WorkBuddy
d9aeff2114
- app.py lifespan:APP_ENV=production 时 logger.warning 提醒 进程内 rate-limit/KeyRing 仅单进程有效(不 sys.exit、不引入 Redis) - key_ring.py docstring 补充多 worker 不共享的后果说明 - deps.py 注释已具备(单进程有效,生产前置 Nginx),无需改动
147 lines
5.4 KiB
Python
147 lines
5.4 KiB
Python
"""API Key 轮换环:多 key 自动切换,遇到限流(429)自动跳过已冷却 key。
|
|
|
|
设计:
|
|
- 进程内纯内存状态(限速是短时状态,无需持久化;D7: 多 worker 部署时各进程
|
|
独立计数、不共享,上游限速额度应按 worker 数分摊,或前置网关统一管理)
|
|
- 单 key 场景零开销:直接透传
|
|
- 多 key 场景:429 时把当前 key 标记冷却(默认 60s),轮转到下一个可用 key
|
|
- 全部 key 都在冷却时:使用最早冷却的那个 key 并等待(退化到单 key 重试)
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import time
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DEFAULT_COOLDOWN = 60.0 # 单个 key 被限流后的冷却时间(秒)
|
|
|
|
|
|
class KeyRing:
|
|
"""多 key 轮换环。在 async 单线程事件循环下无需加锁。"""
|
|
|
|
def __init__(self, keys: list[str], cooldown_seconds: float = DEFAULT_COOLDOWN) -> None:
|
|
self._keys: list[str] = [k.strip() for k in keys if k and k.strip()]
|
|
self._cooldown = cooldown_seconds
|
|
# key → 冷却过期时间戳(时刻);不在表中表示可用
|
|
self._blocked_until: dict[str, float] = {}
|
|
self._index = 0 # 当前轮转位置
|
|
|
|
@property
|
|
def has_multiple(self) -> bool:
|
|
return len(self._keys) > 1
|
|
|
|
@property
|
|
def all_keys(self) -> list[str]:
|
|
return list(self._keys)
|
|
|
|
@property
|
|
def active_key(self) -> str | None:
|
|
"""当前指向的 key(即使正在冷却也返回,用于上报)。"""
|
|
if not self._keys:
|
|
return None
|
|
return self._keys[self._index]
|
|
|
|
def get(self) -> str | None:
|
|
"""获取一个可用 key:优先选不在冷却中的;全部冷却则选最早过期的。"""
|
|
if not self._keys:
|
|
return None
|
|
if len(self._keys) == 1:
|
|
return self._keys[0]
|
|
|
|
now = time.monotonic()
|
|
n = len(self._keys)
|
|
# 从当前 index 开始找一圈,找一个可用的
|
|
for offset in range(n):
|
|
idx = (self._index + offset) % n
|
|
key = self._keys[idx]
|
|
expire = self._blocked_until.get(key, 0.0)
|
|
if now >= expire:
|
|
# 可用:把指针移到这里
|
|
self._index = idx
|
|
# 清理已过期的冷却记录
|
|
if key in self._blocked_until:
|
|
del self._blocked_until[key]
|
|
return key
|
|
|
|
# 全部在冷却中:选最早过期的那个,并等待到它过期
|
|
earliest_key = min(self._keys, key=lambda k: self._blocked_until.get(k, 0.0))
|
|
self._index = self._keys.index(earliest_key)
|
|
return earliest_key
|
|
|
|
def report_rate_limited(self, key: str | None = None) -> str | None:
|
|
"""上报某个 key 被限流(429)。默认是当前 key。返回切换后的新 key。"""
|
|
target = key or self.active_key
|
|
if target and len(self._keys) > 1:
|
|
until = time.monotonic() + self._cooldown
|
|
self._blocked_until[target] = until
|
|
logger.warning(
|
|
"bzzoiro key 被限流(429),冷却 %.0fs: %s", self._cooldown, _mask(target),
|
|
)
|
|
# 轮转到下一个(即使只有一个 key 也做一次 get,保持行为一致)
|
|
return self.get()
|
|
|
|
def wait_if_all_blocked(self) -> float:
|
|
"""如果所有 key 都在冷却中,返回需要等待的秒数;否则返回 0。"""
|
|
if len(self._keys) <= 1:
|
|
return 0.0
|
|
now = time.monotonic()
|
|
remaining = [self._blocked_until.get(k, 0.0) - now for k in self._keys]
|
|
if all(r > -0.001 for r in remaining) and any(r > 0.001 for r in remaining):
|
|
# 全部仍在冷却中
|
|
return max(remaining)
|
|
return 0.0
|
|
|
|
def stats(self) -> dict:
|
|
"""当前 key 环状态(用于管理后台展示)。"""
|
|
now = time.monotonic()
|
|
return {
|
|
"total": len(self._keys),
|
|
"keys": [
|
|
{
|
|
"masked": _mask(k),
|
|
"blocked_remaining": max(0.0, round(self._blocked_until.get(k, 0.0) - now, 1)),
|
|
}
|
|
for k in self._keys
|
|
],
|
|
"active_index": self._index,
|
|
}
|
|
|
|
|
|
def _mask(key: str) -> str:
|
|
"""脱敏:只显示前 4 位和后 4 位。"""
|
|
if len(key) <= 10:
|
|
return key[:2] + "***"
|
|
return key[:4] + "..." + key[-4:]
|
|
|
|
|
|
# ── 全局单例(进程级,按 base URL 隔离) ──────────────────────────
|
|
_RINGS: dict[str, KeyRing] = {}
|
|
|
|
|
|
def parse_keys(value: str | None) -> list[str]:
|
|
"""解析 key 配置值:支持逗号、分号、换行分隔的多个 key。"""
|
|
if not value:
|
|
return []
|
|
# 统一替换分隔符为逗号后拆分
|
|
normalized = value.replace("\n", ",").replace(";", ",")
|
|
return [k.strip() for k in normalized.split(",") if k.strip()]
|
|
|
|
|
|
def get_key_ring(base: str, raw_keys: str | None, cooldown_seconds: float = DEFAULT_COOLDOWN) -> KeyRing:
|
|
"""获取(或创建)某 base URL 对应的 KeyRing。"""
|
|
key = base
|
|
ring = _RINGS.get(key)
|
|
parsed = parse_keys(raw_keys)
|
|
if ring is None:
|
|
ring = KeyRing(parsed, cooldown_seconds)
|
|
_RINGS[key] = ring
|
|
else:
|
|
# 热更新 key 列表(增删 key 无需重启)
|
|
if set(ring.all_keys) != set(parsed):
|
|
ring._keys = parsed
|
|
ring._blocked_until.clear()
|
|
ring._index = 0
|
|
ring._cooldown = cooldown_seconds
|
|
return ring
|