332 lines
12 KiB
Python
332 lines
12 KiB
Python
import asyncio
|
|
import json
|
|
import os
|
|
import signal
|
|
|
|
import aio_pika
|
|
from aio_pika.exceptions import AMQPException
|
|
|
|
import metrics
|
|
from db import db_manager
|
|
from device_token_repository import database_lookup_by_user_id, database_lookup_by_uuid, remove_inactive_push_token
|
|
from secret_manager import cleanup_secret_files, return_credentials
|
|
from secret_manager.fernet_encryption import decrypt_token
|
|
from send_notification import send_notification
|
|
from simple_logger_handler import setup_logger
|
|
|
|
logger = setup_logger(__name__)
|
|
|
|
|
|
rmq_username = return_credentials("/etc/secrets/rmq_username")
|
|
rmq_password = return_credentials("/etc/secrets/rmq_password")
|
|
rmq_host = os.getenv("BACKEND_PN_RMQ_HOST", "localhost:30672")
|
|
rmq_vhost = os.getenv("BACKEND_PN_RMQ_VHOST", "app_notifications")
|
|
rmq_exchange = os.getenv("BACKEND_PN_RMQ_EXCHANGE", "app_notifications")
|
|
|
|
RABBITMQ_URL = f"amqp://{rmq_username}:{rmq_password}@{rmq_host}/{rmq_vhost}"
|
|
|
|
|
|
class RabbitMQConsumer:
|
|
def __init__(self, url=RABBITMQ_URL, db_manager=db_manager, exchange_name=rmq_exchange,rmq_host=rmq_host):
|
|
self.url = url
|
|
self.rmq_host = rmq_host
|
|
self.exchange_name = exchange_name
|
|
self.connection: aio_pika.RobustConnection | None = None
|
|
self.channel: aio_pika.RobustChannel | None = None
|
|
self.exchange: aio_pika.Exchange | None = None
|
|
self.queue: aio_pika.Queue | None = None
|
|
self._closing = False
|
|
self.db_manager = db_manager
|
|
self.queue_name = "notifications"
|
|
self.retry_queue_name = "notifications_retry"
|
|
self.dlq_queue_name = "notifications_dlq"
|
|
self.max_retries = 5
|
|
self.retry_ttl = 5 * 60 * 10000
|
|
logger.debug(f"[RabbitMQ] Initialized consumer for exchange {self.exchange_name}")
|
|
|
|
|
|
async def connect(self):
|
|
logger.info(f"[RabbitMQ] Connecting to {self.rmq_host}...")
|
|
self.connection = await aio_pika.connect_robust(self.url)
|
|
self.channel = await self.connection.channel()
|
|
self.exchange = await self.channel.declare_exchange(
|
|
self.exchange_name, aio_pika.ExchangeType.TOPIC, durable=True
|
|
)
|
|
self.queue = await self.channel.declare_queue(self.queue_name, durable=True)
|
|
await self.queue.bind(self.exchange, routing_key="notify.user.*")
|
|
logger.debug(f"[RabbitMQ] Bound queue '{self.queue_name}' to exchange '{self.exchange_name}'")
|
|
|
|
retry_queue_args = {
|
|
"x-message-ttl": self.retry_ttl,
|
|
"x-dead-letter-exchange": self.exchange_name,
|
|
"x-dead-letter-routing-key": "notify.user.retry",
|
|
}
|
|
logger.debug(f"[RabbitMQ] Declaring retry queue '{self.retry_queue_name}' with args: {retry_queue_args}")
|
|
|
|
self.retry_queue = await self.channel.declare_queue(
|
|
self.retry_queue_name, durable=True, arguments=retry_queue_args
|
|
)
|
|
logger.debug(f"[RabbitMQ] Retry queue '{self.retry_queue_name}' declared.")
|
|
|
|
await self.channel.declare_queue(self.dlq_queue_name, durable=True)
|
|
logger.debug(f"[RabbitMQ] DLQ queue '{self.dlq_queue_name}' declared.")
|
|
|
|
metrics.start_metrics_server(port=9000)
|
|
logger.debug("[Metrics] Prometheus metrics server started on port 9000.")
|
|
|
|
self._metric_tasks = [
|
|
asyncio.create_task(metrics.update_queue_gauge(self.channel, self.queue_name)),
|
|
asyncio.create_task(metrics.update_queue_gauge(self.channel, self.retry_queue_name)),
|
|
asyncio.create_task(metrics.update_queue_gauge(self.channel, self.dlq_queue_name)),
|
|
]
|
|
logger.debug(f"[Metrics] Gauge update tasks created for queues: {self.queue_name}, {self.retry_queue_name}, {self.dlq_queue_name}")
|
|
|
|
logger.info("[RabbitMQ] Connected, exchange, retry, and DLQ queues ready.")
|
|
|
|
async def send_message_to_dlq(self, uuid: str, message: aio_pika.IncomingMessage):
|
|
data = json.loads(message.body.decode())
|
|
data["uuid"] = uuid
|
|
await self.channel.default_exchange.publish(
|
|
aio_pika.Message(
|
|
body=json.dumps(data).encode(),
|
|
content_type="application/json",
|
|
delivery_mode=aio_pika.DeliveryMode.PERSISTENT
|
|
),
|
|
routing_key=self.dlq_queue_name
|
|
)
|
|
logger.warning(f"[RabbitMQ] Message sent to DLQ: {data['uuid']}")
|
|
metrics.MSG_FAILED.inc()
|
|
|
|
async def send_message_to_retry_queue(self, uuid: str, message: aio_pika.IncomingMessage):
|
|
"""
|
|
Publishes a message to the retry queue with an incremented retry counter.
|
|
|
|
Behaviour:
|
|
- Each message is retried up to `max_retries` times.
|
|
- After exceeding the limit, the message is forwarded to the DLQ.
|
|
- The UUID of the intended device is embedded in the message payload.
|
|
|
|
Args:
|
|
uuid: Unique identifier for the device or token associated with the message.
|
|
message: The original RabbitMQ message being retried.
|
|
"""
|
|
data = json.loads(message.body.decode())
|
|
data["uuid"] = uuid
|
|
retry_count = data.get("retry_count", 0) + 1
|
|
data["retry_count"] = retry_count
|
|
|
|
metrics.MSG_RETRY.labels(
|
|
queue_name=self.retry_queue_name,
|
|
uuid=uuid,
|
|
retry_count=retry_count
|
|
).inc()
|
|
logger.debug(f"[RabbitMQ] Metrics updated for retry: uuid={uuid}, retry_count={retry_count}")
|
|
|
|
if retry_count > self.max_retries:
|
|
logger.debug(f"[RabbitMQ] Retry count exceeded max_retries={self.max_retries}, sending to DLQ")
|
|
await self.send_message_to_dlq(uuid, message)
|
|
return
|
|
await self.channel.default_exchange.publish(
|
|
aio_pika.Message(
|
|
body=json.dumps(data).encode(),
|
|
content_type="application/json",
|
|
delivery_mode=aio_pika.DeliveryMode.PERSISTENT
|
|
),
|
|
routing_key=self.retry_queue_name
|
|
)
|
|
logger.info(f"[RabbitMQ] Message requeued for retry #{retry_count}")
|
|
|
|
async def handle_message(self, message: aio_pika.IncomingMessage):
|
|
"""
|
|
Processes a single RabbitMQ message containing push notification data.
|
|
|
|
Message routing:
|
|
- Standard messages use the routing key: notify.user.<user_id>
|
|
- Retry messages use the routing key: notify.user.retry and include a device UUID
|
|
|
|
Behaviour:
|
|
1. Parses and validates the message payload.
|
|
2. Fetches encrypted push tokens from the database by UUID or user ID.
|
|
3. Decrypts tokens and sends notifications via the Expo Push API.
|
|
4. Delegates delivery result handling to `validate_delivery`.
|
|
|
|
Error handling:
|
|
- Invalid JSON → message discarded (not requeued)
|
|
- RabbitMQ transport errors → message requeued
|
|
- Unexpected exceptions → raised for shutdown or external handling
|
|
"""
|
|
if self._closing:
|
|
logger.debug("[RabbitMQ] Skipping message because consumer is closing")
|
|
return
|
|
|
|
async with message.process():
|
|
try:
|
|
data = json.loads(message.body.decode())
|
|
logger.debug(f"[RabbitMQ] Received message: {data}")
|
|
|
|
if message.routing_key in [self.dlq_queue_name, "notify.dlq"]:
|
|
logger.info(f"[RabbitMQ] Message in DLQ, skipping processing: {data}")
|
|
return
|
|
|
|
uuid = data.get("uuid")
|
|
|
|
if uuid:
|
|
encrypted_tokens = await database_lookup_by_uuid(uuid, db_manager)
|
|
else:
|
|
encrypted_tokens = await database_lookup_by_user_id(message.routing_key, db_manager)
|
|
|
|
if not encrypted_tokens:
|
|
logger.warning(f"[RabbitMQ] No push tokens found for routing key {message.routing_key}")
|
|
return
|
|
|
|
token_map = {row["uuid"]: row["token"].decode() for row in encrypted_tokens}
|
|
for uuid, token in token_map.items():
|
|
decrypted_token = decrypt_token(token)
|
|
token_map[uuid] = decrypted_token
|
|
response = await send_notification(message=data,push_tokens=token_map)
|
|
|
|
# Validate delivery - any errors here should NOT prevent acknowledgment
|
|
# because notifications were already sent. Instead, we handle retries/DLQ
|
|
# by publishing new messages.
|
|
try:
|
|
await self.validate_delivery(response,message)
|
|
except Exception as validation_error:
|
|
logger.error(f"[RabbitMQ] Error during delivery validation: {validation_error}", exc_info=True)
|
|
metrics.MSG_FAILED.inc()
|
|
|
|
except json.JSONDecodeError as e:
|
|
logger.error(f"[RabbitMQ] Bad message, discarding: {e}", exc_info=True)
|
|
await message.nack(requeue=False)
|
|
except AMQPException as e:
|
|
logger.error(f"[RabbitMQ] RabbitMQ error: {e}", exc_info=True)
|
|
await message.nack(requeue=True)
|
|
except Exception as e:
|
|
logger.critical(f"[RabbitMQ] Fatal error: {e}", exc_info=True)
|
|
metrics.MSG_FAILED.inc()
|
|
|
|
async def validate_delivery(self,response,message: aio_pika.IncomingMessage):
|
|
"""
|
|
Validates delivery results returned by the Expo Push API after a message is sent.
|
|
|
|
This function inspects both the transport-level and API-level delivery statuses and
|
|
takes appropriate follow-up actions:
|
|
|
|
- Successful deliveries are logged and counted as published.
|
|
- API-level errors trigger token expiration or redirection to the DLQ.
|
|
- Failures (e.g., network or unknown issues) are requeued for retry.
|
|
|
|
Right now the <push-receipt-id> is not stored for further validation but needs to be implemented if the userbase grows.
|
|
|
|
Expected response structure:
|
|
{
|
|
"<uuid>": {
|
|
"status": "ok" | "error" | "failure",
|
|
"data": {
|
|
"data": [
|
|
{
|
|
"status": "ok" | "error",
|
|
"id": "<push-receipt-id>",
|
|
"details": { "error": "<ErrorMessage>" } #Only exists if there is an error
|
|
}
|
|
]
|
|
}
|
|
}
|
|
}
|
|
|
|
See: https://docs.expo.dev/push-notifications/sending-notifications/
|
|
"""
|
|
for uuid, result in response.items():
|
|
try:
|
|
status = result.get("status")
|
|
data_list = result.get("data", {}).get("data", [])
|
|
|
|
api_status = None
|
|
api_error = None
|
|
if data_list and isinstance(data_list, list) and len(data_list) > 0:
|
|
first_item = data_list[0]
|
|
if isinstance(first_item, dict):
|
|
api_status = first_item.get("status")
|
|
api_error = first_item.get("details", {}).get("error")
|
|
|
|
logger.debug(f"[RabbitMQ] Processing delivery for uuid={uuid}, status={status}, api_status={api_status}")
|
|
|
|
if status == "ok" and api_status == "ok":
|
|
logger.info(f"[RabbitMQ] Notification delivered successfully to {uuid}")
|
|
metrics.MSG_PUBLISHED.inc()
|
|
logger.debug(f"[RabbitMQ] Metrics updated for published message: uuid={uuid}")
|
|
|
|
|
|
elif status == "ok" and api_status == "error":
|
|
if api_error == "DeviceNotRegistered":
|
|
expired = await remove_inactive_push_token(uuid, self.db_manager)
|
|
if expired:
|
|
logger.info(f"[RabbitMQ] Device no longer registered for uuid {uuid}, marked expired")
|
|
else:
|
|
logger.error(f"[RabbitMQ] Failed expiring token for uuid: {uuid}")
|
|
else:
|
|
await self.send_message_to_dlq(uuid, message)
|
|
logger.debug(f"[RabbitMQ] Message sent to DLQ due to API error: uuid={uuid}, error={api_error}")
|
|
|
|
|
|
elif status == "error":
|
|
await self.send_message_to_dlq(uuid, message)
|
|
logger.debug(f"[RabbitMQ] Message sent to DLQ due to status=error: uuid={uuid}")
|
|
|
|
elif status == "failure":
|
|
await self.send_message_to_retry_queue(uuid, message)
|
|
logger.debug(f"[RabbitMQ] Message requeued for retry due to status=failure: uuid={uuid}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"[RabbitMQ] Error processing delivery validation for uuid={uuid}: {e}", exc_info=True)
|
|
try:
|
|
await self.send_message_to_dlq(uuid,message)
|
|
except Exception as dlq_error:
|
|
logger.error(f"[RabbitMQ] Failed to send to DLQ for uuid={uuid}: {dlq_error}")
|
|
|
|
|
|
async def consume(self):
|
|
if self._closing:
|
|
logger.debug("[RabbitMQ] Consumer is closing, not starting consume")
|
|
return
|
|
if not self.queue:
|
|
raise RuntimeError("[RabbitMQ] Queue not initialized")
|
|
await self.queue.consume(self.handle_message, no_ack=False)
|
|
logger.info("[RabbitMQ] Consuming messages...")
|
|
|
|
async def close(self):
|
|
self._closing = True
|
|
if self.connection:
|
|
await self.connection.close()
|
|
logger.info("[RabbitMQ] Connection closed")
|
|
|
|
|
|
|
|
async def main():
|
|
logger.info("[APP] Starting application")
|
|
await db_manager.connect()
|
|
consumer = RabbitMQConsumer(db_manager=db_manager)
|
|
await consumer.connect()
|
|
await consumer.consume()
|
|
stop_event = asyncio.Event()
|
|
|
|
SECRET_PATHS = frozenset({
|
|
"/etc/secrets/encryption_key",
|
|
"/etc/secrets/db_username",
|
|
"/etc/secrets/db_password",
|
|
"/etc/secrets/rmq_username",
|
|
"/etc/secrets/rmq_password"
|
|
})
|
|
cleanup_secret_files(SECRET_PATHS)
|
|
|
|
for sig in (signal.SIGINT, signal.SIGTERM):
|
|
asyncio.get_running_loop().add_signal_handler(sig, stop_event.set)
|
|
|
|
await stop_event.wait()
|
|
logger.info("[APP] Stopping application")
|
|
await consumer.close()
|
|
await db_manager.close()
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|
|
|