Merge #59293

59293: sql: fix emitting maximum memory usage ComponentStat r=yuzefovich a=asubiotto While working on some frontend changes, I realized that we weren't displaying max memory usage correctly. This PR fixes that. Please see individual commits for details. Release note: None (no user-visible change) Co-authored-by: Alfonso Subiotto Marques <[email protected]>
cockroachdb · Jan 26, 2021 · 97d14bf · 97d14bf
2 parents d86781c + cb7ee1c
commit 97d14bf
Show file tree

Hide file tree

Showing 18 changed files with 247 additions and 159 deletions.
diff --git a/pkg/sql/colflow/vectorized_flow.go b/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),
 	)

diff --git a/pkg/sql/colflow/vectorized_flow_test.go b/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{},
 	)

diff --git a/pkg/sql/distsql/inbound_test.go b/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

diff --git a/pkg/sql/distsql/server.go b/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 {

diff --git a/pkg/sql/execinfra/flow_context.go b/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.

diff --git a/pkg/sql/execinfrapb/component_stats.go b/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"