client/server: throttle and single-flight health probes

Status(GetFullPeerStatus=true) RPCs trigger a full health probe
(network round-trips to management, signal and the relays). The
desktop UI issues these frequently and concurrently, and a burst of
parallel Get() calls each fired its own probe — the lastProbe guard
was unprotected against concurrent access and only advanced when every
component was healthy, so a sustained unhealthy state (e.g. relay down)
disabled the throttle entirely and let every call re-probe.

Extract the throttle/single-flight policy into probeThrottle:
  - single-flight: only one probe runs at a time; concurrent callers
    that piled up while it ran share its result instead of each
    launching another, even when that probe failed.
  - throttle: lastOK only advances on a fully successful probe, so
    while anything is unhealthy callers keep probing frequently and
    notice recovery quickly (preserved from the original design).

RunHealthProbes now takes a context so a caller that gives up (e.g. a
Status RPC whose client disconnected) cancels the in-flight STUN/TURN
probe instead of letting it run to its per-component timeout. The
engine's own lifetime ctx still applies independently.

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)
 		}
 	}
 

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)
 

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)
 			}
 		}
 	}

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
+	}
+}

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)
+	}
+}

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.
+// stream so both paths return identical snapshots. ctx scopes the health
-func (s *Server) buildStatusResponse(msg *proto.StatusRequest) (*proto.StatusResponse, error) {
+// 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 {
+	s.probeThrottle.Run(ctx, engine, s.statusRecorder, waitForProbeResult)
-		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)
-		}
-	}
 }
 
 // GetConfig of the daemon.

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