diff --git a/agent/agent.go b/agent/agent.go index c3e7c2e..1e66f4a 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -189,22 +189,6 @@ func (a *Agent) connectWS() error { return err } - // 设置读取超时(比服务端 ping 间隔长一些) - conn.SetReadDeadline(time.Now().Add(60 * time.Second)) - - // 设置 Ping 处理器,收到 Ping 时重置读取超时 - conn.SetPingHandler(func(appData string) error { - conn.SetReadDeadline(time.Now().Add(60 * time.Second)) - // 回复 Pong - return conn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(5*time.Second)) - }) - - // 设置 Pong 处理器,收到 Pong 时重置读取超时 - conn.SetPongHandler(func(appData string) error { - conn.SetReadDeadline(time.Now().Add(60 * time.Second)) - return nil - }) - a.wsMu.Lock() a.wsConn = conn a.wsStopCh = make(chan struct{}) @@ -244,12 +228,10 @@ func (a *Agent) readWS() { _, message, err := conn.ReadMessage() if err != nil { + log.Warnf("WebSocket 读取错误: %v", err) return } - // 收到消息,重置读取超时 - conn.SetReadDeadline(time.Now().Add(60 * time.Second)) - var msg WSMessage if err := json.Unmarshal(message, &msg); err != nil { continue diff --git a/internal/controllers/agent_controller.go b/internal/controllers/agent_controller.go index 9eba6cf..973644b 100644 --- a/internal/controllers/agent_controller.go +++ b/internal/controllers/agent_controller.go @@ -410,7 +410,7 @@ func (c *AgentController) wsReadPump(ac *services.AgentConnection, agent *models if r := recover(); r != nil { logger.Errorf("[AgentWS] Agent #%d wsReadPump panic: %v", agent.ID, r) } - c.wsManager.Unregister(agent.ID) + c.wsManager.Unregister(agent.ID, ac) }() // 检查连接是否有效(可能是旧连接被新连接替换) diff --git a/internal/services/agent_ws_service.go b/internal/services/agent_ws_service.go index 98dccf0..4bf3668 100644 --- a/internal/services/agent_ws_service.go +++ b/internal/services/agent_ws_service.go @@ -161,12 +161,13 @@ func (m *AgentWSManager) Register(agentID uint, conn *websocket.Conn, ip string) return ac } -// Unregister 注销连接 -func (m *AgentWSManager) Unregister(agentID uint) { +// Unregister 注销连接(只注销指定的连接实例) +func (m *AgentWSManager) Unregister(agentID uint, ac *AgentConnection) { m.mu.Lock() defer m.mu.Unlock() - if conn, exists := m.connections[agentID]; exists { + // 只有当前连接和 map 中的连接是同一个实例时才删除 + if conn, exists := m.connections[agentID]; exists && conn == ac { // 减少 IP 连接计数 if conn.IP != "" { if count, ok := m.ipConnections[conn.IP]; ok && count > 0 {