Restore robust consumers after broker cancellation

README.md

@@ -33,6 +33,21 @@ If you are a newcomer to RabbitMQ, please start with the [adopted official Rabbi
 * Complete type-hints coverage.
 
 
+## Consumer recovery
+
+With `connect_robust`, consumers registered with `robust=True` on robust
+queues recover after a broker `basic.cancel`, including when a classic queue's
+node restarts. Recovery retries with a delay and restores the queue's bindings
+before registering the cancelled consumers again. Other consumers stay active.
+
+A queue declared normally is recreated if it was deleted by the broker or
+another client. A queue declared with `passive=True` waits until it exists again.
+Explicit `queue.cancel()`, deleting the queue through its own channel with
+`queue.delete()` or `channel.queue_delete()`, and closing the channel or
+connection stop recovery. Consumers or queues registered with `robust=False`
+do not recover automatically.
+
+
 ## Installation
 
 ```shell

aio_pika/robust_channel.py

@@ -71,6 +71,7 @@ def __init__(
         self.reopen_callbacks = CallbackCollection(self)
         self.__restore_lock = asyncio.Lock()
         self.__restored = asyncio.Event()
+        self._close_called = False
 
         self.close_callbacks.remove(self._set_closed_callback)
 
@@ -93,6 +94,7 @@ async def restore(self, channel: Any = None) -> None:
         async with self.__restore_lock:
             if (
                 self.__restored.is_set()
+                or self._close_called
                 or self._connection.close_called
                 or self._connection.is_closed
             ):
@@ -137,13 +139,14 @@ async def _on_close(self, closing: asyncio.Future) -> BaseException | None:
 
     async def close(
         self,
-        exc: Optional[aiormq.abc.ExceptionType] = None,
+        exc: aiormq.abc.ExceptionType | None = None,
     ) -> None:
-        # Avoid recovery when channel is explicitely closed using this method
+        self._close_called = True
         self.__restored.clear()
         await super().close(exc)
 
     async def reopen(self) -> None:
+        self._close_called = False
         await super().reopen()
         await self.reopen_callbacks()
 
@@ -154,6 +157,7 @@ async def _on_open(self) -> None:
         exchanges = self._exchanges.values()
         queues = tuple(chain(*self._queues.values()))
         channel = await self.get_underlay_channel()
+        channel.on_consumer_cancel_callbacks.add(self._on_consumer_cancel)
 
         await channel.basic_qos(
             prefetch_count=self._prefetch_count,
@@ -172,6 +176,78 @@ async def _on_open(self) -> None:
 
         self.__restored.set()
 
+    def _on_consumer_cancel(self, frame: aiormq.spec.Basic.Cancel) -> None:
+        from .robust_connection import RobustConnection
+
+        if isinstance(self._connection, RobustConnection) and any(
+            frame.consumer_tag in queue._consumers
+            for queue in chain(*self._queues.values())
+        ):
+            self._connection._request_recovery()
+
+    async def _restore_consumers(self) -> bool:
+        if self._close_called or self._closed.done():
+            return False
+        try:
+            await self.restore()
+        except Exception:
+            log.debug("Channel recovery failed; retrying", exc_info=True)
+            return True
+        channel = cast(aiormq.Channel, await self.get_underlay_channel())
+        recovery = asyncio.create_task(self._restore_missing_consumers(channel))
+        try:
+            await asyncio.wait(
+                (recovery, self._closed, channel.closing),
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+            if self._close_called or self._closed.done():
+                return False
+            if channel.closing.done():
+                return True
+            return await recovery
+        finally:
+            recovery.cancel()
+            await asyncio.gather(recovery, return_exceptions=True)
+
+    async def _restore_missing_consumers(self, channel: aiormq.Channel) -> bool:
+        retry = False
+        async with self.__restore_lock:
+            for queue in tuple(chain(*self._queues.values())):
+                tags = queue._consumers.keys() - channel.consumers.keys()
+                if not tags:
+                    continue
+                try:
+                    probe = await channel.connection.channel(
+                        publisher_confirms=False,
+                        timeout=5,
+                    )
+                    try:
+                        await queue._probe(probe)
+                    finally:
+                        await probe.close()
+                    if not self.__restored.is_set() or channel.is_closed:
+                        return False
+                    for tag in tags:
+                        kwargs = queue._consumers.get(tag)
+                        if kwargs is None or tag in channel.consumers:
+                            continue
+                        await Queue.consume(
+                            queue,
+                            consumer_tag=tag,
+                            timeout=5,
+                            **kwargs,
+                        )
+                        if queue._consumers.get(tag) is not kwargs:
+                            await channel.basic_cancel(tag)
+                except Exception:
+                    log.debug(
+                        "Consumer recovery failed for %r; retrying",
+                        queue,
+                        exc_info=True,
+                    )
+                    retry = True
+        return retry
+
     async def set_qos(
         self,
         prefetch_count: int = 0,
@@ -301,14 +377,16 @@ async def queue_delete(
         nowait: bool = False,
     ) -> aiormq.spec.Queue.DeleteOk:
         await self.ready()
+        queues = self._queues.pop(queue_name, set())
+        for queue in queues:
+            queue._consumers.clear()
         result = await super().queue_delete(
             queue_name=queue_name,
             timeout=timeout,
             if_unused=if_unused,
             if_empty=if_empty,
             nowait=nowait,
         )
-        self._queues.pop(queue_name, None)
         return result
 
 

aio_pika/robust_connection.py

@@ -60,14 +60,18 @@ def __init__(
             self.__fail_fast_future.set_result(None)
 
         self.__channels: WeakSet[AbstractRobustChannel] = WeakSet()
-        self.__connection_close_event = asyncio.Event()
+        self.__restore_event = asyncio.Event()
         self.__connect_timeout: Optional[TimeoutType] = None
         self.__reconnection_task: Optional[asyncio.Task] = None
 
         self._reconnect_lock = asyncio.Lock()
         self.reconnect_callbacks = CallbackCollection(self)
 
-        self.__connection_close_event.set()
+        self.__restore_event.set()
+
+    def _request_recovery(self) -> None:
+        if not self.is_closed and not self._close_called:
+            self.__restore_event.set()
 
     @property
     def reconnecting(self) -> bool:
@@ -97,7 +101,7 @@ async def _on_connection_close(self, closing: asyncio.Future) -> None:
             self.reconnect_interval,
         )
 
-        self.__connection_close_event.set()
+        self.__restore_event.set()
 
     async def _on_connected(self) -> None:
         await super()._on_connected()
@@ -141,30 +145,29 @@ async def _on_connected(self) -> None:
             await self.reconnect_callbacks()
 
         self.connection_attempt += 1
-        self.__connection_close_event.clear()
 
     async def __connection_factory(self) -> None:
         log.debug("Starting connection factory for %r", self)
         while not self.is_closed and not self._close_called:
-            log.debug("Waiting for connection close event for %r", self)
-            await self.__connection_close_event.wait()
+            log.debug("Waiting for recovery event for %r", self)
+            await self.__restore_event.wait()
+            self.__restore_event.clear()
 
             if self.is_closed or self._close_called:
                 return
 
             # noinspection PyBroadException
             try:
-                self.transport = None
-                self.connected.clear()
-
-                log.debug("Connection attempt for %r", self)
-                await Connection.connect(self, self.__connect_timeout)
-
-                if not self.__fail_fast_future.done():
-                    self.__fail_fast_future.set_result(None)
-
-                log.debug("Connection made on %r", self)
+                transport = self.transport
+                if transport is not None and not transport.connection.is_closed:
+                    for channel in tuple(self.__channels):
+                        if isinstance(channel, RobustChannel):
+                            if await channel._restore_consumers():
+                                self.__restore_event.set()
+                else:
+                    await self.__connect_transport()
             except CONNECTION_EXCEPTIONS as e:
+                self.__restore_event.set()
                 if not self.__fail_fast_future.done():
                     self.__fail_fast_future.set_exception(e)
                     return
@@ -177,6 +180,7 @@ async def __connection_factory(self) -> None:
                     self.reconnect_interval,
                 )
             except asyncio.CancelledError:
+                self.__restore_event.set()
                 if self._close_called or self.is_closed:
                     raise
                 log.warning(
@@ -186,6 +190,7 @@ async def __connection_factory(self) -> None:
                     self.reconnect_interval,
                 )
             except Exception as exc:
+                self.__restore_event.set()
                 log.error(
                     "Reconnect attempt failed %s due to %s: %s. "
                     "Retrying after %r seconds.",
@@ -201,6 +206,15 @@ async def __connection_factory(self) -> None:
 
             await asyncio.sleep(self.reconnect_interval)
 
+    async def __connect_transport(self) -> None:
+        self.transport = None
+        self.connected.clear()
+        log.debug("Connection attempt for %r", self)
+        await Connection.connect(self, self.__connect_timeout)
+        if not self.__fail_fast_future.done():
+            self.__fail_fast_future.set_result(None)
+        log.debug("Connection made on %r", self)
+
     async def connect(self, timeout: TimeoutType = None) -> None:
         self.__connect_timeout = timeout
 
@@ -249,10 +263,9 @@ def channel(
 
     async def close(
         self,
-        exc: Optional[aiormq.abc.ExceptionType] = asyncio.CancelledError,
+        exc: aiormq.abc.ExceptionType | None = asyncio.CancelledError,
     ) -> None:
         self._close_called = True
-
         if self.__reconnection_task is not None:
             self.__reconnection_task.cancel()
             await asyncio.gather(

aio_pika/robust_queue.py

@@ -15,6 +15,7 @@
     AbstractRobustQueue,
     ConsumerTag,
     TimeoutType,
+    get_exchange_name,
 )
 from .exchange import ExchangeParamType
 from .log import get_logger
@@ -72,6 +73,40 @@ async def restore(self, channel: Any = None) -> None:
         for consumer_tag, kwargs in consumers:
             await self.consume(consumer_tag=consumer_tag, **kwargs)
 
+    async def _probe(self, channel: aiormq.abc.AbstractChannel) -> None:
+        await channel.queue_declare(
+            queue=self.name,
+            durable=self.durable,
+            exclusive=self.exclusive,
+            auto_delete=self.auto_delete,
+            arguments=self.arguments,
+            passive=self.passive,
+            timeout=5,
+        )
+        for (exchange, routing_key), kwargs in tuple(self._bindings.items()):
+            await channel.queue_bind(
+                queue=self.name,
+                exchange=get_exchange_name(exchange),
+                routing_key=routing_key,
+                timeout=5,
+                **kwargs,
+            )
+
+    async def delete(
+        self,
+        *,
+        if_unused: bool = True,
+        if_empty: bool = True,
+        timeout: TimeoutType = None,
+    ) -> aiormq.spec.Queue.DeleteOk:
+        self._consumers.clear()
+        return await self.channel.queue_delete(
+            self.name,
+            if_unused=if_unused,
+            if_empty=if_empty,
+            timeout=timeout,
+        )
+
     async def bind(
         self,
         exchange: ExchangeParamType,