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:
1 parent
cf053cb720
commit
787fa6918d
11 files changed
+861
-80
No files matched your search
@@ -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 {
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in new issue
Block a user