Diuse/zhenxun_bot
1
0
Fork 0
mirror of https://github.com/zhenxun-org/zhenxun_bot.git synced 2026-09-28 16:20:56 +08:00
Code Issues Actions 11 Packages Projects Releases Wiki Activity

性能优化 (#2126)

Browse Source 
* 性能优化

* 代码改进

* 优化浏览器代际切换逻辑

* 统一缓存与生命周期

* 添加aiomysql依赖

* 优化插件路径处理逻辑，简化条件判断；在虚拟环境包管理器中添加编码和错误处理参数以增强稳定性

* 🚨 auto fix by pre-commit hooks

* 优化Windows下的关闭逻辑

* 代码优化

* bugfix:修复配置重载问题

* bugfix:修复插件加载启动竞态问题

* 收敛事件入口和权限上下文

* 优化 Windows launcher 关闭重启兜底

---------

Co-authored-by: HibiKier <775757368@qq.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Copaan
2026-04-26 15:50:15 +08:00
committed by GitHub
co-authored by HibiKier pre-commit-ci[bot]
parent 24c316cd2c
commit 5d92ccd3b0
56 changed files with 3092 additions and 806 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pyproject.toml
+1
View File
@@ -45,6 +45,7 @@ dependencies = [
"alibabacloud-devops20210625>=5.0.2,<6.0.0",
"uvloop>=0.21.0; sys_platform != 'win32'",
"pytest-timeout>=2.4.0",
"aiomysql>=0.3.2",
]
[project.scripts]
uv.lock
Generated
+23
View File
@@ -158,6 +158,18 @@ wheels = [
{ url = "https://mirrors.aliyun.com/pypi/packages/62/29/2f8418269e46454a26171bfdd6a055d74febf32234e474930f2f60a17145/aiohttp-3.13.5-cp314-cp314t-win_amd64.whl", hash = "sha256:18a2f6c1182c51baa1d28d68fea51513cb2a76612f038853c0ad3c145423d3d9" },
]
[[package]]
name = "aiomysql"
version = "0.3.2"
source = { registry = "https://mirrors.aliyun.com/pypi/simple/" }
dependencies = [
{ name = "pymysql" },
]
sdist = { url = "https://mirrors.aliyun.com/pypi/packages/29/e0/302aeffe8d90853556f47f3106b89c16cc2ec2a4d269bdfd82e3f4ae12cc/aiomysql-0.3.2.tar.gz", hash = "sha256:72d15ef5cfc34c03468eb41e1b90adb9fd9347b0b589114bd23ead569a02ac1a" }
wheels = [
{ url = "https://mirrors.aliyun.com/pypi/packages/4c/af/aae0153c3e28712adaf462328f6c7a3c196a1c1c27b491de4377dd3e6b52/aiomysql-0.3.2-py3-none-any.whl", hash = "sha256:c82c5ba04137d7afd5c693a258bea8ead2aad77101668044143a991e04632eb2" },
]
[[package]]
name = "aiosignal"
version = "1.4.0"
@@ -2598,6 +2610,15 @@ wheels = [
{ url = "https://mirrors.aliyun.com/pypi/packages/f7/27/a2fc51a4a122dfd1015e921ae9d22fee3d20b0b8080d9a704578bf9deece/pymdown_extensions-10.21.2-py3-none-any.whl", hash = "sha256:5c0fd2a2bea14eb39af8ff284f1066d898ab2187d81b889b75d46d4348c01638" },
]
[[package]]
name = "pymysql"
version = "1.1.2"
source = { registry = "https://mirrors.aliyun.com/pypi/simple/" }
sdist = { url = "https://mirrors.aliyun.com/pypi/packages/f5/ae/1fe3fcd9f959efa0ebe200b8de88b5a5ce3e767e38c7ac32fb179f16a388/pymysql-1.1.2.tar.gz", hash = "sha256:4961d3e165614ae65014e361811a724e2044ad3ea3739de9903ae7c21f539f03" }
wheels = [
{ url = "https://mirrors.aliyun.com/pypi/packages/7c/4c/ad33b92b9864cbde84f259d5df035a6447f91891f5be77788e2a3892bce3/pymysql-1.1.2-py3-none-any.whl", hash = "sha256:e6b1d89711dd51f8f74b1631fe08f039e7d76cf67a42a323d3178f0f25762ed9" },
]
[[package]]
name = "pypika-tortoise"
version = "0.1.6"
@@ -4157,6 +4178,7 @@ source = { editable = "." }
dependencies = [
{ name = "aiocache", extra = ["redis"] },
{ name = "aiofiles" },
{ name = "aiomysql" },
{ name = "alibabacloud-devops20210625" },
{ name = "asyncpg" },
{ name = "beautifulsoup4" },
@@ -4212,6 +4234,7 @@ dev = [
requires-dist = [
{ name = "aiocache", extras = ["redis"], specifier = ">=0.12.3" },
{ name = "aiofiles", specifier = ">=23.2.1" },
{ name = "aiomysql", specifier = ">=0.3.2" },
{ name = "alibabacloud-devops20210625", specifier = ">=5.0.2,<6.0.0" },
{ name = "asyncpg", specifier = ">=0.20.0" },
{ name = "beautifulsoup4", specifier = ">=4.12.3,<5.0.0" },
zhenxun/builtin_plugins/admin/plugin_switch/data_source.py
+1 -4
View File
@@ -1,10 +1,9 @@
from nonebot.adapters import Bot
from zhenxun.models.group_console import GroupConsole
from zhenxun.services.cache import CacheRoot
from zhenxun.services.cache.runtime_cache import GroupMemoryCache
from zhenxun.utils.common_utils import CommonUtils
from zhenxun.utils.enum import BlockType, CacheType
from zhenxun.utils.enum import BlockType
from zhenxun.utils.platform import PlatformUtils
from .strategy import get_strategy
@@ -134,7 +133,6 @@ class PluginManager:
await GroupConsole.bulk_update(
update_list, [norm_field, su_field], batch_size=500
)
await CacheRoot.clear(CacheType.GROUPS)
for group in update_list:
await GroupMemoryCache.upsert_from_model(group)
@@ -318,7 +316,6 @@ class PluginManager:
status=False
)
await CacheRoot.clear(CacheType.GROUPS)
await GroupMemoryCache.refresh()
action_str = "醒来" if status else "休眠"
zhenxun/builtin_plugins/admin/plugin_switch/strategy.py
+1 -3
View File
@@ -4,12 +4,11 @@ from typing import Any, cast
from zhenxun.models.group_console import GroupConsole
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.models.task_info import TaskInfo
from zhenxun.services.cache import CacheRoot
from zhenxun.services.cache.runtime_cache import (
PluginInfoMemoryCache,
TaskInfoMemoryCache,
)
from zhenxun.utils.enum import BlockType, CacheType, PluginType
from zhenxun.utils.enum import BlockType, PluginType
class SwitchStrategy(ABC):
@@ -135,7 +134,6 @@ class PluginStrategy(SwitchStrategy):
await self.refresh_cache()
async def refresh_cache(self) -> None:
await CacheRoot.invalidate_cache(CacheType.PLUGINS)
await PluginInfoMemoryCache.refresh()
zhenxun/builtin_plugins/chat_history/chat_message.py
+27 -8
View File
@@ -10,6 +10,7 @@ from nonebot_plugin_uninfo import Uninfo
from zhenxun.configs.config import Config
from zhenxun.configs.utils import PluginExtraData, RegisterConfig
from zhenxun.models.chat_history import ChatHistory
from zhenxun.services.db_context import with_db_timeout
from zhenxun.services.log import logger
from zhenxun.services.message_load import is_overloaded, should_pause_tasks
from zhenxun.utils.enum import PluginType
@@ -47,6 +48,9 @@ _HISTORY_QUEUE: asyncio.Queue[ChatHistory] = asyncio.Queue(maxsize=5000)
_DROP_COUNT = 0
_LAST_DROP_LOG = 0.0
_DROP_LOG_INTERVAL = 10.0
_FLUSH_BATCH_SIZE = 200
_FLUSH_MAX_PER_TICK = 1000
_FLUSH_DB_TIMEOUT = 5.0
@chat_history.handle()
@@ -80,19 +84,34 @@ async def _(message: UniMsg, session: Uninfo):
@scheduler.scheduled_job(
"interval",
minutes=1,
max_instances=1,
coalesce=True,
)
async def _():
try:
if should_pause_tasks():
return
message_list: list[ChatHistory] = []
while True:
try:
message_list.append(_HISTORY_QUEUE.get_nowait())
except asyncio.QueueEmpty:
flushed = 0
while flushed < _FLUSH_MAX_PER_TICK:
message_list: list[ChatHistory] = []
limit = min(_FLUSH_BATCH_SIZE, _FLUSH_MAX_PER_TICK - flushed)
for _ in range(limit):
try:
message_list.append(_HISTORY_QUEUE.get_nowait())
except asyncio.QueueEmpty:
break
if not message_list:
break
if message_list:
await ChatHistory.bulk_create(message_list)
logger.debug(f"批量添加聊天记录 {len(message_list)} 条", "定时任务")
await with_db_timeout(
ChatHistory.bulk_create(message_list, _FLUSH_BATCH_SIZE),
timeout=_FLUSH_DB_TIMEOUT,
operation=f"ChatHistory.bulk_create[{len(message_list)}]",
source="chat_history",
)
flushed += len(message_list)
if flushed:
backlog = _HISTORY_QUEUE.qsize()
suffix = f"，剩余队列 {backlog} 条" if backlog else ""
logger.debug(f"批量添加聊天记录 {flushed} 条{suffix}", "定时任务")
except Exception as e:
logger.warning("存储聊天记录失败", "chat_history", e=e)
zhenxun/builtin_plugins/hooks/auth/auth_admin.py
+13 -4
View File
@@ -7,9 +7,10 @@ from zhenxun.models.level_user import LevelUser
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.services.cache.runtime_cache import LevelUserMemoryCache, LevelUserSnapshot
from zhenxun.services.log import logger
from zhenxun.utils.utils import get_entity_ids
from zhenxun.utils.utils import EntityIDs, get_entity_ids
from .config import LOGGER_COMMAND, WARNING_THRESHOLD
from .context import PermissionContext
from .exception import SkipPluginException
@@ -20,6 +21,9 @@ async def auth_admin(
LevelUser | LevelUserSnapshot | None, LevelUser | LevelUserSnapshot | None
]
| None = None,
*,
context: PermissionContext | None = None,
entity: EntityIDs | None = None,
):
"""管理员命令 个人权限
@@ -33,7 +37,12 @@ async def auth_admin(
return
try:
entity = get_entity_ids(session)
if context is not None:
entity = context.entity
if cached_levels is None:
cached_levels = context.admin_levels
if entity is None:
entity = get_entity_ids(session)
global_user: LevelUser | LevelUserSnapshot | None = None
group_users: LevelUser | LevelUserSnapshot | None = None
@@ -42,7 +51,7 @@ async def auth_admin(
global_user, group_users = cached_levels
else:
global_user, group_users = await LevelUserMemoryCache.get_levels(
session.user.id, entity.group_id
entity.user_id, entity.group_id
)
user_level = global_user.user_level if global_user else 0
@@ -53,7 +62,7 @@ async def auth_admin(
raise SkipPluginException(
f"{plugin.name}({plugin.module}) 管理员权限不足...",
tip_message=[
At(flag="user", target=session.user.id),
At(flag="user", target=entity.user_id),
f"你的权限不足喔，该功能需要的权限等级: {plugin.admin_level}",
],
tip_check_tag=entity.user_id,
zhenxun/builtin_plugins/hooks/auth/auth_ban.py
+7 -48
View File
@@ -1,4 +1,3 @@
import asyncio
import time
from nonebot.matcher import Matcher
@@ -8,7 +7,6 @@ from nonebot_plugin_uninfo import Uninfo
from zhenxun.configs.config import Config
from zhenxun.models.ban_console import BanConsole
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.services.cache.cache_containers import CacheDict
from zhenxun.services.cache.runtime_cache import BanMemoryCache
from zhenxun.services.db_context import DB_TIMEOUT_SECONDS
from zhenxun.services.log import logger
@@ -16,6 +14,7 @@ from zhenxun.utils.enum import PluginType
from zhenxun.utils.utils import EntityIDs, get_entity_ids
from .config import LOGGER_COMMAND, WARNING_THRESHOLD
from .context import PermissionContext
from .exception import SkipPluginException
from .utils import freq
@@ -25,37 +24,6 @@ Config.add_plugin_config(
"才不会给你发消息.",
help="对被ban用户发送的消息",
)
BAN_CACHE_TTL = 2
BAN_CACHE_TTL_POSITIVE = 30
BAN_CACHE_TTL_NEGATIVE = 5
BAN_CACHE = (
CacheDict("AUTH_BAN_CACHE", expire=0)
if max(BAN_CACHE_TTL_POSITIVE, BAN_CACHE_TTL_NEGATIVE) > 0
else None
)
def _ban_cache_key(user_id: str | None, group_id: str | None) -> str:
return f"{user_id or ''}:{group_id or ''}"
def _ban_cache_get(key: str) -> int | None:
if not BAN_CACHE:
return None
try:
return BAN_CACHE[key]
except KeyError:
return None
def _ban_cache_set(key: str, value: int) -> None:
if not BAN_CACHE:
return
ttl = BAN_CACHE_TTL_POSITIVE if value else BAN_CACHE_TTL_NEGATIVE
if ttl <= 0:
return
BAN_CACHE.set(key, value, expire=ttl)
async def calculate_ban_time(ban_record: BanConsole | None) -> int:
@@ -214,6 +182,7 @@ async def auth_ban(
session: Uninfo,
plugin: PluginInfo,
*,
context: PermissionContext | None = None,
entity: EntityIDs | None = None,
is_superuser: bool = False,
) -> None:
@@ -229,28 +198,18 @@ async def auth_ban(
return
if not matcher.plugin_name:
return
if context is not None:
entity = context.entity
is_superuser = context.is_superuser
if entity is None:
entity = get_entity_ids(session)
if is_superuser:
return
if entity.group_id:
try:
await asyncio.wait_for(
group_handle(entity.group_id), timeout=DB_TIMEOUT_SECONDS
)
except asyncio.TimeoutError:
logger.error(f"群组ban检查超时: {entity.group_id}", LOGGER_COMMAND)
# 超时时不阻塞，继续执行
await group_handle(entity.group_id)
if entity.user_id:
try:
await asyncio.wait_for(
user_handle(plugin, entity, session),
timeout=DB_TIMEOUT_SECONDS,
)
except asyncio.TimeoutError:
logger.error(f"用户ban检查超时: {entity.user_id}", LOGGER_COMMAND)
# 超时时不阻塞，继续执行
await user_handle(plugin, entity, session)
finally:
# 记录总执行时间
elapsed = time.time() - start_time
zhenxun/builtin_plugins/hooks/auth/auth_bot.py
+6
View File
@@ -7,6 +7,7 @@ from zhenxun.services.log import logger
from zhenxun.utils.common_utils import CommonUtils
from .config import LOGGER_COMMAND, WARNING_THRESHOLD
from .context import PermissionContext
from .exception import SkipPluginException
@@ -16,6 +17,8 @@ async def auth_bot(
bot_data: BotConsole | BotSnapshot | None = None,
skip_fetch: bool = False,
allow_sleep_bypass: bool = False,
*,
context: PermissionContext | None = None,
):
"""bot层面的权限检查
@@ -30,6 +33,9 @@ async def auth_bot(
start_time = time.time()
try:
if context is not None:
bot_id = context.event.bot_id
bot_data = context.bot_data
bot: BotConsole | BotSnapshot | None = bot_data
if bot is None and not skip_fetch:
bot = await BotMemoryCache.get(bot_id)
zhenxun/builtin_plugins/hooks/auth/auth_cost.py
+8 -1
View File
@@ -7,13 +7,18 @@ from zhenxun.models.user_console import UserConsole
from zhenxun.services.log import logger
from .config import LOGGER_COMMAND, WARNING_THRESHOLD
from .context import PermissionContext
from .exception import SkipPluginException
DEFAULT_GOLD = 100
async def auth_cost(
user: UserConsole | None, plugin: PluginInfo, session: Uninfo
user: UserConsole | None,
plugin: PluginInfo,
session: Uninfo,
*,
context: PermissionContext | None = None,
) -> int:
"""检测是否满足金币条件
@@ -28,6 +33,8 @@ async def auth_cost(
start_time = time.time()
try:
if context is not None and user is None:
user = context.user
user_gold = user.gold if user else DEFAULT_GOLD
if user_gold < plugin.cost_gold:
"""插件消耗金币不足"""
zhenxun/builtin_plugins/hooks/auth/auth_group.py
+8
View File
@@ -7,6 +7,7 @@ from zhenxun.services.cache.runtime_cache import GroupSnapshot
from zhenxun.services.log import logger
from .config import LOGGER_COMMAND, WARNING_THRESHOLD, SwitchEnum
from .context import PermissionContext
from .exception import SkipPluginException
_GROUP_WAKE_PATTERN = re.compile(r"^醒来$", re.IGNORECASE)
@@ -34,6 +35,8 @@ async def auth_group(
group: GroupConsole | GroupSnapshot | None,
text: str | None,
group_id: str | None,
*,
context: PermissionContext | None = None,
):
"""群黑名单检测 群总开关检测
@@ -42,6 +45,11 @@ async def auth_group(
group: GroupConsole
message: UniMsg
"""
if context is not None:
group = context.group or group
text = context.plain_text
group_id = context.group_id
if not group_id:
return
zhenxun/builtin_plugins/hooks/auth/auth_limit.py
+31 -20
View File
@@ -19,9 +19,10 @@ from zhenxun.utils.limiters import CountLimiter, FreqLimiter, UserBlockLimiter
from zhenxun.utils.manager.priority_manager import PriorityLifecycle
from zhenxun.utils.message import MessageUtils
from zhenxun.utils.time_utils import TimeUtils
from zhenxun.utils.utils import get_entity_ids
from zhenxun.utils.utils import EntityIDs, get_entity_ids
from .config import LOGGER_COMMAND, WARNING_THRESHOLD
from .context import PermissionContext
from .exception import SkipPluginException
driver = nonebot.get_driver()
@@ -31,7 +32,7 @@ _LIMIT_NOTICE_LIMITER = FreqLimiter(_LIMIT_NOTICE_CD)
_LIMIT_NOTICE_TASKS: set[asyncio.Task] = set()
@PriorityLifecycle.on_startup(priority=5)
@PriorityLifecycle.on_startup(priority=7)
async def _():
"""初始化限制"""
await LimitManager.init_limit()
@@ -117,6 +118,7 @@ class LimitManager:
cls.cd_limit = {}
cls.block_limit = {}
cls.count_limit = {}
cls.module_limit_cache.clear()
# 添加新数据
for limit in limit_list:
cls.add_limit(limit)
@@ -137,22 +139,22 @@ class LimitManager:
"""
if limit.module not in cls.add_module:
cls.add_module.append(limit.module)
if limit.limit_type == PluginLimitType.BLOCK:
cls.block_limit[limit.module] = Limit(
limit=limit, limiter=UserBlockLimiter()
)
elif limit.limit_type == PluginLimitType.CD:
cd_value = int(limit.cd or 0)
cls.cd_limit[limit.module] = Limit(
limit=limit, limiter=FreqLimiter(cd_value)
)
elif limit.limit_type == PluginLimitType.COUNT:
max_count = int(limit.max_count or 0)
if max_count <= 0:
return
cls.count_limit[limit.module] = Limit(
limit=limit, limiter=CountLimiter(max_count)
)
if limit.limit_type == PluginLimitType.BLOCK:
cls.block_limit[limit.module] = Limit(
limit=limit, limiter=UserBlockLimiter()
)
elif limit.limit_type == PluginLimitType.CD:
cd_value = int(limit.cd or 0)
cls.cd_limit[limit.module] = Limit(
limit=limit, limiter=FreqLimiter(cd_value)
)
elif limit.limit_type == PluginLimitType.COUNT:
max_count = int(limit.max_count or 0)
if max_count <= 0:
return
cls.count_limit[limit.module] = Limit(
limit=limit, limiter=CountLimiter(max_count)
)
@classmethod
def unblock(
@@ -322,14 +324,23 @@ class LimitManager:
limiter.increase(key_type)
async def auth_limit(plugin: PluginInfo, session: Uninfo):
async def auth_limit(
plugin: PluginInfo,
session: Uninfo,
*,
context: PermissionContext | None = None,
entity: EntityIDs | None = None,
):
"""插件限制
参数:
plugin: PluginInfo
session: Uninfo
"""
entity = get_entity_ids(session)
if context is not None:
entity = context.entity
if entity is None:
entity = get_entity_ids(session)
try:
await asyncio.wait_for(
LimitManager.check(
zhenxun/builtin_plugins/hooks/auth/auth_plugin.py
+16 -7
View File
@@ -10,6 +10,7 @@ from zhenxun.services.log import logger
from zhenxun.utils.enum import BlockType
from .config import LOGGER_COMMAND, WARNING_THRESHOLD
from .context import PermissionContext
from .exception import IsSuperuserException, SkipPluginException
from .utils import freq, is_poke
@@ -107,11 +108,16 @@ class GroupCheck:
class PluginCheck:
def __init__(
self, group: GroupConsole | GroupSnapshot | None, session: Uninfo, is_poke: bool
self,
group: GroupConsole | GroupSnapshot | None,
session: Uninfo,
is_poke: bool,
user_id: str | None,
):
self.session = session
self.is_poke = is_poke
self.group_data = group
self.user_id = user_id or session.user.id
self.group_id = None
if group:
self.group_id = group.group_id
@@ -126,13 +132,11 @@ class PluginCheck:
IgnoredException: 忽略插件
"""
if plugin.block_type == BlockType.PRIVATE:
should_tip = freq.is_send_limit_message(
plugin, self.session.user.id, self.is_poke
)
should_tip = freq.is_send_limit_message(plugin, self.user_id, self.is_poke)
raise SkipPluginException(
f"{plugin.name}({plugin.module}) 该插件在私聊中已被禁用...",
tip_message="该功能在私聊中已被禁用..." if should_tip else None,
tip_check_tag=self.session.user.id if should_tip else None,
tip_check_tag=self.user_id if should_tip else None,
tip_background=should_tip,
)
@@ -153,7 +157,7 @@ class PluginCheck:
if self.group_data and self.group_data.is_super:
raise IsSuperuserException()
sid = self.group_id or self.session.user.id
sid = self.group_id or self.user_id
should_tip = freq.is_send_limit_message(plugin, sid, self.is_poke)
raise SkipPluginException(
f"{plugin.name}({plugin.module}) 全局未开启此功能...",
@@ -176,7 +180,9 @@ async def auth_plugin(
session: Uninfo,
event: Event,
*,
context: PermissionContext | None = None,
skip_group_block: bool = False,
user_id: str | None = None,
):
"""插件状态
@@ -187,8 +193,11 @@ async def auth_plugin(
"""
start_time = time.time()
try:
if context is not None:
group = context.group or group
user_id = context.user_id
is_poke_event = is_poke(event)
user_check = PluginCheck(group, session, is_poke_event)
user_check = PluginCheck(group, session, is_poke_event, user_id)
if group:
block_set, super_block_set = _get_group_block_sets(group)
zhenxun/builtin_plugins/hooks/auth/bot_filter.py
+13 -4
View File
@@ -3,6 +3,7 @@ from nonebot_plugin_uninfo import Uninfo
from zhenxun.configs.config import Config
from .context import PermissionContext
from .exception import SkipPluginException
Config.add_plugin_config(
@@ -15,7 +16,12 @@ Config.add_plugin_config(
)
def bot_filter(session: Uninfo):
def bot_filter(
session: Uninfo,
*,
context: PermissionContext | None = None,
user_id: str | None = None,
):
"""过滤bot调用bot
参数:
@@ -26,10 +32,13 @@ def bot_filter(session: Uninfo):
"""
if not Config.get_config("hook", "FILTER_BOT"):
return
if context is not None:
user_id = context.user_id
bot_ids = list(nonebot.get_bots().keys())
if session.user.id == session.self_id:
checked_user_id = user_id or session.user.id
if checked_user_id == session.self_id:
return
if session.user.id in bot_ids:
if checked_user_id in bot_ids:
raise SkipPluginException(
f"bot:{session.self_id} 尝试调用 bot:{session.user.id}"
f"bot:{session.self_id} 尝试调用 bot:{checked_user_id}"
)
zhenxun/builtin_plugins/hooks/auth/context.py
+321
View File
@@ -0,0 +1,321 @@
from __future__ import annotations
import asyncio
import contextlib
from dataclasses import dataclass, field
from typing import Any
from nonebot.adapters import Bot, Event
from nonebot_plugin_alconna import UniMsg
from nonebot_plugin_uninfo import Uninfo
from zhenxun.services.cache.cache_containers import CacheDict
from zhenxun.utils.platform import PlatformUtils
from zhenxun.utils.utils import EntityIDs, get_entity_ids
AUTH_EVENT_CACHE_TTL = 5
STATE_EVENT_CONTEXT = "_zx_event_context"
STATE_PERMISSION_CONTEXT = "_zx_permission_context"
STATE_ENTITY = "_zx_entity"
STATE_EVENT_CACHE = "_zx_event_cache"
STATE_PLAIN_TEXT = "_zx_plain_text"
STATE_ROUTE_MODULES = "_zx_route_modules"
STATE_IS_SUPERUSER = "_zx_is_superuser"
STATE_PERMISSION_SIDE_EFFECTS = "_zx_permission_side_effects"
EVENT_CACHE_PERMISSION_SIDE_EFFECTS = "permission_side_effects"
EVENT_CACHE = (
CacheDict("AUTH_EVENT_CACHE", expire=AUTH_EVENT_CACHE_TTL)
if AUTH_EVENT_CACHE_TTL > 0
else None
)
@dataclass
class EventContext:
bot_id: str
platform: str
event_type: str
message_id: str | int | None
entity: EntityIDs
plain_text: str = ""
route_modules: set[str] = field(default_factory=set)
route_modules_loaded: bool = False
is_superuser: bool = False
event_cache: dict[str, Any] | None = None
@property
def user_id(self) -> str:
return self.entity.user_id
@property
def group_id(self) -> str | None:
return self.entity.group_id
@property
def channel_id(self) -> str | None:
return self.entity.channel_id
@dataclass
class PermissionSideEffectCache:
auth_results: dict[str, tuple[bool, str | None]] = field(default_factory=dict)
module_locks: dict[str, asyncio.Lock] = field(default_factory=dict)
def lock_for(self, module: str) -> asyncio.Lock:
lock = self.module_locks.get(module)
if lock is None:
lock = asyncio.Lock()
self.module_locks[module] = lock
return lock
@dataclass
class PermissionContext:
event: EventContext
module: str
plugin: Any = None
user: Any = None
group: Any = None
bot_data: Any = None
admin_levels: Any = None
@property
def entity(self) -> EntityIDs:
return self.event.entity
@property
def user_id(self) -> str:
return self.event.user_id
@property
def group_id(self) -> str | None:
return self.event.group_id
@property
def channel_id(self) -> str | None:
return self.event.channel_id
@property
def plain_text(self) -> str:
return self.event.plain_text
@property
def is_superuser(self) -> bool:
return self.event.is_superuser
def resolve_actor_user_id(event: Event, fallback_user_id: str | None) -> str:
"""优先使用事件发起者 ID，避免 notice 场景 session.user 指向 bot 自身。"""
event_user_id = getattr(event, "user_id", None)
if event_user_id is None:
return fallback_user_id or ""
resolved = str(event_user_id)
return resolved or fallback_user_id or ""
def resolve_event_group_id(event: Event, fallback_group_id: str | None) -> str | None:
"""notice 场景 session.group 可能缺失，回退到事件上的 group_id。"""
event_group_id = getattr(event, "group_id", None)
if event_group_id is None:
return fallback_group_id
resolved = str(event_group_id)
return resolved or fallback_group_id
def resolve_event_channel_id(
event: Event, fallback_channel_id: str | None
) -> str | None:
"""频道场景回退到事件上的 channel_id。"""
event_channel_id = getattr(event, "channel_id", None)
if event_channel_id is None:
return fallback_channel_id
resolved = str(event_channel_id)
return resolved or fallback_channel_id
def resolve_entity_ids(event: Event, session: Uninfo) -> EntityIDs:
entity = get_entity_ids(session)
entity.user_id = resolve_actor_user_id(event, entity.user_id)
entity.group_id = resolve_event_group_id(event, entity.group_id)
entity.channel_id = resolve_event_channel_id(event, entity.channel_id)
return entity
def extract_plain_text(message: UniMsg | None, event: Event) -> str:
if message is not None:
with contextlib.suppress(Exception):
return message.extract_plain_text()
with contextlib.suppress(Exception):
plain = event.get_plaintext()
if plain:
return plain.strip()
return ""
def _event_message_id(event: Event) -> str | int | None:
msg_id = getattr(event, "message_id", None)
if msg_id is None:
msg_id = getattr(event, "id", None)
return msg_id
def event_cache_key(
event: Event,
*,
bot_id: str,
platform: str,
entity: EntityIDs,
) -> str:
msg_id = _event_message_id(event)
if msg_id is None:
msg_id = id(event)
group_id = entity.group_id or ""
channel_id = entity.channel_id or ""
return f"{platform}:{bot_id}:{entity.user_id}:{group_id}:{channel_id}:{msg_id}"
def get_event_cache(
event: Event,
*,
bot_id: str,
platform: str,
entity: EntityIDs,
) -> dict[str, Any] | None:
if not EVENT_CACHE:
return None
key = event_cache_key(event, bot_id=bot_id, platform=platform, entity=entity)
try:
return EVENT_CACHE[key]
except KeyError:
cache: dict[str, Any] = {}
EVENT_CACHE[key] = cache
return cache
def _sync_context_state(state: dict[str, Any], context: EventContext) -> None:
state[STATE_EVENT_CONTEXT] = context
state[STATE_ENTITY] = context.entity
state[STATE_EVENT_CACHE] = context.event_cache
state[STATE_PLAIN_TEXT] = context.plain_text
state[STATE_ROUTE_MODULES] = context.route_modules
state[STATE_IS_SUPERUSER] = context.is_superuser
get_permission_side_effect_cache(state=state, event_cache=context.event_cache)
def get_permission_side_effect_cache(
*,
state: dict[str, Any] | None = None,
event_cache: dict[str, Any] | None = None,
) -> PermissionSideEffectCache:
side_effects = None
if state is not None:
side_effects = state.get(STATE_PERMISSION_SIDE_EFFECTS)
if (
not isinstance(side_effects, PermissionSideEffectCache)
and event_cache is not None
):
side_effects = event_cache.get(EVENT_CACHE_PERMISSION_SIDE_EFFECTS)
if not isinstance(side_effects, PermissionSideEffectCache):
side_effects = PermissionSideEffectCache()
if state is not None:
state[STATE_PERMISSION_SIDE_EFFECTS] = side_effects
if event_cache is not None:
event_cache[EVENT_CACHE_PERMISSION_SIDE_EFFECTS] = side_effects
return side_effects
def get_event_context(state: dict[str, Any] | None) -> EventContext | None:
if state is None:
return None
context = state.get(STATE_EVENT_CONTEXT)
return context if isinstance(context, EventContext) else None
def get_or_create_event_context(
bot: Bot,
event: Event,
session: Uninfo,
state: dict[str, Any],
*,
message: UniMsg | None = None,
) -> EventContext:
context = get_event_context(state)
if context is not None:
_sync_context_state(state, context)
return context
entity = state.get(STATE_ENTITY)
if not isinstance(entity, EntityIDs):
entity = resolve_entity_ids(event, session)
platform = PlatformUtils.get_platform(session)
bot_id = str(bot.self_id)
event_cache = state.get(STATE_EVENT_CACHE)
if not isinstance(event_cache, dict):
event_cache = get_event_cache(
event,
bot_id=bot_id,
platform=platform,
entity=entity,
)
text = state.get(STATE_PLAIN_TEXT)
if not isinstance(text, str):
cached_text = event_cache.get("plain_text") if event_cache is not None else None
text = (
cached_text
if isinstance(cached_text, str)
else extract_plain_text(message, event)
)
if event_cache is not None:
event_cache["plain_text"] = text
route_modules_loaded = STATE_ROUTE_MODULES in state
route_modules = state.get(STATE_ROUTE_MODULES)
if not isinstance(route_modules, set):
cached_routes = (
event_cache.get("route_modules") if event_cache is not None else None
)
route_modules = cached_routes if isinstance(cached_routes, set) else set()
route_modules_loaded = isinstance(cached_routes, set)
is_superuser = state.get(STATE_IS_SUPERUSER)
if not isinstance(is_superuser, bool):
is_superuser = entity.user_id in bot.config.superusers
context = EventContext(
bot_id=bot_id,
platform=platform,
event_type=event.get_type(),
message_id=_event_message_id(event),
entity=entity,
plain_text=text,
route_modules=route_modules,
route_modules_loaded=route_modules_loaded,
is_superuser=is_superuser,
event_cache=event_cache,
)
_sync_context_state(state, context)
return context
def set_route_modules(
state: dict[str, Any] | None,
context: EventContext,
route_modules: set[str],
) -> None:
context.route_modules = route_modules
context.route_modules_loaded = True
if context.event_cache is not None:
context.event_cache["route_modules"] = route_modules
if state is not None:
_sync_context_state(state, context)
def store_permission_context(
state: dict[str, Any] | None, context: PermissionContext
) -> None:
if state is not None:
state[STATE_PERMISSION_CONTEXT] = context
zhenxun/builtin_plugins/hooks/auth_checker.py
+427 -181
View File
@@ -1,7 +1,7 @@
import asyncio
from collections.abc import Awaitable, Callable
import contextlib
import os
import importlib
import re
import time
from typing import cast
@@ -11,7 +11,6 @@ from nonebot.adapters import Bot, Event
from nonebot.exception import IgnoredException
from nonebot.matcher import Matcher
import nonebot.message as nb_message
from nonebot_plugin_alconna import UniMsg
from nonebot_plugin_uninfo import Uninfo
from zhenxun.configs.utils import PluginExtraData
@@ -33,7 +32,6 @@ from zhenxun.services.message_load import is_overloaded
from zhenxun.utils.enum import BlockType, GoldHandle, PluginType
from zhenxun.utils.exception import InsufficientGold
from zhenxun.utils.platform import PlatformUtils
from zhenxun.utils.utils import get_entity_ids
from .auth.auth_admin import auth_admin
from .auth.auth_ban import auth_ban
@@ -44,6 +42,16 @@ from .auth.auth_limit import LimitManager, auth_limit
from .auth.auth_plugin import auth_plugin
from .auth.bot_filter import bot_filter
from .auth.config import LOGGER_COMMAND, WARNING_THRESHOLD
from .auth.context import (
EVENT_CACHE,
STATE_PLAIN_TEXT,
EventContext,
PermissionContext,
get_event_context,
get_permission_side_effect_cache,
set_route_modules,
store_permission_context,
)
from .auth.exception import (
IsSuperuserException,
PermissionExemption,
@@ -53,7 +61,6 @@ from .auth.utils import send_message
AUTH_HOOKS_CONCURRENCY_LIMIT = 5
AUTH_DB_CONCURRENCY_LIMIT = 6
AUTH_EVENT_CACHE_TTL = 5 # 增加到5秒，减少缓存抖动
# 超时设置（秒）
@@ -74,13 +81,6 @@ CIRCUIT_RESET_TIME = 300 # 5分钟
HOOKS_CONCURRENCY_LIMIT = AUTH_HOOKS_CONCURRENCY_LIMIT
DB_CONCURRENCY_LIMIT = AUTH_DB_CONCURRENCY_LIMIT
EVENT_CACHE_TTL = AUTH_EVENT_CACHE_TTL
EVENT_CACHE = (
CacheDict("AUTH_EVENT_CACHE", expire=EVENT_CACHE_TTL)
if EVENT_CACHE_TTL > 0
else None
)
# 路由索引缓存
_ROUTE_INDEX_LOCK = asyncio.Lock()
_ROUTE_INDEX_READY = False
@@ -91,21 +91,17 @@ MATCHER_ROUTE_PREFILTER_TTL = 2
PREFILTER_STATS_LOG_INTERVAL = 10.0
CACHE_SWEEP_INTERVAL = 1.0
CPU_COUNT = os.cpu_count() or 4
COMMAND_MATCHER_CONCURRENCY = max(8, min(48, CPU_COUNT * 4))
HEAVY_COMMAND_CONCURRENCY = max(1, min(3, CPU_COUNT // 2))
HEAVY_COMMAND_MODULES = frozenset({"shop", "sign_in"})
# 全局信号量与计数器
HOOKS_ACTIVE_COUNT = 0
HOOKS_SEMAPHORE = asyncio.Semaphore(HOOKS_CONCURRENCY_LIMIT)
COMMAND_MATCHER_SEMAPHORE = asyncio.Semaphore(COMMAND_MATCHER_CONCURRENCY)
HEAVY_COMMAND_SEMAPHORE = asyncio.Semaphore(HEAVY_COMMAND_CONCURRENCY)
DB_SEMAPHORE = asyncio.Semaphore(DB_CONCURRENCY_LIMIT)
DB_ACTIVE_COUNT = 0
_CHECK_MATCHER_PATCHED = False
_ORIGINAL_CHECK_AND_RUN_MATCHER: Callable[..., Awaitable[None]] | None = None
_HANDLE_EVENT_PATCHED = False
_ORIGINAL_HANDLE_EVENT: Callable[..., Awaitable[None]] | None = None
_ORIGINAL_ADAPTER_HANDLE_EVENTS: dict[object, object] = {}
_MATCHER_COMMAND_TYPE_CACHE: dict[type[Matcher], bool] = {}
_MATCHER_COMMAND_LITERAL_CACHE: dict[type[Matcher], tuple[str, ...] | None] = {}
_MATCHER_ALCONNA_SHORTCUT_CACHE: dict[type[Matcher], bool] = {}
@@ -115,6 +111,10 @@ _CHECK_MATCHER_ROUTE_CACHE = CacheDict(
_PREFILTER_STATS = {
"checked": 0,
"skipped": 0,
"before_task_checked": 0,
"before_task_skipped": 0,
"inside_task_checked": 0,
"inside_task_skipped": 0,
"type_miss": 0,
"route_miss": 0,
"command_miss": 0,
@@ -163,33 +163,6 @@ def _debug_log(message: str, *args, **kwargs) -> None:
logger.debug(message, *args, **kwargs)
def _event_cache_key(event: Event, session: Uninfo, entity) -> str:
msg_id = getattr(event, "message_id", None)
if msg_id is None:
msg_id = getattr(event, "id", None)
if msg_id is None:
msg_id = id(event)
platform = PlatformUtils.get_platform(session)
group_id = entity.group_id or ""
channel_id = entity.channel_id or ""
return (
f"{platform}:{session.self_id}:{entity.user_id}:"
f"{group_id}:{channel_id}:{msg_id}"
)
def _get_event_cache(event: Event, session: Uninfo, entity):
if not EVENT_CACHE:
return None
key = _event_cache_key(event, session, entity)
try:
return EVENT_CACHE[key]
except KeyError:
cache = {}
EVENT_CACHE[key] = cache
return cache
def _normalize_command(command: str) -> str:
text = command.strip()
if not text:
@@ -488,6 +461,9 @@ def _event_plain_text(event: Event) -> str:
def _state_plain_text(state: dict | None) -> str:
if state is None:
return ""
context = get_event_context(state)
if context is not None:
return context.plain_text.strip()
text = state.get("_zx_plain_text")
if isinstance(text, str):
return text.strip()
@@ -496,6 +472,9 @@ def _state_plain_text(state: dict | None) -> str:
def _get_route_modules_for_event(event: Event, state: dict | None = None) -> set[str]:
if state is not None:
context = get_event_context(state)
if context is not None and context.route_modules_loaded:
return context.route_modules
route_modules = state.get("_zx_route_modules")
if isinstance(route_modules, set):
return route_modules
@@ -506,15 +485,67 @@ def _get_route_modules_for_event(event: Event, state: dict | None = None) -> set
route_modules = _match_route_modules(_event_plain_text(event))
_CHECK_MATCHER_ROUTE_CACHE[key] = route_modules
if state is not None:
state["_zx_route_modules"] = route_modules
context = get_event_context(state)
if context is not None:
set_route_modules(state, context, route_modules)
else:
state["_zx_route_modules"] = route_modules
return route_modules
def _record_prefilter_stats(skipped: bool, reason: str | None) -> None:
def _prepare_handle_event_state(event: Event, state: dict) -> None:
get_permission_side_effect_cache(state=state)
if event.get_type() != "message":
return
if _state_plain_text(state):
return
text = _event_plain_text(event)
if text:
state[STATE_PLAIN_TEXT] = text
def _build_matcher_state(base_state: dict) -> dict:
get_permission_side_effect_cache(state=base_state)
matcher_state = base_state.copy()
get_permission_side_effect_cache(state=matcher_state)
return matcher_state
async def _run_selected_matcher(
matcher: type[Matcher],
bot: Bot,
event: Event,
state: dict,
stack,
dependency_cache,
) -> None:
await nb_message.check_and_run_matcher(
matcher,
bot,
event,
state,
stack,
dependency_cache,
)
def _record_prefilter_stats(
skipped: bool,
reason: str | None,
stage: str = "inside_task",
) -> None:
global _PREFILTER_LAST_LOG
_PREFILTER_STATS["checked"] += 1
if skipped:
_PREFILTER_STATS["skipped"] += 1
if stage == "before_task":
_PREFILTER_STATS["before_task_checked"] += 1
if skipped:
_PREFILTER_STATS["before_task_skipped"] += 1
else:
_PREFILTER_STATS["inside_task_checked"] += 1
if skipped:
_PREFILTER_STATS["inside_task_skipped"] += 1
if reason == "type_miss":
_PREFILTER_STATS["type_miss"] += 1
elif reason == "route_miss":
@@ -537,6 +568,10 @@ def _record_prefilter_stats(skipped: bool, reason: str | None) -> None:
"matcher prefilter stats: "
f"checked={_PREFILTER_STATS['checked']} "
f"skipped={_PREFILTER_STATS['skipped']} "
f"before_task={_PREFILTER_STATS['before_task_skipped']}/"
f"{_PREFILTER_STATS['before_task_checked']} "
f"inside_task={_PREFILTER_STATS['inside_task_skipped']}/"
f"{_PREFILTER_STATS['inside_task_checked']} "
f"type_miss={_PREFILTER_STATS['type_miss']} "
f"route_miss={_PREFILTER_STATS['route_miss']} "
f"command_miss={_PREFILTER_STATS['command_miss']} "
@@ -643,15 +678,6 @@ def _matcher_has_alconna_shortcuts(matcher_cls: type[Matcher]) -> bool:
return has_shortcuts
def _is_heavy_command_module(module: str) -> bool:
normalized = module.strip().lower()
if not normalized:
return False
if normalized in HEAVY_COMMAND_MODULES:
return True
return any(normalized.endswith(f".{name}") for name in HEAVY_COMMAND_MODULES)
async def _check_matcher_prefilter(
matcher_cls: type[Matcher], event: Event, state: dict | None = None
) -> tuple[bool, str | None]:
@@ -686,6 +712,17 @@ async def _check_matcher_prefilter(
if not module:
return False, None
command_matched = False
matcher_commands = _extract_matcher_command_literals(matcher_cls)
if matcher_commands:
for command in matcher_commands:
if _command_matches(text, command):
command_matched = True
break
else:
if not _matcher_has_alconna_shortcuts(matcher_cls):
return True, "command_miss"
ai_route_modules = _collect_ai_route_modules(event, state)
ai_route_heads = _collect_ai_route_heads(event, state)
if ai_route_modules and module not in ai_route_modules:
@@ -696,25 +733,85 @@ async def _check_matcher_prefilter(
await _ensure_route_index()
if module not in _ROUTE_MODULES_WITH_COMMANDS:
matcher_commands = _extract_matcher_command_literals(matcher_cls)
if matcher_commands:
for command in matcher_commands:
if _command_matches(text, command):
return False, None
if _matcher_has_alconna_shortcuts(matcher_cls):
return False, None
return True, "command_miss"
return False, None
route_modules = _get_route_modules_for_event(event, state)
if module not in route_modules:
if command_matched:
return False, None
if _matcher_has_alconna_shortcuts(matcher_cls):
return False, None
return True, "route_miss"
return False, None
_MATCHER_SEMAPHORE_TIMEOUT = 8.0
def _check_matcher_prefilter_before_task(
matcher_cls: type[Matcher], event: Event, state: dict | None = None
) -> tuple[bool, str | None]:
"""Conservative selector before creating matcher task.
This mirrors the async matcher prefilter but never performs IO or route-index
rebuild. If anything is uncertain, let the existing check_and_run_matcher
patch handle it inside the task.
"""
event_type = event.get_type()
matcher_type = getattr(matcher_cls, "type", "") or ""
if isinstance(matcher_type, str) and matcher_type and matcher_type != event_type:
return True, "type_miss"
if event_type != "message":
return False, None
if getattr(matcher_cls, "temp", False):
return False, None
if not _is_command_matcher_class(matcher_cls):
return False, None
text = _state_plain_text(state)
if not text:
text = _event_plain_text(event)
if state is not None and text:
state["_zx_plain_text"] = text
if not text:
return True, "empty_text"
module = _matcher_module_name(matcher_cls)
if not module:
return False, None
command_matched = False
matcher_commands = _extract_matcher_command_literals(matcher_cls)
has_alconna_shortcuts = _matcher_has_alconna_shortcuts(matcher_cls)
if matcher_commands:
for command in matcher_commands:
if _command_matches(text, command):
command_matched = True
break
else:
if not has_alconna_shortcuts:
return True, "command_miss"
ai_route_modules = _collect_ai_route_modules(event, state)
ai_route_heads = _collect_ai_route_heads(event, state)
if ai_route_modules and module not in ai_route_modules:
if not _matcher_matches_ai_route_heads(matcher_cls, ai_route_heads):
return True, "route_miss"
if not _ROUTE_INDEX_READY:
return False, None
if module not in _ROUTE_MODULES_WITH_COMMANDS:
return False, None
route_modules = _get_route_modules_for_event(event, state)
if module not in route_modules:
if command_matched or has_alconna_shortcuts:
return False, None
return True, "route_miss"
return False, None
_MAX_MATCHER_CACHE = 512
@@ -729,7 +826,7 @@ async def _patched_check_and_run_matcher(
skip, reason = await _check_matcher_prefilter(
Matcher, event, state if isinstance(state, dict) else None
)
_record_prefilter_stats(skip, reason)
_record_prefilter_stats(skip, reason, "inside_task")
if skip:
return
@@ -744,28 +841,6 @@ async def _patched_check_and_run_matcher(
"stack": stack,
"dependency_cache": dependency_cache,
}
if _is_command_matcher_class(Matcher):
module = _matcher_module_name(Matcher)
sem = (
HEAVY_COMMAND_SEMAPHORE
if _is_heavy_command_module(module)
else COMMAND_MATCHER_SEMAPHORE
)
try:
await asyncio.wait_for(sem.acquire(), timeout=_MATCHER_SEMAPHORE_TIMEOUT)
except asyncio.TimeoutError:
logger.warning(
f"matcher semaphore acquire timeout for {module}, "
"executing without concurrency limit",
LOGGER_COMMAND,
)
await original(**kwargs)
return
try:
await original(**kwargs)
finally:
sem.release()
return
await original(**kwargs)
@@ -788,27 +863,137 @@ def _uninstall_matcher_prefilter() -> None:
_ORIGINAL_CHECK_AND_RUN_MATCHER = None
def _get_message_text(
message: UniMsg | None,
event_cache: dict | None,
event: Event | None = None,
) -> str:
if event_cache is not None:
cached = event_cache.get("plain_text")
if isinstance(cached, str):
return cached
async def _patched_handle_event(bot: Bot, event: Event) -> None:
show_log = True
escape_tag = getattr(nb_message, "escape_tag")
logger_ = getattr(nb_message, "logger")
no_log_exception = getattr(nb_message, "NoLogException")
text = ""
if message is not None:
with contextlib.suppress(Exception):
text = message.extract_plain_text()
if not text and event is not None:
with contextlib.suppress(Exception):
text = (event.get_plaintext() or "").strip()
log_msg = f"<m>{escape_tag(bot.type)} {escape_tag(bot.self_id)}</m> | "
try:
log_msg += event.get_log_string()
except no_log_exception:
show_log = False
if show_log:
logger_.opt(colors=True).success(log_msg)
if event_cache is not None:
event_cache["plain_text"] = text
return text
state = {}
dependency_cache = {}
async_exit_stack = getattr(nb_message, "AsyncExitStack")
apply_event_preprocessors = getattr(nb_message, "_apply_event_preprocessors")
apply_event_postprocessors = getattr(nb_message, "_apply_event_postprocessors")
trie_rule = getattr(nb_message, "TrieRule")
matchers = getattr(nb_message, "matchers")
catch = getattr(nb_message, "catch")
stop_propagation = getattr(nb_message, "StopPropagation")
handle_exception = getattr(nb_message, "_handle_exception")
anyio_mod = getattr(nb_message, "anyio")
run_coro_with_shield = getattr(nb_message, "run_coro_with_shield")
async with async_exit_stack() as stack:
if not await apply_event_preprocessors(
bot=bot,
event=event,
state=state,
stack=stack,
dependency_cache=dependency_cache,
):
return
try:
trie_rule.get_value(bot, event, state)
except Exception as e:
logger_.opt(colors=True, exception=e).warning(
"Error while parsing command for event"
)
_prepare_handle_event_state(event, state)
break_flag = False
def _handle_stop_propagation(_exc_group) -> None:
nonlocal break_flag
break_flag = True
logger_.debug("Stop event propagation")
for priority in sorted(matchers.keys()):
if break_flag:
break
if show_log:
logger_.debug(f"Checking for matchers in priority {priority}...")
if not (priority_matchers := matchers[priority]):
continue
with catch(
{
stop_propagation: _handle_stop_propagation,
Exception: handle_exception(
"<r><bg #f8bbd0>Error when checking Matcher.</bg #f8bbd0></r>"
),
}
):
async with anyio_mod.create_task_group() as tg:
for matcher in priority_matchers:
skip, reason = _check_matcher_prefilter_before_task(
matcher,
event,
state,
)
_record_prefilter_stats(skip, reason, "before_task")
if skip:
continue
matcher_state = _build_matcher_state(state)
tg.start_soon(
run_coro_with_shield,
_run_selected_matcher(
matcher,
bot,
event,
matcher_state,
stack,
dependency_cache,
),
)
if show_log:
logger_.debug("Checking for matchers completed")
await apply_event_postprocessors(bot, event, state, stack, dependency_cache)
def _install_handle_event_selector() -> None:
global _HANDLE_EVENT_PATCHED, _ORIGINAL_HANDLE_EVENT
if _HANDLE_EVENT_PATCHED:
return
_ORIGINAL_HANDLE_EVENT = nb_message.handle_event
nb_message.handle_event = _patched_handle_event # type: ignore[assignment]
for module_name in (
"nonebot.adapters.onebot.v11.bot",
"nonebot.adapters.onebot.v12.bot",
"onebug.mixin.process",
):
with contextlib.suppress(Exception):
module = importlib.import_module(module_name)
current = getattr(module, "handle_event", None)
if current is not None:
_ORIGINAL_ADAPTER_HANDLE_EVENTS[module] = current
setattr(module, "handle_event", _patched_handle_event)
_HANDLE_EVENT_PATCHED = True
def _uninstall_handle_event_selector() -> None:
global _HANDLE_EVENT_PATCHED, _ORIGINAL_HANDLE_EVENT
if not _HANDLE_EVENT_PATCHED:
return
if _ORIGINAL_HANDLE_EVENT is not None:
nb_message.handle_event = _ORIGINAL_HANDLE_EVENT # type: ignore[assignment]
for module, original in list(_ORIGINAL_ADAPTER_HANDLE_EVENTS.items()):
with contextlib.suppress(Exception):
setattr(module, "handle_event", original)
_ORIGINAL_ADAPTER_HANDLE_EVENTS.clear()
_HANDLE_EVENT_PATCHED = False
_ORIGINAL_HANDLE_EVENT = None
async def _get_route_context(text: str, event_cache: dict | None) -> set[str]:
@@ -843,12 +1028,14 @@ async def start_auth_runtime_tasks() -> None:
global _CACHE_SWEEP_TASK
await _ensure_route_index()
_install_matcher_prefilter()
_install_handle_event_selector()
if _CACHE_SWEEP_TASK is None or _CACHE_SWEEP_TASK.done():
_CACHE_SWEEP_TASK = asyncio.create_task(_cache_sweep_loop())
async def stop_auth_runtime_tasks() -> None:
global _CACHE_SWEEP_TASK
_uninstall_handle_event_selector()
_uninstall_matcher_prefilter()
task = _CACHE_SWEEP_TASK
_CACHE_SWEEP_TASK = None
@@ -873,6 +1060,12 @@ async def _has_limits_cached(module: str, event_cache: dict | None) -> bool:
@contextlib.asynccontextmanager
async def _db_section():
global DB_ACTIVE_COUNT
if DB_SEMAPHORE.locked():
logger.warning(
"db semaphore saturated, allowing permission check to continue",
LOGGER_COMMAND,
)
raise PermissionExemption("db semaphore saturated, allow pass")
await DB_SEMAPHORE.acquire()
DB_ACTIVE_COUNT += 1
try:
@@ -918,7 +1111,9 @@ def _group_has_plugin_block(group, module: str) -> bool:
)
def _needs_auth_plugin(plugin: PluginInfo, group, entity) -> bool:
def _needs_auth_plugin(plugin: PluginInfo, context: PermissionContext) -> bool:
group = context.group
entity = context.entity
if plugin.block_type == BlockType.ALL and not plugin.status:
if group and getattr(group, "is_super", False):
return False
@@ -953,11 +1148,11 @@ async def _get_bot_data_cached(
async def _get_admin_levels_cached(
session: Uninfo, entity, event_cache
entity, event_cache
) -> tuple[tuple[LevelUserSnapshot | None, LevelUserSnapshot | None] | None, bool]:
if event_cache is not None and "admin_levels" in event_cache:
return event_cache.get("admin_levels"), event_cache.get("admin_timeout", False)
levels = await LevelUserMemoryCache.get_levels(session.user.id, entity.group_id)
levels = await LevelUserMemoryCache.get_levels(entity.user_id, entity.group_id)
if event_cache is not None:
event_cache["admin_levels"] = levels
event_cache["admin_timeout"] = False
@@ -1074,12 +1269,18 @@ async def get_plugin_and_user(
if user_id in user_cache:
user = user_cache[user_id]
else:
async with _db_section():
user = await _fetch_user_readonly(user_dao, user_id)
try:
async with _db_section():
user = await _fetch_user_readonly(user_dao, user_id)
except PermissionExemption:
user = None
user_cache[user_id] = user
else:
async with _db_section():
user = await _fetch_user_readonly(user_dao, user_id)
try:
async with _db_section():
user = await _fetch_user_readonly(user_dao, user_id)
except PermissionExemption:
user = None
return plugin, user
@@ -1089,7 +1290,7 @@ async def get_plugin_cost(
plugin: PluginInfo,
session: Uninfo,
*,
is_superuser: bool = False,
context: PermissionContext | None = None,
) -> int:
"""获取插件费用
@@ -1106,7 +1307,10 @@ async def get_plugin_cost(
返回:
int: 调用插件金币费用
"""
cost_gold = await with_timeout(auth_cost(user, plugin, session), name="auth_cost")
cost_gold = await with_timeout(
auth_cost(user, plugin, session, context=context), name="auth_cost"
)
is_superuser = context.is_superuser if context is not None else False
if is_superuser:
if plugin.plugin_type == PluginType.SUPERUSER:
raise IsSuperuserException()
@@ -1124,7 +1328,7 @@ async def reduce_gold(user_id: str, module: str, cost_gold: int, session: Uninfo
cost_gold: 消耗金币
session: Uninfo
"""
user_dao = DataAccess(UserConsole)
should_clear_cache = False
try:
await with_timeout(
UserConsole.reduce_gold(
@@ -1141,14 +1345,16 @@ async def reduce_gold(user_id: str, module: str, cost_gold: int, session: Uninfo
u.gold = 0
await u.save(update_fields=["gold"])
except asyncio.TimeoutError:
should_clear_cache = True
logger.error(
f"扣除金币超时，用户: {user_id}, 金币: {cost_gold}",
LOGGER_COMMAND,
session=session,
)
# 清除缓存，使下次查询时从数据库获取最新数据
await user_dao.clear_cache(user_id=user_id)
# 正常写入路径由 UserConsole.save() 统一失效缓存；超时状态不确定时兜底清理。
if should_clear_cache:
await DataAccess(UserConsole).clear_cache(user_id=user_id)
logger.debug(f"调用功能花费金币: {cost_gold}", LOGGER_COMMAND, session=session)
@@ -1174,16 +1380,15 @@ async def time_hook(coro, name, recorder: HookTraceRecorder | None = None):
async def _enter_hooks_section():
"""尝试获取全局信号量并更新计数器，超时则抛出 PermissionExemption。"""
"""尝试获取全局信号量并更新计数器，饱和时快速放行。"""
global HOOKS_ACTIVE_COUNT
try:
await asyncio.wait_for(HOOKS_SEMAPHORE.acquire(), timeout=TIMEOUT_SECONDS)
except asyncio.TimeoutError:
if HOOKS_SEMAPHORE.locked():
logger.warning(
"hooks semaphore acquire timeout, allowing pass",
"hooks semaphore saturated, allowing pass",
LOGGER_COMMAND,
)
raise PermissionExemption("hooks semaphore timeout, allow pass")
raise PermissionExemption("hooks semaphore saturated, allow pass")
await HOOKS_SEMAPHORE.acquire()
HOOKS_ACTIVE_COUNT += 1
@@ -1197,14 +1402,7 @@ async def _leave_hooks_section():
async def route_precheck(
matcher: Matcher,
event: Event,
session: Uninfo,
message: UniMsg | None,
*,
entity=None,
event_cache: dict | None = None,
text: str | None = None,
route_modules: set[str] | None = None,
context: EventContext,
) -> bool:
module = matcher.plugin_name or ""
if not module:
@@ -1213,19 +1411,20 @@ async def route_precheck(
return False
if not _is_command_matcher_class(type(matcher)):
return False
if entity is None:
entity = get_entity_ids(session)
if event_cache is None:
event_cache = _get_event_cache(event, session, entity)
if text is None:
text = _get_message_text(message, event_cache, event)
route_modules = context.route_modules if context.route_modules_loaded else None
if route_modules is None:
route_modules = await _get_route_context(text, event_cache)
route_modules = await _get_route_context(
context.plain_text,
context.event_cache,
)
set_route_modules(None, context, route_modules)
if module in _ROUTE_MODULES_WITH_COMMANDS and module not in route_modules:
if _matcher_has_alconna_shortcuts(type(matcher)):
return False
if event_cache is not None:
event_cache["route_skip"] = True
if context.event_cache is not None:
context.event_cache["route_skip"] = True
return True
return False
@@ -1235,14 +1434,10 @@ async def auth(
event: Event,
bot: Bot,
session: Uninfo,
message: UniMsg | None,
*,
context: EventContext,
skip_ban: bool = False,
entity=None,
event_cache: dict | None = None,
text: str | None = None,
route_modules: set[str] | None = None,
is_superuser: bool = False,
state: dict | None = None,
):
"""权限检查
@@ -1251,20 +1446,28 @@ async def auth(
event: Event
bot: bot
session: Uninfo
message: UniMsg
context: EventContext
"""
start_time = time.time()
cost_gold = 0
ignore_flag = False
if entity is None:
entity = get_entity_ids(session)
entity = context.entity
event_cache = context.event_cache
text = context.plain_text
is_superuser = context.is_superuser
route_modules = context.route_modules if context.route_modules_loaded else None
module = matcher.plugin_name or ""
is_command_matcher = _is_command_matcher_class(type(matcher))
if event_cache is None:
event_cache = _get_event_cache(event, session, entity)
auth_allowed = None
auth_result_cache = None
admin_checked_pre = False
permission_context: PermissionContext | None = None
side_effect_cache = get_permission_side_effect_cache(
state=state,
event_cache=event_cache,
)
side_effect_lock = None
entered_side_effect_lock = False
# 仅在慢请求时记录 hook 明细，避免热路径高频构造字符串
hook_recorder = HookTraceRecorder(start_time)
@@ -1278,14 +1481,17 @@ async def auth(
auth_allowed = True
return
if event_cache is not None:
auth_result_cache = event_cache.setdefault("auth_result", {})
cached_result = auth_result_cache.get(module)
if cached_result is not None:
allowed, reason = cached_result
if not allowed:
raise SkipPluginException(reason or "auth cached skip")
return
side_effect_lock = side_effect_cache.lock_for(module)
await side_effect_lock.acquire()
entered_side_effect_lock = True
auth_result_cache = side_effect_cache.auth_results
cached_result = auth_result_cache.get(module)
if cached_result is not None:
allowed, reason = cached_result
if not allowed:
raise SkipPluginException(reason or "auth cached skip")
return
if _is_hidden_plugin(matcher):
auth_allowed = True
@@ -1293,10 +1499,9 @@ async def auth(
if event_cache is not None and event_cache.get("ban_state") is True:
raise SkipPluginException("user or group banned (cached)")
if text is None:
text = _get_message_text(message, event_cache, event)
if route_modules is None:
route_modules = await _get_route_context(text, event_cache)
set_route_modules(state, context, route_modules)
route_skip_checks = (
is_command_matcher
and module in _ROUTE_MODULES_WITH_COMMANDS
@@ -1310,7 +1515,7 @@ async def auth(
auth_allowed = True
return
platform = PlatformUtils.get_platform(session)
platform = context.platform
# 获取插件和用户数据
plugin_user_start = time.time()
try:
@@ -1336,6 +1541,14 @@ async def auth(
auth_allowed = True
return
permission_context = PermissionContext(
event=context,
module=module,
plugin=plugin,
user=user,
)
store_permission_context(state, permission_context)
if not route_skip_checks and _needs_admin_check(plugin):
if plugin.plugin_type in {
PluginType.SUPERUSER,
@@ -1356,13 +1569,18 @@ async def auth(
admin_timeout = False
if event_cache is not None:
admin_levels, admin_timeout = await _get_admin_levels_cached(
session, entity, event_cache
entity, event_cache
)
permission_context.admin_levels = admin_levels
if admin_timeout:
hook_recorder.set("auth_admin", "timeout")
else:
admin_start = time.time()
await auth_admin(plugin, session, cached_levels=admin_levels)
await auth_admin(
plugin,
session,
context=permission_context,
)
hook_recorder.set(
"auth_admin", f"{time.time() - admin_start:.3f}s(pre)"
)
@@ -1386,8 +1604,7 @@ async def auth(
matcher,
session,
plugin,
entity=entity,
is_superuser=is_superuser,
context=permission_context,
)
hook_recorder.set("auth_ban", f"{time.time() - ban_start:.3f}s")
if event_cache is not None:
@@ -1407,7 +1624,7 @@ async def auth(
user,
plugin,
session,
is_superuser=is_superuser,
context=permission_context,
),
name="get_plugin_cost",
)
@@ -1421,7 +1638,7 @@ async def auth(
hook_recorder.set("cost_gold", "skipped")
# 执行 bot_filter
bot_filter(session)
bot_filter(session, context=permission_context)
group = await _get_group_cached(entity, event_cache)
@@ -1439,13 +1656,23 @@ async def auth(
and not route_skip_checks
):
admin_levels, admin_timeout = await _get_admin_levels_cached(
session, entity, event_cache
entity, event_cache
)
permission_context.group = group
permission_context.bot_data = bot_data
if admin_levels is not None:
permission_context.admin_levels = admin_levels
store_permission_context(state, permission_context)
# 并行执行所有 hook 检查，并记录执行时间
hooks_start = time.time()
allow_sleep_bypass = _is_bot_wake_command(module, text)
# 先进入 hooks 并行检查区域；饱和时快速放行，避免创建并积压协程。
await _enter_hooks_section()
entered_hooks = True
# 创建所有 hook 任务
hook_tasks = []
if event_cache is None:
@@ -1455,6 +1682,7 @@ async def auth(
plugin,
bot.self_id,
allow_sleep_bypass=allow_sleep_bypass,
context=permission_context,
),
"auth_bot",
hook_recorder,
@@ -1472,6 +1700,7 @@ async def auth(
bot_data=bot_data,
skip_fetch=True,
allow_sleep_bypass=allow_sleep_bypass,
context=permission_context,
),
"auth_bot",
hook_recorder,
@@ -1483,7 +1712,13 @@ async def auth(
else:
hook_tasks.append(
time_hook(
auth_group(plugin, group, text, entity.group_id),
auth_group(
plugin,
group,
text,
entity.group_id,
context=permission_context,
),
"auth_group",
hook_recorder,
)
@@ -1492,7 +1727,11 @@ async def auth(
if not route_skip_checks and plugin.admin_level and not admin_checked_pre:
if event_cache is None:
hook_tasks.append(
time_hook(auth_admin(plugin, session), "auth_admin", hook_recorder)
time_hook(
auth_admin(plugin, session, context=permission_context),
"auth_admin",
hook_recorder,
)
)
else:
if admin_timeout:
@@ -1500,7 +1739,11 @@ async def auth(
else:
hook_tasks.append(
time_hook(
auth_admin(plugin, session, cached_levels=admin_levels),
auth_admin(
plugin,
session,
context=permission_context,
),
"auth_admin",
hook_recorder,
)
@@ -1510,7 +1753,7 @@ async def auth(
if is_superuser:
hook_recorder.set("auth_plugin", "superuser")
elif not route_skip_checks and _needs_auth_plugin(plugin, group, entity):
elif not route_skip_checks and _needs_auth_plugin(plugin, permission_context):
hook_tasks.append(
time_hook(
auth_plugin(
@@ -1518,6 +1761,7 @@ async def auth(
group,
session,
event,
context=permission_context,
skip_group_block=is_superuser,
),
"auth_plugin",
@@ -1531,18 +1775,17 @@ async def auth(
has_limits = await _has_limits_cached(module, event_cache)
if has_limits:
hook_tasks.append(
time_hook(auth_limit(plugin, session), "auth_limit", hook_recorder)
time_hook(
auth_limit(plugin, session, context=permission_context),
"auth_limit",
hook_recorder,
)
)
else:
hook_recorder.set("auth_limit", "skipped")
else:
hook_recorder.set("auth_limit", "skipped")
if hook_tasks:
# 进入 hooks 并行检查区域（会在高并发时排队）
await _enter_hooks_section()
entered_hooks = True
# 使用 gather 并行执行所有 hook，但添加总体超时控制
try:
await with_timeout(
@@ -1599,6 +1842,9 @@ async def auth(
)
if auth_result_cache is not None and auth_allowed is not None:
auth_result_cache[module] = (auth_allowed, None)
if entered_side_effect_lock and side_effect_lock is not None:
with contextlib.suppress(Exception):
side_effect_lock.release()
# 扣除金币
if not ignore_flag and cost_gold > 0:
gold_start = time.time()
zhenxun/builtin_plugins/hooks/auth_hook.py
+47 -98
View File
@@ -1,4 +1,3 @@
import contextlib
import time
from nonebot import get_driver
@@ -12,14 +11,20 @@ from nonebot_plugin_uninfo import Uninfo
from zhenxun.services.cache.runtime_cache import is_cache_ready
from zhenxun.services.log import logger
from zhenxun.services.message_load import is_overloaded
from zhenxun.services.message_load import is_overloaded, mark_activity
from zhenxun.services.runtime_bootstrap import register_runtime_bootstrap
from zhenxun.utils.utils import get_entity_ids
from .auth.config import LOGGER_COMMAND
from .auth.context import (
get_event_context,
get_or_create_event_context,
resolve_actor_user_id,
resolve_event_channel_id,
resolve_event_group_id,
set_route_modules,
)
from .auth_checker import (
LimitManager,
_get_event_cache,
_get_route_context,
auth,
route_precheck,
@@ -41,17 +46,6 @@ async def _mark_bot_connected(bot: Bot):
_BOT_CONNECT_TS = time.time()
def _extract_plain_text(message: UniMsg | None, event: Event) -> str:
if message is not None:
with contextlib.suppress(Exception):
return message.extract_plain_text()
with contextlib.suppress(Exception):
plain = event.get_plaintext()
if plain:
return plain.strip()
return ""
@driver.on_startup
async def _start_auth_runtime_tasks():
await start_auth_runtime_tasks()
@@ -72,37 +66,9 @@ def _skip_auth_for_plugin(matcher: Matcher) -> bool:
return "chat_history" in module_name
def _resolve_actor_user_id(event: Event, fallback_user_id: str) -> str:
"""优先使用事件发起者ID，避免 notice 场景 session.user 指向 bot 自身。"""
event_user_id = getattr(event, "user_id", None)
if event_user_id is None:
return fallback_user_id
event_user_id = str(event_user_id)
return event_user_id or fallback_user_id
def _resolve_event_group_id(event: Event, fallback_group_id: str | None) -> str | None:
"""notice 场景 session.group 可能缺失，回退到事件上的 group_id。"""
event_group_id = getattr(event, "group_id", None)
if event_group_id is None:
return fallback_group_id
resolved = str(event_group_id)
return resolved or fallback_group_id
def _resolve_event_channel_id(
event: Event, fallback_channel_id: str | None
) -> str | None:
"""频道场景回退到事件上的 channel_id。"""
event_channel_id = getattr(event, "channel_id", None)
if event_channel_id is None:
return fallback_channel_id
resolved = str(event_channel_id)
return resolved or fallback_channel_id
@event_preprocessor
async def _drop_message_before_cache_ready(event: Event):
mark_activity()
if event.get_type() != "message":
return
if not is_cache_ready():
@@ -130,46 +96,22 @@ async def _auth_preprocessor(
return
start_time = time.time()
entity = state.get("_zx_entity")
if entity is None:
entity = get_entity_ids(session)
entity.user_id = _resolve_actor_user_id(event, entity.user_id)
entity.group_id = _resolve_event_group_id(event, entity.group_id)
entity.channel_id = _resolve_event_channel_id(event, entity.channel_id)
state["_zx_entity"] = entity
event_cache = state.get("_zx_event_cache")
if event_cache is None:
event_cache = _get_event_cache(event, session, entity)
state["_zx_event_cache"] = event_cache
text = state.get("_zx_plain_text")
if text is None:
text = _extract_plain_text(message, event)
state["_zx_plain_text"] = text
if event_cache is not None:
event_cache["plain_text"] = text
route_modules = state.get("_zx_route_modules")
if route_modules is None:
route_modules = await _get_route_context(text, event_cache)
state["_zx_route_modules"] = route_modules
is_superuser = state.get("_zx_is_superuser")
if is_superuser is None:
is_superuser = entity.user_id in bot.config.superusers
state["_zx_is_superuser"] = is_superuser
if await route_precheck(
matcher,
event_context = get_or_create_event_context(
bot,
event,
session,
message,
entity=entity,
event_cache=event_cache,
text=text,
route_modules=route_modules,
):
state,
message=message,
)
if not event_context.route_modules_loaded:
route_modules = await _get_route_context(
event_context.plain_text,
event_context.event_cache,
)
set_route_modules(state, event_context, route_modules)
if await route_precheck(matcher, event_context):
return
try:
@@ -178,13 +120,9 @@ async def _auth_preprocessor(
event,
bot,
session,
message,
context=event_context,
skip_ban=False,
entity=entity,
event_cache=event_cache,
text=text,
route_modules=route_modules,
is_superuser=is_superuser,
state=state,
)
except IgnoredException:
raise
@@ -203,16 +141,27 @@ async def _auth_preprocessor(
@run_postprocessor
async def _unblock_after_matcher(matcher: Matcher, session: Uninfo, event: Event):
user_id = _resolve_actor_user_id(event, session.user.id)
group_id = _resolve_event_group_id(event, None)
channel_id = _resolve_event_channel_id(event, None)
if session.group:
if session.group.parent:
group_id = session.group.parent.id
channel_id = session.group.id
else:
group_id = session.group.id
async def _unblock_after_matcher(
matcher: Matcher,
session: Uninfo,
event: Event,
state: T_State,
):
context = get_event_context(state)
if context is not None:
user_id = context.user_id
group_id = context.group_id
channel_id = context.channel_id
else:
user_id = resolve_actor_user_id(event, session.user.id)
group_id = resolve_event_group_id(event, None)
channel_id = resolve_event_channel_id(event, None)
if session.group:
if session.group.parent:
group_id = session.group.parent.id
channel_id = session.group.id
else:
group_id = session.group.id
if user_id and matcher.plugin:
module = matcher.plugin.name
LimitManager.unblock(module, user_id, group_id, channel_id)
zhenxun/builtin_plugins/hooks/chkdsk_hook.py
+24 -14
View File
@@ -16,13 +16,7 @@ from zhenxun.services.log import logger
from zhenxun.utils.enum import PluginType
from zhenxun.utils.message import MessageUtils
malicious_check_time = Config.get_config("hook", "MALICIOUS_CHECK_TIME")
malicious_ban_count = Config.get_config("hook", "MALICIOUS_BAN_COUNT")
if not malicious_check_time:
raise ValueError("模块: [hook], 配置项: [MALICIOUS_CHECK_TIME] 为空或小于0")
if not malicious_ban_count:
raise ValueError("模块: [hook], 配置项: [MALICIOUS_BAN_COUNT] 为空或小于0")
from .auth.context import resolve_actor_user_id, resolve_event_group_id
class BanCheckLimiter:
@@ -36,6 +30,10 @@ class BanCheckLimiter:
self.default_check_time = default_check_time
self.default_count = default_count
def configure(self, check_time: float, count: int) -> None:
self.default_check_time = check_time
self.default_count = count
def add(self, key: str | float):
if self.mint[key] == 1:
self.mtime[key] = time.time()
@@ -59,11 +57,22 @@ class BanCheckLimiter:
_blmt = BanCheckLimiter(
malicious_check_time,
malicious_ban_count,
5,
4,
)
def _get_positive_config(key: str, cast_type: type[int] | type[float]) -> int | float:
value = Config.get_config("hook", key)
try:
parsed_value = cast_type(value)
except (TypeError, ValueError) as e:
raise ValueError(f"模块: [hook], 配置项: [{key}] 不是有效数字") from e
if parsed_value <= 0:
raise ValueError(f"模块: [hook], 配置项: [{key}] 为空或小于0")
return parsed_value
# 恶意触发命令检测
@run_preprocessor
async def _(
@@ -88,11 +97,12 @@ async def _(
else:
return
user_id = session.id1
group_id = session.id3 or session.id2
malicious_ban_time = Config.get_config("hook", "MALICIOUS_BAN_TIME")
if not malicious_ban_time:
raise ValueError("模块: [hook], 配置项: [MALICIOUS_BAN_TIME] 为空或小于0")
user_id = resolve_actor_user_id(event, session.id1)
group_id = resolve_event_group_id(event, session.id3 or session.id2)
malicious_check_time = float(_get_positive_config("MALICIOUS_CHECK_TIME", float))
malicious_ban_count = int(_get_positive_config("MALICIOUS_BAN_COUNT", int))
malicious_ban_time = int(_get_positive_config("MALICIOUS_BAN_TIME", int))
_blmt.configure(malicious_check_time, malicious_ban_count)
if user_id and module:
if _blmt.check(f"{user_id}__{module}"):
await BanConsole.ban(
zhenxun/builtin_plugins/init/__init_cache.py
-1
View File
@@ -29,7 +29,6 @@ def register_cache_types():
GroupPluginSetting,
key_format="{group_id}_{plugin_name}_{key}",
)
CacheRegistry.register(CacheType.GROUP_PLUGIN_SETTINGS_VIEW, dict)
CacheRegistry.register(
CacheType.LEVEL, LevelUser, key_format="{user_id}_{group_id}"
)
zhenxun/builtin_plugins/init/init_plugin.py
+1 -1
View File
@@ -88,7 +88,7 @@ async def _handle_setting(
)
@PriorityLifecycle.on_startup(priority=5)
@PriorityLifecycle.on_startup(priority=4)
async def _():
"""
初始化插件数据配置
zhenxun/builtin_plugins/plugin_store/data_source.py
+26 -23
View File
@@ -3,12 +3,12 @@ from pathlib import Path
import random
import shutil
from aiocache import cached
import ujson as json
from zhenxun.builtin_plugins.plugin_store.models import StorePluginInfo
from zhenxun.configs.path_config import TEMP_PATH
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.services.cache.bounded_ttl import BoundedTTLCache
from zhenxun.services.log import logger
from zhenxun.services.plugin_init import PluginInitManager
from zhenxun.utils.enum import PluginType
@@ -26,6 +26,14 @@ from .config import (
)
from .exceptions import PluginStoreException
_PLUGIN_STORE_DATA_CACHE = BoundedTTLCache[
str, tuple[list[StorePluginInfo], list[StorePluginInfo]]
](
"PLUGIN_STORE_DATA",
ttl_seconds=60,
max_items=1,
)
def row_style(column: str, text: str) -> RowStyle:
"""被动技能文本风格
@@ -56,20 +64,9 @@ class StoreManager:
relative_parts = [part for part in plugin_info.module_path.split(".") if part]
relative_path = Path(*relative_parts) if relative_parts else Path(plugin_name)
path = BASE_PATH.parent / relative_path
if plugin_info.is_dir:
return path
return path.parent / f"{plugin_name}.py"
return path if plugin_info.is_dir else path.parent / f"{plugin_name}.py"
@classmethod
def _is_plugin_installed(
cls, plugin_info: StorePluginInfo, *, is_external: bool
) -> bool:
return cls._resolve_local_plugin_path(
plugin_info, is_external=is_external
).exists()
@classmethod
@cached(60)
async def get_data(cls) -> tuple[list[StorePluginInfo], list[StorePluginInfo]]:
"""获取插件信息数据
@@ -77,15 +74,22 @@ class StoreManager:
tuple[list[StorePluginInfo], list[StorePluginInfo]]:
原生插件信息数据，第三方插件信息数据
"""
cache_key = "plugins_json"
if cached_data := await _PLUGIN_STORE_DATA_CACHE.get(cache_key):
return cached_data
plugins = await RepoFileManager.get_file_content(
DEFAULT_GITHUB_URL, "plugins.json"
)
extra_plugins = await RepoFileManager.get_file_content(
EXTRA_GITHUB_URL, "plugins.json", "index"
)
return [StorePluginInfo(**plugin) for plugin in json.loads(plugins)], [
StorePluginInfo(**plugin) for plugin in json.loads(extra_plugins)
]
result = (
[StorePluginInfo(**plugin) for plugin in json.loads(plugins)],
[StorePluginInfo(**plugin) for plugin in json.loads(extra_plugins)],
)
await _PLUGIN_STORE_DATA_CACHE.set(cache_key, result)
return result
@classmethod
def version_check(cls, plugin_info: StorePluginInfo, suc_plugin: dict[str, str]):
@@ -330,13 +334,12 @@ class StoreManager:
source: 源
"""
repo_type = RepoType.GITHUB if is_external else None
if source == "ali":
if (
source != "ali" and source != "git" and plugin_info.ali_url
) or source == "ali":
repo_type = RepoType.ALIYUN
elif source == "git":
repo_type = RepoType.GITHUB
else:
if plugin_info.ali_url:
repo_type = RepoType.ALIYUN
module_path = plugin_info.module_path
is_dir = plugin_info.is_dir
github_url = plugin_info.github_url
@@ -380,7 +383,7 @@ class StoreManager:
requirement_file = target_dir / requirement_path.path
if requirement_file.exists():
is_install_req = True
await VirtualEnvPackageManager.install_requirement(requirement_file)
await VirtualEnvPackageManager.add_requirement(requirement_file)
if not is_install_req:
# 从仓库根目录查找文件
@@ -401,13 +404,13 @@ class StoreManager:
f"开始安装插件 {module_path} 依赖文件: {requirement_path}",
LOG_COMMAND,
)
await VirtualEnvPackageManager.install_requirement(requirement_path)
await VirtualEnvPackageManager.add_requirement(requirement_path)
if requirements_path.exists():
logger.info(
f"开始安装插件 {module_path} 依赖文件: {requirements_path}",
LOG_COMMAND,
)
await VirtualEnvPackageManager.install_requirement(requirements_path)
await VirtualEnvPackageManager.add_requirement(requirements_path)
@classmethod
async def remove_plugin(cls, index_or_module: str) -> str:
zhenxun/builtin_plugins/shop/_data_source.py
+7 -5
View File
@@ -19,9 +19,9 @@ from zhenxun.models.friend_user import FriendUser
from zhenxun.models.goods_info import GoodsInfo
from zhenxun.models.group_member_info import GroupInfoUser
from zhenxun.models.user_console import UserConsole
from zhenxun.models.user_gold_log import UserGoldLog
from zhenxun.models.user_props_log import UserPropsLog
from zhenxun.services import avatar_service
from zhenxun.services.buffered_writers import append_user_gold_log
from zhenxun.services.log import logger
from zhenxun.ui.models import ImageCell, TextCell
from zhenxun.utils.enum import GoldHandle, PropHandle
@@ -480,10 +480,6 @@ class ShopManage:
).count()
if goods.daily_limit and count >= goods.daily_limit:
return "今天的购买已达限制了喔!"
await UserGoldLog.create(user_id=user_id, gold=price, handle=GoldHandle.BUY)
await UserPropsLog.create(
user_id=user_id, uuid=goods.uuid, gold=price, num=num, handle=PropHandle.BUY
)
logger.info(
f"花费 {price} 金币购买 {goods.goods_name} ×{num} 成功！",
"购买道具",
@@ -494,6 +490,12 @@ class ShopManage:
user.props[goods.uuid] = 0
user.props[goods.uuid] += num
await user.save(update_fields=["gold", "props"])
await append_user_gold_log(
user_id=user_id, gold=int(price), handle=GoldHandle.BUY
)
await UserPropsLog.create(
user_id=user_id, uuid=goods.uuid, gold=price, num=num, handle=PropHandle.BUY
)
return f"花费 {price} 金币购买 {goods.goods_name} ×{num} 成功！"
@classmethod
zhenxun/builtin_plugins/sign_in/_data_source.py
+28 -3
View File
@@ -9,13 +9,17 @@ import pytz
from zhenxun import ui
from zhenxun.configs.path_config import IMAGE_PATH
from zhenxun.models.friend_user import FriendUser
from zhenxun.models.goods_info import GoodsInfo
from zhenxun.models.group_member_info import GroupInfoUser
from zhenxun.models.sign_log import SignLog
from zhenxun.models.sign_user import SignUser
from zhenxun.models.user_console import UserConsole
from zhenxun.services.avatar_service import avatar_service
from zhenxun.services.buffered_writers import append_user_gold_log
from zhenxun.services.log import logger
from zhenxun.ui.models import ImageCell, TextCell
from zhenxun.utils.enum import GoldHandle
from zhenxun.utils.exception import GoodsNotFound
from zhenxun.utils.platform import PlatformUtils
from ._random_event import random_event
@@ -182,11 +186,32 @@ class SignManage:
gift = random_event(float(user.impression))
if isinstance(gift, int):
gold += gift
await UserConsole.add_gold(user.user_id, gold + gift, "sign_in", platform)
user_console = await UserConsole.get_user(user.user_id, platform)
user_console.gold += gold
await user_console.save(update_fields=["gold"])
await append_user_gold_log(
user_id=user.user_id,
gold=gold,
handle=GoldHandle.GET,
source="sign_in",
)
gift = f"额外金币 +{gift}"
else:
await UserConsole.add_gold(user.user_id, gold, "sign_in", platform)
await UserConsole.add_props_by_name(user.user_id, gift, 1, platform)
goods = await GoodsInfo.get_or_none(goods_name=gift)
if not goods:
raise GoodsNotFound("未找到商品...")
user_console = await UserConsole.get_user(user.user_id, platform)
user_console.gold += gold
if goods.uuid not in user_console.props:
user_console.props[goods.uuid] = 0
user_console.props[goods.uuid] += 1
await user_console.save(update_fields=["gold", "props"])
await append_user_gold_log(
user_id=user.user_id,
gold=gold,
handle=GoldHandle.GET,
source="sign_in",
)
gift += " + 1"
logger.info(
f"签到成功. score: {user.impression:.2f} "
zhenxun/builtin_plugins/statistics/statistics_hook.py
+38 -7
View File
@@ -1,5 +1,7 @@
import asyncio
from datetime import datetime
from nonebot import get_driver
from nonebot.adapters import Bot, Event
from nonebot.adapters.onebot.v11 import PokeNotifyEvent
from nonebot.matcher import Matcher
@@ -25,7 +27,35 @@ __plugin_meta__ = PluginMetadata(
).to_dict(),
)
TEMP_LIST = []
STATS_BUFFER_FLUSH_SIZE = 5000
STATS_BUFFER_MAX_RETAIN = 10000
TEMP_LIST: list[Statistics] = []
_STATS_FLUSH_LOCK = asyncio.Lock()
driver = get_driver()
async def _flush_statistics_buffer(reason: str) -> int:
async with _STATS_FLUSH_LOCK:
call_list = TEMP_LIST.copy()
TEMP_LIST.clear()
if not call_list:
return 0
try:
await Statistics.bulk_create(call_list)
except Exception as e:
logger.error(f"{reason}批量添加调用记录失败", "定时任务", e=e)
retain_count = max(STATS_BUFFER_MAX_RETAIN - len(TEMP_LIST), 0)
if retain_count:
TEMP_LIST[:0] = call_list[-retain_count:]
return 0
logger.debug(f"{reason}批量添加调用记录 {len(call_list)} 条", "定时任务")
return len(call_list)
async def _append_statistics(record: Statistics) -> None:
TEMP_LIST.append(record)
if len(TEMP_LIST) >= STATS_BUFFER_FLUSH_SIZE and not _STATS_FLUSH_LOCK.locked():
await _flush_statistics_buffer("缓冲区触发")
@run_postprocessor
@@ -50,7 +80,7 @@ async def _(
if plugin_type == PluginType.NORMAL:
entity = get_entity_ids(session)
logger.debug(f"提交调用记录: {matcher.plugin_name}...", session=session)
TEMP_LIST.append(
await _append_statistics(
Statistics(
user_id=entity.user_id,
group_id=entity.group_id,
@@ -66,10 +96,11 @@ async def _():
try:
if should_pause_tasks():
return
call_list = TEMP_LIST.copy()
TEMP_LIST.clear()
if call_list:
await Statistics.bulk_create(call_list)
logger.debug(f"批量添加调用记录 {len(call_list)} 条", "定时任务")
await _flush_statistics_buffer("定时")
except Exception as e:
logger.error("定时批量添加调用记录", "定时任务", e=e)
@driver.on_shutdown
async def _flush_statistics_on_shutdown():
await _flush_statistics_buffer("关闭")
zhenxun/builtin_plugins/superuser/reload_setting.py
+68 -10
View File
@@ -1,3 +1,5 @@
import contextlib
from nonebot.permission import SUPERUSER
from nonebot.plugin import PluginMetadata
from nonebot.rule import to_me
@@ -11,8 +13,11 @@ from zhenxun.services.llm.config.providers import get_llm_config
from zhenxun.services.llm.manager import clear_model_cache
from zhenxun.services.log import logger
from zhenxun.utils.enum import PluginType
from zhenxun.utils.manager.priority_manager import PriorityLifecycle
from zhenxun.utils.message import MessageUtils
AUTO_RELOAD_JOB_ID = "zhenxun.reload_setting.auto_reload"
__plugin_meta__ = PluginMetadata(
name="重载配置",
description="重新加载config.yaml",
@@ -53,22 +58,75 @@ _matcher = on_alconna(
)
@_matcher.handle()
async def _(session: EventSession, arparma: Arparma):
def _get_auto_reload_interval() -> int:
value = Config.get_config("reload_setting", "AUTO_RELOAD_TIME", 180)
try:
seconds = int(value)
except (TypeError, ValueError):
logger.warning(
f"AUTO_RELOAD_TIME 配置无效: {value!r}，已使用默认值 180 秒",
"重载配置",
)
return 180
if seconds <= 0:
logger.warning(
f"AUTO_RELOAD_TIME 配置小于等于 0: {seconds}，已使用默认值 180 秒",
"重载配置",
)
return 180
return seconds
def _reschedule_auto_reload_job() -> None:
seconds = _get_auto_reload_interval()
if scheduler.get_job(AUTO_RELOAD_JOB_ID):
scheduler.reschedule_job(
AUTO_RELOAD_JOB_ID,
trigger="interval",
seconds=seconds,
)
else:
scheduler.add_job(
_auto_reload_config,
"interval",
seconds=seconds,
id=AUTO_RELOAD_JOB_ID,
replace_existing=True,
)
logger.debug(f"自动重载配置任务间隔已设置为 {seconds} 秒", "重载配置")
async def _reload_plugin_limit_config() -> None:
from zhenxun.builtin_plugins.hooks.auth.auth_limit import LimitManager
from zhenxun.builtin_plugins.init.manager import manager
manager.init()
await manager.load_to_db()
await LimitManager.update_limits()
async def _reload_runtime_config() -> None:
Config.reload()
get_llm_config.cache_clear()
clear_model_cache()
await _reload_plugin_limit_config()
with contextlib.suppress(Exception):
_reschedule_auto_reload_job()
@PriorityLifecycle.on_startup(priority=1)
def _init_auto_reload_job() -> None:
_reschedule_auto_reload_job()
@_matcher.handle()
async def _(session: EventSession, arparma: Arparma):
await _reload_runtime_config()
logger.debug("自动重载配置文件", arparma.header_result, session=session)
await MessageUtils.build_message("重载完成!").send(reply_to=True)
@scheduler.scheduled_job(
"interval",
seconds=Config.get_config("reload_setting", "AUTO_RELOAD_TIME", 180),
)
async def _():
async def _auto_reload_config() -> None:
if Config.get_config("reload_setting", "AUTO_RELOAD"):
Config.reload()
get_llm_config.cache_clear()
clear_model_cache()
await _reload_runtime_config()
logger.debug("已自动重载配置文件...")
zhenxun/builtin_plugins/web_ui/__init__.py
+1 -23
View File
@@ -1,20 +1,17 @@
import asyncio
import secrets
from fastapi import APIRouter, FastAPI
import nonebot
from nonebot.log import default_filter, default_format
from nonebot.plugin import PluginMetadata
from zhenxun.configs.config import Config as gConfig
from zhenxun.configs.utils import PluginExtraData, RegisterConfig
from zhenxun.services.log import logger, logger_
from zhenxun.services.log import logger
from zhenxun.utils.enum import PluginType
from zhenxun.utils.manager.priority_manager import PriorityLifecycle
from .api.configure import router as configure_router
from .api.logs import router as ws_log_routes
from .api.logs.log_manager import LOG_STORAGE
from .api.menu import router as menu_router
from .api.tabs.dashboard import router as dashboard_router
from .api.tabs.database import router as database_router
@@ -95,25 +92,6 @@ WsApiRouter.include_router(chat_routes)
@PriorityLifecycle.on_startup(priority=0)
async def _():
try:
# 存储任务引用的列表，防止任务被垃圾回收
_tasks = []
async def log_sink(message: str):
loop = None
if not loop:
try:
loop = asyncio.get_running_loop()
except Exception as e:
logger.warning("Web Ui log_sink", e=e)
if not loop:
loop = asyncio.new_event_loop()
# 存储任务引用到外部列表中
_tasks.append(loop.create_task(LOG_STORAGE.add(message.rstrip("\n"))))
logger_.add(
log_sink, colorize=True, filter=default_filter, format=default_format
)
app: FastAPI = nonebot.get_app()
app.include_router(BaseApiRouter)
app.include_router(WsApiRouter)
zhenxun/builtin_plugins/web_ui/api/logs/log_manager.py
+76 -12
View File
@@ -1,33 +1,97 @@
import asyncio
from collections import deque
from collections.abc import Awaitable, Callable
from typing import Generic, TypeVar
import contextlib
_T = TypeVar("_T")
LogListener = Callable[[_T], Awaitable[None]]
from nonebot.log import default_filter, default_format
from zhenxun.services.log import logger_
LogListener = Callable[[str], Awaitable[None]]
DEFAULT_MAX_LOGS = 1000
DEFAULT_MAX_LISTENERS = 16
class LogStorage(Generic[_T]):
class LogStorage:
"""
日志存储
"""
def __init__(self, rotation: float = 5 * 60):
def __init__(
self,
rotation: float = 5 * 60,
max_logs: int = DEFAULT_MAX_LOGS,
max_listeners: int = DEFAULT_MAX_LISTENERS,
):
self.count, self.rotation = 0, rotation
self.max_logs = max_logs
self.max_listeners = max_listeners
self.logs: dict[int, str] = {}
self.listeners: set[LogListener[str]] = set()
self._order: deque[int] = deque()
self.listeners: set[LogListener] = set()
async def add(self, log: str):
seq = self.count = self.count + 1
self.logs[seq] = log
self._order.append(seq)
self._trim()
asyncio.get_running_loop().call_later(self.rotation, self.remove, seq)
await asyncio.gather(
*(listener(log) for listener in self.listeners),
return_exceptions=True,
)
listeners = tuple(self.listeners)
if listeners:
results = await asyncio.gather(
*(listener(log) for listener in listeners),
return_exceptions=True,
)
for listener, result in zip(listeners, results, strict=False):
if isinstance(result, BaseException):
self.listeners.discard(listener)
return seq
def add_listener(self, listener: LogListener) -> bool:
if len(self.listeners) >= self.max_listeners:
return False
self.listeners.add(listener)
return True
def remove_listener(self, listener: LogListener) -> None:
self.listeners.discard(listener)
def remove(self, seq: int):
del self.logs[seq]
self.logs.pop(seq, None)
with contextlib.suppress(ValueError):
self._order.remove(seq)
def _trim(self) -> None:
while self._order and self._order[0] not in self.logs:
self._order.popleft()
while len(self.logs) > self.max_logs and self._order:
self.logs.pop(self._order.popleft(), None)
LOG_STORAGE: LogStorage[str] = LogStorage[str]()
LOG_STORAGE = LogStorage()
_LOG_SINK_ID: int | None = None
async def ensure_log_sink_started() -> None:
global _LOG_SINK_ID
if _LOG_SINK_ID is not None:
return
async def log_sink(message: str) -> None:
await LOG_STORAGE.add(message.rstrip("\n"))
_LOG_SINK_ID = logger_.add(
log_sink,
colorize=True,
filter=default_filter,
format=default_format,
)
def stop_log_sink_if_idle() -> None:
global _LOG_SINK_ID
if LOG_STORAGE.listeners or _LOG_SINK_ID is None:
return
logger_.remove(_LOG_SINK_ID)
_LOG_SINK_ID = None
zhenxun/builtin_plugins/web_ui/api/logs/logs.py
+9 -3
View File
@@ -3,7 +3,7 @@ from loguru import logger
from nonebot.utils import escape_tag
from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState
from .log_manager import LOG_STORAGE
from .log_manager import LOG_STORAGE, ensure_log_sink_started, stop_log_sink_if_idle
router = APIRouter()
@@ -11,11 +11,16 @@ router = APIRouter()
@router.websocket("/logs")
async def system_logs_realtime(websocket: WebSocket):
await websocket.accept()
await ensure_log_sink_started()
async def log_listener(log: str):
await websocket.send_text(log)
LOG_STORAGE.listeners.add(log_listener)
if not LOG_STORAGE.add_listener(log_listener):
await websocket.send_text("日志连接数已达上限，请稍后再试。")
await websocket.close()
stop_log_sink_if_idle()
return
try:
while websocket.client_state == WebSocketState.CONNECTED:
recv = await websocket.receive()
@@ -26,4 +31,5 @@ async def system_logs_realtime(websocket: WebSocket):
except WebSocketDisconnect:
pass
finally:
LOG_STORAGE.listeners.remove(log_listener)
LOG_STORAGE.remove_listener(log_listener)
stop_log_sink_if_idle()
zhenxun/builtin_plugins/web_ui/api/tabs/main/__init__.py
+59 -6
View File
@@ -34,6 +34,31 @@ run_time = time.time()
ws_router = APIRouter()
router = APIRouter(prefix="/main")
_SYSTEM_STATUS_CONNECTIONS: set[WebSocket] = set()
_SYSTEM_STATUS_STOPPING = False
async def _close_system_status_websocket(websocket: WebSocket) -> None:
with contextlib.suppress(Exception):
if websocket.client_state == WebSocketState.CONNECTED:
await asyncio.wait_for(
websocket.close(code=1001, reason="server shutdown"),
timeout=2,
)
@driver.on_shutdown
async def _close_system_status_websockets() -> None:
global _SYSTEM_STATUS_STOPPING
_SYSTEM_STATUS_STOPPING = True
websockets = list(_SYSTEM_STATUS_CONNECTIONS)
if not websockets:
return
await asyncio.gather(
*(_close_system_status_websocket(websocket) for websocket in websockets),
return_exceptions=True,
)
_SYSTEM_STATUS_CONNECTIONS.clear()
@router.get(
@@ -243,11 +268,39 @@ async def _(param: BotManageUpdateParam):
@ws_router.websocket("/system_status")
async def system_logs_realtime(websocket: WebSocket, sleep: int = 5):
await websocket.accept()
_SYSTEM_STATUS_CONNECTIONS.add(websocket)
logger.debug("ws system_status is connect")
with contextlib.suppress(
WebSocketDisconnect, ConnectionClosedError, ConnectionClosedOK
):
while websocket.client_state == WebSocketState.CONNECTED:
disconnect_event = asyncio.Event()
async def _watch_disconnect() -> None:
try:
while websocket.client_state == WebSocketState.CONNECTED:
await websocket.receive()
except (WebSocketDisconnect, ConnectionClosedError, ConnectionClosedOK):
pass
except Exception as e:
logger.debug(f"ws system_status receive stopped: {type(e).__name__}")
finally:
disconnect_event.set()
receive_task = asyncio.create_task(_watch_disconnect())
try:
while (
websocket.client_state == WebSocketState.CONNECTED
and not _SYSTEM_STATUS_STOPPING
):
system_status = await get_system_status()
await websocket.send_text(system_status.json())
await asyncio.sleep(sleep)
await asyncio.wait_for(websocket.send_text(system_status.json()), timeout=5)
try:
await asyncio.wait_for(disconnect_event.wait(), timeout=max(sleep, 1))
except TimeoutError:
pass
except (WebSocketDisconnect, ConnectionClosedError, ConnectionClosedOK):
pass
finally:
_SYSTEM_STATUS_CONNECTIONS.discard(websocket)
receive_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await receive_task
await _close_system_status_websocket(websocket)
zhenxun/builtin_plugins/web_ui/api/tabs/plugin_manage/__init__.py
+13 -27
View File
@@ -52,34 +52,20 @@ async def _(
async def _() -> Result[PluginCount]:
try:
plugin_count = PluginCount()
plugin_count.normal = len(
await DbPluginInfo.get_plugins(
plugin_type=PluginType.NORMAL,
load_status=True,
filter_parent=False,
)
)
plugin_count.admin = len(
await DbPluginInfo.get_plugins(
plugin_type__in=[PluginType.ADMIN, PluginType.SUPER_AND_ADMIN],
load_status=True,
filter_parent=False,
)
)
plugin_count.superuser = len(
await DbPluginInfo.get_plugins(
plugin_type__in=[PluginType.SUPERUSER, PluginType.SUPER_AND_ADMIN],
load_status=True,
filter_parent=False,
)
)
plugin_count.other = len(
await DbPluginInfo.get_plugins(
plugin_type__in=[PluginType.HIDDEN, PluginType.DEPENDANT],
load_status=True,
filter_parent=False,
)
plugins = await DbPluginInfo.get_plugins(
load_status=True,
filter_parent=False,
)
for plugin in plugins:
plugin_type = plugin.plugin_type
if plugin_type == PluginType.NORMAL:
plugin_count.normal += 1
if plugin_type in {PluginType.ADMIN, PluginType.SUPER_AND_ADMIN}:
plugin_count.admin += 1
if plugin_type in {PluginType.SUPERUSER, PluginType.SUPER_AND_ADMIN}:
plugin_count.superuser += 1
if plugin_type in {PluginType.HIDDEN, PluginType.DEPENDANT}:
plugin_count.other += 1
return Result.ok(plugin_count, "拿到信息啦!")
except Exception as e:
logger.error(f"{router.prefix}/get_plugin_count 调用错误", "WebUi", e=e)
zhenxun/builtin_plugins/web_ui/api/tabs/plugin_manage/data_source.py
+10 -1
View File
@@ -111,10 +111,19 @@ class ApiDataSource:
other_update_fields = set()
updated_count = 0
errors = []
modules = [item.module for item in params.updates]
plugin_records = await DbPluginInfo.get_plugins(
module__in=modules,
load_status=None,
filter_parent=False,
)
plugin_map = {plugin.module: plugin for plugin in plugin_records}
for item in params.updates:
try:
db_plugin = await DbPluginInfo.get(module=item.module)
db_plugin = plugin_map.get(item.module)
if db_plugin is None:
raise DoesNotExist()
plugin_changed_other = False
plugin_changed_block = False
zhenxun/cli.py
+129 -13
View File
@@ -8,12 +8,27 @@
from __future__ import annotations
import atexit
import importlib.metadata
import os
from pathlib import Path
import signal
import subprocess
import sys
import time
GRACEFUL_SHUTDOWN_TIMEOUT = 15
WORKER_POLL_INTERVAL = 0.1
RESTART_POLL_INTERVAL = 0.5
WORKER_SOFT_EXIT_TIMEOUT = 15.0
WORKER_TERMINATE_TIMEOUT = 5.0
WORKER_KILL_TIMEOUT = 5.0
def _launcher_log(message: str) -> None:
sys.stderr.write(f"[zx launcher] {message}\n")
sys.stderr.flush()
def _print_version() -> None:
try:
@@ -94,32 +109,69 @@ def _run_worker() -> None:
nonebot.logger.info(f"加载第三方插件目录: {ext}")
nonebot.load_plugins(ext)
nonebot.run()
nonebot.run(timeout_graceful_shutdown=GRACEFUL_SHUTDOWN_TIMEOUT)
def _build_worker_command() -> list[str]:
return [sys.executable, "-m", "zhenxun.cli", "run-worker"]
def _get_worker_creationflags() -> int:
if os.name == "nt":
return getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0)
return 0
def _wait_worker_exit(proc: subprocess.Popen, timeout_seconds: float) -> bool:
deadline = time.monotonic() + timeout_seconds
while time.monotonic() < deadline:
if proc.poll() is not None:
return True
time.sleep(0.1)
time.sleep(WORKER_POLL_INTERVAL)
return proc.poll() is not None
def _terminate_worker(proc: subprocess.Popen) -> None:
if proc.poll() is not None:
return
if _wait_worker_exit(proc, 8.0):
return
proc.terminate()
if _wait_worker_exit(proc, 5.0):
_launcher_log(f"stopping worker pid={proc.pid}")
if os.name == "nt":
ctrl_break_event = getattr(signal, "CTRL_BREAK_EVENT", None)
if ctrl_break_event is not None:
try:
_launcher_log(f"sending CTRL_BREAK_EVENT to worker pid={proc.pid}")
proc.send_signal(ctrl_break_event)
except Exception as e:
_launcher_log(f"failed to send CTRL_BREAK_EVENT: {e!r}")
else:
if _wait_worker_exit(proc, WORKER_SOFT_EXIT_TIMEOUT):
_launcher_log(
f"worker pid={proc.pid} exited after CTRL_BREAK_EVENT "
f"with code {proc.returncode}"
)
return
_launcher_log(
f"worker pid={proc.pid} did not exit after "
f"{WORKER_SOFT_EXIT_TIMEOUT:.0f}s"
)
if _wait_worker_exit(proc, 1.0):
return
try:
_launcher_log(f"terminating worker pid={proc.pid}")
proc.terminate()
except Exception as e:
_launcher_log(f"failed to terminate worker: {e!r}")
else:
if _wait_worker_exit(proc, WORKER_TERMINATE_TIMEOUT):
_launcher_log(
f"worker pid={proc.pid} exited after terminate with code "
f"{proc.returncode}"
)
return
_launcher_log(f"worker pid={proc.pid} did not exit after terminate timeout")
_launcher_log(f"killing worker pid={proc.pid}")
proc.kill()
proc.wait(timeout=5)
proc.wait(timeout=WORKER_KILL_TIMEOUT)
def _run_launcher() -> None:
@@ -130,19 +182,83 @@ def _run_launcher() -> None:
)
clear_launcher_restart_signal()
while True:
worker = subprocess.Popen(_build_worker_command(), cwd=str(cwd))
current_worker: subprocess.Popen | None = None
stop_requested = False
stop_signal: int | None = None
def _cleanup_current_worker() -> None:
if current_worker is not None:
_terminate_worker(current_worker)
atexit.register(_cleanup_current_worker)
def _handle_launcher_signal(signum, _frame) -> None:
nonlocal stop_requested, stop_signal
if stop_requested:
_launcher_log(f"received signal {signum} while stopping, exiting launcher")
raise SystemExit(128 + int(signum))
stop_requested = True
stop_signal = int(signum)
_launcher_log(f"received signal {signum}, scheduling worker shutdown")
handled_signals = [signal.SIGINT]
if hasattr(signal, "SIGTERM"):
handled_signals.append(signal.SIGTERM)
if hasattr(signal, "SIGBREAK"):
handled_signals.append(signal.SIGBREAK)
for sig in handled_signals:
try:
return_code = worker.wait()
signal.signal(sig, _handle_launcher_signal)
except Exception:
pass
while True:
if stop_requested:
raise SystemExit(128 + int(stop_signal or signal.SIGINT))
worker_env = os.environ.copy()
worker_env["ZHENXUN_LAUNCHER_PID"] = str(os.getpid())
worker = subprocess.Popen(
_build_worker_command(),
cwd=str(cwd),
creationflags=_get_worker_creationflags(),
env=worker_env,
)
current_worker = worker
restart_requested = False
return_code: int | None = None
next_restart_check = 0.0
try:
while True:
return_code = worker.poll()
if return_code is not None:
break
if stop_requested:
clear_launcher_restart_signal()
_terminate_worker(worker)
raise SystemExit(128 + int(stop_signal or signal.SIGINT))
now = time.monotonic()
if now >= next_restart_check:
next_restart_check = now + RESTART_POLL_INTERVAL
if consume_launcher_restart_signal():
restart_requested = True
_launcher_log(
"detected restart request, stopping current worker"
)
_terminate_worker(worker)
return_code = worker.poll()
break
time.sleep(WORKER_POLL_INTERVAL)
except KeyboardInterrupt:
clear_launcher_restart_signal()
_terminate_worker(worker)
return
finally:
if current_worker is worker:
current_worker = None
should_restart = consume_launcher_restart_signal()
if should_restart:
if restart_requested or consume_launcher_restart_signal():
continue
raise SystemExit(return_code)
raise SystemExit(return_code if return_code is not None else 1)
def main() -> None:
zhenxun/configs/utils/__init__.py
+108 -19
View File
@@ -4,12 +4,12 @@ from pathlib import Path
from typing import Any, TypeVar
import cattrs
from nonebot.log import logger as _nonebot_logger
from pydantic import BaseModel, Field
from ruamel.yaml import YAML
from ruamel.yaml.scanner import ScannerError
from zhenxun.configs.path_config import DATA_PATH
from zhenxun.services.log import logger
from zhenxun.utils.pydantic_compat import (
_dump_pydantic_obj,
_is_pydantic_type,
@@ -38,6 +38,33 @@ _yaml.indent = 2
_yaml.allow_unicode = True
T = TypeVar("T")
_MISSING = object()
class _ConfigLogger:
@staticmethod
def _emit(level: str, info: str, *, e: Exception | None = None) -> None:
logger = _nonebot_logger.opt(exception=e) if e else _nonebot_logger
getattr(logger, level)(info)
@classmethod
def debug(cls, info: str, *_, e: Exception | None = None, **__) -> None:
cls._emit("debug", info, e=e)
@classmethod
def info(cls, info: str, *_, e: Exception | None = None, **__) -> None:
cls._emit("info", info, e=e)
@classmethod
def warning(cls, info: str, *_, e: Exception | None = None, **__) -> None:
cls._emit("warning", info, e=e)
@classmethod
def error(cls, info: str, *_, e: Exception | None = None, **__) -> None:
cls._emit("error", info, e=e)
logger = _ConfigLogger()
class NoSuchConfig(Exception):
@@ -114,21 +141,83 @@ class ConfigsManager:
self._simple_data: dict = {}
self._simple_file = DATA_PATH / "config.yaml"
self.add_module = []
_yaml = YAML()
if file:
file.parent.mkdir(exist_ok=True, parents=True)
self.file = file
self.load_data()
if self._simple_file.exists():
try:
with self._simple_file.open(encoding="utf8") as f:
self._simple_data = _yaml.load(f)
except ScannerError as e:
raise ScannerError(
f"{e}\n**********************************************\n"
f"****** 可能为config.yaml配置文件填写不规范 ******\n"
f"**********************************************"
) from e
self._load_simple_data(raise_on_error=True)
self._apply_simple_data(warn_unknown=False)
def _load_simple_data(self, *, raise_on_error: bool = False) -> None:
if not self._simple_file.exists():
self._simple_data = {}
return
try:
with self._simple_file.open(encoding="utf8") as f:
simple_data = _yaml.load(f) or {}
except ScannerError as e:
message = (
f"{e}\n**********************************************\n"
f"****** 可能为config.yaml配置文件填写不规范 ******\n"
f"**********************************************"
)
if raise_on_error:
raise ScannerError(message) from e
logger.warning(f"读取config.yaml失败，已跳过本次重载: {message}", e=e)
return
except Exception as e:
if raise_on_error:
raise RuntimeError(f"读取config.yaml失败: {e}") from e
logger.warning(f"读取config.yaml失败，已跳过本次重载: {e}", e=e)
return
if not isinstance(simple_data, dict):
message = "config.yaml 顶层必须为字典，已忽略当前内容。"
if raise_on_error:
raise ValueError(message)
logger.warning(message)
self._simple_data = {}
return
self._simple_data = simple_data
@staticmethod
def _find_mapping_key(data: dict, key: str) -> str | None:
if key in data:
return key
upper_key = key.upper()
for raw_key in data:
if str(raw_key).upper() == upper_key:
return raw_key
return None
def _get_simple_config_value(self, module: str, key: str) -> Any:
module_data = self._simple_data.get(module)
if not isinstance(module_data, dict):
return _MISSING
simple_key = self._find_mapping_key(module_data, key.upper())
if simple_key is None:
return _MISSING
return module_data[simple_key]
def _apply_simple_data(self, *, warn_unknown: bool) -> None:
for module, module_data in self._simple_data.items():
if not isinstance(module_data, dict):
if warn_unknown:
logger.warning(f"配置组 {module} 不是字典，已跳过。")
continue
config_group = self._data.get(module)
if not config_group:
if warn_unknown:
logger.warning(f"未知配置组 {module}，已跳过。")
continue
for raw_key, value in module_data.items():
key = str(raw_key).upper()
config_key = self._find_mapping_key(config_group.configs, key)
if config_key is None:
if warn_unknown:
logger.warning(f"未知配置项 {module}.{raw_key}，已跳过。")
continue
config_group.configs[config_key].value = value
def set_name(self, module: str, name: str):
"""设置插件配置中文名出
@@ -223,7 +312,11 @@ class ConfigsManager:
if module in self._data and (config := self._data[module].configs.get(key)):
existing_value = config.value
processed_value = self._normalize_config_data(value, existing_value)
simple_value = self._get_simple_config_value(module, key)
if simple_value is _MISSING:
processed_value = self._normalize_config_data(value, existing_value)
else:
processed_value = self._normalize_config_data(value, simple_value)
processed_default_value = self._normalize_config_data(default_value)
self.add_module.append(f"{module}:{key}".lower())
@@ -231,7 +324,7 @@ class ConfigsManager:
config.help = help
config.arg_parser = arg_parser
config.type = type
if _override:
if simple_value is not _MISSING or _override:
config.value = processed_value
config.default_value = processed_default_value
else:
@@ -371,12 +464,8 @@ class ConfigsManager:
def reload(self):
"""重新加载配置文件"""
if self._simple_file.exists():
with open(self._simple_file, encoding="utf8") as f:
self._simple_data = _yaml.load(f)
for key in self._simple_data.keys():
for k in self._simple_data[key].keys():
self._data[key].configs[k].value = self._simple_data[key][k]
self._load_simple_data()
self._apply_simple_data(warn_unknown=True)
self.save()
def load_data(self):
zhenxun/models/group_member_info.py
+52 -1
View File
@@ -2,6 +2,7 @@ from typing import ClassVar
from tortoise import fields
from zhenxun.configs.config import BotConfig
from zhenxun.services.db_context import Model
@@ -51,4 +52,54 @@ class GroupInfoUser(Model):
@classmethod
async def _run_script(cls):
return ["ALTER TABLE group_info_users DROP COLUMN nickname;"]
db_type = (BotConfig.get_sql_type() or "").lower()
scripts = ["ALTER TABLE group_info_users DROP COLUMN nickname;"]
if "postgres" in db_type:
scripts.extend(
[
(
"ALTER TABLE group_info_users ADD COLUMN IF NOT EXISTS "
"platform character varying(255);"
),
(
"ALTER TABLE group_info_users ALTER COLUMN user_id "
"TYPE character varying(255) USING user_id::character varying;"
),
(
"ALTER TABLE group_info_users ALTER COLUMN group_id "
"TYPE character varying(255) USING group_id::character varying;"
),
(
"ALTER TABLE group_info_users ALTER COLUMN platform "
"TYPE character varying(255) USING platform::character varying;"
),
]
)
elif "mysql" in db_type:
scripts.extend(
[
(
"ALTER TABLE group_info_users ADD COLUMN "
"platform VARCHAR(255) NULL;"
),
(
"ALTER TABLE group_info_users MODIFY COLUMN "
"user_id VARCHAR(255) NOT NULL;"
),
(
"ALTER TABLE group_info_users MODIFY COLUMN "
"group_id VARCHAR(255) NOT NULL;"
),
(
"ALTER TABLE group_info_users MODIFY COLUMN "
"platform VARCHAR(255) NULL;"
),
]
)
elif "sqlite" in db_type:
scripts.append(
"ALTER TABLE group_info_users ADD COLUMN platform VARCHAR(255);"
)
return scripts
zhenxun/models/user_console.py
+18 -8
View File
@@ -5,12 +5,11 @@ from tortoise import fields
from tortoise.exceptions import IntegrityError
from zhenxun.models.goods_info import GoodsInfo
from zhenxun.services.buffered_writers import append_user_gold_log
from zhenxun.services.db_context import Model
from zhenxun.utils.enum import CacheType, GoldHandle
from zhenxun.utils.exception import GoodsNotFound, InsufficientGold
from .user_gold_log import UserGoldLog
class UserConsole(Model):
id = fields.IntField(pk=True, generated=True, auto_increment=True)
@@ -77,6 +76,17 @@ class UserConsole(Model):
user, _ = await cls.get_or_create_user(user_id=user_id, platform=platform)
return user
@classmethod
async def _get_user_for_write(
cls, user_id: str, platform: str | None = None
) -> "UserConsole":
"""获取写入用用户；已有用户不走 get_or_create，避免重复清理缓存。"""
user = await cls.get_or_none(user_id=user_id)
if user is not None:
return user
user, _ = await cls.get_or_create_user(user_id=user_id, platform=platform)
return user
@classmethod
async def get_new_uid(cls) -> int:
"""获取最新uid
@@ -103,10 +113,10 @@ class UserConsole(Model):
source: 来源
platform: 平台.
"""
user, _ = await cls.get_or_create_user(user_id=user_id, platform=platform)
user = await cls._get_user_for_write(user_id=user_id, platform=platform)
user.gold += gold
await user.save(update_fields=["gold"])
await UserGoldLog.create(
await append_user_gold_log(
user_id=user_id, gold=gold, handle=GoldHandle.GET, source=source
)
@@ -131,12 +141,12 @@ class UserConsole(Model):
异常:
InsufficientGold: 金币不足
"""
user, _ = await cls.get_or_create_user(user_id=user_id, platform=platform)
user = await cls._get_user_for_write(user_id=user_id, platform=platform)
if user.gold < gold:
raise InsufficientGold()
user.gold -= gold
await user.save(update_fields=["gold"])
await UserGoldLog.create(
await append_user_gold_log(
user_id=user_id, gold=gold, handle=handle, source=plugin_module
)
@@ -152,7 +162,7 @@ class UserConsole(Model):
num: 道具数量.
platform: 平台.
"""
user, _ = await cls.get_or_create_user(user_id=user_id, platform=platform)
user = await cls._get_user_for_write(user_id=user_id, platform=platform)
if goods_uuid not in user.props:
user.props[goods_uuid] = 0
user.props[goods_uuid] += num
@@ -186,7 +196,7 @@ class UserConsole(Model):
num: 道具数量.
platform: 平台.
"""
user, _ = await cls.get_or_create_user(user_id=user_id, platform=platform)
user = await cls._get_user_for_write(user_id=user_id, platform=platform)
if goods_uuid not in user.props or user.props[goods_uuid] < num:
raise GoodsNotFound("未找到商品或道具数量不足...")
zhenxun/services/avatar_service.py
+5
View File
@@ -63,6 +63,11 @@ class AvatarService:
identifier = str(identifier)
return self.cache_path / platform / f"{identifier}.png"
def clear_memory_cache(self) -> int:
size = len(self._memory_cache)
self._memory_cache.clear()
return size
async def get_avatar_path(
self, platform: str, identifier: str, force_refresh: bool = False
) -> Path | None:
zhenxun/services/buffered_writers.py
+128
View File
@@ -0,0 +1,128 @@
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.utils.enum import GoldHandle
from zhenxun.utils.manager.priority_manager import PriorityLifecycle
LOG_COMMAND = "BufferedWriters"
_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())
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,
)
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)
async def append_user_gold_log(
user_id: str,
gold: int,
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("缓冲区触发")
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
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()
zhenxun/services/cache/__init__.py
Vendored
+10 -6
View File
@@ -74,6 +74,7 @@ from .config import (
)
__all__ = [
"BoundedTTLCache",
"Cache",
"CacheDict",
"CacheManager",
@@ -82,6 +83,7 @@ __all__ = [
]
from . import runtime_cache as _runtime_cache # noqa: F401
from .bounded_ttl import BoundedTTLCache
T = TypeVar("T")
U = TypeVar("U")
@@ -400,9 +402,10 @@ class CacheManager:
"""清除缓存
参数:
cache_type: 缓存类型，为None时清除所有缓存。
注意：受 aiocache 限制，无法按类型精确删除，
指定 cache_type 时仅清除整个 backend（行为与不指定相同）。
cache_type: 缓存类型。为 None 时清除整个 backend。
指定 cache_type 时不再退化为清除整个 backend，避免误删其他类型缓存。
需要刷新模型运行态缓存时，应调用对应
RuntimeCache.refresh/upsert/remove。
返回:
bool: 是否成功
@@ -413,11 +416,12 @@ class CacheManager:
try:
if cache_type:
logger.debug(
f"清除缓存类型 {cache_type}"
"（aiocache 不支持按前缀删除，清除整个 backend）",
logger.warning(
f"拒绝清除缓存类型 {cache_type}: "
"当前后端不支持可靠的按类型清理，已避免清除整个 backend",
LOG_COMMAND,
)
return False
await self.cache_backend.clear() # type: ignore
return True
except Exception as e:
zhenxun/services/cache/bounded_ttl.py
Vendored
+218
View File
@@ -0,0 +1,218 @@
from __future__ import annotations
import asyncio
from collections import OrderedDict
from collections.abc import Callable
from dataclasses import dataclass
import sys
import time
from typing import Generic, TypeVar
import weakref
K = TypeVar("K")
V = TypeVar("V")
def _default_sizeof(value: object) -> int:
if isinstance(value, bytes | bytearray | memoryview):
return len(value)
return 0
@dataclass(frozen=True)
class BoundedTTLCacheStats:
name: str
items: int
max_items: int
total_bytes: int
max_total_bytes: int | None
hits: int
misses: int
sets: int
evictions: int
def to_dict(self) -> dict[str, int | str | None]:
return {
"name": self.name,
"items": self.items,
"max_items": self.max_items,
"total_bytes": self.total_bytes,
"max_total_bytes": self.max_total_bytes,
"hits": self.hits,
"misses": self.misses,
"sets": self.sets,
"evictions": self.evictions,
}
class BoundedTTLCache(Generic[K, V]):
"""Small async TTL/LRU cache with optional total-byte limit."""
_instances: weakref.WeakSet["BoundedTTLCache"] = weakref.WeakSet()
def __init__(
self,
name: str,
ttl_seconds: float,
max_items: int,
max_total_bytes: int | None = None,
sizeof: Callable[[V], int] | None = None,
) -> None:
self.name = name.upper()
self._ttl_seconds = max(ttl_seconds, 0.0)
self._max_items = max(max_items, 1)
self._max_total_bytes = (
max_total_bytes
if isinstance(max_total_bytes, int) and max_total_bytes > 0
else None
)
self._sizeof = sizeof or _default_sizeof
self._cache: OrderedDict[K, tuple[float, V, int]] = OrderedDict()
self._total_bytes = 0
self._hits = 0
self._misses = 0
self._sets = 0
self._evictions = 0
self._lock = asyncio.Lock()
self.__class__._instances.add(self)
def _expire_at(self, now: float) -> float:
if self._ttl_seconds <= 0:
return sys.float_info.max
return now + self._ttl_seconds
def _value_size(self, value: V) -> int:
try:
return max(0, int(self._sizeof(value)))
except Exception:
return 0
def _remove_key_nolock(self, key: K) -> bool:
item = self._cache.pop(key, None)
if item is None:
return False
self._total_bytes -= item[2]
if self._total_bytes < 0:
self._total_bytes = 0
return True
def _pop_oldest_nolock(self) -> bool:
if not self._cache:
return False
_, (_, _, size) = self._cache.popitem(last=False)
self._total_bytes -= size
if self._total_bytes < 0:
self._total_bytes = 0
self._evictions += 1
return True
def _cleanup_nolock(self, now: float) -> None:
expired_keys = [
key for key, (expire_at, _, _) in self._cache.items() if expire_at <= now
]
for key in expired_keys:
if self._remove_key_nolock(key):
self._evictions += 1
while len(self._cache) > self._max_items:
self._pop_oldest_nolock()
if self._max_total_bytes is not None:
while self._total_bytes > self._max_total_bytes and self._cache:
self._pop_oldest_nolock()
async def get(self, key: K) -> V | None:
now = time.monotonic()
async with self._lock:
self._cleanup_nolock(now)
item = self._cache.get(key)
if item is None:
self._misses += 1
return None
expire_at, value, _ = item
if expire_at <= now:
self._remove_key_nolock(key)
self._misses += 1
return None
self._cache.move_to_end(key)
self._hits += 1
return value
async def set(self, key: K, value: V) -> bool:
value_size = self._value_size(value)
if self._max_total_bytes is not None and value_size > self._max_total_bytes:
return False
now = time.monotonic()
async with self._lock:
self._remove_key_nolock(key)
self._cache[key] = (self._expire_at(now), value, value_size)
self._total_bytes += value_size
self._sets += 1
self._cache.move_to_end(key)
self._cleanup_nolock(now)
return key in self._cache
async def delete(self, key: K) -> bool:
async with self._lock:
return self._remove_key_nolock(key)
async def clear(self) -> int:
async with self._lock:
size = len(self._cache)
self._cache.clear()
self._total_bytes = 0
return size
async def stats(self) -> BoundedTTLCacheStats:
now = time.monotonic()
async with self._lock:
self._cleanup_nolock(now)
return BoundedTTLCacheStats(
name=self.name,
items=len(self._cache),
max_items=self._max_items,
total_bytes=self._total_bytes,
max_total_bytes=self._max_total_bytes,
hits=self._hits,
misses=self._misses,
sets=self._sets,
evictions=self._evictions,
)
@classmethod
async def clear_all(cls) -> dict[str, int]:
result: dict[str, int] = {}
for cache in list(cls._instances):
size = await cache.clear()
if size:
result[cache.name] = result.get(cache.name, 0) + size
return result
@classmethod
async def stats_all(cls) -> dict[str, dict[str, int | str | None]]:
result: dict[str, dict[str, int | str | None]] = {}
for cache in list(cls._instances):
stats = await cache.stats()
if not stats.items:
continue
if cache.name not in result:
result[cache.name] = stats.to_dict()
continue
current = result[cache.name]
for key in (
"items",
"max_items",
"total_bytes",
"hits",
"misses",
"sets",
"evictions",
):
current[key] = int(current.get(key) or 0) + int(
getattr(stats, key) or 0
)
current_max_bytes = current.get("max_total_bytes")
if current_max_bytes is not None or stats.max_total_bytes is not None:
current["max_total_bytes"] = int(current_max_bytes or 0) + int(
stats.max_total_bytes or 0
)
return result
zhenxun/services/cache/cache_containers.py
Vendored
+84 -2
View File
@@ -1,9 +1,12 @@
from dataclasses import dataclass
import time
from typing import Any, Generic, TypeVar
import weakref
T = TypeVar("T")
DEFAULT_CACHE_MAX_ITEMS = 10000
@dataclass
class CacheData(Generic[T]):
@@ -16,16 +19,21 @@ class CacheData(Generic[T]):
class CacheDict(Generic[T]):
"""缓存字典类，提供类似普通字典的接口，数据只存储在内存中"""
def __init__(self, name: str, expire: int = 0):
_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: dict[str, CacheData[T]] = {}
self.__class__._instances.add(self)
def expire_time(self, key: str) -> float:
"""获取字典项的过期时间"""
@@ -62,6 +70,7 @@ class CacheDict(Generic[T]):
"""
expire_time = time.time() + self.expire if self.expire > 0 else 0
self._data[key] = CacheData(value=value, expire_time=expire_time)
self._enforce_limit()
def __delitem__(self, key: str) -> None:
"""删除字典项
@@ -122,6 +131,7 @@ class CacheDict(Generic[T]):
expire_time = time.time() + self.expire
self._data[key] = CacheData(value=value, expire_time=expire_time)
self._enforce_limit()
def pop(self, key: str, default: Any = None) -> T:
"""删除并返回字典项
@@ -146,6 +156,32 @@ class CacheDict(Generic[T]):
"""清空字典"""
self._data.clear()
def stats(self) -> dict[str, int]:
"""返回当前缓存条目统计。"""
self._clean_expired()
return {"items": len(self._data), "max_items": self.max_items}
@classmethod
def stats_all(cls) -> dict[str, dict[str, int]]:
"""返回所有 CacheDict 实例的条目统计。"""
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]:
"""清空所有 CacheDict，返回各缓存清理的条目数。"""
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 keys(self) -> list[str]:
"""获取所有键
@@ -187,6 +223,12 @@ class CacheDict(Generic[T]):
for key in expired_keys:
del self._data[key]
def _enforce_limit(self) -> None:
if self.max_items <= 0:
return
while len(self._data) > self.max_items:
self._data.pop(next(iter(self._data)))
def __len__(self) -> int:
"""获取字典长度
@@ -211,17 +253,22 @@ class CacheDict(Generic[T]):
class CacheList(Generic[T]):
"""缓存列表类，提供类似普通列表的接口，数据只存储在内存中"""
def __init__(self, name: str, expire: int = 0):
_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:
@@ -303,6 +350,7 @@ class CacheList(Generic[T]):
self.clear()
self._data.append(CacheData(value=value))
self._enforce_limit()
# 更新过期时间
self._update_expire_time()
@@ -318,6 +366,7 @@ class CacheList(Generic[T]):
self.clear()
self._data.extend([CacheData(value=v) for v in values])
self._enforce_limit()
# 更新过期时间
self._update_expire_time()
@@ -334,6 +383,7 @@ class CacheList(Generic[T]):
self.clear()
self._data.insert(index, CacheData(value=value))
self._enforce_limit()
# 更新过期时间
self._update_expire_time()
@@ -389,6 +439,32 @@ class CacheList(Generic[T]):
# 重置过期时间
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:
"""查找值的索引
@@ -438,6 +514,12 @@ class CacheList(Generic[T]):
"""更新过期时间"""
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:
"""字符串表示
zhenxun/services/cache/runtime_cache.py
Vendored
+169 -42
View File
@@ -10,7 +10,13 @@ import uuid
from zhenxun.services.cache.config import CacheMode
from zhenxun.services.log import logger
from zhenxun.utils.enum import LimitCheckType, LimitWatchType, PluginLimitType
from zhenxun.utils.enum import (
BlockType,
LimitCheckType,
LimitWatchType,
PluginLimitType,
PluginType,
)
from zhenxun.utils.manager.priority_manager import PriorityLifecycle
if TYPE_CHECKING:
@@ -18,24 +24,6 @@ if TYPE_CHECKING:
LOG_COMMAND = "RuntimeCache"
PLUGININFO_MEM_REFRESH_INTERVAL = 1800 # 30分钟 - 插件信息很少变化
BAN_MEM_REFRESH_INTERVAL = 60
BAN_MEM_CLEAN_INTERVAL = 60
BAN_MEM_CLEANUP_DB = True
BAN_MEM_NEGATIVE_TTL = 5
BOT_MEM_REFRESH_INTERVAL = 300 # 5分钟
BOT_MEM_NEGATIVE_TTL = 60
GROUP_MEM_REFRESH_INTERVAL = 300 # 5分钟 - 群组信息很少变化
GROUP_MEM_NEGATIVE_TTL = 60
LEVEL_MEM_REFRESH_INTERVAL = 300 # 5分钟 - 用户等级很少变化
LEVEL_MEM_NEGATIVE_TTL = 60
TASK_MEM_REFRESH_INTERVAL = 900
TASK_MEM_NEGATIVE_TTL = 60
LIMIT_MEM_REFRESH_INTERVAL = 300 # 5分钟
LIMIT_MEM_NEGATIVE_TTL = 30
RUNTIME_CACHE_SYNC_ENABLED = True
RUNTIME_CACHE_SYNC_CHANNEL = "ZHENXUN_RUNTIME_CACHE_SYNC"
def _coerce_int(value, default: int) -> int:
try:
@@ -45,6 +33,27 @@ def _coerce_int(value, default: int) -> int:
return value_int if value_int >= 0 else default
# RuntimeCache 以模型 save/delete 主动失效为主，周期 refresh 只做兜底。
# 这些默认值避免低压力运行时频繁全量扫表。
PLUGININFO_MEM_REFRESH_INTERVAL = 3600 # 60分钟
BAN_MEM_REFRESH_INTERVAL = 300
BAN_MEM_CLEAN_INTERVAL = 60
BAN_MEM_CLEANUP_DB = True
BAN_MEM_NEGATIVE_TTL = 5
BOT_MEM_REFRESH_INTERVAL = 900 # 15分钟
BOT_MEM_NEGATIVE_TTL = 60
GROUP_MEM_REFRESH_INTERVAL = 900 # 15分钟
GROUP_MEM_NEGATIVE_TTL = 60
LEVEL_MEM_REFRESH_INTERVAL = 900 # 15分钟
LEVEL_MEM_NEGATIVE_TTL = 60
TASK_MEM_REFRESH_INTERVAL = 1800
TASK_MEM_NEGATIVE_TTL = 60
LIMIT_MEM_REFRESH_INTERVAL = 900 # 15分钟
LIMIT_MEM_NEGATIVE_TTL = 30
RUNTIME_CACHE_SYNC_ENABLED = True
RUNTIME_CACHE_SYNC_CHANNEL = "ZHENXUN_RUNTIME_CACHE_SYNC"
INSTANCE_ID = uuid.uuid4().hex
_CACHE_READY_EVENT = asyncio.Event()
@@ -89,6 +98,89 @@ def _parse_block_modules(value: str) -> frozenset[str]:
return frozenset(items)
@dataclass(frozen=True)
class PluginInfoSnapshot:
id: int
module: str
module_path: str
name: str
status: bool
block_type: BlockType | None
load_status: bool
author: str | None
version: str | None
level: int
default_status: bool
limit_superuser: bool
menu_type: str
plugin_type: PluginType | None
cost_gold: int
admin_level: int | None
ignore_prompt: bool
is_delete: bool
parent: str | None
is_show: bool
ignore_statistics: bool
impression: float
@classmethod
def from_model(cls, model) -> "PluginInfoSnapshot":
return cls(
id=int(getattr(model, "id", 0) or 0),
module=str(getattr(model, "module", "") or ""),
module_path=str(getattr(model, "module_path", "") or ""),
name=str(getattr(model, "name", "") or ""),
status=bool(getattr(model, "status", True)),
block_type=getattr(model, "block_type", None),
load_status=bool(getattr(model, "load_status", True)),
author=getattr(model, "author", None),
version=getattr(model, "version", None),
level=int(getattr(model, "level", 0) or 0),
default_status=bool(getattr(model, "default_status", True)),
limit_superuser=bool(getattr(model, "limit_superuser", False)),
menu_type=str(getattr(model, "menu_type", "") or ""),
plugin_type=getattr(model, "plugin_type", None),
cost_gold=int(getattr(model, "cost_gold", 0) or 0),
admin_level=getattr(model, "admin_level", None),
ignore_prompt=bool(getattr(model, "ignore_prompt", False)),
is_delete=bool(getattr(model, "is_delete", False)),
parent=getattr(model, "parent", None),
is_show=bool(getattr(model, "is_show", True)),
ignore_statistics=bool(getattr(model, "ignore_statistics", False)),
impression=float(getattr(model, "impression", 0) or 0),
)
def to_model(self):
from zhenxun.models.plugin_info import PluginInfo
plugin = PluginInfo(
id=self.id,
module=self.module,
module_path=self.module_path,
name=self.name,
status=self.status,
block_type=self.block_type,
load_status=self.load_status,
author=self.author,
version=self.version,
level=self.level,
default_status=self.default_status,
limit_superuser=self.limit_superuser,
menu_type=self.menu_type,
plugin_type=self.plugin_type,
cost_gold=self.cost_gold,
admin_level=self.admin_level,
ignore_prompt=self.ignore_prompt,
is_delete=self.is_delete,
parent=self.parent,
is_show=self.is_show,
ignore_statistics=self.ignore_statistics,
impression=self.impression,
)
plugin._saved_in_db = True
return plugin
@dataclass(frozen=True)
class BanEntry:
user_id: str | None
@@ -465,6 +557,20 @@ class RuntimeCacheSync:
@classmethod
async def stop(cls) -> None:
cls._ready = False
if cls._publish_tasks:
tasks = list(cls._publish_tasks)
try:
await asyncio.wait_for(
asyncio.gather(*tasks, return_exceptions=True),
timeout=1.0,
)
except asyncio.TimeoutError:
for task in tasks:
if not task.done():
task.cancel()
finally:
cls._publish_tasks.difference_update(tasks)
if cls._task and not cls._task.done():
cls._task.cancel()
cls._task = None
@@ -480,7 +586,6 @@ class RuntimeCacheSync:
except Exception:
pass
cls._redis = None
cls._ready = False
@classmethod
def publish_event(cls, cache_type: str, action: str, data: dict[str, Any]) -> None:
@@ -559,25 +664,43 @@ class RuntimeCacheSync:
class PluginInfoMemoryCache:
_lock: ClassVar[asyncio.Lock] = asyncio.Lock()
_by_module: ClassVar[dict[str, "PluginInfo"]] = {}
_by_module_path: ClassVar[dict[str, "PluginInfo"]] = {}
_by_module: ClassVar[dict[str, PluginInfoSnapshot]] = {}
_by_module_path: ClassVar[dict[str, PluginInfoSnapshot]] = {}
_loaded: ClassVar[bool] = False
_refresh_task: ClassVar[asyncio.Task | None] = None
_last_refresh: ClassVar[float] = 0.0
@classmethod
def _to_model(cls, snapshot: PluginInfoSnapshot | None) -> "PluginInfo | None":
return snapshot.to_model() if snapshot else None
@classmethod
def _store_snapshot(cls, snapshot: PluginInfoSnapshot) -> None:
if snapshot.module:
old = cls._by_module.get(snapshot.module)
if old and old.module_path != snapshot.module_path:
cls._by_module_path.pop(old.module_path, None)
cls._by_module[snapshot.module] = snapshot
if snapshot.module_path:
old = cls._by_module_path.get(snapshot.module_path)
if old and old.module != snapshot.module:
cls._by_module.pop(old.module, None)
cls._by_module_path[snapshot.module_path] = snapshot
@classmethod
async def refresh(cls) -> None:
from zhenxun.models.plugin_info import PluginInfo
async with cls._lock:
plugins = await PluginInfo.all()
by_module: dict[str, "PluginInfo"] = {}
by_module_path: dict[str, "PluginInfo"] = {}
by_module: dict[str, PluginInfoSnapshot] = {}
by_module_path: dict[str, PluginInfoSnapshot] = {}
for plugin in plugins:
if plugin.module:
by_module[plugin.module] = plugin
if plugin.module_path:
by_module_path[plugin.module_path] = plugin
snapshot = PluginInfoSnapshot.from_model(plugin)
if snapshot.module:
by_module[snapshot.module] = snapshot
if snapshot.module_path:
by_module_path[snapshot.module_path] = snapshot
cls._by_module = by_module
cls._by_module_path = by_module_path
cls._loaded = True
@@ -596,42 +719,42 @@ class PluginInfoMemoryCache:
async def get_by_module(cls, module: str) -> "PluginInfo | None":
if not cls._loaded:
await cls.ensure_loaded()
return cls._by_module.get(module)
return cls._to_model(cls._by_module.get(module))
@classmethod
async def get_all(cls) -> dict[str, "PluginInfo"]:
if not cls._loaded:
await cls.ensure_loaded()
return dict(cls._by_module)
return {
module: snapshot.to_model() for module, snapshot in cls._by_module.items()
}
@classmethod
def get_by_module_path(cls, module_path: str) -> "PluginInfo | None":
return cls._by_module_path.get(module_path)
return cls._to_model(cls._by_module_path.get(module_path))
@classmethod
def set_plugin(cls, plugin) -> None:
if not plugin:
return
if plugin.module:
cls._by_module[plugin.module] = plugin
if getattr(plugin, "module_path", None):
cls._by_module_path[plugin.module_path] = plugin
snapshot = PluginInfoSnapshot.from_model(plugin)
cls._store_snapshot(snapshot)
cls._loaded = True
cls._last_refresh = time.time()
@classmethod
def remove_by_module(cls, module: str) -> None:
cls._by_module.pop(module, None)
snapshot = cls._by_module.pop(module, None)
if snapshot and snapshot.module_path:
cls._by_module_path.pop(snapshot.module_path, None)
@classmethod
async def upsert_from_model(cls, plugin) -> None:
if not plugin:
return
async with cls._lock:
if getattr(plugin, "module", None):
cls._by_module[plugin.module] = plugin
if getattr(plugin, "module_path", None):
cls._by_module_path[plugin.module_path] = plugin
snapshot = PluginInfoSnapshot.from_model(plugin)
cls._store_snapshot(snapshot)
cls._loaded = True
cls._last_refresh = time.time()
@@ -643,9 +766,13 @@ class PluginInfoMemoryCache:
return
async with cls._lock:
if module:
cls._by_module.pop(module, None)
snapshot = cls._by_module.pop(module, None)
if snapshot and snapshot.module_path:
cls._by_module_path.pop(snapshot.module_path, None)
if module_path:
cls._by_module_path.pop(module_path, None)
snapshot = cls._by_module_path.pop(module_path, None)
if snapshot and snapshot.module:
cls._by_module.pop(snapshot.module, None)
@classmethod
async def _refresh_loop(cls, interval: int) -> None:
zhenxun/services/data_access.py
+19 -58
View File
@@ -10,7 +10,11 @@ T = TypeVar("T", bound=Model)
class DataAccess(Generic[T]):
"""数据访问层，根据配置决定是否使用缓存
"""数据访问兼容层，根据配置保留单点缓存读取和清理能力
新的高频运行态路径应优先使用 RuntimeCache 或 BoundedTTLCache。
这里不再把 filter/all/create/update_or_create 结果写入通用缓存，
create/update_or_create 只负责清理旧缓存，避免旧值残留。
使用示例:
```python
@@ -395,34 +399,25 @@ class DataAccess(Generic[T]):
return COMPOSITE_KEY_SEPARATOR.join(key_parts)
async def _cache_items(self, data_list: list[T]) -> None:
"""将数据列表存入缓存
参数:
data_list: 数据列表
"""
if (
not data_list
or not self.cache_type
or cache_config.cache_mode == CacheMode.NONE
):
async def _invalidate_item_cache(self, item: T, action: str) -> None:
if not self.cache_type or cache_config.cache_mode == CacheMode.NONE:
return
try:
# 遍历数据列表，将每条数据存入缓存
cached_count = 0
for item in data_list:
cache_key = self._build_cache_key_for_item(item)
if cache_key is not None:
await self.cache.set(cache_key, item)
cached_count += 1
self._cache_stats[self.cache_type]["sets"] += 1
cache_key = self._build_cache_key_for_item(item)
if cache_key is None:
return
await self.cache.delete(cache_key)
self._cache_stats[self.cache_type]["deletes"] += 1
logger.debug(
f"{self.model_cls.__name__} 批量缓存: {cached_count}/{len(data_list)}项"
f"{self.model_cls.__name__} {action}: 已失效兼容缓存: {cache_key}"
)
except Exception as e:
logger.error(f"{self.model_cls.__name__} 批量缓存失败", e=e)
logger.error(
f"{self.model_cls.__name__} {action}: 更新兼容缓存失败",
e=e,
)
async def filter(self, *args, **kwargs) -> list[T]:
"""筛选数据
@@ -441,9 +436,6 @@ class DataAccess(Generic[T]):
f"{self.model_cls.__name__} filter: 查询结果数量: {len(data_list)}"
)
# 将数据存入缓存
await self._cache_items(data_list)
return data_list
async def all(self) -> list[T]:
@@ -457,9 +449,6 @@ class DataAccess(Generic[T]):
data_list = await self.model_cls.all()
logger.debug(f"{self.model_cls.__name__} all: 查询结果数量: {len(data_list)}")
# 将数据存入缓存
await self._cache_items(data_list)
return data_list
async def count(self, *args, **kwargs) -> int:
@@ -501,24 +490,7 @@ class DataAccess(Generic[T]):
logger.debug(f"{self.model_cls.__name__} create: 创建数据, 参数: {kwargs}")
data = await self.model_cls.create(**kwargs)
# 如果有缓存类型，将数据存入缓存
if self.cache_type and cache_config.cache_mode != CacheMode.NONE:
try:
# 生成缓存键
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
logger.debug(
f"{self.model_cls.__name__} create: "
f"新创建的数据已存入缓存: {cache_key}"
)
except Exception as e:
logger.error(
f"{self.model_cls.__name__} create: 存入缓存失败，参数: {kwargs}",
e=e,
)
await self._invalidate_item_cache(data, "create")
return data
@@ -539,18 +511,7 @@ class DataAccess(Generic[T]):
defaults=defaults, **kwargs
)
# 如果有缓存类型，将数据存入缓存
if self.cache_type and cache_config.cache_mode != CacheMode.NONE:
try:
# 生成缓存键
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
logger.debug(f"更新或创建的数据已存入缓存: {cache_key}")
except Exception as e:
logger.error(f"存入缓存失败，参数: {kwargs}", e=e)
await self._invalidate_item_cache(data, "update_or_create")
return data, created
zhenxun/services/db_context/base_model.py
+1 -1
View File
@@ -150,7 +150,7 @@ class Model(TortoiseModel):
obj = await cls.filter(**kwargs).using_db(connection).get()
result = (obj, False)
if cache_type := cls.get_cache_type():
if result[1] and (cache_type := cls.get_cache_type()):
await CacheRoot.invalidate_cache(cache_type, cls.get_cache_key(result[0]))
return result
zhenxun/services/group_settings_service.py
+6 -3
View File
@@ -5,10 +5,9 @@ import ujson as json
from zhenxun.configs.config import Config
from zhenxun.models.group_plugin_setting import GroupPluginSetting
from zhenxun.services.cache import Cache
from zhenxun.services.cache import BoundedTTLCache
from zhenxun.services.data_access import DataAccess
from zhenxun.services.log import logger
from zhenxun.utils.enum import CacheType
from zhenxun.utils.pydantic_compat import model_dump, model_validate, parse_as
T = TypeVar("T", bound=BaseModel)
@@ -22,7 +21,11 @@ class GroupSettingsService:
def __init__(self):
self.dao = DataAccess(GroupPluginSetting)
self._cache = Cache[dict[str, Any]](CacheType.GROUP_PLUGIN_SETTINGS_VIEW)
self._cache = BoundedTTLCache[str, dict[str, Any]](
"GROUP_PLUGIN_SETTINGS_VIEW",
ttl_seconds=600,
max_items=10000,
)
@staticmethod
def _build_cache_key(group_id: str, plugin_name: str) -> str:
zhenxun/services/memory_governor.py
+307
View File
@@ -0,0 +1,307 @@
from __future__ import annotations
import asyncio
import contextlib
import gc
import inspect
import sys
import time
from typing import Any
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.log import logger
from zhenxun.services.message_load import idle_seconds, is_overloaded
LOG_COMMAND = "MemoryGovernor"
IDLE_CHECK_INTERVAL_SECONDS = 60
IDLE_RECLAIM_SECONDS = 600
RECLAIM_COOLDOWN_SECONDS = 3 * 60 * 60
RECLAIM_TIMEOUT_SECONDS = 10
_task: asyncio.Task | None = None
_reclaim_lock = asyncio.Lock()
_last_reclaim_at = 0.0
def _cooldown_left(now: float | None = None) -> float:
now = time.monotonic() if now is None else now
return max(0.0, _last_reclaim_at + RECLAIM_COOLDOWN_SECONDS - now)
async def start_memory_governor() -> None:
global _task
if _task is not None and not _task.done():
return
if IDLE_CHECK_INTERVAL_SECONDS <= 0 or IDLE_RECLAIM_SECONDS <= 0:
logger.info("idle memory governor disabled", LOG_COMMAND)
return
_task = asyncio.create_task(_idle_reclaim_loop())
async def stop_memory_governor() -> None:
global _task
task = _task
_task = None
if task is not None:
task.cancel()
with contextlib.suppress(BaseException):
await task
async def _idle_reclaim_loop() -> None:
while True:
await asyncio.sleep(IDLE_CHECK_INTERVAL_SECONDS)
if not await _should_reclaim():
continue
if _reclaim_lock.locked():
continue
async with _reclaim_lock:
if not await _should_reclaim():
continue
try:
await asyncio.wait_for(
_run_reclaim(),
timeout=max(RECLAIM_TIMEOUT_SECONDS, 1),
)
except asyncio.TimeoutError:
logger.warning("idle memory reclaim timed out", LOG_COMMAND)
except Exception as exc:
logger.warning("idle memory reclaim failed", LOG_COMMAND, e=exc)
async def _should_reclaim() -> bool:
if _cooldown_left() > 0:
return False
if idle_seconds() < IDLE_RECLAIM_SECONDS:
return False
if is_overloaded():
return False
if await _has_active_auth_work():
return False
return not await _has_active_render_work()
async def _has_active_auth_work() -> bool:
module = sys.modules.get("zhenxun.builtin_plugins.hooks.auth_checker")
if module is None:
return False
hooks_active = int(getattr(module, "HOOKS_ACTIVE_COUNT", 0) or 0)
db_active = int(getattr(module, "DB_ACTIVE_COUNT", 0) or 0)
return hooks_active > 0 or db_active > 0
async def _has_active_render_work() -> bool:
module = sys.modules.get("zhenxun.services.renderer.engine")
if module is None:
return False
manager = getattr(module, "engine_manager", None)
engine = getattr(manager, "_instance", None)
if engine is None:
return False
try:
snapshot = await asyncio.wait_for(engine.get_runtime_snapshot(), timeout=1.0)
except Exception:
return True
if snapshot.get("active_renders", 0):
return True
if snapshot.get("htmlrender_active_tasks", 0):
return True
active_generation = snapshot.get("active_generation")
if isinstance(active_generation, dict) and active_generation.get(
"active_leases", 0
):
return True
retiring = snapshot.get("retiring_generations", [])
if isinstance(retiring, list):
return any(
isinstance(item, dict) and item.get("active_leases", 0) for item in retiring
)
return False
async def _run_reclaim() -> None:
global _last_reclaim_at
start = time.monotonic()
before_rss = _get_total_rss()
cleared: dict[str, Any] = {}
cache_stats_before = {
"cache_dict": CacheDict.stats_all(),
"cache_list": CacheList.stats_all(),
"bounded_ttl": await BoundedTTLCache.stats_all(),
}
cleared["statistics"] = await _flush_statistics_buffer()
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()
cleared["renderer_runtime"] = await _clear_renderer_runtime_caches()
cleared["message_manager"] = _clear_message_manager_cache()
cleared["aiocache_memory"] = await _clear_simple_memory_backend()
collected = gc.collect(2)
malloc_trimmed = _malloc_trim()
after_rss = _get_total_rss()
_last_reclaim_at = time.monotonic()
logger.info(
"idle memory reclaim completed: "
f"cost={time.monotonic() - start:.3f}s "
f"rss_before={_format_bytes(before_rss)} "
f"rss_after={_format_bytes(after_rss)} "
f"gc={collected} malloc_trim={malloc_trimmed} "
f"cleared={cleared} cache_stats_before={cache_stats_before}",
LOG_COMMAND,
)
async def _flush_statistics_buffer() -> int:
module = sys.modules.get("zhenxun.builtin_plugins.statistics.statistics_hook")
if module is None:
return 0
flush = getattr(module, "_flush_statistics_buffer", None)
if flush is None:
return 0
result = await flush("内存回收")
return int(result or 0)
async def _flush_user_gold_log_buffer() -> int:
module = sys.modules.get("zhenxun.services.buffered_writers")
if module is None:
return 0
flush = getattr(module, "flush_user_gold_log_buffer", None)
if flush is None:
return 0
result = await flush("内存回收")
return int(result or 0)
async def _clear_simple_memory_backend() -> bool:
backend = getattr(CacheRoot, "_cache_backend", None)
if not isinstance(backend, SimpleMemoryCache):
return False
await backend.clear()
return True
def _clear_runtime_negative_caches() -> dict[str, int]:
module = sys.modules.get("zhenxun.services.cache.runtime_cache")
if module is None:
return {}
result: dict[str, int] = {}
for name in (
"BotMemoryCache",
"GroupMemoryCache",
"LevelUserMemoryCache",
"TaskInfoMemoryCache",
"PluginLimitMemoryCache",
"BanMemoryCache",
):
cache_cls = getattr(module, name, None)
negative = getattr(cache_cls, "_negative", None)
if isinstance(negative, dict) and negative:
result[name] = len(negative)
negative.clear()
return result
def _clear_auth_local_caches() -> dict[str, int]:
module = sys.modules.get("zhenxun.builtin_plugins.hooks.auth_checker")
if module is None:
return {}
result: dict[str, int] = {}
for name in (
"_MATCHER_COMMAND_TYPE_CACHE",
"_MATCHER_COMMAND_LITERAL_CACHE",
"_MATCHER_ALCONNA_SHORTCUT_CACHE",
):
cache = getattr(module, name, None)
if isinstance(cache, dict) and cache:
result[name] = len(cache)
cache.clear()
return result
def _clear_avatar_memory_cache() -> int:
module = sys.modules.get("zhenxun.services.avatar_service")
if module is None:
return 0
service = getattr(module, "avatar_service", None)
clear = getattr(service, "clear_memory_cache", None)
if not callable(clear):
return 0
result = clear()
return result if isinstance(result, int) and result > 0 else 0
async def _clear_renderer_runtime_caches() -> dict[str, int]:
module = sys.modules.get("zhenxun.services.renderer.service")
if module is None:
return {}
service = getattr(module, "renderer_service", None)
clear = getattr(service, "clear_runtime_caches", None)
if not callable(clear):
return {}
result = clear()
if inspect.isawaitable(result):
result = await result
if not isinstance(result, dict):
return {}
return {
str(key): int(value)
for key, value in result.items()
if isinstance(value, int) and value > 0
}
def _clear_message_manager_cache() -> int:
module = sys.modules.get("zhenxun.utils.manager.message_manager")
if module is None:
return 0
manager_cls = getattr(module, "MessageManager", None)
clear = getattr(manager_cls, "clear_all", None)
if not callable(clear):
return 0
result = clear()
return result if isinstance(result, int) and result > 0 else 0
def _get_total_rss() -> int | None:
try:
import psutil
process = psutil.Process()
total = process.memory_info().rss
for child in process.children(recursive=True):
with contextlib.suppress(Exception):
total += child.memory_info().rss
return int(total)
except Exception:
return None
def _malloc_trim() -> bool:
if sys.platform.startswith(("win", "darwin")):
return False
try:
import ctypes
libc = ctypes.CDLL("libc.so.6")
return bool(libc.malloc_trim(0))
except Exception:
return False
def _format_bytes(value: int | None) -> str:
if value is None:
return "unknown"
return f"{value / 1024 / 1024:.2f}MiB"
zhenxun/services/message_load.py
+11
View File
@@ -3,6 +3,17 @@ from __future__ import annotations
import time
_OVERLOAD_UNTIL = 0.0
_LAST_ACTIVITY = time.monotonic()
def mark_activity() -> None:
"""Record lightweight runtime activity for idle-only maintenance jobs."""
global _LAST_ACTIVITY
_LAST_ACTIVITY = time.monotonic()
def idle_seconds() -> float:
return max(0.0, time.monotonic() - _LAST_ACTIVITY)
def signal_overload(duration: float = 5.0) -> None:
zhenxun/services/renderer/engine.py
+62 -23
View File
@@ -20,6 +20,12 @@ from zhenxun.services.log import logger
from .types import BaseScreenshotEngine
_PLAYWRIGHT_DISCONNECT_ERROR = "Connection closed while reading from the driver"
_PLAYWRIGHT_TARGET_CLOSED_ERROR_MARKERS = (
"TargetClosedError",
"Target page, context or browser has been closed",
"browser has been closed",
"BrowserContext.new_page",
)
_UNRETRIEVED_FUTURE_MESSAGE = "Future exception was never retrieved"
_LOOP_EXCEPTION_FILTER_STATE_ATTR = "_zhenxun_playwright_exception_filter_state"
_DISCONNECT_SUPPRESSION_WINDOW_SECONDS = 10.0
@@ -135,6 +141,14 @@ def _is_ignorable_playwright_disconnect(ctx: dict[str, Any]) -> bool:
)
def _is_playwright_target_closed_error(exc: Exception) -> bool:
exc_name = type(exc).__name__
if exc_name == "TargetClosedError":
return True
message = str(exc)
return any(marker in message for marker in _PLAYWRIGHT_TARGET_CLOSED_ERROR_MARKERS)
def _get_loop_exception_filter_state(
loop: asyncio.AbstractEventLoop,
) -> dict[str, Any] | None:
@@ -1103,29 +1117,54 @@ class PlaywrightEngine(BaseScreenshotEngine):
template_path: str,
render_options: dict[str, Any],
) -> bytes:
generation, context = await self._acquire_context()
page = None
broken = False
try:
page = await context.new_page()
page_options = self._build_page_options(render_options, pooled=True)
viewport = page_options.get("viewport")
if isinstance(viewport, dict):
width = viewport.get("width")
height = viewport.get("height")
if isinstance(width, int) and isinstance(height, int):
await page.set_viewport_size({"width": width, "height": height})
return await self._render_with_page(
page, html, template_path, render_options
)
except Exception:
broken = True
raise
finally:
if page is not None:
with contextlib.suppress(Exception):
await page.close()
await self._release_context(generation, context, broken=broken)
last_error: Exception | None = None
for attempt in range(2):
generation, context = await self._acquire_context()
page = None
broken = False
try:
page = await context.new_page()
page_options = self._build_page_options(render_options, pooled=True)
viewport = page_options.get("viewport")
if isinstance(viewport, dict):
width = viewport.get("width")
height = viewport.get("height")
if isinstance(width, int) and isinstance(height, int):
await page.set_viewport_size({"width": width, "height": height})
return await self._render_with_page(
page, html, template_path, render_options
)
except Exception as e:
broken = True
last_error = e
if attempt == 0:
if _is_playwright_target_closed_error(e):
logger.warning(
"截图引擎浏览器上下文代已失效，切换新代后重试一次。",
"PlaywrightEngine",
e=e,
)
try:
await self._swap_generation("target_closed")
except Exception:
raise e
else:
logger.warning(
"截图引擎上下文已失效，丢弃后重试一次。",
"PlaywrightEngine",
e=e,
)
continue
raise
finally:
if page is not None:
with contextlib.suppress(Exception):
await page.close()
await self._release_context(generation, context, broken=broken)
if last_error is not None:
raise last_error
raise RuntimeError("截图引擎上下文池渲染失败。")
async def _render_html(
self,
zhenxun/services/renderer/result_cache.py
+9 -55
View File
@@ -1,11 +1,9 @@
from __future__ import annotations
import asyncio
from collections import OrderedDict
import hashlib
import time
from typing import Any
from zhenxun.services.cache.bounded_ttl import BoundedTTLCache
from zhenxun.utils.pydantic_compat import dump_json_safely
@@ -23,9 +21,12 @@ class RenderResultMemoryCache:
if isinstance(max_total_bytes, int) and max_total_bytes > 0
else None
)
self._cache: OrderedDict[str, tuple[float, bytes]] = OrderedDict()
self._total_bytes = 0
self._lock = asyncio.Lock()
self._cache = BoundedTTLCache[str, bytes](
"RENDER_RESULT",
ttl_seconds=self._ttl_seconds,
max_items=self._max_items,
max_total_bytes=self._max_total_bytes,
)
@staticmethod
def build_key(payload: Any) -> str:
@@ -37,55 +38,8 @@ class RenderResultMemoryCache:
)
return hashlib.sha256(payload_text.encode("utf-8")).hexdigest()
def _pop_oldest(self) -> None:
if not self._cache:
return
_, (_, value) = self._cache.popitem(last=False)
self._total_bytes -= len(value)
if self._total_bytes < 0:
self._total_bytes = 0
def _cleanup(self, now: float) -> None:
while self._cache:
expire_at, _ = next(iter(self._cache.values()))
if expire_at > now:
break
self._pop_oldest()
while len(self._cache) > self._max_items:
self._pop_oldest()
if self._max_total_bytes is not None:
while self._total_bytes > self._max_total_bytes and self._cache:
self._pop_oldest()
async def get(self, key: str) -> bytes | None:
now = time.monotonic()
async with self._lock:
self._cleanup(now)
item = self._cache.get(key)
if item is None:
return None
expire_at, value = item
if expire_at <= now:
removed = self._cache.pop(key, None)
if removed:
self._total_bytes -= len(removed[1])
if self._total_bytes < 0:
self._total_bytes = 0
return None
self._cache.move_to_end(key)
return value
return await self._cache.get(key)
async def set(self, key: str, value: bytes) -> None:
value_size = len(value)
if self._max_total_bytes is not None and value_size > self._max_total_bytes:
return
now = time.monotonic()
async with self._lock:
if old := self._cache.pop(key, None):
self._total_bytes -= len(old[1])
if self._total_bytes < 0:
self._total_bytes = 0
self._cache[key] = (now + self._ttl_seconds, value)
self._total_bytes += value_size
self._cache.move_to_end(key)
self._cleanup(now)
await self._cache.set(key, value)
zhenxun/services/renderer/service.py
+12
View File
@@ -475,6 +475,18 @@ class RendererService:
raise RuntimeError("ThemeManager尚未初始化。")
return self._theme_manager.list_available_themes()
def clear_runtime_caches(self) -> dict[str, int]:
cleared: dict[str, int] = {}
if self._theme_manager:
cleared.update(self._theme_manager.clear_runtime_caches())
if self._template_engine and self._template_engine.env.cache:
jinja_cache = self._template_engine.env.cache
cache_size = len(jinja_cache)
jinja_cache.clear()
if cache_size:
cleared["jinja_env"] = cache_size
return cleared
async def switch_theme(self, theme_name: str) -> str:
"""
切换UI主题，加载新主题并持久化配置。
zhenxun/services/renderer/theme.py
+18 -1
View File
@@ -51,8 +51,10 @@ class ManifestRegistry:
self._manifest_cache: dict[str, TemplateManifest] = {}
self._lock = asyncio.Lock()
def clear_cache(self):
def clear_cache(self) -> int:
size = len(self._manifest_cache)
self._manifest_cache.clear()
return size
async def get_manifest(
self, component_path: str, skin: str | None = None
@@ -362,6 +364,21 @@ class ThemeManager:
tuple[type, str, str | None], ComponentDependency
] = OrderedDict()
def clear_runtime_caches(self) -> dict[str, int]:
cleared = {
"asset_resolution": len(self._asset_resolution_cache),
"global_template": len(self._global_template_cache),
"component_dependency": len(self._component_dependency_cache),
}
self._asset_resolution_cache.clear()
self._global_template_cache.clear()
self._component_dependency_cache.clear()
if self.manifest_registry:
manifest_count = self.manifest_registry.clear_cache()
if manifest_count:
cleared["manifest"] = manifest_count
return {key: value for key, value in cleared.items() if value}
@staticmethod
def _get_lru_entry(cache: OrderedDict, key: Any) -> Any:
value = cache.get(key)
zhenxun/services/runtime_bootstrap.py
+63 -3
View File
@@ -2,10 +2,17 @@ import asyncio
from concurrent.futures import ThreadPoolExecutor
import contextlib
import os
import signal
import anyio.to_thread
from nonebot.drivers import Driver
from zhenxun.services.log import logger
from zhenxun.services.memory_governor import (
start_memory_governor,
stop_memory_governor,
)
from zhenxun.services.send_queue import start_send_queue, stop_send_queue
from zhenxun.services.uninfo_patch import apply_uninfo_onebot11_patch
from zhenxun.utils.manager.priority_manager import PriorityLifecycle
DEFAULT_EXECUTOR_MIN_WORKERS = 16
@@ -14,6 +21,7 @@ DEFAULT_ANYIO_MIN_TOKENS = 32
DEFAULT_ANYIO_MAX_TOKENS = 128
_thread_executor: ThreadPoolExecutor | None = None
_launcher_watchdog_task: asyncio.Task[None] | None = None
_runtime_hooks_registered = False
_alconna_patch_applied = False
@@ -57,14 +65,60 @@ def _apply_alconna_conflict_patch() -> None:
_alconna_patch_applied = True
def register_runtime_bootstrap(driver: Driver) -> None:
async def _launcher_watchdog_loop(launcher_pid: int) -> None:
try:
import psutil
except Exception:
return
current_pid = os.getpid()
while True:
await asyncio.sleep(2)
if psutil.pid_exists(launcher_pid):
continue
logger.warning(
f"检测到 launcher 进程 {launcher_pid} 已退出，worker 将主动结束...",
"RuntimeBootstrap",
)
with contextlib.suppress(Exception):
os.kill(current_pid, signal.SIGTERM)
return
def _start_launcher_watchdog() -> None:
global _launcher_watchdog_task
if _launcher_watchdog_task is not None and not _launcher_watchdog_task.done():
return
launcher_pid_text = os.getenv("ZHENXUN_LAUNCHER_PID", "").strip()
if not launcher_pid_text:
return
with contextlib.suppress(ValueError):
launcher_pid = int(launcher_pid_text)
if launcher_pid > 0:
_launcher_watchdog_task = asyncio.create_task(
_launcher_watchdog_loop(launcher_pid)
)
async def _stop_launcher_watchdog() -> None:
global _launcher_watchdog_task
task = _launcher_watchdog_task
_launcher_watchdog_task = None
if task is None or task.done():
return
task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await task
def register_runtime_bootstrap(_driver) -> None:
_apply_alconna_conflict_patch()
apply_uninfo_onebot11_patch()
global _runtime_hooks_registered
if _runtime_hooks_registered:
return
_runtime_hooks_registered = True
@driver.on_startup
@PriorityLifecycle.on_startup(priority=-100)
async def _setup_runtime_concurrency() -> None:
global _thread_executor
workers = _get_executor_workers()
@@ -77,10 +131,16 @@ def register_runtime_bootstrap(driver: Driver) -> None:
with contextlib.suppress(Exception):
limiter = anyio.to_thread.current_default_thread_limiter()
limiter.total_tokens = _get_anyio_tokens(workers)
_start_launcher_watchdog()
await start_send_queue()
await start_memory_governor()
@PriorityLifecycle.on_shutdown(priority=50)
async def _shutdown_runtime_concurrency() -> None:
global _thread_executor
await _stop_launcher_watchdog()
await stop_send_queue()
await stop_memory_governor()
executor = _thread_executor
_thread_executor = None
if executor is not None:
zhenxun/services/send_queue.py
+91 -13
View File
@@ -2,21 +2,29 @@ import asyncio
import time
from typing import Any
import nonebot
from nonebot.adapters import Bot
from zhenxun.services.log import logger
_SEND_APIS = {"send_msg", "send_like"}
_QUEUE: asyncio.Queue[tuple[Bot, str, dict[str, Any], asyncio.Future]] = asyncio.Queue()
_SEND_APIS = {"send_msg", "send_group_msg", "send_private_msg", "send_like"}
_WORKERS = 3
_MIN_INTERVAL = 0.05
_QUEUE_MAXSIZE = 2000
_SHUTDOWN_DRAIN_TIMEOUT_SECONDS = 3.0
_QUEUE_PRESSURE_LOG_INTERVAL = 10.0
_QUEUE: asyncio.Queue[tuple[Bot, str, dict[str, Any], asyncio.Future[Any]]] = (
asyncio.Queue(maxsize=_QUEUE_MAXSIZE)
)
_SEND_LOCK = asyncio.Lock()
_LAST_SEND_TS = 0.0
_API_SEMAPHORE = asyncio.Semaphore(3)
_ORIG_CALL_API = Bot.call_api
_PATCHED = False
_WORKER_TASKS: list[asyncio.Task] = []
_QUEUE_TIMEOUT_COUNT = 0
_SEND_LIKE_DROP_COUNT = 0
_LAST_QUEUE_PRESSURE_LOG = 0.0
_STOPPING = False
async def _rate_limit():
@@ -29,15 +37,36 @@ async def _rate_limit():
_LAST_SEND_TS = time.monotonic()
def _log_queue_pressure(reason: str) -> None:
global _LAST_QUEUE_PRESSURE_LOG
now = time.monotonic()
if now - _LAST_QUEUE_PRESSURE_LOG < _QUEUE_PRESSURE_LOG_INTERVAL:
return
_LAST_QUEUE_PRESSURE_LOG = now
logger.warning(
f"{reason}; qsize={_QUEUE.qsize()}/{_QUEUE_MAXSIZE} "
f"timeouts={_QUEUE_TIMEOUT_COUNT} dropped_like={_SEND_LIKE_DROP_COUNT}",
"SendQueue",
)
async def _direct_call_api(bot: Bot, api: str, data: dict[str, Any]) -> Any:
await _rate_limit()
async with _API_SEMAPHORE:
return await _ORIG_CALL_API(bot, api, **data)
async def _worker(worker_id: int):
while True:
bot, api, data, future = await _QUEUE.get()
try:
await _rate_limit()
async with _API_SEMAPHORE:
result = await _ORIG_CALL_API(bot, api, **data)
result = await _direct_call_api(bot, api, data)
if not future.done():
future.set_result(result)
except asyncio.CancelledError:
if not future.done():
future.set_exception(RuntimeError("send queue worker cancelled"))
raise
except Exception as exc:
if not future.done():
future.set_exception(exc)
@@ -54,12 +83,41 @@ async def _worker(worker_id: int):
async def _queued_call_api(self: Bot, api: str, **data: Any):
if api not in _SEND_APIS:
return await _ORIG_CALL_API(self, api, **data)
if _STOPPING:
return await _direct_call_api(self, api, data)
loop = asyncio.get_running_loop()
future: asyncio.Future = loop.create_future()
await _QUEUE.put((self, api, data, future))
future: asyncio.Future[Any] = loop.create_future()
queue_item = (self, api, data, future)
try:
_QUEUE.put_nowait(queue_item)
except asyncio.QueueFull:
if api == "send_like":
global _SEND_LIKE_DROP_COUNT
_SEND_LIKE_DROP_COUNT += 1
_log_queue_pressure("send_like dropped because send queue is full")
return None
global _QUEUE_TIMEOUT_COUNT
_QUEUE_TIMEOUT_COUNT += 1
_log_queue_pressure(f"{api} fallback to direct send because queue is full")
return await _direct_call_api(self, api, data)
return await future
def _drain_pending_futures(reason: str) -> int:
drained = 0
while True:
try:
_, _, _, future = _QUEUE.get_nowait()
except asyncio.QueueEmpty:
break
if not future.done():
future.set_exception(RuntimeError(reason))
_QUEUE.task_done()
drained += 1
return drained
def patch_send_queue() -> None:
global _PATCHED
if _PATCHED:
@@ -68,21 +126,41 @@ def patch_send_queue() -> None:
_PATCHED = True
driver = nonebot.get_driver()
def unpatch_send_queue() -> None:
global _PATCHED
if not _PATCHED:
return
Bot.call_api = _ORIG_CALL_API # type: ignore[assignment]
_PATCHED = False
@driver.on_startup
async def _start_send_queue():
async def start_send_queue() -> None:
global _STOPPING
patch_send_queue()
_STOPPING = False
if _WORKER_TASKS:
return
for idx in range(_WORKERS):
_WORKER_TASKS.append(asyncio.create_task(_worker(idx)))
@driver.on_shutdown
async def _stop_send_queue():
async def stop_send_queue() -> None:
global _STOPPING
_STOPPING = True
try:
await asyncio.wait_for(_QUEUE.join(), timeout=_SHUTDOWN_DRAIN_TIMEOUT_SECONDS)
except asyncio.TimeoutError:
drained = _drain_pending_futures("send queue shutdown before drain completed")
logger.warning(
f"send queue shutdown timed out, dropped pending futures={drained}, "
f"qsize={_QUEUE.qsize()}",
"SendQueue",
)
tasks = _WORKER_TASKS.copy()
_WORKER_TASKS.clear()
for task in tasks:
task.cancel()
if tasks:
await asyncio.gather(*tasks, return_exceptions=True)
unpatch_send_queue()
_STOPPING = False
zhenxun/services/uninfo_patch.py
+149
View File
@@ -0,0 +1,149 @@
import asyncio
from collections.abc import Awaitable, Callable
import contextlib
from typing import Any, cast
from nonebot.adapters import Bot, Event
from nonebot.adapters.onebot.v11.event import GroupMessageEvent
from nonebot.log import logger
_PATCHED = False
_ORIGINAL_FETCH: Callable[..., Awaitable[Any]] | None = None
_ORIGINAL_ONEBOT11_GROUP_MESSAGE: Callable[..., Awaitable[dict[str, Any]]] | None = None
def _sender_value(sender: Any, key: str, default: Any = None) -> Any:
value = getattr(sender, key, default)
return default if value is None else value
def _event_value(event: Event, key: str, default: Any = None) -> Any:
value = getattr(event, key, default)
return default if value is None else value
def _event_group_name(event: Event) -> str | None:
group_name = _event_value(event, "group_name")
if isinstance(group_name, str) and group_name:
return group_name
group = _event_value(event, "group")
if group is not None:
name = _sender_value(group, "name") or _sender_value(group, "group_name")
if isinstance(name, str) and name:
return name
return None
def _has_compatible_onebot11_sender(event: Event) -> bool:
if getattr(event, "_zx_uninfo_full_fetch", False):
return False
sender = _event_value(event, "sender")
if sender is None:
return False
return (
_event_value(event, "user_id") is not None
and _event_value(event, "group_id") is not None
and _sender_value(sender, "nickname") is not None
and _sender_value(sender, "role") is not None
)
async def _fast_onebot11_group_message(bot: Bot, event: Event) -> dict[str, Any]:
"""Build Uninfo session data from OneBot v11 group message event fields.
nonebot-plugin-uninfo's default OneBot v11 fetcher always calls
get_group_info and get_group_member_info for group messages. For normal
matcher rule checks, event-provided sender fields are enough and avoid
multiplying protocol API calls by the number of candidate matchers.
"""
original = _ORIGINAL_ONEBOT11_GROUP_MESSAGE
if not _has_compatible_onebot11_sender(event):
if original is not None:
return await original(bot, event)
logger.debug("Uninfo OneBot11 fast fetch fallback unavailable")
sender = _event_value(event, "sender")
user_id = str(_event_value(event, "user_id", ""))
group_id = str(_event_value(event, "group_id", ""))
nickname = _sender_value(sender, "nickname", "")
card = _sender_value(sender, "card", "") or nickname
return {
"group_id": group_id,
"group_name": _event_group_name(event),
"user_id": user_id,
"name": nickname,
"nickname": card,
"card": card,
"role": _sender_value(sender, "role", "member"),
"join_time": _event_value(event, "join_time"),
"gender": _sender_value(sender, "sex", "unknown") or "unknown",
}
async def _singleflight_fetch(self: Any, bot: Bot, event: Event) -> Any:
original = _ORIGINAL_FETCH
if original is None:
return None
try:
sess_id = self.get_session_id(event)
except ValueError:
return await original(self, bot, event)
session_cache = getattr(self, "session_cache", None)
if isinstance(session_cache, dict) and sess_id in session_cache:
return session_cache[sess_id]
inflight = getattr(self, "_zx_fetch_inflight", None)
if not isinstance(inflight, dict):
inflight = {}
setattr(self, "_zx_fetch_inflight", inflight)
key = (str(getattr(bot, "self_id", "")), event.__class__, sess_id)
task = inflight.get(key)
if task is None or task.done():
task = asyncio.ensure_future(original(self, bot, event))
inflight[key] = task
try:
return await task
finally:
if inflight.get(key) is task and task.done():
inflight.pop(key, None)
def apply_uninfo_onebot11_patch() -> None:
global _ORIGINAL_FETCH, _ORIGINAL_ONEBOT11_GROUP_MESSAGE, _PATCHED
if _PATCHED:
return
with contextlib.suppress(Exception):
from nonebot_plugin_uninfo.adapters.onebot11.main import fetcher
original_endpoint = fetcher.endpoint.get(GroupMessageEvent)
if not getattr(original_endpoint, "__zhenxun_fast_onebot11__", False):
_ORIGINAL_ONEBOT11_GROUP_MESSAGE = cast(
Callable[..., Awaitable[dict[str, Any]]] | None,
original_endpoint,
)
setattr(_fast_onebot11_group_message, "__zhenxun_fast_onebot11__", True)
fetcher.endpoint[GroupMessageEvent] = _fast_onebot11_group_message
try:
from nonebot_plugin_uninfo.fetch import InfoFetcher
except Exception as e:
logger.warning("Uninfo patch skipped", e=e)
return
original_fetch = getattr(InfoFetcher, "fetch", None)
if getattr(original_fetch, "__zhenxun_singleflight__", False):
_PATCHED = True
return
if original_fetch is None:
return
_ORIGINAL_FETCH = cast(Callable[..., Awaitable[Any]], original_fetch)
setattr(_singleflight_fetch, "__zhenxun_singleflight__", True)
setattr(InfoFetcher, "fetch", _singleflight_fetch)
_PATCHED = True
logger.debug("Uninfo OneBot11 fast fetch and singleflight patch applied")
zhenxun/utils/enum.py
-2
View File
@@ -55,8 +55,6 @@ class CacheType(StrEnum):
"""全局全部群组"""
GROUP_PLUGIN_SETTINGS = "GROUP_PLUGIN_SETTINGS"
"""插件分群配置"""
GROUP_PLUGIN_SETTINGS_VIEW = "GROUP_PLUGIN_SETTINGS_VIEW"
"""插件分群配置视图缓存（聚合 dict）"""
USERS = "GLOBAL_ALL_USERS"
"""全部用户"""
BAN = "GLOBAL_ALL_BAN"
zhenxun/utils/http_utils.py
+12 -34
View File
@@ -1,5 +1,4 @@
import asyncio
from collections import OrderedDict
from collections.abc import AsyncGenerator, Awaitable, Callable, Sequence
from contextlib import asynccontextmanager
import os
@@ -21,6 +20,7 @@ from rich.progress import (
import ujson as json
from zhenxun.configs.config import BotConfig
from zhenxun.services.cache.bounded_ttl import BoundedTTLCache
from zhenxun.services.log import logger
from zhenxun.utils.decorator.retry import Retry
from zhenxun.utils.exception import AllURIsFailedError
@@ -168,7 +168,12 @@ class AsyncHttpx:
_CONTENT_CACHE_TTL: ClassVar[float] = 3.0
_CONTENT_CACHE_MAX_ITEMS: ClassVar[int] = 256
_CONTENT_CACHE_MAX_BYTES: ClassVar[int] = 2 * 1024 * 1024
_content_cache: ClassVar[OrderedDict[str, tuple[float, bytes]]] = OrderedDict()
_content_cache: ClassVar[BoundedTTLCache[str, bytes]] = BoundedTTLCache(
"HTTP_IMAGE_CONTENT",
ttl_seconds=_CONTENT_CACHE_TTL,
max_items=_CONTENT_CACHE_MAX_ITEMS,
max_total_bytes=_CONTENT_CACHE_MAX_BYTES,
)
_content_inflight: ClassVar[dict[str, asyncio.Task[Response]]] = {}
_content_cache_lock: ClassVar[asyncio.Lock] = asyncio.Lock()
@@ -191,29 +196,6 @@ class AsyncHttpx:
return True
return "qpic.cn" in lower_url or "qlogo.cn" in lower_url
@classmethod
def _get_cached_content_nolock(cls, key: str) -> bytes | None:
entry = cls._content_cache.get(key)
if not entry:
return None
expire_at, content = entry
if expire_at <= time.monotonic():
cls._content_cache.pop(key, None)
return None
cls._content_cache.move_to_end(key)
return content
@classmethod
def _cleanup_content_cache_nolock(cls) -> None:
now = time.monotonic()
while cls._content_cache:
expire_at, _ = next(iter(cls._content_cache.values()))
if expire_at > now:
break
cls._content_cache.popitem(last=False)
while len(cls._content_cache) > cls._CONTENT_CACHE_MAX_ITEMS:
cls._content_cache.popitem(last=False)
@classmethod
async def _try_cache_content(cls, key: str, response: Response) -> None:
content = response.content
@@ -223,13 +205,7 @@ class AsyncHttpx:
is_image = content_type.startswith("image/") or cls._is_probably_image_url(key)
if not is_image:
return
async with cls._content_cache_lock:
cls._content_cache[key] = (
time.monotonic() + cls._CONTENT_CACHE_TTL,
content,
)
cls._content_cache.move_to_end(key)
cls._cleanup_content_cache_nolock()
await cls._content_cache.set(key, content)
@classmethod
def _prepare_temporary_client_config(cls, client_kwargs: dict) -> dict:
@@ -450,9 +426,11 @@ class AsyncHttpx:
return res.content
cache_key = url
if cached := await cls._content_cache.get(cache_key):
return cached
async with cls._content_cache_lock:
cached = cls._get_cached_content_nolock(cache_key)
if cached is not None:
if cached := await cls._content_cache.get(cache_key):
return cached
task = cls._content_inflight.get(cache_key)
if task is None:
zhenxun/utils/manager/message_manager.py
+60 -8
View File
@@ -1,25 +1,77 @@
from collections import OrderedDict
import time
from typing import ClassVar
class MessageManager:
data: ClassVar[dict[str, list[str]]] = {}
_MAX_USERS: ClassVar[int] = 4096
_MAX_MESSAGES_PER_USER: ClassVar[int] = 200
_TRIM_MESSAGES_TO: ClassVar[int] = 100
_USER_TTL_SECONDS: ClassVar[float] = 6 * 60 * 60
data: ClassVar[OrderedDict[str, tuple[float, list[str]]]] = OrderedDict()
@classmethod
def _prune(cls, now: float | None = None) -> None:
now = time.monotonic() if now is None else now
stale_before = now - cls._USER_TTL_SECONDS
stale_uids = [
uid for uid, (last_seen, _) in cls.data.items() if last_seen <= stale_before
]
for uid in stale_uids:
cls.data.pop(uid, None)
while len(cls.data) > cls._MAX_USERS:
cls.data.popitem(last=False)
@classmethod
def _touch(cls, uid: str, messages: list[str], now: float | None = None) -> None:
now = time.monotonic() if now is None else now
cls.data[uid] = (now, messages)
cls.data.move_to_end(uid)
@classmethod
def add(cls, uid: str, msg_id: str):
if uid not in cls.data:
cls.data[uid] = []
cls.data[uid].append(msg_id)
now = time.monotonic()
cls._prune(now)
_, messages = cls.data.get(uid, (now, []))
messages.append(msg_id)
cls._touch(uid, messages, now)
cls.remove_check(uid)
cls._prune(now)
@classmethod
def check(cls, uid: str, msg_id: str) -> bool:
return msg_id in cls.data.get(uid, [])
now = time.monotonic()
cls._prune(now)
entry = cls.data.get(uid)
if entry is None:
return False
_, messages = entry
cls._touch(uid, messages, now)
return msg_id in messages
@classmethod
def remove_check(cls, uid: str):
if len(cls.data[uid]) > 200:
cls.data[uid] = cls.data[uid][100:]
entry = cls.data.get(uid)
if entry is None:
return
_, messages = entry
if len(messages) > cls._MAX_MESSAGES_PER_USER:
messages = messages[-cls._TRIM_MESSAGES_TO :]
cls._touch(uid, messages)
@classmethod
def get(cls, uid: str) -> list[str]:
return cls.data[uid] if uid in cls.data else []
now = time.monotonic()
cls._prune(now)
entry = cls.data.get(uid)
if entry is None:
return []
_, messages = entry
cls._touch(uid, messages, now)
return list(messages)
@classmethod
def clear_all(cls) -> int:
size = len(cls.data)
cls.data.clear()
return size
zhenxun/utils/manager/virtual_env_package_manager.py
+41
View File
@@ -8,6 +8,7 @@ from zhenxun.configs.config import Config
from zhenxun.services.log import logger
LOG_COMMAND = "VirtualEnvPackageManager"
PROJECT_ROOT = Path(__file__).resolve().parents[3]
Config.add_plugin_config(
"virtualenv",
@@ -191,6 +192,46 @@ class VirtualEnvPackageManager:
)
return stderr
@classmethod
async def add_requirement(cls, requirement_file: Path):
"""将依赖文件写入项目依赖并同步环境
插件商店安装依赖需要持久化到 pyproject.toml/uv.lock，避免重建环境后丢失。
"""
if not requirement_file.exists():
raise FileNotFoundError(f"依赖文件 {requirement_file} 不存在", LOG_COMMAND)
cls._clean_requirements_file(requirement_file)
try:
command = [
"uv",
"add",
"--requirements",
str(requirement_file.absolute()),
]
logger.info(f"执行项目依赖添加指令: {command}", LOG_COMMAND)
result = await asyncio.to_thread(
subprocess.run,
command,
cwd=PROJECT_ROOT,
check=True,
capture_output=True,
text=True,
encoding="utf-8",
errors="replace",
)
logger.debug(
f"项目依赖添加指令执行完成: {result.stdout}",
LOG_COMMAND,
)
return result.stdout
except (CalledProcessError, FileNotFoundError) as e:
stderr = e.stderr if isinstance(e, CalledProcessError) else str(e)
logger.error(
f"项目依赖添加指令执行失败: {stderr}.",
LOG_COMMAND,
)
return stderr
@classmethod
async def list(cls) -> str:
"""列出已安装的依赖包"""