Added acknowledgable messages. (#144)

Author: s3rius

docs/examples/extending/broker.py

@@ -1,6 +1,6 @@
-from typing import AsyncGenerator
+from typing import AsyncGenerator, Union
 
-from taskiq import AsyncBroker, BrokerMessage
+from taskiq import AckableMessage, AsyncBroker, BrokerMessage
 
 
 class MyBroker(AsyncBroker):
@@ -23,7 +23,7 @@ async def kick(self, message: BrokerMessage) -> None:
         # Send a message.message.
         pass
 
-    async def listen(self) -> AsyncGenerator[bytes, None]:
+    async def listen(self) -> AsyncGenerator[Union[bytes, AckableMessage], None]:
         while True:
             # Get new message.
             new_message: bytes = ...  # type: ignore

docs/extending-taskiq/broker.md

@@ -23,7 +23,28 @@ As a broker developer, please send only raw bytes from the `message` field of a
 :::
 
 
-The `listen` method should yield raw bytes that were sent over the network.
+## Acknowledgement
+
+The `listen` method should yield raw bytes of a message.
+But if your broker supports acking or rejecting messages, the broker should return `taskiq.AckableMessage`
+with required fields.
+
+For example:
+
+```python
+
+async def listen(self) -> AsyncGenerator[AckableMessage, None]:
+   for message in self.my_channel:
+      yield AckableMessage(
+         data=message.bytes,
+         # Ack is a function that takes no parameters.
+         # So you either set here method of a message,
+         # or you can make a closure.
+         ack=message.ack
+         # Can be set to None if broker doesn't support it.
+         reject=message.reject
+      )
+```
 
 ## Conventions
 

taskiq/__init__.py

@@ -6,12 +6,22 @@
 from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.abc.result_backend import AsyncResultBackend
 from taskiq.abc.schedule_source import ScheduleSource
+from taskiq.acks import AckableMessage
 from taskiq.brokers.inmemory_broker import InMemoryBroker
 from taskiq.brokers.shared_broker import async_shared_broker
 from taskiq.brokers.zmq_broker import ZeroMQBroker
 from taskiq.context import Context
 from taskiq.events import TaskiqEvents
-from taskiq.exceptions import TaskiqError
+from taskiq.exceptions import (
+    NoResultError,
+    RejectError,
+    ResultGetError,
+    ResultIsReadyError,
+    SecurityError,
+    SendTaskError,
+    TaskiqError,
+    TaskiqResultTimeoutError,
+)
 from taskiq.funcs import gather
 from taskiq.message import BrokerMessage, TaskiqMessage
 from taskiq.middlewares.prometheus_middleware import PrometheusMiddleware
@@ -26,23 +36,31 @@
     "Context",
     "AsyncBroker",
     "TaskiqError",
+    "RejectError",
     "TaskiqState",
     "TaskiqResult",
     "ZeroMQBroker",
     "TaskiqEvents",
+    "SecurityError",
     "TaskiqMessage",
     "BrokerMessage",
+    "ResultGetError",
     "ScheduledTask",
     "TaskiqDepends",
+    "NoResultError",
+    "SendTaskError",
+    "AckableMessage",
     "InMemoryBroker",
     "ScheduleSource",
     "TaskiqScheduler",
     "TaskiqFormatter",
     "AsyncTaskiqTask",
     "TaskiqMiddleware",
+    "ResultIsReadyError",
     "AsyncResultBackend",
     "async_shared_broker",
-    "AsyncTaskiqDecoratedTask",
-    "SimpleRetryMiddleware",
     "PrometheusMiddleware",
+    "SimpleRetryMiddleware",
+    "AsyncTaskiqDecoratedTask",
+    "TaskiqResultTimeoutError",
 ]

taskiq/abc/broker.py

@@ -24,6 +24,7 @@
 from typing_extensions import ParamSpec, Self, TypeAlias
 
 from taskiq.abc.middleware import TaskiqMiddleware
+from taskiq.acks import AckableMessage
 from taskiq.decor import AsyncTaskiqDecoratedTask
 from taskiq.events import TaskiqEvents
 from taskiq.formatters.json_formatter import JSONFormatter
@@ -185,13 +186,19 @@ async def kick(
         """
 
     @abstractmethod
-    def listen(self) -> AsyncGenerator[bytes, None]:
+    def listen(self) -> AsyncGenerator[Union[bytes, AckableMessage], None]:
         """
         This function listens to new messages and yields them.
 
         This it the main point for workers.
         This function is used to get new tasks from the network.
 
+        If your broker support acknowledgement, then you
+        should wrap your message in AckableMessage dataclass.
+
+        If your messages was wrapped in AckableMessage dataclass,
+        taskiq will call ack when finish processing message.
+
         :yield: incoming messages.
         :return: nothing.
         """

taskiq/acks.py

@@ -0,0 +1,21 @@
+import dataclasses
+from typing import Awaitable, Callable, Optional, Union
+
+
+@dataclasses.dataclass
+class AckableMessage:
+    """
+    Message that can be acknowledged.
+
+    If your broker support message acknowledgement,
+    please return this type of message, so we'll be
+    able to mark this message as acknowledged after
+    the function will be executed.
+
+    It adds more reliability to brokers and system
+    as a whole.
+    """
+
+    data: bytes
+    ack: Callable[[], Union[None, Awaitable[None]]]
+    reject: Optional[Callable[[], Union[None, Awaitable[None]]]] = None

taskiq/exceptions.py

@@ -36,3 +36,7 @@ class SecurityError(TaskiqError):
 
 class NoResultError(TaskiqError):
     """Error if user does not want to set result."""
+
+
+class RejectError(TaskiqError):
+    """Error is thrown if message should be rejected."""

taskiq/receiver/receiver.py

@@ -3,15 +3,15 @@
 from concurrent.futures import Executor
 from logging import getLogger
 from time import time
-from typing import Any, Callable, Dict, Optional, Set, get_type_hints
+from typing import Any, Callable, Dict, Optional, Set, Union, get_type_hints
 
 import anyio
 from taskiq_dependencies import DependencyGraph
 
-from taskiq.abc.broker import AsyncBroker
+from taskiq.abc.broker import AckableMessage, AsyncBroker
 from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.context import Context
-from taskiq.exceptions import NoResultError
+from taskiq.exceptions import NoResultError, RejectError
 from taskiq.message import TaskiqMessage
 from taskiq.receiver.params_parser import parse_params
 from taskiq.result import TaskiqResult
@@ -69,9 +69,9 @@ def __init__(  # noqa: WPS211
             )
         self.sem_prefetch = asyncio.Semaphore(max_prefetch)
 
-    async def callback(  # noqa: C901, WPS213
+    async def callback(  # noqa: C901, WPS213, WPS217
         self,
-        message: bytes,
+        message: Union[bytes, AckableMessage],
         raise_err: bool = False,
     ) -> None:
         """
@@ -86,12 +86,16 @@ async def callback(  # noqa: C901, WPS213
         :param raise_err: raise an error if cannot save result in
             result_backend.
         """
+        if isinstance(message, AckableMessage):
+            message_data = message.data
+        else:
+            message_data = message
         try:
-            taskiq_msg = self.broker.formatter.loads(message=message)
+            taskiq_msg = self.broker.formatter.loads(message=message_data)
         except Exception as exc:
             logger.warning(
                 "Cannot parse message: %s. Skipping execution.\n %s",
-                message,
+                message_data,
                 exc,
                 exc_info=True,
             )
@@ -124,9 +128,20 @@ async def callback(  # noqa: C901, WPS213
             target=self.broker.available_tasks[taskiq_msg.task_name].original_func,
             message=taskiq_msg,
         )
+
+        # If broker has an ability to ack or reject messages.
+        if isinstance(message, AckableMessage):
+            # If we received an error for negative acknowledgement.
+            if message.reject is not None and isinstance(result.error, RejectError):
+                await maybe_awaitable(message.reject())
+            # Otherwise we positively acknowledge the message.
+            else:
+                await maybe_awaitable(message.ack())
+
         for middleware in self.broker.middlewares:
             if middleware.__class__.post_execute != TaskiqMiddleware.post_execute:
                 await maybe_awaitable(middleware.post_execute(taskiq_msg, result))
+
         try:
             if not isinstance(result.error, NoResultError):
                 await self.broker.result_backend.set_result(taskiq_msg.task_id, result)
@@ -255,13 +270,16 @@ async def listen(self) -> None:  # pragma: no cover
         """
         await self.broker.startup()
         logger.info("Listening started.")
-        queue: asyncio.Queue[bytes] = asyncio.Queue()
+        queue: "asyncio.Queue[Union[bytes, AckableMessage]]" = asyncio.Queue()
 
         async with anyio.create_task_group() as gr:
             gr.start_soon(self.prefetcher, queue)
             gr.start_soon(self.runner, queue)
 
-    async def prefetcher(self, queue: "asyncio.Queue[Any]") -> None:
+    async def prefetcher(
+        self,
+        queue: "asyncio.Queue[Union[bytes, AckableMessage]]",
+    ) -> None:
         """
         Prefetch tasks data.
 
@@ -280,7 +298,10 @@ async def prefetcher(self, queue: "asyncio.Queue[Any]") -> None:
 
         await queue.put(QUEUE_DONE)
 
-    async def runner(self, queue: "asyncio.Queue[bytes]") -> None:
+    async def runner(
+        self,
+        queue: "asyncio.Queue[Union[bytes, AckableMessage]]",
+    ) -> None:
         """
         Run tasks.
 

taskiq/utils.py

@@ -1,11 +1,11 @@
 import inspect
-from typing import Any, Coroutine, TypeVar, Union
+from typing import Any, Awaitable, Coroutine, TypeVar, Union
 
 _T = TypeVar("_T")  # noqa: WPS111
 
 
 async def maybe_awaitable(
-    possible_coroutine: "Union[_T, Coroutine[Any, Any, _T]]",
+    possible_coroutine: "Union[_T, Coroutine[Any, Any, _T], Awaitable[_T]]",
 ) -> _T:
     """
     Awaits coroutine if needed.