fix #291

Author: cunla

scheduler/helpers/queues/getters.py

@@ -46,6 +46,13 @@ 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

@@ -91,6 +91,15 @@ def __init__(self, connection: ConnectionType, name: str, is_async: bool = True)
         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
+
     def __len__(self) -> int:
         return self.count
 

scheduler/worker/worker.py

@@ -24,6 +24,7 @@
 from scheduler.types import ConnectionType, TimeoutErrorTypes, ConnectionErrorTypes, WatchErrorTypes, ResponseErrorTypes
 from .commands import WorkerCommandsChannelListener
 from .scheduler import WorkerScheduler, SchedulerStatus
+from ..helpers.queues.getters import refresh_queue_connection
 from ..redis_models.lock import QueueLock
 from ..redis_models.worker import WorkerStatus
 
@@ -116,14 +117,14 @@ def __init__(
         self.job_monitoring_interval: int = job_monitoring_interval
         self.maintenance_interval = maintenance_interval
 
-        connection = self._set_connection(connection)
-        self.connection = connection
+        self.connection = self._set_connection(connection)
 
         self.queues = [
             (Queue(name=q, connection=connection) if isinstance(q, str) else q) for q in _ensure_list(queues)
         ]
         self.name: str = name
-        self._validate_name_uniqueness()
+        if model is None:
+            self._validate_name_uniqueness()
         self._ordered_queues = self.queues[:]
 
         self._is_job_execution_process: bool = False
@@ -155,6 +156,9 @@ def __init__(
 
     @property
     def _pid(self) -> int:
+        if self._model is None:
+            logger.debug(f"[Worker {self.name}]: Worker model is None, returning 0 as PID")
+            return 0
         return self._model.pid
 
     def should_run_maintenance_tasks(self) -> bool:
@@ -600,16 +604,19 @@ def fork_job_execution_process(self, job: JobModel, queue: Queue) -> None:
         os.environ["SCHEDULER_JOB_NAME"] = job.name
         if child_pid == 0:  # Child process/Job executor process to run the job
             os.setsid()
-            self._model.job_execution_process_pid = os.getpid()
-            self._model.save(connection=self.connection)
-            self.execute_in_separate_process(job, queue)
+            refresh_queue_connection(queue)
+            self.connection = queue.connection
+            self._model.set_field("job_execution_process_pid", os.getpid(), connection=queue.connection)
+            worker = Worker.from_model(self._model)
+            worker.execute_in_separate_process(job, queue)
             os._exit(0)  # just in case
         else:  # Parent worker process
             logger.debug(
                 f"[Worker {self.name}/{self._pid}]: Forking job execution process, job_execution_process_pid={child_pid}"
             )
+            refresh_queue_connection(queue)
             self._model.job_execution_process_pid = child_pid
-            self._model.save(connection=self.connection)
+            self._model.save(connection=queue.connection)
             self.procline(f"Forked {child_pid} at {time.time()}")
 
     def get_heartbeat_ttl(self, job: JobModel) -> int:
@@ -732,6 +739,7 @@ def execute_in_separate_process(self, job: JobModel, queue: Queue) -> None:
         random.seed()
         self.setup_job_execution_process_signals()
         self._is_job_execution_process = True
+        job = JobModel.get(job.name, self.connection)
         try:
             self.perform_job(job, queue)
         except:  # noqa
@@ -756,7 +764,7 @@ def worker_before_execution(self, job: JobModel, connection: ConnectionType) ->
         self._model.current_job_working_time = 0
         self._model.job_execution_process_pid = current_pid
         heartbeat_ttl = self.get_heartbeat_ttl(job)
-        self._model.heartbeat(self.connection, heartbeat_ttl)
+        self._model.heartbeat(connection, heartbeat_ttl)
         self.procline(
             f"[Worker {self.name}/{self._pid}]: Processing {job.func_name} from {job.queue_name} since {time.time()}"
         )
@@ -792,7 +800,7 @@ def handle_job_success(self, job: JobModel, return_value: Any, queue: Queue) ->
                     self._model.completed_jobs += 1
                     if job.started_at is not None and job.ended_at is not None:
                         self._model.total_working_time_ms += (job.ended_at - job.started_at).microseconds / 1000.0
-                    self._model.save(connection=self.connection)
+                    self._model.save(connection=pipeline)
 
                     job.expire(job.success_ttl, connection=pipeline)
                     logger.debug(f"[Worker {self.name}/{self._pid}]: Removing job {job.name} from active_job_registry")