distsql: fix recently introduced leak of CancelFunc

In 27a65c9 we introduced the logic to
make sure that remote flows respect the quiesce signal. However, we
consciusly ignored the returned CancelFunc since the lifetime of the
context is non-trivial. This introduced a leak of that CancelFunc since
we keep the reference to it in the stopper even when the flow exits, so
every time a remote flow runs on a node, the leak would slowly grow.

This leak is now fixed by ensuring the cancellation function is always
called. This is achieved by passing the function as another thing to do
on the "flow cleanup" (we already always need to close the reserved
memory account, so we have the necessary infrastructure set up).

Release note (bug fix): In 25.1.8, 25.2.1, 25.2.2, and 25.3 betas, a slow
memory leak was introduced that would accumulate whenever a node
executes a part of the distributed plan (the gateway node of the plan is
not affected). The leak can only be mitigated by restarting the node and
is now fixed.

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

@@ -167,14 +167,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.
@@ -187,6 +195,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()
@@ -205,6 +214,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.
@@ -301,18 +311,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(
@@ -580,7 +593,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 */
 	)
 }
 
@@ -640,10 +653,10 @@ 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)
+	// Ensure that the flow respects the node being shut down. We can only call
+	// the cancellation function once the flow exits.
+	var cancel context.CancelFunc
+	ctx, cancel = 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.
@@ -655,7 +668,7 @@ func (ds *ServerImpl) SetupFlow(
 		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