Merge pull request #8 from bigjools/max_concurrency

Add max_concurrency to tasks

Author: NicolasLM

README.rst

@@ -14,6 +14,7 @@ Distinctive features:
 
 - At-least-once or at-most-once delivery per task
 - Periodic tasks without an additional process
+- Concurrency limits on queued jobs
 - Scheduling of tasks in batch
 - Integrations with `Flask, Django, Logging, Sentry and Datadog
   <https://spinach.readthedocs.io/en/stable/user/integrations.html>`_

doc/hacking/contributing.rst

@@ -32,14 +32,24 @@ The code base follows `pep8 <https://www.python.org/dev/peps/pep-0008/>`_
 guidelines with lines wrapping at the 79th character. You can verify that the
 code follows the conventions with::
 
-    $ pep8 spinach tests
+    $ pycodestyle --ignore=E252,W503,W504 spinach tests
 
 Running tests is an invaluable help when adding a new feature or when
 refactoring. Try to add the proper test cases in ``tests/`` together with your
 patch. The test suite can be run with pytest::
 
     $ pytest tests
 
+Because the Redis broker tests require a running Redis server, there is also a
+convenience `tox.ini` that runs all the tests and pep8 checks for you after
+starting Redis in a container via docker-compose.  Simply running::
+
+    $ tox
+
+will build a virtualenv, install Spinach and its dependencies into it,
+start the Redis server in the container, and run tests and pycodestyle,
+tearing down the Redis server container when done.
+
 Compatibility
 -------------
 

doc/index.rst

@@ -9,6 +9,7 @@ Distinctive features:
 
 - At-least-once or at-most-once delivery per task
 - Periodic tasks without an additional process
+- Concurrency limits on queued jobs
 - Scheduling of tasks in batch
 - Embeddable workers for easier testing
 - Integrations with :ref:`Flask, Django, Logging, Sentry and Datadog

doc/user/tasks.rst

@@ -109,6 +109,22 @@ A task can also raise a :class:`AbortException` for short-circuit behavior:
 
 .. autoclass:: spinach.task.AbortException
 
+Limiting task concurrency
+-------------------------
+
+If a task is idempotent it may also have a limit on the number of
+concurrent jobs spawned across all workers. These types of tasks are
+defined with a positive `max_concurrency` value::
+
+    @tasks.task(name='foo', max_retries=10, max_concurrency=1)
+    def foo(a, b):
+        pass
+
+With this definition, no more than one instance of the Task will ever be
+spawned as a running Job, no matter how many are queued and waiting to
+run.
+
+
 Periodic tasks
 --------------
 

setup.py

@@ -67,7 +67,9 @@
             'flush.lua',
             'get_jobs_from_queue.lua',
             'move_future_jobs.lua',
-            'register_periodic_tasks.lua'
+            'register_periodic_tasks.lua',
+            'remove_job_from_running.lua',
+            'set_concurrency_keys.lua',
         ],
     },
 )

spinach/brokers/base.py

@@ -85,6 +85,18 @@ def _to_namespaced(self, value: str) -> str:
     def register_periodic_tasks(self, tasks: Iterable[Task]):
         """Register tasks that need to be scheduled periodically."""
 
+    @abstractmethod
+    def set_concurrency_keys(self, tasks: Iterable[Task]):
+        """Register concurrency data for Tasks.
+
+        Set up anything in the Broker that is required to track
+        concurrency on Tasks, where a Task defines max_concurrency.
+        """
+
+    @abstractmethod
+    def is_queue_empty(self, queue: str) -> bool:
+        """Return True if the provided queue is empty."""
+
     @abstractmethod
     def inspect_periodic_tasks(self) -> List[Tuple[int, str]]:
         """Get the next periodic task schedule.
@@ -93,7 +105,7 @@ def inspect_periodic_tasks(self) -> List[Tuple[int, str]]:
         """
 
     @abstractmethod
-    def enqueue_jobs(self, jobs: Iterable[Job]):
+    def enqueue_jobs(self, jobs: Iterable[Job], from_failure: bool):
         """Enqueue a batch of jobs."""
 
     @abstractmethod

spinach/brokers/memory.py

@@ -23,6 +23,8 @@ def __init__(self):
         self._future_jobs = list()
         self._running_jobs = list()
         self._scheduler = sched.scheduler()
+        self._max_concurrency_keys = dict()
+        self._cur_concurrency_keys = dict()
 
     def _get_queue(self, queue_name: str):
         queue_name = self._to_namespaced(queue_name)
@@ -34,9 +36,16 @@ def _get_queue(self, queue_name: str):
                 self._queues[queue_name] = queue
                 return queue
 
-    def enqueue_jobs(self, jobs: Iterable[Job]):
+    def enqueue_jobs(self, jobs: Iterable[Job], from_failure: bool=False):
         """Enqueue a batch of jobs."""
         for job in jobs:
+            with self._lock:
+                if from_failure:
+                    max_concurrency = self._max_concurrency_keys[
+                        job.task_name
+                    ]
+                    if max_concurrency is not None:
+                        self._cur_concurrency_keys[job.task_name] -= 1
             if job.should_start:
                 job.status = JobStatus.QUEUED
                 queue = self._get_queue(job.queue)
@@ -71,6 +80,11 @@ def move_future_jobs(self) -> int:
 
         return num_jobs_moved
 
+    def set_concurrency_keys(self, tasks: Iterable[Task]):
+        for task in tasks:
+            self._max_concurrency_keys[task.name] = task.max_concurrency
+            self._cur_concurrency_keys[task.name] = 0
+
     def register_periodic_tasks(self, tasks: Iterable[Task]):
         """Register tasks that need to be scheduled periodically."""
         for task in tasks:
@@ -121,18 +135,46 @@ def _get_next_future_job(self) -> Optional[Job]:
             except IndexError:
                 return None
 
+    def is_queue_empty(self, queue: str):
+        return self._get_queue(queue).qsize() == 0
+
     def get_jobs_from_queue(self, queue: str, max_jobs: int) -> List[Job]:
         """Get jobs from a queue."""
         rv = list()
-        while len(rv) < max_jobs:
-            try:
-                job_json_string = self._get_queue(queue).get(block=False)
-            except Empty:
-                break
-
-            job = Job.deserialize(job_json_string)
-            job.status = JobStatus.RUNNING
-            rv.append(job)
+        jobs_to_re_add = list()
+        with self._lock:
+            while len(rv) < max_jobs:
+                try:
+                    job_json_string = self._get_queue(queue).get(block=False)
+                except Empty:
+                    break
+
+                job = Job.deserialize(job_json_string)
+                max_concurrency = self._max_concurrency_keys.get(job.task_name)
+                cur_concurrency = self._cur_concurrency_keys.get(job.task_name)
+                if (
+                    max_concurrency is not None and
+                    cur_concurrency >= max_concurrency
+                ):
+                    jobs_to_re_add.append(job_json_string)
+
+                else:
+                    job.status = JobStatus.RUNNING
+                    rv.append(job)
+                    if max_concurrency is not None:
+                        self._cur_concurrency_keys[job.task_name] += 1
+
+            # Re-add jobs that could not be run due to max_concurrency
+            # limits. Queue does not have a way to insert at the front, so
+            # sadly they go straight to the back again. Given that
+            # MemoryBroker is generally only used for testing, this should
+            # not be a great hardship.
+            logger.debug(
+                "Re-adding %s jobs due to concurrency limits",
+                len(jobs_to_re_add)
+            )
+            for job in jobs_to_re_add:
+                self._get_queue(queue).put(job)
 
         return rv
 
@@ -154,5 +196,13 @@ def remove_job_from_running(self, job: Job):
 
         Easy, the memory broker doesn't track running jobs. If the broker dies
         there is nothing we can do.
+
+        We still need to decrement the current_concurrency count,
+        however, if it exists.
         """
+        with self._lock:
+            max_concurrency = self._max_concurrency_keys[job.task_name]
+            if max_concurrency is not None:
+                self._cur_concurrency_keys[job.task_name] -= 1
+
         self._something_happened.set()

spinach/brokers/redis.py

@@ -17,7 +17,8 @@
 from ..const import (
     FUTURE_JOBS_KEY, NOTIFICATIONS_KEY, RUNNING_JOBS_KEY,
     PERIODIC_TASKS_HASH_KEY, PERIODIC_TASKS_QUEUE_KEY,
-    DEFAULT_ENQUEUE_JOB_RETRIES, ALL_BROKERS_HASH_KEY, ALL_BROKERS_ZSET_KEY
+    DEFAULT_ENQUEUE_JOB_RETRIES, ALL_BROKERS_HASH_KEY, ALL_BROKERS_ZSET_KEY,
+    MAX_CONCURRENCY_KEY, CURRENT_CONCURRENCY_KEY,
 )
 from ..utils import run_forever, call_with_retry
 
@@ -57,10 +58,16 @@ def __init__(self, redis: Optional[StrictRedis]=None,
         self._get_jobs_from_queue = self._load_script(
             'get_jobs_from_queue.lua'
         )
+        self._remove_job_from_running = self._load_script(
+            'remove_job_from_running.lua'
+        )
         self._move_future_jobs = self._load_script('move_future_jobs.lua')
         self._register_periodic_tasks = self._load_script(
             'register_periodic_tasks.lua'
         )
+        self._set_concurrency_keys = self._load_script(
+            'set_concurrency_keys.lua'
+        )
         self._reset()
 
     def _reset(self):
@@ -107,7 +114,11 @@ def _run_script(self, script: Script, *args):
 
         return rv
 
-    def enqueue_jobs(self, jobs: Iterable[Job]):
+    def is_queue_empty(self, queue: str) -> bool:
+        """Return True if the provided queue is empty."""
+        return self._r.llen(self._to_namespaced(queue)) == 0
+
+    def enqueue_jobs(self, jobs: Iterable[Job], from_failure: bool=False):
         """Enqueue a batch of jobs."""
         jobs_to_queue = list()
         for job in jobs:
@@ -124,6 +135,9 @@ def enqueue_jobs(self, jobs: Iterable[Job]):
                 self._to_namespaced(RUNNING_JOBS_KEY.format(self._id)),
                 self.namespace,
                 self._to_namespaced(FUTURE_JOBS_KEY),
+                self._to_namespaced(MAX_CONCURRENCY_KEY),
+                self._to_namespaced(CURRENT_CONCURRENCY_KEY),
+                1 if from_failure else 0,
                 *jobs_to_queue
             )
 
@@ -188,7 +202,9 @@ def get_jobs_from_queue(self, queue: str, max_jobs: int) -> List[Job]:
             self._to_namespaced(queue),
             self._to_namespaced(RUNNING_JOBS_KEY.format(self._id)),
             JobStatus.RUNNING.value,
-            max_jobs
+            max_jobs,
+            self._to_namespaced(MAX_CONCURRENCY_KEY),
+            self._to_namespaced(CURRENT_CONCURRENCY_KEY),
         )
 
         jobs = json.loads(jobs_json_string.decode())
@@ -198,10 +214,14 @@ def get_jobs_from_queue(self, queue: str, max_jobs: int) -> List[Job]:
 
     def remove_job_from_running(self, job: Job):
         if job.max_retries > 0:
-            self._r.hdel(
+            self._run_script(
+                self._remove_job_from_running,
                 self._to_namespaced(RUNNING_JOBS_KEY.format(self._id)),
-                str(job.id)
+                self._to_namespaced(MAX_CONCURRENCY_KEY),
+                self._to_namespaced(CURRENT_CONCURRENCY_KEY),
+                job.serialize(),
             )
+
         self._something_happened.set()
 
     def _subscriber_func(self):
@@ -272,19 +292,38 @@ def enqueue_jobs_from_dead_broker(self, dead_broker_id: uuid.UUID) -> int:
             self._to_namespaced(ALL_BROKERS_HASH_KEY),
             self._to_namespaced(ALL_BROKERS_ZSET_KEY),
             self.namespace,
-            self._to_namespaced(NOTIFICATIONS_KEY)
+            self._to_namespaced(NOTIFICATIONS_KEY),
+            self._to_namespaced(MAX_CONCURRENCY_KEY),
+            self._to_namespaced(CURRENT_CONCURRENCY_KEY),
         )
 
     def register_periodic_tasks(self, tasks: Iterable[Task]):
         """Register tasks that need to be scheduled periodically."""
-        tasks = [task.serialize() for task in tasks]
-        self._number_periodic_tasks = len(tasks)
+        _tasks = [task.serialize() for task in tasks]
+        self._number_periodic_tasks = len(_tasks)
         self._run_script(
             self._register_periodic_tasks,
             math.ceil(datetime.now(timezone.utc).timestamp()),
             self._to_namespaced(PERIODIC_TASKS_HASH_KEY),
             self._to_namespaced(PERIODIC_TASKS_QUEUE_KEY),
-            *tasks
+            *_tasks
+        )
+
+    def set_concurrency_keys(self, tasks: Iterable[Task]):
+        """For each Task, set up its concurrency keys.
+
+        The Lua script handles the logic of:
+            - removing dead keys where a Task was removed
+            - only setting keys where max_concurrency > 0
+        """
+        _tasks = [task.serialize() for task in tasks]
+        if not _tasks:
+            return
+        self._run_script(
+            self._set_concurrency_keys,
+            self._to_namespaced(MAX_CONCURRENCY_KEY),
+            self._to_namespaced(CURRENT_CONCURRENCY_KEY),
+            *_tasks,
         )
 
     def inspect_periodic_tasks(self) -> List[Tuple[int, str]]:

spinach/brokers/redis_scripts/enqueue_job.lua

@@ -4,16 +4,27 @@ local notifications = ARGV[2]
 local running_jobs_key = ARGV[3]
 local namespace = ARGV[4]
 local future_jobs = ARGV[5]
--- jobs starting at ARGV[6]
+local max_concurrency_key = ARGV[6]
+local current_concurrency_key = ARGV[7]
+local from_failure = ARGV[8]
+
+-- jobs starting at ARGV[9]
 
 if not redis.call('set', idempotency_token, 'true', 'EX', 3600, 'NX') then
     redis.log(redis.LOG_WARNING, "Not reprocessing script")
     return -1
 end
 
-for i=6, #ARGV do
+for i=9, #ARGV do
     local job_json = ARGV[i]
     local job = cjson.decode(job_json)
+    if tonumber(from_failure) == 1 then
+        -- job is being requeued after a failure, decrement its concurrency
+        local max_concurrency = tonumber(redis.call('hget', max_concurrency_key, job['task_name']))
+        if max_concurrency ~= nil and max_concurrency ~= -1 then
+            redis.call('hincrby', current_concurrency_key, job['task_name'], -1)
+        end
+    end
     if job["status"] == 2 then
         -- job status is queued
         local queue = string.format("%s/%s", namespace, job["queue"])

spinach/brokers/redis_scripts/enqueue_jobs_from_dead_broker.lua

@@ -4,6 +4,8 @@ local all_brokers_hash_key = ARGV[3]
 local all_brokers_zset_key = ARGV[4]
 local namespace = ARGV[5]
 local notifications = ARGV[6]
+local max_concurrency_key = ARGV[7]
+local current_concurrency_key = ARGV[8]
 
 local num_enqueued_jobs = 0
 
@@ -26,6 +28,13 @@ for _, job_json in ipairs(jobs_json) do
         -- Serialize the job so that it can be put in the queue
         local job_json = cjson.encode(job)
 
+        -- Decrement the current concurrency if we are tracking
+        -- concurrency on the Task.
+        local max_concurrency = tonumber(redis.call('hget', max_concurrency_key, job['task_name']))
+        if max_concurrency ~= nil and max_concurrency ~= -1 then
+            redis.call('hincrby', current_concurrency_key, job['task_name'], -1)
+        end
+
         -- Enqueue the job
         local queue = string.format("%s/%s", namespace, job["queue"])
         redis.call('rpush', queue, job_json)