OpenFlare/internal/apps/agent/heartbeat/cycle.go

// Package heartbeat implements the periodic heartbeat cycle executed by the agent,
// including payload preparation, config sync, WAF IP group application, and observability buffering.
package heartbeat

import (
	"context"
	"log/slog"
	"strings"
	"time"

	"github.com/Rain-kl/Wavelet/internal/apps/agent/config"
	"github.com/Rain-kl/Wavelet/internal/apps/agent/observability"
	"github.com/Rain-kl/Wavelet/internal/apps/agent/protocol"
	"github.com/Rain-kl/Wavelet/internal/apps/agent/state"
	"github.com/Rain-kl/Wavelet/internal/apps/agent/updater"
	edgeheartbeat "github.com/Rain-kl/Wavelet/internal/apps/edge/heartbeat"
	"github.com/Rain-kl/Wavelet/internal/apps/edge/nodeip"
)

// SyncService is the interface used by Cycle to sync active configuration and WAF IP groups.
type SyncService interface {
	SyncOnStartup(ctx context.Context, target *protocol.ActiveConfigMeta) error
	SyncOnce(ctx context.Context, target *protocol.ActiveConfigMeta) error
	WAFIPGroupChecksums() (map[string]string, error)
	ApplyWAFIPGroups(ctx context.Context, groups []protocol.WAFIPGroup) error
}

// SettingsApplier is the interface used by Cycle to apply agent settings received from the server.
type SettingsApplier interface {
	Apply(settings *protocol.AgentSettings) (intervalChanged bool)
	RestartOpenrestyIfNeeded(ctx context.Context)
}

// Cycle holds the dependencies required to execute a single agent heartbeat cycle.
type Cycle struct {
	Config              *config.Config
	StateStore          *state.Store
	ObservabilityBuffer *state.ObservabilityBufferStore
	Heartbeat           API
	Sync                SyncService
	Updater             *updater.Service
	RecordSyncError     func(err error)
}

// Perform executes one complete heartbeat cycle: sends the heartbeat, syncs config, and applies settings.
func (c *Cycle) Perform(ctx context.Context, nodeID string, startup bool, settings SettingsApplier) (bool, error) {
	payload, ackWindows := c.PrepareHeartbeatPayload(ctx, nodeID)
	heartbeatResult, err := c.Heartbeat.Heartbeat(ctx, payload)
	if err != nil {
		return false, err
	}
	c.AckObservabilityWindows(ackWindows)
	if heartbeatResult == nil {
		heartbeatResult = &protocol.HeartbeatResult{}
	}
	mode := "periodic"
	if startup {
		mode = "startup"
	}
	slog.Debug("agent heartbeat succeeded", "mode", mode, "node_id", nodeID)

	var changed bool
	if settings != nil {
		changed = settings.Apply(heartbeatResult.AgentSettings)
	}
	c.ApplyWAFIPGroups(ctx, heartbeatResult.WAFIPGroups)
	if startup {
		if err = c.Sync.SyncOnStartup(ctx, heartbeatResult.ActiveConfig); err != nil {
			c.recordSyncError(err)
			slog.Error("agent startup sync failed", "error", err)
		} else {
			slog.Debug("agent startup sync completed")
		}
	} else if err = c.Sync.SyncOnce(ctx, heartbeatResult.ActiveConfig); err != nil {
		c.recordSyncError(err)
		slog.Error("agent sync failed", "error", err)
	}
	if settings != nil {
		settings.RestartOpenrestyIfNeeded(ctx)
	}
	edgeheartbeat.TryAutoUpdate(ctx, c.Updater, agentSettingsToAutoUpdate(heartbeatResult.AgentSettings), "agent")
	return changed, nil
}

// NodePayload builds and returns the full NodePayload to be sent in a heartbeat request.
func (c *Cycle) NodePayload(ctx context.Context, nodeID string) protocol.NodePayload {
	snapshot, _ := c.StateStore.Load()
	openrestyStatus := strings.TrimSpace(snapshot.OpenrestyStatus)
	if openrestyStatus == "" {
		openrestyStatus = protocol.OpenrestyStatusUnknown
	}
	profile := observability.BuildProfile(c.Config, c.StateStore)
	managedOpenRestyMetrics := observability.CollectManagedOpenRestyMetrics(ctx, c.Config)
	trafficReport, accessLogs, fallbackMetrics := observability.BuildTrafficObservability(c.Config, c.StateStore, managedOpenRestyMetrics)
	if managedOpenRestyMetrics == nil {
		managedOpenRestyMetrics = fallbackMetrics
	}
	metricSnapshot := observability.BuildSnapshot(c.Config, c.StateStore)
	openrestyObservation := observability.BuildOpenrestyObservation(managedOpenRestyMetrics)
	healthEvents := observability.BuildHealthEvents(snapshot)

	ip := c.Config.NodeIP
	if !c.Config.NodeIPConfigured {
		ip = nodeip.DetectWithContext(ctx)
	}

	payload := protocol.NodePayload{
		NodeID:               nodeID,
		Name:                 c.Config.NodeName,
		IP:                   ip,
		Version:              c.Config.Version,
		ExtVersion:           c.Config.ExtVersion,
		CurrentVersion:       snapshot.CurrentVersion,
		LastError:            snapshot.LastError,
		OpenrestyStatus:      openrestyStatus,
		OpenrestyMessage:     snapshot.OpenrestyMessage,
		Profile:              profile,
		Snapshot:             metricSnapshot,
		OpenrestyObservation: openrestyObservation,
		TrafficReport:        trafficReport,
		AccessLogs:           accessLogs,
		HealthEvents:         healthEvents,
	}
	if c.Sync != nil {
		checksums, err := c.Sync.WAFIPGroupChecksums()
		if err != nil {
			slog.Debug("load local waf ip group checksums failed", "error", err)
		} else if len(checksums) > 0 {
			payload.WAFIPGroupChecksums = checksums
		}
	}
	return payload
}

// PrepareHeartbeatPayload constructs the heartbeat payload with buffered observability records and returns the window timestamps to acknowledge.
func (c *Cycle) PrepareHeartbeatPayload(ctx context.Context, nodeID string) (protocol.NodePayload, []int64) {
	payload := c.NodePayload(ctx, nodeID)
	if c.ObservabilityBuffer == nil || (payload.Snapshot == nil && payload.TrafficReport == nil && len(payload.AccessLogs) == 0) {
		return payload, nil
	}
	now := time.Now().UTC()
	retainAfterUnix := now.Add(-time.Duration(c.Config.ObservabilityReplayMinutes) * time.Minute).Unix()
	windowStartedAtUnix := state.ObservabilityWindowStartedAt(payload.Snapshot, payload.OpenrestyObservation, payload.TrafficReport)
	if windowStartedAtUnix <= 0 {
		return payload, nil
	}

	record := state.ObservabilityBufferRecord{
		WindowStartedAtUnix:  windowStartedAtUnix,
		Snapshot:             payload.Snapshot,
		OpenrestyObservation: payload.OpenrestyObservation,
		TrafficReport:        payload.TrafficReport,
		AccessLogs:           payload.AccessLogs,
		QueuedAtUnix:         now.Unix(),
	}
	if err := c.ObservabilityBuffer.Upsert(record, retainAfterUnix); err != nil {
		slog.Error("upsert observability buffer failed", "error", err)
		return payload, nil
	}

	records, err := c.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,
			OpenrestyObservation: item.OpenrestyObservation,
			TrafficReport:        item.TrafficReport,
			AccessLogs:           item.AccessLogs,
		})
		ackWindows = append(ackWindows, item.WindowStartedAtUnix)
	}
	payload.BufferedObservability = buffered
	ackWindows = append(ackWindows, windowStartedAtUnix)
	return payload, ackWindows
}

// AckObservabilityWindows acknowledges the given observability window timestamps in the buffer store.
func (c *Cycle) AckObservabilityWindows(windowStartedAtUnix []int64) {
	if c.ObservabilityBuffer == nil || len(windowStartedAtUnix) == 0 {
		return
	}
	retainAfterUnix := time.Now().UTC().Add(-time.Duration(c.Config.ObservabilityReplayMinutes) * time.Minute).Unix()
	if err := c.ObservabilityBuffer.Ack(windowStartedAtUnix, retainAfterUnix); err != nil {
		slog.Error("ack observability buffer failed", "error", err)
	}
}

// ApplyWAFIPGroups applies the WAF IP groups received from the server via the SyncService.
func (c *Cycle) ApplyWAFIPGroups(ctx context.Context, groups []protocol.WAFIPGroup) {
	if len(groups) == 0 || c.Sync == nil {
		return
	}
	if err := c.Sync.ApplyWAFIPGroups(ctx, groups); err != nil {
		c.recordSyncError(err)
		slog.Error("agent apply waf ip groups failed", "error", err)
	}
}

func (c *Cycle) recordSyncError(err error) {
	if c.RecordSyncError != nil {
		c.RecordSyncError(err)
	}
}

// AgentSettingsToAutoUpdate converts AgentSettings to an AutoUpdateSettings value used by the edge heartbeat updater.
func AgentSettingsToAutoUpdate(settings *protocol.AgentSettings) *edgeheartbeat.AutoUpdateSettings {
	if settings == nil {
		return nil
	}
	return &edgeheartbeat.AutoUpdateSettings{
		AutoUpdate:    settings.AutoUpdate,
		UpdateNow:     settings.UpdateNow,
		UpdateRepo:    settings.UpdateRepo,
		UpdateChannel: settings.UpdateChannel,
		UpdateTag:     settings.UpdateTag,
	}
}

func agentSettingsToAutoUpdate(settings *protocol.AgentSettings) *edgeheartbeat.AutoUpdateSettings {
	return AgentSettingsToAutoUpdate(settings)
}