mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
1d9c77522a
- 删除全部 Python 源码 (src/) 及 Alembic 迁移脚本，归档至 _deprecated_py_src/ - 重构 Rust gateway ai_pipeline: 拆分 planner/finalize 模块，新增 contracts/adaptation 层 - 重组 handlers 模块为 admin/public/proxy/internal/shared 子模块结构 - 新增 executor 模块，引入 Rust 原生数据库迁移 (aether-data/migrations) - 简化 CI/Docker 构建流程，移除 base image 二级构建，统一为单一 app image - 移除 Python 相关基础设施文件 (entrypoint.sh, gunicorn_conf.py, Dockerfile.base)
695 lines
22 KiB
Python
695 lines
22 KiB
Python
"""Redis operations for the Account Pool (provider-agnostic).
|
|
|
|
All pool transient state is stored in Redis. This module centralises key
|
|
naming, Lua scripts, and graceful fallbacks so that the rest of the pool
|
|
layer is free of Redis specifics.
|
|
|
|
Key schema
|
|
----------
|
|
ap:{pid}:sticky:{session_uuid} STRING -> key_id (TTL: config)
|
|
ap:{pid}:lru ZSET member=key_id, score=unix_ts
|
|
ap:{pid}:cooldown:{key_id} STRING -> reason (TTL: error-specific)
|
|
ap:{pid}:cost:{key_id} ZSET member=req_id, score=unix_ts
|
|
ap:{pid}:latency:{key_id} ZSET member=req_id:ttfb_ms, score=unix_ts
|
|
provider_oauth_token_cache:{key_id} STRING -> access_token (TTL: expires - 60)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
import uuid
|
|
from typing import TYPE_CHECKING
|
|
|
|
from src.clients.redis_client import get_redis_client
|
|
from src.core.logger import logger
|
|
|
|
if TYPE_CHECKING:
|
|
import redis.asyncio as aioredis
|
|
|
|
PREFIX = "ap"
|
|
|
|
|
|
def _sticky_key(provider_id: str, session_uuid: str) -> str:
|
|
return f"{PREFIX}:{provider_id}:sticky:{session_uuid}"
|
|
|
|
|
|
def _lru_key(provider_id: str) -> str:
|
|
return f"{PREFIX}:{provider_id}:lru"
|
|
|
|
|
|
def _cooldown_key(provider_id: str, key_id: str) -> str:
|
|
return f"{PREFIX}:{provider_id}:cooldown:{key_id}"
|
|
|
|
|
|
def _cooldown_index_key(provider_id: str) -> str:
|
|
"""SET tracking which keys are in cooldown (for O(1) count queries)."""
|
|
return f"{PREFIX}:{provider_id}:cooldown_idx"
|
|
|
|
|
|
def _cost_key(provider_id: str, key_id: str) -> str:
|
|
return f"{PREFIX}:{provider_id}:cost:{key_id}"
|
|
|
|
|
|
def _latency_key(provider_id: str, key_id: str) -> str:
|
|
return f"{PREFIX}:{provider_id}:latency:{key_id}"
|
|
|
|
|
|
def _oauth_cache_key(key_id: str) -> str:
|
|
return f"provider_oauth_token_cache:{key_id}"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Lua scripts
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# Sticky-select: GET binding, verify it's not in cooldown, refresh TTL.
|
|
# KEYS[1] = sticky key, KEYS[2] = cooldown key prefix (ap:{pid}:cooldown:)
|
|
# ARGV[1] = ttl
|
|
# Returns: key_id or nil
|
|
_STICKY_SELECT_LUA = """
|
|
local binding = redis.call("GET", KEYS[1])
|
|
if not binding then
|
|
return nil
|
|
end
|
|
-- Check cooldown for the bound key
|
|
local cooldown_key = KEYS[2] .. binding
|
|
local in_cooldown = redis.call("EXISTS", cooldown_key)
|
|
if in_cooldown == 1 then
|
|
redis.call("DEL", KEYS[1])
|
|
return nil
|
|
end
|
|
redis.call("EXPIRE", KEYS[1], tonumber(ARGV[1]))
|
|
return binding
|
|
"""
|
|
|
|
# Cost window sum (read-only, no cleanup on read path for performance).
|
|
# KEYS[1] = cost zset key, ARGV[1] = window_start timestamp
|
|
# Returns total token count within the window.
|
|
_COST_WINDOW_SUM_LUA = """
|
|
local key = KEYS[1]
|
|
local window_start = tonumber(ARGV[1])
|
|
local members = redis.call("ZRANGEBYSCORE", key, window_start, "+inf")
|
|
local total = 0
|
|
for _, m in ipairs(members) do
|
|
local colon = string.find(m, ":", 1, true)
|
|
if colon then
|
|
local n = tonumber(string.sub(m, colon + 1))
|
|
if n then total = total + n end
|
|
end
|
|
end
|
|
return total
|
|
"""
|
|
|
|
# Latency window average (read-only, no cleanup on read path for performance).
|
|
# KEYS[1] = latency zset key, ARGV[1] = window_start timestamp
|
|
# Returns nil when there are no samples, or avg(ms) as number.
|
|
_LATENCY_WINDOW_AVG_LUA = """
|
|
local key = KEYS[1]
|
|
local window_start = tonumber(ARGV[1])
|
|
local members = redis.call("ZRANGEBYSCORE", key, window_start, "+inf")
|
|
local total = 0
|
|
local count = 0
|
|
for _, m in ipairs(members) do
|
|
local colon = string.find(m, ":", 1, true)
|
|
if colon then
|
|
local n = tonumber(string.sub(m, colon + 1))
|
|
if n then
|
|
total = total + n
|
|
count = count + 1
|
|
end
|
|
end
|
|
end
|
|
if count == 0 then
|
|
return nil
|
|
end
|
|
return total / count
|
|
"""
|
|
|
|
|
|
async def _get_redis() -> "aioredis.Redis | None":
|
|
return await get_redis_client(require_redis=False)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Sticky session
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def get_sticky_binding(provider_id: str, session_uuid: str, ttl: int) -> str | None:
|
|
"""Get and refresh sticky session binding. Returns key_id or None."""
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return None
|
|
try:
|
|
result = await redis.eval(
|
|
_STICKY_SELECT_LUA,
|
|
2,
|
|
_sticky_key(provider_id, session_uuid),
|
|
f"{PREFIX}:{provider_id}:cooldown:",
|
|
str(ttl),
|
|
)
|
|
if result:
|
|
return result.decode() if isinstance(result, bytes) else str(result)
|
|
return None
|
|
except Exception:
|
|
logger.debug("Pool: sticky GET failed for session {}", session_uuid[:8])
|
|
return None
|
|
|
|
|
|
async def set_sticky_binding(provider_id: str, session_uuid: str, key_id: str, ttl: int) -> None:
|
|
"""Create or update sticky session binding."""
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return
|
|
try:
|
|
await redis.setex(_sticky_key(provider_id, session_uuid), ttl, key_id)
|
|
except Exception:
|
|
logger.debug("Pool: sticky SET failed for session {}", session_uuid[:8])
|
|
|
|
|
|
async def delete_sticky_binding(provider_id: str, session_uuid: str) -> None:
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return
|
|
try:
|
|
await redis.delete(_sticky_key(provider_id, session_uuid))
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# LRU
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def get_lru_scores(provider_id: str, key_ids: list[str]) -> dict[str, float]:
|
|
"""Batch-fetch LRU timestamps. Missing keys get score 0 (highest priority)."""
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return {}
|
|
try:
|
|
lru_k = _lru_key(provider_id)
|
|
scores = await redis.zmscore(lru_k, key_ids)
|
|
result: dict[str, float] = {}
|
|
for kid, score in zip(key_ids, scores):
|
|
result[kid] = float(score) if score is not None else 0.0
|
|
return result
|
|
except Exception:
|
|
logger.debug("Pool: LRU ZMSCORE failed for provider {}", provider_id[:8])
|
|
return {}
|
|
|
|
|
|
async def touch_lru(provider_id: str, key_id: str) -> None:
|
|
"""Update last-used timestamp."""
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return
|
|
try:
|
|
await redis.zadd(_lru_key(provider_id), {key_id: time.time()})
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Cooldown
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def set_cooldown(provider_id: str, key_id: str, reason: str, ttl: int) -> None:
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return
|
|
try:
|
|
pipe = redis.pipeline()
|
|
pipe.setex(_cooldown_key(provider_id, key_id), ttl, reason)
|
|
# Track in index set for O(1) count queries.
|
|
idx_key = _cooldown_index_key(provider_id)
|
|
pipe.sadd(idx_key, key_id)
|
|
# Keep index alive at least as long as the longest cooldown entry.
|
|
# Each set_cooldown call refreshes the TTL so the SET won't expire
|
|
# while there are still active cooldowns.
|
|
pipe.expire(idx_key, ttl + 60)
|
|
await pipe.execute()
|
|
logger.info(
|
|
"Pool: key {} cooldown set: {} ({}s)",
|
|
key_id[:8],
|
|
reason,
|
|
ttl,
|
|
)
|
|
except Exception:
|
|
logger.debug("Pool: cooldown SET failed for key {}", key_id[:8])
|
|
|
|
|
|
async def get_cooldown(provider_id: str, key_id: str) -> str | None:
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return None
|
|
try:
|
|
val = await redis.get(_cooldown_key(provider_id, key_id))
|
|
if val:
|
|
return val.decode() if isinstance(val, bytes) else str(val)
|
|
return None
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
async def clear_cooldown(provider_id: str, key_id: str) -> None:
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return
|
|
try:
|
|
pipe = redis.pipeline()
|
|
pipe.delete(_cooldown_key(provider_id, key_id))
|
|
pipe.srem(_cooldown_index_key(provider_id), key_id)
|
|
await pipe.execute()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
async def batch_get_cooldowns(
|
|
provider_id: str,
|
|
key_ids: list[str],
|
|
*,
|
|
include_ttl: bool = False,
|
|
) -> dict[str, str | None] | dict[str, tuple[str | None, int | None]]:
|
|
"""Batch check cooldown status for multiple keys.
|
|
|
|
When *include_ttl* is ``True``, each value is a ``(reason, ttl_seconds)``
|
|
tuple instead of a plain reason string. The TTL commands are batched in
|
|
the same pipeline so there is no extra round-trip.
|
|
"""
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
if include_ttl:
|
|
return {k: (None, None) for k in key_ids}
|
|
return {k: None for k in key_ids}
|
|
try:
|
|
pipe = redis.pipeline()
|
|
for kid in key_ids:
|
|
ck = _cooldown_key(provider_id, kid)
|
|
pipe.get(ck)
|
|
if include_ttl:
|
|
pipe.ttl(ck)
|
|
results = await pipe.execute()
|
|
|
|
if include_ttl:
|
|
out_ttl: dict[str, tuple[str | None, int | None]] = {}
|
|
# results interleave GET/TTL: [val0, ttl0, val1, ttl1, ...]
|
|
for i, kid in enumerate(key_ids):
|
|
val = results[i * 2]
|
|
ttl_val = results[i * 2 + 1]
|
|
reason: str | None = None
|
|
if val:
|
|
reason = val.decode() if isinstance(val, bytes) else str(val)
|
|
ttl_sec: int | None = None
|
|
if isinstance(ttl_val, int) and ttl_val > 0:
|
|
ttl_sec = ttl_val
|
|
out_ttl[kid] = (reason, ttl_sec)
|
|
return out_ttl
|
|
|
|
out: dict[str, str | None] = {}
|
|
for kid, val in zip(key_ids, results):
|
|
if val:
|
|
out[kid] = val.decode() if isinstance(val, bytes) else str(val)
|
|
else:
|
|
out[kid] = None
|
|
return out
|
|
except Exception:
|
|
if include_ttl:
|
|
return {k: (None, None) for k in key_ids}
|
|
return {k: None for k in key_ids}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Cost tracking
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def add_cost_entry(provider_id: str, key_id: str, tokens: int, window_seconds: int) -> None:
|
|
"""Record a cost entry (tokens used) with automatic window expiry."""
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return
|
|
try:
|
|
now = time.time()
|
|
cost_k = _cost_key(provider_id, key_id)
|
|
member = f"{uuid.uuid4().hex}:{tokens}"
|
|
window_start = now - max(int(window_seconds), 1)
|
|
pipe = redis.pipeline()
|
|
pipe.zadd(cost_k, {member: now})
|
|
# Prune expired entries on the write path (moved from read Lua script).
|
|
pipe.zremrangebyscore(cost_k, "-inf", window_start)
|
|
# Set a TTL slightly larger than the window to auto-clean abandoned keys.
|
|
pipe.expire(cost_k, window_seconds + 600)
|
|
await pipe.execute()
|
|
except Exception:
|
|
logger.debug("Pool: cost ADD failed for key {}", key_id[:8])
|
|
|
|
|
|
async def get_cost_window_total(provider_id: str, key_id: str, window_seconds: int) -> int:
|
|
"""Sum tokens used within the rolling window (single key)."""
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return 0
|
|
try:
|
|
now = time.time()
|
|
window_start = now - window_seconds
|
|
cost_k = _cost_key(provider_id, key_id)
|
|
result = await redis.eval(_COST_WINDOW_SUM_LUA, 1, cost_k, str(window_start))
|
|
return int(result) if result else 0
|
|
except Exception:
|
|
logger.debug("Pool: cost SUM failed for key {}", key_id[:8])
|
|
return 0
|
|
|
|
|
|
async def batch_get_cost_totals(
|
|
provider_id: str, key_ids: list[str], window_seconds: int
|
|
) -> dict[str, int]:
|
|
"""Batch-fetch cost totals for multiple keys using pipeline + Lua."""
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return {k: 0 for k in key_ids}
|
|
try:
|
|
now = time.time()
|
|
window_start = now - window_seconds
|
|
pipe = redis.pipeline()
|
|
for kid in key_ids:
|
|
cost_k = _cost_key(provider_id, kid)
|
|
pipe.eval(_COST_WINDOW_SUM_LUA, 1, cost_k, str(window_start))
|
|
results = await pipe.execute()
|
|
out: dict[str, int] = {}
|
|
for kid, val in zip(key_ids, results):
|
|
out[kid] = int(val) if val else 0
|
|
return out
|
|
except Exception:
|
|
logger.debug("Pool: batch cost SUM failed for provider {}", provider_id[:8])
|
|
return {k: 0 for k in key_ids}
|
|
|
|
|
|
async def record_latency(
|
|
provider_id: str,
|
|
key_id: str,
|
|
ttfb_ms: int,
|
|
window_seconds: int,
|
|
sample_limit: int,
|
|
) -> None:
|
|
"""Record one TTFB sample with rolling-window cleanup."""
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return
|
|
try:
|
|
now = time.time()
|
|
latency_k = _latency_key(provider_id, key_id)
|
|
sample = max(int(ttfb_ms), 0)
|
|
member = f"{uuid.uuid4().hex}:{sample}"
|
|
pipe = redis.pipeline()
|
|
pipe.zadd(latency_k, {member: now})
|
|
window_start = now - max(int(window_seconds), 1)
|
|
pipe.zremrangebyscore(latency_k, "-inf", window_start)
|
|
capped_limit = max(int(sample_limit), 1)
|
|
pipe.zremrangebyrank(latency_k, 0, -(capped_limit + 1))
|
|
pipe.expire(latency_k, max(int(window_seconds), 1) + 600)
|
|
await pipe.execute()
|
|
except Exception:
|
|
logger.debug("Pool: latency ADD failed for key {}", key_id[:8])
|
|
|
|
|
|
async def batch_get_latency_avgs(
|
|
provider_id: str,
|
|
key_ids: list[str],
|
|
window_seconds: int,
|
|
) -> dict[str, float]:
|
|
"""Batch-fetch latency averages (ms) for keys in a rolling window."""
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return {}
|
|
try:
|
|
now = time.time()
|
|
window_start = now - max(int(window_seconds), 1)
|
|
pipe = redis.pipeline()
|
|
for kid in key_ids:
|
|
pipe.eval(_LATENCY_WINDOW_AVG_LUA, 1, _latency_key(provider_id, kid), str(window_start))
|
|
results = await pipe.execute()
|
|
out: dict[str, float] = {}
|
|
for kid, val in zip(key_ids, results):
|
|
if val is None:
|
|
continue
|
|
try:
|
|
out[kid] = float(val)
|
|
except Exception:
|
|
continue
|
|
return out
|
|
except Exception:
|
|
logger.debug("Pool: batch latency AVG failed for provider {}", provider_id[:8])
|
|
return {}
|
|
|
|
|
|
async def clear_cost(provider_id: str, key_id: str) -> None:
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return
|
|
try:
|
|
await redis.delete(_cost_key(provider_id, key_id))
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# OAuth token cache
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def cache_oauth_token(key_id: str, token: str, ttl: int) -> None:
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return
|
|
try:
|
|
if ttl > 0:
|
|
await redis.setex(_oauth_cache_key(key_id), ttl, token)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
async def get_cached_oauth_token(key_id: str) -> str | None:
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return None
|
|
try:
|
|
val = await redis.get(_oauth_cache_key(key_id))
|
|
if val:
|
|
return val.decode() if isinstance(val, bytes) else str(val)
|
|
return None
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
async def invalidate_oauth_token_cache(key_id: str) -> None:
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return
|
|
try:
|
|
await redis.delete(_oauth_cache_key(key_id))
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pool status query (admin)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def get_sticky_session_count(provider_id: str) -> int:
|
|
"""Approximate count of active sticky sessions (via SCAN, for admin display only)."""
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return 0
|
|
try:
|
|
pattern = f"{PREFIX}:{provider_id}:sticky:*"
|
|
count = 0
|
|
async for _ in redis.scan_iter(match=pattern, count=100):
|
|
count += 1
|
|
return count
|
|
except Exception:
|
|
return 0
|
|
|
|
|
|
async def get_key_sticky_count(provider_id: str, key_id: str) -> int:
|
|
"""Count sticky sessions bound to a specific key (admin only).
|
|
|
|
Uses batched SCAN + pipeline MGET to reduce Redis round-trips.
|
|
"""
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return 0
|
|
try:
|
|
pattern = f"{PREFIX}:{provider_id}:sticky:*"
|
|
count = 0
|
|
batch: list[bytes | str] = []
|
|
async for k in redis.scan_iter(match=pattern, count=200):
|
|
batch.append(k)
|
|
if len(batch) >= 200:
|
|
vals = await redis.mget(batch)
|
|
for val in vals:
|
|
if val:
|
|
bound_id = val.decode() if isinstance(val, bytes) else str(val)
|
|
if bound_id == key_id:
|
|
count += 1
|
|
batch.clear()
|
|
if batch:
|
|
vals = await redis.mget(batch)
|
|
for val in vals:
|
|
if val:
|
|
bound_id = val.decode() if isinstance(val, bytes) else str(val)
|
|
if bound_id == key_id:
|
|
count += 1
|
|
return count
|
|
except Exception:
|
|
return 0
|
|
|
|
|
|
async def batch_get_key_sticky_counts(
|
|
provider_id: str,
|
|
key_ids: list[str],
|
|
) -> dict[str, int]:
|
|
"""Count sticky sessions for multiple keys in a single scan (admin only)."""
|
|
if not key_ids:
|
|
return {}
|
|
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return {kid: 0 for kid in key_ids}
|
|
|
|
target_ids = set(key_ids)
|
|
counts: dict[str, int] = {kid: 0 for kid in key_ids}
|
|
|
|
try:
|
|
pattern = f"{PREFIX}:{provider_id}:sticky:*"
|
|
batch: list[bytes | str] = []
|
|
async for key in redis.scan_iter(match=pattern, count=200):
|
|
batch.append(key)
|
|
if len(batch) >= 200:
|
|
vals = await redis.mget(batch)
|
|
for val in vals:
|
|
if not val:
|
|
continue
|
|
bound_id = val.decode() if isinstance(val, bytes) else str(val)
|
|
if bound_id in target_ids:
|
|
counts[bound_id] = counts.get(bound_id, 0) + 1
|
|
batch.clear()
|
|
|
|
if batch:
|
|
vals = await redis.mget(batch)
|
|
for val in vals:
|
|
if not val:
|
|
continue
|
|
bound_id = val.decode() if isinstance(val, bytes) else str(val)
|
|
if bound_id in target_ids:
|
|
counts[bound_id] = counts.get(bound_id, 0) + 1
|
|
|
|
return counts
|
|
except Exception:
|
|
return {kid: 0 for kid in key_ids}
|
|
|
|
|
|
async def get_cooldown_ttl(provider_id: str, key_id: str) -> int | None:
|
|
"""Get remaining cooldown TTL in seconds. None = no cooldown."""
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return None
|
|
try:
|
|
ttl = await redis.ttl(_cooldown_key(provider_id, key_id))
|
|
return ttl if ttl > 0 else None
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
async def batch_get_cooldown_ttls(provider_id: str, key_ids: list[str]) -> dict[str, int | None]:
|
|
"""Batch-fetch cooldown TTLs for multiple keys using pipeline."""
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return {k: None for k in key_ids}
|
|
try:
|
|
pipe = redis.pipeline()
|
|
for kid in key_ids:
|
|
pipe.ttl(_cooldown_key(provider_id, kid))
|
|
results = await pipe.execute()
|
|
out: dict[str, int | None] = {}
|
|
for kid, ttl in zip(key_ids, results):
|
|
out[kid] = int(ttl) if isinstance(ttl, int) and ttl > 0 else None
|
|
return out
|
|
except Exception:
|
|
return {k: None for k in key_ids}
|
|
|
|
|
|
async def batch_count_provider_cooldowns(provider_ids: list[str]) -> dict[str, int]:
|
|
"""Count cooldown entries per provider using the cooldown index set.
|
|
|
|
Uses ``SCARD`` on the ``ap:{pid}:cooldown_idx`` set for O(1) count
|
|
instead of scanning the key-space. The index set is maintained by
|
|
:func:`set_cooldown` / :func:`clear_cooldown`.
|
|
|
|
Note: the index set may contain stale entries (expired cooldowns whose
|
|
TTL elapsed before an explicit clear). This over-count is acceptable
|
|
for admin display purposes -- precision is not critical here.
|
|
"""
|
|
if not provider_ids:
|
|
return {}
|
|
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return {pid: 0 for pid in provider_ids}
|
|
|
|
try:
|
|
pipe = redis.pipeline()
|
|
for pid in provider_ids:
|
|
pipe.scard(_cooldown_index_key(pid))
|
|
results = await pipe.execute()
|
|
counts: dict[str, int] = {}
|
|
for pid, val in zip(provider_ids, results):
|
|
counts[pid] = max(int(val or 0), 0)
|
|
return counts
|
|
except Exception:
|
|
return {pid: 0 for pid in provider_ids}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Stream timeout counter
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_STREAM_TIMEOUT_KEY_FMT = f"{PREFIX}:{{}}:stream_timeout:{{}}"
|
|
|
|
|
|
def _stream_timeout_key(provider_id: str, key_id: str) -> str:
|
|
return _STREAM_TIMEOUT_KEY_FMT.format(provider_id, key_id)
|
|
|
|
|
|
async def incr_stream_timeout_count(
|
|
provider_id: str,
|
|
key_id: str,
|
|
window_seconds: int,
|
|
) -> int:
|
|
"""Increment stream timeout counter and return count within the window.
|
|
|
|
Uses a ZSET with timestamps as scores. Old entries beyond the window
|
|
are pruned on each call. Returns the count of timeouts in the window.
|
|
"""
|
|
redis = await _get_redis()
|
|
if redis is None:
|
|
return 0
|
|
try:
|
|
now = time.time()
|
|
window_start = now - window_seconds
|
|
key = _stream_timeout_key(provider_id, key_id)
|
|
member = f"{uuid.uuid4().hex}"
|
|
pipe = redis.pipeline()
|
|
pipe.zremrangebyscore(key, "-inf", window_start)
|
|
pipe.zadd(key, {member: now})
|
|
pipe.zcard(key)
|
|
pipe.expire(key, window_seconds + 60)
|
|
results = await pipe.execute()
|
|
count = int(results[2]) if results[2] else 0
|
|
return count
|
|
except Exception:
|
|
logger.debug("Pool: stream timeout INCR failed for key {}", key_id[:8])
|
|
return 0
|