asyncio.wait replaced with Semaphore, fixed compatibility issues (#93)

Author: vvanglro

.flake8

@@ -1,5 +1,5 @@
 [flake8]
-max-complexity = 6
+max-complexity = 7
 inline-quotes = double
 max-line-length = 88
 extend-ignore = E203
@@ -90,6 +90,8 @@ ignore =
   N802,
   ; Do not perform function calls in argument defaults.
   B008,
+  ; Found too many public instance attributes
+  WPS230,
 
   ; all init files
   __init__.py:

taskiq/__main__.py

@@ -7,7 +7,7 @@
 from taskiq.abc.cmd import TaskiqCMD
 
 
-def main() -> None:  # noqa: C901, WPS210  # pragma: no cover
+def main() -> None:  # noqa: WPS210  # pragma: no cover
     """
     Main entrypoint of the taskiq.
 

taskiq/abc/broker.py

@@ -30,7 +30,7 @@
 from taskiq.message import BrokerMessage
 from taskiq.result_backends.dummy import DummyResultBackend
 from taskiq.state import TaskiqState
-from taskiq.utils import maybe_awaitable
+from taskiq.utils import maybe_awaitable, remove_suffix
 
 if TYPE_CHECKING:  # pragma: no cover
     from taskiq.abc.formatter import TaskiqFormatter
@@ -57,7 +57,7 @@ def default_id_generator() -> str:
     return uuid4().hex
 
 
-class AsyncBroker(ABC):  # noqa: WPS230
+class AsyncBroker(ABC):
     """
     Async broker.
 
@@ -235,11 +235,7 @@ def inner(
                     fmodule = func.__module__
                     if fmodule == "__main__":  # pragma: no cover
                         fmodule = ".".join(
-                            sys.argv[0]
-                            .removesuffix(
-                                ".py",
-                            )
-                            .split(
+                            remove_suffix(sys.argv[0], ".py").split(
                                 os.path.sep,
                             ),
                         )

taskiq/cli/utils.py

@@ -6,6 +6,8 @@
 from pathlib import Path
 from typing import Any, Generator, List
 
+from taskiq.utils import remove_suffix
+
 logger = getLogger("taskiq.worker")
 
 
@@ -81,6 +83,8 @@ def import_tasks(modules: List[str], pattern: str, fs_discover: bool) -> None:
     """
     if fs_discover:
         for path in Path(".").rglob(pattern):
-            modules.append(str(path).removesuffix(".py").replace(os.path.sep, "."))
+            modules.append(
+                remove_suffix(str(path), ".py").replace(os.path.sep, "."),
+            )
 
     import_from_modules(modules)

taskiq/cli/watcher.py

@@ -24,7 +24,7 @@ def __init__(
             self.gitignore = parse_gitignore(gpath)
         self.callback_kwargs = callback_kwargs
 
-    def dispatch(self, event: FileSystemEvent) -> None:  # noqa: C901
+    def dispatch(self, event: FileSystemEvent) -> None:
         """
         React to event.
 

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 = 100
+    max_async_tasks: int = 10
 
     @classmethod
     def from_cli(  # noqa: WPS213
@@ -128,9 +128,8 @@ def from_cli(  # noqa: WPS213
             "--max-async-tasks",
             type=int,
             dest="max_async_tasks",
-            default=100,
-            help="Maximum simultaneous async tasks per worker process. "
-            + "Infinite if less than 1",
+            default=10,
+            help="Maximum simultaneous async tasks per worker process. ",
         )
 
         namespace = parser.parse_args(args)

taskiq/cli/worker/async_task_runner.py

@@ -39,11 +39,4 @@ async def async_listen_messages(
         # 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)
-
-        # If we have finite number of maximum simultanious tasks,
-        # we await them when we reached the limit.
-        # But we don't await all of them, we await only first completed task,
-        # and then continue.
-        if 1 <= cli_args.max_async_tasks <= len(tasks):
-            _, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
-            tasks = pending
+        logger.debug("Received {0} tasks".format(len(tasks)))

taskiq/cli/worker/receiver.py

@@ -50,6 +50,7 @@ def __init__(self, broker: AsyncBroker, cli_args: WorkerArgs) -> None:
         self.executor = ThreadPoolExecutor(
             max_workers=cli_args.max_threadpool_threads,
         )
+        self.sem = asyncio.Semaphore(cli_args.max_async_tasks)
 
     async def callback(  # noqa: C901, WPS213
         self,
@@ -68,61 +69,62 @@ async def callback(  # noqa: C901, WPS213
         :param raise_err: raise an error if cannot save result in
             result_backend.
         """
-        logger.debug(f"Received message: {message}")
-        if message.task_name not in self.broker.available_tasks:
-            logger.warning(
-                'task "%s" is not found. Maybe you forgot to import it?',
+        async with self.sem:
+            logger.debug(f"Received message: {message}")
+            if message.task_name not in self.broker.available_tasks:
+                logger.warning(
+                    'task "%s" is not found. Maybe you forgot to import it?',
+                    message.task_name,
+                )
+                return
+            logger.debug(
+                "Function for task %s is resolved. Executing...",
                 message.task_name,
             )
-            return
-        logger.debug(
-            "Function for task %s is resolved. Executing...",
-            message.task_name,
-        )
-        try:
-            taskiq_msg = self.broker.formatter.loads(message=message)
-        except Exception as exc:
-            logger.warning(
-                "Cannot parse message: %s. Skipping execution.\n %s",
-                message,
-                exc,
-                exc_info=True,
-            )
-            return
-        for middleware in self.broker.middlewares:
-            if middleware.__class__.pre_execute != TaskiqMiddleware.pre_execute:
-                taskiq_msg = await maybe_awaitable(
-                    middleware.pre_execute(
-                        taskiq_msg,
-                    ),
+            try:
+                taskiq_msg = self.broker.formatter.loads(message=message)
+            except Exception as exc:
+                logger.warning(
+                    "Cannot parse message: %s. Skipping execution.\n %s",
+                    message,
+                    exc,
+                    exc_info=True,
                 )
+                return
+            for middleware in self.broker.middlewares:
+                if middleware.__class__.pre_execute != TaskiqMiddleware.pre_execute:
+                    taskiq_msg = await maybe_awaitable(
+                        middleware.pre_execute(
+                            taskiq_msg,
+                        ),
+                    )
 
-        logger.info(
-            "Executing task %s with ID: %s",
-            taskiq_msg.task_name,
-            taskiq_msg.task_id,
-        )
-        result = await self.run_task(
-            target=self.broker.available_tasks[message.task_name].original_func,
-            message=taskiq_msg,
-        )
-        for middleware in self.broker.middlewares:
-            if middleware.__class__.post_execute != TaskiqMiddleware.post_execute:
-                await maybe_awaitable(middleware.post_execute(taskiq_msg, result))
-        try:
-            await self.broker.result_backend.set_result(message.task_id, result)
-        except Exception as exc:
-            logger.exception(
-                "Can't set result in result backend. Cause: %s",
-                exc,
-                exc_info=True,
+            logger.info(
+                "Executing task %s with ID: %s",
+                taskiq_msg.task_name,
+                taskiq_msg.task_id,
+            )
+            result = await self.run_task(
+                target=self.broker.available_tasks[message.task_name].original_func,
+                message=taskiq_msg,
             )
-            if raise_err:
-                raise exc
+            for middleware in self.broker.middlewares:
+                if middleware.__class__.post_execute != TaskiqMiddleware.post_execute:
+                    await maybe_awaitable(middleware.post_execute(taskiq_msg, result))
+            try:
+                await self.broker.result_backend.set_result(message.task_id, result)
+            except Exception as exc:
+                logger.exception(
+                    "Can't set result in result backend. Cause: %s",
+                    exc,
+                    exc_info=True,
+                )
+                if raise_err:
+                    raise exc
 
-        for middleware in self.broker.middlewares:
-            if middleware.__class__.post_save != TaskiqMiddleware.post_save:
-                await maybe_awaitable(middleware.post_save(taskiq_msg, result))
+            for middleware in self.broker.middlewares:
+                if middleware.__class__.post_save != TaskiqMiddleware.post_save:
+                    await maybe_awaitable(middleware.post_save(taskiq_msg, result))
 
     async def run_task(  # noqa: C901, WPS210
         self,

taskiq/cli/worker/run.py

@@ -47,7 +47,7 @@ async def shutdown_broker(broker: AsyncBroker, timeout: float) -> None:
         )
 
 
-def start_listen(args: WorkerArgs) -> None:  # noqa: C901, WPS213
+def start_listen(args: WorkerArgs) -> None:  # noqa: WPS213
     """
     This function starts actual listening process.
 

taskiq/decor.py

@@ -60,15 +60,15 @@ def __call__(  # noqa: D102
         return self.original_func(*args, **kwargs)
 
     @overload
-    async def kiq(  # noqa: D102
+    async def kiq(
         self: "AsyncTaskiqDecoratedTask[_FuncParams, Coroutine[Any, Any, _T]]",
         *args: _FuncParams.args,
         **kwargs: _FuncParams.kwargs,
     ) -> AsyncTaskiqTask[_T]:
         ...
 
     @overload
-    async def kiq(  # noqa: D102
+    async def kiq(
         self: "AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]",
         *args: _FuncParams.args,
         **kwargs: _FuncParams.kwargs,