mirror of
https://github.com/zhenxun-org/zhenxun_bot.git
synced 2026-10-08 21:30:01 +08:00
添加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:
co-authored by
pre-commit-ci[bot]
ManyManyTomato
AkashiCoin
parent
74bf912d04
commit
8b16126e40
@@ -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
|
||||
|
||||
Vendored
+20
-301
@@ -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
@@ -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,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)
|
||||
|
||||
# 收集所有字段值
|
||||
|
||||
@@ -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}")
|
||||
|
||||
@@ -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
@@ -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()
|
||||
|
||||
@@ -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"] = (
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user