Graceful shutdown without cancelling via SIGUSR1 (#374)

CHANGELOG.md

@@ -2,7 +2,11 @@
 All notable changes to this project will be documented in this file.
 This project adheres to [Semantic Versioning](http://semver.org/).
 
-## [XX.Y.Z] - TODO
+## [23.6.2] - 2019-07-26
+
+### Added
+- Support for graceful shutdown without cancelling using SIGUSR1
+
 ### Removed
 - Support for old `application-services-r` workerType
 

scriptworker/test/test_worker.py

@@ -92,6 +92,38 @@ async def async_main(internal_context, _):
         assert run_tasks_cancelled.result()
 
 
+@pytest.mark.parametrize('running', (True, False))
+def test_main_running_sigusr1(mocker, context, event_loop, running):
+    """Test that sending SIGUSR1 causes the main loop to stop after the next
+    call to async_main without cancelling the task."""
+    run_tasks_cancelled = event_loop.create_future()
+
+    class MockRunTasks:
+        @staticmethod
+        def cancel():
+            run_tasks_cancelled.set_result(True)
+
+    async def async_main(internal_context, _):
+        # scriptworker reads context from a file, so we have to modify the
+        # context given here instead of the variable from the fixture
+        if running:
+            internal_context.running_tasks = MockRunTasks()
+        # Send SIGUSR1 to ourselves so that we stop
+        os.kill(os.getpid(), signal.SIGUSR1)
+
+    _, tmp = tempfile.mkstemp()
+    try:
+        with open(tmp, "w") as fh:
+            json.dump(context.config, fh)
+        mocker.patch.object(worker, 'async_main', new=async_main)
+        mocker.patch.object(sys, 'argv', new=['x', tmp])
+        worker.main(event_loop=event_loop)
+    finally:
+        os.remove(tmp)
+
+    assert not run_tasks_cancelled.done()
+
+
 # async_main {{{1
 @pytest.mark.asyncio
 async def test_async_main(context, mocker, tmpdir):

scriptworker/worker.py

@@ -242,7 +242,14 @@ async def _handle_sigterm():
         if context.running_tasks is not None:
             await context.running_tasks.cancel()
 
+    async def _handle_sigusr1():
+        """Stop accepting new tasks."""
+        log.info("SIGUSR1 received; no more tasks will be taken")
+        nonlocal done
+        done = True
+
     context.event_loop.add_signal_handler(signal.SIGTERM, lambda: asyncio.ensure_future(_handle_sigterm()))
+    context.event_loop.add_signal_handler(signal.SIGUSR1, lambda: asyncio.ensure_future(_handle_sigusr1()))
 
     while not done:
         try:

version.json

@@ -2,7 +2,7 @@
     "version": [
         23,
         6,
-        1
+        2
     ],
-    "version_string": "23.6.1"
+    "version_string": "23.6.2"
 }