RPC框架编写实践-RPC常见限流方法的实现

import asyncio

# 假设容量只有10
import time

leaky_bucket: asyncio.Queue = asyncio.Queue(10)
loop: asyncio.AbstractEventLoop = asyncio.get_event_loop()


async def demo_request(cnt: int) -> None:
    """模仿请求"""
    msg: str = f"I'm mock request:{cnt}"
    future: asyncio.Future = asyncio.Future()
    try:
        leaky_bucket.put_nowait(future)
    except asyncio.QueueFull:
        # 代表桶满了， 溢出来， 该请求要提前抛弃
        print(f"Fail Request:{msg}")
    else:
        # 等待放行
        await future
        print(time.time(), msg)


def timer() -> None:
    """定时放行请求"""
    try:
        # 放行该请求
        future: asyncio.Future = leaky_bucket.get_nowait()
        future.set_result(True)
    except asyncio.QueueEmpty:
        pass
    # 一秒执行一次
    loop.call_later(1, timer)

timer()
# 模拟并发12个请求
loop.run_until_complete(asyncio.gather(*[demo_request(i) for i in range(12)]))

from typing import Any, Coroutine, Union

from rap.server.plugin.processor.limit.rule import Rule


class BaseLimitBackend(object):
    def can_requests(self, key: str, rule: Rule, token_num: int = 1) -> Union[bool, Coroutine[Any, Any, bool]]:
        raise NotImplementedError

    def expected_time(self, key: str, rule: Rule) -> Union[float, Coroutine[Any, Any, float]]:
        raise NotImplementedError

class BaseRedisBackend(BaseLimitBackend, ABC):
    def __init__(self, redis: Union[StrictRedis, StrictRedisCluster]):
        # 初始化Redis模块
        self._redis: "Union[StrictRedis, StrictRedisCluster]" = redis

    async def _block_time_handle(self, key: str, rule: Rule, func: Callable[..., Awaitable[bool]]) -> bool:
        """处理block_time逻辑"""
        block_time_key: str = f"{key}:block_time"
        bucket_block_time: Optional[int] = rule.block_time

        if bucket_block_time is not None and await self._redis.exists(block_time_key):
            # 启用block_time逻辑， 且key已经存在， 那么直接返回False告诉该请求应该被拒绝 
            return False

        # 执行正真的判断是否限流逻辑
        can_requests: bool = await func()

        if not can_requests and bucket_block_time is not None:
            # 启用block_time逻辑且被限流时, 正式启用block time逻辑 
            await self._redis.set(block_time_key, bucket_block_time, ex=bucket_block_time)

        return can_requests

class BaseRedisCellBackend(BaseRedisBackend):
    """
    use redis-cell module
    learn more:https://github.com/brandur/redis-cell

    input: CL.THROTTLE user123 15 30 60 1
        # param  |  desc
        # user123 key
        # 15 maxburst
        # 30 token
        # 60 seconds
        # 1 apply 1token
    output:
        1) (integer) 0        # is allowed
        2) (integer) 16       # total bucket num
        3) (integer) 15       # the remaining limit of the key.
        4) (integer) -1       # the number of seconds until the user should retry,
                              #   and always -1 if the action was allowed.
        5) (integer) 2        # The number of seconds until the limit will reset to its maximum capacity
    """

    async def _call_cell(self, key: str, rule: Rule, token_num: int = 1) -> List[int]:
        """调用redis_cell"""
        result: List[int] = await self._redis.execute_command(
            "CL.THROTTLE", key, rule.max_token - 1, rule.gen_token, int(rule.total_second), token_num
        )
        return result

    def expected_time(self, key: str, rule: Rule) -> Union[float, Coroutine[Any, Any, float]]:
        """获取下次可请求时间"""
        async def _expected_time() -> float:
            block_time_key: str = key + ":block_time"
            block_time = await self._redis.get(block_time_key)
            if block_time:
                return await self._redis.ttl(block_time_key)

            result: List[int] = await self._call_cell(key, rule, 0)
            return float(max(result[3], 0))

        return _expected_time()

class RedisCellBackend(BaseRedisCellBackend):

    def can_requests(self, key: str, rule: Rule, token_num: int = 1) -> Union[bool, Coroutine[Any, Any, bool]]:
        """通过redis-cell判断是否可以请求，以及是否需要休眠等待， 如果需要则休眠固定的时间后再放行"""
        async def _can_requests() -> bool:
            result: List[int] = await self._call_cell(key, rule, token_num)
            can_requests: bool = result[0] == 0
            if can_requests and result[4]:
                await asyncio.sleep(result[4])
            return can_requests

        return self._block_time_handle(key, rule, _can_requests)

class LimitProcessor(BaseProcessor):
    def __init__(self, backend: BaseLimitBackend, rule_list: List[Tuple[RULE_FUNC_TYPE, Rule]]):
        """初始化规则和算法逻辑， 这里的规则之所以是使用传参的方式是仅供参考， 后续整个框架的配置都会抽离成一个config， 供其它组件调用"""
        self._backend: BaseLimitBackend = backend
        self._rule_list: List[Tuple[RULE_FUNC_TYPE, Rule]] = rule_list

    async def process_request(self, request: Request) -> Request:
        # not limit client event
        if request.msg_type == constant.CLIENT_EVENT:
            # 屏蔽event请求
            return request

        for func, rule in self._rule_list:
            # 获取该请求的key
            if inspect.iscoroutinefunction(func):
                key, is_ignore_limit = await func(request)  # type: ignore
            else:
                key, is_ignore_limit = func(request)
            if is_ignore_limit:
                # 如果该请求不应该限流， 直接跳过限流逻辑
                return request
            if key:
                # 匹配到key， 进入限流逻辑
                break
        else:
            raise TooManyRequest()

        # 通过backend判断是否限流
        key = f"rap:processor:{self.__class__.__name__}:{key}"
        can_requests: Union[bool, Awaitable[bool]] = self._backend.can_requests(key, rule)
        if inspect.isawaitable(can_requests):
            can_requests = await can_requests  # type: ignore
        if not can_requests:
            # 如果被限流， 返回异常， 并告知要何时后才可以再次请求
            expected_time: Union[float, Awaitable[float]] = self._backend.expected_time(key, rule)
            if inspect.isawaitable(expected_time):
                expected_time = await expected_time  # type: ignore
            raise TooManyRequest(extra_msg=f"expected time: {expected_time}")
        return request