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 AsyncIterator, 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_error, 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
        self.future: Future = asyncio.Future()

    # =======================================================================
    # =======================================================================
    # ============== 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.future.set_result(True)
        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 = asyncio.run_coroutine_threadsafe(
                self.amq_data_message_handler.send_rpc(route, message),
                loop=self.backpressure_handler.loop,
            )
            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 command(self, service_name: str, fq_method_name: str, **kwargs) -> bool:
        """
        Send a unidirectional Command 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:
            m: DataMessage = self.data_message_factory.create_rpc_message(
                fq_method_name=fq_method_name,
                swarm_id=self.amq_configuration.amq_adapter.swarm_task_id,
                args=kwargs,
            )
            asyncio.run_coroutine_threadsafe(
                self.amq_data_message_handler.send_command(route, m),
                loop=self.backpressure_handler.loop,
            )
            return True
        return False

    async def data_message_generator(self, route: AMQRoute) -> AsyncIterator[DataMessage]:
        """
        An async iterator that blocks until messages arrive, processes them,
        and continues in a loop.

        :param queue_name: The name of the queue to consume messages from
        :yield: The processed DataResponse for each message
        """
        # Get a channel from the connection
        channel = self.rabbit_mq_client.channel
        queue_args = (
            self.rabbit_mq_client.DLQ_ARGS if self.amq_configuration.dispatch.use_dlq else None
        )
        queue = await channel.declare_queue(
            route.queue,
            durable=True,
            exclusive=False,
            auto_delete=False,
            arguments=queue_args,
        )
        await queue.bind(route.exchange, route.key)
        # Start consuming without a callback
        async with queue.iterator() as queue_iter:
            async for message in queue_iter:
                try:
                    # Process the message using the same logic as in inbound_data_message_callback_no_thread
                    data_message: DataMessage = (
                        await self.amq_data_message_handler.reconstitute_data_message(message)
                    )

                    success = yield data_message
                    if not success:
                        logging_warning(f"Message processing failed for {data_message.id()}.")
                        # Nack the message without requeuing in case of failure
                        await message.nack(requeue=False)
                    else:
                        await message.ack()
                except Exception as e:
                    logging_error(f"Error processing message: {e}")
                    # Nack the message without requeuing in case of error
                    await message.nack(requeue=False)

    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)