mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
afbb1b9a5d
- 号池 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 预计算移出循环
132 lines
4.8 KiB
Python
132 lines
4.8 KiB
Python
import json
|
|
from fnmatch import fnmatch
|
|
|
|
import pytest
|
|
|
|
from src.services.scheduling.affinity_manager import CacheAffinityManager
|
|
|
|
|
|
class _FakeRedis:
|
|
def __init__(self, store: dict[str, str]):
|
|
self.store = dict(store)
|
|
self.keys_called = 0
|
|
self.scan_calls = 0
|
|
self.unlink_batches: list[tuple[str, ...]] = []
|
|
self.delete_batches: list[tuple[str, ...]] = []
|
|
|
|
async def scan( # type: ignore[override]
|
|
self, cursor: int = 0, match: str | None = None, count: int = 10
|
|
) -> tuple[int, list[str]]:
|
|
self.scan_calls += 1
|
|
pattern = match or "*"
|
|
matched = [key for key in sorted(self.store) if fnmatch(key, pattern)]
|
|
start = int(cursor)
|
|
batch = matched[start : start + count]
|
|
next_cursor = 0 if start + count >= len(matched) else start + count
|
|
return next_cursor, batch
|
|
|
|
async def get(self, key: str) -> str | None:
|
|
return self.store.get(key)
|
|
|
|
async def mget(self, keys: list[str]) -> list[str | None]:
|
|
return [self.store.get(k) for k in keys]
|
|
|
|
async def unlink(self, *keys: str) -> int:
|
|
self.unlink_batches.append(tuple(keys))
|
|
deleted = 0
|
|
for key in keys:
|
|
if key in self.store:
|
|
deleted += 1
|
|
self.store.pop(key, None)
|
|
return deleted
|
|
|
|
async def delete(self, *keys: str) -> int:
|
|
self.delete_batches.append(tuple(keys))
|
|
deleted = 0
|
|
for key in keys:
|
|
if key in self.store:
|
|
deleted += 1
|
|
self.store.pop(key, None)
|
|
return deleted
|
|
|
|
async def keys(self, _pattern: str) -> list[str]:
|
|
self.keys_called += 1
|
|
raise AssertionError("clear operations should not call Redis KEYS")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_clear_all_uses_scan_and_clears_l1_cache() -> None:
|
|
redis = _FakeRedis(
|
|
{
|
|
"cache_affinity:user-1:openai:model-a": json.dumps({"provider_id": "provider-a"}),
|
|
"cache_affinity:user-2:openai:model-b": json.dumps({"provider_id": "provider-b"}),
|
|
"cache_affinity:user-3:claude:model-c": json.dumps({"provider_id": "provider-c"}),
|
|
"other:key": json.dumps({"provider_id": "provider-x"}),
|
|
}
|
|
)
|
|
manager = CacheAffinityManager(redis_client=redis)
|
|
|
|
await manager._set_l1_entry(
|
|
"cache_affinity:user-1:openai:model-a", {"provider_id": "provider-a"}
|
|
)
|
|
await manager._set_l1_entry("other:key", {"provider_id": "provider-x"})
|
|
|
|
deleted_count = await manager.clear_all()
|
|
|
|
assert deleted_count == 3
|
|
assert redis.keys_called == 0
|
|
assert redis.scan_calls >= 1
|
|
assert redis.unlink_batches == [
|
|
(
|
|
"cache_affinity:user-1:openai:model-a",
|
|
"cache_affinity:user-2:openai:model-b",
|
|
"cache_affinity:user-3:claude:model-c",
|
|
)
|
|
]
|
|
assert "other:key" in redis.store
|
|
assert all(not key.startswith("cache_affinity:") for key in redis.store)
|
|
assert await manager._get_l1_entry("cache_affinity:user-1:openai:model-a") is None
|
|
assert await manager._get_l1_entry("other:key") == {"provider_id": "provider-x"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_invalidate_all_for_provider_uses_scan_not_keys() -> None:
|
|
redis = _FakeRedis(
|
|
{
|
|
"cache_affinity:user-1:openai:model-a": json.dumps({"provider_id": "provider-a"}),
|
|
"cache_affinity:user-2:openai:model-b": json.dumps({"provider_id": "provider-b"}),
|
|
"cache_affinity:user-3:claude:model-c": json.dumps({"provider_id": "provider-a"}),
|
|
}
|
|
)
|
|
manager = CacheAffinityManager(redis_client=redis)
|
|
|
|
await manager._set_l1_entry(
|
|
"cache_affinity:user-1:openai:model-a", {"provider_id": "provider-a"}
|
|
)
|
|
await manager._set_l1_entry(
|
|
"cache_affinity:user-2:openai:model-b", {"provider_id": "provider-b"}
|
|
)
|
|
await manager._set_l1_entry(
|
|
"cache_affinity:user-3:claude:model-c", {"provider_id": "provider-a"}
|
|
)
|
|
|
|
deleted_count = await manager.invalidate_all_for_provider("provider-a")
|
|
|
|
assert deleted_count == 2
|
|
assert redis.keys_called == 0
|
|
assert redis.scan_calls >= 1
|
|
# Pipeline path: batch MGET + batch UNLINK (not per-key DELETE).
|
|
assert len(redis.unlink_batches) == 1
|
|
assert set(redis.unlink_batches[0]) == {
|
|
"cache_affinity:user-1:openai:model-a",
|
|
"cache_affinity:user-3:claude:model-c",
|
|
}
|
|
assert redis.delete_batches == []
|
|
assert "cache_affinity:user-2:openai:model-b" in redis.store
|
|
assert "cache_affinity:user-1:openai:model-a" not in redis.store
|
|
assert "cache_affinity:user-3:claude:model-c" not in redis.store
|
|
assert await manager._get_l1_entry("cache_affinity:user-1:openai:model-a") is None
|
|
assert await manager._get_l1_entry("cache_affinity:user-2:openai:model-b") == {
|
|
"provider_id": "provider-b"
|
|
}
|