sidetransport/sender: replace updatesBuf condVar broadcast by signals

This commit replaces the updatesBuf condVar broadcast by multiple signals().
At this point, this should have the exact same behaviour as before. However,
this opens the door to pacing the condVar signalling, to soften the spike
in the number of runnable goroutines after each updatesBuf publish event.

We added two condVars, and two waiters counters. This helps avoid the
condition where we signal to the same goroutine more than once.

With two condVars, and two counters, when we publish a message, we can
atomically swap the condVars and counters such that future goroutines wait
on the other condVar, while we are signalling to all goroutines that were
waiting on the first condVar.

Fixes: #148211

Release note: None

Author: iskettaneh

pkg/kv/kvserver/closedts/sidetransport/sender.go

@@ -539,8 +539,29 @@ func (s *Sender) GetSnapshot() *ctpb.Update {
 type updatesBuf struct {
 	mu struct {
 		syncutil.Mutex
-		// updated is signaled when a new item is inserted.
-		updated sync.Cond
+
+		// We use two condition variables that we atomically swap to avoid signaling
+		// the same goroutine multiple times. Without this, a goroutine could:
+		// 1. Wake up from a signal.
+		// 2. Process its message.
+		// 3. Start waiting again on the same condition variable.
+		// 4. Get signaled again before other waiting goroutines get a chance.
+		//
+		// By swapping to a new condition variable when publishing a message, we
+		// ensure that goroutines can only wait on the new one. This lets us safely
+		// signal all waiters on the old condition variable without racing against
+		// goroutines re-queuing on it.
+		updated1    sync.Cond
+		numWaiters1 int
+
+		// updated2 is signaled when a new item is inserted.
+		updated2    sync.Cond
+		numWaiters2 int
+
+		// activeCondVar is 0 or 1, indicating which conditional variable new
+		// goroutines should be waiting on. This is flipped after each push.
+		activeCondVar int
+
 		// data contains pointers to the Updates.
 		data []*ctpb.Update
 		// head points to the earliest update in the buffer. If the buffer is empty,
@@ -561,15 +582,36 @@ const updatesBufSize = 50
 
 func newUpdatesBuf() *updatesBuf {
 	buf := &updatesBuf{}
-	buf.mu.updated.L = &buf.mu
+	buf.mu.updated1.L = &buf.mu
+	buf.mu.updated2.L = &buf.mu
 	buf.mu.data = make([]*ctpb.Update, updatesBufSize)
 	return buf
 }
 
 // Push adds a new update to the back of the buffer.
 func (b *updatesBuf) Push(ctx context.Context, update *ctpb.Update) {
 	b.mu.Lock()
-	defer b.mu.Unlock()
+
+	// The goal here is that we want to signal all goroutines who were waiting
+	// for the next update. We know that the goroutine that we will signal will
+	// wake up, and then perform some work and wait again for the next update. We
+	// want to avoid the race condition where we can't differentiate between a
+	// goroutine that is waiting for the next update and a goroutine that was
+	// waiting, got woken up, performed some work, and is then waiting for the
+	// next update.
+	var condVar *sync.Cond
+	var numWaiters *int
+	if b.mu.activeCondVar == 0 {
+		condVar = &b.mu.updated1
+		numWaiters = &b.mu.numWaiters1
+	} else {
+		condVar = &b.mu.updated2
+		numWaiters = &b.mu.numWaiters2
+	}
+
+	// At this point, we know that any goroutine that we wake up will wait on the
+	// next conditional variable.
+	b.mu.activeCondVar = (b.mu.activeCondVar + 1) % 2
 
 	// If the buffer is not empty, sanity check the seq num.
 	if b.sizeLocked() != 0 {
@@ -589,7 +631,21 @@ func (b *updatesBuf) Push(ctx context.Context, update *ctpb.Update) {
 
 	// Notify everybody who might have been waiting for this message - we expect
 	// all the connections to be blocked waiting.
-	b.mu.updated.Broadcast()
+	b.mu.Unlock()
+	b.PaceBroadcastUpdate(condVar, numWaiters)
+}
+
+// PaceBroadcastUpdate paces the conditional variable signaling to avoid
+// overloading the system.
+func (b *updatesBuf) PaceBroadcastUpdate(condVar *sync.Cond, numWaiters *int) {
+	// TODO(ibrahim): We can just hook this with the taskPacer.
+	b.mu.Lock()
+	originalNumWaiters := *numWaiters
+	for originalNumWaiters > 0 {
+		condVar.Signal()
+		originalNumWaiters--
+	}
+	b.mu.Unlock()
 }
 
 func (b *updatesBuf) lastIdxLocked() int {
@@ -612,6 +668,23 @@ func (b *updatesBuf) GetBySeq(ctx context.Context, seqNum ctpb.SeqNum) (*ctpb.Up
 	b.mu.Lock()
 	defer b.mu.Unlock()
 
+	var condVar *sync.Cond
+	var numWaiters *int
+	if b.mu.activeCondVar == 0 {
+		condVar = &b.mu.updated1
+		numWaiters = &b.mu.numWaiters1
+	} else {
+		condVar = &b.mu.updated2
+		numWaiters = &b.mu.numWaiters2
+	}
+
+	// Increment the number of waiters on the active conditional variable, and
+	// decrement the same counter when we return.
+	*numWaiters++
+	defer func() {
+		*numWaiters--
+	}()
+
 	// Loop until the requested seqNum is added to the buffer.
 	for {
 		if b.mu.closed {
@@ -630,7 +703,7 @@ func (b *updatesBuf) GetBySeq(ctx context.Context, seqNum ctpb.SeqNum) (*ctpb.Up
 		}
 		// If the requested msg has not been produced yet, block.
 		if seqNum == lastSeq+1 {
-			b.mu.updated.Wait()
+			condVar.Wait()
 			continue
 		}
 		if seqNum > lastSeq+1 {
@@ -666,7 +739,8 @@ func (b *updatesBuf) Close() {
 	b.mu.Lock()
 	defer b.mu.Unlock()
 	b.mu.closed = true
-	b.mu.updated.Broadcast()
+	b.mu.updated1.Broadcast()
+	b.mu.updated2.Broadcast()
 }
 
 // connFactory is capable of creating new connections to specific nodes.

pkg/kv/kvserver/closedts/sidetransport/sender_test.go

@@ -9,6 +9,7 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -884,3 +885,45 @@ func TestRPCConnStopOnClose(t *testing.T) {
 		return nil
 	})
 }
+
+// TestAllUpdatesBufAreSignalled creates a bunch of goroutines that wait on
+// updatesBuf, and then publishes multiple messages to the updatesBuf. It
+// verifies that all goroutines are signalled for all messages, and no goroutine
+// misses any message.
+func TestAllUpdatesBufAreSignalled(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+	connFactory := &mockConnFactory{}
+	st := cluster.MakeTestingClusterSettings()
+	s, stopper := newMockSenderWithSt(connFactory, st)
+	defer stopper.Stop(ctx)
+
+	const numGoroutines = 100
+	const numMsgs = 100
+	var counter atomic.Int64
+
+	// Start 100 goroutines that will wait on sequence numbers 1-100.
+	wg := sync.WaitGroup{}
+	wg.Add(numGoroutines)
+	for i := 0; i < numGoroutines; i++ {
+		go func() {
+			for seqNum := ctpb.SeqNum(1); seqNum <= numMsgs; seqNum++ {
+				_, ok := s.buf.GetBySeq(ctx, seqNum)
+				if ok {
+					counter.Add(1)
+				}
+			}
+			wg.Done()
+		}()
+	}
+
+	// Publish 10 messages.
+	for i := 0; i < numMsgs; i++ {
+		s.publish(ctx)
+	}
+
+	// Verify that eventually all goroutines received all messages.
+	wg.Wait()
+	require.Equal(t, int64(numGoroutines*numMsgs), counter.Load())
+}