添加uv支持 (#2119)

* bugfix:修复内存泄露和信号量饥饿问题

* 修复图片渲染按高度截断问题

* 优化图片渲染速度

* 权限检查去掉无效引用代码

* 添加uv支持

* 🚨 auto fix by pre-commit hooks

* bugfix:修改gitignore换行

* bugfix:修复测试没有新生成uv.lock

* 修复导入错误

* bugfix:移除重复调用

* 🚨 auto fix by pre-commit hooks

* 清理残余poetry引用

* 更新uv安装方式

* 修复阿里云获取问题

* 增加资源下载提示

* 🚨 auto fix by pre-commit hooks

* 修改资源下载为流式

* 🚨 auto fix by pre-commit hooks

* 提高启动速度

* 移除bot.py支持

* 🚨 auto fix by pre-commit hooks

* 优化win脚本逻辑

* 🚨 auto fix by pre-commit hooks

* 清理残余无效逻辑

* 代码改进

* 🚨 auto fix by pre-commit hooks

* 增加数据库迁移存在性检查

* 🚨 auto fix by pre-commit hooks

* chore(test): 添加pytest超时控制和优雅关闭机制

- 在GitHub Actions工作流中添加作业级和步骤级超时限制,防止测试无限期挂起
- 添加pytest-timeout依赖并配置全局超时为120秒
- 在send_queue服务添加关闭钩子,确保worker任务正确取消
- 在priority_manager添加on_shutdown钩子,支持优先级生命周期的关闭阶段

* chore(lint): 禁用超长行的lint警告

* Modify restart logic for Windows platform

* 🚨 auto fix by pre-commit hooks

* bugfix:修复sys导入问题

* 清理无效结构

* bugfix:修复路径问题

* bugfix:修复shell语法传递给git导致资源获取失败问题

* 优化关闭显示

* bugfix:修复路径问题

* bugfix:增加路径安全

* bugfix:修复orm绕过问题

* 放宽numpy版本限制

* 修改重启方案

* bugfix:修复循环导入

* 优化逻辑

* Enhance disconnect function with error handling

Added error handling for disconnect function and imported ConfigurationError.

* Implement emergency restart mechanism

Added emergency restart mechanism using atexit to ensure process restart even on severe exceptions during shutdown.

* 🚨 auto fix by pre-commit hooks

* 重启行为归一化

* 修复测试检测问题

* bugfix:修复测试侧类型报错问题

* 引入launcher机制

* 移除重启测试

* 收紧缓存调用路径

* 类型注解收敛

* 优化浏览器回收行为

* 优化浏览器渲染

* bugfix:解决重复关闭浏览器问题

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: ManyManyTomato <93612024+ATTomatoo@users.noreply.github.com>
Co-authored-by: AkashiCoin <l1040186796@gmail.com>
This commit is contained in:
Copaan
2026-04-18 23:42:10 +08:00
committed by GitHub
co-authored by pre-commit-ci[bot] ManyManyTomato AkashiCoin
parent 74bf912d04
commit 8b16126e40
83 changed files with 15394 additions and 7062 deletions
-9
View File
@@ -9,15 +9,6 @@ Zhenxun Bot - 核心服务模块
- 定时任务调度器 (scheduler): 提供持久化的、可管理的定时任务服务。
"""
from nonebot import require
require("nonebot_plugin_apscheduler")
require("nonebot_plugin_alconna")
require("nonebot_plugin_session")
require("nonebot_plugin_htmlrender")
require("nonebot_plugin_uninfo")
require("nonebot_plugin_waiter")
from .avatar_service import avatar_service
from .db_context import Model, disconnect, with_db_timeout
from .group_settings_service import group_settings_service
+20 -301
View File
@@ -19,7 +19,7 @@ users = await level_cache.get({"user_id": "123", "group_id": "456"})
await level_cache.set({"user_id": "123", "group_id": "456"}, users)
```
2. 使用CacheDict作为全局字典
2. 使用CacheDict作为内存字典缓存
```python
from zhenxun.services.cache.cache_containers import CacheDict
@@ -29,51 +29,18 @@ config_dict = CacheDict("global_config")
# 创建有过期时间的缓存字典(1小时后过期)
temp_dict = CacheDict("temp_config", expire=3600)
# 使用字典操作
config_dict["key"] = "value"
value = config_dict["key"]
# 保存缓存数据(可选)
await config_dict.save()
value = config_dict.get("key")
```
3. 使用CacheList作为全局列表
```python
from zhenxun.services.cache.cache_containers import CacheList
# 创建缓存列表(默认永不过期)
message_list = CacheList("recent_messages")
# 创建有过期时间的缓存列表(30分钟后过期)
temp_list = CacheList("temp_messages", expire=1800)
# 使用列表操作
message_list.append("新消息")
message = message_list[0]
# 保存缓存数据(可选)
await message_list.save()
```
4. 使用CacheManager的类型化缓存方法
3. 使用CacheRoot直接操作缓存后端
```python
from zhenxun.services.cache import CacheRoot
# 获取字符串类型的缓存字典(向后兼容)
str_cache = CacheRoot.cache_dict("string_cache")
# 获取类型化的缓存字典(推荐)
int_cache = CacheRoot.cache_dict_typed("int_cache", value_type=int)
user_cache = CacheRoot.cache_dict_typed("user_cache", value_type=User)
# 获取类型化的缓存列表
message_list = CacheRoot.cache_list_typed("messages", value_type=str)
user_list = CacheRoot.cache_list_typed("users", value_type=User)
# 使用类型化的缓存
int_cache["count"] = 42 # 类型安全
user_cache["user1"] = User(name="Alice") # 类型安全
message_list.append("Hello") # 类型安全
# 获取/设置缓存后端数据(需先通过 CacheRegistry.register 注册类型)
await CacheRoot.get(cache_type, key)
await CacheRoot.set(cache_type, key, value)
await CacheRoot.invalidate_cache(cache_type, key)
```
"""
@@ -95,7 +62,7 @@ from pydantic import BaseModel
from zhenxun.services.log import logger
from .cache_containers import CacheDict, CacheList
from .cache_containers import CacheDict
from .config import (
CACHE_KEY_PREFIX,
CACHE_KEY_SEPARATOR,
@@ -108,9 +75,7 @@ from .config import (
__all__ = [
"Cache",
"CacheData",
"CacheDict",
"CacheList",
"CacheManager",
"CacheRegistry",
"CacheRoot",
@@ -168,129 +133,12 @@ class CacheModel(BaseModel):
arbitrary_types_allowed = True
"""
CacheData类是缓存系统的核心组件,它负责管理单个缓存项的数据和生命周期。
设计思路:
1. 每个CacheData实例代表一个具名的缓存项,如"用户列表"、"配置数据"等
2. 它提供了数据的懒加载、自动过期和持久化等功能
3. 可以通过func参数提供一个获取数据的函数,在数据不存在或过期时自动调用
4. 支持直接设置_data属性,方便外部直接操作数据
主要用途:
1. 作为CacheDict和CacheList的后端存储
2. 被CacheManager管理,实现统一的缓存生命周期控制
3. 提供数据过期和自动刷新机制
通常情况下,用户不需要直接使用CacheData,而是通过Cache、CacheDict或CacheList来操作缓存。
"""
class CacheData:
"""缓存数据类"""
def __init__(
self,
name: str,
func: Callable,
expire: int = DEFAULT_EXPIRE,
lazy_load: bool = True,
cache: BaseCache | AioCache | None = None,
):
"""初始化缓存数据
参数:
name: 缓存名称
func: 获取数据的函数
expire: 过期时间(秒)
lazy_load: 是否延迟加载
cache: 缓存后端
"""
self.name = name.upper()
self.func = func
self.expire = expire
self.lazy_load = lazy_load
self.cache = cache
self._data = None
self._last_update = 0
# 如果不是延迟加载,立即加载数据
if not lazy_load:
import asyncio
try:
loop = asyncio.get_event_loop()
if not loop.is_running():
loop.run_until_complete(self.get_data())
except Exception:
pass
async def get_data(self) -> Any:
"""获取数据
返回:
Any: 缓存数据
"""
# 检查是否需要更新
now = datetime.now().timestamp()
if self._data is None or (
self.expire > 0 and now - self._last_update > self.expire
):
# 更新数据
try:
self._data = await self.func()
self._last_update = now
except Exception as e:
logger.error(f"获取缓存数据 {self.name} 失败", LOG_COMMAND, e=e)
return self._data
async def set_data(self, data: Any) -> bool:
"""设置数据
参数:
data: 缓存数据
返回:
bool: 是否成功
"""
try:
self._data = data
self._last_update = datetime.now().timestamp()
# 如果有缓存后端,保存到缓存
if self.cache and cache_config.cache_mode != CacheMode.NONE:
await self.cache.set(self.name, data, ttl=self.expire) # type: ignore
return True
except Exception as e:
logger.error(f"设置缓存数据 {self.name} 失败", LOG_COMMAND, e=e)
return False
async def clear(self) -> bool:
"""清除数据
返回:
bool: 是否成功
"""
try:
self._data = None
self._last_update = 0
# 如果有缓存后端,清除缓存
if self.cache and cache_config.cache_mode != CacheMode.NONE:
await self.cache.delete(self.name) # type: ignore
return True
except Exception as e:
logger.error(f"清除缓存数据 {self.name} 失败", LOG_COMMAND, e=e)
return False
class CacheManager:
"""缓存管理器"""
_instance: ClassVar["CacheManager | None"] = None
_cache_backend: BaseCache | AioCache | None = None
_registry: ClassVar[dict[str, CacheModel]] = {}
_data: ClassVar[dict[str, CacheData]] = {}
_list_caches: ClassVar[dict[str, "CacheList"]] = {}
_dict_caches: ClassVar[dict[str, "CacheDict"]] = {}
_enabled = False # 缓存启用标记
@@ -336,105 +184,6 @@ class CacheManager:
self._dict_caches[cache_type] = CacheDict[value_type](cache_type, expire)
return self._dict_caches[cache_type]
def cache_list(
self, cache_type: str, expire: int = 0, value_type: type[U] = str
) -> CacheList[U]:
"""获取缓存列表
参数:
cache_type: 缓存类型
expire: 过期时间(秒)
value_type: 值类型
返回:
CacheList: 缓存列表
"""
if cache_type not in self._list_caches:
self._list_caches[cache_type] = CacheList[value_type](cache_type, expire)
return self._list_caches[cache_type]
def listener(self, cache_type: str):
"""缓存监听器装饰器
在方法调用后自动刷新缓存数据
参数:
cache_type: 缓存类型
返回:
Callable: 装饰器
"""
def decorator(func: Callable):
@wraps(func)
async def wrapper(cls, *args, **kwargs):
# 执行原函数
result = await func(cls, *args, **kwargs)
obj = None
# 如果启用了缓存,自动刷新缓存
if cache_config.cache_mode != CacheMode.NONE:
# 根据返回值类型处理
if isinstance(result, tuple) and len(result) > 0:
# 处理返回元组的情况,如 update_or_create 返回 (obj, created)
obj = result[0]
else:
# 处理返回单个对象的情况
obj = result
# 获取缓存键并刷新缓存
if (
obj
and hasattr(cls, "get_cache_key")
and hasattr(obj, cls.get_cache_key_field())
):
key = cls.get_cache_key(obj)
if key is not None:
await self.invalidate_cache(cache_type, key)
return result
return wrapper
return decorator
async def get_cache(self, cache_type: str) -> Any:
"""获取指定类型的缓存对象
此方法返回一个简单的缓存对象,具有 update 方法
参数:
cache_type: 缓存类型
返回:
Any: 缓存对象
"""
class CacheAdapter:
"""缓存适配器"""
def __init__(self, cache_manager: CacheManager, cache_type: str):
self.cache_manager = cache_manager
self.cache_type = cache_type
async def update(self, key: Any, value: Any) -> None:
"""更新缓存
参数:
key: 缓存键
value: 缓存值
"""
# 先清除旧缓存
await self.cache_manager.invalidate_cache(self.cache_type, key)
# 如果需要,可以在这里添加重新设置缓存的逻辑
# 目前我们只清除缓存,让下次查询时自动重建
return (
CacheAdapter(self, cache_type)
if cache_config.cache_mode != CacheMode.NONE
else None
)
@property
def cache_backend(self) -> BaseCache | AioCache:
"""获取缓存后端"""
@@ -479,35 +228,6 @@ class CacheManager:
)
return self._cache_backend
@property
def _cache(self) -> BaseCache | AioCache:
"""获取缓存后端(别名)"""
return self.cache_backend
async def get_cache_data(self, name: str) -> Any:
"""获取缓存数据
参数:
name: 缓存名称
返回:
Any: 缓存数据
"""
name = name.upper()
# 检查是否存在缓存数据
if name in self._data:
return await self._data[name].get_data()
# 尝试从缓存后端获取
if cache_config.cache_mode != CacheMode.NONE:
try:
data = await self.cache_backend.get(name) # type: ignore
if data is not None:
return data
except Exception as e:
logger.error(f"从缓存后端获取数据 {name} 失败", LOG_COMMAND, e=e)
return None
async def invalidate_cache(
self, cache_type: str, key: str | dict[str, Any] | None = None
) -> bool:
@@ -680,29 +400,28 @@ class CacheManager:
"""清除缓存
参数:
cache_type: 缓存类型,为None时清除所有缓存
cache_type: 缓存类型,为None时清除所有缓存。
注意:受 aiocache 限制,无法按类型精确删除,
指定 cache_type 时仅清除整个 backend(行为与不指定相同)。
返回:
bool: 是否成功
"""
# 如果缓存被禁用或缓存模式为NONE,直接返回False
# 如果缓存被禁用或缓存模式为NONE,直接返回True(无需操作)
if not self.enabled or cache_config.cache_mode == CacheMode.NONE:
return False
return True
try:
if cache_type:
# 清除指定类型的缓存
# pattern = f"{cache_type.upper()}{CACHE_KEY_SEPARATOR}*"
# 由于aiocache可能没有delete_pattern方法,使用其他方式清除
# 这里简化处理,直接清除所有缓存
await self.cache_backend.clear() # type: ignore
else:
# 清除所有缓存
await self.cache_backend.clear() # type: ignore
logger.debug(
f"清除缓存类型 {cache_type}"
"(aiocache 不支持按前缀删除,清除整个 backend)",
LOG_COMMAND,
)
await self.cache_backend.clear() # type: ignore
return True
except Exception as e:
if f"缓存类型 {cache_type} 不存在" not in str(e):
logger.warning("清除缓存失败", LOG_COMMAND, e=e)
logger.warning("清除缓存失败", LOG_COMMAND, e=e)
return False
async def close(self):
+69 -35
View File
@@ -18,20 +18,20 @@ if TYPE_CHECKING:
LOG_COMMAND = "RuntimeCache"
PLUGININFO_MEM_REFRESH_INTERVAL = 300
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 = 60
BOT_MEM_REFRESH_INTERVAL = 300 # 5分钟
BOT_MEM_NEGATIVE_TTL = 60
GROUP_MEM_REFRESH_INTERVAL = 60
GROUP_MEM_REFRESH_INTERVAL = 300 # 5分钟 - 群组信息很少变化
GROUP_MEM_NEGATIVE_TTL = 60
LEVEL_MEM_REFRESH_INTERVAL = 120
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 = 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"
@@ -368,35 +368,47 @@ class PluginLimitSnapshot:
@dataclass(frozen=True)
class TaskInfoSnapshot:
id: int
module: str
name: str
status: bool
load_status: bool
default_status: bool
run_time: str | None
@classmethod
def from_model(cls, model) -> "TaskInfoSnapshot":
return cls(
id=int(getattr(model, "id", 0) or 0),
module=str(model.module),
name=str(getattr(model, "name", "") or ""),
status=bool(getattr(model, "status", True)),
load_status=bool(getattr(model, "load_status", True)),
default_status=bool(getattr(model, "default_status", True)),
run_time=getattr(model, "run_time", None),
)
def to_payload(self) -> dict[str, Any]:
return {
"id": self.id,
"module": self.module,
"name": self.name,
"status": self.status,
"load_status": self.load_status,
"default_status": self.default_status,
"run_time": self.run_time,
}
@classmethod
def from_payload(cls, payload: dict[str, Any]) -> "TaskInfoSnapshot":
return cls(
id=int(payload.get("id", 0) or 0),
module=str(payload.get("module", "")),
name=str(payload.get("name", "") or ""),
status=bool(payload.get("status", True)),
load_status=bool(payload.get("load_status", True)),
default_status=bool(payload.get("default_status", True)),
run_time=payload.get("run_time"),
)
@@ -1208,6 +1220,7 @@ class LevelUserMemoryCache:
class TaskInfoMemoryCache:
_lock: ClassVar[asyncio.Lock] = asyncio.Lock()
_by_module: ClassVar[dict[str, TaskInfoSnapshot]] = {}
_by_name: ClassVar[dict[str, TaskInfoSnapshot]] = {}
_negative: ClassVar[dict[str, float]] = {}
_loaded: ClassVar[bool] = False
_refresh_task: ClassVar[asyncio.Task | None] = None
@@ -1246,7 +1259,15 @@ class TaskInfoMemoryCache:
async with cls._lock:
records = await TaskInfo.all()
cls._by_module = {r.module: TaskInfoSnapshot.from_model(r) for r in records}
by_module: dict[str, TaskInfoSnapshot] = {}
by_name: dict[str, TaskInfoSnapshot] = {}
for record in records:
entry = TaskInfoSnapshot.from_model(record)
by_module[entry.module] = entry
if entry.name:
by_name[entry.name] = entry
cls._by_module = by_module
cls._by_name = by_name
cls._negative = {}
cls._loaded = True
logger.debug(
@@ -1275,6 +1296,21 @@ class TaskInfoMemoryCache:
cls._mark_negative(module)
return None
@classmethod
async def get_by_name(cls, name: str | None) -> TaskInfoSnapshot | None:
name = (name or "").strip()
if not name:
return None
if not cls._loaded:
await cls.ensure_loaded()
return cls._by_name.get(name)
@classmethod
async def get_all(cls) -> list[TaskInfoSnapshot]:
if not cls._loaded:
await cls.ensure_loaded()
return sorted(cls._by_module.values(), key=lambda item: (item.id, item.module))
@classmethod
async def is_disabled(cls, module: str | None) -> bool:
entry = await cls.get(module)
@@ -1287,6 +1323,8 @@ class TaskInfoMemoryCache:
entry = TaskInfoSnapshot.from_model(record)
async with cls._lock:
cls._by_module[entry.module] = entry
if entry.name:
cls._by_name[entry.name] = entry
cls._negative.pop(entry.module, None)
RuntimeCacheSync.publish_event("task", "upsert", entry.to_payload())
@@ -1297,6 +1335,8 @@ class TaskInfoMemoryCache:
return
async with cls._lock:
cls._by_module[entry.module] = entry
if entry.name:
cls._by_name[entry.name] = entry
cls._negative.pop(entry.module, None)
@classmethod
@@ -1305,7 +1345,11 @@ class TaskInfoMemoryCache:
if not module:
return
async with cls._lock:
cls._by_module.pop(module, None)
removed = cls._by_module.pop(module, None)
if removed and removed.name:
current = cls._by_name.get(removed.name)
if current and current.module == removed.module:
cls._by_name.pop(removed.name, None)
RuntimeCacheSync.publish_event("task", "delete", {"module": module})
@classmethod
@@ -1837,37 +1881,27 @@ class BanMemoryCache:
await cls.refresh()
async def _safe_refresh(cache_cls: type, label: str) -> None:
"""安全地刷新单个缓存,异常不影响其他缓存。"""
try:
await cache_cls.refresh()
except Exception as exc:
logger.error(f"{label} cache init failed", LOG_COMMAND, e=exc)
@PriorityLifecycle.on_startup(priority=6)
async def _init_runtime_cache():
await RuntimeCacheSync.start()
try:
await PluginInfoMemoryCache.refresh()
except Exception as exc:
logger.error("plugin cache init failed", LOG_COMMAND, e=exc)
try:
await BotMemoryCache.refresh()
except Exception as exc:
logger.error("bot cache init failed", LOG_COMMAND, e=exc)
try:
await GroupMemoryCache.refresh()
except Exception as exc:
logger.error("group cache init failed", LOG_COMMAND, e=exc)
try:
await LevelUserMemoryCache.refresh()
except Exception as exc:
logger.error("level cache init failed", LOG_COMMAND, e=exc)
try:
await TaskInfoMemoryCache.refresh()
except Exception as exc:
logger.error("task info cache init failed", LOG_COMMAND, e=exc)
try:
await PluginLimitMemoryCache.refresh()
except Exception as exc:
logger.error("plugin limit cache init failed", LOG_COMMAND, e=exc)
try:
await BanMemoryCache.refresh()
except Exception as exc:
logger.error("ban cache init failed", LOG_COMMAND, e=exc)
# 并发刷新所有缓存,互不依赖
await asyncio.gather(
_safe_refresh(PluginInfoMemoryCache, "plugin"),
_safe_refresh(BotMemoryCache, "bot"),
_safe_refresh(GroupMemoryCache, "group"),
_safe_refresh(LevelUserMemoryCache, "level"),
_safe_refresh(TaskInfoMemoryCache, "task info"),
_safe_refresh(PluginLimitMemoryCache, "plugin limit"),
_safe_refresh(BanMemoryCache, "ban"),
)
PluginInfoMemoryCache.start_refresh_task()
BotMemoryCache.start_tasks()
GroupMemoryCache.start_tasks()
+1 -4
View File
@@ -1,3 +1,4 @@
import re
from typing import Any, ClassVar, Generic, TypeVar, cast
from zhenxun.services.cache import Cache, CacheRoot, cache_config
@@ -7,8 +8,6 @@ from zhenxun.services.log import logger
T = TypeVar("T", bound=Model)
cache = CacheRoot.cache_dict("DB_TEST_BAN", 10, int)
class DataAccess(Generic[T]):
"""数据访问层,根据配置决定是否使用缓存
@@ -387,8 +386,6 @@ class DataAccess(Generic[T]):
# 构建键参数字典
key_parts = []
# 从格式字符串中提取所需的字段名
import re
field_names = re.findall(r"{([^}]+)}", cache_model.key_format)
# 收集所有字段值
+114 -30
View File
@@ -1,5 +1,8 @@
import asyncio
import hashlib
import json
from pathlib import Path
import re
from urllib.parse import urlparse
import aiofiles
@@ -7,7 +10,7 @@ import nonebot
from nonebot.utils import is_coroutine_callable
from tortoise import Tortoise
from tortoise.connection import connections
from tortoise.exceptions import OperationalError
from tortoise.exceptions import ConfigurationError, OperationalError
from zhenxun.configs.config import BotConfig
from zhenxun.services.log import logger
@@ -44,6 +47,8 @@ __all__ = [
driver = nonebot.get_driver()
_SCRIPT_HASH_FILE = Path() / "data" / ".db_script_hash"
def get_config() -> dict:
"""获取数据库配置"""
@@ -121,7 +126,6 @@ async def init():
config=get_config(),
)
if db_model.script_method:
db = Tortoise.get_connection("default")
logger.debug(
"即将运行SCRIPT_METHOD方法, 合计 "
f"<u><y>{len(db_model.script_method)}</y></u> 个..."
@@ -134,34 +138,108 @@ async def init():
sql_list += sql
except Exception as e:
logger.debug(f"{module} 执行SCRIPT_METHOD方法出错...", e=e)
for sql in sql_list:
logger.debug(f"执行SQL: {sql}")
try:
await asyncio.wait_for(
db.execute_query_dict(sql), timeout=DB_TIMEOUT_SECONDS
)
except OperationalError as e:
err_str = str(e).lower()
if any(
x in err_str
for x in [
"already exists",
"duplicate column",
"已经存在",
"已存在",
]
):
pass
elif any(
x in err_str for x in ["does not exist", "check that", "不存在"]
) and ("drop" in sql.lower() or "rename" in sql.lower()):
pass
else:
logger.warning(f"执行SQL警告: {sql} || {e}")
except Exception as e:
logger.debug(f"执行SQL: {sql} 错误...", e=e)
if sql_list:
logger.debug("SCRIPT_METHOD方法执行完毕!")
fingerprint = hashlib.md5(
json.dumps(sorted(sql_list), ensure_ascii=False).encode()
).hexdigest()
need_run = not (
_SCRIPT_HASH_FILE.exists()
and _SCRIPT_HASH_FILE.read_text(encoding="utf-8").strip()
== fingerprint
)
if need_run:
db = Tortoise.get_connection("default")
async def table_exists(table_name: str) -> bool:
"""检查表是否存在"""
try:
# PostgreSQL
result = await db.execute_query_dict(
"SELECT to_regclass($1) IS NOT NULL as exists",
[table_name],
)
if result:
return result[0]["exists"]
except Exception:
pass
try:
# MySQL
result = await db.execute_query_dict(
"SELECT COUNT(*) as count FROM information_schema.tables " # noqa: E501
"WHERE table_name = %s",
[table_name],
)
if result:
return result[0]["count"] > 0
except Exception:
pass
try:
# SQLite
result = await db.execute_query_dict(
"SELECT name FROM sqlite_master WHERE type='table' AND name=?", # noqa: E501
[table_name],
)
return len(result) > 0
except Exception:
pass
return True # 如果检查失败,假设表存在,让SQL自己报错
for sql in sql_list:
# 对于 ALTER TABLE 操作,先检查表是否存在
if sql.strip().upper().startswith("ALTER TABLE"):
match = re.match(
r"ALTER\s+TABLE\s+(\w+)", sql, re.IGNORECASE
)
if match:
table_name = match.group(1)
if not await table_exists(table_name):
logger.debug(f"跳过SQL(表不存在): {sql}")
continue
logger.debug(f"执行SQL: {sql}")
try:
await asyncio.wait_for(
db.execute_query_dict(sql),
timeout=DB_TIMEOUT_SECONDS,
)
except OperationalError as e:
err_str = str(e).lower()
sql_lower = sql.lower()
if any(
x in err_str
for x in [
"already exists",
"duplicate column",
"已经存在",
"已存在",
]
):
pass
elif any(
x in err_str
for x in [
"does not exist",
"check that",
"不存在",
"no such column",
]
) and ("drop" in sql_lower or "rename" in sql_lower):
pass
elif "syntax error" in err_str and (
"alter column" in sql_lower
or "drop not null" in sql_lower
):
# SQLite 不支持 PostgreSQL 的 ALTER COLUMN 语法
pass
else:
logger.warning(f"执行SQL警告: {sql} || {e}")
except Exception as e:
logger.debug(f"执行SQL: {sql} 错误...", e=e)
logger.debug("SCRIPT_METHOD方法执行完毕!")
_SCRIPT_HASH_FILE.parent.mkdir(parents=True, exist_ok=True)
_SCRIPT_HASH_FILE.write_text(fingerprint, encoding="utf-8")
else:
logger.debug("迁移脚本无变化,跳过执行")
logger.debug("开始生成数据库表结构...")
await Tortoise.generate_schemas()
logger.debug("数据库表结构生成完毕!")
@@ -170,5 +248,11 @@ async def init():
raise DbConnectError(f"数据库连接错误... e:{e}") from e
@PriorityLifecycle.on_shutdown(priority=100)
async def disconnect():
await connections.close_all()
try:
await connections.close_all()
except ConfigurationError:
logger.debug("数据库连接未初始化,跳过关闭")
except Exception as e:
logger.error(f"关闭数据库连接时发生意外错误: {e}")
+6 -2
View File
@@ -30,7 +30,11 @@ class PluginData(BaseModel):
async def _get_plugins_by_types(plugin_types: list[PluginType]) -> list[PluginData]:
"""根据指定的插件类型列表获取插件数据"""
plugin_list = await PluginInfo.filter(plugin_type__in=plugin_types).all()
plugin_list = await PluginInfo.get_plugins(
load_status=None,
filter_parent=False,
plugin_type__in=plugin_types,
)
data_list = []
for plugin in plugin_list:
if _plugin := nonebot.get_plugin_by_module_name(plugin.module_path):
@@ -42,7 +46,7 @@ async def _get_plugins_by_types(plugin_types: list[PluginType]) -> list[PluginDa
async def _get_task_category() -> dict:
"""获取被动技能帮助类别"""
task_items = []
if task_list := await TaskInfo.all():
if task_list := await TaskInfo.get_tasks(load_status=True):
task_names = "\n".join([task.name for task in task_list])
task_items.append(
{
+529 -144
View File
@@ -1,7 +1,8 @@
import asyncio
from collections import OrderedDict
from collections.abc import Awaitable, Callable
from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine
import contextlib
from dataclasses import dataclass, field
import hashlib
import inspect
import json
@@ -9,6 +10,7 @@ from pathlib import Path
import time
from typing import Any, ClassVar, cast
import nonebot_plugin_htmlrender as htmlrender_module
import nonebot_plugin_htmlrender.browser as htmlrender_browser
import psutil
@@ -17,6 +19,98 @@ from zhenxun.services.log import logger
from .types import BaseScreenshotEngine
_PLAYWRIGHT_DISCONNECT_ERROR = "Connection closed while reading from the driver"
_UNRETRIEVED_FUTURE_MESSAGE = "Future exception was never retrieved"
_LOOP_EXCEPTION_FILTER_STATE_ATTR = "_zhenxun_playwright_exception_filter_state"
_DISCONNECT_SUPPRESSION_WINDOW_SECONDS = 10.0
class HtmlrenderTaskTracker:
"""只追踪 htmlrender 渲染任务的轻量运行时。"""
def __init__(self) -> None:
self._lock = asyncio.Lock()
self._idle_event = asyncio.Event()
self._idle_event.set()
self._active_tasks = 0
self._draining = False
self._drain_reason: str | None = None
@property
def active_tasks(self) -> int:
return self._active_tasks
@property
def is_draining(self) -> bool:
return self._draining
async def reset(self) -> None:
async with self._lock:
self._draining = False
self._drain_reason = None
if self._active_tasks == 0:
self._idle_event.set()
async def resume(self) -> None:
async with self._lock:
self._draining = False
self._drain_reason = None
async def mark_draining(self, reason: str) -> None:
async with self._lock:
self._draining = True
self._drain_reason = reason
async def begin(self, owner: str) -> None:
async with self._lock:
if self._draining:
reason = self._drain_reason or "unknown"
message = (
"htmlrender 正在排空,拒绝新的渲染任务: "
f"owner={owner}, reason={reason}"
)
raise RuntimeError(message)
self._active_tasks += 1
self._idle_event.clear()
async def end(self) -> None:
async with self._lock:
self._active_tasks = max(0, self._active_tasks - 1)
if self._active_tasks == 0:
self._idle_event.set()
async def wait_for_idle(self) -> None:
await self._idle_event.wait()
@contextlib.asynccontextmanager
async def track(self, owner: str):
await self.begin(owner)
try:
yield
finally:
await self.end()
_HTMLRENDER_TASK_TRACKER = HtmlrenderTaskTracker()
@dataclass(slots=True)
class ContextGeneration:
generation_id: int
context_pool: asyncio.LifoQueue[Any] = field(default_factory=asyncio.LifoQueue)
all_contexts: set[Any] = field(default_factory=set)
active_leases: int = 0
retiring: bool = False
def snapshot(self) -> dict[str, int | bool]:
return {
"generation_id": self.generation_id,
"pool_size": self.context_pool.qsize(),
"context_count": len(self.all_contexts),
"active_leases": self.active_leases,
"retiring": self.retiring,
}
async def _await_if_needed(value: Any) -> Any:
if inspect.isawaitable(value):
@@ -32,35 +126,95 @@ async def _get_browser_instance() -> Any:
raise RuntimeError("nonebot_plugin_htmlrender.browser 未提供可用浏览器获取函数。")
async def _shutdown_browser_instance() -> None:
for attr_name in (
"shutdown_htmlrender",
"shutdown_browser",
"close_browser",
"close_htmlrender",
):
shutdown_func = getattr(htmlrender_browser, attr_name, None)
if callable(shutdown_func):
try:
await _await_if_needed(shutdown_func())
finally:
with contextlib.suppress(Exception):
setattr(htmlrender_browser, "_browser", None)
with contextlib.suppress(Exception):
setattr(htmlrender_browser, "_playwright", None)
def _is_ignorable_playwright_disconnect(ctx: dict[str, Any]) -> bool:
exc = ctx.get("exception")
return (
ctx.get("message") == _UNRETRIEVED_FUTURE_MESSAGE
and isinstance(exc, Exception)
and _PLAYWRIGHT_DISCONNECT_ERROR in str(exc)
)
def _get_loop_exception_filter_state(
loop: asyncio.AbstractEventLoop,
) -> dict[str, Any] | None:
state = getattr(loop, _LOOP_EXCEPTION_FILTER_STATE_ATTR, None)
if isinstance(state, dict):
return state
return None
def _ensure_loop_exception_filter(
loop: asyncio.AbstractEventLoop,
) -> dict[str, Any]:
state = _get_loop_exception_filter_state(loop)
if state is not None:
return state
state = {
"original_handler": loop.get_exception_handler(),
"suppress_until": 0.0,
}
def _filter(lp: asyncio.AbstractEventLoop, ctx: dict[str, Any]) -> None:
if _is_ignorable_playwright_disconnect(ctx):
suppress_until = float(state.get("suppress_until", 0.0))
if suppress_until >= time.monotonic():
return
original_handler = state.get("original_handler")
if callable(original_handler):
original_handler(lp, ctx)
return
lp.default_exception_handler(ctx)
loop.set_exception_handler(_filter)
setattr(loop, _LOOP_EXCEPTION_FILTER_STATE_ATTR, state)
return state
def _arm_disconnect_exception_suppression(
loop: asyncio.AbstractEventLoop,
*,
seconds: float = _DISCONNECT_SUPPRESSION_WINDOW_SECONDS,
) -> None:
state = _ensure_loop_exception_filter(loop)
deadline = time.monotonic() + max(seconds, 0.0)
state["suppress_until"] = max(float(state.get("suppress_until", 0.0)), deadline)
async def _shutdown_browser_instance() -> None:
loop = asyncio.get_running_loop()
_arm_disconnect_exception_suppression(loop)
browser_obj = getattr(htmlrender_browser, "_browser", None)
playwright_obj = getattr(htmlrender_browser, "_playwright", None)
if browser_obj is not None:
is_connected_fn = getattr(browser_obj, "is_connected", None)
if callable(is_connected_fn) and not is_connected_fn():
with contextlib.suppress(Exception):
setattr(htmlrender_browser, "_browser", None)
with contextlib.suppress(Exception):
setattr(htmlrender_browser, "_playwright", None)
return
if browser_obj is None and playwright_obj is None:
return
close_func = getattr(browser_obj, "close", None) if browser_obj else None
if callable(close_func):
with contextlib.suppress(Exception):
try:
await _await_if_needed(close_func())
except Exception as e:
if _PLAYWRIGHT_DISCONNECT_ERROR not in str(e):
logger.debug(f"关闭浏览器实例时忽略异常: {e}")
playwright_obj = getattr(htmlrender_browser, "_playwright", None)
stop_func = getattr(playwright_obj, "stop", None) if playwright_obj else None
if callable(stop_func):
with contextlib.suppress(Exception):
try:
await _await_if_needed(stop_func())
except Exception as e:
if _PLAYWRIGHT_DISCONNECT_ERROR not in str(e):
logger.debug(f"关闭 Playwright 实例时忽略异常: {e}")
with contextlib.suppress(Exception):
setattr(htmlrender_browser, "_browser", None)
@@ -68,15 +222,55 @@ async def _shutdown_browser_instance() -> None:
setattr(htmlrender_browser, "_playwright", None)
if callable(close_func) or callable(stop_func):
await asyncio.sleep(0)
def _patch_htmlrender_task_tracking() -> None:
if getattr(htmlrender_browser, "_zhenxun_task_tracking_patched", False):
return
logger.debug(
"未找到 htmlrender 浏览器关闭函数,跳过 shutdown。",
"PlaywrightEngine",
)
try:
import nonebot_plugin_htmlrender.data_source as htmlrender_data_source
except Exception as e:
logger.warning("导入 htmlrender.data_source 失败,跳过任务追踪补丁。", e=e)
return
original_get_new_page = getattr(htmlrender_browser, "get_new_page", None)
if not callable(original_get_new_page):
logger.warning("htmlrender 未提供 get_new_page,跳过任务追踪补丁。")
return
original_get_new_page = cast(Callable[..., Any], original_get_new_page)
@contextlib.asynccontextmanager
async def _tracked_get_new_page(*args: Any, **kwargs: Any) -> AsyncIterator[Any]:
async with _HTMLRENDER_TASK_TRACKER.track("htmlrender"):
page_context = cast(Any, original_get_new_page(*args, **kwargs))
async with page_context as page:
yield page
setattr(htmlrender_browser, "get_new_page", _tracked_get_new_page)
setattr(htmlrender_module, "get_new_page", _tracked_get_new_page)
setattr(htmlrender_data_source, "get_new_page", _tracked_get_new_page)
setattr(htmlrender_browser, "_zhenxun_task_tracking_patched", True)
def _patch_htmlrender_shutdown() -> None:
if getattr(htmlrender_browser, "_zhenxun_shutdown_patched", False):
return
async def _patched_shutdown_browser() -> None:
if _HTMLRENDER_TASK_TRACKER.is_draining:
await _HTMLRENDER_TASK_TRACKER.wait_for_idle()
await _shutdown_browser_instance()
setattr(htmlrender_browser, "shutdown_browser", _patched_shutdown_browser)
setattr(htmlrender_module, "shutdown_browser", _patched_shutdown_browser)
setattr(htmlrender_browser, "_zhenxun_shutdown_patched", True)
def _patch_playwright_env_check_once() -> None:
_patch_htmlrender_task_tracking()
_patch_htmlrender_shutdown()
if getattr(htmlrender_browser, "_zhenxun_check_once_patched", False):
return
@@ -155,9 +349,9 @@ def _patch_playwright_env_check_once() -> None:
class PlaywrightEngine(BaseScreenshotEngine):
"""使用 nonebot-plugin-htmlrender 实现的截图引擎。"""
_MAX_CONCURRENT_RENDER = 2
_CONTEXT_POOL_SIZE = 2
_PREWARM_CONTEXT_COUNT = 1
_MAX_CONCURRENT_RENDER = 4
_CONTEXT_POOL_SIZE = 4
_PREWARM_CONTEXT_COUNT = 2
_SET_CONTENT_WAIT_UNTIL = "domcontentloaded"
_READY_STATE_TIMEOUT_MS = 2_000
_IMAGE_READY_TIMEOUT_MS = 1_800
@@ -183,7 +377,6 @@ class PlaywrightEngine(BaseScreenshotEngine):
_IDLE_CHECK_INTERVAL_SECONDS = 15
_IDLE_RECYCLE_SECONDS = 180
_POOL_UNSAFE_OPTION_KEYS: ClassVar[set[str]] = {
"device_scale_factor",
"color_scheme",
"extra_http_headers",
"forced_colors",
@@ -209,6 +402,7 @@ class PlaywrightEngine(BaseScreenshotEngine):
"timezone_id",
"user_agent",
}
_POOL_DEVICE_SCALE_FACTOR = 2
def __init__(self):
_patch_playwright_env_check_once()
@@ -224,8 +418,9 @@ class PlaywrightEngine(BaseScreenshotEngine):
self._rss_baseline_bytes: int | None = None
self._recent_results: OrderedDict[str, tuple[float, bytes]] = OrderedDict()
self._inflight_tasks: dict[str, asyncio.Task[bytes]] = {}
self._context_pool: asyncio.LifoQueue[Any] = asyncio.LifoQueue()
self._all_contexts: set[Any] = set()
self._generation_counter = 0
self._active_generation: ContextGeneration | None = None
self._retiring_generations: list[ContextGeneration] = []
self._idle_recycle_task: asyncio.Task[None] | None = None
self._closing = False
self._process = psutil.Process()
@@ -311,16 +506,62 @@ class PlaywrightEngine(BaseScreenshotEngine):
if current_rss >= threshold:
self._recycle_pending = True
async def get_runtime_snapshot(self) -> dict[str, Any]:
async with self._state_lock:
active_generation = (
self._active_generation.snapshot()
if self._active_generation is not None
else None
)
retiring_generations = [
generation.snapshot() for generation in self._retiring_generations
]
return {
"closing": self._closing,
"active_renders": self._active_renders,
"render_count": self._render_count,
"recycle_pending": self._recycle_pending,
"last_recycle_at": self._last_recycle_at,
"generation_counter": self._generation_counter,
"active_generation": active_generation,
"retiring_generations": retiring_generations,
"retiring_generation_count": len(retiring_generations),
"inflight_task_count": len(self._inflight_tasks),
"recent_result_count": len(self._recent_results),
"htmlrender_active_tasks": _HTMLRENDER_TASK_TRACKER.active_tasks,
"htmlrender_draining": _HTMLRENDER_TASK_TRACKER.is_draining,
}
async def _log_runtime_snapshot(self, reason: str) -> None:
snapshot = await self.get_runtime_snapshot()
logger.trace(
f"截图引擎状态快照[{reason}]: {snapshot}",
)
def _create_generation_nolock(self) -> ContextGeneration:
self._generation_counter += 1
return ContextGeneration(generation_id=self._generation_counter)
def _ensure_active_generation_nolock(self) -> ContextGeneration:
if self._active_generation is None:
self._active_generation = self._create_generation_nolock()
return self._active_generation
async def initialize(self) -> None:
async with self._state_lock:
if self._idle_recycle_task and not self._idle_recycle_task.done():
return
self._closing = False
self._generation_counter = 0
self._active_generation = None
self._retiring_generations.clear()
self._last_render_finished_at = time.monotonic()
if current_rss := self._get_total_rss():
self._rss_baseline_bytes = current_rss
self._idle_recycle_task = asyncio.create_task(self._idle_recycle_loop())
await self._prewarm_browser_and_pool()
await _HTMLRENDER_TASK_TRACKER.reset()
await self._log_runtime_snapshot("initialize")
# 浏览器在首次 _acquire_context 时按需启动,无需预热
async def close(self) -> None:
idle_task: asyncio.Task[None] | None = None
@@ -328,21 +569,34 @@ class PlaywrightEngine(BaseScreenshotEngine):
self._closing = True
idle_task = self._idle_recycle_task
self._idle_recycle_task = None
for task in self._inflight_tasks.values():
task.cancel()
self._inflight_tasks.clear()
self._recent_results.clear()
self._recycle_pending = False
await _HTMLRENDER_TASK_TRACKER.mark_draining("engine_close")
if idle_task:
idle_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await idle_task
await _HTMLRENDER_TASK_TRACKER.wait_for_idle()
async with self._state_lock:
inflight_tasks = list(self._inflight_tasks.values())
if inflight_tasks:
await asyncio.gather(*inflight_tasks, return_exceptions=True)
async with self._state_lock:
self._inflight_tasks.clear()
await self._log_runtime_snapshot("close:before_dispose")
await self._dispose_context_pool()
await _shutdown_browser_instance()
await self._log_runtime_snapshot("close:after_shutdown")
async def _on_render_begin(self) -> None:
await _HTMLRENDER_TASK_TRACKER.begin("zhenxun_renderer")
async with self._state_lock:
self._active_renders += 1
@@ -354,10 +608,11 @@ class PlaywrightEngine(BaseScreenshotEngine):
now = time.monotonic()
self._last_render_finished_at = now
self._mark_recycle_if_needed_nolock(now)
if self._recycle_pending and self._active_renders == 0:
if self._recycle_pending and _HTMLRENDER_TASK_TRACKER.active_tasks == 0:
self._recycle_pending = False
self._last_recycle_at = now
should_recycle = True
await _HTMLRENDER_TASK_TRACKER.end()
if should_recycle:
await self._recycle_browser("active")
@@ -378,6 +633,7 @@ class PlaywrightEngine(BaseScreenshotEngine):
options.pop("disable_animations", None)
if pooled:
options.pop("base_url", None)
options.pop("device_scale_factor", None)
return options
@staticmethod
@@ -413,6 +669,9 @@ class PlaywrightEngine(BaseScreenshotEngine):
for key in cls._POOL_UNSAFE_OPTION_KEYS:
if key in render_options:
return False
dsf = render_options.get("device_scale_factor")
if dsf is not None and dsf != cls._POOL_DEVICE_SCALE_FACTOR:
return False
return True
async def _render_with_page(
@@ -424,7 +683,7 @@ class PlaywrightEngine(BaseScreenshotEngine):
) -> bytes:
if self._debug_console_log:
page.on("console", lambda msg: logger.debug(f"浏览器控制台: {msg.text}"))
await page.goto(template_path, wait_until="domcontentloaded")
await page.goto(template_path, wait_until="commit")
await page.set_content(html, wait_until=self._SET_CONTENT_WAIT_UNTIL)
if bool(render_options.get("disable_animations", False)):
await self._disable_page_animations(page)
@@ -477,13 +736,7 @@ class PlaywrightEngine(BaseScreenshotEngine):
current_height = int(viewport.get("height") or 0)
if target_height > current_height:
await page.set_viewport_size(
{
"width": width,
"height": min(
target_height,
self._FULL_PAGE_VIEWPORT_MAX_HEIGHT,
),
}
{"width": width, "height": target_height}
)
if clip_padding <= 0:
@@ -510,18 +763,76 @@ class PlaywrightEngine(BaseScreenshotEngine):
return await element.screenshot(**element_screenshot_options)
async def _wait_for_visual_stability(self, page: Any) -> None:
# 先做一次快速预检,判断页面是否有外部图片和自定义字体
resource_hints: dict[str, bool] | None = None
with contextlib.suppress(Exception):
resource_hints = await page.evaluate(
"""
() => {
const imgs = document.images || [];
let hasUnloadedImages = false;
for (let i = 0; i < imgs.length; i++) {
if (!imgs[i].complete) { hasUnloadedImages = true; break; }
}
const hasCustomFonts = !!(
document.fonts && document.fonts.size > 0
);
return {
ready: document.readyState === 'complete',
images: hasUnloadedImages,
fonts: hasCustomFonts,
};
}
"""
)
# 如果预检已知全部就绪,直接返回
if (
isinstance(resource_hints, dict)
and resource_hints.get("ready") is True
and resource_hints.get("images") is not True
and resource_hints.get("fonts") is not True
):
return
# 对需要的等待项并行执行
waiters: list[Coroutine[Any, Any, None]] = []
need_ready = not (
isinstance(resource_hints, dict) and resource_hints.get("ready") is True
)
need_images = (
not isinstance(resource_hints, dict) or resource_hints.get("images") is True
)
need_fonts = (
not isinstance(resource_hints, dict) or resource_hints.get("fonts") is True
)
if need_ready:
waiters.append(self._wait_ready_state(page))
if need_images:
waiters.append(self._wait_images_loaded(page))
if need_fonts:
waiters.append(self._wait_fonts_ready(page))
if waiters:
await asyncio.gather(*waiters)
async def _wait_ready_state(self, page: Any) -> None:
with contextlib.suppress(Exception):
await page.wait_for_function(
"() => document.readyState === 'complete'",
timeout=self._READY_STATE_TIMEOUT_MS,
)
async def _wait_images_loaded(self, page: Any) -> None:
with contextlib.suppress(Exception):
await page.wait_for_function(
"() => Array.from(document.images || []).every(img => img.complete)",
timeout=self._IMAGE_READY_TIMEOUT_MS,
)
async def _wait_fonts_ready(self, page: Any) -> None:
with contextlib.suppress(Exception):
await page.evaluate(
"""
@@ -595,15 +906,12 @@ class PlaywrightEngine(BaseScreenshotEngine):
content_height, int
):
return
if (
content_width < 10
or content_height < 10
or content_width > self._FULL_PAGE_VIEWPORT_MAX_WIDTH
or content_height > self._FULL_PAGE_VIEWPORT_MAX_HEIGHT
):
if content_width < 10 or content_height < 10:
return
target_width = max(width, content_width)
target_width = min(
max(width, content_width), self._FULL_PAGE_VIEWPORT_MAX_WIDTH
)
target_height = max(height, content_height)
await page.set_viewport_size(
{"width": target_width, "height": target_height}
@@ -627,17 +935,116 @@ class PlaywrightEngine(BaseScreenshotEngine):
with contextlib.suppress(Exception):
await page.close()
async def _acquire_context(self) -> Any:
try:
return self._context_pool.get_nowait()
except asyncio.QueueEmpty:
pass
async def _dispose_generation(self, generation: ContextGeneration) -> None:
contexts = list(generation.all_contexts)
generation.all_contexts.clear()
while True:
try:
generation.context_pool.get_nowait()
except asyncio.QueueEmpty:
break
for context in contexts:
with contextlib.suppress(Exception):
await context.close()
async def _cleanup_retiring_generations(self) -> None:
disposable: list[ContextGeneration] = []
async with self._state_lock:
if len(self._all_contexts) < self._CONTEXT_POOL_SIZE:
create_new = True
else:
create_new = False
remaining: list[ContextGeneration] = []
for generation in self._retiring_generations:
if generation.active_leases <= 0:
disposable.append(generation)
else:
remaining.append(generation)
self._retiring_generations = remaining
for generation in disposable:
await self._dispose_generation(generation)
async def _dispose_context_pool(self) -> None:
async with self._state_lock:
generations: list[ContextGeneration] = []
if self._active_generation is not None:
generations.append(self._active_generation)
self._active_generation = None
generations.extend(self._retiring_generations)
self._retiring_generations = []
for generation in generations:
await self._dispose_generation(generation)
async def _build_generation(self) -> ContextGeneration:
async with self._state_lock:
generation = self._create_generation_nolock()
if self._closing:
return generation
try:
browser = await _get_browser_instance()
except Exception as e:
logger.warning("截图引擎浏览器预热失败。", "PlaywrightEngine", e=e)
return generation
for _ in range(self._PREWARM_CONTEXT_COUNT):
if self._closing:
break
if len(generation.all_contexts) >= self._CONTEXT_POOL_SIZE:
break
context = None
try:
context = await browser.new_context(
viewport={"width": 800, "height": 10},
device_scale_factor=2,
)
page = await context.new_page()
await page.goto("about:blank", wait_until="domcontentloaded")
await page.set_content(
"<html><body></body></html>",
wait_until="domcontentloaded",
)
await page.close()
except Exception as e:
logger.warning("截图引擎上下文预热失败。", "PlaywrightEngine", e=e)
if context is not None:
with contextlib.suppress(Exception):
await context.close()
break
generation.all_contexts.add(context)
generation.context_pool.put_nowait(context)
return generation
async def _swap_generation(self, reason: str) -> None:
new_generation = await self._build_generation()
async with self._state_lock:
old_generation = self._active_generation
if old_generation is not None:
old_generation.retiring = True
self._retiring_generations.append(old_generation)
self._active_generation = new_generation
await self._cleanup_retiring_generations()
logger.debug(
f"截图引擎触发代际切换({reason}),新代={new_generation.generation_id}",
"PlaywrightEngine",
)
await self._log_runtime_snapshot(f"swap_generation:{reason}")
async def _acquire_context(self) -> tuple[ContextGeneration, Any]:
generation: ContextGeneration | None = None
create_new = False
async with self._state_lock:
generation = self._ensure_active_generation_nolock()
try:
context = generation.context_pool.get_nowait()
generation.active_leases += 1
return generation, context
except asyncio.QueueEmpty:
create_new = len(generation.all_contexts) < self._CONTEXT_POOL_SIZE
if create_new:
browser = await _get_browser_instance()
@@ -646,57 +1053,57 @@ class PlaywrightEngine(BaseScreenshotEngine):
device_scale_factor=2,
)
async with self._state_lock:
self._all_contexts.add(context)
return context
return await self._context_pool.get()
async def _release_context(self, context: Any, broken: bool = False) -> None:
if broken:
await self._discard_context(context)
return
target_generation = generation
if target_generation.retiring and self._active_generation is not None:
target_generation = self._active_generation
target_generation.all_contexts.add(context)
target_generation.active_leases += 1
return target_generation, context
context = await generation.context_pool.get()
async with self._state_lock:
if self._closing:
broken = True
elif context not in self._all_contexts:
broken = True
else:
self._context_pool.put_nowait(context)
return
generation.active_leases += 1
return generation, context
if broken:
await self._discard_context(context)
async def _discard_context(self, context: Any) -> None:
async def _release_context(
self,
generation: ContextGeneration,
context: Any,
broken: bool = False,
) -> None:
should_discard = broken
async with self._state_lock:
existed = context in self._all_contexts
generation.active_leases = max(0, generation.active_leases - 1)
if self._closing or generation.retiring:
should_discard = True
elif context not in generation.all_contexts:
should_discard = True
elif not should_discard:
generation.context_pool.put_nowait(context)
if should_discard:
await self._discard_context(generation, context)
await self._cleanup_retiring_generations()
async def _discard_context(
self, generation: ContextGeneration, context: Any
) -> None:
async with self._state_lock:
existed = context in generation.all_contexts
if existed:
self._all_contexts.remove(context)
generation.all_contexts.remove(context)
if existed:
with contextlib.suppress(Exception):
await context.close()
async def _dispose_context_pool(self) -> None:
async with self._state_lock:
contexts = list(self._all_contexts)
self._all_contexts.clear()
while True:
try:
self._context_pool.get_nowait()
except asyncio.QueueEmpty:
break
for context in contexts:
with contextlib.suppress(Exception):
await context.close()
async def _render_with_context_pool(
self,
html: str,
template_path: str,
render_options: dict[str, Any],
) -> bytes:
context = await self._acquire_context()
generation, context = await self._acquire_context()
page = None
broken = False
try:
@@ -718,7 +1125,7 @@ class PlaywrightEngine(BaseScreenshotEngine):
if page is not None:
with contextlib.suppress(Exception):
await page.close()
await self._release_context(context, broken=broken)
await self._release_context(generation, context, broken=broken)
async def _render_html(
self,
@@ -735,66 +1142,35 @@ class PlaywrightEngine(BaseScreenshotEngine):
async def _recycle_browser(self, reason: str) -> None:
async with self._recycle_lock:
try:
await self._dispose_context_pool()
await _shutdown_browser_instance()
await self._swap_generation(reason)
current_rss = self._get_total_rss()
if current_rss is not None:
self._update_rss_baseline_nolock(current_rss)
await self._prewarm_browser_and_pool()
logger.debug(
f"截图引擎触发回收({reason}),已重建浏览器实例。",
"PlaywrightEngine",
)
await self._log_runtime_snapshot(f"recycle:{reason}")
except Exception as e:
logger.warning("浏览器实例重建失败。", "PlaywrightEngine", e=e)
async def _prewarm_browser_and_pool(self) -> None:
if self._closing:
return
try:
browser = await _get_browser_instance()
except Exception as e:
logger.warning("截图引擎浏览器预热失败。", "PlaywrightEngine", e=e)
async with self._state_lock:
has_active_generation = self._active_generation is not None
if has_active_generation:
return
for _ in range(self._PREWARM_CONTEXT_COUNT):
async with self._state_lock:
if self._closing:
return
if len(self._all_contexts) >= self._CONTEXT_POOL_SIZE:
return
if self._context_pool.qsize() >= self._PREWARM_CONTEXT_COUNT:
return
context = None
try:
context = await browser.new_context(
viewport={"width": 800, "height": 10},
device_scale_factor=2,
)
page = await context.new_page()
await page.goto("about:blank", wait_until="domcontentloaded")
await page.set_content(
"<html><body></body></html>",
wait_until="domcontentloaded",
)
await page.close()
except Exception as e:
logger.warning("截图引擎上下文预热失败。", "PlaywrightEngine", e=e)
if context is not None:
with contextlib.suppress(Exception):
await context.close()
generation = await self._build_generation()
dispose_generation = False
async with self._state_lock:
if self._closing:
dispose_generation = True
elif self._active_generation is None:
self._active_generation = generation
return
else:
dispose_generation = True
async with self._state_lock:
if self._closing:
with contextlib.suppress(Exception):
await context.close()
return
if context in self._all_contexts:
continue
self._all_contexts.add(context)
self._context_pool.put_nowait(context)
if dispose_generation:
await self._dispose_generation(generation)
async def _idle_recycle_loop(self) -> None:
while True:
@@ -804,7 +1180,7 @@ class PlaywrightEngine(BaseScreenshotEngine):
if self._closing:
return
now = time.monotonic()
if self._active_renders > 0:
if _HTMLRENDER_TASK_TRACKER.active_tasks > 0:
continue
if now - self._last_recycle_at < self._RECYCLE_COOLDOWN_SECONDS:
continue
@@ -850,6 +1226,9 @@ class PlaywrightEngine(BaseScreenshotEngine):
return result
async def render(self, html: str, base_url_path: Path, **render_options) -> bytes:
if self._closing or _HTMLRENDER_TASK_TRACKER.is_draining:
raise RuntimeError("截图引擎正在排空/关闭,暂不接受新的渲染任务。")
base_url_for_browser = self._normalize_base_url(base_url_path)
final_render_options = {
@@ -909,6 +1288,12 @@ class EngineManager:
await self._instance.initialize()
return self._instance
async def get_runtime_snapshot(self) -> dict[str, Any]:
engine = await self.get_engine()
if isinstance(engine, PlaywrightEngine):
return await engine.get_runtime_snapshot()
return {"engine": type(engine).__name__}
async def close(self):
if self._instance:
await self._instance.close()
+31 -10
View File
@@ -2,7 +2,7 @@ from collections.abc import Callable
import os
from pathlib import Path
import re
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, ClassVar
from jinja2 import (
ChoiceLoader,
@@ -258,6 +258,34 @@ class ComponentRenderStrategy(RenderStrategy):
class TemplateFileRenderStrategy(RenderStrategy):
"""独立模板文件渲染策略。"""
_env_cache: ClassVar[dict[Path, "RelativePathEnvironment"]] = {}
_ENV_CACHE_MAX = 32
@classmethod
def _get_or_create_env(
cls,
template_dir: Path,
base_loader: Any,
) -> "RelativePathEnvironment":
env = cls._env_cache.get(template_dir)
if env is not None:
return env
temp_loader = FileSystemLoader(str(template_dir))
temp_env_loader = (
ChoiceLoader([temp_loader, base_loader]) if base_loader else temp_loader
)
env = RelativePathEnvironment(
loader=temp_env_loader,
enable_async=True,
autoescape=select_autoescape(["html", "xml"]),
)
if len(cls._env_cache) >= cls._ENV_CACHE_MAX:
cls._env_cache.pop(next(iter(cls._env_cache)))
cls._env_cache[template_dir] = env
return env
async def render(self, context: "RenderContext") -> RenderResult:
component = context.component
template_path = getattr(component, "template_path")
@@ -265,16 +293,9 @@ class TemplateFileRenderStrategy(RenderStrategy):
logger.debug(f"正在渲染独立模板: '{template_path}'", "RendererService")
template_dir = template_path.parent
temp_loader = FileSystemLoader(str(template_dir))
base_loader = context.template_engine.env.loader
temp_env_loader = (
ChoiceLoader([temp_loader, base_loader]) if base_loader else temp_loader
)
temp_env = RelativePathEnvironment(
loader=temp_env_loader,
enable_async=True,
autoescape=select_autoescape(["html", "xml"]),
)
temp_env = self._get_or_create_env(template_dir, base_loader)
temp_env.globals.update(context.template_engine.env.globals)
temp_env.filters.update(context.template_engine.env.filters)
temp_env.globals["asset"] = (
+3 -1
View File
@@ -6,6 +6,8 @@ import os
import anyio.to_thread
from nonebot.drivers import Driver
from zhenxun.utils.manager.priority_manager import PriorityLifecycle
DEFAULT_EXECUTOR_MIN_WORKERS = 16
DEFAULT_EXECUTOR_MAX_WORKERS = 64
DEFAULT_ANYIO_MIN_TOKENS = 32
@@ -76,7 +78,7 @@ def register_runtime_bootstrap(driver: Driver) -> None:
limiter = anyio.to_thread.current_default_thread_limiter()
limiter.total_tokens = _get_anyio_tokens(workers)
@driver.on_shutdown
@PriorityLifecycle.on_shutdown(priority=50)
async def _shutdown_runtime_concurrency() -> None:
global _thread_executor
executor = _thread_executor
+10
View File
@@ -76,3 +76,13 @@ async def _start_send_queue():
patch_send_queue()
for idx in range(_WORKERS):
_WORKER_TASKS.append(asyncio.create_task(_worker(idx)))
@driver.on_shutdown
async def _stop_send_queue():
tasks = _WORKER_TASKS.copy()
_WORKER_TASKS.clear()
for task in tasks:
task.cancel()
if tasks:
await asyncio.gather(*tasks, return_exceptions=True)