fix：客户端断开连接时，服务端处理websocket异常

backend/controller/common.py

@@ -1,4 +1,9 @@
+import asyncio
+
+from starlette.websockets import WebSocket
+
 from core.security import generate_token, verify_password
+from core.utils import get_system_info
 from dbhelper.user import get_user
 from schemas import LoginForm, LoginResult, Response
 
@@ -18,3 +23,14 @@ async def login(auth_data: LoginForm) -> Response[LoginResult]:
 async def about() -> Response:
     """关于"""
     pass
+
+
+async def websocket(ws: WebSocket):
+    await ws.accept()
+    try:
+        while True:
+            await asyncio.sleep(1)
+            await ws.send_json(get_system_info())
+    except Exception as e:
+        print("断开了链接", e)
+        await ws.close()

backend/controller/ws.py

@@ -1,14 +0,0 @@
-from fastapi import FastAPI
-from fastapi.websockets import WebSocket
-
-from core.utils import get_system_info
-
-# websocket app
-ws_app = FastAPI()
-
-
-@ws_app.websocket("/ws")
-async def ws_func(ws: WebSocket):
-    await ws.accept()
-    while True:
-        await ws.send_json(get_system_info())

backend/core/utils.py

@@ -1,5 +1,4 @@
 import random
-import time
 
 
 def list_to_tree(
@@ -32,7 +31,6 @@ def list_to_tree(
 
 def get_system_info():
     """获取系统信息"""
-    time.sleep(1)
     return {
         "usage": {
             "cpu": f"{random.random() * 100: .2}",

backend/main.py

@@ -1,6 +1,5 @@
 from fastapi import FastAPI
 
-from controller.ws import ws_app
 from core.events import close_orm, init_orm
 from core.exceptions import exception_handlers
 from core.log import logger
@@ -15,14 +14,13 @@ app = FastAPI(
     exception_handlers=exception_handlers,
 )
 
-app.mount("/", ws_app)
 
 if __name__ == "__main__":
     import uvicorn
-    from fastapi.routing import Mount
+    from fastapi.routing import APIWebSocketRoute
 
     for i in app.routes:
-        if not isinstance(i, Mount):
+        if not isinstance(i, APIWebSocketRoute):
             logger.info(
                 f"{i.path}, {i.methods}, {i.__dict__.get('summary')}, {i.endpoint}"
             )

backend/router/url.py

@@ -2,7 +2,7 @@ from typing import Any, Callable, get_type_hints
 
 from fastapi import Depends, routing
 
-from controller.common import about, login
+from controller.common import about, login, websocket
 from controller.menu import menu_add, menu_arr, menu_del, menu_put
 from controller.role import (
     role_add,
@@ -172,6 +172,7 @@ routes = [
     Route.put(
         "/menu/{pk}", endpoint=menu_put, tags=["菜单管理"], summary="菜单更新", **has_perm
     ),
+    routing.APIWebSocketRoute("/ws", endpoint=websocket),
 ]
 
 __all__ = [routes]

frontend/src/views/main/setting/about/about.vue

@@ -40,7 +40,6 @@ const { systemUsage, performance } = toRefs(data)
 
 onUnmounted(() => {
   ws.close()
-  console.log('关闭socket 连接')
 })
 </script>