From 7cd2960fda7877c6c8c9424a038667020e01c23b Mon Sep 17 00:00:00 2001 From: Stanislav Hejny Date: Tue, 15 Jul 2025 16:40:02 +0100 Subject: [PATCH] feat/58 - improve Backpressure algorithm --- amqp/adapter/backpressure_handler.py | 287 ++++++++++----------- amqp/config/amq_configuration.py | 2 +- amqp/service/amq_message_handler.py | 12 +- amqp/service/amq_service.py | 21 +- tests/adapter/test_backpressure_handler.py | 129 +++++---- tests/service/test_amq_service.py | 26 +- 6 files changed, 230 insertions(+), 247 deletions(-) diff --git a/amqp/adapter/backpressure_handler.py b/amqp/adapter/backpressure_handler.py index 8033f54..f566e52 100644 --- a/amqp/adapter/backpressure_handler.py +++ b/amqp/adapter/backpressure_handler.py @@ -3,6 +3,7 @@ import json import time from asyncio import AbstractEventLoop from threading import Thread +from typing import Optional import psutil from aio_pika import Message @@ -94,9 +95,8 @@ class RunningAverage: class BackpressureHandler: # Track the number of messages currently being processed - current_parallel_executions = 0 + current_load = 0 # helps detect IDLE condition - last_data_message_time = 0 # helps prevent flooding the system with backpressure events last_backpressure_event_time = 0 last_backpressure_event = ScalingRequestAlert.UPDATE @@ -116,93 +116,74 @@ class BackpressureHandler: self.exchange = None self.swarm_service_id = self.config.amq_adapter.swarm_service_id self.swarm_task_id = self.config.amq_adapter.swarm_task_id - self.parallel_workers = self.config.backpressure.threshold_threads + self.maximum_load = self.config.backpressure.threshold_load 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: - """ - Check if the loop should continue running. - Positive value of do_loop indicates the number of iterations left. - Negative value indicates the loop should run indefinitely. - """ - _val = self.do_loop != 0 - if self.do_loop > 0: - self.do_loop -= 1 - return _val - - def increase_parallel_executions(self): - """Increase the number of parallel executions""" + def increase_current_load(self): + """Increase the number of parallel executions, or current load""" logging_info( "Backpressure: Increase parallel executions, current=%s", - self.current_parallel_executions, + self.current_load, ) - self.current_parallel_executions += 1 + self.current_load += 1 - def decrease_parallel_executions(self): - """Decrease the number of parallel executions""" + def decrease_current_load(self): + """Decrease the number of parallel executions, or current load""" logging_info( "Backpressure: Decrease parallel executions, current=%s", - self.current_parallel_executions, + self.current_load, ) - if self.current_parallel_executions > 0: - self.current_parallel_executions -= 1 + if self.current_load > 0: + self.current_load -= 1 - def update_parallel_executions(self, count: int): - """Update the number of parallel executions""" - self.current_parallel_executions = count + def update_current_load(self, count: int): + """Update the number of parallel executions, or current load""" + self.current_load = count def update_last_data_message_time(self): - logging_info("Backpressure: Update last data message time") + # logging_info("Backpressure: Update last data message time") """Update the last data message time""" - self.last_data_message_time = time.time() + self.last_backpressure_event_time = time.time() async def update_backpressure_value(self, current: int, maximum: int): """ Update the current backpressure value and check for overload conditions. This method is called by the AMQService.backpressure() method to update the current parallel executions count and check for overload conditions. - + Args: current: Current value of the backpressure metric maximum: Maximum value of the backpressure metric """ logging_info( - "Backpressure: Updating backpressure value, current=%s, maximum=%s", - current, - maximum + "Backpressure: Updating backpressure value, current=%s, maximum=%s", current, maximum ) - + if maximum > 0 and maximum > self.maximum_load: + self.maximum_load = maximum + # Update the current parallel executions count - self.update_parallel_executions(current) - - # Update the last data message time to indicate activity - self.update_last_data_message_time() - + self.update_current_load(current) + # Check for overload conditions await self.check_overload_condition() - + + # Update the last data message time to indicate activity + self.update_last_data_message_time() + # If maximum has changed, update the parallel workers threshold - if maximum != self.parallel_workers and maximum > 0: + if maximum != self.maximum_load and maximum > 0: logging_info( "Backpressure: Updating maximum parallel workers from %s to %s", - self.parallel_workers, - maximum + self.maximum_load, + maximum, ) - self.parallel_workers = maximum + self.maximum_load = maximum - def start_backpressure_monitor(self) -> Thread: - # Check if CPU monitoring is enabled in configuration - cpu_monitoring_enabled = self.config.backpressure.cpu_monitoring_enabled - - if not cpu_monitoring_enabled: - logging_info("Backpressure CPU monitoring is disabled by configuration") - return None - + def start_backpressure_monitor(self) -> Optional[Thread]: # Start the Backpressure monitor loop self.thread = Thread(target=self.backpressure_monitor_loop) self.thread.daemon = True # This makes it a daemon thread @@ -215,50 +196,54 @@ class BackpressureHandler: or if the service has been idle for too long (IDLE). """ current_time = time.time() - self.update_last_data_message_time() - + last_event_delta = current_time - self.last_backpressure_event_time logging_info( "Backpressure: Check conditions, current=%s, max=%s, last_activity=%s", - self.current_parallel_executions, - self.parallel_workers, - current_time - self.last_data_message_time, + self.current_load, + self.maximum_load, + last_event_delta, ) - + # Check for OVERLOAD condition - if self.current_parallel_executions >= self.parallel_workers - 1: + if self.current_load >= round(0.8 * self.maximum_load): # Check if the last backpressure event was not an overload if self.last_backpressure_event != ScalingRequestAlert.OVERLOAD: # Trigger the overload event await self.handle_backpressure_overload_event() + # update / reset time-window so that the OVERLOAD is not sent too often self.last_backpressure_event_time = current_time self.last_backpressure_event = ScalingRequestAlert.OVERLOAD - + # Check for IDLE condition - elif self.current_parallel_executions == 0: + elif self.current_load == 0: idle_duration = self.config.backpressure.idle_duration # Check if service has been idle for longer than the configured duration - if (current_time - self.last_data_message_time > idle_duration and - (current_time - self.last_backpressure_event_time > self.config.backpressure.time_window or - self.last_backpressure_event != ScalingRequestAlert.IDLE)): - + if ( + last_event_delta > idle_duration + and self.last_backpressure_event == ScalingRequestAlert.IDLE + ): + logging_info( "Backpressure: Service has been idle for %s seconds (threshold: %s seconds)", - current_time - self.last_data_message_time, - idle_duration + last_event_delta, + idle_duration, ) - + # Trigger the idle event await self._handle_backpressure_idle_event() + # update / reset time-window so that the IDLE is not sent too often self.last_backpressure_event_time = current_time + else: + # Don't send IDLE right away when the usage drops to zero / below threshold, but wait for a while self.last_backpressure_event = ScalingRequestAlert.IDLE - + # If neither OVERLOAD nor IDLE, and it's time for an update, send UPDATE event - elif (current_time - self.last_backpressure_event_time > self.config.backpressure.time_window): + elif last_event_delta > self.config.backpressure.time_window: _scaling_request: ScaleRequestV1 = ScaleRequestV1( self.swarm_service_id, self.swarm_task_id, - self.parallel_workers, - self.parallel_workers - self.current_parallel_executions, + self.maximum_load, + self.maximum_load - self.current_load, ScalingRequestAlert.UPDATE, ) # Address the message to any adapter capable of supporting BACKPRESSURE request @@ -266,8 +251,12 @@ class BackpressureHandler: self.last_backpressure_event_time = current_time self.last_backpressure_event = ScalingRequestAlert.UPDATE + self.update_last_data_message_time() + async def _backpressure_monitor(self): """Monitor the backpressure conditions and trigger events accordingly""" + # Check if CPU monitoring is enabled in configuration + cpu_monitoring_enabled = self.config.backpressure.cpu_monitoring_enabled _time_window_millis = self.config.backpressure.time_window _idle_duration_millis = self.config.backpressure.idle_duration _overload_duration_millis = self.config.backpressure.cpu_overload_duration @@ -281,73 +270,76 @@ class BackpressureHandler: _cpu_usage_state = CPU_ACTIVE _cpu_usage_changed = time.time() while self._do_loop(): - _current_cpu_usage = psutil.cpu_percent(interval=None) - self.average_cpu_usage.add(_current_cpu_usage) - _average_cpu_usage = self.average_cpu_usage.get_average() - _now = time.time() - # logging_info( - # "Backpressure: current CPU[pct]=%s, avg=%s, state=%s, last dataMsg=%s, last event=%s, last eventTime=%s", - # _current_cpu_usage, - # _average_cpu_usage, - # _cpu_usage_state, - # self.last_data_message_time, - # self.last_backpressure_event, - # self.last_backpressure_event_time, - # ) - if _average_cpu_usage < IDLE_THRESHOLD: - if _cpu_usage_state != CPU_IDLE: - _cpu_usage_changed = _now - _cpu_usage_state = CPU_IDLE - elif _average_cpu_usage > OVERLOAD_THRESHOLD: - if _cpu_usage_state != CPU_OVERLOAD: - _cpu_usage_changed = _now - _cpu_usage_state = CPU_OVERLOAD - else: - if _cpu_usage_state != CPU_ACTIVE: - _cpu_usage_changed = _now - _cpu_usage_state = CPU_ACTIVE + if cpu_monitoring_enabled: + _current_cpu_usage = psutil.cpu_percent(interval=None) + self.average_cpu_usage.add(_current_cpu_usage) + _average_cpu_usage = self.average_cpu_usage.get_average() + _now = time.time() + # logging_info( + # "Backpressure: current CPU[pct]=%s, avg=%s, state=%s, last dataMsg=%s, last event=%s, last eventTime=%s", + # _current_cpu_usage, + # _average_cpu_usage, + # _cpu_usage_state, + # self.last_backpressure_event, + # self.last_backpressure_event_time, + # ) + if _average_cpu_usage < IDLE_THRESHOLD: + if _cpu_usage_state != CPU_IDLE: + _cpu_usage_changed = _now + _cpu_usage_state = CPU_IDLE + elif _average_cpu_usage > OVERLOAD_THRESHOLD: + if _cpu_usage_state != CPU_OVERLOAD: + _cpu_usage_changed = _now + _cpu_usage_state = CPU_OVERLOAD + else: + if _cpu_usage_state != CPU_ACTIVE: + _cpu_usage_changed = _now + _cpu_usage_state = CPU_ACTIVE - # Check if the current time exceeds the last backpressure event time - if ( - _cpu_usage_state == CPU_OVERLOAD - and _now - _cpu_usage_changed > _overload_duration_millis - and ( - _now - self.last_backpressure_event_time > _time_window_millis - or self.last_backpressure_event != ScalingRequestAlert.OVERLOAD - ) - ): - # Trigger the overload event - await self.handle_backpressure_overload_event() - self.last_backpressure_event_time = _now - self.last_backpressure_event = ScalingRequestAlert.OVERLOAD - elif ( - _cpu_usage_state == CPU_IDLE - and _now - _cpu_usage_changed > _idle_duration_millis - and _now - self.last_data_message_time > _idle_duration_millis - and ( - _now - self.last_backpressure_event_time > _time_window_millis - or self.last_backpressure_event != ScalingRequestAlert.IDLE - ) - ): - # Trigger the idle event - await self._handle_backpressure_idle_event() - self.last_backpressure_event = ScalingRequestAlert.IDLE - self.last_backpressure_event_time = _now - else: - # Trigger update event - if _now - self.last_backpressure_event_time > _time_window_millis: - # Trigger the update event - _scaling_request: ScaleRequestV1 = ScaleRequestV1( - self.swarm_service_id, - self.swarm_task_id, - 100, - _average_cpu_usage, - ScalingRequestAlert.UPDATE, + # Check if the current time exceeds the last backpressure event time + if ( + _cpu_usage_state == CPU_OVERLOAD + and _now - _cpu_usage_changed > _overload_duration_millis + and ( + _now - self.last_backpressure_event_time > _time_window_millis + or self.last_backpressure_event != ScalingRequestAlert.OVERLOAD ) - # Address the message to any adapter capable of supporting BACKPRESSURE request - await self.publish_backpressure_request(_scaling_request) + ): + # Trigger the overload event + await self.handle_backpressure_overload_event() self.last_backpressure_event_time = _now - self.last_backpressure_event = ScalingRequestAlert.UPDATE + self.last_backpressure_event = ScalingRequestAlert.OVERLOAD + elif ( + _cpu_usage_state == CPU_IDLE + and _now - _cpu_usage_changed > _idle_duration_millis + and ( + _now - self.last_backpressure_event_time > _time_window_millis + or self.last_backpressure_event != ScalingRequestAlert.IDLE + ) + ): + # Trigger the idle event + await self._handle_backpressure_idle_event() + self.last_backpressure_event = ScalingRequestAlert.IDLE + self.last_backpressure_event_time = _now + else: + # Trigger update event + if _now - self.last_backpressure_event_time > _time_window_millis: + # Trigger the update event + _scaling_request: ScaleRequestV1 = ScaleRequestV1( + self.swarm_service_id, + self.swarm_task_id, + 100, + _average_cpu_usage, + ScalingRequestAlert.UPDATE, + ) + # Address the message to any adapter capable of supporting BACKPRESSURE request + await self.publish_backpressure_request(_scaling_request) + self.last_backpressure_event_time = _now + self.last_backpressure_event = ScalingRequestAlert.UPDATE + else: + # If CPU monitoring is disabled, just check the current load + await self.check_overload_condition() + await asyncio.sleep(_monitor_interval) def backpressure_monitor_loop(self): @@ -360,14 +352,12 @@ class BackpressureHandler: scaling_request: ScaleRequestV1 = ScaleRequestV1( self.swarm_service_id, self.swarm_task_id, - self.parallel_workers, - self.parallel_workers - self.current_parallel_executions, + self.maximum_load, + self.maximum_load - self.current_load, ScalingRequestAlert.OVERLOAD, ) # Address the message to any management service instance capable of supporting BACKPRESSURE request await self.publish_backpressure_request(scaling_request) - # update / reset time-window so that the OVERLOAD is not sent too often - self.last_data_message_time = time.time() async def _handle_backpressure_idle_event(self): logging_warning("Backpressure: Service is idle.") @@ -375,14 +365,12 @@ class BackpressureHandler: scaling_request: ScaleRequestV1 = ScaleRequestV1( self.swarm_service_id, self.swarm_task_id, - self.parallel_workers, - self.parallel_workers - self.current_parallel_executions, + self.maximum_load, + self.maximum_load - self.current_load, ScalingRequestAlert.IDLE, ) # Address the message to any adapter capable of supporting BACKPRESSURE request await self.publish_backpressure_request(scaling_request) - # update / reset time-window so that the IDLE is not sent too often - self.last_data_message_time = time.time() async def publish_backpressure_request(self, scaling_request: ScaleRequestV1): # Publish the backpressure request to the management service @@ -429,3 +417,14 @@ class BackpressureHandler: if BackpressureHandler._callback_list is not None: BackpressureHandler._callback_list["_wrap_rabbit_mq_api"] = 1 await await_future(asyncio.run_coroutine_threadsafe(_wrap_rabbit_mq_api(), self.loop)) + + def _do_loop(self) -> bool: + """ + Helper function for unit tests to perform several loops only. Check if the loop should continue running. + Positive value of do_loop indicates the number of iterations left. + Negative value indicates the loop should run indefinitely. + """ + _val = self.do_loop != 0 + if self.do_loop > 0: + self.do_loop -= 1 + return _val diff --git a/amqp/config/amq_configuration.py b/amqp/config/amq_configuration.py index 3f96fb3..1ed3113 100644 --- a/amqp/config/amq_configuration.py +++ b/amqp/config/amq_configuration.py @@ -150,7 +150,7 @@ class Backpressure: :param config: config parser to read configuration values """ - self.threshold_threads = config.getint( + self.threshold_load = config.getint( "CleverMicro-AMQ", self.PREFIX + "threshold", fallback=int(os.getenv("PARALLEL_WORKERS", 1)), diff --git a/amqp/service/amq_message_handler.py b/amqp/service/amq_message_handler.py index b30429c..630f1ed 100644 --- a/amqp/service/amq_message_handler.py +++ b/amqp/service/amq_message_handler.py @@ -124,10 +124,10 @@ class DataMessageHandler: :return: DataResponse """ # Increment the counter for parallel executions and check resource availability - self.backpressure_handler.increase_parallel_executions() + self.backpressure_handler.increase_current_load() await self.backpressure_handler.check_overload_condition() logging_info( - f"PARALLEL: Current Capacity [{self.backpressure_handler.current_parallel_executions}] of [{self.backpressure_handler.parallel_workers}]" + f"PARALLEL: Current Capacity [{self.backpressure_handler.current_load}] of [{self.backpressure_handler.maximum_load}]" ) if message.reply_to: self.reply_to = message.reply_to @@ -155,7 +155,7 @@ class DataMessageHandler: except Exception as e: logging_error(f"Request handler: {e}") - self.backpressure_handler.decrease_parallel_executions() + self.backpressure_handler.decrease_current_load() await self.backpressure_handler.check_overload_condition() async def reconstitute_data_message( @@ -208,12 +208,12 @@ class DataMessageHandler: # 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() + global last_backpressure_event_time + last_backpressure_event_time = time.time() logging_info( "######[%s] Done, single ACK: %sms.", delivery_tag, - last_data_message_time - start, + last_backpressure_event_time - start, ) async def send_reply(self, reply_to: str, response: DataResponse): diff --git a/amqp/service/amq_service.py b/amqp/service/amq_service.py index a400057..e9cd7ed 100644 --- a/amqp/service/amq_service.py +++ b/amqp/service/amq_service.py @@ -83,24 +83,27 @@ class AMQService: logging.info("***********************************************") self.loop.run_forever() - def backpressure(self, current: int, maximum: int = -1): + def backpressure(self, current_availability: 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 + based on the current_availability and maximum values. + - OVERLOAD event is triggered immediately when the current_availability 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 + - IDLE event is triggered when the current_availability 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'. + :param current_availability: Currently available caopacity. + :param maximum: Maximum value of the backpressure metric. If undefined or set to lt 0 (e.g. -1), + the maximum is than hgdetermined from max value of current_availability seen over the time. """ 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) + maximum = self.amq_configuration.backpressure.threshold_load + asyncio.run_coroutine_threadsafe( + self.backpressure_handler.update_backpressure_value(current_availability, maximum), + loop=self.backpressure_handler.loop, + ) def rpc(self, service_name: str, fq_method_name: str, **kwargs) -> Future: """ diff --git a/tests/adapter/test_backpressure_handler.py b/tests/adapter/test_backpressure_handler.py index ea977a7..fafb42e 100644 --- a/tests/adapter/test_backpressure_handler.py +++ b/tests/adapter/test_backpressure_handler.py @@ -13,9 +13,6 @@ 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, @@ -289,38 +286,38 @@ class TestBackpressureHandler: assert handler.config == mock_config assert handler.swarm_service_id == "clever-amqp-service" assert handler.swarm_task_id == "python-adapter" - assert handler.parallel_workers == mock_config.backpressure.threshold_threads + assert handler.maximum_load == mock_config.backpressure.threshold_load assert handler.average_cpu_usage is None - def test_increase_parallel_executions(self, mock_channel, mock_loop, mock_config): + def test_increase_current_load(self, mock_channel, mock_loop, mock_config): """ - Test the increase_parallel_executions method. + Test the increase_current_load method. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - initial_executions = handler.current_parallel_executions - handler.increase_parallel_executions() - assert handler.current_parallel_executions == initial_executions + 1 + initial_load = handler.current_load + handler.increase_current_load() + assert handler.current_load == initial_load + 1 - def test_decrease_parallel_executions(self, mock_channel, mock_loop, mock_config): + def test_decrease_current_load(self, mock_channel, mock_loop, mock_config): """ - Test the decrease_parallel_executions method. + Test the decrease_current_load method. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - handler.current_parallel_executions = 5 - handler.decrease_parallel_executions() - assert handler.current_parallel_executions == 4 + handler.current_load = 5 + handler.decrease_current_load() + assert handler.current_load == 4 - handler.current_parallel_executions = 0 - handler.decrease_parallel_executions() - assert handler.current_parallel_executions == 0 # Should not go below 0 + handler.current_load = 0 + handler.decrease_current_load() + assert handler.current_load == 0 # Should not go below 0 - def test_update_parallel_executions(self, mock_channel, mock_loop, mock_config): + def test_update_current_load(self, mock_channel, mock_loop, mock_config): """ - Test the update_parallel_executions method. + Test the update_current_load method. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - handler.update_parallel_executions(100) - assert handler.current_parallel_executions == 100 + handler.update_current_load(100) + assert handler.current_load == 100 def test_update_last_data_message_time(self, mock_channel, mock_loop, mock_config): """ @@ -330,7 +327,7 @@ class TestBackpressureHandler: with patch("time.time") as mock_time: mock_time.return_value = 12345 handler.update_last_data_message_time() - assert handler.last_data_message_time == 12345 + assert handler.last_backpressure_event_time == 12345 def test_start_backpressure_monitor(self, mock_channel, mock_loop, mock_config): """ @@ -345,19 +342,25 @@ class TestBackpressureHandler: @pytest.mark.asyncio async def test_check_overload_condition_no_overload(self, mock_channel, mock_loop, mock_config): """ - Test the check_overload_condition method when there is no overload. + Test the check_overload_condition method when there is no overload and only short time elapsed from last message """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - handler.current_parallel_executions = 1 - handler.parallel_workers = 5 + handler.current_load = 1 + handler.maximum_load = 5 handler.last_backpressure_event = ScalingRequestAlert.UPDATE handler.publish_backpressure_request = AsyncMock() + handler.last_backpressure_event_time = 12340 with patch("time.time") as mock_time: mock_time.return_value = 12345 - with patch.object(handler, "handle_backpressure_overload_event") as mock_handle_overload: + with patch.object( + handler, "handle_backpressure_overload_event" + ) as mock_handle_overload: await handler.check_overload_condition() mock_handle_overload.assert_not_called() - if 12345 - handler.last_backpressure_event_time > mock_config.backpressure.time_window: + if ( + 12345 - handler.last_backpressure_event_time + > mock_config.backpressure.time_window + ): handler.publish_backpressure_request.assert_called_once() assert handler.last_backpressure_event == ScalingRequestAlert.UPDATE assert handler.last_backpressure_event_time == 12345 @@ -370,7 +373,7 @@ class TestBackpressureHandler: Test the check_overload_condition method when there is an overload. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - handler.current_parallel_executions = 9 # Assuming parallel_workers is 10 + handler.current_load = 9 # Assuming maximum_load is 10 handler.last_backpressure_event = ScalingRequestAlert.UPDATE handler.publish_backpressure_request = AsyncMock() with patch.object(handler, "handle_backpressure_overload_event") as mock_handle_overload: @@ -390,7 +393,7 @@ class TestBackpressureHandler: Test the check_overload_condition method when the last event was already an overload. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - handler.current_parallel_executions = 10 + handler.current_load = 10 handler.last_backpressure_event = ScalingRequestAlert.OVERLOAD handler.publish_backpressure_request = AsyncMock() with patch.object(handler, "handle_backpressure_overload_event") as mock_handle_overload: @@ -418,7 +421,7 @@ class TestBackpressureHandler: # 1. Simulate an overload condition: handler.last_backpressure_event_time = 0 # Force event trigger - handler.last_data_message_time = 0 + handler.last_backpressure_event_time = 0 handler.do_loop = 1 handler.config.backpressure.cpu_overload_duration = ( -1 @@ -432,7 +435,7 @@ class TestBackpressureHandler: ) # Mock the RunningAverage in IDLE state handler.average_cpu_usage.is_filled = True # Together with 0 values ensures IDLE state handler.last_backpressure_event_time = 0 - handler.last_data_message_time = 0 + handler.last_backpressure_event_time = 0 handler.do_loop = 1 handler.config.backpressure.idle_duration = ( -1 @@ -445,7 +448,7 @@ class TestBackpressureHandler: handler._handle_backpressure_idle_event.call_count = 0 handler.average_cpu_usage = RunningAverage(100) # Mock the RunningAverage in pritine state handler.last_backpressure_event_time = 0 - handler.last_data_message_time = 0 + handler.last_backpressure_event_time = 0 handler.do_loop = 1 await handler._backpressure_monitor() assert ( @@ -465,7 +468,6 @@ class TestBackpressureHandler: mock_time.return_value = 12345 await handler.handle_backpressure_overload_event() assert handler.publish_backpressure_request.call_count == 1 - assert handler.last_data_message_time == 12345 @pytest.mark.asyncio async def test_handle_backpressure_idle_event(self, mock_channel, mock_loop, mock_config): @@ -478,12 +480,11 @@ class TestBackpressureHandler: mock_time.return_value = 12345 await handler._handle_backpressure_idle_event() 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): + async def test_update_backpressure_value_overload(self, mock_channel, mock_loop, mock_config): """ - Test the check_resource_availability method when resource usage is in OVERLOAD state. + Test the update_backpressure_value method when resource usage is in OVERLOAD state. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler.handle_backpressure_overload_event = AsyncMock() @@ -492,17 +493,15 @@ class TestBackpressureHandler: # 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 + handler.last_backpressure_event_time = 10 # time of last message + handler._resource_usage_changed = 10 # time of last change 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) + await handler.update_backpressure_value(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}") @@ -513,12 +512,11 @@ class TestBackpressureHandler: 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): + async def test_update_backpressure_value_idle(self, mock_channel, mock_loop, mock_config): """ - Test the check_resource_availability method when resource usage is in IDLE state. + Test the update_backpressure_value method when resource usage is in IDLE state. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler._handle_backpressure_idle_event = AsyncMock() @@ -528,17 +526,16 @@ class TestBackpressureHandler: # 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) + await handler.update_backpressure_value(0, 10) # 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) + await handler.update_backpressure_value(0, 10) # Verify IDLE event was triggered handler._handle_backpressure_idle_event.assert_called_once() @@ -548,9 +545,9 @@ class TestBackpressureHandler: ) @pytest.mark.asyncio - async def test_check_resource_availability_update(self, mock_channel, mock_loop, mock_config): + async def test_update_backpressure_value_update(self, mock_channel, mock_loop, mock_config): """ - Test the check_resource_availability method when resource usage is in ACTIVE state. + Test the update_backpressure_value method when resource usage is in ACTIVE state. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler.publish_backpressure_request = AsyncMock() @@ -564,13 +561,12 @@ class TestBackpressureHandler: 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) + await handler.update_backpressure_value(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( @@ -622,28 +618,28 @@ class TestBackpressureHandler: Test the update_backpressure_value method. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - handler.update_parallel_executions = MagicMock() + handler.update_current_load = MagicMock() handler.update_last_data_message_time = MagicMock() handler.check_overload_condition = AsyncMock() # Test with current value and same maximum - await handler.update_backpressure_value(5, handler.parallel_workers) - handler.update_parallel_executions.assert_called_once_with(5) + await handler.update_backpressure_value(5, handler.maximum_load) + handler.update_current_load.assert_called_once_with(5) handler.update_last_data_message_time.assert_called_once() handler.check_overload_condition.assert_called_once() # Reset mocks - handler.update_parallel_executions.reset_mock() + handler.update_current_load.reset_mock() handler.update_last_data_message_time.reset_mock() handler.check_overload_condition.reset_mock() # Test with current value and different maximum - original_workers = handler.parallel_workers + original_workers = handler.maximum_load await handler.update_backpressure_value(3, original_workers + 5) - handler.update_parallel_executions.assert_called_once_with(3) + handler.update_current_load.assert_called_once_with(3) handler.update_last_data_message_time.assert_called_once() handler.check_overload_condition.assert_called_once() - assert handler.parallel_workers == original_workers + 5 + assert handler.maximum_load == original_workers + 5 @pytest.mark.asyncio async def test_check_overload_condition_idle(self, mock_channel, mock_loop, mock_config): @@ -654,8 +650,8 @@ class TestBackpressureHandler: handler._handle_backpressure_idle_event = AsyncMock() # Set up initial state for IDLE condition - handler.current_parallel_executions = 0 - handler.last_backpressure_event = ScalingRequestAlert.UPDATE + handler.current_load = 0 + handler.last_backpressure_event = ScalingRequestAlert.IDLE with patch("time.time") as mock_time: # Set current time @@ -663,11 +659,8 @@ class TestBackpressureHandler: mock_time.return_value = current_time # Set last data message time to be older than idle_duration - handler.last_data_message_time = ( - current_time - mock_config.backpressure.idle_duration - 10 - ) handler.last_backpressure_event_time = ( - current_time - mock_config.backpressure.time_window - 10 + current_time - mock_config.backpressure.idle_duration - 10 ) # Call the method @@ -687,8 +680,8 @@ class TestBackpressureHandler: handler.publish_backpressure_request = AsyncMock() # Set up initial state for UPDATE condition - handler.current_parallel_executions = 2 - handler.parallel_workers = 10 + handler.current_load = 2 + handler.maximum_load = 10 handler.last_backpressure_event = ScalingRequestAlert.IDLE # Coming from IDLE state with patch("time.time") as mock_time: @@ -720,8 +713,8 @@ class TestBackpressureHandler: thread = handler.start_backpressure_monitor() # Verify no thread was started - assert thread is None - assert not hasattr(handler, "thread") or handler.thread is None + assert thread is not None + assert hasattr(handler, "thread") or handler.thread is not None def test_start_backpressure_monitor_enabled(self, mock_channel, mock_loop, mock_config): """ diff --git a/tests/service/test_amq_service.py b/tests/service/test_amq_service.py index 7dc4ad6..31bcac9 100644 --- a/tests/service/test_amq_service.py +++ b/tests/service/test_amq_service.py @@ -1,8 +1,7 @@ import unittest -from unittest.mock import AsyncMock, MagicMock, Mock, patch +from unittest.mock import AsyncMock, Mock, patch from amqp.config.amq_configuration import AMQAdapter, AMQConfiguration -from amqp.model.model import DataMessage from amqp.service.amq_service import AMQService @@ -14,6 +13,11 @@ class TestAMQService(unittest.IsolatedAsyncioTestCase): self.mock_adapter.route_mapping = {} self.mock_adapter.swarm_service_id = "test-swarm-id" self.mock_adapter.swarm_task_id = "test-task-id" + self.mock_adapter.consul_host = "consul" + self.mock_adapter.consul_port = 4500 + self.mock_adapter.consul_counter_key = "services/ids/counter" + self.mock_adapter.consul_max_retries = 5 + self.mock_adapter.consul_initial_retry_delay = 0.2 self.mock_config = Mock(spec=AMQConfiguration) self.mock_config.amq_adapter = self.mock_adapter @@ -29,7 +33,7 @@ class TestAMQService(unittest.IsolatedAsyncioTestCase): # Add backpressure configuration self.mock_backpressure = Mock() - self.mock_backpressure.threshold_threads = 10 + self.mock_backpressure.threshold_load = 10 self.mock_config.backpressure = self.mock_backpressure # Create service instance with mocked configuration @@ -145,22 +149,6 @@ class TestAMQService(unittest.IsolatedAsyncioTestCase): # Verify that client was reinitialized mock_client.init.assert_called_once() - def test_send_message(self): - # Create mock message and future - mock_message = Mock(spec=DataMessage) - mock_message.id.return_value = "test-id" - mock_future = MagicMock() - - # Setup mock message handler - self.service.amq_data_message_handler = Mock() - self.service.amq_data_message_handler.outstanding = {} - - # Send message - self.service.send_message(mock_message, mock_future) - - # Verify message was added to outstanding - self.assertEqual(self.service.amq_data_message_handler.outstanding["test-id"], mock_future) - def test_remove_outstanding_future(self): # Setup mock message handler with outstanding message self.service.amq_data_message_handler = Mock()