Stanislav Hejny
144 lines
6.3 KiB
Python
144 lines
6.3 KiB
Python
import logging
|
|
import time
|
|
|
|
from typing import Callable, Dict
|
|
from concurrent.futures import Future
|
|
from opentelemetry import trace
|
|
from opentelemetry.propagate import get_global_textmap, set_global_textmap
|
|
from opentelemetry.propagators.tracecontext import TraceContextTextMapPropagator
|
|
from opentelemetry.trace import Span, Tracer, TraceState
|
|
|
|
from amqp.adapter.amq_message_factory import AMQMessageFactory
|
|
from amqp.model.model import AMQResponse
|
|
from amqp.router.serializer import map_as_string
|
|
|
|
|
|
class AMQServiceMessageHandler:
|
|
def __init__(self, http_adapter, tracer):
|
|
self.http_adapter = http_adapter
|
|
self.tracer = tracer
|
|
self.outstanding = {}
|
|
|
|
def on_delivery_callback_http(self, channel, consumerTag, delivery):
|
|
start = time.time()
|
|
deliveryTag = delivery.envelope.delivery_tag
|
|
debug = f" [DBG] ######[{deliveryTag}] (AMQ MESSAGE), env: {delivery.envelope}, props: {delivery.properties}"
|
|
logging.info(debug)
|
|
|
|
amq_message = AMQMessageFactory.from_bytes(delivery.body)
|
|
logging.info(" [*] ######[%s] AMQ Message received, ctag=%s, msgId=%s, traceInfo=%s",
|
|
deliveryTag, consumerTag, amq_message.id, map_as_string(amq_message.trace_info))
|
|
|
|
propagator = TraceContextTextMapPropagator.get_instance()
|
|
context = propagator.extract(trace.get_current_context(), amq_message.trace_info , self.get_text_getter())
|
|
parent_span = Span.from_context(context)
|
|
context = context.with_span(parent_span)
|
|
|
|
with context.activate():
|
|
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 = self.http_adapter.handle(amq_message)
|
|
future.add_done_callback(
|
|
lambda f: self.on_http_response(delivery, deliveryTag, channel, f.result())
|
|
)
|
|
else:
|
|
logging.error(" [E] Unknown MAGIC value: [%s], don't know how to interpret the message",
|
|
amq_message.magic())
|
|
|
|
duration = time.time() - start
|
|
if "___NACK___" in str(amq_message.body()):
|
|
logging.info(" [x] ######[%s] Done, single NACK: %sms.", deliveryTag, duration)
|
|
channel.basic_nack(deliveryTag, False, False)
|
|
else:
|
|
logging.info(" [x] ######[%s] Done, single ACK: %sms.", deliveryTag, duration)
|
|
channel.basic_ack(deliveryTag, False)
|
|
|
|
def on_http_response(self, delivery, deliveryTag, channel, http_response):
|
|
response = AMQResponse(
|
|
delivery.properties.correlation_id,
|
|
http_response.status_code,
|
|
http_response.headers.get("Content-Type"),
|
|
http_response.body,
|
|
http_response.status_message,
|
|
"",
|
|
self.collect_trace_info(trace.get_current_context())
|
|
)
|
|
try:
|
|
self.send_reply(delivery.properties.reply_to, deliveryTag, channel, response)
|
|
except Exception as e:
|
|
raise RuntimeError(e)
|
|
|
|
|
|
class AMQServiceMessageHandler:
|
|
def __init__(self, http_adapter, tracer):
|
|
self.http_adapter = http_adapter
|
|
self.tracer = tracer
|
|
self.outstanding: Dict[str, 'ResponseHolder'] = {}
|
|
|
|
def collect_trace_info(self, context):
|
|
map = {}
|
|
TraceContextTextMapPropagator.get_instance().inject(
|
|
context if context else trace.get_current_context(),
|
|
map, self.set_text
|
|
)
|
|
return map
|
|
|
|
def send_reply(self, reply_to, delivery_tag, channel, response):
|
|
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(
|
|
RouterBase.AMQ_REPLY_TO_EXCHANGE,
|
|
reply_to,
|
|
reply_props,
|
|
AMQResponseFactory.serialize(response)
|
|
)
|
|
|
|
def on_reply_consumer(self, channel):
|
|
def handle_delivery(consumerTag, envelope, properties, body):
|
|
reply_id = properties.correlation_id
|
|
content_type = properties.content_type
|
|
delivery_tag = envelope.delivery_tag
|
|
debug = f" [DBG] ######[{delivery_tag}] (RPC REPLY) env: {envelope}, props: {properties}, BODY: {body.decode()}"
|
|
logging.info(debug)
|
|
print(f"OUTSTANDING RECV: {self.outstanding}, size={len(self.outstanding)}")
|
|
|
|
with self.outstanding.lock:
|
|
reply = 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()
|
|
reply.response().set_result(Response.ok(result.body(), content_type).build())
|
|
|
|
channel.basic_ack(delivery_tag, False)
|
|
|
|
return handle_delivery
|
|
|
|
def register_response_holder(self, id, response_holder):
|
|
self.outstanding[id] = response_holder
|
|
logging.info(" [*] AMQMessage Sent, OUTSTANDING response holder registered: %s, size=%s",
|
|
self.outstanding, len(self.outstanding))
|
|
|
|
def get_default_trace_state(self):
|
|
return TraceState(entries={})
|
|
|
|
def get_text_getter(self):
|
|
return lambda map, key: map.get(key)
|
|
|
|
def set_text(self, carrier, key, value):
|
|
carrier[key] = value
|