163 lines
4.8 KiB
Python
163 lines
4.8 KiB
Python
# core/cache.py
|
||
import json as _stdlib_json
|
||
from contextlib import asynccontextmanager
|
||
from typing import Any, AsyncGenerator, Optional, Union
|
||
|
||
import logging
|
||
logger = logging.getLogger(__name__)
|
||
|
||
from redis.asyncio import Redis
|
||
|
||
# ===== 尝试使用你的 fff，否则回退到 stdlib =====
|
||
try:
|
||
import orjson as json
|
||
except ImportError:
|
||
json = _stdlib_json # type: ignore
|
||
|
||
|
||
# ===== Redis 客户端单例（可替换为你自己的）=====
|
||
class RedisClient:
|
||
_instance: Optional["RedisClient"] = None
|
||
_redis: Optional[Redis] = None
|
||
|
||
def __new__(cls):
|
||
if cls._instance is None:
|
||
cls._instance = super().__new__(cls)
|
||
return cls._instance
|
||
|
||
async def init_redis(self, url: str = "redis://localhost:6379/0") -> None:
|
||
if self._redis is None:
|
||
self._redis = Redis.from_url(url, decode_responses=False)
|
||
|
||
@property
|
||
def client(self) -> Redis:
|
||
if self._redis is None:
|
||
raise RuntimeError("Redis not initialized. Call init_redis() first.")
|
||
return self._redis
|
||
|
||
async def close(self) -> None:
|
||
if self._redis:
|
||
await self._redis.close()
|
||
self._redis = None
|
||
|
||
|
||
redis_client = RedisClient()
|
||
|
||
|
||
# ===== FastAPI 依赖注入 =====
|
||
async def get_redis() -> Redis:
|
||
"""FastAPI 依赖：获取 Redis 客户端"""
|
||
return redis_client.client
|
||
|
||
async def invalidate_cache(key: str) -> bool:
|
||
"""
|
||
主动删除缓存键。
|
||
返回是否成功删除（Redis delete 返回被删除的 key 数量）。
|
||
"""
|
||
try:
|
||
result = await redis_client.client.delete(key)
|
||
return result > 0
|
||
except Exception as e:
|
||
logger.exception(f"Failed to invalidate cache key: {key}", exc_info=e)
|
||
return False
|
||
|
||
# ===== 缓存上下文管理器 =====
|
||
class CacheResult:
|
||
__slots__ = ("key", "hit", "value", "_to_set", "_set_called")
|
||
|
||
def __init__(self, key: str) -> None:
|
||
self.key = key
|
||
self.hit = False
|
||
self.value: Any = None
|
||
self._to_set: Any = None
|
||
self._set_called = False
|
||
|
||
def set(self, value: Any) -> None:
|
||
"""标记要缓存的值（可为 None）"""
|
||
self._to_set = value
|
||
self._set_called = True
|
||
|
||
|
||
@asynccontextmanager
|
||
async def cache_if(
|
||
key: str,
|
||
ttl: int = 3600,
|
||
redis: Optional[Redis] = None,
|
||
) -> AsyncGenerator[CacheResult, None]:
|
||
"""
|
||
异步缓存上下文管理器，支持 None 值缓存（防穿透）。
|
||
|
||
Args:
|
||
key: 缓存键
|
||
ttl: 正常值缓存时间（秒）
|
||
redis: 可选 Redis 客户端（用于测试或自定义）
|
||
|
||
Usage:
|
||
async with cache_if("report:123") as cache:
|
||
if cache.hit:
|
||
return cache.value
|
||
result = await compute()
|
||
cache.set(result) # result 可为 None
|
||
"""
|
||
result = CacheResult(key)
|
||
client = redis or redis_client.client
|
||
|
||
# 尝试读缓存
|
||
try:
|
||
cached_val = await client.get(key)
|
||
if cached_val is not None:
|
||
# 解码
|
||
if cached_val == b"__NULL__":
|
||
result.hit = True
|
||
result.value = None
|
||
else:
|
||
result.hit = True
|
||
result.value = json.loads(cached_val)
|
||
except Exception:
|
||
# Redis 不可用，降级（不中断主流程）
|
||
logger.exception(f"Redis get error for key: {key}")
|
||
pass
|
||
|
||
yield result
|
||
|
||
# 写缓存（仅当调用了 set()）
|
||
if result._set_called:
|
||
try:
|
||
if result._to_set is None:
|
||
val = b"__NULL__"
|
||
ex = 60 # 空值短 TTL
|
||
else:
|
||
# 注意：orjson.dumps 返回 bytes，stdlib 返回 str → 统一转 bytes
|
||
serialized = json.dumps(result._to_set)
|
||
val = serialized if isinstance(serialized, bytes) else serialized.encode("utf-8")
|
||
ex = ttl
|
||
await client.setex(key, ex, val)
|
||
except Exception:
|
||
# 写缓存失败，不影响主流程
|
||
pass
|
||
|
||
|
||
# ===== 装饰器版（可选补充）=====
|
||
from functools import wraps
|
||
import asyncio
|
||
import hashlib
|
||
|
||
def cached(ttl: int = 3600):
|
||
"""函数缓存装饰器（使用 cache_if）"""
|
||
def decorator(func):
|
||
@wraps(func)
|
||
async def wrapper(*args, **kwargs):
|
||
# 生成 key（简单版，可替换为更 robust 的）
|
||
key_data = str(args) + str(sorted(kwargs.items()))
|
||
key = f"cached:{func.__name__}:{hashlib.md5(key_data.encode()).hexdigest()}"
|
||
logger.debug(f"Cache key: {key}")
|
||
|
||
async with cache_if(key, ttl) as cache:
|
||
if cache.hit:
|
||
logger.debug(f"Cache hit for key: {key}")
|
||
return cache.value
|
||
result = await func(*args, **kwargs)
|
||
cache.set(result)
|
||
return result
|
||
return wrapper
|
||
return decorator |