fix:registries do not keep connections

Author: cunla

scheduler/helpers/queues/getters.py

@@ -46,13 +46,6 @@ def _get_connection(config: QueueConfiguration, use_strict_broker: bool = False)
     )
 
 
-def refresh_queue_connection(queue: Queue) -> None:
-    """Refreshes the connection of a given Queue"""
-    queue_settings = get_queue_configuration(queue.name)
-    connection = _get_connection(queue_settings)
-    queue.refresh_connection(connection)
-
-
 def get_queue(name: str = "default") -> Queue:
     """Returns an DjangoQueue using parameters defined in `SCHEDULER_QUEUES`"""
     queue_settings = get_queue_configuration(name)

scheduler/helpers/queues/queue_logic.py

@@ -82,21 +82,12 @@ def __init__(self, connection: ConnectionType, name: str, is_async: bool = True)
         self.connection: ConnectionType = connection
         self.name = name
         self._is_async = is_async
-        self.queued_job_registry = QueuedJobRegistry(connection=self.connection, name=self.name)
-        self.active_job_registry = ActiveJobRegistry(connection=self.connection, name=self.name)
-        self.failed_job_registry = FailedJobRegistry(connection=self.connection, name=self.name)
-        self.finished_job_registry = FinishedJobRegistry(connection=self.connection, name=self.name)
-        self.scheduled_job_registry = ScheduledJobRegistry(connection=self.connection, name=self.name)
-        self.canceled_job_registry = CanceledJobRegistry(connection=self.connection, name=self.name)
-
-    def refresh_connection(self, connection: ConnectionType) -> None:
-        self.connection = connection
-        self.queued_job_registry.connection = connection
-        self.active_job_registry.connection = connection
-        self.failed_job_registry.connection = connection
-        self.finished_job_registry.connection = connection
-        self.scheduled_job_registry.connection = connection
-        self.canceled_job_registry.connection = connection
+        self.queued_job_registry = QueuedJobRegistry(name=self.name)
+        self.active_job_registry = ActiveJobRegistry(name=self.name)
+        self.failed_job_registry = FailedJobRegistry(name=self.name)
+        self.finished_job_registry = FinishedJobRegistry(name=self.name)
+        self.scheduled_job_registry = ScheduledJobRegistry(name=self.name)
+        self.canceled_job_registry = CanceledJobRegistry(name=self.name)
 
     def __len__(self) -> int:
         return self.count
@@ -114,7 +105,7 @@ def clean_registries(self, timestamp: Optional[float] = None) -> None:
         Removed jobs are added to the global failed job queue.
         """
         before_score = timestamp or current_timestamp()
-        self.queued_job_registry.compact()
+        self.queued_job_registry.compact(self.connection)
         started_jobs: List[Tuple[str, float]] = self.active_job_registry.get_job_names_before(
             self.connection, before_score
         )
@@ -142,7 +133,7 @@ def clean_registries(self, timestamp: Optional[float] = None) -> None:
                 getattr(self, registry).cleanup(connection=self.connection, timestamp=before_score)
 
     def first_queued_job_name(self) -> Optional[str]:
-        return self.queued_job_registry.get_first()
+        return self.queued_job_registry.get_first(self.connection)
 
     @property
     def count(self) -> int:
@@ -160,12 +151,12 @@ def get_registry(self, name: str) -> JobNamesRegistry:
 
     def get_all_job_names(self) -> List[str]:
         all_job_names = list()
-        all_job_names.extend(self.queued_job_registry.all())
-        all_job_names.extend(self.finished_job_registry.all())
-        all_job_names.extend(self.active_job_registry.all())
-        all_job_names.extend(self.failed_job_registry.all())
-        all_job_names.extend(self.scheduled_job_registry.all())
-        all_job_names.extend(self.canceled_job_registry.all())
+        all_job_names.extend(self.queued_job_registry.all(self.connection))
+        all_job_names.extend(self.finished_job_registry.all(self.connection))
+        all_job_names.extend(self.active_job_registry.all(self.connection))
+        all_job_names.extend(self.failed_job_registry.all(self.connection))
+        all_job_names.extend(self.scheduled_job_registry.all(self.connection))
+        all_job_names.extend(self.canceled_job_registry.all(self.connection))
         res = list(filter(lambda job_name: JobModel.exists(job_name, self.connection), all_job_names))
         return res
 
@@ -307,7 +298,7 @@ def dequeue_any(
         while True:
             registries = [q.queued_job_registry for q in queues]
             for registry in registries:
-                registry.compact()
+                registry.compact(connection)
 
             registry_key, job_name = QueuedJobRegistry.pop(connection, registries, timeout)
             if job_name is None:
@@ -416,7 +407,7 @@ def enqueue_job(
             if at_front:
                 score = current_timestamp()
             else:
-                score = self.queued_job_registry.get_last_timestamp() or current_timestamp()
+                score = self.queued_job_registry.get_last_timestamp(self.connection) or current_timestamp()
             self.scheduled_job_registry.delete(connection=pipe, job_name=job_model.name)
             self.queued_job_registry.add(connection=pipe, score=score, job_name=job_model.name)
             pipe.execute()

scheduler/management/commands/delete_failed_executions.py

@@ -16,7 +16,7 @@ def add_arguments(self, parser: CommandParser) -> None:
 
     def handle(self, *args, **options):
         queue = get_queue(options.get("queue", "default"))
-        job_names = queue.failed_job_registry.all()
+        job_names = queue.failed_job_registry.all(queue.connection)
         jobs = JobModel.get_many(job_names, connection=queue.connection)
         func_name = options.get("func", None)
         if func_name is not None:

scheduler/models/task.py

@@ -186,12 +186,14 @@ def is_scheduled(self) -> bool:
         if self.job_name is None:  # no job_id => is not scheduled
             return False
         # check whether job_id is in scheduled/queued/active jobs
-        res = (
-            (self.job_name in self.rqueue.scheduled_job_registry.all())
-            or (self.job_name in self.rqueue.queued_job_registry.all())
-            or (self.job_name in self.rqueue.active_job_registry.all())
-        )
-        # If the job_id is not scheduled/queued/started,
+        with self.rqueue.connection.pipeline() as pipeline:
+            self.rqueue.scheduled_job_registry.exists(pipeline, self.job_name)
+            self.rqueue.queued_job_registry.exists(pipeline, self.job_name)
+            self.rqueue.active_job_registry.exists(pipeline, self.job_name)
+            results = pipeline.execute()
+            res = any([item is not None for item in results])
+
+        # If the job_name is not scheduled/queued/started,
         # update the job_id to None. (The job_id belongs to a previous run which is completed)
         if not res:
             self.job_name = None

scheduler/redis_models/registry/base_registry.py

@@ -32,54 +32,52 @@ def delete(self, connection: ConnectionType, job_name: str) -> None:
         logger.debug(f"[registry {self._key}] Deleting {job_name}")
         connection.zrem(self._key, job_name)
 
+    def exists(self, connection: ConnectionType, job_name: str) -> bool:
+        return connection.zrank(self._key, job_name) is not None
+
 
 class JobNamesRegistry(ZSetModel):
     _element_key_template: ClassVar[str] = ":registry:{}"
 
-    def __init__(self, connection: ConnectionType, name: str) -> None:
+    def __init__(self, name: str) -> None:
         super().__init__(name=name)
-        self.connection = connection
-
-    def __len__(self) -> int:
-        return self.count(self.connection)
-
-    def __contains__(self, item: str) -> bool:
-        return self.connection.zrank(self._key, item) is not None
 
-    def all(self, start: int = 0, end: int = -1) -> List[str]:
+    def all(self, connection: ConnectionType, start: int = 0, end: int = -1) -> List[str]:
         """Returns a list of all job names.
 
+        :param connection: Broker connection
         :param start: Start score/timestamp, default to 0.
         :param end: End score/timestamp, default to -1 (i.e., no max score).
         :returns: Returns a list of all job names with timestamp from start to end
         """
-        self.cleanup(self.connection)
-        res = [as_str(job_name) for job_name in self.connection.zrange(self._key, start, end)]
-        logger.debug(f"Getting jobs for registry {self._key}: {len(res)} found.")
+        self.cleanup(connection)
+        res = [as_str(job_name) for job_name in connection.zrange(self._key, start, end)]
+        logger.debug(f"Getting jobs for registry {self.key}: {len(res)} found.")
         return res
 
-    def all_with_timestamps(self, start: int = 0, end: int = -1) -> List[Tuple[str, float]]:
+    def all_with_timestamps(self, connection: ConnectionType, start: int = 0, end: int = -1) -> List[Tuple[str, float]]:
         """Returns a list of all job names with their timestamps.
 
+        :param connection: Broker connection
         :param start: Start score/timestamp, default to 0.
         :param end: End score/timestamp, default to -1 (i.e., no max score).
         :returns: Returns a list of all job names with timestamp from start to end
         """
-        self.cleanup(self.connection)
-        res = self.connection.zrange(self._key, start, end, withscores=True)
+        self.cleanup(connection)
+        res = connection.zrange(self._key, start, end, withscores=True)
         logger.debug(f"Getting jobs for registry {self._key}: {len(res)} found.")
         return [(as_str(job_name), timestamp) for job_name, timestamp in res]
 
-    def get_first(self) -> Optional[str]:
+    def get_first(self, connection: ConnectionType) -> Optional[str]:
         """Returns the first job in the registry."""
-        self.cleanup(self.connection)
-        first_job = self.connection.zrange(self._key, 0, 0)
+        self.cleanup(connection)
+        first_job = connection.zrange(self._key, 0, 0)
         return first_job[0].decode() if first_job else None
 
-    def get_last_timestamp(self) -> Optional[int]:
+    def get_last_timestamp(self, connection: ConnectionType) -> Optional[int]:
         """Returns the latest timestamp in the registry."""
-        self.cleanup(self.connection)
-        last_timestamp = self.connection.zrange(self._key, -1, -1, withscores=True)
+        self.cleanup(connection)
+        last_timestamp = connection.zrange(self._key, -1, -1, withscores=True)
         return int(last_timestamp[0][1]) if last_timestamp else None
 
     @property
@@ -88,7 +86,7 @@ def key(self) -> str:
 
     @classmethod
     def pop(
-        cls, connection: ConnectionType, registries: Sequence[Self], timeout: Optional[int]
+            cls, connection: ConnectionType, registries: Sequence[Self], timeout: Optional[int]
     ) -> Tuple[Optional[str], Optional[str]]:
         """Helper method to abstract away from some Redis API details
 

scheduler/redis_models/registry/queue_registries.py

@@ -17,18 +17,18 @@ def cleanup(self, connection: ConnectionType, timestamp: Optional[float] = None)
         and `all()` methods implemented in JobIdsRegistry."""
         pass
 
-    def compact(self) -> None:
+    def compact(self, connection: ConnectionType) -> None:
         """Removes all "dead" jobs from the queue by cycling through it, while guaranteeing FIFO semantics."""
-        jobs_with_ts = self.all_with_timestamps()
+        jobs_with_ts = self.all_with_timestamps(connection)
         for job_name, timestamp in jobs_with_ts:
-            if not JobModel.exists(job_name, self.connection):
-                self.delete(connection=self.connection, job_name=job_name)
+            if not JobModel.exists(job_name, connection):
+                self.delete(connection=connection, job_name=job_name)
 
-    def empty(self) -> None:
-        queued_jobs_count = self.count(connection=self.connection)
-        with self.connection.pipeline() as pipe:
+    def empty(self, connection: ConnectionType) -> None:
+        queued_jobs_count = self.count(connection=connection)
+        with connection.pipeline() as pipe:
             for offset in range(0, queued_jobs_count, 1000):
-                job_names = self.all(offset, 1000)
+                job_names = self.all(connection, offset, 1000)
                 for job_name in job_names:
                     self.delete(connection=pipe, job_name=job_name)
                 JobModel.delete_many(job_names, connection=pipe)
@@ -76,24 +76,26 @@ def schedule(self, connection: ConnectionType, job_name: str, scheduled_datetime
         timestamp = scheduled_datetime.timestamp()
         return self.add(connection=connection, job_name=job_name, score=timestamp)
 
-    def get_jobs_to_schedule(self, timestamp: int, chunk_size: int = 1000) -> List[str]:
+    def get_jobs_to_schedule(self, connection: ConnectionType, timestamp: int, chunk_size: int = 1000) -> List[str]:
         """Gets a list of job names that should be scheduled.
 
+        :param connection: Broker connection
         :param timestamp: timestamp/score of jobs in SortedSet.
         :param chunk_size: Max results to return.
         :returns: A list of job names
         """
-        jobs_to_schedule = self.connection.zrangebyscore(self._key, 0, max=timestamp, start=0, num=chunk_size)
+        jobs_to_schedule = connection.zrangebyscore(self._key, 0, max=timestamp, start=0, num=chunk_size)
         return [as_str(job_name) for job_name in jobs_to_schedule]
 
-    def get_scheduled_time(self, job_name: str) -> Optional[datetime]:
+    def get_scheduled_time(self, connection: ConnectionType, job_name: str) -> Optional[datetime]:
         """Returns datetime (UTC) at which job is scheduled to be enqueued
 
+        :param connection: Broker connection
         :param job_name: Job name
         :returns: The scheduled time as datetime object, or None if job is not found
         """
 
-        score: Optional[float] = self.connection.zscore(self._key, job_name)
+        score: Optional[float] = connection.zscore(self._key, job_name)
         if not score:
             return None
 

scheduler/templatetags/scheduler_tags.py

@@ -68,4 +68,4 @@ def job_runtime(job: JobModel):
 
 @register.filter
 def job_scheduled_time(job: JobModel, queue: Queue):
-    return queue.scheduled_job_registry.get_scheduled_time(job.name)
+    return queue.scheduled_job_registry.get_scheduled_time(queue.connection, job.name)

scheduler/tests/test_job_decorator.py

@@ -100,7 +100,7 @@ def test_job_decorator_different_queue(self):
 
     def _assert_job_with_func_and_props(self, queue_name, expected_func, expected_result_ttl, expected_timeout):
         queue = get_queue(queue_name)
-        jobs = JobModel.get_many(queue.queued_job_registry.all(), queue.connection)
+        jobs = JobModel.get_many(queue.queued_job_registry.all(queue.connection), queue.connection)
         self.assertEqual(1, len(jobs))
 
         j = jobs[0]

scheduler/tests/test_mgmt_commands/test_delete_failed_executions.py

@@ -12,9 +12,9 @@ def test_delete_failed_executions__delete_jobs(self):
         queue = get_queue("default")
         call_command("delete_failed_executions", queue="default")
         queue.create_and_enqueue_job(failing_job)
-        self.assertEqual(1, len(queue.queued_job_registry))
+        self.assertEqual(1, queue.queued_job_registry.count(queue.connection))
         worker = create_worker("default", burst=True)
         worker.work()
-        self.assertEqual(1, len(queue.failed_job_registry))
+        self.assertEqual(1, queue.failed_job_registry.count(queue.connection))
         call_command("delete_failed_executions", queue="default")
-        self.assertEqual(0, len(queue.failed_job_registry))
+        self.assertEqual(0, queue.failed_job_registry.count(queue.connection))

scheduler/tests/test_mgmt_commands/test_run_job.py

@@ -10,11 +10,11 @@
 class RunJobTest(TestCase):
     def test_run_job__should_schedule_job(self):
         queue = get_queue("default")
-        queue.queued_job_registry.empty()
+        queue.queued_job_registry.empty(queue.connection)
         func_name = f"{test_job.__module__}.{test_job.__name__}"
         # act
         call_command("run_job", func_name, queue="default")
         # assert
-        job_list = JobModel.get_many(queue.queued_job_registry.all(), queue.connection)
+        job_list = JobModel.get_many(queue.queued_job_registry.all(queue.connection), queue.connection)
         self.assertEqual(1, len(job_list))
         self.assertEqual(func_name + "()", job_list[0].get_call_string())