Added simples retries logic.

s3rius · s3rius · commit 0a7c5956fef8 · 2022-08-08T12:38:48.000+04:00
Signed-off-by: Pavel Kirilin &lt;win10@list.ru&gt;
diff --git a/taskiq/__init__.py b/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",
 ]
diff --git a/taskiq/abc/broker.py b/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(
diff --git a/taskiq/abc/formatter.py b/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.
         """
 
diff --git a/taskiq/abc/middleware.py b/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,24 +62,23 @@ 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(
diff --git a/taskiq/cli/async_task_runner.py b/taskiq/cli/async_task_runner.py
@@ -311,7 +311,6 @@ 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
@@ -326,10 +325,7 @@ async def async_listen_messages(  # noqa: C901, WPS210, WPS213
             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:
diff --git a/taskiq/formatters/__init__.py b/taskiq/formatters/__init__.py
@@ -0,0 +1 @@
+"""Taskiq formatters."""
diff --git a/taskiq/formatters/json_formatter.py b/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",
             },
         )
diff --git a/taskiq/kicker.py b/taskiq/kicker.py
@@ -11,7 +11,6 @@
     Union,
     overload,
 )
-from uuid import uuid4
 
 from pydantic import BaseModel
 from typing_extensions import ParamSpec
@@ -37,14 +36,7 @@ def __init__(
         self,
         task_name: str,
         broker: "AsyncBroker",
-        labels: Dict[
-            str,
-            Union[
-                str,
-                int,
-                float,
-            ],
-        ],
+        labels: Dict[str, Any],
     ) -> None:
         self.task_name = task_name
         self.broker = broker
@@ -118,17 +110,17 @@ async def kiq(  # noqa: C901
         )
         message = self._prepare_message(*args, **kwargs)
         for middleware in self.broker.middlewares:
-            pre_send_res = middleware.pre_send(message, self.labels)
+            pre_send_res = middleware.pre_send(message)
             if isawaitable(pre_send_res):
                 message = await pre_send_res
             else:
                 message = pre_send_res  # type: ignore
         try:
-            await self.broker.kick(self.broker.formatter.dumps(message, self.labels))
+            await self.broker.kick(self.broker.formatter.dumps(message))
         except Exception as exc:
             raise SendTaskError() from exc
         for middleware in self.broker.middlewares:
-            post_send_res = middleware.post_send(message, self.labels)
+            post_send_res = middleware.post_send(message)
             if isawaitable(post_send_res):
                 await post_send_res
         return AsyncTaskiqTask(
@@ -167,17 +159,18 @@ def _prepare_message(  # noqa: WPS210
         """
         formatted_args = []
         formatted_kwargs = {}
+        labels = {}
         for arg in args:
             formatted_args.append(self._prepare_arg(arg))
         for kwarg_name, kwarg_val in kwargs.items():
             formatted_kwargs[kwarg_name] = self._prepare_arg(kwarg_val)
-
-        task_id = uuid4().hex
+        for label, label_val in self.labels.items():
+            labels[label] = str(label_val)
 
         return TaskiqMessage(
-            task_id=task_id,
+            task_id=self.broker.id_generator,
             task_name=self.task_name,
-            meta=self.labels,
+            labels=labels,
             args=formatted_args,
             kwargs=formatted_kwargs,
         )
diff --git a/taskiq/message.py b/taskiq/message.py
@@ -14,7 +14,7 @@ class TaskiqMessage(BaseModel):
 
     task_id: str
     task_name: str
-    meta: Dict[str, str]
+    labels: Dict[str, str]
     args: List[Any]
     kwargs: Dict[str, Any]
 
@@ -25,4 +25,4 @@ class BrokerMessage(BaseModel):
     task_id: str
     task_name: str
     message: str
-    headers: Dict[str, Any]
+    labels: Dict[str, str]
diff --git a/taskiq/middlewares/__init__.py b/taskiq/middlewares/__init__.py
@@ -0,0 +1 @@
+"""Taskiq middlewares."""
diff --git a/taskiq/middlewares/retry_middleware.py b/taskiq/middlewares/retry_middleware.py
@@ -0,0 +1,53 @@
+from logging import getLogger
+from typing import Any
+
+from taskiq.abc.middleware import TaskiqMiddleware
+from taskiq.message import TaskiqMessage
+from taskiq.result import TaskiqResult
+
+logger = getLogger("taskiq.retry_middleware")
+
+
+class SimpleRetryMiddleware(TaskiqMiddleware):
+    """Middleware to add retries."""
+
+    def __init__(
+        self,
+        default_retry_count: int = 3,
+    ) -> None:
+        self.default_retry_count = default_retry_count
+
+    async def on_error(
+        self,
+        message: "TaskiqMessage",
+        result: "TaskiqResult[Any]",
+        exception: Exception,
+    ) -> None:
+        """
+        Retry on error.
+
+        This middleware is used to retry
+        tasks on errors.
+
+        If error is found during the execution
+        this function is invoked.
+
+        :param message: Message that caused the error.
+        :param result: execution result.
+        :param exception: found exception.
+        """
+        retry_on_error = message.labels.get("retry_on_error")
+        # Check if retrying is enabled for the task.
+        if retry_on_error != "True":
+            return
+
+        logger.info(f"Task '{message.task_name}' invocation failed. Retrying.")
+
+        # Getting number of previous retries.
+        retries = int(message.labels.get("_retries", 0)) + 1
+        message.labels["_retries"] = str(retries)
+        max_retries = int(message.labels.get("max_retries", self.default_retry_count))
+        if retries < max_retries:
+            message.task_id = self.broker.id_generator()
+            broker_message = self.broker.formatter.dumps(message=message)
+            await self.broker.kick(broker_message)
diff --git a/taskiq/plugins/__init__.py b/taskiq/plugins/__init__.py
