kvserver/rangefeed: move active streams to stream muxer

This patch moves the existing active stream management to StreamMuxer
without changing any existing behaviour. The main purpose is to make future
commits cleaner.

Part of: cockroachdb#126561
Epic: none
Release note: none

pkg/kv/kvserver/rangefeed/stream_muxer.go

@@ -12,8 +12,11 @@ package rangefeed
 
 import (
 	"context"
+	"sync"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 )
@@ -29,7 +32,10 @@ type severStreamSender interface {
 type StreamMuxer struct {
 	// Note that lockedMuxStream wraps the underlying grpc server stream, ensuring
 	// thread safety.
-	sender  severStreamSender
+	sender severStreamSender
+
+	// streamID -> streamInfo for active rangefeeds
+	activeStreams sync.Map
 
 	// 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
@@ -49,17 +55,68 @@ func NewStreamMuxer(sender severStreamSender) *StreamMuxer {
 	}
 }
 
-func (sm *StreamMuxer) AppendMuxError(e *kvpb.MuxRangeFeedEvent) {
+// AddStream registers a server rangefeed stream with the muxer. It remains
+// active until DisconnectRangefeedWithError 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 {
+		log.Fatalf(context.Background(), "stream %d already exists", streamID)
+	}
+
+}
+
+// transformRangefeedErrToClientError converts a rangefeed error to a client
+// error to be sent back to client. This also handles nil values, preventing nil
+// pointer dereference.
+func transformRangefeedErrToClientError(err *kvpb.Error) *kvpb.Error {
+	if err == nil {
+		// When processor is stopped when it no longer has any registrations, it
+		// would attempt to close all feeds again with a nil error. Theoretically,
+		// this should never happen as processor would always stop with a reason if
+		// feeds are active.
+		return kvpb.NewError(kvpb.NewRangeFeedRetryError(kvpb.RangeFeedRetryError_REASON_RANGEFEED_CLOSED))
+	}
+	return err
+}
+
+func (sm *StreamMuxer) appendMuxError(e *kvpb.MuxRangeFeedEvent) {
 	sm.mu.Lock()
 	defer sm.mu.Unlock()
 	sm.mu.muxErrors = append(sm.mu.muxErrors, e)
-	// Note that notify is unblocking.
+	// Note that notifyMuxError is non-blocking.
 	select {
 	case sm.notifyMuxError <- struct{}{}:
 	default:
 	}
 }
 
+// DisconnectRangefeedWithError disconnects a stream with an error. Safe to call
+// repeatedly for the same stream, but subsequent errors are ignored. It ensures
+// 1. cancel stream context 2. send error back to client once for the stream.
+//
+// Note that this function can be called by the processor worker while holding
+// raftMu, so it is important that this function doesn't block IO. It does so by
+// delegating the responsibility of sending mux error to muxer.Run.
+func (sm *StreamMuxer) DisconnectRangefeedWithError(
+	streamID int64, rangeID roachpb.RangeID, err *kvpb.Error,
+) {
+	if cancelFunc, ok := sm.activeStreams.LoadAndDelete(streamID); ok {
+		if f, ok := cancelFunc.(context.CancelFunc); ok {
+			f()
+		}
+		clientErrorEvent := transformRangefeedErrToClientError(err)
+		ev := &kvpb.MuxRangeFeedEvent{
+			StreamID: streamID,
+			RangeID:  rangeID,
+		}
+		ev.SetValue(&kvpb.RangeFeedError{
+			Error: *clientErrorEvent,
+		})
+		sm.appendMuxError(ev)
+	}
+}
+
 func (sm *StreamMuxer) detachMuxErrors() []*kvpb.MuxRangeFeedEvent {
 	sm.mu.Lock()
 	defer sm.mu.Unlock()
@@ -89,7 +146,8 @@ func (sm *StreamMuxer) Run(ctx context.Context, stopper *stop.Stopper) {
 	for {
 		select {
 		case <-sm.notifyMuxError:
-			for _, clientErr := range sm.detachMuxErrors() {
+			toSend := sm.detachMuxErrors()
+			for _, clientErr := range toSend {
 				if err := sm.sender.Send(clientErr); err != nil {
 					return
 				}

pkg/kv/kvserver/rangefeed/stream_muxer_test.go

@@ -0,0 +1,188 @@
+// Copyright 2024 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package rangefeed
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	"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"
+	"github.com/cockroachdb/cockroach/pkg/util/stop"
+	"github.com/cockroachdb/errors"
+	"github.com/stretchr/testify/require"
+)
+
+// TestStreamMuxerOnStop tests that the StreamMuxer stops when the context is cancelled.
+func TestStreamMuxerOnContextCancel(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	stopper := stop.NewStopper()
+
+	testServerStream := newTestServerStream()
+	muxer, cleanUp := NewTestStreamMuxer(t, ctx, stopper, testServerStream)
+	defer cleanUp()
+	defer stopper.Stop(ctx)
+
+	cancel()
+	expectedErrEvent := &kvpb.MuxRangeFeedEvent{
+		StreamID: 0,
+		RangeID:  1,
+	}
+	expectedErrEvent.MustSetValue(&kvpb.RangeFeedError{
+		Error: *kvpb.NewError(context.Canceled),
+	})
+	muxer.appendMuxError(expectedErrEvent)
+	time.Sleep(10 * time.Millisecond)
+	require.False(t, testServerStream.hasEvent(expectedErrEvent))
+	require.Equal(t, 0, testServerStream.totalEventsSent())
+}
+
+func TestStreamMuxer(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	stopper := stop.NewStopper()
+
+	testServerStream := newTestServerStream()
+	muxer, cleanUp := NewTestStreamMuxer(t, ctx, stopper, testServerStream)
+	defer cleanUp()
+
+	// Note that this also tests that the StreamMuxer stops when the stopper is
+	// stopped. If not, the test will hang.
+	defer stopper.Stop(ctx)
+
+	t.Run("nil handling", func(t *testing.T) {
+		const streamID = 0
+		const rangeID = 1
+		streamCtx, cancel := context.WithCancel(context.Background())
+		muxer.AddStream(0, cancel)
+		// Note that kvpb.NewError(nil) == nil.
+		muxer.DisconnectRangefeedWithError(streamID, rangeID, kvpb.NewError(nil))
+		require.Equal(t, context.Canceled, streamCtx.Err())
+		expectedErrEvent := &kvpb.MuxRangeFeedEvent{
+			StreamID: streamID,
+			RangeID:  rangeID,
+		}
+		expectedErrEvent.MustSetValue(&kvpb.RangeFeedError{
+			Error: *kvpb.NewError(kvpb.NewRangeFeedRetryError(kvpb.RangeFeedRetryError_REASON_RANGEFEED_CLOSED)),
+		})
+		time.Sleep(10 * time.Millisecond)
+		require.Equal(t, 1, testServerStream.totalEventsSent())
+		require.True(t, testServerStream.hasEvent(expectedErrEvent))
+
+		// Repeat closing the stream does nothing.
+		muxer.DisconnectRangefeedWithError(streamID, rangeID,
+			kvpb.NewError(kvpb.NewRangeFeedRetryError(kvpb.RangeFeedRetryError_REASON_RANGEFEED_CLOSED)))
+		time.Sleep(10 * time.Millisecond)
+		require.Equal(t, 1, testServerStream.totalEventsSent())
+	})
+
+	t.Run("send rangefeed completion error", func(t *testing.T) {
+		testRangefeedCompletionErrors := []struct {
+			streamID int64
+			rangeID  roachpb.RangeID
+			Error    error
+		}{
+			{0, 1, kvpb.NewRangeFeedRetryError(kvpb.RangeFeedRetryError_REASON_RANGEFEED_CLOSED)},
+			{1, 1, context.Canceled},
+			{2, 2, &kvpb.NodeUnavailableError{}},
+		}
+
+		for _, muxError := range testRangefeedCompletionErrors {
+			muxer.AddStream(muxError.streamID, func() {})
+		}
+
+		var wg sync.WaitGroup
+		for _, muxError := range testRangefeedCompletionErrors {
+			wg.Add(1)
+			func(streamID int64, rangeID roachpb.RangeID, err error) {
+				muxer.DisconnectRangefeedWithError(streamID, rangeID, kvpb.NewError(err))
+				wg.Done()
+			}(muxError.streamID, muxError.rangeID, muxError.Error)
+		}
+		wg.Wait()
+
+		for _, muxError := range testRangefeedCompletionErrors {
+			testutils.SucceedsSoon(t, func() error {
+				ev := &kvpb.MuxRangeFeedEvent{
+					StreamID: muxError.streamID,
+					RangeID:  muxError.rangeID,
+				}
+				ev.MustSetValue(&kvpb.RangeFeedError{
+					Error: *kvpb.NewError(muxError.Error),
+				})
+				if testServerStream.hasEvent(ev) {
+					return nil
+				}
+				return errors.Newf("expected error %v not found", muxError)
+			})
+		}
+	})
+}
+
+func TestStreamMuxerOnBlockingIO(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	stopper := stop.NewStopper()
+
+	testServerStream := newTestServerStream()
+	muxer, cleanUp := NewTestStreamMuxer(t, ctx, stopper, testServerStream)
+	defer cleanUp()
+	defer stopper.Stop(ctx)
+
+	const streamID = 0
+	const rangeID = 1
+	_, cancel := context.WithCancel(context.Background())
+	muxer.AddStream(0, cancel)
+	ev := &kvpb.MuxRangeFeedEvent{
+		StreamID: streamID,
+		RangeID:  rangeID,
+	}
+	ev.SetValue(
+		*rangeFeedCheckpoint(
+			roachpb.Span{Key: roachpb.Key("a"), EndKey: roachpb.Key("m")},
+			hlc.Timestamp{WallTime: 1}),
+	)
+	err := muxer.sender.Send(ev)
+	require.NoError(t, err)
+	require.True(t, testServerStream.hasEvent(ev))
+
+	// Block the stream.
+	unblock := testServerStream.BlockSend()
+
+	// Although stream is blocked, we should be able to disconnect the stream
+	// without blocking.
+	muxer.DisconnectRangefeedWithError(streamID, rangeID,
+		kvpb.NewError(kvpb.NewRangeFeedRetryError(kvpb.RangeFeedRetryError_REASON_NO_LEASEHOLDER)))
+	expectedErrEvent := &kvpb.MuxRangeFeedEvent{
+		StreamID: streamID,
+		RangeID:  rangeID,
+	}
+	expectedErrEvent.MustSetValue(&kvpb.RangeFeedError{
+		Error: *kvpb.NewError(kvpb.NewRangeFeedRetryError(kvpb.RangeFeedRetryError_REASON_NO_LEASEHOLDER)),
+	})
+	unblock()
+	time.Sleep(10 * time.Millisecond)
+	// Receive the event after unblocking.
+	require.True(t, testServerStream.hasEvent(expectedErrEvent))
+}