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

@@ -29,6 +29,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"
 )
 
@@ -328,7 +329,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)
@@ -373,21 +374,27 @@ func (g *routineGenerator) startInternal(ctx context.Context, txn *kv.Txn) (err
 			params := runParams{ctx, g.p.ExtendedEvalContext(), g.p}
 			var builder *sqlstats.RecordedStatementStatsBuilder
 			if stmt.FingerprintId != 0 {
+				// TODO: What other builder options should we set here?
 				builder = sqlstats.NewRecordedStatementStatsBuilder(appstatspb.StmtFingerprintID(stmt.FingerprintId), g.p.extendedEvalCtx.SessionID, stmt.DbName, stmt, stmt.AppName).
 					PlanFlags(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.StatementError(err)
 				return err
 			}
+			// TODO: What other stats can we record here?
 			builder.QueryLevelStats(queryStats)
 			forwardInnerQueryStats(g.p.routineMetadataForwarder, queryStats)
 			if openCursor {

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

@@ -13,6 +13,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"
 )
 
@@ -46,7 +47,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
@@ -433,3 +434,62 @@ func (s RoutineBodyStmt) GetQueryID() clusterunique.ID {
 func (s RoutineBodyStmt) StatementType() StatementType {
 	return s.StmtType
 }
+
+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()
+}