import asyncio import aio_pika 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[OutgoingMessage] = asyncio.Queue() self._flush_task: asyncio.Task | None = None self._closing = False self._ready = asyncio.Event() async def connect(self): username = return_credentials("/etc/secrets/rmq_username") password = return_credentials("/etc/secrets/rmq_password") url = ( f"amqp://{quote(username, safe='')}:{quote(password, safe='')}" f"@{self.rmq_host}/{quote(self.rmq_vhost, safe='')}" ) logger.info(f"[RabbitMQ] Connecting to {self.rmq_host}...") self.connection = await aio_pika.connect_robust(url) self.channel = await self.connection.channel(publisher_confirms=True) self.exchange = await self.channel.declare_exchange( 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 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: 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(outgoing_message.message).encode(), content_type="application/json", delivery_mode=aio_pika.DeliveryMode.PERSISTENT ), routing_key=outgoing_message.routing_key, mandatory=True ) 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 {outgoing_message.routing_key}: {e}. Requeuing message.", exc_info=True) await asyncio.sleep(2) 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() async def close(self): 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 except asyncio.CancelledError: logger.debug("[RabbitMQ] Flush task cancelled") if self.connection: await self.connection.close() logger.info("[RabbitMQ] Connection closed")