Track which worker executed a task (#136)

Author: RealOrangeOne

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

@@ -23,6 +23,7 @@
 from django_tasks.exceptions import InvalidTaskBackendError
 from django_tasks.signals import task_finished, task_started
 from django_tasks.task import DEFAULT_QUEUE_NAME
+from django_tasks.utils import get_random_id
 
 package_logger = logging.getLogger("django_tasks")
 logger = logging.getLogger("django_tasks.backends.database.db_worker")
@@ -38,6 +39,7 @@ def __init__(
         backend_name: str,
         startup_delay: bool,
         max_tasks: Optional[int],
+        worker_id: str,
     ):
         self.queue_names = queue_names
         self.process_all_queues = "*" in queue_names
@@ -51,6 +53,8 @@ def __init__(
         self.running_task = False
         self._run_tasks = 0
 
+        self.worker_id = worker_id
+
     def shutdown(self, signum: int, frame: Optional[FrameType]) -> None:
         if not self.running:
             logger.warning(
@@ -82,11 +86,17 @@ def reset_signals(self) -> None:
         if hasattr(signal, "SIGQUIT"):
             signal.signal(signal.SIGQUIT, signal.SIG_DFL)
 
-    def start(self) -> None:
-        logger.info("Starting worker for queues=%s", ",".join(self.queue_names))
+    def run(self) -> None:
+        self.configure_signals()
+
+        logger.info(
+            "Starting worker worker_id=%s 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
+            # Add a random small delay before starting to avoid a thundering herd
             time.sleep(random.random())
 
         while self.running:
@@ -109,19 +119,24 @@ 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)
 
             if self.batch and task_result is None:
                 # If we're running in "batch" mode, terminate the loop (and thus the worker)
-                logger.info("No more tasks to run - exiting gracefully.")
+                logger.info(
+                    "No more tasks to run for worker_id=%s - exiting gracefully.",
+                    self.worker_id,
+                )
                 return None
 
             if self.max_tasks is not None and self._run_tasks >= self.max_tasks:
                 logger.info(
-                    "Run maximum tasks (%d) - exiting gracefully.", self._run_tasks
+                    "Run maximum tasks (%d) on worker=%s - exiting gracefully.",
+                    self._run_tasks,
+                    self.worker_id,
                 )
                 return None
 
@@ -199,6 +214,14 @@ def valid_max_tasks(val: str) -> int:
     return num
 
 
+def validate_worker_id(val: str) -> str:
+    if not val:
+        raise ArgumentTypeError("Worker id must not be empty")
+    if len(val) > 64:
+        raise ArgumentTypeError("Worker ids must be shorter than 64 characters")
+    return val
+
+
 class Command(BaseCommand):
     help = "Run a database background worker"
 
@@ -249,6 +272,13 @@ def add_arguments(self, parser: ArgumentParser) -> None:
             type=valid_max_tasks,
             help="If provided, the maximum number of tasks the worker will execute before exiting.",
         )
+        parser.add_argument(
+            "--worker-id",
+            nargs="?",
+            type=validate_worker_id,
+            help="Worker id. MUST be unique across worker pool (default: auto-generate)",
+            default=get_random_id(),
+        )
 
     def configure_logging(self, verbosity: int) -> None:
         if verbosity == 0:
@@ -274,6 +304,7 @@ def handle(
         startup_delay: bool,
         reload: bool,
         max_tasks: Optional[int],
+        worker_id: str,
         **options: dict,
     ) -> None:
         self.configure_logging(verbosity)
@@ -291,14 +322,15 @@ def handle(
             backend_name=backend_name,
             startup_delay=startup_delay,
             max_tasks=max_tasks,
+            worker_id=worker_id,
         )
 
         if reload:
             if os.environ.get(DJANGO_AUTORELOAD_ENV) == "true":
                 # Only the child process should configure its signals
                 worker.configure_signals()
 
-            run_with_reloader(worker.start)
+            run_with_reloader(worker.run)
         else:
             worker.configure_signals()
-            worker.start()
+            worker.run()

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

@@ -0,0 +1,20 @@
+# Generated by Django 5.1.6 on 2025-02-16 15:55
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        (
+            "django_tasks_database",
+            "0017_remove_dbtaskresult_django_task_new_ordering_idx_and_more",
+        ),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="dbtaskresult",
+            name="worker_id",
+            field=models.CharField(default="", max_length=64, verbose_name="worker id"),
+        ),
+    ]

django_tasks/backends/database/migrations/0019_remove_dbtaskresult_worker_id_and_more.py

@@ -0,0 +1,21 @@
+# Generated by Django 4.2 on 2025-06-20 12:03
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("django_tasks_database", "0018_dbtaskresult_worker_id"),
+    ]
+
+    operations = [
+        migrations.RemoveField(
+            model_name="dbtaskresult",
+            name="worker_id",
+        ),
+        migrations.AddField(
+            model_name="dbtaskresult",
+            name="worker_ids",
+            field=models.JSONField(default=list, verbose_name="worker id"),
+        ),
+    ]

django_tasks/backends/database/models.py

@@ -99,6 +99,7 @@ class DBTaskResult(GenericBase[P, T], models.Model):
     priority = models.IntegerField(_("priority"), default=DEFAULT_PRIORITY)
 
     task_path = models.TextField(_("task path"))
+    worker_ids = models.JSONField(_("worker id"), default=list)
 
     queue_name = models.CharField(
         _("queue name"), default=DEFAULT_QUEUE_NAME, max_length=32
@@ -177,6 +178,7 @@ def task_result(self) -> "TaskResult[T]":
             kwargs=self.args_kwargs["kwargs"],
             backend=self.backend_name,
             errors=[],
+            worker_ids=self.worker_ids,
         )
 
         if self.status == ResultStatus.FAILED:
@@ -206,13 +208,14 @@ def task_name(self) -> str:
             return self.task_path
 
     @retry(backoff_delay=0)
-    def claim(self) -> None:
+    def claim(self, worker_id: str) -> None:
         """
         Mark as job as being run
         """
         self.status = ResultStatus.RUNNING
         self.started_at = timezone.now()
-        self.save(update_fields=["status", "started_at"])
+        self.worker_ids = [*self.worker_ids, worker_id]
+        self.save(update_fields=["status", "started_at", "worker_ids"])
 
     @retry()
     def set_succeeded(self, return_value: Any) -> None:

django_tasks/backends/dummy.py

@@ -52,6 +52,7 @@ def enqueue(
             kwargs=kwargs,
             backend=self.alias,
             errors=[],
+            worker_ids=[],
         )
 
         if self._get_enqueue_on_commit_for_task(task) is not False:

django_tasks/backends/immediate.py

@@ -29,6 +29,11 @@
 class ImmediateBackend(BaseTaskBackend):
     supports_async_task = True
 
+    def __init__(self, alias: str, params: dict):
+        super().__init__(alias, params)
+
+        self.worker_id = get_random_id()
+
     def _execute_task(self, task_result: TaskResult) -> None:
         """
         Execute the task for the given `TaskResult`, mutating it with the outcome
@@ -45,6 +50,7 @@ def _execute_task(self, task_result: TaskResult) -> None:
         object.__setattr__(task_result, "status", ResultStatus.RUNNING)
         object.__setattr__(task_result, "started_at", timezone.now())
         object.__setattr__(task_result, "last_attempted_at", timezone.now())
+        task_result.worker_ids.append(self.worker_id)
         task_started.send(sender=type(self), task_result=task_result)
 
         try:
@@ -98,6 +104,7 @@ def enqueue(
             kwargs=kwargs,
             backend=self.alias,
             errors=[],
+            worker_ids=[],
         )
 
         if self._get_enqueue_on_commit_for_task(task) is not False:

django_tasks/backends/rq.py

@@ -53,6 +53,10 @@ class Job(BaseJob):
     def perform(self) -> Any:
         task_result = self.into_task_result()
 
+        assert self.worker_name is not None
+        self.meta.setdefault("_django_tasks_worker_ids", []).append(self.worker_name)
+        self.save_meta()  # type: ignore[no-untyped-call]
+
         task_started.send(type(task_result.task.get_backend()), task_result=task_result)
 
         return super().perform()
@@ -103,6 +107,7 @@ def into_task_result(self) -> TaskResult:
             kwargs=self.kwargs,
             backend=self.meta["backend_name"],
             errors=[],
+            worker_ids=self.meta.get("_django_tasks_worker_ids", []),
         )
 
         exception_classes = self.meta.get("_django_tasks_exceptions", []).copy()
@@ -188,6 +193,7 @@ def enqueue(
             kwargs=kwargs,
             backend=self.alias,
             errors=[],
+            worker_ids=[],
         )
 
         job = queue.create_job(
@@ -209,7 +215,6 @@ def save_result() -> None:
                 job = queue.schedule_job(job, task.run_after)
 
             object.__setattr__(task_result, "enqueued_at", job.enqueued_at)
-
             task_enqueued.send(type(self), task_result=task_result)
 
         if self._get_enqueue_on_commit_for_task(task):

django_tasks/task.py

@@ -42,6 +42,7 @@
     "last_attempted_at",
     "status",
     "enqueued_at",
+    "worker_ids",
 }
 
 
@@ -283,6 +284,9 @@ class TaskResult(Generic[T]):
     errors: list[TaskError]
     """The errors raised when running the task"""
 
+    worker_ids: list[str]
+    """The workers which have processed the task"""
+
     _return_value: Optional[T] = field(init=False, default=None)
 
     @property
@@ -307,12 +311,7 @@ def is_finished(self) -> bool:
 
     @property
     def attempts(self) -> int:
-        attempts = len(self.errors)
-
-        if self.status == ResultStatus.SUCCEEDED:
-            attempts += 1
-
-        return attempts
+        return len(self.worker_ids)
 
     def refresh(self) -> None:
         """

django_tasks/utils.py

@@ -66,7 +66,7 @@ def get_exception_traceback(exc: BaseException) -> str:
 
 def get_random_id() -> str:
     """
-    Return a random string for use as a task id.
+    Return a random string for use as a task or worker id.
 
     Whilst 64 characters is the max, just use 32 as a sensible middle-ground.