From 18cca9a44ee4048f8faabe0870bb50b223d7d3b9 Mon Sep 17 00:00:00 2001 From: Johannes Maron Date: Mon, 10 Aug 2026 18:25:39 +0200 Subject: [PATCH 1/2] Fix AppRegistryNotReady error for tasks w/ model access --- tests/test_executor.py | 24 ++++++++++++++++++++++++ tests/testapp/settings.py | 2 +- tests/testapp/tasks.py | 6 ++++++ threadmill/executor.py | 2 ++ 4 files changed, 33 insertions(+), 1 deletion(-) diff --git a/tests/test_executor.py b/tests/test_executor.py index 456e376..7009381 100644 --- a/tests/test_executor.py +++ b/tests/test_executor.py @@ -22,6 +22,7 @@ boom_retry_raises, boom_retry_thrice, boom_with_retry, + count_users, echo, ) from threadmill.backends.base import Broker @@ -158,6 +159,29 @@ def test_run__processes_enqueued_tasks_end_to_end(self): assert {r.id for r in results} == {r.id for r in enqueued} assert all(r.status == TaskResultStatus.SUCCESSFUL for r in results) + def test_run__executes_model_task_in_spawned_worker(self): + """run() executes a model-accessing task in a spawned worker process.""" + original_start_method = multiprocessing.get_start_method() + multiprocessing.set_start_method("spawn", force=True) + try: + enqueued = default_task_backend.enqueue(count_users) + executor = TaskExecutor( + backend=default_task_backend, + workers=1, + threads=1, + queues=("default",), + ) + run_thread = threading.Thread(target=executor.run, daemon=True) + run_thread.start() + time.sleep(3) + executor.shutdown() + run_thread.join(timeout=5) + assert not run_thread.is_alive() + result = default_task_backend.get_result(enqueued.id) + assert result.status == TaskResultStatus.SUCCESSFUL + finally: + multiprocessing.set_start_method(original_start_method, force=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]) diff --git a/tests/testapp/settings.py b/tests/testapp/settings.py index 1d4e2fd..dd4767c 100644 --- a/tests/testapp/settings.py +++ b/tests/testapp/settings.py @@ -82,7 +82,7 @@ DATABASES = { "default": { "ENGINE": "django.db.backends.sqlite3", - "NAME": ":memory:", + "NAME": BASE_DIR / "db.sqlite3", } } diff --git a/tests/testapp/tasks.py b/tests/testapp/tasks.py index 0be41ac..af578f6 100644 --- a/tests/testapp/tasks.py +++ b/tests/testapp/tasks.py @@ -22,6 +22,12 @@ def boom(): raise ValueError("boom") +@task() +def count_users(): + """Count all users in the database (tests model access in workers).""" + from django.contrib.auth.models import User # noqa + + @task(queue_name="compute") def compute_workload(): """Calculate the first 1000 prime numbers.""" diff --git a/threadmill/executor.py b/threadmill/executor.py index cfd4d64..f9490b3 100644 --- a/threadmill/executor.py +++ b/threadmill/executor.py @@ -15,6 +15,7 @@ from queue import Empty from traceback import format_exception +import django from django.tasks import TaskResult, task_backends from django.tasks.base import TaskContext, TaskError, TaskResultStatus from django.tasks.signals import task_finished, task_started @@ -146,6 +147,7 @@ def __init__( def run(self) -> None: """Start consumer execution inside this process.""" + django.setup() logger.info("Starting worker process %s", self.name) self.lock = threading.Lock() self.expired = threading.Event() From ccbeffc85bae938057ecd7503fc6f35768e9441a Mon Sep 17 00:00:00 2001 From: Johannes Maron Date: Wed, 2 Sep 2026 21:10:03 +0200 Subject: [PATCH 2/2] Close old database connections around task execution Register a receiver for the task_started and task_finished signals that calls close_old_connections(), so worker processes reuse database connections the same way request handlers do, honoring CONN_MAX_AGE. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- tests/test_executor.py | 4 ++ tests/test_signals.py | 85 ++++++++++++++++++++++++++++++++++++++++++ threadmill/apps.py | 8 ++++ threadmill/signals.py | 10 +++++ 4 files changed, 107 insertions(+) create mode 100644 tests/test_signals.py create mode 100644 threadmill/apps.py create mode 100644 threadmill/signals.py diff --git a/tests/test_executor.py b/tests/test_executor.py index a2729eb..b7af974 100644 --- a/tests/test_executor.py +++ b/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( diff --git a/tests/test_signals.py b/tests/test_signals.py new file mode 100644 index 0000000..e94b289 --- /dev/null +++ b/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() diff --git a/threadmill/apps.py b/threadmill/apps.py new file mode 100644 index 0000000..1b0a367 --- /dev/null +++ b/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 diff --git a/threadmill/signals.py b/threadmill/signals.py new file mode 100644 index 0000000..57bb710 --- /dev/null +++ b/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()