feat: signal handler to wait for task completion before shutting down (#345)

* feat: k8s signal handler to wait for task completion before shutting down

* refactor: signal handler to be more generic

* refactor: rename to job_completion_wait, split signal handler to a new
function
tests: add test for new signal handler flow

* fix: use asyncio.TimeoutError to support python3.7

* Update arq/worker.py

Co-authored-by: Samuel Colvin <[email protected]>

Author: JonasKs

arq/worker.py

@@ -1,4 +1,5 @@
 import asyncio
+import contextlib
 import inspect
 import logging
 import signal
@@ -153,6 +154,10 @@ class Worker:
     :param after_job_end: coroutine function to run after job has ended and results have been recorded
     :param handle_signals: default true, register signal handlers,
       set to false when running inside other async framework
+    :param job_completion_wait: time to wait before cancelling tasks after a signal.
+      Useful together with ``terminationGracePeriodSeconds`` in kubernetes,
+      when you want to make the pod complete jobs before shutting down.
+      The worker will not pick new tasks while waiting for shut down.
     :param max_jobs: maximum number of jobs to run at a time
     :param job_timeout: default job timeout (max run time)
     :param keep_result: default duration to keep job results for
@@ -192,6 +197,7 @@ def __init__(
         on_job_end: Optional['StartupShutdown'] = None,
         after_job_end: Optional['StartupShutdown'] = None,
         handle_signals: bool = True,
+        job_completion_wait: int = 0,
         max_jobs: int = 10,
         job_timeout: 'SecondsTimedelta' = 300,
         keep_result: 'SecondsTimedelta' = 3600,
@@ -263,13 +269,19 @@ def __init__(
         self._last_health_check: float = 0
         self._last_health_check_log: Optional[str] = None
         self._handle_signals = handle_signals
+        self._job_completion_wait = job_completion_wait
         if self._handle_signals:
-            self._add_signal_handler(signal.SIGINT, self.handle_sig)
-            self._add_signal_handler(signal.SIGTERM, self.handle_sig)
+            if self._job_completion_wait:
+                self._add_signal_handler(signal.SIGINT, self.handle_sig_wait_for_completion)
+                self._add_signal_handler(signal.SIGTERM, self.handle_sig_wait_for_completion)
+            else:
+                self._add_signal_handler(signal.SIGINT, self.handle_sig)
+                self._add_signal_handler(signal.SIGTERM, self.handle_sig)
         self.on_stop: Optional[Callable[[Signals], None]] = None
         # whether or not to retry jobs on Retry and CancelledError
         self.retry_jobs = retry_jobs
         self.allow_abort_jobs = allow_abort_jobs
+        self.allow_pick_jobs: bool = True
         self.aborting_tasks: Set[str] = set()
         self.max_burst_jobs = max_burst_jobs
         self.job_serializer = job_serializer
@@ -361,23 +373,23 @@ async def _poll_iteration(self) -> None:
             if burst_jobs_remaining < 1:
                 return
             count = min(burst_jobs_remaining, count)
+        if self.allow_pick_jobs:
+            async with self.sem:  # don't bother with zrangebyscore until we have "space" to run the jobs
+                now = timestamp_ms()
+                job_ids = await self.pool.zrangebyscore(
+                    self.queue_name, min=float('-inf'), start=self._queue_read_offset, num=count, max=now
+                )
 
-        async with self.sem:  # don't bother with zrangebyscore until we have "space" to run the jobs
-            now = timestamp_ms()
-            job_ids = await self.pool.zrangebyscore(
-                self.queue_name, min=float('-inf'), start=self._queue_read_offset, num=count, max=now
-            )
-
-        await self.start_jobs(job_ids)
+            await self.start_jobs(job_ids)
 
-        if self.allow_abort_jobs:
-            await self._cancel_aborted_jobs()
+            if self.allow_abort_jobs:
+                await self._cancel_aborted_jobs()
 
-        for job_id, t in list(self.tasks.items()):
-            if t.done():
-                del self.tasks[job_id]
-                # required to make sure errors in run_job get propagated
-                t.result()
+            for job_id, t in list(self.tasks.items()):
+                if t.done():
+                    del self.tasks[job_id]
+                    # required to make sure errors in run_job get propagated
+                    t.result()
 
         await self.heart_beat()
 
@@ -757,6 +769,55 @@ def _add_signal_handler(self, signum: Signals, handler: Callable[[Signals], None
     def _jobs_started(self) -> int:
         return self.jobs_complete + self.jobs_retried + self.jobs_failed + len(self.tasks)
 
+    async def _sleep_until_tasks_complete(self) -> None:
+        """
+        Sleeps until all tasks are done. Used together with asyncio.wait_for()
+        """
+        while len(self.tasks):
+            await asyncio.sleep(0.1)
+
+    async def _wait_for_tasks_to_complete(self, signum: Signals) -> None:
+        """
+        Wait for tasks to complete, until `wait_for_job_completion_on_signal_second` has been reached.
+        """
+        with contextlib.suppress(asyncio.TimeoutError):
+            await asyncio.wait_for(
+                self._sleep_until_tasks_complete(),
+                self._job_completion_wait,
+            )
+        logger.info(
+            'shutdown on %s, wait complete ◆ %d jobs complete ◆ %d failed ◆ %d retries ◆ %d ongoing to cancel',
+            signum.name,
+            self.jobs_complete,
+            self.jobs_failed,
+            self.jobs_retried,
+            sum(not t.done() for t in self.tasks.values()),
+        )
+        for t in self.tasks.values():
+            if not t.done():
+                t.cancel()
+        self.main_task and self.main_task.cancel()
+        self.on_stop and self.on_stop(signum)
+
+    def handle_sig_wait_for_completion(self, signum: Signals) -> None:
+        """
+        Alternative signal handler that allow tasks to complete within a given time before shutting down the worker.
+        Time can be configured using `wait_for_job_completion_on_signal_second`.
+        The worker will stop picking jobs when signal has been received.
+        """
+        sig = Signals(signum)
+        logger.info('Setting allow_pick_jobs to `False`')
+        self.allow_pick_jobs = False
+        logger.info(
+            'shutdown on %s ◆ %d jobs complete ◆ %d failed ◆ %d retries ◆ %d to be completed',
+            sig.name,
+            self.jobs_complete,
+            self.jobs_failed,
+            self.jobs_retried,
+            len(self.tasks),
+        )
+        self.loop.create_task(self._wait_for_tasks_to_complete(signum=sig))
+
     def handle_sig(self, signum: Signals) -> None:
         sig = Signals(signum)
         logger.info(

tests/test_worker.py

@@ -110,6 +110,34 @@ async def test_handle_no_sig(caplog):
     assert worker.tasks[1].cancel.call_count == 1
 
 
+async def test_worker_signal_completes_job_before_shutting_down(caplog, arq_redis: ArqRedis, worker):
+    caplog.set_level(logging.INFO)
+
+    async def sleep_job(ctx, time):
+        await asyncio.sleep(time)
+
+    await arq_redis.enqueue_job('sleep_job', 0.2, _job_id='short_sleep')  # should be cancelled
+    await arq_redis.enqueue_job('sleep_job', 5, _job_id='long_sleep')  # should be cancelled
+    worker = worker(
+        functions=[func(sleep_job, name='sleep_job', max_tries=1)],
+        job_completion_wait=0.4,
+        job_timeout=10,
+    )
+    assert worker.jobs_complete == 0
+    asyncio.create_task(worker.main())
+    await asyncio.sleep(0.1)
+    worker.handle_sig_wait_for_completion(signal.SIGINT)
+    assert worker.allow_pick_jobs is False
+    await asyncio.sleep(0.5)
+    logs = [rec.message for rec in caplog.records]
+    assert 'shutdown on SIGINT ◆ 0 jobs complete ◆ 0 failed ◆ 0 retries ◆ 2 to be completed' in logs
+    assert 'shutdown on SIGINT, wait complete ◆ 1 jobs complete ◆ 0 failed ◆ 0 retries ◆ 1 ongoing to cancel' in logs
+    assert 'long_sleep:sleep_job cancelled, will be run again' in logs[-1]
+    assert worker.jobs_complete == 1
+    assert worker.jobs_retried == 1
+    assert worker.jobs_failed == 0
+
+
 async def test_job_successful(arq_redis: ArqRedis, worker, caplog):
     caplog.set_level(logging.INFO)
     await arq_redis.enqueue_job('foobar', _job_id='testing')