From 7fd1c62596451c1a2aab05cff7e262d4f7f90ee4 Mon Sep 17 00:00:00 2001 From: Stanislav Hejny Date: Sat, 28 Jun 2025 10:20:55 +0200 Subject: [PATCH] Feat #58 - fix backpressure logic and update tests --- .gitignore | 1 + amqp/adapter/backpressure_handler.py | 8 ++ amqp/adapter/data_message_factory.py | 29 ++++++- amqp/config/amq_configuration.py | 5 ++ amqp/config/application.properties | 6 +- amqp/model/model.py | 3 + amqp/rabbitmq/rabbit_mq_client.py | 5 +- amqp/router/route_database.py | 9 +- amqp/router/router_base.py | 5 +- amqp/service/amq_message_handler.py | 11 ++- amqp/service/amq_service.py | 72 ++++++++++++++-- pyproject.toml | 2 +- tests/adapter/test_backpressure_handler.py | 95 ++++++++++++++++++++++ 13 files changed, 234 insertions(+), 17 deletions(-) diff --git a/.gitignore b/.gitignore index 9f169a6..8644835 100644 --- a/.gitignore +++ b/.gitignore @@ -60,3 +60,4 @@ unused-code/ amq_adapter.egg-info/ build/ coverage.xml +.aider* diff --git a/amqp/adapter/backpressure_handler.py b/amqp/adapter/backpressure_handler.py index 97fd5a0..ea0cf72 100644 --- a/amqp/adapter/backpressure_handler.py +++ b/amqp/adapter/backpressure_handler.py @@ -13,6 +13,10 @@ from amqp.config.amq_configuration import AMQConfiguration from amqp.model.model import ScalingRequestAlert from amqp.router.utils import await_future, await_result +CPU_IDLE, CPU_ACTIVE, CPU_OVERLOAD = 0, 1, 2 +IDLE_THRESHOLD, OVERLOAD_THRESHOLD = 0.15, 0.90 +RESOURCE_UNSET, RESOURCE_IDLE, RESOURCE_ACTIVE, RESOURCE_OVERLOAD = -1, 0, 1, 2 + class ScaleRequestV1: @@ -115,6 +119,10 @@ class BackpressureHandler: self.parallel_workers = self.config.backpressure.threshold_threads self.average_cpu_usage = None self.do_loop = -1 + self._resource_usage_state = RESOURCE_UNSET + self._resource_usage_changed = 0 + self._resource_average_value = 0 + self._last_resource_max_value = 100 # Default max value for CPU usage def _do_loop(self) -> bool: """ diff --git a/amqp/adapter/data_message_factory.py b/amqp/adapter/data_message_factory.py index b42f011..02a17ad 100644 --- a/amqp/adapter/data_message_factory.py +++ b/amqp/adapter/data_message_factory.py @@ -7,10 +7,11 @@ import json import re from datetime import datetime, timezone from io import BytesIO -from typing import Dict, List +from typing import Any, Dict, List import amqp.adapter.file_handler from amqp.adapter.file_handler import FileHandler, StreamingFileHandler +from amqp.adapter.pydantic_serializer import python_type_to_json from amqp.adapter.serializer import ( map_as_string, map_with_list_as_string, @@ -112,6 +113,32 @@ class DataMessageFactory: payload, ) + def create_rpc_message( + self, + fq_method_name: str, + swarm_id: str, + args: Dict[str, Any] = None, + ) -> DataMessage: + """ + Create a RPC message with the given parameters. + :param fq_method_name: fully qualified method name (e.g. 'com.example.ServiceClass.method_name') + :param args: arguments whose serialized form will form the request body + :return: DataMessage object + """ + package, class_name, method = fq_method_name.rsplit(".", 2) + headers = {"Content-Type": ["application/octet-stream"], "Return-Address": [swarm_id]} + trace_info = {} + return DataMessage( + DataMessageMagicByte.RPC_REQUEST.value, + self.generate_snowflake_id(), + package, + class_name, + method, + headers, + trace_info, + python_type_to_json(args) if args else b"", + ) + @staticmethod def amq_message_routing_key(message: DataMessage) -> str: """ diff --git a/amqp/config/amq_configuration.py b/amqp/config/amq_configuration.py index 383c7ef..c593c2b 100644 --- a/amqp/config/amq_configuration.py +++ b/amqp/config/amq_configuration.py @@ -40,6 +40,11 @@ class AMQAdapter: self.PREFIX + "require-authenticated-user", fallback=False, ) + self.backpressure_enabled = config.getboolean( + "CleverMicro-AMQ", + self.PREFIX + "default-backpressure-enabled", + fallback=False, + ) # these 2 shall be provided by Docker swarm, format is a random alphanumeric string self.swarm_service_id = os.getenv("SWARM_SERVICE_ID", default="clever-amqp-service") self.swarm_task_id = os.getenv("SWARM_TASK_ID", default="python-adapter") diff --git a/amqp/config/application.properties b/amqp/config/application.properties index fdfa5a8..1a70c64 100644 --- a/amqp/config/application.properties +++ b/amqp/config/application.properties @@ -14,6 +14,8 @@ cm.amq-adapter.generator-id=1 cm.amq-adapter.route-mapping= # Indicate if AMQ Adapter should respond with HTTP 401 Unauthorized if the request does not contain valid JWT token cm.amq-adapter.require-authenticated-user=false +# Indicate if AMQ Adapter should use inbuilt backpressure monitor (True) or rely only on the service backpressure updates (False) +cm.amq-adapter.default-backpressure-enabled=false # CleverMicro AMQ Adapter dispatcher, settings related to RabbitMQ connection and optional REST forwarding cm.dispatch.use-dlq=true @@ -27,7 +29,8 @@ cm.dispatch.download-dir=/tmp/downloaded_files #cm.dispatch.service-port=8080 # Backpressure monitor settings -# Define maximum number of messages of being processed in parallel before triggering backpressure OVERLOAD alert +# Define maximum number of resources that service can use before triggering backpressure OVERLOAD alert +# 'resource' is a generic term that can mean CPU, memory, parallel requests, etc. cm.backpressure.threshold=5 # Define maximum CPU usage (%) before triggering backpressure OVERLOAD alert cm.backpressure.threshold-cpu-overload=90 @@ -39,4 +42,3 @@ cm.backpressure.time-window=10 cm.backpressure.cpu-overload-duration=30 # Define the time window length, in seconds, during which, if no request is made to the Service, the system will trigger IDLE alert cm.backpressure.cpu-idle-duration=30 - diff --git a/amqp/model/model.py b/amqp/model/model.py index fe31fc4..08742b3 100644 --- a/amqp/model/model.py +++ b/amqp/model/model.py @@ -91,6 +91,9 @@ class CleverMicroMessage: def content_type(self) -> str: return self.headers().get("Content-Type", self.DEFAULT_CONTENT_TYPE)[0] + def return_address(self) -> str: + return self.headers().get("Return-Address", "")[0] + class DataMessage(CleverMicroMessage): """ diff --git a/amqp/rabbitmq/rabbit_mq_client.py b/amqp/rabbitmq/rabbit_mq_client.py index 3804a84..d839ce6 100644 --- a/amqp/rabbitmq/rabbit_mq_client.py +++ b/amqp/rabbitmq/rabbit_mq_client.py @@ -17,7 +17,6 @@ from amqp.router.router_base import ( AMQ_RPC_EXCHANGE, DLQ_EXCHANGE, ) -from amqp.router.router_consumer import RouterConsumer from amqp.router.utils import sanitize_as_rabbitmq_name @@ -25,7 +24,7 @@ class RabbitMQClient: PREFETCH_COUNT = 1 DLQ_ARGS = {"x-dead-letter-exchange": DLQ_EXCHANGE} - def __init__(self, configuration): + def __init__(self, configuration, router): self.channel: aio_pika.RobustChannel | None = None self.connection: aio_pika.RobustConnection | None = None self.rpc_exchange: aio_pika.RobustExchange | None = None @@ -37,7 +36,7 @@ class RabbitMQClient: self.amq_configuration.dispatch.amq_host, ) logging.info("*******************************************") - self.router = RouterConsumer(configuration) + self.router = router # this.meterRegistry = new PrometheusMeterRegistry(PrometheusConfig.DEFAULT); self.reply_to_exchange: AbstractRobustExchange | None = ( diff --git a/amqp/router/route_database.py b/amqp/router/route_database.py index 4f36d68..56406c6 100644 --- a/amqp/router/route_database.py +++ b/amqp/router/route_database.py @@ -1,5 +1,5 @@ import re -from typing import Set +from typing import Optional, Set from amqp.adapter.logging_utils import logging_info from amqp.model.model import AMQRoute @@ -55,6 +55,13 @@ class RouteDatabase: return route return None + def find_route_by_service(self, service_name: str) -> Optional[AMQRoute]: + logging_info(f"findRoute service_name:{service_name} in:{self.wrap_routes()}") + for route in self.routes: + if route.component_name == service_name: + return route + return None + def add_route(self, route: AMQRoute) -> bool: if route not in self.routes: self.routes.add(route) diff --git a/amqp/router/router_base.py b/amqp/router/router_base.py index dd760f1..f7cc4ab 100644 --- a/amqp/router/router_base.py +++ b/amqp/router/router_base.py @@ -1,4 +1,4 @@ -from typing import Callable, List, Set +from typing import Callable, List, Optional, Set from amqp.adapter.amq_route_factory import NULL_ROUTE, AMQRouteFactory from amqp.adapter.logging_utils import logging_debug, logging_error, logging_info @@ -49,6 +49,9 @@ class RouterBase: def find_route_by_message(self, message: DataMessage) -> AMQRoute | None: return self.route_database.find_route(message.domain(), message.path()) + def find_route_by_service(self, service_name: str) -> Optional[AMQRoute]: + return self.route_database.find_route_by_service(service_name) + def add_routes(self, message: str) -> None: try: amq_route_list = self.unwrap_route_list(message) diff --git a/amqp/service/amq_message_handler.py b/amqp/service/amq_message_handler.py index af67d0d..790a330 100644 --- a/amqp/service/amq_message_handler.py +++ b/amqp/service/amq_message_handler.py @@ -123,8 +123,12 @@ class DataMessageHandler: :param loop: Event loop :return: DataResponse """ - # Increment the counter for parallel executions + # Increment the counter for parallel executions and check resource availability self.backpressure_handler.increase_parallel_executions() + await self.backpressure_handler.check_resource_availability( + self.backpressure_handler.current_parallel_executions, + self.backpressure_handler.parallel_workers, + ) logging_info( f"PARALLEL: Current Capacity [{self.backpressure_handler.current_parallel_executions}] of [{self.backpressure_handler.parallel_workers}]" ) @@ -132,7 +136,6 @@ class DataMessageHandler: self.reply_to = message.reply_to try: - await self.backpressure_handler.check_overload_condition() _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}") @@ -156,6 +159,10 @@ class DataMessageHandler: logging_error(f"Request handler: {e}") self.backpressure_handler.decrease_parallel_executions() + await self.backpressure_handler.check_resource_availability( + self.backpressure_handler.current_parallel_executions, + self.backpressure_handler.parallel_workers, + ) async def reconstitute_data_message( self, message: AbstractIncomingMessage diff --git a/amqp/service/amq_service.py b/amqp/service/amq_service.py index 0e161ee..acd9672 100644 --- a/amqp/service/amq_service.py +++ b/amqp/service/amq_service.py @@ -14,8 +14,9 @@ from amqp.adapter.data_message_factory import DataMessageFactory from amqp.adapter.logging_utils import logging_info, logging_warning from amqp.adapter.service_message_factory import ServiceMessageFactory from amqp.config.amq_configuration import AMQConfiguration -from amqp.model.model import DataMessage +from amqp.model.model import AMQRoute, DataMessage from amqp.rabbitmq.rabbit_mq_client import RabbitMQClient +from amqp.router.router_consumer import RouterConsumer from amqp.service.amq_message_handler import DataMessageHandler from amqp.service.tracing import TracingConfig, initialize_telemetry @@ -45,7 +46,8 @@ class AMQService: self.data_message_factory = DataMessageFactory.get_instance( self.amq_configuration.amq_adapter.generator_id ) - self.rabbit_mq_client = RabbitMQClient(self.amq_configuration) + self.router = RouterConsumer(self.amq_configuration) + self.rabbit_mq_client = RabbitMQClient(self.amq_configuration, self.router) self.tracer = initialize_telemetry( tracing_config=TracingConfig( service_name=amq_configuration.amq_adapter.service_name, @@ -57,8 +59,17 @@ class AMQService: self.amq_configuration.amq_adapter.service_name ) self.backpressure_thread: Optional[Thread] = None + self.backpressure_handler: Optional[BackpressureHandler] = None + # ======================================================================= + # ======================================================================= + # ====================== A P I M e t h o d s ========================= + # ======================================================================= + # ======================================================================= def run(self): + """ + Start the AMQ service, initializing RabbitMQ connection and message handlers. + """ # self.init(loop) will initialize RabbitMQ connection asyncio.set_event_loop(self.loop) self.loop.run_until_complete(self.init(self.loop, self.service_adapter)) @@ -68,9 +79,57 @@ class AMQService: logging.info("***********************************************") self.loop.run_forever() + def backpressure(self, current: int, maximum: int = -1): + """ + Send Backpressure event to the Management Service. The Backpressure logic determines the type of the event + based on the current and maximum values. + - OVERLOAD event is triggered immediately when the current value exceeds the maximum threshold, which is + typically set to 80% of the maximum value. + - IDLE event is triggered when the current value drops below 20% of the maximum value and remains there for + a IDLE_WINDOW period of time (see BackpressureHandler). + - UPDATE event is sent if neither OVERLOAD nor IDLE conditions are met, indicating that the backpressure is + within acceptable limits. + :param current: Current value of the backpressure metric. + :param maximum: Maximum value of the backpressure metric. If undefined or set to lt 0, the maximum is set + according to the configuration value 'cm.backpressure.threshold'. + """ + if self.backpressure_handler is not None: + if maximum < 0: + maximum = self.amq_configuration.backpressure.threshold_threads + self.backpressure_handler.update_backpressure_value(current, maximum) + + def rpc(self, service_name: str, fq_method_name: str, **kwargs) -> Future: + """ + Send a Remote Procedure Call (RPC) message to the AMQ service. + :param service_name: The service name to receive the RPC message. + :param fq_method_name: The fully qualified method name to execute the RPC message. + :param kwargs: Additional keyword arguments to be passed to the RPC method. + :return: A Future that resolves when the RPC response is received. + """ + route: Optional[AMQRoute] = self.router.find_route_by_service(service_name) + if route is not None: + message: DataMessage = self.data_message_factory.create_rpc_message( + fq_method_name=fq_method_name, + swarm_id=self.amq_configuration.amq_adapter.swarm_task_id, + **kwargs, + ) + result: Future = self.amq_data_message_handler.send_rpc(route, message) + self.set_outstanding_future(message, result) + return result + + result = asyncio.Future() + result.set_exception( + NameError(f"No route exists. Service '{service_name}' not found in routing table.") + ) + return result + + # ======================================================================= + # ====================== Internal Methods ============================ + # ======================================================================= + async def init(self, loop, service_adapter): _reply_to_exchange: AbstractRobustExchange = await self.rabbit_mq_client.init(loop) - backpressure_handler = BackpressureHandler( + self.backpressure_handler = BackpressureHandler( channel=self.rabbit_mq_client.channel, loop=loop, config=self.amq_configuration, @@ -84,11 +143,12 @@ class AMQService: self.rabbit_mq_client, self.amq_configuration, loop=loop, - backpressure_handler=backpressure_handler, + backpressure_handler=self.backpressure_handler, ) await self.register_routes() await self.rabbit_mq_client.request_routes() - self.backpressure_thread = backpressure_handler.start_backpressure_monitor() + logging_info("Starting backpressure monitor thread") + self.backpressure_thread = self.backpressure_handler.start_backpressure_monitor() def cleanup(self): logging_info("AMQService.cleanup() rmqc=%s", self.rabbit_mq_client) @@ -122,7 +182,7 @@ class AMQService: ) return None - def send_message(self, message: DataMessage, future: Future): + def set_outstanding_future(self, message: DataMessage, future: Future): self.amq_data_message_handler.outstanding[str(message.id())] = future def remove_outstanding_future(self, message_id: str): diff --git a/pyproject.toml b/pyproject.toml index 79de392..34e3e43 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "amq_adapter" -version = "0.2.33" +version = "0.2.34" description = "The Python implementation of the AMQ Adapter for CleverMicro" readme = "README.md" authors = [ diff --git a/tests/adapter/test_backpressure_handler.py b/tests/adapter/test_backpressure_handler.py index 42b7219..d2b75e0 100644 --- a/tests/adapter/test_backpressure_handler.py +++ b/tests/adapter/test_backpressure_handler.py @@ -13,6 +13,9 @@ from aio_pika import Exchange from aio_pika.abc import AbstractRobustChannel from amqp.adapter.backpressure_handler import ( + RESOURCE_ACTIVE, + RESOURCE_OVERLOAD, + RESOURCE_UNSET, BackpressureHandler, RunningAverage, ScaleRequestV1, @@ -463,6 +466,98 @@ class TestBackpressureHandler: assert handler.publish_backpressure_request.call_count == 1 assert handler.last_data_message_time == 12345 + @pytest.mark.asyncio + async def test_check_resource_availability_overload(self, mock_channel, mock_loop, mock_config): + """ + Test the check_resource_availability method when resource usage is in OVERLOAD state. + """ + handler = BackpressureHandler(mock_channel, mock_loop, mock_config) + handler.handle_backpressure_overload_event = AsyncMock() + handler._handle_backpressure_idle_event = AsyncMock() + handler.publish_backpressure_request = AsyncMock() + + # Set up initial state + handler.last_backpressure_event = ScalingRequestAlert.UPDATE + handler.last_backpressure_event_time = 0 + handler._resource_usage_state = RESOURCE_ACTIVE # Start with ACTIVE state + handler._resource_usage_changed = 0 + + with patch("time.time") as mock_time: + mock_time.return_value = 100 + # Call with resource usage above OVERLOAD threshold (90% of maximum) + await handler.check_resource_availability(95, 100) + + # Debug output + print(f"Resource usage state: {handler._resource_usage_state}") + print(f"Last backpressure event: {handler.last_backpressure_event}") + print(f"Last backpressure event time: {handler.last_backpressure_event_time}") + print(f"Time window: {mock_config.backpressure.time_window}") + print("Resource usage: 95, Maximum: 100") + print(f"Overload threshold calculation: {round(0.9 * 100)}") + + # Verify OVERLOAD event was triggered + handler.handle_backpressure_overload_event.assert_called_once() + assert handler.last_backpressure_event == ScalingRequestAlert.OVERLOAD + assert handler.last_backpressure_event_time == 100 + assert handler._resource_usage_state == RESOURCE_OVERLOAD + + @pytest.mark.asyncio + async def test_check_resource_availability_idle(self, mock_channel, mock_loop, mock_config): + """ + Test the check_resource_availability method when resource usage is in IDLE state. + """ + handler = BackpressureHandler(mock_channel, mock_loop, mock_config) + handler._handle_backpressure_idle_event = AsyncMock() + handler.handle_backpressure_overload_event = AsyncMock() + handler.publish_backpressure_request = AsyncMock() + + # Set up initial state + handler.last_backpressure_event = ScalingRequestAlert.UPDATE + handler.last_backpressure_event_time = 0 + handler._resource_usage_state = RESOURCE_UNSET # Start with UNSET state + handler._resource_usage_changed = 0 + + with patch("time.time") as mock_time: + # First call to set state to IDLE + mock_time.return_value = 50 + await handler.check_resource_availability(10, 100) + + # Second call after idle_duration has passed + mock_time.return_value = 50 + mock_config.backpressure.idle_duration + 1 + await handler.check_resource_availability(10, 100) + + # Verify IDLE event was triggered + handler._handle_backpressure_idle_event.assert_called_once() + assert handler.last_backpressure_event == ScalingRequestAlert.IDLE + assert ( + handler.last_backpressure_event_time == 50 + mock_config.backpressure.idle_duration + 1 + ) + + @pytest.mark.asyncio + async def test_check_resource_availability_update(self, mock_channel, mock_loop, mock_config): + """ + Test the check_resource_availability method when resource usage is in ACTIVE state. + """ + handler = BackpressureHandler(mock_channel, mock_loop, mock_config) + handler.publish_backpressure_request = AsyncMock() + handler.handle_backpressure_overload_event = AsyncMock() + handler._handle_backpressure_idle_event = AsyncMock() + + # Set up initial state + handler.last_backpressure_event = ScalingRequestAlert.IDLE + handler.last_backpressure_event_time = 0 + + with patch("time.time") as mock_time: + mock_time.return_value = 100 + # Call with resource usage between thresholds + await handler.check_resource_availability(50, 100) + + # Verify UPDATE event was triggered + assert handler.publish_backpressure_request.call_count == 1 + assert handler.last_backpressure_event == ScalingRequestAlert.UPDATE + assert handler.last_backpressure_event_time == 100 + assert handler._resource_usage_state == RESOURCE_ACTIVE + @pytest.mark.asyncio async def test_publish_backpressure_request_exchange_exists( self, mock_channel, mock_loop, mock_config