From a4f1de358f2146edf60c94b71250d619380fdbf4 Mon Sep 17 00:00:00 2001 From: Selwin Ong Date: Thu, 26 Jul 2012 09:36:50 +0700 Subject: [PATCH] Raise a StopException when Control+C is pressed. --- rq/worker.py | 18 +++++++++++++++--- tests/fixtures.py | 4 ++++ 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/rq/worker.py b/rq/worker.py index 9fb34019..cb0c466b 100644 --- a/rq/worker.py +++ b/rq/worker.py @@ -28,6 +28,10 @@ red = make_colorizer('darkred') blue = make_colorizer('darkblue') +class StopRequested(Exception): + pass + + def iterable(x): return hasattr(x, '__iter__') @@ -246,10 +250,16 @@ class Worker(object): self.log.debug('Ignoring signal %s.' % signal_name(signum)) return - msg = 'Warm shut down. Press Ctrl+C again for a cold shutdown.' + msg = 'Warm shut down requested.' self.log.warning(msg) - self._stopped = True - self.log.debug('Stopping after current horse is finished.') + # If shutdown is requested in the middle of a job, wait until finish + # before shutting down + if self.state == 'busy': + self._stopped = 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) @@ -286,6 +296,8 @@ class Worker(object): connection=self.connection) if result is None: break + except StopRequested: + break except UnpickleError as e: msg = '*** Ignoring unpickleable data on %s.' % \ green(e.queue.name) diff --git a/tests/fixtures.py b/tests/fixtures.py index 3d45ab60..da3c13e7 100644 --- a/tests/fixtures.py +++ b/tests/fixtures.py @@ -56,3 +56,7 @@ with Connection(): @job(queue='default') def decorated_job(x, y): return x + y + + +def long_running_job(): + time.sleep(10)