sql: emit max memory usage on the gateway node

Previously, the max memory usage execution stat was only emitted by the last
outbox to drain metadata in a flow. However, there are cases where no outbox
is present in the flow (e.g. single-node queries). This commit adds a Gateway
field to the FlowCtx which specifies whether this flow is running on the
gateway node or not and changes outboxes to only emit this stat if they are
not run on the gateway node. Gateway node max memory usage is now emitted in
the trace during flow.Cleanup.

This commit also adds a NodeID field to the ComponentID in ComponentStats used
to uniquely describe a flow-level ComponentStat. The FlowID was used
previously, which is not unique across nodes.

Release note: None

pkg/sql/colflow/vectorized_flow.go

@@ -195,6 +195,7 @@ func (f *vectorizedFlow) Setup(
 		helper,
 		vectorizedRemoteComponentCreator{},
 		recordingStats,
+		f.Gateway,
 		f.GetWaitGroup(),
 		f.GetSyncFlowConsumer(),
 		flowCtx.Cfg.NodeDialer,
@@ -438,6 +439,7 @@ type vectorizedFlowCreator struct {
 	streamIDToInputOp              map[execinfrapb.StreamID]opDAGWithMetaSources
 	streamIDToSpecIdx              map[execinfrapb.StreamID]int
 	recordingStats                 bool
+	isGatewayNode                  bool
 	vectorizedStatsCollectorsQueue []colexec.VectorizedStatsCollector
 	waitGroup                      *sync.WaitGroup
 	syncFlowConsumer               execinfra.RowReceiver
@@ -505,6 +507,7 @@ func newVectorizedFlowCreator(
 	helper flowCreatorHelper,
 	componentCreator remoteComponentCreator,
 	recordingStats bool,
+	isGatewayNode bool,
 	waitGroup *sync.WaitGroup,
 	syncFlowConsumer execinfra.RowReceiver,
 	nodeDialer *nodedialer.Dialer,
@@ -948,17 +951,13 @@ func (s *vectorizedFlowCreator) setupOutput(
 						// Start a separate recording so that GetRecording will return
 						// the recordings for only the child spans containing stats.
 						ctx, span := tracing.ChildSpanRemote(ctx, "")
-						if atomic.AddInt32(&s.numOutboxesDrained, 1) == atomic.LoadInt32(&s.numOutboxes) {
+						if atomic.AddInt32(&s.numOutboxesDrained, 1) == atomic.LoadInt32(&s.numOutboxes) && !s.isGatewayNode {
 							// At the last outbox, we can accurately retrieve stats for the
 							// whole flow from parent monitors. These stats are added to a
 							// flow-level span.
 							span.SetTag(execinfrapb.FlowIDTagKey, flowCtx.ID)
 							span.SetSpanStats(&execinfrapb.ComponentStats{
-								Component: execinfrapb.ComponentID{
-									Type:   execinfrapb.ComponentID_FLOW,
-									FlowID: flowCtx.ID,
-									// TODO(radu): the node ID should be part of the ComponentID.
-								},
+								Component: execinfrapb.FlowComponentID(outputStream.OriginNodeID, flowCtx.ID),
 								FlowStats: execinfrapb.FlowStats{
 									MaxMemUsage: optional.MakeUint(uint64(flowCtx.EvalCtx.Mon.MaximumBytes())),
 								},
@@ -1369,7 +1368,7 @@ func ConvertToVecTree(
 		fuseOpt = flowinfra.FuseAggressively
 	}
 	creator := newVectorizedFlowCreator(
-		newNoopFlowCreatorHelper(), vectorizedRemoteComponentCreator{}, false,
+		newNoopFlowCreatorHelper(), vectorizedRemoteComponentCreator{}, false, false,
 		nil, &execinfra.RowChannel{}, nil, execinfrapb.FlowID{}, colcontainer.DiskQueueCfg{},
 		flowCtx.Cfg.VecFDSemaphore, flowCtx.TypeResolverFactory.NewTypeResolver(flowCtx.EvalCtx.Txn),
 	)

pkg/sql/colflow/vectorized_flow_test.go

@@ -222,7 +222,7 @@ func TestDrainOnlyInputDAG(t *testing.T) {
 	}
 	var wg sync.WaitGroup
 	vfc := newVectorizedFlowCreator(
-		&vectorizedFlowCreatorHelper{f: f}, componentCreator, false, &wg, &execinfra.RowChannel{},
+		&vectorizedFlowCreatorHelper{f: f}, componentCreator, false, false, &wg, &execinfra.RowChannel{},
 		nil /* nodeDialer */, execinfrapb.FlowID{}, colcontainer.DiskQueueCfg{},
 		nil /* fdSemaphore */, descs.DistSQLTypeResolver{},
 	)

pkg/sql/distsql/inbound_test.go

@@ -95,7 +95,7 @@ func TestOutboxInboundStreamIntegration(t *testing.T) {
 	}
 
 	streamID := execinfrapb.StreamID(1)
-	outbox := flowinfra.NewOutbox(&flowCtx, execinfra.StaticNodeID, streamID, nil /* numOutboxes */)
+	outbox := flowinfra.NewOutbox(&flowCtx, execinfra.StaticNodeID, streamID, nil /* numOutboxes */, false /* isGatewayNode */)
 	outbox.Init(rowenc.OneIntCol)
 
 	// WaitGroup for the outbox and inbound stream. If the WaitGroup is done, no

pkg/sql/distsql/server.go

@@ -306,7 +306,9 @@ func (ds *ServerImpl) setupFlow(
 	}
 
 	// Create the FlowCtx for the flow.
-	flowCtx := ds.NewFlowContext(ctx, req.Flow.FlowID, evalCtx, req.TraceKV, localState)
+	flowCtx := ds.NewFlowContext(
+		ctx, req.Flow.FlowID, evalCtx, req.TraceKV, localState, req.Flow.Gateway == roachpb.NodeID(ds.NodeID.SQLInstanceID()),
+	)
 
 	// req always contains the desired vectorize mode, regardless of whether we
 	// have non-nil localState.EvalContext. We don't want to update EvalContext
@@ -392,6 +394,7 @@ func (ds *ServerImpl) NewFlowContext(
 	evalCtx *tree.EvalContext,
 	traceKV bool,
 	localState LocalState,
+	isGatewayNode bool,
 ) execinfra.FlowCtx {
 	// TODO(radu): we should sanity check some of these fields.
 	flowCtx := execinfra.FlowCtx{
@@ -402,6 +405,7 @@ func (ds *ServerImpl) NewFlowContext(
 		NodeID:         ds.ServerConfig.NodeID,
 		TraceKV:        traceKV,
 		Local:          localState.IsLocal,
+		Gateway:        isGatewayNode,
 	}
 
 	if localState.IsLocal && localState.Collection != nil {

pkg/sql/execinfra/flow_context.go

@@ -62,6 +62,9 @@ type FlowCtx struct {
 	// Local is true if this flow is being run as part of a local-only query.
 	Local bool
 
+	// Gateway is true if this flow is being run on the gateway node.
+	Gateway bool
+
 	// TypeResolverFactory is used to construct transaction bound TypeResolvers
 	// to resolve type references during flow setup. It is not safe for concurrent
 	// use and is intended to be used only during flow setup and initialization.

pkg/sql/execinfrapb/component_stats.go

@@ -15,6 +15,7 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/util/humanizeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/optional"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
@@ -42,6 +43,15 @@ func StreamComponentID(flowID FlowID, streamID StreamID) ComponentID {
 	}
 }
 
+// FlowComponentID returns a ComponentID for the given flow.
+func FlowComponentID(nodeID roachpb.NodeID, flowID FlowID) ComponentID {
+	return ComponentID{
+		FlowID: flowID,
+		Type:   ComponentID_FLOW,
+		NodeID: nodeID,
+	}
+}
+
 // FlowIDTagKey is the key used for flow id tags in tracing spans.
 const (
 	FlowIDTagKey = tracing.TagPrefix + "flowid"