Merge pull request #135 from FlorianPerucki/master

add pause/resume queue methods

Author: ggueret

docs/command-line.md

@@ -56,6 +56,7 @@ You can pass additional configuration flags:
  - `--report_file`: Filepath of a json dump of the worker status. Disabled if none.
  - `--subqueues_refresh_interval`: Seconds between worker refreshes of the known subqueues.
  - `--subqueues_delimiter`: Delimiter between main queue and subqueue names.
+ - `--paused_queues_refresh_interval`: Seconds between worker refreshes of the paused queues list.
  - `--admin_port`: Start an admin server on this port, if provided. Incompatible with --processes. Defaults to **0**
  - `--admin_ip`: IP for the admin server to listen on. Use "0.0.0.0" to allow access from outside. Defaults to **127.0.0.1**.
  - `--local_ip`: Overwrite the local IP, to be displayed in the dashboard.
@@ -103,4 +104,3 @@ $ mrq-run tasks.mylib.myfile.MyTask '{"param1": 1, "param2": True}'
 # Shorter syntax which casts all values as strings (equivalent to '{"param1": "1", "param2": "ok"}')
 $ mrq-run tasks.mylib.myfile.MyTask param1 1 param2 ok
 ```
-

mrq/config.py

@@ -336,11 +336,18 @@ def add_parser_args(parser, config_type):
 
         parser.add_argument(
             '--subqueues_refresh_interval',
-            default=60,
+            default=10,
             action='store',
             type=float,
             help="Seconds between worker refreshes of the known subqueues")
 
+        parser.add_argument(
+            '--paused_queues_refresh_interval',
+            default=10,
+            action='store',
+            type=float,
+            help="Seconds between worker refreshes of the paused queues list")
+
         parser.add_argument(
             '--subqueues_delimiter',
             default='/',

mrq/queue.py

@@ -25,6 +25,7 @@ class Queue(object):
     # This is a mutable type so it is shared by all instances
     # of Queue in the current process
     known_queues = {}
+    paused_queues = set()
 
     def __init__(self, queue_id, add_to_known_queues=False):
 
@@ -76,6 +77,11 @@ def redis_key_started(cls):
         """ Returns the global redis key used to store started job ids """
         return "%s:s:started" % context.get_current_config()["redis_prefix"]
 
+    @classmethod
+    def redis_key_paused_queues(cls):
+        """ Returns the redis key used to store this queue. """
+        return "%s:s:paused" % (context.get_current_config()["redis_prefix"])
+
     @classmethod
     def redis_key_known_queues(cls):
         """ Returns the global redis key used to store started job ids """
@@ -112,6 +118,11 @@ def redis_known_queues(cls):
             for value, score in context.connections.redis.zrange(cls.redis_key_known_queues(), 0, -1, withscores=True)
         }
 
+    @classmethod
+    def redis_paused_queues(cls):
+        """ Returns the set of currently paused queues """
+        return context.connections.redis.smembers(cls.redis_key_paused_queues())
+
     def redis_known_subqueues(self):
         """ Return the known subqueues of this queue as Queue objects. """
         delimiter = context.get_current_config()["subqueues_delimiter"]
@@ -147,6 +158,37 @@ def unserialize_job_ids(self, job_ids):
         else:
             return [x.encode('hex') for x in job_ids]
 
+    def _get_pausable_id(self):
+      """
+          Get the queue id (either id or root_id) that should be used to pause/unpause the current queue
+          TODO: handle subqueues with more than one level, e.g. "queue/subqueue/"
+      """
+      queue = self.id
+      delimiter = context.get_current_config().get("subqueues_delimiter")
+      if delimiter is not None and self.id.endswith(delimiter):
+          queue = self.root_id
+      return queue
+
+    def pause(self):
+        """ Adds this queue to the set of paused queues """
+        context.connections.redis.sadd(Queue.redis_key_paused_queues(), self._get_pausable_id())
+
+    def is_paused(self):
+        """
+            Returns wether the queue is paused or not.
+            Warning: this does NOT ensure that the queue was effectively added to
+            the set of paused queues. See the 'paused_queues_refresh_interval' option.
+        """
+        root_is_paused = False
+        if self.root_id != self.id:
+            root_is_paused = context.connections.redis.sismember(Queue.redis_key_paused_queues(), self.root_id)
+
+        return root_is_paused or context.connections.redis.sismember(Queue.redis_key_paused_queues(), self.id)
+
+    def resume(self):
+        """ Resumes a paused queue """
+        context.connections.redis.srem(Queue.redis_key_paused_queues(), self._get_pausable_id())
+
     def size(self):
         """ Returns the total number of jobs on the queue """
 
@@ -472,6 +514,7 @@ def dequeue_jobs(self, max_jobs=1, job_class=None, worker=None):
             for j in job_data:
                 j["status"] = "started"
                 j["queue"] = retry_queue
+                j["raw_queue"] = self.id
                 if worker:
                     j["worker"] = worker.id
 

mrq/worker.py

@@ -227,6 +227,14 @@ def greenlet_subqueues(self):
 
             time.sleep(self.config["subqueues_refresh_interval"])
 
+    def greenlet_paused_queues(self):
+
+      while True:
+
+          # Update the process-local list of paused queues
+          Queue.paused_queues = Queue.redis_paused_queues()
+          time.sleep(self.config["paused_queues_refresh_interval"])
+
     def get_memory(self):
         mmaps = self.process.get_memory_maps()
         mem = {
@@ -416,7 +424,13 @@ def work_init(self):
 
         self.status = "started"
 
-        self.greenlets["subqueues"] = gevent.spawn(self.greenlet_subqueues)
+        # An interval of 0 disables the refresh
+        if self.config["subqueues_refresh_interval"] > 0:
+            self.greenlets["subqueues"] = gevent.spawn(self.greenlet_subqueues)
+
+        # An interval of 0 disables the refresh
+        if self.config["paused_queues_refresh_interval"] > 0:
+            self.greenlets["paused_queues"] = gevent.spawn(self.greenlet_paused_queues)
 
         self.greenlets["report"] = gevent.spawn(self.greenlet_report)
 
@@ -470,9 +484,15 @@ def work_loop(self, max_jobs=None):
 
                 jobs = []
 
-                for queue_i in xrange(len(self.queues)):
+                available_queues = [
+                    queue for queue in self.queues
+                    if queue.root_id not in Queue.paused_queues and
+                    queue.id not in Queue.paused_queues
+                ]
+
+                for queue_i in xrange(len(available_queues)):
 
-                    queue = self.queues[(queue_i + queue_offset) % len(self.queues)]
+                    queue = available_queues[(queue_i + queue_offset) % len(available_queues)]
 
                     max_jobs_per_queue = free_pool_slots - len(jobs)
 
@@ -481,7 +501,7 @@ def work_loop(self, max_jobs=None):
                         break
 
                     if self.config["dequeue_strategy"] == "parallel":
-                        max_jobs_per_queue = max(1, int(max_jobs_per_queue / (len(self.queues) - queue_i)))
+                        max_jobs_per_queue = max(1, int(max_jobs_per_queue / (len(available_queues) - queue_i)))
 
                     jobs += queue.dequeue_jobs(
                         max_jobs=max_jobs_per_queue,

tests/test_pause.py

@@ -0,0 +1,130 @@
+from mrq.job import Job
+import pytest
+from mrq.queue import Queue, send_task
+import time
+from mrq.context import set_current_config, get_config
+
+
+def test_pause_resume(worker):
+
+    worker.start(flags="--paused_queues_refresh_interval=0.1")
+
+    Queue("high").pause()
+
+    assert Queue("high").is_paused()
+
+    # wait for the paused_queues list to be refreshed
+    time.sleep(2)
+
+    job_id1 = send_task(
+        "tests.tasks.general.MongoInsert", {"a": 41},
+        queue="high")
+
+    job_id2 = send_task(
+        "tests.tasks.general.MongoInsert", {"a": 43},
+        queue="low")
+
+    time.sleep(5)
+
+    job1 = Job(job_id1).fetch().data
+    job2 = Job(job_id2).fetch().data
+
+    assert job1["status"] == "queued"
+
+    assert job2["status"] == "success"
+    assert job2["result"] == {"a": 43}
+
+    assert worker.mongodb_jobs.tests_inserts.count() == 1
+
+    Queue("high").resume()
+
+    Job(job_id1).wait(poll_interval=0.01)
+
+    job1 = Job(job_id1).fetch().data
+
+    assert job1["status"] == "success"
+    assert job1["result"] == {"a": 41}
+
+    assert worker.mongodb_jobs.tests_inserts.count() == 2
+
+    worker.stop()
+
+
+def test_pause_refresh_interval(worker):
+
+    """ Tests that a refresh interval of 0 disables the pause functionnality """
+
+    worker.start(flags="--paused_queues_refresh_interval=0")
+
+    Queue("high").pause()
+
+    assert Queue("high").is_paused()
+
+    # wait for the paused_queues list to be refreshed
+    time.sleep(2)
+
+    job_id1 = send_task(
+        "tests.tasks.general.MongoInsert", {"a": 41},
+        queue="high")
+
+    time.sleep(5)
+
+    job1 = Job(job_id1).fetch().data
+
+    assert job1["status"] == "success"
+    assert job1["result"] == {"a": 41}
+
+    worker.stop()
+
+
+def test_pause_subqueue(worker):
+
+    # set config in current context in order to have a subqueue delimiter
+    set_current_config(get_config(config_type="worker"))
+
+    worker.start(queues="high high/", flags="--subqueues_refresh_interval=1 --paused_queues_refresh_interval=1")
+
+    Queue("high").pause()
+
+    assert Queue("high/").is_paused()
+
+    # wait for the paused_queues list to be refreshed
+    time.sleep(2)
+
+    job_id1 = send_task(
+        "tests.tasks.general.MongoInsert", {"a": 41},
+        queue="high")
+
+    job_id2 = send_task(
+        "tests.tasks.general.MongoInsert", {"a": 43},
+        queue="high/subqueue")
+
+    # wait a bit to make sure the jobs status will still be queued
+    time.sleep(5)
+
+    job1 = Job(job_id1).fetch().data
+    job2 = Job(job_id2).fetch().data
+
+    assert job1["status"] == "queued"
+    assert job2["status"] == "queued"
+
+    assert worker.mongodb_jobs.tests_inserts.count() == 0
+
+    Queue("high/").resume()
+
+    Job(job_id1).wait(poll_interval=0.01)
+
+    Job(job_id2).wait(poll_interval=0.01)
+
+    job1 = Job(job_id1).fetch().data
+    job2 = Job(job_id2).fetch().data
+
+    assert job1["status"] == "success"
+    assert job1["result"] == {"a": 41}
+
+    assert job2["status"] == "success"
+    assert job2["result"] == {"a": 43}
+
+    assert worker.mongodb_jobs.tests_inserts.count() == 2
+
+    worker.stop()

tests/test_subqueues.py

@@ -53,3 +53,24 @@ def test_custom_delimiters(worker, delimiter):
     job_id = worker.send_task("tests.tasks.general.GetTime", {}, queue=subqueue, block=False)
     Job(job_id).wait(poll_interval=0.01)
     worker.stop()
+
+
+def test_refresh_interval(worker):
+
+    """ Tests that a refresh interval of 0 disables the subqueue detection """
+
+    worker.start(queues="test/", flags="--subqueues_refresh_interval=0")
+
+    time.sleep(2)
+
+    job_id1 = worker.send_task(
+        "tests.tasks.general.GetTime", {"a": 41},
+        queue="test/subqueue", block=False)
+
+    time.sleep(5)
+
+    job1 = Job(job_id1).fetch().data
+
+    assert job1["status"] == "queued"
+
+    worker.stop()