perf: Redis 操作优化,移除号池 key 列表的 Usage 聚合查询

- 号池 key 列表移除 Usage 表聚合查询(total_tokens/total_cost_usd),消除慢 SQL
- cooldown 计数从 SCAN 改为 SCARD (O(1)),通过 cooldown_idx SET 维护索引
- 读路径 Lua 脚本移除 ZREMRANGEBYSCORE,清理移至写路径减少开销
- affinity 清理从 KEYS 改为 SCAN 分批删除,invalidate_all_for_provider 改用 pipeline 批量 MGET+UNLINK
- 缓存监控页添加清除按钮 loading 状态,Redis SCAN 并发限制为 4
- 提取 redis_utils 模块统一 SCAN+批量删除逻辑
- 调度热路径跳过 cooldown TTL 查询,account_state 预计算移出循环
This commit is contained in:
fawney19
2026-03-09 14:54:45 +08:00
parent 9516619b92
commit afbb1b9a5d
11 changed files with 352 additions and 187 deletions

View File

@@ -20,6 +20,7 @@ from src.api.base.pipeline import ApiRequestPipeline
from src.clients.redis_client import get_redis_client_sync
from src.core.crypto import crypto_service
from src.core.logger import logger
from src.core.redis_utils import scan_delete_pattern
from src.database import get_db
from src.models.database import ApiKey, User
from src.services.scheduling.affinity_manager import get_affinity_manager
@@ -28,6 +29,8 @@ from src.services.system.config import SystemConfigService
router = APIRouter(prefix="/api/admin/monitoring/cache", tags=["Admin - Monitoring: Cache"])
pipeline = ApiRequestPipeline()
REDIS_SCAN_BATCH_SIZE = 200
REDIS_DELETE_BATCH_SIZE = 500
def mask_api_key(api_key: str | None, prefix_len: int = 8, suffix_len: int = 4) -> str | None:
@@ -1649,13 +1652,16 @@ class AdminRedisCacheCategoriesAdapter(AdminApiAdapter):
"data": {"available": False, "message": "Redis 未启用"},
}
async def _count_keys(pattern: str) -> int:
count = 0
async for _ in redis.scan_iter(match=pattern, count=500):
count += 1
return count
semaphore = asyncio.Semaphore(4)
# 并行扫描所有分类的 key 数量,避免串行 20 次 SCAN
async def _count_keys(pattern: str) -> int:
async with semaphore:
count = 0
async for _ in redis.scan_iter(match=pattern, count=500):
count += 1
return count
# 限制扫描并发,避免刷新监控页时对 Redis 造成瞬时压力
counts = await asyncio.gather(
*[_count_keys(pattern) for _, _, pattern, _ in _CACHE_CATEGORIES]
)
@@ -1719,16 +1725,7 @@ class AdminClearRedisCacheCategoryAdapter(AdminApiAdapter):
if not redis:
raise HTTPException(status_code=503, detail="Redis 未启用")
keys_to_delete: list[str] = []
async for key in redis.scan_iter(match=pattern, count=200):
keys_to_delete.append(key)
deleted_count = 0
# 分批删除,避免单次 DELETE 命令阻塞 Redis 事件循环
batch_size = 1000
for i in range(0, len(keys_to_delete), batch_size):
batch = keys_to_delete[i : i + batch_size]
deleted_count += await redis.delete(*batch)
deleted_count = await scan_delete_pattern(redis, pattern)
logger.warning(
"已清除 Redis 缓存分类(管理员操作): {} ({}), pattern={}, deleted={}",
@@ -1767,17 +1764,8 @@ class AdminClearAllModelMappingCacheAdapter(AdminApiAdapter):
if not redis:
raise HTTPException(status_code=503, detail="Redis 未启用")
deleted_count = 0
# 删除所有模型相关的缓存键
keys_to_delete = []
async for key in redis.scan_iter(match="model:*", count=100):
keys_to_delete.append(key)
async for key in redis.scan_iter(match="global_model:*", count=100):
keys_to_delete.append(key)
if keys_to_delete:
deleted_count = await redis.delete(*keys_to_delete)
deleted_count = await scan_delete_pattern(redis, "model:*")
deleted_count += await scan_delete_pattern(redis, "global_model:*")
logger.warning(f"已清除所有模型映射缓存(管理员操作): {deleted_count} 个键")
context.add_audit_metadata(

View File

@@ -29,7 +29,7 @@ from src.core.crypto import crypto_service
from src.core.exceptions import NotFoundException
from src.core.logger import logger
from src.database import get_db
from src.models.database import Provider, ProviderAPIKey, Usage
from src.models.database import Provider, ProviderAPIKey
from src.services.provider.fingerprint import generate_fingerprint
from src.services.provider.pool import redis_ops as pool_redis
from src.services.provider.pool.account_state import resolve_pool_account_state
@@ -634,7 +634,6 @@ class AdminListPoolKeysAdapter(AdminApiAdapter):
count_query_ms = 0.0
keys_query_ms = 0.0
redis_state_ms = 0.0
usage_stats_ms = 0.0
serialize_ms = 0.0
db = context.db
@@ -782,37 +781,6 @@ class AdminListPoolKeysAdapter(AdminApiAdapter):
{},
)
usage_stats_by_key: dict[str, dict[str, Any]] = {}
if key_ids:
usage_stats_started_at = time.perf_counter()
usage_rows = (
db.query(
Usage.provider_api_key_id.label("key_id"),
func.count(Usage.id).label("request_count"),
func.coalesce(func.sum(Usage.total_tokens), 0).label("total_tokens"),
func.coalesce(func.sum(Usage.total_cost_usd), 0.0).label("total_cost_usd"),
func.max(Usage.created_at).label("last_used_at"),
)
.filter(
Usage.provider_id == pid,
Usage.provider_api_key_id.in_(key_ids),
Usage.status.notin_(["pending", "streaming"]),
)
.group_by(Usage.provider_api_key_id)
.all()
)
usage_stats_ms = (time.perf_counter() - usage_stats_started_at) * 1000.0
usage_stats_by_key = {
str(row.key_id): {
"request_count": int(row.request_count or 0),
"total_tokens": int(row.total_tokens or 0),
"total_cost_usd": float(row.total_cost_usd or 0.0),
"last_used_at": getattr(row, "last_used_at", None),
}
for row in usage_rows
if getattr(row, "key_id", None)
}
key_details: list[PoolKeyDetail] = []
serialize_started_at = time.perf_counter()
for k in keys:
@@ -895,15 +863,8 @@ class AdminListPoolKeysAdapter(AdminApiAdapter):
if isinstance(getattr(k, "api_formats", None), list)
else []
)
key_usage_stats = usage_stats_by_key.get(kid, {})
key_request_count = int(
key_usage_stats.get("request_count") or getattr(k, "request_count", 0) or 0
)
key_total_tokens = int(key_usage_stats.get("total_tokens") or 0)
key_total_cost_usd = float(key_usage_stats.get("total_cost_usd") or 0.0)
key_last_used_at = getattr(k, "last_used_at", None) or key_usage_stats.get(
"last_used_at"
)
key_request_count = int(getattr(k, "request_count", 0) or 0)
key_last_used_at = getattr(k, "last_used_at", None)
oauth_auth_config = _extract_oauth_auth_config(k)
key_details.append(
@@ -962,8 +923,6 @@ class AdminListPoolKeysAdapter(AdminApiAdapter):
cost_window_usage=cost_usage,
cost_limit=cost_limit,
request_count=key_request_count,
total_tokens=key_total_tokens,
total_cost_usd=key_total_cost_usd,
sticky_sessions=sticky_counts.get(kid, 0),
lru_score=lru_scores.get(kid),
created_at=(
@@ -980,7 +939,7 @@ class AdminListPoolKeysAdapter(AdminApiAdapter):
total_ms = (time.perf_counter() - started_at) * 1000.0
logger.info(
"[POOL_KEYS_TIMING] provider={} page={} page_size={} status={} search={} total={} count_ms={:.2f} fetch_ms={:.2f} redis_ms={:.2f} usage_ms={:.2f} serialize_ms={:.2f} total_ms={:.2f}",
"[POOL_KEYS_TIMING] provider={} page={} page_size={} status={} search={} total={} count_ms={:.2f} fetch_ms={:.2f} redis_ms={:.2f} serialize_ms={:.2f} total_ms={:.2f}",
pid[:8],
self.page,
self.page_size,
@@ -990,7 +949,6 @@ class AdminListPoolKeysAdapter(AdminApiAdapter):
count_query_ms,
keys_query_ms,
redis_state_ms,
usage_stats_ms,
serialize_ms,
total_ms,
)

View File

@@ -103,8 +103,6 @@ class PoolKeyDetail(BaseModel):
cost_window_usage: int = 0
cost_limit: int | None = None
request_count: int = 0
total_tokens: int = 0
total_cost_usd: float = 0.0
sticky_sessions: int = 0
lru_score: float | None = None
created_at: str | None = None

54
src/core/redis_utils.py Normal file
View File

@@ -0,0 +1,54 @@
"""Redis batch-delete utilities.
Provides SCAN-based pattern deletion with UNLINK preference to minimise
blocking on the Redis server. Used by cache monitoring and affinity
manager modules.
"""
from __future__ import annotations
from typing import Any
from src.core.logger import logger
# Defaults matching existing usage across the codebase.
DEFAULT_SCAN_BATCH_SIZE = 200
DEFAULT_DELETE_BATCH_SIZE = 500
async def delete_redis_keys(redis: Any, keys: list[str]) -> int:
"""Batch-delete Redis keys, preferring UNLINK over DELETE."""
if not keys:
return 0
try:
unlink = getattr(redis, "unlink", None)
if callable(unlink):
return int(await unlink(*keys))
except Exception as exc:
logger.debug("Redis UNLINK failed, falling back to DELETE: {}", exc)
return int(await redis.delete(*keys))
async def scan_delete_pattern(
redis: Any,
pattern: str,
*,
scan_batch_size: int = DEFAULT_SCAN_BATCH_SIZE,
delete_batch_size: int = DEFAULT_DELETE_BATCH_SIZE,
) -> int:
"""SCAN + batch-delete keys matching *pattern* without blocking Redis."""
deleted_count = 0
cursor: int | str = 0
while True:
cursor, keys = await redis.scan(cursor=cursor, match=pattern, count=scan_batch_size)
if keys:
for i in range(0, len(keys), delete_batch_size):
batch = keys[i : i + delete_batch_size]
deleted_count += await delete_redis_keys(redis, batch)
if int(cursor) == 0:
break
return deleted_count

View File

@@ -121,7 +121,9 @@ class PoolManager:
all_key_ids = [str(c.key.id) for c in candidates]
# Fire independent Redis queries concurrently.
_cooldown_coro = redis_ops.batch_get_cooldowns(pid, all_key_ids, include_ttl=True)
# Only fetch reason (no TTL) on the scheduling hot path -- TTL is only
# used for trace display and costs an extra pipeline command per key.
_cooldown_coro = redis_ops.batch_get_cooldowns(pid, all_key_ids, include_ttl=False)
_cost_coro = (
redis_ops.batch_get_cost_totals(pid, all_key_ids, self.config.cost_window_seconds)
if (
@@ -154,17 +156,7 @@ class PoolManager:
gathered = await asyncio.gather(*coros)
cooldowns_raw = gathered[0]
# cooldowns_raw: dict[str, tuple[str | None, int | None]]
cooldowns: dict[str, str | None] = {}
cooldown_ttls: dict[str, int | None] = {}
for kid, val in cooldowns_raw.items():
if isinstance(val, tuple):
cooldowns[kid] = val[0]
cooldown_ttls[kid] = val[1]
else:
cooldowns[kid] = val
cooldown_ttls[kid] = None
cooldowns: dict[str, str | None] = gathered[0]
# Cost check
cost_exhausted: set[str] = set()
@@ -225,6 +217,17 @@ class PoolManager:
pass
# --- 3. Classify candidates -----------------------------------
# Pre-compute account states to avoid repeated dict parsing inside loop.
account_states: dict[str, Any] = {}
for c in candidates:
kid = str(c.key.id)
if kid not in account_states:
account_states[kid] = resolve_pool_account_state(
provider_type=provider_type,
upstream_metadata=getattr(c.key, "upstream_metadata", None),
oauth_invalid_reason=getattr(c.key, "oauth_invalid_reason", None),
)
sticky_candidate: ProviderCandidate | None = None
available: list[ProviderCandidate] = []
skipped: list[ProviderCandidate] = []
@@ -246,12 +249,8 @@ class PoolManager:
trace.candidate_traces[kid] = ct
continue
# Cooldown?
account_state = resolve_pool_account_state(
provider_type=provider_type,
upstream_metadata=getattr(c.key, "upstream_metadata", None),
oauth_invalid_reason=getattr(c.key, "oauth_invalid_reason", None),
)
# Account blocked?
account_state = account_states[kid]
if account_state.blocked:
c.is_skipped = True
skip_reason = account_state.reason or account_state.label or "account blocked"
@@ -274,7 +273,7 @@ class PoolManager:
ct.skipped = True
ct.skip_type = "cooldown"
ct.cooldown_reason = cd_reason
ct.cooldown_ttl = cooldown_ttls.get(kid)
ct.cooldown_ttl = None # TTL skipped on hot path for perf
_attach_pool_extra(c, ct)
trace.candidate_traces[kid] = ct
continue
@@ -578,18 +577,24 @@ class PoolManager:
pass
# --- 3. Classify keys -------------------------------------------------
# Pre-compute account states to avoid repeated dict parsing inside loop.
account_states_sk: dict[str, Any] = {}
for k in keys:
kid = str(k.id)
if kid not in account_states_sk:
account_states_sk[kid] = resolve_pool_account_state(
provider_type=self.provider_type,
upstream_metadata=getattr(k, "upstream_metadata", None),
oauth_invalid_reason=getattr(k, "oauth_invalid_reason", None),
)
sticky_key: ProviderAPIKey | None = None
available: list[ProviderAPIKey] = []
for k in keys:
kid = str(k.id)
account_state = resolve_pool_account_state(
provider_type=self.provider_type,
upstream_metadata=getattr(k, "upstream_metadata", None),
oauth_invalid_reason=getattr(k, "oauth_invalid_reason", None),
)
if account_state.blocked:
if account_states_sk[kid].blocked:
continue
if cooldowns.get(kid) is not None:

View File

@@ -16,10 +16,9 @@ provider_oauth_token_cache:{key_id} STRING -> access_token (TTL: expires
from __future__ import annotations
import asyncio
import time
import uuid
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING
from src.clients.redis_client import get_redis_client
from src.core.logger import logger
@@ -42,6 +41,11 @@ 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}"
@@ -78,13 +82,12 @@ redis.call("EXPIRE", KEYS[1], tonumber(ARGV[1]))
return binding
"""
# Cost window cleanup + sum tokens in a single round-trip.
# 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])
redis.call("ZREMRANGEBYSCORE", key, "-inf", window_start)
local members = redis.call("ZRANGEBYSCORE", key, window_start, "+inf")
local total = 0
for _, m in ipairs(members) do
@@ -97,13 +100,12 @@ end
return total
"""
# Latency window cleanup + average in a single round-trip.
# 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])
redis.call("ZREMRANGEBYSCORE", key, "-inf", window_start)
local members = redis.call("ZRANGEBYSCORE", key, window_start, "+inf")
local total = 0
local count = 0
@@ -218,7 +220,16 @@ async def set_cooldown(provider_id: str, key_id: str, reason: str, ttl: int) ->
if redis is None:
return
try:
await redis.setex(_cooldown_key(provider_id, key_id), ttl, reason)
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],
@@ -247,7 +258,10 @@ async def clear_cooldown(provider_id: str, key_id: str) -> None:
if redis is None:
return
try:
await redis.delete(_cooldown_key(provider_id, key_id))
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
@@ -320,8 +334,11 @@ async def add_cost_entry(provider_id: str, key_id: str, tokens: int, window_seco
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()
@@ -605,11 +622,15 @@ async def batch_get_cooldown_ttls(provider_id: str, key_ids: list[str]) -> dict[
async def batch_count_provider_cooldowns(provider_ids: list[str]) -> dict[str, int]:
"""Count cooldown entries per provider using per-provider SCAN.
"""Count cooldown entries per provider using the cooldown index set.
Each provider's cooldown keys are scanned independently with a
targeted pattern ``ap:{pid}:cooldown:*``, avoiding full-keyspace traversal.
Multiple providers are scanned concurrently via ``asyncio.gather``.
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 {}
@@ -618,25 +639,14 @@ async def batch_count_provider_cooldowns(provider_ids: list[str]) -> dict[str, i
if redis is None:
return {pid: 0 for pid in provider_ids}
async def _count_one(r: Any, pid: str) -> tuple[str, int]:
pattern = f"{PREFIX}:{pid}:cooldown:*"
count = 0
async for _key in r.scan_iter(match=pattern, count=200):
count += 1
return pid, count
try:
results = await asyncio.gather(
*[_count_one(redis, pid) for pid in provider_ids],
return_exceptions=True,
)
counts: dict[str, int] = {}
for r in results:
if isinstance(r, Exception):
continue
counts[r[0]] = r[1]
pipe = redis.pipeline()
for pid in provider_ids:
counts.setdefault(pid, 0)
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}

View File

@@ -29,6 +29,7 @@ from typing import Any, NamedTuple
from src.config.constants import CacheTTL
from src.core.logger import logger
from src.core.redis_utils import delete_redis_keys, scan_delete_pattern
class CacheAffinity(NamedTuple):
@@ -76,6 +77,8 @@ class CacheAffinityManager:
# 默认缓存TTL- 使用统一常量
DEFAULT_CACHE_TTL = CacheTTL.CACHE_AFFINITY
REDIS_SCAN_BATCH_SIZE = 200
REDIS_DELETE_BATCH_SIZE = 500
def __init__(
self, redis_client: Any | None = None, default_ttl: int = DEFAULT_CACHE_TTL
@@ -261,6 +264,29 @@ class CacheAffinityManager:
await self._set_l1_entry(cache_key, None)
async def _delete_redis_keys(self, keys: list[str]) -> int:
if self._is_memory_backend() or not keys:
return 0
return await delete_redis_keys(self.redis, keys)
async def _scan_delete_pattern(self, pattern: str) -> int:
if self._is_memory_backend():
return 0
return await scan_delete_pattern(
self.redis,
pattern,
scan_batch_size=self.REDIS_SCAN_BATCH_SIZE,
delete_batch_size=self.REDIS_DELETE_BATCH_SIZE,
)
async def _clear_l1_entries_by_prefix(self, prefix: str) -> int:
"""清理匹配前缀的 L1 本地缓存。"""
async with self._l1_lock:
keys_to_remove = [key for key in self._l1_cache if key.startswith(prefix)]
for key in keys_to_remove:
self._l1_cache.pop(key, None)
return len(keys_to_remove)
async def _snapshot_memory_items(self) -> dict[str, dict[str, Any]]:
"""复制内存存储内容(仅内存模式使用)"""
lock = self._get_memory_lock()
@@ -494,20 +520,49 @@ class CacheAffinityManager:
invalidated_count = 0
if not self._is_memory_backend():
pattern = "cache_affinity:*"
keys = await self.redis.keys(pattern)
cursor: int | str = 0
while True:
cursor, scan_keys = await self.redis.scan(
cursor=cursor,
match="cache_affinity:*",
count=self.REDIS_SCAN_BATCH_SIZE,
)
if scan_keys:
# Pipeline batch GET to reduce round-trips.
values = await self.redis.mget(scan_keys)
keys_to_delete: list[str] = []
for key, raw in zip(scan_keys, values):
if not raw:
continue
try:
data = json.loads(raw)
except Exception:
continue
if data.get("provider_id") == provider_id:
keys_to_delete.append(key)
if keys_to_delete:
deleted = await self._delete_redis_keys(keys_to_delete)
invalidated_count += deleted
self._stats["cache_invalidations"] += deleted
# Clear L1 for deleted keys.
for key in keys_to_delete:
await self._set_l1_entry(key, None)
if int(cursor) == 0:
break
else:
keys = list((await self._snapshot_memory_items()).keys())
for key in keys:
affinity_dict = await self._load_affinity_dict(key)
if not affinity_dict:
continue
for key in keys:
affinity_dict = await self._load_affinity_dict(key)
if not affinity_dict:
continue
if affinity_dict.get("provider_id") == provider_id:
await self._delete_affinity_key(key)
invalidated_count += 1
self._stats["cache_invalidations"] += 1
if affinity_dict.get("provider_id") == provider_id:
await self._delete_affinity_key(key)
invalidated_count += 1
self._stats["cache_invalidations"] += 1
if invalidated_count > 0:
logger.debug(
@@ -530,17 +585,17 @@ class CacheAffinityManager:
"""
try:
if not self._is_memory_backend():
keys = await self.redis.keys("cache_affinity:*")
if keys:
await self.redis.delete(*keys)
logger.debug(f"清除所有Redis缓存亲和性: {len(keys)}")
return len(keys)
return 0
count = await self._scan_delete_pattern("cache_affinity:*")
await self._clear_l1_entries_by_prefix("cache_affinity:")
if count:
logger.debug(f"清除所有Redis缓存亲和性: {count}")
return count
lock = self._get_memory_lock()
async with lock:
count = len(self._memory_store)
self._memory_store.clear()
await self._clear_l1_entries_by_prefix("cache_affinity:")
if count:
logger.debug(f"清除所有内存缓存亲和性: {count}")
return count