improved preventing new jobs starting when the worker is stopping

Author: samuelcolvin

HISTORY.rst

@@ -3,6 +3,10 @@
 History
 -------
 
+v0.3.2 (TBD)
+...................
+* improved solution for preventing new jobs starting when the worker is about to stop
+
 v0.3.1 (2017-01-20)
 ...................
 * fix main process signal handling so the worker shuts down when just the main process receives a signal

arq/version.py

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

arq/worker.py

@@ -36,7 +36,7 @@ class BadJob(Exception):
 
 
 # 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
+SIG_PROXY = signal.SIGRTMIN + 7
 
 
 class BaseWorker(RedisMixin):
@@ -84,7 +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)
+        signal.signal(SIG_PROXY, self.handle_proxy_signal)
         super().__init__(**kwargs)  # type: ignore # TODO
         self._closing_lock = asyncio.Lock(loop=self.loop)
 
@@ -160,7 +160,7 @@ async def run(self, reuse=False):
         finally:
             if reuse:
                 work_logger.info('waiting for %d jobs to finish', len(self._pending_tasks))
-                await asyncio.gather(*self._pending_tasks, loop=self.loop)
+                await asyncio.wait(self._pending_tasks, loop=self.loop)
             else:
                 await self.close()
             if self._task_exception:
@@ -186,28 +186,28 @@ async def work(self):
                     work_logger.debug('got job from the quit queue, stopping')
                     break
                 queue = queue_lookup[raw_queue]
-                await self.schedule_job(data, queue, raw_queue)
+                self.schedule_job(data, queue)
+                await self.below_concurrency_limit()
 
-    async def schedule_job(self, data, queue, raw_queue):
+    def schedule_job(self, data, queue):
         work_logger.debug('scheduling job from queue %s', queue)
         job = self.job_class(queue, data)
 
-        pt_cnt = len(self._pending_tasks)
-        if pt_cnt >= self.max_concurrent_tasks:
-            work_logger.debug('%d pending tasks, waiting for one to finish before creating task for %s', pt_cnt, job)
-            _, 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)
         self._pending_tasks.add(task)
 
+    async def below_concurrency_limit(self):
+        pt_cnt = len(self._pending_tasks)
+        while True:
+            if pt_cnt < self.max_concurrent_tasks:
+                return
+            work_logger.debug('%d pending tasks, waiting for one to finish', pt_cnt)
+            _, self._pending_tasks = await asyncio.wait(self._pending_tasks, loop=self.loop,
+                                                        return_when=asyncio.FIRST_COMPLETED)
+            pt_cnt = len(self._pending_tasks)
+
     def cancel_job(self, task, job):
         if not task.cancel():
             return
@@ -290,9 +290,9 @@ async def close(self):
             await super().close()
             self._closed = True
 
-    def handle_supervisor_signal(self, signum, frame):
+    def handle_proxy_signal(self, signum, frame):
         self.running = False
-        work_logger.warning('pid=%d, got shutdown signal from main process, stopping...', os.getpid())
+        work_logger.warning('pid=%d, got signal proxied 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)
@@ -302,7 +302,7 @@ def handle_supervisor_signal(self, signum, frame):
     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(SIG_PROXY, 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)
@@ -392,11 +392,11 @@ def handle_sig(self, signum, frame):
         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)
-        # sleep to make sure handle_sig above has executed if it's going to and detached handle_supervisor_signal
+        # sleep to make sure worker.handle_sig above has executed if it's going to and detached handle_proxy_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)
+            os.kill(self.process.pid, SIG_PROXY)
 
     def handle_sig_force(self, signum, frame):
         work_logger.error('got signal: %s again, forcing exit', Signals(signum).name)

setup.cfg

@@ -1,6 +1,7 @@
 [tool:pytest]
 testpaths = tests
 addopts = --isort
+timeout = 5
 
 [flake8]
 max-complexity = 10

tests/example.py

@@ -20,8 +20,8 @@ class WorkerSignalQuit(Worker):
     """
     max_concurrent_tasks = 1
 
-    async def schedule_job(self, *args):
-        await super().schedule_job(*args)
+    def schedule_job(self, *args):
+        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, *args):
-        await super().schedule_job(*args)
+    def schedule_job(self, *args):
+        super().schedule_job(*args)
         if self.jobs_complete >= 2:
             self.handle_sig_force(2, None)

tests/requirements.txt

@@ -11,6 +11,7 @@ pytest-cov==2.4.0
 pytest-isort==0.1.0
 pytest-mock==1.5.0
 pytest-sugar==0.8.0
+pytest-timeout==1.2.0
 pytest-toolbox==0.1
 pytz==2016.10
 Sphinx==1.5.1

tests/test_cli.py

@@ -56,7 +56,7 @@ def test_main_process_sigint(tmpworkdir, monkeypatch, caplog):
     work_runner = arq.worker.RunWorkerProcess('test.py', 'Worker')
     work_runner.handle_sig(signal.SIGINT, None)
     assert 'got signal: SIGINT, waiting for worker pid=123 to finish...' in caplog
-    os_kill.assert_called_once_with(123, arq.worker.SIG_SUPERVISOR)
+    os_kill.assert_called_once_with(123, arq.worker.SIG_PROXY)
 
 
 def test_main_process_sigint_worker_stopped(tmpworkdir, monkeypatch, caplog):

tests/test_worker.py

@@ -126,8 +126,7 @@ async def test_run_quit(tmpworkdir, redis_conn, actor, caplog):
     await worker.run()
     # 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 '1 pending tasks, waiting for one to finish' in caplog
     assert 2 == await redis_conn.llen(b'arq:q:dft')
 
 
@@ -177,7 +176,7 @@ 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):
+async def test_run_proxy_signal(actor, monkeypatch):
     mock_signal_signal = MagicMock()
     monkeypatch.setattr(arq.worker.signal, 'signal', mock_signal_signal)
     mock_signal_alarm = MagicMock()
@@ -189,7 +188,7 @@ async def test_run_supervisor_signal(actor, monkeypatch):
     assert mock_signal_alarm.call_count == 0
 
     with pytest.raises(arq.worker.HandledExit):
-        worker.handle_supervisor_signal(arq.worker.SIG_SUPERVISOR, None)
+        worker.handle_proxy_signal(arq.worker.SIG_PROXY, None)
 
     assert worker.running is False
     assert mock_signal_signal.call_count == 6