sql: add cumulative KV time to TraceAnalyzer

This commit adds the total time spent waiting for KV requests to
TraceAnalyzer.

Release note: None.

pkg/sql/execstats/BUILD.bazel

@@ -20,9 +20,10 @@ go_test(
     srcs = [
         "main_test.go",
         "traceanalyzer_test.go",
+        "utils_test.go",
     ],
+    embed = [":execstats"],
     deps = [
-        ":execstats",
         "//pkg/base",
         "//pkg/roachpb",
         "//pkg/security",
@@ -37,7 +38,10 @@ go_test(
         "//pkg/testutils/testcluster",
         "//pkg/util/leaktest",
         "//pkg/util/log",
+        "//pkg/util/optional",
         "//pkg/util/tracing",
+        "//pkg/util/uuid",
+        "//vendor/github.com/stretchr/testify/assert",
         "//vendor/github.com/stretchr/testify/require",
     ],
 )

pkg/sql/execstats/traceanalyzer.go

@@ -12,6 +12,7 @@ package execstats
 
 import (
 	"strconv"
+	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
@@ -21,6 +22,7 @@ import (
 	"github.com/gogo/protobuf/types"
 )
 
+// processorStats contains stats for a specific processor extracted from a trace.
 type processorStats struct {
 	nodeID roachpb.NodeID
 	stats  *execinfrapb.ComponentStats
@@ -89,6 +91,7 @@ type NodeLevelStats struct {
 	MaxMemoryUsageGroupedByNode   map[roachpb.NodeID]int64
 	KVBytesReadGroupedByNode      map[roachpb.NodeID]int64
 	KVRowsReadGroupedByNode       map[roachpb.NodeID]int64
+	KVTimeGroupedByNode           map[roachpb.NodeID]time.Duration
 }
 
 // QueryLevelStats returns all the query level stats that correspond to the given traces and flow metadata.
@@ -97,6 +100,7 @@ type QueryLevelStats struct {
 	MaxMemUsage      int64
 	KVBytesRead      int64
 	KVRowsRead       int64
+	KVTime           time.Duration
 }
 
 // TraceAnalyzer is a struct that helps calculate top-level statistics from a
@@ -191,6 +195,7 @@ func (a *TraceAnalyzer) ProcessStats() error {
 		MaxMemoryUsageGroupedByNode:   make(map[roachpb.NodeID]int64),
 		KVBytesReadGroupedByNode:      make(map[roachpb.NodeID]int64),
 		KVRowsReadGroupedByNode:       make(map[roachpb.NodeID]int64),
+		KVTimeGroupedByNode:           make(map[roachpb.NodeID]time.Duration),
 	}
 	var errs error
 
@@ -201,6 +206,7 @@ func (a *TraceAnalyzer) ProcessStats() error {
 		}
 		a.nodeLevelStats.KVBytesReadGroupedByNode[stats.nodeID] += int64(stats.stats.KV.BytesRead.Value())
 		a.nodeLevelStats.KVRowsReadGroupedByNode[stats.nodeID] += int64(stats.stats.KV.TuplesRead.Value())
+		a.nodeLevelStats.KVTimeGroupedByNode[stats.nodeID] += stats.stats.KV.KVTime.Value()
 	}
 
 	// Process streamStats.
@@ -255,6 +261,7 @@ func (a *TraceAnalyzer) ProcessStats() error {
 		MaxMemUsage:      int64(0),
 		KVBytesRead:      int64(0),
 		KVRowsRead:       int64(0),
+		KVTime:           time.Duration(0),
 	}
 
 	for _, bytesSentByNode := range a.nodeLevelStats.NetworkBytesSentGroupedByNode {
@@ -274,6 +281,10 @@ func (a *TraceAnalyzer) ProcessStats() error {
 	for _, kvRowsRead := range a.nodeLevelStats.KVRowsReadGroupedByNode {
 		a.queryLevelStats.KVRowsRead += kvRowsRead
 	}
+
+	for _, kvTime := range a.nodeLevelStats.KVTimeGroupedByNode {
+		a.queryLevelStats.KVTime += kvTime
+	}
 	return errs
 }
 

pkg/sql/execstats/traceanalyzer_test.go

@@ -13,7 +13,9 @@ package execstats_test
 import (
 	"context"
 	"fmt"
+	"reflect"
 	"testing"
+	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -27,7 +29,10 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/optional"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
+	"github.com/cockroachdb/cockroach/pkg/util/uuid"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -166,3 +171,39 @@ func TestTraceAnalyzer(t *testing.T) {
 		require.Equal(t, int64(30*8), queryLevelStats.KVBytesRead)
 	}
 }
+
+func TestTraceAnalyzerProcessStats(t *testing.T) {
+	a := &execstats.TraceAnalyzer{FlowMetadata: &execstats.FlowMetadata{}}
+	a.AddComponentStats(
+		1, /* nodeID */
+		&execinfrapb.ComponentStats{
+			Component: execinfrapb.ProcessorComponentID(
+				execinfrapb.FlowID{UUID: uuid.MakeV4()},
+				1, /* processorID */
+			),
+			KV: execinfrapb.KVStats{
+				KVTime: optional.MakeTimeValue(3 * time.Second),
+			},
+		},
+	)
+
+	a.AddComponentStats(
+		2, /* nodeID */
+		&execinfrapb.ComponentStats{
+			Component: execinfrapb.ProcessorComponentID(
+				execinfrapb.FlowID{UUID: uuid.MakeV4()},
+				2, /* processorID */
+			),
+			KV: execinfrapb.KVStats{
+				KVTime: optional.MakeTimeValue(5 * time.Second),
+			},
+		},
+	)
+
+	expected := execstats.QueryLevelStats{KVTime: 8 * time.Second}
+
+	assert.NoError(t, a.ProcessStats())
+	if got := a.GetQueryLevelStats(); !reflect.DeepEqual(got, expected) {
+		t.Errorf("ProcessStats() = %v, want %v", got, expected)
+	}
+}

pkg/sql/execstats/utils_test.go

@@ -0,0 +1,52 @@
+// Copyright 2020 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package execstats
+
+import (
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
+)
+
+// Modifies TraceAnalyzer internal state to add stats for the processor/stream/flow specified
+// in stats.ComponentID and the given node ID.
+func (a *TraceAnalyzer) AddComponentStats(
+	nodeID roachpb.NodeID, stats *execinfrapb.ComponentStats,
+) {
+	switch stats.Component.Type {
+	case execinfrapb.ComponentID_PROCESSOR:
+		processorStat := &processorStats{
+			nodeID: nodeID,
+			stats:  stats,
+		}
+		if a.FlowMetadata.processorStats == nil {
+			a.FlowMetadata.processorStats = make(map[execinfrapb.ProcessorID]*processorStats)
+		}
+		a.FlowMetadata.processorStats[execinfrapb.ProcessorID(stats.Component.ID)] = processorStat
+	case execinfrapb.ComponentID_STREAM:
+		streamStat := &streamStats{
+			originNodeID: nodeID,
+			stats:        stats,
+		}
+		if a.FlowMetadata.streamStats == nil {
+			a.FlowMetadata.streamStats = make(map[execinfrapb.StreamID]*streamStats)
+		}
+		a.FlowMetadata.streamStats[execinfrapb.StreamID(stats.Component.ID)] = streamStat
+	default:
+		flowStat := &flowStats{
+			nodeID: nodeID,
+		}
+		flowStat.stats = append(flowStat.stats, stats)
+		if a.FlowMetadata.flowStats == nil {
+			a.FlowMetadata.flowStats = make(map[execinfrapb.FlowID]*flowStats)
+		}
+		a.FlowMetadata.flowStats[stats.Component.FlowID] = flowStat
+	}
+}