diff --git a/docs/command-line.md b/docs/command-line.md index f5669cec..786ad2e0 100644 --- a/docs/command-line.md +++ b/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 ``` - diff --git a/mrq/config.py b/mrq/config.py index ad772fab..044ce7cf 100644 --- a/mrq/config.py +++ b/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='/', diff --git a/mrq/queue.py b/mrq/queue.py index 61a04197..9803fad5 100644 --- a/mrq/queue.py +++ b/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 diff --git a/mrq/worker.py b/mrq/worker.py index f959a2bd..51bcaeae 100644 --- a/mrq/worker.py +++ b/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, diff --git a/tests/test_pause.py b/tests/test_pause.py new file mode 100644 index 00000000..9ee0dd41 --- /dev/null +++ b/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() diff --git a/tests/test_subqueues.py b/tests/test_subqueues.py index 34b65f33..4c58561f 100644 --- a/tests/test_subqueues.py +++ b/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()