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")