From 2c641b0f64565d985ad61e9371266db96e9eddfd Mon Sep 17 00:00:00 2001 From: Stanislav Hejny Date: Sat, 12 Apr 2025 10:50:35 +0100 Subject: [PATCH] Fix the multithread/waiting for result --- amqp/adapter/cleverthis_service_adapter.py | 64 ++++++++++++++++------ amqp/adapter/file_uploader.py | 3 + amqp/service/amq_message_handler.py | 62 +++++++++++---------- 3 files changed, 83 insertions(+), 46 deletions(-) diff --git a/amqp/adapter/cleverthis_service_adapter.py b/amqp/adapter/cleverthis_service_adapter.py index 36c4d35..45caf3b 100644 --- a/amqp/adapter/cleverthis_service_adapter.py +++ b/amqp/adapter/cleverthis_service_adapter.py @@ -25,6 +25,10 @@ from amqp.model.model import DataMessage, DataResponse from amqp.adapter.multipart_parser import CleverMultiPartParser, process_form_data from amqp.adapter.file_uploader import publish_file_to_rabbitmq +from IPython.core import debugger + +debug = debugger.Pdb().set_trace + @dataclass class AMQMessage(DataMessage): @@ -32,6 +36,7 @@ class AMQMessage(DataMessage): Expand the core DataMessage type to include the channel, which is required to create a dedicated queue in case the response should initiate file download. """ + def __init__(self, data_message: DataMessage, channel: AbstractRobustChannel): # Copy all fields from DataMessage self.magic = data_message.magic @@ -113,9 +118,10 @@ async def convert_file_response(obj: FileResponse, channel: AbstractRobustChanne """ Converts a FileResponse object to a JSON string containing the filename and stream. """ - _future: Future = asyncio.run_coroutine_threadsafe( - publish_file_to_rabbitmq(channel, obj.path), loop=channel.channel.loop) - return json.dumps({'files':[ + _future: Future[str] = asyncio.Future(loop=channel.channel.loop) + asyncio.run_coroutine_threadsafe( + publish_file_to_rabbitmq(channel, obj.path, _future), loop=channel.channel.loop) + return json.dumps({'files': [ {'filename': obj.filename, 'mediaType': 'application/octet-stream', 'streamName': _future.result() @@ -124,10 +130,10 @@ async def convert_file_response(obj: FileResponse, channel: AbstractRobustChanne async def call_cleverthis_api( - message: AMQMessage, - args: Any, - api_method: Callable[[Any, Any], Coroutine], - success_code: int + message: AMQMessage, + args: Any, + api_method: Callable[[Any, Any], Coroutine], + success_code: int ) -> DataResponse: """ Handles POST and PUT requests, which typically contain possibly complex payload data, passed as a dictionary. @@ -143,7 +149,7 @@ async def call_cleverthis_api( _verified_args = {} _sig = inspect.signature(api_method) _required_args = list(_sig.parameters.values()) - #_required_args = getattr(api_method, '__annotations__', {}) + # _required_args = getattr(api_method, '__annotations__', {}) for arg in _required_args: if arg.name in args: _verified_args[arg.name] = args[arg.name] @@ -153,20 +159,27 @@ async def call_cleverthis_api( f" [*] INVOKE ENDPOINT: {api_method.__name__}, ARGS: {', '.join(f'{k}={v}' for k, v in _verified_args.items())}") _result: Any = await api_method(**_verified_args) logging.info(f" [*] ENDPOINT RESULT: {api_method.__name__}, RESULT: {str(_result)}") - _json_result, _content_type = await convert_file_response( + + if isinstance(_result, FileResponse): + _json_result, _content_type = await convert_file_response( _result, message.channel - ), 'application/octet-stream' if isinstance(_result, FileResponse) else _result, 'application/json' + ), 'application/octet-stream' + else: + _json_result = _result + _content_type = 'application/json' + _binary_result = _json_result.encode('utf-8') if hasattr( _json_result, 'encode' ) else _json_result.body if hasattr( _json_result, 'body' ) else _json_result.read() if hasattr(_json_result, 'read') else _json_result + return DataResponseFactory.of(message.id, success_code, _content_type, _binary_result, message.trace_info) except HTTPException as http_error: _tb = traceback.extract_tb(sys.exc_info()[2])[-1] logging.error(f" [E] Error in {_tb.filename}, line {_tb.lineno}, function '{_tb.name}': {http_error}") return DataResponseFactory.of(message.id, http_error.status_code, "text/plain", - str(http_error.detail).encode('utf-8'), message.trace_info) + str(http_error.detail).encode('utf-8'), message.trace_info) except Exception as error: _tb = traceback.extract_tb(sys.exc_info()[2])[-1] logging.error(f" [E] Error in {_tb.filename}, line {_tb.lineno}, function '{_tb.name}': {error}") @@ -200,6 +213,7 @@ class CleverThisServiceAdapter: IMPORTANT: This class is not intended to be used directly. It should be subclassed for each CleverThis service and override the methods to handle the specific API calls. """ + def __init__(self, amq_configuration: AMQConfiguration, session: ClientSession = None): """ Initializes the CleverThisServiceAdapter with the given AMQ configuration and optional HTTP session. @@ -210,7 +224,8 @@ class CleverThisServiceAdapter: self.session = session self.service_host = self.amq_configuration.dispatch.service_host self.service_port = self.amq_configuration.dispatch.service_port - self.require_authenticated_user = self.amq_configuration.amq_adapter.require_authenticated_user + 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 def get_content_type(self, message_headers: Dict[str, List[str]]) -> str: """ @@ -223,7 +238,7 @@ class CleverThisServiceAdapter: return values[0].split(';')[0] return 'application/json' - async def get_current_user(self, token: str): + async def get_current_user(self, token: str) -> object: """ To Be Overriden in subclass for given CleverThis Service, to return the current user based on the token, in format appropriate for the service. @@ -232,7 +247,7 @@ class CleverThisServiceAdapter: """ return None - async def on_message(self, message: AMQMessage) -> DataResponse: + async def on_message(self, message: AMQMessage, future: Future[DataResponse] = None) -> DataResponse: """ Called when a message is received from the AMQP. :param message: @@ -241,16 +256,27 @@ class CleverThisServiceAdapter: _auth_user: Any | None = None amq_response: Any | None = None + logging.warn( + f'Got here with: _auth_user={_auth_user} and require_authenticated_user={self.require_authenticated_user}') if self.require_authenticated_user: _auth: str = message.headers.get('Authorization', '') + logging.info(f"Auth: {_auth}") + if isinstance(_auth, List): + _auth = _auth[0] if 'Bearer' in _auth: try: token = _auth.split(' ')[1] + logging.warn(f'Token: {token}') _auth_user = await self.get_current_user(token) except Exception as error: logging.error(f"Error getting user: {error}") - if _auth_user or not self.require_authenticated_user: + # bearer_in_auth = 'Bearer' in _auth + # logging.warn(f'Got here2 with: Bearer_in_auth={bearer_in_auth}, _auth_user={_auth_user} and require_authenticated_user={self.require_authenticated_user}') + logging.warn( + f'Got here2 with: _auth_user={_auth_user} and require_authenticated_user={self.require_authenticated_user}') + + if _auth_user or not self.require_authenticated_user or message.path.endswith('/login'): method = message.method.lower() if method == 'get': amq_response = await self.get(message, _auth_user) @@ -267,6 +293,8 @@ class CleverThisServiceAdapter: 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 # ================================================================================================== @@ -289,10 +317,10 @@ class CleverThisServiceAdapter: if self.session: _http_resp: ClientResponse = await self.session.get(url, headers=headers) return DataResponseFactory.of(message.id, _http_resp.status, _http_resp.content_type, - await _http_resp.read(), - message.trace_info) + await _http_resp.read(), + message.trace_info) return DataResponseFactory.of(message.id, 404, "text/plain", - str("Destination location does not exists").encode('utf-8'), message.trace_info) + str("Destination location does not exists").encode('utf-8'), message.trace_info) async def put(self, message: DataMessage, auth_user: Any) -> DataResponse: """ diff --git a/amqp/adapter/file_uploader.py b/amqp/adapter/file_uploader.py index b21694f..d8ce6d8 100644 --- a/amqp/adapter/file_uploader.py +++ b/amqp/adapter/file_uploader.py @@ -12,6 +12,7 @@ from amqp.router.router_base import AMQ_REPLY_TO_EXCHANGE async def publish_file_to_rabbitmq( channel: AbstractRobustChannel, file_path: Path, + future: asyncio.Future[str] = None, 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 @@ -70,4 +71,6 @@ async def publish_file_to_rabbitmq( logging.debug(f"Published chunk {_offset}/{_file_size} to {exchange_name}:{_queue_name}") print(f"File {file_path} published to RabbitMQ exchange {exchange_name}, queue {_queue_name}") + if future: + future.set_result(_queue_name) return _queue_name diff --git a/amqp/service/amq_message_handler.py b/amqp/service/amq_message_handler.py index 24d6dae..d7b073d 100644 --- a/amqp/service/amq_message_handler.py +++ b/amqp/service/amq_message_handler.py @@ -136,7 +136,7 @@ class DataMessageHandler: def sync_on_message(self, message: AbstractIncomingMessage, amq_message: DataMessage, - loop: AbstractEventLoop) -> DataResponse: + 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. @@ -179,35 +179,39 @@ class DataMessageHandler: # 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}") - response: DataResponse = asyncio.run( - self.service_adapter.on_message(a_message) - ) - try: - logging.info(" [*] ######[%s] AMQ Message Reply.1, To=%s, msgId=%s", - message.delivery_tag, message.reply_to, response.id) - future: Future = asyncio.run_coroutine_threadsafe(self.send_reply(message.reply_to, response), loop=loop) - logging.info(" [*] ######[%s] AMQ Message Reply.2, To=%s, future=%s, done=%s", - message.delivery_tag, message.reply_to, future, future.done()) - 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) + ftr: Future[DataResponse] = asyncio.run_coroutine_threadsafe(self.service_adapter.on_message(a_message, future), loop) + print(f"WILL AWAIT RESP {ftr} -> {ftr.done()}") - 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}") + 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) - with lock: - current_parallel_executions -= 1 + 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}") - return response + 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): @@ -259,7 +263,9 @@ class DataMessageHandler: # map_as_string(amq_message.trace_info), context_with_span, context) if amq_message.magic == DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST: - future = executor.submit(self.sync_on_message, message, amq_message, self.loop) + 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: @@ -304,7 +310,7 @@ class DataMessageHandler: message=pika_message, routing_key=reply_to ) - logging.info(" [*] ###### Sending Reply.2 To=%s, res=%s", reply_to, _res) + 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): """