Fix IsBenched / notification inconsistency

Author: aaronbuchwald

snow/networking/benchlist/benchlist.go

@@ -158,14 +158,6 @@ func (b *benchlist) observe(nodeID ids.NodeID, v float64) {
 		b.nodes[nodeID] = n
 	}
 
-	// If the bench duration has expired, clear the benched state and reset
-	// the failure probability so that the node gets a clean slate rather
-	// than being immediately re-benched due to stale EWMA history.
-	if n.isBenched && b.benchExpired(n) {
-		n.isBenched = false
-		n.failureProbability = math.NewUninitializedAverager(b.halflife)
-	}
-
 	n.failureProbability.Observe(v, b.clock.Time())
 
 	p := n.failureProbability.Read()
@@ -183,18 +175,13 @@ func (b *benchlist) observe(nodeID ids.NodeID, v float64) {
 	}
 }
 
-// benchExpired reports whether n's bench duration has elapsed.
-func (b *benchlist) benchExpired(n *node) bool {
-	return !n.benchedAt.IsZero() && b.clock.Time().After(n.benchedAt.Add(b.benchDuration))
-}
-
 // IsBenched returns true if messages to nodeID should immediately fail.
 func (b *benchlist) IsBenched(nodeID ids.NodeID) bool {
 	b.lock.RLock()
 	defer b.lock.RUnlock()
 
 	n, ok := b.nodes[nodeID]
-	return ok && n.isBenched && !b.benchExpired(n)
+	return ok && n.isBenched
 }
 
 // TODO: Close goroutine within run during shutdown.
@@ -218,17 +205,23 @@ func (b *benchlist) run() {
 			)
 
 			if j.bench {
+				// The producer path (observe) is the source of truth for
+				// b.nodes transitions triggered by EWMA. By the time a bench
+				// job is emitted, n.isBenched was already set under b.lock.
+				// This consumer branch only drives callbacks/metrics and tracks
+				// callback-visible deadlines.
 				// Always update the timeout deadline. If the node is
-				// already benched from the consumer's perspective
-				// (e.g., observe detected expiration and immediately
-				// re-benched before the timer fired), this extends
-				// the deadline without a duplicate Benched notification.
+				// already benched from the consumer's perspective, this
+				// extends the deadline without a duplicate Benched
+				// notification.
 				timeoutHeap.Push(j.nodeID, time.Now().Add(b.benchDuration))
 				if !benched.Contains(j.nodeID) {
 					benched.Add(j.nodeID)
 					b.benchable.Benched(b.ctx.ChainID, j.nodeID)
 				}
 			} else if benched.Contains(j.nodeID) {
+				// Same as above: EWMA unbench jobs don't update b.nodes here
+				// because observe already flipped n.isBenched before enqueuing.
 				// Guard: only unbench if the consumer still considers
 				// the node benched, avoiding duplicate Unbenched calls
 				// when both EWMA and timeout race to unbench.
@@ -246,6 +239,16 @@ func (b *benchlist) run() {
 				}
 				timeoutHeap.Pop()
 				benched.Remove(nodeID)
+				b.lock.Lock()
+				// Timeout unbenching originates in this consumer goroutine
+				// (not in observe), so we must synchronize b.nodes here to
+				// keep IsBenched consistent with callback-visible state.
+				// Timeout-based unbench gives the node a clean slate.
+				if n, exists := b.nodes[nodeID]; exists {
+					n.isBenched = false
+					n.failureProbability = math.NewUninitializedAverager(b.halflife)
+				}
+				b.lock.Unlock()
 
 				b.ctx.Log.Debug("unbenching node due to timeout",
 					zap.Stringer("nodeID", nodeID),

snow/networking/benchlist/benchlist_test.go

@@ -122,77 +122,6 @@ func TestBenchlist(t *testing.T) {
 	requireBenched()
 }
 
-// Regression test: when observe() detects that a bench has expired via the
-// mock clock (before the consumer's real-time timer fires) and the EWMA
-// immediately re-benches the node, the consumer must not send a duplicate
-// Benched notification. The benchable.Benched() assertion
-// (require.NotContains) would fail if the guard were missing.
-func TestBenchlistNoDuplicateBench(t *testing.T) {
-	require := require.New(t)
-
-	snowCtx := snowtest.Context(t, snowtest.CChainID)
-	ctx := snowtest.ConsensusContext(snowCtx)
-	vdrs := validators.NewManager()
-	vdrID := ids.GenerateTestNodeID()
-
-	require.NoError(vdrs.AddStaker(ctx.SubnetID, vdrID, nil, ids.Empty, 1))
-
-	benchable := &benchable{
-		t:           t,
-		wantChainID: ctx.ChainID,
-		updated:     make(chan struct{}, 1),
-	}
-	b, err := newBenchlist(
-		ctx,
-		benchable,
-		vdrs,
-		Config{
-			Halflife:           DefaultHalflife,
-			UnbenchProbability: DefaultUnbenchProbability,
-			BenchProbability:   DefaultBenchProbability,
-			BenchDuration:      10 * time.Second,
-		},
-		prometheus.NewRegistry(),
-	)
-	require.NoError(err)
-
-	now := time.Now()
-	b.clock.Set(now)
-
-	// Bench the node: p = 2/3 > 0.5
-	b.RegisterResponse(vdrID)
-	b.RegisterFailure(vdrID)
-	b.RegisterFailure(vdrID)
-	<-benchable.updated
-	require.True(b.IsBenched(vdrID))
-
-	// Advance the mock clock past the bench duration. The real-time timer
-	// has NOT fired, so the consumer still considers the node benched.
-	now = now.Add(11 * time.Second)
-	b.clock.Set(now)
-	require.False(b.IsBenched(vdrID))
-
-	// Register a failure. observe() detects the bench expired, clears
-	// isBenched, resets the EWMA to a clean slate, then observes the
-	// failure (p = 1.0 > 0.5) and immediately re-benches.
-	// This sends a bench job for a node the consumer already considers
-	// benched. Without the duplicate-bench guard, the consumer would call
-	// benchable.Benched() again, failing the require.NotContains assertion.
-	b.RegisterFailure(vdrID)
-	require.True(b.IsBenched(vdrID))
-
-	// Drive the probability below unbenchProbability to trigger an EWMA
-	// unbench. Since jobs are FIFO, when we receive the unbench
-	// notification we know the re-bench job above was already processed
-	// (without a duplicate Benched call).
-	for range 20 {
-		b.RegisterResponse(vdrID)
-	}
-	<-benchable.updated
-	require.False(b.IsBenched(vdrID))
-	require.Empty(benchable.benched)
-}
-
 func TestBenchlistTimeout(t *testing.T) {
 	require := require.New(t)
 
@@ -232,15 +161,10 @@ func TestBenchlistTimeout(t *testing.T) {
 	<-benchable.updated
 	require.True(b.IsBenched(vdrID))
 
-	// Advance the mock clock past the bench duration so IsBenched returns
-	// false immediately.
-	now = now.Add(100 * time.Millisecond)
-	b.clock.Set(now)
-	require.False(b.IsBenched(vdrID))
-
 	// The consumer goroutine's timer also fires (real time ≥ 50ms) and calls
 	// Unbenched on the benchable.
 	<-benchable.updated
+	require.False(b.IsBenched(vdrID))
 	require.Empty(benchable.benched)
 }
 
@@ -271,7 +195,7 @@ func TestBenchlistTimeoutCleansSlate(t *testing.T) {
 			Halflife:           DefaultHalflife,
 			UnbenchProbability: DefaultUnbenchProbability,
 			BenchProbability:   DefaultBenchProbability,
-			BenchDuration:      10 * time.Second,
+			BenchDuration:      50 * time.Millisecond,
 		},
 		prometheus.NewRegistry(),
 	)
@@ -287,9 +211,8 @@ func TestBenchlistTimeoutCleansSlate(t *testing.T) {
 	<-benchable.updated
 	require.True(b.IsBenched(vdrID))
 
-	// Advance past the bench duration so the bench expires.
-	now = now.Add(11 * time.Second)
-	b.clock.Set(now)
+	// Wait for timeout-based unbench.
+	<-benchable.updated
 	require.False(b.IsBenched(vdrID))
 
 	// Register a response followed by a failure. With a clean slate, p = 1/2