fix:registries do not keep connections

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/tests/test_worker/test_scheduler.py

@@ -27,7 +27,7 @@ def test_scheduler_schedules_tasks(self):
             self.assertFalse(task.rqueue.queued_job_registry.exists(task.rqueue.connection, task.job_name))
             self.assertTrue(task.rqueue.scheduled_job_registry.exists(task.rqueue.connection, task.job_name))
 
-            scheduler = WorkerScheduler([task.rqueue], worker_name="fake-worker", connection=task.rqueue.connection)
+            scheduler = WorkerScheduler([task.rqueue], worker_name="fake-worker")
 
             # act
             traveller.move_to(50)

scheduler/worker/scheduler.py

@@ -3,18 +3,19 @@
 import traceback
 from datetime import datetime
 from enum import Enum
+from logging import DEBUG, INFO
 from threading import Thread
 from typing import List, Set, Optional, Sequence, Dict
 
 import django
 
 from scheduler.helpers.queues import Queue
 from scheduler.helpers.queues import get_queue
+from scheduler.helpers.queues.getters import get_queue_connection
 from scheduler.helpers.utils import current_timestamp
 from scheduler.models import Task
 from scheduler.redis_models import SchedulerLock, JobModel, ScheduledJobRegistry
 from scheduler.settings import SCHEDULER_CONFIG, logger
-from scheduler.types import ConnectionType
 
 
 class SchedulerStatus(str, Enum):
@@ -31,22 +32,15 @@ def _reschedule_tasks() -> None:
 
 
 class WorkerScheduler:
-    def __init__(
-        self,
-        queues: Sequence[Queue],
-        connection: ConnectionType,
-        worker_name: str,
-        interval: Optional[int] = None,
-    ) -> None:
-        interval = interval or SCHEDULER_CONFIG.SCHEDULER_INTERVAL
+    def __init__(self, queues: Sequence[Queue], worker_name: str, interval: Optional[int] = None) -> None:
         self._queues = queues
+        if len(queues) == 0:
+            raise ValueError("At least one queue must be provided to WorkerScheduler")
         self._scheduled_job_registries: List[ScheduledJobRegistry] = []
         self.lock_acquisition_time: Optional[datetime] = None
-        self._pool_class = connection.connection_pool.connection_class
-        self._pool_kwargs = connection.connection_pool.connection_kwargs.copy()
         self._locks: Dict[str, SchedulerLock] = dict()
-        self.connection = connection
-        self.interval = interval
+        self.connection = get_queue_connection(queues[0].name)
+        self.interval = interval or SCHEDULER_CONFIG.SCHEDULER_INTERVAL
         self._stop_requested = False
         self.status = SchedulerStatus.STOPPED
         self._thread: Optional[Thread] = None
@@ -57,6 +51,9 @@ def __init__(
     def pid(self) -> Optional[int]:
         return self._pid
 
+    def log(self, level: int, message: str, *args, **kwargs) -> None:
+        logger.log(level, f"[Scheduler {self.worker_name}/{self._pid}]: {message}", *args, **kwargs)
+
     def _should_reacquire_locks(self) -> bool:
         """Returns True if lock_acquisition_time is longer than 10 minutes ago"""
         if not self.lock_acquisition_time:
@@ -70,12 +67,10 @@ def _acquire_locks(self) -> Set[str]:
         if self.pid is None:
             self._pid = os.getpid()
         queue_names = [queue.name for queue in self._queues]
-        logger.debug(
-            f"""[Scheduler {self.worker_name}/{self.pid}] Trying to acquire locks for {", ".join(queue_names)}"""
-        )
+        self.log(DEBUG, f"""Trying to acquire locks for {", ".join(queue_names)}""")
         for queue in self._queues:
             lock = SchedulerLock(queue.name)
-            if lock.acquire(self.pid, connection=queue.connection, expire=self.interval + 60):
+            if lock.acquire(self.pid, connection=self.connection, expire=self.interval + 60):
                 self._locks[queue.name] = lock
                 successful_locks.add(queue.name)
 
@@ -85,7 +80,7 @@ def _acquire_locks(self) -> Set[str]:
         for queue_name in self._locks:
             queue = get_queue(queue_name)
             self._scheduled_job_registries.append(queue.scheduled_job_registry)
-        logger.debug(f"[Scheduler {self.worker_name}/{self.pid}] Locks acquired for {', '.join(self._locks.keys())}")
+        self.log(DEBUG, f"Locks acquired for {', '.join(self._locks.keys())}")
         return successful_locks
 
     def start(self) -> None:
@@ -98,24 +93,22 @@ def start(self) -> None:
 
     def request_stop_and_wait(self) -> None:
         """Toggle self._stop_requested that's checked on every loop"""
-        logger.debug(f"[Scheduler {self.worker_name}/{self.pid}] Stop Scheduler requested")
+        self.log(DEBUG, f"Stop Scheduler requested")
         self._stop_requested = True
         if self._thread is not None:
             self._thread.join()
 
     def heartbeat(self) -> None:
         """Updates the TTL on scheduler keys and the locks"""
         lock_keys = ", ".join(self._locks.keys())
-        logger.debug(f"[Scheduler {self.worker_name}/{self.pid}] Scheduler updating lock for queue {lock_keys}")
+        self.log(DEBUG, f"Scheduler updating lock for queue {lock_keys}")
         with self.connection.pipeline() as pipeline:
             for lock in self._locks.values():
                 lock.expire(self.connection, expire=self.interval + 60)
             pipeline.execute()
 
     def stop(self) -> None:
-        logger.info(
-            f"[Scheduler {self.worker_name}/{self.pid}] Stopping scheduler, releasing locks for {', '.join(self._locks.keys())}..."
-        )
+        self.log(INFO, f"Stopping scheduler, releasing locks for {', '.join(self._locks.keys())}...")
         self.release_locks()
         self.status = SchedulerStatus.STOPPED
 
@@ -128,7 +121,7 @@ def release_locks(self) -> None:
 
     def work(self) -> None:
         queue_names = [queue.name for queue in self._queues]
-        logger.info(f"""[Scheduler {self.worker_name}/{self.pid}] Scheduler for {", ".join(queue_names)} started""")
+        self.log(INFO, f"""Scheduler for {", ".join(queue_names)} started""")
         django.setup()
 
         while True:

scheduler/worker/worker.py

@@ -97,16 +97,16 @@ def from_model(cls, model: WorkerModel) -> Self:
         return res
 
     def __init__(
-            self,
-            queues: Iterable[Union[str, Queue]],
-            name: str,
-            maintenance_interval: int = SCHEDULER_CONFIG.DEFAULT_MAINTENANCE_TASK_INTERVAL,
-            job_monitoring_interval: int = SCHEDULER_CONFIG.DEFAULT_JOB_MONITORING_INTERVAL,
-            dequeue_strategy: DequeueStrategy = DequeueStrategy.DEFAULT,
-            fork_job_execution: bool = True,
-            with_scheduler: bool = True,
-            burst: bool = False,
-            model: Optional[WorkerModel] = None,
+        self,
+        queues: Iterable[Union[str, Queue]],
+        name: str,
+        maintenance_interval: int = SCHEDULER_CONFIG.DEFAULT_MAINTENANCE_TASK_INTERVAL,
+        job_monitoring_interval: int = SCHEDULER_CONFIG.DEFAULT_JOB_MONITORING_INTERVAL,
+        dequeue_strategy: DequeueStrategy = DequeueStrategy.DEFAULT,
+        fork_job_execution: bool = True,
+        with_scheduler: bool = True,
+        burst: bool = False,
+        model: Optional[WorkerModel] = None,
     ) -> None:
         self.fork_job_execution = fork_job_execution
         self.job_monitoring_interval: int = job_monitoring_interval
@@ -290,7 +290,7 @@ def bootstrap(self) -> None:
         self.log(INFO, f"Worker {self.name} started with PID {os.getpid()}")
         self._command_listener.start()
         if self.with_scheduler:
-            self.scheduler = WorkerScheduler(self.queues, worker_name=self.name, connection=self.connection)
+            self.scheduler = WorkerScheduler(self.queues, worker_name=self.name)
             self.scheduler.start()
             self._model.has_scheduler = True
             self._model.save(connection=self.connection)
@@ -345,7 +345,7 @@ def run_maintenance_tasks(self) -> None:
         self._model.save(connection=self.connection)
 
     def dequeue_job_and_maintain_ttl(
-            self, timeout: Optional[int], max_idle_time: Optional[int] = None
+        self, timeout: Optional[int], max_idle_time: Optional[int] = None
     ) -> Tuple[Optional[JobModel], Optional[Queue]]:
         """Dequeues a job while maintaining the TTL.
         :param timeout: The timeout for the dequeue operation.
@@ -504,7 +504,7 @@ def reorder_queues(self, reference_queue: Queue) -> None:
             return
         if self._dequeue_strategy == DequeueStrategy.ROUND_ROBIN:
             pos = self._ordered_queues.index(reference_queue)
-            self._ordered_queues = self._ordered_queues[pos + 1:] + self._ordered_queues[: pos + 1]
+            self._ordered_queues = self._ordered_queues[pos + 1 :] + self._ordered_queues[: pos + 1]
             return
         if self._dequeue_strategy == DequeueStrategy.RANDOM:
             shuffle(self._ordered_queues)
@@ -584,7 +584,7 @@ def monitor_job_execution_process(self, job: JobModel, queue: Queue) -> None:
         while True:
             try:
                 with SCHEDULER_CONFIG.DEATH_PENALTY_CLASS(
-                        self.job_monitoring_interval, JobExecutionMonitorTimeoutException
+                    self.job_monitoring_interval, JobExecutionMonitorTimeoutException
                 ):
                     retpid, ret_val = self._wait_for_job_execution_process()
                 break
@@ -801,7 +801,7 @@ class RoundRobinWorker(Worker):
 
     def reorder_queues(self, reference_queue: Queue) -> None:
         pos = self._ordered_queues.index(reference_queue)
-        self._ordered_queues = self._ordered_queues[pos + 1:] + self._ordered_queues[: pos + 1]
+        self._ordered_queues = self._ordered_queues[pos + 1 :] + self._ordered_queues[: pos + 1]
 
 
 class RandomWorker(Worker):