Added plugin to create delay messages.

Author: chandr-andr

.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"

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,30 +227,47 @@ 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)
         else:
-            rmq_msg.expiration = timedelta(seconds=delay)
-            await self.write_channel.default_exchange.publish(
-                rmq_msg,
-                routing_key=self._delay_queue_name,
-            )
+            if 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_message.expiration = timedelta(seconds=delay)
+                await self.write_channel.default_exchange.publish(
+                    rmq_message,
+                    routing_key=self._delay_queue_name,
+                )
 
     async def listen(self) -> AsyncGenerator[bytes, None]:
         """
@@ -232,15 +287,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,
+        )

tests/conftest.py

@@ -185,3 +185,33 @@ async def test_delayed_message(
 
     # Check that we can get the message.
     await main_queue.get()
+
+
+@pytest.mark.anyio
+async def test_delayed_message_with_plugin(
+    broker_with_delayed_message_plugin: AioPikaBroker,
+    test_channel: Channel,
+    queue_name: str,
+) -> None:
+    """Test that we can send delayed messages with plugin.
+
+    :param broker_with_delayed_message_plugin: broker with
+    turned on plugin integration.
+    :param test_channel: amqp channel for tests.
+    :param queue_name: test queue name.
+    """
+    main_queue = await test_channel.get_queue(queue_name)
+    broker_msg = BrokerMessage(
+        task_id="1",
+        task_name="name",
+        message=b"message",
+        labels={"delay": "2"},
+    )
+
+    await broker_with_delayed_message_plugin.kick(broker_msg)
+    with pytest.raises(QueueEmpty):
+        await main_queue.get(no_ack=True)
+
+    await asyncio.sleep(2)
+
+    assert await main_queue.get()