Merge #146270

146270: kvserver: remove UnregisterFromReplica callback r=tbg a=stevendanna

The UnregisterFromReplica callback required obtaining the rangefeedMu mutex from inside the rangefeed scheduler worker. This is a problem because the current holder of the rangefeedMu (specifically (*ScheduledProcessor).Register) may require a response from the very same scheduler worker to make progress and release rangefeedMu.

Here, we solve that by removing this UnregisterFromReplica callback completely. This callback was responsible for removing the processor from the replica when the processor had been shut down.

But, nearly every code path that calls Stop() already removes the processor from the replica itself. The only case where this wasn't true is when the processor stops itself because it has no more active registrations.

Not removing the processor from the replica in that case has two consequences:

1. We may hold onto memory related to the ScheduledProcessor struct,

2. The replica does extra work because of the rangefeed processor is set.

Here, we choose to ignore (1) as a minor problem since a most ranges that have a rangefeed started once will have one started again in the future.

We solve (2) by making the `stopped` state variable accessible from outside the processor and consulting it before using the rangefeed processor on the replica. The assumption here is that the atomic load is cheap in comparison to the work we need to do when the processor is present.

Fixes #144828

Epic: none

Release note (bug fix): Fix a bug that could lead to a node stall.

Co-authored-by: Steven Danna <[email protected]>

Author: craig[bot]

pkg/kv/kvserver/rangefeed/processor.go

@@ -104,11 +104,6 @@ type Config struct {
 	// for low-volume system ranges, since the worker pool is small (default 2).
 	// Only has an effect when Scheduler is used.
 	Priority bool
-
-	// UnregisterFromReplica is a callback provided from the
-	// replica that this processor can call when shutting down to
-	// remove itself from the replica.
-	UnregisterFromReplica func(Processor)
 }
 
 // SetDefaults initializes unset fields in Config to values
@@ -166,6 +161,9 @@ type Processor interface {
 	// It is not valid to restart a processor after it has been stopped.
 	StopWithErr(pErr *kvpb.Error)
 
+	// Returns true if a stop event has already been processed by this processor.
+	Stopping() bool
+
 	// Lifecycle of registrations.
 
 	// Register registers the stream over the specified span of keys.

pkg/kv/kvserver/rangefeed/scheduled_processor.go

@@ -57,9 +57,20 @@ type ScheduledProcessor struct {
 
 	requestQueue chan request
 	eventC       chan *event
-	// If true, processor is not processing data anymore and waiting for registrations
-	// to be complete.
-	stopping bool
+
+	// stopping, when true, indicates that a stop request has been processed.
+	// After this is set, all subsequent registration requests will fail fast and
+	// any events will be ignored. This is only ever set via stopInternal, but it
+	// is an atomic to allow it to be read outside this processor. This is used in
+	// the replica to handle the case where we have stopped ourselves but are
+	// still attached to the replica. It should be preferred over stoppedC for all
+	// uses other than unblocking request waiters since stoppedC is not closed
+	// until the processsor is fully closed up.
+	stopping atomic.Bool
+	// stoppedC is closed during the final steps of processor shutdown. This is
+	// used to unblock anyone who might be waiting on a response. Callsites that
+	// need to decide whether or not to process a request should consult stopping
+	// instead.
 	stoppedC chan struct{}
 
 	// stopper passed by start that is used for firing up async work from scheduler.
@@ -185,7 +196,7 @@ func (p *ScheduledProcessor) processEvents(ctx context.Context) {
 	for max := len(p.eventC); max > 0; max-- {
 		select {
 		case e := <-p.eventC:
-			if !p.stopping {
+			if !p.stopping.Load() {
 				// If we are stopping, there's no need to forward any remaining
 				// data since registrations already have errors set.
 				p.consumeEvent(ctx, e)
@@ -255,10 +266,6 @@ func (p *ScheduledProcessor) cleanup() {
 	p.taskCancel()
 	close(p.stoppedC)
 	p.MemBudget.Close(ctx)
-	if p.UnregisterFromReplica != nil {
-		p.UnregisterFromReplica(p)
-	}
-
 }
 
 // Stop shuts down the processor and closes all registrations. Safe to call on
@@ -278,6 +285,10 @@ func (p *ScheduledProcessor) StopWithErr(pErr *kvpb.Error) {
 	p.sendStop(pErr)
 }
 
+func (p *ScheduledProcessor) Stopping() bool {
+	return p.stopping.Load()
+}
+
 // DisconnectSpanWithErr disconnects all rangefeed registrations that overlap
 // the given span with the given error.
 func (p *ScheduledProcessor) DisconnectSpanWithErr(span roachpb.Span, pErr *kvpb.Error) {
@@ -299,7 +310,7 @@ func (p *ScheduledProcessor) stopInternal(ctx context.Context, pErr *kvpb.Error)
 	p.reg.DisconnectAllOnShutdown(ctx, pErr)
 	// First set stopping flag to ensure that once all registrations are removed
 	// processor should stop.
-	p.stopping = true
+	p.stopping.Store(true)
 	p.scheduler.StopProcessor()
 }
 
@@ -349,7 +360,7 @@ func (p *ScheduledProcessor) Register(
 	}
 
 	filter := runRequest(p, func(ctx context.Context, p *ScheduledProcessor) *Filter {
-		if p.stopping {
+		if p.stopping.Load() {
 			return nil
 		}
 		if !p.Span.AsRawSpanWithNoLocals().Contains(r.getSpan()) {

pkg/kv/kvserver/replica.go

@@ -1046,6 +1046,9 @@ type Replica struct {
 		// Requires Replica.raftMu be held when providing logical ops and
 		//  informing the processor of closed timestamp updates. This properly
 		//  synchronizes updates that are linearized and driven by the Raft log.
+		//
+		// proc should only be accessed via getRangefeedProcessorAndFilter or
+		// getRangefeedProcessor in nearly all cases.
 		proc rangefeed.Processor
 		// opFilter is a best-effort filter that informs the raft processing
 		// goroutine of which logical operations the rangefeed processor is

pkg/kv/kvserver/replica_rangefeed.go

@@ -348,7 +348,18 @@ func (r *Replica) RangeFeed(
 func (r *Replica) getRangefeedProcessorAndFilter() (rangefeed.Processor, *rangefeed.Filter) {
 	r.rangefeedMu.RLock()
 	defer r.rangefeedMu.RUnlock()
-	return r.rangefeedMu.proc, r.rangefeedMu.opFilter
+	p := r.rangefeedMu.proc
+	if p != nil && p.Stopping() {
+		// This is here only to try to preserve existing behaviour when fixing
+		// #144828. Nearly all call paths that stop the processor immediately remove
+		// the processor from the replica. The only call path where this isn't true
+		// is when the processor stops itself after find all of its registrations
+		// have been removed. Thus, we check stopping here to avoid doing work on
+		// this stop-but-not-yet-removed processor.
+		return nil, r.rangefeedMu.opFilter
+	} else {
+		return p, r.rangefeedMu.opFilter
+	}
 }
 
 func (r *Replica) getRangefeedProcessor() rangefeed.Processor {
@@ -485,21 +496,20 @@ func (r *Replica) registerWithRangefeedRaftMuLocked(
 	desc := r.Desc()
 	tp := rangefeedTxnPusher{ir: r.store.intentResolver, r: r, span: desc.RSpan()}
 	cfg := rangefeed.Config{
-		AmbientContext:        r.AmbientContext,
-		Clock:                 r.Clock(),
-		Stopper:               r.store.stopper,
-		Settings:              r.store.ClusterSettings(),
-		RangeID:               r.RangeID,
-		Span:                  desc.RSpan(),
-		TxnPusher:             &tp,
-		PushTxnsAge:           r.store.TestingKnobs().RangeFeedPushTxnsAge,
-		EventChanCap:          defaultEventChanCap,
-		EventChanTimeout:      defaultEventChanTimeout,
-		Metrics:               r.store.metrics.RangeFeedMetrics,
-		MemBudget:             feedBudget,
-		Scheduler:             r.store.getRangefeedScheduler(),
-		Priority:              isSystemSpan, // only takes effect when Scheduler != nil
-		UnregisterFromReplica: r.unsetRangefeedProcessor,
+		AmbientContext:   r.AmbientContext,
+		Clock:            r.Clock(),
+		Stopper:          r.store.stopper,
+		Settings:         r.store.ClusterSettings(),
+		RangeID:          r.RangeID,
+		Span:             desc.RSpan(),
+		TxnPusher:        &tp,
+		PushTxnsAge:      r.store.TestingKnobs().RangeFeedPushTxnsAge,
+		EventChanCap:     defaultEventChanCap,
+		EventChanTimeout: defaultEventChanTimeout,
+		Metrics:          r.store.metrics.RangeFeedMetrics,
+		MemBudget:        feedBudget,
+		Scheduler:        r.store.getRangefeedScheduler(),
+		Priority:         isSystemSpan, // only takes effect when Scheduler != nil
 	}
 	p = rangefeed.NewProcessor(cfg)