267 lines
7.0 KiB
Python
267 lines
7.0 KiB
Python
|
"""缓存模块
|
|||
|
|
|
|||
|
|
提供内存缓存和 TTL 管理
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import asyncio
|
|||
|
|
import hashlib
|
|||
|
|
import time
|
|||
|
|
from dataclasses import dataclass, field
|
|||
|
|
from typing import Any, Generic, TypeVar
|
|||
|
|
|
|||
|
|
from minenasai.core.logging import get_logger
|
|||
|
|
|
|||
|
|
logger = get_logger(__name__)
|
|||
|
|
|
|||
|
|
T = TypeVar("T")
|
|||
|
|
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class CacheEntry(Generic[T]):
|
|||
|
|
"""缓存条目"""
|
|||
|
|
|
|||
|
|
key: str
|
|||
|
|
value: T
|
|||
|
|
created_at: float
|
|||
|
|
expires_at: float
|
|||
|
|
hits: int = 0
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def is_expired(self) -> bool:
|
|||
|
|
"""是否过期"""
|
|||
|
|
return time.time() > self.expires_at
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def ttl_remaining(self) -> float:
|
|||
|
|
"""剩余 TTL（秒）"""
|
|||
|
|
return max(0, self.expires_at - time.time())
|
|||
|
|
|
|||
|
|
|
|||
|
|
class MemoryCache(Generic[T]):
|
|||
|
|
"""内存缓存
|
|||
|
|
|
|||
|
|
支持 TTL、最大容量、LRU 淘汰
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
max_size: int = 1000,
|
|||
|
|
default_ttl: float = 300.0, # 5分钟
|
|||
|
|
cleanup_interval: float = 60.0, # 1分钟清理一次
|
|||
|
|
) -> None:
|
|||
|
|
self._cache: dict[str, CacheEntry[T]] = {}
|
|||
|
|
self._max_size = max_size
|
|||
|
|
self._default_ttl = default_ttl
|
|||
|
|
self._cleanup_interval = cleanup_interval
|
|||
|
|
self._cleanup_task: asyncio.Task[None] | None = None
|
|||
|
|
|
|||
|
|
# 统计
|
|||
|
|
self._hits = 0
|
|||
|
|
self._misses = 0
|
|||
|
|
|
|||
|
|
async def start(self) -> None:
|
|||
|
|
"""启动后台清理任务"""
|
|||
|
|
if self._cleanup_task is None:
|
|||
|
|
self._cleanup_task = asyncio.create_task(self._cleanup_loop())
|
|||
|
|
logger.info("cache_cleanup_started", interval=self._cleanup_interval)
|
|||
|
|
|
|||
|
|
async def stop(self) -> None:
|
|||
|
|
"""停止后台清理任务"""
|
|||
|
|
if self._cleanup_task:
|
|||
|
|
self._cleanup_task.cancel()
|
|||
|
|
try:
|
|||
|
|
await self._cleanup_task
|
|||
|
|
except asyncio.CancelledError:
|
|||
|
|
pass
|
|||
|
|
self._cleanup_task = None
|
|||
|
|
|
|||
|
|
async def _cleanup_loop(self) -> None:
|
|||
|
|
"""清理循环"""
|
|||
|
|
while True:
|
|||
|
|
await asyncio.sleep(self._cleanup_interval)
|
|||
|
|
self._cleanup_expired()
|
|||
|
|
|
|||
|
|
def _cleanup_expired(self) -> int:
|
|||
|
|
"""清理过期条目"""
|
|||
|
|
expired_keys = [k for k, v in self._cache.items() if v.is_expired]
|
|||
|
|
for key in expired_keys:
|
|||
|
|
del self._cache[key]
|
|||
|
|
|
|||
|
|
if expired_keys:
|
|||
|
|
logger.debug("cache_cleanup", removed=len(expired_keys))
|
|||
|
|
|
|||
|
|
return len(expired_keys)
|
|||
|
|
|
|||
|
|
def _evict_lru(self) -> None:
|
|||
|
|
"""LRU 淘汰"""
|
|||
|
|
if len(self._cache) < self._max_size:
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# 按命中次数和创建时间排序，淘汰最少使用的
|
|||
|
|
sorted_entries = sorted(
|
|||
|
|
self._cache.items(),
|
|||
|
|
key=lambda x: (x[1].hits, x[1].created_at),
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 淘汰 10% 的条目
|
|||
|
|
evict_count = max(1, len(sorted_entries) // 10)
|
|||
|
|
for key, _ in sorted_entries[:evict_count]:
|
|||
|
|
del self._cache[key]
|
|||
|
|
|
|||
|
|
logger.debug("cache_eviction", evicted=evict_count)
|
|||
|
|
|
|||
|
|
def get(self, key: str) -> T | None:
|
|||
|
|
"""获取缓存值"""
|
|||
|
|
entry = self._cache.get(key)
|
|||
|
|
|
|||
|
|
if entry is None:
|
|||
|
|
self._misses += 1
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
if entry.is_expired:
|
|||
|
|
del self._cache[key]
|
|||
|
|
self._misses += 1
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
entry.hits += 1
|
|||
|
|
self._hits += 1
|
|||
|
|
return entry.value
|
|||
|
|
|
|||
|
|
def set(self, key: str, value: T, ttl: float | None = None) -> None:
|
|||
|
|
"""设置缓存值"""
|
|||
|
|
if ttl is None:
|
|||
|
|
ttl = self._default_ttl
|
|||
|
|
|
|||
|
|
# 检查容量
|
|||
|
|
if len(self._cache) >= self._max_size:
|
|||
|
|
self._evict_lru()
|
|||
|
|
|
|||
|
|
now = time.time()
|
|||
|
|
self._cache[key] = CacheEntry(
|
|||
|
|
key=key,
|
|||
|
|
value=value,
|
|||
|
|
created_at=now,
|
|||
|
|
expires_at=now + ttl,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def delete(self, key: str) -> bool:
|
|||
|
|
"""删除缓存"""
|
|||
|
|
if key in self._cache:
|
|||
|
|
del self._cache[key]
|
|||
|
|
return True
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
def clear(self) -> int:
|
|||
|
|
"""清空缓存"""
|
|||
|
|
count = len(self._cache)
|
|||
|
|
self._cache.clear()
|
|||
|
|
return count
|
|||
|
|
|
|||
|
|
def exists(self, key: str) -> bool:
|
|||
|
|
"""检查 key 是否存在且未过期"""
|
|||
|
|
entry = self._cache.get(key)
|
|||
|
|
if entry is None:
|
|||
|
|
return False
|
|||
|
|
if entry.is_expired:
|
|||
|
|
del self._cache[key]
|
|||
|
|
return False
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
def get_stats(self) -> dict[str, Any]:
|
|||
|
|
"""获取统计信息"""
|
|||
|
|
total = self._hits + self._misses
|
|||
|
|
hit_rate = self._hits / total if total > 0 else 0.0
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
"size": len(self._cache),
|
|||
|
|
"max_size": self._max_size,
|
|||
|
|
"hits": self._hits,
|
|||
|
|
"misses": self._misses,
|
|||
|
|
"hit_rate": round(hit_rate * 100, 2),
|
|||
|
|
"default_ttl": self._default_ttl,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def make_cache_key(*args: Any, **kwargs: Any) -> str:
|
|||
|
|
"""生成缓存 key"""
|
|||
|
|
key_parts = [str(arg) for arg in args]
|
|||
|
|
key_parts.extend(f"{k}={v}" for k, v in sorted(kwargs.items()))
|
|||
|
|
key_string = ":".join(key_parts)
|
|||
|
|
return hashlib.md5(key_string.encode()).hexdigest()
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 全局缓存实例
|
|||
|
|
_response_cache: MemoryCache[dict[str, Any]] | None = None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_response_cache() -> MemoryCache[dict[str, Any]]:
|
|||
|
|
"""获取响应缓存"""
|
|||
|
|
global _response_cache
|
|||
|
|
if _response_cache is None:
|
|||
|
|
_response_cache = MemoryCache(
|
|||
|
|
max_size=500,
|
|||
|
|
default_ttl=300.0, # 5分钟
|
|||
|
|
)
|
|||
|
|
return _response_cache
|
|||
|
|
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class RateLimiter:
|
|||
|
|
"""速率限制器
|
|||
|
|
|
|||
|
|
令牌桶算法实现
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
rate: float # 每秒允许的请求数
|
|||
|
|
burst: int # 突发容量
|
|||
|
|
|
|||
|
|
_tokens: float = field(init=False)
|
|||
|
|
_last_update: float = field(init=False)
|
|||
|
|
|
|||
|
|
def __post_init__(self) -> None:
|
|||
|
|
self._tokens = float(self.burst)
|
|||
|
|
self._last_update = time.time()
|
|||
|
|
|
|||
|
|
def _refill(self) -> None:
|
|||
|
|
"""补充令牌"""
|
|||
|
|
now = time.time()
|
|||
|
|
elapsed = now - self._last_update
|
|||
|
|
self._tokens = min(self.burst, self._tokens + elapsed * self.rate)
|
|||
|
|
self._last_update = now
|
|||
|
|
|
|||
|
|
def acquire(self, tokens: int = 1) -> bool:
|
|||
|
|
"""尝试获取令牌"""
|
|||
|
|
self._refill()
|
|||
|
|
|
|||
|
|
if self._tokens >= tokens:
|
|||
|
|
self._tokens -= tokens
|
|||
|
|
return True
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
async def wait(self, tokens: int = 1) -> None:
|
|||
|
|
"""等待获取令牌"""
|
|||
|
|
while not self.acquire(tokens):
|
|||
|
|
# 计算需要等待的时间
|
|||
|
|
needed = tokens - self._tokens
|
|||
|
|
wait_time = needed / self.rate
|
|||
|
|
await asyncio.sleep(min(wait_time, 0.1))
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def available_tokens(self) -> float:
|
|||
|
|
"""可用令牌数"""
|
|||
|
|
self._refill()
|
|||
|
|
return self._tokens
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 全局速率限制器
|
|||
|
|
_rate_limiters: dict[str, RateLimiter] = {}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_rate_limiter(name: str, rate: float = 10.0, burst: int = 20) -> RateLimiter:
|
|||
|
|
"""获取或创建速率限制器"""
|
|||
|
|
if name not in _rate_limiters:
|
|||
|
|
_rate_limiters[name] = RateLimiter(rate=rate, burst=burst)
|
|||
|
|
return _rate_limiters[name]
|