feat/#1 - refactor code for better design and testability

This commit was merged in pull request #11.
This commit is contained in:
Stanislav Hejny
2025-04-22 08:50:16 +01:00
parent 1de7afe4f5
commit 7f10905791
17 changed files with 1297 additions and 826 deletions
+29 -114
View File
@@ -1,11 +1,9 @@
import asyncio
import logging
import sys
import threading
import time
import concurrent.futures
import os
import traceback
from asyncio import Future, AbstractEventLoop
from typing import Dict
@@ -16,38 +14,26 @@ from opentelemetry import trace
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from opentelemetry.trace import Tracer, TraceState
from amqp.adapter.backpressure_handler import BackpressureHandler
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.data_response_factory import DataResponseFactory
from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter, AMQMessage
from amqp.adapter.file_downloader import on_message
from amqp.adapter.logging_utils import logging_error
from amqp.adapter.file_handler import StreamingFileHandler, FileHandler
from amqp.adapter.logging_utils import logging_error, logging_debug
from amqp.adapter.service_message_factory import ServiceMessageFactory
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import (
DataResponse,
DataMessage,
MultipartDataMessage,
ScalingRequest,
ScalingRequestAlert,
)
from amqp.model.service_message_type import ServiceMessageType
from amqp.rabbitmq.rabbit_mq_client import RabbitMQClient
from amqp.router.router_base import ANY_RECIPIENT
from amqp.adapter.serializer import map_as_string
# Track the number of messages currently being processed
current_parallel_executions = 0
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
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 + 2)
executor = concurrent.futures.ThreadPoolExecutor(
max_workers=BackpressureHandler.parallel_workers + 2
)
def collect_trace_info():
@@ -77,6 +63,7 @@ class DataMessageHandler:
rabbit_mq_client: RabbitMQClient,
amq_configuration: AMQConfiguration,
loop,
file_handler: FileHandler = StreamingFileHandler(),
):
self.service_adapter: CleverThisServiceAdapter = service_adapter
self.tracer = tracer
@@ -94,8 +81,10 @@ class DataMessageHandler:
self.backpressure_monitor_window = (
amq_configuration.backpressure.time_window / 1000
) # convert to seconds
global parallel_workers
self.file_handler: FileHandler = file_handler
self.backpressure_handler = BackpressureHandler(
channel=rabbit_mq_client.channel, loop=loop
)
_backpressure_treshold = amq_configuration.backpressure.threshold
"""
Determine the number of parallel workers based on the backpressure threshold and PARALLEL_WORKERS env var.
@@ -103,54 +92,8 @@ class DataMessageHandler:
as the number of parallel workers.
"""
if _backpressure_treshold > 0:
if parallel_workers <= 0:
parallel_workers = _backpressure_treshold
async def monitor_backpressure(self):
while True:
global last_data_message_time
global last_backpressure_event_time
global last_backpressure_event
now = time.time()
event: ScalingRequestAlert = ScalingRequestAlert.UPDATE
if now - last_backpressure_event_time >= self.backpressure_monitor_window:
if (
current_parallel_executions == 0
and now - last_data_message_time > self.backpressure_monitor_window
):
event = ScalingRequestAlert.IDLE
elif current_parallel_executions > parallel_workers - 2:
event = ScalingRequestAlert.OVERLOAD
last_backpressure_event_time = now
last_backpressure_event = event
last_backpressure_event_time = now
logging.info(
f"Backpressure event: {event}, threads: {current_parallel_executions} of {parallel_workers}"
)
scaling_request: ScalingRequest = ScalingRequest(
swarm_service_id,
swarm_task_id,
parallel_workers,
parallel_workers - current_parallel_executions,
event,
)
# Address the message to any adapter capable of supporting BACKPRESSURE request
scale_up_request = self.service_message_factory.of(
ServiceMessageType.BACKPRESSURE,
scaling_request.to_json(),
self.reply_to if self.reply_to else ANY_RECIPIENT,
)
asyncio.run_coroutine_threadsafe(
self.rabbit_mq_client.router.publish_service_message(
scale_up_request,
exchange=self.rabbit_mq_client.router.cm_response_exchange,
),
loop=self.loop,
)
await asyncio.sleep(self.backpressure_monitor_window) # sleep in seconds
if BackpressureHandler.parallel_workers <= 0:
BackpressureHandler.parallel_workers = _backpressure_treshold
async def inbound_data_message_callback(self, message: AbstractIncomingMessage):
threading.Thread(
@@ -185,38 +128,11 @@ class DataMessageHandler:
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(
swarm_service_id,
swarm_task_id,
parallel_workers,
parallel_workers - current_parallel_executions,
ScalingRequestAlert.OVERLOAD,
)
# Address the message to any adapter capable of supporting BACKPRESSURE request
scale_up_request = self.service_message_factory.of(
ServiceMessageType.BACKPRESSURE,
scaling_request.to_json(),
self.reply_to if self.reply_to else ANY_RECIPIENT,
)
asyncio.run_coroutine_threadsafe(
self.rabbit_mq_client.router.publish_service_message(
scale_up_request,
exchange=self.rabbit_mq_client.router.cm_response_exchange,
),
loop=_loop,
)
# update / reset time-window so that the OVERLOAD is not sent too often
global last_data_message_time
last_datadeliveryTag_message_time = time.time()
async def run_http_request_magic(
self,
message: AbstractIncomingMessage,
amq_message: DataMessage,
_loop: AbstractEventLoop,
loop: AbstractEventLoop,
) -> None:
"""
Synchronous version of on_message method. This is a wrapper around the async on_message method.
@@ -226,30 +142,31 @@ class DataMessageHandler:
:param loop: Event loop
:return: DataResponse
"""
global current_parallel_executions
# Increment the counter for parallel executions
current_parallel_executions += 1
BackpressureHandler.current_parallel_executions += 1
logging.info(
f"PARALLEL: Current Capacity [{current_parallel_executions}] of [{parallel_workers}]"
f"PARALLEL: Current Capacity [{BackpressureHandler.current_parallel_executions}] of [{BackpressureHandler.parallel_workers}]"
)
if message.reply_to:
self.reply_to = message.reply_to
if current_parallel_executions > parallel_workers - 2:
self.handle_backpressure(_loop)
if (
BackpressureHandler.current_parallel_executions
> BackpressureHandler.parallel_workers - 2
):
await self.backpressure_handler.handle_backpressure_overload_event()
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}")
_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}")
try:
logging.info(
" [*] ######[%s] AMQ Message Reply.1, To=%s, msgId=%s",
message.delivery_tag,
message.reply_to,
response.id,
_response.id,
)
await self.send_reply(message.reply_to, response)
await self.send_reply(message.reply_to, _response)
logging.info(
" [*] ######[%s] AMQ Message Reply.2, To=%s",
message.delivery_tag,
@@ -257,9 +174,7 @@ class DataMessageHandler:
)
except Exception as e:
logging_error(f"Processing message: {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
@@ -271,7 +186,7 @@ class DataMessageHandler:
os.remove(filename)
except OSError as e:
logging_error(f"Deleting file {filename}: {e}")
current_parallel_executions -= 1
BackpressureHandler.current_parallel_executions -= 1
return
except Exception as e:
logging_error(f"Request handler: {e}")
@@ -293,7 +208,7 @@ class DataMessageHandler:
f" [DBG] ######[{delivery_tag}] AMQ Message from stream: {amq_message}"
)
if amq_message is not None:
extra_data = await on_message(
extra_data = await self.file_handler.on_message(
message=message,
json_data=amq_message.body(),
rabbitmq_url=self.rabbit_mq_url,