Feat #58 - fix backpressure logic and update tests
Unit test coverage / pytest (push) Failing after 1m17s
/ build-and-push (push) Successful in 1m29s

This commit is contained in:
Stanislav Hejny
2025-06-28 10:20:55 +02:00
parent a2d76aca2f
commit 7fd1c62596
13 changed files with 234 additions and 17 deletions
+1
View File
@@ -60,3 +60,4 @@ unused-code/
amq_adapter.egg-info/
build/
coverage.xml
.aider*
+8
View File
@@ -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:
"""
+28 -1
View File
@@ -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:
"""
+5
View File
@@ -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")
+4 -2
View File
@@ -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
+3
View File
@@ -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):
"""
+2 -3
View File
@@ -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 = (
+8 -1
View File
@@ -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)
+4 -1
View File
@@ -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)
+9 -2
View File
@@ -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
+66 -6
View File
@@ -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):
+1 -1
View File
@@ -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 = [
@@ -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