[release-18.0] vtexplain: Ensure memory topo is set up for throttler (#15279)

go/cmd/vtexplain/cli/vtexplain.go

@@ -24,6 +24,7 @@ import (
 	"vitess.io/vitess/go/acl"
 	"vitess.io/vitess/go/vt/logutil"
 	"vitess.io/vitess/go/vt/servenv"
+	"vitess.io/vitess/go/vt/topo/memorytopo"
 	"vitess.io/vitess/go/vt/vtexplain"
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
 
@@ -175,7 +176,9 @@ func parseAndRun() error {
 		Target:          dbName,
 	}
 
-	vte, err := vtexplain.Init(context.Background(), vschema, schema, ksShardMap, opts)
+	ctx := context.Background()
+	ts := memorytopo.NewServer(ctx, vtexplain.Cell)
+	vte, err := vtexplain.Init(context.Background(), ts, vschema, schema, ksShardMap, opts)
 	if err != nil {
 		return err
 	}

go/vt/srvtopo/watch.go

@@ -23,6 +23,7 @@ import (
 	"time"
 
 	"vitess.io/vitess/go/stats"
+	"vitess.io/vitess/go/timer"
 	"vitess.io/vitess/go/vt/log"
 	"vitess.io/vitess/go/vt/topo"
 )
@@ -204,8 +205,11 @@ func (entry *watchEntry) onErrorLocked(ctx context.Context, err error, init bool
 			entry.value = nil
 		}
 	} else {
-		entry.lastError = fmt.Errorf("ResilientWatch stream failed for %v: %w", entry.key, err)
-		log.Errorf("%v", entry.lastError)
+		if !topo.IsErrType(err, topo.Interrupted) {
+			// No need to log if we're explicitly interrupted.
+			entry.lastError = fmt.Errorf("ResilientWatch stream failed for %v: %w", entry.key, err)
+			log.Errorf("%v", entry.lastError)
+		}
 
 		// Even though we didn't get a new value, update the lastValueTime
 		// here since the watch was successfully running before and we want
@@ -224,8 +228,7 @@ func (entry *watchEntry) onErrorLocked(ctx context.Context, err error, init bool
 
 	if len(entry.listeners) > 0 && !topo.IsErrType(err, topo.Interrupted) {
 		go func() {
-			time.Sleep(entry.rw.cacheRefreshInterval)
-
+			_ = timer.SleepContext(ctx, entry.rw.cacheRefreshInterval)
 			entry.mutex.Lock()
 			entry.ensureWatchingLocked(ctx)
 			entry.mutex.Unlock()

go/vt/srvtopo/watch_srvkeyspace.go

@@ -40,7 +40,7 @@ func (k *srvKeyspaceKey) String() string {
 func NewSrvKeyspaceWatcher(ctx context.Context, topoServer *topo.Server, counts *stats.CountersWithSingleLabel, cacheRefresh, cacheTTL time.Duration) *SrvKeyspaceWatcher {
 	watch := func(entry *watchEntry) {
 		key := entry.key.(*srvKeyspaceKey)
-		requestCtx, requestCancel := context.WithCancel(context.Background())
+		requestCtx, requestCancel := context.WithCancel(ctx)
 		defer requestCancel()
 
 		current, changes, err := topoServer.WatchSrvKeyspace(requestCtx, key.cell, key.keyspace)

go/vt/vtadmin/api.go

@@ -37,6 +37,7 @@ import (
 	"vitess.io/vitess/go/vt/concurrency"
 	"vitess.io/vitess/go/vt/log"
 	"vitess.io/vitess/go/vt/topo"
+	"vitess.io/vitess/go/vt/topo/memorytopo"
 	"vitess.io/vitess/go/vt/topo/topoproto"
 	"vitess.io/vitess/go/vt/vtadmin/cluster"
 	"vitess.io/vitess/go/vt/vtadmin/cluster/dynamic"
@@ -2148,7 +2149,8 @@ func (api *API) VTExplain(ctx context.Context, req *vtadminpb.VTExplainRequest)
 		return nil, er.Error()
 	}
 
-	vte, err := vtexplain.Init(ctx, srvVSchema, schema, shardMap, &vtexplain.Options{ReplicationMode: "ROW"})
+	ts := memorytopo.NewServer(ctx, vtexplain.Cell)
+	vte, err := vtexplain.Init(ctx, ts, srvVSchema, schema, shardMap, &vtexplain.Options{ReplicationMode: "ROW"})
 	if err != nil {
 		return nil, fmt.Errorf("error initilaizing vtexplain: %w", err)
 	}

go/vt/vtexplain/vtexplain.go

@@ -31,6 +31,7 @@ import (
 
 	"vitess.io/vitess/go/vt/discovery"
 	"vitess.io/vitess/go/vt/servenv"
+	"vitess.io/vitess/go/vt/topo"
 	"vitess.io/vitess/go/vt/vtgate"
 
 	"vitess.io/vitess/go/jsonutil"
@@ -54,7 +55,7 @@ func init() {
 }
 
 const (
-	vtexplainCell = "explainCell"
+	Cell = "explainCell"
 
 	// ModeMulti is the default mode with autocommit implemented at vtgate
 	ModeMulti = "multi"
@@ -181,7 +182,7 @@ type TabletActions struct {
 }
 
 // Init sets up the fake execution environment
-func Init(ctx context.Context, vSchemaStr, sqlSchema, ksShardMapStr string, opts *Options) (*VTExplain, error) {
+func Init(ctx context.Context, ts *topo.Server, vSchemaStr, sqlSchema, ksShardMapStr string, opts *Options) (*VTExplain, error) {
 	// Verify options
 	if opts.ReplicationMode != "ROW" && opts.ReplicationMode != "STATEMENT" {
 		return nil, fmt.Errorf("invalid replication mode \"%s\"", opts.ReplicationMode)
@@ -201,7 +202,7 @@ func Init(ctx context.Context, vSchemaStr, sqlSchema, ksShardMapStr string, opts
 		Autocommit:   true,
 	}}
 	vte.setGlobalTabletEnv(tabletEnv)
-	err = vte.initVtgateExecutor(ctx, vSchemaStr, ksShardMapStr, opts)
+	err = vte.initVtgateExecutor(ctx, ts, vSchemaStr, ksShardMapStr, opts)
 	if err != nil {
 		return nil, fmt.Errorf("initVtgateExecutor: %v", err.Error())
 	}

go/vt/vtexplain/vtexplain_test.go

@@ -29,6 +29,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"vitess.io/vitess/go/test/utils"
+	"vitess.io/vitess/go/vt/topo/memorytopo"
 
 	"vitess.io/vitess/go/vt/key"
 	"vitess.io/vitess/go/vt/proto/topodata"
@@ -49,7 +50,7 @@ type testopts struct {
 	shardmap map[string]map[string]*topo.ShardInfo
 }
 
-func initTest(ctx context.Context, mode string, opts *Options, topts *testopts, t *testing.T) *VTExplain {
+func initTest(ctx context.Context, ts *topo.Server, mode string, opts *Options, topts *testopts, t *testing.T) *VTExplain {
 	schema, err := os.ReadFile("testdata/test-schema.sql")
 	require.NoError(t, err)
 
@@ -65,7 +66,7 @@ func initTest(ctx context.Context, mode string, opts *Options, topts *testopts,
 	}
 
 	opts.ExecutionMode = mode
-	vte, err := Init(ctx, string(vSchema), string(schema), shardmap, opts)
+	vte, err := Init(ctx, ts, string(vSchema), string(schema), shardmap, opts)
 	require.NoError(t, err, "vtexplain Init error\n%s", string(schema))
 	return vte
 }
@@ -88,7 +89,8 @@ func runTestCase(testcase, mode string, opts *Options, topts *testopts, t *testi
 	t.Run(testcase, func(t *testing.T) {
 		ctx := utils.LeakCheckContext(t)
 
-		vte := initTest(ctx, mode, opts, topts, t)
+		ts := memorytopo.NewServer(ctx, Cell)
+		vte := initTest(ctx, ts, mode, opts, topts, t)
 		defer vte.Stop()
 
 		sqlFile := fmt.Sprintf("testdata/%s-queries.sql", testcase)
@@ -154,8 +156,8 @@ func TestExplain(t *testing.T) {
 
 func TestErrors(t *testing.T) {
 	ctx := utils.LeakCheckContext(t)
-
-	vte := initTest(ctx, ModeMulti, defaultTestOpts(), &testopts{}, t)
+	ts := memorytopo.NewServer(ctx, Cell)
+	vte := initTest(ctx, ts, ModeMulti, defaultTestOpts(), &testopts{}, t)
 	defer vte.Stop()
 
 	tests := []struct {
@@ -194,8 +196,8 @@ func TestErrors(t *testing.T) {
 
 func TestJSONOutput(t *testing.T) {
 	ctx := utils.LeakCheckContext(t)
-
-	vte := initTest(ctx, ModeMulti, defaultTestOpts(), &testopts{}, t)
+	ts := memorytopo.NewServer(ctx, Cell)
+	vte := initTest(ctx, ts, ModeMulti, defaultTestOpts(), &testopts{}, t)
 	defer vte.Stop()
 	sql := "select 1 from user where id = 1"
 	explains, err := vte.Run(sql)
@@ -344,7 +346,8 @@ func TestInit(t *testing.T) {
   }
 }`
 	schema := "create table table_missing_primary_vindex (id int primary key)"
-	_, err := Init(ctx, vschema, schema, "", defaultTestOpts())
+	ts := memorytopo.NewServer(ctx, Cell)
+	_, err := Init(ctx, ts, vschema, schema, "", defaultTestOpts())
 	require.Error(t, err)
 	require.Contains(t, err.Error(), "missing primary col vindex")
 }

go/vt/vtexplain/vtexplain_vtgate.go

@@ -22,26 +22,23 @@ package vtexplain
 import (
 	"context"
 	"fmt"
+	"path"
 	"sort"
 	"strings"
 
 	"vitess.io/vitess/go/cache/theine"
-	"vitess.io/vitess/go/vt/vtgate/logstats"
-	"vitess.io/vitess/go/vt/vtgate/vindexes"
-
-	"vitess.io/vitess/go/vt/topo"
-	"vitess.io/vitess/go/vt/topo/memorytopo"
-
-	"vitess.io/vitess/go/vt/vterrors"
-
 	"vitess.io/vitess/go/json2"
 	"vitess.io/vitess/go/streamlog"
 	"vitess.io/vitess/go/vt/discovery"
 	"vitess.io/vitess/go/vt/key"
 	"vitess.io/vitess/go/vt/log"
 	"vitess.io/vitess/go/vt/srvtopo"
+	"vitess.io/vitess/go/vt/topo"
+	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vtgate"
 	"vitess.io/vitess/go/vt/vtgate/engine"
+	"vitess.io/vitess/go/vt/vtgate/logstats"
+	"vitess.io/vitess/go/vt/vtgate/vindexes"
 	"vitess.io/vitess/go/vt/vttablet/queryservice"
 
 	querypb "vitess.io/vitess/go/vt/proto/query"
@@ -50,14 +47,14 @@ import (
 	vtgatepb "vitess.io/vitess/go/vt/proto/vtgate"
 )
 
-func (vte *VTExplain) initVtgateExecutor(ctx context.Context, vSchemaStr, ksShardMapStr string, opts *Options) error {
+func (vte *VTExplain) initVtgateExecutor(ctx context.Context, ts *topo.Server, vSchemaStr, ksShardMapStr string, opts *Options) error {
 	vte.explainTopo = &ExplainTopo{NumShards: opts.NumShards}
-	vte.explainTopo.TopoServer = memorytopo.NewServer(ctx, vtexplainCell)
+	vte.explainTopo.TopoServer = ts
 	vte.healthCheck = discovery.NewFakeHealthCheck(nil)
 
-	resolver := vte.newFakeResolver(ctx, opts, vte.explainTopo, vtexplainCell)
+	resolver := vte.newFakeResolver(ctx, opts, vte.explainTopo, Cell)
 
-	err := vte.buildTopology(ctx, opts, vSchemaStr, ksShardMapStr, opts.NumShards)
+	err := vte.buildTopology(ctx, ts, opts, vSchemaStr, ksShardMapStr, opts.NumShards)
 	if err != nil {
 		return err
 	}
@@ -75,7 +72,7 @@ func (vte *VTExplain) initVtgateExecutor(ctx context.Context, vSchemaStr, ksShar
 	var schemaTracker vtgate.SchemaInfo // no schema tracker for these tests
 	queryLogBufferSize := 10
 	plans := theine.NewStore[vtgate.PlanCacheKey, *engine.Plan](4*1024*1024, false)
-	vte.vtgateExecutor = vtgate.NewExecutor(ctx, vte.explainTopo, vtexplainCell, resolver, opts.Normalize, false, streamSize, plans, schemaTracker, false, opts.PlannerVersion)
+	vte.vtgateExecutor = vtgate.NewExecutor(ctx, vte.explainTopo, Cell, resolver, opts.Normalize, false, streamSize, plans, schemaTracker, false, opts.PlannerVersion)
 	vte.vtgateExecutor.SetQueryLogger(streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize))
 
 	return nil
@@ -95,7 +92,7 @@ func (vte *VTExplain) newFakeResolver(ctx context.Context, opts *Options, serv s
 	return vtgate.NewResolver(srvResolver, serv, cell, sc)
 }
 
-func (vte *VTExplain) buildTopology(ctx context.Context, opts *Options, vschemaStr string, ksShardMapStr string, numShardsPerKeyspace int) error {
+func (vte *VTExplain) buildTopology(ctx context.Context, ts *topo.Server, opts *Options, vschemaStr string, ksShardMapStr string, numShardsPerKeyspace int) error {
 	vte.explainTopo.Lock.Lock()
 	defer vte.explainTopo.Lock.Unlock()
 
@@ -120,6 +117,10 @@ func (vte *VTExplain) buildTopology(ctx context.Context, opts *Options, vschemaS
 		return err
 	}
 
+	conn, err := ts.ConnForCell(ctx, Cell)
+	if err != nil {
+		return err
+	}
 	vte.explainTopo.TabletConns = make(map[string]*explainTablet)
 	vte.explainTopo.KeyspaceShards = make(map[string]map[string]*topodatapb.ShardReference)
 	for ks, vschema := range vte.explainTopo.Keyspaces {
@@ -130,6 +131,32 @@ func (vte *VTExplain) buildTopology(ctx context.Context, opts *Options, vschemaS
 
 		vte.explainTopo.KeyspaceShards[ks] = make(map[string]*topodatapb.ShardReference)
 
+		srvPath := path.Join(topo.KeyspacesPath, ks, topo.SrvKeyspaceFile)
+		srvKeyspace := &topodatapb.SrvKeyspace{
+			Partitions: []*topodatapb.SrvKeyspace_KeyspacePartition{
+				{
+					ServedType:      topodatapb.TabletType_PRIMARY,
+					ShardReferences: shards,
+				},
+				{
+					ServedType:      topodatapb.TabletType_REPLICA,
+					ShardReferences: shards,
+				},
+				{
+					ServedType:      topodatapb.TabletType_RDONLY,
+					ShardReferences: shards,
+				},
+			},
+		}
+		data, err := srvKeyspace.MarshalVT()
+		if err != nil {
+			return err
+		}
+		_, err = conn.Update(ctx, srvPath, data, nil)
+		if err != nil {
+			return err
+		}
+
 		for _, shard := range shards {
 			// If the topology is in the middle of a reshard, there can be two shards covering the same key range (e.g.
 			// both source shard 80- and target shard 80-c0 cover the keyrange 80-c0). For the purposes of explain, we
@@ -142,14 +169,13 @@ func (vte *VTExplain) buildTopology(ctx context.Context, opts *Options, vschemaS
 			hostname := fmt.Sprintf("%s/%s", ks, shard.Name)
 			log.Infof("registering test tablet %s for keyspace %s shard %s", hostname, ks, shard.Name)
 
-			tablet := vte.healthCheck.AddFakeTablet(vtexplainCell, hostname, 1, ks, shard.Name, topodatapb.TabletType_PRIMARY, true, 1, nil, func(t *topodatapb.Tablet) queryservice.QueryService {
-				return vte.newTablet(ctx, opts, t)
+			tablet := vte.healthCheck.AddFakeTablet(Cell, hostname, 1, ks, shard.Name, topodatapb.TabletType_PRIMARY, true, 1, nil, func(t *topodatapb.Tablet) queryservice.QueryService {
+				return vte.newTablet(ctx, opts, t, ts)
 			})
 			vte.explainTopo.TabletConns[hostname] = tablet.(*explainTablet)
 			vte.explainTopo.KeyspaceShards[ks][shard.Name] = shard
 		}
 	}
-
 	return err
 }
 

go/vt/vtexplain/vtexplain_vttablet.go

@@ -24,6 +24,7 @@ import (
 	"sync"
 
 	"vitess.io/vitess/go/vt/sidecardb"
+	"vitess.io/vitess/go/vt/topo"
 
 	"vitess.io/vitess/go/mysql"
 	"vitess.io/vitess/go/mysql/collations"
@@ -34,7 +35,6 @@ import (
 	"vitess.io/vitess/go/vt/log"
 	"vitess.io/vitess/go/vt/mysqlctl"
 	"vitess.io/vitess/go/vt/sqlparser"
-	"vitess.io/vitess/go/vt/topo/memorytopo"
 	"vitess.io/vitess/go/vt/topo/topoproto"
 	"vitess.io/vitess/go/vt/vtgate/evalengine"
 
@@ -102,7 +102,7 @@ type explainTablet struct {
 
 var _ queryservice.QueryService = (*explainTablet)(nil)
 
-func (vte *VTExplain) newTablet(ctx context.Context, opts *Options, t *topodatapb.Tablet) *explainTablet {
+func (vte *VTExplain) newTablet(ctx context.Context, opts *Options, t *topodatapb.Tablet, ts *topo.Server) *explainTablet {
 	db := fakesqldb.New(nil)
 	sidecardb.AddSchemaInitQueries(db, true)
 
@@ -117,7 +117,7 @@ func (vte *VTExplain) newTablet(ctx context.Context, opts *Options, t *topodatap
 	config.EnableTableGC = false
 
 	// XXX much of this is cloned from the tabletserver tests
-	tsv := tabletserver.NewTabletServer(ctx, topoproto.TabletAliasString(t.Alias), config, memorytopo.NewServer(ctx, ""), t.Alias)
+	tsv := tabletserver.NewTabletServer(ctx, topoproto.TabletAliasString(t.Alias), config, ts, t.Alias)
 
 	tablet := explainTablet{db: db, tsv: tsv, vte: vte}
 	db.Handler = &tablet

go/vt/vtexplain/vtexplain_vttablet_test.go

@@ -20,10 +20,12 @@ import (
 	"context"
 	"encoding/json"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"vitess.io/vitess/go/vt/topo/memorytopo"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/schema"
 
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
@@ -70,7 +72,8 @@ create table t2 (
 
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
-	vte, err := Init(ctx, testVSchema, testSchema, "", opts)
+	ts := memorytopo.NewServer(ctx, Cell)
+	vte, err := Init(ctx, ts, testVSchema, testSchema, "", opts)
 	require.NoError(t, err)
 	defer vte.Stop()
 
@@ -129,17 +132,22 @@ create table test_partitioned (
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	vte := initTest(ctx, ModeMulti, defaultTestOpts(), &testopts{}, t)
+	ts := memorytopo.NewServer(ctx, Cell)
+	vte := initTest(ctx, ts, ModeMulti, defaultTestOpts(), &testopts{}, t)
 	defer vte.Stop()
 
 	tabletEnv, _ := newTabletEnvironment(ddls, defaultTestOpts())
 	vte.setGlobalTabletEnv(tabletEnv)
 
 	tablet := vte.newTablet(ctx, defaultTestOpts(), &topodatapb.Tablet{
-		Keyspace: "test_keyspace",
+		Keyspace: "ks_sharded",
 		Shard:    "-80",
-		Alias:    &topodatapb.TabletAlias{},
-	})
+		Alias: &topodatapb.TabletAlias{
+			Cell: Cell,
+		},
+	}, ts)
+
+	time.Sleep(10 * time.Millisecond)
 	se := tablet.tsv.SchemaEngine()
 	tables := se.GetSchema()