sql: Add latency recordings to udf sub statement calls

Adds RoutineStatementLatencyRecorder struct that implements
sqlstats.StatementLatencyRecorder. Now statement latencies
are recorded

Epic: CRDB-55081
Release note: None

Author: kyle-a-wong

pkg/sql/opt/exec/execbuilder/BUILD.bazel

@@ -57,6 +57,7 @@ go_library(
         "//pkg/util/metamorphic",
         "//pkg/util/timeutil",
         "//pkg/util/treeprinter",
+        "@com_github_cockroachdb_crlib//crtime",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_cockroachdb_redact//:redact",
     ],

pkg/sql/opt/exec/execbuilder/scalar.go

@@ -28,6 +28,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util/errorutil"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/crlib/crtime"
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/redact"
 )
@@ -883,7 +884,9 @@ func (b *Builder) buildSubquery(
 			if err != nil {
 				return err
 			}
-			err = fn(plan, tree.RoutineBodyStmt{} /* stmtForDistSQLDiagram */, true /* isFinalPlan */)
+			recorder := tree.NewStatementLatencyRecorder()
+			recorder.RecordPhase(tree.RoutineStatementStarted, crtime.NowMono())
+			err = fn(plan, tree.RoutineBodyStmt{} /* stmt */, recorder, true /* isFinalPlan */)
 			if err != nil {
 				return err
 			}
@@ -1175,6 +1178,9 @@ func (b *Builder) buildRoutinePlanGenerator(
 
 		format := tree.FmtHideConstants | tree.FmtFlags(tree.QueryFormattingForFingerprintsMask.Get(&b.evalCtx.Settings.SV))
 		for i := range stmts {
+			latencyRecorder := tree.NewStatementLatencyRecorder()
+			latencyRecorder.RecordPhase(tree.RoutineStatementStarted, crtime.NowMono())
+			latencyRecorder.RecordPhase(tree.RoutineStatementStartParsing, crtime.NowMono())
 			var bodyStmt tree.RoutineBodyStmt
 			stmt := stmts[i]
 			props := stmtProps[i]
@@ -1208,7 +1214,8 @@ func (b *Builder) buildRoutinePlanGenerator(
 				AppName:        appName,
 				FingerprintId:  uint64(fpId),
 			}
-
+			latencyRecorder.RecordPhase(tree.RoutineStatementEndParsing, crtime.NowMono())
+			latencyRecorder.RecordPhase(tree.RoutineStatementStartPlanning, crtime.NowMono())
 			o.Init(ctx, b.evalCtx, b.catalog)
 			f := o.Factory()
 
@@ -1312,8 +1319,9 @@ func (b *Builder) buildRoutinePlanGenerator(
 					return expectedLazyRoutineError("subquery")
 				}
 			}
+			latencyRecorder.RecordPhase(tree.RoutineStatementEndPlanning, crtime.NowMono())
 			incrementRoutineStmtCounter(b.evalCtx.StartedRoutineStatementCounters, dbName, appName, tag)
-			err = fn(plan, bodyStmt, isFinalPlan)
+			err = fn(plan, bodyStmt, latencyRecorder, isFinalPlan)
 			if err != nil {
 				return err
 			}

pkg/sql/routine.go

@@ -30,6 +30,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/metamorphic"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
+	"github.com/cockroachdb/crlib/crtime"
 	"github.com/cockroachdb/errors"
 )
 
@@ -329,7 +330,7 @@ func (g *routineGenerator) startInternal(ctx context.Context, txn *kv.Txn) (err
 	rrw := NewRowResultWriter(&g.rch)
 	var cursorHelper *plpgsqlCursorHelper
 	err = g.expr.ForEachPlan(ctx, ef, rrw, g.args,
-		func(plan tree.RoutinePlan, stmt tree.RoutineBodyStmt, isFinalPlan bool) error {
+		func(plan tree.RoutinePlan, stmt tree.RoutineBodyStmt, latencyRecorder *tree.RoutineStatementLatencyRecorder, isFinalPlan bool) error {
 			stmtIdx++
 			opName := "routine-stmt-" + g.expr.Name + "-" + strconv.Itoa(stmtIdx)
 			ctx, sp := tracing.ChildSpan(ctx, opName)
@@ -372,26 +373,30 @@ func (g *routineGenerator) startInternal(ctx context.Context, txn *kv.Txn) (err
 
 			// Run the plan.
 			params := runParams{ctx, g.p.ExtendedEvalContext(), g.p}
-			var builder sqlstats.RecordedStatementStatsBuilder[sqlstats.StatementLatencyRecorder, topLevelQueryStats, PlanInfo]
+			var builder sqlstats.RecordedStatementStatsBuilder[*tree.RoutineStatementLatencyRecorder, topLevelQueryStats, PlanInfo]
 			if stmt.FingerprintId != 0 {
-				builder = sqlstats.NewRecordedStatementStatsBuilder[sqlstats.StatementLatencyRecorder, topLevelQueryStats, PlanInfo](
+				builder = sqlstats.NewRecordedStatementStatsBuilder[*tree.RoutineStatementLatencyRecorder, topLevelQueryStats, PlanInfo](
 					appstatspb.StmtFingerprintID(stmt.FingerprintId),
 					g.p.extendedEvalCtx.SessionID,
 					stmt.DbName,
 					stmt.FingerprintStr,
 					stmt.SummaryStr,
 					clusterunique.ID{},
-					stmt.StmtType, stmt.AppName,
-					PlanInfo{
+					stmt.StmtType,
+					stmt.AppName, PlanInfo{
 						planFlags: g.p.curPlan.flags, // TODO: Is this right?
 						planGist:  explain.PlanGist{},
 					})
 				defer func() {
-					g.p.ExtendedEvalContext().sqlStatsIngester.RecordStatement(builder.Build())
+					latencyRecorder.RecordPhase(tree.RoutineStatementEndExec, crtime.NowMono())
+					// TODO: only record if statement stats recording is enabled. We can also avoid building the
+					// 	tree.RoutineBodyStmt struct passed into this function if stats recording is disabled.
+					g.p.ExtendedEvalContext().sqlStatsIngester.RecordStatement(builder.LatencyRecorder(latencyRecorder).Build())
 				}()
 			}
+			latencyRecorder.RecordPhase(tree.RoutineStatementStartExec, crtime.NowMono())
 			queryStats, err := runPlanInsidePlan(ctx, params, plan.(*planComponents), w, g, stmt.StmtString)
-
+			latencyRecorder.RecordPhase(tree.RoutineStatementEndExec, crtime.NowMono())
 			if err != nil {
 				builder = builder.StatementError(err)
 				return err

pkg/sql/sem/tree/BUILD.bazel

@@ -189,6 +189,7 @@ go_library(
         "//pkg/util/uuid",
         "//pkg/util/vector",
         "@com_github_cockroachdb_apd_v3//:apd",
+        "@com_github_cockroachdb_crlib//crtime",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_cockroachdb_redact//:redact",
         "@com_github_google_go_cmp//cmp",

pkg/sql/sem/tree/routine.go

@@ -12,6 +12,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
+	"github.com/cockroachdb/crlib/crtime"
 	"github.com/cockroachdb/errors"
 )
 
@@ -45,7 +46,7 @@ type RoutineBodyStmt struct {
 // to specify the SQL stmt corresponding to the plan.
 // - isFinalPlan is true if no more plans will be generated after the current
 // plan.
-type RoutinePlanGeneratedFunc func(plan RoutinePlan, stmt RoutineBodyStmt, isFinalPlan bool) error
+type RoutinePlanGeneratedFunc func(plan RoutinePlan, stmt RoutineBodyStmt, latencyRecorder *RoutineStatementLatencyRecorder, isFinalPlan bool) error
 
 // RoutinePlan represents a plan for a statement in a routine. It currently maps
 // to exec.Plan. We use the empty interface here rather than exec.Plan to avoid
@@ -412,3 +413,62 @@ func (node *TxnControlExpr) Walk(v Visitor) Expr {
 	// Cannot walk into a TxnOp, so this is a no-op.
 	return node
 }
+
+type RoutineStatementPhase int
+
+const (
+	RoutineStatementStarted RoutineStatementPhase = iota
+	RoutineStatementStartParsing
+	RoutineStatementEndParsing
+	RoutineStatementStartPlanning
+	RoutineStatementEndPlanning
+	RoutineStatementStartExec
+	RoutineStatementEndExec
+	RoutineStatementEnd
+	RoutineStatementNumPhases
+)
+
+type RoutineStatementLatencyRecorder struct {
+	times [RoutineStatementNumPhases]crtime.Mono
+}
+
+func NewStatementLatencyRecorder() *RoutineStatementLatencyRecorder {
+	return &RoutineStatementLatencyRecorder{}
+}
+
+func (r *RoutineStatementLatencyRecorder) RecordPhase(
+	phase RoutineStatementPhase, time crtime.Mono,
+) {
+	r.times[phase] = time
+}
+
+func (r *RoutineStatementLatencyRecorder) RunLatency() time.Duration {
+	return r.times[RoutineStatementEndExec].Sub(r.times[RoutineStatementStartExec])
+}
+
+// IdleLatency There is no idle latency for routine statements
+func (r *RoutineStatementLatencyRecorder) IdleLatency() time.Duration {
+	return 0
+}
+
+func (r *RoutineStatementLatencyRecorder) ServiceLatency() time.Duration {
+	return r.times[RoutineStatementEndExec].Sub(r.times[RoutineStatementStarted])
+}
+func (r *RoutineStatementLatencyRecorder) ParsingLatency() time.Duration {
+	return r.times[RoutineStatementEndParsing].Sub(r.times[RoutineStatementStartParsing])
+}
+func (r *RoutineStatementLatencyRecorder) PlanningLatency() time.Duration {
+	return r.times[RoutineStatementEndPlanning].Sub(r.times[RoutineStatementStartPlanning])
+}
+func (r *RoutineStatementLatencyRecorder) ProcessingLatency() time.Duration {
+	return r.ServiceLatency() // TODO: is there a difference between processing latency and service latency?
+}
+func (r *RoutineStatementLatencyRecorder) ExecOverheadLatency() time.Duration {
+	return r.ServiceLatency() - r.ProcessingLatency() // Will be 0 since processing and service latency are the same
+}
+func (r *RoutineStatementLatencyRecorder) StartTime() time.Time {
+	return r.times[RoutineStatementStarted].ToUTC()
+}
+func (r *RoutineStatementLatencyRecorder) EndTime() time.Time {
+	return r.times[RoutineStatementEnd].ToUTC()
+}

pkg/sql/sqlstats/ssprovider.go

@@ -211,6 +211,10 @@ func NewRecordedStatementStatsBuilder[L StatementLatencyRecorder, Q QueryStats,
 func (b RecordedStatementStatsBuilder[L, Q, P]) LatencyRecorder(
 	recorder L,
 ) RecordedStatementStatsBuilder[L, Q, P] {
+	if b.stmtStats == nil {
+		return b
+	}
+
 	b.stmtStats.RunLatencySec = recorder.RunLatency().Seconds()
 	b.stmtStats.IdleLatencySec = recorder.IdleLatency().Seconds()
 	b.stmtStats.ServiceLatencySec = recorder.ServiceLatency().Seconds()
@@ -234,6 +238,9 @@ func (b RecordedStatementStatsBuilder[L, Q, P]) QueryLevelStats(
 func (b RecordedStatementStatsBuilder[L, Q, P]) ExecStats(
 	execStats *execstats.QueryLevelStats,
 ) RecordedStatementStatsBuilder[L, Q, P] {
+	if b.stmtStats == nil {
+		return b
+	}
 	if execStats == nil {
 		return b
 	}
@@ -254,13 +261,19 @@ func (b RecordedStatementStatsBuilder[L, Q, P]) ExecStats(
 func (b RecordedStatementStatsBuilder[L, Q, P]) Indexes(
 	indexes []string,
 ) RecordedStatementStatsBuilder[L, Q, P] {
+	if b.stmtStats == nil {
+		return b
+	}
 	b.stmtStats.Indexes = indexes
 	return b
 }
 
 func (b RecordedStatementStatsBuilder[L, Q, P]) StatementError(
 	stmtErr error,
 ) RecordedStatementStatsBuilder[L, Q, P] {
+	if b.stmtStats == nil {
+		return b
+	}
 	if stmtErr == nil {
 		return b
 	}
@@ -272,6 +285,9 @@ func (b RecordedStatementStatsBuilder[L, Q, P]) StatementError(
 func (b RecordedStatementStatsBuilder[L, Q, P]) AutoRetry(
 	autoRetryCount int, autoRetryReason error,
 ) RecordedStatementStatsBuilder[L, Q, P] {
+	if b.stmtStats == nil {
+		return b
+	}
 	b.stmtStats.AutoRetryCount = autoRetryCount
 	b.stmtStats.AutoRetryReason = autoRetryReason
 	return b
@@ -280,24 +296,36 @@ func (b RecordedStatementStatsBuilder[L, Q, P]) AutoRetry(
 func (b RecordedStatementStatsBuilder[L, Q, P]) RowsAffected(
 	rowsAffected int,
 ) RecordedStatementStatsBuilder[L, Q, P] {
+	if b.stmtStats == nil {
+		return b
+	}
 	b.stmtStats.RowsAffected = rowsAffected
 	return b
 }
 
 func (b RecordedStatementStatsBuilder[L, Q, P]) IndexRecommendations(
 	idxRecommendations []string,
 ) RecordedStatementStatsBuilder[L, Q, P] {
+	if b.stmtStats == nil {
+		return b
+	}
 	b.stmtStats.IndexRecommendations = idxRecommendations
 	return b
 }
 
 func (b RecordedStatementStatsBuilder[L, Q, P]) UnderOuterTxn() RecordedStatementStatsBuilder[L, Q, P] {
+	if b.stmtStats == nil {
+		return b
+	}
 	b.stmtStats.UnderOuterTxn = true
 	return b
 }
 func (b RecordedStatementStatsBuilder[L, Q, P]) QueryTags(
 	queryTags []sqlcommenter.QueryTag,
 ) RecordedStatementStatsBuilder[L, Q, P] {
+	if b.stmtStats == nil {
+		return b
+	}
 	b.stmtStats.QueryTags = queryTags
 	return b
 }