feat: add redis six-gate rate limiter backend

2026-07-21 00:27:16 -04:00
parent acc1bcc18a
commit aae739cebe
4 changed files with 498 additions and 14 deletions
@@ -0,0 +1,6 @@
+"""Redis 治理后端(M2): 与内存版同一契约的跨进程实现(D3 双后端)。"""
+
+from polygateway.backends.redis.breaker import RedisGate
+from polygateway.backends.redis.limiter import RedisLimiter
+
+__all__ = ["RedisGate", "RedisLimiter"]
@@ -0,0 +1,307 @@
+"""跨进程限流后端: CHS 六道闸 Lua 逐字移植(D3 双后端;M2 设计 §2)。
+
+语义蓝本 `reference/CHSAnalyzer/app/coordination/{scripts,limiter}.py`:
+读-判-占在单条 Lua 内原子完成;并发 = ZSET 租约(score=过期时刻 ms,惰性
+清理 + PEXPIRE 双保险);RPM/TPM = 分钟固定窗口计数,窗口 id 用 **Redis
+服务器时钟**(多进程口径统一);TPM 预扣 settle 多退少补且退款落 acquire
+时的窗口(INCRBY 可为负,存储保留负值保证跨窗口结算正确)。
+
+对 CHS 的已声明偏离(设计 §9 勘误): ① 限额 0 = 该闸不启用(M1 契约冻结;
+CHS Lua 无此守卫,0 在其语义下是全拒),Lua 判据显式加 `limit > 0`;
+② key 前缀 `cclimit:` → `pgw:limit:`;③ `LimiterError` → `GovernanceBackendError`;
+④ `source_stats` 读侧 clamp ≥0(展示口径与内存版统一,存储不 clamp)。
+契约量纲为秒,毫秒换算是本后端私事(ports.py 模块 docstring)。
+"""
+
+from __future__ import annotations
+
+import asyncio
+import uuid
+from typing import TYPE_CHECKING
+
+from loguru import logger
+from redis.exceptions import RedisError
+
+from polygateway.errors import GovernanceBackendError
+from polygateway.types import GlobalLimits, SourceConfig, SourceStats
+
+if TYPE_CHECKING:
+    from redis.asyncio import Redis
+
+_WINDOW_TTL_S = 120  # 分钟窗口兜底过期(CHS limiter.py:16)
+_PROGRESS_TTL_S = 3600  # 远大于任何 stall_window,防进度键过期造成假停滞(CHS :21)
+
+# KEYS: g_lease s_lease g_rpm s_rpm g_tpm s_tpm(rpm/tpm key 已含 window 后缀)
+# ARGV: g_conc s_conc g_rpm s_rpm g_tpm s_tpm est lease_ttl_ms window_ttl_s lease_id
+# 拒绝零副作用(仅清过期 lease);闸判据 CHS scripts.py:18-23 + 库版 limit>0 守卫
+ACQUIRE = """
+local t = redis.call('TIME')
+local now = tonumber(t[1]) * 1000 + math.floor(tonumber(t[2]) / 1000)
+redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, now)
+redis.call('ZREMRANGEBYSCORE', KEYS[2], 0, now)
+local gi = redis.call('ZCARD', KEYS[1])
+local si = redis.call('ZCARD', KEYS[2])
+local gr = tonumber(redis.call('GET', KEYS[3]) or '0')
+local sr = tonumber(redis.call('GET', KEYS[4]) or '0')
+local gt = tonumber(redis.call('GET', KEYS[5]) or '0')
+local st = tonumber(redis.call('GET', KEYS[6]) or '0')
+local est = tonumber(ARGV[7])
+if tonumber(ARGV[1]) > 0 and gi >= tonumber(ARGV[1]) then return 0 end
+if tonumber(ARGV[2]) > 0 and si >= tonumber(ARGV[2]) then return 0 end
+if tonumber(ARGV[3]) > 0 and gr >= tonumber(ARGV[3]) then return 0 end
+if tonumber(ARGV[4]) > 0 and sr >= tonumber(ARGV[4]) then return 0 end
+if tonumber(ARGV[5]) > 0 and gt + est > tonumber(ARGV[5]) then return 0 end
+if tonumber(ARGV[6]) > 0 and st + est > tonumber(ARGV[6]) then return 0 end
+local expire_at = now + tonumber(ARGV[8])
+redis.call('ZADD', KEYS[1], expire_at, ARGV[10])
+redis.call('ZADD', KEYS[2], expire_at, ARGV[10])
+redis.call('PEXPIRE', KEYS[1], ARGV[8])
+redis.call('PEXPIRE', KEYS[2], ARGV[8])
+redis.call('INCR', KEYS[3]); redis.call('EXPIRE', KEYS[3], ARGV[9])
+redis.call('INCR', KEYS[4]); redis.call('EXPIRE', KEYS[4], ARGV[9])
+redis.call('INCRBY', KEYS[5], est); redis.call('EXPIRE', KEYS[5], ARGV[9])
+redis.call('INCRBY', KEYS[6], est); redis.call('EXPIRE', KEYS[6], ARGV[9])
+return 1
+"""
+
+# KEYS: g_lease s_lease ; ARGV: lease_id(ZREM 幂等,CHS scripts.py:37-41)
+RELEASE = """
+redis.call('ZREM', KEYS[1], ARGV[1])
+redis.call('ZREM', KEYS[2], ARGV[1])
+return 1
+"""
+
+# KEYS: g_tpm s_tpm ; ARGV: delta window_ttl_s(delta 可负=退款;落 acquire 窗口)
+SETTLE = """
+redis.call('INCRBY', KEYS[1], ARGV[1]); redis.call('EXPIRE', KEYS[1], ARGV[2])
+redis.call('INCRBY', KEYS[2], ARGV[1]); redis.call('EXPIRE', KEYS[2], ARGV[2])
+return 1
+"""
+
+# KEYS: s_lease s_rpm s_tpm(CHS scripts.py:51-58)
+STATS = """
+local t = redis.call('TIME')
+local now = tonumber(t[1]) * 1000 + math.floor(tonumber(t[2]) / 1000)
+redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, now)
+return {redis.call('ZCARD', KEYS[1]),
+        tonumber(redis.call('GET', KEYS[2]) or '0'),
+        tonumber(redis.call('GET', KEYS[3]) or '0')}
+"""
+
+# KEYS: progress_key ; ARGV: ttl_s(CHS scripts.py:61-67)
+PROGRESS_MARK = """
+local t = redis.call('TIME')
+local now = tonumber(t[1]) * 1000 + math.floor(tonumber(t[2]) / 1000)
+redis.call('SET', KEYS[1], now)
+redis.call('EXPIRE', KEYS[1], ARGV[1])
+return 1
+"""
+
+# KEYS: progress_key(-1=从未成功;时钟回拨 clamp 0;CHS scripts.py:70-78)
+PROGRESS_AGE = """
+local v = redis.call('GET', KEYS[1])
+if not v then return -1 end
+local t = redis.call('TIME')
+local now = tonumber(t[1]) * 1000 + math.floor(tonumber(t[2]) / 1000)
+local d = now - tonumber(v)
+if d < 0 then return 0 end
+return d
+"""
+
+
+class _RedisPermit:
+    """入场许可: 持有并发租约,负责释放与按 acquire 窗口结算 TPM(幂等)。"""
+
+    def __init__(
+        self, limiter: RedisLimiter, source: str, lease_id: str, est: int, window: int
+    ) -> None:
+        self._limiter = limiter
+        self._source = source
+        self._lease_id = lease_id
+        self._est = est
+        self._window = window
+        self._released = False
+        self._settled = False
+
+    async def release(self) -> None:
+        """归还并发租约;释放侧 Redis 失败降级 warning(设计 §10),不冒泡。"""
+        if self._released:
+            return
+        self._released = True
+        try:
+            await self._limiter._release_lease(self._source, self._lease_id)
+        except GovernanceBackendError as exc:
+            logger.warning("permit release 降级(租约将由 TTL 回收): {}", exc)
+
+    async def settle(self, actual_tokens: int) -> None:
+        """按实际 usage 结算 TPM 差额;释放侧失败降级 warning,不冒泡。"""
+        if self._settled:
+            return
+        self._settled = True
+        delta = actual_tokens - self._est
+        if delta == 0:
+            return
+        try:
+            await self._limiter._settle_tpm(self._source, delta, self._window)
+        except GovernanceBackendError as exc:
+            logger.warning("permit settle 降级(窗口计数将随 TTL 过期): {}", exc)
+
+
+class RedisLimiter:
+    """六道闸的跨进程实现;同 scope 的多进程 worker 天然共享全局限额。"""
+
+    def __init__(
+        self,
+        *,
+        scope: str,
+        sources: dict[str, SourceConfig],
+        global_limits: GlobalLimits,
+        redis: Redis,
+        lease_ttl_s: float = 1500.0,
+        poll_interval_s: float = 0.05,
+    ) -> None:
+        if lease_ttl_s <= 0 or poll_interval_s <= 0:
+            raise ValueError("lease_ttl_s 与 poll_interval_s 必须 > 0")
+        self._scope = scope.lower()
+        self._sources = dict(sources)
+        self._global = global_limits
+        self._lease_ttl_ms = int(lease_ttl_s * 1000)
+        self._poll_interval_s = poll_interval_s
+        self._redis = redis
+        self._owns_client = False
+        self._acquire_lua = redis.register_script(ACQUIRE)
+        self._release_lua = redis.register_script(RELEASE)
+        self._settle_lua = redis.register_script(SETTLE)
+        self._stats_lua = redis.register_script(STATS)
+        self._progress_mark_lua = redis.register_script(PROGRESS_MARK)
+        self._progress_age_lua = redis.register_script(PROGRESS_AGE)
+
+    @classmethod
+    def from_url(cls, url: str, **kwargs) -> RedisLimiter:
+        """自建并持有 Redis 客户端(aclose 时代关);共享后端请直接注入 redis。"""
+        import redis.asyncio as aioredis
+
+        limiter = cls(redis=aioredis.from_url(url), **kwargs)
+        limiter._owns_client = True
+        return limiter
+
+    # —— key 布局(CHS limiter.py:100-115,前缀改 pgw:limit:)——
+
+    def _cfg(self, source_key: str) -> SourceConfig:
+        cfg = self._sources.get(source_key)
+        if cfg is None:
+            raise GovernanceBackendError(f"未知源 {source_key!r}(scope={self._scope})")
+        return cfg
+
+    def _lease_keys(self, source_key: str) -> tuple[str, str]:
+        return (
+            f"pgw:limit:GLOBAL:{self._scope}:lease",
+            f"pgw:limit:{self._scope}:{source_key}:lease",
+        )
+
+    def _window_keys(self, source_key: str, window: int) -> dict[str, str]:
+        gl = f"GLOBAL:{self._scope}"
+        return {
+            "g_rpm": f"pgw:limit:{gl}:rpm:{window}",
+            "s_rpm": f"pgw:limit:{self._scope}:{source_key}:rpm:{window}",
+            "g_tpm": f"pgw:limit:{gl}:tpm:{window}",
+            "s_tpm": f"pgw:limit:{self._scope}:{source_key}:tpm:{window}",
+        }
+
+    def _progress_key(self) -> str:
+        return f"pgw:limit:GLOBAL:{self._scope}:progress_ms"
+
+    async def _window_id(self) -> int:
+        """Redis 服务器时钟生成分钟窗口 id(CHS limiter.py:95-98)。"""
+        sec, _ = await self._redis.time()
+        return int(sec) // 60
+
+    # —— 契约方法 ——
+
+    async def try_acquire(self, source_key: str, est_tokens: int) -> _RedisPermit | None:
+        """非阻塞准入: 六闸单 Lua 原子判定;拒绝零副作用;Redis 失败报错不放行。"""
+        cfg = self._cfg(source_key)
+        lease_id = uuid.uuid4().hex
+        try:
+            window = await self._window_id()
+            gl, sl = self._lease_keys(source_key)
+            wk = self._window_keys(source_key, window)
+            ok = await self._acquire_lua(
+                keys=[gl, sl, wk["g_rpm"], wk["s_rpm"], wk["g_tpm"], wk["s_tpm"]],
+                args=[
+                    self._global.max_concurrency,
+                    cfg.max_concurrency,
+                    self._global.rpm,
+                    cfg.rpm,
+                    self._global.tpm,
+                    cfg.tpm,
+                    est_tokens,
+                    self._lease_ttl_ms,
+                    _WINDOW_TTL_S,
+                    lease_id,
+                ],
+            )
+        except RedisError as exc:
+            raise GovernanceBackendError(f"限流后端 try_acquire 失败: {exc}") from exc
+        if ok != 1:
+            return None
+        return _RedisPermit(self, source_key, lease_id, est_tokens, window)
+
+    async def acquire(self, source_key: str, est_tokens: int) -> _RedisPermit:
+        """阻塞准入: 轮询直到拿到 permit;sleep 可被取消(取消穿透铁律)。"""
+        while True:
+            permit = await self.try_acquire(source_key, est_tokens)
+            if permit is not None:
+                return permit
+            await asyncio.sleep(self._poll_interval_s)
+
+    async def _release_lease(self, source_key: str, lease_id: str) -> None:
+        gl, sl = self._lease_keys(source_key)
+        try:
+            await self._release_lua(keys=[gl, sl], args=[lease_id])
+        except RedisError as exc:
+            raise GovernanceBackendError(f"限流后端 release 失败: {exc}") from exc
+
+    async def _settle_tpm(self, source_key: str, delta: int, window: int) -> None:
+        wk = self._window_keys(source_key, window)
+        try:
+            await self._settle_lua(keys=[wk["g_tpm"], wk["s_tpm"]], args=[delta, _WINDOW_TTL_S])
+        except RedisError as exc:
+            raise GovernanceBackendError(f"限流后端 settle 失败: {exc}") from exc
+
+    async def source_stats(self, source_key: str) -> SourceStats:
+        """当前窗口快照;读侧 clamp ≥0(展示口径,存储保留负值)。"""
+        self._cfg(source_key)
+        try:
+            window = await self._window_id()
+            _, sl = self._lease_keys(source_key)
+            wk = self._window_keys(source_key, window)
+            res = await self._stats_lua(keys=[sl, wk["s_rpm"], wk["s_tpm"]])
+        except RedisError as exc:
+            raise GovernanceBackendError(f"限流后端 source_stats 失败: {exc}") from exc
+        return SourceStats(
+            inflight=int(res[0]),
+            rpm_used=max(0, int(res[1])),
+            tpm_used=max(0, int(res[2])),
+        )
+
+    async def mark_progress(self) -> None:
+        """记录 scope 级"最近出餐"时刻;失败报错,降级责任在中间件记账侧(T5)。"""
+        try:
+            await self._progress_mark_lua(keys=[self._progress_key()], args=[_PROGRESS_TTL_S])
+        except RedisError as exc:
+            raise GovernanceBackendError(f"限流后端 mark_progress 失败: {exc}") from exc
+
+    async def progress_age_s(self) -> float:
+        """距上次全局成功的秒数;仅键缺失(-1)= 从未进展 → inf(CHS limiter.py:208)。"""
+        try:
+            res = await self._progress_age_lua(keys=[self._progress_key()])
+        except RedisError as exc:
+            raise GovernanceBackendError(f"限流后端 progress_age_s 失败: {exc}") from exc
+        return float("inf") if int(res) == -1 else int(res) / 1000.0
+
+    async def aclose(self) -> None:
+        """幂等释放自建客户端;注入的客户端归注入方管理。"""
+        if self._owns_client:
+            self._owns_client = False
+            await self._redis.aclose()