Updated broker, so it uses different channels for reading and writing.

Signed-off-by: Pavel Kirilin <[email protected]>

Author: s3rius

taskiq_aio_pika/broker.py

@@ -2,9 +2,8 @@
 from logging import getLogger
 from typing import Any, AsyncGenerator, Callable, Optional, TypeVar
 
-from aio_pika import Channel, ExchangeType, Message, connect_robust
+from aio_pika import ExchangeType, Message, connect_robust
 from aio_pika.abc import AbstractChannel, AbstractRobustConnection
-from aio_pika.pool import Pool
 from taskiq.abc.broker import AsyncBroker
 from taskiq.abc.result_backend import AsyncResultBackend
 from taskiq.message import BrokerMessage
@@ -57,48 +56,49 @@ def __init__(  # noqa: WPS211
         """
         super().__init__(result_backend, task_id_generator)
 
-        async def _get_rmq_connection() -> AbstractRobustConnection:
-            return await connect_robust(
-                url,
-                loop=loop,
-                **connection_kwargs,
-            )
-
-        self._connection_pool: Pool[AbstractRobustConnection] = Pool(
-            _get_rmq_connection,
-            max_size=max_connection_pool_size,
-            loop=loop,
-        )
-
-        async def get_channel() -> AbstractChannel:
-            async with self._connection_pool.acquire() as connection:
-                return await connection.channel()
-
-        self._channel_pool: Pool[Channel] = Pool(
-            get_channel,
-            max_size=max_channel_pool_size,
-            loop=loop,
-        )
-
+        self.url = url
+        self._loop = loop
+        self._conn_kwargs = connection_kwargs
         self._exchange_name = exchange_name
         self._exchange_type = exchange_type
         self._qos = qos
         self._declare_exchange = declare_exchange
         self._queue_name = queue_name
         self._routing_key = routing_key
+        self.read_conn: Optional[AbstractRobustConnection] = None
+        self.write_conn: Optional[AbstractRobustConnection] = None
+        self.write_channel: Optional[AbstractChannel] = None
+        self.read_channel: Optional[AbstractChannel] = None
 
-    async def startup(self) -> None:
+    async def startup(self) -> None:  # noqa: WPS217
         """Create exchange and queue on startup."""
-        async with self._channel_pool.acquire() as channel:
-            if self._declare_exchange:
-                exchange = await channel.declare_exchange(
-                    self._exchange_name,
-                    type=self._exchange_type,
-                )
-            else:
-                exchange = await channel.get_exchange(self._exchange_name, ensure=False)
-            queue = await channel.declare_queue(self._queue_name)
-            await queue.bind(exchange=exchange, routing_key=self._routing_key)
+        self.write_conn = await connect_robust(
+            self.url,
+            loop=self._loop,
+            **self._conn_kwargs,
+        )
+        self.write_channel = await self.write_conn.channel()
+
+        if self.is_worker_process:
+            self.read_conn = await connect_robust(
+                self.url,
+                loop=self._loop,
+                **self._conn_kwargs,
+            )
+            self.read_channel = await self.read_conn.channel()
+
+        if self._declare_exchange:
+            exchange = await self.write_channel.declare_exchange(
+                self._exchange_name,
+                type=self._exchange_type,
+            )
+        else:
+            exchange = await self.write_channel.get_exchange(
+                self._exchange_name,
+                ensure=False,
+            )
+        queue = await self.write_channel.declare_queue(self._queue_name)
+        await queue.bind(exchange=exchange, routing_key=self._routing_key)
 
     async def kick(self, message: BrokerMessage) -> None:
         """
@@ -111,6 +111,8 @@ 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 awaited.
         :param message: message to send.
         """
         rmq_msg = Message(
@@ -121,9 +123,13 @@ async def kick(self, message: BrokerMessage) -> None:
                 **message.labels,
             },
         )
-        async with self._channel_pool.acquire() as channel:
-            exchange = await channel.get_exchange(self._exchange_name, ensure=False)
-            await exchange.publish(rmq_msg, routing_key=message.task_name)
+        if self.write_channel is None:
+            raise ValueError("Please run startup before kicking.")
+        exchange = await self.write_channel.get_exchange(
+            self._exchange_name,
+            ensure=False,
+        )
+        await exchange.publish(rmq_msg, routing_key=message.task_name)
 
     async def listen(self) -> AsyncGenerator[BrokerMessage, None]:
         """
@@ -132,29 +138,38 @@ async def listen(self) -> AsyncGenerator[BrokerMessage, None]:
         This function listens to queue and yields
         new messages.
 
+        :raises ValueError: if startup wasn't called.
         :yield: parsed broker messages.
         """
-        async with self._channel_pool.acquire() as channel:
-            await channel.set_qos(prefetch_count=self._qos)
-            queue = await channel.get_queue(self._queue_name, ensure=False)
-            async with queue.iterator() as queue_iter:
-                async for rmq_message in queue_iter:
-                    async with rmq_message.process():
-                        try:
-                            yield BrokerMessage(
-                                task_id=rmq_message.headers.pop("task_id"),
-                                task_name=rmq_message.headers.pop("task_name"),
-                                message=rmq_message.body,
-                                labels=rmq_message.headers,
-                            )
-                        except (ValueError, LookupError) as exc:
-                            logger.debug(
-                                "Cannot read broker message %s",
-                                exc,
-                                exc_info=True,
-                            )
+        if self.read_channel is None:
+            raise ValueError("Call startup before starting listening.")
+        await self.read_channel.set_qos(prefetch_count=0)
+        queue = await self.read_channel.get_queue(self._queue_name, ensure=False)
+        async with queue.iterator() as queue_iter:
+            async for rmq_message in queue_iter:
+                async with rmq_message.process():
+                    try:
+                        yield BrokerMessage(
+                            task_id=rmq_message.headers.pop("task_id"),
+                            task_name=rmq_message.headers.pop("task_name"),
+                            message=rmq_message.body,
+                            labels=rmq_message.headers,
+                        )
+                    except (ValueError, LookupError) as exc:
+                        logger.debug(
+                            "Cannot read broker message %s",
+                            exc,
+                            exc_info=True,
+                        )
 
     async def shutdown(self) -> None:
         """Close all connections on shutdown."""
         await super().shutdown()
-        await self._connection_pool.close()
+        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()