From 71ec05978bb2b1eab36e1e57cf799001aa08bb05 Mon Sep 17 00:00:00 2001 From: Florian Date: Sun, 11 Oct 2026 12:32:05 +0200 Subject: [PATCH] Added individual pull queues for each user, retry and error logic and descriptions --- src/rabbitmq_handler.py | 117 ++++++++++++++++++++++++++++++++-------- 1 file changed, 95 insertions(+), 22 deletions(-) diff --git a/src/rabbitmq_handler.py b/src/rabbitmq_handler.py index eb93117..3ea9b95 100644 --- a/src/rabbitmq_handler.py +++ b/src/rabbitmq_handler.py @@ -1,24 +1,48 @@ import asyncio import aio_pika -from aio_pika.exceptions import AMQPException +from aio_pika.exceptions import AMQPException, DeliveryError from secret_manager import return_credentials import os from simple_logger_handler import setup_logger import json from urllib.parse import quote +from typing import NamedTuple +from enum import IntEnum logger = setup_logger(__name__) +MAX_UNROUTABLE_ATTEMPTS = 3 +DRAIN_TIMEOUT_SECONDS = 10 # how long close() waits for buffered messages + +class Kind(IntEnum): + PUSH = 0 + PULL = 1 + +class OutgoingMessage(NamedTuple): + user_id: int + routing_key: str + message: dict + attempts: int + kind: Kind class RabbitMQProducer: + """ + Buffered RabbitMQ producer. publish() only enqueues in memory; a background + task sends each message to the exchange, once with the push key and once with + the pull key per user. + """ def __init__(self): self.rmq_host = os.getenv("BACKEND_API_INTERNAL_RMQ_HOST", "localhost") self.rmq_vhost = os.getenv("BACKEND_API_INTERNAL_RMQ_VHOST", "app_notifications") self.rmq_exchange = os.getenv("BACKEND_API_INTERNAL_RMQ_EXCHANGE", "app_notifications") + self.pull_routing_prefix = os.getenv("BACKEND_PULL_NOTIFICATIONS_RMQ_ROUTING_PREFIX", "notify.pull") + self.push_routing_prefix = os.getenv("BACKEND_PUSH_NOTIFICATIONS_RMQ_ROUTING_PREFIX", "notify.user") + self.pull_queue_prefix = os.getenv("BACKEND_PULL_NOTIFICATIONS_RMQ_QUEUE_PREFIX", "notifications-pull") + self.push_queue_name = os.getenv("BACKEND_PUSH_NOTIFICATIONS_RMQ_QUEUE", "notifications") self.connection: aio_pika.RobustConnection | None = None self.channel: aio_pika.RobustChannel | None = None self.exchange: aio_pika.Exchange | None = None - self._queue: asyncio.Queue[tuple[str, dict]] = asyncio.Queue() + self._queue: asyncio.Queue[OutgoingMessage] = asyncio.Queue() self._flush_task: asyncio.Task | None = None self._closing = False self._ready = asyncio.Event() @@ -38,37 +62,88 @@ class RabbitMQProducer: self.rmq_exchange, aio_pika.ExchangeType.TOPIC, durable=True ) logger.info(f"[RabbitMQ] Connected and exchange '{self.rmq_exchange}' ready.") + push_queue = await self.channel.declare_queue(self.push_queue_name, durable=True) + await push_queue.bind(self.exchange, routing_key=f"{self.push_routing_prefix}.*") self._ready.set() self._flush_task = asyncio.create_task(self._flush_queue_loop()) logger.debug("[RabbitMQ] Flush queue loop started") - async def publish(self, routing_key: int, message: dict): - await self._queue.put((f"notify.user.{routing_key}", message)) - await self._queue.put((f"notify.windows.{routing_key}", message)) - logger.debug(f"[RabbitMQ] Queued message for {routing_key}") + async def ensure_user_queue(self, user_id: int) -> aio_pika.Queue: + """ + Declare the user's durable queue (with expiry and message TTL) and bind it to + the exchange with the user's routing key. Safe to call repeatedly, creates it if it expired, and resets its + idle-expiry timer. + """ + logger.debug("[RabbitMQ] Declaring user pull queue for user:{user_id}") + queue = await self.channel.declare_queue( + f"{self.pull_queue_prefix}.{user_id}", + durable=True, + arguments={ + "x-expires": 30 * 24 * 3600 * 1000, + "x-message-ttl": 7 * 24 * 3600 * 1000, + }, + ) + await queue.bind(self.exchange, routing_key=f"{self.pull_routing_prefix}.{user_id}") + return queue + + async def publish(self, user_id: str | int, message: dict): + """ + Buffer a message for a user under two routing keys (push and pull). Works while RabbitMQ is down. + Raises: + - RuntimeError if the producer is closing + """ + if self._closing: + raise RuntimeError("[RabbitMQ] Producer is closing") + user_id = int(user_id) + await self._queue.put(OutgoingMessage(user_id, f"{self.push_routing_prefix}.{user_id}", message, 0, Kind.PUSH)) + await self._queue.put(OutgoingMessage(user_id, f"{self.pull_routing_prefix}.{user_id}", message, 0, Kind.PULL)) + logger.debug(f"[RabbitMQ] Queued message for {user_id}") async def _flush_queue_loop(self): + """Background task: wait until connected, then take buffered messages one at a + time and publish them as persistent JSON. + + - Declares the user's pull queue first when the key is a pull key. + - Returned/unroutable messages (DeliveryError) are retried up to + MAX_UNROUTABLE_ATTEMPTS times, then dropped with an error log. + - Connection errors are retried without limit after a 2 second pause. + - Any other error (e.g. non-serializable payload) drops the message so the + loop keeps running. + """ logger.debug("[RabbitMQ] Waiting for ready signal before flushing queue") await self._ready.wait() while True: - routing_key, message = await self._queue.get() - logger.debug(f"[RabbitMQ] Publishing message to routing key {routing_key}: {message}") + outgoing_message = await self._queue.get() + logger.debug(f"[RabbitMQ] Publishing message to routing key {outgoing_message.routing_key}: {outgoing_message.message}") try: + if outgoing_message.kind == Kind.PULL: + await self.ensure_user_queue(outgoing_message.user_id) await self.exchange.publish( aio_pika.Message( - body=json.dumps(message).encode(), + body=json.dumps(outgoing_message.message).encode(), content_type="application/json", delivery_mode=aio_pika.DeliveryMode.PERSISTENT ), - routing_key=routing_key, + routing_key=outgoing_message.routing_key, mandatory=True ) - logger.debug(f"[RabbitMQ] Successfully published to notify.user.{routing_key}") + logger.debug(f"[RabbitMQ] Successfully published to {outgoing_message.routing_key}") + + except DeliveryError as e: + attempts = outgoing_message.attempts + 1 + if attempts >= MAX_UNROUTABLE_ATTEMPTS: + logger.error(f"[RabbitMQ] Dropping message for {outgoing_message.routing_key} after {attempts} attempts: {e}") + else: + await asyncio.sleep(2) + await self._queue.put(outgoing_message._replace(attempts=attempts)) except (AMQPException, OSError) as e: - logger.warning(f"[RabbitMQ] Failed to publish to {routing_key}: {e}. Requeuing message.", exc_info=True) + logger.warning(f"[RabbitMQ] Failed to publish to {outgoing_message.routing_key}: {e}. Requeuing message.", exc_info=True) await asyncio.sleep(2) - await self._queue.put((routing_key, message)) + await self._queue.put(outgoing_message) + + except Exception as e: + logger.error(f"[RabbitMQ] Dropping unpublishable message for {outgoing_message.routing_key}: {e}", exc_info=True) finally: self._queue.task_done() @@ -76,20 +151,18 @@ class RabbitMQProducer: logger.info("[RabbitMQ] Closing producer...") self._closing = True if self._flush_task and not self._flush_task.done(): + try: + await asyncio.wait_for(self._queue.join(), timeout=DRAIN_TIMEOUT_SECONDS) + except asyncio.TimeoutError: + logger.warning(f"[RabbitMQ] {self._queue.qsize()} message(s) still buffered, dropping on shutdown") + self._flush_task.cancel() try: await self._flush_task - logger.debug("[RabbitMQ] Flush task cancelled successfully") except asyncio.CancelledError: - logger.debug("[RabbitMQ] Flush task cancellation caught") + logger.debug("[RabbitMQ] Flush task cancelled") + if self.connection: await self.connection.close() logger.info("[RabbitMQ] Connection closed") - -async def main(): - producer = RabbitMQProducer() - await producer.connect() - await producer.publish("1", {"type": "notification", "content": "Test message"}) - await producer.close() -