kvserver/rangefeed: add unbuffered registration

pkg/kv/kvserver/rangefeed/BUILD.bazel

@@ -20,6 +20,7 @@ go_library(
         "stream_manager.go",
         "task.go",
         "test_helpers.go",
+        "unbuffered_registration.go",
         "unbuffered_sender.go",
     ],
     importpath = "github.com/cockroachdb/cockroach/pkg/kv/kvserver/rangefeed",
@@ -77,6 +78,7 @@ go_test(
         "sender_helper_test.go",
         "stream_manager_test.go",
         "task_test.go",
+        "unbuffered_registration_test.go",
         "unbuffered_sender_test.go",
     ],
     embed = [":rangefeed"],

pkg/kv/kvserver/rangefeed/buffered_registration.go

@@ -350,3 +350,15 @@ var overflowLogEvery = log.Every(5 * time.Second)
 func (br *bufferedRegistration) shouldLogOverflow(checkpointSent bool) bool {
 	return (!checkpointSent) || log.V(1) || overflowLogEvery.ShouldLog()
 }
+
+// Used for testing only.
+func (br *bufferedRegistration) getBuf() chan *sharedEvent {
+	return br.buf
+}
+
+// Used for testing only.
+func (br *bufferedRegistration) getOverflowed() bool {
+	br.mu.Lock()
+	defer br.mu.Unlock()
+	return br.mu.overflowed
+}

pkg/kv/kvserver/rangefeed/buffered_sender_test.go

@@ -71,7 +71,8 @@ func TestBufferedSenderDisconnectStream(t *testing.T) {
 		require.NoError(t, bs.waitForEmptyBuffer(ctx))
 		sm.DisconnectStream(int64(streamID), err)
 		require.NoError(t, bs.waitForEmptyBuffer(ctx))
-		require.Equal(t, 1, testServerStream.totalEventsSent())
+		require.Equalf(t, 1, testServerStream.totalEventsSent(),
+			"expected only 1 error event in %s", testServerStream.String())
 		testRangefeedCounter.waitForRangefeedCount(t, 0)
 	})
 }

pkg/kv/kvserver/rangefeed/processor_helpers_test.go

@@ -23,6 +23,24 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+type testBufferedStream struct {
+	Stream
+}
+
+var _ BufferedStream = (*testBufferedStream)(nil)
+
+func (tb *testBufferedStream) SendBuffered(
+	e *kvpb.RangeFeedEvent, _ *SharedBudgetAllocation,
+) error {
+	// In production code, buffered stream would be responsible for properly using
+	// and releasing the alloc. We just ignore memory accounting here.
+	return tb.SendUnbuffered(e)
+}
+
+func (tb *testBufferedStream) Disconnect(err *kvpb.Error) {
+	tb.Stream.SendError(err)
+}
+
 func makeLogicalOp(val interface{}) enginepb.MVCCLogicalOp {
 	var op enginepb.MVCCLogicalOp
 	op.MustSetValue(val)
@@ -189,11 +207,12 @@ const testProcessorEventCCap = 16
 const testProcessorEventCTimeout = 10 * time.Millisecond
 
 type processorTestHelper struct {
-	span         roachpb.RSpan
-	rts          *resolvedTimestamp
-	syncEventC   func()
-	sendSpanSync func(*roachpb.Span)
-	scheduler    *ClientScheduler
+	span                     roachpb.RSpan
+	rts                      *resolvedTimestamp
+	syncEventC               func()
+	sendSpanSync             func(*roachpb.Span)
+	scheduler                *ClientScheduler
+	toBufferedStreamIfNeeded func(s Stream) Stream
 }
 
 // syncEventAndRegistrations waits for all previously sent events to be
@@ -237,16 +256,23 @@ type rangefeedTestType bool
 
 var (
 	scheduledProcessorWithUnbufferedSender rangefeedTestType
+	scheduledProcessorWithBufferedSender   rangefeedTestType
 )
 
 var testTypes = []rangefeedTestType{
 	scheduledProcessorWithUnbufferedSender,
+	scheduledProcessorWithBufferedSender,
 }
 
-// NB: When adding new types, please keep make sure existing
-// benchmarks will keep their old name.
 func (t rangefeedTestType) String() string {
-	return "scheduled"
+	switch t {
+	case scheduledProcessorWithUnbufferedSender:
+		return "scheduled"
+	case scheduledProcessorWithBufferedSender:
+		return "scheduled/registration=buffered"
+	default:
+		panic("unknown rangefeed test type")
+	}
 }
 
 type testConfig struct {
@@ -426,6 +452,18 @@ func newTestProcessor(
 			p.syncSendAndWait(&syncEvent{c: make(chan struct{}), testRegCatchupSpan: span})
 		}
 		h.scheduler = &p.scheduler
+		switch cfg.feedType {
+		case scheduledProcessorWithUnbufferedSender:
+			h.toBufferedStreamIfNeeded = func(s Stream) Stream {
+				return s
+			}
+		case scheduledProcessorWithBufferedSender:
+			h.toBufferedStreamIfNeeded = func(s Stream) Stream {
+				return &testBufferedStream{Stream: s}
+			}
+		default:
+			panic("unknown rangefeed test type")
+		}
 	default:
 		panic("unknown processor type")
 	}