colexec, colflow: add network latency stat to streams

pkg/server/server_sql.go

@@ -355,6 +355,10 @@ func newSQLServer(ctx context.Context, cfg sqlServerArgs) (*sqlServer, error) {
 		RPCContext:     cfg.rpcContext,
 		Stopper:        cfg.stopper,
 
+		LatencyGetter: &serverpb.LatencyGetter{
+			NodesStatusServer: &cfg.nodesStatusServer,
+		},
+
 		TempStorage:     tempEngine,
 		TempStoragePath: cfg.TempStorageConfig.Path,
 		TempFS:          tempFS,

pkg/server/serverpb/BUILD.bazel

@@ -31,6 +31,7 @@ go_library(
         "//pkg/util/errorutil",
         "//pkg/util/log",
         "//pkg/util/metric",
+        "//pkg/util/timeutil",
         "//pkg/util/uuid",
         "//vendor/github.com/gogo/protobuf/proto",
         "//vendor/github.com/gogo/protobuf/sortkeys",

pkg/server/serverpb/status.go

@@ -12,8 +12,11 @@ package serverpb
 
 import (
 	context "context"
+	"time"
 
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/util/errorutil"
+	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 )
 
 // SQLStatusServer is a smaller version of the serverpb.StatusInterface which
@@ -61,3 +64,49 @@ func (s *OptionalNodesStatusServer) OptionalNodesStatusServer(
 	}
 	return v.(NodesStatusServer), nil
 }
+
+// LatencyGetter stores the map of latencies obtained from the NodesStatusServer.
+// These latencies are displayed on the streams of EXPLAIN ANALYZE diagrams.
+// This struct is put here to avoid import cycles.
+type LatencyGetter struct {
+	latencyMap        map[roachpb.NodeID]map[roachpb.NodeID]int64
+	lastUpdatedTime   time.Time
+	NodesStatusServer *OptionalNodesStatusServer
+}
+
+const updateThreshold = 5 * time.Second
+
+// GetLatency is a helper function that updates the latencies between nodes
+// if the time since the last update exceeds the updateThreshold. This function
+// returns the latency between the origin node and the target node.
+func (lg *LatencyGetter) GetLatency(
+	ctx context.Context, originNodeID roachpb.NodeID, targetNodeID roachpb.NodeID,
+) int64 {
+	if timeutil.Since(lg.lastUpdatedTime) < updateThreshold {
+		return lg.latencyMap[originNodeID][targetNodeID]
+	}
+	// Update latencies in latencyMap.
+	ss, err := lg.NodesStatusServer.OptionalNodesStatusServer(errorutil.FeatureNotAvailableToNonSystemTenantsIssue)
+	if err != nil {
+		// When latency is 0, it is not shown on EXPLAIN ANALYZE diagrams.
+		return 0
+	}
+	if lg.latencyMap == nil {
+		lg.latencyMap = make(map[roachpb.NodeID]map[roachpb.NodeID]int64)
+	}
+	response, _ := ss.Nodes(ctx, &NodesRequest{})
+	for _, sendingNode := range response.Nodes {
+		sendingNodeID := sendingNode.Desc.NodeID
+		if lg.latencyMap[sendingNodeID] == nil {
+			lg.latencyMap[sendingNodeID] = make(map[roachpb.NodeID]int64)
+		}
+		for _, receivingNode := range response.Nodes {
+			receivingNodeID := receivingNode.Desc.NodeID
+			if sendingNodeID != receivingNodeID {
+				lg.latencyMap[sendingNodeID][receivingNodeID] = sendingNode.Activity[receivingNodeID].Latency
+			}
+		}
+	}
+	lg.lastUpdatedTime = timeutil.Now()
+	return lg.latencyMap[originNodeID][targetNodeID]
+}

pkg/sql/colexec/execpb/stats.go

@@ -21,6 +21,8 @@ import (
 
 var _ tracing.SpanStats = &VectorizedStats{}
 var _ execinfrapb.DistSQLSpanStats = &VectorizedStats{}
+var _ tracing.SpanStats = &VectorizedInboxStats{}
+var _ execinfrapb.DistSQLSpanStats = &VectorizedInboxStats{}
 
 const (
 	batchesOutputTagSuffix     = "output.batches"
@@ -31,6 +33,7 @@ const (
 	maxVecDiskBytesTagSuffix   = "disk.vectorized.max"
 	bytesReadTagSuffix         = "bytes.read"
 	rowsReadTagSuffix          = "rows.read"
+	networkLatencyTagSuffix    = "network.latency"
 )
 
 // Stats is part of SpanStats interface.
@@ -57,6 +60,26 @@ func (vs *VectorizedStats) Stats() map[string]string {
 	return stats
 }
 
+// Stats is part of SpanStats interface.
+func (vs *VectorizedInboxStats) Stats() map[string]string {
+	stats := map[string]string{
+		batchesOutputTagSuffix: fmt.Sprintf("%d", vs.BaseVectorizedStats.NumBatches),
+		// TODO(cathymw): Have inbox collect its own deserialization time with a
+		// timer and display deserialization time instead of ioTime.
+		ioTimeTagSuffix: fmt.Sprintf("%v", vs.BaseVectorizedStats.Time.Round(time.Microsecond)),
+	}
+	if vs.BaseVectorizedStats.BytesRead != 0 {
+		stats[bytesReadTagSuffix] = humanizeutil.IBytes(vs.BaseVectorizedStats.BytesRead)
+	}
+	if vs.BaseVectorizedStats.RowsRead != 0 {
+		stats[rowsReadTagSuffix] = fmt.Sprintf("%d", vs.BaseVectorizedStats.RowsRead)
+	}
+	if vs.NetworkLatency != 0 {
+		stats[networkLatencyTagSuffix] = fmt.Sprintf("%v", time.Duration(vs.NetworkLatency).Round(time.Microsecond))
+	}
+	return stats
+}
+
 const (
 	batchesOutputQueryPlanSuffix     = "batches output"
 	tuplesOutputQueryPlanSuffix      = "tuples output"
@@ -66,6 +89,7 @@ const (
 	maxVecDiskBytesQueryPlanSuffix   = "max vectorized disk allocated"
 	bytesReadQueryPlanSuffix         = "bytes read"
 	rowsReadQueryPlanSuffix          = "rows read"
+	networkLatencyQueryPlanSuffix    = "network latency"
 )
 
 // StatsForQueryPlan is part of DistSQLSpanStats interface.
@@ -98,3 +122,23 @@ func (vs *VectorizedStats) StatsForQueryPlan() []string {
 	}
 	return stats
 }
+
+// StatsForQueryPlan is part of DistSQLSpanStats interface.
+func (vs *VectorizedInboxStats) StatsForQueryPlan() []string {
+	stats := []string{
+		fmt.Sprintf("%s: %d", batchesOutputQueryPlanSuffix, vs.BaseVectorizedStats.NumBatches),
+		fmt.Sprintf("%s: %v", ioTimeQueryPlanSuffix, vs.BaseVectorizedStats.Time.Round(time.Microsecond)),
+	}
+	if vs.BaseVectorizedStats.BytesRead != 0 {
+		stats = append(stats,
+			fmt.Sprintf("%s: %s", bytesReadQueryPlanSuffix, humanizeutil.IBytes(vs.BaseVectorizedStats.BytesRead)))
+	}
+	if vs.BaseVectorizedStats.RowsRead != 0 {
+		stats = append(stats, fmt.Sprintf("%s: %d", rowsReadQueryPlanSuffix, vs.BaseVectorizedStats.RowsRead))
+	}
+	if vs.NetworkLatency != 0 {
+		stats = append(stats,
+			fmt.Sprintf("%s: %v", networkLatencyQueryPlanSuffix, time.Duration(vs.NetworkLatency).Round(time.Microsecond)))
+	}
+	return stats
+}