feat: WebSocket 消息可靠投递(纯内存事件队列)

- agent 数据写入增加 10s 超时,写失败立即断连重连,消除无超时写卡死期间快照被覆盖的静默丢失
- 事件类消息(SSH登录/防篡改/指令响应/监控检测)改走内存 outbox:seq + 服务端累计 ack + 断线重放,上限 1024 条溢出丢最旧,重启清空(可接受)
- 注册请求携带 BootID 会话标识:探针重启后 seq 从 1 重新分配,服务端据此重置确认位点,避免新消息被旧位点误判为重复而丢弃
- server 读循环与处理解耦(per-agent 有序队列 + 串行 worker),慢处理不再阻塞 TCP 读取,队列满断连由探针重放兜底
- 修复探针快速重连时旧连接误注销新连接的竞态(此前会闪断新连接)
- IsAgentEnabled 加 3s TTL 缓存,禁用立即失效;pong 在线状态写库去抖(首次 + 5 分钟心跳刷新)
- VM 写超时默认 30s -> 10s,快速失败配合重放兜底

新旧版本双向兼容:注册响应 Reliable 字段协商,旧 server 下探针自动降级直发
This commit is contained in:
dushixiang committed 2026-08-30 22:30:41 +08:00
1 parent cf053cb720
commit 787fa6918d
11 files changed
+861 -80

No files matched your search

+11 -9
View File
@@ -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 {
+35 -12
View File
@@ -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 &registerReq, 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,