distsql: improve SetupFlow RPC handling by server when canceled

This commit makes the server-side of `SetupFlow` RPC react better to the
context cancellation. In particular, after setting up the flow on the
remote node but before starting it in a new goroutine we now check
whether the context of the RPC is canceled or not; if it is, then we
don't actually need to run the flow at all, so we now short-circuit.

This improves the shutdown protocol in some cases when the context is
canceled on the gateway (client-side of the RPC) after we have dialed
the node in `runnerRequest.run` but before the new goroutine handling
the RPC on the client side has been spun up. This scenario is pretty
rare currently but might occur more often due to the changes in the
following commit. Note that there is no correctness issue here - even
before this change, the flow would be shutdown properly (when the outbox
would fail to execute `FlowStream` RPC due to the inbox having exited).

Release note: None

pkg/sql/distsql/server.go

@@ -621,8 +621,9 @@ func (ds *ServerImpl) SetupFlow(
 	_, rpcSpan := ds.setupSpanForIncomingRPC(ctx, req)
 	defer rpcSpan.Finish()
 
+	rpcCtx := ctx
 	// Note: the passed context will be canceled when this RPC completes, so we
-	// can't associate it with the flow.
+	// can't associate it with the flow since it outlives the RPC.
 	ctx = ds.AnnotateCtx(context.Background())
 	if err := func() error {
 		// Reserve some memory for this remote flow which is a poor man's
@@ -637,7 +638,18 @@ func (ds *ServerImpl) SetupFlow(
 			ctx, rpcSpan, ds.memMonitor, &reserved, req, nil, /* rowSyncFlowConsumer */
 			nil /* batchSyncFlowConsumer */, LocalState{},
 		)
+		// Check whether the RPC context has been canceled indicating that we
+		// actually don't need to run this flow. This can happen when the
+		// context is canceled after Dial() but before issuing SetupFlow RPC in
+		// runnerRequest.run().
+		if err == nil {
+			err = rpcCtx.Err()
+		}
 		if err != nil {
+			// Make sure to clean up the flow if it was created.
+			if f != nil {
+				f.Cleanup(ctx)
+			}
 			return err
 		}
 		return ds.remoteFlowRunner.RunFlow(ctx, f)