Added individual pull queues for each user, retry and error logic and descriptions

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