diff --git a/backend/common/core/sqlbot_cache.py b/backend/common/core/sqlbot_cache.py index 02cd2d82..cce07c06 100644 --- a/backend/common/core/sqlbot_cache.py +++ b/backend/common/core/sqlbot_cache.py @@ -1,50 +1,12 @@ -import contextvars from fastapi_cache import FastAPICache from functools import partial, wraps from typing import Optional, Any, Dict, Tuple from inspect import signature -from contextlib import asynccontextmanager -import asyncio -import random -from collections import defaultdict from common.core.config import settings from common.utils.utils import SQLBotLogUtil from fastapi_cache.backends.inmemory import InMemoryBackend -# 使用contextvar来跟踪当前线程已持有的锁 -_held_locks = contextvars.ContextVar('held_locks', default=set()) -# 高效锁管理器 -class LockManager: - _locks = defaultdict(asyncio.Lock) - - @classmethod - def get_lock(cls, key: str) -> asyncio.Lock: - return cls._locks[key] -@asynccontextmanager -async def _get_cache_lock(key: str): - # 获取当前已持有的锁集合 - current_locks = _held_locks.get() - - # 如果已经持有这个锁,直接yield(锁传递) - if key in current_locks: - yield - return - - # 否则获取锁并添加到当前上下文中 - lock = LockManager.get_lock(key) - try: - await lock.acquire() - # 更新当前持有的锁集合 - new_locks = current_locks | {key} - token = _held_locks.set(new_locks) - - yield - - finally: - # 恢复之前的锁集合 - _held_locks.reset(token) - if lock.locked(): - lock.release() +from fastapi_cache.decorator import cache as original_cache def custom_key_builder( func: Any, @@ -91,7 +53,6 @@ def cache( *, cacheName: str, # 必须提供cacheName keyExpression: Optional[str] = None, - jitter: int = 60, # 默认抖动60秒 ): def decorator(func): # 预先生成key builder @@ -108,30 +69,17 @@ def cache( # 生成缓存键 cache_key = used_key_builder( func=func, - namespace=namespace, + namespace=str(namespace) if namespace else "", args=args, kwargs=kwargs ) - # 防击穿锁 - async with _get_cache_lock(cache_key): - backend = FastAPICache.get_backend() - # 双重检查 - if (cached := await backend.get(cache_key)) is not None: - SQLBotLogUtil.debug(f"Cache hit: {cache_key}") - return cached - - # 执行函数并缓存结果 - result = await func(*args, **kwargs) - - actual_expire = expire + random.randint(-jitter, jitter) - if await backend.get(cache_key): - await backend.clear(cache_key) - await backend.set(cache_key, result, actual_expire) - - SQLBotLogUtil.debug(f"Cache set: {cache_key} (expire: {actual_expire}s)") - return result - + return await original_cache( + expire=expire, + namespace=str(namespace) if namespace else "", + key_builder=lambda *_, **__: cache_key + )(func)(*args, **kwargs) + return wrapper return decorator @@ -141,7 +89,6 @@ def clear_cache( cacheName: str, keyExpression: Optional[str] = None, ): - """精确清除单个缓存项的装饰器""" def decorator(func): @wraps(func) async def wrapper(*args, **kwargs): @@ -149,24 +96,21 @@ def clear_cache( return await func(*args, **kwargs) cache_key = custom_key_builder( func=func, - namespace=namespace, + namespace=str(namespace) if namespace else "", args=args, kwargs=kwargs, cacheName=cacheName, keyExpression=keyExpression, ) - - # 加锁防止竞争 - async with _get_cache_lock(cache_key): - backend = FastAPICache.get_backend() - result = None - if await backend.get(cache_key): - await backend.clear(cache_key) - result = await func(*args, **kwargs) - if await backend.get(cache_key): - await backend.clear(cache_key) - SQLBotLogUtil.debug(f"Cache cleared: {cache_key}") - return result + backend = FastAPICache.get_backend() + if await backend.get(cache_key): + if settings.CACHE_TYPE.lower() == "redis": + redis = backend.redis + await redis.delete(cache_key) + else: + await backend.clear(key=cache_key) + SQLBotLogUtil.debug(f"Cache cleared: {cache_key}") + return await func(*args, **kwargs) return wrapper return decorator