import asyncio import os import threading import time from asyncio import AbstractEventLoop, Future from typing import 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 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 from amqp.adapter.service_message_factory import ServiceMessageFactory from amqp.config.amq_configuration import AMQConfiguration from amqp.model.model import ( DataMessage, 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 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 == DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST: await self.run_http_request_magic(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 self.backpressure_handler.increase_parallel_executions() logging_info( f"PARALLEL: Current Capacity [{self.backpressure_handler.current_parallel_executions}] of [{self.backpressure_handler.parallel_workers}]" ) if message.reply_to: self.reply_to = message.reply_to try: await self.backpressure_handler.check_overload_condition() _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_parallel_executions() async def reconstitute_data_message( self, message: AbstractIncomingMessage ) -> DataMessage | None: 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_data_message_time last_data_message_time = time.time() logging_info( "######[%s] Done, single ACK: %sms.", delivery_tag, last_data_message_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 reply for DataMessage that we sent out earlier - must match the received response with the original request and respond 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)