Locked jobs

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))

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

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)
 

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/

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])

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

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"