135 lines
4.8 KiB
Python
135 lines
4.8 KiB
Python
import asyncio
|
|
import aio_pika
|
|
from aio_pika.exceptions import AMQPException
|
|
from secret_manager import return_credentials
|
|
import os
|
|
from simple_logger_handler import setup_logger
|
|
import json
|
|
import signal
|
|
from urllib.parse import quote
|
|
import weakref
|
|
|
|
logger = setup_logger(__name__)
|
|
|
|
class RabbitMQConsumer:
|
|
"""Pull-based RabbitMQ consumer. Messages are fetched on demand, one queue per user."""
|
|
def __init__(self):
|
|
self.rmq_host = os.getenv("BACKEND_PUSH_NOTIFICATIONS_RMQ_HOST", "localhost")
|
|
self.rmq_vhost = os.getenv("BACKEND_PUSH_NOTIFICATIONS_RMQ_VHOST", "app_notifications")
|
|
self.rmq_exchange = os.getenv("BACKEND_PUSH_NOTIFICATIONS_RMQ_EXCHANGE", "app_notifications")
|
|
self.pull_queue_prefix = os.getenv("BACKEND_PULL_NOTIFICATIONS_RMQ_QUEUE_PREFIX", "notifications-pull")
|
|
self.pull_routing_prefix = os.getenv("BACKEND_PULL_NOTIFICATIONS_RMQ_ROUTING_PREFIX", "notify.pull")
|
|
|
|
self._user_locks: weakref.WeakValueDictionary[int, asyncio.Lock] = weakref.WeakValueDictionary()
|
|
self.connection: aio_pika.RobustConnection | None = None
|
|
self.channel: aio_pika.RobustChannel | None = None
|
|
self.exchange: aio_pika.Exchange | None = None
|
|
self._closing = False
|
|
logger.debug(f"[RabbitMQ] Initialized consumer for exchange {self.rmq_exchange}")
|
|
|
|
def _lock_for(self, user_id: int) -> asyncio.Lock:
|
|
"""
|
|
Return the asyncio lock for this user, creating it if needed. Locks are held
|
|
weakly, so they are cleaned up automatically once nobody is using them.
|
|
"""
|
|
lock = self._user_locks.get(user_id)
|
|
if lock is None:
|
|
lock = asyncio.Lock()
|
|
self._user_locks[user_id] = lock
|
|
return lock
|
|
|
|
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()
|
|
self.exchange = await self.channel.declare_exchange(
|
|
self.rmq_exchange, aio_pika.ExchangeType.TOPIC, durable=True
|
|
)
|
|
logger.info("[RabbitMQ] Connected to RabbitMQ")
|
|
|
|
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, recreates it if it expired, and resets its
|
|
idle-expiry timer.
|
|
"""
|
|
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 fetch_messages(self, user_id: int, max_messages: int = 100) -> list[dict]:
|
|
"""Drain up to max_messages queued messages for this user and return their
|
|
decoded JSON bodies.
|
|
|
|
- Only one fetch per user runs at a time; different users run concurrently.
|
|
- Messages are acked after the whole batch has been read, so an error mid-read
|
|
puts the batch back on the queue and re-raises.
|
|
- Undecodable messages are rejected without requeue (discarded).
|
|
- Returns an empty list if nothing is queued.
|
|
- Delivery is at-most-once towards the client: messages are acked before the
|
|
HTTP response is sent.
|
|
"""
|
|
|
|
if self._closing:
|
|
raise RuntimeError("[RabbitMQ] Consumer is closing")
|
|
if not self.channel:
|
|
raise RuntimeError("[RabbitMQ] Queue not initialized, call connect() first")
|
|
|
|
user_id = int(user_id)
|
|
|
|
async with self._lock_for(user_id): #one drain at a time
|
|
queue = await self.ensure_user_queue(user_id)
|
|
results: list[dict] = []
|
|
to_ack: list[aio_pika.IncomingMessage] = []
|
|
try:
|
|
while len(to_ack) < max_messages:
|
|
message = await queue.get(no_ack=False, fail=False, timeout=5)
|
|
if message is None:
|
|
break
|
|
|
|
try:
|
|
data = json.loads(message.body.decode())
|
|
except (json.JSONDecodeError, UnicodeDecodeError) as e:
|
|
logger.error(f"[RabbitMQ] Bad message, discarding: {e}", exc_info=True)
|
|
await message.reject(requeue=False)
|
|
continue
|
|
|
|
results.append(data)
|
|
to_ack.append(message)
|
|
|
|
for message in to_ack:
|
|
await message.ack()
|
|
|
|
logger.debug(f"[RabbitMQ] Fetched {len(results)} message(s)")
|
|
return results
|
|
|
|
except (AMQPException, OSError) as e:
|
|
logger.error(f"[RabbitMQ] Error while fetching, requeueing batch: {e}", exc_info=True)
|
|
for message in to_ack:
|
|
try:
|
|
await message.nack(requeue=True)
|
|
except Exception:
|
|
pass # channel is gone; the broker requeues unacked messages itself
|
|
raise
|
|
|
|
async def close(self):
|
|
self._closing = True
|
|
if self.connection:
|
|
await self.connection.close()
|
|
logger.info("[RabbitMQ] Connection closed")
|
|
|