rangefeed: defer unregistration until stream is removed from manager

Unregistering the registration from the processor has the side-effect of
closing the underlying memory budget. We don't want that to happen until
all of the events have been cleared from the queue.

Here, we pass the unregistration callback up via the Disconnector and
then the stream manager calls it when removing the stream.

Epic: none
Release note: None

Author: stevendanna

pkg/kv/kvserver/rangefeed/buffered_registration.go

@@ -158,6 +158,13 @@ func (br *bufferedRegistration) IsDisconnected() bool {
 	return br.mu.disconnected
 }
 
+// Unregister implements Disconnector.
+//
+// The bufferedRegistration unregisters itself via Disconnect because it is
+// responsible for all of its buffered memory and thus there is no reason to
+// delay unregistration.
+func (br *bufferedRegistration) Unregister() {}
+
 // Disconnect cancels the output loop context for the registration and passes an
 // error to the output error stream for the registration.
 // Safe to run multiple times, but subsequent errors would be discarded.

pkg/kv/kvserver/rangefeed/processor_test.go

@@ -1618,18 +1618,112 @@ func TestIntentScannerOnError(t *testing.T) {
 	}
 	err := s.Start(stopper, erroringScanConstructor)
 	require.ErrorContains(t, err, "scanner error")
+}
+
+// TestProcessorMemoryAccountingOnError tests that when a
+// buffered sender disconnects because of an error, the memory budget continues
+// to account for any previously buffered events until they are actually sent.
+//
+// Note, this tests the case where the error is a memory overflow, but any error
+// that disconnects our registration could have been used.
+func TestProcessorMemoryAccountingOnError(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	ctx := context.Background()
+	stopper := stop.NewStopper()
+	defer stopper.Stop(ctx)
+
+	queueCap := int64(10)
+	streamID := int64(1)
+
+	st := cluster.MakeTestingClusterSettings()
+	RangefeedSingleBufferedSenderQueueMaxPerReg.Override(ctx, &st.SV, queueCap)
+
+	fb := newTestBudget(math.MaxInt64)
+	testServerStream := newTestServerStream()
+	bs := NewBufferedSender(testServerStream, st, NewBufferedSenderMetrics())
+
+	smMetrics := NewStreamManagerMetrics()
+	sm := NewStreamManager(bs, smMetrics)
+	require.NoError(t, sm.Start(ctx, stopper))
+	defer sm.Stop(ctx)
+
+	// Create a processor with our budget.
+	p, h, pStopper := newTestProcessor(t,
+		withBudget(fb),
+		withRangefeedTestType(scheduledProcessorWithBufferedSender))
+	defer pStopper.Stop(ctx)
+
+	// Block the sender so the buffer will fill up.
+	unblock := testServerStream.BlockSend()
+	defer func() {
+		if unblock != nil {
+			unblock()
+		}
+	}()
+
+	startTime := hlc.Timestamp{WallTime: 1}
+	sm.RegisteringStream(streamID)
+	registered, d, _ := p.Register(
+		ctx,
+		roachpb.RSpan{Key: roachpb.RKey("a"), EndKey: roachpb.RKey("z")},
+		startTime,
+		nil,   /* catchUpIter */
+		false, /* withDiff */
+		false, /* withFiltering */
+		false, /* withOmitRemote */
+		noBulkDelivery,
+		sm.NewStream(streamID, 1 /* rangeID */),
+	)
+	require.True(t, registered)
+	sm.AddStream(streamID, d)
 
-	// The processor should be stopped eventually.
-	p := (s).(*ScheduledProcessor)
+	// Overflow the queue.
+	for i := range queueCap + 1 {
+		v := writeValueOpWithKV(roachpb.Key("k"), hlc.Timestamp{WallTime: startTime.WallTime + i + 1}, []byte("val"))
+		require.True(t, p.ConsumeLogicalOps(ctx, v))
+	}
+
+	// Once all events have been sent to the registration, we should be overflowed
+	// and disconnection.
+	h.syncEventC()
 	testutils.SucceedsSoon(t, func() error {
-		select {
-		case <-p.stoppedC:
-			_, ok := sch.shards[shardIndex(p.ID(), len(sch.shards), p.Priority)].procs[p.ID()]
-			require.False(t, ok)
-			require.False(t, sch.priorityIDs.Contains(p.ID()))
+		if d.IsDisconnected() {
 			return nil
-		default:
-			return errors.New("processor not stopped")
 		}
+		return errors.New("waiting for registration to disconnect")
+	})
+
+	// At this point, the registration should be disconnected but the buffered
+	// sender still has events in its queue. Assert that the memory budget still
+	// accounts for the memory in that queue.
+	//
+	// NB: This could be racy if change the structure of the code in the future.
+	// Namely, perhaps it isn't zero, now, but perhaps it becomes zero at some
+	// time in the future. We try to defend against that here by sending 2 sync
+	// events to help ensure we've definitely processed any processor requests.
+	//
+	// At the time this test was written, this test caught the bug on every run.
+	h.syncEventC()
+	h.syncEventC()
+
+	fb.mu.Lock()
+	budgetUsed := fb.mu.memBudget.Used()
+	fb.mu.Unlock()
+	require.Greater(t, budgetUsed, int64(0),
+		"memory budget should still account for events in buffered sender after overflow")
+
+	// Unblocking the sender should drain the queue and free everything from the
+	// memory budget.
+	unblock()
+	unblock = nil
+
+	testutils.SucceedsSoon(t, func() error {
+		fb.mu.Lock()
+		defer fb.mu.Unlock()
+		if used := fb.mu.memBudget.Used(); used != 0 {
+			return errors.Errorf("budget still has %d bytes allocated", used)
+		}
+		return nil
 	})
 }

pkg/kv/kvserver/rangefeed/registry.go

@@ -29,6 +29,9 @@ type Disconnector interface {
 	// Disconnected is a permanent state; once IsDisconnected returns true, it
 	// always returns true
 	IsDisconnected() bool
+	// Unregister is called when an error has finally been delivered to the
+	// underlying stream.
+	Unregister()
 }
 
 // registration defines an interface for registration that can be added to a

pkg/kv/kvserver/rangefeed/sender_helper_test.go

@@ -208,3 +208,5 @@ func (c *cancelCtxDisconnector) IsDisconnected() bool {
 	defer c.mu.Unlock()
 	return c.mu.disconnected
 }
+
+func (c *cancelCtxDisconnector) Unregister() {}

pkg/kv/kvserver/rangefeed/stream_manager.go

@@ -123,6 +123,7 @@ func (sm *StreamManager) OnError(streamID int64) {
 		defer sm.streams.Unlock()
 		if d, ok := sm.streams.m[streamID]; ok {
 			assertTrue(d.IsDisconnected(), "OnError called on connected registration")
+			d.Unregister()
 			delete(sm.streams.m, streamID)
 			sm.metrics.ActiveMuxRangeFeed.Dec(1)
 		}
@@ -163,6 +164,13 @@ func (sm *StreamManager) AddStream(streamID int64, d Disconnector) {
 	if d.IsDisconnected() {
 		// If the stream is already disconnected, we don't add it to streams. The
 		// registration will have already sent an error to the client.
+		//
+		// TODO(ssd): Technically this error event might live in the buffer still
+		// and unregistering now may close the underlying memory budget related to
+		// that event. At the moment, there isn't a better place to do this however
+		// because the whole point of IsDisconnected() is that the error event might have
+		// raced us and already be at the client.
+		d.Unregister()
 		return
 	}
 	if _, ok := sm.streams.m[streamID]; ok {
@@ -218,6 +226,9 @@ func (sm *StreamManager) Stop(ctx context.Context) {
 		// sent to the client after shutdown, but the gRPC stream will still
 		// terminate.
 		disconnector.Disconnect(rangefeedClosedErr)
+		// At this point the sender has been cleaned up so any memory allocations it
+		// had should already be gone.
+		disconnector.Unregister()
 	}
 	sm.streams.m = nil
 }

pkg/kv/kvserver/rangefeed/unbuffered_registration.go

@@ -181,7 +181,9 @@ func (ubr *unbufferedRegistration) disconnectLocked(pErr *kvpb.Error) {
 	}
 	ubr.mu.disconnected = true
 	ubr.stream.SendError(pErr)
-	ubr.removeRegFromProcessor(ubr)
+	// NB: The unbuffered registration does not unregister itself on Disconnect
+	// because it still has memory in the buffered sender and we do not want to
+	// free any underlying memory budgets until that has been cleared.
 }
 
 // IsDisconnected returns true if the registration is disconnected.
@@ -191,6 +193,12 @@ func (ubr *unbufferedRegistration) IsDisconnected() bool {
 	return ubr.mu.disconnected
 }
 
+// Unregister implements Disconnector.
+func (ubr *unbufferedRegistration) Unregister() {
+	assertTrue(ubr.IsDisconnected(), "connected registration in Unregister")
+	ubr.removeRegFromProcessor(ubr)
+}
+
 // runOutputLoop is run in a goroutine. It is responsible for running the
 // catch-up scan, and then publishing any events buffered in catchUpBuf to the
 // sender (or discarding catch-up buffer in the case of an error).