amq-adapter-python/amqp/adapter/backpressure_handler.py

import asyncio
import json
import time
from asyncio import AbstractEventLoop
from threading import Thread
from typing import Optional

from aio_pika import Message
from aio_pika.abc import AbstractRobustChannel

from amqp.adapter.logging_utils import logging_info, logging_warning
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import ScalingRequestAlert
from amqp.router.utils import await_future, await_result


class ScaleRequestV1:

    def __init__(
        self,
        serviceId: str,
        taskId: str,
        max_availability: int,
        current_availability: int,
        requestType: ScalingRequestAlert,
    ):
        self.version = 1  # Version of the request, currently 1
        self.serviceId = serviceId
        self.taskId = taskId
        self.max_availability = max_availability
        self.current_availability = current_availability
        self.requestType = requestType
        self.thread = None

    def to_json(self) -> str:
        """Converts the object to a JSON string matching the Java structure"""
        return json.dumps(
            {
                "version": self.version,
                "serviceId": self.serviceId,
                "taskId": self.taskId,
                "maxAvailability": self.max_availability,
                "currentAvailability": self.current_availability,
                "requestType": self.requestType.value,  # Using .value for Enum serialization
            },
            indent=2,
        )


class BackpressureHandler:
    # Track the number of messages currently being processed
    current_availability = 0
    # helps detect IDLE condition
    # helps prevent flooding the system with backpressure events
    last_backpressure_event_time = 0
    last_backpressure_event = ScalingRequestAlert.UPDATE
    # _callback_list is used in unit tests to record the invoked callbacks
    _callback_list = None

    def __init__(
        self,
        channel: AbstractRobustChannel,
        loop: AbstractEventLoop,
        config: AMQConfiguration,
    ):
        self.lock = None
        self.channel = channel
        self.loop = loop
        self.config = config
        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.do_loop = -1
        self._resource_usage_changed = 0
        self._resource_average_value = 0
        self._last_resource_max_value = 100  # Default max value for CPU usage
        self.max_availability = -1
        self.current_availability = 1
        self.max_load = 1
        self.current_load = 0

    def increase_current_load(self):
        """Increase the number of parallel executions, or current_availability load"""
        logging_info(
            "Backpressure: Increase current load (%s) by 1",
            self.current_load,
        )
        self.current_load += 1

    def decrease_current_load(self):
        """Decrease the number of parallel executions, or current load"""
        logging_info(
            "Backpressure: Decrease current load(%s) by 1",
            self.current_load,
        )
        if self.current_load > 0:
            self.current_load -= 1

    def update_current_availability(self, count: int):
        """Update the number of parallel executions, or current load"""
        self.current_availability = count

    def update_last_backpressure_event_time(self):
        # logging_info("Backpressure: Update last data message time")
        """Update the last data message time"""
        self.last_backpressure_event_time = time.time()

    async def update_backpressure_value(self, current_availability: 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_availability: Current value of the backpressure metric
            maximum: Maximum value of the backpressure metric
        """
        logging_info(
            "Backpressure: Updating backpressure value, current_availability=%s, maximum=%s",
            current_availability,
            maximum,
        )
        if maximum > 0 and maximum > self.max_availability:
            self.max_availability = maximum
        # Update the current_availability / parallel executions count
        self.update_current_availability(current_availability)
        # Check for overload conditions
        await self.check_overload_condition()

    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
        self.thread.start()
        return self.thread

    async def check_overload_condition(self):
        """
        Check if the current_availability availability is too low (OVERLOAD)
        or if the service has been idle for too long with high availability (IDLE).

        Note: current_availability represents available capacity, not used capacity.
        - Low availability (close to 0) means OVERLOAD
        - High availability (close to maximum) with no activity means IDLE
        """
        current_time = time.time()
        last_event_delta: float = current_time - self.last_backpressure_event_time

        logging_info(
            "Backpressure: Check conditions, current_availability=%s, max=%s, last_activity=%s",
            self.current_availability,
            self.max_availability,
            last_event_delta,
        )

        # Check for OVERLOAD condition - low availability (less than 20% of maximum)
        if self.current_availability <= round(0.1 * self.max_availability):
            # Check if the last backpressure event was not an overload
            if (
                self.last_backpressure_event != ScalingRequestAlert.OVERLOAD
                or last_event_delta > self.config.backpressure.idle_duration
            ):
                # 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 - high availability (more than 80% of maximum) with no activity
        elif self.current_availability >= round(0.8 * self.max_availability):
            idle_duration = self.config.backpressure.idle_duration
            # Check if service has been idle for longer than the configured duration
            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)",
                    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 availability increases, but wait for a while, send UPDATE instead
                self.last_backpressure_event = ScalingRequestAlert.IDLE
                _scaling_request: ScaleRequestV1 = ScaleRequestV1(
                    self.swarm_service_id,
                    self.swarm_task_id,
                    self.max_availability,
                    self.current_availability,  # Current availability is passed directly
                    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 = current_time

        # If neither OVERLOAD nor IDLE, and it's time for an update, send UPDATE event
        elif last_event_delta > self.config.backpressure.time_window:
            _scaling_request: ScaleRequestV1 = ScaleRequestV1(
                self.swarm_service_id,
                self.swarm_task_id,
                self.max_availability,
                self.current_availability,  # Current availability is passed directly
                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 = current_time
            self.last_backpressure_event = ScalingRequestAlert.UPDATE

        self.update_last_backpressure_event_time()

    async def _backpressure_monitor(self):
        """Periodically monitor the backpressure conditions and trigger events accordingly"""
        _monitor_interval = 0.5  # Monitor every 500ms
        while self._do_loop():
            await self.check_overload_condition()
            await asyncio.sleep(_monitor_interval)

    def backpressure_monitor_loop(self):
        _loop = asyncio.new_event_loop()
        _loop.run_until_complete(self._backpressure_monitor())

    async def handle_backpressure_overload_event(self):
        logging_warning("Backpressure: Capacity close to depleted!")
        # Send an Overload event to the Management service
        scaling_request: ScaleRequestV1 = ScaleRequestV1(
            self.swarm_service_id,
            self.swarm_task_id,
            self.max_availability,
            self.current_availability,  # Current availability is passed directly
            ScalingRequestAlert.OVERLOAD,
        )
        # Address the message to any management service instance capable of supporting BACKPRESSURE request
        await self.publish_backpressure_request(scaling_request)

    async def _handle_backpressure_idle_event(self):
        logging_warning("Backpressure: Service is idle.")
        # Send an Idle event to the Management service
        scaling_request: ScaleRequestV1 = ScaleRequestV1(
            self.swarm_service_id,
            self.swarm_task_id,
            self.max_availability,
            self.current_availability,  # Current availability is passed directly
            ScalingRequestAlert.IDLE,
        )
        # Address the message to any adapter capable of supporting BACKPRESSURE request
        await self.publish_backpressure_request(scaling_request)

    async def publish_backpressure_request(self, scaling_request: ScaleRequestV1):
        # Publish the backpressure request to the management service
        logging_info(
            f"Publishing backpressure for {scaling_request.serviceId} with request type {scaling_request.requestType}"
        )
        if not self.exchange:

            async def _wrap_rabbit_mq_api_init(channel):
                _exchange = await channel.get_exchange(name="cleverthis.clevermicro.management")
                return _exchange

            if BackpressureHandler._callback_list is not None:
                BackpressureHandler._callback_list["_wrap_rabbit_mq_api_init"] = 1
            self.exchange = await await_result(
                asyncio.run_coroutine_threadsafe(_wrap_rabbit_mq_api_init(self.channel), self.loop)
            )

        if self.exchange:

            async def _wrap_rabbit_mq_api():
                if not self.channel.is_closed:
                    binary_content: bytes = scaling_request.to_json().encode("utf-8")
                    pika_message: Message = Message(
                        body=binary_content,
                        content_encoding="utf-8",
                        delivery_mode=2,
                        content_type="application/octet-stream",
                        headers=None,
                        priority=0,
                        correlation_id=None,
                    )
                    await self.exchange.publish(
                        message=pika_message, routing_key="backpressure-scaling-v1"
                    )
                    logging_info(
                        "Service Message Published to %s, msg: %s",
                        self.exchange.name,
                        str(binary_content),
                    )
                    return True
                return False

            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