Store the worker processing a task on the task itself

Author: RealOrangeOne

django_tasks/backends/database/management/commands/db_worker.py

@@ -4,6 +4,7 @@
 import signal
 import sys
 import time
+import uuid
 from argparse import ArgumentParser, ArgumentTypeError
 from types import FrameType
 from typing import Optional
@@ -34,6 +35,7 @@ def __init__(
         batch: bool,
         backend_name: str,
         startup_delay: bool,
+        worker_id: uuid.UUID,
     ):
         self.queue_names = queue_names
         self.process_all_queues = "*" in queue_names
@@ -42,6 +44,8 @@ def __init__(
         self.backend_name = backend_name
         self.startup_delay = startup_delay
 
+        self.worker_id = worker_id
+
         self.running = True
         self.running_task = False
 
@@ -72,7 +76,11 @@ def configure_signals(self) -> None:
     def start(self) -> None:
         self.configure_signals()
 
-        logger.info("Starting worker for queues=%s", ",".join(self.queue_names))
+        logger.info(
+            "Starting worker worker_id=%s for queues=%s",
+            self.worker_id,
+            ",".join(self.queue_names),
+        )
 
         if self.startup_delay and self.interval:
             # Add a random small delay before starting the loop to avoid a thundering herd
@@ -101,7 +109,7 @@ def start(self) -> None:
 
                     if task_result is not None:
                         # "claim" the task, so it isn't run by another worker process
-                        task_result.claim()
+                        task_result.claim(self.worker_id)
 
                 if task_result is not None:
                     self.run_task(task_result)
@@ -150,7 +158,10 @@ def run_task(self, db_task_result: DBTaskResult) -> None:
                 sender = type(db_task_result.task.get_backend())
                 task_result = db_task_result.task_result
             except (ModuleNotFoundError, SuspiciousOperation):
-                logger.exception("Task id=%s failed unexpectedly", db_task_result.id)
+                logger.exception(
+                    "Task id=%s failed unexpectedly",
+                    db_task_result.id,
+                )
             else:
                 task_finished.send(
                     sender=sender,
@@ -214,6 +225,13 @@ def add_arguments(self, parser: ArgumentParser) -> None:
             dest="startup_delay",
             help="Don't add a small delay at startup.",
         )
+        parser.add_argument(
+            "--worker-id",
+            nargs="?",
+            type=uuid.UUID,
+            help="Worker id. MUST be unique across worker pool (default: auto-generate)",
+            default=uuid.uuid4(),
+        )
 
     def configure_logging(self, verbosity: int) -> None:
         if verbosity == 0:
@@ -239,6 +257,7 @@ def handle(
         batch: bool,
         backend_name: str,
         startup_delay: bool,
+        worker_id: uuid.UUID,
         **options: dict,
     ) -> None:
         self.configure_logging(verbosity)
@@ -249,9 +268,12 @@ def handle(
             batch=batch,
             backend_name=backend_name,
             startup_delay=startup_delay,
+            worker_id=worker_id,
         )
 
         worker.start()
 
         if batch:
-            logger.info("No more tasks to run - exiting gracefully.")
+            logger.info(
+                "No more tasks to run for worker_id=%s - exiting gracefully.", worker_id
+            )

django_tasks/backends/database/migrations/0015_dbtaskresult_worker_id.py

@@ -0,0 +1,17 @@
+# Generated by Django 5.1.5 on 2025-02-01 12:16
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("django_tasks_database", "0014_remove_dbtaskresult_exception_data"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="dbtaskresult",
+            name="worker_id",
+            field=models.UUIDField(null=True, verbose_name="worker id"),
+        ),
+    ]

django_tasks/backends/database/models.py

@@ -89,6 +89,8 @@ class DBTaskResult(GenericBase[P, T], models.Model):
     started_at = models.DateTimeField(_("started at"), null=True)
     finished_at = models.DateTimeField(_("finished at"), null=True)
 
+    worker_id = models.UUIDField(_("worker id"), null=True)
+
     args_kwargs = models.JSONField(_("args kwargs"))
 
     priority = models.IntegerField(_("priority"), default=DEFAULT_PRIORITY)
@@ -163,6 +165,7 @@ def task_result(self) -> "TaskResult[T]":
             args=self.args_kwargs["args"],
             kwargs=self.args_kwargs["kwargs"],
             backend=self.backend_name,
+            worker_id=None if self.worker_id is None else str(self.worker_id),
         )
 
         object.__setattr__(task_result, "_exception_class", exception_class)
@@ -186,13 +189,14 @@ def task_name(self) -> str:
             return self.task_path
 
     @retry(backoff_delay=0)
-    def claim(self) -> None:
+    def claim(self, worker_id: uuid.UUID) -> None:
         """
-        Mark as job as being run
+        Mark as job as being run by a worker
         """
         self.status = ResultStatus.RUNNING
         self.started_at = timezone.now()
-        self.save(update_fields=["status", "started_at"])
+        self.worker_id = worker_id
+        self.save(update_fields=["status", "started_at", "worker_id"])
 
     @retry()
     def set_succeeded(self, return_value: Any) -> None:
@@ -201,13 +205,15 @@ def set_succeeded(self, return_value: Any) -> None:
         self.return_value = return_value
         self.exception_class_path = ""
         self.traceback = ""
+        self.worker_id = None
         self.save(
             update_fields=[
                 "status",
                 "return_value",
                 "finished_at",
                 "exception_class_path",
                 "traceback",
+                "worker_id",
             ]
         )
 
@@ -218,12 +224,14 @@ def set_failed(self, exc: BaseException) -> None:
         self.exception_class_path = get_module_path(type(exc))
         self.traceback = get_exception_traceback(exc)
         self.return_value = None
+        self.worker_id = None
         self.save(
             update_fields=[
                 "status",
                 "return_value",
                 "finished_at",
                 "exception_class_path",
                 "traceback",
+                "worker_id",
             ]
         )

django_tasks/backends/dummy.py

@@ -47,6 +47,7 @@ def enqueue(
             args=args,
             kwargs=kwargs,
             backend=self.alias,
+            worker_id=None,
         )
 
         if self._get_enqueue_on_commit_for_task(task) is not False:

django_tasks/backends/immediate.py

@@ -80,6 +80,7 @@ def enqueue(
             args=args,
             kwargs=kwargs,
             backend=self.alias,
+            worker_id=None,
         )
 
         if self._get_enqueue_on_commit_for_task(task) is not False:

django_tasks/task.py

@@ -41,6 +41,7 @@
     "started_at",
     "status",
     "enqueued_at",
+    "worker_id",
 }
 
 
@@ -249,6 +250,9 @@ class TaskResult(Generic[T]):
     backend: str
     """The name of the backend the task will run on"""
 
+    worker_id: Optional[str]
+    """The id of the worker running the task"""
+
     _exception_class: Optional[type[BaseException]] = field(init=False, default=None)
     _traceback: Optional[str] = field(init=False, default=None)
 

tests/tests/test_database_backend.py

@@ -415,6 +415,8 @@ def test_enqueue_logs(self) -> None:
     }
 )
 class DatabaseBackendWorkerTestCase(TransactionTestCase):
+    worker_id = uuid.uuid4()
+
     run_worker = staticmethod(
         partial(
             call_command,
@@ -423,6 +425,7 @@ class DatabaseBackendWorkerTestCase(TransactionTestCase):
             batch=True,
             interval=0,
             startup_delay=False,
+            worker_id=worker_id,
         )
     )
 
@@ -454,6 +457,7 @@ def test_run_enqueued_task(self) -> None:
                 self.assertGreaterEqual(result.started_at, result.enqueued_at)  # type:ignore[arg-type,misc]
                 self.assertGreaterEqual(result.finished_at, result.started_at)  # type:ignore[arg-type,misc]
                 self.assertEqual(result.status, ResultStatus.SUCCEEDED)
+                self.assertIsNone(result.worker_id)
 
                 self.assertEqual(DBTaskResult.objects.ready().count(), 0)
 
@@ -629,6 +633,15 @@ def test_fractional_interval(self) -> None:
 
         self.assertEqual(worker_class.mock_calls[0].kwargs["interval"], 0.1)
 
+    def test_invalid_worker_id(self) -> None:
+        output = StringIO()
+        with redirect_stderr(output):
+            with self.assertRaises(SystemExit):
+                execute_from_command_line(
+                    ["django-admin", "db_worker", "--worker-id", "123"]
+                )
+        self.assertIn("invalid UUID value", output.getvalue())
+
     def test_run_after(self) -> None:
         result = test_tasks.noop_task.using(
             run_after=timezone.now() + timedelta(hours=10)
@@ -702,10 +715,10 @@ def test_verbose_logging(self) -> None:
         self.assertEqual(
             stdout.getvalue().splitlines(),
             [
-                "Starting worker for queues=default",
+                f"Starting worker worker_id={self.worker_id} for queues=default",
                 f"Task id={result.id} path=tests.tasks.noop_task state=RUNNING",
                 f"Task id={result.id} path=tests.tasks.noop_task state=SUCCEEDED",
-                "No more tasks to run - exiting gracefully.",
+                f"No more tasks to run for worker_id={self.worker_id} - exiting gracefully.",
             ],
         )
 
@@ -1333,6 +1346,7 @@ def test_interrupt_signals(self) -> None:
 
                 result.refresh()
                 self.assertEqual(result.status, ResultStatus.RUNNING)
+                self.assertIsNotNone(result.worker_id)
 
                 process.send_signal(sig)
 
@@ -1343,6 +1357,7 @@ def test_interrupt_signals(self) -> None:
                 result.refresh()
 
                 self.assertEqual(result.status, ResultStatus.SUCCEEDED)
+                self.assertIsNone(result.worker_id)
 
     @skipIf(sys.platform == "win32", "Cannot emulate CTRL-C on Windows")
     def test_repeat_ctrl_c(self) -> None:
@@ -1355,6 +1370,7 @@ def test_repeat_ctrl_c(self) -> None:
 
         result.refresh()
         self.assertEqual(result.status, ResultStatus.RUNNING)
+        self.assertIsNotNone(result.worker_id)
 
         process.send_signal(signal.SIGINT)
 
@@ -1373,6 +1389,7 @@ def test_repeat_ctrl_c(self) -> None:
         result.refresh()
         self.assertEqual(result.status, ResultStatus.FAILED)
         self.assertEqual(result.exception_class, SystemExit)
+        self.assertIsNone(result.worker_id)
 
     @skipIf(sys.platform == "win32", "Windows doesn't support SIGKILL")
     def test_kill(self) -> None:
@@ -1411,6 +1428,7 @@ def test_system_exit_task(self) -> None:
         result.refresh()
         self.assertEqual(result.status, ResultStatus.FAILED)
         self.assertEqual(result.exception_class, SystemExit)
+        self.assertIsNone(result.worker_id)
 
     def test_keyboard_interrupt_task(self) -> None:
         result = test_tasks.failing_task_keyboard_interrupt.enqueue()
@@ -1423,6 +1441,7 @@ def test_keyboard_interrupt_task(self) -> None:
         result.refresh()
         self.assertEqual(result.status, ResultStatus.FAILED)
         self.assertEqual(result.exception_class, KeyboardInterrupt)
+        self.assertIsNone(result.worker_id)
 
     def test_multiple_workers(self) -> None:
         results = [test_tasks.sleep_for.enqueue(0.1) for _ in range(10)]
@@ -1439,6 +1458,7 @@ def test_multiple_workers(self) -> None:
         for result in results:
             result.refresh()
             self.assertEqual(result.status, ResultStatus.SUCCEEDED)
+            self.assertIsNone(result.worker_id)
 
         all_output = ""