45073486a7
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.
174 lines
6.6 KiB
Python
174 lines
6.6 KiB
Python
"""进程内限流后端: 与 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
|