From bc93ac7eaad337a29624f89d995a11e4a028c42d Mon Sep 17 00:00:00 2001
From: shanth96 <shanth.sathiyaseelan@shopify.com>
Date: Thu, 8 Aug 2024 11:14:23 -0400
Subject: [PATCH] more tracing

---
 go/vt/vtgate/executor.go                    | 2 ++
 go/vt/vtgate/plan_execute.go                | 2 +-
 go/vt/vttablet/tabletserver/query_engine.go | 2 +-
 3 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/go/vt/vtgate/executor.go b/go/vt/vtgate/executor.go
index 0854f74631a..8898f5edeb6 100644
--- a/go/vt/vtgate/executor.go
+++ b/go/vt/vtgate/executor.go
@@ -1084,6 +1084,7 @@ func (e *Executor) getPlan(
 		return nil, err
 	}
 
+	normalSpan, ctx := trace.NewSpan(ctx, "getPlan.normalize")
 	// Normalize if possible
 	shouldNormalize := e.canNormalizeStatement(stmt, setVarComment)
 	parameterize := allowParameterization && shouldNormalize
@@ -1103,6 +1104,7 @@ func (e *Executor) getPlan(
 	if err != nil {
 		return nil, err
 	}
+	normalSpan.Finish()
 	stmt = rewriteASTResult.AST
 	bindVarNeeds := rewriteASTResult.BindVarNeeds
 	if shouldNormalize {
diff --git a/go/vt/vtgate/plan_execute.go b/go/vt/vtgate/plan_execute.go
index 6a1642eb454..aea9ba78014 100644
--- a/go/vt/vtgate/plan_execute.go
+++ b/go/vt/vtgate/plan_execute.go
@@ -92,6 +92,7 @@ func (e *Executor) newExecute(
 	var lastVSchemaCreated time.Time
 	vs := e.VSchema()
 	lastVSchemaCreated = vs.GetCreated()
+	step2Span.Finish()
 	for try := 0; try < MaxBufferingRetries; try++ {
 		if try > 0 && !vs.GetCreated().After(lastVSchemaCreated) {
 			// There is a race due to which the executor's vschema may not have been updated yet.
@@ -105,7 +106,6 @@ func (e *Executor) newExecute(
 		if err != nil {
 			return err
 		}
-		step2Span.Finish()
 
 		// 3: Create a plan for the query
 		// If we are retrying, it is likely that the routing rules have changed and hence we need to
diff --git a/go/vt/vttablet/tabletserver/query_engine.go b/go/vt/vttablet/tabletserver/query_engine.go
index 80f9e0b71be..a216257e011 100644
--- a/go/vt/vttablet/tabletserver/query_engine.go
+++ b/go/vt/vttablet/tabletserver/query_engine.go
@@ -379,7 +379,7 @@ func (qe *QueryEngine) getPlan(curSchema *currentSchema, sql string) (*TabletPla
 
 // GetPlan returns the TabletPlan that for the query. Plans are cached in an LRU cache.
 func (qe *QueryEngine) GetPlan(ctx context.Context, logStats *tabletenv.LogStats, sql string, skipQueryPlanCache bool) (*TabletPlan, error) {
-	span, _ := trace.NewSpan(ctx, "QueryEngine.GetPlan")
+	span, ctx := trace.NewSpan(ctx, "QueryEngine.GetPlan")
 	defer span.Finish()
 
 	var plan *TabletPlan