Merge pull request #25 from taskiq-python/feature/zmq

Added ZeroMQ broker, fixed multiprocessing.

Author: s3rius

README.md

@@ -154,6 +154,8 @@ Or you can let taskiq find all python modules named tasks in current directory r
 taskiq test_project.broker:broker -fsd
 ```
 
+If you have uvloop installed, taskiq will automatically install new policies to event loop.
+
 You can always run `--help` to see all possible options.
 
 
@@ -262,3 +264,48 @@ async def main():
 asyncio.run(main())
 
 ```
+
+
+# Available Brokers
+
+Taskiq has several brokers out of the box:
+* InMemoryBroker
+* ZeroMQBroker
+
+## InMemoryBroker
+
+This broker is created for development purpose.
+You can easily use it without setting up workers for your project.
+
+It works the same as real brokers, but with some limitations.
+
+1. It cannot use `pre_execute` and `post_execute` hooks in middlewares.
+2. You cannot use it in multiprocessing applications without real result_backend.
+
+This broker is sutable for local development only.
+
+InMemoryBroker parameters:
+* `sync_tasks_pool_size` - number of threads for threadpool executor.
+    All sync functions are executed in threadpool.
+* `logs_format` - format which is used to collect logs from task execution.
+* `max_stored_results` - maximum number of results that is stored in memory. This number is used only if no custom result backend is provided.
+* `cast_types` - whether to use agressive type cast for types.
+* `result_backend` - custom result backend. By default
+    it uses `InmemoryResultBackend`.
+* `task_id_generator` - custom function to generate task ids.
+
+## ZeroMQBroker
+
+ZeroMQ is not available by default. To enable it. Please install [pyzmq](https://pypi.org/project/pyzmq/),
+or you can install `taskiq[zmq]`.
+
+This broker doesn't have limitations, but it requires you to set up a worker using taskiq CLI.
+Also please ensure that worker process is started after your application is ready.
+
+Because ZMQ publishes messages in socket and all worker processes will connect to it.
+
+ZeroMQBroker parameters:
+* `zmq_pub_host` - host that used to publish messages.
+* `zmq_sub_host` - host to subscribe to publisher.
+* `result_backend` - custom result backend.
+* `task_id_generator` - custom function to generate task ids.

poetry.lock

@@ -27,6 +27,8 @@ keywords = ["taskiq", "tasks", "distributed", "async"]
 python = "^3.7"
 typing-extensions = ">=3.10.0.0"
 pydantic = "^1.6.2"
+pyzmq = { version = "^23.2.0", optional = true }
+uvloop = { version = "^0.16.0", optional = true }
 
 [tool.poetry.dev-dependencies]
 pytest = "^7.1.2"
@@ -41,6 +43,10 @@ wemake-python-styleguide = "^0.16.1"
 coverage = "^6.4.2"
 pytest-cov = "^3.0.0"
 
+[tool.poetry.extras]
+zmq = ["pyzmq"]
+uv = ["uvloop"]
+
 [tool.poetry.scripts]
 taskiq = "taskiq.__main__:main"
 
@@ -54,6 +60,7 @@ show_error_codes = true
 implicit_reexport = true
 allow_untyped_decorators = true
 warn_return_any = false
+warn_unused_ignores = false
 
 [tool.isort]
 profile = "black"

pyproject.toml

@@ -3,7 +3,9 @@
 from taskiq.abc.formatter import TaskiqFormatter
 from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.abc.result_backend import AsyncResultBackend
+from taskiq.brokers.inmemory_broker import InMemoryBroker
 from taskiq.brokers.shared_broker import async_shared_broker
+from taskiq.brokers.zmq_broker import ZeroMQBroker
 from taskiq.exceptions import TaskiqError
 from taskiq.message import BrokerMessage, TaskiqMessage
 from taskiq.result import TaskiqResult
@@ -13,8 +15,10 @@
     "AsyncBroker",
     "TaskiqError",
     "TaskiqResult",
+    "ZeroMQBroker",
     "TaskiqMessage",
     "BrokerMessage",
+    "InMemoryBroker",
     "TaskiqFormatter",
     "AsyncTaskiqTask",
     "TaskiqMiddleware",

taskiq/__init__.py

@@ -57,6 +57,7 @@ class AsyncBroker(ABC):
     """
 
     available_tasks: Dict[str, AsyncTaskiqDecoratedTask[Any, Any]] = {}
+    is_worker_process: bool = False
 
     def __init__(
         self,
@@ -69,7 +70,6 @@ def __init__(
             task_id_generator = default_id_generator
         self.middlewares: "List[TaskiqMiddleware]" = []
         self.result_backend = result_backend
-        self.is_worker_process = False
         self.decorator_class = AsyncTaskiqDecoratedTask
         self.formatter: "TaskiqFormatter" = JSONFormatter()
         self.id_generator = task_id_generator
@@ -101,6 +101,7 @@ async def shutdown(self) -> None:
             middleware_shutdown = middleware.shutdown()
             if inspect.isawaitable(middleware_shutdown):
                 await middleware_shutdown
+        await self.result_backend.shutdown()
 
     @abstractmethod
     async def kick(

taskiq/abc/broker.py

@@ -0,0 +1,69 @@
+from typing import AsyncGenerator, Callable, Optional, TypeVar
+
+from taskiq.abc.broker import AsyncBroker
+from taskiq.abc.result_backend import AsyncResultBackend
+from taskiq.message import BrokerMessage
+
+try:
+    import zmq  # noqa: WPS433
+    from zmq.asyncio import Context  # noqa: WPS433
+except ImportError:
+    zmq = None  # type: ignore
+
+_T = TypeVar("_T")  # noqa: WPS111
+
+
+class ZeroMQBroker(AsyncBroker):
+    """
+    ZeroMQ broker.
+
+    This broker starts a socket ON A CLIENT SIDE,
+    and all workers connect to this socket using sub_host.
+
+    If you're using this socket you have to be sure,
+    that your workers start after the client is ready.
+    """
+
+    def __init__(
+        self,
+        zmq_pub_host: str = "tcp://0.0.0.0:5555",
+        zmq_sub_host: str = "tcp://localhost:5555",
+        result_backend: "Optional[AsyncResultBackend[_T]]" = None,
+        task_id_generator: Optional[Callable[[], str]] = None,
+    ) -> None:
+        if zmq is None:
+            raise RuntimeError(
+                "To use ZMQ broker please install pyzmq lib or taskiq[zmq].",
+            )
+        super().__init__(result_backend, task_id_generator)
+        self.context = Context()
+        self.pub_host = zmq_pub_host
+        self.sub_host = zmq_sub_host
+        if self.is_worker_process:
+            self.socket = self.context.socket(zmq.SUB)
+            self.socket.setsockopt(zmq.SUBSCRIBE, b"")
+        else:
+            self.socket = self.context.socket(zmq.PUB)
+            self.socket.bind(self.pub_host)
+
+    async def kick(self, message: BrokerMessage) -> None:
+        """
+        Kicking message.
+
+        This method is used to publish message
+        via socket.
+
+        :param message: message to publish.
+        """
+        with self.socket.connect(self.sub_host) as sock:
+            await sock.send_string(message.json())
+
+    async def listen(self) -> AsyncGenerator[BrokerMessage, None]:
+        """
+        Start accepting new messages.
+
+        :yield: received broker message
+        """
+        while True:  # noqa: WPS457
+            with self.socket.connect(self.sub_host) as sock:
+                yield BrokerMessage.parse_raw(await sock.recv_string())

taskiq/brokers/zmq_broker.py

@@ -1,12 +1,10 @@
 import asyncio
 import inspect
 import io
-import signal
-import sys
 from concurrent.futures import Executor, ThreadPoolExecutor
 from logging import getLogger
 from time import time
-from typing import Any, Callable, Dict, List, NoReturn, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 from pydantic import parse_obj_as
 
@@ -180,81 +178,6 @@ async def run_task(  # noqa: C901, WPS210, WPS211
     return result
 
 
-def exit_process(task: "asyncio.Task[Any]") -> NoReturn:
-    """
-    This function exits from the current process.
-
-    It receives asyncio Task of broker.shutdown().
-    We check if there were an exception or returned value.
-
-    If the function raised an exception, we print it with stack trace.
-    If it returned a value, we log it.
-
-    After this, we cancel all current tasks in the loop
-    and exits.
-
-    :param task: broker.shutdown task.
-    """
-    exitcode = 0
-    try:
-        result = task.result()
-        if result is not None:
-            logger.info("Broker returned value on shutdown: '%s'", str(result))
-    except Exception as exc:
-        logger.warning("Exception was found while shutting down!")
-        logger.warning(exc, exc_info=True)
-        exitcode = 1
-
-    loop = asyncio.get_event_loop()
-    for running_task in asyncio.all_tasks(loop):
-        running_task.cancel()
-
-    logger.info("Worker process killed.")
-    sys.exit(exitcode)
-
-
-def signal_handler(broker: AsyncBroker) -> Callable[[int, Any], None]:
-    """
-    Signal handler.
-
-    This function is used to generate
-    real signal handler using closures.
-
-    It takes current broker as an argument
-    and returns function that shuts it down.
-
-    :param broker: current broker.
-    :returns: signal handler function.
-    """
-
-    def _handler(signum: int, _frame: Any) -> None:
-        """
-        Exit signal handler.
-
-        This signal handler
-        calls shutdown for broker and after
-        the task is done it exits process with 0 status code.
-
-        :param signum: received signal.
-        :param _frame: current execution frame.
-        """
-        if getattr(broker, "_is_shutting_down", False):
-            # We're already shutting down the broker.
-            return
-
-        # We set this flag to not call this method twice.
-        # Since we add an asynchronous task in loop
-        # It can wait for execution for some time.
-        # We want to execute shutdown only once. Otherwise
-        # it would give us Undefined Behaviour.
-        broker._is_shutting_down = True  # type: ignore  # noqa: WPS437
-        logger.info(f"Got {signum} signal. Shutting down worker process.")
-        task = asyncio.create_task(broker.shutdown())
-        task.add_done_callback(exit_process)
-
-    return _handler
-
-
 async def async_listen_messages(  # noqa: C901, WPS210, WPS213
     broker: AsyncBroker,
     cli_args: TaskiqArgs,
@@ -268,15 +191,6 @@ async def async_listen_messages(  # noqa: C901, WPS210, WPS213
     :param broker: broker to listen to.
     :param cli_args: CLI arguments for worker.
     """
-    signal.signal(
-        signal.SIGTERM,
-        signal_handler(broker),
-    )
-    signal.signal(
-        signal.SIGINT,
-        signal_handler(broker),
-    )
-
     logger.info("Runing startup event.")
     await broker.startup()
     executor = ThreadPoolExecutor(