import asyncio import logging import logging.config import os from asyncio import Future from threading import Thread from typing import Optional from aio_pika.abc import AbstractRobustExchange from amqp.adapter.amq_route_factory import AMQRouteFactory from amqp.adapter.backpressure_handler import BackpressureHandler from amqp.adapter.consul_global_id_generator import get_unique_instance_id from amqp.adapter.data_message_factory import DataMessageFactory from amqp.adapter.logging_utils import logging_info, logging_warning from amqp.adapter.service_message_factory import ServiceMessageFactory from amqp.config.amq_configuration import AMQConfiguration from amqp.model.model import AMQRoute, DataMessage from amqp.rabbitmq.rabbit_mq_client import RabbitMQClient from amqp.router.router_consumer import RouterConsumer from amqp.service.amq_message_handler import DataMessageHandler from amqp.service.tracing import TracingConfig, initialize_telemetry logging.basicConfig(level=logging.DEBUG) logging.info(f"CWD = {os.getcwd()}") logging_conf_path = os.path.join(os.getcwd(), "amqp", "service", "logging.conf") if os.path.exists(logging_conf_path): logging.config.fileConfig(logging_conf_path) else: logging_conf_path = os.path.join(os.getcwd(), "logging.conf") if os.path.exists(logging_conf_path): logging.config.fileConfig(logging_conf_path) class AMQService: CLEVER_AMQP_CLIENT_VERSION = os.getenv("CLEVER_AMQP_CLIENT_VERSION") MAX_WAIT = 300 # defaults to 5 minutes (300 seconds) def __init__(self, amq_configuration: AMQConfiguration, service_adapter=None): logging.info("***********************************************") logging.info("************** AMQ Service ***************") logging.info("***********************************************") self.loop = asyncio.new_event_loop() self.amq_configuration = amq_configuration self.service_adapter = service_adapter self.unique_id = get_unique_instance_id(amq_configuration.amq_adapter) amq_configuration.amq_adapter.generator_id = self.unique_id self.data_message_factory = DataMessageFactory.get_instance( self.unique_id & DataMessageFactory.MACHINE_ID_MASK ) self.router = RouterConsumer(self.amq_configuration) self.rabbit_mq_client = RabbitMQClient(self.amq_configuration, self.router) self.tracer = initialize_telemetry( tracing_config=TracingConfig( service_name=amq_configuration.amq_adapter.service_name, service_version=os.getenv("SERVICE_VERSION", "unknown"), ) ) self.amq_data_message_handler = None self.service_message_factory = ServiceMessageFactory( self.amq_configuration.amq_adapter.service_name ) self.backpressure_thread: Optional[Thread] = None self.backpressure_handler: Optional[BackpressureHandler] = None # ======================================================================= # ======================================================================= # ============== P u b l i c A P I M e t h o d s =================== # ======================================================================= # ======================================================================= def run(self): """ Start the AMQ service, initializing RabbitMQ connection and message handlers. """ # self.init(loop) will initialize RabbitMQ connection asyncio.set_event_loop(self.loop) self.loop.run_until_complete(self.init(self.loop, self.service_adapter)) logging.info("***********************************************") logging.info("****** AMQ Service Init Complete *********") logging.info("***********************************************") self.loop.run_forever() def backpressure(self, current: int, maximum: int = -1): """ Send Backpressure event to the Management Service. The Backpressure logic determines the type of the event based on the current and maximum values. - OVERLOAD event is triggered immediately when the current value exceeds the maximum threshold, which is typically set to 80% of the maximum value. - IDLE event is triggered when the current value drops below 20% of the maximum value and remains there for a IDLE_WINDOW period of time (see BackpressureHandler). - UPDATE event is sent if neither OVERLOAD nor IDLE conditions are met, indicating that the backpressure is within acceptable limits. :param current: Current value of the backpressure metric. :param maximum: Maximum value of the backpressure metric. If undefined or set to lt 0, the maximum is set according to the configuration value 'cm.backpressure.threshold'. """ if self.backpressure_handler is not None: if maximum < 0: maximum = self.amq_configuration.backpressure.threshold_threads self.backpressure_handler.update_backpressure_value(current, maximum) def rpc(self, service_name: str, fq_method_name: str, **kwargs) -> Future: """ Send a Remote Procedure Call (RPC) message to the AMQ service. :param service_name: The service name to receive the RPC message. :param fq_method_name: The fully qualified method name to execute the RPC message. :param kwargs: Additional keyword arguments to be passed to the RPC method. :return: A Future that resolves when the RPC response is received. """ route: Optional[AMQRoute] = self.router.find_route_by_service(service_name) if route is not None: message: DataMessage = self.data_message_factory.create_rpc_message( fq_method_name=fq_method_name, swarm_id=self.amq_configuration.amq_adapter.swarm_task_id, **kwargs, ) result: Future = self.amq_data_message_handler.send_rpc(route, message) self.set_outstanding_future(message, result) return result result = asyncio.Future() result.set_exception( NameError(f"No route exists. Service '{service_name}' not found in routing table.") ) return result def get_unique_instance_id(self) -> int: """ Provides cluster-wide unique instance ID for this AMQ service instance. This ID is used to identify the service instance in the cluster and is unique across all instances. IDs are managed by shared Consul KV store, ensuring that even if multiple instances of the service are running, they will each have unique ID. """ return self.unique_id # ======================================================================= # ====================== Internal Methods ============================ # ======================================================================= async def init(self, loop, service_adapter): _reply_to_exchange: AbstractRobustExchange = await self.rabbit_mq_client.init(loop) self.backpressure_handler = BackpressureHandler( channel=self.rabbit_mq_client.channel, loop=loop, config=self.amq_configuration, ) self.amq_data_message_handler = DataMessageHandler( service_adapter, self.tracer, self.data_message_factory, self.service_message_factory, _reply_to_exchange, self.rabbit_mq_client, self.amq_configuration, loop=loop, backpressure_handler=self.backpressure_handler, ) await self.register_routes() await self.rabbit_mq_client.request_routes() logging_info("Starting backpressure monitor thread") self.backpressure_thread = self.backpressure_handler.start_backpressure_monitor() def cleanup(self): logging_info("AMQService.cleanup() rmqc=%s", self.rabbit_mq_client) if self.rabbit_mq_client: self.rabbit_mq_client.cleanup() async def reinitialize_if_disconnected(self): if not self.rabbit_mq_client.channel or self.rabbit_mq_client.channel.is_closed: self.rabbit_mq_client = RabbitMQClient(self.amq_configuration, self.router) await self.rabbit_mq_client.init(loop=self.loop) await self.register_routes() await self.rabbit_mq_client.request_routes() async def register_routes(self) -> AbstractRobustExchange | None: route_mapping = self.amq_configuration.amq_adapter.route_mapping if AMQRouteFactory.validate(route_mapping): try: await self.rabbit_mq_client.register_inbound_routes( route_mapping, self.amq_data_message_handler.inbound_data_message_callback_no_thread, ) return await self.rabbit_mq_client.register_data_reply_callback( self.amq_data_message_handler.reply_received_callback ) except Exception as e: raise RuntimeError(e) elif route_mapping: logging_warning( "route [%s] failed validation, no routes were registered.", route_mapping, ) return None def set_outstanding_future(self, message: DataMessage, future: Future): self.amq_data_message_handler.outstanding[str(message.id())] = future def remove_outstanding_future(self, message_id: str): self.amq_data_message_handler.outstanding.pop(message_id, None)