[vtadmin] tracing refactor

go/vt/vtadmin/api.go

@@ -313,12 +313,18 @@ func (api *API) GetKeyspaces(ctx context.Context, req *vtadminpb.GetKeyspacesReq
 				return
 			}
 
-			resp, err := c.Vtctld.GetKeyspaces(ctx, &vtctldatapb.GetKeyspacesRequest{})
+			getKeyspacesSpan, getKeyspacesCtx := trace.NewSpan(ctx, "Cluster.GetKeyspaces")
+			cluster.AnnotateSpan(c, getKeyspacesSpan)
+
+			resp, err := c.Vtctld.GetKeyspaces(getKeyspacesCtx, &vtctldatapb.GetKeyspacesRequest{})
 			if err != nil {
 				er.RecordError(fmt.Errorf("GetKeyspaces(cluster = %s): %w", c.ID, err))
+				getKeyspacesSpan.Finish()
 				return
 			}
 
+			getKeyspacesSpan.Finish()
+
 			kss := make([]*vtadminpb.Keyspace, 0, len(resp.Keyspaces))
 
 			var (
@@ -332,6 +338,13 @@ func (api *API) GetKeyspaces(ctx context.Context, req *vtadminpb.GetKeyspacesReq
 				// Find all shards for each keyspace in the cluster, in parallel
 				go func(c *cluster.Cluster, ks *vtctldatapb.Keyspace) {
 					defer kwg.Done()
+
+					span, ctx := trace.NewSpan(ctx, "Cluster.FindAllShardsInKeyspace")
+					defer span.Finish()
+
+					cluster.AnnotateSpan(c, span)
+					span.Annotate("keyspace", ks.Name)
+
 					sr, err := c.Vtctld.FindAllShardsInKeyspace(ctx, &vtctldatapb.FindAllShardsInKeyspaceRequest{
 						Keyspace: ks.Name,
 					})
@@ -462,11 +475,17 @@ func (api *API) getSchemas(ctx context.Context, c *cluster.Cluster, tablets []*v
 		return nil, err
 	}
 
-	resp, err := c.Vtctld.GetKeyspaces(ctx, &vtctldatapb.GetKeyspacesRequest{})
+	getKeyspacesSpan, getKeyspacesCtx := trace.NewSpan(ctx, "Cluster.GetKeyspaces")
+	cluster.AnnotateSpan(c, getKeyspacesSpan)
+
+	resp, err := c.Vtctld.GetKeyspaces(getKeyspacesCtx, &vtctldatapb.GetKeyspacesRequest{})
 	if err != nil {
+		getKeyspacesSpan.Finish()
 		return nil, err
 	}
 
+	getKeyspacesSpan.Finish()
+
 	var (
 		schemas []*vtadminpb.Schema
 		wg      sync.WaitGroup
@@ -712,7 +731,7 @@ func (api *API) GetVSchemas(ctx context.Context, req *vtadminpb.GetVSchemasReque
 		go func(c *cluster.Cluster) {
 			defer wg.Done()
 
-			span, ctx := trace.NewSpan(ctx, "API.getVSchemasForCluster")
+			span, ctx := trace.NewSpan(ctx, "Cluster.GetVSchemas")
 			defer span.Finish()
 
 			cluster.AnnotateSpan(c, span)
@@ -722,12 +741,18 @@ func (api *API) GetVSchemas(ctx context.Context, req *vtadminpb.GetVSchemasReque
 				return
 			}
 
-			keyspaces, err := c.Vtctld.GetKeyspaces(ctx, &vtctldatapb.GetKeyspacesRequest{})
+			getKeyspacesSpan, getKeyspacesCtx := trace.NewSpan(ctx, "Cluster.GetKeyspaces")
+			cluster.AnnotateSpan(c, getKeyspacesSpan)
+
+			keyspaces, err := c.Vtctld.GetKeyspaces(getKeyspacesCtx, &vtctldatapb.GetKeyspacesRequest{})
 			if err != nil {
 				rec.RecordError(fmt.Errorf("GetKeyspaces(cluster = %s): %w", c.ID, err))
+				getKeyspacesSpan.Finish()
 				return
 			}
 
+			getKeyspacesSpan.Finish()
+
 			var (
 				clusterM        sync.Mutex
 				clusterWG       sync.WaitGroup
@@ -740,7 +765,6 @@ func (api *API) GetVSchemas(ctx context.Context, req *vtadminpb.GetVSchemasReque
 
 				go func(keyspace *vtctldatapb.Keyspace) {
 					defer clusterWG.Done()
-
 					vschema, err := c.GetVSchema(ctx, keyspace.Name)
 					if err != nil {
 						clusterRec.RecordError(fmt.Errorf("GetVSchema(keyspace = %s): %w", keyspace.Name, err))
@@ -799,13 +823,18 @@ func (api *API) VTExplain(ctx context.Context, req *vtadminpb.VTExplainRequest)
 		return nil, fmt.Errorf("%w: %s", errors.ErrUnsupportedCluster, req.Cluster)
 	}
 
+	span.Annotate("keyspace", req.Keyspace)
+	cluster.AnnotateSpan(c, span)
+
 	tablet, err := c.FindTablet(ctx, func(t *vtadminpb.Tablet) bool {
 		return t.Tablet.Keyspace == req.Keyspace && topo.IsInServingGraph(t.Tablet.Type) && t.Tablet.Type != topodatapb.TabletType_MASTER && t.State == vtadminpb.Tablet_SERVING
 	})
 	if err != nil {
 		return nil, fmt.Errorf("cannot find serving, non-primary tablet in keyspace=%s: %w", req.Keyspace, err)
 	}
 
+	span.Annotate("tablet_alias", topoproto.TabletAliasString(tablet.Tablet.Alias))
+
 	if err := c.Vtctld.Dial(ctx); err != nil {
 		return nil, err
 	}
@@ -848,6 +877,12 @@ func (api *API) VTExplain(ctx context.Context, req *vtadminpb.VTExplainRequest)
 	go func(c *cluster.Cluster) {
 		defer wg.Done()
 
+		span, ctx := trace.NewSpan(ctx, "Cluster.GetSrvVSchema")
+		defer span.Finish()
+
+		span.Annotate("cell", tablet.Tablet.Alias.Cell)
+		cluster.AnnotateSpan(c, span)
+
 		res, err := c.Vtctld.GetSrvVSchema(ctx, &vtctldatapb.GetSrvVSchemaRequest{
 			Cell: tablet.Tablet.Alias.Cell,
 		})
@@ -876,6 +911,12 @@ func (api *API) VTExplain(ctx context.Context, req *vtadminpb.VTExplainRequest)
 	go func(c *cluster.Cluster) {
 		defer wg.Done()
 
+		span, ctx := trace.NewSpan(ctx, "Cluster.FindAllShardsInKeyspace")
+		defer span.Finish()
+
+		span.Annotate("keyspace", req.Keyspace)
+		cluster.AnnotateSpan(c, span)
+
 		ksm, err := c.Vtctld.FindAllShardsInKeyspace(ctx, &vtctldatapb.FindAllShardsInKeyspaceRequest{
 			Keyspace: req.Keyspace,
 		})

go/vt/vtadmin/cluster/cluster.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"strings"
 	"time"
 
 	"vitess.io/vitess/go/trace"
@@ -201,6 +202,15 @@ func (c *Cluster) parseTablet(rows *sql.Rows) (*vtadminpb.Tablet, error) {
 
 // GetTablets returns all tablets in the cluster.
 func (c *Cluster) GetTablets(ctx context.Context) ([]*vtadminpb.Tablet, error) {
+	span, ctx := trace.NewSpan(ctx, "Cluster.GetTablets")
+	defer span.Finish()
+
+	AnnotateSpan(c, span)
+
+	return c.getTablets(ctx)
+}
+
+func (c *Cluster) getTablets(ctx context.Context) ([]*vtadminpb.Tablet, error) {
 	if err := c.DB.Dial(ctx, ""); err != nil {
 		return nil, err
 	}
@@ -227,10 +237,22 @@ func (c *Cluster) GetTablets(ctx context.Context) ([]*vtadminpb.Tablet, error) {
 // bunch of tablets once and make a series of GetSchema calls without Cluster
 // refetching the tablet list each time.
 func (c *Cluster) GetSchema(ctx context.Context, req *vtctldatapb.GetSchemaRequest, tablet *vtadminpb.Tablet) (*vtadminpb.Schema, error) {
+	span, ctx := trace.NewSpan(ctx, "Cluster.GetSchema")
+	defer span.Finish()
+
+	AnnotateSpan(c, span)
+
 	// Copy the request to not mutate the caller's request object.
 	r := *req
 	r.TabletAlias = tablet.Tablet.Alias
 
+	span.Annotate("tablet_alias", topoproto.TabletAliasString(r.TabletAlias))
+	span.Annotate("exclude_tables", strings.Join(r.ExcludeTables, ","))
+	span.Annotate("tables", strings.Join(r.Tables, ","))
+	span.Annotate("include_views", r.IncludeViews)
+	span.Annotate("table_names_only", r.TableNamesOnly)
+	span.Annotate("table_sizes_only", r.TableSizesOnly)
+
 	schema, err := c.Vtctld.GetSchema(ctx, &r)
 	if err != nil {
 		return nil, err
@@ -274,7 +296,12 @@ func (c *Cluster) GetVSchema(ctx context.Context, keyspace string) (*vtadminpb.V
 
 // FindTablet returns the first tablet in a given cluster that satisfies the filter function.
 func (c *Cluster) FindTablet(ctx context.Context, filter func(*vtadminpb.Tablet) bool) (*vtadminpb.Tablet, error) {
-	tablets, err := c.FindTablets(ctx, filter, 1)
+	span, ctx := trace.NewSpan(ctx, "Cluster.FindTablet")
+	defer span.Finish()
+
+	AnnotateSpan(c, span)
+
+	tablets, err := c.findTablets(ctx, filter, 1)
 	if err != nil {
 		return nil, err
 	}
@@ -290,6 +317,17 @@ func (c *Cluster) FindTablet(ctx context.Context, filter func(*vtadminpb.Tablet)
 // the filter function. If N = -1, then all matching tablets are returned.
 // Ordering is not guaranteed, and callers should write their filter functions accordingly.
 func (c *Cluster) FindTablets(ctx context.Context, filter func(*vtadminpb.Tablet) bool, n int) ([]*vtadminpb.Tablet, error) {
+	span, ctx := trace.NewSpan(ctx, "Cluster.FindTablets")
+	defer span.Finish()
+
+	AnnotateSpan(c, span)
+
+	return c.findTablets(ctx, filter, n)
+}
+
+func (c *Cluster) findTablets(ctx context.Context, filter func(*vtadminpb.Tablet) bool, n int) ([]*vtadminpb.Tablet, error) {
+	span, _ := trace.FromContext(ctx)
+
 	tablets, err := c.GetTablets(ctx)
 	if err != nil {
 		return nil, err
@@ -299,6 +337,10 @@ func (c *Cluster) FindTablets(ctx context.Context, filter func(*vtadminpb.Tablet
 		n = len(tablets)
 	}
 
+	if span != nil {
+		span.Annotate("max_result_length", n) // this is a bad name; I didn't want just "n", but it's more like, "requested result length".
+	}
+
 	results := make([]*vtadminpb.Tablet, 0, n)
 	for _, t := range tablets {
 		if len(results) >= n {

go/vt/vtadmin/cluster/discovery/discovery_static_file.go

@@ -26,6 +26,7 @@ import (
 
 	"github.com/spf13/pflag"
 
+	"vitess.io/vitess/go/trace"
 	vtadminpb "vitess.io/vitess/go/vt/proto/vtadmin"
 )
 
@@ -139,7 +140,14 @@ func (d *StaticFileDiscovery) parseConfig(bytes []byte) error {
 
 // DiscoverVTGate is part of the Discovery interface.
 func (d *StaticFileDiscovery) DiscoverVTGate(ctx context.Context, tags []string) (*vtadminpb.VTGate, error) {
-	gates, err := d.DiscoverVTGates(ctx, tags)
+	span, ctx := trace.NewSpan(ctx, "StaticFileDiscovery.DiscoverVTGate")
+	defer span.Finish()
+
+	return d.discoverVTGate(ctx, tags)
+}
+
+func (d *StaticFileDiscovery) discoverVTGate(ctx context.Context, tags []string) (*vtadminpb.VTGate, error) {
+	gates, err := d.discoverVTGates(ctx, tags)
 	if err != nil {
 		return nil, err
 	}
@@ -155,6 +163,9 @@ func (d *StaticFileDiscovery) DiscoverVTGate(ctx context.Context, tags []string)
 
 // DiscoverVTGateAddr is part of the Discovery interface.
 func (d *StaticFileDiscovery) DiscoverVTGateAddr(ctx context.Context, tags []string) (string, error) {
+	span, ctx := trace.NewSpan(ctx, "StaticFileDiscovery.DiscoverVTGateAddr")
+	defer span.Finish()
+
 	gate, err := d.DiscoverVTGate(ctx, tags)
 	if err != nil {
 		return "", err
@@ -165,6 +176,13 @@ func (d *StaticFileDiscovery) DiscoverVTGateAddr(ctx context.Context, tags []str
 
 // DiscoverVTGates is part of the Discovery interface.
 func (d *StaticFileDiscovery) DiscoverVTGates(ctx context.Context, tags []string) ([]*vtadminpb.VTGate, error) {
+	span, ctx := trace.NewSpan(ctx, "StaticFileDiscovery.DiscoverVTGates")
+	defer span.Finish()
+
+	return d.discoverVTGates(ctx, tags)
+}
+
+func (d *StaticFileDiscovery) discoverVTGates(ctx context.Context, tags []string) ([]*vtadminpb.VTGate, error) {
 	if len(tags) == 0 {
 		results := []*vtadminpb.VTGate{}
 		for _, g := range d.gates.byName {
@@ -204,7 +222,14 @@ func (d *StaticFileDiscovery) DiscoverVTGates(ctx context.Context, tags []string
 
 // DiscoverVtctld is part of the Discovery interface.
 func (d *StaticFileDiscovery) DiscoverVtctld(ctx context.Context, tags []string) (*vtadminpb.Vtctld, error) {
-	vtctlds, err := d.DiscoverVtctlds(ctx, tags)
+	span, ctx := trace.NewSpan(ctx, "StaticFileDiscovery.DiscoverVtctld")
+	defer span.Finish()
+
+	return d.discoverVtctld(ctx, tags)
+}
+
+func (d *StaticFileDiscovery) discoverVtctld(ctx context.Context, tags []string) (*vtadminpb.Vtctld, error) {
+	vtctlds, err := d.discoverVtctlds(ctx, tags)
 	if err != nil {
 		return nil, err
 	}
@@ -220,7 +245,10 @@ func (d *StaticFileDiscovery) DiscoverVtctld(ctx context.Context, tags []string)
 
 // DiscoverVtctldAddr is part of the Discovery interface.
 func (d *StaticFileDiscovery) DiscoverVtctldAddr(ctx context.Context, tags []string) (string, error) {
-	vtctld, err := d.DiscoverVtctld(ctx, tags)
+	span, ctx := trace.NewSpan(ctx, "StaticFileDiscovery.DiscoverVtctldAddr")
+	defer span.Finish()
+
+	vtctld, err := d.discoverVtctld(ctx, tags)
 	if err != nil {
 		return "", err
 	}
@@ -230,6 +258,13 @@ func (d *StaticFileDiscovery) DiscoverVtctldAddr(ctx context.Context, tags []str
 
 // DiscoverVtctlds is part of the Discovery interface.
 func (d *StaticFileDiscovery) DiscoverVtctlds(ctx context.Context, tags []string) ([]*vtadminpb.Vtctld, error) {
+	span, ctx := trace.NewSpan(ctx, "StaticFileDiscovery.DiscoverVtctlds")
+	defer span.Finish()
+
+	return d.discoverVtctlds(ctx, tags)
+}
+
+func (d *StaticFileDiscovery) discoverVtctlds(ctx context.Context, tags []string) ([]*vtadminpb.Vtctld, error) {
 	if len(tags) == 0 {
 		results := []*vtadminpb.Vtctld{}
 		for _, v := range d.vtctlds.byName {

go/vt/vtadmin/cluster/trace.go

@@ -16,10 +16,14 @@ limitations under the License.
 
 package cluster
 
-import "vitess.io/vitess/go/trace"
+import (
+	"vitess.io/vitess/go/trace"
+	"vitess.io/vitess/go/vt/vtadmin/vtadminproto"
+)
 
 // AnnotateSpan adds the cluster_id and cluster_name to a span.
 func AnnotateSpan(c *Cluster, span trace.Span) {
-	span.Annotate("cluster_id", c.ID)
-	span.Annotate("cluster_name", c.Name)
+	vtadminproto.AnnotateClusterSpan(c.ToProto(), span)
+	// (TODO:@ajm188) add support for discovery impls to add annotations to a
+	// span, like `discovery_impl` and any parameters that might be relevant.
 }

go/vt/vtadmin/vtadminproto/trace.go

@@ -0,0 +1,29 @@
+/*
+Copyright 2021 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package vtadminproto
+
+import (
+	"vitess.io/vitess/go/trace"
+
+	vtadminpb "vitess.io/vitess/go/vt/proto/vtadmin"
+)
+
+// AnnotateClusterSpan adds the cluster_id and cluster_name to a span.
+func AnnotateClusterSpan(c *vtadminpb.Cluster, span trace.Span) {
+	span.Annotate("cluster_id", c.Id)
+	span.Annotate("cluster_name", c.Name)
+}