Add max_jobs test and fix #103

Author: sylvinus

README.md

@@ -4,8 +4,6 @@
 
 Full documentation is available on [readthedocs](http://mrq.readthedocs.org/en/latest/)
 
-/!\ MRQ is not yet ready for public use. Soon!
-
 # Why?
 
 MRQ is an opinionated task queue. It aims to be simple and beautiful like [RQ](http://python-rq.org) while having performances close to [Celery](http://celeryproject.org)

docs/command-line.md

@@ -40,12 +40,12 @@ The following general flags can be passed as command-line arguments to either **
 
 `mrq-worker` starts a new worker and takes one argument list:
 
- - `queues`: The queues to listen on.Defaults to **default** , which will listen on all queues. 
+ - `queues`: The queues to listen on.Defaults to **default** , which will listen on all queues.
 
 You can pass additional configuration flags:
 
- - `--max_jobs`: Gevent:max number of jobs to do before quitting. Use as a temporary workaround for memory leaks in your tasks. Defaults to **0**
- - `--max_memory`: Max memory (in Mb) after which the process will be shut down. Use with `--processes [1-N]` 
+ - `--max_jobs`: Gevent:max number of jobs to do before quitting. Use as a workaround for memory leaks in your tasks. Defaults to **0**
+ - `--max_memory`: Max memory (in Mb) after which the process will be shut down. Use with `--processes [1-N]`
                   to have supervisord automatically respawn the worker when this happens. Defaults to **1**
  - `--grenlets, --gevent, --g`: Max number of greenlets to use. Defaults to **1**.
  - `--processes, --p`: Number of processes to launch with supervisord. Defaults to **0** (no supervisord).
@@ -87,11 +87,11 @@ To do this we added a ```--add_network_latency=0.1``` config option that will ad
 ## mrq-run
 
 `mrq-run` runs a one-off task. If you add the `--queue` option that will enqueue it to be later ran by a worker.
- 
+
  - `taskpath`: Task to run.
  - `taskargs`: JSON-encoded arguments, or "key value" pairs.
  - `--queue`: Queue the task on this queue instead of running it right away.
- 
+
 Typical usage is:
 
 ```

docs/configuration.md

@@ -50,13 +50,13 @@ DEFAULT_JOB_CANCEL_TTL = 86400 #Seconds the tasks are kept in MongoDB when statu
 DEFAULT_JOB_TIMEOUT = 3600 #In seconds, delay before interrupting the job.Defaults to 3600 seconds which is 1 hour.
 DEFAULT_JOB_MAX_RETRIES = 3 #Set the status to "maxretries" after retrying that many times.Defaults to 3 seconds.
 DEFAULT_JOB_RETRY_DELAY = 3 #Seconds before a job in retry status is requeued again.Defaults to 3 seconds.
-USE_LARGE_JOB_IDS = False #Do not use compacted job IDs in Redis. For compatibility with 0.1.x only. Defaults to 
+USE_LARGE_JOB_IDS = False #Do not use compacted job IDs in Redis. For compatibility with 0.1.x only. Defaults to
 
 """ mrq-worker settings
 """
-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. Temp workaround for memory leaks.Defaults to 0
-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 
+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_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.
 SUPERVISORD_TEMPLATE = "supervisord_templates/default.conf" #Path of supervisord template to use. Defaults to supervisord_templates/default.conf.

mrq/worker.py

@@ -411,6 +411,12 @@ def work_loop(self, max_jobs=None):
 
                     free_pool_slots = self.gevent_pool.free_count()
 
+                    if max_jobs:
+                        total_started = (self.pool_size - free_pool_slots) + self.done_jobs
+                        free_pool_slots = min(free_pool_slots, max_jobs - total_started)
+                        if free_pool_slots == 0:
+                            break
+
                     if free_pool_slots > 0:
                         self.status = "wait"
                         break
@@ -423,6 +429,9 @@ def work_loop(self, max_jobs=None):
 
                     max_jobs_per_queue = free_pool_slots - len(jobs)
 
+                    if max_jobs_per_queue <= 0:
+                        break
+
                     if self.config["dequeue_strategy"] == "parallel":
                         max_jobs_per_queue = 1
 
@@ -432,9 +441,6 @@ def work_loop(self, max_jobs=None):
                         worker=self
                     )
 
-                    if len(jobs) >= free_pool_slots:
-                        break
-
                 for job in jobs:
 
                     # TODO investigate spawn_raw?

tests/tasks/redis.py

@@ -1,5 +1,6 @@
 from mrq.task import Task
 from mrq.context import connections, subpool_map
+import gevent
 
 
 class MultiRedis(Task):
@@ -19,7 +20,10 @@ class Disconnections(Task):
 
     def run(self, params):
 
+        get_clients = lambda: [c for c in connections.redis.client_list() if c.get("cmd") != "client"]
+
         def inner(i):
+            print "Greenlet #%s, %s clients so far" % (id(gevent.getcurrent()), len(get_clients()))
             return connections.redis.get("test")
 
         if params["subpool_size"]:

tests/test_general.py

@@ -143,3 +143,4 @@ def test_general_exception_status(worker):
     assert job1["status"] == "failed"
     assert "raise" in job1["traceback"]
     assert "xyz" in job1["traceback"]
+

tests/test_interrupts.py

@@ -304,3 +304,18 @@ def test_interrupt_redis_started_jobs(worker):
     assert Queue("xxx").size() == 2
     assert Queue("default").size() == 0
     assert connections.redis.zcard(Queue.redis_key_started()) == 0
+
+
+def test_interrupt_maxjobs(worker):
+
+    # The worker will stop after doing 5 jobs
+    worker.start(flags="--max_jobs 5 --greenlets 2", queues="test1 default")
+
+    worker.send_tasks("tests.tasks.general.Add", [
+        {"a": i, "b": 1, "sleep": 0}
+        for i in range(12)
+    ], block=False)
+
+    time.sleep(2)
+
+    assert Queue("default").size() == 7