Close old database connections around task execution

tests/test_executor.py

@@ -7,6 +7,7 @@
 import time
 import uuid
 
+import pytest
 from django.tasks import (
     TaskContext,
     TaskResult,
@@ -182,6 +183,7 @@ def test_run__executes_model_task_in_spawned_worker(self):
         finally:
             multiprocessing.set_start_method(original_start_method, force=True)
 
+    @pytest.mark.django_db(transaction=True)
     def test_worker_acquires_updates_and_acknowledges(self):
         """Worker acquires, executes, and acknowledges via its own backend."""
         enqueued = default_task_backend.enqueue(echo, args=[42])
@@ -288,6 +290,8 @@ def test_shutdown_requested__is_settable(self):
 class TestWorkerThread:
     """Tests for the WorkerThread class."""
 
+    pytestmark = pytest.mark.django_db(transaction=True)
+
     def test_execute_task_result__successful_execution(self):
         """execute_task_result runs a task and returns SUCCESSFUL result."""
         result = WorkerThread(

tests/test_signals.py

@@ -0,0 +1,85 @@
+import time
+import uuid
+from unittest import mock
+
+import pytest
+from django.db import DEFAULT_DB_ALIAS, connections
+from django.tasks import TaskResult, TaskResultStatus
+from django.tasks.backends.immediate import ImmediateBackend
+from django.tasks.signals import task_finished, task_started
+from django.utils import timezone
+
+from tests.testapp.tasks import echo
+from threadmill.executor import TaskExecutor, WorkerProcess, WorkerThread
+
+
+def _task_result(task, *args) -> TaskResult:
+    """Build a READY `TaskResult` without touching the backend."""
+    return TaskResult(
+        task=task,
+        id=str(uuid.uuid4()),
+        status=TaskResultStatus.READY,
+        enqueued_at=timezone.now(),
+        started_at=None,
+        finished_at=None,
+        last_attempted_at=None,
+        args=list(args),
+        kwargs={},
+        backend="default",
+        errors=[],
+        worker_ids=[],
+    )
+
+
+@pytest.fixture
+def stale_connection():
+    """Open a database connection and mark it as expired."""
+    connection = connections[DEFAULT_DB_ALIAS]
+    connection.ensure_connection()
+    connection.close_at = time.monotonic() - 1
+    yield connection
+    connection.close_at = None
+    connection.close()
+
+
+@pytest.mark.django_db(transaction=True)
+class TestCloseTaskDatabaseConnection:
+    def test_task_started__closes_stale_connection(self, stale_connection):
+        """A stale connection is closed when a task starts."""
+        with mock.patch.object(
+            stale_connection, "close", wraps=stale_connection.close
+        ) as spy:
+            task_started.send(TaskExecutor, task_result=_task_result(echo))
+
+        spy.assert_called_once()
+
+    def test_task_finished__closes_stale_connection(self, stale_connection):
+        """A stale connection is closed when a task finished."""
+        with mock.patch.object(
+            stale_connection, "close", wraps=stale_connection.close
+        ) as spy:
+            task_finished.send(TaskExecutor, task_result=_task_result(echo))
+
+        spy.assert_called_once()
+
+    def test_task_started__ignores_other_senders(self, stale_connection):
+        """Signals from other executors do not close the connection."""
+        with mock.patch.object(stale_connection, "close") as spy:
+            task_started.send(ImmediateBackend, task_result=_task_result(echo))
+
+        spy.assert_not_called()
+
+    def test_execute_task_result__closes_stale_connection(self, stale_connection):
+        """Executing a task closes stale connections via lifecycle signals."""
+        with mock.patch.object(
+            stale_connection, "close", wraps=stale_connection.close
+        ) as spy:
+            worker = WorkerProcess(
+                thread_count=1, backend_alias="default", queues=("default",)
+            )
+            thread = WorkerThread(worker=worker, index=0, backend=None)
+
+            result = thread.execute_task_result(_task_result(echo, 42))
+
+        assert result.status is TaskResultStatus.SUCCESSFUL
+        spy.assert_called()

threadmill/apps.py

@@ -0,0 +1,8 @@
+from django.apps import AppConfig
+
+
+class ThreadmillConfig(AppConfig):
+    name = "threadmill"
+
+    def ready(self) -> None:
+        from . import signals  # noqa: F401

threadmill/signals.py

@@ -0,0 +1,10 @@
+from django.db import close_old_connections
+from django.dispatch import receiver
+from django.tasks.signals import task_finished, task_started
+
+from .executor import TaskExecutor
+
+
+@receiver([task_started, task_finished], sender=TaskExecutor)
+def close_task_database_connection(sender, task_result, **kwargs):
+    close_old_connections()