from fastapi import WebSocket, APIRouter, WebSocketDisconnect, FastAPI from typing import Dict, Any, Optional import datetime import asyncio import json from contextlib import asynccontextmanager # 创建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秒）""" 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 # 存储所有已连接的客户端（key：客户端IP、value：ClientConnection对象） 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)}") # 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: 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__}）") finally: if client_ip in connected_clients: del connected_clients[client_ip] print(f"[{current_time:%Y-%m-%d %H:%M:%S}] 客户端 {client_ip} 已从连接列表移除") else: print(f"[{current_time:%Y-%m-%d %H:%M:%S}] 无超时客户端、心跳检查完成") # 3. 等待30秒后进行下一轮检查 await asyncio.sleep(30) @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 # 应用运行中 if heartbeat_task and not heartbeat_task.done(): heartbeat_task.cancel() try: await heartbeat_task 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)}") async def send_heartbeat_ack(client_ip: str, client_timestamp: Any) -> bool: """向客户端回复心跳确认（严格遵循 {"timestamp":xxxxx, "type":"heartbeat"} 格式）""" 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, "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)}") 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} 因心跳回复失败被移除") 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 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) else: print(f"[{recv_time:%Y-%m-%d %H:%M:%S}] 非心跳消息（类型：{message.get('type')}）、不回复") 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: # 无论何种退出原因、确保客户端从列表中移除 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)}")