Files
PolyGateway/src/polygateway/backends/memory/limiter.py
T
iomgaa 45073486a7 fix: reparent governance backend failures under GatewayUnavailableError (issue #7)
A fail-closed limiter or breaker backend means the scope cannot emit a
single request, which is exactly scope-level unavailability. But the error
sat directly under PolyGatewayError, so a caller writing only
`except GatewayUnavailableError` dropped it into the catch-all branch:
Redis blips once and a backlog of tasks burns its business failure budget
into the dead letter queue, over a fault a restart would clear.

Three gate paths leak to callers rather than being absorbed by
_record_quietly (try_acquire, try_enter, progress_age_s); each is now
pinned by a test, since none of them had one before.

The two unknown-source sites move to SourceNotConfiguredError instead of
following along. They report a misconfigured source name, not an outage,
and letting them into the retryable family would be the mirror of the bug
being fixed here: the task would retry forever and never surface.
2026-08-06 04:53:52 -04:00

174 lines
6.6 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""进程内限流后端: 与 Redis 版同一契约的六道闸实现(D3 双后端)。
语义蓝本 CHS `app/coordination/limiter.py`: 并发 = 带 TTL 的租约(持有者
死亡后过期回收);RPM/TPM = 分钟**固定窗口**计数(`int(now/60)`,与 CHS
Lua 同款口径);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 SourceNotConfiguredError
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 SourceNotConfiguredError(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