mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
Merge branch 'main' of https://github.com/dataease/SQLBot
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -398,6 +398,7 @@ defineExpose({
|
||||
<el-scrollbar v-loading="copyLoading" class="custom-tree">
|
||||
<el-tree
|
||||
ref="resourceListTree"
|
||||
style="overflow-x: hidden"
|
||||
menu
|
||||
:empty-text="t('dashboard.no_dashboard')"
|
||||
:default-expanded-keys="expandedArray"
|
||||
|
||||
Reference in New Issue
Block a user