WS方案

ws/ws.py

@@ -1,223 +1,157 @@
-from fastapi import WebSocket, APIRouter, WebSocketDisconnect, FastAPI
-from typing import Dict, Any, Optional
-import datetime
 import asyncio
+import datetime
 import json
 from contextlib import asynccontextmanager
+from typing import Dict, Optional
 
-# 创建WebSocket路由
+import numpy as np
+from fastapi import WebSocket, APIRouter, WebSocketDisconnect, FastAPI
+
+# -------------------------- 配置常量（简化硬编码） --------------------------
+HEARTBEAT_INTERVAL = 30  # 心跳检查间隔（秒）
+HEARTBEAT_TIMEOUT = 60  # 客户端超时阈值（秒）
+WS_ENDPOINT = "/ws"  # WebSocket端点路径
+
+# -------------------------- 核心数据结构与全局变量 --------------------------
 ws_router = APIRouter()
 
 
-# 客户端连接信息数据结构
+# 客户端连接封装（仅保留核心属性和方法）
 class ClientConnection:
     def __init__(self, websocket: WebSocket, client_ip: str):
         self.websocket = websocket
         self.client_ip = client_ip
-        self.last_heartbeat = datetime.datetime.now()  # 初始心跳时间为连接时间
-
-    def update_heartbeat(self):
-        """更新心跳时间为当前时间"""
         self.last_heartbeat = datetime.datetime.now()
-        print(f"[{datetime.datetime.now():%Y-%m-%d %H:%M:%S}] 客户端 {self.client_ip} 心跳时间已更新")
 
-    def is_alive(self, timeout_seconds: int = 60) -> bool:
-        """检查客户端是否活跃（心跳超时阈值：60秒）"""
+    # 更新心跳时间
+    def update_heartbeat(self):
+        self.last_heartbeat = datetime.datetime.now()
+
+    # 检查是否存活（超时返回False）
+    def is_alive(self) -> bool:
         timeout = (datetime.datetime.now() - self.last_heartbeat).total_seconds()
-        print(f"[{datetime.datetime.now():%Y-%m-%d %H:%M:%S}] 客户端 {self.client_ip} 心跳检查："
-              f"上次心跳距今 {timeout:.1f} 秒（阈值：{timeout_seconds}秒）")
-        return timeout < timeout_seconds
+        return timeout < HEARTBEAT_TIMEOUT
 
 
-# 存储所有已连接的客户端（key：客户端IP、value：ClientConnection对象）
+# 全局连接管理（IP -> 连接实例）
 connected_clients: Dict[str, ClientConnection] = {}
-
-# 心跳检查任务引用（全局变量、用于应用关闭时取消任务）
+# 心跳任务（全局引用，用于关闭时清理）
 heartbeat_task: Optional[asyncio.Task] = None
 
 
+# -------------------------- 心跳检查逻辑（精简日志） --------------------------
 async def heartbeat_checker():
-    """定期检查客户端心跳（每30秒一次）、超时直接剔除（不发通知）"""
     while True:
-        current_time = datetime.datetime.now()
-        print(f"\n[{current_time:%Y-%m-%d %H:%M:%S}] === 开始新一轮心跳检查 ===")
-        print(f"[{current_time:%Y-%m-%d %H:%M:%S}] 当前在线客户端总数：{len(connected_clients)}")
+        now = datetime.datetime.now()
+        # 1. 筛选超时客户端（避免遍历中修改字典）
+        timeout_ips = [ip for ip, conn in connected_clients.items() if not conn.is_alive()]
 
-        # 1. 收集超时客户端IP（避免遍历中修改字典）
-        timeout_clients = []
-        for client_ip, connection in connected_clients.items():
-            if not connection.is_alive():
-                timeout_clients.append(client_ip)
-
-        # 2. 处理超时客户端（关闭连接+移除记录）
-        if timeout_clients:
-            print(f"[{current_time:%Y-%m-%d %H:%M:%S}] 发现超时客户端：{timeout_clients}（共{len(timeout_clients)}个）")
-            for client_ip in timeout_clients:
+        # 2. 处理超时连接（关闭+移除）
+        if timeout_ips:
+            print(f"[{now:%H:%M:%S}] 心跳检查：{len(timeout_ips)}个客户端超时（{timeout_ips}）")
+            for ip in timeout_ips:
                 try:
-                    connection = connected_clients[client_ip]
-                    await connection.websocket.close(code=1008, reason="心跳超时（>60秒）")
-                    print(f"[{current_time:%Y-%m-%d %H:%M:%S}] 客户端 {client_ip} 已关闭（超时）")
-                except Exception as e:
-                    print(
-                        f"[{current_time:%Y-%m-%d %H:%M:%S}] 关闭客户端 {client_ip} 失败：{str(e)}（错误类型：{type(e).__name__}）")
+                    await connected_clients[ip].websocket.close(code=1008, reason="心跳超时")
                 finally:
-                    if client_ip in connected_clients:
-                        del connected_clients[client_ip]
-                        print(f"[{current_time:%Y-%m-%d %H:%M:%S}] 客户端 {client_ip} 已从连接列表移除")
+                    connected_clients.pop(ip, None)
         else:
-            print(f"[{current_time:%Y-%m-%d %H:%M:%S}] 无超时客户端、心跳检查完成")
+            print(f"[{now:%H:%M:%S}] 心跳检查：{len(connected_clients)}个客户端在线，无超时")
 
-        # 3. 等待30秒后进行下一轮检查
-        await asyncio.sleep(30)
+        # 3. 等待下一轮检查
+        await asyncio.sleep(HEARTBEAT_INTERVAL)
 
 
+# -------------------------- 应用生命周期（简化异常处理） --------------------------
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    """应用生命周期管理：启动时创建心跳任务、关闭时取消任务"""
     global heartbeat_task
+    # 启动心跳任务
     heartbeat_task = asyncio.create_task(heartbeat_checker())
-    print(f"[{datetime.datetime.now():%Y-%m-%d %H:%M:%S}] 心跳检查任务已启动（任务ID：{id(heartbeat_task)}）")
-    yield  # 应用运行中
+    print(f"[{datetime.datetime.now():%H:%M:%S}] 心跳任务启动（ID：{id(heartbeat_task)}）")
+    yield
+    # 关闭时取消心跳任务
     if heartbeat_task and not heartbeat_task.done():
         heartbeat_task.cancel()
         try:
             await heartbeat_task
+            print(f"[{datetime.datetime.now():%H:%M:%S}] 心跳任务已取消")
         except asyncio.CancelledError:
-            print(f"[{datetime.datetime.now():%Y-%m-%d %H:%M:%S}] 心跳检查任务已正常取消")
-        except Exception as e:
-            print(f"[{datetime.datetime.now():%Y-%m-%d %H:%M:%S}] 取消心跳任务时出错：{str(e)}")
+            pass
 
 
-async def send_heartbeat_ack(client_ip: str, client_timestamp: Any) -> bool:
-    """向客户端回复心跳确认（严格遵循 {"timestamp":xxxxx, "type":"heartbeat"} 格式）"""
+# -------------------------- 消息处理（合并冗余逻辑） --------------------------
+async def send_heartbeat_ack(client_ip: str):
+    """回复心跳确认"""
     if client_ip not in connected_clients:
-        print(f"[{datetime.datetime.now():%Y-%m-%d %H:%M:%S}] 回复心跳失败：客户端 {client_ip} 不在连接列表中")
         return False
-
-    # 服务端当前格式化时间戳（字符串类型、与日志时间格式匹配）
-    server_latest_timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    ack_msg = {
-        "timestamp": server_latest_timestamp,
+    try:
+        ack = {
+            "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
             "type": "heartbeat"
         }
-
-    try:
-        connection = connected_clients[client_ip]
-        await connection.websocket.send_json(ack_msg)
-        print(
-            f"[{datetime.datetime.now():%Y-%m-%d %H:%M:%S}] 已向客户端 {client_ip} 回复心跳：{json.dumps(ack_msg, ensure_ascii=False)}")
+        await connected_clients[client_ip].websocket.send_json(ack)
         return True
-    except Exception as e:
-        print(
-            f"[{datetime.datetime.now():%Y-%m-%d %H:%M:%S}] 回复客户端 {client_ip} 心跳失败：{str(e)}（错误类型：{type(e).__name__}）")
-        if client_ip in connected_clients:
-            del connected_clients[client_ip]
-            print(f"[{datetime.datetime.now():%Y-%m-%d %H:%M:%S}] 客户端 {client_ip} 因心跳回复失败被移除")
+    except Exception:
+        connected_clients.pop(client_ip, None)
         return False
 
 
-# 新增：向指定客户端发送消息的方法
-async def send_message_to_client(client_ip: str, message: Dict[str, Any]) -> bool:
-    """
-    向指定客户端IP发送JSON消息（需确保客户端已在连接列表中）
-
-    :param client_ip: 目标客户端IP（必须与`connected_clients`中的key完全一致，如"192.168.1.100"）
-    :param message: 待发送的JSON消息（需为Dict类型，FastAPI会自动序列化为JSON字符串）
-    :return: 发送成功返回True，失败（客户端不存在/发送异常）返回False
-    """
-    current_time = datetime.datetime.now()
-    # 1. 校验客户端是否在线
-    if client_ip not in connected_clients:
-        print(
-            f"[{current_time:%Y-%m-%d %H:%M:%S}] 发送消息失败：客户端 {client_ip} 不在连接列表中（当前在线数：{len(connected_clients)}）")
-        return False
-
+async def handle_text_msg(client_ip: str, text: str, conn: ClientConnection):
+    """处理文本消息（核心：心跳+JSON解析）"""
     try:
-        # 2. 获取客户端连接实例并发送消息
-        client_conn = connected_clients[client_ip]
-        await client_conn.websocket.send_json(message)
-
-        # 3. 打印发送成功日志（格式化消息便于查看）
-        formatted_msg = json.dumps(message, ensure_ascii=False, indent=2)
-        print(f"[{current_time:%Y-%m-%d %H:%M:%S}] 已向客户端 {client_ip} 发送消息：\n{formatted_msg}")
-        return True
-
-    except Exception as e:
-        # 4. 处理发送异常（如网络中断、客户端已离线但未被清理）
-        error_msg = f"[{current_time:%Y-%m-%d %H:%M:%S}] 向客户端 {client_ip} 发送消息失败：{str(e)}（错误类型：{type(e).__name__}）"
-        print(error_msg)
-
-        # 5. 清理无效连接（避免僵尸连接残留）
-        if client_ip in connected_clients:
-            del connected_clients[client_ip]
-            print(f"[{current_time:%Y-%m-%d %H:%M:%S}] 客户端 {client_ip} 因发送异常已从连接列表移除")
-
-        return False
-
-
-@ws_router.websocket("/ws")
-async def websocket_endpoint(websocket: WebSocket):
-    """WebSocket核心端点：处理连接建立/消息接收/连接关闭"""
-    current_time = datetime.datetime.now()
-    # 1. 接受客户端连接请求
-    await websocket.accept()
-    # 获取客户端IP（作为唯一标识）
-    client_ip = websocket.client.host if websocket.client else "unknown_ip"
-    print(f"\n[{current_time:%Y-%m-%d %H:%M:%S}] 客户端 {client_ip} 连接请求已接受（WebSocket握手成功）")
-
-    try:
-        # 2. 处理"同一IP重复连接"场景：关闭旧连接、保留新连接
-        if client_ip in connected_clients:
-            old_connection = connected_clients[client_ip]
-            await old_connection.websocket.close(code=1008, reason="同一IP新连接已建立")
-            del connected_clients[client_ip]
-            print(f"[{current_time:%Y-%m-%d %H:%M:%S}] 已关闭客户端 {client_ip} 的旧连接（新连接已建立）")
-
-        # 3. 注册新客户端到连接列表
-        new_connection = ClientConnection(websocket, client_ip)
-        connected_clients[client_ip] = new_connection
-        print(
-            f"[{current_time:%Y-%m-%d %H:%M:%S}] 客户端 {client_ip} 已注册到连接列表、当前在线数：{len(connected_clients)}")
-
-        # 4. 循环接收客户端消息（持续监听）
-        while True:
-            raw_data = await websocket.receive_text()
-            recv_time = datetime.datetime.now()
-            print(f"\n[{recv_time:%Y-%m-%d %H:%M:%S}] 收到客户端 {client_ip} 的消息：{raw_data}")
-
-            # 尝试解析JSON消息
-            try:
-                message = json.loads(raw_data)
-                print(
-                    f"[{recv_time:%Y-%m-%d %H:%M:%S}] 消息解析成功：{json.dumps(message, ensure_ascii=False, indent=2)}")
-
-                # 5. 区分消息类型：仅处理心跳、其他消息不回复
-                if message.get("type") == "heartbeat":
-                    client_timestamp = message.get("timestamp")
-                    if client_timestamp is None:
-                        print(f"[{recv_time:%Y-%m-%d %H:%M:%S}] 警告：客户端 {client_ip} 发送的心跳缺少'timestamp'字段")
-                        continue
-
-                    new_connection.update_heartbeat()
-                    await send_heartbeat_ack(client_ip, client_timestamp)
+        msg = json.loads(text)
+        # 仅处理心跳类型消息
+        if msg.get("type") == "heartbeat":
+            conn.update_heartbeat()
+            await send_heartbeat_ack(client_ip)
         else:
-                    print(f"[{recv_time:%Y-%m-%d %H:%M:%S}] 非心跳消息（类型：{message.get('type')}）、不回复")
+            print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：收到消息：{msg}")
+    except json.JSONDecodeError:
+        print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：无效JSON消息")
 
-            except json.JSONDecodeError as e:
-                print(f"[{recv_time:%Y-%m-%d %H:%M:%S}] 客户端 {client_ip} 消息格式错误：无效JSON（错误：{str(e)}）")
-            except Exception as e:
-                print(
-                    f"[{recv_time:%Y-%m-%d %H:%M:%S}] 处理客户端 {client_ip} 消息时出错：{str(e)}（错误类型：{type(e).__name__}）")
 
-    except WebSocketDisconnect as e:
-        print(
-            f"\n[{datetime.datetime.now():%Y-%m-%d %H:%M:%S}] 客户端 {client_ip} 主动断开连接（代码：{e.code}、原因：{e.reason}）")
-    except Exception as e:
-        print(
-            f"\n[{datetime.datetime.now():%Y-%m-%d %H:%M:%S}] 客户端 {client_ip} 连接异常：{str(e)}（错误类型：{type(e).__name__}）")
-    finally:
-        # 无论何种退出原因、确保客户端从列表中移除
+async def handle_binary_msg(client_ip: str, data: bytes):
+    """处理二进制消息（保留扩展入口）"""
+    print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：收到{len(data)}字节二进制数据")
+    # 将二进制转化为 numpy 数组
+    data_ndarray = np.frombuffer(data, dtype=np.uint8)
+    # 进行检测
+
+# -------------------------- WebSocket核心端点 --------------------------
+@ws_router.websocket(WS_ENDPOINT)
+async def websocket_endpoint(websocket: WebSocket):
+    # 接受连接 + 获取客户端IP
+    await websocket.accept()
+    client_ip = websocket.client.host if websocket.client else "unknown"
+    now = datetime.datetime.now()
+    print(f"[{now:%H:%M:%S}] 客户端{client_ip}：连接成功")
+
+    try:
+        # 处理重复连接（关闭旧连接）
         if client_ip in connected_clients:
-            del connected_clients[client_ip]
-            print(
-                f"[{datetime.datetime.now():%Y-%m-%d %H:%M:%S}] 客户端 {client_ip} 已从连接列表移除、当前在线数：{len(connected_clients)}")
+            await connected_clients[client_ip].websocket.close(code=1008, reason="同一IP新连接")
+            connected_clients.pop(client_ip)
+            print(f"[{now:%H:%M:%S}] 客户端{client_ip}：关闭旧连接")
+
+        # 注册新连接
+        new_conn = ClientConnection(websocket, client_ip)
+        connected_clients[client_ip] = new_conn
+        print(f"[{now:%H:%M:%S}] 客户端{client_ip}：注册成功，当前在线{len(connected_clients)}个")
+
+        # 循环接收消息
+        while True:
+            data = await websocket.receive()
+            if "text" in data:
+                await handle_text_msg(client_ip, data["text"], new_conn)
+            elif "bytes" in data:
+                await handle_binary_msg(client_ip, data["bytes"])
+
+    # 异常处理（断开/错误）
+    except WebSocketDisconnect as e:
+        print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：主动断开（代码：{e.code}）")
+    except Exception as e:
+        print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：连接异常（{str(e)[:50]}）")
+    finally:
+        # 清理连接
+        connected_clients.pop(client_ip, None)
+        print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：连接已清理，当前在线{len(connected_clients)}个")