性能优化 (#2126)

* 性能优化

* 代码改进

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

* 统一缓存与生命周期

* 添加aiomysql依赖

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

* 🚨 auto fix by pre-commit hooks

* 优化Windows下的关闭逻辑

* 代码优化

* bugfix:修复配置重载问题

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

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

* 优化 Windows launcher 关闭重启兜底

---------

Co-authored-by: HibiKier <775757368@qq.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

zhenxun/builtin_plugins/web_ui/api/logs/log_manager.py

@@ -1,33 +1,97 @@
 import asyncio
+from collections import deque
 from collections.abc import Awaitable, Callable
-from typing import Generic, TypeVar
+import contextlib
 
-_T = TypeVar("_T")
-LogListener = Callable[[_T], Awaitable[None]]
+from nonebot.log import default_filter, default_format
+
+from zhenxun.services.log import logger_
+
+LogListener = Callable[[str], Awaitable[None]]
+DEFAULT_MAX_LOGS = 1000
+DEFAULT_MAX_LISTENERS = 16
 
 
-class LogStorage(Generic[_T]):
+class LogStorage:
     """
     日志存储
     """
 
-    def __init__(self, rotation: float = 5 * 60):
+    def __init__(
+        self,
+        rotation: float = 5 * 60,
+        max_logs: int = DEFAULT_MAX_LOGS,
+        max_listeners: int = DEFAULT_MAX_LISTENERS,
+    ):
         self.count, self.rotation = 0, rotation
+        self.max_logs = max_logs
+        self.max_listeners = max_listeners
         self.logs: dict[int, str] = {}
-        self.listeners: set[LogListener[str]] = set()
+        self._order: deque[int] = deque()
+        self.listeners: set[LogListener] = set()
 
     async def add(self, log: str):
         seq = self.count = self.count + 1
         self.logs[seq] = log
+        self._order.append(seq)
+        self._trim()
         asyncio.get_running_loop().call_later(self.rotation, self.remove, seq)
-        await asyncio.gather(
-            *(listener(log) for listener in self.listeners),
-            return_exceptions=True,
-        )
+        listeners = tuple(self.listeners)
+        if listeners:
+            results = await asyncio.gather(
+                *(listener(log) for listener in listeners),
+                return_exceptions=True,
+            )
+            for listener, result in zip(listeners, results, strict=False):
+                if isinstance(result, BaseException):
+                    self.listeners.discard(listener)
         return seq
 
+    def add_listener(self, listener: LogListener) -> bool:
+        if len(self.listeners) >= self.max_listeners:
+            return False
+        self.listeners.add(listener)
+        return True
+
+    def remove_listener(self, listener: LogListener) -> None:
+        self.listeners.discard(listener)
+
     def remove(self, seq: int):
-        del self.logs[seq]
+        self.logs.pop(seq, None)
+        with contextlib.suppress(ValueError):
+            self._order.remove(seq)
+
+    def _trim(self) -> None:
+        while self._order and self._order[0] not in self.logs:
+            self._order.popleft()
+        while len(self.logs) > self.max_logs and self._order:
+            self.logs.pop(self._order.popleft(), None)
 
 
-LOG_STORAGE: LogStorage[str] = LogStorage[str]()
+LOG_STORAGE = LogStorage()
+
+_LOG_SINK_ID: int | None = None
+
+
+async def ensure_log_sink_started() -> None:
+    global _LOG_SINK_ID
+    if _LOG_SINK_ID is not None:
+        return
+
+    async def log_sink(message: str) -> None:
+        await LOG_STORAGE.add(message.rstrip("\n"))
+
+    _LOG_SINK_ID = logger_.add(
+        log_sink,
+        colorize=True,
+        filter=default_filter,
+        format=default_format,
+    )
+
+
+def stop_log_sink_if_idle() -> None:
+    global _LOG_SINK_ID
+    if LOG_STORAGE.listeners or _LOG_SINK_ID is None:
+        return
+    logger_.remove(_LOG_SINK_ID)
+    _LOG_SINK_ID = None

zhenxun/builtin_plugins/web_ui/api/logs/logs.py

@@ -3,7 +3,7 @@ from loguru import logger
 from nonebot.utils import escape_tag
 from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState
 
-from .log_manager import LOG_STORAGE
+from .log_manager import LOG_STORAGE, ensure_log_sink_started, stop_log_sink_if_idle
 
 router = APIRouter()
 
@@ -11,11 +11,16 @@ router = APIRouter()
 @router.websocket("/logs")
 async def system_logs_realtime(websocket: WebSocket):
     await websocket.accept()
+    await ensure_log_sink_started()
 
     async def log_listener(log: str):
         await websocket.send_text(log)
 
-    LOG_STORAGE.listeners.add(log_listener)
+    if not LOG_STORAGE.add_listener(log_listener):
+        await websocket.send_text("日志连接数已达上限，请稍后再试。")
+        await websocket.close()
+        stop_log_sink_if_idle()
+        return
     try:
         while websocket.client_state == WebSocketState.CONNECTED:
             recv = await websocket.receive()
@@ -26,4 +31,5 @@ async def system_logs_realtime(websocket: WebSocket):
     except WebSocketDisconnect:
         pass
     finally:
-        LOG_STORAGE.listeners.remove(log_listener)
+        LOG_STORAGE.remove_listener(log_listener)
+        stop_log_sink_if_idle()

zhenxun/builtin_plugins/web_ui/api/tabs/main/__init__.py

@@ -34,6 +34,31 @@ run_time = time.time()
 
 ws_router = APIRouter()
 router = APIRouter(prefix="/main")
+_SYSTEM_STATUS_CONNECTIONS: set[WebSocket] = set()
+_SYSTEM_STATUS_STOPPING = False
+
+
+async def _close_system_status_websocket(websocket: WebSocket) -> None:
+    with contextlib.suppress(Exception):
+        if websocket.client_state == WebSocketState.CONNECTED:
+            await asyncio.wait_for(
+                websocket.close(code=1001, reason="server shutdown"),
+                timeout=2,
+            )
+
+
+@driver.on_shutdown
+async def _close_system_status_websockets() -> None:
+    global _SYSTEM_STATUS_STOPPING
+    _SYSTEM_STATUS_STOPPING = True
+    websockets = list(_SYSTEM_STATUS_CONNECTIONS)
+    if not websockets:
+        return
+    await asyncio.gather(
+        *(_close_system_status_websocket(websocket) for websocket in websockets),
+        return_exceptions=True,
+    )
+    _SYSTEM_STATUS_CONNECTIONS.clear()
 
 
 @router.get(
@@ -243,11 +268,39 @@ async def _(param: BotManageUpdateParam):
 @ws_router.websocket("/system_status")
 async def system_logs_realtime(websocket: WebSocket, sleep: int = 5):
     await websocket.accept()
+    _SYSTEM_STATUS_CONNECTIONS.add(websocket)
     logger.debug("ws system_status is connect")
-    with contextlib.suppress(
-        WebSocketDisconnect, ConnectionClosedError, ConnectionClosedOK
-    ):
-        while websocket.client_state == WebSocketState.CONNECTED:
+
+    disconnect_event = asyncio.Event()
+
+    async def _watch_disconnect() -> None:
+        try:
+            while websocket.client_state == WebSocketState.CONNECTED:
+                await websocket.receive()
+        except (WebSocketDisconnect, ConnectionClosedError, ConnectionClosedOK):
+            pass
+        except Exception as e:
+            logger.debug(f"ws system_status receive stopped: {type(e).__name__}")
+        finally:
+            disconnect_event.set()
+
+    receive_task = asyncio.create_task(_watch_disconnect())
+    try:
+        while (
+            websocket.client_state == WebSocketState.CONNECTED
+            and not _SYSTEM_STATUS_STOPPING
+        ):
             system_status = await get_system_status()
-            await websocket.send_text(system_status.json())
-            await asyncio.sleep(sleep)
+            await asyncio.wait_for(websocket.send_text(system_status.json()), timeout=5)
+            try:
+                await asyncio.wait_for(disconnect_event.wait(), timeout=max(sleep, 1))
+            except TimeoutError:
+                pass
+    except (WebSocketDisconnect, ConnectionClosedError, ConnectionClosedOK):
+        pass
+    finally:
+        _SYSTEM_STATUS_CONNECTIONS.discard(websocket)
+        receive_task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await receive_task
+        await _close_system_status_websocket(websocket)

zhenxun/builtin_plugins/web_ui/api/tabs/plugin_manage/__init__.py

@@ -52,34 +52,20 @@ async def _(
 async def _() -> Result[PluginCount]:
     try:
         plugin_count = PluginCount()
-        plugin_count.normal = len(
-            await DbPluginInfo.get_plugins(
-                plugin_type=PluginType.NORMAL,
-                load_status=True,
-                filter_parent=False,
-            )
-        )
-        plugin_count.admin = len(
-            await DbPluginInfo.get_plugins(
-                plugin_type__in=[PluginType.ADMIN, PluginType.SUPER_AND_ADMIN],
-                load_status=True,
-                filter_parent=False,
-            )
-        )
-        plugin_count.superuser = len(
-            await DbPluginInfo.get_plugins(
-                plugin_type__in=[PluginType.SUPERUSER, PluginType.SUPER_AND_ADMIN],
-                load_status=True,
-                filter_parent=False,
-            )
-        )
-        plugin_count.other = len(
-            await DbPluginInfo.get_plugins(
-                plugin_type__in=[PluginType.HIDDEN, PluginType.DEPENDANT],
-                load_status=True,
-                filter_parent=False,
-            )
+        plugins = await DbPluginInfo.get_plugins(
+            load_status=True,
+            filter_parent=False,
         )
+        for plugin in plugins:
+            plugin_type = plugin.plugin_type
+            if plugin_type == PluginType.NORMAL:
+                plugin_count.normal += 1
+            if plugin_type in {PluginType.ADMIN, PluginType.SUPER_AND_ADMIN}:
+                plugin_count.admin += 1
+            if plugin_type in {PluginType.SUPERUSER, PluginType.SUPER_AND_ADMIN}:
+                plugin_count.superuser += 1
+            if plugin_type in {PluginType.HIDDEN, PluginType.DEPENDANT}:
+                plugin_count.other += 1
         return Result.ok(plugin_count, "拿到信息啦!")
     except Exception as e:
         logger.error(f"{router.prefix}/get_plugin_count 调用错误", "WebUi", e=e)

zhenxun/builtin_plugins/web_ui/api/tabs/plugin_manage/data_source.py

@@ -111,10 +111,19 @@ class ApiDataSource:
         other_update_fields = set()
         updated_count = 0
         errors = []
+        modules = [item.module for item in params.updates]
+        plugin_records = await DbPluginInfo.get_plugins(
+            module__in=modules,
+            load_status=None,
+            filter_parent=False,
+        )
+        plugin_map = {plugin.module: plugin for plugin in plugin_records}
 
         for item in params.updates:
             try:
-                db_plugin = await DbPluginInfo.get(module=item.module)
+                db_plugin = plugin_map.get(item.module)
+                if db_plugin is None:
+                    raise DoesNotExist()
                 plugin_changed_other = False
                 plugin_changed_block = False