Stanislav Hejny
187 lines
8.2 KiB
Python
187 lines
8.2 KiB
Python
import asyncio
|
|
import logging
|
|
import time
|
|
import concurrent.futures
|
|
import os
|
|
from asyncio import Future
|
|
from typing import Dict
|
|
|
|
from aio_pika import IncomingMessage, Message
|
|
from aio_pika.abc import AbstractIncomingMessage, AbstractRobustExchange
|
|
|
|
from opentelemetry import trace
|
|
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
|
|
from opentelemetry.trace import Tracer, TraceState
|
|
|
|
from amqp.adapter.amq_message_factory import AMQMessageFactory
|
|
from amqp.adapter.amq_response_factory import AMQResponseFactory
|
|
from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter
|
|
from amqp.model.model import AMQResponse, AMQMessage
|
|
from amqp.router.serializer import map_as_string
|
|
|
|
parallel_workers = os.getenv("PARALLEL_WORKERS", default=10)
|
|
|
|
# Shared thread pool for parallel execution
|
|
executor = concurrent.futures.ThreadPoolExecutor(max_workers=parallel_workers)
|
|
|
|
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: CleverThisServiceAdapter,
|
|
tracer: Tracer, message_factory: AMQMessageFactory,
|
|
reply_to_exchange: AbstractRobustExchange):
|
|
self.service_adapter: CleverThisServiceAdapter = service_adapter
|
|
self.tracer = tracer
|
|
self.message_factory: AMQMessageFactory = message_factory
|
|
self.reply_to_exchange: AbstractRobustExchange = reply_to_exchange
|
|
# Dictionary to store outstanding Futures
|
|
self.outstanding: Dict[str, asyncio.Future] = {}
|
|
|
|
async def inbound_message_callback(self, message: AbstractIncomingMessage):
|
|
"""
|
|
Receives the message as bytes from RabbitMQ queue and deserializes it to the AMQMessage
|
|
Ensures Jaeger trace-id propagation.
|
|
Invokes custom handler/mapper/adapter method that corresponds to this message
|
|
:param message:
|
|
:return:
|
|
"""
|
|
start = time.time()
|
|
delivery_tag = message.delivery_tag
|
|
debug = f" [DBG] ######[{delivery_tag}] (AMQ MESSAGE), env: {message.headers}, props: {message.properties}"
|
|
logging.info(debug)
|
|
|
|
amq_message: AMQMessage = AMQMessageFactory.from_bytes(message.body)
|
|
logging.info(" [*] ######[%s] AMQ Message received, ctag=%s, msgId=%s, traceInfo=%s",
|
|
delivery_tag, message.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:
|
|
future = executor.submit(self.service_adapter.sync_on_message, amq_message)
|
|
# response: AMQResponse = asyncio.run(self.http_adapter.on_message(amq_message))
|
|
# Wait for the result and send it back to the reply queue
|
|
try:
|
|
response: AMQResponse = future.result() # Block until the result is available
|
|
await self.send_reply(message.reply_to, delivery_tag, response)
|
|
except Exception as e:
|
|
print(f"[E] Error processing message: {e}")
|
|
await message.nack(requeue=False)
|
|
return
|
|
else:
|
|
logging.error(" [E] Unknown MAGIC value: [%s], don't know how to interpret the message",
|
|
amq_message.magic)
|
|
await message.nack(requeue=False)
|
|
return
|
|
|
|
duration = time.time() - start
|
|
logging.info(" [x] ######[%s] Done, single ACK: %sms.", delivery_tag, duration)
|
|
await message.ack(multiple=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, response)
|
|
except Exception as e:
|
|
raise RuntimeError(e)
|
|
|
|
async def send_reply(self, reply_to: str, delivery_tag: int | None, 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 response:
|
|
:return:
|
|
"""
|
|
if reply_to:
|
|
logging.info(" [*] ######[%s] AMQ Message Reply, To=%s, msgId=%s",
|
|
delivery_tag, reply_to, response.id)
|
|
pika_message: Message = Message(
|
|
body=AMQResponseFactory.serialize(response),
|
|
correlation_id=response.id,
|
|
content_type=response.content_type[0]
|
|
if isinstance(response.content_type, list) else response.content_type
|
|
)
|
|
await self.reply_to_exchange.publish(
|
|
message=pika_message,
|
|
routing_key=reply_to
|
|
)
|
|
|
|
async def reply_received_callback(self, message: IncomingMessage):
|
|
"""
|
|
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 message:
|
|
:return:
|
|
"""
|
|
reply_id = message.correlation_id
|
|
content_type = message.content_type
|
|
delivery_tag = message.delivery_tag
|
|
debug = f" [DBG] ######[{delivery_tag}] (RPC REPLY) env: {message.headers}, props: {message.properties}, BODY: {message.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(message.body)
|
|
# Complete the Future with the response data
|
|
future.set_result(response.body())
|
|
|
|
await message.ack(multiple=False)
|