Separation of the receiver from the CLI (#90)

Author: hugiron

taskiq/brokers/inmemory_broker.py

@@ -1,17 +1,17 @@
 import asyncio
 import inspect
 from collections import OrderedDict
+from concurrent.futures import ThreadPoolExecutor
 from typing import Any, AsyncGenerator, Callable, Optional, Set, TypeVar, get_type_hints
 
 from taskiq_dependencies import DependencyGraph
 
 from taskiq.abc.broker import AsyncBroker
 from taskiq.abc.result_backend import AsyncResultBackend, TaskiqResult
-from taskiq.cli.worker.args import WorkerArgs
-from taskiq.cli.worker.receiver import Receiver
 from taskiq.events import TaskiqEvents
 from taskiq.exceptions import TaskiqError
 from taskiq.message import BrokerMessage
+from taskiq.receiver import Receiver
 from taskiq.utils import maybe_awaitable
 
 _ReturnType = TypeVar("_ReturnType")
@@ -91,11 +91,11 @@ class InMemoryBroker(AsyncBroker):
     def __init__(  # noqa: WPS211
         self,
         sync_tasks_pool_size: int = 4,
-        logs_format: Optional[str] = None,
         max_stored_results: int = 100,
         cast_types: bool = True,
         result_backend: Optional[AsyncResultBackend[Any]] = None,
         task_id_generator: Optional[Callable[[], str]] = None,
+        max_async_tasks: int = 30,
     ) -> None:
         if result_backend is None:
             result_backend = InmemoryResultBackend(
@@ -105,15 +105,12 @@ def __init__(  # noqa: WPS211
             result_backend=result_backend,
             task_id_generator=task_id_generator,
         )
+        self.executor = ThreadPoolExecutor(sync_tasks_pool_size)
         self.receiver = Receiver(
-            self,
-            WorkerArgs(
-                broker="",
-                modules=[],
-                max_threadpool_threads=sync_tasks_pool_size,
-                no_parse=not cast_types,
-                log_collector_format=logs_format or WorkerArgs.log_collector_format,
-            ),
+            broker=self,
+            executor=self.executor,
+            validate_params=cast_types,
+            max_async_tasks=max_async_tasks,
         )
         self._running_tasks: "Set[asyncio.Task[Any]]" = set()
 
@@ -170,3 +167,4 @@ async def shutdown(self) -> None:
         for event in (TaskiqEvents.CLIENT_SHUTDOWN, TaskiqEvents.WORKER_SHUTDOWN):
             for handler in self.event_handlers.get(event, []):
                 await maybe_awaitable(handler(self.state))
+        self.executor.shutdown()

taskiq/cli/worker/args.py

@@ -23,7 +23,7 @@ class WorkerArgs:
     shutdown_timeout: float = 5
     reload: bool = False
     no_gitignore: bool = False
-    max_async_tasks: int = 10
+    max_async_tasks: int = 100
 
     @classmethod
     def from_cli(  # noqa: WPS213
@@ -128,7 +128,7 @@ def from_cli(  # noqa: WPS213
             "--max-async-tasks",
             type=int,
             dest="max_async_tasks",
-            default=10,
+            default=100,
             help="Maximum simultaneous async tasks per worker process. ",
         )
 

taskiq/cli/worker/async_task_runner.py

@@ -1,15 +1,16 @@
 import asyncio
 import logging
 import signal
+from concurrent.futures import ThreadPoolExecutor
 from typing import Any
 
 from watchdog.observers import Observer
 
 from taskiq.abc.broker import AsyncBroker
 from taskiq.cli.utils import import_object, import_tasks
 from taskiq.cli.worker.args import WorkerArgs
-from taskiq.cli.worker.async_task_runner import async_listen_messages
 from taskiq.cli.worker.process_manager import ProcessManager
+from taskiq.receiver import Receiver
 
 try:
     import uvloop  # noqa: WPS433
@@ -102,7 +103,15 @@ def interrupt_handler(signum: int, _frame: Any) -> None:
 
     loop = asyncio.get_event_loop()
     try:
-        loop.run_until_complete(async_listen_messages(broker, args))
+        logger.debug("Initialize receiver.")
+        with ThreadPoolExecutor(args.max_threadpool_threads) as pool:
+            receiver = Receiver(
+                broker=broker,
+                executor=pool,
+                validate_params=not args.no_parse,
+                max_async_tasks=args.max_async_tasks,
+            )
+            loop.run_until_complete(receiver.listen())
     except KeyboardInterrupt:
         logger.warning("Worker process interrupted.")
         loop.run_until_complete(shutdown_broker(broker, args.shutdown_timeout))

taskiq/cli/worker/run.py

@@ -0,0 +1,6 @@
+"""Package for message receiver."""
+from taskiq.receiver.receiver import Receiver
+
+__all__ = [
+    "Receiver",
+]

taskiq/receiver/__init__.py

@@ -1,18 +1,17 @@
 import asyncio
 import inspect
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import Executor
 from logging import getLogger
 from time import time
-from typing import Any, Callable, Dict, get_type_hints
+from typing import Any, Callable, Dict, Optional, get_type_hints
 
 from taskiq_dependencies import DependencyGraph
 
 from taskiq.abc.broker import AsyncBroker
 from taskiq.abc.middleware import TaskiqMiddleware
-from taskiq.cli.worker.args import WorkerArgs
-from taskiq.cli.worker.params_parser import parse_params
 from taskiq.context import Context
 from taskiq.message import BrokerMessage, TaskiqMessage
+from taskiq.receiver.params_parser import parse_params
 from taskiq.result import TaskiqResult
 from taskiq.state import TaskiqState
 from taskiq.utils import maybe_awaitable
@@ -37,20 +36,24 @@ def _run_sync(target: Callable[..., Any], message: TaskiqMessage) -> Any:
 class Receiver:
     """Class that uses as a callback handler."""
 
-    def __init__(self, broker: AsyncBroker, cli_args: WorkerArgs) -> None:
+    def __init__(
+        self,
+        broker: AsyncBroker,
+        executor: Optional[Executor] = None,
+        validate_params: bool = True,
+        max_async_tasks: int = 20,
+    ) -> None:
         self.broker = broker
-        self.cli_args = cli_args
+        self.executor = executor
+        self.validate_params = validate_params
         self.task_signatures: Dict[str, inspect.Signature] = {}
         self.task_hints: Dict[str, Dict[str, Any]] = {}
         self.dependency_graphs: Dict[str, DependencyGraph] = {}
         for task in self.broker.available_tasks.values():
             self.task_signatures[task.task_name] = inspect.signature(task.original_func)
             self.task_hints[task.task_name] = get_type_hints(task.original_func)
             self.dependency_graphs[task.task_name] = DependencyGraph(task.original_func)
-        self.executor = ThreadPoolExecutor(
-            max_workers=cli_args.max_threadpool_threads,
-        )
-        self.sem = asyncio.Semaphore(cli_args.max_async_tasks)
+        self.sem = asyncio.Semaphore(max_async_tasks)
 
     async def callback(  # noqa: C901, WPS213
         self,
@@ -152,10 +155,10 @@ async def run_task(  # noqa: C901, WPS210
         loop = asyncio.get_running_loop()
         returned = None
         found_exception = None
-        signature = self.task_signatures.get(message.task_name)
+        signature = None
+        if self.validate_params:
+            signature = self.task_signatures.get(message.task_name)
         dependency_graph = self.dependency_graphs.get(message.task_name)
-        if self.cli_args.no_parse:
-            signature = None
         parse_params(signature, self.task_hints.get(message.task_name) or {}, message)
 
         dep_ctx = None
@@ -221,3 +224,25 @@ async def run_task(  # noqa: C901, WPS210
                     )
 
         return result
+
+    async def listen(self) -> None:  # pragma: no cover
+        """
+        This function iterates over tasks asynchronously.
+
+        It uses listen() method of an AsyncBroker
+        to get new messages from queues.
+        """
+        logger.debug("Runing startup event.")
+        await self.broker.startup()
+        logger.info("Listening started.")
+        tasks = set()
+        async for message in self.broker.listen():
+            task = asyncio.create_task(self.callback(message=message, raise_err=False))
+            tasks.add(task)
+
+            # We want the task to remove itself from the set when it's done.
+            #
+            # Because python's GC can silently cancel task
+            # and it considered to be Hisenbug.
+            # https://textual.textualize.io/blog/2023/02/11/the-heisenbug-lurking-in-your-async-code/
+            task.add_done_callback(tasks.discard)

taskiq/cli/worker/params_parser.py renamed to taskiq/receiver/params_parser.py

@@ -5,8 +5,8 @@
 import pytest
 from pydantic import BaseModel
 
-from taskiq.cli.worker.params_parser import parse_params
 from taskiq.message import TaskiqMessage
+from taskiq.receiver.params_parser import parse_params
 
 
 class _TestPydanticClass(BaseModel):

taskiq/cli/worker/receiver.py renamed to taskiq/receiver/receiver.py

@@ -1,4 +1,5 @@
 import asyncio
+from concurrent.futures import ThreadPoolExecutor
 from typing import Any, Optional
 
 import pytest
@@ -7,16 +8,15 @@
 from taskiq.abc.broker import AsyncBroker
 from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.brokers.inmemory_broker import InMemoryBroker
-from taskiq.cli.worker.args import WorkerArgs
-from taskiq.cli.worker.receiver import Receiver
 from taskiq.message import BrokerMessage, TaskiqMessage
+from taskiq.receiver import Receiver
 from taskiq.result import TaskiqResult
 
 
 def get_receiver(
     broker: Optional[AsyncBroker] = None,
     no_parse: bool = False,
-    cli_args: Optional[WorkerArgs] = None,
+    max_async_tasks: int = 10,
 ) -> Receiver:
     """
     Returns receiver with custom broker and args.
@@ -28,15 +28,11 @@ def get_receiver(
     """
     if broker is None:
         broker = InMemoryBroker()
-    if cli_args is None:
-        cli_args = WorkerArgs(
-            broker="",
-            modules=[],
-            no_parse=no_parse,
-        )
     return Receiver(
         broker,
-        cli_args,
+        executor=ThreadPoolExecutor(max_workers=10),
+        validate_params=not no_parse,
+        max_async_tasks=max_async_tasks,
     )
 
 
@@ -261,13 +257,7 @@ async def task_sem() -> int:
         await asyncio.sleep(1)
         return 1
 
-    cli_args = WorkerArgs(
-        broker="",
-        modules=[],
-        no_parse=False,
-        max_async_tasks=3,
-    )
-    receiver = get_receiver(broker, cli_args=cli_args)
+    receiver = get_receiver(broker, max_async_tasks=3)
 
     broker_message = broker.formatter.dumps(
         TaskiqMessage(
@@ -279,5 +269,6 @@ async def task_sem() -> int:
         ),
     )
     tasks = [asyncio.create_task(receiver.callback(broker_message)) for _ in range(5)]
-    await asyncio.sleep(0)
+    await asyncio.sleep(0.3)
     assert sem_num == 3
+    await asyncio.gather(*tasks)