348 lines
17 KiB
Python
348 lines
17 KiB
Python
import asyncio
|
|
import logging
|
|
import sys
|
|
import threading
|
|
import time
|
|
import concurrent.futures
|
|
import os
|
|
import traceback
|
|
from asyncio import Future, AbstractEventLoop
|
|
from typing import Dict
|
|
|
|
from aio_pika import 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.data_message_factory import DataMessageFactory
|
|
from amqp.adapter.data_response_factory import DataResponseFactory
|
|
from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter, AMQMessage
|
|
from amqp.adapter.file_downloader import on_message
|
|
from amqp.adapter.service_message_factory import ServiceMessageFactory
|
|
from amqp.config.amq_configuration import AMQConfiguration
|
|
from amqp.model.model import DataResponse, DataMessage, MultipartDataMessage, ScalingRequest, ScalingRequestAlert
|
|
from amqp.model.service_message_type import ServiceMessageType
|
|
from amqp.rabbitmq.rabbit_mq_client import RabbitMQClient
|
|
from amqp.router.router_base import ANY_RECIPIENT
|
|
from amqp.adapter.serializer import map_as_string
|
|
|
|
# Track the number of messages currently being processed
|
|
current_parallel_executions = 0
|
|
last_data_message_time = 0 # helps detect IDLE condition
|
|
last_backpressure_event_time = 0 # helps prevent flooding the system with backpressure events
|
|
last_backpressure_event = ScalingRequestAlert.IDLE
|
|
|
|
# Global lock to synchronize access to the current_parallel_executions variable
|
|
lock = threading.Lock()
|
|
parallel_workers = int(os.getenv("PARALLEL_WORKERS", default=0))
|
|
swarm_service_id = os.getenv("SWARM_SERVICE_ID", default="clever-amqp-service")
|
|
swarm_task_id = os.getenv("SWARM_TASK_ID", default="python-adapter")
|
|
# 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 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):
|
|
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
|
|
# Dictionary to store outstanding Futures
|
|
self.outstanding: Dict[str, asyncio.Future] = {}
|
|
self.reply_to: str | None = None
|
|
self.backpressure_monitor_window = amq_configuration.backpressure.time_window / 1000 # convert to seconds
|
|
|
|
global parallel_workers
|
|
_backpressure_treshold = amq_configuration.backpressure.threshold
|
|
"""
|
|
Determine the number of parallel workers based on the backpressure threshold and PARALLEL_WORKERS env var.
|
|
If env var PARALLEL_WORKERS is not set (or set to 0), use the backpressure threshold config
|
|
as the number of parallel workers.
|
|
"""
|
|
if _backpressure_treshold > 0:
|
|
if parallel_workers <= 0:
|
|
parallel_workers = _backpressure_treshold
|
|
|
|
async def monitor_backpressure(self):
|
|
while True:
|
|
global last_data_message_time
|
|
global last_backpressure_event_time
|
|
global last_backpressure_event
|
|
|
|
now = time.time()
|
|
event: ScalingRequestAlert = ScalingRequestAlert.UPDATE
|
|
if now - last_backpressure_event_time >= self.backpressure_monitor_window:
|
|
if current_parallel_executions == 0 and now - last_data_message_time > self.backpressure_monitor_window:
|
|
event = ScalingRequestAlert.IDLE
|
|
elif current_parallel_executions > parallel_workers - 2:
|
|
event = ScalingRequestAlert.OVERLOAD
|
|
|
|
last_backpressure_event_time = now
|
|
last_backpressure_event = event
|
|
last_backpressure_event_time = now
|
|
logging.info(f"Backpressure event: {event}, threads: {current_parallel_executions} of {parallel_workers}")
|
|
|
|
scaling_request: ScalingRequest = ScalingRequest(
|
|
swarm_service_id, swarm_task_id,
|
|
parallel_workers, parallel_workers - current_parallel_executions,
|
|
event
|
|
)
|
|
# Address the message to any adapter capable of supporting BACKPRESSURE request
|
|
scale_up_request = self.service_message_factory.of(
|
|
ServiceMessageType.BACKPRESSURE,
|
|
scaling_request.to_json(),
|
|
self.reply_to if self.reply_to else ANY_RECIPIENT
|
|
)
|
|
asyncio.run_coroutine_threadsafe(
|
|
self.rabbit_mq_client.router.publish_service_message(
|
|
scale_up_request, exchange=self.rabbit_mq_client.router.cm_response_exchange
|
|
),
|
|
loop=self.loop
|
|
)
|
|
await asyncio.sleep(self.backpressure_monitor_window) # sleep in seconds
|
|
|
|
def sync_on_message(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
|
|
"""
|
|
global current_parallel_executions
|
|
|
|
# Increment the counter for parallel executions
|
|
with lock:
|
|
current_parallel_executions += 1
|
|
logging.info(f"PARALLEL: Current Capacity [{current_parallel_executions}] of [{parallel_workers}]")
|
|
if message.reply_to:
|
|
self.reply_to = message.reply_to
|
|
if current_parallel_executions > parallel_workers - 2:
|
|
logging.warning("Warning: Capacity close to depleted!")
|
|
# Send a scaling request to the Management service
|
|
scaling_request: ScalingRequest = ScalingRequest(
|
|
swarm_service_id, swarm_task_id,
|
|
parallel_workers, parallel_workers - current_parallel_executions,
|
|
ScalingRequestAlert.OVERLOAD
|
|
)
|
|
# Address the message to any adapter capable of supporting BACKPRESSURE request
|
|
scale_up_request = self.service_message_factory.of(
|
|
ServiceMessageType.BACKPRESSURE,
|
|
scaling_request.to_json(),
|
|
self.reply_to if self.reply_to else ANY_RECIPIENT
|
|
)
|
|
asyncio.run_coroutine_threadsafe(
|
|
self.rabbit_mq_client.router.publish_service_message(
|
|
scale_up_request, exchange=self.rabbit_mq_client.router.cm_response_exchange
|
|
),
|
|
loop=loop
|
|
)
|
|
# update / reset time-window so that the OVERLOAD is not sent too often
|
|
global last_data_message_time
|
|
last_data_message_time = time.time()
|
|
|
|
# Call the async on_message method. This shall also ACK the message after processing.
|
|
a_message = AMQMessage(amq_message, self.rabbit_mq_client.channel)
|
|
future: Future[DataResponse] = asyncio.Future(loop=loop)
|
|
print(f" [*] ######[{message.delivery_tag}] about to call asyncio.run(on_message) with AMQ Message: {a_message}")
|
|
ftr: Future[DataResponse] = asyncio.run_coroutine_threadsafe(self.service_adapter.on_message(a_message, future), loop)
|
|
print(f"WILL AWAIT RESP {ftr} -> {ftr.done()}")
|
|
|
|
async def wait_for_response(future: Future):
|
|
print(f"AWAITINMG RESPONSE %s -> %s", future, future.done())
|
|
response: DataResponse = future.result()
|
|
try:
|
|
logging.info(" [*] ######[%s] AMQ Message Reply.1, To=%s, msgId=%s",
|
|
message.delivery_tag, message.reply_to, response.id)
|
|
await self.send_reply(message.reply_to, response)
|
|
logging.info(" [*] ######[%s] AMQ Message Reply.2, To=%s",
|
|
message.delivery_tag, message.reply_to)
|
|
except Exception as e:
|
|
logging.info(f"[E] Error processing message: {e}")
|
|
tb = traceback.extract_tb(sys.exc_info()[2])[-1]
|
|
logging.error(f"[E] Error in {tb.filename}, line {tb.lineno}, function '{tb.name}': {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"Error deleting file {filename}: {e}")
|
|
|
|
with lock:
|
|
current_parallel_executions -= 1
|
|
print(f"KICK OFF wait proc {future}")
|
|
asyncio.run_coroutine_threadsafe(wait_for_response(future), loop)
|
|
|
|
|
|
async def inbound_data_message_callback(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()
|
|
delivery_tag = message.delivery_tag
|
|
debug = f" [DBG] ######[{delivery_tag}] (AMQ MESSAGE), env: {message.headers}, props: {message.properties}"
|
|
logging.debug(debug)
|
|
addressing: int = message.headers.get('clevermicro.addressing', 0)
|
|
|
|
amq_message: DataMessage | None = None
|
|
if addressing == 0:
|
|
amq_message = DataMessageFactory.from_bytes(message.body)
|
|
else:
|
|
loop = asyncio.get_event_loop()
|
|
amq_message = await DataMessageFactory.from_stream(message.body, rabbitmq_url=self.rabbit_mq_url, loop=loop)
|
|
logging.info(f" [DBG] ######[{delivery_tag}] AMQ Message from stream: {amq_message}")
|
|
if amq_message is not None:
|
|
extra_data = await on_message(
|
|
message=message,
|
|
json_data=amq_message.body(),
|
|
rabbitmq_url=self.rabbit_mq_url,
|
|
loop=loop,
|
|
output_dir=self.output_dir)
|
|
logging.info(f" [x] ######[{delivery_tag}] Multipart download result: {extra_data}")
|
|
amq_message = MultipartDataMessage(amq_message, extra_data)
|
|
|
|
if amq_message is not None:
|
|
logging.info(" [*] ######[%s] AMQ Message received, ctag=%s, msgId=%s, %s:%s, traceInfo=%s",
|
|
delivery_tag, message.consumer_tag, amq_message.id,
|
|
amq_message.method, amq_message.path, 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 == DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST:
|
|
future = concurrent.futures.ThreadPoolExecutor().submit(
|
|
self.sync_on_message, message, amq_message, self.loop
|
|
)
|
|
# response: DataResponse = asyncio.run(self.http_adapter.on_message(amq_message))
|
|
# Wait and send it back to the reply queue
|
|
#try:
|
|
# response: DataResponse = future.result() # Block until the result is available
|
|
# await self.send_reply(message.reply_to, delivery_tag, response)
|
|
#except Exception as e:
|
|
# logging.info(f"[E] Error processing message: {e}")
|
|
# tb = traceback.extract_tb(sys.exc_info()[2])[-1]
|
|
# logging.error(f"[E] Error in {tb.filename}, line {tb.lineno}, function '{tb.name}': {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
|
|
else:
|
|
logging.error(" [E] ######[%s] No valid AMQ Message present.", delivery_tag)
|
|
|
|
global last_data_message_time
|
|
last_data_message_time = time.time()
|
|
logging.info(" [x] ######[%s] Done, single ACK: %sms.", delivery_tag, last_data_message_time - start)
|
|
await message.ack(multiple=False)
|
|
|
|
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:
|
|
: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
|
|
)
|
|
logging.info(" [*] ###### Sending Reply.1 To=%s, msgId=%s", reply_to, response.id)
|
|
_res = await self.reply_to_exchange.publish(
|
|
message=pika_message,
|
|
routing_key=reply_to
|
|
)
|
|
logging.info(" [*] ###### DONE EXCH PUB(%s) Sending Reply.2 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
|
|
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 = 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)
|