Merge cockroachdb#126491

126491: kvserver/rangefeed: add rangeID as part of activeStreams map r=nvanbenschoten a=wenyihu6

This patch adds rangeID into the activeStreams map without changing any existing
behavior. The main purpose is to make future commits cleaner. RangeID will be
needed when we enable streamMuxer to shut down all streams.


Part of: cockroachdb#126560
Release note: none

Co-authored-by: Wenyi Hu <[email protected]>

pkg/kv/kvserver/rangefeed/stream_muxer.go

@@ -115,8 +115,8 @@ type StreamMuxer struct {
 	// metrics is used to record rangefeed metrics for the node.
 	metrics RangefeedMetricsRecorder
 
-	// streamID -> context.CancelFunc for active rangefeeds
-	activeStreams syncutil.Map[int64, context.CancelFunc]
+	// streamID -> streamInfo for active rangefeeds
+	activeStreams syncutil.Map[int64, streamInfo]
 
 	// notifyMuxError is a buffered channel of size 1 used to signal the presence
 	// of muxErrors. Additional signals are dropped if the channel is already full
@@ -142,12 +142,24 @@ func NewStreamMuxer(sender ServerStreamSender, metrics RangefeedMetricsRecorder)
 	}
 }
 
+// streamInfo contains the rangeID and cancel function for an active rangefeed.
+// It should be treated as immutable.
+type streamInfo struct {
+	rangeID roachpb.RangeID
+	cancel  context.CancelFunc
+}
+
 // AddStream registers a server rangefeed stream with the StreamMuxer. It
 // remains active until DisconnectStreamWithError is called with the same
 // streamID. Caller must ensure no duplicate stream IDs are added without
 // disconnecting the old one first.
-func (sm *StreamMuxer) AddStream(streamID int64, cancel context.CancelFunc) {
-	if _, loaded := sm.activeStreams.LoadOrStore(streamID, &cancel); loaded {
+func (sm *StreamMuxer) AddStream(
+	streamID int64, rangeID roachpb.RangeID, cancel context.CancelFunc,
+) {
+	if _, loaded := sm.activeStreams.LoadOrStore(streamID, &streamInfo{
+		rangeID: rangeID,
+		cancel:  cancel,
+	}); loaded {
 		log.Fatalf(context.Background(), "stream %d already exists", streamID)
 	}
 	sm.metrics.UpdateMetricsOnRangefeedConnect()
@@ -201,8 +213,9 @@ func (sm *StreamMuxer) appendMuxError(e *kvpb.MuxRangeFeedEvent) {
 func (sm *StreamMuxer) DisconnectStreamWithError(
 	streamID int64, rangeID roachpb.RangeID, err *kvpb.Error,
 ) {
-	if cancelFunc, ok := sm.activeStreams.LoadAndDelete(streamID); ok {
-		(*cancelFunc)()
+	if stream, ok := sm.activeStreams.LoadAndDelete(streamID); ok {
+		// Fine to skip nil checking here since that would be a programming error.
+		stream.cancel()
 		clientErrorEvent := transformRangefeedErrToClientError(err)
 		ev := &kvpb.MuxRangeFeedEvent{
 			StreamID: streamID,

pkg/kv/kvserver/rangefeed/stream_muxer_test.go

@@ -48,7 +48,7 @@ func TestStreamMuxer(t *testing.T) {
 		const streamID = 0
 		const rangeID = 1
 		streamCtx, cancel := context.WithCancel(context.Background())
-		muxer.AddStream(0, cancel)
+		muxer.AddStream(streamID, rangeID, cancel)
 		// Note that kvpb.NewError(nil) == nil.
 		require.Equal(t, testRangefeedCounter.get(), int32(1))
 		muxer.DisconnectStreamWithError(streamID, rangeID, kvpb.NewError(nil))
@@ -86,7 +86,7 @@ func TestStreamMuxer(t *testing.T) {
 		require.Equal(t, testRangefeedCounter.get(), int32(0))
 
 		for _, muxError := range testRangefeedCompletionErrors {
-			muxer.AddStream(muxError.streamID, func() {})
+			muxer.AddStream(muxError.streamID, muxError.rangeID, func() {})
 		}
 
 		require.Equal(t, testRangefeedCounter.get(), int32(3))
@@ -140,7 +140,7 @@ func TestStreamMuxerOnBlockingIO(t *testing.T) {
 	const streamID = 0
 	const rangeID = 1
 	streamCtx, streamCancel := context.WithCancel(context.Background())
-	muxer.AddStream(0, streamCancel)
+	muxer.AddStream(0, rangeID, streamCancel)
 
 	ev := &kvpb.MuxRangeFeedEvent{
 		StreamID: streamID,

pkg/server/node.go

@@ -2027,7 +2027,7 @@ func (n *Node) MuxRangeFeed(stream kvpb.Internal_MuxRangeFeedServer) error {
 				streamID: req.StreamID,
 				wrapped:  streamMuxer,
 			}
-			streamMuxer.AddStream(req.StreamID, cancel)
+			streamMuxer.AddStream(req.StreamID, req.RangeID, cancel)
 
 			// Rangefeed attempts to register rangefeed a request over the specified
 			// span. If registration fails, it returns an error. Otherwise, it returns