RPC框架编写实践--熔断器

# 引用的库
import random
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple

from rap.client.model import Request, Response
from rap.client.processor.base import BaseProcessor
from rap.client.types import CLIENT_EVENT_FN
from rap.common.collect_statistics import WindowStatistics
from rap.common.exceptions import BaseRapError
from rap.common.utils import Constant, EventEnum

if TYPE_CHECKING:
    from rap.client.core import BaseClient


class CircuitBreakerExc(Exception):
    """熔断时的异常类型"""
    pass


class BaseCircuitBreakerProcessor(BaseProcessor):
    """基础熔断封装"""

    exc: Exception = NotImplementedError()

    def __init__(
        self,
        k: float = 2.0,  # google sre default
        expire: int = 180,
        interval: int = 120,  # google sre default
        prefix: str = "circuit_breaker",
        window_statistics: Optional[WindowStatistics] = None,
    ):
        """
        k 公式中k的值
        expire 每个指标数据在内存中的过期时间， 每次请求都会刷新
        interval 统计窗口时间
        prefix 以及 window statistics是内部存储实现
        """
        self._prefix: str = prefix
        self._expire: int = expire
        self._interval: int = interval
        self._window_statistics: WindowStatistics = window_statistics or WindowStatistics()
        self._probability_dict: Dict[str, float] = {}

        def upload_probability(stats_dict: Dict[Any, int]) -> None:
            # 内部存储实现， 每隔一段时间会刷新数据
            # 收到统计好的数据后， 根据公式统计每个指标的概率
            _dict: Dict[str, Dict[str, int]] = {}
            for key, value in stats_dict.items():
                if key.startswith(self._prefix):
                    _, index, type_ = key.split("|")
                    if index not in _dict:
                        _dict[index] = {}
                    _dict[index][type_] = value
            for index, metric_dict in _dict.items():
                total: int = metric_dict.get("total", 0)
                error_cnt: int = metric_dict.get("error", 0)
                self._probability_dict[index] = max(0.0, (total - k * (total - error_cnt)) / (total + 1))

        self._window_statistics.add_priority_callback(upload_probability)
        self.event_dict: Dict[EventEnum, List[CLIENT_EVENT_FN]] = {
            EventEnum.after_start: [self.start_event_handle],
            EventEnum.before_end: [self.stop_event_handle],
        }

    def get_index_from_request(self, request: Request) -> str:
        """从request获取指标的值"""
        raise NotImplementedError

    def get_index_from_response(self, response: Response) -> str:
        """从response获取指标的值"""
        raise NotImplementedError

    async def process_request(self, request: Request) -> Request:
        if request.msg_type == Constant.CLIENT_EVENT:
            # 忽略掉系统探测的请求
            return request
        index: str = self.get_index_from_request(request)
        total_key: str = f"{self._prefix}|{index}|total"
        # 记录请求数
        self._window_statistics.set_gauge_value(total_key, self._expire, self._interval)
        if random.randint(0, 100) < self._probability_dict.get(index, 0.0) * 100:
            # 判断是否可以放行， 如果不可放行， 则记录失败信息
            error_key: str = f"{self._prefix}|{index}|error"
            # 内部存储的方法， 类似于Prometheus的gauge
            self._window_statistics.set_gauge_value(error_key, self._expire, self._interval)
            raise self.exc
        return request

    async def process_exc(self, response: Response, exc: Exception) -> Tuple[Response, Exception]:
        # 如果是Python的异常或者是自定义中状态码大于等于500的异常，则记录失败信息
        if request.msg_type != Constant.CLIENT_EVENT \
            and not (isinstance(exc, BaseRapError) and exc.status_code < 500):
            error_key: str = f"{self._prefix}|{self.get_index_from_response(response)}|error"
            # 内部存储的方法， 类似于Prometheus的gauge
            self._window_statistics.set_gauge_value(error_key, self._expire, self._interval)
        return response, exc
    

class HostCircuitBreakerProcessor(BaseCircuitBreakerProcessor):
    """以host为维度的异常"""

    exc: Exception = CircuitBreakerExc("Service Unavailable")

    def get_index_from_request(self, request: Request) -> str:
        return request.header["host"][0]

    def get_index_from_response(self, response: Response) -> str:
        return response.conn.peer_tuple[0]


class FuncCircuitBreakerProcessor(BaseCircuitBreakerProcessor):
    """以函数名为维度的异常"""

    exc: Exception = CircuitBreakerExc("Service's func Unavailable")

    def get_index_from_request(self, request: Request) -> str:
        return request.target

    def get_index_from_response(self, response: Response) -> str:
        return response.target