feat: add in-memory limiter and breaker satisfying backend contracts

2026-07-20 06:54:07 -04:00
parent 454a8b5e0f
commit 4a176b6220
9 changed files with 645 additions and 0 deletions
@@ -0,0 +1,170 @@
+"""进程内限流后端: 与 Redis 版同一契约的六道闸实现(D3 双后端)。
+
+语义蓝本 CHS `app/coordination/limiter.py`: 并发 = 带 TTL 的租约(持有者
+死亡后过期回收);RPM/TPM = 分钟滑动窗口计数;TPM 入场按 est 预扣,settle
+按实际结算多退少补且退款落 acquire 时的窗口。检查-占用在单次同步段内完成
+(无 await 穿插),单事件循环下天然原子;本实现不跨进程,是单进程部署的
+正确答案(多 worker 用 M2 Redis 后端)。
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+import uuid
+from typing import TYPE_CHECKING
+
+from polygateway.errors import GovernanceBackendError
+from polygateway.types import GlobalLimits, SourceConfig, SourceStats
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+_GLOBAL = "__global__"
+_WINDOW_S = 60.0
+
+
+class _MemoryPermit:
+    """入场许可;release/settle 幂等(CHS _RedisPermit 同款 flag 语义)。"""
+
+    def __init__(self, limiter: InMemoryLimiter, 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:
+        if self._released:
+            return
+        self._released = True
+        self._limiter._release_lease(self._source, self._lease_id)
+
+    async def settle(self, actual_tokens: int) -> None:
+        if self._settled:
+            return
+        self._settled = True
+        delta = actual_tokens - self._est
+        if delta:
+            self._limiter._settle_tpm(self._source, delta, self._window)
+
+
+class InMemoryLimiter:
+    """六道闸: 全局/单源 × 并发/RPM/TPM;限额 0 = 该闸不启用。"""
+
+    def __init__(
+        self,
+        *,
+        scope: str,
+        sources: dict[str, SourceConfig],
+        global_limits: GlobalLimits,
+        lease_ttl_s: float = 1500.0,
+        now: Callable[[], float] = time.monotonic,
+        sleep: Callable[[float], object] = asyncio.sleep,
+        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
+        self._sources = dict(sources)
+        self._global = global_limits
+        self._lease_ttl_s = lease_ttl_s
+        self._now = now
+        self._sleep = sleep
+        self._poll_interval_s = poll_interval_s
+        # source → lease_id → 过期时刻
+        self._leases: dict[str, dict[str, float]] = {name: {} for name in self._sources}
+        # (source|__global__, window_id) → 计数
+        self._rpm: dict[tuple[str, int], int] = {}
+        self._tpm: dict[tuple[str, int], int] = {}
+        self._progress_at: float | None = None
+
+    # —— 内部原语(同步,单事件循环下原子)——
+
+    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 _window(self) -> int:
+        return int(self._now() / _WINDOW_S)
+
+    def _purge_leases(self) -> None:
+        now = self._now()
+        for leases in self._leases.values():
+            expired = [lid for lid, expiry in leases.items() if expiry <= now]
+            for lid in expired:
+                del leases[lid]
+
+    def _inflight(self, source_key: str) -> int:
+        return len(self._leases[source_key])
+
+    def _inflight_total(self) -> int:
+        return sum(len(leases) for leases in self._leases.values())
+
+    def _release_lease(self, source_key: str, lease_id: str) -> None:
+        self._leases[source_key].pop(lease_id, None)
+
+    def _settle_tpm(self, source_key: str, delta: int, window: int) -> None:
+        # 退款/补记落 acquire 时的窗口,计数下限 0(CHS Lua 同款)
+        for key in (source_key, _GLOBAL):
+            slot = (key, window)
+            self._tpm[slot] = max(0, self._tpm.get(slot, 0) + delta)
+
+    def _gates_pass(self, cfg: SourceConfig, est_tokens: int, window: int) -> bool:
+        checks = (
+            (self._global.max_concurrency, self._inflight_total() + 1),
+            (cfg.max_concurrency, self._inflight(cfg.name) + 1),
+            (self._global.rpm, self._rpm.get((_GLOBAL, window), 0) + 1),
+            (cfg.rpm, self._rpm.get((cfg.name, window), 0) + 1),
+            (self._global.tpm, self._tpm.get((_GLOBAL, window), 0) + est_tokens),
+            (cfg.tpm, self._tpm.get((cfg.name, window), 0) + est_tokens),
+        )
+        return all(limit <= 0 or would_be <= limit for limit, would_be in checks)
+
+    # —— 契约方法 ——
+
+    async def try_acquire(self, source_key: str, est_tokens: int) -> _MemoryPermit | None:
+        """非阻塞准入: 六闸全过才占用;任一满则返回 None 且零副作用。"""
+        cfg = self._cfg(source_key)
+        self._purge_leases()
+        window = self._window()
+        if not self._gates_pass(cfg, est_tokens, window):
+            return None
+        lease_id = uuid.uuid4().hex
+        self._leases[source_key][lease_id] = self._now() + self._lease_ttl_s
+        for key in (source_key, _GLOBAL):
+            self._rpm[(key, window)] = self._rpm.get((key, window), 0) + 1
+            if est_tokens:
+                self._tpm[(key, window)] = self._tpm.get((key, window), 0) + est_tokens
+        return _MemoryPermit(self, source_key, lease_id, est_tokens, window)
+
+    async def acquire(self, source_key: str, est_tokens: int) -> _MemoryPermit:
+        """阻塞准入: 轮询直到拿到 permit;sleep 可被取消(取消穿透铁律)。"""
+        while True:
+            permit = await self.try_acquire(source_key, est_tokens)
+            if permit is not None:
+                return permit
+            await self._sleep(self._poll_interval_s)
+
+    async def source_stats(self, source_key: str) -> SourceStats:
+        self._cfg(source_key)
+        self._purge_leases()
+        window = self._window()
+        return SourceStats(
+            inflight=self._inflight(source_key),
+            rpm_used=self._rpm.get((source_key, window), 0),
+            tpm_used=self._tpm.get((source_key, window), 0),
+        )
+
+    async def mark_progress(self) -> None:
+        """记录"最近一次出餐"时刻,供背压 stall 判定(M2)读取。"""
+        self._progress_at = self._now()
+
+    async def progress_age_s(self) -> float:
+        if self._progress_at is None:
+            return float("inf")
+        return self._now() - self._progress_at