Support running coroutines in new asyncio worker

This commit adds a new type of worker, in addition to the existing
threaded worker, that can run tasks defined as asyncio coroutines.

This makes Spinach compatible with the whole asyncio ecosystem while
providing better concurrency for tasks that are heavily IO bound.

The implementation swaps the queue used between the engine and the
workers with one that works both with sync and async code. Then the
existing `Workers` class is split into two main parts: the threaded
workers (using the sync part of the queue) and the asyncio workers
(using the async part of the queue).

There are two main shortcomings to this implementation:

1. Scheduling jobs is still blocking, so when creating a job from an
   asyncio task, care must be taken to wrap the call in
   `asyncio.to_thread` to prevent the event loop from blocking.

2. The compatibility with 3rd-party integrations defined in the contrib
   package is not guarantied when running asyncio tasks.

Author: NicolasLM

README.rst

@@ -14,14 +14,15 @@ Redis task queue for Python 3 heavily inspired by Celery and RQ.
 
 Distinctive features:
 
+- Threaded and asyncio workers
 - 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>`_
 - Embeddable workers for easier testing
-- Python 3, threaded, explicit... see `design choices
+- See `design choices
   <https://spinach.readthedocs.io/en/stable/user/design.html>`_ for more
   details
 

doc/index.rst

@@ -7,15 +7,15 @@ Spinach is a Redis task queue for Python 3 heavily inspired by Celery and RQ.
 
 Distinctive features:
 
+- Threaded and asyncio workers
 - 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
   <integrations>`
-- Python 3, threaded, explicit... see :ref:`design choices <design>` for more
-  details
+- See :ref:`design choices <design>` for more details
 
 Installation::
 
@@ -53,6 +53,7 @@ Getting started with spinach:
     user/jobs
     user/engine
     user/queues
+    user/asyncio
     user/integrations
     user/signals
     user/production

doc/user/asyncio.rst

@@ -0,0 +1,51 @@
+.. _asyncio:
+
+Asyncio
+=======
+
+Spinach allows to define and run tasks as asyncio coroutines. In this mode the worker is a single
+thread that runs all tasks asynchronously. This allows for greater concurrency as well as
+compatibility with the asyncio ecosystem.
+
+Creating async tasks
+--------------------
+
+To define an asynchronous task, just prefix its definition with the ``async`` keyword::
+
+    @spin.task(name='compute')
+    async def compute(a, b):
+        await asyncio.sleep(1)
+        print('Computed {} + {} = {}'.format(a, b, a + b))
+
+To run the workers in asynchronous mode, pass the ``AsyncioWorkers`` class to ``start_workers``::
+
+    from spinach import AsyncioWorkers
+
+    spin.start_workers(number=256, workers_class=AsyncioWorkers)
+
+When using the asyncio workers, the ``number`` argument can be set quite high because each worker
+is just a coroutine, consuming a negligible amount of resources.
+
+Scheduling jobs
+---------------
+
+Because internally only workers are asyncio aware, jobs are still sent to Redis using a blocking
+socket. This means that to schedule jobs from asynchronous code, care must be taken to send jobs
+from outside the event loop. This can be achieve using `asyncio.to_thread
+<https://docs.python.org/3/library/asyncio-task.html#asyncio.to_thread>`_::
+
+    await asyncio.to_thread(spin.schedule, compute, 2, 4)
+
+Code scheduling a lot of jobs should use :ref:`batches <batch>` to improve performance.
+
+Example
+-------
+
+.. literalinclude:: ../../examples/asyncio_workers.py
+
+
+.. note:: If an application defines both sync and async tasks, each kind of task should go in its
+          own :ref:`queue <queues>` so that sync tasks are picked by threaded workers and async
+          tasks by asyncio workers.
+
+.. note:: Not all contrib :ref:`integrations <integrations>` may work with asynchronous workers.

doc/user/design.rst

@@ -11,33 +11,22 @@ be summed up as: explicit is better than implicit. Spinach makes sure that it
 does not provide any convenient feature that can backfire in more complex
 usages.
 
-Threaded workers
-----------------
+Threaded & asynchronous workers
+-------------------------------
 
-Spinach workers are threaded while other task queues like Celery or RQ rely on
-processes by default.
+Spinach workers are either threaded or asynchronous while other task queues
+like Celery or RQ rely on processes by default.
 
-Threaded workers work best with IO bound tasks: tasks that make requests to
-other services, query a database or read files. If your task are CPU bound,
-meaning that you do heavy computations in Python, a process based worker will
-be more efficient.
+Threaded and asynchronous workers work best with IO bound tasks: tasks that
+make requests to other services, query a database or read files. If your tasks
+are CPU bound, meaning that you do heavy computations in Python, a process
+based worker will be more efficient.
 
 Tasks in a typical web application are more often than not IO bound. The choice
-of threads as unit of concurrency is a sensible one.
-
-Threads also have the advantage of being lighter than processes, a system can
-handle more threads than processes before resources get exhausted.
-
-Thread safety
-~~~~~~~~~~~~~
-
-As Spinach workers are threads, care must be taken to make sure that the
-application is thread-safe. The good news is that your application is probably
-already thread-safe: web frameworks are often run threaded as well, so they
-take care of most of the heavy work for you.
+of threads or coroutines as unit of concurrency is a sensible one.
 
-You can read an article I wrote for an `introduction to thread-safety
-<https://lemanchet.fr/articles/learning-python-3-threading-module.html>`_.
+Threads and coroutines also have the advantage of being lighter than processes,
+a system can handle more threads than processes before resources get exhausted.
 
 Fork
 ~~~~
@@ -191,4 +180,4 @@ the ability to process jobs.
 
 Because worker processes can die unexpectedly (power loss, OOM killed, extended
 network outage...), Spinach tries to detect dead workers and reschedule
-the jobs that were running on them if the jobs are safe to be retried.
+the jobs that were running on them if the jobs are safe to be retried.

doc/user/production.rst

@@ -8,9 +8,9 @@ Advices to read before deploying an application using Spinach to production.
 Spinach
 -------
 
-Since Spinach relies heavily on threads the user's code MUST be thread-safe.
-This is usually quite easy to achieve on a traditional web application because
-frameworks like Flask or Django make that obvious.
+Since by default Spinach executes jobs in a separate threads, the user's code
+must be thread-safe. This is usually quite easy to achieve on a traditional web
+application because frameworks like Flask or Django make that straightforward.
 
 Tasks should not store state in the process between invocations. Instead all
 state must be stored in an external system, like a database or a cache. This
@@ -70,7 +70,7 @@ Spinach:
 - Task `args` and `kwargs` are JSON serializable and small in size
 - Jobs are sent in :class:`Batch` to the broker when multiple jobs are to be
   scheduled at once
-- The user's code is thread-safe
+- The user's code is thread-safe when using the default threaded workers
 - Tasks do not store state in the process between invocations
 - Logging is configured and exceptions are sent to Sentry, see
   :doc:`integrations`

doc/user/tasks.rst

@@ -173,6 +173,8 @@ directly on the :class:`Engine` using::
 .. autoclass:: spinach.task.Tasks
     :members:
 
+.. _batch:
+
 Batch
 -----
 

examples/asyncio_workers.py

@@ -0,0 +1,30 @@
+import aiohttp
+from spinach import Engine, MemoryBroker, Batch, AsyncioWorkers
+
+spin = Engine(MemoryBroker())
+
+
+@spin.task(name='get_pokemon_name')
+async def get_pokemon_name(pokemon_id: int):
+    """Call an HTTP API to retrieve a pokemon name by its ID."""
+    url = f'https://pokeapi.co/api/v2/pokemon/{pokemon_id}'
+    async with aiohttp.ClientSession() as session:
+        async with session.get(url) as response:
+            pokemon = await response.json()
+
+    print(f'Pokemon #{pokemon_id} is {pokemon["name"]}')
+
+
+# Schedule a batch of 150 tasks to retrieve the name of the
+# first 150 pokemons.
+batch = Batch()
+for pokemon_id in range(1, 151):
+    batch.schedule(get_pokemon_name, pokemon_id)
+spin.schedule_batch(batch)
+
+# Start the asyncio workers and process the tasks
+spin.start_workers(
+    number=256,
+    workers_class=AsyncioWorkers,
+    stop_when_queue_empty=True
+)

spinach/__init__.py

@@ -3,5 +3,6 @@
 from .const import VERSION
 from .engine import Engine
 from .task import Tasks, Batch, RetryException, AbortException
+from .worker import ThreadWorkers, AsyncioWorkers
 
 __version__ = VERSION

spinach/engine.py

@@ -1,13 +1,14 @@
 from datetime import datetime, timezone
 from logging import getLogger
 import threading
+from typing import Type
 
 from .task import Tasks, Batch, Schedulable
 from .utils import run_forever, handle_sigterm
 from .job import Job, JobStatus, advance_job_status
 from .brokers.base import Broker
 from .const import DEFAULT_QUEUE, DEFAULT_NAMESPACE, DEFAULT_WORKER_NUMBER
-from .worker import Workers
+from .worker import BaseWorkers, ThreadWorkers
 from . import exc
 
 
@@ -185,17 +186,21 @@ def _arbiter_func(self, stop_when_queue_empty=False):
 
         logger.debug('Arbiter terminated')
 
-    def start_workers(self, number: int=DEFAULT_WORKER_NUMBER,
-                      queue=DEFAULT_QUEUE, block=True,
-                      stop_when_queue_empty=False):
+    def start_workers(self, number: int = DEFAULT_WORKER_NUMBER,
+                      queue: str = DEFAULT_QUEUE, block: bool = True,
+                      stop_when_queue_empty=False,
+                      workers_class: Type[BaseWorkers] = ThreadWorkers):
         """Start the worker threads.
 
-        :arg number: number of worker threads to launch
-        :arg queue: name of the queue to consume, see :doc:`queues`
+        :arg number: number of workers to launch, each job running uses one
+                     worker.
+        :arg queue: name of the queue to consume, see :doc:`queues`.
         :arg block: whether to block the calling thread until a signal arrives
-             and workers get terminated
+             and workers get terminated.
         :arg stop_when_queue_empty: automatically stop the workers when the
              queue is empty. Useful mostly for one-off scripts and testing.
+        :arg worker_class: Class to change the behavior of workers,
+             defaults to threaded workers
         """
         if self._arbiter or self._workers:
             raise RuntimeError('Workers are already running')
@@ -213,7 +218,7 @@ def start_workers(self, number: int=DEFAULT_WORKER_NUMBER,
         self._broker.start()
 
         # Start workers
-        self._workers = Workers(
+        self._workers = workers_class(
             num_workers=number,
             namespace=self.namespace,
         )

spinach/queuey.py

@@ -0,0 +1,115 @@
+import asyncio
+from collections import deque
+from concurrent.futures import Future
+import threading
+from typing import Tuple, Optional, Any, Deque
+
+
+class Queuey:
+    """Hybrid queue allowing to interface sync and async(io) code.
+
+    It is widely inspired by a talk by David Beazley on the subject:
+    https://www.youtube.com/watch?v=x1ndXuw7S0s
+
+    One big difference with a normal queue is that even with a maxsize
+    set to a fixed number, this queue can still end up taking an
+    infinite amount of memory since pending get/put operation are kept
+    as futures.
+
+    It is an alternative to the 3rd-party Janus library which had
+    shortcomings when used in Spinach:
+    - Janus queues have to be created in an asyncio coroutine, turning
+      the creation of the queue in the Workers class into a strange dance.
+    - It was not obvious to me how to implement showing the queue as full
+      if there are unfinished tasks.
+    - It adds a few dependencies only needed by a fractions of users, adds a
+      ton of code for something that should be simple.
+    """
+
+    def __init__(self, maxsize: int):
+        self.maxsize = maxsize
+        self._mutex = threading.Lock()
+        self._items: Deque[Any] = deque()
+        self._getters: Deque[Future] = deque()
+        self._putters: Deque[Tuple[Any, Future]] = deque()
+        self._unfinished_tasks = 0
+
+    def _get_noblock(self) -> Tuple[Optional[Any], Optional[Future]]:
+        with self._mutex:
+            if self._items:
+                if self._putters:
+                    # About to remove one item from the queue which means
+                    # that a new spot will be available. Since there are
+                    # putters waiting, wake up one and take its item.
+                    item, put_fut = self._putters.popleft()
+                    self._items.append(item)
+                    put_fut.set_result(True)
+                return self._items.popleft(), None
+
+            else:
+                fut = Future()
+                self._getters.append(fut)
+                return None, fut
+
+    def _put_noblock(self, item: Any) -> Optional[Future]:
+        with self._mutex:
+            if len(self._items) < self.maxsize:
+                self._items.append(item)
+                self._unfinished_tasks += 1
+                if self._getters:
+                    self._getters.popleft().set_result(self._items.popleft())
+            else:
+                fut = Future()
+                self._putters.append((item, fut))
+                return fut
+
+    def get_sync(self) -> Any:
+        item, fut = self._get_noblock()
+        if fut:
+            item = fut.result()
+        return item
+
+    def put_sync(self, item: Any) -> None:
+        fut = self._put_noblock(item)
+        if fut is None:
+            return
+
+        fut.result()
+
+    async def get_async(self) -> Any:
+        item, fut = self._get_noblock()
+        if fut:
+            item = await asyncio.wait_for(asyncio.wrap_future(fut), None)
+        return item
+
+    async def put_async(self, item: Any) -> None:
+        fut = self._put_noblock(item)
+        if fut is None:
+            return
+
+        await asyncio.wait_for(asyncio.wrap_future(fut), None)
+
+    def task_done(self) -> None:
+        """Indicate that a formerly enqueued task is complete.
+
+        Raises a ValueError if called more times than there were items
+        placed in the queue.
+        """
+        with self._mutex:
+            unfinished = self._unfinished_tasks - 1
+            if unfinished < 0:
+                raise ValueError('task_done() called too many times')
+
+            self._unfinished_tasks = unfinished
+
+    def empty(self) -> bool:
+        with self._mutex:
+            return self._unfinished_tasks == 0
+
+    def full(self) -> bool:
+        with self._mutex:
+            return self.maxsize <= self._unfinished_tasks
+
+    def available_slots(self) -> int:
+        with self._mutex:
+            return self.maxsize - self._unfinished_tasks