From 0720c5963fac2ffac3d424b03886fa3289994719 Mon Sep 17 00:00:00 2001 From: Stanislav Hejny Date: Wed, 16 Apr 2025 00:24:17 +0100 Subject: [PATCH] restructure to fit around python event-loop --- amqp/adapter/cleverthis_service_adapter.py | 49 +++-- amqp/adapter/file_uploader.py | 35 ++-- amqp/service/amq_message_handler.py | 201 ++++++++++----------- amqp/service/amq_service.py | 11 +- pyproject.toml | 2 +- 5 files changed, 153 insertions(+), 145 deletions(-) diff --git a/amqp/adapter/cleverthis_service_adapter.py b/amqp/adapter/cleverthis_service_adapter.py index 4c7c040..b48dc96 100644 --- a/amqp/adapter/cleverthis_service_adapter.py +++ b/amqp/adapter/cleverthis_service_adapter.py @@ -9,12 +9,12 @@ import logging import sys import traceback import pathlib -from asyncio import Future +from asyncio import Future, AbstractEventLoop from typing import Dict, List, Tuple, Callable, Coroutine, Any from urllib.parse import parse_qs, urlparse from xml.etree import ElementTree -from aio_pika.abc import AbstractRobustChannel +from aio_pika.abc import AbstractRobustChannel, AbstractRobustConnection, AbstractRobustExchange, AbstractExchange from aiohttp import ClientSession, ClientResponse from dataclasses import dataclass from fastapi import HTTPException @@ -28,6 +28,8 @@ from amqp.adapter.file_uploader import publish_file_to_rabbitmq from IPython.core import debugger +from amqp.router.router_base import AMQ_REPLY_TO_EXCHANGE + debug = debugger.Pdb().set_trace @@ -38,7 +40,7 @@ class AMQMessage(DataMessage): in case the response should initiate file download. """ - def __init__(self, data_message: DataMessage, channel: AbstractRobustChannel): + def __init__(self, data_message: DataMessage, connection: AbstractRobustConnection): # Copy all fields from DataMessage self.magic = data_message.magic self.id = data_message.id @@ -49,7 +51,7 @@ class AMQMessage(DataMessage): self.trace_info = data_message.trace_info self.base64_body = data_message.base64_body # Add the new field - self.channel = channel + self.connection = connection def parse_request_body(message: DataMessage): @@ -115,11 +117,35 @@ def parse_request_body(message: DataMessage): raise ValueError(f"Unsupported Content-Type: {content_type}") -async def convert_file_response(obj: FileResponse, channel: AbstractRobustChannel) -> str: +async def convert_file_response( + obj: FileResponse, + connection: AbstractRobustConnection, + loop: AbstractEventLoop | None +) -> str: """ Converts a FileResponse object to a JSON string containing the filename and stream. """ - (_stream_name, _file_size) = await publish_file_to_rabbitmq(channel, obj.path) + async def wrap_connection_channel(connection: AbstractRobustConnection) -> tuple[AbstractRobustChannel, AbstractExchange, str]: + """ Wrap call to channel() to make it Awaitable """ + _channel = await connection.channel() + exchange_name: str = AMQ_REPLY_TO_EXCHANGE + _exchange = await _channel.get_exchange(name=exchange_name, ensure=True) + _queue = await _channel.declare_queue(exclusive=True) # Declare a unique, exclusive queue + _queue_name = _queue.name + await _queue.bind( + exchange=exchange_name, + routing_key=_queue_name, # Use queue name as routing key + ) + + logging.debug(f" [*] Publishing to exchange: {exchange_name}, q: {_queue_name}") + return _channel, _exchange, _queue_name + + _future = asyncio.run_coroutine_threadsafe(wrap_connection_channel(connection), loop) + while not _future.done(): + await asyncio.sleep(0.1) # allow coroutine to execute + _channel, _exchange, _routing_key = _future.result() + (_stream_name, _file_size) = await publish_file_to_rabbitmq(_exchange, obj.path, _routing_key, loop) + await _channel.close() return json.dumps({'files': [ {'filename': obj.filename, 'mediaType': 'application/octet-stream', @@ -133,7 +159,8 @@ async def call_cleverthis_api( message: AMQMessage, args: Any, api_method: Callable[[Any, Any], Coroutine], - success_code: int + success_code: int, + loop: AbstractEventLoop | None = None ) -> DataResponse: """ Handles POST and PUT requests, which typically contain possibly complex payload data, passed as a dictionary. @@ -162,7 +189,7 @@ async def call_cleverthis_api( if isinstance(_result, FileResponse): _json_result, _content_type = await convert_file_response( - _result, message.channel + _result, message.connection, loop ), 'application/octet-stream' else: _json_result = _result @@ -224,6 +251,7 @@ class CleverThisServiceAdapter: self.session = session self.service_host = self.amq_configuration.dispatch.service_host self.service_port = self.amq_configuration.dispatch.service_port + self.loop: AbstractEventLoop | None = None self.require_authenticated_user = eval( self.amq_configuration.amq_adapter.require_authenticated_user.capitalize()) if self.amq_configuration.amq_adapter.require_authenticated_user.strip() else True @@ -247,7 +275,7 @@ class CleverThisServiceAdapter: """ return None - async def on_message(self, message: AMQMessage, future: Future[DataResponse] = None) -> DataResponse: + async def on_message(self, message: AMQMessage) -> DataResponse: """ Called when a message is received from the AMQP. :param message: @@ -292,9 +320,6 @@ class CleverThisServiceAdapter: raise ValueError(f"Unexpected HTTP method: {message.method}") else: amq_response = DataResponseFactory.of(message.id, 401, "text/plain", b"Unauthorized", message.trace_info) - - if future: - future.set_result(amq_response) return amq_response # ================================================================================================== diff --git a/amqp/adapter/file_uploader.py b/amqp/adapter/file_uploader.py index 3bd729a..19148b8 100644 --- a/amqp/adapter/file_uploader.py +++ b/amqp/adapter/file_uploader.py @@ -1,10 +1,12 @@ import logging +from asyncio import AbstractEventLoop + import aio_pika import asyncio import os from pathlib import Path -from aio_pika.abc import AbstractRobustChannel +from aio_pika.abc import AbstractRobustChannel, AbstractRobustExchange from aio_pika import Message, DeliveryMode from amqp.adapter.file_downloader import IN_PROGRESS, END_OF_FILE @@ -12,10 +14,11 @@ from amqp.router.router_base import AMQ_REPLY_TO_EXCHANGE async def publish_file_to_rabbitmq( - channel: AbstractRobustChannel, + exchange: AbstractRobustExchange, file_path: Path, + routing_key: str, + loop: AbstractEventLoop, max_chunk_size: int = 2**20, # 1MB default chunk size - exchange_name: str = AMQ_REPLY_TO_EXCHANGE, content_type: str = 'application/octet-stream', # Default content type delivery_mode: DeliveryMode = DeliveryMode.PERSISTENT, ) -> tuple[str, int]: @@ -25,10 +28,12 @@ async def publish_file_to_rabbitmq( exchange, and uses the queue name as the routing key. Args: - channel: The aio_pika RobustChannel to use for communication with RabbitMQ. + exchange: The aio_pika RobustChannel to use for communication with RabbitMQ. file_path: The path to the file to publish. max_chunk_size: The maximum size of each chunk (in bytes). - exchange_name: The name of the RabbitMQ exchange to publish to. + routing_key: The name of the RabbitMQ exchange routing key to publish to correct queue. + loop: correct event loop + max_chunk_size: chunk size to send content_type: The content type of the file data. delivery_mode: The delivery mode (e.g., PERSISTENT or TRANSIENT). @@ -38,15 +43,6 @@ async def publish_file_to_rabbitmq( if not file_path.is_file(): raise FileNotFoundError(f"File not found: {file_path}") _file_size = os.path.getsize(file_path) - _exchange = await channel.get_exchange(name=exchange_name, ensure=True) - _queue = await channel.declare_queue(exclusive=True) # Declare a unique, exclusive queue - _queue_name = _queue.name - await _queue.bind( - exchange=exchange_name, - routing_key=_queue_name, # Use queue name as routing key - ) - - logging.info(f"Publishing file: {file_path} (size: {_file_size}) to exchg: {exchange_name}, q: {_queue_name}") with open(file_path, 'rb') as f: _offset = 0 @@ -67,9 +63,12 @@ async def publish_file_to_rabbitmq( 'eof': END_OF_FILE if max_chunk_size + _offset >= _file_size else IN_PROGRESS } ) - await _exchange.publish(message, routing_key=_queue_name) + _future = asyncio.run_coroutine_threadsafe(exchange.publish(message, routing_key=routing_key), loop) + while not _future.done(): + # sleep allows coroutine to execute. must wait for completion to ensure all chunks written in right order + await asyncio.sleep(0.02) _offset += len(_chunk) - logging.debug(f"Published chunk {_offset}/{_file_size} to {exchange_name}:{_queue_name}") + logging.debug(f"Published chunk {_offset}/{_file_size} to {exchange.name}:{routing_key}") - print(f"File {file_path} published to RabbitMQ exchange {exchange_name}, queue {_queue_name}") - return _queue_name, _file_size + print(f"File {file_path} published to RabbitMQ exchange {exchange.name}, queue {routing_key}") + return routing_key, _file_size diff --git a/amqp/service/amq_message_handler.py b/amqp/service/amq_message_handler.py index d7b073d..1a3887b 100644 --- a/amqp/service/amq_message_handler.py +++ b/amqp/service/amq_message_handler.py @@ -34,13 +34,11 @@ 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) +executor = concurrent.futures.ThreadPoolExecutor(max_workers=parallel_workers + 2) def collect_trace_info(): @@ -79,6 +77,7 @@ class DataMessageHandler: 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 @@ -133,10 +132,61 @@ class DataMessageHandler: ) await asyncio.sleep(self.backpressure_monitor_window) # sleep in seconds - def sync_on_message(self, + async def inbound_data_message_callback(self, message: AbstractIncomingMessage): + threading.Thread(target=lambda: asyncio.run(self.run_callback_thread(message))).start() + + async def run_callback_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 == DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST: + await self.run_http_request_magic(message, amq_message, self.loop) + else: + logging.error(" [E] 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(" [E] ######[%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) + + def handle_backpressure(self, _loop: AbstractEventLoop): + 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_datadeliveryTag_message_time = time.time() + + async def run_http_request_magic(self, message: AbstractIncomingMessage, amq_message: DataMessage, - loop: AbstractEventLoop) -> None: + _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. @@ -148,45 +198,17 @@ class DataMessageHandler: 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() + 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: + self.handle_backpressure(_loop) - # 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: + a_message = AMQMessage(amq_message, self.rabbit_mq_client.connection) + response: DataResponse = await self.service_adapter.on_message(a_message) + print(f"Result in thread: {response}") try: logging.info(" [*] ######[%s] AMQ Message Reply.1, To=%s, msgId=%s", message.delivery_tag, message.reply_to, response.id) @@ -197,7 +219,7 @@ class DataMessageHandler: 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) + asyncio.run_coroutine_threadsafe(message.nack(requeue=False), loop=_loop) if isinstance(amq_message, MultipartDataMessage): # remove the downloaded files from the output_dir @@ -207,94 +229,59 @@ class DataMessageHandler: 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) + current_parallel_executions -= 1 + return + except Exception as e: + print(f"Error in thread: {e}") + return - 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) - + 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: - loop = asyncio.get_event_loop() - amq_message = await DataMessageFactory.from_stream(message.body, rabbitmq_url=self.rabbit_mq_url, loop=loop) + amq_message = await DataMessageFactory.from_stream(message.body, rabbitmq_url=self.rabbit_mq_url, loop=self.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, + loop=self.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: + if amq_message: logging.info(" [*] ######[%s] AMQ Message received, ctag=%s, msgId=%s, %s:%s, traceInfo=%s", - delivery_tag, message.consumer_tag, amq_message.id, + 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 - 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) + 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(" [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 reply_to:connection.channel :param response: :return: """ @@ -306,10 +293,10 @@ class DataMessageHandler: 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( + _res = asyncio.run_coroutine_threadsafe(self.reply_to_exchange.publish( message=pika_message, routing_key=reply_to - ) + ), self.loop) 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): diff --git a/amqp/service/amq_service.py b/amqp/service/amq_service.py index 0adcedd..0ff04fc 100644 --- a/amqp/service/amq_service.py +++ b/amqp/service/amq_service.py @@ -38,7 +38,7 @@ class AMQService: self.loop = asyncio.get_event_loop() self.amq_configuration = amq_configuration - self.message_factory = DataMessageFactory.get_instance(self.amq_configuration.amq_adapter.generator_id) + self.data_message_factory = DataMessageFactory.get_instance(self.amq_configuration.amq_adapter.generator_id) self.rabbit_mq_client = RabbitMQClient(self.amq_configuration) self.tracer = initialize_trace(name=amq_configuration.amq_adapter.service_name) self.amq_data_message_handler = None @@ -47,9 +47,6 @@ class AMQService: ) # self.init(loop) will initialize RabbitMQ connection self.loop.run_until_complete(self.init(self.loop, service_adapter)) - # loop.run_until_complete(consume(loop)) - # self.consumer_thread = asyncio.create_task(self.rabbit_mq_client.channel.start_consuming()) - # self.rabbit_mq_client.request_routes() logging.info("***********************************************") logging.info("****** AMQ Service Init Complete *********") @@ -57,13 +54,13 @@ class AMQService: self.loop.run_forever() async def init(self, loop, service_adapter): - exchange: AbstractRobustExchange = await self.rabbit_mq_client.init(loop) + _reply_to_exchange: AbstractRobustExchange = await self.rabbit_mq_client.init(loop) self.amq_data_message_handler = DataMessageHandler( service_adapter, self.tracer, - self.message_factory, + self.data_message_factory, self.service_message_factory, - exchange, + _reply_to_exchange, self.rabbit_mq_client, self.amq_configuration, loop=loop diff --git a/pyproject.toml b/pyproject.toml index 516b809..c7c6deb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ packages = ["amqp","amqp.adapter","amqp.config","amqp.model","amqp.rabbitmq","am [project] name = "amq_adapter" -version = "0.2.18" +version = "0.2.19" description = "The Python implementation of the AMQ Adaptor for CleverMicro" readme = "README.md" license = { text = "Apache License 2.0" }