Added broker for shared tasks.

Signed-off-by: Pavel Kirilin <[email protected]>

Author: s3rius

taskiq/abc/broker.py

@@ -31,6 +31,7 @@ def __init__(
             result_backend = DummyResultBackend()
         self.result_backend = result_backend
         self.is_worker_process = False
+        self.decorator_class = AsyncTaskiqDecoratedTask
 
     async def startup(self) -> None:
         """Do something when starting broker."""
@@ -133,7 +134,7 @@ def inner(
                 wrapper = wraps(func)
 
                 decorated_task = wrapper(
-                    AsyncTaskiqDecoratedTask(
+                    self.decorator_class(
                         broker=self,
                         original_func=func,
                         labels=inner_labels,

taskiq/brokers/shared_broker.py

@@ -0,0 +1,76 @@
+from typing import AsyncGenerator, Optional, TypeVar
+
+from taskiq.abc.broker import AsyncBroker
+from taskiq.decor import AsyncTaskiqDecoratedTask
+from taskiq.exceptions import TaskiqError
+from taskiq.kicker import AsyncKicker
+from taskiq.message import TaskiqMessage
+from taskiq.types_helpers import ReturnType_
+
+Params_ = TypeVar("Params_")  # noqa: WPS120
+
+
+class SharedDecoratedTask(AsyncTaskiqDecoratedTask[Params_, ReturnType_]):
+    """Decorator that is used with shared broker."""
+
+    def kicker(self) -> AsyncKicker[Params_, ReturnType_]:
+        """
+        This method updates getting default kicker.
+
+        In this method we want to get default broker from
+        our shared broker and send task to it, instead
+        of shared_broker.
+
+        :raises TaskiqError: if _default_broker is not set.
+        :return: new kicker.
+        """
+        broker = getattr(self.broker, "_default_broker", None)
+        if broker is None:
+            raise TaskiqError(
+                "You cannot use kiq directly on shared task "
+                "without setting the default_broker.",
+            )
+        return AsyncKicker(
+            task_name=self.task_name,
+            broker=broker,
+            labels=self.labels,
+        )
+
+
+class AsyncSharedBroker(AsyncBroker):
+    """Broker for creating shared tasks."""
+
+    def __init__(self) -> None:
+        super().__init__(None)
+        self._default_broker: Optional[AsyncBroker] = None
+        self.decorator_class = SharedDecoratedTask
+
+    async def kick(self, message: TaskiqMessage) -> None:
+        """
+        Shared broker cannot kick tasks.
+
+        :param message: message to send.
+        :raises TaskiqError: if called.
+        """
+        raise TaskiqError("Shared broker cannot kick tasks.")
+
+    def default_broker(self, new_broker: AsyncBroker) -> None:
+        """
+        Updates default broker.
+
+        :param new_broker: new async broker to kick tasks with.
+        """
+        self._default_broker = new_broker
+
+    async def listen(self) -> AsyncGenerator[TaskiqMessage, None]:  # type: ignore
+        """
+        Shared broker cannot listen to tasks.
+
+        This method will throw an exception.
+
+        :raises TaskiqError: if called.
+        """
+        raise TaskiqError("Shared broker cannot listen")
+
+
+async_shared_broker = AsyncSharedBroker()

taskiq/cli/async_task_runner.py

@@ -64,7 +64,7 @@ def parse_params(  # noqa: C901
         argnum += 1
         annot = params_type.annotation
         value = None
-        logger.debug("Trying to parse %s as %s" % (param_name, params_type.annotation))
+        logger.debug("Trying to parse %s as %s", param_name, params_type.annotation)
         if argnum >= len(message.args):
             value = message.kwargs.get(param_name)
             if value is None:
@@ -146,7 +146,8 @@ async def run_task(  # noqa: WPS210
         except Exception as exc:
             is_err = True
             logger.error(
-                "Exception found while executing function: %s" % exc,
+                "Exception found while executing function: %s",
+                exc,
                 exc_info=True,
             )
         execution_time = time() - start_time
@@ -180,7 +181,7 @@ def exit_process(task: asyncio.Task[Any]) -> NoReturn:
     try:
         result = task.result()
         if result is not None:
-            logger.info("Broker returned value on shutdown: '%s'" % str(result))
+            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)

taskiq/cli/worker.py

@@ -133,15 +133,15 @@ def watch_workers_restarts(args: TaskiqArgs) -> None:
             if worker.is_alive():
                 continue
             if worker.exitcode is not None and worker.exitcode > 0 and restart_workers:
-                logger.info("Trying to restart the worker-%s" % worker_id)
+                logger.info("Trying to restart the worker-%s", worker_id)
                 worker_processes[worker_id] = Process(
                     target=start_listen,
                     kwargs={"args": args},
                     name=f"worker-{worker_id}",
                 )
                 worker_processes[worker_id].start()
             else:
-                logger.info("Worker-%s has finished." % worker_id)
+                logger.info("Worker-%s has finished.", worker_id)
                 worker.join()
                 process_to_remove.append(worker)
 
@@ -162,7 +162,7 @@ def run_worker(args: TaskiqArgs) -> None:
         level=getLevelName(args.log_level),
         format=("[%(asctime)s][%(levelname)-7s][%(processName)s] %(message)s"),
     )
-    logger.info("Starting %s worker processes." % args.workers)
+    logger.info("Starting %s worker processes.", args.workers)
 
     global worker_processes  # noqa: WPS420
 
@@ -174,11 +174,9 @@ def run_worker(args: TaskiqArgs) -> None:
         )
         work_proc.start()
         logger.debug(
-            "Started process worker-%d with pid %s "
-            % (
-                process,
-                work_proc.pid,
-            ),
+            "Started process worker-%d with pid %s ",
+            process,
+            work_proc.pid,
         )
         worker_processes.append(work_proc)
 

taskiq/result_backends/dummy.py

@@ -1,6 +1,7 @@
 from typing import Any, TypeVar
 
 from taskiq.abc.result_backend import AsyncResultBackend
+from taskiq.result import TaskiqResult
 
 _ReturnType = TypeVar("_ReturnType")
 
@@ -20,23 +21,31 @@ async def set_result(self, task_id: str, result: Any) -> None:
         """
         return await super().set_result(task_id, result)
 
-    async def is_result_ready(self, _id: str) -> bool:
+    async def is_result_ready(self, task_id: str) -> bool:
         """
         Checks whether task is completed.
 
         Result is always ready,
         since it doesn't care about tasks.
 
-        :param _id: task's id.
+        :param task_id: task's id.
         :return: true.
         """
         return True
 
-    def get_result(self, _id: str, _with_logs: bool = False) -> Any:
+    async def get_result(self, task_id: str, with_logs: bool = False) -> Any:
         """
         Returns None.
 
         This result doesn't care about passed parameters.
 
-        :param _id: task's id.
+        :param task_id: task's id.
+        :param with_logs: wether to fetch logs.
+        :returns: TaskiqResult.
         """
+        return TaskiqResult(
+            is_err=False,
+            log=None,
+            return_value=None,
+            execution_time=0,
+        )