wip

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
 	}
 )
 

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 {

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

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 {

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
+	}
 }