netbird/client/server/probe_throttle.go

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