add more spans to query exec

Signed-off-by: Austen Lacy <[email protected]>

go/vt/vtgate/engine/limit.go

@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"io"
 
+	"vitess.io/vitess/go/trace"
 	"vitess.io/vitess/go/vt/vtgate/evalengine"
 
 	"vitess.io/vitess/go/sqltypes"
@@ -54,6 +55,8 @@ func (l *Limit) GetTableName() string {
 
 // TryExecute satisfies the Primitive interface.
 func (l *Limit) TryExecute(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool) (*sqltypes.Result, error) {
+	span, ctx := trace.NewSpan(ctx, "Limit.TryExecute")
+	defer span.Finish()
 	count, offset, err := l.getCountAndOffset(vcursor, bindVars)
 	if err != nil {
 		return nil, err

go/vt/vtgate/engine/route.go

@@ -25,6 +25,7 @@ import (
 	"time"
 
 	"vitess.io/vitess/go/mysql/collations"
+	"vitess.io/vitess/go/trace"
 	"vitess.io/vitess/go/vt/vtgate/evalengine"
 
 	"vitess.io/vitess/go/mysql"
@@ -173,6 +174,9 @@ func (route *Route) SetTruncateColumnCount(count int) {
 
 // TryExecute performs a non-streaming exec.
 func (route *Route) TryExecute(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool) (*sqltypes.Result, error) {
+	span, ctx := trace.NewSpan(ctx, "Route.executePlan")
+	defer span.Finish()
+
 	if route.QueryTimeout != 0 {
 		var cancel context.CancelFunc
 		ctx, cancel = context.WithTimeout(ctx, time.Duration(route.QueryTimeout)*time.Millisecond)

go/vt/vtgate/engine/rows.go

@@ -20,6 +20,7 @@ import (
 	"context"
 
 	"vitess.io/vitess/go/sqltypes"
+	"vitess.io/vitess/go/trace"
 	querypb "vitess.io/vitess/go/vt/proto/query"
 )
 
@@ -55,7 +56,9 @@ func (r *Rows) GetTableName() string {
 }
 
 // TryExecute implements the Primitive interface
-func (r *Rows) TryExecute(context.Context, VCursor, map[string]*querypb.BindVariable, bool) (*sqltypes.Result, error) {
+func (r *Rows) TryExecute(ctx context.Context, v VCursor, vars map[string]*querypb.BindVariable, wantFields bool) (*sqltypes.Result, error) {
+	span, _ := trace.NewSpan(ctx, "Rows.executePlan")
+	defer span.Finish()
 	return &sqltypes.Result{
 		Fields:   r.fields,
 		InsertID: 0,

go/vt/vtgate/engine/send.go

@@ -20,6 +20,7 @@ import (
 	"context"
 
 	"vitess.io/vitess/go/sqltypes"
+	"vitess.io/vitess/go/trace"
 	"vitess.io/vitess/go/vt/key"
 	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
 	"vitess.io/vitess/go/vt/srvtopo"
@@ -86,6 +87,8 @@ func (s *Send) GetTableName() string {
 
 // TryExecute implements Primitive interface
 func (s *Send) TryExecute(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool) (*sqltypes.Result, error) {
+	span, ctx := trace.NewSpan(ctx, "Send.TryExecute")
+	defer span.Finish()
 	rss, _, err := vcursor.ResolveDestinations(ctx, s.Keyspace.Name, nil, []key.Destination{s.TargetDestination})
 	if err != nil {
 		return nil, err

go/vt/vtgate/executor.go

@@ -381,6 +381,9 @@ func saveSessionStats(safeSession *SafeSession, stmtType sqlparser.StatementType
 }
 
 func (e *Executor) execute(ctx context.Context, safeSession *SafeSession, sql string, bindVars map[string]*querypb.BindVariable, logStats *logstats.LogStats) (sqlparser.StatementType, *sqltypes.Result, error) {
+	span, ctx := trace.NewSpan(ctx, "Executor.execute")
+	defer span.Finish()
+
 	var err error
 	var qr *sqltypes.Result
 	var stmtType sqlparser.StatementType
@@ -968,6 +971,9 @@ type iQueryOption interface {
 // getPlan computes the plan for the given query. If one is in
 // the cache, it reuses it.
 func (e *Executor) getPlan(ctx context.Context, vcursor *vcursorImpl, sql string, comments sqlparser.MarginComments, bindVars map[string]*querypb.BindVariable, qo iQueryOption, logStats *logstats.LogStats) (*engine.Plan, error) {
+	span, ctx := trace.NewSpan(ctx, "Executor.getPlan")
+	defer span.Finish()
+
 	if e.VSchema() == nil {
 		return nil, errors.New("vschema not initialized")
 	}

go/vt/vtgate/plan_execute.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"time"
 
+	"vitess.io/vitess/go/trace"
 	"vitess.io/vitess/go/vt/vtgate/logstats"
 
 	"vitess.io/vitess/go/sqltypes"
@@ -42,6 +43,9 @@ func (e *Executor) newExecute(
 	execPlan planExec, // used when there is a plan to execute
 	recResult txResult, // used when it's something simple like begin/commit/rollback/savepoint
 ) error {
+	span, ctx := trace.NewSpan(ctx, "Executor.newExecute")
+	defer span.Finish()
+
 	// 1: Prepare before planning and execution
 
 	// Start an implicit transaction if necessary.
@@ -105,6 +109,8 @@ func (e *Executor) newExecute(
 
 // handleTransactions deals with transactional queries: begin, commit, rollback and savepoint management
 func (e *Executor) handleTransactions(ctx context.Context, safeSession *SafeSession, plan *engine.Plan, logStats *logstats.LogStats, vcursor *vcursorImpl) (*sqltypes.Result, error) {
+	span, ctx := trace.NewSpan(ctx, "Executor.handleTransactions")
+	defer span.Finish()
 	// We need to explicitly handle errors, and begin/commit/rollback, since these control transactions. Everything else
 	// will fall through and be handled through planning
 	switch plan.Type {
@@ -200,6 +206,9 @@ func (e *Executor) executePlan(
 	execStart time.Time,
 ) (*sqltypes.Result, error) {
 
+	span, ctx := trace.NewSpan(ctx, "Executor.executePlan")
+	defer span.Finish()
+
 	// 4: Execute!
 	qr, err := vcursor.ExecutePrimitive(ctx, plan.Instructions, bindVars, true)
 

go/vt/vtgate/vcursor_impl.go

@@ -23,6 +23,7 @@ import (
 	"strings"
 	"sync/atomic"
 
+	"vitess.io/vitess/go/trace"
 	"vitess.io/vitess/go/vt/vtgate/logstats"
 
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
@@ -392,6 +393,8 @@ func (vc *vcursorImpl) TargetString() string {
 const MaxBufferingRetries = 3
 
 func (vc *vcursorImpl) ExecutePrimitive(ctx context.Context, primitive engine.Primitive, bindVars map[string]*querypb.BindVariable, wantfields bool) (*sqltypes.Result, error) {
+	span, ctx := trace.NewSpan(ctx, "vcursorImpl.ExecutePrimitive")
+	defer span.Finish()
 	for try := 0; try < MaxBufferingRetries; try++ {
 		res, err := primitive.TryExecute(ctx, vc, bindVars, wantfields)
 		if err != nil && vterrors.RootCause(err) == buffer.ShardMissingError {

go/vt/vttablet/tabletserver/connpool/dbconn.go

@@ -157,6 +157,8 @@ func (dbc *DBConn) Exec(ctx context.Context, query string, maxrows int, wantfiel
 }
 
 func (dbc *DBConn) execOnce(ctx context.Context, query string, maxrows int, wantfields bool) (*sqltypes.Result, error) {
+	span, ctx := trace.NewSpan(ctx, "DBConn.execOnce")
+	defer span.Finish()
 	dbc.current.Set(query)
 	defer dbc.current.Set("")
 
@@ -185,6 +187,8 @@ func (dbc *DBConn) execOnce(ctx context.Context, query string, maxrows int, want
 
 // ExecOnce executes the specified query, but does not retry on connection errors.
 func (dbc *DBConn) ExecOnce(ctx context.Context, query string, maxrows int, wantfields bool) (*sqltypes.Result, error) {
+	span, ctx := trace.NewSpan(ctx, "DBConn.ExecOnce")
+	defer span.Finish()
 	return dbc.execOnce(ctx, query, maxrows, wantfields)
 }