112 lines
4.7 KiB
Python
112 lines
4.7 KiB
Python
|
"""限流原子计数共享 helper。
|
|||
|
|
|
|||
|
|
封装 ``CachePort.incr + expire`` 原子限流范式，消除 ``get/set`` 非原子
|
|||
|
|
读改写导致的高并发超额放行（C-C1）。供 ``rate_limit_stage`` 与
|
|||
|
|
``admin_message_service`` 复用，避免重复实现错误范式。
|
|||
|
|
|
|||
|
|
故障语义（与 ``health_check_service._checkProbeRateLimit`` 对齐）：
|
|||
|
|
|
|||
|
|
- ``incr`` 失败（``DependencyError``）→ fail-closed：请求尚未计数，翻译
|
|||
|
|
为 ``RateLimitError`` 拒绝，避免 Redis 故障期间限流失效放行超额请求。
|
|||
|
|
- ``expire`` 失败（``DependencyError``）→ fail-fast：请求已计数，直接
|
|||
|
|
上抛 ``DependencyError`` 暴露基础设施故障（INV-7），不在此处吞异常——
|
|||
|
|
窗口未设置会导致计数永久累积，限流失效。
|
|||
|
|
- 首次递增（``count == 1``）设置窗口 TTL（对应 Redis ``EXPIRE``，只设
|
|||
|
|
TTL 不修改值，消除"首次 ``set`` 覆盖并发 ``incr`` 计数"的竞态）。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
from yuxi.channels.contract.errors import DependencyError, RateLimitError
|
|||
|
|
from yuxi.channels.contract.ports.driven.cache_port import CachePort
|
|||
|
|
from yuxi.channels.contract.ports.driven.logger_port import LoggerPort
|
|||
|
|
|
|||
|
|
__all__ = ["RateLimitChecker"]
|
|||
|
|
|
|||
|
|
|
|||
|
|
class RateLimitChecker:
|
|||
|
|
"""原子限流计数器，封装 ``incr + expire`` 范式。
|
|||
|
|
|
|||
|
|
单一职责：基于 ``CachePort`` 原子操作校验单 key 计数是否超限。不负责
|
|||
|
|
key 构造与阈值配置（由调用方传入），``resource`` 字段标识限流作用域。
|
|||
|
|
|
|||
|
|
用法::
|
|||
|
|
|
|||
|
|
checker = RateLimitChecker(cache_port, logger)
|
|||
|
|
await checker.check(
|
|||
|
|
resource=f"control-plane:{user_id}:{ip}",
|
|||
|
|
cache_key=f"rate_limit:{user_id}:{ip}",
|
|||
|
|
max_requests=100,
|
|||
|
|
window_seconds=60,
|
|||
|
|
retry_after_ms=60_000,
|
|||
|
|
trace_id=trace_id,
|
|||
|
|
)
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, cache_port: CachePort, logger: LoggerPort | None = None) -> None:
|
|||
|
|
"""初始化限流计数器。
|
|||
|
|
|
|||
|
|
参数：
|
|||
|
|
cache_port: 缓存被驱动端口，提供 ``incr`` / ``expire`` 原子操作。
|
|||
|
|
logger: 日志端口（可选），``incr`` 故障 fail-closed 时记录告警。
|
|||
|
|
"""
|
|||
|
|
self._cache = cache_port
|
|||
|
|
self._logger = logger
|
|||
|
|
|
|||
|
|
async def check(
|
|||
|
|
self,
|
|||
|
|
*,
|
|||
|
|
resource: str,
|
|||
|
|
cache_key: str,
|
|||
|
|
max_requests: int,
|
|||
|
|
window_seconds: int,
|
|||
|
|
retry_after_ms: int,
|
|||
|
|
trace_id: str | None = None,
|
|||
|
|
) -> None:
|
|||
|
|
"""校验限流计数是否超限，超限或缓存故障时抛出。
|
|||
|
|
|
|||
|
|
参数：
|
|||
|
|
resource: 限流资源标识（写入 ``RateLimitError.resource``）。
|
|||
|
|
cache_key: 缓存计数 key。
|
|||
|
|
max_requests: 窗口内最大请求数（``count > max_requests`` 拒绝）。
|
|||
|
|
window_seconds: 限流窗口秒数（首次递增时设置 TTL）。
|
|||
|
|
retry_after_ms: 建议重试等待毫秒。
|
|||
|
|
trace_id: 追踪 ID（写入异常上下文与日志）。
|
|||
|
|
|
|||
|
|
抛出：
|
|||
|
|
RateLimitError: 请求频率超限，或 ``incr`` 缓存故障时 fail-closed。
|
|||
|
|
DependencyError: ``expire`` 缓存故障时 fail-fast（请求已计数）。
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
count = await self._cache.incr(cache_key)
|
|||
|
|
except DependencyError as exc:
|
|||
|
|
# incr 失败时请求尚未计数，fail-closed 翻译为 RateLimitError
|
|||
|
|
# 拒绝（与 health_check_service._checkProbeRateLimit 一致，INV-7），
|
|||
|
|
# 避免 Redis 故障期间限流失效放行超额请求。
|
|||
|
|
if self._logger is not None:
|
|||
|
|
await self._logger.error(
|
|||
|
|
"rate_limit cache failure, fail-closed reject",
|
|||
|
|
trace_id=trace_id,
|
|||
|
|
cache_key=cache_key,
|
|||
|
|
error=str(exc),
|
|||
|
|
)
|
|||
|
|
raise RateLimitError(
|
|||
|
|
resource=resource,
|
|||
|
|
retry_after_ms=retry_after_ms,
|
|||
|
|
trace_id=trace_id,
|
|||
|
|
) from exc
|
|||
|
|
# 首次递增（count == 1）设置 TTL 窗口。expire 只设过期时间不修改值，
|
|||
|
|
# 避免 set 覆盖高并发期间其他请求的 incr 计数。
|
|||
|
|
#
|
|||
|
|
# 故障策略与 incr 不同：expire 失败时请求已计数，fail-fast 直接抛
|
|||
|
|
# DependencyError 暴露基础设施故障（INV-7），不吞异常——窗口未设置
|
|||
|
|
# 会导致计数永久累积，限流失效。
|
|||
|
|
if count == 1:
|
|||
|
|
await self._cache.expire(cache_key, window_seconds)
|
|||
|
|
if count > max_requests:
|
|||
|
|
raise RateLimitError(
|
|||
|
|
resource=resource,
|
|||
|
|
retry_after_ms=retry_after_ms,
|
|||
|
|
trace_id=trace_id,
|
|||
|
|
)
|