Merge branch 'master' into save_safety

Author: cunla

docs/changelog.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## v4.0.8 🌈
+
+### 🐛 Bug Fixes
+
+- fix:registries do not keep connections
+
 ## v4.0.7 🌈
 
 ### 🐛 Bug Fixes

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "django-tasks-scheduler"
-version = "4.0.7"
+version = "4.0.8"
 description = "An async job scheduler for django using redis/valkey brokers"
 authors = [{ name = "Daniel Moran", email = "[email protected]" }]
 requires-python = ">=3.10"

scheduler/helpers/queues/__init__.py

@@ -3,8 +3,8 @@
     "InvalidJobOperation",
     "get_queue",
     "get_all_workers",
-    "perform_job",
+    "queue_perform_job",
 ]
 
 from .getters import get_queue, get_all_workers
-from .queue_logic import Queue, InvalidJobOperation, perform_job
+from .queue_logic import Queue, InvalidJobOperation, queue_perform_job

scheduler/helpers/queues/getters.py

@@ -46,11 +46,10 @@ 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)
+def get_queue_connection(queue_name: str) -> ConnectionType:
+    queue_settings = get_queue_configuration(queue_name)
     connection = _get_connection(queue_settings)
-    queue.refresh_connection(connection)
+    return connection
 
 
 def get_queue(name: str = "default") -> Queue:

scheduler/helpers/queues/queue_logic.py

@@ -34,7 +34,7 @@ class NoSuchRegistryError(Exception):
     pass
 
 
-def perform_job(job_model: JobModel, connection: ConnectionType) -> Any:  # noqa
+def queue_perform_job(job_model: JobModel, connection: ConnectionType) -> Any:  # noqa
     """The main execution method. Invokes the job function with the job arguments.
 
     :returns: The job's return value
@@ -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
 
@@ -280,7 +271,7 @@ def run_sync(self, job: JobModel) -> JobModel:
         """Run a job synchronously, meaning on the same process the method was called."""
         job.prepare_for_execution("sync", self.active_job_registry, self.connection)
         try:
-            result = perform_job(job, self.connection)
+            result = queue_perform_job(job, self.connection)
             self.job_handle_success(job, result=result, job_info_ttl=job.job_info_ttl, result_ttl=job.success_ttl)
         except Exception as e:  # noqa
             logger.warning(f"Job {job.name} failed with exception: {e}")
@@ -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/helpers/sentry_integration.py

@@ -32,7 +32,7 @@ def setup_once() -> None:
         version = parse_version(scheduler.__version__)
         _check_minimum_version(SentryIntegration, version)
 
-        old_perform_job = Worker.perform_job
+        old_perform_job = Worker.worker_perform_job
 
         @ensure_integration_enabled(SentryIntegration, old_perform_job)
         def sentry_patched_perform_job(self: Any, job_model: JobModel, *args: Queue, **kwargs: Any) -> bool:
@@ -65,7 +65,7 @@ def sentry_patched_perform_job(self: Any, job_model: JobModel, *args: Queue, **k
 
             return rv
 
-        Worker.perform_job = sentry_patched_perform_job  # type: ignore[method-assign]
+        Worker.worker_perform_job = sentry_patched_perform_job  # type: ignore[method-assign]
 
         old_handle_exception = Worker.handle_exception
 

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

@@ -25,59 +25,59 @@ def count(self, connection: ConnectionType) -> int:
         return connection.zcard(self._key)
 
     def add(self, connection: ConnectionType, job_name: str, score: float, update_existing_only: bool = False) -> int:
+        logger.debug(f"[registry {self._key}] Adding {job_name} / {score}")
         return connection.zadd(self._key, {job_name: float(score)}, xx=update_existing_only)
 
     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]:
-        """Returns the last timestamp in the registry."""
-        self.cleanup(self.connection)
-        last_timestamp = self.connection.zrange(self._key, -1, -1, withscores=True)
+    def get_last_timestamp(self, connection: ConnectionType) -> Optional[int]:
+        """Returns the latest timestamp in the registry."""
+        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

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