diff --git a/mrq/basetasks/cleaning.py b/mrq/basetasks/cleaning.py index 532a98a1..6130ac70 100644 --- a/mrq/basetasks/cleaning.py +++ b/mrq/basetasks/cleaning.py @@ -10,6 +10,8 @@ class RequeueInterruptedJobs(Task): """ Requeue jobs that were marked as status=interrupt when a worker got a SIGTERM. """ + max_concurrency = 1 + def run(self, params): return run_task("mrq.basetasks.utils.JobAction", { "status": "interrupt", @@ -21,6 +23,8 @@ class RequeueRetryJobs(Task): """ Requeue jobs that were marked as retry. """ + max_concurrency = 1 + def run(self, params): return run_task("mrq.basetasks.utils.JobAction", { "status": "retry", @@ -37,6 +41,8 @@ class RequeueStartedJobs(Task): The timeout parameter of this task is in addition to the task's own timeout. """ + max_concurrency = 1 + def run(self, params): additional_timeout = params.get("timeout", 300) @@ -76,6 +82,8 @@ class RequeueRedisStartedJobs(Task): redis.lpop and mongodb.update """ + max_concurrency = 1 + def run(self, params): redis_key_started = Queue.redis_key_started() @@ -124,6 +132,8 @@ class RequeueLostJobs(Task): They could have been lost by a Redis flush or another severe issue """ + max_concurrency = 1 + def run(self, params): # If there are more than this much items on the queue, we don't try to check if our mongodb @@ -186,6 +196,9 @@ class MigrateKnownQueues(Task): """ Migrate known_queues from old set format to new zset """ + + max_concurrency = 1 + def run(self, params): key = "%s:known_queues" % get_current_config()["redis_prefix"] for queue in connections.redis.smembers(key): @@ -202,6 +215,8 @@ class CleanKnownQueues(Task): - be empty """ + max_concurrency = 1 + def run(self, params): max_age = int(params.get("max_age") or (7 * 86400)) diff --git a/mrq/exceptions.py b/mrq/exceptions.py index d4217778..87c3537c 100644 --- a/mrq/exceptions.py +++ b/mrq/exceptions.py @@ -50,3 +50,7 @@ class StopRequested(GreenletExit): class JobInterrupt(GreenletExit): """ Interrupts that stop a job in its execution, e.g. when responding to a SIGTERM. """ pass + + +class MaxConcurrencyInterrupt(_MrqInterrupt): + pass diff --git a/mrq/job.py b/mrq/job.py index d00492f4..5bcf3b8a 100644 --- a/mrq/job.py +++ b/mrq/job.py @@ -1,7 +1,8 @@ import datetime from bson import ObjectId +from redis.exceptions import LockError import time -from .exceptions import RetryInterrupt, MaxRetriesInterrupt, AbortInterrupt +from .exceptions import RetryInterrupt, MaxRetriesInterrupt, AbortInterrupt, MaxConcurrencyInterrupt from .utils import load_class_by_path, group_iter import gevent import objgraph @@ -70,6 +71,12 @@ def __init__(self, job_id, queue=None, start=False, fetch=False): elif fetch: self.fetch(start=False, full_data=False) + @property + def redis_max_concurrency_key(self): + """ Returns the global redis key used to store started job ids """ + return "%s:c:%s" % (context.get_current_config()["redis_prefix"], + self.data["path"]) + def exists(self): """ Returns True if a job with the current _id exists in MongoDB. """ return bool(self.collection.find_one({"_id": self.id}, projection={"_id": 1})) @@ -274,7 +281,27 @@ def perform(self): self.task.is_main_task = True - result = self.task.run_wrapped(self.data["params"]) + try: + lock = None + + if self.task.max_concurrency: + + if self.task.max_concurrency > 1: + raise NotImplementedError() + + # TODO: implement a semaphore + lock = context.connections.redis.lock(self.redis_max_concurrency_key, timeout=self.timeout + 5) + if not lock.acquire(blocking=True, blocking_timeout=0): + raise MaxConcurrencyInterrupt() + + result = self.task.run_wrapped(self.data["params"]) + + finally: + if lock: + try: + lock.release() + except LockError: + pass self.save_success(result) diff --git a/mrq/task.py b/mrq/task.py index 22dd0202..2aea0711 100644 --- a/mrq/task.py +++ b/mrq/task.py @@ -3,6 +3,7 @@ class Task(object): # Are we the first task that a Job called? is_main_task = False + max_concurrency = 0 # Default write concern values when setting status=success # http://docs.mongodb.org/manual/reference/write-concern/ diff --git a/mrq/worker.py b/mrq/worker.py index 51bcaeae..5f4e01d9 100644 --- a/mrq/worker.py +++ b/mrq/worker.py @@ -16,7 +16,7 @@ from .job import Job from .exceptions import (TimeoutInterrupt, StopRequested, JobInterrupt, AbortInterrupt, - RetryInterrupt, MaxRetriesInterrupt) + RetryInterrupt, MaxRetriesInterrupt, MaxConcurrencyInterrupt) from .context import (set_current_worker, set_current_job, get_current_job, get_current_config, connections, enable_greenlet_tracing) from .queue import Queue @@ -617,6 +617,10 @@ def perform_job(self, job): try: job.perform() + except MaxConcurrencyInterrupt: + self.log.error("Max concurrency reached") + job._save_status("maxconcurrency", exception=True) + except RetryInterrupt: self.log.error("Caught retry") job.save_retry(sys.exc_info()[1]) diff --git a/tests/tasks/concurrency.py b/tests/tasks/concurrency.py new file mode 100644 index 00000000..5e4ca174 --- /dev/null +++ b/tests/tasks/concurrency.py @@ -0,0 +1,18 @@ +import time +from mrq.task import Task +from mrq.context import log +from .general import Add + +class LockedAdd(Add): + + max_concurrency = 1 + + def run(self, params): + log.info("adding", params) + res = params.get("a", 0) + params.get("b", 0) + + if params.get("sleep", 0): + log.info("sleeping", params.get("sleep", 0)) + time.sleep(params.get("sleep", 0)) + + return res diff --git a/tests/test_interrupts.py b/tests/test_interrupts.py index 5404658a..135924ca 100644 --- a/tests/test_interrupts.py +++ b/tests/test_interrupts.py @@ -319,3 +319,25 @@ def test_interrupt_maxjobs(worker): time.sleep(2) assert Queue("default").size() == 7 + + +def test_interrupt_maxconcurrency(worker): + + # The worker will raise a maxconcurrency on the second job + worker.start(flags="--greenlets=2") + + job_ids = worker.send_tasks("tests.tasks.concurrency.LockedAdd", [ + {"a": i, "b": 1, "sleep": 2} + for i in range(2) + ], block=False) + + worker.wait_for_tasks_results(job_ids, accept_statuses=["success", "failed", "maxconcurrency"]) + job_statuses = [Job(job_id).fetch().data["status"] for job_id in job_ids] + assert job_statuses == ["success", "maxconcurrency"] + + # the job concurrency key must be equal to 0 + last_job_id = worker.send_task("tests.tasks.concurrency.LockedAdd", + {"a": 1, "b": 1, "sleep": 2}, block=False + ) + last_job = Job(last_job_id).wait(poll_interval=0.01) + assert last_job.get("status") == "success"