Merge pull request binarymatt#105 from heynemann/master

Implementing timeout for workers

Author: binarymatt

pyres/exceptions.py

@@ -1,2 +1,5 @@
 class NoQueueError(Exception):
     pass
+
+class TimeoutError(RuntimeError):
+    pass

pyres/scripts.py

@@ -106,6 +106,7 @@ def pyres_worker():
     parser.add_option('-l', '--log-level', dest='log_level', default='info', help='log level.  Valid values are "debug", "info", "warning", "error", "critical", in decreasing order of verbosity. Defaults to "info" if parameter not specified.')
     parser.add_option('-f', dest='logfile', help='If present, a logfile will be used.  "stderr", "stdout", and "syslog" are all special values.')
     parser.add_option('-p', dest='pidfile', help='If present, a pidfile will be used.')
+    parser.add_option("-t", '--timeout', dest='timeout', default=None, help='the timeout in seconds for this worker')
     (options,args) = parser.parse_args()
 
     if len(args) != 1:
@@ -120,6 +121,8 @@ def pyres_worker():
     if interval is not None:
         interval = int(interval)
 
+    timeout = options.timeout is None and options.timeout or int(options.timeout)
+
     queues = args[0].split(',')
     server = '%s:%s' % (options.host,options.port)
-    Worker.run(queues, server, interval)
+    Worker.run(queues, server, interval, timeout=timeout)

pyres/worker.py

@@ -6,7 +6,7 @@
 import commands
 import random
 
-from pyres.exceptions import NoQueueError
+from pyres.exceptions import NoQueueError, TimeoutError
 from pyres.job import Job
 from pyres import ResQ, Stat, __version__
 
@@ -25,13 +25,14 @@ class and passes a comma-separated list of queues to listen on.::
 
     job_class = Job
 
-    def __init__(self, queues=(), server="localhost:6379", password=None):
+    def __init__(self, queues=(), server="localhost:6379", password=None, timeout=None):
         self.queues = queues
         self.validate_queues()
         self._shutdown = False
         self.child = None
         self.pid = os.getpid()
         self.hostname = os.uname()[1]
+        self.timeout = timeout
 
         if isinstance(server, basestring):
             self.resq = ResQ(server=server, password=password)
@@ -153,13 +154,34 @@ def work(self, interval=5):
                                                       datetime.datetime.now()))
 
                     try:
-                        os.waitpid(self.child, 0)
+                        start = datetime.datetime.now()
+
+                        # waits for the result or times out
+                        while True:
+                            result = os.waitpid(self.child, os.WNOHANG)
+                            if result != (0, 0):
+                                break
+                            time.sleep(0.5)
+
+                            now = datetime.datetime.now()
+                            if self.timeout and ((now - start).seconds > self.timeout):
+                                os.kill(self.child, signal.SIGKILL)
+                                os.waitpid(-1, os.WNOHANG)
+                                raise TimeoutError("Timed out after %d seconds" % self.timeout)
+
                     except OSError as ose:
                         import errno
 
                         if ose.errno != errno.EINTR:
                             raise ose
-                    #os.wait()
+
+                    except TimeoutError as e:
+                        exceptionType, exceptionValue, exceptionTraceback = sys.exc_info()
+                        logger.exception("%s timed out: %s" % (job, e))
+                        job.fail(exceptionTraceback)
+                        self.failed()
+                        self.done_working()
+
                     logger.debug('done waiting')
                 else:
                     self._setproctitle("Processing %s since %s" %
@@ -283,8 +305,8 @@ def worker_pids(self):
                                        grep pyres_worker").split("\n"))
 
     @classmethod
-    def run(cls, queues, server="localhost:6379", interval=None):
-        worker = cls(queues=queues, server=server)
+    def run(cls, queues, server="localhost:6379", interval=None, timeout=None):
+        worker = cls(queues=queues, server=server, timeout=timeout)
         if interval is not None:
             worker.work(interval)
         else:

tests/__init__.py

@@ -37,6 +37,14 @@ def perform(fail_until):
         else:
             return True
 
+class TimeoutJob(object):
+    queue = 'basic'
+
+    @staticmethod
+    def perform(wait_for):
+        import time
+        time.sleep(wait_for)
+        return "Done Sleeping"
 
 class TestProcess(object):
     queue = 'high'

tests/test_worker.py

@@ -1,4 +1,4 @@
-from tests import PyResTests, Basic, TestProcess, ErrorObject, RetryOnExceptionJob
+from tests import PyResTests, Basic, TestProcess, ErrorObject, RetryOnExceptionJob, TimeoutJob
 from pyres import ResQ
 from pyres.job import Job
 from pyres.scheduler import Scheduler
@@ -202,6 +202,23 @@ def test_retry_on_exception(self):
         assert True == worker.process()
         assert worker.get_failed() == 0
 
+    def test_kills_stale_workers_after_timeout(self):
+        import signal
+        timeout = 1
+
+        worker = Worker(['basic'], timeout=timeout)
+        self.resq.enqueue(TimeoutJob, timeout + 1)
+
+        child = os.fork()
+        if child:
+            assert worker.get_failed() == 0
+            time.sleep(timeout + 2)
+            os.kill(child, signal.SIGKILL)
+            os.waitpid(-1, os.WNOHANG)
+            assert worker.get_failed() == 1
+        else:
+            worker.work()
+
     def test_retries_give_up_eventually(self):
         now = datetime.datetime.now()
         self.set_current_time(now)