Review cleanup: remove redundant nGranted, internalize test helpers

- Remove nGranted field; len(holders) is the single source of truth
- Move IsLocked/NGranted to test-only helpers (isIdle/nGranted)
- Refactor lockedAdvanceFirstWaiting to accept elem param, absorbing
  the equality guard from all 4 callers
- Rename inFlightAfterCancel → nGrantedAfterCancel in race test

AI disclosure: Claude Code assisted with development. Every line of code was either written by or carefully reviewed by me :)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Brett Wines <bwines@slack-corp.com>

Author: bgwines

go/vt/vttablet/tabletserver/loadshed/codelq.go

@@ -194,9 +194,7 @@ func (q *CoDelQueue) lockedPeek() *Request {
 		if req.signaledValue == nil || req.signaledValue == grantSentinel {
 			return req
 		}
-		if q.firstWaiting == front {
-			q.lockedAdvanceFirstWaiting()
-		}
+		q.lockedAdvanceFirstWaiting(front)
 		q.queue.Remove(front)
 		req.codelqElem = nil
 		if req.isDroppable() {
@@ -217,9 +215,7 @@ func (q *CoDelQueue) lockedPeek() *Request {
 // updating dropping state if appropriate.
 func (q *CoDelQueue) lockedPopElem(elem *list.Element, err error) *Request {
 	req := elem.Value.(*Request)
-	if q.firstWaiting == elem {
-		q.lockedAdvanceFirstWaiting()
-	}
+	q.lockedAdvanceFirstWaiting(elem)
 	q.queue.Remove(elem)
 	req.codelqElem = nil
 
@@ -242,9 +238,7 @@ func (q *CoDelQueue) lockedRemove(r *Request) {
 	if r.codelqElem == nil {
 		return
 	}
-	if q.firstWaiting == r.codelqElem {
-		q.lockedAdvanceFirstWaiting()
-	}
+	q.lockedAdvanceFirstWaiting(r.codelqElem)
 	q.queue.Remove(r.codelqElem)
 	r.codelqElem = nil
 
@@ -264,18 +258,17 @@ func (q *CoDelQueue) lockedOnGrant(r *Request) {
 			q.dropping = false
 		}
 	}
-	if q.firstWaiting == r.codelqElem {
-		q.lockedAdvanceFirstWaiting()
-	}
+	q.lockedAdvanceFirstWaiting(r.codelqElem)
 }
 
-// lockedAdvanceFirstWaiting moves the firstWaiting pointer to the next
-// unsignaled request in the queue, or nil if none remain.
-func (q *CoDelQueue) lockedAdvanceFirstWaiting() {
-	if q.firstWaiting == nil {
+// lockedAdvanceFirstWaiting advances the firstWaiting pointer past elem if
+// elem is the current firstWaiting. elem is a queue entry that is no longer
+// waiting — either because it was granted, removed, or dropped.
+func (q *CoDelQueue) lockedAdvanceFirstWaiting(elem *list.Element) {
+	if q.firstWaiting != elem {
 		return
 	}
-	for e := q.firstWaiting.Next(); e != nil; e = e.Next() {
+	for e := elem.Next(); e != nil; e = e.Next() {
 		if e.Value.(*Request).signaledValue == nil {
 			q.firstWaiting = e
 			return

go/vt/vttablet/tabletserver/loadshed/snake.go

@@ -46,7 +46,6 @@ type (
 		mu sync.Mutex
 
 		q              *SelfContentionAwareCoDelQueue
-		nGranted       int
 		holders        map[*Request]struct{}
 		maxAgeTimers   map[*Request]*time.Timer
 		dropTimer      *time.Timer
@@ -91,7 +90,7 @@ func (s *Snake) capacity() int {
 }
 
 func (s *Snake) hasCapacity() bool {
-	return s.nGranted < s.capacity()
+	return len(s.holders) < s.capacity()
 }
 
 // Acquire acquires a slot. It blocks until a slot is granted, the request
@@ -147,20 +146,6 @@ func (s *Snake) Acquire(ctx context.Context, valveID string) (*SafeUnlock, error
 	}
 }
 
-// IsLocked reports whether any slot is currently held.
-func (s *Snake) IsLocked() bool {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	return s.nGranted > 0
-}
-
-// InFlight reports the number of currently held slots.
-func (s *Snake) InFlight() int {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	return s.nGranted
-}
-
 // IsHealthy reports whether the CoDel queue is healthy.
 func (s *Snake) IsHealthy() bool {
 	s.mu.Lock()
@@ -193,7 +178,6 @@ func (s *Snake) release(req *Request, excValue error) error {
 	s.lockedStopMaxAgeTimer(req)
 	delete(s.holders, req)
 	s.q.lockedComplete(req)
-	s.nGranted--
 	s.lockedTryGrantOne()
 	s.mu.Unlock()
 
@@ -206,13 +190,11 @@ func (s *Snake) releaseOnCancel(req *Request) {
 	s.lockedStopMaxAgeTimer(req)
 	delete(s.holders, req)
 	s.q.lockedComplete(req)
-	s.nGranted--
 	s.lockedTryGrantOne()
 	s.mu.Unlock()
 }
 
 func (s *Snake) lockedGrant(req *Request) {
-	s.nGranted++
 	s.holders[req] = struct{}{}
 	s.q.lockedOnGrant(req)
 	s.lockedStartMaxAgeTimer(req)

go/vt/vttablet/tabletserver/loadshed/snake_nholder_test.go

@@ -44,12 +44,12 @@ func TestSnake_NHolder_ConcurrentGrants(t *testing.T) {
 				unlocks[i] = u
 			}
 
-			assert.Equal(t, cap, s.InFlight())
+			assert.Equal(t, cap, s.nGranted())
 
 			for _, u := range unlocks {
 				u.Release()
 			}
-			assert.Equal(t, 0, s.InFlight())
+			assert.Equal(t, 0, s.nGranted())
 		})
 	}
 }
@@ -116,7 +116,7 @@ func TestSnake_NHolder_ParallelThroughput(t *testing.T) {
 				if err != nil {
 					continue
 				}
-				cur := int64(s.InFlight())
+				cur := int64(s.nGranted())
 				for {
 					old := maxConcurrent.Load()
 					if cur <= old || maxConcurrent.CompareAndSwap(old, cur) {
@@ -191,7 +191,7 @@ func TestSnake_NHolder_DynamicCapacityDecrease(t *testing.T) {
 	}
 
 	cap.Store(2)
-	assert.Equal(t, 4, s.InFlight(), "existing holders are not evicted")
+	assert.Equal(t, 4, s.nGranted(), "existing holders are not evicted")
 
 	acquired := make(chan struct{})
 	go func() {
@@ -235,7 +235,7 @@ func TestSnake_NHolder_ValveSerialization(t *testing.T) {
 	require.NoError(t, err)
 	u3, err := s.Acquire(t.Context(), "valve-c")
 	require.NoError(t, err)
-	assert.Equal(t, 3, s.InFlight())
+	assert.Equal(t, 3, s.nGranted())
 
 	results := make(chan string, 6)
 	var wg sync.WaitGroup
@@ -266,12 +266,12 @@ func TestSnake_NHolder_ValveSerialization(t *testing.T) {
 		count++
 	}
 	assert.Equal(t, 6, count, "all valve-serialized requests should complete")
-	assert.Equal(t, 0, s.InFlight())
+	assert.Equal(t, 0, s.nGranted())
 }
 
-// --- N-holder: InFlight accuracy under concurrency ---
+// --- N-holder: NGranted accuracy under concurrency ---
 
-func TestSnake_NHolder_InFlightAccuracy(t *testing.T) {
+func TestSnake_NHolder_NGrantedAccuracy(t *testing.T) {
 	const capacity = 5
 	cfg := defaultSnakeConfig()
 	cfg.Capacity = func() int { return capacity }
@@ -288,16 +288,16 @@ func TestSnake_NHolder_InFlightAccuracy(t *testing.T) {
 			if err != nil {
 				return
 			}
-			if s.InFlight() > capacity {
+			if s.nGranted() > capacity {
 				violations.Add(1)
 			}
 			u.Release()
 		}()
 	}
 
 	wg.Wait()
-	assert.Zero(t, violations.Load(), "InFlight must never exceed capacity")
-	assert.Equal(t, 0, s.InFlight())
+	assert.Zero(t, violations.Load(), "NGranted must never exceed capacity")
+	assert.Equal(t, 0, s.nGranted())
 }
 
 // --- N-holder: context cancel frees slot for waiter ---
@@ -369,7 +369,7 @@ func TestSnake_NHolder_MaxAge_MultipleHolders(t *testing.T) {
 	u2.Release()
 
 	assert.Eventually(t, func() bool {
-		return s.InFlight() == 0
+		return s.nGranted() == 0
 	}, 1*time.Second, 5*time.Millisecond)
 }
 
@@ -419,7 +419,6 @@ func TestSnake_NHolder_MemoryCleanup(t *testing.T) {
 
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	assert.Equal(t, 0, s.nGranted)
 	assert.Empty(t, s.holders)
 	assert.Equal(t, 0, s.q.lockedLen())
 	assert.Empty(t, s.q.pendingRequests)

go/vt/vttablet/tabletserver/loadshed/snake_overload_test.go

@@ -140,7 +140,7 @@ func TestSnake_Overload_MassCancel(t *testing.T) {
 
 	// Release original holder
 	unlock.Release()
-	assert.False(t, s.IsLocked())
+	assert.True(t, s.isIdle())
 
 	// Verify the lock is still usable
 	u, err := s.Acquire(t.Context(), "mass-cancel")
@@ -376,7 +376,6 @@ func TestSnake_Memory_CleanBaseline(t *testing.T) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
-	assert.Equal(t, 0, s.nGranted, "nGranted should be 0 after all releases")
 	assert.Empty(t, s.holders, "holders map should be empty after all releases")
 	assert.Equal(t, 0, s.q.lockedLen(), "queue should be empty")
 	assert.Empty(t, s.q.pendingRequests, "pendingRequests map should be empty")
@@ -475,7 +474,7 @@ func TestSnake_MaxAge_VsContextCancel_Race(t *testing.T) {
 		cancel()
 	}
 
-	assert.False(t, s.IsLocked(), "lock must not be stuck after max-age/cancel races")
+	assert.True(t, s.isIdle(), "lock must not be stuck after max-age/cancel races")
 }
 
 // --- Double-signal panic invariant ---
@@ -506,5 +505,5 @@ func TestSnake_NoPanicOnConcurrentCancel(t *testing.T) {
 	unlock.Release()
 	wg.Wait()
 
-	assert.False(t, s.IsLocked())
+	assert.True(t, s.isIdle())
 }

go/vt/vttablet/tabletserver/loadshed/snake_race_test.go

@@ -74,10 +74,10 @@ func TestSnake_CancelVsGrant_Race(t *testing.T) {
 		if acquireErr != nil {
 			// B was cancelled. Verify the lock isn't stuck.
 			s.mu.Lock()
-			inFlightAfterCancel := s.nGranted
+			nGrantedAfterCancel := len(s.holders)
 			s.mu.Unlock()
 
-			if inFlightAfterCancel > 0 {
+			if nGrantedAfterCancel > 0 {
 				leaked.Add(1)
 			}