hurui200320
45941de6b1
Remove the loop for monitoring the CPU usage, adopt changes in the management service for the reporting message format, update the timing for sending out report, similar to java adapter. ISSUES RELATED: clevermicro/amq-adapter-python#58
278 lines
12 KiB
Python
278 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
|
|
await 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:
|
|
_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}")
|
|
# recover availability
|
|
await 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)
|