Merge pull request #15 from taskiq-python/issue-14

Author: s3rius

.github/workflows/test.yml

@@ -29,7 +29,7 @@ jobs:
   pytest:
     services:
       rabbit:
-        image: rabbitmq:3.9.16-alpine
+        image: heidiks/rabbitmq-delayed-message-exchange:latest
         env:
           RABBITMQ_DEFAULT_USER: "guest"
           RABBITMQ_DEFAULT_PASS: "guest"

README.md

@@ -20,9 +20,13 @@ You can send delayed messages and set priorities to messages using labels.
 
 ## Delays
 
+### **Default retries**
+
 To send delayed message, you have to specify
 delay label. You can do it with `task` decorator,
-or by using kicker. For example:
+or by using kicker.  
+In this type of delay we are using additional queue with `expiration` parameter and after with time message will be deleted from `delay` queue and sent to the main taskiq queue.
+For example:
 
 ```python
 broker = AioPikaBroker()
@@ -48,6 +52,36 @@ async def main():
     # have to wait delay period before message is going to be sent.
 ```
 
+### **Retries with `rabbitmq-delayed-message-exchange` plugin**
+
+To send delayed message you can install `rabbitmq-delayed-message-exchange`
+plugin https://github.com/rabbitmq/rabbitmq-delayed-message-exchange.
+
+And you need to configure you broker.
+There is `delayed_message_exchange_plugin` `AioPikaBroker` parameter and it must be `True` to turn on delayed message functionality.  
+
+The delay plugin can handle tasks with different delay times well, and the delay based on dead letter queue is suitable for tasks with the same delay time.  
+For example:
+
+```python
+broker = AioPikaBroker(
+    delayed_message_exchange_plugin=True,
+)
+
+@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()
+```
 
 ## Priorities
 

poetry.lock

@@ -53,6 +53,7 @@ def __init__(  # noqa: WPS211
         routing_key: str = "#",
         exchange_type: ExchangeType = ExchangeType.TOPIC,
         max_priority: Optional[int] = None,
+        delayed_message_exchange_plugin: bool = False,
         **connection_kwargs: Any,
     ) -> None:
         """
@@ -79,6 +80,8 @@ def __init__(  # noqa: WPS211
         :param exchange_type: type of the exchange.
             Used only if `declare_exchange` is True.
         :param max_priority: maximum priority value for messages.
+        :param delayed_message_exchange_plugin: turn on or disable
+            delayed-message-exchange rabbitmq plugin.
         :param connection_kwargs: additional keyword arguments,
             for connect_robust method of aio-pika.
         """
@@ -95,6 +98,7 @@ def __init__(  # noqa: WPS211
         self._queue_name = queue_name
         self._routing_key = routing_key
         self._max_priority = max_priority
+        self._delayed_message_exchange_plugin = delayed_message_exchange_plugin
 
         self._dead_letter_queue_name = f"{queue_name}.dead_letter"
         if dead_letter_queue_name:
@@ -104,6 +108,8 @@ def __init__(  # noqa: WPS211
         if delay_queue_name:
             self._delay_queue_name = delay_queue_name
 
+        self._delay_plugin_exchange_name = f"{exchange_name}.plugin_delay"
+
         self.read_conn: Optional[AbstractRobustConnection] = None
         self.write_conn: Optional[AbstractRobustConnection] = None
         self.write_channel: Optional[AbstractChannel] = None
@@ -132,9 +138,31 @@ async def startup(self) -> None:  # noqa: WPS217
                 self._exchange_name,
                 type=self._exchange_type,
             )
+
+        if self._delayed_message_exchange_plugin:
+            await self.write_channel.declare_exchange(
+                self._delay_plugin_exchange_name,
+                type=ExchangeType.X_DELAYED_MESSAGE,
+                arguments={
+                    "x-delayed-type": "direct",
+                },
+            )
+
         if self._declare_queues:
             await self.declare_queues(self.write_channel)
 
+    async def shutdown(self) -> None:
+        """Close all connections on shutdown."""
+        await super().shutdown()
+        if self.write_channel:
+            await self.write_channel.close()
+        if self.read_channel:
+            await self.read_channel.close()
+        if self.write_conn:
+            await self.write_conn.close()
+        if self.read_conn:
+            await self.read_conn.close()
+
     async def declare_queues(
         self,
         channel: AbstractChannel,
@@ -163,14 +191,24 @@ async def declare_queues(
             self._queue_name,
             arguments=args,
         )
-        await channel.declare_queue(
-            self._delay_queue_name,
-            arguments={
-                "x-dead-letter-exchange": "",
-                "x-dead-letter-routing-key": self._queue_name,
-            },
+        if self._delayed_message_exchange_plugin:
+            await queue.bind(
+                exchange=self._delay_plugin_exchange_name,
+                routing_key=self._routing_key,
+            )
+        else:
+            await channel.declare_queue(
+                self._delay_queue_name,
+                arguments={
+                    "x-dead-letter-exchange": "",
+                    "x-dead-letter-routing-key": self._queue_name,
+                },
+            )
+
+        await queue.bind(
+            exchange=self._exchange_name,
+            routing_key=self._routing_key,
         )
-        await queue.bind(exchange=self._exchange_name, routing_key=self._routing_key)
         return queue
 
     async def kick(self, message: BrokerMessage) -> None:
@@ -189,28 +227,44 @@ async def kick(self, message: BrokerMessage) -> None:
         """
         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,
-            headers={
+
+        message_base_params: dict[str, Any] = {
+            "body": message.message,
+            "headers": {
                 "task_id": message.task_id,
                 "task_name": message.task_name,
                 **message.labels,
             },
-            delivery_mode=DeliveryMode.PERSISTENT,
-            priority=priority,
+            "delivery_mode": DeliveryMode.PERSISTENT,
+        }
+
+        message_base_params["priority"] = parse_val(
+            int,
+            message.labels.get("priority"),
         )
-        delay = parse_val(int, message.labels.get("delay"))
+
+        delay: Optional[int] = parse_val(int, message.labels.get("delay"))
+        rmq_message: Message = Message(**message_base_params)
+
         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)
+            await exchange.publish(rmq_message, routing_key=message.task_name)
+        elif self._delayed_message_exchange_plugin:
+            rmq_message.headers["x-delay"] = delay * 1000
+            exchange = await self.write_channel.get_exchange(
+                self._delay_plugin_exchange_name,
+            )
+            await exchange.publish(
+                rmq_message,
+                routing_key=self._routing_key,
+            )
         else:
-            rmq_msg.expiration = timedelta(seconds=delay)
+            rmq_message.expiration = timedelta(seconds=delay)
             await self.write_channel.default_exchange.publish(
-                rmq_msg,
+                rmq_message,
                 routing_key=self._delay_queue_name,
             )
 
@@ -232,15 +286,3 @@ async def listen(self) -> AsyncGenerator[bytes, None]:
             async for message in iterator:
                 async with message.process():
                     yield message.body
-
-    async def shutdown(self) -> None:
-        """Close all connections on shutdown."""
-        await super().shutdown()
-        if self.write_channel:
-            await self.write_channel.close()
-        if self.read_channel:
-            await self.read_channel.close()
-        if self.write_conn:
-            await self.write_conn.close()
-        if self.read_conn:
-            await self.read_conn.close()

taskiq_aio_pika/broker.py

@@ -42,6 +42,16 @@ def queue_name() -> str:
     return uuid4().hex
 
 
+@pytest.fixture
+def routing_key() -> str:
+    """
+    Generated routing key.
+
+    :return: random routing key.
+    """
+    return uuid4().hex
+
+
 @pytest.fixture
 def delay_queue_name() -> str:
     """
@@ -155,3 +165,67 @@ async def broker(
             if_empty=False,
             if_unused=False,
         )
+
+
+@pytest.fixture
+async def broker_with_delayed_message_plugin(
+    amqp_url: str,
+    queue_name: str,
+    delay_queue_name: str,
+    dead_queue_name: str,
+    exchange_name: str,
+    routing_key: str,
+    test_channel: Channel,
+) -> AsyncGenerator[AioPikaBroker, None]:
+    """
+    Yields new broker instance.
+
+    This function is used to
+    create broker, run startup,
+    and shutdown after test.
+
+    :param amqp_url: current rabbitmq connection string.
+    :param test_channel: amqp channel for tests.
+    :param queue_name: test queue name.
+    :param delay_queue_name: test delay queue name.
+    :param dead_queue_name: test dead letter queue name.
+    :param exchange_name: test exchange name.
+    :param routing_key: routing_key.
+    :yield: broker.
+    """
+    broker = AioPikaBroker(
+        url=amqp_url,
+        declare_exchange=True,
+        exchange_name=exchange_name,
+        dead_letter_queue_name=dead_queue_name,
+        queue_name=queue_name,
+        delayed_message_exchange_plugin=True,
+        routing_key=routing_key,
+    )
+    broker.is_worker_process = True
+
+    await broker.startup()
+
+    yield broker
+
+    await broker.shutdown()
+
+    exchange = await test_channel.get_exchange(exchange_name)
+    await exchange.delete(
+        timeout=1,
+        if_unused=False,
+    )
+    plugin_exchange = await test_channel.get_exchange(
+        broker._delay_plugin_exchange_name,
+    )
+    await plugin_exchange.delete(
+        timeout=1,
+        if_unused=False,
+    )
+    for i_queue_name in (queue_name, delay_queue_name, dead_queue_name):
+        queue = await test_channel.get_queue(i_queue_name, ensure=False)
+        await queue.delete(
+            timeout=1,
+            if_empty=False,
+            if_unused=False,
+        )