From f1aa806f42095ab96490c352066f267cdad11418 Mon Sep 17 00:00:00 2001
From: Andres Taylor <andres@planetscale.com>
Date: Thu, 12 Sep 2024 11:36:25 +0200
Subject: [PATCH] wip

---
 go/logstats/logstats.go                  |  6 ------
 go/vt/vtgate/engine/fake_vcursor_test.go |  8 ++++++--
 go/vt/vtgate/engine/primitive.go         |  9 ++++++++-
 go/vt/vtgate/engine/trace.go             |  5 +++--
 go/vt/vtgate/vcursor_impl.go             | 14 +++++++++-----
 5 files changed, 26 insertions(+), 16 deletions(-)

diff --git a/go/logstats/logstats.go b/go/logstats/logstats.go
index 76f41497e55..fb45e8a5cc2 100644
--- a/go/logstats/logstats.go
+++ b/go/logstats/logstats.go
@@ -52,12 +52,6 @@ type (
 		SessionUUID    string
 		CachedPlan     bool
 		ActiveKeyspace string // ActiveKeyspace is the selected keyspace `use ks`
-		PrimitiveStats map[int]PrimitiveStats
-	}
-
-	PrimitiveStats struct {
-		NoOfCalls int
-		Rows      []int
 	}
 )
 
diff --git a/go/vt/vtgate/engine/fake_vcursor_test.go b/go/vt/vtgate/engine/fake_vcursor_test.go
index c9682bd4895..4c03c4bd99b 100644
--- a/go/vt/vtgate/engine/fake_vcursor_test.go
+++ b/go/vt/vtgate/engine/fake_vcursor_test.go
@@ -132,7 +132,9 @@ func (t *noopVCursor) UnresolvedTransactions(ctx context.Context, keyspace strin
 	panic("implement me")
 }
 
-func (t *noopVCursor) EnableOperatorTracing() {}
+func (t *noopVCursor) EnableOperatorTracing() func() map[int]PrimitiveStats {
+	panic("implement me")
+}
 
 func (t *noopVCursor) SetExec(ctx context.Context, name string, value string) error {
 	panic("implement me")
@@ -876,7 +878,9 @@ func (f *loggingVCursor) UnresolvedTransactions(_ context.Context, _ string) ([]
 	return f.transactionStatusOutput, nil
 }
 
-func (f *loggingVCursor) EnableOperatorTracing() {}
+func (f *loggingVCursor) EnableOperatorTracing() func() map[int]PrimitiveStats {
+	panic("implement me")
+}
 
 // SQLParser implements VCursor
 func (t *loggingVCursor) SQLParser() *sqlparser.Parser {
diff --git a/go/vt/vtgate/engine/primitive.go b/go/vt/vtgate/engine/primitive.go
index 8d20e776208..2eae55344d9 100644
--- a/go/vt/vtgate/engine/primitive.go
+++ b/go/vt/vtgate/engine/primitive.go
@@ -141,7 +141,14 @@ type (
 		// UnresolvedTransactions reads the state of all the unresolved atomic transactions in the given keyspace.
 		UnresolvedTransactions(ctx context.Context, keyspace string) ([]*querypb.TransactionMetadata, error)
 
-		EnableOperatorTracing()
+		// EnableOperatorTracing enables operator tracing for the current session
+		// the returned function should be called after the query is done to get the operator stats
+		EnableOperatorTracing() func() map[int]PrimitiveStats
+	}
+
+	PrimitiveStats struct {
+		NoOfCalls int
+		Rows      []int
 	}
 
 	// SessionActions gives primitives ability to interact with the session state
diff --git a/go/vt/vtgate/engine/trace.go b/go/vt/vtgate/engine/trace.go
index e7264068342..7e958a050d8 100644
--- a/go/vt/vtgate/engine/trace.go
+++ b/go/vt/vtgate/engine/trace.go
@@ -43,8 +43,9 @@ func (t *Trace) GetTableName() string {
 }
 
 func (t *Trace) GetFields(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
-	vcursor.EnableOperatorTracing()
-	return t.Inner.GetFields(ctx, vcursor, bindVars)
+	getter := vcursor.EnableOperatorTracing()
+	fields, err := t.Inner.GetFields(ctx, vcursor, bindVars)
+	return fields, err
 }
 
 func (t *Trace) NeedsTransaction() bool {
diff --git a/go/vt/vtgate/vcursor_impl.go b/go/vt/vtgate/vcursor_impl.go
index 75e87d1e646..0b3ff8448cf 100644
--- a/go/vt/vtgate/vcursor_impl.go
+++ b/go/vt/vtgate/vcursor_impl.go
@@ -132,6 +132,7 @@ type vcursorImpl struct {
 	warmingReadsChannel chan bool
 
 	resultsObserver resultsObserver
+	primitiveStats  map[int]engine.PrimitiveStats
 }
 
 // newVcursorImpl creates a vcursorImpl. Before creating this object, you have to separate out any marginComments that came with
@@ -520,12 +521,12 @@ func (vc *vcursorImpl) ExecutePrimitive(ctx context.Context, primitive engine.Pr
 }
 
 func (vc *vcursorImpl) logOpTraffic(primitive engine.Primitive, res *sqltypes.Result) {
-	if vc.logStats.LogOpStats() {
+	if vc.primitiveStats != nil {
 		key := int(primitive.GetID())
-		stats := vc.logStats.PrimitiveStats[key]
+		stats := vc.primitiveStats[key]
 		stats.NoOfCalls++
 		stats.Rows = append(stats.Rows, len(res.Rows))
-		vc.logStats.PrimitiveStats[key] = stats
+		vc.primitiveStats[key] = stats
 	}
 }
 
@@ -1511,6 +1512,9 @@ func (vc *vcursorImpl) GetForeignKeyChecksState() *bool {
 	return vc.fkChecksState
 }
 
-func (vc *vcursorImpl) EnableOperatorTracing() {
-	vc.logStats.EnableOpStats()
+func (vc *vcursorImpl) EnableOperatorTracing() func() map[int]engine.PrimitiveStats {
+	vc.primitiveStats = map[int]engine.PrimitiveStats{}
+	return func() map[int]engine.PrimitiveStats {
+		return vc.primitiveStats
+	}
 }