mirror of
https://github.com/zhenxun-org/zhenxun_bot.git
synced 2026-10-07 12:50:03 +08:00
✨ feat(hook): 增强认证钩子和运行时缓存管理 (#2106)
* ✨ feat(hook): 增强认证钩子和运行时缓存管理 ``` ♻ refactor(hook): 移除未使用的配置项并优化缓存设置 移除 AUTH_HOOKS_CONCURRENCY_LIMIT 配置项，该配置项不再使用 ✨feat(auth_ban): 简化缓存配置并添加实体参数支持 将 BAN_CACHE_TTL 相关配置从动态配置改为常量定义， 移除复杂的 TTL 值转换逻辑，并为 auth_ban 函数添加可选的 entity 参数以支持外部传入实体信息 ♻ refactor(auth_limit): 移除未使用的配置依赖 移除 AUTH_LIMIT_NOTICE_CD 配置项，直接使用常量值 2 作为限制通知冷却时间 📦 依赖更新: update playwright dependency to version 1.57.0 in pyproject.toml and requirements.txt ``` * ✨ feat(auth_checker): 增强插件模块处理和预过滤逻辑，支持用户插件兼容性 * ✨ feat(http_utils): 添加内容缓存机制以优化HTTP响应处理 * ✨ feat(cache): 添加群组插件设置视图缓存类型并更新相关逻辑 * ✨ feat(renderer): 优化渲染引擎，增加内存缓存管理和HTML文档处理逻辑 * ✨ feat(sign_in): 添加好感度排行和好感度总排行命令 * ✨ feat(renderer): 增强浏览器实例管理和模板预处理，支持历史 include 语法兼容 * ✨ feat(renderer): 优化 Playwright 环境检查逻辑，增加结果缓存以提高性能 * ✨ feat(renderer): 增强模板文件渲染策略，优化资产加载路径处理 * 🚨 auto fix by pre-commit hooks * ✨ feat(log): 增加日志内容安全序列化，避免超长 base64 等污染日志 ✨ feat(log_sanitizer): 添加对嵌入超长 base64/data URI 的清理功能 ✨ feat(auth_checker): 添加 Alconna 快捷方式缓存检查，优化路由匹配逻辑 * feat(renderer): 添加渲染结果内存缓存功能以优化性能 feat(help): 实现帮助菜单图像缓存机制 feat(sign_in): 更新HTML卡片生成以支持动画禁用和剪裁 feat(superuser): 在启动时预热超级用户帮助缓存 feat(theme): 优化主题管理器的资源解析缓存 * 🚨 auto fix by pre-commit hooks * feat(renderer): 增加全页面视口最大宽度限制并优化内容尺寸计算 --------- Co-authored-by: ATTomatoo <1126160939@qq.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: HibiKier <45528451+HibiKier@users.noreply.github.com>
This commit is contained in:
ManyManyTomato
co-authored by
ATTomatoo
pre-commit-ci[bot]
HibiKier
ATTomatoo
pre-commit-ci[bot]
HibiKier
parent
51f4773e14
commit
ce94f63d9a
@@ -1,17 +1,856 @@
|
||||
import asyncio
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Awaitable, Callable
|
||||
import contextlib
|
||||
import hashlib
|
||||
import inspect
|
||||
import json
|
||||
from pathlib import Path
|
||||
import time
|
||||
from typing import Any, ClassVar, cast
|
||||
|
||||
from nonebot_plugin_htmlrender import html_to_pic
|
||||
import nonebot_plugin_htmlrender.browser as htmlrender_browser
|
||||
import psutil
|
||||
|
||||
from zhenxun.configs.config import Config
|
||||
from zhenxun.services.log import logger
|
||||
|
||||
from .types import BaseScreenshotEngine
|
||||
|
||||
|
||||
async def _await_if_needed(value: Any) -> Any:
|
||||
if inspect.isawaitable(value):
|
||||
return await cast(Awaitable[Any], value)
|
||||
return value
|
||||
|
||||
|
||||
async def _get_browser_instance() -> Any:
|
||||
for attr_name in ("get_browser", "get_new_browser"):
|
||||
browser_getter = getattr(htmlrender_browser, attr_name, None)
|
||||
if callable(browser_getter):
|
||||
return await _await_if_needed(browser_getter())
|
||||
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)
|
||||
return
|
||||
|
||||
browser_obj = getattr(htmlrender_browser, "_browser", None)
|
||||
close_func = getattr(browser_obj, "close", None) if browser_obj else None
|
||||
if callable(close_func):
|
||||
with contextlib.suppress(Exception):
|
||||
await _await_if_needed(close_func())
|
||||
|
||||
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):
|
||||
await _await_if_needed(stop_func())
|
||||
|
||||
with contextlib.suppress(Exception):
|
||||
setattr(htmlrender_browser, "_browser", None)
|
||||
with contextlib.suppress(Exception):
|
||||
setattr(htmlrender_browser, "_playwright", None)
|
||||
|
||||
if callable(close_func) or callable(stop_func):
|
||||
return
|
||||
|
||||
logger.debug(
|
||||
"未找到 htmlrender 浏览器关闭函数，跳过 shutdown。",
|
||||
"PlaywrightEngine",
|
||||
)
|
||||
|
||||
|
||||
def _patch_playwright_env_check_once() -> None:
|
||||
if getattr(htmlrender_browser, "_zhenxun_check_once_patched", False):
|
||||
return
|
||||
|
||||
original_check: Callable[..., Awaitable[Any]] | None = None
|
||||
check_attr_name = ""
|
||||
for attr_name in ("check_playwright_env", "check_browser_env"):
|
||||
candidate = getattr(htmlrender_browser, attr_name, None)
|
||||
if callable(candidate):
|
||||
original_check = cast(Callable[..., Awaitable[Any]], candidate)
|
||||
check_attr_name = attr_name
|
||||
break
|
||||
|
||||
if original_check is None:
|
||||
logger.debug(
|
||||
"未找到 htmlrender 环境检查函数，跳过 check_once 补丁。",
|
||||
"PlaywrightEngine",
|
||||
)
|
||||
setattr(htmlrender_browser, "_zhenxun_check_once_patched", True)
|
||||
return
|
||||
|
||||
check_func = original_check
|
||||
state: dict[str, Any] = {"checked": False, "result": None}
|
||||
check_lock: asyncio.Lock | None = None
|
||||
|
||||
def _is_browser_usable(browser_obj: Any) -> bool:
|
||||
if browser_obj is None:
|
||||
return False
|
||||
is_connected = getattr(browser_obj, "is_connected", None)
|
||||
if callable(is_connected):
|
||||
with contextlib.suppress(Exception):
|
||||
return bool(is_connected())
|
||||
# 无法判断连接状态时，保守认为可用
|
||||
return True
|
||||
|
||||
def _get_current_browser_candidate() -> Any:
|
||||
current = state["result"]
|
||||
if _is_browser_usable(current):
|
||||
return current
|
||||
fallback = getattr(htmlrender_browser, "_browser", None)
|
||||
if _is_browser_usable(fallback):
|
||||
return fallback
|
||||
return None
|
||||
|
||||
async def _check_once(**kwargs: Any) -> Any:
|
||||
nonlocal check_lock
|
||||
if state["checked"]:
|
||||
cached_browser = _get_current_browser_candidate()
|
||||
if cached_browser is not None:
|
||||
return cached_browser
|
||||
state["checked"] = False
|
||||
state["result"] = None
|
||||
|
||||
if check_lock is None:
|
||||
check_lock = asyncio.Lock()
|
||||
async with check_lock:
|
||||
if state["checked"]:
|
||||
cached_browser = _get_current_browser_candidate()
|
||||
if cached_browser is not None:
|
||||
return cached_browser
|
||||
state["checked"] = False
|
||||
state["result"] = None
|
||||
|
||||
result = await check_func(**kwargs)
|
||||
state["checked"] = True
|
||||
state["result"] = result
|
||||
|
||||
browser = _get_current_browser_candidate()
|
||||
if browser is not None:
|
||||
return browser
|
||||
return result
|
||||
|
||||
setattr(htmlrender_browser, check_attr_name, _check_once)
|
||||
setattr(htmlrender_browser, "_zhenxun_check_once_patched", True)
|
||||
|
||||
|
||||
class PlaywrightEngine(BaseScreenshotEngine):
|
||||
"""使用 nonebot-plugin-htmlrender 实现的截图引擎。"""
|
||||
|
||||
_MAX_CONCURRENT_RENDER = 2
|
||||
_CONTEXT_POOL_SIZE = 2
|
||||
_PREWARM_CONTEXT_COUNT = 1
|
||||
_SET_CONTENT_WAIT_UNTIL = "domcontentloaded"
|
||||
_READY_STATE_TIMEOUT_MS = 2_000
|
||||
_IMAGE_READY_TIMEOUT_MS = 1_800
|
||||
_FONT_READY_TIMEOUT_MS = 1_200
|
||||
_FULL_PAGE_VIEWPORT_MAX_HEIGHT = 4_096
|
||||
_FULL_PAGE_VIEWPORT_MAX_WIDTH = 4_096
|
||||
_CLIP_PADDING_DEFAULT = 0
|
||||
_DISABLE_ANIMATIONS_STYLE = """
|
||||
*, *::before, *::after {
|
||||
animation: none !important;
|
||||
transition: none !important;
|
||||
caret-color: transparent !important;
|
||||
scroll-behavior: auto !important;
|
||||
}
|
||||
"""
|
||||
_RECENT_RESULT_TTL_SECONDS = 1.5
|
||||
_RECENT_RESULT_MAX_ITEMS = 64
|
||||
_RSS_RECYCLE_MIN_THRESHOLD_BYTES = 700 * 1024 * 1024
|
||||
_RSS_RECYCLE_MAX_THRESHOLD_BYTES = 1200 * 1024 * 1024
|
||||
_RSS_RECYCLE_HEADROOM_BYTES = 224 * 1024 * 1024
|
||||
_RECYCLE_COOLDOWN_SECONDS = 300
|
||||
_RECYCLE_CHECK_EVERY = 8
|
||||
_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",
|
||||
"geolocation",
|
||||
"has_touch",
|
||||
"http_credentials",
|
||||
"ignore_https_errors",
|
||||
"is_mobile",
|
||||
"java_script_enabled",
|
||||
"locale",
|
||||
"permissions",
|
||||
"proxy",
|
||||
"record_har_content",
|
||||
"record_har_mode",
|
||||
"record_har_omit_content",
|
||||
"record_har_path",
|
||||
"record_video_dir",
|
||||
"record_video_size",
|
||||
"reduced_motion",
|
||||
"screen",
|
||||
"service_workers",
|
||||
"storage_state",
|
||||
"timezone_id",
|
||||
"user_agent",
|
||||
}
|
||||
|
||||
def __init__(self):
|
||||
_patch_playwright_env_check_once()
|
||||
self._render_semaphore = asyncio.Semaphore(self._MAX_CONCURRENT_RENDER)
|
||||
self._debug_console_log = bool(Config.get_config("UI", "DEBUG_MODE", False))
|
||||
self._state_lock = asyncio.Lock()
|
||||
self._recycle_lock = asyncio.Lock()
|
||||
self._active_renders = 0
|
||||
self._render_count = 0
|
||||
self._recycle_pending = False
|
||||
self._last_recycle_at = 0.0
|
||||
self._last_render_finished_at = time.monotonic()
|
||||
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._idle_recycle_task: asyncio.Task[None] | None = None
|
||||
self._closing = False
|
||||
self._process = psutil.Process()
|
||||
|
||||
@staticmethod
|
||||
def _normalize_base_url(path: Path) -> str:
|
||||
base_url = path.absolute().as_uri()
|
||||
if not base_url.endswith("/"):
|
||||
base_url += "/"
|
||||
return base_url
|
||||
|
||||
@staticmethod
|
||||
def _build_render_key(
|
||||
html: str, template_path: str, render_options: dict[str, Any]
|
||||
) -> str:
|
||||
options_json = json.dumps(render_options, sort_keys=True, default=str)
|
||||
hasher = hashlib.sha256()
|
||||
hasher.update(template_path.encode("utf-8"))
|
||||
hasher.update(b"\0")
|
||||
hasher.update(options_json.encode("utf-8"))
|
||||
hasher.update(b"\0")
|
||||
hasher.update(html.encode("utf-8", errors="ignore"))
|
||||
return hasher.hexdigest()
|
||||
|
||||
def _cleanup_recent_results_nolock(self, now: float) -> None:
|
||||
while self._recent_results:
|
||||
expire_at, _ = next(iter(self._recent_results.values()))
|
||||
if expire_at > now:
|
||||
break
|
||||
self._recent_results.popitem(last=False)
|
||||
while len(self._recent_results) > self._RECENT_RESULT_MAX_ITEMS:
|
||||
self._recent_results.popitem(last=False)
|
||||
|
||||
def _get_recent_result_nolock(self, key: str, now: float) -> bytes | None:
|
||||
entry = self._recent_results.get(key)
|
||||
if not entry:
|
||||
return None
|
||||
expire_at, result = entry
|
||||
if expire_at <= now:
|
||||
self._recent_results.pop(key, None)
|
||||
return None
|
||||
self._recent_results.move_to_end(key)
|
||||
return result
|
||||
|
||||
def _get_total_rss(self) -> int | None:
|
||||
try:
|
||||
total_rss = self._process.memory_info().rss
|
||||
for child in self._process.children(recursive=True):
|
||||
with contextlib.suppress(Exception):
|
||||
total_rss += child.memory_info().rss
|
||||
return total_rss
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def _update_rss_baseline_nolock(self, current_rss: int) -> None:
|
||||
if self._rss_baseline_bytes is None or current_rss < self._rss_baseline_bytes:
|
||||
self._rss_baseline_bytes = current_rss
|
||||
return
|
||||
|
||||
threshold = self._rss_baseline_bytes + self._RSS_RECYCLE_HEADROOM_BYTES * 2
|
||||
if current_rss >= threshold:
|
||||
self._rss_baseline_bytes = int(
|
||||
self._rss_baseline_bytes * 0.9 + current_rss * 0.1
|
||||
)
|
||||
|
||||
def _get_dynamic_threshold_nolock(self, current_rss: int) -> int:
|
||||
self._update_rss_baseline_nolock(current_rss)
|
||||
baseline = self._rss_baseline_bytes or current_rss
|
||||
dynamic = baseline + self._RSS_RECYCLE_HEADROOM_BYTES
|
||||
dynamic = max(dynamic, self._RSS_RECYCLE_MIN_THRESHOLD_BYTES)
|
||||
dynamic = min(dynamic, self._RSS_RECYCLE_MAX_THRESHOLD_BYTES)
|
||||
return dynamic
|
||||
|
||||
def _mark_recycle_if_needed_nolock(self, now: float) -> None:
|
||||
if self._render_count % self._RECYCLE_CHECK_EVERY != 0:
|
||||
return
|
||||
if now - self._last_recycle_at < self._RECYCLE_COOLDOWN_SECONDS:
|
||||
return
|
||||
current_rss = self._get_total_rss()
|
||||
if current_rss is None:
|
||||
return
|
||||
threshold = self._get_dynamic_threshold_nolock(current_rss)
|
||||
if current_rss >= threshold:
|
||||
self._recycle_pending = True
|
||||
|
||||
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._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()
|
||||
|
||||
async def close(self) -> None:
|
||||
idle_task: asyncio.Task[None] | None = None
|
||||
async with self._state_lock:
|
||||
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
|
||||
|
||||
if idle_task:
|
||||
idle_task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await idle_task
|
||||
|
||||
await self._dispose_context_pool()
|
||||
await _shutdown_browser_instance()
|
||||
|
||||
async def _on_render_begin(self) -> None:
|
||||
async with self._state_lock:
|
||||
self._active_renders += 1
|
||||
|
||||
async def _on_render_end(self) -> None:
|
||||
should_recycle = False
|
||||
async with self._state_lock:
|
||||
self._active_renders = max(0, self._active_renders - 1)
|
||||
self._render_count += 1
|
||||
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:
|
||||
self._recycle_pending = False
|
||||
self._last_recycle_at = now
|
||||
should_recycle = True
|
||||
if should_recycle:
|
||||
await self._recycle_browser("active")
|
||||
|
||||
@staticmethod
|
||||
def _build_page_options(
|
||||
render_options: dict[str, Any], *, pooled: bool
|
||||
) -> dict[str, Any]:
|
||||
options = render_options.copy()
|
||||
options.pop("wait", None)
|
||||
options.pop("type", None)
|
||||
options.pop("quality", None)
|
||||
options.pop("scale", None)
|
||||
options.pop("screenshot_scale", None)
|
||||
options.pop("screenshot_timeout", None)
|
||||
options.pop("full_page", None)
|
||||
options.pop("clip_selector", None)
|
||||
options.pop("clip_padding", None)
|
||||
options.pop("disable_animations", None)
|
||||
if pooled:
|
||||
options.pop("base_url", None)
|
||||
return options
|
||||
|
||||
@staticmethod
|
||||
def _build_screenshot_options(render_options: dict[str, Any]) -> dict[str, Any]:
|
||||
scale = render_options.get("screenshot_scale", render_options.get("scale"))
|
||||
if scale not in ("css", "device"):
|
||||
scale = None
|
||||
return {
|
||||
"full_page": bool(render_options.get("full_page", True)),
|
||||
"type": render_options.get("type", "png"),
|
||||
"quality": render_options.get("quality"),
|
||||
"scale": scale,
|
||||
"timeout": render_options.get("screenshot_timeout", 30_000),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _get_wait_timeout(render_options: dict[str, Any]) -> int:
|
||||
wait = render_options.get("wait", 0)
|
||||
if isinstance(wait, int):
|
||||
return max(wait, 0)
|
||||
return 0
|
||||
|
||||
@staticmethod
|
||||
def _coerce_non_negative_int(value: Any, default: int = 0) -> int:
|
||||
try:
|
||||
value_int = int(value)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
return value_int if value_int >= 0 else default
|
||||
|
||||
@classmethod
|
||||
def _should_use_context_pool(cls, render_options: dict[str, Any]) -> bool:
|
||||
for key in cls._POOL_UNSAFE_OPTION_KEYS:
|
||||
if key in render_options:
|
||||
return False
|
||||
return True
|
||||
|
||||
async def _render_with_page(
|
||||
self,
|
||||
page: Any,
|
||||
html: str,
|
||||
template_path: str,
|
||||
render_options: dict[str, Any],
|
||||
) -> 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.set_content(html, wait_until=self._SET_CONTENT_WAIT_UNTIL)
|
||||
if bool(render_options.get("disable_animations", False)):
|
||||
await self._disable_page_animations(page)
|
||||
await self._wait_for_visual_stability(page)
|
||||
if wait_ms := self._get_wait_timeout(render_options):
|
||||
await page.wait_for_timeout(wait_ms)
|
||||
screenshot_options = self._build_screenshot_options(render_options)
|
||||
clip_selector = render_options.get("clip_selector")
|
||||
if isinstance(clip_selector, str) and clip_selector.strip():
|
||||
if image_bytes := await self._capture_by_selector(
|
||||
page,
|
||||
selector=clip_selector.strip(),
|
||||
screenshot_options=screenshot_options,
|
||||
clip_padding=self._coerce_non_negative_int(
|
||||
render_options.get("clip_padding"),
|
||||
self._CLIP_PADDING_DEFAULT,
|
||||
),
|
||||
):
|
||||
return image_bytes
|
||||
await self._optimize_full_page_capture(page, screenshot_options)
|
||||
return await page.screenshot(**screenshot_options)
|
||||
|
||||
async def _disable_page_animations(self, page: Any) -> None:
|
||||
with contextlib.suppress(Exception):
|
||||
await page.add_style_tag(content=self._DISABLE_ANIMATIONS_STYLE)
|
||||
|
||||
async def _capture_by_selector(
|
||||
self,
|
||||
page: Any,
|
||||
selector: str,
|
||||
screenshot_options: dict[str, Any],
|
||||
clip_padding: int,
|
||||
) -> bytes | None:
|
||||
element = await page.query_selector(selector)
|
||||
if element is None:
|
||||
return None
|
||||
|
||||
element_screenshot_options = {
|
||||
"type": screenshot_options.get("type", "png"),
|
||||
"quality": screenshot_options.get("quality"),
|
||||
"timeout": screenshot_options.get("timeout", 30_000),
|
||||
}
|
||||
with contextlib.suppress(Exception):
|
||||
box = await element.bounding_box()
|
||||
if box and clip_padding > 0:
|
||||
viewport = page.viewport_size or {}
|
||||
width = int(viewport.get("width") or 0)
|
||||
if width > 0:
|
||||
target_height = int(box["y"] + box["height"] + clip_padding)
|
||||
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,
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
if clip_padding <= 0:
|
||||
return await element.screenshot(**element_screenshot_options)
|
||||
|
||||
with contextlib.suppress(Exception):
|
||||
clip_box = await element.bounding_box()
|
||||
if clip_box is None:
|
||||
return await element.screenshot(**element_screenshot_options)
|
||||
clip = {
|
||||
"x": max(clip_box["x"] - clip_padding, 0),
|
||||
"y": max(clip_box["y"] - clip_padding, 0),
|
||||
"width": clip_box["width"] + clip_padding * 2,
|
||||
"height": clip_box["height"] + clip_padding * 2,
|
||||
}
|
||||
page_options = {
|
||||
"type": screenshot_options.get("type", "png"),
|
||||
"quality": screenshot_options.get("quality"),
|
||||
"timeout": screenshot_options.get("timeout", 30_000),
|
||||
"clip": clip,
|
||||
}
|
||||
return await page.screenshot(**page_options)
|
||||
|
||||
return await element.screenshot(**element_screenshot_options)
|
||||
|
||||
async def _wait_for_visual_stability(self, page: Any) -> None:
|
||||
with contextlib.suppress(Exception):
|
||||
await page.wait_for_function(
|
||||
"() => document.readyState === 'complete'",
|
||||
timeout=self._READY_STATE_TIMEOUT_MS,
|
||||
)
|
||||
|
||||
with contextlib.suppress(Exception):
|
||||
await page.wait_for_function(
|
||||
"() => Array.from(document.images || []).every(img => img.complete)",
|
||||
timeout=self._IMAGE_READY_TIMEOUT_MS,
|
||||
)
|
||||
|
||||
with contextlib.suppress(Exception):
|
||||
await page.evaluate(
|
||||
"""
|
||||
async (timeoutMs) => {
|
||||
if (!document.fonts || !document.fonts.ready) return;
|
||||
await Promise.race([
|
||||
document.fonts.ready,
|
||||
new Promise(resolve => setTimeout(resolve, timeoutMs)),
|
||||
]);
|
||||
}
|
||||
""",
|
||||
self._FONT_READY_TIMEOUT_MS,
|
||||
)
|
||||
|
||||
async def _optimize_full_page_capture(
|
||||
self, page: Any, screenshot_options: dict[str, Any]
|
||||
) -> None:
|
||||
if not bool(screenshot_options.get("full_page")):
|
||||
return
|
||||
|
||||
viewport = page.viewport_size or {}
|
||||
width = viewport.get("width")
|
||||
height = viewport.get("height")
|
||||
if (
|
||||
not isinstance(width, int)
|
||||
or width <= 0
|
||||
or not isinstance(height, int)
|
||||
or height <= 0
|
||||
):
|
||||
return
|
||||
|
||||
with contextlib.suppress(Exception):
|
||||
content_size = await page.evaluate(
|
||||
"""
|
||||
() => {
|
||||
const body = document.body;
|
||||
const doc = document.documentElement;
|
||||
const bodyWidth = body ? Math.max(
|
||||
body.scrollWidth,
|
||||
body.offsetWidth,
|
||||
body.clientWidth
|
||||
) : 0;
|
||||
const bodyHeight = body ? Math.max(
|
||||
body.scrollHeight,
|
||||
body.offsetHeight,
|
||||
body.clientHeight
|
||||
) : 0;
|
||||
const docWidth = doc ? Math.max(
|
||||
doc.scrollWidth,
|
||||
doc.offsetWidth,
|
||||
doc.clientWidth
|
||||
) : 0;
|
||||
const docHeight = doc ? Math.max(
|
||||
doc.scrollHeight,
|
||||
doc.offsetHeight,
|
||||
doc.clientHeight
|
||||
) : 0;
|
||||
return {
|
||||
width: Math.ceil(Math.max(bodyWidth, docWidth, 10)),
|
||||
height: Math.ceil(Math.max(bodyHeight, docHeight, 10)),
|
||||
};
|
||||
}
|
||||
"""
|
||||
)
|
||||
if not isinstance(content_size, dict):
|
||||
return
|
||||
|
||||
content_width = content_size.get("width")
|
||||
content_height = content_size.get("height")
|
||||
if not isinstance(content_width, int) or not isinstance(
|
||||
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
|
||||
):
|
||||
return
|
||||
|
||||
target_width = max(width, content_width)
|
||||
target_height = max(height, content_height)
|
||||
await page.set_viewport_size(
|
||||
{"width": target_width, "height": target_height}
|
||||
)
|
||||
screenshot_options["full_page"] = False
|
||||
|
||||
async def _render_with_oneoff_page(
|
||||
self,
|
||||
html: str,
|
||||
template_path: str,
|
||||
render_options: dict[str, Any],
|
||||
) -> bytes:
|
||||
browser = await _get_browser_instance()
|
||||
page_options = self._build_page_options(render_options, pooled=False)
|
||||
page = await browser.new_page(**page_options)
|
||||
try:
|
||||
return await self._render_with_page(
|
||||
page, html, template_path, render_options
|
||||
)
|
||||
finally:
|
||||
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 with self._state_lock:
|
||||
if len(self._all_contexts) < self._CONTEXT_POOL_SIZE:
|
||||
create_new = True
|
||||
else:
|
||||
create_new = False
|
||||
|
||||
if create_new:
|
||||
browser = await _get_browser_instance()
|
||||
context = await browser.new_context(
|
||||
viewport={"width": 800, "height": 10},
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
if broken:
|
||||
await self._discard_context(context)
|
||||
|
||||
async def _discard_context(self, context: Any) -> None:
|
||||
async with self._state_lock:
|
||||
existed = context in self._all_contexts
|
||||
if existed:
|
||||
self._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()
|
||||
page = None
|
||||
broken = False
|
||||
try:
|
||||
page = await context.new_page()
|
||||
page_options = self._build_page_options(render_options, pooled=True)
|
||||
viewport = page_options.get("viewport")
|
||||
if isinstance(viewport, dict):
|
||||
width = viewport.get("width")
|
||||
height = viewport.get("height")
|
||||
if isinstance(width, int) and isinstance(height, int):
|
||||
await page.set_viewport_size({"width": width, "height": height})
|
||||
return await self._render_with_page(
|
||||
page, html, template_path, render_options
|
||||
)
|
||||
except Exception:
|
||||
broken = True
|
||||
raise
|
||||
finally:
|
||||
if page is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
await page.close()
|
||||
await self._release_context(context, broken=broken)
|
||||
|
||||
async def _render_html(
|
||||
self,
|
||||
html: str,
|
||||
template_path: str,
|
||||
render_options: dict[str, Any],
|
||||
) -> bytes:
|
||||
if self._should_use_context_pool(render_options):
|
||||
return await self._render_with_context_pool(
|
||||
html, template_path, render_options
|
||||
)
|
||||
return await self._render_with_oneoff_page(html, template_path, render_options)
|
||||
|
||||
async def _recycle_browser(self, reason: str) -> None:
|
||||
async with self._recycle_lock:
|
||||
try:
|
||||
await self._dispose_context_pool()
|
||||
await _shutdown_browser_instance()
|
||||
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",
|
||||
)
|
||||
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)
|
||||
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()
|
||||
return
|
||||
|
||||
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)
|
||||
|
||||
async def _idle_recycle_loop(self) -> None:
|
||||
while True:
|
||||
await asyncio.sleep(self._IDLE_CHECK_INTERVAL_SECONDS)
|
||||
should_recycle = False
|
||||
async with self._state_lock:
|
||||
if self._closing:
|
||||
return
|
||||
now = time.monotonic()
|
||||
if self._active_renders > 0:
|
||||
continue
|
||||
if now - self._last_recycle_at < self._RECYCLE_COOLDOWN_SECONDS:
|
||||
continue
|
||||
idle_for = now - self._last_render_finished_at
|
||||
if idle_for < self._IDLE_RECYCLE_SECONDS:
|
||||
continue
|
||||
current_rss = self._get_total_rss()
|
||||
if current_rss is None:
|
||||
continue
|
||||
threshold = self._get_dynamic_threshold_nolock(current_rss)
|
||||
if current_rss >= threshold:
|
||||
self._last_recycle_at = now
|
||||
should_recycle = True
|
||||
if should_recycle:
|
||||
await self._recycle_browser("idle")
|
||||
|
||||
async def _render_and_store_result(
|
||||
self,
|
||||
key: str,
|
||||
html: str,
|
||||
base_url_for_browser: str,
|
||||
render_options: dict[str, Any],
|
||||
) -> bytes:
|
||||
async with self._render_semaphore:
|
||||
await self._on_render_begin()
|
||||
try:
|
||||
result = await self._render_html(
|
||||
html,
|
||||
base_url_for_browser,
|
||||
render_options,
|
||||
)
|
||||
finally:
|
||||
await self._on_render_end()
|
||||
|
||||
async with self._state_lock:
|
||||
now = time.monotonic()
|
||||
self._recent_results[key] = (
|
||||
now + self._RECENT_RESULT_TTL_SECONDS,
|
||||
result,
|
||||
)
|
||||
self._recent_results.move_to_end(key)
|
||||
self._cleanup_recent_results_nolock(now)
|
||||
return result
|
||||
|
||||
async def render(self, html: str, base_url_path: Path, **render_options) -> bytes:
|
||||
base_url_for_browser = base_url_path.absolute().as_uri()
|
||||
if not base_url_for_browser.endswith("/"):
|
||||
base_url_for_browser += "/"
|
||||
base_url_for_browser = self._normalize_base_url(base_url_path)
|
||||
|
||||
final_render_options = {
|
||||
"viewport": {"width": 800, "height": 10},
|
||||
@@ -19,12 +858,40 @@ class PlaywrightEngine(BaseScreenshotEngine):
|
||||
"base_url": base_url_for_browser,
|
||||
}
|
||||
|
||||
return await html_to_pic(
|
||||
html=html,
|
||||
template_path=base_url_for_browser,
|
||||
**final_render_options,
|
||||
dedupe_key = self._build_render_key(
|
||||
html,
|
||||
base_url_for_browser,
|
||||
final_render_options,
|
||||
)
|
||||
|
||||
owner = False
|
||||
async with self._state_lock:
|
||||
now = time.monotonic()
|
||||
self._cleanup_recent_results_nolock(now)
|
||||
if cached_result := self._get_recent_result_nolock(dedupe_key, now):
|
||||
return cached_result
|
||||
|
||||
task = self._inflight_tasks.get(dedupe_key)
|
||||
if task is None:
|
||||
task = asyncio.create_task(
|
||||
self._render_and_store_result(
|
||||
dedupe_key,
|
||||
html,
|
||||
base_url_for_browser,
|
||||
final_render_options,
|
||||
)
|
||||
)
|
||||
self._inflight_tasks[dedupe_key] = task
|
||||
owner = True
|
||||
|
||||
try:
|
||||
return await task
|
||||
finally:
|
||||
if owner:
|
||||
async with self._state_lock:
|
||||
if self._inflight_tasks.get(dedupe_key) is task:
|
||||
self._inflight_tasks.pop(dedupe_key, None)
|
||||
|
||||
|
||||
class EngineManager:
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user