backend-api-internal/src/rabbitmq_handler.py

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")