rangefeed: pop multiple events from queue in buffered sender

This is a small optimization that is perhaps a step toward something
like #152359.

Epic: none
Release note: None

Author: stevendanna

pkg/kv/kvserver/rangefeed/buffered_sender.go

@@ -73,6 +73,9 @@ type BufferedSender struct {
 	// Note that lockedMuxStream wraps the underlying grpc server stream, ensuring
 	// thread safety.
 	sender ServerStreamSender
+	// sendBufSize is the number of items we pop from the queue at a time. Exposed for
+	// testing.
+	sendBufSize int
 
 	// queueMu protects the buffer queue.
 	queueMu struct {
@@ -96,6 +99,8 @@ type BufferedSender struct {
 	metrics *BufferedSenderMetrics
 }
 
+const defaultSendBufSize = 64
+
 type streamState int64
 
 const (
@@ -119,8 +124,9 @@ func NewBufferedSender(
 	sender ServerStreamSender, settings *cluster.Settings, bsMetrics *BufferedSenderMetrics,
 ) *BufferedSender {
 	bs := &BufferedSender{
-		sender:  sender,
-		metrics: bsMetrics,
+		sendBufSize: defaultSendBufSize,
+		sender:      sender,
+		metrics:     bsMetrics,
 	}
 	bs.notifyDataC = make(chan struct{}, 1)
 	bs.queueMu.buffer = newEventQueue()
@@ -219,6 +225,8 @@ func (bs *BufferedSender) sendUnbuffered(ev *kvpb.MuxRangeFeedEvent) error {
 func (bs *BufferedSender) run(
 	ctx context.Context, stopper *stop.Stopper, onError func(streamID int64),
 ) error {
+	eventsBuf := make([]sharedMuxEvent, 0, bs.sendBufSize)
+
 	for {
 		select {
 		case <-ctx.Done():
@@ -231,39 +239,47 @@ func (bs *BufferedSender) run(
 			return nil
 		case <-bs.notifyDataC:
 			for {
-				e, success := bs.popFront()
-				if !success {
+				eventsBuf = bs.popEvents(eventsBuf[:0], bs.sendBufSize)
+				if len(eventsBuf) == 0 {
 					break
 				}
 
-				bs.metrics.BufferedSenderQueueSize.Dec(1)
-				err := bs.sender.Send(e.ev)
-				e.alloc.Release(ctx)
-				if e.ev.Error != nil {
-					onError(e.ev.StreamID)
-				}
-				if err != nil {
-					return err
+				bs.metrics.BufferedSenderQueueSize.Dec(int64(len(eventsBuf)))
+				for _, evt := range eventsBuf {
+					// TODO(ssd): This might be another location where we could transform
+					// multiple events into BulkEvents. We can't just throw them all in a
+					// bulk event though since we are processing events for different
+					// streams here.
+					err := bs.sender.Send(evt.ev)
+					evt.alloc.Release(ctx)
+					if evt.ev.Error != nil {
+						onError(evt.ev.StreamID)
+					}
+					if err != nil {
+						return err
+					}
 				}
+				clear(eventsBuf) // Clear so referenced MuxRangeFeedEvents can be GC'd.
 			}
 		}
 	}
 }
 
-// popFront pops the front event from the buffer queue. It returns the event and
-// a boolean indicating if the event was successfully popped.
-func (bs *BufferedSender) popFront() (e sharedMuxEvent, success bool) {
+// popEvents appends up to eventsToPop events into dest, returning the appended slice.
+func (bs *BufferedSender) popEvents(dest []sharedMuxEvent, eventsToPop int) []sharedMuxEvent {
 	bs.queueMu.Lock()
 	defer bs.queueMu.Unlock()
-	event, ok := bs.queueMu.buffer.popFront()
-	if ok {
+	dest = bs.queueMu.buffer.popFrontInto(dest, eventsToPop)
+
+	// Update accounting for everything we popped.
+	for _, event := range dest {
 		state, streamFound := bs.queueMu.byStream[event.ev.StreamID]
 		if streamFound {
 			state.queueItems--
 			bs.queueMu.byStream[event.ev.StreamID] = state
 		}
 	}
-	return event, ok
+	return dest
 }
 
 // addStream initializes the per-stream tracking for the given streamID.

pkg/kv/kvserver/rangefeed/buffered_sender_test.go

@@ -198,12 +198,14 @@ func TestBufferedSenderOnOverflow(t *testing.T) {
 		require.NoError(t, bs.sendBuffered(muxEv, nil))
 	}
 	require.Equal(t, queueCap, int64(bs.len()))
-	e, success := bs.popFront()
+	dest := bs.popEvents(nil, 1)
+	require.Equal(t, 1, len(dest))
+	e := dest[0]
 	require.Equal(t, sharedMuxEvent{
 		ev:    muxEv,
 		alloc: nil,
 	}, e)
-	require.True(t, success)
+
 	require.Equal(t, queueCap-1, int64(bs.len()))
 	require.NoError(t, bs.sendBuffered(muxEv, nil))
 	require.Equal(t, queueCap, int64(bs.len()))
@@ -229,6 +231,7 @@ func TestBufferedSenderOnOverflowMultiStream(t *testing.T) {
 	queueCap := int64(24)
 	RangefeedSingleBufferedSenderQueueMaxPerReg.Override(ctx, &st.SV, queueCap)
 	bs := NewBufferedSender(testServerStream, st, NewBufferedSenderMetrics())
+	bs.sendBufSize = 1 // Test requires knowing exactly when the queue will fill.
 	require.Equal(t, queueCap, bs.queueMu.perStreamCapacity)
 
 	sm := NewStreamManager(bs, smMetrics)

pkg/kv/kvserver/rangefeed/event_queue.go

@@ -79,21 +79,38 @@ func (q *eventQueue) pushBack(e sharedMuxEvent) {
 	q.size++
 }
 
-func (q *eventQueue) popFront() (sharedMuxEvent, bool) {
-	if q.size == 0 {
-		return sharedMuxEvent{}, false
+// popFrontInto appends up to eventsToPop events into dest.
+func (q *eventQueue) popFrontInto(dest []sharedMuxEvent, eventsToPop int) []sharedMuxEvent {
+	if eventsToPop == 0 || q.size == 0 {
+		return dest
 	}
+
 	if q.read == eventQueueChunkSize {
+		assertTrue(q.first.nextChunk != nil, "nextChunk should be non-nil")
 		removed := q.first
 		q.first = q.first.nextChunk
 		putPooledQueueChunk(removed)
 		q.read = 0
 	}
-	res := q.first.data[q.read]
-	q.first.data[q.read] = sharedMuxEvent{}
-	q.read++
-	q.size--
-	return res, true
+
+	// We only read out of the current chunk. We could loop until we've reached
+	// eventsToPop.
+	availableInChunk := eventQueueChunkSize - q.read
+	if q.first == q.last {
+		// Last chunk - only up to write position.
+		availableInChunk = q.write - q.read
+	}
+
+	if eventsToPop > availableInChunk {
+		eventsToPop = availableInChunk
+	}
+
+	dest = append(dest, q.first.data[q.read:q.read+eventsToPop]...)
+	clear(q.first.data[q.read : q.read+eventsToPop])
+
+	q.read += eventsToPop
+	q.size -= eventsToPop
+	return dest
 }
 
 // free drops references held by the queue.