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

import asyncio
import importlib
import inspect
import os
import threading
import time
from asyncio import AbstractEventLoop, Future
from typing import Any, Dict

from aio_pika import Message
from aio_pika.abc import AbstractIncomingMessage, AbstractRobustExchange
from opentelemetry import trace
from opentelemetry.trace import Tracer, TraceState
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from pika.exchange_type import ExchangeType

from amqp.adapter.backpressure_handler import BackpressureHandler
from amqp.adapter.cleverthis_service_adapter import AMQMessage, CleverThisServiceAdapter
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.data_response_factory import DataResponseFactory
from amqp.adapter.file_handler import FileHandler, StreamingFileHandler
from amqp.adapter.logging_utils import logging_debug, logging_error, logging_info
from amqp.adapter.serializer import map_as_string, parse_map_string
from amqp.adapter.service_message_factory import ServiceMessageFactory
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import (
    DataMessage,
    DataMessageMagicByte,
    DataResponse,
    MultipartDataMessage,
)
from amqp.rabbitmq.rabbit_mq_client import RabbitMQClient


def collect_trace_info():
    trace_map = {}
    TraceContextTextMapPropagator().inject(
        carrier=trace_map, context=trace.context_api.get_current()
    )
    return trace_map


def get_default_trace_state():
    return TraceState(entries={})


def set_text(carrier, key, value):
    carrier[key] = value


async def invoke_command(
    package_name: str, class_name: str, function_name: str, msg_id: str = "", kwargs=None
) -> Any:
    if kwargs is None:
        kwargs = {}
    try:
        # 1. Dynamically import the package
        logging_info(f"Attempting to import package: {package_name}")
        package = importlib.import_module(package_name)
        logging_info(f"Package '{package_name}' imported successfully.")

        target_callable = None

        if class_name == "__global__":
            # 2. Handle global function invocation
            logging_info(f"Attempting to find global function: {function_name} in {package_name}")
            target_callable = getattr(package, function_name)
            logging_info(f"Global function '{function_name}' found.")
        else:
            # 3. Handle class method invocation
            logging_info(f"Attempting to find class: {class_name} in {package_name}")
            target_class = getattr(package, class_name)
            logging_info(f"Class '{class_name}' found. Instantiating...")

            # Instantiate the class. Assuming a default constructor or one that doesn't
            # require arguments from kwargs for initialization.
            # If the class __init__ needs specific arguments, this part might need adjustment
            # based on how those arguments are provided.
            instance = target_class()
            logging_info(f"Instance of '{class_name}' created.")

            logging_info(f"Attempting to find method: {function_name} in class {class_name}")
            target_callable = getattr(instance, function_name)
            logging_info(f"Method '{function_name}' found.")

        # 4. Invoke the function/method
        logging_info(f"Invoking '{function_name}' with arguments: {kwargs}")
        if inspect.iscoroutinefunction(target_callable):
            result = await target_callable(**kwargs)
            logging_info(f"Asynchronous function '{function_name}' invoked. Result: {result}")
        else:
            result = target_callable(**kwargs)
            logging_info(f"Synchronous function '{function_name}' invoked. Result: {result}")

        response: DataResponse = DataResponseFactory.of(
            id=msg_id,
            response_code=200,
            content_type="text/plain",
            body=result.encode("utf-8"),
            trace_info=collect_trace_info(),
            error=None,
            error_cause=None,
        )
        response._magic = DataMessageMagicByte.RPC_REQUEST.value
        return response

    except ModuleNotFoundError as e:
        logging_error(f"Error: Package '{package_name}' not found. {e}")
        raise
    except AttributeError as e:
        logging_error(
            f"Error: Class '{class_name}' or function/method '{function_name}' not found in '{package_name}'. {e}"
        )
        raise
    except Exception as e:
        logging_error(f"An unexpected error occurred during invocation: {e}")
        raise


class DataMessageHandler:
    def __init__(
        self,
        service_adapter: CleverThisServiceAdapter,
        tracer: Tracer,
        message_factory: DataMessageFactory,
        service_message_factory: ServiceMessageFactory,
        reply_to_exchange: AbstractRobustExchange,
        rabbit_mq_client: RabbitMQClient,
        amq_configuration: AMQConfiguration,
        loop,
        backpressure_handler: BackpressureHandler,
        file_handler: FileHandler = StreamingFileHandler(),
    ):
        self.service_adapter: CleverThisServiceAdapter = service_adapter
        self.tracer = tracer
        self.message_factory: DataMessageFactory = message_factory
        self.service_message_factory: ServiceMessageFactory = service_message_factory
        self.reply_to_exchange: AbstractRobustExchange = reply_to_exchange
        self.rabbit_mq_client: RabbitMQClient = rabbit_mq_client
        self.rabbit_mq_url = amq_configuration.dispatch.rabbit_mq_url
        self.output_dir = amq_configuration.dispatch.download_dir
        self.loop = loop
        self.service_adapter.loop = loop
        # Dictionary to store outstanding Futures
        self.outstanding: Dict[str, asyncio.Future] = {}
        self.reply_to: str | None = None
        self.file_handler: FileHandler = file_handler
        self.backpressure_handler: BackpressureHandler = backpressure_handler

    async def inbound_data_message_callback_as_thread(self, message: AbstractIncomingMessage):
        """
        The callback method that is called when a DataMessage is received from RabbitMQ, starts new Thread.
        """
        threading.Thread(
            target=lambda: asyncio.run(self.inbound_data_message_callback_no_thread(message))
        ).start()

    async def inbound_data_message_callback_no_thread(self, message: AbstractIncomingMessage):
        """
        Receives the DataMessage as bytes from RabbitMQ queue and deserializes it to the DataMessage object
        Ensures Jaeger trace-id propagation.
        Invokes custom handler/mapper/adapter method that corresponds to this message
        :param message:
        :return:
        """
        start = time.time()
        amq_message = await self.reconstitute_data_message(message)
        if amq_message is not None:
            self.ensure_trace_info(amq_message)
            if amq_message.magic() == DataMessageMagicByte.HTTP_REQUEST.value:
                await self.run_http_request_magic(message, amq_message, self.loop)
            elif amq_message.magic() == DataMessageMagicByte.RPC_REQUEST.value:
                await self.run_rpc_request(message, amq_message, self.loop)
            else:
                logging_error(
                    "Unknown MAGIC value: [%s], don't know how to interpret the message",
                    amq_message.magic(),
                )
                asyncio.run_coroutine_threadsafe(message.nack(requeue=False), self.loop)
                return
        else:
            logging_error("######[%s] No valid AMQ Message present.", message.delivery_tag)
        self.record_duration(start, message.delivery_tag)
        asyncio.run_coroutine_threadsafe(message.ack(multiple=False), self.loop)

    async def run_http_request_magic(
        self,
        message: AbstractIncomingMessage,
        amq_message: DataMessage,
        loop: AbstractEventLoop,
    ) -> None:
        """
        Synchronous version of on_message method. This is a wrapper around the async on_message method.
        It also counts parallel executions and handles backpressure.
        :param message: RabbitMQ raw message
        :param amq_message: DataMessage
        :param loop: Event loop
        :return: DataResponse
        """
        # Increment the counter for parallel executions and check resource availability
        self.backpressure_handler.increase_current_load()
        await self.backpressure_handler.check_overload_condition()
        logging_info(
            f"PARALLEL: Current Capacity [{self.backpressure_handler.current_availability}] of [{self.backpressure_handler.max_availability}]"
        )
        if message.reply_to:
            self.reply_to = message.reply_to

        try:
            _a_message = AMQMessage(amq_message, self.rabbit_mq_client.connection)
            _response: DataResponse = await self.service_adapter.on_message(_a_message)
            logging_debug(f"Result in thread: {_response}")
            try:
                await self.send_reply(message.reply_to, _response)
            except Exception as e:
                logging_error(f"Processing message: {e}")
                asyncio.run_coroutine_threadsafe(message.nack(requeue=False), loop=loop)

            if isinstance(amq_message, MultipartDataMessage):
                # remove the downloaded files from the output_dir
                for file_id, filename in amq_message.extra_data.items():
                    try:
                        logging_info(
                            f"[{message.delivery_tag}] Deleting file: {file_id} in {filename}"
                        )
                        os.remove(filename)
                    except OSError as e:
                        logging_error(f"Deleting file {filename}: {e}")
        except Exception as e:
            logging_error(f"Request handler: {e}")

        self.backpressure_handler.decrease_current_load()

    async def run_rpc_request(
        self,
        message: AbstractIncomingMessage,
        amq_message: DataMessage,
        loop: AbstractEventLoop,
    ) -> None:
        """
        Synchronous version of on_message method. This is a wrapper around the async on_message method.
        It also counts parallel executions and handles backpressure.
        :param message: RabbitMQ raw message
        :param amq_message: DataMessage
        :param loop: Event loop
        :return: DataResponse
        """
        # Increment the counter for parallel executions and check resource availability
        self.backpressure_handler.increase_current_load()
        await self.backpressure_handler.check_overload_condition()
        logging_info(
            f"PARALLEL: Current Capacity [{self.backpressure_handler.current_availability}] of [{self.backpressure_handler.max_availability}]"
        )
        if message.reply_to:
            self.reply_to = message.reply_to

        try:
            _a_message = AMQMessage(amq_message, self.rabbit_mq_client.connection)
            _response: DataResponse = await self.on_rpc(_a_message)
            logging_debug(f"Result in thread: {_response}")
            try:
                await self.send_reply(message.reply_to, _response)
            except Exception as e:
                logging_error(f"Processing message: {e}")
                asyncio.run_coroutine_threadsafe(message.nack(requeue=False), loop=loop)

            if isinstance(amq_message, MultipartDataMessage):
                # remove the downloaded files from the output_dir
                for file_id, filename in amq_message.extra_data.items():
                    try:
                        logging_info(
                            f"[{message.delivery_tag}] Deleting file: {file_id} in {filename}"
                        )
                        os.remove(filename)
                    except OSError as e:
                        logging_error(f"Deleting file {filename}: {e}")
        except Exception as e:
            logging_error(f"Request handler: {e}")

        self.backpressure_handler.decrease_current_load()

    async def on_rpc(self, a_message: AMQMessage) -> DataResponse:
        """
        Dynamically invokes a function or a class method from a specified package.

        Args:
            a_message (AMQMessage): The message containing the package, class, and function details.
        Args encoded in the AMQMessage:
            package_name (str): The name of the Python package to import.
            class_name (str): The name of the class within the package.
                              Use '__global__' if the function_name refers to a global function
                              in the package.
            function_name (str): The name of the function or method to invoke.
            kwargs (Dict[str, Any]): A dictionary of keyword arguments to pass to the function/method (in the body).

        Returns:
            Any: The return value of the invoked function or method.

        Raises:
            ModuleNotFoundError: If the specified package cannot be imported.
            AttributeError: If the specified class or function/method does not exist.
            Exception: For any other errors during invocation.
        """
        package_name: str = a_message.method()
        class_name: str = a_message.domain()
        function_name: str = a_message.path()
        kwargs: Dict[str, Any] = parse_map_string(a_message._base64body.decode("utf-8"))

        return await invoke_command(
            package_name, class_name, function_name, a_message.id().as_string(), kwargs
        )

    async def reconstitute_data_message(
        self, message: AbstractIncomingMessage
    ) -> DataMessage | None:
        """
        DataMessage may be serialized as string in separate queue (addressing == 1) or be a MultiPart message with
        files streamed over individual queues, one per file.
        This function uses 'addressing' mode from message headers to determine how to deserialize (reconstitute)
        the original message.
        """
        amq_message: DataMessage | None = None
        delivery_tag = message.delivery_tag
        addressing: int = message.headers.get("clevermicro.addressing", 0)
        if addressing == 0:
            amq_message = DataMessageFactory.from_bytes(message.body)
        else:
            amq_message = await DataMessageFactory.from_stream(
                message.body, rabbitmq_url=self.rabbit_mq_url, loop=self.loop
            )
            logging_info(f"######[{delivery_tag}] AMQ Message from stream: {amq_message}")
            if amq_message is not None:
                extra_data = await self.file_handler.on_message(
                    message=message,
                    json_data=amq_message.body(),
                    rabbitmq_url=self.rabbit_mq_url,
                    loop=self.loop,
                    output_dir=self.output_dir,
                )
                logging_info(f"######[{delivery_tag}] Multipart download result: {extra_data}")
                amq_message = MultipartDataMessage(amq_message, extra_data)
        if amq_message:
            logging_info(
                "######[%s] AMQ Message received, ctag=%s, msgId=%s, %s:%s, traceInfo=%s",
                message.delivery_tag,
                message.consumer_tag,
                amq_message.id(),
                amq_message.method(),
                amq_message.path(),
                map_as_string(amq_message.trace_info()),
            )
        return amq_message

    def ensure_trace_info(self, amq_message: DataMessage):
        propagator = TraceContextTextMapPropagator()
        context = propagator.extract(
            carrier=amq_message.trace_info(), context=trace.context_api.get_current()
        )
        self.tracer.start_span(name="AMQ-Adapter-Consumer", context=context)
        # logging_info("CTX=%s", context)
        # current_span = self.tracer.start_span("AMQ-Adapter.Processing.Request", parent=parent_span)
        # context_with_span = trace.set_span_in_context(current_span, context)
        # logging.info(" [*] CTX2=%s", context_with_span)
        # propagator.inject(context_with_span, amq_message.trace_info, self.set_text)
        # logging.info(" [***] NEW TI=%s CTX=%s SCOPE=%s",
        #             map_as_string(amq_message.trace_info()), context_with_span, context)

    def record_duration(self, start, delivery_tag):
        global last_backpressure_event_time
        last_backpressure_event_time = time.time()
        logging_info(
            "######[%s] Done, single ACK: %sms.",
            delivery_tag,
            last_backpressure_event_time - start,
        )

    async def send_reply(self, reply_to: str, response: DataResponse):
        """
        The Service side code. When the CleverXXX service has output ready, in DataResponse object,
        call this method to return the reply back to the caller
        :param reply_to:connection.channel
        :param response:
        :return:
        """
        if reply_to:
            pika_message: Message = Message(
                body=DataResponseFactory.serialize(response),
                correlation_id=response.id(),
                content_type=(
                    response.content_type()[0]
                    if isinstance(response.content_type(), list)
                    else response.content_type()
                ),
            )
            _res = asyncio.run_coroutine_threadsafe(
                self.reply_to_exchange.publish(message=pika_message, routing_key=reply_to),
                self.loop,
            )
            logging_debug(
                "###### DONE Reply Published to exchg:(%s) To=%s, res=%s",
                self.reply_to_exchange.name,
                reply_to,
                _res,
            )

    async def reply_received_callback(self, message: AbstractIncomingMessage):
        """
        The handler for reply for DataMessage that we sent out earlier -
        the code matches the received response with the original request and passes the response to the caller.
        :param message:
        :return:
        """
        reply_id = message.correlation_id
        delivery_tag = message.delivery_tag
        debug = f"######[{delivery_tag}] (RPC REPLY) env: {message.headers}, props: {message.properties}, BODY: {message.body.decode()}"
        logging_debug(debug)
        # TODO - figure out implementation of this part, as this is sending reply back either to a service,
        #        or to the user
        logging_info(f"OUTSTANDING RECV: {self.outstanding}, size={len(self.outstanding)}")

        future: Future = self.outstanding.pop(reply_id, None)
        # if reply is None:
        # logging.warn(" [W] ######[%s] IGNORING REPLY. No outstanding request for ID %s",
        #             delivery_tag, reply_id)
        # else:
        # result = DataResponseFactory.from_bytes(body)
        # reply.tracing_span().end()
        # TODO - figure out the right response format (not needed for Clever Swarm v0.1)
        # that currently runs Java version, this method here is for the future compatibility to enable
        # internal, service-2-service communication, and the proper response type will be decided than
        # reply.response().set_result(Response.ok(result.body, content_type).build())
        if future:
            response: DataResponse = DataResponseFactory.from_bytes(message.body)
            # Complete the Future with the response data
            future.set_result(response.body())

        await message.ack(multiple=False)

    async def send_rpc(self, route, message: DataMessage) -> Future:
        """
        Sends an RPC message to the specified route and returns a Future that will be completed
        when the response is received.

        :param route: AMQRoute object containing exchange and routing information
        :param message: DataMessage to be sent
        :return: Future that will be completed with the response payload
        """
        # Create a Future to be completed when the response is received
        future = asyncio.Future()

        # Store the future in the outstanding dictionary using the message ID as the key
        message_id = message.id().as_string()
        self.outstanding[message_id] = future

        # Create a Pika message with the correlation ID set to the message ID
        pika_message = Message(
            body=DataMessageFactory.serialize(message),
            correlation_id=message_id,
            reply_to=self.rabbit_mq_client.router.get_reply_to_queue_name(),
            content_type=(
                message.content_type()[0]
                if isinstance(message.content_type(), list)
                else message.content_type()
            ),
        )

        # Get the exchange from the route
        rpc_exchange_name = route.exchange
        exchange: AbstractRobustExchange = await self.rabbit_mq_client.channel.declare_exchange(
            name=rpc_exchange_name, type=ExchangeType.topic
        )
        # Publish the message to the exchange with the component name as the routing key
        await exchange.publish(message=pika_message, routing_key=route.component_name)

        logging_debug(
            "###### RPC Request Published to exchg:(%s) routing_key=%s, msgId=%s",
            rpc_exchange_name,
            route.component_name,
            message_id,
        )

        return future

    async def send_command(self, route, message: DataMessage) -> bool:
        """
        Sends an RPC message to the specified route and returns a Future that will be completed
        when the response is received.

        :param route: AMQRoute object containing exchange and routing information
        :param message: DataMessage to be sent
        :return: Future that will be completed with the response payload
        """
        rmq_body = DataMessageFactory.serialize(message)
        cid = message.id().as_string()
        ctype = (
            message.content_type()[0]
            if isinstance(message.content_type(), list)
            else message.content_type()
        )
        print(
            f"Sending RPC command to route: {route}, message ID: {cid} with body: {rmq_body} correlation ID: {cid}, content type: {ctype}"
        )
        # Create a Pika message with the correlation ID set to the message ID
        pika_message = Message(
            body=rmq_body,
            correlation_id=cid,
            content_type=ctype,
        )
        print(f"SENDING RPC COMMAND: {pika_message.body} with ID: {message.id().as_string()}")

        # Get the exchange from the route
        rpc_exchange_name = route.exchange
        exchange: AbstractRobustExchange = await self.rabbit_mq_client.channel.declare_exchange(
            name=rpc_exchange_name, type=ExchangeType.topic
        )
        # Publish the message to the exchange with the component name as the routing key
        await exchange.publish(message=pika_message, routing_key=route.component_name)

        logging_info(
            "###### RPC Request Published to exchg:(%s) routing_key=%s, msgId=%s",
            rpc_exchange_name,
            route.component_name,
            message.id().as_string(),
        )

        return True