Merge pull request #23 from taskiq-python/feature/retries

Feature/retries

Author: s3rius

taskiq/__init__.py

@@ -1,13 +1,24 @@
 """Distributed task manager."""
 from taskiq.abc.broker import AsyncBroker, AsyncTaskiqDecoratedTask
+from taskiq.abc.formatter import TaskiqFormatter
+from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.abc.result_backend import AsyncResultBackend
-from taskiq.message import TaskiqMessage
+from taskiq.brokers.shared_broker import async_shared_broker
+from taskiq.exceptions import TaskiqError
+from taskiq.message import BrokerMessage, TaskiqMessage
+from taskiq.result import TaskiqResult
 from taskiq.task import AsyncTaskiqTask
 
 __all__ = [
-    "TaskiqMessage",
     "AsyncBroker",
-    "AsyncTaskiqDecoratedTask",
-    "AsyncResultBackend",
+    "TaskiqError",
+    "TaskiqResult",
+    "TaskiqMessage",
+    "BrokerMessage",
+    "TaskiqFormatter",
     "AsyncTaskiqTask",
+    "TaskiqMiddleware",
+    "AsyncResultBackend",
+    "async_shared_broker",
+    "AsyncTaskiqDecoratedTask",
 ]

taskiq/abc/broker.py

@@ -1,3 +1,4 @@
+import inspect
 from abc import ABC, abstractmethod
 from functools import wraps
 from logging import getLogger
@@ -13,12 +14,13 @@
     Union,
     overload,
 )
+from uuid import uuid4
 
 from typing_extensions import ParamSpec
 
 from taskiq.decor import AsyncTaskiqDecoratedTask
+from taskiq.formatters.json_formatter import JSONFormatter
 from taskiq.message import BrokerMessage
-from taskiq.plugins.json_formatter import JSONFormatter
 from taskiq.result_backends.dummy import DummyResultBackend
 
 if TYPE_CHECKING:
@@ -33,6 +35,18 @@
 logger = getLogger("taskiq")
 
 
+def default_id_generator() -> str:
+    """
+    Default task_id generator.
+
+    This function is used to generate id's
+    for tasks.
+
+    :return: new task_id.
+    """
+    return uuid4().hex
+
+
 class AsyncBroker(ABC):
     """
     Async broker.
@@ -47,14 +61,18 @@ class AsyncBroker(ABC):
     def __init__(
         self,
         result_backend: "Optional[AsyncResultBackend[_T]]" = None,
+        task_id_generator: Optional[Callable[[], str]] = None,
     ) -> None:
         if result_backend is None:
             result_backend = DummyResultBackend()
+        if task_id_generator is None:
+            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
 
     def add_middlewares(self, middlewares: "List[TaskiqMiddleware]") -> None:
         """
@@ -79,6 +97,10 @@ async def shutdown(self) -> None:
         This method is called,
         when broker is closig.
         """
+        for middleware in self.middlewares:
+            middleware_shutdown = middleware.shutdown()
+            if inspect.isawaitable(middleware_shutdown):
+                await middleware_shutdown
 
     @abstractmethod
     async def kick(

taskiq/abc/formatter.py

@@ -1,5 +1,4 @@
 from abc import ABC, abstractmethod
-from typing import Any, Dict
 
 from taskiq.message import BrokerMessage, TaskiqMessage
 
@@ -8,12 +7,11 @@ class TaskiqFormatter(ABC):
     """Custom formatter for brokers."""
 
     @abstractmethod
-    def dumps(self, message: TaskiqMessage, labels: Dict[str, Any]) -> BrokerMessage:
+    def dumps(self, message: TaskiqMessage) -> BrokerMessage:
         """
         Dump message to broker message instance.
 
         :param message: message to send.
-        :param labels: task's labels.
         :return: message for brokers.
         """
 

taskiq/abc/middleware.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Coroutine, Dict, Union
+from typing import TYPE_CHECKING, Any, Coroutine, Union
 
 if TYPE_CHECKING:
     from taskiq.abc.broker import AsyncBroker
@@ -20,10 +20,12 @@ def set_broker(self, broker: "AsyncBroker") -> None:
         """
         self.broker = broker
 
+    def shutdown(self) -> Union[None, Coroutine[Any, Any, None]]:
+        """This function is used to do some work on shutdown."""
+
     def pre_send(
         self,
         message: "TaskiqMessage",
-        labels: Dict[str, Any],
     ) -> "Union[TaskiqMessage, Coroutine[Any, Any, TaskiqMessage]]":
         """
         Hook that executes before sending the task to worker.
@@ -32,15 +34,13 @@ def pre_send(
         the message is sent to broker.
 
         :param message: message to send.
-        :param labels: task's labels.
         :return: modified message.
         """
         return message
 
     def post_send(
         self,
         message: "TaskiqMessage",
-        labels: Dict[str, Any],
     ) -> "Union[None, Coroutine[Any, Any, None]]":
         """
         This hook is executed right after the task is sent.
@@ -49,13 +49,11 @@ def post_send(
         after the messages is kicked in broker.
 
         :param message: kicked message.
-        :param labels: labels for a message.
         """
 
     def pre_execute(
         self,
         message: "TaskiqMessage",
-        labels: Dict[str, Any],
     ) -> "Union[TaskiqMessage, Coroutine[Any, Any, TaskiqMessage]]":
         """
         This hook is called before executing task.
@@ -64,22 +62,35 @@ def pre_execute(
         executes in the worker process.
 
         :param message: incoming parsed taskiq message.
-        :param labels: task's labels without user-supplied lables.
         :return: modified message.
         """
         return message
 
     def post_execute(
         self,
+        message: "TaskiqMessage",
         result: "TaskiqResult[Any]",
-        labels: Dict[str, Any],
     ) -> "Union[None, Coroutine[Any, Any, None]]":
         """
         This hook executes after task is complete.
 
         This is a worker-side hook. It's called
         in worker process.
 
+        :param message: incoming message.
         :param result: result of execution for current task.
-        :param labels: task's labels. Without user-supplied labels.
+        """
+
+    def on_error(
+        self,
+        message: "TaskiqMessage",
+        result: "TaskiqResult[Any]",
+        exception: Exception,
+    ) -> "Union[None, Coroutine[Any, Any, None]]":
+        """
+        This function is called when exception is found.
+
+        :param message: incoming message.
+        :param result: returned value.
+        :param exception: found exception.
         """

taskiq/brokers/inmemory_broker.py

@@ -125,6 +125,7 @@ async def kick(self, message: BrokerMessage) -> None:
             message=taskiq_message,
             log_collector_format=self.logs_format,
             executor=self.executor,
+            middlewares=self.middlewares,
         )
         await self.result_backend.set_result(message.task_id, result)
 

taskiq/cli/async_task_runner.py

@@ -6,11 +6,12 @@
 from concurrent.futures import Executor, ThreadPoolExecutor
 from logging import getLogger
 from time import time
-from typing import Any, Callable, Dict, NoReturn, Optional
+from typing import Any, Callable, Dict, List, NoReturn, Optional
 
 from pydantic import parse_obj_as
 
 from taskiq.abc.broker import AsyncBroker
+from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.cli.args import TaskiqArgs
 from taskiq.cli.log_collector import log_collector
 from taskiq.message import TaskiqMessage
@@ -97,12 +98,13 @@ def run_sync(target: Callable[..., Any], message: TaskiqMessage) -> Any:
     return target(*message.args, **message.kwargs)
 
 
-async def run_task(  # noqa: WPS210
+async def run_task(  # noqa: C901, WPS210, WPS211
     target: Callable[..., Any],
     signature: Optional[inspect.Signature],
     message: TaskiqMessage,
     log_collector_format: str,
     executor: Optional[Executor] = None,
+    middlewares: Optional[List[TaskiqMiddleware]] = None,
 ) -> TaskiqResult[Any]:
     """
     This function actually executes functions.
@@ -123,12 +125,16 @@ async def run_task(  # noqa: WPS210
     :param message: received message.
     :param log_collector_format: Log format in wich logs are collected.
     :param executor: executor to run sync tasks.
+    :param middlewares: list of broker's middlewares in case of errors.
     :return: result of execution.
     """
+    if middlewares is None:
+        middlewares = []
+
     loop = asyncio.get_running_loop()
     logs = io.StringIO()
-    is_err = False
     returned = None
+    found_exception = None
     # Captures function's logs.
     parse_params(signature, message)
     with log_collector(logs, log_collector_format):
@@ -144,7 +150,7 @@ async def run_task(  # noqa: WPS210
                     message,
                 )
         except Exception as exc:
-            is_err = True
+            found_exception = exc
             logger.error(
                 "Exception found while executing function: %s",
                 exc,
@@ -154,12 +160,23 @@ async def run_task(  # noqa: WPS210
 
     raw_logs = logs.getvalue()
     logs.close()
-    return TaskiqResult(
-        is_err=is_err,
+    result: "TaskiqResult[Any]" = TaskiqResult(
+        is_err=found_exception is not None,
         log=raw_logs,
         return_value=returned,
         execution_time=execution_time,
     )
+    if found_exception is not None:
+        for middleware in middlewares:
+            err_handler = middleware.on_error(
+                message,
+                result,
+                found_exception,
+            )
+            if inspect.isawaitable(err_handler):
+                await err_handler
+
+    return result
 
 
 def exit_process(task: "asyncio.Task[Any]") -> NoReturn:
@@ -294,24 +311,21 @@ async def async_listen_messages(  # noqa: C901, WPS210, WPS213
         for middleware in broker.middlewares:
             pre_ex_res = middleware.pre_execute(
                 taskiq_msg,
-                broker.available_tasks[message.task_name].labels,
             )
             if inspect.isawaitable(pre_ex_res):
                 taskiq_msg = await pre_ex_res
             else:
                 taskiq_msg = pre_ex_res  # type: ignore
         result = await run_task(
-            broker.available_tasks[message.task_name].original_func,
-            task_signatures.get(message.task_name),
-            taskiq_msg,
-            cli_args.log_collector_format,
-            executor,
+            target=broker.available_tasks[message.task_name].original_func,
+            signature=task_signatures.get(message.task_name),
+            message=taskiq_msg,
+            log_collector_format=cli_args.log_collector_format,
+            executor=executor,
+            middlewares=broker.middlewares,
         )
         for middleware in broker.middlewares:
-            post_ex_res = middleware.post_execute(
-                result,
-                broker.available_tasks[message.task_name].labels,
-            )
+            post_ex_res = middleware.post_execute(taskiq_msg, result)
             if inspect.isawaitable(post_ex_res):
                 await post_ex_res
         try:

taskiq/formatters/__init__.py

@@ -0,0 +1 @@
+"""Taskiq formatters."""

taskiq/plugins/json_formatter.py renamed to taskiq/formatters/json_formatter.py

@@ -1,25 +1,22 @@
-from typing import Any, Dict
-
 from taskiq.abc.formatter import TaskiqFormatter
 from taskiq.message import BrokerMessage, TaskiqMessage
 
 
 class JSONFormatter(TaskiqFormatter):
     """Default taskiq formatter."""
 
-    def dumps(self, message: TaskiqMessage, labels: Dict[str, Any]) -> BrokerMessage:
+    def dumps(self, message: TaskiqMessage) -> BrokerMessage:
         """
         Dumps taskiq message to some broker message format.
 
         :param message: message to send.
-        :param labels: message's labels.
         :return: Dumped message.
         """
         return BrokerMessage(
             task_id=message.task_id,
             task_name=message.task_name,
             message=message.json(),
-            headers={
+            labels={
                 "content_type": "application/json",
             },
         )