Merge pull request #587 from samuelcolvin/signal-changes

Signal changes

.travis.yml

@@ -15,6 +15,6 @@ install:
   - pip install coveralls --use-mirrors
   #- pip install pytest  # installed by Travis by default already
 script:
-  - py.test --cov rq
+  - RUN_SLOW_TESTS_TOO=1 py.test --cov rq
 after_success:
   - coveralls

rq/worker.py

@@ -323,7 +323,10 @@ class Worker(object):
         gracefully.
         """
 
-        def request_force_stop(signum, frame):
+        signal.signal(signal.SIGINT, self.request_stop)
+        signal.signal(signal.SIGTERM, self.request_stop)
+
+    def request_force_stop(self, signum, frame):
         """Terminates the application (cold shutdown).
         """
         self.log.warning('Cold shut down')
@@ -341,14 +344,14 @@ class Worker(object):
                     raise
         raise SystemExit()
 
-        def request_stop(signum, frame):
+    def request_stop(self, signum, frame):
         """Stops the current worker loop but waits for child processes to
         end gracefully (warm shutdown).
         """
         self.log.debug('Got signal {0}'.format(signal_name(signum)))
 
-            signal.signal(signal.SIGINT, request_force_stop)
+        signal.signal(signal.SIGINT, self.request_force_stop)
-            signal.signal(signal.SIGTERM, request_force_stop)
+        signal.signal(signal.SIGTERM, self.request_force_stop)
 
         msg = 'Warm shut down requested'
         self.log.warning(msg)
@@ -362,9 +365,6 @@ class Worker(object):
         else:
             raise StopRequested()
 
-        signal.signal(signal.SIGINT, request_stop)
-        signal.signal(signal.SIGTERM, request_stop)
-
     def check_for_suspension(self, burst):
         """Check to see if workers have been suspended by `rq suspend`"""
 

run_tests

@@ -17,9 +17,9 @@ else
     safe_rg=cat
 fi
 
-export ONLY_RUN_FAST_TESTS=1
+export RUN_SLOW_TESTS_TOO=1
 if [ "$1" = '-f' ]; then  # Poor man's argparse
-    unset ONLY_RUN_FAST_TESTS
+    unset RUN_SLOW_TESTS_TOO
     shift 1
 fi
 

tests/__init__.py

@@ -32,7 +32,7 @@ def slow(f):
 
     @wraps(f)
     def _inner(*args, **kwargs):
-        if os.environ.get('ONLY_RUN_FAST_TESTS'):
+        if os.environ.get('RUN_SLOW_TESTS_TOO'):
             f(*args, **kwargs)
 
     return _inner

tests/test_worker.py

@@ -5,6 +5,9 @@ from __future__ import (absolute_import, division, print_function,
 import os
 from datetime import timedelta
 from time import sleep
+import signal
+import time
+from multiprocessing import Process
 
 from tests import RQTestCase, slow
 from tests.fixtures import (create_file, create_file_after_timeout,
@@ -468,3 +471,74 @@ class TestWorker(RQTestCase):
         worker = Worker(queue, connection=self.testconn)
         worker.work(burst=True)
         self.assertEqual(self.testconn.zcard(registry.key), 0)
+
+
+def kill_worker(pid, double_kill):
+    # wait for the worker to be started over on the main process
+    time.sleep(0.5)
+    os.kill(pid, signal.SIGTERM)
+    if double_kill:
+        # give the worker time to switch signal handler
+        time.sleep(0.5)
+        os.kill(pid, signal.SIGTERM)
+
+
+class TestWorkerShutdown(RQTestCase):
+    def setUp(self):
+        # we want tests to fail if signal are ignored and the work remain running,
+        # so set a signal to kill them after 5 seconds
+        signal.signal(signal.SIGALRM, self._timeout)
+        signal.alarm(5)
+
+    def _timeout(self, signal, frame):
+        raise AssertionError("test still running after 5 seconds, "
+                             "likely the worker wasn't shutdown correctly")
+
+    @slow
+    def test_idle_worker_warm_shutdown(self):
+        """worker with no ongoing job receiving single SIGTERM signal and shutting down"""
+        w = Worker('foo')
+        self.assertFalse(w._stop_requested)
+        p = Process(target=kill_worker, args=(os.getpid(), False))
+        p.start()
+
+        w.work()
+
+        p.join(1)
+        self.assertFalse(w._stop_requested)
+
+    @slow
+    def test_working_worker_warm_shutdown(self):
+        """worker with an ongoing job receiving single SIGTERM signal, allowing job to finish then shutting down"""
+        fooq = Queue('foo')
+        w = Worker(fooq)
+
+        sentinel_file = '/tmp/.rq_sentinel_warm'
+        fooq.enqueue(create_file_after_timeout, sentinel_file, 2)
+        self.assertFalse(w._stop_requested)
+        p = Process(target=kill_worker, args=(os.getpid(), False))
+        p.start()
+
+        w.work()
+
+        p.join(2)
+        self.assertTrue(w._stop_requested)
+        self.assertTrue(os.path.exists(sentinel_file))
+
+    @slow
+    def test_working_worker_cold_shutdown(self):
+        """worker with an ongoing job receiving double SIGTERM signal and shutting down immediately"""
+        fooq = Queue('foo')
+        w = Worker(fooq)
+        sentinel_file = '/tmp/.rq_sentinel_cold'
+        fooq.enqueue(create_file_after_timeout, sentinel_file, 2)
+        self.assertFalse(w._stop_requested)
+        p = Process(target=kill_worker, args=(os.getpid(), True))
+        p.start()
+
+        self.assertRaises(SystemExit, w.work)
+
+        p.join(1)
+        self.assertTrue(w._stop_requested)
+        self.assertFalse(os.path.exists(sentinel_file))
+