perf: Enable user cache

backend/common/core/sqlbot_cache.py

@@ -1,35 +1,50 @@
+import contextvars
 from fastapi_cache import FastAPICache
-from fastapi_cache.decorator import cache as original_cache
 from functools import partial, wraps
-from typing import Optional, Set, Any, Dict, Tuple
-from inspect import Parameter, signature
+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.utils.utils import SQLBotLogUtil
 
-
-# 锁管理
-_cache_locks = {}
+# 使用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):
-    lock = _cache_locks.setdefault(key, asyncio.Lock())
-    async with lock:
-        try:
-            yield
-        finally:
-            if key in _cache_locks and not lock.locked():
-                del _cache_locks[key]
-
-def should_skip_param(param: Parameter) -> bool:
-    """判断参数是否应该被忽略(依赖注入参数)"""
-    return (
-        param.kind == Parameter.VAR_KEYWORD or  # **kwargs
-        param.kind == Parameter.VAR_POSITIONAL or  # *args
-        hasattr(param.annotation, "__module__") and 
-        param.annotation.__module__.startswith(('fastapi', 'starlette', "sqlmodel.orm.session"))
-    )
+    # 获取当前已持有的锁集合
+    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()
 
 def custom_key_builder(
     func: Any,
@@ -37,137 +52,116 @@ def custom_key_builder(
     *,
     args: Tuple[Any, ...] = (),
     kwargs: Dict[str, Any],
-    additional_skip_args: Optional[Set[str]] = None,
-    cacheName: Optional[str] = None,
+    cacheName: str,
     keyExpression: Optional[str] = None,
 ) -> str:
-    """完全兼容FastAPICache的键生成器"""
-    if cacheName: 
+    try:
         base_key = f"{namespace}:{cacheName}:"
         
         if keyExpression:
-            try:
-                sig = signature(func)
-                bound_args = sig.bind_partial(*args, **kwargs)
-                bound_args.apply_defaults()
-                
-                if keyExpression.startswith("args["):
-                    import re
-                    match = re.match(r"args\[(\d+)\]", keyExpression)
-                    if match:
-                        index = int(match.group(1))
-                        value = bound_args.args[index]
-                        base_key += f"{value}:"
-                else:
-                    parts = keyExpression.split('.')
-                    value = bound_args.arguments[parts[0]]
-                    for part in parts[1:]:
-                        value = getattr(value, part)
-                    base_key += f"{value}:"
-                
-            except (IndexError, KeyError, AttributeError) as e:
-                SQLBotLogUtil.warning(f"Failed to evaluate keyExpression '{keyExpression}': {str(e)}")
+            sig = signature(func)
+            bound_args = sig.bind_partial(*args, **kwargs)
+            bound_args.apply_defaults()
+            
+            # 支持args[0]格式
+            if keyExpression.startswith("args["):
+                import re
+                if match := re.match(r"args\[(\d+)\]", keyExpression):
+                    index = int(match.group(1))
+                    value = bound_args.args[index]
+                    return f"{base_key}{value}"
+            
+            # 支持属性路径格式
+            parts = keyExpression.split('.')
+            value = bound_args.arguments[parts[0]]
+            for part in parts[1:]:
+                value = getattr(value, part)
+            return f"{base_key}{value}"
         
-        return base_key
-    
-    sig = signature(func)
-    auto_skip_args = {
-        name for name, param in sig.parameters.items()
-        if should_skip_param(param)
-    }
-    skip_args = auto_skip_args.union(additional_skip_args or set())
-    filtered_kwargs = {
-        k: v for k, v in kwargs.items() if k not in skip_args
-    }
-    
-    bound_args = sig.bind_partial(*args, **kwargs)
-    bound_args.apply_defaults()
-    
-    filtered_args = []
-    for i, (name, value) in enumerate(bound_args.arguments.items()):
-        if i < len(args) and name not in skip_args:
-            filtered_args.append(value)
-    filtered_args = tuple(filtered_args)
-    
-    default_key_builder = FastAPICache.get_key_builder()
-    return default_key_builder(
-        func=func,
-        namespace=namespace,
-        args=filtered_args,
-        kwargs=filtered_kwargs,
-    )
+        # 默认使用第一个参数作为key
+        return f"{base_key}{args[0] if args else 'default'}"
+        
+    except Exception as e:
+        SQLBotLogUtil.error(f"Key builder error: {str(e)}")
+        raise ValueError(f"Invalid cache key generation: {e}") from e
 
 def cache(
-    expire: Optional[int] = 60 * 60 * 24,
-    namespace: Optional[str] = None,
-    key_builder: Optional[Any] = None,
+    expire: int = 60 * 60 * 24,
+    namespace: str = "",
     *,
-    additional_skip_args: Optional[Set[str]] = None,
-    cacheName: Optional[str] = None,
+    cacheName: str,  # 必须提供cacheName
     keyExpression: Optional[str] = None,
+    jitter: int = 60,  # 默认抖动60秒
 ):
-    """完全兼容的缓存装饰器"""
     def decorator(func):
-        if key_builder is None:
-            used_key_builder = partial(
-                custom_key_builder,
-                additional_skip_args=additional_skip_args,
-                cacheName=cacheName,
-                keyExpression=keyExpression
-            )
-        else:
-            used_key_builder = key_builder
-            
+        # 预先生成key builder
+        used_key_builder = partial(
+            custom_key_builder,
+            cacheName=cacheName,
+            keyExpression=keyExpression
+        )
+        
         @wraps(func)
         async def wrapper(*args, **kwargs):
+            # 生成缓存键
             cache_key = used_key_builder(
                 func=func,
-                namespace=namespace or "",
+                namespace=namespace,
                 args=args,
                 kwargs=kwargs
             )
             
+            # 防击穿锁
             async with _get_cache_lock(cache_key):
-                SQLBotLogUtil.info(f"Using cache key: {cache_key}")
                 backend = FastAPICache.get_backend()
-                cached_value = await backend.get(cache_key)
-                if cached_value is not None:
-                    SQLBotLogUtil.info(f"Cache hit for key: {cache_key}, the value is: {cached_value}")
-                    return cached_value
                 
+                # 双重检查
+                if (cached := await backend.get(cache_key)) is not None:
+                    SQLBotLogUtil.debug(f"Cache hit: {cache_key}")
+                    return cached
+                
+                # 执行函数并缓存结果
                 result = await func(*args, **kwargs)
-                await backend.set(cache_key, result, expire)
-                SQLBotLogUtil.info(f"Cache miss for key: {cache_key}, result cached.")
+                
+                actual_expire = expire + random.randint(-jitter, jitter)
+                await backend.set(cache_key, result, actual_expire)
+                
+                SQLBotLogUtil.debug(f"Cache set: {cache_key} (expire: {actual_expire}s)")
                 return result
                 
         return wrapper
     return decorator
 
 def clear_cache(
-    namespace: Optional[str] = None,
-    cacheName: Optional[str] = None,
+    namespace: str = "",
+    *,
+    cacheName: str,
     keyExpression: Optional[str] = None,
 ):
-    """清除缓存的装饰器"""
+    """精确清除单个缓存项的装饰器"""
     def decorator(func):
         @wraps(func)
         async def wrapper(*args, **kwargs):
             cache_key = custom_key_builder(
                 func=func,
-                namespace=namespace or "",
+                namespace=namespace,
                 args=args,
                 kwargs=kwargs,
                 cacheName=cacheName,
                 keyExpression=keyExpression,
             )
             
+            # 加锁防止竞争
             async with _get_cache_lock(cache_key):
-                if await FastAPICache.get_backend().get(cache_key):
-                    await FastAPICache.clear(key=cache_key)
+                backend = FastAPICache.get_backend()
+                result = None
+                if await backend.get(cache_key):
+                    await backend.clear(cache_key)
                     result = await func(*args, **kwargs)
-                    SQLBotLogUtil.info(f"Clearing cache for key: {cache_key}")
-                    return result
-                return await func(*args, **kwargs)
+                    if await backend.get(cache_key):
+                        await backend.clear(cache_key)
+                    SQLBotLogUtil.info(f"Cache cleared: {cache_key}")
+                return result
         
         return wrapper
     return decorator