kvserver/rangefeed: add unbuffered registration

This patch adds unbufferedRegistration. UnbufferedRegistration is like BufferedRegistration but uses BufferedSender to buffer live raft updates instead of a using fixed size channel and having a dedicated per-registration goroutine to volley events to underlying gRPC stream. Instead, there is only one BufferedSender for each incoming node.MuxRangefeed gRPC call. BufferedSender is responsible for buffering and sending its updates to the underlying gRPC stream in a dedicated goroutine O(node). Resolved: cockroachdb#110432 Release note: A new cluster setting `kv.rangefeed.buffered_stream_sender.enabled` can now be used to allow rangefeed to use buffered sender for all rangefeed feeds instead of buffering events separately per client per range. Co-authored-by: Steven Danna <[email protected]>
wenyihu6 · Dec 4, 2024 · 7351476 · 7351476
1 parent 7da80b1
commit 7351476
Show file tree

Hide file tree

Showing 11 changed files with 1,655 additions and 470 deletions.
diff --git a/pkg/kv/kvserver/rangefeed/BUILD.bazel b/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"],

diff --git a/pkg/kv/kvserver/rangefeed/buffered_registration.go b/pkg/kv/kvserver/rangefeed/buffered_registration.go
@@ -347,3 +347,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
+}
diff --git a/pkg/kv/kvserver/rangefeed/buffered_sender_test.go b/pkg/kv/kvserver/rangefeed/buffered_sender_test.go
@@ -7,12 +7,14 @@ package rangefeed
 
 import (
 	"context"
+	"fmt"
 	"sync"
 	"sync/atomic"
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -71,7 +73,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)
 	})
 }
@@ -156,3 +159,106 @@ func TestBufferedSenderChaosWithStop(t *testing.T) {
 		require.Equal(t, 0, bs.len())
 	})
 }
+
+func TestBufferedSenderErrorHandling(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+
+	testutils.RunValues(t, "feed type", testTypes, func(t *testing.T, rt rangefeedTestType) {
+		testServerStream := newTestServerStream()
+		testRangefeedCounter := newTestRangefeedCounter()
+		var s sender
+		switch rt {
+		case scheduledProcessorWithUnbufferedSender:
+			s = NewUnbufferedSender(testServerStream)
+		case scheduledProcessorWithBufferedSender:
+			s = NewBufferedSender(testServerStream)
+		default:
+			t.Fatalf("unknown rangefeed test type %v", rt)
+		}
+
+		sm := NewStreamManager(s, testRangefeedCounter)
+		stopper := stop.NewStopper()
+		defer stopper.Stop(ctx)
+		require.NoError(t, sm.Start(ctx, stopper))
+		const sID, rID = int64(0), 1
+		disconnectErr := kvpb.NewError(fmt.Errorf("disconnection error"))
+
+		expectErrorHandlingInvariance := func(p Processor) {
+			testRangefeedCounter.waitForRangefeedCount(t, 0)
+			testutils.SucceedsSoon(t, func() error {
+				if p.Len() == 0 {
+					return nil
+				}
+				return errors.Newf("expected 0 registrations, found %d", p.Len())
+			})
+			testServerStream.waitForEvent(t, makeMuxRangefeedErrorEvent(sID, rID, disconnectErr))
+			require.Equalf(t, 1, testServerStream.totalEventsFilterBy(
+				func(e *kvpb.MuxRangeFeedEvent) bool {
+					return e.Error != nil
+				}), "expected only 1 error event in %s", testServerStream.String())
+		}
+		t.Run("Fail to register rangefeed with the processor", func(t *testing.T) {
+			p, _, stopper := newTestProcessor(t, withRangefeedTestType(rt))
+			defer stopper.Stop(ctx)
+			sm.NewStream(sID, rID)
+			// We mock failed registration by not calling p.Register.
+			// node.MuxRangefeed would call sendBuffered with error event.
+			require.NoError(t, sm.sender.sendBuffered(makeMuxRangefeedErrorEvent(sID, rID, disconnectErr), nil))
+			expectErrorHandlingInvariance(p)
+			testServerStream.reset()
+		})
+		t.Run("Disconnect stream after registration with processor but before adding to stream manager",
+			func(t *testing.T) {
+				p, h, stopper := newTestProcessor(t, withRangefeedTestType(rt))
+				defer stopper.Stop(ctx)
+				stream := sm.NewStream(sID, rID)
+				registered, d, _ := p.Register(ctx, h.span, hlc.Timestamp{}, nil, /* catchUpIter */
+					false /* withDiff */, false /* withFiltering */, false, /* withOmitRemote */
+					stream)
+				require.True(t, registered)
+				go p.StopWithErr(disconnectErr)
+				require.Equal(t, 0, testRangefeedCounter.get())
+				sm.AddStream(sID, d)
+				expectErrorHandlingInvariance(p)
+				testServerStream.reset()
+			})
+		t.Run("Disconnect stream after registration with processor and stream manager", func(t *testing.T) {
+			stream := sm.NewStream(sID, rID)
+			p, h, stopper := newTestProcessor(t, withRangefeedTestType(rt))
+			defer stopper.Stop(ctx)
+			registered, d, _ := p.Register(ctx, h.span, hlc.Timestamp{}, nil, /* catchUpIter */
+				false /* withDiff */, false /* withFiltering */, false, /* withOmitRemote */
+				stream)
+			require.True(t, registered)
+			sm.AddStream(sID, d)
+			require.Equal(t, 1, p.Len())
+			require.Equal(t, 1, testRangefeedCounter.get())
+			sm.DisconnectStream(sID, disconnectErr)
+			expectErrorHandlingInvariance(p)
+			testServerStream.reset()
+		})
+		t.Run("Stream manager disconnects everything", func(t *testing.T) {
+			stream := sm.NewStream(sID, rID)
+			p, h, stopper := newTestProcessor(t, withRangefeedTestType(rt))
+			defer stopper.Stop(ctx)
+			registered, d, _ := p.Register(ctx, h.span, hlc.Timestamp{}, nil, /* catchUpIter */
+				false /* withDiff */, false /* withFiltering */, false, /* withOmitRemote */
+				stream)
+			require.True(t, registered)
+			sm.AddStream(sID, d)
+			require.Equal(t, 1, testRangefeedCounter.get())
+			require.Equal(t, 1, p.Len())
+			sm.Stop(ctx)
+			// No disconnect events should be sent during Stop().
+			testRangefeedCounter.waitForRangefeedCount(t, 0)
+			testutils.SucceedsSoon(t, func() error {
+				if p.Len() == 0 {
+					return nil
+				}
+				return errors.Newf("expected 0 registrations, found %d", p.Len())
+			})
+		})
+	})
+}
diff --git a/pkg/kv/kvserver/rangefeed/processor_helpers_test.go b/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")
 	}