[client, management] Recover from stuck proof collections and keep re-tracked accounts on the challenge refresher (#7890)

Co-authored-by: Viktor Liu <viktor@netbird.io>

client/internal/certproof/collector.go

@@ -2,7 +2,7 @@ package certproof
 
 import (
 	"context"
-	"sync/atomic"
+	"sync"
 	"time"
 
 	log "github.com/sirupsen/logrus"
@@ -11,28 +11,51 @@ import (
 	"github.com/netbirdio/netbird/shared/management/proto"
 )
 
-const collectTimeout = 45 * time.Second
+const (
+	collectTimeout = 45 * time.Second
 
-// collecting is the single-flight flag Collectors share by default. It outlives the
-// engine, since a collection abandoned in a token, TPM or keychain call keeps running
-// after the engine that started it has stopped, and the next engine must not start
-// another one on top of it.
-var collecting atomic.Bool
+	// lostAfter is how many deadlines an abandoned collection is waited for before a new
+	// one may start beside it. A store that is merely slow finishes well within that and
+	// releases its slot; only one wedged for good is worked around.
+	lostAfter = 10
 
-// Collector runs CollectProofs with a deadline and at most one collection at a time in
-// the process. Token, TPM and keychain calls cannot be interrupted, so a collection that
-// overruns is abandoned rather than awaited, and a new one is refused until it has
-// finished. The zero value is ready to use.
+	// maxInFlight bounds how many abandoned collections may pile up. A store that never
+	// answers would otherwise leak a goroutine on every collection for the life of the
+	// daemon.
+	maxInFlight = 2
+)
+
+// collectSlots tracks the collections running in the process. Collectors share one by
+// default: a collection abandoned in a token, TPM or keychain call keeps running after
+// the engine that started it has stopped, and the next engine must account for it.
+type collectSlots struct {
+	mu   sync.Mutex
+	next uint64
+	// running holds the start time of each collection still running, by slot id.
+	running map[uint64]time.Time
+}
+
+var processSlots collectSlots
+
+// Collector runs CollectProofs with a deadline and, in the normal case, one collection
+// at a time in the process. Token, TPM and keychain calls cannot be interrupted, and on
+// Linux there is not even a process to kill, so a collection that overruns is abandoned
+// rather than awaited. A single wedged call would then refuse every later collection for
+// the life of the daemon, so once an abandoned collection has been running for lostAfter
+// deadlines another may start beside it, up to maxInFlight. The zero value is ready to
+// use.
 type Collector struct {
-	// busy overrides the process-wide single-flight flag when set.
-	busy *atomic.Bool
+	// slots overrides the process-wide slots when set.
+	slots *collectSlots
 	// timeout overrides collectTimeout when set.
 	timeout time.Duration
+	// now overrides the clock when set.
+	now func() time.Time
 }
 
 // Collect answers the certificate challenges in checks, returning no proofs when there
-// are no challenges, when a previous collection is still running, or when this one does
-// not finish in time. Missing proofs fail the certificate check on management.
+// are no challenges, when previous collections still hold the slots, or when this one
+// does not finish in time. Missing proofs fail the certificate check on management.
 func (c *Collector) Collect(ctx context.Context, checks []*proto.Checks, peerKey []byte, cfg Config) []certposture.Proof {
 	return c.collect(ctx, checks, func(ctx context.Context) []certposture.Proof {
 		return CollectProofs(ctx, checks, peerKey, cfg)
@@ -43,9 +66,9 @@ func (c *Collector) collect(ctx context.Context, checks []*proto.Checks, run fun
 	if len(certificateChallenges(checks)) == 0 {
 		return nil
 	}
-	busy := c.flag()
-	if !busy.CompareAndSwap(false, true) {
-		log.Warnf("certificate posture: previous proof collection is still running, sending no proofs")
+	slots := c.slotsInUse()
+	slot, ok := slots.start(c.clock(), time.Duration(lostAfter)*c.deadline())
+	if !ok {
 		return nil
 	}
 
@@ -57,7 +80,7 @@ func (c *Collector) collect(ctx context.Context, checks []*proto.Checks, run fun
 		// The slot is freed before the result is delivered, so a caller that starts the
 		// next collection right after this one returned is not turned away.
 		proofs := run(ctx)
-		busy.Store(false)
+		slots.finish(slot)
 		done <- proofs
 	}()
 
@@ -70,11 +93,24 @@ func (c *Collector) collect(ctx context.Context, checks []*proto.Checks, run fun
 	}
 }
 
-func (c *Collector) flag() *atomic.Bool {
-	if c.busy != nil {
-		return c.busy
+// Stuck reports whether lost collections hold every slot. No collection starts again
+// until one of them returns, which for a store that never answers means a restart.
+func (c *Collector) Stuck() bool {
+	return c.slotsInUse().allLost(c.clock(), time.Duration(lostAfter)*c.deadline())
+}
+
+func (c *Collector) slotsInUse() *collectSlots {
+	if c.slots != nil {
+		return c.slots
 	}
-	return &collecting
+	return &processSlots
+}
+
+func (c *Collector) clock() time.Time {
+	if c.now != nil {
+		return c.now()
+	}
+	return time.Now()
 }
 
 func (c *Collector) deadline() time.Duration {
@@ -83,3 +119,71 @@ func (c *Collector) deadline() time.Duration {
 	}
 	return collectTimeout
 }
+
+// start claims a slot at now, reporting its id and whether the caller may collect. A
+// collection that has held a slot for lost is taken to be wedged.
+func (s *collectSlots) start(now time.Time, lost time.Duration) (uint64, bool) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if len(s.running) > 0 {
+		oldest := s.oldest()
+		switch {
+		case len(s.running) >= maxInFlight:
+			log.Warnf("certificate posture: %d proof collections are wedged, sending no proofs", len(s.running))
+			return 0, false
+		case now.Sub(oldest) < lost:
+			log.Warnf("certificate posture: previous proof collection is still running, sending no proofs")
+			return 0, false
+		default:
+			log.Warnf("certificate posture: a proof collection has been stuck since %s, starting another", oldest.Format(time.RFC3339))
+		}
+	}
+	if s.running == nil {
+		s.running = make(map[uint64]time.Time, maxInFlight)
+	}
+	s.next++
+	s.running[s.next] = now
+	return s.next, true
+}
+
+// finish releases the slot with the given id.
+func (s *collectSlots) finish(slot uint64) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	delete(s.running, slot)
+}
+
+// allLost reports whether every slot is held by a collection that has run for lost.
+func (s *collectSlots) allLost(now time.Time, lost time.Duration) bool {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if len(s.running) < maxInFlight {
+		return false
+	}
+	for _, started := range s.running {
+		if now.Sub(started) < lost {
+			return false
+		}
+	}
+	return true
+}
+
+// idle reports whether no collection is running.
+func (s *collectSlots) idle() bool {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return len(s.running) == 0
+}
+
+// oldest returns the start time of the longest-running collection. The caller holds mu
+// and has checked that one is running.
+func (s *collectSlots) oldest() time.Time {
+	var oldest time.Time
+	for _, started := range s.running {
+		if oldest.IsZero() || started.Before(oldest) {
+			oldest = started
+		}
+	}
+	return oldest
+}

client/internal/certproof/collector_test.go

@@ -40,7 +40,7 @@ func TestCollector_ReturnsProofs(t *testing.T) {
 }
 
 func TestCollector_AbandonsStuckCollection(t *testing.T) {
-	c := Collector{timeout: 50 * time.Millisecond, busy: new(atomic.Bool)}
+	c := Collector{timeout: 50 * time.Millisecond, slots: &collectSlots{}}
 	release := make(chan struct{})
 	finished := make(chan struct{})
 
@@ -66,7 +66,7 @@ func TestCollector_AbandonsStuckCollection(t *testing.T) {
 
 	close(release)
 	<-finished
-	require.Eventually(t, func() bool { return !c.busy.Load() }, time.Second, 5*time.Millisecond, "the collector frees up once the stuck call returns")
+	require.Eventually(t, func() bool { return c.slots.idle() }, time.Second, 5*time.Millisecond, "the collector frees up once the stuck call returns")
 
 	want := []certposture.Proof{{Nonce: []byte("nonce")}}
 	assert.Equal(t, want, c.collect(context.Background(), challengeChecks, func(context.Context) []certposture.Proof { return want }),
@@ -74,7 +74,7 @@ func TestCollector_AbandonsStuckCollection(t *testing.T) {
 }
 
 func TestCollector_CancelsContextAtDeadline(t *testing.T) {
-	c := Collector{timeout: 20 * time.Millisecond, busy: new(atomic.Bool)}
+	c := Collector{timeout: 20 * time.Millisecond, slots: &collectSlots{}}
 	cancelled := make(chan struct{})
 
 	c.collect(context.Background(), challengeChecks, func(ctx context.Context) []certposture.Proof {
@@ -94,9 +94,9 @@ func TestCollector_CancelsContextAtDeadline(t *testing.T) {
 // collection stuck in a token call outlives the engine that started it, so the next
 // engine's Collector must not start another one until it has finished.
 func TestCollector_SingleFlightAcrossCollectors(t *testing.T) {
-	shared := new(atomic.Bool)
-	first := Collector{timeout: 20 * time.Millisecond, busy: shared}
-	second := Collector{timeout: time.Second, busy: shared}
+	shared := &collectSlots{}
+	first := Collector{timeout: 20 * time.Millisecond, slots: shared}
+	second := Collector{timeout: time.Second, slots: shared}
 
 	release := make(chan struct{})
 	stuck := func(context.Context) []certposture.Proof {
@@ -114,7 +114,7 @@ func TestCollector_SingleFlightAcrossCollectors(t *testing.T) {
 	assert.False(t, called, "a second collection does not start on top of the abandoned one")
 
 	close(release)
-	require.Eventually(t, func() bool { return !shared.Load() }, time.Second, 5*time.Millisecond)
+	require.Eventually(t, func() bool { return shared.idle() }, time.Second, 5*time.Millisecond)
 	assert.Len(t, second.collect(context.Background(), challengeChecks, func(context.Context) []certposture.Proof { return []certposture.Proof{{}} }), 1,
 		"collections resume once the abandoned one finished")
 }
@@ -122,5 +122,86 @@ func TestCollector_SingleFlightAcrossCollectors(t *testing.T) {
 // TestCollector_ZeroValuesShareTheProcessFlag: the zero value uses the process-wide flag.
 func TestCollector_ZeroValuesShareTheProcessFlag(t *testing.T) {
 	var a, b Collector
-	assert.Same(t, a.flag(), b.flag(), "separate Collectors share one single-flight flag")
+	assert.Same(t, a.slotsInUse(), b.slotsInUse(), "separate Collectors share the process slots")
+}
+
+// TestCollector_LostCollectionDoesNotBlockForever: a store call that never answers, such
+// as a wedged token on Linux where there is no process to kill, must not refuse every
+// later collection. Once it has run for lostAfter deadlines another starts beside it,
+// and no more than maxInFlight ever run.
+func TestCollector_LostCollectionDoesNotBlockForever(t *testing.T) {
+	now := time.Now()
+	slots := &collectSlots{}
+	c := Collector{timeout: 10 * time.Millisecond, slots: slots, now: func() time.Time { return now }}
+
+	release := make(chan struct{})
+	defer close(release)
+	var started atomic.Int32
+	wedged := func(context.Context) []certposture.Proof {
+		started.Add(1)
+		<-release
+		return nil
+	}
+
+	c.collect(context.Background(), challengeChecks, wedged)
+	waitStarted(t, &started, 1)
+	c.collect(context.Background(), challengeChecks, wedged)
+	assert.Equal(t, int32(1), started.Load(), "a second collection is refused while the first may still finish")
+	assert.False(t, c.Stuck(), "one slow collection is not stuck yet")
+
+	now = now.Add(time.Duration(lostAfter) * c.deadline())
+	c.collect(context.Background(), challengeChecks, wedged)
+	waitStarted(t, &started, 2)
+	assert.False(t, c.Stuck(), "the second collection may still return, so the collector is not stuck yet")
+
+	now = now.Add(time.Duration(lostAfter) * c.deadline())
+	c.collect(context.Background(), challengeChecks, wedged)
+	assert.Equal(t, int32(2), started.Load(), "no more than maxInFlight collections run")
+	assert.True(t, c.Stuck(), "with every slot held by a lost collection the collector reports itself stuck")
+}
+
+// TestCollector_LostWindowFollowsTheOldestRunningCollection: A wedges, B starts beside it
+// once A counts as lost, then A returns. B is recent, so a third collection must wait for
+// B to be lost in its own right rather than inherit A's start time.
+func TestCollector_LostWindowFollowsTheOldestRunningCollection(t *testing.T) {
+	now := time.Now()
+	slots := &collectSlots{}
+	c := Collector{timeout: 10 * time.Millisecond, slots: slots, now: func() time.Time { return now }}
+	lost := time.Duration(lostAfter) * c.deadline()
+
+	releaseA := make(chan struct{})
+	releaseB := make(chan struct{})
+	defer close(releaseB)
+	var started atomic.Int32
+	blockOn := func(release chan struct{}) func(context.Context) []certposture.Proof {
+		return func(context.Context) []certposture.Proof {
+			started.Add(1)
+			<-release
+			return nil
+		}
+	}
+
+	c.collect(context.Background(), challengeChecks, blockOn(releaseA))
+	waitStarted(t, &started, 1)
+	now = now.Add(lost)
+	c.collect(context.Background(), challengeChecks, blockOn(releaseB))
+	waitStarted(t, &started, 2)
+
+	close(releaseA)
+	require.Eventually(t, func() bool {
+		slots.mu.Lock()
+		defer slots.mu.Unlock()
+		return len(slots.running) == 1
+	}, time.Second, 5*time.Millisecond, "A frees its slot once it returns")
+
+	c.collect(context.Background(), challengeChecks, blockOn(make(chan struct{})))
+	assert.Equal(t, int32(2), started.Load(), "a third collection waits while B is recent")
+}
+
+// waitStarted waits for the collections the test launched to have started. collect may
+// return at its deadline before the goroutine running the store call was scheduled.
+func waitStarted(t *testing.T, started *atomic.Int32, n int32) {
+	t.Helper()
+	require.Eventually(t, func() bool { return started.Load() == n }, time.Second, time.Millisecond,
+		"%d collections should have started", n)
 }

client/internal/engine_certposture.go

@@ -57,6 +57,9 @@ type certPostureState struct {
 	hasDelivery bool
 	cached      []certposture.Proof
 	cachedFor   string
+	// stuck is whether the collector was last seen with every slot held by a lost
+	// collection.
+	stuck bool
 }
 
 // record stores the outcome of a collection for the challenges identified by
@@ -123,8 +126,17 @@ func (s *certPostureState) needsCollection(challengesKey, userContext string, no
 	}
 }
 
+// setStuck records whether the collector is stuck and reports whether it just became so.
+func (s *certPostureState) setStuck(stuck bool) bool {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	became := stuck && !s.stuck
+	s.stuck = stuck
+	return became
+}
+
 // cachedProofs returns the cached proofs management still accepts for the challenges in
-// checks: those whose nonce is current, or from the window before, for one of them. Proofs
+// checks: those whose nonce is current, or from an adjacent window, for one of them. Proofs
 // signed for the previous nonce bridge the time until the watcher has signed the new one.
 func (s *certPostureState) cachedProofs(checks []*mgmProto.Checks) []certposture.Proof {
 	s.mu.Lock()
@@ -230,6 +242,19 @@ func (e *Engine) publishCertificatePostureEvent(proven bool) {
 			"Access to some resources may be blocked until one is available.", nil)
 }
 
+// publishCertificateStoreStuckEvent tells the user that reading the certificate store
+// stopped answering. Collection resumes only once a stuck call returns, which a store
+// that hangs for good never does short of a restart.
+func (e *Engine) publishCertificateStoreStuckEvent() {
+	if e.statusRecorder == nil {
+		return
+	}
+	e.statusRecorder.PublishEvent(cProto.SystemEvent_WARNING, cProto.SystemEvent_SYSTEM,
+		"certificate posture: the certificate store stopped responding",
+		"NetBird cannot read the certificates required by your organization's device policy because "+
+			"the certificate store stopped responding. If this persists, restart the NetBird service.", nil)
+}
+
 // watchCertificatePosture owns certificate proof collection until ctx is done. It
 // collects when woken by new checks or a sync, and on every tick when the cached proofs
 // went stale, such as after a login following an autostart.
@@ -303,6 +328,9 @@ func (e *Engine) refreshCertificateProofs() error {
 	if e.certState.record(key, userContext, proofs, time.Now()) {
 		e.publishCertificatePostureEvent(len(proofs) > 0)
 	}
+	if e.certState.setStuck(e.certProofs.Stuck()) {
+		e.publishCertificateStoreStuckEvent()
+	}
 
 	if e.certState.sameAsDelivered(proofs) {
 		log.Debugf("certificate posture: proofs unchanged since the last meta sync, not sending")

client/internal/engine_certposture_test.go

@@ -39,6 +39,17 @@ func TestCertPostureState_Record(t *testing.T) {
 	assert.True(t, unproven.record("k", "", nil, now), "losing the proof again is reported")
 }
 
+// TestCertPostureState_ReportsStuckOnce: a store that stopped answering is reported to
+// the user once when it wedges the collector, not on every collection that follows, and
+// again if it wedges anew after recovering.
+func TestCertPostureState_ReportsStuckOnce(t *testing.T) {
+	var s certPostureState
+	assert.True(t, s.setStuck(true), "becoming stuck is reported")
+	assert.False(t, s.setStuck(true), "staying stuck is not reported again")
+	assert.False(t, s.setStuck(false), "recovering is not a stuck report")
+	assert.True(t, s.setStuck(true), "getting stuck again is reported again")
+}
+
 func TestCertPostureState_NeedsCollection(t *testing.T) {
 	now := time.Now()
 	one := []certposture.Proof{{}}