Periodic task tests, ready/finished runner events

Author: dcwatson

src/dbtasks/periodic.py

@@ -1,17 +1,24 @@
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import Callable
 
-from .schedule import Crontab, Schedule
+from .schedule import Crontab, Every, Schedule
 
 
 class Periodic:
+    schedule: Schedule
+
     def __init__(
         self,
-        spec: Schedule | str,
+        spec: Schedule | str | int | timedelta,
         args: list | tuple | Callable[[], list | tuple] | None = None,
         kwargs: dict | Callable[[], dict] | None = None,
     ):
-        self.schedule = spec if isinstance(spec, Schedule) else Crontab(spec)
+        if isinstance(spec, Schedule):
+            self.schedule = spec
+        elif isinstance(spec, str) and len(spec.split()) == 5:
+            self.schedule = Crontab(spec)
+        else:
+            self.schedule = Every(spec)
         self._args = args
         self._kwargs = kwargs
 

src/dbtasks/runner.py

@@ -43,7 +43,6 @@ def run_task(pk: str) -> TaskResultStatus:
         # Seems like this might be a welcome addition:
         # https://discuss.python.org/t/adding-finalizer-to-the-threading-library/54186
         connection.close()
-        pass
 
 
 @task
@@ -63,6 +62,7 @@ def __init__(
         worker_id: str | None = None,
         backend: str = DEFAULT_TASK_BACKEND_ALIAS,
         loop_delay: float = 0.5,
+        init_periodic: bool = True,
     ):
         self.workers = workers
         self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=workers)
@@ -78,15 +78,20 @@ def __init__(
         self.backend = task_backends[backend]
         if not isinstance(self.backend, DatabaseBackend):
             raise ImproperlyConfigured("Backend must be a `DatabaseBackend`")
+        # Signaled when the runner is ready and processing tasks.
+        self.ready = threading.Event()
         # Signaled when the runner should stop.
         self.stopsign = threading.Event()
-        # Signaled when the queue is empty (no READY tasks).
+        # Signaled when the runner is finished stopping.
+        self.finished = threading.Event()
+        # Signaled each time the queue is empty (no READY tasks).
         self.empty = threading.Event()
         # Covers `self.tasks`, `self.seen_modules`, and `self.processed` access.
         self.lock = threading.Lock()
         # Allows callers to block on a single task being completed.
         self.waiting: dict[str, threading.Event] = {}
         self.periodic: dict[str, Periodic] = {}
+        self.should_init_periodic = init_periodic
         if retain := self.backend.options.get("retain"):
             # If the task backend specifies a retention period, schedule a periodic task
             # to delete finished tasks older than that period.
@@ -156,20 +161,50 @@ def task_done(
         except Exception as ex:
             logger.info(f"Task {task_path} ({pk}) raised {ex}")
 
+        if was_periodic and (schedule := self.periodic.get(task_path)):
+            after = timezone.make_aware(schedule.next())
+            # Since this can run in the task's thread, we need to clean up the
+            # connection afterwards since it may not be closed at the end of `run`.
+            with connection.temporary_connection():
+                t = ScheduledTask.objects.create(
+                    task_path=task_path,
+                    args=schedule.args,
+                    kwargs=schedule.kwargs,
+                    backend=self.backend.alias,
+                    run_after=after,
+                    periodic=True,
+                )
+                logger.info(f"Re-scheduled {t} for {after}")
+
+        # If anyone is waiting on this task, wake them up.
         if event := self.waiting.get(pk):
             event.set()
 
-        if was_periodic and (schedule := self.periodic.get(task_path)):
-            after = timezone.make_aware(schedule.next())
-            t = ScheduledTask.objects.create(
-                task_path=task_path,
-                args=schedule.args,
-                kwargs=schedule.kwargs,
-                backend=self.backend.alias,
-                run_after=after,
-                periodic=True,
-            )
-            logger.info(f"Re-scheduled {t} for {after}")
+    def submit_task(self, task: ScheduledTask, start: bool = True) -> TaskResult:
+        """
+        Submits a `ScheduledTask` for execution, marking it as RUNNING and setting its
+        `started_at` timestamp if `start=True`.
+        """
+        if start:
+            task.status = TaskResultStatus.RUNNING
+            task.started_at = timezone.now()
+            task.worker_ids.append(self.worker_id)
+            task.save(update_fields=["status", "started_at", "worker_ids"])
+        logger.debug(f"Submitting {task} for execution")
+        f = self.executor.submit(run_task, task.task_id)
+        with self.lock:
+            # Keep track of task modules we've seen, so we can reload them.
+            self.seen_modules.add(task.task_path.rsplit(".", 1)[0])
+            self.tasks[task.task_id] = f
+        f.add_done_callback(
+            functools.partial(
+                self.task_done,
+                task.task_id,
+                task.task_path,
+                task.periodic,
+            ),
+        )
+        return task.result
 
     def schedule_tasks(self) -> float:
         """
@@ -193,20 +228,8 @@ def schedule_tasks(self) -> float:
         self.empty.clear()
 
         for t in tasks:
-            logger.debug(f"Submitting {t} for execution")
-            f = self.executor.submit(run_task, t.task_id)
-            with self.lock:
-                # Keep track of task modules we've seen, so we can reload them.
-                self.seen_modules.add(t.task_path.rsplit(".", 1)[0])
-                self.tasks[t.task_id] = f
-            f.add_done_callback(
-                functools.partial(
-                    self.task_done,
-                    t.task_id,
-                    t.task_path,
-                    t.periodic,
-                ),
-            )
+            # get_tasks starts all of the returned tasks atomically, no need to here.
+            self.submit_task(t, start=False)
 
         if len(tasks) >= available:
             # We got a full batch, try again immediately.
@@ -244,7 +267,12 @@ def run(self):
         """
         logger.info(f"Starting task runner with {self.workers} workers")
         self.processed = 0
-        self.init_periodic()
+        if self.should_init_periodic:
+            with transaction.atomic(durable=True):
+                self.init_periodic()
+                transaction.on_commit(self.ready.set)
+        else:
+            self.ready.set()
         try:
             while not self.stopsign.is_set():
                 delay = self.schedule_tasks()
@@ -254,6 +282,7 @@ def run(self):
         finally:
             self.executor.shutdown()
             connection.close()
+        self.finished.set()
 
     def wait_for(self, result: TaskResult, timeout: float | None = None) -> bool:
         """

src/dbtasks/testing.py

@@ -15,13 +15,18 @@ def setUpClass(cls):
         super().setUpClass()
         # Run with a very short loop delay to speed up tests. None of our test tasks
         # take very long, so there's not much point in waiting aside from not flooding
-        # the database with queries for new tasks.
-        cls.runner = Runner(workers=os.cpu_count() - 1, loop_delay=0.01)
-        cls.runner_thread = threading.Thread(target=cls.runner.run)
-        cls.runner_thread.start()
+        # the database with queries for new tasks. Also we don't initialize periodic
+        # tasks - individual tests can call `self.runner.init_periodic()`.
+        cls.runner = Runner(
+            workers=max(1, (os.cpu_count() or 4) - 1),
+            loop_delay=0.01,
+            init_periodic=False,
+        )
+        threading.Thread(target=cls.runner.run).start()
+        cls.runner.ready.wait()
 
     @classmethod
     def tearDownClass(cls):
         super().tearDownClass()
         cls.runner.stop()
-        cls.runner_thread.join()
+        cls.runner.finished.wait()

tests/settings.py

@@ -1,4 +1,3 @@
-import datetime
 import os
 
 INSTALLED_APPS = ["dbtasks", "tests"]
@@ -45,8 +44,9 @@
         "BACKEND": "dbtasks.backend.DatabaseBackend",
         "OPTIONS": {
             "immediate": False,
-            "retain": datetime.timedelta(days=7),
-            "periodic": {},
+            "periodic": {
+                "tests.tasks.maintenance": "1h",
+            },
         },
     },
 }

tests/tasks.py

@@ -9,3 +9,8 @@
 def send_mail(to: str, message: str):
     logger.info(f"Sending mail to {to}: {message}")
     return {"sent": True}
+
+
+@task
+def maintenance():
+    pass

tests/tests/test_tasks.py

@@ -1,6 +1,8 @@
 from django.tasks import TaskResult, TaskResultStatus
 
-from ..tasks import send_mail
+from dbtasks.models import ScheduledTask
+
+from ..tasks import maintenance, send_mail
 from ..utils import LoggedRunnerTestCase
 
 
@@ -23,3 +25,41 @@ def test_lots_of_tasks(self):
         self.runner.wait()
         self.assertEqual(self.runner.processed, 100)
         self.assertEqual(set(self.task_logs["INFO"]), expected)
+
+    def test_periodic(self):
+        self.assertIn(maintenance.module_path, self.runner.periodic)
+        self.runner.init_periodic()
+        first = ScheduledTask.objects.get(
+            task_path=maintenance.module_path,
+            status=TaskResultStatus.READY,
+            periodic=True,
+        )
+        # Run the initial scheduled task manually.
+        result = self.runner.submit_task(first)
+        self.assertTrue(self.runner.wait_for(result))
+        first.refresh_from_db()
+        self.assertEqual(result.status, TaskResultStatus.SUCCESSFUL)
+        self.assertEqual(first.status, TaskResultStatus.SUCCESSFUL)
+        # Make sure a second periodic task got scheduled when the first completed.
+        second = ScheduledTask.objects.get(
+            task_path=maintenance.module_path,
+            status=TaskResultStatus.READY,
+            periodic=True,
+        )
+        self.assertNotEqual(first.id, second.id)
+        self.assertGreater(second.enqueued_at, first.enqueued_at)
+        # Now run the maintenance task manually, not periodically.
+        result: TaskResult = maintenance.enqueue()
+        self.assertTrue(self.runner.wait_for(result))
+        manual = ScheduledTask.objects.get(pk=result.id)
+        # Make sure the manual run was not marked periodic, and that no new tasks were
+        # automatically scheduled afterwards.
+        self.assertFalse(manual.periodic)
+        self.assertEqual(
+            ScheduledTask.objects.filter(
+                task_path=maintenance.module_path, status=TaskResultStatus.READY
+            ).count(),
+            1,
+        )
+        second.refresh_from_db()
+        self.assertEqual(second.status, TaskResultStatus.READY)

tests/utils.py

@@ -7,9 +7,10 @@ class LoggedRunnerTestCase(RunnerTestCase):
     def setUp(self):
         with stashlock:
             stash.clear()
+        super().setUp()
 
     def tearDown(self):
-        pass
+        super().tearDown()
 
     @property
     def logs(self) -> dict: