fix(scheduler): remove registry entries whose job model is gone

Author: javiermunooz

scheduler/helpers/queues/queue_logic.py

@@ -131,7 +131,10 @@ def clean_registries(self, timestamp: float | None = None) -> None:
 
         for job_name, job_score in started_jobs:
             job = JobModel.get(job_name, connection=self.connection)
-            if job is None or not job.has_failure_callback or job_score + job.timeout > before_score:
+            if job is None:
+                self.active_job_registry.delete(connection=self.connection, job_name=job_name)
+                continue
+            if not job.has_failure_callback or job_score + job.timeout > before_score:
                 continue
 
             logger.debug(f"Running failure callbacks for {job.name}")

scheduler/tests/test_internals.py

@@ -6,6 +6,8 @@
 from django.utils import timezone
 
 from scheduler.helpers.callback import Callback, CallbackSetupError
+from scheduler.helpers.queues import get_queue
+from scheduler.helpers.utils import current_timestamp
 from scheduler.models import TaskType, get_next_cron_time, get_scheduled_task
 from scheduler.tests.testtools import SchedulerBaseCase, task_factory
 
@@ -47,6 +49,16 @@ def test_callback_bad_arguments(self):
         self.assertEqual(str(cm.exception), "Callback `func` must be a string or function, received 1")
 
 
+class TestCleanRegistries(SchedulerBaseCase):
+    def test_active_registry_entry_without_job_is_removed(self):
+        queue = get_queue("default")
+        queue.active_job_registry.add(queue.connection, "orphan-job-name", current_timestamp() - 3600)
+
+        queue.clean_registries()
+
+        self.assertFalse(queue.active_job_registry.exists(queue.connection, "orphan-job-name"))
+
+
 class TestConfSettings(SchedulerBaseCase):
     @override_settings(SCHEDULER_CONFIG=[])
     def test_conf_settings__bad_scheduler_config(self):

scheduler/tests/test_worker/test_scheduler.py

@@ -3,7 +3,9 @@
 import time_machine
 from django.utils import timezone
 
+from scheduler.helpers.utils import current_timestamp
 from scheduler.models import TaskType
+from scheduler.redis_models import JobModel
 from scheduler.settings import SCHEDULER_CONFIG
 from scheduler.tests.testtools import SchedulerBaseCase, task_factory
 from scheduler.worker import WorkerScheduler, create_worker
@@ -37,3 +39,32 @@ def test_scheduler_schedules_tasks(self):
             self.assertIsNotNone(task.job_name)
             self.assertTrue(task.rqueue.queued_job_registry.exists(task.rqueue.connection, task.job_name))
             self.assertFalse(task.rqueue.scheduled_job_registry.exists(task.rqueue.connection, task.job_name))
+
+    def test_scheduler_removes_scheduled_registry_entry_without_job(self):
+        # arrange
+        task = task_factory(TaskType.CRON)
+        job_name = task.job_name
+        self.assertIsNotNone(job_name)
+        connection = task.rqueue.connection
+        registry = task.rqueue.scheduled_job_registry
+        connection.delete(JobModel.key_for(job_name))
+        registry.add(connection, job_name, current_timestamp() - 10)
+
+        scheduler = WorkerScheduler([task.rqueue], worker_name="fake-worker")
+        scheduler._acquire_locks()
+
+        # act
+        scheduler.enqueue_scheduled_jobs()
+
+        # assert
+        self.assertFalse(registry.exists(connection, job_name))
+        self.assertFalse(task.rqueue.queued_job_registry.exists(connection, job_name))
+
+        # act: the next pass schedules the task again
+        scheduler.enqueue_scheduled_jobs()
+
+        # assert
+        task.refresh_from_db()
+        self.assertIsNotNone(task.job_name)
+        self.assertNotEqual(job_name, task.job_name)
+        self.assertTrue(registry.exists(connection, task.job_name))

scheduler/worker/scheduler.py

@@ -148,9 +148,11 @@ def enqueue_scheduled_jobs(self) -> None:
             queue = get_queue(registry.name)
             jobs = JobModel.get_many(job_names, connection=self.connection)
             with self.connection.pipeline() as pipeline:
-                for job in jobs:
+                for job_name, job in zip(job_names, jobs):
                     if job is not None:
                         queue.enqueue_job(job, pipeline=pipeline, at_front=job.at_front)
+                    else:
+                        registry.delete(connection=pipeline, job_name=job_name)
                 pipeline.execute()
         self.status = SchedulerStatus.STARTED