Stanislav Hejny
279 lines
12 KiB
Python
279 lines
12 KiB
Python
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,
|
|
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
|
|
|
|
|
|
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)
|
|
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 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_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 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)
|