diff --git a/client/internal/connect.go b/client/internal/connect.go index d038a6ef2..eff2c9489 100644 --- a/client/internal/connect.go +++ b/client/internal/connect.go @@ -403,6 +403,7 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan StateManager: stateManager, UpdateManager: c.updateManager, ClientMetrics: c.clientMetrics, + MetricsCtx: c.ctx, }, mobileDependency) engine.SetSyncResponsePersistence(c.persistSyncResponse) c.engine = engine diff --git a/client/internal/engine.go b/client/internal/engine.go index 1ee4a2936..fb1d08f5e 100644 --- a/client/internal/engine.go +++ b/client/internal/engine.go @@ -172,6 +172,7 @@ type EngineServices struct { StateManager *statemanager.Manager UpdateManager *updater.Manager ClientMetrics *metrics.ClientMetrics + MetricsCtx context.Context } // Engine is a mechanism responsible for reacting on Signal and Management stream events and managing connections to the remote peers. @@ -264,6 +265,7 @@ type Engine struct { // clientMetrics collects and pushes metrics clientMetrics *metrics.ClientMetrics + metricsCtx context.Context jobExecutor *jobexec.Executor jobExecutorWG sync.WaitGroup @@ -316,6 +318,7 @@ func NewEngine( probeStunTurn: relay.NewStunTurnProbe(relay.DefaultCacheTTL), jobExecutor: jobexec.NewExecutor(), clientMetrics: services.ClientMetrics, + metricsCtx: services.MetricsCtx, updateManager: services.UpdateManager, syncStoreDir: config.StateDir, } @@ -1073,7 +1076,7 @@ func (e *Engine) handleMetricsUpdate(config *mgmProto.MetricsConfig) { return } log.Infof("received metrics configuration from management: enabled=%v", config.GetEnabled()) - e.clientMetrics.UpdatePushFromMgm(e.ctx, config.GetEnabled()) + e.clientMetrics.UpdatePushFromMgm(e.metricsCtx, config.GetEnabled()) } func toFlowLoggerConfig(config *mgmProto.FlowConfig) (*nftypes.FlowConfig, error) { diff --git a/client/internal/metrics/metrics.go b/client/internal/metrics/metrics.go index d07a99a2d..cfe477107 100644 --- a/client/internal/metrics/metrics.go +++ b/client/internal/metrics/metrics.go @@ -7,6 +7,7 @@ import ( "fmt" "io" "sync" + "sync/atomic" "time" log "github.com/sirupsen/logrus" @@ -75,7 +76,7 @@ type ClientMetrics struct { agentInfo AgentInfo mu sync.RWMutex - push *Push + push atomic.Pointer[Push] pushMu sync.Mutex wg sync.WaitGroup pushCancel context.CancelFunc @@ -167,10 +168,7 @@ func (c *ClientMetrics) UpdateAgentInfo(agentInfo AgentInfo, publicKey string) { c.agentInfo = agentInfo c.mu.Unlock() - c.pushMu.Lock() - push := c.push - c.pushMu.Unlock() - if push != nil { + if push := c.push.Load(); push != nil { push.SetPeerID(agentInfo.peerID) } } @@ -194,7 +192,7 @@ func (c *ClientMetrics) StartPush(ctx context.Context, config PushConfig) { c.pushMu.Lock() defer c.pushMu.Unlock() - if c.push != nil { + if c.push.Load() != nil { log.Warnf("metrics push already running") return } @@ -230,13 +228,13 @@ func (c *ClientMetrics) UpdatePushFromMgm(ctx context.Context, enabled bool) { defer c.pushMu.Unlock() if enabled { - if c.push != nil { + if c.push.Load() != nil { return } log.Infof("enabled metrics push by management") c.startPushLocked(ctx, PushConfigFromEnv()) } else { - if c.push == nil { + if c.push.Load() == nil { return } log.Infof("disabled metrics push by management") @@ -261,22 +259,23 @@ func (c *ClientMetrics) startPushLocked(ctx context.Context, config PushConfig) ctx, cancel := context.WithCancel(ctx) c.pushCancel = cancel + c.push.Store(push) c.wg.Add(1) go func() { defer c.wg.Done() push.Start(ctx) + c.push.CompareAndSwap(push, nil) }() - c.push = push } // stopPushLocked stops push. Caller must hold pushMu. func (c *ClientMetrics) stopPushLocked() { - if c.push == nil { + if c.push.Load() == nil { return } c.pushCancel() c.wg.Wait() - c.push = nil + c.push.Store(nil) }