django-lambda-tasks 0.4.2__tar.gz → 0.4.3__tar.gz

{django_lambda_tasks-0.4.2 → django_lambda_tasks-0.4.3}/.kiro/steering/product.md

@@ -136,10 +136,11 @@ class SQSLambdaTaskMessage(BaseModel):
 - Returns `{"batchItemFailures": [...]}` for partial-batch failure reporting
 - Only pre-execution failures (malformed message, import error, misconfiguration) are reported as `batchItemFailures` — task logic failures are caught and recorded as `FAILED` TaskRecords without raising
 - Recommended SQS queue settings: `maxReceiveCount=1` with a DLQ configured; automatic retries are not useful since task failures are not re-driven by design
-- Cold-start sequence runs inside the handler on the first invocation (not at module import time) to avoid Lambda init-duration timeouts: a temporary `StreamHandler` is attached to the `lambda_tasks` logger, then `resolve_environment()` → `resolve_secrets_into_env()` (handler removed) → conditional `django.setup()`
+- Cold-start sequence runs inside the handler on the first invocation (not at module import time) to avoid Lambda init-duration timeouts: memory limit is set, a temporary `StreamHandler` is attached to the `lambda_tasks` logger, then `resolve_environment()` → `resolve_secrets_into_env()` (handler removed) → conditional `django.setup()`
 - 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
 
 ## Environment Loader
 

{django_lambda_tasks-0.4.2 → django_lambda_tasks-0.4.3}/.kiro/steering/structure.md

@@ -42,7 +42,7 @@ django-lambda-tasks/
 - `decorators.py` — defines `@lambda_task`; enforces kwargs-only at decoration time
 - `models.py` — `TaskRecord` (Django ORM), `SQSLambdaTaskMessage` (Pydantic, SQS schema + execution logic), `SQSLambdaTask` (Pydantic, holds message + routing; `_execute()` publishes to SQS, executes eagerly, or submits to the local process pool; `execute_on_commit()` registers `_execute` with `transaction.on_commit`)
 - `local_executor.py` — `ProcessPoolExecutor`-based async local execution; `get_pool()` lazily creates a module-level pool; `submit_task()` fire-and-forget submission; `_execute_in_worker()` deserializes and runs the task in a child process; `_pool_initializer()` calls `django.setup()` per worker
-- `handler.py` — Lambda entry point; cold-start init (temporary log handler → `resolve_environment()` → `resolve_secrets_into_env()` → handler removed → conditional `django.setup()`) runs inside the handler on first invocation, guarded by `_cold_start_done` sentinel; processes SQS records independently; returns `batchItemFailures`
+- `handler.py` — Lambda entry point; cold-start init (memory limit → temporary log handler → `resolve_environment()` → `resolve_secrets_into_env()` → handler removed → conditional `django.setup()`) runs inside the handler on first invocation, guarded by `_cold_start_done` sentinel; processes SQS records independently; returns `batchItemFailures`
 - `environment_loader.py` — loads env vars from a Secrets Manager secret at cold start; validates flat JSON format; idempotent via `_loaded` sentinel
 - `secret_loader.py` — resolves `LAMBDA_TASKS_SECRET_*` env vars from Secrets Manager before Django starts; validates format, detects conflicts, batches API calls; idempotent via `_loaded` sentinel
 - `logging.py` — `task_logger` singleton; `message_id` set/cleared around each task execution

{django_lambda_tasks-0.4.2 → django_lambda_tasks-0.4.3}/PKG-INFO

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

{django_lambda_tasks-0.4.2 → django_lambda_tasks-0.4.3}/lambda_tasks/handler.py

@@ -14,6 +14,7 @@ warm invocations skip it.
 
 import logging
 import os
+import resource
 
 import django
 from django.apps import apps
@@ -26,6 +27,36 @@ logger = logging.getLogger(__name__)
 _cold_start_done: bool = False
 
 
+_MEMORY_RESERVED_MB = 128
+_MEMORY_MINIMUM_LIMIT_MB = 64
+
+
+def _set_memory_limit() -> None:
+    """Set RLIMIT_DATA from AWS_LAMBDA_FUNCTION_MEMORY_SIZE if available.
+
+    When running in Lambda, this causes Python to raise MemoryError on
+    excessive allocation instead of being killed by the OOM killer.
+
+    128 MB is reserved for the Python runtime, shared libraries, and OS
+    overhead. The limit is floored at 64 MB so that even small Lambdas
+    get some protection.
+    """
+    memory_mb = os.environ.get("AWS_LAMBDA_FUNCTION_MEMORY_SIZE")
+
+    if memory_mb is None:
+        return
+
+    limit_bytes = (
+        max(int(memory_mb) - _MEMORY_RESERVED_MB, _MEMORY_MINIMUM_LIMIT_MB)
+        * 1024
+        * 1024
+    )
+    resource.setrlimit(resource.RLIMIT_AS, (limit_bytes, limit_bytes))
+    logger.info(
+        f"Set RLIMIT_AS to {limit_bytes} bytes ({memory_mb} MB - {_MEMORY_RESERVED_MB} MB reserved)"
+    )
+
+
 def _perform_cold_start() -> None:
     """Run one-time initialisation: env loading, secrets, Django setup.
 
@@ -49,6 +80,8 @@ def _perform_cold_start() -> None:
     lambda_tasks_logger.addHandler(boot_handler)
     lambda_tasks_logger.setLevel(logging.DEBUG)
 
+    _set_memory_limit()
+
     try:
         resolve_environment()
         resolve_secrets_into_env()

{django_lambda_tasks-0.4.2 → django_lambda_tasks-0.4.3}/pyproject.toml

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

django_lambda_tasks-0.4.3/tests/test_memory_limit.py

@@ -0,0 +1,131 @@
+"""
+Unit tests for memory limit enforcement in the Lambda handler.
+
+The handler sets resource.RLIMIT_AS during cold start so that runaway
+memory allocation raises MemoryError instead of triggering the Lambda OOM
+killer.
+"""
+
+import resource
+from unittest.mock import patch
+
+import lambda_tasks.handler as handler_module
+
+
+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.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_rlimit_as_not_set_when_env_var_missing(self, monkeypatch):
+        """When AWS_LAMBDA_FUNCTION_MEMORY_SIZE is not set, RLIMIT_AS is unchanged."""
+        monkeypatch.delenv("DJANGO_SETTINGS_MODULE", raising=False)
+        monkeypatch.delenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", raising=False)
+        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)
+
+        assert calls == []
+
+    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")
+
+        call_order: list[str] = []
+
+        def fake_setrlimit(which: int, limits: tuple[int, int]) -> None:
+            call_order.append("setrlimit")
+
+        def fake_resolve_environment() -> None:
+            call_order.append("resolve_environment")
+
+        def fake_resolve_secrets() -> None:
+            call_order.append("resolve_secrets_into_env")
+
+        monkeypatch.setattr(resource, "setrlimit", fake_setrlimit)
+        monkeypatch.setattr(
+            "lambda_tasks.handler.resolve_environment", fake_resolve_environment
+        )
+        monkeypatch.setattr(
+            "lambda_tasks.handler.resolve_secrets_into_env", fake_resolve_secrets
+        )
+
+        handler_module._cold_start_done = False
+        handler_module.handler(event={"Records": []}, context=None)
+
+        assert call_order == [
+            "setrlimit",
+            "resolve_environment",
+            "resolve_secrets_into_env",
+        ]
+
+    def test_memory_limit_is_logged(self, monkeypatch, caplog):
+        """Setting the memory limit emits an info log."""
+        import logging
+
+        monkeypatch.delenv("DJANGO_SETTINGS_MODULE", raising=False)
+        monkeypatch.setenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", "1024")
+        monkeypatch.setattr("lambda_tasks.handler.resolve_environment", lambda: None)
+        monkeypatch.setattr(
+            "lambda_tasks.handler.resolve_secrets_into_env", lambda: None
+        )
+        monkeypatch.setattr(resource, "setrlimit", lambda *args: None)
+
+        handler_module._cold_start_done = False
+
+        with caplog.at_level(logging.INFO, logger="lambda_tasks.handler"):
+            handler_module.handler(event={"Records": []}, context=None)
+
+        assert any("RLIMIT_AS" in msg for msg in caplog.messages)
+
+    def test_rlimit_floored_at_minimum_when_memory_too_small(self, monkeypatch):
+        """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.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 = 64 * 1024 * 1024
+        assert calls == [(resource.RLIMIT_AS, (expected, expected))]