Merge pull request #6 from taskiq-python/feature/delays

Added delayed and prioritized messages.

Author: s3rius

.flake8

@@ -91,6 +91,11 @@ per-file-ignores =
   WPS609,
   ; Found protected attribute usage
   WPS437,
+  ; Found too many await expressions:
+  WPS217,
+  ; Found too many arguments
+  WPS211,
+
 
 exclude =
   ./.git,

README.md

@@ -14,6 +14,73 @@ async def test() -> None:
 
 ```
 
+## Non-obvious things
+
+You can send delayed messages and set priorities to messages using labels.
+
+## Delays
+
+To send delayed message, you have to specify
+delay label. You can do it with `task` decorator,
+or by using kicker. For example:
+
+```python
+broker = AioPikaBroker()
+
+@broker.task(delay=3)
+async def delayed_task() -> int:
+    return 1
+
+async def main():
+    await broker.startup()
+    # This message will be received by workers
+    # After 3 seconds delay.
+    await delayed_task.kiq()
+
+    # This message is going to be received after the delay in 4 seconds.
+    # Since we overriden the `delay` label using kicker.
+    await delayed_task.kicker().with_labels(delay=4).kiq()
+
+    # This message is going to be send immediately. Since we deleted the label.
+    await delayed_task.kicker().with_labels(delay=None).kiq()
+
+    # Of course the delay is managed by rabbitmq, so you don't
+    # have to wait delay period before message is going to be sent.
+```
+
+
+## Priorities
+
+You can define priorities for messages using `priority` label.
+Messages with higher priorities are delivered faster.
+But to use priorities you need to define `max_priority` of the main queue, by passing `max_priority` parameter in broker's init.
+This parameter sets maximum priority for the queue and
+declares it as the prority queue.
+
+Before doing so please read the [documentation](https://www.rabbitmq.com/priority.html#behaviour) about what
+downsides you get by using prioritized queues.
+
+
+```python
+broker = AioPikaBroker(max_priority=10)
+
+# We can define default priority for tasks.
+@broker.task(priority=2)
+async def prio_task() -> int:
+    return 1
+
+async def main():
+    await broker.startup()
+    # This message has priority = 2.
+    await prio_task.kiq()
+
+    # This message is going to have priority 4.
+    await prio_task.kicker().with_labels(priority=4).kiq()
+
+    # This message is going to have priority 0.
+    await prio_task.kicker().with_labels(priority=None).kiq()
+
+```
 
 ## Configuration
 
@@ -26,4 +93,9 @@ AioPikaBroker parameters:
 * `queue_name` - queue that used to get incoming messages.
 * `routing_key` - that used to bind that queue to the exchange.
 * `declare_exchange` - whether you want to declare new exchange if it doesn't exist.
+* `max_priority` - maximum priority for messages.
+* `delay_queue_name` - custom delay queue name.
+    This queue is used to deliver messages with delays.
+* `dead_letter_queue_name` - custom dead letter queue name.
+    This queue is used to receive negatively acknowleged messages from the main queue.
 * `qos` - number of messages that worker can prefetch.

taskiq_aio_pika/broker.py

@@ -1,8 +1,9 @@
 import asyncio
+from datetime import timedelta
 from logging import getLogger
-from typing import Any, Callable, Coroutine, Optional, TypeVar
+from typing import Any, Callable, Coroutine, Dict, Optional, TypeVar
 
-from aio_pika import ExchangeType, Message, connect_robust
+from aio_pika import DeliveryMode, ExchangeType, Message, connect_robust
 from aio_pika.abc import (
     AbstractChannel,
     AbstractIncomingMessage,
@@ -17,6 +18,26 @@
 logger = getLogger("taskiq.aio_pika_broker")
 
 
+def parse_val(
+    parse_func: Callable[[str], _T],
+    target: Optional[str] = None,
+) -> Optional[_T]:
+    """
+    Parse string to some value.
+
+    :param parse_func: function to use if value is present.
+    :param target: value to parse, defaults to None
+    :return: Optional value.
+    """
+    if target is None:
+        return None
+
+    try:
+        return parse_func(target)
+    except ValueError:
+        return None
+
+
 class AioPikaBroker(AsyncBroker):
     """Broker that works with RabbitMQ."""
 
@@ -29,9 +50,12 @@ def __init__(  # noqa: WPS211
         loop: Optional[asyncio.AbstractEventLoop] = None,
         exchange_name: str = "taskiq",
         queue_name: str = "taskiq",
+        dead_letter_queue_name: Optional[str] = None,
+        delay_queue_name: Optional[str] = None,
         declare_exchange: bool = True,
         routing_key: str = "#",
         exchange_type: ExchangeType = ExchangeType.TOPIC,
+        max_priority: Optional[int] = None,
         **connection_kwargs: Any,
     ) -> None:
         """
@@ -46,11 +70,16 @@ def __init__(  # noqa: WPS211
         :param loop: specific even loop.
         :param exchange_name: name of exchange that used to send messages.
         :param queue_name: queue that used to get incoming messages.
+        :param dead_letter_queue_name: custom name for dead-letter queue.
+            by default it set to {queue_name}.dead_letter.
+        :param delay_queue_name: custom name for queue that used to
+            deliver messages with delays.
         :param declare_exchange: whether you want to declare new exchange
             if it doesn't exist.
         :param routing_key: that used to bind that queue to the exchange.
         :param exchange_type: type of the exchange.
             Used only if `declare_exchange` is True.
+        :param max_priority: maximum priority value for messages.
         :param connection_kwargs: additional keyword arguments,
             for connect_robust method of aio-pika.
         """
@@ -65,6 +94,16 @@ def __init__(  # noqa: WPS211
         self._declare_exchange = declare_exchange
         self._queue_name = queue_name
         self._routing_key = routing_key
+        self._max_priority = max_priority
+
+        self._dead_letter_queue_name = f"{queue_name}.dead_letter"
+        if dead_letter_queue_name:
+            self._dead_letter_queue_name = dead_letter_queue_name
+
+        self._delay_queue_name = f"{queue_name}.delay"
+        if delay_queue_name:
+            self._delay_queue_name = delay_queue_name
+
         self.read_conn: Optional[AbstractRobustConnection] = None
         self.write_conn: Optional[AbstractRobustConnection] = None
         self.write_channel: Optional[AbstractChannel] = None
@@ -97,7 +136,26 @@ async def startup(self) -> None:  # noqa: WPS217
                 self._exchange_name,
                 ensure=False,
             )
-        queue = await self.write_channel.declare_queue(self._queue_name)
+        await self.write_channel.declare_queue(
+            self._dead_letter_queue_name,
+        )
+        args: "Dict[str, Any]" = {
+            "x-dead-letter-exchange": "",
+            "x-dead-letter-routing-key": self._dead_letter_queue_name,
+        }
+        if self._max_priority is not None:
+            args["x-max-priority"] = self._max_priority
+        queue = await self.write_channel.declare_queue(
+            self._queue_name,
+            arguments=args,
+        )
+        await self.write_channel.declare_queue(
+            self._delay_queue_name,
+            arguments={
+                "x-dead-letter-exchange": "",
+                "x-dead-letter-routing-key": self._queue_name,
+            },
+        )
         await queue.bind(exchange=exchange, routing_key=self._routing_key)
 
     async def kick(self, message: BrokerMessage) -> None:
@@ -111,25 +169,35 @@ async def kick(self, message: BrokerMessage) -> None:
         in headers. And message's routing key is the same
         as the task_name.
 
-
         :raises ValueError: if startup wasn't called.
         :param message: message to send.
         """
         if self.write_channel is None:
             raise ValueError("Please run startup before kicking.")
+        priority = parse_val(int, message.labels.get("priority"))
         rmq_msg = Message(
             body=message.message.encode(),
             headers={
                 "task_id": message.task_id,
                 "task_name": message.task_name,
                 **message.labels,
             },
+            delivery_mode=DeliveryMode.PERSISTENT,
+            priority=priority,
         )
-        exchange = await self.write_channel.get_exchange(
-            self._exchange_name,
-            ensure=False,
-        )
-        await exchange.publish(rmq_msg, routing_key=message.task_name)
+        delay = parse_val(int, message.labels.get("delay"))
+        if delay is None:
+            exchange = await self.write_channel.get_exchange(
+                self._exchange_name,
+                ensure=False,
+            )
+            await exchange.publish(rmq_msg, routing_key=message.task_name)
+        else:
+            rmq_msg.expiration = timedelta(seconds=delay)
+            await self.write_channel.default_exchange.publish(
+                rmq_msg,
+                routing_key=self._delay_queue_name,
+            )
 
     async def listen(
         self,
@@ -178,15 +246,18 @@ async def process_message(self, message: AbstractIncomingMessage) -> None:
         :param message: received message.
         """
         async with message.process():
+            headers = {}
+            for header_name, header_value in message.headers.items():
+                headers[header_name] = str(header_value)
             try:
                 broker_message = BrokerMessage(
-                    task_id=message.headers.pop("task_id"),
-                    task_name=message.headers.pop("task_name"),
+                    task_id=headers.pop("task_id"),
+                    task_name=headers.pop("task_name"),
                     message=message.body,
-                    labels=message.headers,
+                    labels=headers,
                 )
             except (ValueError, LookupError) as exc:
-                logger.debug(
+                logger.warning(
                     "Cannot read broker message %s",
                     exc,
                     exc_info=True,