Fix spurious per-result waiter count and consolidate counter management

When the global ConsolidatorQueryWaiterCap was hit by a caller of a
different query, the current query's per-result waiter count was left
spuriously positive, causing unnecessary proto3 row caching. Fix by
decrementing the per-result count alongside the global count for all
non-original callers.

Also refactors counter management: AddWaiterCounter now increments both
the per-result and global counters (removing AddPerResultWaiterCounter),
and TotalWaiterCount is a read method on the Consolidator interface
instead of being read through AddWaiterCounter's return value.

Note: the two counter increments in AddWaiterCounter are not jointly
atomic. This is benign — the leader checks HasWaiters() before
Broadcast(), so it always sees the pre-decrement state. The only effect
is momentary soft-limit imprecision on ConsolidatorQueryWaiterCap, which
is a soft cap anyway.

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

Co-Authored-By: Claude <svc-devxp-claude@slack-corp.com>
Signed-off-by: Brett Wines <bwines@slack-corp.com>

Author: bgwines

go/sync2/consolidator.go

@@ -30,6 +30,7 @@ type Consolidator interface {
 	Create(string) (PendingResult, bool)
 	Items() []ConsolidatorCacheItem
 	Record(query string)
+	TotalWaiterCount() int64
 }
 
 // PendingResult is a wrapper for result of a query.
@@ -41,7 +42,7 @@ type PendingResult interface {
 	Result() *sqltypes.Result
 	Wait()
 	HasWaiters() bool
-	AddWaiterCounter(int64) *int64
+	AddWaiterCounter(int64)
 }
 
 type consolidator struct {
@@ -64,12 +65,12 @@ type pendingResult struct {
 	// executing is used to block additional requests.
 	// The original request holds a write lock while additional ones are blocked
 	// on acquiring a read lock (see Wait() below.)
-	executing            sync.RWMutex
-	consolidator         *consolidator
-	query                string
-	result               *sqltypes.Result
-	err                  error
-	perResultWaiterCount atomic.Int64
+	executing    sync.RWMutex
+	consolidator *consolidator
+	query        string
+	result       *sqltypes.Result
+	err          error
+	waiterCount  atomic.Int64
 }
 
 // Create adds a query to currently executing queries and acquires a
@@ -80,7 +81,6 @@ func (co *consolidator) Create(query string) (PendingResult, bool) {
 	defer co.mu.Unlock()
 	var r *pendingResult
 	if r, ok := co.queries[query]; ok {
-		r.perResultWaiterCount.Add(1)
 		r.AddWaiterCounter(1)
 		return r, false
 	}
@@ -121,7 +121,7 @@ func (rs *pendingResult) SetResult(res *sqltypes.Result) {
 }
 
 func (rs *pendingResult) HasWaiters() bool {
-	return rs.perResultWaiterCount.Load() > 0
+	return rs.waiterCount.Load() > 0
 }
 
 // Wait waits for the original query to complete execution. Wait should
@@ -131,9 +131,11 @@ func (rs *pendingResult) Wait() {
 	rs.executing.RLock()
 }
 
-func (rs *pendingResult) AddWaiterCounter(c int64) *int64 {
+func (rs *pendingResult) AddWaiterCounter(c int64) {
+	// Non-atomic pair is benign: ConsolidatorQueryWaiterCap is a soft limit and
+	// the per-waiter count is only checked before Broadcast().
+	rs.waiterCount.Add(c)
 	atomic.AddInt64(rs.consolidator.totalWaiterCount, c)
-	return rs.consolidator.totalWaiterCount
 }
 
 // ConsolidatorCache is a thread-safe object used for counting how often recent
@@ -145,6 +147,10 @@ type ConsolidatorCache struct {
 	totalWaiterCount *int64
 }
 
+func (cc *ConsolidatorCache) TotalWaiterCount() int64 {
+	return atomic.LoadInt64(cc.totalWaiterCount)
+}
+
 // NewConsolidatorCache creates a new cache with the given capacity.
 func NewConsolidatorCache(capacity int64) *ConsolidatorCache {
 	return &ConsolidatorCache{cache.NewLRUCache[*ccount](capacity), new(int64)}

go/sync2/consolidator_test.go

@@ -49,8 +49,8 @@ func TestAddWaiterCount(t *testing.T) {
 	wgAdd.Wait()
 	wgSub.Wait()
 
-	if *pr.AddWaiterCounter(0) != 0 {
-		t.Fatalf("Expect 0 totalWaiterCount but got: %v", *pr.AddWaiterCounter(0))
+	if con.TotalWaiterCount() != 0 {
+		t.Fatalf("Expect 0 totalWaiterCount but got: %d", con.TotalWaiterCount())
 	}
 }
 

go/sync2/fake_consolidator.go

@@ -30,7 +30,8 @@ type FakeConsolidator struct {
 	// CreateReturnCreated pre-configures the return value of Create calls.
 	CreateReturn *FakeConsolidatorCreateReturn
 	// RecordCalls can be usd to inspect Record calls.
-	RecordCalls []string
+	RecordCalls      []string
+	totalWaiterCount int64
 }
 
 // FakeConsolidatorCreateReturn wraps the two return values of a call to
@@ -53,10 +54,10 @@ type FakePendingResult struct {
 	WaitCalls int
 	// AddWaiterCounterCalls can be used to inspect AddWaiterCounter calls.
 	AddWaiterCounterCalls []int64
-	WaiterCount int64
-	PerResultWaiterCount int64
-	err                  error
-	result               *sqltypes.Result
+	WaiterCount           int64
+	Consolidator          *FakeConsolidator
+	err                   error
+	result                *sqltypes.Result
 }
 
 var (
@@ -87,6 +88,14 @@ func (fc *FakeConsolidator) Items() []ConsolidatorCacheItem {
 	return nil
 }
 
+func (fc *FakeConsolidator) TotalWaiterCount() int64 {
+	return fc.totalWaiterCount
+}
+
+func (fc *FakeConsolidator) SetTotalWaiterCount(count int64) {
+	fc.totalWaiterCount = count
+}
+
 // Broadcast records the Broadcast call for later verification.
 func (fr *FakePendingResult) Broadcast() {
 	fr.BroadcastCalls++
@@ -118,12 +127,12 @@ func (fr *FakePendingResult) Wait() {
 }
 
 func (fr *FakePendingResult) HasWaiters() bool {
-	return fr.PerResultWaiterCount > 0
+	return fr.WaiterCount > 0
 }
 
 // AddWaiterCounter records the call and simulates waiter count changes.
-func (fr *FakePendingResult) AddWaiterCounter(delta int64) *int64 {
+func (fr *FakePendingResult) AddWaiterCounter(delta int64) {
 	fr.AddWaiterCounterCalls = append(fr.AddWaiterCounterCalls, delta)
 	fr.WaiterCount += delta
-	return &fr.WaiterCount
+	fr.Consolidator.totalWaiterCount += delta
 }

go/vt/vttablet/tabletserver/query_executor.go

@@ -761,7 +761,7 @@ func (qre *QueryExecutor) execSelect() (*sqltypes.Result, error) {
 			}
 		} else {
 			waiterCap := qre.tsv.config.ConsolidatorQueryWaiterCap
-			if waiterCap == 0 || *q.AddWaiterCounter(0) <= waiterCap {
+			if waiterCap == 0 || qre.tsv.qe.consolidator.TotalWaiterCount() <= waiterCap {
 				qre.logStats.QuerySources |= tabletenv.QuerySourceConsolidator
 				startTime := time.Now()
 				q.Wait()

go/vt/vttablet/tabletserver/query_executor_test.go

@@ -1423,7 +1423,7 @@ func TestQueryExecutorShouldConsolidate(t *testing.T) {
 
 			// Set up consolidator pre-conditions.
 
-			fakePendingResult := &sync2.FakePendingResult{}
+			fakePendingResult := &sync2.FakePendingResult{Consolidator: fakeConsolidator}
 			fakePendingResult.SetResult(result)
 			fakeConsolidator.CreateReturn = &sync2.FakeConsolidatorCreateReturn{
 				Created:       !tcase.consolidatorHasIdenticalQuery,
@@ -1494,10 +1494,11 @@ func TestQueryExecutorConsolidatorWaiterCapFallback(t *testing.T) {
 	}
 
 	// Set up consolidator to simulate an identical query already running (Created=false)
-	fakePendingResult := &sync2.FakePendingResult{}
+	fakePendingResult := &sync2.FakePendingResult{Consolidator: fakeConsolidator}
 	fakePendingResult.SetResult(result)
 	// Start with waiter count above the cap (2 > 1), so the condition fails
 	fakePendingResult.WaiterCount = 2
+	fakeConsolidator.SetTotalWaiterCount(2)
 
 	fakeConsolidator.CreateReturn = &sync2.FakeConsolidatorCreateReturn{
 		Created:       false, // Simulate identical query already running
@@ -1528,10 +1529,9 @@ func TestQueryExecutorConsolidatorWaiterCapFallback(t *testing.T) {
 	// Verify we did NOT broadcast (because we're not the original)
 	require.Equal(t, 0, fakePendingResult.BroadcastCalls)
 
-	// Verify AddWaiterCounter was called: once with 0 (to check count), once with -1 (cleanup)
-	require.Len(t, fakePendingResult.AddWaiterCounterCalls, 2)
-	require.Equal(t, int64(0), fakePendingResult.AddWaiterCounterCalls[0])  // Check current count
-	require.Equal(t, int64(-1), fakePendingResult.AddWaiterCounterCalls[1]) // Decrement
+	// Verify AddWaiterCounter was called once with -1 (cleanup)
+	require.Len(t, fakePendingResult.AddWaiterCounterCalls, 1)
+	require.Equal(t, int64(-1), fakePendingResult.AddWaiterCounterCalls[0])
 
 	// Verify fallback executed the query independently
 	require.Equal(t, 1, db.GetQueryCalledNum(input))