♻️🚑️🤖 RabbitMQ: bound the logs queue to prevent unbounded broker memory growth

packages/service-library/src/servicelib/rabbitmq/_client.py

@@ -1,16 +1,20 @@
 import asyncio
+import datetime
 import logging
 from dataclasses import dataclass, field
 from functools import partial
-from typing import Annotated, Final
+from typing import Annotated, Any, Final
 from uuid import uuid4
 
 import aio_pika
 from aiormq import ChannelInvalidStateError
+from aiormq.exceptions import ChannelNotFoundEntity
 from annotated_types import doc
+from common_library.async_tools import cancel_wait_task
 from common_library.logging.logging_errors import create_troubleshooting_log_kwargs
-from pydantic import NonNegativeInt
+from pydantic import NonNegativeInt, PositiveInt
 
+from ..background_task import create_periodic_task
 from ..logging_utils import log_catch, log_context
 from ._client_base import RabbitMQClientBase
 from ._models import (
@@ -34,6 +38,9 @@
 _DEFAULT_RABBITMQ_EXECUTION_TIMEOUT_S: Final[int] = 5
 _HEADER_X_DEATH: Final[str] = "x-death"
 
+_BACKLOG_MONITOR_INTERVAL: Final[datetime.timedelta] = datetime.timedelta(seconds=60)
+_BACKLOG_MONITOR_CONSECUTIVE_GROWTH_TO_WARN: Final[int] = 3
+
 _DEFAULT_UNEXPECTED_ERROR_RETRY_DELAY_S: Final[float] = 1
 _DEFAULT_UNEXPECTED_ERROR_MAX_ATTEMPTS: Final[NonNegativeInt] = 15
 
@@ -56,7 +63,20 @@
     message_handler: MessageHandler,
     max_retries_upon_error: int,
     message: aio_pika.abc.AbstractIncomingMessage,
+    *,
+    dead_letter_requeue_enabled: bool,
 ) -> None:
+    if not dead_letter_requeue_enabled:
+        # NOTE: no dead-letter-exchange configured on this queue: nacking here drops the message
+        # for good, there is no retry to report on
+        _logger.warning(
+            "Handler '%s' failed for message_id='%s'; dropping it (dead-letter requeue disabled)",
+            message_handler,
+            message.message_id,
+        )
+        await message.nack(requeue=False)
+        return
+
     count = _get_x_death_count(message)
     _logger.debug(
         "Nacking message '%s' from handler '%s', death count %s, max retries %s",
@@ -88,6 +108,8 @@
     message_handler: MessageHandler,
     max_retries_upon_error: int,
     message: aio_pika.abc.AbstractIncomingMessage,
+    *,
+    dead_letter_requeue_enabled: bool,
 ) -> None:
     log_error_context = {
         "message_id": message.message_id,
@@ -113,7 +135,12 @@
                                 logging.DEBUG,
                                 msg=f"Nack message {message.exchange=}, {message.routing_key=}",
                             ):
-                                await _nack_message(message_handler, max_retries_upon_error, message)
+                                await _nack_message(
+                                    message_handler,
+                                    max_retries_upon_error,
+                                    message,
+                                    dead_letter_requeue_enabled=dead_letter_requeue_enabled,
+                                )
                 except Exception as exc:  # pylint: disable=broad-exception-caught
                     _logger.exception(
                         **create_troubleshooting_log_kwargs(
@@ -125,7 +152,12 @@
                         )
                     )
                     with log_catch(_logger, reraise=False):
-                        await _nack_message(message_handler, max_retries_upon_error, message)
+                        await _nack_message(
+                            message_handler,
+                            max_retries_upon_error,
+                            message,
+                            dead_letter_requeue_enabled=dead_letter_requeue_enabled,
+                        )
 
     except ChannelInvalidStateError as exc:
         # NOTE: this error can happen as can be seen in aio-pika code
@@ -146,6 +178,7 @@
 class RabbitMQClient(RabbitMQClientBase):
     _connection_pool: aio_pika.pool.Pool | None = field(init=False, default=None)
     _channel_pool: aio_pika.pool.Pool | None = field(init=False, default=None)
+    _backlog_monitor_tasks: dict[QueueName, asyncio.Task] = field(init=False, default_factory=dict)
 
     def __post_init__(self) -> None:
         # recommendations are 1 connection per process
@@ -175,6 +208,14 @@
             logging.INFO,
             msg=f"{self.client_name} closing connection to RabbitMQ",
         ):
+            with log_catch(_logger, reraise=False):
+                await asyncio.gather(
+                    *(
+                        cancel_wait_task(task, max_delay=_DEFAULT_RABBITMQ_EXECUTION_TIMEOUT_S)
+                        for task in self._backlog_monitor_tasks.values()
+                    )
+                )
+            self._backlog_monitor_tasks.clear()
             assert self._channel_pool  # nosec
             await self._channel_pool.close()
             assert self._connection_pool  # nosec
@@ -188,10 +229,61 @@
             channel.close_callbacks.add(self._channel_close_callback)
             return channel
 
+    def _start_backlog_monitor(self, queue_name: QueueName, exchange_name: ExchangeName, max_length: int) -> None:
+        """Warns if a bounded queue's ready-message count keeps growing, i.e. the
+        consumer cannot keep up with the publish rate (see `subscribe`'s `max_length`)."""
+        previous_message_count: int | None = None
+        consecutive_growth = 0
+
+        async def _check_backlog() -> None:
+            nonlocal previous_message_count, consecutive_growth
+            assert self._connection_pool  # nosec
+            # NOTE: uses a dedicated, short-lived channel (not the shared `_channel_pool`) since
+            # that pool can hand out a channel that is concurrently in use to actively deliver
+            # an unacked message to a consumer, causing protocol interleaving on that channel
+            async with self._connection_pool.acquire() as connection:
+                channel = await connection.channel()
+                try:
+                    declared = await channel.declare_queue(
+                        queue_name, passive=True, timeout=_DEFAULT_RABBITMQ_EXECUTION_TIMEOUT_S
+                    )
+                    message_count = declared.declaration_result.message_count
+                except ChannelNotFoundEntity:
+                    # queue is gone (e.g. unsubscribe() raced this check): stop monitoring for good
+                    _logger.debug("Queue '%s' no longer exists, stopping its backlog monitor", queue_name)
+                    task = self._backlog_monitor_tasks.pop(queue_name, None)
+                    if task is not None:
+                        task.cancel()
+                    return
+                finally:
+                    await channel.close()
+            consecutive_growth = (
+                consecutive_growth + 1
+                if previous_message_count is not None and message_count and message_count > previous_message_count
+                else 0
+            )
+            if consecutive_growth >= _BACKLOG_MONITOR_CONSECUTIVE_GROWTH_TO_WARN:
+                _logger.warning(
+                    "Queue '%s' (exchange '%s') backlog kept growing for %s consecutive checks "
+                    "(now %s/%s ready messages): the consumer may not be keeping up with the publish rate",
+                    queue_name,
+                    exchange_name,
+                    consecutive_growth,
+                    message_count,
+                    max_length,
+                )
+            previous_message_count = message_count
+
+        self._backlog_monitor_tasks[queue_name] = create_periodic_task(
+            _check_backlog,
+            interval=_BACKLOG_MONITOR_INTERVAL,
+            task_name=f"rabbitmq_backlog_monitor_{queue_name}",
+        )
+
     async def _create_consumer_tag(self, exchange_name) -> ConsumerTag:
         return ConsumerTag(f"{get_rabbitmq_client_unique_name(self.client_name)}_{exchange_name}_{uuid4()}")
 
-    async def subscribe(
+    async def subscribe(  # noqa: PLR0913 # pylint: disable=too-many-arguments
         self,
         exchange_name: ExchangeName,
         message_handler: Annotated[
@@ -232,17 +324,57 @@
             NonNegativeInt,
             doc(
                 "Also acts as a soft timeout: if `message_handler` does not finish processing "
-                "the message before this is reached, the message will be redelivered"
+                "the message before this is reached, the message will be redelivered (or, if "
+                "`enable_dead_letter_requeue=False`, simply dropped)"
             ),
         ] = RABBIT_QUEUE_MESSAGE_DEFAULT_TTL_MS,
+        max_length: Annotated[
+            NonNegativeInt | None,
+            doc(
+                "Caps the queue depth. Once reached, the OLDEST ready messages are silently "
+                "dropped (`x-overflow: drop-head`) to make room for new ones, protecting the "
+                "broker from unbounded memory growth if `message_handler` ever falls behind "
+                "the publish rate. None (default) leaves the queue unbounded; only set this "
+                "for exchanges where losing old messages is preferable to broker instability. "
+                "NOTE: RabbitMQ dead-letters messages dropped this way too (reason `maxlen`), same "
+                "as nacked/expired ones; pair this with `enable_dead_letter_requeue=False`, or the "
+                "dropped messages will bounce forever between this queue and its delay queue until "
+                "RabbitMQ's own dead-letter-cycle detector catches it"
+            ),
+        ] = None,
+        prefetch_count: Annotated[
+            PositiveInt | None,
+            doc(
+                "Maximum number of messages delivered to (and awaiting ack from) `message_handler` "
+                "concurrently. None (default) uses 1 for shared (`exclusive_queue=False`) queues, or "
+                f"{_DEFAULT_PREFETCH_VALUE} for exclusive ones. Raise this for lightweight, I/O-bound "
+                "handlers on high-throughput exchanges, where the default is otherwise the throughput "
+                "ceiling regardless of how fast `message_handler` actually runs"
+            ),
+        ] = None,
         unexpected_error_retry_delay_s: Annotated[
             float,
-            doc("Time to wait between each retry when `message_handler` raised or returned `False`"),
+            doc(
+                "Time to wait between each retry when `message_handler` raised or returned `False`. "
+                "Has no effect when `enable_dead_letter_requeue=False`: such messages are dropped "
+                "immediately instead of being retried"
+            ),
         ] = _DEFAULT_UNEXPECTED_ERROR_RETRY_DELAY_S,
         unexpected_error_max_attempts: Annotated[
             int,
             doc("Maximum amount of retries when `message_handler` raised or returned `False`"),
         ] = _DEFAULT_UNEXPECTED_ERROR_MAX_ATTEMPTS,
+        enable_dead_letter_requeue: Annotated[
+            bool,
+            doc(
+                "When True (default), messages that are nacked or that expire after sitting "
+                "`message_ttl` in the queue are bounced through a delay queue and re-published into "
+                "THIS SAME exchange for a retry, up to `unexpected_error_max_attempts` times. Set False "
+                "for best-effort/fire-and-forget exchanges (e.g. live UI notifications) where a stale "
+                "message has no value: expired/nacked messages are then simply dropped instead of "
+                "generating more publish traffic back into an exchange that may already be backlogged"
+            ),
+        ] = True,
     ) -> Annotated[
         tuple[QueueName, ConsumerTag],
         doc("Returns the queue name and consumer tag of the subscription"),
@@ -256,7 +388,11 @@
 
         assert self._channel_pool  # nosec
         async with self._channel_pool.acquire() as channel:
-            qos_value = 1 if exclusive_queue is False else _DEFAULT_PREFETCH_VALUE
+            qos_value = (
+                prefetch_count
+                if prefetch_count is not None
+                else (1 if exclusive_queue is False else _DEFAULT_PREFETCH_VALUE)
+            )
             await channel.set_qos(qos_value)
 
             exchange = await channel.declare_exchange(
@@ -271,41 +407,55 @@
             # exclusive means that the queue is only available for THIS very client
             # and will be deleted when the client disconnects
             # NOTE what is a dead letter exchange, see https://www.rabbitmq.com/dlx.html
-            delayed_exchange_name = _DELAYED_EXCHANGE_NAME.format(exchange_name=exchange_name)
+            queue_arguments: dict[str, Any] = {}
+            if enable_dead_letter_requeue:
+                delayed_exchange_name = _DELAYED_EXCHANGE_NAME.format(exchange_name=exchange_name)
+                queue_arguments["x-dead-letter-exchange"] = delayed_exchange_name
+            if max_length is not None:
+                queue_arguments["x-max-length"] = max_length
+                queue_arguments["x-overflow"] = "drop-head"
             queue = await declare_queue(
                 channel,
                 self.client_name,
                 non_exclusive_queue_name or exchange_name,
                 exclusive_queue=exclusive_queue,
                 message_ttl=message_ttl,
-                arguments={"x-dead-letter-exchange": delayed_exchange_name},
+                arguments=queue_arguments,
             )
             if topics is None:
                 await queue.bind(exchange, routing_key="")
             else:
                 await asyncio.gather(*(queue.bind(exchange, routing_key=topic) for topic in topics))
 
-            delayed_exchange = await channel.declare_exchange(
-                delayed_exchange_name, aio_pika.ExchangeType.FANOUT, durable=True
-            )
-            delayed_queue_name = _DELAYED_QUEUE_NAME.format(queue_name=non_exclusive_queue_name or exchange_name)
-
-            delayed_queue = await declare_queue(
-                channel,
-                self.client_name,
-                delayed_queue_name,
-                exclusive_queue=exclusive_queue,
-                message_ttl=int(unexpected_error_retry_delay_s * 1000),
-                arguments={"x-dead-letter-exchange": exchange.name},
-            )
-            await delayed_queue.bind(delayed_exchange)
+            if enable_dead_letter_requeue:
+                delayed_exchange = await channel.declare_exchange(
+                    delayed_exchange_name, aio_pika.ExchangeType.FANOUT, durable=True
+                )
+                delayed_queue_name = _DELAYED_QUEUE_NAME.format(queue_name=non_exclusive_queue_name or exchange_name)
+
+                delayed_queue = await declare_queue(
+                    channel,
+                    self.client_name,
+                    delayed_queue_name,
+                    exclusive_queue=exclusive_queue,
+                    message_ttl=int(unexpected_error_retry_delay_s * 1000),
+                    arguments={"x-dead-letter-exchange": exchange.name},
+                )
+                await delayed_queue.bind(delayed_exchange)
 
             consumer_tag = await self._create_consumer_tag(exchange_name)
             await queue.consume(
-                partial(_on_message, message_handler, unexpected_error_max_attempts),
+                partial(
+                    _on_message,
+                    message_handler,
+                    unexpected_error_max_attempts,
+                    dead_letter_requeue_enabled=enable_dead_letter_requeue,
+                ),
                 exclusive=exclusive_queue,
                 consumer_tag=consumer_tag,
             )
+            if max_length is not None:
+                self._start_backlog_monitor(queue.name, exchange_name, max_length)
             return queue.name, consumer_tag
 
     async def add_topics(
@@ -323,7 +473,7 @@
                 self.client_name,
                 exchange_name,
                 exclusive_queue=True,
-                arguments={"x-dead-letter-exchange": _DELAYED_EXCHANGE_NAME.format(exchange_name=exchange_name)},
+                passive=True,
             )
 
             await asyncio.gather(*(queue.bind(exchange, routing_key=topic) for topic in topics))
@@ -342,7 +492,7 @@
                 self.client_name,
                 exchange_name,
                 exclusive_queue=True,
-                arguments={"x-dead-letter-exchange": _DELAYED_EXCHANGE_NAME.format(exchange_name=exchange_name)},
+                passive=True,
             )
 
             await asyncio.gather(
@@ -363,6 +513,9 @@
             queue = await channel.get_queue(queue_name)
             # NOTE: we force delete here
             await queue.delete(if_unused=False, if_empty=False)
+        backlog_monitor_task = self._backlog_monitor_tasks.pop(queue_name, None)
+        if backlog_monitor_task is not None:
+            await cancel_wait_task(backlog_monitor_task, max_delay=_DEFAULT_RABBITMQ_EXECUTION_TIMEOUT_S)
 
     async def publish(self, exchange_name: ExchangeName, message: RabbitMessage) -> None:
         """publish message in the exchange exchange_name.

packages/service-library/src/servicelib/rabbitmq/_utils.py

@@ -74,14 +74,28 @@ async def declare_queue(
     exclusive_queue: bool,
     arguments: dict[str, Any] | None = None,
     message_ttl: NonNegativeInt = RABBIT_QUEUE_MESSAGE_DEFAULT_TTL_MS,
+    passive: bool = False,
 ) -> aio_pika.abc.AbstractRobustQueue:
+    """Declares (or, if `passive=True`, just looks up) the queue derived from `queue_name`.
+
+    `passive=False` (default): asks the broker to create the queue if it does not exist, or
+    otherwise assert that it exists AND has the exact same `durable`/`exclusive`/`arguments` as
+    passed here; a mismatch raises `ChannelPreconditionFailed` and closes the channel. Use this
+    only where this call is meant to define the queue's shape (e.g. the original `subscribe()`).
+
+    `passive=True`: only checks that the queue already exists (raises if not) and returns a handle
+    to it; the broker does not create anything and does not compare arguments, so this is safe to
+    use to just re-attach to (e.g. to bind/unbind topics on) a queue defined elsewhere, even if the
+    `arguments` passed here don't match what that queue was actually created with.
+    """
     default_arguments = {"x-message-ttl": message_ttl}
     if arguments is not None:
         default_arguments.update(arguments)
     queue_parameters: dict[str, Any] = {
         "durable": not exclusive_queue,
         "exclusive": exclusive_queue,
         "arguments": default_arguments,
+        "passive": passive,
         "name": f"{get_rabbitmq_client_unique_name(client_name)}_{queue_name}_exclusive",
     }
     if not exclusive_queue: