fixing signal handling

Author: samuelcolvin

HISTORY.rst

@@ -3,6 +3,11 @@
 History
 -------
 
+v0.3.1 (2017-01-20)
+...................
+* fix main process signal handling so the worker shuts down when just the main process receives a signal
+* re-enqueue un-started jobs popped from the queue if the worker is about to exit
+
 v0.3.0 (2017-01-19)
 ...................
 * rename settings class to ``RedisSettings`` and simplify significantly

arq/version.py

@@ -2,4 +2,4 @@
 
 __all__ = ['VERSION']
 
-VERSION = StrictVersion('0.3')
+VERSION = StrictVersion('0.3.1')

arq/worker.py

@@ -35,6 +35,10 @@ class BadJob(Exception):
     pass
 
 
+# special signal sent by the main process in case the worker process hasn't received a signal (eg. SIGTERM or SIGINT)
+SIG_SUPERVISOR = signal.SIGRTMIN + 7
+
+
 class BaseWorker(RedisMixin):
     """
     Base class for Workers to inherit from.
@@ -80,6 +84,7 @@ def __init__(self, *,
         self.job_class = None  # type: type # TODO
         signal.signal(signal.SIGINT, self.handle_sig)
         signal.signal(signal.SIGTERM, self.handle_sig)
+        signal.signal(SIG_SUPERVISOR, self.handle_supervisor_signal)
         super().__init__(**kwargs)  # type: ignore # TODO
         self._closing_lock = asyncio.Lock(loop=self.loop)
 
@@ -176,15 +181,15 @@ async def work(self):
                 msg = await redis.blpop(*redis_queues, timeout=1)
                 if msg is None:
                     continue
-                _queue, data = msg
-                if self._burst_mode and _queue == quit_queue:
+                raw_queue, data = msg
+                if self._burst_mode and raw_queue == quit_queue:
                     work_logger.debug('got job from the quit queue, stopping')
                     break
-                queue = queue_lookup[_queue]
-                work_logger.debug('scheduling job from queue %s', queue)
-                await self.schedule_job(queue, data)
+                queue = queue_lookup[raw_queue]
+                await self.schedule_job(data, queue, raw_queue)
 
-    async def schedule_job(self, queue, data):
+    async def schedule_job(self, data, queue, raw_queue):
+        work_logger.debug('scheduling job from queue %s', queue)
         job = self.job_class(queue, data)
 
         pt_cnt = len(self._pending_tasks)
@@ -193,6 +198,11 @@ async def schedule_job(self, queue, data):
             _, self._pending_tasks = await asyncio.wait(self._pending_tasks, loop=self.loop,
                                                         return_when=asyncio.FIRST_COMPLETED)
 
+        if not self.running:
+            work_logger.warning('job popped from queue, but exit is imminent, re-queueing the job')
+            async with await self.get_redis_conn() as redis:
+                await redis.lpush(raw_queue, data)
+            return
         task = self.loop.create_task(self.run_job(job))
         task.add_done_callback(self.job_callback)
         self.loop.call_later(self.timeout_seconds, self.cancel_job, task, job)
@@ -280,9 +290,19 @@ async def close(self):
             await super().close()
             self._closed = True
 
+    def handle_supervisor_signal(self, signum, frame):
+        self.running = False
+        work_logger.warning('pid=%d, got shutdown signal from main process, stopping...', os.getpid())
+        signal.signal(signal.SIGINT, self.handle_sig_force)
+        signal.signal(signal.SIGTERM, self.handle_sig_force)
+        signal.signal(signal.SIGALRM, self.handle_sig_force)
+        signal.alarm(self.shutdown_delay)
+        raise HandledExit()
+
     def handle_sig(self, signum, frame):
         self.running = False
         work_logger.warning('pid=%d, got signal: %s, stopping...', os.getpid(), Signals(signum).name)
+        signal.signal(SIG_SUPERVISOR, signal.SIG_IGN)
         signal.signal(signal.SIGINT, self.handle_sig_force)
         signal.signal(signal.SIGTERM, self.handle_sig_force)
         signal.signal(signal.SIGALRM, self.handle_sig_force)
@@ -365,17 +385,18 @@ def run_worker(self, worker_path, worker_class, burst):
         work_logger.critical('worker process %s exited badly with exit code %s',
                              self.process.pid, self.process.exitcode)
         sys.exit(3)
-        # could restart worker here, but better to leave it up to the real manager
+        # could restart worker here, but better to leave it up to the real manager eg. docker restart: always
 
     def handle_sig(self, signum, frame):
         signal.signal(signal.SIGINT, self.handle_sig_force)
         signal.signal(signal.SIGTERM, self.handle_sig_force)
         work_logger.warning('got signal: %s, waiting for worker pid=%s to finish...', Signals(signum).name,
                             self.process and self.process.pid)
-        for i in range(100):  # pragma: no branch
-            if not self.process or not self.process.is_alive():
-                return
-            time.sleep(0.1)
+        # sleep to make sure handle_sig above has executed if it's going to and detached handle_supervisor_signal
+        time.sleep(0.01)
+        if self.process and self.process.is_alive():
+            work_logger.debug("sending custom shutdown signal to worker in case it didn't receive the signal")
+            os.kill(self.process.pid, SIG_SUPERVISOR)
 
     def handle_sig_force(self, signum, frame):
         work_logger.error('got signal: %s again, forcing exit', Signals(signum).name)

tests/example.py

@@ -20,8 +20,8 @@ class WorkerSignalQuit(Worker):
     """
     max_concurrent_tasks = 1
 
-    async def schedule_job(self, queue, data):
-        await super().schedule_job(queue, data)
+    async def schedule_job(self, *args):
+        await super().schedule_job(*args)
         if self.jobs_complete >= 2:
             self.handle_sig(2, None)
 
@@ -30,7 +30,7 @@ class WorkerSignalTwiceQuit(Worker):
     """
     worker which simulates receiving sigint twice after 2 jobs
     """
-    async def schedule_job(self, queue, data):
-        await super().schedule_job(queue, data)
+    async def schedule_job(self, *args):
+        await super().schedule_job(*args)
         if self.jobs_complete >= 2:
             self.handle_sig_force(2, None)

tests/test_cli.py

@@ -45,15 +45,34 @@ def test_worker_exited_badly(tmpworkdir, monkeypatch):
 def test_main_process_sigint(tmpworkdir, monkeypatch, caplog):
     monkeypatch.setattr(arq.worker.Process, 'start', MagicMock())
     monkeypatch.setattr(arq.worker.Process, 'join', MagicMock())
-    is_alive = MagicMock(side_effect=[True, True, False, False])
-    monkeypatch.setattr(arq.worker.Process, 'is_alive', is_alive)
+    monkeypatch.setattr(arq.worker.Process, 'is_alive', MagicMock(return_value=True))
     monkeypatch.setattr(arq.worker.Process, 'exitcode', 0)
     monkeypatch.setattr(arq.worker.Process, 'pid', 123)
+
+    os_kill = MagicMock()
+    monkeypatch.setattr(arq.worker.os, 'kill', os_kill)
+
     tmpworkdir.join('test.py').write(EXAMPLE_FILE)
     work_runner = arq.worker.RunWorkerProcess('test.py', 'Worker')
     work_runner.handle_sig(signal.SIGINT, None)
-    assert is_alive.call_count == 3
     assert 'got signal: SIGINT, waiting for worker pid=123 to finish...' in caplog
+    os_kill.assert_called_once_with(123, arq.worker.SIG_SUPERVISOR)
+
+
+def test_main_process_sigint_worker_stopped(tmpworkdir, monkeypatch, caplog):
+    monkeypatch.setattr(arq.worker.Process, 'start', MagicMock())
+    monkeypatch.setattr(arq.worker.Process, 'join', MagicMock())
+    monkeypatch.setattr(arq.worker.Process, 'is_alive', MagicMock(return_value=False))
+    monkeypatch.setattr(arq.worker.Process, 'exitcode', 0)
+    monkeypatch.setattr(arq.worker.Process, 'pid', 123)
+
+    os_kill = MagicMock()
+    monkeypatch.setattr(arq.worker.os, 'kill', os_kill)
+
+    tmpworkdir.join('test.py').write(EXAMPLE_FILE)
+    work_runner = arq.worker.RunWorkerProcess('test.py', 'Worker')
+    work_runner.handle_sig(signal.SIGINT, None)
+    assert os_kill.called is False
 
 
 def test_main_process_sigint_twice(tmpworkdir, monkeypatch, caplog):
@@ -88,6 +107,5 @@ def test_main_process_sigint_twice_worker_running(tmpworkdir, monkeypatch, caplo
     with pytest.raises(arq.worker.ImmediateExit):
         work_runner.handle_sig_force(signal.SIGINT, None)
     assert is_alive.call_count == 1
-    assert os_kill.called
-    assert os_kill.call_args == ((123, signal.SIGTERM),)
+    os_kill.assert_called_once_with(123, signal.SIGTERM)
     assert 'got signal: SIGINT again, forcing exit' in caplog

tests/test_worker.py

@@ -2,9 +2,11 @@
 import logging
 import re
 from multiprocessing import Process
+from unittest.mock import MagicMock
 
 import pytest
 
+import arq.worker
 from arq.testing import RaiseWorker
 from arq.worker import import_string, start_worker
 
@@ -117,12 +119,16 @@ async def test_run_quit(tmpworkdir, redis_conn, actor, caplog):
     await actor.save_slow(2, 0.1)
     await actor.save_slow(3, 0.1)
     await actor.save_slow(4, 0.1)
+    assert 4 == await redis_conn.llen(b'arq:q:dft')
 
     assert not tmpworkdir.join('save_slow').exists()
     worker = WorkerQuit(loop=actor.loop)
     await worker.run()
-    assert tmpworkdir.join('save_slow').read() == '3'
+    # the third job should be remove from the queue and readded
+    assert tmpworkdir.join('save_slow').read() == '2'
     assert '1 pending tasks, waiting for one to finish before creating task for DemoActor.save_slow(2, 0.1)' in caplog
+    assert 'job popped from queue, but exit is imminent, re-queueing the job' in caplog
+    assert 2 == await redis_conn.llen(b'arq:q:dft')
 
 
 async def test_task_exc(redis_conn, actor, caplog):
@@ -171,6 +177,25 @@ def test_run_sigint_twice(tmpworkdir, redis_conn, loop, caplog):
     assert 'Worker exiting after an unhandled error: ImmediateExit' in caplog
 
 
+async def test_run_supervisor_signal(actor, monkeypatch):
+    mock_signal_signal = MagicMock()
+    monkeypatch.setattr(arq.worker.signal, 'signal', mock_signal_signal)
+    mock_signal_alarm = MagicMock()
+    monkeypatch.setattr(arq.worker.signal, 'alarm', mock_signal_alarm)
+
+    worker = Worker(burst=True, loop=actor.loop)
+    assert worker.running is True
+    assert mock_signal_signal.call_count == 3
+    assert mock_signal_alarm.call_count == 0
+
+    with pytest.raises(arq.worker.HandledExit):
+        worker.handle_supervisor_signal(arq.worker.SIG_SUPERVISOR, None)
+
+    assert worker.running is False
+    assert mock_signal_signal.call_count == 6
+    assert mock_signal_alarm.call_count == 1
+
+
 async def test_non_existent_function(redis_conn, actor, caplog):
     await actor.enqueue_job('doesnt_exist')
     worker = Worker(burst=True, loop=actor.loop)