amq-adapter-python/amqp/service/amq_service.py

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.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.data_message_factory = DataMessageFactory.get_instance(
            self.amq_configuration.amq_adapter.generator_id
        )
        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

    # =======================================================================
    # =======================================================================
    # ====================== 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 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

    # =======================================================================
    # ======================  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)