diff --git a/go.mod b/go.mod index e61f322..d443cc3 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,7 @@ module github.com/dushixiang/pika go 1.25.5 require ( + github.com/boltdb/bolt v1.3.1 github.com/coreos/go-oidc/v3 v3.17.0 github.com/fsnotify/fsnotify v1.9.0 github.com/go-errors/errors v1.5.1 diff --git a/go.sum b/go.sum index 618b0be..9fb5bfa 100644 --- a/go.sum +++ b/go.sum @@ -3,6 +3,7 @@ aead.dev/minisign v0.3.0 h1:8Xafzy5PEVZqYDNP60yJHARlW1eOQtsKNp/Ph2c0vRA= aead.dev/minisign v0.3.0/go.mod h1:NLvG3Uoq3skkRMDuc3YHpWUTMTrSExqm+Ij73W13F6Y= filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA= filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4= +github.com/boltdb/bolt v1.3.1/go.mod h1:clJnj/oiGkjum5o1McbSZDSLxVThjynRyGBgiAx27Ps= github.com/coreos/go-oidc/v3 v3.17.0 h1:hWBGaQfbi0iVviX4ibC7bk8OKT5qNr4klBaCHVNvehc= github.com/coreos/go-oidc/v3 v3.17.0/go.mod h1:wqPbKFrVnE90vty060SB40FCJ8fTHTxSwyXJqZH+sI8= github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= diff --git a/internal/handler/agent_ws_handler.go b/internal/handler/agent_ws_handler.go index 7c8c679..ad69b14 100644 --- a/internal/handler/agent_ws_handler.go +++ b/internal/handler/agent_ws_handler.go @@ -154,7 +154,7 @@ func (h *AgentHandler) handleMetricsMessage(ctx context.Context, agentID string, if err != nil { return err } - return h.metricService.HandleMetricData(ctx, agentID, string(metricsWrapper.Type), json.RawMessage(metricsData)) + return h.metricService.HandleMetricData(ctx, agentID, string(metricsWrapper.Type), metricsData, metricsWrapper.Timestamp) } func (h *AgentHandler) handleCommandResponseMessage(ctx context.Context, agentID string, data json.RawMessage) error { diff --git a/internal/protocol/message.go b/internal/protocol/message.go index b476375..78f003c 100644 --- a/internal/protocol/message.go +++ b/internal/protocol/message.go @@ -39,8 +39,9 @@ type AgentInfo struct { // MetricsPayload 指标数据包装，发送端/接收端统一使用 type MetricsPayload struct { - Type MetricType `json:"type"` - Data interface{} `json:"data"` + Type MetricType `json:"type"` + Data interface{} `json:"data"` + Timestamp int64 `json:"timestamp,omitempty"` // 客户端采集时间(毫秒) } type MessageType string diff --git a/internal/service/metric_service.go b/internal/service/metric_service.go index 2afa171..43e5681 100644 --- a/internal/service/metric_service.go +++ b/internal/service/metric_service.go @@ -51,8 +51,10 @@ func NewMetricService(logger *zap.Logger, db *gorm.DB, propertyService *Property } // HandleMetricData 处理指标数据 -func (s *MetricService) HandleMetricData(ctx context.Context, agentID string, metricType string, data json.RawMessage) error { - now := time.Now().UnixMilli() +func (s *MetricService) HandleMetricData(ctx context.Context, agentID string, metricType string, data json.RawMessage, timestamp int64) error { + if timestamp == 0 { + timestamp = time.Now().UnixMilli() + } // 更新内存缓存 latestMetrics, ok := s.latestCache.Get(agentID) @@ -69,7 +71,7 @@ func (s *MetricService) HandleMetricData(ctx context.Context, agentID string, me return err } latestMetrics.CPU = &cpuData - metrics := s.convertToMetrics(agentID, metricType, &cpuData, now) + metrics := s.convertToMetrics(agentID, metricType, &cpuData, timestamp) return s.vmClient.Write(ctx, metrics) case protocol.MetricTypeMemory: @@ -78,7 +80,7 @@ func (s *MetricService) HandleMetricData(ctx context.Context, agentID string, me return err } latestMetrics.Memory = &memData - metrics := s.convertToMetrics(agentID, metricType, &memData, now) + metrics := s.convertToMetrics(agentID, metricType, &memData, timestamp) return s.vmClient.Write(ctx, metrics) case protocol.MetricTypeDisk: @@ -104,7 +106,7 @@ func (s *MetricService) HandleMetricData(ctx context.Context, agentID string, me Used: totalUsed, Free: totalFree, } - metrics := s.convertToMetrics(agentID, metricType, diskDataList, now) + metrics := s.convertToMetrics(agentID, metricType, diskDataList, timestamp) return s.vmClient.Write(ctx, metrics) case protocol.MetricTypeNetwork: @@ -134,7 +136,7 @@ func (s *MetricService) HandleMetricData(ctx context.Context, agentID string, me zap.String("agentId", agentID), zap.Error(err)) } - metrics := s.convertToMetrics(agentID, metricType, networkDataList, now) + metrics := s.convertToMetrics(agentID, metricType, networkDataList, timestamp) return s.vmClient.Write(ctx, metrics) case protocol.MetricTypeNetworkConnection: @@ -143,7 +145,7 @@ func (s *MetricService) HandleMetricData(ctx context.Context, agentID string, me return err } latestMetrics.NetworkConnection = &connData - metrics := s.convertToMetrics(agentID, metricType, &connData, now) + metrics := s.convertToMetrics(agentID, metricType, &connData, timestamp) return s.vmClient.Write(ctx, metrics) case protocol.MetricTypeDiskIO: @@ -151,7 +153,7 @@ func (s *MetricService) HandleMetricData(ctx context.Context, agentID string, me if err := json.Unmarshal(data, &diskIODataList); err != nil { return err } - metrics := s.convertToMetrics(agentID, metricType, diskIODataList, now) + metrics := s.convertToMetrics(agentID, metricType, diskIODataList, timestamp) return s.vmClient.Write(ctx, metrics) case protocol.MetricTypeHost: @@ -170,7 +172,7 @@ func (s *MetricService) HandleMetricData(ctx context.Context, agentID string, me Uptime: hostData.Uptime, BootTime: hostData.BootTime, Procs: hostData.Procs, - Timestamp: now, + Timestamp: timestamp, } latestMetrics.Host = hostMetric return s.metricRepo.SaveHostMetric(ctx, hostMetric) @@ -182,7 +184,7 @@ func (s *MetricService) HandleMetricData(ctx context.Context, agentID string, me } // 更新缓存 latestMetrics.GPU = gpuDataList - metrics := s.convertToMetrics(agentID, metricType, gpuDataList, now) + metrics := s.convertToMetrics(agentID, metricType, gpuDataList, timestamp) return s.vmClient.Write(ctx, metrics) case protocol.MetricTypeTemperature: @@ -192,7 +194,7 @@ func (s *MetricService) HandleMetricData(ctx context.Context, agentID string, me } // 更新缓存 latestMetrics.Temp = tempDataList - metrics := s.convertToMetrics(agentID, metricType, tempDataList, now) + metrics := s.convertToMetrics(agentID, metricType, tempDataList, timestamp) return s.vmClient.Write(ctx, metrics) case protocol.MetricTypeMonitor: @@ -206,10 +208,10 @@ func (s *MetricService) HandleMetricData(ctx context.Context, agentID string, me // 更新缓存 latestMetrics.Monitors = monitorDataList for _, monitorData := range monitorDataList { - s.updateMonitorCache(agentID, &monitorData, now) + s.updateMonitorCache(agentID, &monitorData, timestamp) } - metrics := s.convertToMetrics(agentID, metricType, monitorDataList, now) + metrics := s.convertToMetrics(agentID, metricType, monitorDataList, timestamp) return s.vmClient.Write(ctx, metrics) default: diff --git a/pkg/agent/collector/manager.go b/pkg/agent/collector/manager.go index cd04adf..5b8256c 100644 --- a/pkg/agent/collector/manager.go +++ b/pkg/agent/collector/manager.go @@ -1,6 +1,8 @@ package collector import ( + "time" + "github.com/dushixiang/pika/internal/protocol" "github.com/dushixiang/pika/pkg/agent/config" ) @@ -177,8 +179,9 @@ func (m *Manager) sendMetrics(conn WebSocketWriter, metricType protocol.MetricTy return conn.WriteJSON(protocol.OutboundMessage{ Type: protocol.MessageTypeMetrics, Data: protocol.MetricsPayload{ - Type: metricType, - Data: data, + Type: metricType, + Data: data, + Timestamp: time.Now().UnixMilli(), }, }) } diff --git a/pkg/agent/service/agent.go b/pkg/agent/service/agent.go index b30b61e..5963c4c 100644 --- a/pkg/agent/service/agent.go +++ b/pkg/agent/service/agent.go @@ -70,6 +70,7 @@ type Agent struct { activeConn *safeConn collectorMu sync.RWMutex collectorManager *collector.Manager + metricsBuffer *metricsBuffer tamperProtector *tamper.Protector sshMonitor *sshmonitor.Monitor } @@ -77,10 +78,12 @@ type Agent struct { // New 创建 Agent 实例 func New(cfg *config.Config) *Agent { return &Agent{ - cfg: cfg, - idMgr: id.NewManager(), - tamperProtector: tamper.NewProtector(), - sshMonitor: sshmonitor.NewMonitor(), + cfg: cfg, + idMgr: id.NewManager(), + collectorManager: collector.NewManager(cfg), + metricsBuffer: newMetricsBuffer(), + tamperProtector: tamper.NewProtector(), + sshMonitor: sshmonitor.NewMonitor(), } } @@ -90,6 +93,8 @@ func (a *Agent) Start(ctx context.Context) error { ctx, cancel := context.WithCancel(ctx) a.cancel = cancel + go a.metricsLoop(ctx) + // 启动探针主循环 b := &backoff.Backoff{ Min: 5 * time.Second, @@ -183,13 +188,8 @@ func (a *Agent) runOnce(ctx context.Context, onConnected func()) error { slog.Info("探针注册成功，开始监控...") - // 创建采集器管理器 - collectorManager := collector.NewManager(a.cfg) - a.setActiveConn(conn) - a.setCollectorManager(collectorManager) defer func() { - a.setCollectorManager(nil) a.setActiveConn(nil) }() @@ -220,16 +220,6 @@ func (a *Agent) runOnce(ctx context.Context, onConnected func()) error { } }) - // 启动指标采集循环 - wg.Go(func() { - if err := a.metricsLoop(ctx, conn, collectorManager, done); err != nil { - select { - case errChan <- fmt.Errorf("数据采集失败: %w", err): - default: - } - } - }) - // 启动防篡改事件监控 wg.Go(func() { a.tamperEventLoop(ctx, conn, done) @@ -455,9 +445,15 @@ func (a *Agent) getCollectorManager() *collector.Manager { } // metricsLoop 指标采集循环 -func (a *Agent) metricsLoop(ctx context.Context, conn *safeConn, manager *collector.Manager, done chan struct{}) error { +func (a *Agent) metricsLoop(ctx context.Context) { + manager := a.getCollectorManager() + if manager == nil { + manager = collector.NewManager(a.cfg) + a.setCollectorManager(manager) + } + // 立即采集一次动态数据 - if err := a.collectAndSendAllMetrics(conn, manager); err != nil { + if err := a.collectAndSendAllMetrics(manager); err != nil { slog.Warn("初始数据采集失败", "error", err) } @@ -469,73 +465,93 @@ func (a *Agent) metricsLoop(ctx context.Context, conn *safeConn, manager *collec select { case <-ticker.C: // 采集并发送各种动态指标 - if err := a.collectAndSendAllMetrics(conn, manager); err != nil { - return fmt.Errorf("数据采集失败: %w", err) + if err := a.collectAndSendAllMetrics(manager); err != nil { + slog.Warn("数据采集失败", "error", err) } - case <-done: - return nil case <-ctx.Done(): - return nil + return } } } // collectAndSendAllMetrics 采集并发送所有动态指标 -func (a *Agent) collectAndSendAllMetrics(conn *safeConn, manager *collector.Manager) error { +func (a *Agent) collectAndSendAllMetrics(manager *collector.Manager) error { + if manager == nil { + return fmt.Errorf("采集器未初始化") + } + + conn := a.getActiveConn() + if conn != nil { + if sent, err := a.metricsBuffer.Flush(conn); err != nil { + slog.Warn("发送缓存指标失败", "error", err) + } else if sent > 0 { + slog.Info("已发送缓存指标", "count", sent) + } + } + + writer := newMetricsWriter(conn, a.metricsBuffer) var hasError bool // CPU 动态指标 - if err := manager.CollectAndSendCPU(conn); err != nil { + if err := manager.CollectAndSendCPU(writer); err != nil { slog.Warn("发送CPU指标失败", "error", err) hasError = true } // 内存动态指标 - if err := manager.CollectAndSendMemory(conn); err != nil { + if err := manager.CollectAndSendMemory(writer); err != nil { slog.Warn("发送内存指标失败", "error", err) hasError = true } // 磁盘指标 - if err := manager.CollectAndSendDisk(conn); err != nil { + if err := manager.CollectAndSendDisk(writer); err != nil { slog.Warn("发送磁盘指标失败", "error", err) hasError = true } // 磁盘 IO 指标 - if err := manager.CollectAndSendDiskIO(conn); err != nil { + if err := manager.CollectAndSendDiskIO(writer); err != nil { slog.Warn("发送磁盘IO指标失败", "error", err) hasError = true } // 网络指标 - if err := manager.CollectAndSendNetwork(conn); err != nil { + if err := manager.CollectAndSendNetwork(writer); err != nil { slog.Warn("发送网络指标失败", "error", err) hasError = true } // 网络连接统计 - if err := manager.CollectAndSendNetworkConnection(conn); err != nil { + if err := manager.CollectAndSendNetworkConnection(writer); err != nil { slog.Warn("发送网络连接统计失败", "error", err) hasError = true } // 主机信息 - if err := manager.CollectAndSendHost(conn); err != nil { + if err := manager.CollectAndSendHost(writer); err != nil { slog.Warn("发送主机信息失败", "error", err) hasError = true } // GPU 信息（可选） - if err := manager.CollectAndSendGPU(conn); err != nil { + if err := manager.CollectAndSendGPU(writer); err != nil { slog.Info("发送GPU信息失败", "error", err) } // 温度信息（可选） - if err := manager.CollectAndSendTemperature(conn); err != nil { + if err := manager.CollectAndSendTemperature(writer); err != nil { slog.Info("发送温度信息失败", "error", err) } + if writer.buffered { + if conn == nil { + slog.Info("当前连接不可用，指标已写入缓存") + } else if writer.sendErr != nil { + slog.Warn("发送指标失败，已写入缓存", "error", writer.sendErr) + } + } + if hasError { return fmt.Errorf("部分指标采集失败") } diff --git a/pkg/agent/service/metrics_buffer.go b/pkg/agent/service/metrics_buffer.go new file mode 100644 index 0000000..1567e7e --- /dev/null +++ b/pkg/agent/service/metrics_buffer.go @@ -0,0 +1,183 @@ +package service + +import ( + "encoding/binary" + "encoding/json" + "fmt" + "log/slog" + "os" + "path/filepath" + "sync" + "time" + + "github.com/boltdb/bolt" + "github.com/dushixiang/pika/internal/protocol" + "github.com/dushixiang/pika/pkg/agent/collector" + "github.com/dushixiang/pika/pkg/agent/utils" +) + +const ( + metricsBufferDBName = "metrics_buffer.db" + metricsBufferBucket = "metrics_buffer" + metricsBufferTimeout = 2 * time.Second +) + +type metricsBuffer struct { + path string + mu sync.Mutex +} + +func newMetricsBuffer() *metricsBuffer { + path := filepath.Join(utils.GetSafeHomeDir(), ".pika", metricsBufferDBName) + return &metricsBuffer{path: path} +} + +func (b *metricsBuffer) Append(v interface{}) error { + payload, err := json.Marshal(v) + if err != nil { + return fmt.Errorf("序列化指标缓存失败: %w", err) + } + + b.mu.Lock() + defer b.mu.Unlock() + + db, err := b.openDB() + if err != nil { + return err + } + defer db.Close() + + if err := db.Update(func(tx *bolt.Tx) error { + bucket, err := tx.CreateBucketIfNotExists([]byte(metricsBufferBucket)) + if err != nil { + return fmt.Errorf("创建指标缓存桶失败: %w", err) + } + + seq, err := bucket.NextSequence() + if err != nil { + return fmt.Errorf("获取指标缓存序列失败: %w", err) + } + + key := make([]byte, 8) + binary.BigEndian.PutUint64(key, seq) + if err := bucket.Put(key, payload); err != nil { + return fmt.Errorf("写入指标缓存失败: %w", err) + } + + return nil + }); err != nil { + return err + } + + return nil +} + +func (b *metricsBuffer) Flush(writer collector.WebSocketWriter) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + + db, err := b.openDB() + if err != nil { + if os.IsNotExist(err) { + return 0, nil + } + return 0, err + } + defer db.Close() + + var ( + sent int + sendErr error + ) + + if err := db.Update(func(tx *bolt.Tx) error { + bucket := tx.Bucket([]byte(metricsBufferBucket)) + if bucket == nil { + return nil + } + + cursor := bucket.Cursor() + for k, v := cursor.First(); k != nil; k, v = cursor.Next() { + if sendErr != nil { + break + } + + var msg protocol.OutboundMessage + if err := json.Unmarshal(v, &msg); err != nil { + slog.Warn("缓存指标解析失败，已跳过", "error", err) + if err := cursor.Delete(); err != nil { + return fmt.Errorf("删除损坏缓存失败: %w", err) + } + continue + } + + if err := writer.WriteJSON(msg); err != nil { + sendErr = err + break + } + + if err := cursor.Delete(); err != nil { + return fmt.Errorf("删除已发送缓存失败: %w", err) + } + sent++ + } + + return nil + }); err != nil { + return sent, err + } + + if sendErr != nil { + return sent, sendErr + } + + return sent, nil +} + +func (b *metricsBuffer) openDB() (*bolt.DB, error) { + if err := os.MkdirAll(filepath.Dir(b.path), 0755); err != nil { + return nil, fmt.Errorf("创建指标缓存目录失败: %w", err) + } + + db, err := bolt.Open(b.path, 0600, &bolt.Options{Timeout: metricsBufferTimeout}) + if err != nil { + return nil, fmt.Errorf("打开指标缓存数据库失败: %w", err) + } + + return db, nil +} + +type metricsWriter struct { + conn *safeConn + buffer *metricsBuffer + buffered bool + sendErr error +} + +func newMetricsWriter(conn *safeConn, buffer *metricsBuffer) *metricsWriter { + return &metricsWriter{ + conn: conn, + buffer: buffer, + } +} + +func (w *metricsWriter) WriteJSON(v interface{}) error { + if w.conn == nil { + if err := w.buffer.Append(v); err != nil { + return err + } + w.buffered = true + return nil + } + + if err := w.conn.WriteJSON(v); err != nil { + if bufferErr := w.buffer.Append(v); bufferErr != nil { + return fmt.Errorf("写入指标缓存失败: %w", bufferErr) + } + w.buffered = true + w.sendErr = err + return nil + } + + return nil +}