django-lambda-tasks 0.6.3__tar.gz → 0.6.5__tar.gz

{django_lambda_tasks-0.6.3 → django_lambda_tasks-0.6.5}/PKG-INFO

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

{django_lambda_tasks-0.6.3 → django_lambda_tasks-0.6.5}/lambda_tasks/handler.py

@@ -32,19 +32,27 @@ _cold_start_done: bool = False
 
 _MEMORY_RESERVED_MB = 128
 _MEMORY_MINIMUM_LIMIT_MB = 64
-_CGROUP_MEMORY_MAX_PATH = Path("/sys/fs/cgroup/memory.max")
+_CGROUP_V2_MEMORY_MAX_PATH = Path("/sys/fs/cgroup/memory.max")
+_CGROUP_V1_MEMORY_LIMIT_PATH = Path("/sys/fs/cgroup/memory/memory.limit_in_bytes")
 
 
 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 cgroup."""
     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()
+        value = _CGROUP_V2_MEMORY_MAX_PATH.read_text().strip()
         if value == "max":
             return None
+        else:
+            return int(value) // (1024 * 1024)
+    except (FileNotFoundError, ValueError):
+        pass
+
+    try:
+        value = _CGROUP_V1_MEMORY_LIMIT_PATH.read_text().strip()
         return int(value) // (1024 * 1024)
     except (FileNotFoundError, ValueError):
         return None

{django_lambda_tasks-0.6.3 → django_lambda_tasks-0.6.5}/lambda_tasks/local_executor.py

@@ -2,6 +2,7 @@
 
 import atexit
 import logging
+import multiprocessing
 import signal
 import threading
 import uuid
@@ -23,9 +24,10 @@ def _pool_initializer() -> None:
     atexit as a fallback). This prevents workers from being killed
     mid-operation and leaking semaphores.
 
-    Closes all inherited database connections before setting up Django.
-    On Linux (fork-based spawning), child processes inherit copies of the
-    parent's DB connections which are in an inconsistent state.
+    Closes all database connections before setting up Django. With the
+    forkserver context, workers don't inherit the parent's connections,
+    but AppConfig.ready() hooks may open connections during django.setup()
+    that should be closed before task execution begins.
 
     Then sets up Django for task execution.
     """
@@ -34,10 +36,10 @@ def _pool_initializer() -> None:
     import django
     from django.db import connections
 
-    connections.close_all()
-
     django.setup()
 
+    connections.close_all()
+
 
 def _shutdown_pool() -> None:
     """Shut down the pool, releasing its POSIX semaphores promptly.
@@ -118,12 +120,22 @@ def _install_shutdown_handlers() -> None:
 
 
 def get_pool() -> ProcessPoolExecutor:
-    """Return the shared ProcessPoolExecutor, creating it on first call."""
+    """Return the shared ProcessPoolExecutor, creating it on first call.
+
+    Uses the "forkserver" multiprocessing context so that worker processes
+    are forked from a clean server process rather than from the Django parent.
+    This prevents workers from inheriting the parent's open database
+    connections (and other file descriptors), which would otherwise cause
+    "server closed the connection unexpectedly" errors in containerized
+    environments where close() in the child tears down the shared socket.
+    """
     global _pool
     if _pool is None:
         conf = LambdaTasksSettings()
+        mp_context = multiprocessing.get_context("forkserver")
         _pool = ProcessPoolExecutor(
             max_workers=conf.LOCAL_WORKERS,
+            mp_context=mp_context,
             initializer=_pool_initializer,
         )
         atexit.register(_shutdown_pool)

{django_lambda_tasks-0.6.3 → django_lambda_tasks-0.6.5}/pyproject.toml

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

{django_lambda_tasks-0.6.3 → django_lambda_tasks-0.6.5}/tests/test_local_executor.py

@@ -189,6 +189,25 @@ class TestGetPool:
 
         assert lambda_tasks.local_executor._pool is pool
 
+    @pytest.mark.usefixtures("_reset_pool")
+    def test_get_pool_uses_forkserver_context(self, settings):
+        """get_pool() uses the forkserver multiprocessing context.
+
+        This prevents workers from inheriting the parent's open database
+        connections, which would cause connection errors in containerized
+        environments.
+        """
+        import lambda_tasks.local_executor
+
+        settings.LAMBDA_TASKS_LOCAL_WORKERS = 2
+        settings.LAMBDA_TASKS_EAGER = False
+
+        pool = lambda_tasks.local_executor.get_pool()
+
+        # ProcessPoolExecutor stores the context; forkserver processes
+        # use the ForkServerProcess type
+        assert pool._mp_context.get_start_method() == "forkserver"
+
 
 # ---------------------------------------------------------------------------
 # Property-based tests: get_pool() (Requirement 2)

{django_lambda_tasks-0.6.3 → django_lambda_tasks-0.6.5}/tests/test_memory_limit.py

@@ -56,7 +56,12 @@ class TestMemoryLimitSetDuringColdStart:
         monkeypatch.setattr(
             "lambda_tasks.handler.resolve_secrets_into_env", lambda: None
         )
-        monkeypatch.setattr(handler_module, "_CGROUP_MEMORY_MAX_PATH", fake_path(None))
+        monkeypatch.setattr(
+            handler_module, "_CGROUP_V2_MEMORY_MAX_PATH", fake_path(None)
+        )
+        monkeypatch.setattr(
+            handler_module, "_CGROUP_V1_MEMORY_LIMIT_PATH", fake_path(None)
+        )
 
         calls: list[tuple] = []
 
@@ -70,7 +75,7 @@ class TestMemoryLimitSetDuringColdStart:
 
         assert calls == []
 
-    def test_rlimit_as_set_from_cgroup(self, monkeypatch):
+    def test_rlimit_as_set_from_cgroup_v2(self, monkeypatch):
         """When AWS_LAMBDA_FUNCTION_MEMORY_SIZE is not set, falls back to cgroup v2."""
         monkeypatch.delenv("DJANGO_SETTINGS_MODULE", raising=False)
         monkeypatch.delenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", raising=False)
@@ -81,9 +86,14 @@ class TestMemoryLimitSetDuringColdStart:
         # 2 GB in bytes
         monkeypatch.setattr(
             handler_module,
-            "_CGROUP_MEMORY_MAX_PATH",
+            "_CGROUP_V2_MEMORY_MAX_PATH",
             fake_path(str(2048 * 1024 * 1024)),
         )
+        monkeypatch.setattr(
+            handler_module,
+            "_CGROUP_V1_MEMORY_LIMIT_PATH",
+            fake_path(None),
+        )
 
         calls: list[tuple] = []
 
@@ -106,7 +116,12 @@ class TestMemoryLimitSetDuringColdStart:
         monkeypatch.setattr(
             "lambda_tasks.handler.resolve_secrets_into_env", lambda: None
         )
-        monkeypatch.setattr(handler_module, "_CGROUP_MEMORY_MAX_PATH", fake_path("max"))
+        monkeypatch.setattr(
+            handler_module, "_CGROUP_V2_MEMORY_MAX_PATH", fake_path("max")
+        )
+        monkeypatch.setattr(
+            handler_module, "_CGROUP_V1_MEMORY_LIMIT_PATH", fake_path(None)
+        )
 
         calls: list[tuple] = []
 
@@ -120,6 +135,39 @@ class TestMemoryLimitSetDuringColdStart:
 
         assert calls == []
 
+    def test_rlimit_as_set_from_cgroup_v1(self, monkeypatch):
+        """When cgroup v2 is unavailable, falls back to cgroup v1 (Fargate)."""
+        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
+        )
+        monkeypatch.setattr(
+            handler_module,
+            "_CGROUP_V2_MEMORY_MAX_PATH",
+            fake_path(None),
+        )
+        # 4 GB in bytes
+        monkeypatch.setattr(
+            handler_module,
+            "_CGROUP_V1_MEMORY_LIMIT_PATH",
+            fake_path(str(4096 * 1024 * 1024)),
+        )
+
+        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 = (4096 - 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)