Spaces:
Paused
Paused
| """内存缓存与限流计数，减少对数据库的依赖。 | |
| key 缓存：首次校验后把 key 的校验信息与配额缓存在内存， | |
| 后续请求命中缓存即不查库；后台改 key 时清缓存。 | |
| RPM：用内存滑动窗口计数，不再查 request_logs 表。 | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import time | |
| from dataclasses import dataclass | |
| from typing import Dict, List, Set | |
| import hashlib | |
| class CachedKey: | |
| """key 的缓存校验信息（不含敏感的完整 key）。""" | |
| id: int | |
| name: str | |
| prefix: str | |
| key_hash: str | |
| is_active: bool | |
| expires_at: float | None # epoch 秒 | |
| allowed_models: list[str] | None | |
| rpm_limit: int | None | |
| tpm_limit: int | None | |
| token_quota: int | None | |
| tokens_used: int | |
| # 负载快照时间，用于判断是否需要回库刷新 tokens_used | |
| loaded_at: float | |
| class KeyCache: | |
| """按 prefix 索引的 key 内存缓存。""" | |
| def __init__(self) -> None: | |
| self._by_prefix: Dict[str, CachedKey] = {} | |
| self._by_id: Dict[int, str] = {} # id -> prefix，用于按 id 失效 | |
| self._lock = asyncio.Lock() | |
| def get(self, prefix: str) -> CachedKey | None: | |
| return self._by_prefix.get(prefix) | |
| async def put(self, k: CachedKey) -> None: | |
| async with self._lock: | |
| self._by_prefix[k.prefix] = k | |
| self._by_id[k.id] = k.prefix | |
| async def invalidate_id(self, key_id: int) -> None: | |
| """按 id 失效（后台改/删 key 时调用）。""" | |
| async with self._lock: | |
| prefix = self._by_id.pop(key_id, None) | |
| if prefix: | |
| self._by_prefix.pop(prefix, None) | |
| async def invalidate_prefix(self, prefix: str) -> None: | |
| async with self._lock: | |
| ck = self._by_prefix.pop(prefix, None) | |
| if ck: | |
| self._by_id.pop(ck.id, None) | |
| async def clear(self) -> None: | |
| async with self._lock: | |
| self._by_prefix.clear() | |
| self._by_id.clear() | |
| async def add_usage(self, key_id: int, tokens: int) -> None: | |
| """成功转发后把用量累加到内存缓存（不回库），保持配额判断近实时。""" | |
| async with self._lock: | |
| prefix = self._by_id.get(key_id) | |
| if prefix: | |
| ck = self._by_prefix.get(prefix) | |
| if ck: | |
| ck.tokens_used += tokens | |
| class RPMCounter: | |
| """内存滑动窗口 RPM 计数，按 key_id 维度。""" | |
| def __init__(self, window: float = 60.0) -> None: | |
| self._window = window | |
| self._hits: Dict[int, List[float]] = {} | |
| self._lock = asyncio.Lock() | |
| async def hit_and_check(self, key_id: int, limit: int | None) -> bool: | |
| """记录一次命中并返回是否在限额内（True=放行，False=超限）。""" | |
| if not limit: | |
| return True | |
| now = time.time() | |
| cutoff = now - self._window | |
| async with self._lock: | |
| hits = self._hits.get(key_id) | |
| if hits is None: | |
| hits = [] | |
| self._hits[key_id] = hits | |
| # 清掉过期命中 | |
| hits[:] = [t for t in hits if t > cutoff] | |
| if len(hits) >= limit: | |
| return False | |
| hits.append(now) | |
| return True | |
| key_cache = KeyCache() | |
| rpm_counter = RPMCounter() | |
| def verify_key_hash(raw: str, digest: str) -> bool: | |
| return hashlib.sha256(raw.encode()).hexdigest() == digest | |
| async def cached_key_from(raw: str, db_loader) -> CachedKey | None: | |
| """先查缓存，未命中则调 db_loader(prefix) 回库加载并缓存。""" | |
| prefix = raw[:20] | |
| ck = key_cache.get(prefix) | |
| if ck is not None and verify_key_hash(raw, ck.key_hash): | |
| return ck | |
| # 未命中或哈希不匹配，回库 | |
| db_key = await db_loader(prefix) | |
| if db_key is None: | |
| return None | |
| exp = db_key.expires_at.timestamp() if db_key.expires_at else None | |
| ck = CachedKey( | |
| id=db_key.id, name=db_key.name, prefix=db_key.prefix, key_hash=db_key.key_hash, | |
| is_active=db_key.is_active, expires_at=exp, allowed_models=db_key.allowed_models, | |
| rpm_limit=db_key.rpm_limit, tpm_limit=db_key.tpm_limit, | |
| token_quota=db_key.token_quota, tokens_used=db_key.tokens_used or 0, | |
| loaded_at=time.time(), | |
| ) | |
| await key_cache.put(ck) | |
| return ck | |