freeleaps-service-hub/apps/devops/app/backend/infra/rabbitmq/async_subscriber.py
2025-07-30 10:50:22 +08:00

84 lines
3.2 KiB
Python

from asyncio import AbstractEventLoop
from app.common.log.module_logger import ModuleLogger
import json
import asyncio
from .async_client import AsyncMQClient
class AsyncMQSubscriber(AsyncMQClient):
def __init__(self, channel_name: str) -> None:
super().__init__(channel_name=channel_name)
self.process_callable = None
self.routing_key = self.channel_name
self.consumer_callbacks = {}
self.consumer_callbacks_lock = asyncio.Lock()
self.module_logger = ModuleLogger(sender_id="AsyncMQSubscriber")
async def process_incoming_message(self, message):
"""Processing incoming message from RabbitMQ"""
await message.ack()
body = message.body
if body:
async with self.consumer_callbacks_lock:
for registry_key, callback_info in self.consumer_callbacks.items():
try:
await callback_info["method"](
registry_key, json.loads(body), callback_info["args"]
)
except Exception as err:
await self.module_logger.log_exception(
exception=err,
text=f"Error processing message for consumer '{registry_key}'",
)
async def subscribe(self, max_retries=10, event_loop: AbstractEventLoop = None):
"""Attempts to bind and consume messages, with retry mechanism."""
retries = 0
while retries < max_retries:
try:
await self.bind(max_retries=5, event_loop=event_loop)
await self.queue.consume(
no_ack=False, exclusive=True, callback=self.process_incoming_message
)
break
except Exception as e:
await self.module_logger.log_exception(
exception=e,
text=f"Failed to subscribe at {retries} time, will retry",
)
retries += 1
await asyncio.sleep(5)
else:
await self.module_logger.log_exception(
exception=ConnectionError(
f"Exceeded max retries ({max_retries}) for subscription."
),
text=f"Subscription failed for {self.channel_name} after {max_retries} attempts.",
)
async def register_consumer(
self,
registry_key: str,
callback_method,
args: dict,
):
"""Register a consumer callback with a unique key."""
async with self.consumer_callbacks_lock:
self.consumer_callbacks[registry_key] = {
"method": callback_method,
"args": args,
}
async def unregister_consumer(
self,
registry_key: str,
):
"""Unregister a consumer callback by its key."""
async with self.consumer_callbacks_lock:
if registry_key in self.consumer_callbacks:
del self.consumer_callbacks[registry_key]
async def clear_all_consumers(self):
"""Unregister all consumer callbacks."""
async with self.consumer_callbacks_lock:
self.consumer_callbacks.clear()