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,47 +323,47 @@ class Worker(object):
         gracefully.
         """
 
-        def request_force_stop(signum, frame):
-            """Terminates the application (cold shutdown).
-            """
-            self.log.warning('Cold shut down')
-
-            # Take down the horse with the worker
-            if self.horse_pid:
-                msg = 'Taking down horse {0} with me'.format(self.horse_pid)
-                self.log.debug(msg)
-                try:
-                    os.kill(self.horse_pid, signal.SIGKILL)
-                except OSError as e:
-                    # ESRCH ("No such process") is fine with us
-                    if e.errno != errno.ESRCH:
-                        self.log.debug('Horse already down')
-                        raise
-            raise SystemExit()
-
-        def request_stop(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.SIGTERM, request_force_stop)
-
-            msg = 'Warm shut down requested'
-            self.log.warning(msg)
-
-            # If shutdown is requested in the middle of a job, wait until
-            # finish before shutting down
-            if self.get_state() == 'busy':
-                self._stop_requested = True
-                self.log.debug('Stopping after current horse is finished. '
-                               'Press Ctrl+C again for a cold shutdown.')
-            else:
-                raise StopRequested()
-
-        signal.signal(signal.SIGINT, request_stop)
-        signal.signal(signal.SIGTERM, request_stop)
+        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')
+
+        # Take down the horse with the worker
+        if self.horse_pid:
+            msg = 'Taking down horse {0} with me'.format(self.horse_pid)
+            self.log.debug(msg)
+            try:
+                os.kill(self.horse_pid, signal.SIGKILL)
+            except OSError as e:
+                # ESRCH ("No such process") is fine with us
+                if e.errno != errno.ESRCH:
+                    self.log.debug('Horse already down')
+                    raise
+        raise SystemExit()
+
+    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, self.request_force_stop)
+        signal.signal(signal.SIGTERM, self.request_force_stop)
+
+        msg = 'Warm shut down requested'
+        self.log.warning(msg)
+
+        # If shutdown is requested in the middle of a job, wait until
+        # finish before shutting down
+        if self.get_state() == 'busy':
+            self._stop_requested = True
+            self.log.debug('Stopping after current horse is finished. '
+                           'Press Ctrl+C again for a cold shutdown.')
+        else:
+            raise StopRequested()
 
     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))
+