Improve TopoServer Performance and Efficiency For Keyspace Shards

go/vt/topo/keyspace.go

@@ -19,11 +19,14 @@ package topo
 import (
 	"context"
 	"path"
+	"runtime"
+	"sort"
 	"sync"
 
 	"golang.org/x/sync/errgroup"
 
 	"vitess.io/vitess/go/constants/sidecar"
+	"vitess.io/vitess/go/vt/key"
 	"vitess.io/vitess/go/vt/vterrors"
 
 	"vitess.io/vitess/go/event"
@@ -36,6 +39,12 @@ import (
 
 // This file contains keyspace utility functions
 
+// Default concurrency to use in order to avoid overhwelming the topo server.
+// This uses a heuristic based on the number of vCPUs available -- where it's
+// assumed that as larger machines are used for Vitess deployments they will
+// be able to do more concurrently.
+var defaultConcurrency = runtime.NumCPU()
+
 // KeyspaceInfo is a meta struct that contains metadata to give the
 // data more context and convenience. This is the main way we interact
 // with a keyspace.
@@ -191,9 +200,50 @@ func (ts *Server) FindAllShardsInKeyspace(ctx context.Context, keyspace string,
 		opt.Concurrency = 1
 	}
 
+	// First try to get all shards using List if we can.
+	shardsPath := path.Join(KeyspacesPath, keyspace, ShardsPath)
+	listResults, err := ts.globalCell.List(ctx, shardsPath)
+	if err != nil || len(listResults) == 0 {
+		if IsErrType(err, NoNode) {
+			// The path doesn't exist, let's see if the keyspace exists.
+			_, kerr := ts.GetKeyspace(ctx, keyspace)
+			if kerr == nil {
+				// We simply have no shards.
+				return make(map[string]*ShardInfo, 0), nil
+			}
+			return nil, vterrors.Wrapf(err, "FindAllShardsInKeyspace(%v): List", keyspace)
+		}
+		// Currently the ZooKeeper implementation does not support scans
+		// so we fall back to concurrently fetching the shards one by one.
+		// It is possible that the response is too large in which case we
+		// also fall back to the one by one fetch in that case.
+		if !IsErrType(err, NoImplementation) && !IsErrType(err, ResourceExhausted) {
+			return nil, vterrors.Wrapf(err, "FindAllShardsInKeyspace(%v): List", keyspace)
+		}
+		// Continue on with the shard by shard method.
+	} else {
+		result := make(map[string]*ShardInfo, len(listResults))
+		for _, entry := range listResults {
+			// The key looks like this: /vitess/global/keyspaces/commerce/shards/-80/Shard
+			shardName := path.Base(path.Dir(string(entry.Key))) // The base part of the dir is "-80"
+			shard := &topodatapb.Shard{}
+			if err = shard.UnmarshalVT(entry.Value); err != nil {
+				return nil, vterrors.Wrapf(err, "FindAllShardsInKeyspace(%v): bad shard data", keyspace)
+			}
+			result[shardName] = &ShardInfo{
+				keyspace:  keyspace,
+				shardName: shardName,
+				version:   entry.Version,
+				Shard:     shard,
+			}
+		}
+		return result, nil
+	}
+
+	// Fall back to the shard by shard method.
 	shards, err := ts.GetShardNames(ctx, keyspace)
 	if err != nil {
-		return nil, vterrors.Wrapf(err, "failed to get list of shards for keyspace '%v'", keyspace)
+		return nil, vterrors.Wrapf(err, "failed to get list of shard names for keyspace '%v'", keyspace)
 	}
 
 	// Keyspaces with a large number of shards and geographically distributed
@@ -245,25 +295,28 @@ func (ts *Server) FindAllShardsInKeyspace(ctx context.Context, keyspace string,
 
 // GetServingShards returns all shards where the primary is serving.
 func (ts *Server) GetServingShards(ctx context.Context, keyspace string) ([]*ShardInfo, error) {
-	shards, err := ts.GetShardNames(ctx, keyspace)
+	shards, err := ts.FindAllShardsInKeyspace(ctx, keyspace, &FindAllShardsInKeyspaceOptions{
+		Concurrency: defaultConcurrency, // Limit concurrency to avoid overwhelming the topo server.
+	})
 	if err != nil {
 		return nil, vterrors.Wrapf(err, "failed to get list of shards for keyspace '%v'", keyspace)
 	}
 
 	result := make([]*ShardInfo, 0, len(shards))
 	for _, shard := range shards {
-		si, err := ts.GetShard(ctx, keyspace, shard)
-		if err != nil {
-			return nil, vterrors.Wrapf(err, "GetShard(%v, %v) failed", keyspace, shard)
-		}
-		if !si.IsPrimaryServing {
+		if !shard.IsPrimaryServing {
 			continue
 		}
-		result = append(result, si)
+		result = append(result, shard)
 	}
 	if len(result) == 0 {
 		return nil, vterrors.Errorf(vtrpcpb.Code_FAILED_PRECONDITION, "%v has no serving shards", keyspace)
 	}
+	// Sort the shards by KeyRange for deterministic results.
+	sort.Slice(result, func(i, j int) bool {
+		return key.KeyRangeLess(result[i].KeyRange, result[j].KeyRange)
+	})
+
 	return result, nil
 }