Add worker pings

Author: RealOrangeOne

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

@@ -3,6 +3,7 @@
 import random
 import signal
 import sys
+import threading
 import time
 import uuid
 from argparse import ArgumentParser, ArgumentTypeError
@@ -16,7 +17,7 @@
 
 from django_tasks import DEFAULT_TASK_BACKEND_ALIAS, tasks
 from django_tasks.backends.database.backend import DatabaseBackend
-from django_tasks.backends.database.models import DBTaskResult
+from django_tasks.backends.database.models import DBTaskResult, DBWorkerPing
 from django_tasks.backends.database.utils import exclusive_transaction
 from django_tasks.exceptions import InvalidTaskBackendError
 from django_tasks.signals import task_finished
@@ -46,11 +47,29 @@ def __init__(
 
         self.worker_id = worker_id
 
-        self.running = True
+        self.ping_thread = threading.Thread(target=self.run_ping)
+
+        self.should_exit = threading.Event()
         self.running_task = False
 
+    def run_ping(self) -> None:
+        try:
+            while True:
+                try:
+                    DBWorkerPing.ping(self.worker_id)
+                except Exception:
+                    logger.exception(
+                        "Error sending worker ping worker_id=%s", self.worker_id
+                    )
+                if self.should_exit.wait(timeout=10):
+                    break
+        finally:
+            # Close any connections opened in this thread
+            for conn in connections.all():
+                conn.close()
+
     def shutdown(self, signum: int, frame: Optional[FrameType]) -> None:
-        if not self.running:
+        if self.should_exit.is_set():
             logger.warning(
                 "Received %s - terminating current task.", signal.strsignal(signum)
             )
@@ -60,7 +79,7 @@ def shutdown(self, signum: int, frame: Optional[FrameType]) -> None:
             "Received %s - shutting down gracefully... (press Ctrl+C again to force)",
             signal.strsignal(signum),
         )
-        self.running = False
+        self.should_exit.set()
 
         if not self.running_task:
             # If we're not currently running a task, exit immediately.
@@ -73,8 +92,9 @@ def configure_signals(self) -> None:
         if hasattr(signal, "SIGQUIT"):
             signal.signal(signal.SIGQUIT, self.shutdown)
 
-    def start(self) -> None:
+    def run(self) -> None:
         self.configure_signals()
+        self.ping_thread.start()
 
         logger.info(
             "Starting worker worker_id=%s for queues=%s",
@@ -86,7 +106,7 @@ def start(self) -> None:
             # Add a random small delay before starting the loop to avoid a thundering herd
             time.sleep(random.random())
 
-        while self.running:
+        while not self.should_exit.is_set():
             tasks = DBTaskResult.objects.ready().filter(backend_name=self.backend_name)
             if not self.process_all_queues:
                 tasks = tasks.filter(queue_name__in=self.queue_names)
@@ -117,16 +137,13 @@ def start(self) -> None:
             finally:
                 self.running_task = False
 
-                for conn in connections.all(initialized_only=True):
-                    conn.close()
-
             if self.batch and task_result is None:
                 # If we're running in "batch" mode, terminate the loop (and thus the worker)
                 return None
 
-            # If ctrl-c has just interrupted a task, self.running was cleared,
+            # If ctrl-c has just interrupted a task, self.should_exit was cleared,
             # and we should not sleep, but rather exit immediately.
-            if self.running and not task_result:
+            if not self.should_exit.is_set() and not task_result:
                 # Wait before checking for another task
                 time.sleep(self.interval)
 
@@ -168,6 +185,14 @@ def run_task(self, db_task_result: DBTaskResult) -> None:
                     task_result=task_result,
                 )
 
+        for conn in connections.all(initialized_only=True):
+            conn.close()
+
+    def stop(self) -> None:
+        self.should_exit.set()
+        DBWorkerPing.cleanup_ping(self.worker_id)
+        self.ping_thread.join()
+
 
 def valid_backend_name(val: str) -> str:
     try:
@@ -271,7 +296,10 @@ def handle(
             worker_id=worker_id,
         )
 
-        worker.start()
+        try:
+            worker.run()
+        finally:
+            worker.stop()
 
         if batch:
             logger.info(

django_tasks/backends/database/migrations/0016_dbworkerping.py

@@ -0,0 +1,29 @@
+# Generated by Django 5.1.5 on 2025-02-01 13:08
+
+import uuid
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("django_tasks_database", "0015_dbtaskresult_worker_id"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="DBWorkerPing",
+            fields=[
+                (
+                    "worker_id",
+                    models.UUIDField(
+                        default=uuid.uuid4,
+                        editable=False,
+                        primary_key=True,
+                        serialize=False,
+                    ),
+                ),
+                ("last_ping", models.DateTimeField()),
+            ],
+        ),
+    ]

django_tasks/backends/database/models.py

@@ -1,5 +1,6 @@
 import logging
 import uuid
+from datetime import timedelta
 from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar
 
 import django
@@ -235,3 +236,32 @@ def set_failed(self, exc: BaseException) -> None:
                 "worker_id",
             ]
         )
+
+
+class DBWorkerPingQuerySet(models.QuerySet):
+    STALE_TIME = 600
+
+    def stale(self) -> "DBWorkerPingQuerySet":
+        return self.filter(
+            last_ping__lte=timezone.now() - timedelta(seconds=self.STALE_TIME)
+        )
+
+
+class DBWorkerPing(models.Model):
+    worker_id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
+
+    last_ping = models.DateTimeField()
+
+    objects = DBWorkerPingQuerySet.as_manager()
+
+    @classmethod
+    @retry()
+    def ping(cls, worker_id: str) -> None:
+        cls.objects.update_or_create(
+            worker_id=worker_id, defaults={"last_ping": timezone.now()}
+        )
+
+    @classmethod
+    @retry()
+    def cleanup_ping(cls, worker_id: str) -> None:
+        cls.objects.filter(worker_id=worker_id).delete()

tests/settings.py

@@ -59,11 +59,11 @@
     )
 }
 
-# Set exclusive transactions in 5.1+
-if django.VERSION >= (5, 1) and "sqlite" in DATABASES["default"]["ENGINE"]:
-    DATABASES["default"].setdefault("OPTIONS", {})["transaction_mode"] = "EXCLUSIVE"
-
 if "sqlite" in DATABASES["default"]["ENGINE"]:
+    if django.VERSION >= (5, 1):
+        # Set exclusive transactions in 5.1+
+        DATABASES["default"].setdefault("OPTIONS", {})["transaction_mode"] = "EXCLUSIVE"
+
     DATABASES["default"]["TEST"] = {"NAME": os.path.join(BASE_DIR, "db-test.sqlite3")}
 
 

tests/tests/test_database_backend.py

@@ -447,7 +447,7 @@ def test_run_enqueued_task(self) -> None:
 
                 self.assertEqual(result.status, ResultStatus.NEW)
 
-                with self.assertNumQueries(9 if connection.vendor == "mysql" else 8):
+                with self.assertNumQueries(12 if connection.vendor == "mysql" else 11):
                     self.run_worker()
 
                 self.assertEqual(result.status, ResultStatus.NEW)
@@ -468,28 +468,28 @@ def test_batch_processes_all_tasks(self) -> None:
 
         self.assertEqual(DBTaskResult.objects.ready().count(), 4)
 
-        with self.assertNumQueries(27 if connection.vendor == "mysql" else 23):
+        with self.assertNumQueries(30 if connection.vendor == "mysql" else 26):
             self.run_worker()
 
         self.assertEqual(DBTaskResult.objects.ready().count(), 0)
         self.assertEqual(DBTaskResult.objects.succeeded().count(), 3)
         self.assertEqual(DBTaskResult.objects.failed().count(), 1)
 
     def test_no_tasks(self) -> None:
-        with self.assertNumQueries(3):
+        with self.assertNumQueries(6):
             self.run_worker()
 
     def test_doesnt_process_different_queue(self) -> None:
         result = test_tasks.noop_task.using(queue_name="queue-1").enqueue()
 
         self.assertEqual(DBTaskResult.objects.ready().count(), 1)
 
-        with self.assertNumQueries(3):
+        with self.assertNumQueries(6):
             self.run_worker()
 
         self.assertEqual(DBTaskResult.objects.ready().count(), 1)
 
-        with self.assertNumQueries(9 if connection.vendor == "mysql" else 8):
+        with self.assertNumQueries(12 if connection.vendor == "mysql" else 11):
             self.run_worker(queue_name=result.task.queue_name)
 
         self.assertEqual(DBTaskResult.objects.ready().count(), 0)
@@ -499,12 +499,12 @@ def test_process_all_queues(self) -> None:
 
         self.assertEqual(DBTaskResult.objects.ready().count(), 1)
 
-        with self.assertNumQueries(3):
+        with self.assertNumQueries(6):
             self.run_worker()
 
         self.assertEqual(DBTaskResult.objects.ready().count(), 1)
 
-        with self.assertNumQueries(9 if connection.vendor == "mysql" else 8):
+        with self.assertNumQueries(12 if connection.vendor == "mysql" else 11):
             self.run_worker(queue_name="*")
 
         self.assertEqual(DBTaskResult.objects.ready().count(), 0)
@@ -519,7 +519,7 @@ def test_failing_task(self) -> None:
         with self.assertRaisesMessage(ValueError, "Task has not finished yet"):
             result.traceback  # noqa: B018
 
-        with self.assertNumQueries(9 if connection.vendor == "mysql" else 8):
+        with self.assertNumQueries(12 if connection.vendor == "mysql" else 11):
             self.run_worker()
 
         self.assertEqual(result.status, ResultStatus.NEW)
@@ -553,7 +553,7 @@ def test_complex_exception(self) -> None:
         with self.assertRaisesMessage(ValueError, "Task has not finished"):
             result.traceback  # noqa: B018
 
-        with self.assertNumQueries(9 if connection.vendor == "mysql" else 8):
+        with self.assertNumQueries(12 if connection.vendor == "mysql" else 11):
             self.run_worker()
 
         self.assertEqual(result.status, ResultStatus.NEW)
@@ -577,12 +577,12 @@ def test_doesnt_process_different_backend(self) -> None:
 
         self.assertEqual(DBTaskResult.objects.ready().count(), 1)
 
-        with self.assertNumQueries(3):
+        with self.assertNumQueries(6):
             self.run_worker(backend_name="dummy")
 
         self.assertEqual(DBTaskResult.objects.ready().count(), 1)
 
-        with self.assertNumQueries(9 if connection.vendor == "mysql" else 8):
+        with self.assertNumQueries(12 if connection.vendor == "mysql" else 11):
             self.run_worker(backend_name=result.backend)
 
         self.assertEqual(DBTaskResult.objects.ready().count(), 0)
@@ -650,7 +650,7 @@ def test_run_after(self) -> None:
         self.assertEqual(DBTaskResult.objects.count(), 1)
         self.assertEqual(DBTaskResult.objects.ready().count(), 0)
 
-        with self.assertNumQueries(3):
+        with self.assertNumQueries(6):
             self.run_worker()
 
         self.assertEqual(DBTaskResult.objects.count(), 1)
@@ -661,7 +661,7 @@ def test_run_after(self) -> None:
 
         self.assertEqual(DBTaskResult.objects.ready().count(), 1)
 
-        with self.assertNumQueries(9 if connection.vendor == "mysql" else 8):
+        with self.assertNumQueries(12 if connection.vendor == "mysql" else 11):
             self.run_worker()
 
         self.assertEqual(DBTaskResult.objects.ready().count(), 0)
@@ -711,7 +711,7 @@ def test_verbose_logging(self) -> None:
 
         stdout = StringIO()
         self.run_worker(verbosity=3, stdout=stdout, stderr=stdout)
-
+        self.maxDiff = None
         self.assertEqual(
             stdout.getvalue().splitlines(),
             [