test:suspend-worker

cunla · cunla · commit 3debe11006ff · 2025-04-11T11:59:53.000-04:00
diff --git a/docs/changelog.md b/docs/changelog.md
@@ -1,5 +1,9 @@
 # Changelog
 
+## v4.0.0 🌈
+
+See breaking changes in 4.0.0 beta versions.
+
 ## v4.0.0b3 🌈
 
 Refactor the code to make it more organized and easier to maintain. This includes:
diff --git a/scheduler/tests/test_worker/test_worker_commands.py b/scheduler/tests/test_worker/test_worker_commands.py
@@ -1,84 +1,42 @@
-from time import sleep
-
-from django.test import tag
-
 from scheduler.helpers.queues import get_queue
-from scheduler.redis_models import JobStatus, JobModel, WorkerModel
-from scheduler.tests.jobs import long_job, two_seconds_job
-from .. import testtools
+from scheduler.tests.jobs import test_job
 from ..test_views.base import BaseTestCase
-from ...worker.commands import KillWorkerCommand, send_command, StopJobCommand
+from ...redis_models import JobModel
+from ...worker import create_worker
+from ...worker.commands import send_command
+from ...worker.commands.suspend_worker import SuspendWorkCommand
 
 
-@tag("multiprocess")
 class WorkerCommandsTest(BaseTestCase):
-    def test_kill_job_command__current_job(self):
+    def test_stop_worker_command__green(self):
         # Arrange
-        queue = get_queue("django_tasks_scheduler_test")
-        job = queue.create_and_enqueue_job(long_job)
+        worker_name = "test"
+        queue = get_queue("default")
+        job = queue.create_and_enqueue_job(test_job)
         self.assertTrue(job.is_queued)
-        process, worker_name = testtools.run_worker_in_process("django_tasks_scheduler_test")
-        sleep(0.1)
-        job = JobModel.get(job.name, connection=queue.connection)
-        self.assertEqual(JobStatus.STARTED, job.status)
-
+        worker = create_worker("default", name=worker_name, burst=True, with_scheduler=False)
+        worker.worker_start()
         # Act
-        send_command(queue.connection, StopJobCommand(worker_name=worker_name, job_name=job.name))
+        send_command(queue.connection, SuspendWorkCommand(worker_name=worker_name))
+        worker.work()
 
         # Assert
+        self.assertTrue(job.is_queued)
+        self.assertTrue(worker._model.is_suspended)
 
-        process.terminate()
-        process.join(2)
-        process.kill()
-
-        job = JobModel.get(job.name, connection=queue.connection)
-        worker_model = WorkerModel.get(worker_name, connection=queue.connection)
-        self.assertEqual(job.name, worker_model.stopped_job_name)
-        self.assertEqual(job.name, worker_model.current_job_name)
-        self.assertEqual(0, worker_model.completed_jobs)
-        self.assertEqual(0, worker_model.failed_job_count)
-        self.assertEqual(0, worker_model.successful_job_count)
-        self.assertEqual(JobStatus.STOPPED, job.status)
-        self.assertNotIn(job.name, queue.queued_job_registry.all())
-
-    def test_kill_job_command__different_job(self):
+    def test_stop_worker_command__bad_worker_name(self):
         # Arrange
-        queue = get_queue("django_tasks_scheduler_test")
-        job = queue.create_and_enqueue_job(two_seconds_job)
+        worker_name = "test"
+        queue = get_queue("default")
+        job = queue.create_and_enqueue_job(test_job)
         self.assertTrue(job.is_queued)
-        process, worker_name = testtools.run_worker_in_process("django_tasks_scheduler_test")
-        sleep(0.2)
-        job = JobModel.get(job.name, connection=queue.connection)
-        self.assertEqual(JobStatus.STARTED, job.status)
-
+        worker = create_worker("default", name=worker_name, burst=True, with_scheduler=False)
+        worker.bootstrap()
         # Act
-        send_command(queue.connection, StopJobCommand(worker_name=worker_name, job_name=job.name + "1"))
-        sleep(0.1)
-        process.kill()
-        process.join()
+        send_command(queue.connection, SuspendWorkCommand(worker_name=worker_name + "1"))
+        worker.work()
+
         # Assert
+        self.assertFalse(worker._model.is_suspended)
         job = JobModel.get(job.name, connection=queue.connection)
-        self.assertEqual(JobStatus.STARTED, job.status)
-        self.assertNotIn(job.name, queue.queued_job_registry.all())
-        worker_model = WorkerModel.get(worker_name, connection=queue.connection)
-        self.assertEqual(0, worker_model.completed_jobs)
-        self.assertEqual(0, worker_model.failed_job_count)
-        self.assertEqual(0, worker_model.successful_job_count)
-        self.assertIsNone(worker_model.stopped_job_name)
-        self.assertEqual(job.name, worker_model.current_job_name)
-
-    def test_kill_worker_command(self):
-        queue = get_queue("django_tasks_scheduler_test")
-        process, worker_name = testtools.run_worker_in_process("django_tasks_scheduler_test")
-        sleep(0.1)
-        # act
-        send_command(queue.connection, KillWorkerCommand(worker_name=worker_name))
-        # assert
-        sleep(0.2)
-        process.kill()
-        process.join()
-        worker_model = WorkerModel.get(worker_name, connection=queue.connection)
-        self.assertEqual(0, worker_model.completed_jobs)
-        self.assertEqual(0, worker_model.failed_job_count)
-        self.assertEqual(0, worker_model.successful_job_count)
-        self.assertIsNotNone(worker_model.shutdown_requested_date)
+        self.assertFalse(job.is_queued)
diff --git a/scheduler/tests/test_worker/test_worker_commands_multiprocess.py b/scheduler/tests/test_worker/test_worker_commands_multiprocess.py
@@ -0,0 +1,84 @@
+from time import sleep
+
+from django.test import tag
+
+from scheduler.helpers.queues import get_queue
+from scheduler.redis_models import JobStatus, JobModel, WorkerModel
+from scheduler.tests.jobs import long_job, two_seconds_job
+from .. import testtools
+from ..test_views.base import BaseTestCase
+from ...worker.commands import KillWorkerCommand, send_command, StopJobCommand
+
+
+@tag("multiprocess")
+class WorkerCommandsTest(BaseTestCase):
+    def test_kill_job_command__current_job(self):
+        # Arrange
+        queue = get_queue("django_tasks_scheduler_test")
+        job = queue.create_and_enqueue_job(long_job)
+        self.assertTrue(job.is_queued)
+        process, worker_name = testtools.run_worker_in_process("django_tasks_scheduler_test")
+        sleep(0.1)
+        job = JobModel.get(job.name, connection=queue.connection)
+        self.assertEqual(JobStatus.STARTED, job.status)
+
+        # Act
+        send_command(queue.connection, StopJobCommand(worker_name=worker_name, job_name=job.name))
+
+        # Assert
+
+        process.terminate()
+        process.join(2)
+        process.kill()
+
+        job = JobModel.get(job.name, connection=queue.connection)
+        worker_model = WorkerModel.get(worker_name, connection=queue.connection)
+        self.assertEqual(job.name, worker_model.stopped_job_name)
+        self.assertEqual(job.name, worker_model.current_job_name)
+        self.assertEqual(0, worker_model.completed_jobs)
+        self.assertEqual(0, worker_model.failed_job_count)
+        self.assertEqual(0, worker_model.successful_job_count)
+        self.assertEqual(JobStatus.STOPPED, job.status)
+        self.assertNotIn(job.name, queue.queued_job_registry.all())
+
+    def test_kill_job_command__different_job(self):
+        # Arrange
+        queue = get_queue("django_tasks_scheduler_test")
+        job = queue.create_and_enqueue_job(two_seconds_job)
+        self.assertTrue(job.is_queued)
+        process, worker_name = testtools.run_worker_in_process("django_tasks_scheduler_test")
+        sleep(0.2)
+        job = JobModel.get(job.name, connection=queue.connection)
+        self.assertEqual(JobStatus.STARTED, job.status)
+
+        # Act
+        send_command(queue.connection, StopJobCommand(worker_name=worker_name, job_name=job.name + "1"))
+        sleep(0.1)
+        process.kill()
+        process.join()
+        # Assert
+        job = JobModel.get(job.name, connection=queue.connection)
+        self.assertEqual(JobStatus.STARTED, job.status)
+        self.assertNotIn(job.name, queue.queued_job_registry.all())
+        worker_model = WorkerModel.get(worker_name, connection=queue.connection)
+        self.assertEqual(0, worker_model.completed_jobs)
+        self.assertEqual(0, worker_model.failed_job_count)
+        self.assertEqual(0, worker_model.successful_job_count)
+        self.assertIsNone(worker_model.stopped_job_name)
+        self.assertEqual(job.name, worker_model.current_job_name)
+
+    def test_kill_worker_command(self):
+        queue = get_queue("django_tasks_scheduler_test")
+        process, worker_name = testtools.run_worker_in_process("django_tasks_scheduler_test")
+        sleep(0.1)
+        # act
+        send_command(queue.connection, KillWorkerCommand(worker_name=worker_name))
+        # assert
+        sleep(0.2)
+        process.kill()
+        process.join()
+        worker_model = WorkerModel.get(worker_name, connection=queue.connection)
+        self.assertEqual(0, worker_model.completed_jobs)
+        self.assertEqual(0, worker_model.failed_job_count)
+        self.assertEqual(0, worker_model.successful_job_count)
+        self.assertIsNotNone(worker_model.shutdown_requested_date)
diff --git a/scheduler/worker/commands/suspend_worker.py b/scheduler/worker/commands/suspend_worker.py
@@ -10,7 +10,7 @@ class SuspendWorkCommand(WorkerCommand):
     command_name = "suspend"
 
     def process_command(self, connection: ConnectionType) -> None:
-        logger.debug(f"Received command to suspend worker {self.job_name}")
+        logger.debug(f"Received command to suspend worker {self.worker_name}")
         worker_model = WorkerModel.get(self.worker_name, connection)
         if worker_model is None:
             logger.warning(f"Worker {self.worker_name} not found")
diff --git a/scheduler/worker/worker.py b/scheduler/worker/worker.py
@@ -109,18 +109,18 @@ def from_model(cls, model: WorkerModel) -> Self:
         return res
 
     def __init__(
-        self,
-        queues,
-        name: str,
-        connection: Optional[ConnectionType] = None,
-        maintenance_interval: int = SCHEDULER_CONFIG.DEFAULT_MAINTENANCE_TASK_INTERVAL,
-        job_monitoring_interval=SCHEDULER_CONFIG.DEFAULT_JOB_MONITORING_INTERVAL,
-        dequeue_strategy: DequeueStrategy = DequeueStrategy.DEFAULT,
-        disable_default_exception_handler: bool = False,
-        fork_job_execution: bool = True,
-        with_scheduler: bool = True,
-        burst: bool = False,
-        model: Optional[WorkerModel] = None,
+            self,
+            queues,
+            name: str,
+            connection: Optional[ConnectionType] = None,
+            maintenance_interval: int = SCHEDULER_CONFIG.DEFAULT_MAINTENANCE_TASK_INTERVAL,
+            job_monitoring_interval=SCHEDULER_CONFIG.DEFAULT_JOB_MONITORING_INTERVAL,
+            dequeue_strategy: DequeueStrategy = DequeueStrategy.DEFAULT,
+            disable_default_exception_handler: bool = False,
+            fork_job_execution: bool = True,
+            with_scheduler: bool = True,
+            burst: bool = False,
+            model: Optional[WorkerModel] = None,
     ):  # noqa
         self.fork_job_execution = fork_job_execution
         self.job_monitoring_interval = job_monitoring_interval
@@ -212,9 +212,9 @@ def _install_signal_handlers(self) -> None:
         signal.signal(signal.SIGTERM, self.request_stop)
 
     def work(
-        self,
-        max_jobs: Optional[int] = None,
-        max_idle_time: Optional[int] = None,
+            self,
+            max_jobs: Optional[int] = None,
+            max_idle_time: Optional[int] = None,
     ) -> bool:
         """Starts the work loop.
 
@@ -351,13 +351,13 @@ def _check_for_suspension(self, burst: bool) -> None:
         """Check to see if workers have been suspended by `rq suspend`"""
         before_state = None
         notified = False
-        while self._model.shutdown_requested_date is not None and self._model.is_suspended:
+        while self._model.is_suspended:
             if burst:
                 logger.info(
                     f"[Worker {self.name}/{self._pid}]: Suspended in burst mode, exiting, "
                     f"Note: There could still be unfinished jobs on the queue"
                 )
-                raise StopRequested
+                raise StopRequested()
 
             if not notified:
                 logger.info(f"[Worker {self.name}/{self._pid}]: Worker suspended, trigger ResumeCommand")
@@ -389,7 +389,7 @@ def run_maintenance_tasks(self):
             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[JobModel, Queue]:
         """Dequeues a job while maintaining the TTL.
         :param timeout: The timeout for the dequeue operation.
@@ -564,7 +564,7 @@ def reorder_queues(self, reference_queue: Queue):
             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)
@@ -595,7 +595,7 @@ def refresh(self, update_queues: bool = False):
         """
         self._model = WorkerModel.get(self.name, connection=self.connection)
         if self._model is None:
-            msg = f"[Worker {self.name}/{self._pid}]: Worker {self.name} not found, quitting..."
+            msg = f"[Worker {self.name}/{self._pid}]: Worker broker record for {self.name} not found, quitting..."
             logger.error(msg)
             raise WorkerNotFound(msg)
         if update_queues:
@@ -648,7 +648,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, rusage = self.wait_for_job_execution_process()
                 break
@@ -884,7 +884,7 @@ class RoundRobinWorker(Worker):
 
     def reorder_queues(self, reference_queue):
         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):
