mirror of
https://github.com/Rain-kl/OpenFlare.git
synced 2026-10-04 23:16:37 +08:00
[优化] go 引用调整
This commit is contained in:
ryan
@@ -0,0 +1,699 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user