clevermicro/amq-adapter-python
14
0
Fork 0
Code Issues 16 Pull Requests 4 Actions Packages Projects Releases Wiki Activity

restructure to fit around python event-loop

Browse Source
This commit is contained in:
Stanislav Hejny
2025-04-16 00:24:17 +01:00
parent b550d98eaf
commit 0720c5963f
5 changed files with 153 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files
amqp/adapter/cleverthis_service_adapter.py
+37 -12
View File
@@ -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
# ==================================================================================================
amqp/adapter/file_uploader.py
+17 -18
View File
@@ -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
amqp/service/amq_message_handler.py
+74 -87
View File
@@ -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,27 +132,34 @@ class DataMessageHandler:
)
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
async def inbound_data_message_callback(self, message: AbstractIncomingMessage):
threading.Thread(target=lambda: asyncio.run(self.run_callback_thread(message))).start()
# 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:
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(
@@ -171,22 +177,38 @@ class DataMessageHandler:
self.rabbit_mq_client.router.publish_service_message(
scale_up_request, exchange=self.rabbit_mq_client.router.cm_response_exchange
),
loop=loop
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()
last_datadeliveryTag_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 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
"""
global current_parallel_executions
async def wait_for_response(future: Future):
print(f"AWAITINMG RESPONSE %s -> %s", future, future.done())
response: DataResponse = future.result()
# Increment the counter for parallel executions
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)
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,53 +229,41 @@ 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)
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
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)
@@ -262,39 +272,16 @@ class DataMessageHandler:
# 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 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):
amqp/service/amq_service.py
+4 -7
View File
@@ -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
pyproject.toml
+1 -1
View File
@@ -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" }