Max time (#152)

* add max_time to worker

* add max_time parameter to config

* script limit on max_time

* add tests for max_time

* add config doc for max_time

* review

* refactor max_time tests

Author: mehdigmira

docs/configuration.md

@@ -56,6 +56,7 @@ USE_LARGE_JOB_IDS = False #Do not use compacted job IDs in Redis. For compatibil
 """
 QUEUES = ("default",) # The queues to listen on.Defaults to default , which will listen on all queues.
 MAX_JOBS = 0 #Gevent:max number of jobs to do before quitting. Workaround for memory leaks in your tasks. Defaults to 0
+MAX_TIME = 0 # max number of seconds a worker runs before quitting
 MAX_MEMORY = 1 #Max memory (in Mb) after which the process will be shut down. Use with PROCESS = [1-N] to have supervisord automatically respawn the worker when this happens.Defaults to 1
 GRENLETS = 1 #Max number of greenlets to use.Defaults to 1.
 PROCESSES = 0 #Number of processes to launch with supervisord.Defaults to 0.

mrq/config.py

@@ -267,6 +267,13 @@ def add_parser_args(parser, config_type):
             help='Gevent: max number of jobs to do before quitting.' +
                  ' Temp workaround for memory leaks')
 
+        parser.add_argument(
+            '--max_time',
+            default=0.0,
+            type=float,
+            action='store',
+            help='Max time a worker should run before quitting.')
+
         parser.add_argument(
             '--max_memory',
             default=0,

mrq/worker.py

@@ -15,7 +15,6 @@
 import sys
 import json as json_stdlib
 import ujson as json
-import http.server
 from bson import ObjectId
 from collections import defaultdict
 
@@ -53,6 +52,7 @@ def __init__(self):
 
         self.done_jobs = 0
         self.max_jobs = self.config["max_jobs"]
+        self.max_time = datetime.timedelta(seconds=self.config["max_time"]) or None
 
         self.connected = False  # MongoDB + Redis
 
@@ -419,7 +419,7 @@ def work(self):
         """
         self.work_init()
 
-        self.work_loop(max_jobs=self.max_jobs)
+        self.work_loop(max_jobs=self.max_jobs, max_time=self.max_time)
 
         return self.work_stop()
 
@@ -449,7 +449,7 @@ def work_init(self):
 
         self.install_signal_handlers()
 
-    def work_loop(self, max_jobs=None):
+    def work_loop(self, max_jobs=None, max_time=None):
 
         self.done_jobs = 0
         self.idle_wait_count = 0
@@ -459,6 +459,7 @@ def work_loop(self, max_jobs=None):
         try:
 
             queue_offset = 0
+            max_time_reached = False
 
             while True:
 
@@ -472,6 +473,12 @@ def work_loop(self, max_jobs=None):
 
                 while True:
 
+                    # we put this here to make sure we have a strict limit on max_time
+                    if max_time and datetime.datetime.utcnow() - self.datestarted >= max_time:
+                        self.log.info("Reached max_time=%s" % max_time.seconds)
+                        max_time_reached = True
+                        break
+
                     free_pool_slots = self.gevent_pool.free_count()
 
                     if max_jobs:
@@ -487,6 +494,9 @@ def work_loop(self, max_jobs=None):
                     self.status = "full"
                     gevent.sleep(0.01)
 
+                if max_time_reached:
+                    break
+
                 jobs = []
 
                 available_queues = [
@@ -561,7 +571,6 @@ def work_loop(self, max_jobs=None):
 
                 self.log.debug("Joining the greenlet pool...")
                 self.status = "join"
-
                 self.gevent_pool.join(timeout=None, raise_error=False)
                 self.log.debug("Joined.")
 

tests/test_interrupts.py

@@ -1,7 +1,7 @@
 import time
 import datetime
 from builtins import str
-from mrq.job import Job
+from mrq.job import Job, get_job_result
 from mrq.queue import Queue
 from bson import ObjectId
 import pytest
@@ -322,6 +322,25 @@ def test_interrupt_maxjobs(worker):
     assert Queue("default").size() == 7
 
 
+def test_worker_interrupt_after_max_time(worker):
+    worker.start(flags="--greenlets=2 --max_time=2", queues="test1 default")
+
+    task_ids = worker.send_tasks("tests.tasks.general.Add", [{"a": i, "b": 1, "sleep": 3} for i in range(5)],
+                                 block=False)
+
+    time.sleep(5)
+
+    results = [get_job_result(task_id) for task_id in task_ids]
+
+    queued_tasks = [result for result in results if result['status'] == "queued"]
+    successful_tasks = [(i, result) for i, result in enumerate(results) if result['status'] == "success"]
+
+    assert len(queued_tasks) == 3
+    assert len(successful_tasks) == 2
+    for i, result in successful_tasks:
+        assert result['result'] == i + 1
+
+
 def test_interrupt_maxconcurrency(worker):
 
     # The worker will raise a maxconcurrency on the second job