import asyncio import logging import time from asyncio import Future from typing import Dict import pika from opentelemetry import trace from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator from opentelemetry.trace import Tracer, TraceState from pika.channel import Channel from amqp.adapter.amq_message_factory import AMQMessageFactory from amqp.adapter.amq_response_factory import AMQResponseFactory from amqp.adapter.clever_swarm_adapter import CleverSwarmAdapter from amqp.model.model import AMQResponse, AMQMessage from amqp.router.router_base import AMQ_REPLY_TO_EXCHANGE from amqp.router.serializer import map_as_string 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 AMQServiceMessageHandler: def __init__(self, service_adapter: CleverSwarmAdapter, tracer: Tracer, message_factory: AMQMessageFactory): self.http_adapter: CleverSwarmAdapter = service_adapter self.tracer = tracer self.message_factory: AMQMessageFactory = message_factory # Dictionary to store outstanding Futures self.outstanding: Dict[str, asyncio.Future] = {} def inbound_message_callback_http(self, channel: Channel, method: pika.spec.Basic.Deliver, properties: pika.spec.BasicProperties, body: bytes): """ Receives the message as bytes from RabbitMQ queue and deserializes it to the AMQMessage Ensures Jaeger trace-id propagation. In loose coupling mode, invokes appropriate adapter (in v0.1 only HTTPAdapter is supported). In tight coupling mode, invokes custom handler/mapper/adapter method that :param body: :param properties: :param method: :param channel: :return: """ start = time.time() delivery_tag = method.delivery_tag debug = f" [DBG] ######[{delivery_tag}] (AMQ MESSAGE), env: {method}, props: {properties}" logging.info(debug) amq_message: AMQMessage = AMQMessageFactory.from_bytes(body) logging.info(" [*] ######[%s] AMQ Message received, ctag=%s, msgId=%s, traceInfo=%s", delivery_tag, method.consumer_tag, amq_message.id, map_as_string(amq_message.trace_info)) 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) if amq_message.magic == AMQMessageFactory.MAGIC_BYTE_HTTP_REQUEST: response: AMQResponse = self.http_adapter.on_message(amq_message) else: logging.error(" [E] Unknown MAGIC value: [%s], don't know how to interpret the message", amq_message.magic) duration = time.time() - start logging.info(" [x] ######[%s] Done, single ACK: %sms.", delivery_tag, duration) channel.basic_ack(delivery_tag, False) # if "___NACK___" in str(amq_message.body()): # logging.info(" [x] ######[%s] Done, single NACK: %sms.", delivery_tag, duration) # channel.basic_nack(delivery_tag, False, False) # else: # logging.info(" [x] ######[%s] Done, single ACK: %sms.", delivery_tag, duration) # channel.basic_ack(delivery_tag, False) def on_http_response(self, delivery, deliveryTag, channel, http_response): """ In loose coupling mode, receives the HTTP response to the REST request the HTTPAdapter made to the associated CleverXXX service. Convert the response to AMQResponse message and use RabbitMQ RPC reply-to mechanism to deliver this response back to API Gateway side AMQProducer. :param delivery: :param deliveryTag: :param channel: :param http_response: :return: """ response = AMQResponse( delivery.properties.correlation_id, http_response.status_code, http_response.headers.get("Content-Type"), http_response.body, http_response.status_message, "", collect_trace_info() ) try: self.send_reply(delivery.properties.reply_to, deliveryTag, channel, response) except Exception as e: raise RuntimeError(e) def send_reply(self, reply_to: str, delivery_tag: str, channel: Channel, response: AMQResponse): """ The Service side code. When the CleverXXX service has output ready, in AMQResponse object, call this method to return the reply back to the caller :param reply_to: :param delivery_tag: :param channel: :param response: :return: """ if reply_to: logging.info(" [*] ######[%s] AMQ Message Reply, To=%s, msgId=%s", delivery_tag, reply_to, response.id) reply_props = pika.BasicProperties( correlation_id=response.id, content_type=response.content_type ) channel.basic_publish( exchange=AMQ_REPLY_TO_EXCHANGE, routing_key=reply_to, body=AMQResponseFactory.serialize(response), properties=reply_props, ) def reply_received_callback(self, channel: Channel, method: pika.spec.Basic.Deliver, properties: pika.spec.BasicProperties, body: bytes): """ The producer/client (API Gateway) facing response_message recipient - must match the received response with the original request and produce the HTTP reply. Note: this code runs on the API Gateway side of RabbitMQ :param body: :param properties: :param method: :param channel: :return: """ reply_id = properties.correlation_id content_type = properties.content_type delivery_tag = method.delivery_tag debug = f" [DBG] ######[{delivery_tag}] (RPC REPLY) env: {method}, props: {properties}, BODY: {body.decode()}" logging.info(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 = AMQResponseFactory.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: AMQResponse = AMQResponseFactory.from_bytes(body) # Complete the Future with the response data future.set_result(response.body()) channel.basic_ack(delivery_tag, False)