Merge pull request #149920 from yuzefovich/blathers/backport-release-25.3-149800

release-25.3: distsql: fix recently introduced leak of CancelFunc

Author: yuzefovich

pkg/kv/kvclient/kvstreamer/streamer.go

@@ -532,6 +532,9 @@ func (s *Streamer) Enqueue(ctx context.Context, reqs []kvpb.RequestUnion) (retEr
 			},
 			s.coordinator.mainLoop,
 		); err != nil {
+			// The server is shutting down. It's ok to not call
+			// s.coordinatorCtxCancel in this case.
+			//
 			// The new goroutine wasn't spun up, so mainLoop won't get executed
 			// and we have to decrement the wait group ourselves.
 			s.waitGroup.Done()

pkg/sql/distsql/server.go

@@ -166,14 +166,22 @@ func (ds *ServerImpl) setDraining(drain bool) error {
 	return nil
 }
 
+type onFlowCleanupFn func()
+
+func (f onFlowCleanupFn) Do() {
+	if f != nil {
+		f()
+	}
+}
+
 // setupFlow creates a Flow.
-//
-//   - reserved: specifies the upfront memory reservation that the flow takes
-//     ownership of. This account is already closed if an error is returned or
-//     will be closed through Flow.Cleanup.
-//
-//   - localState: specifies if the flow runs entirely on this node and, if it
-//     does, specifies the txn and other attributes.
+// - reserved: specifies the upfront memory reservation that the flow takes
+// ownership of. This account is already closed if an error is returned or will
+// be closed through Flow.Cleanup.
+// - localState: specifies if the flow runs entirely on this node and, if it
+// does, specifies the txn and other attributes.
+// - onFlowCleanup, if non-nil, will be called at the end of Flow.Cleanup. It'll
+// also be called if this method returns an error.
 //
 // Note: unless an error is returned, the returned context contains a span that
 // must be finished through Flow.Cleanup.
@@ -186,6 +194,7 @@ func (ds *ServerImpl) setupFlow(
 	rowSyncFlowConsumer execinfra.RowReceiver,
 	batchSyncFlowConsumer execinfra.BatchReceiver,
 	localState LocalState,
+	onFlowCleanup onFlowCleanupFn,
 ) (retCtx context.Context, _ flowinfra.Flow, _ execopnode.OpChains, retErr error) {
 	var sp *tracing.Span                       // will be Finish()ed by Flow.Cleanup()
 	var monitor, diskMonitor *mon.BytesMonitor // will be closed in Flow.Cleanup()
@@ -204,6 +213,7 @@ func (ds *ServerImpl) setupFlow(
 				onFlowCleanupEnd(ctx)
 			} else {
 				reserved.Close(ctx)
+				onFlowCleanup.Do()
 			}
 			// We finish the span after performing other cleanup in case that
 			// cleanup accesses the context with the span.
@@ -300,18 +310,21 @@ func (ds *ServerImpl) setupFlow(
 			onFlowCleanupEnd = func(ctx context.Context) {
 				localEvalCtx.Txn = origTxn
 				reserved.Close(ctx)
+				onFlowCleanup.Do()
 			}
 			// Update the Txn field early (before f.SetTxn() below) since some
 			// processors capture the field in their constructor (see #41992).
 			localEvalCtx.Txn = leafTxn
 		} else {
 			onFlowCleanupEnd = func(ctx context.Context) {
 				reserved.Close(ctx)
+				onFlowCleanup.Do()
 			}
 		}
 	} else {
 		onFlowCleanupEnd = func(ctx context.Context) {
 			reserved.Close(ctx)
+			onFlowCleanup.Do()
 		}
 		if localState.IsLocal {
 			return nil, nil, nil, errors.AssertionFailedf(
@@ -579,7 +592,7 @@ func (ds *ServerImpl) SetupLocalSyncFlow(
 ) (context.Context, flowinfra.Flow, execopnode.OpChains, error) {
 	return ds.setupFlow(
 		ctx, tracing.SpanFromContext(ctx), parentMonitor, &mon.BoundAccount{}, /* reserved */
-		req, output, batchOutput, localState,
+		req, output, batchOutput, localState, nil, /* onFlowCleanup */
 	)
 }
 
@@ -639,10 +652,6 @@ func (ds *ServerImpl) SetupFlow(
 	// Note: the passed context will be canceled when this RPC completes, so we
 	// can't associate it with the flow since it outlives the RPC.
 	ctx = ds.AnnotateCtx(context.Background())
-	// Ensure that the flow respects the node being shut down. Note that since
-	// the flow outlives the RPC, we cannot defer the cancel function, so we
-	// simply ignore it.
-	ctx, _ = ds.Stopper.WithCancelOnQuiesce(ctx)
 	if err := func() error {
 		// Reserve some memory for this remote flow which is a poor man's
 		// admission control based on the RAM usage.
@@ -651,10 +660,18 @@ func (ds *ServerImpl) SetupFlow(
 		if err != nil {
 			return err
 		}
+		// Ensure that the flow respects the node being shut down. We can only
+		// call the cancellation function once the flow exits.
+		//
+		// setupFlow will either call 'cancel' if an error is returned, or the
+		// cancellation function is taken over by the flow, and it'll be called
+		// in Flow.Cleanup.
+		var cancel context.CancelFunc
+		ctx, cancel = ds.Stopper.WithCancelOnQuiesce(ctx)
 		var f flowinfra.Flow
 		ctx, f, _, err = ds.setupFlow(
 			ctx, rpcSpan, ds.memMonitor, &reserved, req, nil, /* rowSyncFlowConsumer */
-			nil /* batchSyncFlowConsumer */, LocalState{},
+			nil /* batchSyncFlowConsumer */, LocalState{}, onFlowCleanupFn(cancel),
 		)
 		// Check whether the RPC context has been canceled indicating that we
 		// actually don't need to run this flow. This can happen when the