diff --git a/internal/handler/agent_enabled_test.go b/internal/handler/agent_enabled_test.go index 5209feb..fc52697 100644 --- a/internal/handler/agent_enabled_test.go +++ b/internal/handler/agent_enabled_test.go @@ -23,26 +23,28 @@ func TestHandleWebSocketMessageIgnoresDisabledAgent(t *testing.T) { t.Fatalf("migrate agent: %v", err) } - agent := models.Agent{ID: "agent-disabled", Enabled: true} - if err := database.Create(&agent).Error; err != nil { - t.Fatalf("create agent: %v", err) - } - if err := database.Model(&models.Agent{}).Where("id = ?", agent.ID).Update("enabled", false).Error; err != nil { - t.Fatalf("disable agent: %v", err) - } - h := &AgentHandler{ logger: zap.NewNop(), agentService: &service.AgentService{ AgentRepo: repo.NewAgentRepo(database), }, } + + agent := models.Agent{ID: "agent-disabled", Enabled: true} + if err := database.Create(&agent).Error; err != nil { + t.Fatalf("create agent: %v", err) + } + // 通过服务层切换启用状态，保证 IsAgentEnabled 的缓存被正确失效 + if err := h.agentService.UpdateAgentEnabled(context.Background(), agent.ID, false); err != nil { + t.Fatalf("disable agent: %v", err) + } + invalidMetrics := json.RawMessage(`{`) if err := h.handleWebSocketMessage(context.Background(), agent.ID, string(protocol.MessageTypeMetrics), invalidMetrics); err != nil { t.Fatalf("disabled agent data should be ignored, got error: %v", err) } - if err := database.Model(&models.Agent{}).Where("id = ?", agent.ID).Update("enabled", true).Error; err != nil { + if err := h.agentService.UpdateAgentEnabled(context.Background(), agent.ID, true); err != nil { t.Fatalf("enable agent: %v", err) } if err := h.handleWebSocketMessage(context.Background(), agent.ID, string(protocol.MessageTypeMetrics), invalidMetrics); err == nil { diff --git a/internal/handler/agent_ws_handler.go b/internal/handler/agent_ws_handler.go index 7562d63..7341049 100644 --- a/internal/handler/agent_ws_handler.go +++ b/internal/handler/agent_ws_handler.go @@ -38,11 +38,12 @@ func (h *AgentHandler) HandleWebSocket(c *echo.Context) error { return err } - defer func() { - h.markAgentOffline(agent.ID) - }() + // 同步探针会话：BootID 变化（探针重启）时重置确认位点， + // 再把当前位点随注册响应返回 + h.wsManager.BeginAgentSession(agent.ID, registerReq.BootID) - // 发送注册成功响应 + // 发送注册成功响应：声明可靠投递支持，并回传该探针的累计确认 + // 位点，探针据此跳过已处理消息、只重放未确认部分 if err := h.sendRegisterSuccess(conn, agent.ID); err != nil { h.logger.Error("failed to send register ack", zap.Error(err)) conn.Close() @@ -72,6 +73,12 @@ func (h *AgentHandler) HandleWebSocket(c *echo.Context) error { h.wsManager.Register(client) + // 连接结束后标记离线；若探针已由新连接替换则跳过，避免旧连接的 + // 退出把新连接打成离线 + defer func() { + h.markAgentOffline(client) + }() + // 启动读写协程 go client.WritePump() client.ReadPump(context.Background()) @@ -152,27 +159,41 @@ func (h *AgentHandler) readRegisterRequest(conn *websocket.Conn) (*protocol.Regi return ®isterReq, nil } -func (h *AgentHandler) markAgentOffline(agentID string) { - _ = h.agentService.UpdateAgentStatus(context.Background(), agentID, 0) +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, 256), + Send: make(chan []byte, 512), Manager: h.wsManager, LastActive: time.Now(), } } -func (h *AgentHandler) handleWebSocketPong(agentID string) { +// pongStatusRefreshInterval 在线状态心跳刷新间隔：本连接首次 pong 立即 +// 写库，之后按该间隔周期性刷新，避免每个 pong（10s/探针）都打数据库 +const pongStatusRefreshInterval = 5 * time.Minute + +func (h *AgentHandler) handleWebSocketPong(client *ws.Client) { + if !client.ShouldWriteStatus(pongStatusRefreshInterval) { + return + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - if err := h.agentService.UpdateAgentStatus(ctx, agentID, 1); err != nil { - h.logger.Warn("failed to update agent status on pong", zap.String("agentID", agentID), zap.Error(err)) + if err := h.agentService.UpdateAgentStatus(ctx, client.ID, 1); err != nil { + h.logger.Warn("failed to update agent status on pong", zap.String("agentID", client.ID), zap.Error(err)) + return } + client.MarkStatusWritten() } func (h *AgentHandler) handleMetricsMessage(ctx context.Context, agentID string, data json.RawMessage) error { @@ -260,8 +281,10 @@ func (h *AgentHandler) handleTamperProtectMessage(ctx context.Context, agentID s // sendRegisterSuccess 发送注册成功响应 func (h *AgentHandler) sendRegisterSuccess(conn *websocket.Conn, agentID string) error { resp := protocol.RegisterResponse{ - AgentID: agentID, - Status: "success", + AgentID: agentID, + Status: "success", + Reliable: true, + AckSeq: h.wsManager.AckSeq(agentID), } return conn.WriteJSON(protocol.OutboundMessage{ Type: protocol.MessageTypeRegisterAck, diff --git a/internal/protocol/message.go b/internal/protocol/message.go index b4300d8..37dbdbe 100644 --- a/internal/protocol/message.go +++ b/internal/protocol/message.go @@ -6,18 +6,25 @@ import "encoding/json" type InputMessage struct { Type MessageType `json:"type"` Data json.RawMessage `json:"data"` + Seq uint64 `json:"seq,omitempty"` // 可靠投递序号（agent→server 事件流，0 表示不参与） } // OutboundMessage WebSocket 出站消息结构 type OutboundMessage struct { Type MessageType `json:"type"` Data interface{} `json:"data"` + Seq uint64 `json:"seq,omitempty"` // 可靠投递序号（agent→server 事件流，0 表示不参与） } // RegisterRequest 注册请求 type RegisterRequest struct { AgentInfo AgentInfo `json:"agentInfo"` ApiKey string `json:"apiKey"` + // BootID 探针进程级会话标识（每次进程启动重新生成）。服务端据此 + // 识别探针重启：纯内存事件队列重启后 seq 从 1 重新分配，必须重置 + // 该探针的累计确认位点，否则新消息会被旧位点误判为重复而丢弃。 + // 旧版探针不携带该字段。 + BootID string `json:"bootId,omitempty"` } // RegisterResponse 注册响应 @@ -25,6 +32,16 @@ type RegisterResponse struct { AgentID string `json:"agentId"` Status string `json:"status"` Message string `json:"message,omitempty"` + // Reliable 表示服务端支持可靠投递（seq 去重 + ack 确认）。 + // 旧版服务端不返回该字段，探针应降级为直接发送。 + Reliable bool `json:"reliable,omitempty"` + // AckSeq 服务端已处理该探针事件流的最高序号，探针据此跳过已确认消息 + AckSeq uint64 `json:"ackSeq,omitempty"` +} + +// AckData 可靠投递累计确认：seq 及之前的所有事件流消息已被服务端处理 +type AckData struct { + Seq uint64 `json:"seq"` } // AgentInfo 探针信息 @@ -62,6 +79,8 @@ const ( // 指标消息（批量） MessageTypeMetrics MessageType = "metrics" MessageTypeMonitorConfig MessageType = "monitor_config" + // 可靠投递确认（server→agent，累计确认事件流处理进度） + MessageTypeAck MessageType = "ack" // 防篡改消息 MessageTypeTamperProtect MessageType = "tamper_protect" MessageTypeTamperEvent MessageType = "tamper_event" diff --git a/internal/service/agent_service.go b/internal/service/agent_service.go index 04cbf21..ac052d8 100644 --- a/internal/service/agent_service.go +++ b/internal/service/agent_service.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "sort" + "sync" "time" "github.com/go-orz/orz" @@ -27,8 +28,19 @@ type AgentService struct { apiKeyService *ApiKeyService metricService *MetricService geoipService *GeoIPService + // enabledCache IsAgentEnabled 的短 TTL 缓存。该检查位于每条 + // WebSocket 消息的处理路径上，直连数据库会在探针数多时拖垮处理。 + enabledCache sync.Map // agentID → enabledCacheEntry } +type enabledCacheEntry struct { + value bool + expiresAt time.Time +} + +// enabledCacheTTL 探针启用状态缓存时长：禁用操作最迟该时长后生效 +const enabledCacheTTL = 3 * time.Second + // ErrInvalidAgentOrder 表示提交的排序列表无法完整、唯一地对应当前探针列表。 var ErrInvalidAgentOrder = errors.New("无效的探针排序") @@ -150,9 +162,26 @@ func (s *AgentService) UpdateAgentStatus(ctx context.Context, agentID string, st return s.AgentRepo.UpdateStatus(ctx, agentID, status, time.Now().UnixMilli()) } -// IsAgentEnabled 查询探针是否允许接收并处理数据。 +// IsAgentEnabled 查询探针是否允许接收并处理数据（带短 TTL 缓存， +// 禁用操作会立即失效缓存）。 func (s *AgentService) IsAgentEnabled(ctx context.Context, agentID string) (bool, error) { - return s.AgentRepo.IsEnabled(ctx, agentID) + if cached, ok := s.enabledCache.Load(agentID); ok { + entry := cached.(enabledCacheEntry) + if time.Now().Before(entry.expiresAt) { + return entry.value, nil + } + } + + enabled, err := s.AgentRepo.IsEnabled(ctx, agentID) + if err != nil { + return false, err + } + + s.enabledCache.Store(agentID, enabledCacheEntry{ + value: enabled, + expiresAt: time.Now().Add(enabledCacheTTL), + }) + return enabled, nil } // UpdateAgentEnabled 更新探针启用状态。禁用时立即标记离线，后续状态更新由仓储层原子拦截。 @@ -168,7 +197,13 @@ func (s *AgentService) UpdateAgentEnabled(ctx context.Context, agentID string, e if !enabled { updates["status"] = 0 } - return s.AgentRepo.UpdateColumnsById(ctx, agentID, updates) + if err := s.AgentRepo.UpdateColumnsById(ctx, agentID, updates); err != nil { + return err + } + + // 立即失效缓存，让禁用/启用马上对消息处理路径生效 + s.enabledCache.Delete(agentID) + return nil } // UpdatePublicIP 更新探针的公网 IP 信息 diff --git a/internal/websocket/manager.go b/internal/websocket/manager.go index e5a1a2c..50fab86 100644 --- a/internal/websocket/manager.go +++ b/internal/websocket/manager.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "sync" + "sync/atomic" "time" "github.com/gorilla/websocket" @@ -15,8 +16,42 @@ const ( serverPingInterval = 10 * time.Second serverPongWait = 30 * time.Second serverWriteWait = 10 * time.Second + + // agentQueueSize 每个探针的待处理消息队列长度。队列满时断开连接， + // 由探针重连后重放未确认消息，读循环不会被慢处理（DB/VM 写入）拖死。 + agentQueueSize = 512 + // agentQueueWait 消息入队的最长等待，超时判定为处理能力不足 + agentQueueWait = 2 * time.Second + // serverSendWait 下行消息写入 Send 通道的最长等待 + serverSendWait = 3 * time.Second ) +// agentMessage 从探针收到的一条待处理消息 +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 +} + +func newProcessor() *processor { + return &processor{ + queue: make(chan agentMessage, agentQueueSize), + quit: make(chan struct{}), + } +} + +func (p *processor) stop() { + p.quitOnce.Do(func() { close(p.quit) }) +} + // Client WebSocket客户端 type Client struct { ID string // 探针ID @@ -24,8 +59,12 @@ type Client struct { Send chan []byte // 发送消息通道 Manager *Manager // 管理器引用 LastActive time.Time // 最后活跃时间 - closed bool // 标记channel是否已关闭 - closeMu sync.Mutex // 保护closed字段 + proc *processor // 上行消息的有序处理器 + // lastStatusWrite 上次把在线状态写入数据库的时间（毫秒），用于 + // 对 pong 触发的状态更新去抖 + lastStatusWrite atomic.Int64 + closed bool // 标记channel是否已关闭 + closeMu sync.Mutex } // Manager WebSocket连接管理器 @@ -38,13 +77,18 @@ type Manager struct { logger *zap.Logger // 日志 onMessage MessageHandler // 消息处理器 onPong PongHandler // Pong 处理器 + // ackSeqs 每个探针事件流的累计确认位点（探针重启或服务端重启后 + // 丢失，退化为重放+去重，可能产生少量重复处理） + ackSeqs map[string]uint64 + // sessions 每个探针最近一次注册携带的 BootID，用于识别探针进程重启 + sessions map[string]string } // MessageHandler 消息处理器接口 type MessageHandler func(ctx context.Context, probeID string, messageType string, data json.RawMessage) error // PongHandler Pong 处理器接口 -type PongHandler func(probeID string) +type PongHandler func(client *Client) // NewManager 创建新的WebSocket管理器 func NewManager(logger *zap.Logger) *Manager { @@ -53,6 +97,8 @@ func NewManager(logger *zap.Logger) *Manager { 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, } } @@ -89,8 +135,11 @@ func (m *Manager) Run(ctx context.Context) { } } -// Register 注册客户端（公开方法） +// Register 注册客户端（公开方法）。同步初始化处理器，确保随后的 +// ReadPump 可以立即入队；注册表维护仍走 Run 循环的注册通道。 func (m *Manager) Register(client *Client) { + client.proc = newProcessor() + go m.processLoop(client) m.register <- client } @@ -100,10 +149,11 @@ func (m *Manager) registerClient(client *Client) { defer m.mu.Unlock() // 如果已存在该探针的连接，先关闭旧连接 - if oldClient, exists := m.clients[client.ID]; exists { + 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() } m.clients[client.ID] = client @@ -115,11 +165,14 @@ func (m *Manager) unregisterClient(client *Client) { m.mu.Lock() defer m.mu.Unlock() - if _, exists := m.clients[client.ID]; exists { + // 只注销当前注册的连接：探针快速重连时旧连接的退出路径不能把 + // 新连接从注册表中移除，否则新连接的下行通道会被关闭 + 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))) } + client.stopProcessor() } // broadcastMessage 广播消息 @@ -154,10 +207,10 @@ func (m *Manager) checkInactiveClients() { for _, client := range inactiveClients { // 再次检查客户端是否仍然存在（避免竞态条件） m.mu.RLock() - _, exists := m.clients[client.ID] + cur, exists := m.clients[client.ID] m.mu.RUnlock() - if exists { + if exists && cur == client { m.logger.Warn("agent inactive timeout, disconnecting", zap.String("agentID", client.ID)) client.Conn.Close() m.unregister <- client @@ -165,7 +218,8 @@ func (m *Manager) checkInactiveClients() { } } -// SendToClient 发送消息给指定客户端 +// SendToClient 发送消息给指定客户端。通道满时最多等待 serverSendWait， +// 避免高负载时指令/配置被静默丢弃 func (m *Manager) SendToClient(probeID string, message []byte) error { m.mu.RLock() client, exists := m.clients[probeID] @@ -178,7 +232,7 @@ func (m *Manager) SendToClient(probeID string, message []byte) error { select { case client.Send <- message: return nil - default: + case <-time.After(serverSendWait): return ErrSendTimeout } } @@ -210,11 +264,48 @@ func (m *Manager) ClientCount() int { return len(m.clients) } -// ReadPump 读取客户端消息 +// AckSeq 返回探针事件流的累计确认位点。注册响应携带该值，探针据此 +// 跳过服务端已处理的消息，避免无谓重放。 +func (m *Manager) AckSeq(probeID string) uint64 { + m.mu.RLock() + defer m.mu.RUnlock() + return m.ackSeqs[probeID] +} + +// BeginAgentSession 在注册时同步探针会话：BootID 变化说明探针进程 +// 重启过（纯内存事件队列清空、seq 从 1 重新分配），必须重置确认位点， +// 否则重启后新分配的小序号会被旧位点误判为重复消息而丢弃。 +// 旧版探针不携带 BootID，保持原位点不变。 +func (m *Manager) BeginAgentSession(probeID, bootID string) { + if bootID == "" { + 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)) @@ -222,7 +313,7 @@ func (c *Client) ReadPump(ctx context.Context) { c.Conn.SetReadDeadline(time.Now().Add(serverPongWait)) c.LastActive = time.Now() if c.Manager.onPong != nil { - go c.Manager.onPong(c.ID) + go c.Manager.onPong(c) } return nil }) @@ -233,7 +324,7 @@ func (c *Client) ReadPump(ctx context.Context) { if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) { c.Manager.logger.Error("websocket read error", zap.Error(err), zap.String("agentID", c.ID)) } - break + return } c.LastActive = time.Now() @@ -245,15 +336,92 @@ func (c *Client) ReadPump(ctx context.Context) { continue } - // 处理消息 - if c.Manager.onMessage != nil { - if err := c.Manager.onMessage(ctx, c.ID, string(msg.Type), msg.Data); err != nil { - c.Manager.logger.Error("failed to handle message", zap.Error(err), zap.String("agentID", c.ID), zap.String("type", string(msg.Type))) - } + 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)) + 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) { + 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)) + 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 msg.seq > 0 { + m.advanceAckSeq(c.ID, msg.seq) + m.sendAck(c, msg.seq) + } + + case <-c.proc.quit: + return + } + } +} + +// sendAck 向探针发送累计确认。走下行 Send 通道，满时丢弃：丢失的 +// ack 只会让探针多保留一会儿消息，重连重放时按 seq 去重，无副作用。 +func (m *Manager) sendAck(c *Client, seq uint64) { + payload, err := json.Marshal(protocol.OutboundMessage{ + Type: protocol.MessageTypeAck, + Data: protocol.AckData{Seq: seq}, + }) + if err != nil { + return + } + select { + case c.Send <- payload: + default: + } +} + // WritePump 向客户端写入消息 func (c *Client) WritePump() { ticker := time.NewTicker(serverPingInterval) @@ -286,6 +454,19 @@ func (c *Client) WritePump() { } } +// 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() @@ -299,6 +480,8 @@ func (c *Client) closeChannel() { // 错误定义 var ( - ErrClientNotFound = &websocket.CloseError{Code: 1000, Text: "client not found"} - ErrSendTimeout = &websocket.CloseError{Code: 1001, Text: "send timeout"} + 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"} ) diff --git a/internal/websocket/manager_test.go b/internal/websocket/manager_test.go new file mode 100644 index 0000000..33a8cdf --- /dev/null +++ b/internal/websocket/manager_test.go @@ -0,0 +1,189 @@ +package websocket + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/pika-monitor/pika/internal/protocol" + "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() + return c +} + +func readAck(t *testing.T, c *Client) uint64 { + t.Helper() + select { + case raw := <-c.Send: + var msg protocol.InputMessage + if err := json.Unmarshal(raw, &msg); err != nil { + t.Fatalf("unmarshal ack: %v", err) + } + if msg.Type != protocol.MessageTypeAck { + t.Fatalf("expected ack message, got %s", msg.Type) + } + var ack protocol.AckData + if err := json.Unmarshal(msg.Data, &ack); err != nil { + t.Fatalf("unmarshal ack data: %v", err) + } + return ack.Seq + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for ack") + return 0 + } +} + +func TestProcessLoopDedupAckAndOrder(t *testing.T) { + m := NewManager(zap.NewNop()) + var processed []string + done := make(chan struct{}) + m.SetMessageHandler(func(ctx context.Context, probeID, typ string, data json.RawMessage) error { + processed = append(processed, typ) + if len(processed) == 2 { + close(done) + } + 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"} + + select { + case <-done: + case <-time.After(2 * time.Second): + 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) + } + if got := m.AckSeq("agent-1"); got != 3 { + t.Fatalf("AckSeq = %d, want 3", got) + } + + // 去重跳过的消息也回 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) + } +} + +func TestProcessLoopSeqZeroBypassesAck(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 + return nil + }) + + client := newTestClient(m, "agent-2") + go m.processLoop(client) + defer client.stopProcessor() + + // 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") + } + + select { + case raw := <-client.Send: + t.Fatalf("seq=0 message should not be acked, got %s", raw) + 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) + } + if got := m.AckSeq("missing"); got != 0 { + t.Fatalf("AckSeq(missing) = %d, want 0", got) + } +} + +func TestBeginAgentSessionResetsAckSeqOnReboot(t *testing.T) { + m := NewManager(zap.NewNop()) + + // 探针运行期间位点推进 + m.BeginAgentSession("a", "boot-1") + m.advanceAckSeq("a", 5000) + + // 同一进程重连：位点保留，重放去重继续有效 + m.BeginAgentSession("a", "boot-1") + if got := m.AckSeq("a"); got != 5000 { + t.Fatalf("AckSeq after reconnect = %d, want 5000", got) + } + + // 探针重启（新 BootID，内存队列清空、seq 从 1 重新分配）： + // 位点必须重置，否则重启后的小序号会被误判为重复消息丢弃 + m.BeginAgentSession("a", "boot-2") + if got := m.AckSeq("a"); got != 0 { + t.Fatalf("AckSeq after reboot = %d, want 0", got) + } + + // 旧版探针不携带 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) + } +} + +func TestEnqueueReturnsAfterProcessorStop(t *testing.T) { + m := NewManager(zap.NewNop()) + client := newTestClient(m, "agent-3") + client.stopProcessor() + + if err := client.enqueue(agentMessage{seq: 1, typ: "x"}); err != ErrProcessorStopped { + t.Fatalf("enqueue after stop = %v, want ErrProcessorStopped", err) + } +} + +func TestClientShouldWriteStatusDebounce(t *testing.T) { + m := NewManager(zap.NewNop()) + client := newTestClient(m, "agent-4") + + if !client.ShouldWriteStatus(time.Minute) { + t.Fatal("first check should pass") + } + client.MarkStatusWritten() + if client.ShouldWriteStatus(time.Minute) { + t.Fatal("second check within interval should be debounced") + } + if !client.ShouldWriteStatus(0) { + t.Fatal("zero interval should always pass") + } +} diff --git a/internal/wire.go b/internal/wire.go index eb81418..2236aca 100644 --- a/internal/wire.go +++ b/internal/wire.go @@ -103,17 +103,21 @@ type AppComponents struct { // provideVMClient 提供 VictoriaMetrics 客户端 func provideVMClient(cfg *config.AppConfig, logger *zap.Logger) *vmclient.VMClient { + // 写入超时默认 10s：VM 抖动时快速失败，未确认消息由探针重放 + // 机制兜底，避免 30s 级别的停滞阻塞探针消息的串行处理 + const defaultWriteTimeout = 10 * time.Second + // 检查配置 if cfg.VictoriaMetrics == nil || !cfg.VictoriaMetrics.Enabled { logger.Info("VictoriaMetrics is not enabled, using default configuration") // 返回一个默认配置的客户端（用于本地开发） - return vmclient.NewVMClient("http://localhost:8428", 30*time.Second, 60*time.Second) + return vmclient.NewVMClient("http://localhost:8428", defaultWriteTimeout, 60*time.Second) } // 使用配置创建客户端 writeTimeout := time.Duration(cfg.VictoriaMetrics.WriteTimeout) * time.Second if writeTimeout == 0 { - writeTimeout = 30 * time.Second + writeTimeout = defaultWriteTimeout } queryTimeout := time.Duration(cfg.VictoriaMetrics.QueryTimeout) * time.Second diff --git a/internal/wire_gen.go b/internal/wire_gen.go index 5c441b8..717e981 100644 --- a/internal/wire_gen.go +++ b/internal/wire_gen.go @@ -28,12 +28,8 @@ func InitializeApp(logger *zap.Logger, db *gorm.DB, cfg *config.AppConfig) (*App accountHandler := handler.NewAccountHandler(accountService) apiKeyService := service.NewApiKeyService(logger, db) propertyService := service.NewPropertyService(logger, db) - themeService, err := service.NewThemeService(logger, propertyService, cfg) - if err != nil { - return nil, err - } - notifier := service.NewNotifier(logger) alertRuleService := service.NewAlertRuleService(logger, db, propertyService) + notifier := service.NewNotifier(logger) notificationService := service.NewNotificationService(logger, db, propertyService, alertRuleService, notifier) trafficService := service.NewTrafficService(logger, db, notificationService) vmClient := provideVMClient(cfg, logger) @@ -48,7 +44,6 @@ func InitializeApp(logger *zap.Logger, db *gorm.DB, cfg *config.AppConfig) (*App tamperService := service.NewTamperService(logger, db, manager, notificationService) ddnsService := service.NewDDNSService(logger, db, propertyService, manager) sshLoginService := service.NewSSHLoginService(logger, db, manager, geoIPService, notificationService) - publicIPService := service.NewPublicIPService(logger, db, propertyService, manager) agentHandler := handler.NewAgentHandler(logger, agentService, trafficService, metricService, monitorService, tamperService, ddnsService, sshLoginService, apiKeyService, propertyService, manager) apiKeyHandler := handler.NewApiKeyHandler(logger, apiKeyService) alertService := service.NewAlertService(logger, db, propertyService, alertRuleService, monitorService, notifier) @@ -60,8 +55,13 @@ func InitializeApp(logger *zap.Logger, db *gorm.DB, cfg *config.AppConfig) (*App dnsProviderHandler := handler.NewDNSProviderHandler(logger, propertyService) ddnsHandler := handler.NewDDNSHandler(logger, ddnsService) sshLoginHandler := handler.NewSSHLoginHandler(logger, sshLoginService) + themeService, err := service.NewThemeService(logger, propertyService, cfg) + if err != nil { + return nil, err + } themeHandler := handler.NewThemeHandler(logger, themeService) webHandler := handler.NewWebHandler(themeService, propertyService) + publicIPService := service.NewPublicIPService(logger, db, propertyService, manager) appComponents := &AppComponents{ AccountHandler: accountHandler, AgentHandler: agentHandler, @@ -133,16 +133,19 @@ type AppComponents struct { // provideVMClient 提供 VictoriaMetrics 客户端 func provideVMClient(cfg *config.AppConfig, logger *zap.Logger) *vmclient.VMClient { + // 写入超时默认 10s：VM 抖动时快速失败，未确认消息由探针重放 + // 机制兜底，避免 30s 级别的停滞阻塞探针消息的串行处理 + const defaultWriteTimeout = 10 * time.Second if cfg.VictoriaMetrics == nil || !cfg.VictoriaMetrics.Enabled { logger.Info("VictoriaMetrics is not enabled, using default configuration") - return vmclient.NewVMClient("http://localhost:8428", 30*time.Second, 60*time.Second) + return vmclient.NewVMClient("http://localhost:8428", defaultWriteTimeout, 60*time.Second) } writeTimeout := time.Duration(cfg.VictoriaMetrics.WriteTimeout) * time.Second if writeTimeout == 0 { - writeTimeout = 30 * time.Second + writeTimeout = defaultWriteTimeout } queryTimeout := time.Duration(cfg.VictoriaMetrics.QueryTimeout) * time.Second diff --git a/pkg/agent/service/agent.go b/pkg/agent/service/agent.go index e799144..c26e79c 100644 --- a/pkg/agent/service/agent.go +++ b/pkg/agent/service/agent.go @@ -11,6 +11,7 @@ import ( "runtime" "strings" "sync" + "sync/atomic" "time" "github.com/pika-monitor/pika/internal/protocol" @@ -22,6 +23,7 @@ import ( "github.com/pika-monitor/pika/pkg/agent/tamper" "github.com/pika-monitor/pika/pkg/version" + "github.com/google/uuid" "github.com/gorilla/websocket" "github.com/jpillora/backoff" "github.com/sourcegraph/conc" @@ -32,6 +34,10 @@ const ( agentPongWait = 30 * time.Second agentWriteWait = 5 * time.Second agentCollectTimeout = 30 * time.Second + // agentDataWriteWait 数据帧写入超时。没有超时的写在对端停止读取 + // （网络劣化、服务端处理阻塞）时会无限阻塞并持有连接写锁， + // 期间采集快照被覆盖造成静默丢数据。 + agentDataWriteWait = 10 * time.Second ) // safeConn 线程安全的 WebSocket 连接包装器 @@ -40,17 +46,23 @@ type safeConn struct { mu sync.Mutex } -// WriteJSON 线程安全地写入 JSON 消息 +// WriteJSON 线程安全地写入 JSON 消息（带写超时） func (sc *safeConn) WriteJSON(v interface{}) error { sc.mu.Lock() defer sc.mu.Unlock() + if err := sc.conn.SetWriteDeadline(time.Now().Add(agentDataWriteWait)); err != nil { + return err + } return sc.conn.WriteJSON(v) } -// WriteMessage 线程安全地写入消息 +// WriteMessage 线程安全地写入消息（带写超时） func (sc *safeConn) WriteMessage(messageType int, data []byte) error { sc.mu.Lock() defer sc.mu.Unlock() + if err := sc.conn.SetWriteDeadline(time.Now().Add(agentDataWriteWait)); err != nil { + return err + } return sc.conn.WriteMessage(messageType, data) } @@ -75,12 +87,15 @@ func (sc *safeConn) WriteControl(messageType int, data []byte, deadline time.Tim type Agent struct { cfg *config.Config idMgr *id.Manager + bootID string cancel context.CancelFunc connMu sync.RWMutex activeConn *safeConn collectorMu sync.RWMutex collectorManager *collector.Manager metricsStore *metricsStore + outbox *outbox + reliable atomic.Bool tamperProtector *tamper.Protector sshMonitor *sshmonitor.Monitor } @@ -90,8 +105,10 @@ func New(cfg *config.Config) *Agent { return &Agent{ cfg: cfg, idMgr: id.NewManager(), + bootID: uuid.NewString(), collectorManager: collector.NewManager(cfg), metricsStore: newMetricsStore(), + outbox: newOutbox(), tamperProtector: tamper.NewProtector(), sshMonitor: sshmonitor.NewMonitor(), } @@ -201,7 +218,8 @@ func (a *Agent) runOnce(ctx context.Context, onRegistered func()) error { }) // 发送注册消息 - if err := a.registerAgent(conn); err != nil { + registerResp, err := a.registerAgent(conn) + if err != nil { return fmt.Errorf("register failed: %w", err) } onRegistered() @@ -212,6 +230,12 @@ func (a *Agent) runOnce(ctx context.Context, onRegistered func()) error { // 重置发送游标，使下一个 sendLoop tick 立即发送当前全量快照， // 保证重连瞬间服务端就能拿到最新数据 a.metricsStore.reset() + // 服务端声明支持可靠投递时，事件类消息改走 outbox；注册响应携带的 + // 累计确认位点先修剪队列，避免重放服务端已处理的消息 + a.reliable.Store(registerResp.Reliable) + if registerResp.Reliable && registerResp.AckSeq > 0 { + a.outbox.ack(registerResp.AckSeq) + } defer func() { a.setActiveConn(nil) }() @@ -243,6 +267,13 @@ func (a *Agent) runOnce(ctx context.Context, onRegistered func()) error { } }) + // 事件类消息可靠发送（仅服务端支持时启用） + if registerResp.Reliable { + wg.Go(func() { + a.eventSendLoop(done, conn) + }) + } + // 启动防篡改事件监控 wg.Go(func() { a.tamperEventLoop(ctx, done) @@ -321,6 +352,8 @@ func (a *Agent) readLoop(conn *websocket.Conn, done chan struct{}) error { go a.handleSSHLoginConfig(msg.Data) case protocol.MessageTypeUninstall: go a.handleUninstall() + case protocol.MessageTypeAck: + a.handleAck(msg.Data) default: // 忽略其他类型 } @@ -345,12 +378,14 @@ func (a *Agent) pingLoop(ctx context.Context, conn *safeConn, done chan struct{} } } -// registerAgent 注册探针 -func (a *Agent) registerAgent(conn *safeConn) error { +// registerAgent 注册探针，返回服务端的注册响应（含可靠投递协商结果） +func (a *Agent) registerAgent(conn *safeConn) (protocol.RegisterResponse, error) { + var registerResp protocol.RegisterResponse + // 加载或生成探针 ID agentID, err := a.idMgr.Load() if err != nil { - return fmt.Errorf("load agent ID failed: %w", err) + return registerResp, fmt.Errorf("load agent ID failed: %w", err) } slog.Info("Agent ID", "id", agentID, "path", a.idMgr.GetPath()) @@ -366,7 +401,8 @@ func (a *Agent) registerAgent(conn *safeConn) error { agentName = hostname } - // 构建注册请求 + // 构建注册请求（BootID 标识本次探针进程，服务端据此识别重启并 + // 重置事件流确认位点——纯内存 outbox 重启后 seq 从 1 重新分配） registerReq := protocol.RegisterRequest{ AgentInfo: protocol.AgentInfo{ ID: agentID, @@ -377,41 +413,82 @@ func (a *Agent) registerAgent(conn *safeConn) error { Version: GetVersion(), }, ApiKey: a.cfg.Server.APIKey, + BootID: a.bootID, } if err := conn.WriteJSON(protocol.OutboundMessage{ Type: protocol.MessageTypeRegister, Data: registerReq, }); err != nil { - return fmt.Errorf("send register message failed: %w", err) + return registerResp, fmt.Errorf("send register message failed: %w", err) } // 读取注册响应 var response protocol.InputMessage if err := conn.ReadJSON(&response); err != nil { - return fmt.Errorf("read register response failed: %w", err) + return registerResp, fmt.Errorf("read register response failed: %w", err) } // 检查响应类型 if response.Type == protocol.MessageTypeRegisterErr { var errResp protocol.RegisterResponse if err := json.Unmarshal(response.Data, &errResp); err == nil { - return fmt.Errorf("register failed: %s", errResp.Message) + return registerResp, fmt.Errorf("register failed: %s", errResp.Message) } - return fmt.Errorf("register failed: unknown error") + return registerResp, fmt.Errorf("register failed: unknown error") } if response.Type != protocol.MessageTypeRegisterAck { - return fmt.Errorf("register failed: unexpected response type %s", response.Type) + return registerResp, fmt.Errorf("register failed: unexpected response type %s", response.Type) } - var registerResp protocol.RegisterResponse if err := json.Unmarshal(response.Data, ®isterResp); err != nil { - return fmt.Errorf("parse register response failed: %w", err) + return registerResp, fmt.Errorf("parse register response failed: %w", err) } - slog.Info("注册成功", "agentId", registerResp.AgentID, "status", registerResp.Status) - return nil + slog.Info("注册成功", "agentId", registerResp.AgentID, "status", registerResp.Status, + "reliable", registerResp.Reliable, "serverAckSeq", registerResp.AckSeq) + return registerResp, nil +} + +// handleAck 处理服务端的累计确认，修剪已投递的事件消息 +func (a *Agent) handleAck(data json.RawMessage) { + var ack protocol.AckData + if err := json.Unmarshal(data, &ack); err != nil { + slog.Warn("解析投递确认失败", "error", err) + return + } + a.outbox.ack(ack.Seq) +} + +// eventSendLoop 事件类消息的可靠发送循环（每个连接一个）。 +// 连接建立后先从队头重放全部未确认消息（服务端按 seq 去重），之后增量 +// 发送新入队的消息；写失败直接返回并断开连接，未确认消息留给下一个 +// 连接重放。 +func (a *Agent) eventSendLoop(done chan struct{}, conn *safeConn) { + 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 + } + lastSent = msg.Seq + } + if !a.outbox.wait(done) { + return + } + } +} + +// killConn 写失败时立即关闭连接，让读循环马上感知并触发重连，而不是 +// 等 30 秒 pong 超时。Close 可安全并发调用，会以错误唤醒阻塞中的读写。 +func (a *Agent) killConn(conn *safeConn) { + if conn == nil { + return + } + _ = conn.Close() } func (a *Agent) handleMonitorConfig(data json.RawMessage) { @@ -434,21 +511,16 @@ func (a *Agent) handleMonitorConfig(data json.RawMessage) { slog.Info("收到服务监控配置，立即执行检测", "count", len(payload.Items)) - // 立即执行一次监控检测，结果以单元素 batch 上报 + // 立即执行一次监控检测，结果以单元素 batch 上报。 + // 走 sendOutboundMessage（可靠模式下进 outbox），断线期间不丢 sample := manager.CollectMonitor(payload.Items) - conn := a.getActiveConn() - if conn == nil { - slog.Warn("当前连接不可用，无法上报监控结果") - return - } - - if err := conn.WriteJSON(protocol.OutboundMessage{ + if err := a.sendOutboundMessage(protocol.OutboundMessage{ Type: protocol.MessageTypeMetrics, Data: protocol.MetricsBatch{Samples: []protocol.MetricSample{sample}}, }); err != nil { slog.Warn("发送监控结果失败", "error", err) } else { - slog.Info("服务监控检测完成，已上报监控项结果", "count", len(payload.Items)) + slog.Info("服务监控检测完成，已提交监控项结果", "count", len(payload.Items)) } } @@ -464,13 +536,22 @@ func (a *Agent) getActiveConn() *safeConn { return a.activeConn } +// 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 } @@ -571,7 +652,9 @@ func (a *Agent) sendMetricsOnce() { Data: protocol.MetricsBatch{Samples: samples}, }); err != nil { slog.Warn("发送指标 batch 失败", "error", err) - // 连接写失败仅用于日志，重连由 readLoop/pingLoop 驱动 + // 写失败说明连接已不可用，主动断开让读循环立刻触发重连； + // 未确认样本由 metricsStore 在重连后重发（latest-wins） + a.killConn(conn) return } diff --git a/pkg/agent/service/outbox.go b/pkg/agent/service/outbox.go new file mode 100644 index 0000000..8c2184f --- /dev/null +++ b/pkg/agent/service/outbox.go @@ -0,0 +1,134 @@ +package service + +import ( + "log/slog" + "sync" + "time" + + "github.com/pika-monitor/pika/internal/protocol" +) + +const ( + // outboxMaxEntries 待确认队列上限。溢出时丢弃最旧消息：新事件（尤其是 + // 安全事件）比旧事件更有价值，且队列只增不减通常意味着探针已离线过久。 + outboxMaxEntries = 1024 + + // outboxDropLogInterval 溢出告警的日志节流，避免事件风暴时刷屏 + outboxDropLogInterval = 30 * time.Second +) + +type outboxEntry struct { + seq uint64 + msg protocol.OutboundMessage +} + +// outbox 事件类出站消息的有界待确认队列，实现 at-least-once 投递： +// +// - 每条消息入队时分配单调递增 seq； +// - 网络写成功不代表删除，只有收到服务端累计 ack（或重连时注册响应 +// 携带的 AckSeq）才修剪队列； +// - 每次新连接从队头整体重放未确认消息，由服务端按 seq 去重。 +// +// 纯内存实现：进程重启后清空，seq 从 1 重新分配。服务端通过注册请求 +// 携带的 BootID 识别探针重启并重置确认位点，避免新序号被旧位点误判 +// 为重复消息。 +type outbox struct { + mu sync.Mutex + entries []outboxEntry + nextSeq uint64 + ackedSeq uint64 + dropped uint64 + lastDrop time.Time + notify chan struct{} +} + +func newOutbox() *outbox { + return &outbox{ + notify: make(chan struct{}, 1), + } +} + +// enqueue 分配序号并入队，返回分配的 seq。队列满时丢弃最旧消息。 +func (o *outbox) enqueue(msg protocol.OutboundMessage) uint64 { + o.mu.Lock() + defer o.mu.Unlock() + + o.nextSeq++ + msg.Seq = o.nextSeq + o.entries = append(o.entries, outboxEntry{seq: msg.Seq, msg: msg}) + + if len(o.entries) > outboxMaxEntries { + drop := len(o.entries) - outboxMaxEntries + o.entries = o.entries[drop:] + o.dropped += uint64(drop) + if time.Since(o.lastDrop) >= outboxDropLogInterval { + slog.Warn("事件待确认队列溢出，已丢弃最旧消息", + "droppedTotal", o.dropped, "queueSize", len(o.entries)) + o.lastDrop = time.Now() + } + } + + o.signal() + return msg.Seq +} + +// after 返回 seq 大于 lastSent 的消息快照（按序），供发送协程增量发送或 +// 连接建立后整体重放。ack 修剪只从队头移除，不影响该判定。 +func (o *outbox) after(lastSent uint64) []protocol.OutboundMessage { + o.mu.Lock() + defer o.mu.Unlock() + + msgs := make([]protocol.OutboundMessage, 0, len(o.entries)) + for _, e := range o.entries { + if e.seq > lastSent { + msgs = append(msgs, e.msg) + } + } + return msgs +} + +// ack 记录服务端累计确认并修剪队列，返回本次修剪的消息数。 +func (o *outbox) ack(seq uint64) int { + o.mu.Lock() + defer o.mu.Unlock() + + if seq <= o.ackedSeq { + return 0 + } + o.ackedSeq = seq + + kept := o.entries[:0] + for _, e := range o.entries { + if e.seq > seq { + kept = append(kept, e) + } + } + trimmed := len(o.entries) - len(kept) + o.entries = kept + return trimmed +} + +// stats 返回队列统计（用于日志与观测）。 +func (o *outbox) stats() (pending int, acked uint64, dropped uint64) { + o.mu.Lock() + defer o.mu.Unlock() + return len(o.entries), o.ackedSeq, o.dropped +} + +// signal 唤醒发送协程（非阻塞，协程自带快照逻辑不会漏消息）。 +func (o *outbox) signal() { + select { + case o.notify <- struct{}{}: + default: + } +} + +// wait 阻塞等待新消息或退出信号。 +func (o *outbox) wait(done <-chan struct{}) bool { + select { + case <-o.notify: + return true + case <-done: + return false + } +} diff --git a/pkg/agent/service/outbox_test.go b/pkg/agent/service/outbox_test.go new file mode 100644 index 0000000..ea3becb --- /dev/null +++ b/pkg/agent/service/outbox_test.go @@ -0,0 +1,106 @@ +package service + +import ( + "testing" + + "github.com/pika-monitor/pika/internal/protocol" +) + +func TestOutboxSeqAssignment(t *testing.T) { + o := newOutbox() + for i := 1; i <= 3; i++ { + seq := o.enqueue(protocol.OutboundMessage{Type: protocol.MessageTypeSSHLoginEvent}) + if seq != uint64(i) { + t.Fatalf("enqueue #%d got seq %d", i, seq) + } + } + + msgs := o.after(0) + if len(msgs) != 3 { + t.Fatalf("after(0) returned %d messages, want 3", len(msgs)) + } + for i, msg := range msgs { + if msg.Seq != uint64(i+1) { + t.Fatalf("msgs[%d].Seq = %d, want %d", i, msg.Seq, i+1) + } + } +} + +func TestOutboxAckTrimsPrefix(t *testing.T) { + o := newOutbox() + for i := 0; i < 5; i++ { + o.enqueue(protocol.OutboundMessage{Type: protocol.MessageTypeTamperEvent}) + } + + if trimmed := o.ack(3); trimmed != 3 { + t.Fatalf("ack(3) trimmed %d, want 3", trimmed) + } + if pending, acked, _ := o.stats(); pending != 2 || acked != 3 { + t.Fatalf("after ack(3): pending=%d acked=%d, want pending=2 acked=3", pending, acked) + } + + // 重复/回退的 ack 无效 + if trimmed := o.ack(2); trimmed != 0 { + t.Fatalf("stale ack(2) trimmed %d, want 0", trimmed) + } + + // 剩余消息 seq 保持原值，重放判定不受修剪影响 + msgs := o.after(0) + if len(msgs) != 2 || msgs[0].Seq != 4 || msgs[1].Seq != 5 { + t.Fatalf("remaining messages = %+v, want seq 4,5", msgs) + } +} + +func TestOutboxAfterSkipsSent(t *testing.T) { + o := newOutbox() + for i := 0; i < 4; i++ { + o.enqueue(protocol.OutboundMessage{Type: protocol.MessageTypeCommandResp}) + } + + msgs := o.after(2) + if len(msgs) != 2 || msgs[0].Seq != 3 || msgs[1].Seq != 4 { + t.Fatalf("after(2) = %+v, want seq 3,4", msgs) + } + + // 全部已发送时返回空 + if msgs := o.after(4); len(msgs) != 0 { + t.Fatalf("after(4) = %+v, want empty", msgs) + } +} + +func TestOutboxOverflowDropsOldest(t *testing.T) { + o := newOutbox() + for i := 0; i < outboxMaxEntries+50; i++ { + o.enqueue(protocol.OutboundMessage{Type: protocol.MessageTypeTamperEvent}) + } + + pending, _, dropped := o.stats() + if pending != outboxMaxEntries { + t.Fatalf("pending = %d, want %d", pending, outboxMaxEntries) + } + if dropped != 50 { + t.Fatalf("dropped = %d, want 50", dropped) + } + + // 队列保留的是最新的消息 + msgs := o.after(0) + first := uint64(outboxMaxEntries + 50 - outboxMaxEntries + 1) + if len(msgs) != outboxMaxEntries || msgs[0].Seq != first { + t.Fatalf("oldest retained seq = %d, want %d", msgs[0].Seq, first) + } +} + +func TestOutboxEnqueueSetsMsgSeq(t *testing.T) { + o := newOutbox() + msg := protocol.OutboundMessage{Type: protocol.MessageTypeCommandResp} + o.enqueue(msg) + + // enqueue 内部复制赋值，入队消息带 seq + if msg.Seq != 0 { + t.Fatalf("caller message mutated, seq = %d", msg.Seq) + } + stored := o.after(0) + if len(stored) != 1 || stored[0].Seq != 1 { + t.Fatalf("stored = %+v, want one message with seq 1", stored) + } +}