mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
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:
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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
54
src/core/redis_utils.py
Normal 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
|
||||
@@ -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:
|
||||
|
||||
@@ -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}
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user