feat: add observability buffer for heartbeat data

- Introduced `ObservabilityBufferStore` to manage buffered observability records. - Implemented methods for upserting, replaying, and acknowledging records in the buffer. - Enhanced `AgentNodePayload` and `NodePayload` to include buffered observability data. - Updated tests to cover new functionality for observability buffering. - Modified existing services to persist and handle buffered observability data during heartbeats. - Added configuration options for observability buffer path and replay minutes. - Updated documentation to reflect new observability features and configurations.
2026-10-08 16:46:37 +08:00 · 2026-03-14 17:57:50 +08:00
parent 7628397785
commit e25b41fd75
18 changed files with 798 additions and 64 deletions
@@ -40,12 +40,13 @@ type UpdateOptions struct {
 }

 type Runner struct {
-	Config           *config.Config
-	StateStore       *state.Store
-	HeartbeatService HeartbeatService
-	SyncService      SyncService
-	Updater          Updater
-	RuntimeManager   RuntimeManager
+	Config              *config.Config
+	StateStore          *state.Store
+	ObservabilityBuffer *state.ObservabilityBufferStore
+	HeartbeatService    HeartbeatService
+	SyncService         SyncService
+	Updater             Updater
+	RuntimeManager      RuntimeManager

 	autoUpdate          bool
 	updateNow           bool
@@ -63,10 +64,12 @@ func (r *Runner) Run(ctx context.Context) error {
 	slog.Info("agent runner started", "node_id", nodeID, "node", r.Config.NodeName, "ip", r.Config.NodeIP)
 	if r.hasAgentToken() {
 		r.refreshOpenrestyHealth(ctx)
-		heartbeatResult, hbErr := r.HeartbeatService.Heartbeat(ctx, r.nodePayload(nodeID))
+		payload, ackWindows := r.prepareHeartbeatPayload(nodeID)
+		heartbeatResult, hbErr := r.HeartbeatService.Heartbeat(ctx, payload)
 		if hbErr != nil {
 			slog.Error("agent startup heartbeat failed", "error", hbErr)
 		} else {
+			r.ackObservabilityWindows(ackWindows)
 			if heartbeatResult == nil {
 				heartbeatResult = &protocol.HeartbeatResult{}
 			}
@@ -101,10 +104,12 @@ func (r *Runner) Run(ctx context.Context) error {
 				continue
 			}
 			r.refreshOpenrestyHealth(ctx)
-			heartbeatResult, hbErr := r.HeartbeatService.Heartbeat(ctx, r.nodePayload(nodeID))
+			payload, ackWindows := r.prepareHeartbeatPayload(nodeID)
+			heartbeatResult, hbErr := r.HeartbeatService.Heartbeat(ctx, payload)
 			if hbErr != nil {
 				slog.Error("agent heartbeat failed", "error", hbErr)
 			} else {
+				r.ackObservabilityWindows(ackWindows)
 				if heartbeatResult == nil {
 					heartbeatResult = &protocol.HeartbeatResult{}
 				}
@@ -220,11 +225,13 @@ func (r *Runner) tryRegister(ctx context.Context, nodeID *string) error {
 	*nodeID = response.NodeID
 	slog.Info("agent discovery registration succeeded", "node_id", response.NodeID)
 	r.refreshOpenrestyHealth(ctx)
-	heartbeatResult, heartbeatErr := r.HeartbeatService.Heartbeat(ctx, r.nodePayload(*nodeID))
+	payload, ackWindows := r.prepareHeartbeatPayload(*nodeID)
+	heartbeatResult, heartbeatErr := r.HeartbeatService.Heartbeat(ctx, payload)
 	if heartbeatErr != nil {
 		slog.Error("agent post-register heartbeat failed", "error", heartbeatErr)
 		return nil
 	}
+	r.ackObservabilityWindows(ackWindows)
 	if heartbeatResult == nil {
 		heartbeatResult = &protocol.HeartbeatResult{}
 	}
@@ -332,3 +339,60 @@ func (r *Runner) nodePayload(nodeID string) protocol.NodePayload {
 		HealthEvents:     healthEvents,
 	}
 }
+
+func (r *Runner) prepareHeartbeatPayload(nodeID string) (protocol.NodePayload, []int64) {
+	payload := r.nodePayload(nodeID)
+	if r.ObservabilityBuffer == nil || payload.Snapshot == nil {
+		return payload, nil
+	}
+	now := time.Now().UTC()
+	retainAfterUnix := now.Add(-time.Duration(r.Config.ObservabilityReplayMinutes) * time.Minute).Unix()
+	windowStartedAtUnix := state.ObservabilityWindowStartedAt(payload.Snapshot, payload.TrafficReport)
+	if windowStartedAtUnix <= 0 {
+		return payload, nil
+	}
+
+	record := state.ObservabilityBufferRecord{
+		WindowStartedAtUnix: windowStartedAtUnix,
+		Snapshot:            payload.Snapshot,
+		TrafficReport:       payload.TrafficReport,
+		QueuedAtUnix:        now.Unix(),
+	}
+	if err := r.ObservabilityBuffer.Upsert(record, retainAfterUnix); err != nil {
+		slog.Error("upsert observability buffer failed", "error", err)
+		return payload, nil
+	}
+
+	records, err := r.ObservabilityBuffer.Replayable(windowStartedAtUnix, retainAfterUnix)
+	if err != nil {
+		slog.Error("load replayable observability buffer failed", "error", err)
+		return payload, []int64{windowStartedAtUnix}
+	}
+
+	ackWindows := make([]int64, 0, len(records)+1)
+	buffered := make([]protocol.BufferedObservabilityRecord, 0, len(records))
+	for _, item := range records {
+		if item.WindowStartedAtUnix <= 0 {
+			continue
+		}
+		buffered = append(buffered, protocol.BufferedObservabilityRecord{
+			WindowStartedAtUnix: item.WindowStartedAtUnix,
+			Snapshot:            item.Snapshot,
+			TrafficReport:       item.TrafficReport,
+		})
+		ackWindows = append(ackWindows, item.WindowStartedAtUnix)
+	}
+	payload.BufferedObservability = buffered
+	ackWindows = append(ackWindows, windowStartedAtUnix)
+	return payload, ackWindows
+}
+
+func (r *Runner) ackObservabilityWindows(windowStartedAtUnix []int64) {
+	if r.ObservabilityBuffer == nil || len(windowStartedAtUnix) == 0 {
+		return
+	}
+	retainAfterUnix := time.Now().UTC().Add(-time.Duration(r.Config.ObservabilityReplayMinutes) * time.Minute).Unix()
+	if err := r.ObservabilityBuffer.Ack(windowStartedAtUnix, retainAfterUnix); err != nil {
+		slog.Error("ack observability buffer failed", "error", err)
+	}
+}
@@ -311,7 +311,7 @@ func TestRunnerHeartbeatPayloadIncludesObservabilityExtensions(t *testing.T) {
 	}
 	if err := os.WriteFile(
 		filepath.Join(filepath.Dir(runner.Config.RouteConfigPath), "atsflare_access.log"),
-		[]byte("{\"ts\":\"2026-03-14T10:00:00Z\",\"host\":\"edge.example.com\",\"remote_addr\":\"10.0.0.8\",\"status\":200}\n"),
+		[]byte("{\"ts\":\""+time.Now().UTC().Format(time.RFC3339)+"\",\"host\":\"edge.example.com\",\"remote_addr\":\"10.0.0.8\",\"status\":200}\n"),
 		0o644,
 	); err != nil {
 		t.Fatalf("failed to prepare access log: %v", err)
@@ -343,6 +343,81 @@ func TestRunnerHeartbeatPayloadIncludesObservabilityExtensions(t *testing.T) {
 	}
 }

+func TestRunnerReplaysBufferedObservabilityAfterHeartbeatRecovery(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	tempDir := t.TempDir()
+	stateStore := state.NewStore(filepath.Join(tempDir, "state.json"))
+	bufferStore := state.NewObservabilityBufferStore(filepath.Join(tempDir, "observability-buffer.json"))
+	nowUnix := time.Now().UTC().Unix()
+	bufferWindow := nowUnix - (nowUnix % 60) - 60
+	if err := bufferStore.Upsert(state.ObservabilityBufferRecord{
+		WindowStartedAtUnix: bufferWindow,
+		Snapshot:            &protocol.NodeMetricSnapshot{CapturedAtUnix: bufferWindow + 5, CPUUsagePercent: 30},
+		TrafficReport:       &protocol.NodeTrafficReport{WindowStartedAtUnix: bufferWindow, WindowEndedAtUnix: bufferWindow + 60, RequestCount: 8},
+		QueuedAtUnix:        bufferWindow + 60,
+	}, 0); err != nil {
+		t.Fatalf("failed to seed observability buffer: %v", err)
+	}
+	heartbeatService := &fakeHeartbeatService{
+		heartbeatErrs:    []error{errors.New("server offline"), nil},
+		heartbeatResults: []*protocol.HeartbeatResult{{}, {}},
+		onHeartbeat: func(callCount int) {
+			if callCount >= 2 {
+				cancel()
+			}
+		},
+	}
+	runner := &Runner{
+		Config: &config.Config{
+			AgentToken:                 "agent-token",
+			NodeName:                   "edge-buffer-01",
+			NodeIP:                     "10.0.0.52",
+			AgentVersion:               config.AgentVersion,
+			NginxVersion:               "1.27.1.2",
+			DataDir:                    tempDir,
+			RouteConfigPath:            filepath.Join(tempDir, "conf.d", "atsflare_routes.conf"),
+			HeartbeatInterval:          config.MillisecondDuration(10 * time.Millisecond),
+			ObservabilityReplayMinutes: 15,
+		},
+		StateStore:          stateStore,
+		ObservabilityBuffer: bufferStore,
+		HeartbeatService:    heartbeatService,
+		SyncService:         &fakeSyncService{},
+	}
+	if err := os.MkdirAll(filepath.Dir(runner.Config.RouteConfigPath), 0o755); err != nil {
+		t.Fatalf("failed to prepare route config dir: %v", err)
+	}
+	if err := os.WriteFile(
+		filepath.Join(filepath.Dir(runner.Config.RouteConfigPath), "atsflare_access.log"),
+		[]byte("{\"ts\":\""+time.Now().UTC().Format(time.RFC3339)+"\",\"host\":\"edge.example.com\",\"remote_addr\":\"10.0.0.8\",\"status\":200}\n"),
+		0o644,
+	); err != nil {
+		t.Fatalf("failed to prepare access log: %v", err)
+	}
+
+	runErr := runner.Run(ctx)
+	if runErr != context.Canceled {
+		t.Fatalf("expected run to stop by context cancellation, got %v", runErr)
+	}
+	if len(heartbeatService.heartbeatPayloads) != 2 {
+		t.Fatalf("expected two heartbeat payloads, got %d", len(heartbeatService.heartbeatPayloads))
+	}
+	secondPayload := heartbeatService.heartbeatPayloads[1]
+	if len(secondPayload.BufferedObservability) != 1 {
+		t.Fatalf("expected second heartbeat to replay one buffered observation, got %+v", secondPayload.BufferedObservability)
+	}
+
+	replayable, err := bufferStore.Replayable(0, 0)
+	if err != nil {
+		t.Fatalf("Replayable after recovery failed: %v", err)
+	}
+	if len(replayable) != 0 {
+		t.Fatalf("expected buffer to be acked after successful heartbeat, got %+v", replayable)
+	}
+}
+
 func TestRunnerDiscoveryRegisterUpdatesTokenAndNodeID(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()