diff --git a/client/embed/embed.go b/client/embed/embed.go index 04bc60fb8..59d278fa1 100644 --- a/client/embed/embed.go +++ b/client/embed/embed.go @@ -464,7 +464,7 @@ func (c *Client) Status() (peer.FullStatus, error) { if connect != nil { engine := connect.Engine() if engine != nil { - _ = engine.RunHealthProbes(false) + _ = engine.RunHealthProbes(context.Background(), false) } } diff --git a/client/internal/engine.go b/client/internal/engine.go index 99ae66673..be260418d 100644 --- a/client/internal/engine.go +++ b/client/internal/engine.go @@ -1173,7 +1173,7 @@ func (e *Engine) handleBundle(params *mgmProto.BundleParameters) (*mgmProto.JobR TempDir: e.config.TempDir, ClientMetrics: e.clientMetrics, RefreshStatus: func() { - e.RunHealthProbes(true) + e.RunHealthProbes(e.ctx, true) }, } @@ -2058,7 +2058,20 @@ func (e *Engine) getRosenpassAddr() string { // RunHealthProbes executes health checks for Signal, Management, Relay, and WireGuard services // and updates the status recorder with the latest states. -func (e *Engine) RunHealthProbes(waitForResult bool) bool { +// +// ctx scopes the (potentially slow) STUN/TURN probing: a caller that gives up — +// e.g. a Status RPC whose client disconnected — cancels its ctx and the probe +// returns instead of running to its per-component timeout. The engine's own +// lifetime ctx still applies independently, so an engine shutdown aborts the +// probe even if the caller's ctx is context.Background(). +func (e *Engine) RunHealthProbes(ctx context.Context, waitForResult bool) bool { + // Tie the caller's ctx to the engine lifetime: either cancelling aborts + // the probe below. + ctx, cancel := context.WithCancel(ctx) + defer cancel() + stop := context.AfterFunc(e.ctx, cancel) + defer stop() + e.syncMsgMux.Lock() signalHealthy := e.signal.IsHealthy() @@ -2081,9 +2094,9 @@ func (e *Engine) RunHealthProbes(waitForResult bool) bool { if runtime.GOOS != "js" { var results []relay.ProbeResult if waitForResult { - results = e.probeStunTurn.ProbeAllWaitResult(e.ctx, stuns, turns) + results = e.probeStunTurn.ProbeAllWaitResult(ctx, stuns, turns) } else { - results = e.probeStunTurn.ProbeAll(e.ctx, stuns, turns) + results = e.probeStunTurn.ProbeAll(ctx, stuns, turns) } e.statusRecorder.UpdateRelayStates(results) diff --git a/client/server/debug.go b/client/server/debug.go index 33247db5f..4120b4cd7 100644 --- a/client/server/debug.go +++ b/client/server/debug.go @@ -52,7 +52,10 @@ func (s *Server) DebugBundle(_ context.Context, req *proto.DebugBundleRequest) ( if engine != nil { refreshStatus = func() { log.Debug("refreshing system health status for debug bundle") - engine.RunHealthProbes(true) + // Background ctx: the bundle wants a full, fresh probe regardless + // of the DebugBundle RPC client's lifetime. The engine's own ctx + // still aborts it on shutdown. + engine.RunHealthProbes(context.Background(), true) } } } diff --git a/client/server/probe_throttle.go b/client/server/probe_throttle.go new file mode 100644 index 000000000..ec6137e15 --- /dev/null +++ b/client/server/probe_throttle.go @@ -0,0 +1,88 @@ +package server + +import ( + "context" + "sync" + "time" + + log "github.com/sirupsen/logrus" +) + +// healthProbeRunner runs the full, expensive probe (network round-trips to +// management, signal and the relays) and reports whether every component was +// healthy. ctx cancels the probe when the caller gives up. Satisfied by +// *internal.Engine. +type healthProbeRunner interface { + RunHealthProbes(ctx context.Context, waitForResult bool) bool +} + +// statsRefresher does the cheap WireGuard-stats refresh callers fall back to +// when a fresh probe isn't warranted. Satisfied by *peer.Status. +type statsRefresher interface { + RefreshWireGuardStats() error +} + +// probeThrottle rate-limits and single-flights the daemon's health probes. +// +// Health probes are expensive (network round-trips to management, signal and +// the relays), while Status(GetFullPeerStatus=true) RPCs can arrive frequently +// and concurrently — the desktop UI alone issues one per connect/disconnect. +// probeThrottle keeps that load bounded with two rules: +// +// - Single-flight: only one probe runs at a time. Callers that pile up while +// a probe is in flight share its result instead of each launching another, +// even when that probe failed. A failed probe therefore does not make every +// waiter re-probe in turn; the next, non-overlapping caller can try again. +// - Throttle: after a fully successful probe the result is cached for +// interval. While any component is unhealthy the cache is not advanced, so +// later callers keep probing frequently and notice recovery quickly — the +// intentional "probe often while unhealthy" behaviour from the original +// design. +type probeThrottle struct { + interval time.Duration + + mu sync.Mutex + lastOK time.Time // last fully-successful probe; drives the throttle window + completedAt time.Time // when the most recent probe finished; drives single-flight sharing +} + +func newProbeThrottle(interval time.Duration) *probeThrottle { + return &probeThrottle{interval: interval} +} + +// Run decides whether to run a fresh health probe or serve the most recent +// result. It serialises concurrent callers: at most one runner.RunHealthProbes +// executes at a time and the rest call refresher.RefreshWireGuardStats and read +// the snapshot it produced. +// +// Both calls run while the throttle's lock is held, so a slow probe blocks +// other callers until it completes — that blocking is the single-flight +// guarantee. ctx is forwarded to RunHealthProbes so a caller that gives up +// cancels the in-flight probe (and any caller still queued on the lock falls +// through quickly once it acquires it, since the probe ctx is already done). +func (t *probeThrottle) Run(ctx context.Context, runner healthProbeRunner, refresher statsRefresher, waitForResult bool) { + entered := time.Now() + + t.mu.Lock() + defer t.mu.Unlock() + + // A probe that finished after we entered ran while we were waiting on the + // lock — i.e. a peer in the same burst already probed for us, so share its + // result rather than launch another. This holds even when that probe + // failed, so a failed probe doesn't make every waiter re-probe in turn. + sharedRecentProbe := t.completedAt.After(entered) + throttled := time.Since(t.lastOK) <= t.interval + + if sharedRecentProbe || throttled { + if err := refresher.RefreshWireGuardStats(); err != nil { + log.Debugf("failed to refresh WireGuard stats: %v", err) + } + return + } + + healthy := runner.RunHealthProbes(ctx, waitForResult) + t.completedAt = time.Now() + if healthy { + t.lastOK = t.completedAt + } +} diff --git a/client/server/probe_throttle_test.go b/client/server/probe_throttle_test.go new file mode 100644 index 000000000..cae776fa4 --- /dev/null +++ b/client/server/probe_throttle_test.go @@ -0,0 +1,109 @@ +package server + +import ( + "context" + "sync" + "sync/atomic" + "testing" + "time" +) + +// fakeProber implements both healthProbeRunner and statsRefresher with +// caller-supplied behaviour. +type fakeProber struct { + onProbe func() bool + onRefresh func() +} + +func (f fakeProber) RunHealthProbes(context.Context, bool) bool { + return f.onProbe() +} + +func (f fakeProber) RefreshWireGuardStats() error { + if f.onRefresh != nil { + f.onRefresh() + } + return nil +} + +func TestProbeThrottle_CachesAfterSuccess(t *testing.T) { + pt := newProbeThrottle(time.Minute) + + var probes, refreshes int + prober := fakeProber{ + onProbe: func() bool { probes++; return true }, + onRefresh: func() { refreshes++ }, + } + + pt.Run(context.Background(), prober, prober, false) + pt.Run(context.Background(), prober, prober, false) + + if probes != 1 { + t.Fatalf("expected 1 probe within the throttle window, got %d", probes) + } + if refreshes != 1 { + t.Fatalf("expected the throttled caller to refresh stats once, got %d", refreshes) + } +} + +func TestProbeThrottle_StaysOpenWhileUnhealthy(t *testing.T) { + pt := newProbeThrottle(time.Minute) + + var probes int + prober := fakeProber{onProbe: func() bool { probes++; return false }} // never healthy + + // Sequential, non-overlapping callers must each re-probe while unhealthy: + // a failed probe does not advance the throttle window. + pt.Run(context.Background(), prober, prober, false) + pt.Run(context.Background(), prober, prober, false) + pt.Run(context.Background(), prober, prober, false) + + if probes != 3 { + t.Fatalf("expected every non-overlapping caller to probe while unhealthy, got %d", probes) + } +} + +func TestProbeThrottle_SingleFlightSharesResult(t *testing.T) { + pt := newProbeThrottle(time.Minute) + + var probes int32 + release := make(chan struct{}) + started := make(chan struct{}) + + // First caller blocks inside the probe until released, holding the lock so + // the others pile up behind it. + prober := fakeProber{onProbe: func() bool { + if atomic.AddInt32(&probes, 1) == 1 { + close(started) + <-release + } + return false // unhealthy — the share must happen regardless of result + }} + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + pt.Run(context.Background(), prober, prober, false) + }() + + <-started // ensure the first probe is in flight before the burst arrives + + const waiters = 9 + wg.Add(waiters) + for i := 0; i < waiters; i++ { + go func() { + defer wg.Done() + pt.Run(context.Background(), prober, prober, false) + }() + } + + // Give the waiters time to block on the lock, then let the first finish. + time.Sleep(50 * time.Millisecond) + close(release) + wg.Wait() + + if got := atomic.LoadInt32(&probes); got != 1 { + t.Fatalf("expected a concurrent burst to run exactly 1 probe, got %d", got) + } +} diff --git a/client/server/server.go b/client/server/server.go index 22165fbcf..6b070b1e7 100644 --- a/client/server/server.go +++ b/client/server/server.go @@ -87,7 +87,7 @@ type Server struct { statusRecorder *peer.Status sessionWatcher *internal.SessionWatcher - lastProbe time.Time + probeThrottle *probeThrottle persistSyncResponse bool isSessionActive atomic.Bool @@ -131,6 +131,7 @@ func New(ctx context.Context, logFile string, configFile string, profilesDisable networksDisabled: networksDisabled, jwtCache: newJWTCache(), extendAuthSessionFlow: auth.NewPendingFlow(), + probeThrottle: newProbeThrottle(probeThreshold), } agent := &serverAgent{s} s.sleepHandler = sleephandler.New(agent) @@ -1242,13 +1243,14 @@ func (s *Server) Status( } } - return s.buildStatusResponse(msg) + return s.buildStatusResponse(ctx, msg) } // buildStatusResponse composes a StatusResponse from the current daemon // state. Shared between the unary Status RPC and the SubscribeStatus -// stream so both paths return identical snapshots. -func (s *Server) buildStatusResponse(msg *proto.StatusRequest) (*proto.StatusResponse, error) { +// stream so both paths return identical snapshots. ctx scopes the health +// probe runProbes may trigger — a caller that disconnects cancels it. +func (s *Server) buildStatusResponse(ctx context.Context, msg *proto.StatusRequest) (*proto.StatusResponse, error) { state := internal.CtxGetState(s.rootCtx) status, err := state.Status() if err != nil { @@ -1277,7 +1279,7 @@ func (s *Server) buildStatusResponse(msg *proto.StatusRequest) (*proto.StatusRes s.statusRecorder.UpdateRosenpass(s.config.RosenpassEnabled, s.config.RosenpassPermissive) if msg.GetFullPeerStatus { - s.runProbes(msg.ShouldRunProbes) + s.runProbes(ctx, msg.ShouldRunProbes) fullStatus := s.statusRecorder.GetFullStatus() pbFullStatus := fullStatus.ToProto() pbFullStatus.Events = s.statusRecorder.GetEventHistory() @@ -1707,7 +1709,7 @@ func isUnixRunningDesktop() bool { return os.Getenv("DESKTOP_SESSION") != "" || os.Getenv("XDG_CURRENT_DESKTOP") != "" } -func (s *Server) runProbes(waitForProbeResult bool) { +func (s *Server) runProbes(ctx context.Context, waitForProbeResult bool) { if s.connectClient == nil { return } @@ -1717,15 +1719,7 @@ func (s *Server) runProbes(waitForProbeResult bool) { return } - if time.Since(s.lastProbe) > probeThreshold { - if engine.RunHealthProbes(waitForProbeResult) { - s.lastProbe = time.Now() - } - } else { - if err := s.statusRecorder.RefreshWireGuardStats(); err != nil { - log.Debugf("failed to refresh WireGuard stats: %v", err) - } - } + s.probeThrottle.Run(ctx, engine, s.statusRecorder, waitForProbeResult) } // GetConfig of the daemon. diff --git a/client/server/status_stream.go b/client/server/status_stream.go index aa1dc508c..c6ba547eb 100644 --- a/client/server/status_stream.go +++ b/client/server/status_stream.go @@ -44,7 +44,7 @@ func (s *Server) SubscribeStatus(req *proto.StatusRequest, stream proto.DaemonSe } func (s *Server) sendStatusSnapshot(req *proto.StatusRequest, stream proto.DaemonService_SubscribeStatusServer) error { - resp, err := s.buildStatusResponse(req) + resp, err := s.buildStatusResponse(stream.Context(), req) if err != nil { log.Warnf("build status snapshot for stream: %v", err) return err