From 0fa920e972b3ff3c038c56f26c2b841fc0405f34 Mon Sep 17 00:00:00 2001 From: dushixiang Date: Mon, 31 Aug 2026 21:52:03 +0800 Subject: [PATCH] =?UTF-8?q?refactor:=20=E9=87=8D=E6=9E=84=20WebSocket=20?= =?UTF-8?q?=E5=8F=AF=E9=9D=A0=E6=8A=95=E9=80=92=E4=BC=9A=E8=AF=9D=E6=A8=A1?= =?UTF-8?q?=E5=9E=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/handler/agent_admin_handler.go | 3 +- internal/handler/agent_ws_handler.go | 105 ++-- internal/websocket/manager.go | 772 +++++++++++++++--------- internal/websocket/manager_test.go | 298 ++++++--- pkg/agent/service/agent.go | 96 ++- pkg/agent/service/outbox.go | 5 + pkg/agent/service/outbox_test.go | 98 +++ 7 files changed, 892 insertions(+), 485 deletions(-) diff --git a/internal/handler/agent_admin_handler.go b/internal/handler/agent_admin_handler.go index fd83e70..089aa8e 100644 --- a/internal/handler/agent_admin_handler.go +++ b/internal/handler/agent_admin_handler.go @@ -246,7 +246,7 @@ func (h *AgentHandler) Delete(c *echo.Context) error { } // 断开连接 - client.Conn.Close() + client.Close() } // 删除探针及其所有相关数据 @@ -257,6 +257,7 @@ func (h *AgentHandler) Delete(c *echo.Context) error { zap.Error(err)) return err } + h.wsManager.ForgetAgent(agentID) h.logger.Info("探针已删除", zap.String("agentID", agentID), diff --git a/internal/handler/agent_ws_handler.go b/internal/handler/agent_ws_handler.go index 7341049..a020ae9 100644 --- a/internal/handler/agent_ws_handler.go +++ b/internal/handler/agent_ws_handler.go @@ -3,6 +3,7 @@ package handler import ( "context" "encoding/json" + "errors" "net/http" "time" @@ -38,13 +39,24 @@ func (h *AgentHandler) HandleWebSocket(c *echo.Context) error { return err } - // 同步探针会话：BootID 变化（探针重启）时重置确认位点， - // 再把当前位点随注册响应返回 - h.wsManager.BeginAgentSession(agent.ID, registerReq.BootID) + // 先原子激活连接及其 BootID 会话，再发送注册响应。这样响应中的 + // AckSeq 与当前连接属于同一次会话切换，旧连接无法在其间污染位点。 + client := ws.NewClient(agent.ID, conn, h.wsManager) + ackSeq := h.wsManager.Register(client, registerReq.BootID) + // RegisterAgent 的在线写入发生在连接切换之前，可能被恰好退出的旧连接 + // 覆盖；会话激活后再确认一次，确保数据库状态属于 current client。 + if err := h.agentService.UpdateAgentStatus(context.Background(), agent.ID, 1); err != nil { + h.logger.Warn("failed to confirm agent online status", zap.String("agentID", agent.ID), zap.Error(err)) + } + defer func() { + if h.wsManager.Unregister(client) { + h.markAgentOffline(client.ID) + } + }() // 发送注册成功响应：声明可靠投递支持，并回传该探针的累计确认 // 位点，探针据此跳过已处理消息、只重放未确认部分 - if err := h.sendRegisterSuccess(conn, agent.ID); err != nil { + if err := h.sendRegisterSuccess(conn, agent.ID, ackSeq); err != nil { h.logger.Error("failed to send register ack", zap.Error(err)) conn.Close() return err @@ -68,17 +80,6 @@ func (h *AgentHandler) HandleWebSocket(c *echo.Context) error { } } - // 创建客户端并注册到管理器 - client := h.newClient(agent.ID, conn) - - h.wsManager.Register(client) - - // 连接结束后标记离线；若探针已由新连接替换则跳过，避免旧连接的 - // 退出把新连接打成离线 - defer func() { - h.markAgentOffline(client) - }() - // 启动读写协程 go client.WritePump() client.ReadPump(context.Background()) @@ -159,22 +160,8 @@ func (h *AgentHandler) readRegisterRequest(conn *websocket.Conn) (*protocol.Regi return ®isterReq, nil } -func (h *AgentHandler) markAgentOffline(client *ws.Client) { - // 探针已由新连接替换时跳过：离线状态由新连接的活跃心跳管理 - if cur, ok := h.wsManager.GetClient(client.ID); ok && cur != client { - return - } - _ = h.agentService.UpdateAgentStatus(context.Background(), client.ID, 0) -} - -func (h *AgentHandler) newClient(agentID string, conn *websocket.Conn) *ws.Client { - return &ws.Client{ - ID: agentID, - Conn: conn, - Send: make(chan []byte, 512), - Manager: h.wsManager, - LastActive: time.Now(), - } +func (h *AgentHandler) markAgentOffline(agentID string) { + _ = h.agentService.UpdateAgentStatus(context.Background(), agentID, 0) } // pongStatusRefreshInterval 在线状态心跳刷新间隔：本连接首次 pong 立即 @@ -189,7 +176,13 @@ func (h *AgentHandler) handleWebSocketPong(client *ws.Client) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - if err := h.agentService.UpdateAgentStatus(ctx, client.ID, 1); err != nil { + current, err := h.wsManager.DoIfCurrent(client, func() error { + return h.agentService.UpdateAgentStatus(ctx, client.ID, 1) + }) + if !current { + return + } + if err != nil { h.logger.Warn("failed to update agent status on pong", zap.String("agentID", client.ID), zap.Error(err)) return } @@ -199,36 +192,60 @@ func (h *AgentHandler) handleWebSocketPong(client *ws.Client) { func (h *AgentHandler) handleMetricsMessage(ctx context.Context, agentID string, data json.RawMessage) error { var batch protocol.MetricsBatch if err := json.Unmarshal(data, &batch); err != nil { - return err + return ws.Permanent(err) } + var transientErr error + var permanentErr error for _, sample := range batch.Samples { metricsData, err := json.Marshal(sample.Data) if err != nil { h.logger.Warn("failed to marshal metric sample", zap.Error(err)) + permanentErr = errors.Join(permanentErr, err) continue } if err := h.metricService.HandleMetricData(ctx, agentID, string(sample.Type), metricsData, sample.Timestamp); err != nil { h.logger.Warn("failed to handle metric sample", zap.Error(err), zap.String("type", string(sample.Type))) + if isPayloadError(err) { + permanentErr = errors.Join(permanentErr, err) + } else { + transientErr = errors.Join(transientErr, err) + } continue } } - return nil + if transientErr != nil { + return transientErr + } + return ws.Permanent(permanentErr) } func (h *AgentHandler) handleCommandResponseMessage(ctx context.Context, agentID string, data json.RawMessage) error { var cmdResp protocol.CommandResponse if err := json.Unmarshal(data, &cmdResp); err != nil { - return err + return ws.Permanent(err) } - return h.agentService.HandleCommandResponse(ctx, agentID, &cmdResp) + return classifyMessageError(h.agentService.HandleCommandResponse(ctx, agentID, &cmdResp)) +} + +func isPayloadError(err error) bool { + var syntaxErr *json.SyntaxError + var typeErr *json.UnmarshalTypeError + return errors.As(err, &syntaxErr) || errors.As(err, &typeErr) +} + +func classifyMessageError(err error) error { + if isPayloadError(err) { + return ws.Permanent(err) + } + return err } func (h *AgentHandler) handleTamperEventMessage(ctx context.Context, agentID string, data json.RawMessage) error { var eventData protocol.TamperEventData if err := json.Unmarshal(data, &eventData); err != nil { h.logger.Error("failed to unmarshal tamper event", zap.Error(err)) - return err + return ws.Permanent(err) } return h.tamperService.CreateEvent(ctx, agentID, &eventData) } @@ -237,7 +254,7 @@ func (h *AgentHandler) handleDDNSIPReportMessage(ctx context.Context, agentID st var ipReport protocol.DDNSIPReportData if err := json.Unmarshal(data, &ipReport); err != nil { h.logger.Error("failed to unmarshal ddns ip report", zap.Error(err)) - return err + return ws.Permanent(err) } return h.ddnsService.HandleIPReport(ctx, agentID, &ipReport) } @@ -246,7 +263,7 @@ func (h *AgentHandler) handlePublicIPReportMessage(ctx context.Context, agentID var ipReport protocol.PublicIPReportData if err := json.Unmarshal(data, &ipReport); err != nil { h.logger.Error("failed to unmarshal public ip report", zap.Error(err)) - return err + return ws.Permanent(err) } return h.agentService.UpdatePublicIP(ctx, agentID, ipReport.IPv4, ipReport.IPv6) } @@ -255,7 +272,7 @@ func (h *AgentHandler) handleSSHLoginEventMessage(ctx context.Context, agentID s var eventData protocol.SSHLoginEvent if err := json.Unmarshal(data, &eventData); err != nil { h.logger.Error("failed to unmarshal ssh login event", zap.Error(err)) - return err + return ws.Permanent(err) } return h.sshLoginService.HandleEvent(ctx, agentID, eventData) } @@ -264,7 +281,7 @@ func (h *AgentHandler) handleSSHLoginConfigResultMessage(ctx context.Context, ag var resultData protocol.SSHLoginConfigResult if err := json.Unmarshal(data, &resultData); err != nil { h.logger.Error("failed to unmarshal ssh login config result", zap.Error(err)) - return err + return ws.Permanent(err) } return h.sshLoginService.HandleConfigResult(ctx, agentID, resultData) } @@ -273,18 +290,18 @@ func (h *AgentHandler) handleTamperProtectMessage(ctx context.Context, agentID s var protectResp protocol.TamperProtectResponse if err := json.Unmarshal(data, &protectResp); err != nil { h.logger.Error("failed to unmarshal tamper protect response", zap.Error(err)) - return err + return ws.Permanent(err) } return h.tamperService.HandleConfigResult(ctx, agentID, protectResp) } // sendRegisterSuccess 发送注册成功响应 -func (h *AgentHandler) sendRegisterSuccess(conn *websocket.Conn, agentID string) error { +func (h *AgentHandler) sendRegisterSuccess(conn *websocket.Conn, agentID string, ackSeq uint64) error { resp := protocol.RegisterResponse{ AgentID: agentID, Status: "success", Reliable: true, - AckSeq: h.wsManager.AckSeq(agentID), + AckSeq: ackSeq, } return conn.WriteJSON(protocol.OutboundMessage{ Type: protocol.MessageTypeRegisterAck, diff --git a/internal/websocket/manager.go b/internal/websocket/manager.go index 50fab86..983b04b 100644 --- a/internal/websocket/manager.go +++ b/internal/websocket/manager.go @@ -3,6 +3,7 @@ package websocket import ( "context" "encoding/json" + "errors" "sync" "sync/atomic" "time" @@ -17,103 +18,202 @@ const ( serverPongWait = 30 * time.Second serverWriteWait = 10 * time.Second - // agentQueueSize 每个探针的待处理消息队列长度。队列满时断开连接， - // 由探针重连后重放未确认消息，读循环不会被慢处理（DB/VM 写入）拖死。 - agentQueueSize = 512 - // agentQueueWait 消息入队的最长等待，超时判定为处理能力不足 - agentQueueWait = 2 * time.Second - // serverSendWait 下行消息写入 Send 通道的最长等待 + // 可靠事件必须保持会话内顺序；队列满表示服务端已无法及时消费， + // 断开当前传输后由 agent 保留 outbox 并重连重放。 + reliableQueueSize = 512 + reliableQueueWait = 2 * time.Second + + // 普通遥测是 latest-wins，不应阻塞可靠事件。队列满时丢弃本批， + // agent 下一轮会继续上报更新后的快照。 + bestEffortQueueSize = 128 + serverSendWait = 3 * time.Second + + messageRetryMin = 500 * time.Millisecond + messageRetryMax = 10 * time.Second ) -// agentMessage 从探针收到的一条待处理消息 +// MessageHandler 处理一条 agent 上行消息。 +type MessageHandler func(ctx context.Context, agentID string, messageType string, data json.RawMessage) error + +// PongHandler 处理连接 pong。 +type PongHandler func(client *Client) + +// PermanentMessageError 表示重放也无法成功的坏消息，例如 JSON 结构错误。 +// 可靠处理器会记录并确认这类消息，避免单条毒消息永久阻塞后续事件。 +type PermanentMessageError struct { + err error +} + +func (e *PermanentMessageError) Error() string { return e.err.Error() } +func (e *PermanentMessageError) Unwrap() error { return e.err } + +// Permanent 把不可重试的消息错误标记为永久错误。 +func Permanent(err error) error { + if err == nil { + return nil + } + return &PermanentMessageError{err: err} +} + +func isPermanent(err error) bool { + var target *PermanentMessageError + return errors.As(err, &target) +} + type agentMessage struct { seq uint64 typ string data json.RawMessage } -// processor 单个探针连接的有序处理单元：ReadPump 只负责解析并入队， -// processLoop 串行消费，慢处理不再阻塞 TCP 读取造成反压。 -type processor struct { - queue chan agentMessage - quit chan struct{} - quitOnce sync.Once +// agentSession 是可靠投递状态的唯一所有者。它绑定 agentID + BootID， +// 与具体 WebSocket 连接解耦：同一进程重连复用 session，新进程注册时 +// 原子替换 session，旧 worker 即使仍在返回路径上也无法提交 ACK。 +type agentSession struct { + agentID string + bootID string + + ctx context.Context + cancel context.CancelFunc + + reliableQueue chan agentMessage + bestEffortQueue chan agentMessage + ackSeq atomic.Uint64 } -func newProcessor() *processor { - return &processor{ - queue: make(chan agentMessage, agentQueueSize), - quit: make(chan struct{}), +func newAgentSession(agentID, bootID string) *agentSession { + ctx, cancel := context.WithCancel(context.Background()) + return &agentSession{ + agentID: agentID, + bootID: bootID, + ctx: ctx, + cancel: cancel, + reliableQueue: make(chan agentMessage, reliableQueueSize), + bestEffortQueue: make(chan agentMessage, bestEffortQueueSize), } } -func (p *processor) stop() { - p.quitOnce.Do(func() { close(p.quit) }) -} - -// Client WebSocket客户端 +// Client 只代表一次 WebSocket 传输，不拥有可靠会话状态。 type Client struct { - ID string // 探针ID - Conn *websocket.Conn // WebSocket连接 - Send chan []byte // 发送消息通道 - Manager *Manager // 管理器引用 - LastActive time.Time // 最后活跃时间 - proc *processor // 上行消息的有序处理器 - // lastStatusWrite 上次把在线状态写入数据库的时间（毫秒），用于 - // 对 pong 触发的状态更新去抖 + ID string + conn *websocket.Conn + manager *Manager + session *agentSession + + send chan []byte + done chan struct{} + + closeOnce sync.Once + lastActive atomic.Int64 lastStatusWrite atomic.Int64 - closed bool // 标记channel是否已关闭 - closeMu sync.Mutex } -// Manager WebSocket连接管理器 +// NewClient 创建尚未注册到 Manager 的连接。 +func NewClient(agentID string, conn *websocket.Conn, manager *Manager) *Client { + c := &Client{ + ID: agentID, + conn: conn, + manager: manager, + send: make(chan []byte, 512), + done: make(chan struct{}), + } + c.touch() + return c +} + +func (c *Client) touch() { + c.lastActive.Store(time.Now().UnixMilli()) +} + +func (c *Client) inactiveFor(now time.Time) time.Duration { + last := c.lastActive.Load() + if last == 0 { + return 0 + } + return now.Sub(time.UnixMilli(last)) +} + +func (c *Client) close() { + c.closeOnce.Do(func() { + close(c.done) + if c.conn != nil { + _ = c.conn.Close() + } + }) +} + +// Close 主动关闭本次传输；Handler 的退出路径负责同步注销 current client。 +func (c *Client) Close() { + c.close() +} + +func (c *Client) sendWithTimeout(message []byte, timeout time.Duration) error { + select { + case <-c.done: + return ErrClientNotFound + default: + } + + timer := time.NewTimer(timeout) + defer timer.Stop() + + select { + case <-c.done: + return ErrClientNotFound + case c.send <- message: + return nil + case <-timer.C: + return ErrSendTimeout + } +} + +func (c *Client) trySend(message []byte) bool { + select { + case <-c.done: + return false + default: + } + + select { + case <-c.done: + return false + case c.send <- message: + return true + default: + return false + } +} + +// Manager 同步维护当前连接和每个 agent 的会话状态。 type Manager struct { - clients map[string]*Client // 客户端映射 probeID -> Client - register chan *Client // 注册通道 - unregister chan *Client // 注销通道 - broadcast chan []byte // 广播通道 - mu sync.RWMutex // 读写锁 - logger *zap.Logger // 日志 - onMessage MessageHandler // 消息处理器 - onPong PongHandler // Pong 处理器 - // ackSeqs 每个探针事件流的累计确认位点（探针重启或服务端重启后 - // 丢失，退化为重放+去重，可能产生少量重复处理） - ackSeqs map[string]uint64 - // sessions 每个探针最近一次注册携带的 BootID，用于识别探针进程重启 - sessions map[string]string + mu sync.RWMutex + clients map[string]*Client + sessions map[string]*agentSession + + logger *zap.Logger + onMessage MessageHandler + onPong PongHandler } -// MessageHandler 消息处理器接口 -type MessageHandler func(ctx context.Context, probeID string, messageType string, data json.RawMessage) error - -// PongHandler Pong 处理器接口 -type PongHandler func(client *Client) - -// NewManager 创建新的WebSocket管理器 func NewManager(logger *zap.Logger) *Manager { return &Manager{ - clients: make(map[string]*Client), - register: make(chan *Client, 10), - unregister: make(chan *Client, 10), - broadcast: make(chan []byte, 256), - ackSeqs: make(map[string]uint64), - sessions: make(map[string]string), - logger: logger, + clients: make(map[string]*Client), + sessions: make(map[string]*agentSession), + logger: logger, } } -// SetMessageHandler 设置消息处理器 func (m *Manager) SetMessageHandler(handler MessageHandler) { m.onMessage = handler } -// SetPongHandler 设置 Pong 处理器 func (m *Manager) SetPongHandler(handler PongHandler) { m.onPong = handler } -// Run 启动管理器 +// Run 只负责周期性连接保活检查和进程退出清理；注册/注销本身是同步的， +// 注册响应拿到的 ACK 与刚完成的会话切换因此属于同一个原子操作。 func (m *Manager) Run(ctx context.Context) { ticker := time.NewTicker(30 * time.Second) defer ticker.Stop() @@ -121,135 +221,189 @@ func (m *Manager) Run(ctx context.Context) { for { select { case <-ctx.Done(): + m.shutdown() m.logger.Info("websocket manager stopped") return - case client := <-m.register: - m.registerClient(client) - case client := <-m.unregister: - m.unregisterClient(client) - case message := <-m.broadcast: - m.broadcastMessage(message) case <-ticker.C: m.checkInactiveClients() } } } -// Register 注册客户端（公开方法）。同步初始化处理器，确保随后的 -// ReadPump 可以立即入队；注册表维护仍走 Run 循环的注册通道。 -func (m *Manager) Register(client *Client) { - client.proc = newProcessor() - go m.processLoop(client) - m.register <- client -} +// Register 原子激活连接并返回该 BootID 会话已确认的累计位点。 +// 同 BootID 重连复用处理队列；BootID 变化创建全新 session 并取消旧 session。 +func (m *Manager) Register(client *Client, bootID string) uint64 { + var oldClient *Client + var oldSession *agentSession + var startSession *agentSession -// registerClient 注册客户端 -func (m *Manager) registerClient(client *Client) { m.mu.Lock() - defer m.mu.Unlock() - - // 如果已存在该探针的连接，先关闭旧连接 - if oldClient, exists := m.clients[client.ID]; exists && oldClient != client { - m.logger.Info("agent reconnected, closing old connection", zap.String("agentID", client.ID)) - oldClient.closeChannel() - oldClient.Conn.Close() - oldClient.stopProcessor() + oldClient = m.clients[client.ID] + session := m.sessions[client.ID] + if session == nil || (bootID != "" && session.bootID != bootID) { + oldSession = session + session = newAgentSession(client.ID, bootID) + m.sessions[client.ID] = session + startSession = session } - + client.session = session m.clients[client.ID] = client - m.logger.Info("agent connected", zap.String("agentID", client.ID), zap.Int("totalClients", len(m.clients))) -} + ackSeq := session.ackSeq.Load() + clientCount := len(m.clients) + m.mu.Unlock() -// unregisterClient 注销客户端 -func (m *Manager) unregisterClient(client *Client) { - m.mu.Lock() - defer m.mu.Unlock() - - // 只注销当前注册的连接：探针快速重连时旧连接的退出路径不能把 - // 新连接从注册表中移除，否则新连接的下行通道会被关闭 - if cur, exists := m.clients[client.ID]; exists && cur == client { - delete(m.clients, client.ID) - client.closeChannel() - m.logger.Info("agent disconnected", zap.String("agentID", client.ID), zap.Int("totalClients", len(m.clients))) + // map/session 指针已切换后再关闭旧对象。旧 worker 的提交路径会先 + // 校验 session 指针，所以无法污染新会话。 + if oldSession != nil { + oldSession.cancel() } - client.stopProcessor() + if oldClient != nil && oldClient != client { + oldClient.close() + } + if startSession != nil { + go m.reliableProcessLoop(startSession) + go m.bestEffortProcessLoop(startSession) + } + + m.logger.Info("agent connected", + zap.String("agentID", client.ID), + zap.String("bootID", bootID), + zap.Uint64("ackSeq", ackSeq), + zap.Int("totalClients", clientCount)) + return ackSeq } -// broadcastMessage 广播消息 -func (m *Manager) broadcastMessage(message []byte) { +// Unregister 只移除仍为 current 的连接，返回值用于决定是否标记离线。 +func (m *Manager) Unregister(client *Client) bool { + m.mu.Lock() + current := m.clients[client.ID] == client + if current { + delete(m.clients, client.ID) + } + clientCount := len(m.clients) + m.mu.Unlock() + + client.close() + if current { + m.logger.Info("agent disconnected", + zap.String("agentID", client.ID), + zap.Int("totalClients", clientCount)) + } + return current +} + +// ForgetAgent 删除已移除 agent 的连接与会话状态，终止常驻 worker。 +func (m *Manager) ForgetAgent(agentID string) { + m.mu.Lock() + client := m.clients[agentID] + session := m.sessions[agentID] + delete(m.clients, agentID) + delete(m.sessions, agentID) + m.mu.Unlock() + + if client != nil { + client.close() + } + if session != nil { + session.cancel() + } +} + +func (m *Manager) shutdown() { + m.mu.Lock() + clients := make([]*Client, 0, len(m.clients)) + for _, client := range m.clients { + clients = append(clients, client) + } + sessions := make([]*agentSession, 0, len(m.sessions)) + for _, session := range m.sessions { + sessions = append(sessions, session) + } + m.clients = make(map[string]*Client) + m.sessions = make(map[string]*agentSession) + m.mu.Unlock() + + for _, client := range clients { + client.close() + } + for _, session := range sessions { + session.cancel() + } +} + +func (m *Manager) isCurrentSession(session *agentSession) bool { m.mu.RLock() defer m.mu.RUnlock() - - for _, client := range m.clients { - select { - case client.Send <- message: - default: - // 发送失败，客户端可能已断开 - m.logger.Warn("failed to send message, client may be disconnected", zap.String("agentID", client.ID)) - } - } + return m.sessions[session.agentID] == session +} + +func (m *Manager) currentClientForSession(session *agentSession) *Client { + m.mu.RLock() + defer m.mu.RUnlock() + if m.sessions[session.agentID] != session { + return nil + } + client := m.clients[session.agentID] + if client == nil || client.session != session { + return nil + } + return client } -// checkInactiveClients 检查不活跃的客户端 func (m *Manager) checkInactiveClients() { + now := time.Now() m.mu.RLock() - inactiveClients := make([]*Client, 0) - timeout := 2 * time.Minute - + inactive := make([]*Client, 0) for _, client := range m.clients { - if time.Since(client.LastActive) > timeout { - inactiveClients = append(inactiveClients, client) + if client.inactiveFor(now) > 2*time.Minute { + inactive = append(inactive, client) } } m.mu.RUnlock() - // 断开不活跃的客户端 - for _, client := range inactiveClients { - // 再次检查客户端是否仍然存在（避免竞态条件） + for _, client := range inactive { m.mu.RLock() - cur, exists := m.clients[client.ID] + current := m.clients[client.ID] == client m.mu.RUnlock() - - if exists && cur == client { + if current { m.logger.Warn("agent inactive timeout, disconnecting", zap.String("agentID", client.ID)) - client.Conn.Close() - m.unregister <- client + client.close() } } } -// SendToClient 发送消息给指定客户端。通道满时最多等待 serverSendWait， -// 避免高负载时指令/配置被静默丢弃 -func (m *Manager) SendToClient(probeID string, message []byte) error { +func (m *Manager) SendToClient(agentID string, message []byte) error { m.mu.RLock() - client, exists := m.clients[probeID] + client := m.clients[agentID] m.mu.RUnlock() - - if !exists { + if client == nil { return ErrClientNotFound } - - select { - case client.Send <- message: - return nil - case <-time.After(serverSendWait): - return ErrSendTimeout - } + return client.sendWithTimeout(message, serverSendWait) } -// GetClient 获取客户端 -func (m *Manager) GetClient(probeID string) (*Client, bool) { +func (m *Manager) GetClient(agentID string) (*Client, bool) { m.mu.RLock() defer m.mu.RUnlock() - client, exists := m.clients[probeID] - return client, exists + client, ok := m.clients[agentID] + return client, ok +} + +// DoIfCurrent 在连接仍为 current 的前提下执行状态变更，并用注册锁保证 +// 回调完成前不会切换连接。用于 pong 在线写库，避免旧连接的迟到写入覆盖 +// 新连接失败后的离线状态。 +func (m *Manager) DoIfCurrent(client *Client, fn func() error) (bool, error) { + m.mu.RLock() + defer m.mu.RUnlock() + if m.clients[client.ID] != client { + return false, nil + } + return true, fn() } -// GetAllClients 获取所有客户端ID func (m *Manager) GetAllClients() []string { m.mu.RLock() defer m.mu.RUnlock() - ids := make([]string, 0, len(m.clients)) for id := range m.clients { ids = append(ids, id) @@ -257,158 +411,209 @@ func (m *Manager) GetAllClients() []string { return ids } -// ClientCount 获取客户端数量 func (m *Manager) ClientCount() int { m.mu.RLock() defer m.mu.RUnlock() return len(m.clients) } -// AckSeq 返回探针事件流的累计确认位点。注册响应携带该值，探针据此 -// 跳过服务端已处理的消息，避免无谓重放。 -func (m *Manager) AckSeq(probeID string) uint64 { +func (m *Manager) AckSeq(agentID string) uint64 { m.mu.RLock() - defer m.mu.RUnlock() - return m.ackSeqs[probeID] + session := m.sessions[agentID] + m.mu.RUnlock() + if session == nil { + return 0 + } + return session.ackSeq.Load() } -// BeginAgentSession 在注册时同步探针会话：BootID 变化说明探针进程 -// 重启过（纯内存事件队列清空、seq 从 1 重新分配），必须重置确认位点， -// 否则重启后新分配的小序号会被旧位点误判为重复消息而丢弃。 -// 旧版探针不携带 BootID，保持原位点不变。 -func (m *Manager) BeginAgentSession(probeID, bootID string) { - if bootID == "" { +// enqueueBestEffort/ enqueueReliable 在持有 current-client 读锁时完成入队。 +// Register 必须取得写锁，因此同 BootID 重连时，旧连接已经读出的消息必定 +// 先入队，然后才会切换连接；切换完成后旧连接不能再把较大 seq 插到新连接 +// 重放的较小 seq 前面，累计 ACK 顺序不会被跨连接竞态破坏。 +func (m *Manager) enqueueBestEffort(client *Client, msg agentMessage) bool { + m.mu.RLock() + defer m.mu.RUnlock() + if m.clients[client.ID] != client || m.sessions[client.ID] != client.session { + return false + } + select { + case client.session.bestEffortQueue <- msg: + return true + default: + return false + } +} + +func (m *Manager) enqueueReliable(ctx context.Context, client *Client, msg agentMessage) error { + m.mu.RLock() + defer m.mu.RUnlock() + if m.clients[client.ID] != client || m.sessions[client.ID] != client.session { + return ErrClientNotFound + } + + timer := time.NewTimer(reliableQueueWait) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-client.done: + return ErrClientNotFound + case <-client.session.ctx.Done(): + return ErrClientNotFound + case client.session.reliableQueue <- msg: + return nil + case <-timer.C: + return ErrQueueFull + } +} + +// ReadPump 只解析并按可靠性分流。可靠队列提供反压；best-effort 遥测 +// 队列满时直接丢弃，避免 VM 慢写阻塞安全事件。 +func (c *Client) ReadPump(ctx context.Context) { + defer c.close() + if c.conn == nil { return } - m.mu.Lock() - defer m.mu.Unlock() - - if m.sessions[probeID] != bootID { - m.sessions[probeID] = bootID - m.ackSeqs[probeID] = 0 - } -} - -// advanceAckSeq 推进探针的累计确认位点（只前进不回退） -func (m *Manager) advanceAckSeq(probeID string, seq uint64) { - m.mu.Lock() - defer m.mu.Unlock() - if m.ackSeqs[probeID] < seq { - m.ackSeqs[probeID] = seq - } -} - -// ReadPump 读取客户端消息。只做解析和入队，处理在 processLoop 中进行， -// 保证慢处理不会阻塞读取造成 TCP 反压。 -func (c *Client) ReadPump(ctx context.Context) { - defer func() { - c.Manager.unregister <- c - c.Conn.Close() - c.stopProcessor() - }() - - c.Conn.SetReadDeadline(time.Now().Add(serverPongWait)) - c.Conn.SetPongHandler(func(string) error { - c.Conn.SetReadDeadline(time.Now().Add(serverPongWait)) - c.LastActive = time.Now() - if c.Manager.onPong != nil { - go c.Manager.onPong(c) + _ = c.conn.SetReadDeadline(time.Now().Add(serverPongWait)) + c.conn.SetPongHandler(func(string) error { + if err := c.conn.SetReadDeadline(time.Now().Add(serverPongWait)); err != nil { + return err + } + c.touch() + if c.manager.onPong != nil { + // 同步执行，确保连接退出后的离线写库一定发生在本次 pong + // 状态刷新之后，避免遗留 goroutine 把已断开的连接重新标在线。 + c.manager.onPong(c) } return nil }) for { - _, message, err := c.Conn.ReadMessage() + _, raw, err := c.conn.ReadMessage() if err != nil { if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) { - c.Manager.logger.Error("websocket read error", zap.Error(err), zap.String("agentID", c.ID)) + c.manager.logger.Error("websocket read error", zap.Error(err), zap.String("agentID", c.ID)) } return } + c.touch() - c.LastActive = time.Now() - - // 解析消息 var msg protocol.InputMessage - if err := json.Unmarshal(message, &msg); err != nil { - c.Manager.logger.Error("failed to parse message", zap.Error(err), zap.String("agentID", c.ID)) + if err := json.Unmarshal(raw, &msg); err != nil { + c.manager.logger.Error("failed to parse message", zap.Error(err), zap.String("agentID", c.ID)) continue } - if err := c.enqueue(agentMessage{seq: msg.Seq, typ: string(msg.Type), data: msg.Data}); err != nil { - c.Manager.logger.Warn("agent processing queue stalled, disconnecting for replay", - zap.String("agentID", c.ID), zap.Error(err)) + item := agentMessage{seq: msg.Seq, typ: string(msg.Type), data: msg.Data} + if msg.Seq == 0 { + if !c.manager.enqueueBestEffort(c, item) { + c.manager.logger.Warn("best-effort agent queue full, dropping telemetry", + zap.String("agentID", c.ID), zap.String("type", string(msg.Type))) + } + continue + } + + if err := c.manager.enqueueReliable(ctx, c, item); err != nil { + c.manager.logger.Warn("reliable agent queue stalled, disconnecting for replay", + zap.String("agentID", c.ID), zap.Uint64("seq", msg.Seq), zap.Error(err)) return } } } -// enqueue 将消息放入该连接的有序处理队列。队列满时最多等待 -// agentQueueWait，仍无法入队则返回错误，由调用方断开连接触发探针重放。 -func (c *Client) enqueue(msg agentMessage) error { - // 先检查处理器是否已停止：select 各分支同时就绪时是随机选择， - // 不前置检查会导致停止后的消息仍被入队却永远无人消费 - select { - case <-c.proc.quit: - return ErrProcessorStopped - default: - } - - select { - case c.proc.queue <- msg: - return nil - case <-time.After(agentQueueWait): - return ErrQueueFull - case <-c.proc.quit: - return ErrProcessorStopped - } -} - -// stopProcessor 停止该连接的处理器 -func (c *Client) stopProcessor() { - if c.proc != nil { - c.proc.stop() - } -} - -// processLoop 串行处理单个探针的消息，保证 per-agent 顺序。 -// 带 seq 的消息处理成功后推进确认位点并回累计 ack；重放消息按 -// 确认位点去重。处理失败按丢弃处理并记录（与旧行为一致），避免 -// 毒消息卡死整个事件流。 -func (m *Manager) processLoop(c *Client) { +func (m *Manager) bestEffortProcessLoop(session *agentSession) { for { select { - case msg := <-c.proc.queue: - if msg.seq > 0 && msg.seq <= m.AckSeq(c.ID) { - // 重放去重：该消息已处理过。仍回 ack（当前确认位点 - // 不低于该 seq），让探针能修剪队列 - m.sendAck(c, m.AckSeq(c.ID)) + case <-session.ctx.Done(): + return + case msg := <-session.bestEffortQueue: + if session.ctx.Err() != nil || !m.isCurrentSession(session) { + return + } + if m.onMessage == nil { + continue + } + if err := m.onMessage(session.ctx, session.agentID, msg.typ, msg.data); err != nil { + m.logger.Error("failed to handle best-effort message, dropped", + zap.Error(err), zap.String("agentID", session.agentID), zap.String("type", msg.typ)) + } + } + } +} + +func (m *Manager) reliableProcessLoop(session *agentSession) { + for { + select { + case <-session.ctx.Done(): + return + case msg := <-session.reliableQueue: + if session.ctx.Err() != nil || !m.isCurrentSession(session) { + return + } + if msg.seq <= session.ackSeq.Load() { + m.sendAck(session, session.ackSeq.Load()) continue } - if m.onMessage != nil { - if err := m.onMessage(context.Background(), c.ID, msg.typ, msg.data); err != nil { - m.logger.Error("failed to handle message, dropped", - zap.Error(err), zap.String("agentID", c.ID), - zap.String("type", msg.typ), zap.Uint64("seq", msg.seq)) - } + if !m.processReliableMessage(session, msg) { + return } - - if msg.seq > 0 { - m.advanceAckSeq(c.ID, msg.seq) - m.sendAck(c, msg.seq) + if !m.isCurrentSession(session) { + return } - - case <-c.proc.quit: - return + session.ackSeq.Store(msg.seq) + m.sendAck(session, msg.seq) } } } -// sendAck 向探针发送累计确认。走下行 Send 通道，满时丢弃：丢失的 -// ack 只会让探针多保留一会儿消息，重连重放时按 seq 去重，无副作用。 -func (m *Manager) sendAck(c *Client, seq uint64) { +// processReliableMessage 对瞬时错误原位退避重试，保证后续序号不能越过 +// 失败消息；永久格式错误明确丢弃并允许推进位点。 +func (m *Manager) processReliableMessage(session *agentSession, msg agentMessage) bool { + retryDelay := messageRetryMin + for { + if session.ctx.Err() != nil || !m.isCurrentSession(session) { + return false + } + if m.onMessage == nil { + return true + } + + err := m.onMessage(session.ctx, session.agentID, msg.typ, msg.data) + if err == nil { + return true + } + if isPermanent(err) { + m.logger.Error("permanent reliable message error, dropped", + zap.Error(err), zap.String("agentID", session.agentID), + zap.String("type", msg.typ), zap.Uint64("seq", msg.seq)) + return true + } + + m.logger.Warn("reliable message handling failed, retrying", + zap.Error(err), zap.String("agentID", session.agentID), + zap.String("type", msg.typ), zap.Uint64("seq", msg.seq), + zap.Duration("retryAfter", retryDelay)) + + timer := time.NewTimer(retryDelay) + select { + case <-session.ctx.Done(): + timer.Stop() + return false + case <-timer.C: + } + if retryDelay < messageRetryMax { + retryDelay *= 2 + if retryDelay > messageRetryMax { + retryDelay = messageRetryMax + } + } + } +} + +func (m *Manager) sendAck(session *agentSession, seq uint64) { payload, err := json.Marshal(protocol.OutboundMessage{ Type: protocol.MessageTypeAck, Data: protocol.AckData{Seq: seq}, @@ -416,72 +621,51 @@ func (m *Manager) sendAck(c *Client, seq uint64) { if err != nil { return } - select { - case c.Send <- payload: - default: + if client := m.currentClientForSession(session); client != nil { + client.trySend(payload) } } -// WritePump 向客户端写入消息 func (c *Client) WritePump() { + if c.conn == nil { + return + } ticker := time.NewTicker(serverPingInterval) - defer func() { - ticker.Stop() - c.Conn.Close() - }() + defer ticker.Stop() for { select { - case message, ok := <-c.Send: - c.Conn.SetWriteDeadline(time.Now().Add(serverWriteWait)) - if !ok { - // 通道已关闭 - c.Conn.WriteMessage(websocket.CloseMessage, []byte{}) + case <-c.done: + return + case message := <-c.send: + _ = c.conn.SetWriteDeadline(time.Now().Add(serverWriteWait)) + if err := c.conn.WriteMessage(websocket.TextMessage, message); err != nil { + c.manager.logger.Error("failed to write message", zap.Error(err), zap.String("agentID", c.ID)) + c.close() return } - - if err := c.Conn.WriteMessage(websocket.TextMessage, message); err != nil { - c.Manager.logger.Error("failed to write message", zap.Error(err), zap.String("agentID", c.ID)) - return - } - case <-ticker.C: - c.Conn.SetWriteDeadline(time.Now().Add(serverWriteWait)) - if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil { + _ = c.conn.SetWriteDeadline(time.Now().Add(serverWriteWait)) + if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil { + c.close() return } } } } -// ShouldWriteStatus 报告是否需要把在线状态写入数据库。去抖：本连接 -// 首次调用立即通过，之后每 refreshInterval 通过一次。 func (c *Client) ShouldWriteStatus(refreshInterval time.Duration) bool { now := time.Now().UnixMilli() last := c.lastStatusWrite.Load() return last == 0 || now-last >= refreshInterval.Milliseconds() } -// MarkStatusWritten 记录本次在线状态已成功写库 func (c *Client) MarkStatusWritten() { c.lastStatusWrite.Store(time.Now().UnixMilli()) } -// closeChannel 安全地关闭发送通道 -func (c *Client) closeChannel() { - c.closeMu.Lock() - defer c.closeMu.Unlock() - - if !c.closed { - close(c.Send) - c.closed = true - } -} - -// 错误定义 var ( - ErrClientNotFound = &websocket.CloseError{Code: 1000, Text: "client not found"} - ErrSendTimeout = &websocket.CloseError{Code: 1001, Text: "send timeout"} - ErrQueueFull = &websocket.CloseError{Code: 1013, Text: "processing queue full"} - ErrProcessorStopped = &websocket.CloseError{Code: 1000, Text: "processor stopped"} + ErrClientNotFound = &websocket.CloseError{Code: 1000, Text: "client not found"} + ErrSendTimeout = &websocket.CloseError{Code: 1001, Text: "send timeout"} + ErrQueueFull = &websocket.CloseError{Code: 1013, Text: "reliable queue full"} ) diff --git a/internal/websocket/manager_test.go b/internal/websocket/manager_test.go index 33a8cdf..69d6e66 100644 --- a/internal/websocket/manager_test.go +++ b/internal/websocket/manager_test.go @@ -3,6 +3,9 @@ package websocket import ( "context" "encoding/json" + "errors" + "sync" + "sync/atomic" "testing" "time" @@ -10,20 +13,16 @@ import ( "go.uber.org/zap" ) -func newTestClient(m *Manager, id string) *Client { - c := &Client{ - ID: id, - Send: make(chan []byte, 16), - Manager: m, - } - c.proc = newProcessor() +func newTestClient(m *Manager, id, bootID string) *Client { + c := NewClient(id, nil, m) + m.Register(c, bootID) return c } func readAck(t *testing.T, c *Client) uint64 { t.Helper() select { - case raw := <-c.Send: + case raw := <-c.send: var msg protocol.InputMessage if err := json.Unmarshal(raw, &msg); err != nil { t.Fatalf("unmarshal ack: %v", err) @@ -42,29 +41,28 @@ func readAck(t *testing.T, c *Client) uint64 { } } -func TestProcessLoopDedupAckAndOrder(t *testing.T) { +func TestReliableSessionDedupAckAndOrder(t *testing.T) { m := NewManager(zap.NewNop()) + defer m.shutdown() + + var mu sync.Mutex var processed []string done := make(chan struct{}) - m.SetMessageHandler(func(ctx context.Context, probeID, typ string, data json.RawMessage) error { + m.SetMessageHandler(func(_ context.Context, _, typ string, _ json.RawMessage) error { + mu.Lock() processed = append(processed, typ) if len(processed) == 2 { close(done) } + mu.Unlock() return nil }) - client := newTestClient(m, "agent-1") - go m.processLoop(client) - defer client.stopProcessor() - - // 模拟断线前的确认位点：seq 1 已处理 - m.advanceAckSeq("agent-1", 1) - - // seq 1 是重放（应去重跳过），seq 2、3 是新消息（按序处理） - client.proc.queue <- agentMessage{seq: 1, typ: "replayed"} - client.proc.queue <- agentMessage{seq: 2, typ: "first"} - client.proc.queue <- agentMessage{seq: 3, typ: "second"} + client := newTestClient(m, "agent-1", "boot-1") + client.session.ackSeq.Store(1) + client.session.reliableQueue <- agentMessage{seq: 1, typ: "replayed"} + client.session.reliableQueue <- agentMessage{seq: 2, typ: "first"} + client.session.reliableQueue <- agentMessage{seq: 3, typ: "second"} select { case <-done: @@ -72,110 +70,230 @@ func TestProcessLoopDedupAckAndOrder(t *testing.T) { t.Fatal("timed out waiting for processing") } - if len(processed) != 2 || processed[0] != "first" || processed[1] != "second" { - t.Fatalf("processed = %v, want [first second]", processed) + mu.Lock() + got := append([]string(nil), processed...) + mu.Unlock() + if len(got) != 2 || got[0] != "first" || got[1] != "second" { + t.Fatalf("processed = %v, want [first second]", got) } - if got := m.AckSeq("agent-1"); got != 3 { - t.Fatalf("AckSeq = %d, want 3", got) + if ack := m.AckSeq("agent-1"); ack != 3 { + t.Fatalf("AckSeq = %d, want 3", ack) } - - // 去重跳过的消息也回 ack（位点 1），随后是新消息的 ack（2、3） - if seq := readAck(t, client); seq != 1 { - t.Fatalf("ack for replayed message = %d, want 1", seq) - } - if seq := readAck(t, client); seq != 2 { - t.Fatalf("ack for first message = %d, want 2", seq) - } - if seq := readAck(t, client); seq != 3 { - t.Fatalf("ack for second message = %d, want 3", seq) + for i, want := range []uint64{1, 2, 3} { + if ack := readAck(t, client); ack != want { + t.Fatalf("ack #%d = %d, want %d", i, ack, want) + } } } -func TestProcessLoopSeqZeroBypassesAck(t *testing.T) { +func TestReliableFailureDoesNotAdvanceOrSkip(t *testing.T) { m := NewManager(zap.NewNop()) - handled := make(chan string, 1) - m.SetMessageHandler(func(ctx context.Context, probeID, typ string, data json.RawMessage) error { - handled <- typ + defer m.shutdown() + + var allowFirst atomic.Bool + processed := make(chan string, 4) + m.SetMessageHandler(func(_ context.Context, _, typ string, _ json.RawMessage) error { + if typ == "first" && !allowFirst.Load() { + return errors.New("temporary database failure") + } + processed <- typ return nil }) - client := newTestClient(m, "agent-2") - go m.processLoop(client) - defer client.stopProcessor() + client := newTestClient(m, "agent-2", "boot-1") + client.session.reliableQueue <- agentMessage{seq: 1, typ: "first"} + client.session.reliableQueue <- agentMessage{seq: 2, typ: "second"} - // seq=0 表示不参与可靠投递（如常规指标消息）：正常处理但不 ack - client.proc.queue <- agentMessage{seq: 0, typ: "metrics"} - - select { - case typ := <-handled: - if typ != "metrics" { - t.Fatalf("handled = %s, want metrics", typ) - } - case <-time.After(2 * time.Second): - t.Fatal("timed out waiting for processing") + time.Sleep(100 * time.Millisecond) + if ack := m.AckSeq("agent-2"); ack != 0 { + t.Fatalf("AckSeq during failure = %d, want 0", ack) } - select { - case raw := <-client.Send: - t.Fatalf("seq=0 message should not be acked, got %s", raw) + case typ := <-processed: + t.Fatalf("processed %s before seq 1 recovered", typ) default: } -} -func TestAdvanceAckSeqMonotonic(t *testing.T) { - m := NewManager(zap.NewNop()) - m.advanceAckSeq("a", 5) - m.advanceAckSeq("a", 3) // 回退无效 - if got := m.AckSeq("a"); got != 5 { - t.Fatalf("AckSeq = %d, want 5", got) + allowFirst.Store(true) + for _, want := range []string{"first", "second"} { + select { + case got := <-processed: + if got != want { + t.Fatalf("processed %q, want %q", got, want) + } + case <-time.After(2 * time.Second): + t.Fatalf("timed out waiting for %s", want) + } } - if got := m.AckSeq("missing"); got != 0 { - t.Fatalf("AckSeq(missing) = %d, want 0", got) + if ack := m.AckSeq("agent-2"); ack != 2 { + t.Fatalf("AckSeq after recovery = %d, want 2", ack) } } -func TestBeginAgentSessionResetsAckSeqOnReboot(t *testing.T) { +func TestBootIDReplacementRejectsOldSessionCommit(t *testing.T) { m := NewManager(zap.NewNop()) + defer m.shutdown() - // 探针运行期间位点推进 - m.BeginAgentSession("a", "boot-1") - m.advanceAckSeq("a", 5000) + oldStarted := make(chan struct{}) + releaseOld := make(chan struct{}) + newProcessed := make(chan struct{}) + m.SetMessageHandler(func(_ context.Context, _, typ string, _ json.RawMessage) error { + switch typ { + case "old": + close(oldStarted) + <-releaseOld // 模拟不遵守 context 的旧业务处理器 + case "new": + close(newProcessed) + } + return nil + }) - // 同一进程重连：位点保留，重放去重继续有效 - m.BeginAgentSession("a", "boot-1") - if got := m.AckSeq("a"); got != 5000 { - t.Fatalf("AckSeq after reconnect = %d, want 5000", got) + oldClient := newTestClient(m, "agent-3", "boot-old") + oldSession := oldClient.session + oldSession.reliableQueue <- agentMessage{seq: 100, typ: "old"} + <-oldStarted + + newClient := newTestClient(m, "agent-3", "boot-new") + if oldClient.session == newClient.session { + t.Fatal("BootID change should replace the session") + } + if ack := m.AckSeq("agent-3"); ack != 0 { + t.Fatalf("new session AckSeq = %d, want 0", ack) } - // 探针重启（新 BootID，内存队列清空、seq 从 1 重新分配）： - // 位点必须重置，否则重启后的小序号会被误判为重复消息丢弃 - m.BeginAgentSession("a", "boot-2") - if got := m.AckSeq("a"); got != 0 { - t.Fatalf("AckSeq after reboot = %d, want 0", got) + close(releaseOld) + newClient.session.reliableQueue <- agentMessage{seq: 1, typ: "new"} + select { + case <-newProcessed: + case <-time.After(2 * time.Second): + t.Fatal("new session message was not processed") } - - // 旧版探针不携带 BootID：不影响现有位点 - m.advanceAckSeq("a", 7) - m.BeginAgentSession("a", "") - if got := m.AckSeq("a"); got != 7 { - t.Fatalf("AckSeq after legacy register = %d, want 7", got) + if ack := readAck(t, newClient); ack != 1 { + t.Fatalf("new session ack = %d, want 1", ack) + } + time.Sleep(50 * time.Millisecond) + if ack := m.AckSeq("agent-3"); ack != 1 { + t.Fatalf("old session polluted AckSeq: got %d, want 1", ack) } } -func TestEnqueueReturnsAfterProcessorStop(t *testing.T) { +func TestSameBootReconnectReusesSession(t *testing.T) { m := NewManager(zap.NewNop()) - client := newTestClient(m, "agent-3") - client.stopProcessor() + defer m.shutdown() - if err := client.enqueue(agentMessage{seq: 1, typ: "x"}); err != ErrProcessorStopped { - t.Fatalf("enqueue after stop = %v, want ErrProcessorStopped", err) + first := newTestClient(m, "agent-4", "boot-1") + first.session.ackSeq.Store(9) + second := newTestClient(m, "agent-4", "boot-1") + + if first.session != second.session { + t.Fatal("same BootID reconnect should reuse session") } + if ack := m.AckSeq("agent-4"); ack != 9 { + t.Fatalf("AckSeq after reconnect = %d, want 9", ack) + } + select { + case <-first.done: + default: + t.Fatal("old transport should be closed after replacement") + } + if m.Unregister(first) { + t.Fatal("old transport must not unregister the replacement") + } + if current, ok := m.GetClient("agent-4"); !ok || current != second { + t.Fatal("replacement transport was removed by old unregister") + } +} + +func TestReplacedTransportCannotEnqueueOutOfOrderSequence(t *testing.T) { + m := NewManager(zap.NewNop()) + defer m.shutdown() + processed := make(chan string, 2) + m.SetMessageHandler(func(_ context.Context, _, typ string, _ json.RawMessage) error { + processed <- typ + return nil + }) + + oldClient := newTestClient(m, "agent-order", "boot-1") + newClient := newTestClient(m, "agent-order", "boot-1") + if err := m.enqueueReliable(context.Background(), oldClient, agentMessage{seq: 10, typ: "stale"}); !errors.Is(err, ErrClientNotFound) { + t.Fatalf("old transport enqueue = %v, want ErrClientNotFound", err) + } + if err := m.enqueueReliable(context.Background(), newClient, agentMessage{seq: 1, typ: "current"}); err != nil { + t.Fatalf("new transport enqueue: %v", err) + } + + select { + case typ := <-processed: + if typ != "current" { + t.Fatalf("processed %q, want current", typ) + } + case <-time.After(time.Second): + t.Fatal("current transport message was not processed") + } + select { + case typ := <-processed: + t.Fatalf("processed stale transport message %q", typ) + case <-time.After(50 * time.Millisecond): + } +} + +func TestBestEffortDoesNotBlockReliableEvents(t *testing.T) { + m := NewManager(zap.NewNop()) + defer m.shutdown() + + blockTelemetry := make(chan struct{}) + reliableDone := make(chan struct{}) + m.SetMessageHandler(func(_ context.Context, _, typ string, _ json.RawMessage) error { + if typ == "metrics" { + <-blockTelemetry + } else if typ == "event" { + close(reliableDone) + } + return nil + }) + + client := newTestClient(m, "agent-5", "boot-1") + client.session.bestEffortQueue <- agentMessage{typ: "metrics"} + client.session.reliableQueue <- agentMessage{seq: 1, typ: "event"} + select { + case <-reliableDone: + case <-time.After(500 * time.Millisecond): + t.Fatal("reliable event was blocked by best-effort telemetry") + } + close(blockTelemetry) +} + +func TestPermanentMessageAdvancesAck(t *testing.T) { + m := NewManager(zap.NewNop()) + defer m.shutdown() + m.SetMessageHandler(func(context.Context, string, string, json.RawMessage) error { + return Permanent(errors.New("invalid payload")) + }) + + client := newTestClient(m, "agent-6", "boot-1") + client.session.reliableQueue <- agentMessage{seq: 1, typ: "bad"} + if ack := readAck(t, client); ack != 1 { + t.Fatalf("permanent message ack = %d, want 1", ack) + } +} + +func TestClosedClientRejectsSendWithoutPanic(t *testing.T) { + m := NewManager(zap.NewNop()) + client := newTestClient(m, "agent-7", "boot-1") + client.close() + + if client.trySend([]byte("ack")) { + t.Fatal("closed client accepted non-blocking send") + } + if err := client.sendWithTimeout([]byte("command"), time.Second); !errors.Is(err, ErrClientNotFound) { + t.Fatalf("send to closed client = %v, want ErrClientNotFound", err) + } + m.shutdown() } func TestClientShouldWriteStatusDebounce(t *testing.T) { m := NewManager(zap.NewNop()) - client := newTestClient(m, "agent-4") - + client := NewClient("agent-8", nil, m) if !client.ShouldWriteStatus(time.Minute) { t.Fatal("first check should pass") } diff --git a/pkg/agent/service/agent.go b/pkg/agent/service/agent.go index c26e79c..28feee3 100644 --- a/pkg/agent/service/agent.go +++ b/pkg/agent/service/agent.go @@ -11,7 +11,6 @@ import ( "runtime" "strings" "sync" - "sync/atomic" "time" "github.com/pika-monitor/pika/internal/protocol" @@ -95,7 +94,6 @@ type Agent struct { collectorManager *collector.Manager metricsStore *metricsStore outbox *outbox - reliable atomic.Bool tamperProtector *tamper.Protector sshMonitor *sshmonitor.Monitor } @@ -123,6 +121,10 @@ func (a *Agent) Start(ctx context.Context) error { // 采集与发送解耦：采集永远运行（断线也写快照），发送独立读快照上报 go a.collectLoop(ctx) go a.sendLoop(ctx) + // 安全事件消费属于 Agent 生命周期，而不是单次 WebSocket 连接。 + // 断线和退避期间仍持续写入 outbox，避免上游小缓冲区溢出丢事件。 + go a.tamperEventLoop(ctx) + go a.sshLoginEventLoop(ctx) // 启动探针主循环 b := &backoff.Backoff{ @@ -230,9 +232,8 @@ func (a *Agent) runOnce(ctx context.Context, onRegistered func()) error { // 重置发送游标，使下一个 sendLoop tick 立即发送当前全量快照， // 保证重连瞬间服务端就能拿到最新数据 a.metricsStore.reset() - // 服务端声明支持可靠投递时，事件类消息改走 outbox；注册响应携带的 - // 累计确认位点先修剪队列，避免重放服务端已处理的消息 - a.reliable.Store(registerResp.Reliable) + // 可靠服务端返回累计确认位点；旧服务端没有 ACK，发送器会在每次 + // WebSocket 写成功后本地确认，从而完整支持升级后的降级/回滚。 if registerResp.Reliable && registerResp.AckSeq > 0 { a.outbox.ack(registerResp.AckSeq) } @@ -267,21 +268,10 @@ func (a *Agent) runOnce(ctx context.Context, onRegistered func()) error { } }) - // 事件类消息可靠发送（仅服务端支持时启用） - if registerResp.Reliable { - wg.Go(func() { - a.eventSendLoop(done, conn) - }) - } - - // 启动防篡改事件监控 + // 每个连接只有一个 outbox 发送器。可靠服务端等待累计 ACK；旧服务端 + // 按写成功确认，队列不会因版本回退而被永久搁置。 wg.Go(func() { - a.tamperEventLoop(ctx, done) - }) - - // 启动 SSH 登录事件监控 - wg.Go(func() { - a.sshLoginEventLoop(ctx, done) + a.eventSendLoop(done, conn, registerResp.Reliable) }) // 等待第一个错误或上下文取消 @@ -461,23 +451,36 @@ func (a *Agent) handleAck(data json.RawMessage) { a.outbox.ack(ack.Seq) } -// eventSendLoop 事件类消息的可靠发送循环（每个连接一个）。 -// 连接建立后先从队头重放全部未确认消息（服务端按 seq 去重），之后增量 -// 发送新入队的消息；写失败直接返回并断开连接，未确认消息留给下一个 -// 连接重放。 -func (a *Agent) eventSendLoop(done chan struct{}, conn *safeConn) { +// eventSendLoop 是单连接 outbox 发送器。可靠模式下由服务端 ACK 修剪； +// 兼容旧服务端时以 WebSocket 写成功作为本地确认，至少保证断线期间已经 +// 入队的消息会在降级连接上继续发送。 +func (a *Agent) eventSendLoop(done chan struct{}, conn *safeConn, reliable bool) { + if err := sendOutboxLoop(a.outbox, done, conn, reliable); err != nil { + slog.Warn("事件消息发送失败，等待重连后重放", "error", err) + a.killConn(conn) + } +} + +type jsonWriter interface { + WriteJSON(v interface{}) error +} + +// sendOutboxLoop 是与连接生命周期无关的发送算法，便于独立验证可靠和 +// 旧版降级语义。写失败不修改待确认消息，由调用方关闭连接并重连。 +func sendOutboxLoop(outbox *outbox, done <-chan struct{}, writer jsonWriter, reliable bool) error { var lastSent uint64 for { - for _, msg := range a.outbox.after(lastSent) { - if err := conn.WriteJSON(msg); err != nil { - slog.Warn("事件消息发送失败，等待重连后重放", "error", err, "seq", msg.Seq) - a.killConn(conn) - return + for _, msg := range outbox.after(lastSent) { + if err := writer.WriteJSON(msg); err != nil { + return fmt.Errorf("send outbox seq %d: %w", msg.Seq, err) } lastSent = msg.Seq + if !reliable { + outbox.ack(msg.Seq) + } } - if !a.outbox.wait(done) { - return + if !outbox.wait(done) { + return nil } } } @@ -512,7 +515,7 @@ func (a *Agent) handleMonitorConfig(data json.RawMessage) { slog.Info("收到服务监控配置，立即执行检测", "count", len(payload.Items)) // 立即执行一次监控检测，结果以单元素 batch 上报。 - // 走 sendOutboundMessage（可靠模式下进 outbox），断线期间不丢 + // 走常驻 outbox，断线和版本降级期间也不会绕过队列 sample := manager.CollectMonitor(payload.Items) if err := a.sendOutboundMessage(protocol.OutboundMessage{ Type: protocol.MessageTypeMetrics, @@ -536,25 +539,10 @@ func (a *Agent) getActiveConn() *safeConn { return a.activeConn } -// sendOutboundMessage 发送事件类出站消息。服务端支持可靠投递时入队 -// outbox，由专用协程发送并等待服务端确认，断线期间不丢；否则降级为 -// 直接发送（与旧服务端兼容）。 +// sendOutboundMessage 把事件类消息无条件写入常驻 outbox。网络版本协商 +// 只影响连接级发送器如何确认，不影响事件是否在断线期间被接收。 func (a *Agent) sendOutboundMessage(msg protocol.OutboundMessage) error { - if a.reliable.Load() { - a.outbox.enqueue(msg) - return nil - } - - conn := a.getActiveConn() - if conn == nil { - return fmt.Errorf("connection is unavailable") - } - - if err := conn.WriteJSON(msg); err != nil { - a.killConn(conn) - return err - } - + a.outbox.enqueue(msg) return nil } @@ -798,14 +786,12 @@ func (a *Agent) sendTamperProtectResponse(success bool, message string, paths [] } // tamperEventLoop 防篡改事件监控循环（包含事件和告警） -func (a *Agent) tamperEventLoop(ctx context.Context, done chan struct{}) { +func (a *Agent) tamperEventLoop(ctx context.Context) { eventCh := a.tamperProtector.GetEvents() alertCh := a.tamperProtector.GetAlerts() for { select { - case <-done: - return case <-ctx.Done(): return case event := <-eventCh: @@ -1090,13 +1076,11 @@ func (a *Agent) sendSSHLoginConfigResult(success bool, enabled bool, message str } // sshLoginEventLoop SSH登录事件监控循环 -func (a *Agent) sshLoginEventLoop(ctx context.Context, done chan struct{}) { +func (a *Agent) sshLoginEventLoop(ctx context.Context) { eventCh := a.sshMonitor.GetEvents() for { select { - case <-done: - return case <-ctx.Done(): return case event := <-eventCh: diff --git a/pkg/agent/service/outbox.go b/pkg/agent/service/outbox.go index 8c2184f..ee44e2a 100644 --- a/pkg/agent/service/outbox.go +++ b/pkg/agent/service/outbox.go @@ -92,6 +92,11 @@ func (o *outbox) ack(seq uint64) int { o.mu.Lock() defer o.mu.Unlock() + // 不信任超过本进程已分配范围的 ACK。异常或过期服务端位点不能 + // 把未来序号也标成已确认，否则后续消息会长期无法正常修剪。 + if seq > o.nextSeq { + seq = o.nextSeq + } if seq <= o.ackedSeq { return 0 } diff --git a/pkg/agent/service/outbox_test.go b/pkg/agent/service/outbox_test.go index ea3becb..3087911 100644 --- a/pkg/agent/service/outbox_test.go +++ b/pkg/agent/service/outbox_test.go @@ -1,11 +1,36 @@ package service import ( + "errors" + "sync" "testing" + "time" "github.com/pika-monitor/pika/internal/protocol" ) +type recordingJSONWriter struct { + mu sync.Mutex + messages []protocol.OutboundMessage + err error + written chan struct{} +} + +func (w *recordingJSONWriter) WriteJSON(value interface{}) error { + if w.err != nil { + return w.err + } + msg := value.(protocol.OutboundMessage) + w.mu.Lock() + w.messages = append(w.messages, msg) + w.mu.Unlock() + select { + case w.written <- struct{}{}: + default: + } + return nil +} + func TestOutboxSeqAssignment(t *testing.T) { o := newOutbox() for i := 1; i <= 3; i++ { @@ -104,3 +129,76 @@ func TestOutboxEnqueueSetsMsgSeq(t *testing.T) { t.Fatalf("stored = %+v, want one message with seq 1", stored) } } + +func TestOutboxClampsAckToAllocatedSequence(t *testing.T) { + o := newOutbox() + o.enqueue(protocol.OutboundMessage{Type: protocol.MessageTypeTamperEvent}) + o.enqueue(protocol.OutboundMessage{Type: protocol.MessageTypeSSHLoginEvent}) + + if trimmed := o.ack(999); trimmed != 2 { + t.Fatalf("ack beyond allocated range trimmed %d, want 2", trimmed) + } + if _, acked, _ := o.stats(); acked != 2 { + t.Fatalf("acked seq = %d, want clamp to 2", acked) + } + + seq := o.enqueue(protocol.OutboundMessage{Type: protocol.MessageTypeCommandResp}) + if seq != 3 { + t.Fatalf("next seq = %d, want 3", seq) + } + if pending, _, _ := o.stats(); pending != 1 { + t.Fatalf("future message was incorrectly pre-acked; pending = %d, want 1", pending) + } +} + +func TestSendOutboundMessageQueuesWhileDisconnected(t *testing.T) { + a := &Agent{outbox: newOutbox()} + if err := a.sendOutboundMessage(protocol.OutboundMessage{Type: protocol.MessageTypeSSHLoginEvent}); err != nil { + t.Fatalf("queue disconnected event: %v", err) + } + if pending, _, _ := a.outbox.stats(); pending != 1 { + t.Fatalf("pending = %d, want 1", pending) + } +} + +func TestLegacyOutboxSenderDrainsPendingMessages(t *testing.T) { + o := newOutbox() + o.enqueue(protocol.OutboundMessage{Type: protocol.MessageTypeTamperEvent}) + o.enqueue(protocol.OutboundMessage{Type: protocol.MessageTypeSSHLoginEvent}) + + done := make(chan struct{}) + writer := &recordingJSONWriter{written: make(chan struct{}, 2)} + result := make(chan error, 1) + go func() { + result <- sendOutboxLoop(o, done, writer, false) + }() + + for i := 0; i < 2; i++ { + select { + case <-writer.written: + case <-time.After(time.Second): + t.Fatal("legacy sender did not flush pending messages") + } + } + close(done) + if err := <-result; err != nil { + t.Fatalf("legacy sender: %v", err) + } + if pending, acked, _ := o.stats(); pending != 0 || acked != 2 { + t.Fatalf("after legacy flush: pending=%d acked=%d, want 0/2", pending, acked) + } +} + +func TestOutboxWriteFailureKeepsMessagePending(t *testing.T) { + o := newOutbox() + o.enqueue(protocol.OutboundMessage{Type: protocol.MessageTypeCommandResp}) + writer := &recordingJSONWriter{err: errors.New("network down"), written: make(chan struct{}, 1)} + + err := sendOutboxLoop(o, make(chan struct{}), writer, false) + if err == nil { + t.Fatal("expected write failure") + } + if pending, acked, _ := o.stats(); pending != 1 || acked != 0 { + t.Fatalf("after write failure: pending=%d acked=%d, want 1/0", pending, acked) + } +}