Split run_task() into claim_task() and run_claimed_task()

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

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,15 +410,29 @@ 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,
                 db_task.task_path,
                 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()

tests/test_backend.py

@@ -1,13 +1,15 @@
 """Tests for the database task backend."""
 
 from datetime import timedelta
+from unittest.mock import patch
 
 import pytest
 from django.tasks import task_backends
 from django.tasks.base import TaskResultStatus
 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")