From 8024df1122070cd307fead745a945b24bf724bd6 Mon Sep 17 00:00:00 2001 From: Shinya Okano Date: Fri, 2 Oct 2026 13:56:46 +0900 Subject: [PATCH] Split run_task() into claim_task() and run_claimed_task() run_task() claimed the task and ran it in one call. A worker that fetches tasks in one thread and runs them in another needs the claim in the fetching thread, so that its next fetch does not return the same READY row, and the run elsewhere. The two steps are now public methods and run_task() calls one and then the other; nothing that calls run_task() changes, and the existing tests pass unmodified. Refs #25 --- CHANGELOG.md | 5 +++ django_database_task/backends.py | 64 ++++++++++++++++++++++++-------- tests/test_backend.py | 52 ++++++++++++++++++++++++++ 3 files changed, 105 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e18f093..33371e7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,11 @@ ### Added +- **`claim_task()` and `run_claimed_task()` on the backend.** The two steps + of `DatabaseTaskBackend.run_task()`, which now calls one and then the + other, so that a worker can claim a task in the thread that fetched it and + run it in another. `run_task()` itself is unchanged, and so is everything + that calls it. Existing projects need no changes. - **`--workers N` on `run_database_tasks`.** The command keeps N worker processes running from one invocation. Each worker is the command run again without `--workers`, so every other option reaches it unchanged and diff --git a/django_database_task/backends.py b/django_database_task/backends.py index 1a0d729..965754c 100644 --- a/django_database_task/backends.py +++ b/django_database_task/backends.py @@ -352,14 +352,44 @@ def _db_task_to_result(self, db_task, task): return result + def claim_task(self, db_task, worker_id=None): + """ + Move a READY task to RUNNING and record the attempt on it. + + The first of the two steps of :meth:`run_task`; the second is + :meth:`run_claimed_task`. A worker that fetches tasks in one thread + and runs them in another claims here, in the fetching thread, so + that its next fetch does not hand it the same READY row again. + + The row is locked and its status re-checked inside one transaction, + so of two workers holding the same READY row exactly one gets to + write RUNNING. The worker id is appended to the list stored in the + database rather than to the copy on ``db_task``, so an attempt + recorded after ``db_task`` was loaded is kept. + + Args: + db_task: The :class:`~django_database_task.models.DatabaseTask` + to claim. It must be READY. + worker_id: Optional worker identifier, recorded on the task. + + Returns: + True if this call claimed the task, and ``db_task`` carries what + was written. False if it was no longer READY (another worker + claimed it, or it was deleted); nothing was written and the task + must not be run. + """ + return self._claim_task(db_task, worker_id) + def run_task(self, db_task, worker_id=None): """ Claim a task and execute it (called from the executor and the management command). - The task is claimed with :meth:`_claim_task` first, so two workers + The task is claimed with :meth:`claim_task` first, so two workers that fetched the same READY row cannot both run it: whichever claims - second gets None and runs nothing. + second gets None and runs nothing. The run itself is + :meth:`run_claimed_task`; the two are separate so that a worker can + claim in one thread and run in another. A task whose function cannot be imported (the module was renamed, the function removed, the worker runs older code than the enqueuer) @@ -380,7 +410,7 @@ def run_task(self, db_task, worker_id=None): not run here. The result of a task that could not be started has ``task`` set to None. """ - if not self._claim_task(db_task, worker_id): + if not self.claim_task(db_task, worker_id): logger.info( "Task not run: id=%s path=%s is no longer READY", db_task.id, @@ -388,7 +418,21 @@ def run_task(self, db_task, worker_id=None): extra=task_log_fields(db_task, worker_id), ) return None + return self.run_claimed_task(db_task, worker_id) + def run_claimed_task(self, db_task, worker_id=None): + """ + Execute a task that :meth:`claim_task` has moved to RUNNING. + + The second of the two steps of :meth:`run_task`, which documents + what is recorded, logged and signalled. ``db_task`` must be the + instance the claim was made with, as the claim wrote the attempt + onto it, and ``worker_id`` the one it was claimed with. + + Returns: + TaskResult after execution. The result of a task that could not + be started has ``task`` set to None. + """ # Past the RUNNING write but before the block that records failures, # so an error here would otherwise leave the task RUNNING with no # error, and no other worker would take it. @@ -541,19 +585,7 @@ def _record_error(self, db_task, exc): return error def _claim_task(self, db_task, worker_id): - """ - Move a READY task to RUNNING and record the attempt on it. - - The row is locked and its status re-checked inside one transaction, - so of two workers holding the same READY row exactly one gets to - write RUNNING. The worker id is appended to the list stored in the - database rather than to the copy on ``db_task``, so an attempt - recorded after ``db_task`` was loaded is kept. - - Returns: - True if this call claimed the task, False if it was no longer - READY. On True, ``db_task`` carries what was written. - """ + """The claim behind :meth:`claim_task`, which documents it.""" from .models import DatabaseTask now = timezone.now() diff --git a/tests/test_backend.py b/tests/test_backend.py index 02107e0..6012130 100644 --- a/tests/test_backend.py +++ b/tests/test_backend.py @@ -1,6 +1,7 @@ """Tests for the database task backend.""" from datetime import timedelta +from unittest.mock import patch import pytest from django.tasks import task_backends @@ -8,6 +9,7 @@ from django.tasks.exceptions import TaskResultDoesNotExist from django.utils import timezone +from django_database_task.backends import DatabaseTaskBackend from django_database_task.models import DatabaseTask from . import tasks as test_tasks @@ -433,3 +435,53 @@ def test_async_task_failure(self): db_task.refresh_from_db() assert db_task.status == TaskResultStatus.FAILED + + +@pytest.mark.django_db +class TestClaimAndRunSeparately: + """claim_task() and run_claimed_task() are the two steps of run_task().""" + + def test_claim_then_run(self): + test_tasks.counting_task_runs.clear() + result = counting_task.enqueue() + db_task = DatabaseTask.objects.get(id=result.id) + backend = task_backends["default"] + + assert backend.claim_task(db_task, worker_id="worker-1") is True + claimed = DatabaseTask.objects.get(id=result.id) + assert claimed.status == TaskResultStatus.RUNNING + assert claimed.worker_ids_json == ["worker-1"] + assert test_tasks.counting_task_runs == [] + + final_result = backend.run_claimed_task(db_task, worker_id="worker-1") + + assert final_result.status == TaskResultStatus.SUCCESSFUL + assert test_tasks.counting_task_runs == [1] + assert DatabaseTask.objects.get(id=result.id).status == ( + TaskResultStatus.SUCCESSFUL + ) + + def test_a_lost_claim_returns_false_and_writes_nothing(self): + result = counting_task.enqueue() + first = DatabaseTask.objects.get(id=result.id) + second = DatabaseTask.objects.get(id=result.id) + backend = task_backends["default"] + + assert backend.claim_task(first, worker_id="worker-1") is True + assert backend.claim_task(second, worker_id="worker-2") is False + + assert DatabaseTask.objects.get(id=result.id).worker_ids_json == ["worker-1"] + + def test_run_task_is_the_two_steps(self): + result = simple_task.enqueue(1, 2) + db_task = DatabaseTask.objects.get(id=result.id) + backend = task_backends["default"] + + with ( + patch.object(DatabaseTaskBackend, "claim_task", return_value=True) as claim, + patch.object(DatabaseTaskBackend, "run_claimed_task") as run, + ): + backend.run_task(db_task, worker_id="w") + + claim.assert_called_once_with(db_task, "w") + run.assert_called_once_with(db_task, "w")