Merge #158055

158055: execstats: Extract AdmissionWorkQueueStats from traces r=alyshanjahani-crl a=alyshanjahani-crl

This commit extracts the AdmissionWorkQueueStats that we are recording on traces and adds it to a new field AdmissionWaitTime on the QueryLevelStats object.

The total admission wait time is now displayed in explain analyze.

A follow up PR will further plumb this through to SQL Stats.

Part of: https://cockroachlabs.atlassian.net/browse/CRDB-55925

Release note (sql change): Admission control wait time is now displayed as part of EXPLAIN ANALYZE when non zero.

Co-authored-by: alyshanjahani-crl <[email protected]>

Author: craig[bot]

pkg/sql/execstats/BUILD.bazel

@@ -14,6 +14,7 @@ go_library(
         "//pkg/kv/kvpb",
         "//pkg/sql/execinfrapb",
         "//pkg/util",
+        "//pkg/util/admission/admissionpb",
         "//pkg/util/buildutil",
         "//pkg/util/optional",
         "//pkg/util/protoutil",

pkg/sql/execstats/traceanalyzer.go

@@ -12,6 +12,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
 	"github.com/cockroachdb/cockroach/pkg/util"
+	"github.com/cockroachdb/cockroach/pkg/util/admission/admissionpb"
 	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing/tracingpb"
 	"github.com/cockroachdb/errors"
@@ -123,6 +124,7 @@ type QueryLevelStats struct {
 	ContentionEvents                   []kvpb.ContentionEvent
 	RUEstimate                         float64
 	CPUTime                            time.Duration
+	AdmissionWaitTime                  time.Duration
 	// SQLInstanceIDs is an ordered list of SQL instances that were involved in
 	// query processing.
 	SQLInstanceIDs []int32
@@ -188,6 +190,7 @@ func (s *QueryLevelStats) Accumulate(other QueryLevelStats) {
 	s.ContentionEvents = append(s.ContentionEvents, other.ContentionEvents...)
 	s.RUEstimate += other.RUEstimate
 	s.CPUTime += other.CPUTime
+	s.AdmissionWaitTime += other.AdmissionWaitTime
 	s.SQLInstanceIDs = util.CombineUnique(s.SQLInstanceIDs, other.SQLInstanceIDs)
 	s.KVNodeIDs = util.CombineUnique(s.KVNodeIDs, other.KVNodeIDs)
 	s.Regions = util.CombineUnique(s.Regions, other.Regions)
@@ -397,6 +400,25 @@ func getAllContentionEvents(trace []tracingpb.RecordedSpan) []kvpb.ContentionEve
 	return contentionEvents
 }
 
+// getAdmissionWaitTime returns the total admission wait time accumulated
+// from all AdmissionWorkQueueStats events found in the given trace.
+func getAdmissionWaitTime(trace []tracingpb.RecordedSpan) time.Duration {
+	var totalWaitTime time.Duration
+	var ev admissionpb.AdmissionWorkQueueStats
+	for i := range trace {
+		trace[i].Structured(func(any *pbtypes.Any, _ time.Time) {
+			if !pbtypes.Is(any, &ev) {
+				return
+			}
+			if err := pbtypes.UnmarshalAny(any, &ev); err != nil {
+				return
+			}
+			totalWaitTime += ev.WaitDurationNanos
+		})
+	}
+	return totalWaitTime
+}
+
 // GetQueryLevelStats returns all the top-level stats in a QueryLevelStats
 // struct. GetQueryLevelStats tries to process as many stats as possible. If
 // errors occur while processing stats, GetQueryLevelStats returns the combined
@@ -417,5 +439,6 @@ func GetQueryLevelStats(
 		queryLevelStats.Accumulate(analyzer.GetQueryLevelStats())
 	}
 	queryLevelStats.ContentionEvents = getAllContentionEvents(trace)
+	queryLevelStats.AdmissionWaitTime = getAdmissionWaitTime(trace)
 	return queryLevelStats, errs
 }

pkg/sql/execstats/traceanalyzer_test.go

@@ -244,6 +244,7 @@ func TestQueryLevelStatsAccumulate(t *testing.T) {
 		ContentionTime:                     7 * time.Second,
 		LockWaitTime:                       4 * time.Second,
 		LatchWaitTime:                      3 * time.Second,
+		AdmissionWaitTime:                  1 * time.Second,
 		ContentionEvents:                   []kvpb.ContentionEvent{aEvent},
 		MaxDiskUsage:                       8,
 		RUEstimate:                         9,
@@ -280,6 +281,7 @@ func TestQueryLevelStatsAccumulate(t *testing.T) {
 		ContentionTime:                     14 * time.Second,
 		LockWaitTime:                       10 * time.Second,
 		LatchWaitTime:                      4 * time.Second,
+		AdmissionWaitTime:                  2 * time.Second,
 		ContentionEvents:                   []kvpb.ContentionEvent{bEvent},
 		MaxDiskUsage:                       15,
 		RUEstimate:                         16,
@@ -315,6 +317,7 @@ func TestQueryLevelStatsAccumulate(t *testing.T) {
 		ContentionTime:                     21 * time.Second,
 		LockWaitTime:                       14 * time.Second,
 		LatchWaitTime:                      7 * time.Second,
+		AdmissionWaitTime:                  3 * time.Second,
 		ContentionEvents:                   []kvpb.ContentionEvent{aEvent, bEvent},
 		MaxDiskUsage:                       15,
 		RUEstimate:                         25,

pkg/sql/instrumentation.go

@@ -888,6 +888,9 @@ func (ih *instrumentationHelper) emitExplainAnalyzePlanToOutputBuilder(
 		if queryStats.LatchWaitTime != 0 {
 			ob.AddLatchWaitTime(queryStats.LatchWaitTime)
 		}
+		if queryStats.AdmissionWaitTime != 0 {
+			ob.AddAdmissionWaitTime(queryStats.AdmissionWaitTime)
+		}
 
 		ob.AddMaxMemUsage(queryStats.MaxMemUsage)
 		ob.AddDistSQLNetworkStats(queryStats.DistSQLNetworkMessages, queryStats.DistSQLNetworkBytesSent)

pkg/sql/opt/exec/explain/output.go

@@ -418,6 +418,16 @@ func (ob *OutputBuilder) AddLatchWaitTime(latchWaitTime time.Duration) {
 	)
 }
 
+// AddAdmissionWaitTime adds a top-level field for the cumulative time spent waiting
+// in admission control.
+func (ob *OutputBuilder) AddAdmissionWaitTime(admissionWaitTime time.Duration) {
+	ob.AddFlakyTopLevelField(
+		DeflakeVolatile,
+		"cumulative time spent waiting in admission control",
+		string(humanizeutil.Duration(admissionWaitTime)),
+	)
+}
+
 // AddMaxMemUsage adds a top-level field for the memory used by the query.
 func (ob *OutputBuilder) AddMaxMemUsage(bytes int64) {
 	ob.AddFlakyTopLevelField(