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 self.maximum_availability = -1 # ======================================================================= # ======================================================================= # ============== 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_availability: int, maximum: int = -1): """ Send Backpressure event to the Management Service. The Backpressure logic determines the type of the event based on the current_availability and maximum values. - OVERLOAD event is triggered immediately when the current_availability value is below threshold typically set to 10% of the maximum value, or 0 if no maximum is provided. - IDLE event is triggered when the current_availability value exeeds 80% of the maximum value or is greater than 0 if no maximum is provided. - UPDATE event is sent if neither OVERLOAD nor IDLE conditions are met, indicating that the backpressure is within acceptable limits. :param current_availability: Currently available capacity. :param maximum: Maximum value of the backpressure metric. If undefined or set to lt 0 (e.g. -1), the maximum is than determined from max value of current_availability seen over the time. """ if self.backpressure_handler is not None: # watch maximum & current value to always have some valid maximum reference if maximum < 0: if current_availability > self.maximum_availability: self.maximum_availability = current_availability else: if maximum > self.maximum_availability: self.maximum_availability = maximum asyncio.run_coroutine_threadsafe( self.backpressure_handler.update_backpressure_value( current_availability, maximum if maximum > 0 else self.maximum_availability ), loop=self.backpressure_handler.loop, ) 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)