Diuse/zhenxun_bot
1
0
Fork 0
mirror of https://github.com/zhenxun-org/zhenxun_bot.git synced 2026-10-22 04:30:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

bugfix:修复sqlite部分场景下锁竞争问题；新增插件恶意触发配置文件 (#2144)

Browse Source 
* bugfix:修复sqlite部分场景下锁竞争问题；新增插件恶意触发配置文件

* 文件没同步完
This commit is contained in:
Copaan
2026-06-24 09:11:03 +08:00
committed by GitHub
parent 73cbe2a609
commit f4d2342693
39 changed files with 1577 additions and 829 deletions
Download Patch File Download Diff File Expand all files Collapse all files
zhenxun/services/buffered_writers.py
+32 -93
View File
@@ -1,61 +1,49 @@
from __future__ import annotations
import asyncio
from collections import deque
import contextlib
import time
from zhenxun.models.user_gold_log import UserGoldLog
from zhenxun.services.log import logger
from zhenxun.services.low_priority_writer import (
LowPriorityWriterConfig,
append_low_priority_record,
flush_low_priority_writer,
register_low_priority_writer,
)
from zhenxun.utils.enum import GoldHandle
from zhenxun.utils.manager.priority_manager import PriorityLifecycle
LOG_COMMAND = "BufferedWriters"
_WRITER_NAME = "user_gold_log"
_USER_GOLD_LOG_BUFFER_MAX_RETAIN = 10_000
_USER_GOLD_LOG_FLUSH_TRIGGER_SIZE = 128
_USER_GOLD_LOG_FLUSH_BATCH_SIZE = 500
_USER_GOLD_LOG_FLUSH_INTERVAL_SECONDS = 60.0
_USER_GOLD_LOG_DROP_LOG_INTERVAL_SECONDS = 10.0
_user_gold_log_buffer: deque[UserGoldLog] = deque()
_user_gold_log_buffer_lock = asyncio.Lock()
_user_gold_log_flush_lock = asyncio.Lock()
_user_gold_log_flush_task: asyncio.Task[None] | None = None
_user_gold_log_dropped = 0
_user_gold_log_last_drop_log_at = 0.0
def _ensure_user_gold_log_flush_task() -> None:
global _user_gold_log_flush_task
if _user_gold_log_flush_task is not None and not _user_gold_log_flush_task.done():
return
_user_gold_log_flush_task = asyncio.create_task(_user_gold_log_flush_loop())
async def _write_user_gold_log_batch(
batch: list[UserGoldLog],
reason: str,
) -> None:
from zhenxun.services.db_context import with_db_timeout
def _record_user_gold_log_drop() -> None:
global _user_gold_log_dropped, _user_gold_log_last_drop_log_at
_user_gold_log_dropped += 1
now = time.monotonic()
if now - _user_gold_log_last_drop_log_at < _USER_GOLD_LOG_DROP_LOG_INTERVAL_SECONDS:
return
_user_gold_log_last_drop_log_at = now
logger.warning(
"user_gold_log buffer full, dropped "
f"{_user_gold_log_dropped} records, backlog={len(_user_gold_log_buffer)}",
LOG_COMMAND,
await with_db_timeout(
UserGoldLog.bulk_create(batch, _USER_GOLD_LOG_FLUSH_BATCH_SIZE),
timeout=5.0,
operation=f"UserGoldLog.bulk_create[{len(batch)}]",
source=f"user_gold_log:{reason}",
)
async def _user_gold_log_flush_loop() -> None:
while True:
await asyncio.sleep(_USER_GOLD_LOG_FLUSH_INTERVAL_SECONDS)
try:
await flush_user_gold_log_buffer("定时")
except asyncio.CancelledError:
raise
except Exception as exc:
logger.warning("定时批量写入金币流水失败", LOG_COMMAND, e=exc)
register_config = LowPriorityWriterConfig(
name=_WRITER_NAME,
write_batch=_write_user_gold_log_batch,
batch_size=_USER_GOLD_LOG_FLUSH_BATCH_SIZE,
trigger_size=_USER_GOLD_LOG_FLUSH_TRIGGER_SIZE,
max_retain=_USER_GOLD_LOG_BUFFER_MAX_RETAIN,
flush_interval_seconds=_USER_GOLD_LOG_FLUSH_INTERVAL_SECONDS,
max_items_per_cycle=_USER_GOLD_LOG_FLUSH_BATCH_SIZE,
backoff_base_seconds=30.0,
backoff_max_seconds=600.0,
log_command=LOG_COMMAND,
)
async def append_user_gold_log(
@@ -64,65 +52,16 @@ async def append_user_gold_log(
handle: GoldHandle,
source: str | None = None,
) -> None:
_ensure_user_gold_log_flush_task()
record = UserGoldLog(user_id=user_id, gold=gold, handle=handle, source=source)
async with _user_gold_log_buffer_lock:
if len(_user_gold_log_buffer) >= _USER_GOLD_LOG_BUFFER_MAX_RETAIN:
_user_gold_log_buffer.popleft()
_record_user_gold_log_drop()
_user_gold_log_buffer.append(record)
should_flush = (
len(_user_gold_log_buffer) >= _USER_GOLD_LOG_FLUSH_TRIGGER_SIZE
and not _user_gold_log_flush_lock.locked()
)
if should_flush:
await flush_user_gold_log_buffer("缓冲区触发")
await append_low_priority_record(_WRITER_NAME, record)
async def flush_user_gold_log_buffer(reason: str) -> int:
async with _user_gold_log_flush_lock:
written = 0
while True:
batch: list[UserGoldLog] = []
async with _user_gold_log_buffer_lock:
if not _user_gold_log_buffer:
break
while (
_user_gold_log_buffer
and len(batch) < _USER_GOLD_LOG_FLUSH_BATCH_SIZE
):
batch.append(_user_gold_log_buffer.popleft())
if not batch:
break
try:
await UserGoldLog.bulk_create(batch, _USER_GOLD_LOG_FLUSH_BATCH_SIZE)
except Exception as exc:
async with _user_gold_log_buffer_lock:
retain_count = max(
_USER_GOLD_LOG_BUFFER_MAX_RETAIN - len(_user_gold_log_buffer),
0,
)
for record in reversed(batch[-retain_count:]):
_user_gold_log_buffer.appendleft(record)
logger.error(f"{reason}批量写入金币流水失败", LOG_COMMAND, e=exc)
return written
written += len(batch)
if written:
logger.debug(f"{reason}批量写入金币流水 {written} 条", LOG_COMMAND)
return written
return await flush_low_priority_writer(_WRITER_NAME, reason, force=True)
async def stop_user_gold_log_buffer() -> int:
global _user_gold_log_flush_task
task = _user_gold_log_flush_task
_user_gold_log_flush_task = None
if task is not None:
task.cancel()
with contextlib.suppress(BaseException):
await task
return await flush_user_gold_log_buffer("关闭")
@PriorityLifecycle.on_shutdown(priority=90)
async def _flush_user_gold_log_buffer_on_shutdown() -> None:
await stop_user_gold_log_buffer()
register_low_priority_writer(register_config)
zhenxun/services/cache/__init__.py
Vendored
+9 -21
View File
@@ -52,7 +52,6 @@ from typing import Any, ClassVar, Generic, TypeVar, cast, get_type_hints
from typing_extensions import Self
from aiocache import Cache as AioCache
from aiocache import SimpleMemoryCache
from aiocache.base import BaseCache
from aiocache.serializers import JsonSerializer
import nonebot
@@ -109,6 +108,10 @@ driver = nonebot.get_driver()
cache_config = nonebot.get_plugin_config(Config)
def _redis_cache_enabled() -> bool:
return cache_config.cache_mode == CacheMode.REDIS and bool(cache_config.redis_host)
class CacheException(Exception):
"""缓存相关异常"""
@@ -190,15 +193,10 @@ class CacheManager:
def cache_backend(self) -> BaseCache | AioCache:
"""获取缓存后端"""
if self._cache_backend is None:
ttl = cache_config.redis_expire
if cache_config.cache_mode == CacheMode.NONE:
ttl = 0
logger.info("缓存功能已禁用，使用非持久化内存缓存", LOG_COMMAND)
elif cache_config.cache_mode == CacheMode.REDIS and cache_config.redis_host:
if _redis_cache_enabled():
try:
from aiocache import RedisCache
# 使用Redis缓存
self._cache_backend = RedisCache(
serializer=JsonSerializer(),
namespace=CACHE_KEY_PREFIX,
@@ -215,19 +213,11 @@ class CacheManager:
return self._cache_backend
except ImportError as e:
logger.error(
"导入aiocache[redis]失败，将默认使用内存缓存...",
"导入aiocache[redis]失败，CacheRoot 模型缓存已禁用",
LOG_COMMAND,
e=e,
)
else:
logger.info("使用内存缓存", LOG_COMMAND)
# 默认使用内存缓存
self._cache_backend = SimpleMemoryCache(
serializer=JsonSerializer(),
namespace=CACHE_KEY_PREFIX,
timeout=30,
ttl=ttl,
)
raise CacheException("CacheRoot 后端未启用")
return self._cache_backend
async def invalidate_cache(
@@ -815,11 +805,9 @@ class Cache(Generic[T]):
@driver.on_startup
async def _():
CacheRoot.enabled = cache_config.cache_mode != CacheMode.NONE
CacheRoot.enabled = _redis_cache_enabled()
if CacheRoot.enabled:
logger.info("缓存系统已启用", LOG_COMMAND)
else:
logger.info("缓存系统已禁用", LOG_COMMAND)
logger.info("CacheRoot Redis 模型缓存已启用", LOG_COMMAND)
@driver.on_shutdown
zhenxun/services/cache/cache_containers.py
Vendored
-282
View File
@@ -248,285 +248,3 @@ class CacheDict(Generic[T]):
# 清理过期的键
self._clean_expired()
return f"CacheDict({self.name}, {len(self._data)} items)"
class CacheList(Generic[T]):
"""缓存列表类，提供类似普通列表的接口，数据只存储在内存中"""
_instances: weakref.WeakSet = weakref.WeakSet()
def __init__(self, name: str, expire: int = 0, max_items: int | None = None):
"""初始化缓存列表
参数:
name: 列表名称
expire: 过期时间（秒），默认为0表示永不过期
max_items: 最大缓存项数，None 使用统一默认值，0 表示不限制
"""
self.name = name.upper()
self.expire = expire
self.max_items = DEFAULT_CACHE_MAX_ITEMS if max_items is None else max_items
self._data: list[CacheData[T]] = []
self._expire_time = 0
self.__class__._instances.add(self)
# 如果设置了过期时间，计算整个列表的过期时间
if self.expire > 0:
self._expire_time = time.time() + self.expire
def __getitem__(self, index: int) -> T:
"""获取列表项
参数:
index: 列表索引
返回:
T: 列表值
"""
# 检查整个列表是否过期
if self._is_expired():
self.clear()
raise IndexError(f"列表索引 {index} 超出范围")
if 0 <= index < len(self._data):
return self._data[index].value
raise IndexError(f"列表索引 {index} 超出范围")
def __setitem__(self, index: int, value: T):
"""设置列表项
参数:
index: 列表索引
value: 列表值
"""
# 检查整个列表是否过期
if self._is_expired():
self.clear()
# 确保索引有效
while len(self._data) <= index:
raise IndexError(f"列表索引 {index} 超出范围")
self._data[index] = CacheData(value=value)
# 更新过期时间
self._update_expire_time()
def __delitem__(self, index: int):
"""删除列表项
参数:
index: 列表索引
"""
# 检查整个列表是否过期
if self._is_expired():
self.clear()
raise IndexError(f"列表索引 {index} 超出范围")
if not 0 <= index < len(self._data):
raise IndexError(f"列表索引 {index} 超出范围")
del self._data[index]
# 更新过期时间
self._update_expire_time()
def __len__(self) -> int:
"""获取列表长度
返回:
int: 列表长度
"""
# 检查整个列表是否过期
if self._is_expired():
self.clear()
return len(self._data)
def append(self, value: T):
"""添加列表项
参数:
value: 列表值
"""
# 检查整个列表是否过期
if self._is_expired():
self.clear()
self._data.append(CacheData(value=value))
self._enforce_limit()
# 更新过期时间
self._update_expire_time()
def extend(self, values: list[T]):
"""扩展列表
参数:
values: 要添加的值列表
"""
# 检查整个列表是否过期
if self._is_expired():
self.clear()
self._data.extend([CacheData(value=v) for v in values])
self._enforce_limit()
# 更新过期时间
self._update_expire_time()
def insert(self, index: int, value: T):
"""插入列表项
参数:
index: 插入位置
value: 列表值
"""
# 检查整个列表是否过期
if self._is_expired():
self.clear()
self._data.insert(index, CacheData(value=value))
self._enforce_limit()
# 更新过期时间
self._update_expire_time()
def pop(self, index: int = -1) -> T:
"""删除并返回列表项
参数:
index: 列表索引，默认为最后一项
返回:
Any: 列表值
"""
# 检查整个列表是否过期
if self._is_expired():
self.clear()
raise IndexError("从空列表中弹出")
if not self._data:
raise IndexError("从空列表中弹出")
item = self._data.pop(index)
# 更新过期时间
self._update_expire_time()
return item.value
def remove(self, value: T):
"""删除第一个匹配的列表项
参数:
value: 要删除的值
"""
# 检查整个列表是否过期
if self._is_expired():
self.clear()
raise ValueError(f"{value} 不在列表中")
# 查找匹配的项
for i, item in enumerate(self._data):
if item.value == value:
del self._data[i]
# 更新过期时间
self._update_expire_time()
return
raise ValueError(f"{value} 不在列表中")
def clear(self) -> None:
"""清空列表"""
self._data.clear()
# 重置过期时间
self._update_expire_time()
def stats(self) -> dict[str, int]:
"""返回当前缓存条目统计。"""
if self._is_expired():
self.clear()
return {"items": len(self._data), "max_items": self.max_items}
@classmethod
def stats_all(cls) -> dict[str, dict[str, int]]:
"""返回所有 CacheList 实例的条目统计。"""
result: dict[str, dict[str, int]] = {}
for cache in list(cls._instances):
stats = cache.stats()
if stats["items"]:
result[cache.name] = stats
return result
@classmethod
def clear_all(cls) -> dict[str, int]:
result: dict[str, int] = {}
for cache in list(cls._instances):
size = len(cache._data)
if size:
cache.clear()
result[cache.name] = result.get(cache.name, 0) + size
return result
def index(self, value: T, start: int = 0, end: int | None = None) -> int:
"""查找值的索引
参数:
value: 要查找的值
start: 起始索引
end: 结束索引
返回:
int: 索引位置
"""
# 检查整个列表是否过期
if self._is_expired():
self.clear()
raise ValueError(f"{value} 不在列表中")
end = end if end is not None else len(self._data)
for i in range(start, min(end, len(self._data))):
if self._data[i].value == value:
return i
raise ValueError(f"{value} 不在列表中")
def count(self, value: T) -> int:
"""计算值出现的次数
参数:
value: 要计数的值
返回:
int: 出现次数
"""
# 检查整个列表是否过期
if self._is_expired():
self.clear()
return 0
# sourcery skip: simplify-constant-sum
return sum(1 for item in self._data if item.value == value)
def _is_expired(self) -> bool:
"""检查整个列表是否过期"""
return self._expire_time > 0 and self._expire_time < time.time()
def _update_expire_time(self):
"""更新过期时间"""
self._expire_time = time.time() + self.expire if self.expire > 0 else 0
def _enforce_limit(self) -> None:
if self.max_items <= 0:
return
if len(self._data) > self.max_items:
del self._data[: len(self._data) - self.max_items]
def __str__(self) -> str:
"""字符串表示
返回:
str: 字符串表示
"""
# 检查整个列表是否过期
if self._is_expired():
self.clear()
return f"CacheList({self.name}, {len(self._data)} items)"
zhenxun/services/cache/runtime_cache.py
Vendored
+78 -8
View File
@@ -11,6 +11,7 @@ import uuid
from zhenxun.services.cache.config import CacheMode
from zhenxun.services.log import logger
from zhenxun.services.message_load import is_db_unhealthy
from zhenxun.utils.enum import (
BlockType,
LimitCheckType,
@@ -57,6 +58,7 @@ RUNTIME_CACHE_SYNC_ENABLED = True
RUNTIME_CACHE_SYNC_CHANNEL = "ZHENXUN_RUNTIME_CACHE_SYNC"
RUNTIME_CACHE_LOAD_RETRY_SECONDS = 1.0
RUNTIME_CACHE_STARTUP_REFRESH_SKIP_SECONDS = 5.0
RUNTIME_CACHE_DB_TIMEOUT_SECONDS = 3.0
INSTANCE_ID = uuid.uuid4().hex
@@ -749,6 +751,9 @@ class RuntimeCacheMutation:
@classmethod
async def ensure_loaded(cls, cache_cls: type, label: str) -> None:
if is_db_unhealthy():
cls.mark_error(cache_cls, RuntimeError("database unhealthy"))
return
if getattr(cache_cls, "_loaded", False):
return
now = time.monotonic()
@@ -756,6 +761,9 @@ class RuntimeCacheMutation:
return
lock = cls._load_locks.setdefault(label, asyncio.Lock())
async with lock:
if is_db_unhealthy():
cls.mark_error(cache_cls, RuntimeError("database unhealthy"))
return
if getattr(cache_cls, "_loaded", False):
return
now = time.monotonic()
@@ -768,7 +776,27 @@ class RuntimeCacheMutation:
cls._retry_after[label] = (
time.monotonic() + RUNTIME_CACHE_LOAD_RETRY_SECONDS
)
raise
return
@staticmethod
async def read_db(cache_cls: type, coro, *, operation: str):
from zhenxun.services.db_context import with_db_timeout
if is_db_unhealthy():
RuntimeCacheMutation.mark_error(
cache_cls, RuntimeError("database unhealthy")
)
return None
try:
return await with_db_timeout(
coro,
timeout=RUNTIME_CACHE_DB_TIMEOUT_SECONDS,
operation=operation,
source="runtime_cache",
)
except Exception as exc:
RuntimeCacheMutation.mark_error(cache_cls, exc)
return None
@staticmethod
def mark_refreshed(cache_cls: type) -> None:
@@ -834,7 +862,13 @@ class PluginInfoMemoryCache:
from zhenxun.models.plugin_info import PluginInfo
async with cls._lock:
plugins = await PluginInfo.all()
plugins = await RuntimeCacheMutation.read_db(
cls,
PluginInfo.all(),
operation="PluginInfoMemoryCache.refresh",
)
if plugins is None:
return
by_module: dict[str, PluginInfoSnapshot] = {}
by_module_path: dict[str, PluginInfoSnapshot] = {}
for plugin in plugins:
@@ -1016,7 +1050,13 @@ class BotMemoryCache:
from zhenxun.models.bot_console import BotConsole
async with cls._lock:
records = await BotConsole.all()
records = await RuntimeCacheMutation.read_db(
cls,
BotConsole.all(),
operation="BotMemoryCache.refresh",
)
if records is None:
return
cls._by_id = {str(r.bot_id): BotSnapshot.from_model(r) for r in records}
RuntimeCacheMutation.clear_negative_all(cls)
RuntimeCacheMutation.mark_refreshed(cls)
@@ -1201,7 +1241,13 @@ class GroupMemoryCache:
from zhenxun.models.group_console import GroupConsole
async with cls._lock:
records = await GroupConsole.all()
records = await RuntimeCacheMutation.read_db(
cls,
GroupConsole.all(),
operation="GroupMemoryCache.refresh",
)
if records is None:
return
by_key: dict[tuple[str, str], GroupSnapshot] = {}
for record in records:
entry = GroupSnapshot.from_model(record)
@@ -1374,7 +1420,13 @@ class LevelUserMemoryCache:
from zhenxun.models.level_user import LevelUser
async with cls._lock:
records = await LevelUser.all()
records = await RuntimeCacheMutation.read_db(
cls,
LevelUser.all(),
operation="LevelUserMemoryCache.refresh",
)
if records is None:
return
by_key: dict[tuple[str, str], LevelUserSnapshot] = {}
by_user_max: dict[str, int] = {}
for record in records:
@@ -1607,7 +1659,13 @@ class TaskInfoMemoryCache:
from zhenxun.models.task_info import TaskInfo
async with cls._lock:
records = await TaskInfo.all()
records = await RuntimeCacheMutation.read_db(
cls,
TaskInfo.all(),
operation="TaskInfoMemoryCache.refresh",
)
if records is None:
return
by_module: dict[str, TaskInfoSnapshot] = {}
by_name: dict[str, TaskInfoSnapshot] = {}
for record in records:
@@ -1790,7 +1848,13 @@ class PluginLimitMemoryCache:
from zhenxun.models.plugin_limit import PluginLimit
async with cls._lock:
records = await PluginLimit.filter(status=True).all()
records = await RuntimeCacheMutation.read_db(
cls,
PluginLimit.filter(status=True).all(),
operation="PluginLimitMemoryCache.refresh",
)
if records is None:
return
by_id: dict[int, PluginLimitSnapshot] = {}
by_module: dict[str, list[PluginLimitSnapshot]] = {}
for record in records:
@@ -2015,7 +2079,13 @@ class BanMemoryCache:
async with cls._lock:
now_ts = time.time()
records = await BanConsole.all()
records = await RuntimeCacheMutation.read_db(
cls,
BanConsole.all(),
operation="BanMemoryCache.refresh",
)
if records is None:
return
by_user: dict[str, BanEntry] = {}
by_group: dict[str, BanEntry] = {}
by_user_group: dict[tuple[str, str], BanEntry] = {}
zhenxun/services/data_access.py
+30 -20
View File
@@ -48,6 +48,7 @@ class DataAccess(Generic[T]):
# 添加缓存统计信息
_cache_stats: ClassVar[dict] = {}
_ENABLE_CACHE_STATS: ClassVar[bool] = False
# 空结果标记
_NULL_RESULT = "__NULL_RESULT_PLACEHOLDER__"
# 默认空结果缓存时间（秒）- 设置为5分钟，避免频繁查询数据库
@@ -87,12 +88,10 @@ class DataAccess(Generic[T]):
self.key_field = getattr(model_cls, "cache_key_field", key_field)
self.cache_type = getattr(model_cls, "cache_type", cache_type)
if not self.cache_type:
raise ValueError("缓存类型不能为空")
self.cache = Cache(self.cache_type)
self.cache = Cache(self.cache_type) if self.cache_type else None
# 初始化缓存统计
if self.cache_type not in self._cache_stats:
if self.cache_type and self.cache_type not in self._cache_stats:
self._cache_stats[self.cache_type] = {
"hits": 0, # 缓存命中次数
"misses": 0, # 缓存未命中次数
@@ -137,6 +136,12 @@ class DataAccess(Generic[T]):
stats["null_sets"] = 0
stats["deletes"] = 0
@classmethod
def _bump_cache_stat(cls, cache_type: str | None, key: str) -> None:
if not cls._ENABLE_CACHE_STATS or not cache_type:
return
cls._cache_stats[cache_type][key] += 1
def _build_cache_key_from_kwargs(self, **kwargs) -> str | None:
"""从关键字参数构建缓存键
@@ -191,7 +196,7 @@ class DataAccess(Generic[T]):
# 如果成功构建缓存键，尝试从缓存获取
if cache_key is not None:
data = await self.cache.get(cache_key)
data = await self.cache.get(cache_key) if self.cache else None
logger.debug(
f"{self.model_cls.__name__} key: {cache_key}"
f" 从缓存获取到的数据 {type(data)}: {data}"
@@ -199,7 +204,7 @@ class DataAccess(Generic[T]):
if data == self._NULL_RESULT:
# 空结果缓存命中
self._cache_stats[self.cache_type]["null_hits"] += 1
self._bump_cache_stat(self.cache_type, "null_hits")
logger.debug(
f"{self.model_cls.__name__} 从缓存获取到空结果: {cache_key}"
)
@@ -211,14 +216,14 @@ class DataAccess(Generic[T]):
return None
elif data:
# 缓存命中
self._cache_stats[self.cache_type]["hits"] += 1
self._bump_cache_stat(self.cache_type, "hits")
logger.debug(
f"{self.model_cls.__name__} 从缓存获取数据成功: {cache_key}"
)
return cast(T, data)
else:
# 缓存未命中
self._cache_stats[self.cache_type]["misses"] += 1
self._bump_cache_stat(self.cache_type, "misses")
logger.debug(f"{self.model_cls.__name__} 缓存未命中: {cache_key}")
except Exception as e:
logger.error(f"{self.model_cls.__name__} 从缓存获取数据失败: {kwargs}", e=e)
@@ -234,8 +239,9 @@ class DataAccess(Generic[T]):
cache_key = self._build_cache_key_for_item(data)
if cache_key is not None:
# 存入缓存
await self.cache.set(cache_key, data)
self._cache_stats[self.cache_type]["sets"] += 1
if self.cache:
await self.cache.set(cache_key, data)
self._bump_cache_stat(self.cache_type, "sets")
logger.debug(
f"{self.model_cls.__name__} 数据已存入缓存: {cache_key}"
)
@@ -249,8 +255,8 @@ class DataAccess(Generic[T]):
# 存入空结果缓存，使用较短的过期时间
await self.cache.set(
cache_key, self._NULL_RESULT, expire=self._NULL_RESULT_TTL
)
self._cache_stats[self.cache_type]["null_sets"] += 1
) if self.cache else None
self._bump_cache_stat(self.cache_type, "null_sets")
logger.debug(
f"{self.model_cls.__name__} 空结果已存入缓存: {cache_key},"
f" TTL={self._NULL_RESULT_TTL}秒"
@@ -343,8 +349,9 @@ class DataAccess(Generic[T]):
return False
# 删除缓存
await self.cache.delete(cache_key)
self._cache_stats[self.cache_type]["deletes"] += 1
if self.cache:
await self.cache.delete(cache_key)
self._bump_cache_stat(self.cache_type, "deletes")
logger.debug(f"已清除{self.model_cls.__name__}缓存: {cache_key}")
return True
except Exception as e:
@@ -416,8 +423,9 @@ class DataAccess(Generic[T]):
if cache_key is None:
return
await self.cache.delete(cache_key)
self._cache_stats[self.cache_type]["deletes"] += 1
if self.cache:
await self.cache.delete(cache_key)
self._bump_cache_stat(self.cache_type, "deletes")
logger.debug(
f"{self.model_cls.__name__} {action}: 已失效兼容缓存: {cache_key}"
)
@@ -543,8 +551,9 @@ class DataAccess(Generic[T]):
if cache_key is not None:
# 如果成功构建缓存键，直接删除缓存
await self.cache.delete(cache_key)
self._cache_stats[self.cache_type]["deletes"] += 1
if self.cache:
await self.cache.delete(cache_key)
self._bump_cache_stat(self.cache_type, "deletes")
logger.debug(
f"{self.model_cls.__name__} delete: 已删除缓存: {cache_key}"
)
@@ -558,8 +567,9 @@ class DataAccess(Generic[T]):
for item in items:
item_cache_key = self._build_cache_key_for_item(item)
if item_cache_key is not None:
await self.cache.delete(item_cache_key)
self._cache_stats[self.cache_type]["deletes"] += 1
if self.cache:
await self.cache.delete(item_cache_key)
self._bump_cache_stat(self.cache_type, "deletes")
if items:
logger.debug(
f"{self.model_cls.__name__} delete:"
zhenxun/services/db_context/__init__.py
+1
View File
@@ -18,6 +18,7 @@ from zhenxun.configs.config import BotConfig
from zhenxun.services.log import logger
from zhenxun.utils.manager.priority_manager import PriorityLifecycle
from . import watchdog as _watchdog # noqa: F401
from .base_model import Model
from .config import (
DB_TIMEOUT_SECONDS,
zhenxun/services/db_context/base_model.py
+8 -8
View File
@@ -27,7 +27,7 @@ class Model(TortoiseModel):
"""
sem_data: ClassVar[dict[str, dict[str, asyncio.Semaphore]]] = {}
_current_locks: ClassVar[dict[int, DbLockType]] = {} # 跟踪当前协程持有的锁
_current_locks: ClassVar[dict[tuple[str, int], DbLockType]] = {}
def __init_subclass__(cls, **kwargs):
super().__init_subclass__(**kwargs)
@@ -95,21 +95,21 @@ class Model(TortoiseModel):
@classmethod
def _require_lock(cls, lock_type: DbLockType) -> bool:
"""检查是否需要真正加锁"""
task_id = id(asyncio.current_task())
return cls._current_locks.get(task_id) != lock_type
lock_key = (cls.__name__, id(asyncio.current_task()))
return cls._current_locks.get(lock_key) != lock_type
@classmethod
@contextlib.asynccontextmanager
async def _lock_context(cls, lock_type: DbLockType):
"""带重入检查的锁上下文"""
task_id = id(asyncio.current_task())
lock_key = (cls.__name__, id(asyncio.current_task()))
need_lock = cls._require_lock(lock_type)
if need_lock and (sem := cls.get_semaphore(lock_type)):
cls._current_locks[task_id] = lock_type
cls._current_locks[lock_key] = lock_type
async with sem:
yield
cls._current_locks.pop(task_id, None)
cls._current_locks.pop(lock_key, None)
else:
yield
@@ -171,8 +171,8 @@ class Model(TortoiseModel):
await obj.save()
result = (obj, False)
else:
# 创建时不重复加锁
result = await cls.create(**kwargs, **(defaults or {})), True
obj = await super().create(**kwargs, **(defaults or {}))
result = (obj, True)
if cache_type := cls.get_cache_type():
await CacheRoot.invalidate_cache(
zhenxun/services/db_context/config.py
+3 -1
View File
@@ -27,7 +27,9 @@ MYSQL_CONFIG = {
SQLITE_CONFIG = {
"journal_mode": "WAL", # 提高并发写入性能
"busy_timeout": 30000, # SQLite 锁等待超时，单位毫秒
# SQLite 的底层锁等待应接近业务超时，避免上层放弃后 worker 仍长时间占线。
# Windows bind mount / Docker 场景下 SQLite 不适合高写并发。
"busy_timeout": 5000,
"foreign_keys": "ON",
}
zhenxun/services/db_context/utils.py
+47 -1
View File
@@ -1,7 +1,9 @@
import asyncio
import contextlib
import time
from zhenxun.services.log import logger
from zhenxun.services.message_load import signal_db_unhealthy
from .config import (
DB_TIMEOUT_SECONDS,
@@ -9,6 +11,40 @@ from .config import (
SLOW_QUERY_THRESHOLD,
)
_SQLITE_STALL_UNTIL = 0.0
_SQLITE_STALL_REASON = ""
_DB_UNHEALTHY_TIMEOUT_SECONDS = 30.0
_SQLITE_STALL_TIMEOUT_SECONDS = 60.0
def _is_sqlite_connection() -> bool:
with contextlib.suppress(Exception):
from tortoise import Tortoise
connection = Tortoise.get_connection("default")
capabilities = getattr(connection, "capabilities", None)
dialect = str(getattr(capabilities, "dialect", "") or "").lower()
return dialect.startswith("sqlite")
return False
def _mark_sqlite_stall(reason: str, duration: float) -> None:
global _SQLITE_STALL_REASON, _SQLITE_STALL_UNTIL
until = time.monotonic() + max(duration, 0.0)
if until > _SQLITE_STALL_UNTIL:
_SQLITE_STALL_UNTIL = until
_SQLITE_STALL_REASON = str(reason or "")[:200]
def is_sqlite_stall_suspected() -> bool:
return time.monotonic() < _SQLITE_STALL_UNTIL
def sqlite_stall_reason() -> str:
if not is_sqlite_stall_suspected():
return ""
return _SQLITE_STALL_REASON
async def with_db_timeout(
coro,
@@ -19,13 +55,23 @@ async def with_db_timeout(
"""带超时控制的数据库操作"""
start_time = time.time()
try:
logger.debug(f"开始执行数据库操作: {operation} 来源: {source}")
result = await asyncio.wait_for(coro, timeout=timeout)
elapsed = time.time() - start_time
if elapsed > SLOW_QUERY_THRESHOLD and operation:
logger.warning(f"慢查询: {operation} 耗时 {elapsed:.3f}s", LOG_COMMAND)
return result
except asyncio.TimeoutError:
timeout_reason = f"{operation or 'database_operation'} from {source or '-'}"
unhealthy_duration = _DB_UNHEALTHY_TIMEOUT_SECONDS
if _is_sqlite_connection():
unhealthy_duration = _SQLITE_STALL_TIMEOUT_SECONDS
_mark_sqlite_stall(timeout_reason, unhealthy_duration)
logger.warning(
"SQLite 数据库操作超时，疑似 aiosqlite worker/连接被锁等待卡住；"
"已暂停低优先级数据库任务",
LOG_COMMAND,
)
signal_db_unhealthy(unhealthy_duration, reason=timeout_reason)
if operation:
logger.error(
f"数据库操作超时: {operation} (>{timeout}s) 来源: {source}",
zhenxun/services/db_context/watchdog.py
+123
View File
@@ -0,0 +1,123 @@
from __future__ import annotations
import asyncio
import contextlib
import time
from tortoise import Tortoise
from tortoise.connection import connections
from zhenxun.services.log import logger
from zhenxun.services.low_priority_writer import low_priority_writer_active_count
from zhenxun.services.message_load import (
signal_db_unhealthy,
)
from zhenxun.utils.manager.priority_manager import PriorityLifecycle
from .config import LOG_COMMAND
_CHECK_INTERVAL_SECONDS = 15.0
_CHECK_TIMEOUT_SECONDS = 2.0
_FAIL_THRESHOLD = 3
_UNHEALTHY_SECONDS = 60.0
_RECONNECT_COOLDOWN_SECONDS = 60.0
_RECONNECT_WAIT_IDLE_SECONDS = 3.0
_WATCHDOG_TASK: asyncio.Task[None] | None = None
_RECONNECT_LOCK = asyncio.Lock()
_LAST_RECONNECT_AT = 0.0
def _is_sqlite_connection() -> bool:
with contextlib.suppress(Exception):
connection = Tortoise.get_connection("default")
capabilities = getattr(connection, "capabilities", None)
dialect = str(getattr(capabilities, "dialect", "") or "").lower()
return dialect.startswith("sqlite")
return False
async def _select_one() -> None:
connection = Tortoise.get_connection("default")
await connection.execute_query("SELECT 1")
async def _try_reconnect(reason: str) -> None:
global _LAST_RECONNECT_AT
now = time.monotonic()
if now - _LAST_RECONNECT_AT < _RECONNECT_COOLDOWN_SECONDS:
return
if low_priority_writer_active_count() > 0:
return
async with _RECONNECT_LOCK:
now = time.monotonic()
if now - _LAST_RECONNECT_AT < _RECONNECT_COOLDOWN_SECONDS:
return
await asyncio.sleep(_RECONNECT_WAIT_IDLE_SECONDS)
if low_priority_writer_active_count() > 0:
return
try:
await connections.close_all(discard=True)
# ConnectionHandler lazily recreates default connection from db_config.
Tortoise.get_connection("default")
_LAST_RECONNECT_AT = time.monotonic()
logger.warning(
f"SQLite watchdog rebuilt default connection: {reason}",
LOG_COMMAND,
)
except Exception as exc:
_LAST_RECONNECT_AT = time.monotonic()
signal_db_unhealthy(
_UNHEALTHY_SECONDS,
reason=f"watchdog reconnect:{reason}",
)
logger.warning("SQLite watchdog reconnect failed", LOG_COMMAND, e=exc)
async def _watchdog_loop() -> None:
failures = 0
while True:
await asyncio.sleep(_CHECK_INTERVAL_SECONDS)
if not _is_sqlite_connection():
failures = 0
continue
try:
await asyncio.wait_for(_select_one(), timeout=_CHECK_TIMEOUT_SECONDS)
failures = 0
except asyncio.CancelledError:
raise
except Exception as exc:
failures += 1
reason = f"sqlite watchdog SELECT 1 failed x{failures}: {exc}"
signal_db_unhealthy(_UNHEALTHY_SECONDS, reason=reason)
logger.warning(reason, LOG_COMMAND)
if failures >= _FAIL_THRESHOLD:
await _try_reconnect(reason)
def start_db_watchdog() -> None:
global _WATCHDOG_TASK
if _WATCHDOG_TASK is not None and not _WATCHDOG_TASK.done():
return
_WATCHDOG_TASK = asyncio.create_task(_watchdog_loop())
async def stop_db_watchdog() -> None:
global _WATCHDOG_TASK
task = _WATCHDOG_TASK
_WATCHDOG_TASK = None
if task is None:
return
task.cancel()
with contextlib.suppress(BaseException):
await task
@PriorityLifecycle.on_startup(priority=8)
async def _start_db_watchdog() -> None:
start_db_watchdog()
@PriorityLifecycle.on_shutdown(priority=10)
async def _stop_db_watchdog() -> None:
await stop_db_watchdog()
zhenxun/services/group_settings_service.py
+7 -11
View File
@@ -6,7 +6,6 @@ import ujson as json
from zhenxun.configs.config import Config
from zhenxun.models.group_plugin_setting import GroupPluginSetting
from zhenxun.services.cache import BoundedTTLCache
from zhenxun.services.data_access import DataAccess
from zhenxun.services.log import logger
from zhenxun.utils.pydantic_compat import model_dump, model_validate, parse_as
@@ -20,7 +19,6 @@ class GroupSettingsService:
"""
def __init__(self):
self.dao = DataAccess(GroupPluginSetting)
self._cache = BoundedTTLCache[str, dict[str, Any]](
"GROUP_PLUGIN_SETTINGS_VIEW",
ttl_seconds=600,
@@ -48,13 +46,12 @@ class GroupSettingsService:
settings_dict = model_dump(settings_model)
json_value = json.dumps(settings_dict, ensure_ascii=False)
await self.dao.update_or_create(
await GroupPluginSetting.update_or_create(
defaults={"settings": json_value}, # type: ignore
group_id=group_id,
plugin_name=plugin_name,
)
await self.dao.clear_cache(group_id=group_id, plugin_name=plugin_name)
await self._clear_merged_cache(group_id, plugin_name)
async def set_key_value(
@@ -72,19 +69,19 @@ class GroupSettingsService:
setting_entry.settings[key] = value
await setting_entry.save(update_fields=["settings"])
await self.dao.clear_cache(group_id=group_id, plugin_name=plugin_name)
await self._clear_merged_cache(group_id, plugin_name)
async def reset_key(self, group_id: str, plugin_name: str, key: str) -> bool:
"""重置单个配置项"""
setting = await self.dao.get_or_none(group_id=group_id, plugin_name=plugin_name)
setting = await GroupPluginSetting.safe_get_or_none(
group_id=group_id, plugin_name=plugin_name
)
if setting and isinstance(setting.settings, dict) and key in setting.settings:
del setting.settings[key]
if not setting.settings:
await setting.delete()
else:
await setting.save(update_fields=["settings"])
await self.dao.clear_cache(group_id=group_id, plugin_name=plugin_name)
await self._clear_merged_cache(group_id, plugin_name)
return True
return False
@@ -119,12 +116,11 @@ class GroupSettingsService:
返回:
bool: 如果成功删除了一个条目，则返回 True，否则返回 False。
"""
deleted_count = await self.dao.delete(
deleted_count = await GroupPluginSetting.filter(
group_id=group_id, plugin_name=plugin_name
)
).delete()
if deleted_count > 0:
await self.dao.clear_cache(group_id=group_id, plugin_name=plugin_name)
await self._clear_merged_cache(group_id, plugin_name)
logger.debug(f"已重置插件 '{plugin_name}' 在群组 '{group_id}' 的配置。")
return True
@@ -176,7 +172,7 @@ class GroupSettingsService:
for key in global_config_group.configs.keys()
}
group_setting_entry = await self.dao.get_or_none(
group_setting_entry = await GroupPluginSetting.safe_get_or_none(
group_id=group_id, plugin_name=plugin_name
)
if group_setting_entry:
zhenxun/services/hot_query_cache.py
+97 -29
View File
@@ -9,6 +9,8 @@ from typing import Any, Literal
from tortoise.functions import Count
from zhenxun.services.cache.bounded_ttl import BoundedTTLCache
from zhenxun.services.db_context import with_db_timeout
from zhenxun.services.message_load import is_db_unhealthy
@dataclass(frozen=True, slots=True)
@@ -81,6 +83,26 @@ _CHAT_RANK_LOCKS: dict[str, asyncio.Lock] = {}
_CHAT_FIRST_MSG_LOCKS: dict[str, asyncio.Lock] = {}
_STATISTICS_LOCKS: dict[str, asyncio.Lock] = {}
_MAX_LOCK_POOL_SIZE = 4096
_MEMBER_DB_TIMEOUT = 2.0
_AGGREGATE_DB_TIMEOUT = 3.0
async def _read_or_default(
coro,
*,
timeout: float,
operation: str,
default,
):
try:
return await with_db_timeout(
coro,
timeout=timeout,
operation=operation,
source="hot_query_cache",
)
except TimeoutError:
return default
def _get_lock(pool: dict[str, asyncio.Lock], key: str) -> asyncio.Lock:
@@ -126,13 +148,20 @@ async def get_group_members(
from zhenxun.models.group_member_info import GroupInfoUser
rows = await GroupInfoUser.filter(group_id=group_key).values_list(
"id",
"user_id",
"user_name",
"user_join_time",
"uid",
"platform",
if is_db_unhealthy():
return ()
rows = await _read_or_default(
GroupInfoUser.filter(group_id=group_key).values_list(
"id",
"user_id",
"user_name",
"user_join_time",
"uid",
"platform",
),
timeout=_MEMBER_DB_TIMEOUT,
operation="hot_query_cache.get_group_members",
default=(),
)
members = tuple(
GroupMemberSnapshot(
@@ -192,15 +221,20 @@ async def get_group_member_map(
if missing:
from zhenxun.models.group_member_info import GroupInfoUser
rows = await GroupInfoUser.filter(
group_id=group_key, user_id__in=missing
).values_list(
"id",
"user_id",
"user_name",
"user_join_time",
"uid",
"platform",
if is_db_unhealthy():
return result
rows = await _read_or_default(
GroupInfoUser.filter(group_id=group_key, user_id__in=missing).values_list(
"id",
"user_id",
"user_name",
"user_join_time",
"uid",
"platform",
),
timeout=_MEMBER_DB_TIMEOUT,
operation="hot_query_cache.get_group_member_map",
default=(),
)
found: set[str] = set()
for row in rows:
@@ -258,8 +292,13 @@ async def get_group_user_ids(group_id: str | int | None) -> set[str]:
from zhenxun.models.group_member_info import GroupInfoUser
rows = await GroupInfoUser.filter(group_id=group_key).values_list(
"user_id", flat=True
if is_db_unhealthy():
return set()
rows = await _read_or_default(
GroupInfoUser.filter(group_id=group_key).values_list("user_id", flat=True),
timeout=_MEMBER_DB_TIMEOUT,
operation="hot_query_cache.get_group_user_ids",
default=(),
)
user_ids = tuple(str(user_id) for user_id in rows if user_id)
await _GROUP_USER_IDS_CACHE.set(group_key, user_ids)
@@ -283,8 +322,13 @@ async def get_user_group_ids(user_id: str | int | None) -> list[str]:
from zhenxun.models.group_member_info import GroupInfoUser
rows = await GroupInfoUser.filter(user_id=user_key).values_list(
"group_id", flat=True
if is_db_unhealthy():
return []
rows = await _read_or_default(
GroupInfoUser.filter(user_id=user_key).values_list("group_id", flat=True),
timeout=_MEMBER_DB_TIMEOUT,
operation="hot_query_cache.get_user_group_ids",
default=(),
)
group_ids = tuple(str(group_id) for group_id in rows if group_id)
await _USER_GROUP_CACHE.set(user_key, group_ids)
@@ -314,8 +358,15 @@ async def get_member_names(
if missing:
from zhenxun.models.group_member_info import GroupInfoUser
rows = await GroupInfoUser.filter(user_id__in=missing).values_list(
"user_id", "user_name"
if is_db_unhealthy():
return result
rows = await _read_or_default(
GroupInfoUser.filter(user_id__in=missing).values_list(
"user_id", "user_name"
),
timeout=_MEMBER_DB_TIMEOUT,
operation="hot_query_cache.get_member_names",
default=(),
)
for user_id, user_name in rows:
user_key = str(user_id)
@@ -395,6 +446,8 @@ async def get_chat_history_rank_cached(
if cached is not None:
return list(cached)
if is_db_unhealthy():
return []
order_prefix = "-" if order == "DESC" else ""
query: Any = model.filter(group_id=gid) if gid else model
if date_scope:
@@ -403,12 +456,15 @@ async def get_chat_history_rank_cached(
date_scope[1].isoformat(" "),
)
query = query.filter(create_time__range=filter_scope)
rows = (
await query.annotate(count=Count("user_id"))
rows = await _read_or_default(
query.annotate(count=Count("user_id"))
.order_by(f"{order_prefix}count")
.group_by("user_id")
.limit(limit)
.values_list("user_id", "count")
.values_list("user_id", "count"),
timeout=_AGGREGATE_DB_TIMEOUT,
operation="hot_query_cache.get_chat_history_rank",
default=(),
)
result = tuple((str(user_id), int(count)) for user_id, count in rows)
await _CHAT_RANK_CACHE.set(key, result)
@@ -430,8 +486,15 @@ async def get_chat_history_first_msg_datetime_cached(
if cached is not None:
return cached[0]
if is_db_unhealthy():
return None
query: Any = model.filter(group_id=group_id) if group_id else model.all()
message = await query.order_by("create_time").first()
message = await _read_or_default(
query.order_by("create_time").first(),
timeout=_AGGREGATE_DB_TIMEOUT,
operation="hot_query_cache.get_chat_history_first_msg",
default=None,
)
result = getattr(message, "create_time", None) if message else None
await _CHAT_FIRST_MSG_CACHE.set(key, (result,))
return result
@@ -459,6 +522,8 @@ async def get_statistics_plugin_counts_cached(
if cached is not None:
return list(cached)
if is_db_unhealthy():
return []
from zhenxun.models.statistics import Statistics
query: Any = Statistics
@@ -472,10 +537,13 @@ async def get_statistics_plugin_counts_cached(
query = query.filter(plugin_name=plugin_name)
if start_time:
query = query.filter(create_time__gte=start_time)
rows = (
await query.annotate(count=Count("id"))
rows = await _read_or_default(
query.annotate(count=Count("id"))
.group_by("plugin_name")
.values_list("plugin_name", "count")
.values_list("plugin_name", "count"),
timeout=_AGGREGATE_DB_TIMEOUT,
operation="hot_query_cache.get_statistics_plugin_counts",
default=(),
)
result = tuple((str(plugin), int(count)) for plugin, count in rows)
await _STATISTICS_COUNT_CACHE.set(key, result)
zhenxun/services/low_priority_writer.py
+329
View File
@@ -0,0 +1,329 @@
from __future__ import annotations
import asyncio
from collections import deque
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
import time
from typing import Any
from zhenxun.services.log import logger
from zhenxun.services.message_load import should_pause_tasks, signal_db_unhealthy
from zhenxun.utils.manager.priority_manager import PriorityLifecycle
LOG_COMMAND = "LowPriorityWriter"
WriteBatch = Callable[[list[Any], str], Awaitable[None]]
_POLL_INTERVAL_SECONDS = 1.0
_DB_UNHEALTHY_SECONDS = 30.0
@dataclass(slots=True)
class LowPriorityWriterConfig:
name: str
write_batch: WriteBatch
batch_size: int = 500
trigger_size: int = 500
max_retain: int = 10_000
flush_interval_seconds: float = 60.0
max_items_per_cycle: int = 1_000
backoff_base_seconds: float = 30.0
backoff_max_seconds: float = 600.0
log_command: str = LOG_COMMAND
@dataclass(slots=True)
class _WriterState:
config: LowPriorityWriterConfig
buffer: deque[Any] = field(default_factory=deque)
lock: asyncio.Lock = field(default_factory=asyncio.Lock)
dropped: int = 0
last_drop_log_at: float = 0.0
last_flush_at: float = field(default_factory=time.monotonic)
failures: int = 0
backoff_until: float = 0.0
_WRITERS: dict[str, _WriterState] = {}
_WORKER_TASK: asyncio.Task[None] | None = None
_WAKE_EVENT: asyncio.Event | None = None
_FLUSH_LOCK = asyncio.Lock()
_ACTIVE_FLUSHES = 0
_STOPPING = False
def _wake() -> None:
if _WAKE_EVENT is not None:
_WAKE_EVENT.set()
def _ensure_worker() -> None:
global _STOPPING, _WAKE_EVENT, _WORKER_TASK
if _STOPPING:
_STOPPING = False
try:
loop = asyncio.get_running_loop()
except RuntimeError:
return
if _WAKE_EVENT is None:
_WAKE_EVENT = asyncio.Event()
if _WORKER_TASK is not None and not _WORKER_TASK.done():
return
_WORKER_TASK = loop.create_task(_worker_loop())
def register_low_priority_writer(config: LowPriorityWriterConfig) -> None:
"""Register or update an append-only low-priority DB writer."""
if config.batch_size <= 0:
raise ValueError("batch_size must be positive")
if config.trigger_size <= 0:
raise ValueError("trigger_size must be positive")
if config.max_retain <= 0:
raise ValueError("max_retain must be positive")
state = _WRITERS.get(config.name)
if state is None:
_WRITERS[config.name] = _WriterState(config=config)
else:
state.config = config
_ensure_worker()
async def append_low_priority_record(name: str, record: Any) -> bool:
"""Append a record without doing DB work in the caller's hot path."""
state = _WRITERS.get(name)
if state is None:
raise KeyError(f"low priority writer not registered: {name}")
_ensure_worker()
should_wake = False
async with state.lock:
if len(state.buffer) >= state.config.max_retain:
state.buffer.popleft()
state.dropped += 1
_log_drop_if_needed(state)
state.buffer.append(record)
should_wake = len(state.buffer) >= state.config.trigger_size
if should_wake:
_wake()
return True
def _log_drop_if_needed(state: _WriterState) -> None:
now = time.monotonic()
if now - state.last_drop_log_at < 10.0:
return
state.last_drop_log_at = now
logger.warning(
f"{state.config.name} low priority buffer full, "
f"dropped={state.dropped}, backlog={len(state.buffer)}",
state.config.log_command,
)
async def flush_low_priority_writer(
name: str,
reason: str,
*,
force: bool = False,
) -> int:
state = _WRITERS.get(name)
if state is None:
return 0
async with _FLUSH_LOCK:
return await _flush_state(state, reason, force=force)
async def flush_all_low_priority_writers(
reason: str,
*,
force: bool = False,
) -> int:
total = 0
async with _FLUSH_LOCK:
for state in list(_WRITERS.values()):
total += await _flush_state(state, reason, force=force)
return total
async def _worker_loop() -> None:
while not _STOPPING:
event = _WAKE_EVENT
if event is None:
await asyncio.sleep(_POLL_INTERVAL_SECONDS)
else:
try:
await asyncio.wait_for(
event.wait(),
timeout=_POLL_INTERVAL_SECONDS,
)
except asyncio.TimeoutError:
pass
event.clear()
if should_pause_tasks():
continue
async with _FLUSH_LOCK:
for state in list(_WRITERS.values()):
if _state_due_for_flush(state):
await _flush_state(state, "低优先队列")
def _state_due_for_flush(state: _WriterState) -> bool:
if not state.buffer:
return False
now = time.monotonic()
if now < state.backoff_until:
return False
return (
len(state.buffer) >= state.config.trigger_size
or now - state.last_flush_at >= state.config.flush_interval_seconds
)
async def _flush_state(
state: _WriterState,
reason: str,
*,
force: bool = False,
) -> int:
if not force:
if should_pause_tasks():
return 0
if time.monotonic() < state.backoff_until:
return 0
written = 0
max_items = state.config.max_items_per_cycle if not force else float("inf")
while written < max_items:
batch = await _take_batch(state)
if not batch:
break
try:
await _write_batch(state, batch, reason)
except (TimeoutError, asyncio.TimeoutError) as exc:
# asyncio.wait_for timeout only cancels the awaiter. With SQLite/aiosqlite
# the worker thread may still finish the SQL later, so restoring this
# append-only low priority batch can duplicate rows. Prefer dropping the
# uncertain batch; chat history/statistics/logs are lossy by design here.
_mark_uncertain_timeout(state, reason, exc, len(batch))
break
except Exception as exc:
await _restore_batch(state, batch)
_mark_failure(state, reason, exc)
break
written += len(batch)
state.failures = 0
state.backoff_until = 0.0
state.last_flush_at = time.monotonic()
if written:
logger.debug(
f"{reason}写入 {state.config.name} {written} 条, "
f"backlog={len(state.buffer)}",
state.config.log_command,
)
return written
async def _take_batch(state: _WriterState) -> list[Any]:
batch: list[Any] = []
async with state.lock:
while state.buffer and len(batch) < state.config.batch_size:
batch.append(state.buffer.popleft())
return batch
async def _restore_batch(state: _WriterState, batch: list[Any]) -> None:
if not batch:
return
async with state.lock:
retain_count = max(state.config.max_retain - len(state.buffer), 0)
restore_items = batch[-retain_count:] if retain_count else []
for record in reversed(restore_items):
state.buffer.appendleft(record)
dropped = len(batch) - len(restore_items)
if dropped:
state.dropped += dropped
_log_drop_if_needed(state)
async def _write_batch(
state: _WriterState,
batch: list[Any],
reason: str,
) -> None:
global _ACTIVE_FLUSHES
_ACTIVE_FLUSHES += 1
try:
await state.config.write_batch(batch, reason)
finally:
_ACTIVE_FLUSHES = max(_ACTIVE_FLUSHES - 1, 0)
def _mark_failure(state: _WriterState, reason: str, exc: Exception) -> None:
state.failures += 1
backoff = min(
state.config.backoff_base_seconds * (2 ** (state.failures - 1)),
state.config.backoff_max_seconds,
)
state.backoff_until = time.monotonic() + backoff
signal_db_unhealthy(_DB_UNHEALTHY_SECONDS, reason=f"{state.config.name}:{reason}")
logger.warning(
f"{reason}写入 {state.config.name} 失败, "
f"backoff={backoff:.0f}s, backlog={len(state.buffer)}",
state.config.log_command,
e=exc,
)
def _mark_uncertain_timeout(
state: _WriterState,
reason: str,
exc: BaseException,
batch_size: int,
) -> None:
state.failures += 1
backoff = min(
state.config.backoff_base_seconds * (2 ** (state.failures - 1)),
state.config.backoff_max_seconds,
)
state.backoff_until = time.monotonic() + backoff
state.dropped += batch_size
signal_db_unhealthy(_DB_UNHEALTHY_SECONDS, reason=f"{state.config.name}:{reason}")
log_exc = exc if isinstance(exc, Exception) else None
logger.warning(
f"{reason}写入 {state.config.name} 超时, "
f"dropped_uncertain={batch_size}, backoff={backoff:.0f}s, "
f"backlog={len(state.buffer)}",
state.config.log_command,
e=log_exc,
)
def low_priority_writer_active_count() -> int:
return _ACTIVE_FLUSHES
def low_priority_writer_backlog() -> dict[str, int]:
return {name: len(state.buffer) for name, state in _WRITERS.items()}
async def stop_low_priority_writer() -> int:
global _STOPPING, _WORKER_TASK
_STOPPING = True
task = _WORKER_TASK
_WORKER_TASK = None
if task is not None:
task.cancel()
try:
await task
except BaseException:
pass
return await flush_all_low_priority_writers("关闭", force=True)
@PriorityLifecycle.on_startup(priority=3)
async def _start_low_priority_writer() -> None:
_ensure_worker()
@PriorityLifecycle.on_shutdown(priority=95)
async def _stop_low_priority_writer() -> None:
await stop_low_priority_writer()
zhenxun/services/memory_governor.py
+1 -3
View File
@@ -12,7 +12,7 @@ from aiocache import SimpleMemoryCache
from zhenxun.services.cache import CacheRoot
from zhenxun.services.cache.bounded_ttl import BoundedTTLCache
from zhenxun.services.cache.cache_containers import CacheDict, CacheList
from zhenxun.services.cache.cache_containers import CacheDict
from zhenxun.services.log import logger
from zhenxun.services.message_load import idle_seconds, is_overloaded
@@ -131,7 +131,6 @@ async def _run_reclaim() -> None:
cleared: dict[str, Any] = {}
cache_stats_before = {
"cache_dict": CacheDict.stats_all(),
"cache_list": CacheList.stats_all(),
"bounded_ttl": await BoundedTTLCache.stats_all(),
}
@@ -139,7 +138,6 @@ async def _run_reclaim() -> None:
cleared["user_gold_logs"] = await _flush_user_gold_log_buffer()
cleared["bounded_ttl_clear"] = await BoundedTTLCache.clear_all()
cleared["cache_dict_clear"] = CacheDict.clear_all()
cleared["cache_list_clear"] = CacheList.clear_all()
cleared["runtime_negative"] = _clear_runtime_negative_caches()
cleared["auth_local"] = _clear_auth_local_caches()
cleared["avatar_l1"] = _clear_avatar_memory_cache()
zhenxun/services/message_load.py
+25 -1
View File
@@ -3,6 +3,8 @@ from __future__ import annotations
import time
_OVERLOAD_UNTIL = 0.0
_DB_UNHEALTHY_UNTIL = 0.0
_DB_UNHEALTHY_REASON = ""
_LAST_ACTIVITY = time.monotonic()
@@ -31,5 +33,27 @@ def is_overloaded() -> bool:
return time.monotonic() < _OVERLOAD_UNTIL
def signal_db_unhealthy(duration: float = 30.0, reason: str = "") -> None:
"""Mark database-dependent low-priority tasks as unsafe to run briefly."""
global _DB_UNHEALTHY_REASON, _DB_UNHEALTHY_UNTIL
if duration <= 0:
return
now = time.monotonic()
until = now + duration
if until > _DB_UNHEALTHY_UNTIL:
_DB_UNHEALTHY_UNTIL = until
_DB_UNHEALTHY_REASON = str(reason or "")[:200]
def is_db_unhealthy() -> bool:
return time.monotonic() < _DB_UNHEALTHY_UNTIL
def db_unhealthy_reason() -> str:
if not is_db_unhealthy():
return ""
return _DB_UNHEALTHY_REASON
def should_pause_tasks() -> bool:
return is_overloaded()
return is_overloaded() or is_db_unhealthy()