Added better support for multi-broker setup. (#169)

Author: s3rius

docs/examples/extending/middleware_async.py

@@ -5,6 +5,14 @@
 
 
 class MyMiddleware(TaskiqMiddleware):
+    async def startup(self) -> None:
+        print("RUN STARTUP")
+        await sleep(1)
+
+    async def shutdown(self) -> None:
+        print("RUN SHUTDOWN")
+        await sleep(1)
+
     async def pre_execute(self, message: "TaskiqMessage") -> TaskiqMessage:
         await sleep(1)
         return message

docs/examples/extending/middleware_sync.py

@@ -5,6 +5,14 @@
 
 
 class MyMiddleware(TaskiqMiddleware):
+    def startup(self) -> None:
+        print("RUN STARTUP")
+        sleep(1)
+
+    def shutdown(self) -> None:
+        print("RUN SHUTDOWN")
+        sleep(1)
+
     def pre_execute(self, message: "TaskiqMessage") -> TaskiqMessage:
         sleep(1)
         return message

taskiq/abc/broker.py

@@ -69,10 +69,6 @@ class AsyncBroker(ABC):
     """
 
     available_tasks: Dict[str, AsyncTaskiqDecoratedTask[Any, Any]] = {}
-    # True only if broker runs in worker process.
-    is_worker_process: bool = False
-    # True only if broker runs in scheduler process.
-    is_scheduler_process: bool = False
 
     def __init__(
         self,
@@ -111,6 +107,10 @@ def __init__(
         ] = defaultdict(list)
         self.state = TaskiqState()
         self.custom_dependency_context: Dict[Any, Any] = {}
+        # True only if broker runs in worker process.
+        self.is_worker_process: bool = False
+        # True only if broker runs in scheduler process.
+        self.is_scheduler_process: bool = False
 
     def add_dependency_context(self, new_ctx: Dict[Any, Any]) -> None:
         """
@@ -151,6 +151,10 @@ async def startup(self) -> None:
         for handler in self.event_handlers[event]:
             await maybe_awaitable(handler(self.state))
 
+        for middleware in self.middlewares:
+            if middleware.__class__.startup != TaskiqMiddleware.startup:
+                await maybe_awaitable(middleware.startup)
+
         await self.result_backend.startup()
 
     async def shutdown(self) -> None:
@@ -168,6 +172,10 @@ async def shutdown(self) -> None:
         for handler in self.event_handlers[event]:
             await maybe_awaitable(handler(self.state))
 
+        for middleware in self.middlewares:
+            if middleware.__class__.shutdown != TaskiqMiddleware.shutdown:
+                await maybe_awaitable(middleware.shutdown)
+
         await self.result_backend.shutdown()
 
     @abstractmethod

taskiq/abc/middleware.py

@@ -20,6 +20,26 @@ def set_broker(self, broker: "AsyncBroker") -> None:
         """
         self.broker = broker
 
+    def startup(self) -> "Union[None, Coroutine[Any, Any, None]]":
+        """
+        Startup method to perform various action during startup.
+
+        This function can be either sync or async.
+        Executed during broker's startup.
+
+        :returns nothing.
+        """
+
+    def shutdown(self) -> "Union[None, Coroutine[Any, Any, None]]":
+        """
+        Shutdown method to perform various action during shutdown.
+
+        This function can be either sync or async.
+        Executed during broker's shutdown.
+
+        :returns nothing.
+        """
+
     def pre_send(
         self,
         message: "TaskiqMessage",

taskiq/brokers/zmq_broker.py

@@ -8,7 +8,7 @@
 
 try:
     import zmq  # noqa: WPS433
-    from zmq.asyncio import Context  # noqa: WPS433
+    from zmq.asyncio import Context, Socket  # noqa: WPS433
 except ImportError:
     zmq = None  # type: ignore
 
@@ -43,12 +43,23 @@ def __init__(
         self.context = Context()
         self.pub_host = zmq_pub_host
         self.sub_host = zmq_sub_host
+        self.socket: Socket
+
+    async def startup(self) -> None:
+        """
+        Startup for zmq broker.
+
+        This function creates actual connections to
+        sockets. if current process is worker,
+        it subscribes, otherwise it becomes publisher.
+        """
         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)
+        await super().startup()
 
     async def kick(self, message: BrokerMessage) -> None:
         """

taskiq/cli/scheduler/run.py

@@ -5,7 +5,6 @@
 
 from pycron import is_now
 
-from taskiq.abc.broker import AsyncBroker
 from taskiq.cli.scheduler.args import SchedulerArgs
 from taskiq.cli.utils import import_object, import_tasks
 from taskiq.kicker import AsyncKicker
@@ -71,13 +70,16 @@ async def run_scheduler(args: SchedulerArgs) -> None:  # noqa: C901, WPS210, WPS
 
     :param args: parsed CLI args.
     """
-    AsyncBroker.is_scheduler_process = True
-    scheduler = import_object(args.scheduler)
-    if not isinstance(scheduler, TaskiqScheduler):
+    if isinstance(args.scheduler, str):
+        scheduler = import_object(args.scheduler)
+    else:
+        scheduler = args.scheduler
+    if not isinstance(args.scheduler, TaskiqScheduler):
         print(  # noqa: WPS421
             "Imported scheduler is not a subclass of TaskiqScheduler.",
         )
         exit(1)  # noqa: WPS421
+    scheduler.broker.is_scheduler_process = True
     import_tasks(args.modules, args.tasks_pattern, args.fs_discover)
     basicConfig(
         level=getLevelName(args.log_level),

taskiq/cli/worker/run.py

@@ -118,11 +118,11 @@ def interrupt_handler(signum: int, _frame: Any) -> None:
     # broker is running as a worker.
     # We must set this field before importing tasks,
     # so broker will remember all tasks it's related to.
-    AsyncBroker.is_worker_process = True
     broker = import_object(args.broker)
-    import_tasks(args.modules, args.tasks_pattern, args.fs_discover)
     if not isinstance(broker, AsyncBroker):
         raise ValueError("Unknown broker type. Please use AsyncBroker instance.")
+    broker.is_worker_process = True
+    import_tasks(args.modules, args.tasks_pattern, args.fs_discover)
 
     receiver_type = get_receiver_type(args)
     receiver_args = dict(args.receiver_arg)

taskiq/middlewares/prometheus_middleware.py

@@ -4,7 +4,6 @@
 from tempfile import gettempdir
 from typing import Any, Optional
 
-from taskiq.abc.broker import AsyncBroker
 from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.message import TaskiqMessage
 from taskiq.result import TaskiqResult
@@ -36,9 +35,6 @@ def __init__(
         self.saved_results = None
         self.execution_time = None
 
-        if not AsyncBroker.is_worker_process:
-            return
-
         metrics_path = metrics_path or Path(gettempdir()) / "taskiq_worker"
 
         if not metrics_path.exists():
@@ -52,11 +48,7 @@ def __init__(
         logger.debug("Initializing metrics")
 
         try:
-            from prometheus_client import (  # noqa: WPS433
-                Counter,
-                Histogram,
-                start_http_server,
-            )
+            from prometheus_client import Counter, Histogram  # noqa: WPS433
         except ImportError as exc:
             raise ImportError(
                 "Cannot initialize metrics. Please install 'taskiq[metrics]'.",
@@ -87,10 +79,23 @@ def __init__(
             "Tome of function execution",
             ["task_name"],
         )
-        try:
-            start_http_server(port=server_port, addr=server_addr)
-        except OSError as exc:
-            logger.debug("Cannot start prometheus server: %s", exc)
+        self.server_port = server_port
+        self.server_addr = server_addr
+
+    def startup(self) -> None:
+        """
+        Prometheus startup.
+
+        This function starts prometheus server.
+        It starts it only in case if it's a worker process.
+        """
+        from prometheus_client import start_http_server  # noqa: WPS433
+
+        if self.broker.is_worker_process:
+            try:
+                start_http_server(port=self.server_port, addr=self.server_addr)
+            except OSError as exc:
+                logger.debug("Cannot start prometheus server: %s", exc)
 
     def pre_execute(
         self,

tests/conftest.py

@@ -26,5 +26,3 @@ def reset_broker() -> Generator[None, None, None]:
     """
     yield
     AsyncBroker.available_tasks = {}
-    AsyncBroker.is_worker_process = False
-    AsyncBroker.is_scheduler_process = False