Files
OpenFlare/openflare-agent/internal/agent/runner.go
T
2026-06-06 10:46:40 +08:00

700 lines
22 KiB
Go

package agent
import (
"context"
"encoding/json"
"errors"
"log/slog"
"strings"
"time"
"github.com/rain-kl/openflare/openflare-agent/internal/config"
"github.com/rain-kl/openflare/openflare-agent/internal/observability"
"github.com/rain-kl/openflare/openflare-agent/internal/protocol"
"github.com/rain-kl/openflare/openflare-agent/internal/state"
"github.com/rain-kl/openflare/openflare-agent/internal/wsclient"
)
type HeartbeatService interface {
Register(ctx context.Context, payload protocol.NodePayload) (*protocol.RegisterNodeResponse, error)
Heartbeat(ctx context.Context, payload protocol.NodePayload) (*protocol.HeartbeatResult, error)
SetToken(token string)
}
type SyncService interface {
SyncOnStartup(ctx context.Context, target *protocol.ActiveConfigMeta) error
SyncOnce(ctx context.Context, target *protocol.ActiveConfigMeta) error
ForceSyncOnce(ctx context.Context, target *protocol.ActiveConfigMeta) error
WAFIPGroupChecksums() (map[string]string, error)
ApplyWAFIPGroups(ctx context.Context, groups []protocol.WAFIPGroup) error
}
type Updater interface {
CheckAndUpdate(ctx context.Context, repo string, options UpdateOptions) error
}
type RuntimeManager interface {
CheckHealth(ctx context.Context) error
Restart(ctx context.Context) error
}
type WebSocketService interface {
Connect(ctx context.Context) (protocol.WebSocketConnection, error)
SetToken(token string)
URL() string
}
type UpdateOptions struct {
Channel string
TagName string
Force bool
}
type Runner struct {
Config *config.Config
StateStore *state.Store
ObservabilityBuffer *state.ObservabilityBufferStore
HeartbeatService HeartbeatService
SyncService SyncService
Updater Updater
RuntimeManager RuntimeManager
WebSocketService WebSocketService
autoUpdate bool
updateNow bool
updateRepo string
updateChan string
updateTag string
restartOpenrestyNow bool
websocketUpgradeEnabled bool
}
func (r *Runner) Run(ctx context.Context) error {
nodeID, err := r.StateStore.EnsureNodeID()
if err != nil {
return err
}
slog.Info("agent runner started", "node_id", nodeID, "node", r.Config.NodeName, "ip", r.Config.NodeIP)
if r.hasAccessToken() {
if _, hbErr := r.performHeartbeatCycle(ctx, nodeID, true); hbErr != nil {
slog.Error("agent startup heartbeat failed", "error", hbErr)
}
} else if err = r.tryRegister(ctx, &nodeID); err != nil {
slog.Error("agent initial discovery register failed", "error", err)
}
heartbeatTicker := time.NewTicker(r.Config.HeartbeatInterval.Duration())
defer heartbeatTicker.Stop()
var wsDone <-chan error
wsBackoff := newWebSocketBackoff()
nextWSAttempt := time.Now()
tryStartWebSocket := func() {
if wsDone != nil || !r.shouldUseWebSocket() || time.Now().Before(nextWSAttempt) {
return
}
done, startErr := r.startWebSocket(ctx, nodeID)
if startErr != nil {
delay := wsBackoff.Next()
nextWSAttempt = time.Now().Add(delay)
slog.Debug("agent ws upgrade failed; falling back to http heartbeat",
"enabled", r.websocketUpgradeEnabled,
"url", r.websocketURL(),
"retry_after", delay,
"error", startErr,
)
return
}
wsBackoff.Reset()
wsDone = done
slog.Debug("agent switched to websocket mode", "url", r.websocketURL())
}
tryStartWebSocket()
for {
select {
case <-ctx.Done():
slog.Info("agent runner shutting down", "error", ctx.Err())
return ctx.Err()
case wsErr := <-wsDone:
wsDone = nil
delay := wsBackoff.Next()
nextWSAttempt = time.Now().Add(delay)
slog.Debug("agent ws disconnected; resuming http heartbeat", "retry_after", delay, "error", wsErr)
if r.hasAccessToken() {
if _, hbErr := r.performHeartbeatCycle(ctx, nodeID, false); hbErr != nil {
slog.Error("agent heartbeat after ws disconnect failed", "error", hbErr)
}
}
case <-heartbeatTicker.C:
if wsDone != nil {
continue
}
if !r.hasAccessToken() {
if err = r.tryRegister(ctx, &nodeID); err != nil {
slog.Error("agent discovery register failed", "error", err)
}
continue
}
if changed, hbErr := r.performHeartbeatCycle(ctx, nodeID, false); hbErr != nil {
slog.Error("agent heartbeat failed", "error", hbErr)
} else {
if changed {
heartbeatTicker.Reset(r.Config.HeartbeatInterval.Duration())
}
tryStartWebSocket()
}
}
}
}
func (r *Runner) performHeartbeatCycle(ctx context.Context, nodeID string, startup bool) (bool, error) {
r.refreshOpenrestyHealth(ctx)
payload, ackWindows := r.prepareHeartbeatPayload(nodeID)
heartbeatResult, err := r.HeartbeatService.Heartbeat(ctx, payload)
if err != nil {
return false, err
}
r.ackObservabilityWindows(ackWindows)
if heartbeatResult == nil {
heartbeatResult = &protocol.HeartbeatResult{}
}
mode := "periodic"
if startup {
mode = "startup"
}
slog.Debug("agent heartbeat succeeded", "mode", mode, "node_id", nodeID)
changed := r.applySettings(heartbeatResult.AgentSettings)
r.applyWAFIPGroups(ctx, heartbeatResult.WAFIPGroups)
if startup {
if err = r.SyncService.SyncOnStartup(ctx, heartbeatResult.ActiveConfig); err != nil {
r.recordSyncError(err)
slog.Error("agent startup sync failed", "error", err)
} else {
slog.Debug("agent startup sync completed")
}
} else if err = r.SyncService.SyncOnce(ctx, heartbeatResult.ActiveConfig); err != nil {
r.recordSyncError(err)
slog.Error("agent sync failed", "error", err)
}
r.tryRestartOpenresty(ctx)
r.tryAutoUpdate(ctx)
return changed, nil
}
func (r *Runner) shouldUseWebSocket() bool {
enabled := r.WebSocketService != nil && r.websocketUpgradeEnabled && r.hasAccessToken()
slog.Debug("agent ws upgrade eligibility checked", "enabled", enabled, "server_enabled", r.websocketUpgradeEnabled, "url", r.websocketURL())
return enabled
}
func (r *Runner) websocketURL() string {
if r.WebSocketService == nil {
return ""
}
return r.WebSocketService.URL()
}
func (r *Runner) startWebSocket(ctx context.Context, nodeID string) (<-chan error, error) {
if r.WebSocketService == nil {
return nil, errors.New("websocket service is not configured")
}
conn, err := r.WebSocketService.Connect(ctx)
if err != nil {
return nil, err
}
done := make(chan error, 1)
go func() {
defer func() {
_ = conn.Close()
}()
done <- r.runWebSocket(ctx, nodeID, conn)
}()
return done, nil
}
type agentWSHandler struct {
runner *Runner
conn protocol.WebSocketConnection
nodeID string
statusTicker *time.Ticker
}
func (h *agentWSHandler) OnConnect(ctx context.Context) error {
return h.runner.sendWebSocketStatus(ctx, h.nodeID, h.conn)
}
func (h *agentWSHandler) HandleMessage(ctx context.Context, msg wsclient.WSMessage) error {
var payloadBytes []byte
if msg.Payload != nil {
payloadBytes = []byte(msg.Payload)
}
protoMsg := protocol.WSMessage{
Type: msg.Type,
Payload: payloadBytes,
}
changed, err := h.runner.handleWebSocketMessage(ctx, protoMsg, h.conn)
if err != nil {
return err
}
if changed {
h.statusTicker.Reset(h.runner.Config.HeartbeatInterval.Duration())
}
return nil
}
func (h *agentWSHandler) OnClose(err error) {
slog.Error("agent ws receive failed", "error", err)
}
func (r *Runner) runWebSocket(ctx context.Context, nodeID string, conn protocol.WebSocketConnection) error {
slog.Debug("agent ws connected", "url", conn.URL(), "node_id", nodeID)
statusTicker := time.NewTicker(r.Config.HeartbeatInterval.Duration())
defer statusTicker.Stop()
childCtx, cancel := context.WithCancel(ctx)
defer cancel()
// Start status ticker sender in background
go func() {
for {
select {
case <-childCtx.Done():
return
case <-statusTicker.C:
if err := r.sendWebSocketStatus(childCtx, nodeID, conn); err != nil {
slog.Error("agent ws send status failed", "error", err)
_ = conn.Close()
return
}
}
}
}()
wsConn, ok := conn.(*wsclient.Connection)
if !ok {
return errors.New("invalid websocket connection type")
}
return wsConn.RunReceiveLoop(childCtx, &agentWSHandler{
runner: r,
conn: conn,
nodeID: nodeID,
statusTicker: statusTicker,
})
}
func (r *Runner) sendWebSocketStatus(ctx context.Context, nodeID string, conn protocol.WebSocketConnection) error {
r.refreshOpenrestyHealth(ctx)
payload, ackWindows := r.prepareHeartbeatPayload(nodeID)
if err := conn.SendStatus(payload); err != nil {
return err
}
r.ackObservabilityWindows(ackWindows)
return nil
}
func (r *Runner) handleWebSocketMessage(ctx context.Context, message protocol.WSMessage, conn protocol.WebSocketConnection) (bool, error) {
switch message.Type {
case protocol.WSMessageTypeSettings:
var settings protocol.AgentSettings
if err := json.Unmarshal(message.Payload, &settings); err != nil {
slog.Debug("agent ws settings decode failed", "error", err)
return false, nil
}
changed := r.applySettings(&settings)
r.tryRestartOpenresty(ctx)
r.tryAutoUpdate(ctx)
if !r.websocketUpgradeEnabled {
slog.Debug("agent ws disabled by server settings; falling back to http heartbeat")
return changed, errors.New("websocket upgrade disabled by server")
}
return changed, nil
case protocol.WSMessageTypeActiveConfig:
var target protocol.ActiveConfigMeta
if err := json.Unmarshal(message.Payload, &target); err != nil {
slog.Debug("agent ws active config decode failed", "error", err)
return false, nil
}
slog.Debug("agent ws active config received", "version", target.Version, "checksum", target.Checksum, "trigger_sync", true)
if err := r.SyncService.SyncOnce(ctx, &target); err != nil {
r.recordSyncError(err)
slog.Error("agent ws triggered sync failed", "version", target.Version, "error", err)
}
return false, nil
case protocol.WSMessageTypeForceSyncConfig:
var target protocol.ActiveConfigMeta
if err := json.Unmarshal(message.Payload, &target); err != nil {
slog.Debug("agent ws force sync config decode failed", "error", err)
return false, nil
}
slog.Debug("agent ws force sync config received", "version", target.Version, "checksum", target.Checksum, "trigger_sync", true)
if err := r.SyncService.ForceSyncOnce(ctx, &target); err != nil {
r.recordSyncError(err)
slog.Error("agent ws triggered force sync failed", "version", target.Version, "error", err)
}
return false, nil
case protocol.WSMessageTypeWAFIPGroups:
var groups []protocol.WAFIPGroup
if err := json.Unmarshal(message.Payload, &groups); err != nil {
slog.Debug("agent ws waf ip groups decode failed", "error", err)
return false, nil
}
r.applyWAFIPGroups(ctx, groups)
return false, nil
case protocol.WSMessageTypePing:
slog.Debug("agent ws ping received")
return false, conn.SendPong()
case protocol.WSMessageTypePong:
slog.Debug("agent ws pong received")
return false, nil
default:
slog.Debug("agent ws unsupported message type", "type", message.Type)
return false, nil
}
}
type webSocketBackoff struct {
delays []time.Duration
index int
}
func newWebSocketBackoff() *webSocketBackoff {
return &webSocketBackoff{
delays: []time.Duration{
time.Second,
2 * time.Second,
5 * time.Second,
10 * time.Second,
30 * time.Second,
},
}
}
func (backoff *webSocketBackoff) Next() time.Duration {
if backoff == nil || len(backoff.delays) == 0 {
return 30 * time.Second
}
if backoff.index >= len(backoff.delays) {
return backoff.delays[len(backoff.delays)-1]
}
delay := backoff.delays[backoff.index]
backoff.index++
return delay
}
func (backoff *webSocketBackoff) Reset() {
if backoff != nil {
backoff.index = 0
}
}
func (r *Runner) hasAccessToken() bool {
return strings.TrimSpace(r.Config.AccessToken) != ""
}
func (r *Runner) applySettings(settings *protocol.AgentSettings) bool {
if settings == nil {
return false
}
changed := false
if settings.HeartbeatInterval > 0 {
newInterval := config.MillisecondDuration(time.Duration(settings.HeartbeatInterval) * time.Millisecond)
if newInterval != r.Config.HeartbeatInterval {
slog.Info("agent heartbeat interval updated", "from", r.Config.HeartbeatInterval, "to", newInterval)
r.Config.HeartbeatInterval = newInterval
changed = true
}
}
if settings.WebsocketUpgradeEnabled != r.websocketUpgradeEnabled {
slog.Debug("agent websocket upgrade setting updated", "from", r.websocketUpgradeEnabled, "to", settings.WebsocketUpgradeEnabled)
}
r.websocketUpgradeEnabled = settings.WebsocketUpgradeEnabled
r.autoUpdate = settings.AutoUpdate
r.updateNow = settings.UpdateNow
r.updateRepo = strings.TrimSpace(settings.UpdateRepo)
r.updateChan = strings.TrimSpace(settings.UpdateChannel)
r.updateTag = strings.TrimSpace(settings.UpdateTag)
r.restartOpenrestyNow = settings.RestartOpenrestyNow
return changed
}
func (r *Runner) tryRestartOpenresty(ctx context.Context) {
if !r.restartOpenrestyNow {
return
}
r.restartOpenrestyNow = false
if r.RuntimeManager == nil {
return
}
slog.Info("agent openresty restart requested by server")
if err := r.RuntimeManager.Restart(ctx); err != nil {
slog.Error("agent openresty restart failed", "error", err)
r.recordOpenrestyUnhealthy(err, false)
return
}
slog.Info("agent openresty restart succeeded")
r.recordOpenrestyHealthy()
}
func (r *Runner) tryAutoUpdate(ctx context.Context) {
force := r.updateNow
shouldCheck := r.autoUpdate || force
r.updateNow = false
r.updateTag = strings.TrimSpace(r.updateTag)
if !shouldCheck || r.Updater == nil || r.updateRepo == "" {
return
}
channel := "stable"
if force && r.updateChan != "" {
channel = r.updateChan
}
if err := r.Updater.CheckAndUpdate(ctx, r.updateRepo, UpdateOptions{
Channel: channel,
TagName: r.updateTag,
Force: force,
}); err != nil {
slog.Error("agent update check failed", "error", err)
}
if force {
r.updateTag = ""
r.updateChan = ""
}
}
func (r *Runner) tryRegister(ctx context.Context, nodeID *string) error {
if strings.TrimSpace(r.Config.DiscoveryToken) == "" {
return errors.New("agent_token 为空且未配置 discovery_token")
}
slog.Info("agent discovery registration started")
response, err := r.HeartbeatService.Register(ctx, r.nodePayload(*nodeID))
if err != nil {
return err
}
if response == nil || strings.TrimSpace(response.AccessToken) == "" || strings.TrimSpace(response.NodeID) == "" {
return errors.New("discovery register response 缺少 node_id 或 agent_token")
}
snapshot, err := r.StateStore.Load()
if err != nil {
return err
}
snapshot.NodeID = response.NodeID
if err = r.StateStore.Save(snapshot); err != nil {
return err
}
r.Config.AccessToken = response.AccessToken
r.Config.DiscoveryToken = ""
if err = r.Config.Save(); err != nil {
return err
}
r.HeartbeatService.SetToken(response.AccessToken)
if r.WebSocketService != nil {
r.WebSocketService.SetToken(response.AccessToken)
}
*nodeID = response.NodeID
slog.Info("agent discovery registration succeeded", "node_id", response.NodeID)
r.refreshOpenrestyHealth(ctx)
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{}
}
r.applySettings(heartbeatResult.AgentSettings)
r.applyWAFIPGroups(ctx, heartbeatResult.WAFIPGroups)
if err = r.SyncService.SyncOnStartup(ctx, heartbeatResult.ActiveConfig); err != nil {
r.recordSyncError(err)
slog.Error("agent post-register startup sync failed", "error", err)
} else {
slog.Debug("agent post-register startup sync completed")
}
r.tryRestartOpenresty(ctx)
r.tryAutoUpdate(ctx)
return nil
}
func (r *Runner) recordSyncError(err error) {
if err == nil || r.StateStore == nil {
return
}
snapshot, loadErr := r.StateStore.Load()
if loadErr != nil {
slog.Error("load state before recording sync error failed", "error", loadErr)
return
}
snapshot.LastError = err.Error()
slog.Warn("recording sync error into state", "error", snapshot.LastError)
if saveErr := r.StateStore.Save(snapshot); saveErr != nil {
slog.Error("save state after sync error failed", "error", saveErr)
}
}
func (r *Runner) refreshOpenrestyHealth(ctx context.Context) {
if r.RuntimeManager == nil || r.StateStore == nil {
return
}
if err := r.RuntimeManager.CheckHealth(ctx); err != nil {
if strings.Contains(err.Error(), "openresty config not exists") {
return
}
r.recordOpenrestyUnhealthy(err, true)
return
}
r.recordOpenrestyHealthy()
}
func (r *Runner) recordOpenrestyHealthy() {
if r.StateStore == nil {
return
}
snapshot, err := r.StateStore.Load()
if err != nil {
slog.Error("load state before recording openresty health failed", "error", err)
return
}
if snapshot.OpenrestyStatus == protocol.OpenrestyStatusHealthy && strings.TrimSpace(snapshot.OpenrestyMessage) == "" {
return
}
snapshot.OpenrestyStatus = protocol.OpenrestyStatusHealthy
snapshot.OpenrestyMessage = ""
if err = r.StateStore.Save(snapshot); err != nil {
slog.Error("save state after recording openresty health failed", "error", err)
}
}
func (r *Runner) recordOpenrestyUnhealthy(err error, fallbackOnly bool) {
if err == nil || r.StateStore == nil {
return
}
snapshot, loadErr := r.StateStore.Load()
if loadErr != nil {
slog.Error("load state before recording openresty error failed", "error", loadErr)
return
}
message := strings.TrimSpace(err.Error())
if !fallbackOnly || strings.TrimSpace(snapshot.OpenrestyMessage) == "" {
snapshot.OpenrestyMessage = message
}
snapshot.OpenrestyStatus = protocol.OpenrestyStatusUnhealthy
if saveErr := r.StateStore.Save(snapshot); saveErr != nil {
slog.Error("save state after recording openresty error failed", "error", saveErr)
}
}
func (r *Runner) nodePayload(nodeID string) protocol.NodePayload {
snapshot, _ := r.StateStore.Load()
openrestyStatus := strings.TrimSpace(snapshot.OpenrestyStatus)
if openrestyStatus == "" {
openrestyStatus = protocol.OpenrestyStatusUnknown
}
profile := observability.BuildProfile(r.Config, r.StateStore)
managedOpenRestyMetrics := observability.CollectManagedOpenRestyMetrics(r.Config)
trafficReport, accessLogs, fallbackMetrics := observability.BuildTrafficObservability(r.Config, r.StateStore, managedOpenRestyMetrics)
if managedOpenRestyMetrics == nil {
managedOpenRestyMetrics = fallbackMetrics
}
metricSnapshot := observability.BuildSnapshot(r.Config, r.StateStore)
openrestyObservation := observability.BuildOpenrestyObservation(managedOpenRestyMetrics)
healthEvents := observability.BuildHealthEvents(snapshot)
payload := protocol.NodePayload{
NodeID: nodeID,
Name: r.Config.NodeName,
IP: r.Config.NodeIP,
Version: r.Config.Version,
ExtVersion: r.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 r.SyncService != nil {
checksums, err := r.SyncService.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
}
func (r *Runner) applyWAFIPGroups(ctx context.Context, groups []protocol.WAFIPGroup) {
if len(groups) == 0 || r.SyncService == nil {
return
}
if err := r.SyncService.ApplyWAFIPGroups(ctx, groups); err != nil {
r.recordSyncError(err)
slog.Error("agent apply waf ip groups failed", "error", err)
}
}
func (r *Runner) prepareHeartbeatPayload(nodeID string) (protocol.NodePayload, []int64) {
payload := r.nodePayload(nodeID)
if r.ObservabilityBuffer == nil || (payload.Snapshot == nil && payload.TrafficReport == nil && len(payload.AccessLogs) == 0) {
return payload, nil
}
now := time.Now().UTC()
retainAfterUnix := now.Add(-time.Duration(r.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 := 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,
OpenrestyObservation: item.OpenrestyObservation,
TrafficReport: item.TrafficReport,
AccessLogs: item.AccessLogs,
})
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)
}
}