django-lambda-tasks 0.6.4__tar.gz → 0.6.6__tar.gz

{django_lambda_tasks-0.6.4 → django_lambda_tasks-0.6.6}/.kiro/steering/product.md

@@ -151,7 +151,7 @@ class SQSLambdaTaskMessage(BaseModel):
 - A module-level `_cold_start_done` sentinel ensures the sequence runs only once; subsequent warm invocations skip it
 - Both loaders run unconditionally (outside the `DJANGO_SETTINGS_MODULE` check) — the environment secret may provide that var, and individual secrets may depend on environment-loaded vars
 - A temporary `StreamHandler` is attached to the `lambda_tasks` logger for the duration of the loaders so their log output is visible before Django's `LOGGING` dictConfig has run; it is removed immediately after so that Django's configuration is the sole authority on logging from that point on
-- `resource.setrlimit(RLIMIT_DATA)` is set from `AWS_LAMBDA_FUNCTION_MEMORY_SIZE` so that excessive allocation raises `MemoryError` instead of triggering the OOM killer
+- `resource.setrlimit(RLIMIT_AS)` is set from `AWS_LAMBDA_FUNCTION_MEMORY_SIZE` (Lambda) or the ECS task metadata endpoint (Fargate/Batch) so that excessive allocation raises `MemoryError` instead of triggering the OOM killer
 
 `main()` in `handler.py` — AWS Batch container entry point (`python -m lambda_tasks.handler`):
 - Reads `LAMBDA_TASKS_MESSAGE` env var (the serialized task JSON, set by `submit_batch_job` via container overrides)

{django_lambda_tasks-0.6.4 → django_lambda_tasks-0.6.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: django-lambda-tasks
-Version: 0.6.4
+Version: 0.6.6
 Summary: Run async tasks in a lambda function
 Requires-Python: >=3.10
 Requires-Dist: awslambdaric

{django_lambda_tasks-0.6.4 → django_lambda_tasks-0.6.6}/lambda_tasks/handler.py

@@ -12,12 +12,13 @@ timeouts. The sequence is guarded by a module-level sentinel so subsequent
 warm invocations skip it.
 """
 
+import json
 import logging
 import os
 import resource
 import sys
+import urllib.request
 import uuid
-from pathlib import Path
 
 import django
 from django.apps import apps
@@ -32,29 +33,31 @@ _cold_start_done: bool = False
 
 _MEMORY_RESERVED_MB = 128
 _MEMORY_MINIMUM_LIMIT_MB = 64
-_CGROUP_MEMORY_MAX_PATH = Path("/sys/fs/cgroup/memory.max")
 
 
 def _get_memory_mb() -> int | None:
-    """Get container memory limit from Lambda env var or cgroup v2."""
+    """Get container memory limit from Lambda env var or ECS metadata endpoint."""
     memory_mb_str = os.environ.get("AWS_LAMBDA_FUNCTION_MEMORY_SIZE")
     if memory_mb_str is not None:
         return int(memory_mb_str)
 
-    try:
-        value = _CGROUP_MEMORY_MAX_PATH.read_text().strip()
-        if value == "max":
+    metadata_uri = os.environ.get("ECS_CONTAINER_METADATA_URI_V4")
+    if metadata_uri is not None:
+        try:
+            with urllib.request.urlopen(f"{metadata_uri}/task", timeout=2) as response:
+                data = json.loads(response.read())
+            return int(data["Limits"]["Memory"])
+        except Exception:
             return None
-        return int(value) // (1024 * 1024)
-    except (FileNotFoundError, ValueError):
-        return None
+
+    return None
 
 
 def _set_memory_limit() -> None:
     """Set RLIMIT_AS from the container memory limit.
 
     On Lambda, reads AWS_LAMBDA_FUNCTION_MEMORY_SIZE. On Fargate/Batch,
-    reads the cgroup v2 memory limit. This causes Python to raise
+    queries the ECS task metadata endpoint. This causes Python to raise
     MemoryError on excessive allocation instead of being killed by the
     OOM killer.
 

{django_lambda_tasks-0.6.4 → django_lambda_tasks-0.6.6}/pyproject.toml

@@ -7,7 +7,7 @@ packages = ["lambda_tasks"]
 
 [project]
 name = "django-lambda-tasks"
-version = "0.6.4"
+version = "0.6.6"
 description = "Run async tasks in a lambda function"
 readme = "README.md"
 requires-python = ">=3.10"

{django_lambda_tasks-0.6.4 → django_lambda_tasks-0.6.6}/tests/test_memory_limit.py

@@ -6,30 +6,18 @@ memory allocation raises MemoryError instead of triggering the Lambda OOM
 killer.
 """
 
+import json
 import resource
 
 import lambda_tasks.handler as handler_module
 
 
-class _FakePath:
-    def __init__(self, *, content: str | None) -> None:
-        self._content = content
-
-    def read_text(self) -> str:
-        if self._content is None:
-            raise FileNotFoundError
-        return self._content
-
-
-def fake_path(content: str | None) -> _FakePath:
-    return _FakePath(content=content)
-
-
 class TestMemoryLimitSetDuringColdStart:
     def test_rlimit_as_set_from_env_var(self, monkeypatch):
         """When AWS_LAMBDA_FUNCTION_MEMORY_SIZE is set, RLIMIT_AS is configured with 128 MB reserved."""
         monkeypatch.delenv("DJANGO_SETTINGS_MODULE", raising=False)
         monkeypatch.setenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", "512")
+        monkeypatch.delenv("ECS_CONTAINER_METADATA_URI_V4", raising=False)
         monkeypatch.setattr("lambda_tasks.handler.resolve_environment", lambda: None)
         monkeypatch.setattr(
             "lambda_tasks.handler.resolve_secrets_into_env", lambda: None
@@ -48,15 +36,15 @@ class TestMemoryLimitSetDuringColdStart:
         expected = (512 - 128) * 1024 * 1024
         assert calls == [(resource.RLIMIT_AS, (expected, expected))]
 
-    def test_rlimit_as_not_set_when_env_var_missing_and_no_cgroup(self, monkeypatch):
-        """When AWS_LAMBDA_FUNCTION_MEMORY_SIZE is not set and cgroup is unavailable, RLIMIT_AS is unchanged."""
+    def test_rlimit_as_not_set_when_no_memory_source(self, monkeypatch):
+        """When neither Lambda env var nor ECS metadata is available, RLIMIT_AS is unchanged."""
         monkeypatch.delenv("DJANGO_SETTINGS_MODULE", raising=False)
         monkeypatch.delenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", raising=False)
+        monkeypatch.delenv("ECS_CONTAINER_METADATA_URI_V4", raising=False)
         monkeypatch.setattr("lambda_tasks.handler.resolve_environment", lambda: None)
         monkeypatch.setattr(
             "lambda_tasks.handler.resolve_secrets_into_env", lambda: None
         )
-        monkeypatch.setattr(handler_module, "_CGROUP_MEMORY_MAX_PATH", fake_path(None))
 
         calls: list[tuple] = []
 
@@ -70,20 +58,33 @@ class TestMemoryLimitSetDuringColdStart:
 
         assert calls == []
 
-    def test_rlimit_as_set_from_cgroup(self, monkeypatch):
-        """When AWS_LAMBDA_FUNCTION_MEMORY_SIZE is not set, falls back to cgroup v2."""
+    def test_rlimit_as_set_from_ecs_metadata(self, monkeypatch):
+        """When ECS_CONTAINER_METADATA_URI_V4 is set, queries the task endpoint for memory."""
         monkeypatch.delenv("DJANGO_SETTINGS_MODULE", raising=False)
         monkeypatch.delenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", raising=False)
+        monkeypatch.setenv("ECS_CONTAINER_METADATA_URI_V4", "http://169.254.170.2/v4")
         monkeypatch.setattr("lambda_tasks.handler.resolve_environment", lambda: None)
         monkeypatch.setattr(
             "lambda_tasks.handler.resolve_secrets_into_env", lambda: None
         )
-        # 2 GB in bytes
-        monkeypatch.setattr(
-            handler_module,
-            "_CGROUP_MEMORY_MAX_PATH",
-            fake_path(str(2048 * 1024 * 1024)),
-        )
+
+        task_metadata = json.dumps({"Limits": {"CPU": 2, "Memory": 8192}}).encode()
+
+        class FakeResponse:
+            def read(self):
+                return task_metadata
+
+            def __enter__(self):
+                return self
+
+            def __exit__(self, *args):
+                pass
+
+        def fake_urlopen(url, *, timeout=None):
+            assert url == "http://169.254.170.2/v4/task"
+            return FakeResponse()
+
+        monkeypatch.setattr("urllib.request.urlopen", fake_urlopen)
 
         calls: list[tuple] = []
 
@@ -95,18 +96,23 @@ class TestMemoryLimitSetDuringColdStart:
         handler_module._cold_start_done = False
         handler_module.handler(event={"Records": []}, context=None)
 
-        expected = (2048 - 128) * 1024 * 1024
+        expected = (8192 - 128) * 1024 * 1024
         assert calls == [(resource.RLIMIT_AS, (expected, expected))]
 
-    def test_rlimit_as_not_set_when_cgroup_is_max(self, monkeypatch):
-        """When cgroup memory.max is 'max' (unlimited), RLIMIT_AS is unchanged."""
+    def test_rlimit_as_not_set_when_ecs_metadata_fails(self, monkeypatch):
+        """When ECS metadata endpoint fails, RLIMIT_AS is unchanged."""
         monkeypatch.delenv("DJANGO_SETTINGS_MODULE", raising=False)
         monkeypatch.delenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", raising=False)
+        monkeypatch.setenv("ECS_CONTAINER_METADATA_URI_V4", "http://169.254.170.2/v4")
         monkeypatch.setattr("lambda_tasks.handler.resolve_environment", lambda: None)
         monkeypatch.setattr(
             "lambda_tasks.handler.resolve_secrets_into_env", lambda: None
         )
-        monkeypatch.setattr(handler_module, "_CGROUP_MEMORY_MAX_PATH", fake_path("max"))
+
+        def fake_urlopen(url, *, timeout=None):
+            raise OSError("connection refused")
+
+        monkeypatch.setattr("urllib.request.urlopen", fake_urlopen)
 
         calls: list[tuple] = []
 
@@ -120,10 +126,34 @@ class TestMemoryLimitSetDuringColdStart:
 
         assert calls == []
 
+    def test_lambda_env_var_takes_precedence_over_ecs_metadata(self, monkeypatch):
+        """AWS_LAMBDA_FUNCTION_MEMORY_SIZE is preferred over ECS metadata."""
+        monkeypatch.delenv("DJANGO_SETTINGS_MODULE", raising=False)
+        monkeypatch.setenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", "512")
+        monkeypatch.setenv("ECS_CONTAINER_METADATA_URI_V4", "http://169.254.170.2/v4")
+        monkeypatch.setattr("lambda_tasks.handler.resolve_environment", lambda: None)
+        monkeypatch.setattr(
+            "lambda_tasks.handler.resolve_secrets_into_env", lambda: None
+        )
+
+        calls: list[tuple] = []
+
+        def fake_setrlimit(which: int, limits: tuple[int, int]) -> None:
+            calls.append((which, limits))
+
+        monkeypatch.setattr(resource, "setrlimit", fake_setrlimit)
+
+        handler_module._cold_start_done = False
+        handler_module.handler(event={"Records": []}, context=None)
+
+        expected = (512 - 128) * 1024 * 1024
+        assert calls == [(resource.RLIMIT_AS, (expected, expected))]
+
     def test_memory_limit_set_before_loaders(self, monkeypatch):
         """The memory limit is set before resolve_environment runs."""
         monkeypatch.delenv("DJANGO_SETTINGS_MODULE", raising=False)
         monkeypatch.setenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", "256")
+        monkeypatch.delenv("ECS_CONTAINER_METADATA_URI_V4", raising=False)
 
         call_order: list[str] = []
 
@@ -159,6 +189,7 @@ class TestMemoryLimitSetDuringColdStart:
 
         monkeypatch.delenv("DJANGO_SETTINGS_MODULE", raising=False)
         monkeypatch.setenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", "1024")
+        monkeypatch.delenv("ECS_CONTAINER_METADATA_URI_V4", raising=False)
         monkeypatch.setattr("lambda_tasks.handler.resolve_environment", lambda: None)
         monkeypatch.setattr(
             "lambda_tasks.handler.resolve_secrets_into_env", lambda: None
@@ -176,6 +207,7 @@ class TestMemoryLimitSetDuringColdStart:
         """When AWS_LAMBDA_FUNCTION_MEMORY_SIZE <= reserved, RLIMIT_AS is set to the 64 MB minimum."""
         monkeypatch.delenv("DJANGO_SETTINGS_MODULE", raising=False)
         monkeypatch.setenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", "128")
+        monkeypatch.delenv("ECS_CONTAINER_METADATA_URI_V4", raising=False)
         monkeypatch.setattr("lambda_tasks.handler.resolve_environment", lambda: None)
         monkeypatch.setattr(
             "lambda_tasks.handler.resolve_secrets_into_env", lambda: None