kvserver/rangefeed: add capacity to node level buffered sender

This patch adds capacity to node level buffered sender which will shut down all
registrations if the node level buffer had overflowed.

Part of: #129813
Release note: none

Author: wenyihu6

pkg/kv/kvserver/rangefeed/buffered_sender.go

@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/retry"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
@@ -56,8 +57,10 @@ type BufferedSender struct {
 	// queueMu protects the buffer queue.
 	queueMu struct {
 		syncutil.Mutex
-		stopped bool
-		buffer  *eventQueue
+		stopped  bool
+		buffer   *eventQueue
+		capacity int64
+		overflow bool
 	}
 
 	// notifyDataC is used to notify the BufferedSender.run goroutine that there
@@ -72,14 +75,16 @@ type BufferedSender struct {
 }
 
 func NewBufferedSender(
-	sender ServerStreamSender, bsMetrics *BufferedSenderMetrics,
+	sender ServerStreamSender, bsMetrics *BufferedSenderMetrics, maxQueueSize int64,
 ) *BufferedSender {
 	bs := &BufferedSender{
 		sender:  sender,
 		metrics: bsMetrics,
 	}
 	bs.queueMu.buffer = newEventQueue()
 	bs.notifyDataC = make(chan struct{}, 1)
+	bs.queueMu.buffer = newEventQueue()
+	bs.queueMu.capacity = maxQueueSize
 	return bs
 }
 
@@ -95,6 +100,15 @@ func (bs *BufferedSender) sendBuffered(
 	if bs.queueMu.stopped {
 		return errors.New("stream sender is stopped")
 	}
+	if bs.queueMu.overflow {
+		// Is this too spammy
+		log.Dev.Error(context.Background(), "buffer capacity exceeded")
+		return newRetryErrBufferCapacityExceeded()
+	}
+	if bs.queueMu.buffer.len() >= bs.queueMu.capacity {
+		bs.queueMu.overflow = true
+		return newRetryErrBufferCapacityExceeded()
+	}
 	// TODO(wenyihu6): pass an actual context here
 	alloc.Use(context.Background())
 	bs.queueMu.buffer.pushBack(sharedMuxEvent{ev, alloc})
@@ -130,7 +144,7 @@ func (bs *BufferedSender) run(
 			return nil
 		case <-bs.notifyDataC:
 			for {
-				e, success := bs.popFront()
+				e, success, overflowed, remains := bs.popFront()
 				if !success {
 					break
 				}
@@ -143,18 +157,26 @@ func (bs *BufferedSender) run(
 				if err != nil {
 					return err
 				}
+				if overflowed && remains == int64(0) {
+					return newRetryErrBufferCapacityExceeded()
+				}
 			}
 		}
 	}
 }
 
 // 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) {
+func (bs *BufferedSender) popFront() (
+	e sharedMuxEvent,
+	success bool,
+	overflowed bool,
+	remains int64,
+) {
 	bs.queueMu.Lock()
 	defer bs.queueMu.Unlock()
 	event, ok := bs.queueMu.buffer.popFront()
-	return event, ok
+	return event, ok, bs.queueMu.overflow, bs.queueMu.buffer.len()
 }
 
 // cleanup is called when the sender is stopped. It is expected to free up

pkg/kv/kvserver/rangefeed/buffered_sender_test.go

@@ -34,7 +34,7 @@ func TestBufferedSenderDisconnectStream(t *testing.T) {
 	defer stopper.Stop(ctx)
 	testServerStream := newTestServerStream()
 	smMetrics := NewStreamManagerMetrics()
-	bs := NewBufferedSender(testServerStream, NewBufferedSenderMetrics())
+	bs := NewBufferedSender(testServerStream, NewBufferedSenderMetrics(), 1000)
 	sm := NewStreamManager(bs, smMetrics)
 	require.NoError(t, sm.Start(ctx, stopper))
 	defer sm.Stop(ctx)
@@ -87,7 +87,7 @@ func TestBufferedSenderChaosWithStop(t *testing.T) {
 	testServerStream := newTestServerStream()
 
 	smMetrics := NewStreamManagerMetrics()
-	bs := NewBufferedSender(testServerStream, NewBufferedSenderMetrics())
+	bs := NewBufferedSender(testServerStream, NewBufferedSenderMetrics(), 1000)
 	sm := NewStreamManager(bs, smMetrics)
 	require.NoError(t, sm.Start(ctx, stopper))
 

pkg/kv/kvserver/rangefeed/stream_manager_test.go

@@ -180,7 +180,7 @@ func TestStreamManagerErrorHandling(t *testing.T) {
 		case scheduledProcessorWithUnbufferedSender:
 			s = NewUnbufferedSender(testServerStream)
 		case scheduledProcessorWithBufferedSender:
-			s = NewBufferedSender(testServerStream, NewBufferedSenderMetrics())
+			s = NewBufferedSender(testServerStream, NewBufferedSenderMetrics(), 1000)
 		default:
 			t.Fatalf("unknown rangefeed test type %v", rt)
 		}

pkg/kv/kvserver/rangefeed/unbuffered_registration_test.go

@@ -34,7 +34,7 @@ func TestUnbufferedRegWithStreamManager(t *testing.T) {
 	defer stopper.Stop(ctx)
 	testServerStream := newTestServerStream()
 	smMetrics := NewStreamManagerMetrics()
-	bs := NewBufferedSender(testServerStream, NewBufferedSenderMetrics())
+	bs := NewBufferedSender(testServerStream, NewBufferedSenderMetrics(), 1000)
 	sm := NewStreamManager(bs, smMetrics)
 	require.NoError(t, sm.Start(ctx, stopper))
 
@@ -106,7 +106,7 @@ func TestUnbufferedRegCorrectnessOnDisconnect(t *testing.T) {
 	defer stopper.Stop(ctx)
 	testServerStream := newTestServerStream()
 	smMetrics := NewStreamManagerMetrics()
-	bs := NewBufferedSender(testServerStream, NewBufferedSenderMetrics())
+	bs := NewBufferedSender(testServerStream, NewBufferedSenderMetrics(), 1000)
 	sm := NewStreamManager(bs, smMetrics)
 	require.NoError(t, sm.Start(ctx, stopper))
 	defer sm.Stop(ctx)

pkg/kv/kvserver/replica_rangefeed.go

@@ -93,6 +93,15 @@ var RangefeedUseBufferedSender = settings.RegisterBoolSetting(
 	metamorphic.ConstantWithTestBool("kv.rangefeed.buffered_sender.enabled", true),
 )
 
+var RangefeedSingleBufferedSenderQueueMaxSize = settings.RegisterIntSetting(
+	settings.SystemOnly,
+	"kv.rangefeed.buffered_sender_queue.max_size",
+	"max size of the buffered sender queue",
+	// what should we set here a single buffered sender here for one Rangefeed
+	// call
+	2048,
+)
+
 func init() {
 	// Inject into kvserverbase to allow usage from kvcoord.
 	kvserverbase.RangeFeedRefreshInterval = RangeFeedRefreshInterval

pkg/server/node.go

@@ -2224,8 +2224,8 @@ func (n *Node) muxRangeFeed(muxStream kvpb.RPCInternal_MuxRangeFeedStream) error
 
 	sm := &rangefeed.StreamManager{}
 	if kvserver.RangefeedUseBufferedSender.Get(&n.storeCfg.Settings.SV) {
-		sm = rangefeed.NewStreamManager(rangefeed.NewBufferedSender(lockedMuxStream, n.metrics.BufferedSenderMetrics),
-			n.metrics.StreamManagerMetrics)
+		sm = rangefeed.NewStreamManager(rangefeed.NewBufferedSender(lockedMuxStream, n.metrics.BufferedSenderMetrics,
+			kvserver.RangefeedSingleBufferedSenderQueueMaxSize.Get(&n.storeCfg.Settings.SV)), n.metrics.StreamManagerMetrics)
 	} else {
 		sm = rangefeed.NewStreamManager(rangefeed.NewUnbufferedSender(lockedMuxStream), n.metrics.StreamManagerMetrics)
 	}