From 33dd285af5622f297ca3f0a013ceae43c05bfadc Mon Sep 17 00:00:00 2001 From: Dirkjan Bussink Date: Mon, 21 Aug 2023 15:09:59 +0200 Subject: [PATCH 01/17] Add leak checking for vtgate tests This also introduces context passing in a lot of places so we can properly shut things down. Signed-off-by: Dirkjan Bussink --- go.mod | 2 +- go.sum | 7 +- go/cache/cache.go | 3 + go/cache/lru_cache.go | 4 + go/cache/null.go | 2 + go/cmd/vtexplain/vtexplain.go | 3 +- go/test/utils/noleak.go | 55 + go/vt/discovery/keyspace_events.go | 26 +- go/vt/topo/memorytopo/memorytopo.go | 7 +- go/vt/vtadmin/api.go | 2 +- go/vt/vtexplain/vtexplain.go | 5 +- go/vt/vtexplain/vtexplain_test.go | 25 +- go/vt/vtexplain/vtexplain_vtgate.go | 8 +- go/vt/vtgate/autocommit_test.go | 100 +- go/vt/vtgate/bench_test.go | 110 -- .../evalengine/integration/comparison_test.go | 9 +- go/vt/vtgate/executor.go | 22 +- go/vt/vtgate/executor_ddl_test.go | 8 +- go/vt/vtgate/executor_dml_test.go | 626 ++++++---- go/vt/vtgate/executor_framework_test.go | 53 +- go/vt/vtgate/executor_scatter_stats_test.go | 32 +- go/vt/vtgate/executor_select_test.go | 1040 ++++++++++++----- go/vt/vtgate/executor_set_test.go | 53 +- go/vt/vtgate/executor_stream_test.go | 24 +- go/vt/vtgate/executor_test.go | 566 ++++++--- go/vt/vtgate/executor_vschema_ddl_test.go | 135 ++- go/vt/vtgate/executor_vstream_test.go | 12 +- go/vt/vtgate/legacy_scatter_conn_test.go | 77 +- go/vt/vtgate/mysql_protocol_test.go | 177 --- go/vt/vtgate/planbuilder/plan_test.go | 1 + go/vt/vtgate/plugin_mysql_server.go | 23 +- go/vt/vtgate/plugin_mysql_server_test.go | 25 +- go/vt/vtgate/querylog.go | 27 +- go/vt/vtgate/queryz_test.go | 17 +- go/vt/vtgate/sandbox_test.go | 7 +- go/vt/vtgate/scatter_conn_test.go | 52 +- go/vt/vtgate/tabletgateway.go | 4 +- go/vt/vtgate/tabletgateway_flaky_test.go | 36 +- go/vt/vtgate/tabletgateway_test.go | 29 +- go/vt/vtgate/tx_conn_test.go | 202 +++- go/vt/vtgate/vstream_manager_test.go | 57 +- go/vt/vtgate/vtgate.go | 10 +- 42 files changed, 2400 insertions(+), 1283 deletions(-) create mode 100644 go/test/utils/noleak.go delete mode 100644 go/vt/vtgate/bench_test.go delete mode 100644 go/vt/vtgate/mysql_protocol_test.go diff --git a/go.mod b/go.mod index 7f1b4f4d042..078b4db8922 100644 --- a/go.mod +++ b/go.mod @@ -106,7 +106,7 @@ require ( github.com/spf13/afero v1.9.3 github.com/spf13/jwalterweatherman v1.1.0 github.com/xlab/treeprint v1.2.0 - go.uber.org/goleak v1.1.11 + go.uber.org/goleak v1.2.1 golang.org/x/sync v0.1.0 modernc.org/sqlite v1.20.3 ) diff --git a/go.sum b/go.sum index 4076ec8a33f..5a3fa32201a 100644 --- a/go.sum +++ b/go.sum @@ -634,8 +634,8 @@ go.uber.org/atomic v1.9.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc= go.uber.org/atomic v1.10.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0= go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE= go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0= -go.uber.org/goleak v1.1.11 h1:wy28qYRKZgnJTxGxvye5/wgWr1EKjmUDGYox5mGlRlI= -go.uber.org/goleak v1.1.11/go.mod h1:cwTWslyiVhfpKIDGSZEM2HlOvcqm+tG4zioyIeLoqMQ= +go.uber.org/goleak v1.2.1 h1:NBol2c7O1ZokfZ0LEU9K6Whx/KnwvepVetCUhtKja4A= +go.uber.org/goleak v1.2.1/go.mod h1:qlT2yGI9QafXHhZZLxlSuNsMw3FFLxBr+tBRlmO1xH4= go.uber.org/mock v0.2.0 h1:TaP3xedm7JaAgScZO7tlvlKrqT0p7I6OsdGB5YNSMDU= go.uber.org/mock v0.2.0/go.mod h1:J0y0rp9L3xiff1+ZBfKxlC1fz2+aO16tw0tsDOixfuM= go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0= @@ -692,8 +692,6 @@ golang.org/x/lint v0.0.0-20191125180803-fdd1cda4f05f/go.mod h1:5qLYkcX4OjUUV8bRu golang.org/x/lint v0.0.0-20200130185559-910be7a94367/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY= golang.org/x/lint v0.0.0-20200302205851-738671d3881b/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY= golang.org/x/lint v0.0.0-20201208152925-83fdc39ff7b5/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY= -golang.org/x/lint v0.0.0-20210508222113-6edffad5e616 h1:VLliZ0d+/avPrXXH+OakdXhpJuEoBZuwh1m2j7U6Iug= -golang.org/x/lint v0.0.0-20210508222113-6edffad5e616/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY= golang.org/x/mobile v0.0.0-20190312151609-d3739f865fa6/go.mod h1:z+o9i4GpDbdi3rU15maQ/Ox0txvL9dWGYEHz965HBQE= golang.org/x/mobile v0.0.0-20190719004257-d2bd2a29d028/go.mod h1:E/iHnbuqvinMTCcRqshq8CkpyQDoeVncDDYHnLhea+o= golang.org/x/mod v0.0.0-20190513183733-4bf6d317e70e/go.mod h1:mXi4GBBbnImb6dmsKGUJ2LatrhH/nqhxcFungHvyanc= @@ -936,7 +934,6 @@ golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4f golang.org/x/tools v0.0.0-20210108195828-e2f9c7f1fc8e/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= golang.org/x/tools v0.1.0/go.mod h1:xkSsbof2nBLbhDlRMhhhyNLN/zl3eTqcnHD5viDpcZ0= golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk= -golang.org/x/tools v0.1.5/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk= golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc= golang.org/x/tools v0.4.0/go.mod h1:UE5sM2OK9E/d67R0ANs2xJizIymRP5gJU295PvKXxjQ= golang.org/x/tools v0.8.0 h1:vSDcovVPld282ceKgDimkRSC8kpaH1dgyc9UMzlt84Y= diff --git a/go/cache/cache.go b/go/cache/cache.go index b6466132452..a801d075fde 100644 --- a/go/cache/cache.go +++ b/go/cache/cache.go @@ -38,6 +38,9 @@ type Cache interface { UsedCapacity() int64 MaxCapacity() int64 SetCapacity(int64) + + // Close shuts down this cache and stops any background goroutines. + Close() } type cachedObject interface { diff --git a/go/cache/lru_cache.go b/go/cache/lru_cache.go index 31ceadaf201..8cc89ac55a4 100644 --- a/go/cache/lru_cache.go +++ b/go/cache/lru_cache.go @@ -250,3 +250,7 @@ func (lru *LRUCache) checkCapacity() { lru.evictions++ } } + +func (lru *LRUCache) Close() { + lru.Clear() +} diff --git a/go/cache/null.go b/go/cache/null.go index c99d52eb2ec..2e1eeeb0d2d 100644 --- a/go/cache/null.go +++ b/go/cache/null.go @@ -71,3 +71,5 @@ func (n *nullCache) SetCapacity(_ int64) {} func (n *nullCache) Evictions() int64 { return 0 } + +func (n *nullCache) Close() {} diff --git a/go/cmd/vtexplain/vtexplain.go b/go/cmd/vtexplain/vtexplain.go index b173a7566c0..68ceed51316 100644 --- a/go/cmd/vtexplain/vtexplain.go +++ b/go/cmd/vtexplain/vtexplain.go @@ -17,6 +17,7 @@ limitations under the License. package main import ( + "context" "fmt" "os" @@ -147,7 +148,7 @@ func parseAndRun() error { Target: dbName, } - vte, err := vtexplain.Init(vschema, schema, ksShardMap, opts) + vte, err := vtexplain.Init(context.Background(), vschema, schema, ksShardMap, opts) if err != nil { return err } diff --git a/go/test/utils/noleak.go b/go/test/utils/noleak.go new file mode 100644 index 00000000000..b24fb585579 --- /dev/null +++ b/go/test/utils/noleak.go @@ -0,0 +1,55 @@ +package utils + +import ( + "os" + "os/exec" + "strconv" + "strings" + "testing" + "time" + + "go.uber.org/goleak" +) + +func EnsureNoLeaks(t testing.TB) { + if t.Failed() { + return + } + ensureNoGoroutines(t) + ensureNoOpenSockets(t) +} + +func ensureNoGoroutines(t testing.TB) { + var ignored = []goleak.Option{ + goleak.IgnoreTopFunction("github.com/golang/glog.(*fileSink).flushDaemon"), + goleak.IgnoreTopFunction("github.com/golang/glog.(*loggingT).flushDaemon"), + goleak.IgnoreTopFunction("vitess.io/vitess/go/vt/dbconfigs.init.0.func1"), + goleak.IgnoreTopFunction("vitess.io/vitess/go/vt/vtgate.resetAggregators"), + goleak.IgnoreTopFunction("vitess.io/vitess/go/vt/vtgate.processQueryInfo"), + goleak.IgnoreTopFunction("github.com/patrickmn/go-cache.(*janitor).Run"), + } + + var err error + for i := 0; i < 5; i++ { + err = goleak.Find(ignored...) + if err == nil { + return + } + time.Sleep(100 * time.Millisecond) + } + t.Fatal(err) +} + +func ensureNoOpenSockets(t testing.TB) { + cmd := exec.Command("lsof", "-a", "-p", strconv.Itoa(os.Getpid()), "-i", "-P", "-V") + cmd.Stderr = nil + lsof, err := cmd.Output() + if err == nil { + t.Errorf("found open sockets:\n%s", lsof) + } else { + if strings.Contains(string(lsof), "no Internet files located") { + return + } + t.Errorf("failed to run `lsof`: %v (%q)", err, lsof) + } +} diff --git a/go/vt/discovery/keyspace_events.go b/go/vt/discovery/keyspace_events.go index ec9988b500c..163f240de8c 100644 --- a/go/vt/discovery/keyspace_events.go +++ b/go/vt/discovery/keyspace_events.go @@ -221,7 +221,7 @@ func (kew *KeyspaceEventWatcher) run(ctx context.Context) { if result == nil { return } - kew.processHealthCheck(result) + kew.processHealthCheck(ctx, result) } } }() @@ -234,7 +234,7 @@ func (kew *KeyspaceEventWatcher) run(ctx context.Context) { return } for _, ks := range keyspaces { - kew.getKeyspaceStatus(ks) + kew.getKeyspaceStatus(ctx, ks) } }() } @@ -544,22 +544,22 @@ func (kss *keyspaceState) onSrvVSchema(vs *vschemapb.SrvVSchema, err error) bool // newKeyspaceState allocates the internal state required to keep track of availability incidents // in this keyspace, and starts up a SrvKeyspace watcher on our topology server which will update // our keyspaceState with any topology changes in real time. -func newKeyspaceState(kew *KeyspaceEventWatcher, cell, keyspace string) *keyspaceState { +func newKeyspaceState(ctx context.Context, kew *KeyspaceEventWatcher, cell, keyspace string) *keyspaceState { log.Infof("created dedicated watcher for keyspace %s/%s", cell, keyspace) kss := &keyspaceState{ kew: kew, keyspace: keyspace, shards: make(map[string]*shardState), } - kew.ts.WatchSrvKeyspace(context.Background(), cell, keyspace, kss.onSrvKeyspace) - kew.ts.WatchSrvVSchema(context.Background(), cell, kss.onSrvVSchema) + kew.ts.WatchSrvKeyspace(ctx, cell, keyspace, kss.onSrvKeyspace) + kew.ts.WatchSrvVSchema(ctx, cell, kss.onSrvVSchema) return kss } // processHealthCheck is the callback that is called by the global HealthCheck stream that was initiated // by this KeyspaceEventWatcher. it redirects the TabletHealth event to the corresponding keyspaceState -func (kew *KeyspaceEventWatcher) processHealthCheck(th *TabletHealth) { - kss := kew.getKeyspaceStatus(th.Target.Keyspace) +func (kew *KeyspaceEventWatcher) processHealthCheck(ctx context.Context, th *TabletHealth) { + kss := kew.getKeyspaceStatus(ctx, th.Target.Keyspace) if kss == nil { return } @@ -569,12 +569,12 @@ func (kew *KeyspaceEventWatcher) processHealthCheck(th *TabletHealth) { // getKeyspaceStatus returns the keyspaceState object for the corresponding keyspace, allocating it // if we've never seen the keyspace before. -func (kew *KeyspaceEventWatcher) getKeyspaceStatus(keyspace string) *keyspaceState { +func (kew *KeyspaceEventWatcher) getKeyspaceStatus(ctx context.Context, keyspace string) *keyspaceState { kew.mu.Lock() defer kew.mu.Unlock() kss := kew.keyspaces[keyspace] if kss == nil { - kss = newKeyspaceState(kew, kew.localCell, keyspace) + kss = newKeyspaceState(ctx, kew, kew.localCell, keyspace) kew.keyspaces[keyspace] = kss } if kss.deleted { @@ -596,11 +596,11 @@ func (kew *KeyspaceEventWatcher) getKeyspaceStatus(keyspace string) *keyspaceSta // This is not a fully accurate heuristic, but it's good enough that we'd want to buffer the // request for the given target under the assumption that the reason why it cannot be completed // right now is transitory. -func (kew *KeyspaceEventWatcher) TargetIsBeingResharded(target *querypb.Target) bool { +func (kew *KeyspaceEventWatcher) TargetIsBeingResharded(ctx context.Context, target *querypb.Target) bool { if target.TabletType != topodatapb.TabletType_PRIMARY { return false } - ks := kew.getKeyspaceStatus(target.Keyspace) + ks := kew.getKeyspaceStatus(ctx, target.Keyspace) if ks == nil { return false } @@ -617,11 +617,11 @@ func (kew *KeyspaceEventWatcher) TargetIsBeingResharded(target *querypb.Target) // to determine that there was a serving primary which now became non serving. This is only possible in a DemotePrimary // RPC which are only called from ERS and PRS. So buffering will stop when these operations succeed. // We return the tablet alias of the primary if it is serving. -func (kew *KeyspaceEventWatcher) PrimaryIsNotServing(target *querypb.Target) (*topodatapb.TabletAlias, bool) { +func (kew *KeyspaceEventWatcher) PrimaryIsNotServing(ctx context.Context, target *querypb.Target) (*topodatapb.TabletAlias, bool) { if target.TabletType != topodatapb.TabletType_PRIMARY { return nil, false } - ks := kew.getKeyspaceStatus(target.Keyspace) + ks := kew.getKeyspaceStatus(ctx, target.Keyspace) if ks == nil { return nil, false } diff --git a/go/vt/topo/memorytopo/memorytopo.go b/go/vt/topo/memorytopo/memorytopo.go index 504f1d4bd39..ea7c8112ceb 100644 --- a/go/vt/topo/memorytopo/memorytopo.go +++ b/go/vt/topo/memorytopo/memorytopo.go @@ -236,14 +236,13 @@ func (n *node) PropagateWatchError(err error) { // NewServerAndFactory returns a new MemoryTopo and the backing factory for all // the cells. It will create one cell for each parameter passed in. It will log.Exit out // in case of a problem. -func NewServerAndFactory(cells ...string) (*topo.Server, *Factory) { +func NewServerAndFactory(ctx context.Context, cells ...string) (*topo.Server, *Factory) { f := &Factory{ cells: make(map[string]*node), generation: uint64(rand.Int63n(1 << 60)), } f.cells[topo.GlobalCell] = f.newDirectory(topo.GlobalCell, nil) - ctx := context.Background() ts, err := topo.NewWithFactory(f, "" /*serverAddress*/, "" /*root*/) if err != nil { log.Exitf("topo.NewWithFactory() failed: %v", err) @@ -258,8 +257,8 @@ func NewServerAndFactory(cells ...string) (*topo.Server, *Factory) { } // NewServer returns the new server -func NewServer(cells ...string) *topo.Server { - server, _ := NewServerAndFactory(cells...) +func NewServer(ctx context.Context, cells ...string) *topo.Server { + server, _ := NewServerAndFactory(ctx, cells...) return server } diff --git a/go/vt/vtadmin/api.go b/go/vt/vtadmin/api.go index de973fd283f..92d11ba18ea 100644 --- a/go/vt/vtadmin/api.go +++ b/go/vt/vtadmin/api.go @@ -2148,7 +2148,7 @@ func (api *API) VTExplain(ctx context.Context, req *vtadminpb.VTExplainRequest) return nil, er.Error() } - vte, err := vtexplain.Init(srvVSchema, schema, shardMap, &vtexplain.Options{ReplicationMode: "ROW"}) + vte, err := vtexplain.Init(ctx, srvVSchema, schema, shardMap, &vtexplain.Options{ReplicationMode: "ROW"}) if err != nil { return nil, fmt.Errorf("error initilaizing vtexplain: %w", err) } diff --git a/go/vt/vtexplain/vtexplain.go b/go/vt/vtexplain/vtexplain.go index 74810dc618f..4b3c4c3bf47 100644 --- a/go/vt/vtexplain/vtexplain.go +++ b/go/vt/vtexplain/vtexplain.go @@ -21,6 +21,7 @@ package vtexplain import ( "bytes" + "context" "fmt" "sort" "strings" @@ -180,7 +181,7 @@ type TabletActions struct { } // Init sets up the fake execution environment -func Init(vSchemaStr, sqlSchema, ksShardMapStr string, opts *Options) (*VTExplain, error) { +func Init(ctx context.Context, 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) @@ -200,7 +201,7 @@ func Init(vSchemaStr, sqlSchema, ksShardMapStr string, opts *Options) (*VTExplai Autocommit: true, }} vte.setGlobalTabletEnv(tabletEnv) - err = vte.initVtgateExecutor(vSchemaStr, ksShardMapStr, opts) + err = vte.initVtgateExecutor(ctx, vSchemaStr, ksShardMapStr, opts) if err != nil { return nil, fmt.Errorf("initVtgateExecutor: %v", err.Error()) } diff --git a/go/vt/vtexplain/vtexplain_test.go b/go/vt/vtexplain/vtexplain_test.go index 30fd289b671..29b2cfdf415 100644 --- a/go/vt/vtexplain/vtexplain_test.go +++ b/go/vt/vtexplain/vtexplain_test.go @@ -17,6 +17,7 @@ limitations under the License. package vtexplain import ( + "context" "encoding/json" "fmt" "os" @@ -48,7 +49,7 @@ type testopts struct { shardmap map[string]map[string]*topo.ShardInfo } -func initTest(mode string, opts *Options, topts *testopts, t *testing.T) *VTExplain { +func initTest(ctx context.Context, mode string, opts *Options, topts *testopts, t *testing.T) *VTExplain { schema, err := os.ReadFile("testdata/test-schema.sql") require.NoError(t, err) @@ -64,7 +65,7 @@ func initTest(mode string, opts *Options, topts *testopts, t *testing.T) *VTExpl } opts.ExecutionMode = mode - vte, err := Init(string(vSchema), string(schema), shardmap, opts) + vte, err := Init(ctx, string(vSchema), string(schema), shardmap, opts) require.NoError(t, err, "vtexplain Init error\n%s", string(schema)) return vte } @@ -85,7 +86,10 @@ func testExplain(testcase string, opts *Options, t *testing.T) { func runTestCase(testcase, mode string, opts *Options, topts *testopts, t *testing.T) { t.Run(testcase, func(t *testing.T) { - vte := initTest(mode, opts, topts, t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + vte := initTest(ctx, mode, opts, topts, t) sqlFile := fmt.Sprintf("testdata/%s-queries.sql", testcase) sql, err := os.ReadFile(sqlFile) @@ -171,7 +175,10 @@ func TestExplain(t *testing.T) { } func TestErrors(t *testing.T) { - vte := initTest(ModeMulti, defaultTestOpts(), &testopts{}, t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + vte := initTest(ctx, ModeMulti, defaultTestOpts(), &testopts{}, t) tests := []struct { SQL string @@ -208,7 +215,10 @@ func TestErrors(t *testing.T) { } func TestJSONOutput(t *testing.T) { - vte := initTest(ModeMulti, defaultTestOpts(), &testopts{}, t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + vte := initTest(ctx, ModeMulti, defaultTestOpts(), &testopts{}, t) sql := "select 1 from user where id = 1" explains, err := vte.Run(sql) require.NoError(t, err, "vtexplain error") @@ -344,6 +354,9 @@ func TestUsingKeyspaceShardMap(t *testing.T) { } func TestInit(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + vschema := `{ "ks1": { "sharded": true, @@ -353,7 +366,7 @@ func TestInit(t *testing.T) { } }` schema := "create table table_missing_primary_vindex (id int primary key)" - _, err := Init(vschema, schema, "", defaultTestOpts()) + _, err := Init(ctx, vschema, schema, "", defaultTestOpts()) require.Error(t, err) require.Contains(t, err.Error(), "missing primary col vindex") } diff --git a/go/vt/vtexplain/vtexplain_vtgate.go b/go/vt/vtexplain/vtexplain_vtgate.go index 682919e5569..310fba4060d 100644 --- a/go/vt/vtexplain/vtexplain_vtgate.go +++ b/go/vt/vtexplain/vtexplain_vtgate.go @@ -50,9 +50,9 @@ import ( vtgatepb "vitess.io/vitess/go/vt/proto/vtgate" ) -func (vte *VTExplain) initVtgateExecutor(vSchemaStr, ksShardMapStr string, opts *Options) error { +func (vte *VTExplain) initVtgateExecutor(ctx context.Context, vSchemaStr, ksShardMapStr string, opts *Options) error { vte.explainTopo = &ExplainTopo{NumShards: opts.NumShards} - vte.explainTopo.TopoServer = memorytopo.NewServer(vtexplainCell) + vte.explainTopo.TopoServer = memorytopo.NewServer(ctx, vtexplainCell) vte.healthCheck = discovery.NewFakeHealthCheck(nil) resolver := vte.newFakeResolver(opts, vte.explainTopo, vtexplainCell) @@ -73,10 +73,8 @@ func (vte *VTExplain) initVtgateExecutor(vSchemaStr, ksShardMapStr string, opts streamSize := 10 var schemaTracker vtgate.SchemaInfo // no schema tracker for these tests - vte.vtgateExecutor = vtgate.NewExecutor(context.Background(), vte.explainTopo, vtexplainCell, resolver, opts.Normalize, false, streamSize, cache.DefaultConfig, schemaTracker, false, opts.PlannerVersion) - queryLogBufferSize := 10 - vtgate.SetQueryLogger(streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize)) + vte.vtgateExecutor = vtgate.NewExecutor(ctx, vte.explainTopo, vtexplainCell, resolver, opts.Normalize, false, streamSize, cache.DefaultConfig, schemaTracker, false, opts.PlannerVersion, streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize)) return nil } diff --git a/go/vt/vtgate/autocommit_test.go b/go/vt/vtgate/autocommit_test.go index 06a0ef46619..965bda4326a 100644 --- a/go/vt/vtgate/autocommit_test.go +++ b/go/vt/vtgate/autocommit_test.go @@ -21,6 +21,7 @@ import ( "testing" "github.com/stretchr/testify/require" + "vitess.io/vitess/go/test/utils" "vitess.io/vitess/go/sqltypes" @@ -35,7 +36,11 @@ import ( // TestAutocommitUpdateSharded: instant-commit. func TestAutocommitUpdateSharded(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() _, err := autocommitExec(executor, "update user set a=2 where id = 1") require.NoError(t, err) @@ -52,7 +57,11 @@ func TestAutocommitUpdateSharded(t *testing.T) { // TestAutocommitUpdateLookup: transaction: select before update. func TestAutocommitUpdateLookup(t *testing.T) { - executor, sbc1, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() sbclookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("b|a", "int64|varbinary"), "2|1", @@ -81,7 +90,11 @@ func TestAutocommitUpdateLookup(t *testing.T) { // TestAutocommitUpdateVindexChange: transaction: select & update before final update. func TestAutocommitUpdateVindexChange(t *testing.T) { - executor, sbc, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() sbc.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("id|name|lastname|name_lastname_keyspace_id_map", "int64|int32|varchar|int64"), "1|1|foo|0", @@ -120,7 +133,11 @@ func TestAutocommitUpdateVindexChange(t *testing.T) { // TestAutocommitDeleteSharded: instant-commit. func TestAutocommitDeleteSharded(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() _, err := autocommitExec(executor, "delete from user_extra where user_id = 1") require.NoError(t, err) @@ -137,7 +154,11 @@ func TestAutocommitDeleteSharded(t *testing.T) { // TestAutocommitDeleteLookup: transaction: select before update. func TestAutocommitDeleteLookup(t *testing.T) { - executor, sbc1, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() sbc1.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("id|name|lastname", "int64|int32|varchar"), "1|1|foo", @@ -179,7 +200,11 @@ func TestAutocommitDeleteLookup(t *testing.T) { // TestAutocommitDeleteIn: instant-commit. func TestAutocommitDeleteIn(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() _, err := autocommitExec(executor, "delete from user_extra where user_id in (1, 2)") require.NoError(t, err) @@ -196,7 +221,11 @@ func TestAutocommitDeleteIn(t *testing.T) { // TestAutocommitDeleteMultiShard: instant-commit. func TestAutocommitDeleteMultiShard(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() _, err := autocommitExec(executor, "delete from user_extra where user_id = user_id + 1") require.NoError(t, err) @@ -216,7 +245,11 @@ func TestAutocommitDeleteMultiShard(t *testing.T) { // TestAutocommitDeleteMultiShardAutoCommit: instant-commit. func TestAutocommitDeleteMultiShardAutoCommit(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() _, err := autocommitExec(executor, "delete /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ from user_extra where user_id = user_id + 1") require.NoError(t, err) @@ -236,7 +269,11 @@ func TestAutocommitDeleteMultiShardAutoCommit(t *testing.T) { // TestAutocommitInsertSharded: instant-commit. func TestAutocommitInsertSharded(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() _, err := autocommitExec(executor, "insert into user_extra(user_id, v) values (1, 2)") require.NoError(t, err) @@ -255,7 +292,11 @@ func TestAutocommitInsertSharded(t *testing.T) { // TestAutocommitInsertLookup: transaction: select before update. func TestAutocommitInsertLookup(t *testing.T) { - executor, sbc1, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() _, err := autocommitExec(executor, "insert into user(id, v, name) values (1, 2, 'myname')") require.NoError(t, err) @@ -282,7 +323,11 @@ func TestAutocommitInsertLookup(t *testing.T) { // TestAutocommitInsertShardAutoCommit: instant-commit. func TestAutocommitInsertMultishardAutoCommit(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() _, err := autocommitExec(executor, "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ into user_extra(user_id, v) values (1, 2), (3, 4)") require.NoError(t, err) @@ -305,7 +350,8 @@ func TestAutocommitInsertMultishardAutoCommit(t *testing.T) { }}) testCommitCount(t, "sbc2", sbc2, 0) - executor, sbc1, sbc2, _ = createExecutorEnv() + executor, sbc1, sbc2, _ = createExecutorEnv(ctx) + defer executor.Close() // Make the first shard fail - the second completes anyway sbc1.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 1 _, err = autocommitExec(executor, "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ into user_extra(user_id, v) values (1, 2), (3, 4)") @@ -326,7 +372,11 @@ func TestAutocommitInsertMultishardAutoCommit(t *testing.T) { } func TestAutocommitInsertMultishard(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() _, err := autocommitExec(executor, "insert into user_extra(user_id, v) values (1, 2), (3, 4)") require.NoError(t, err) @@ -352,7 +402,11 @@ func TestAutocommitInsertMultishard(t *testing.T) { // TestAutocommitInsertAutoinc: instant-commit: sequence fetch is not transactional. func TestAutocommitInsertAutoinc(t *testing.T) { - executor, _, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() _, err := autocommitExec(executor, "insert into main1(id, name) values (null, 'myname')") require.NoError(t, err) @@ -371,7 +425,11 @@ func TestAutocommitInsertAutoinc(t *testing.T) { // TestAutocommitTransactionStarted: no instant-commit. func TestAutocommitTransactionStarted(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() session := &vtgatepb.Session{ TargetString: "@primary", @@ -403,7 +461,11 @@ func TestAutocommitTransactionStarted(t *testing.T) { // TestAutocommitDirectTarget: instant-commit. func TestAutocommitDirectTarget(t *testing.T) { - executor, _, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() session := &vtgatepb.Session{ TargetString: "TestUnsharded/0@primary", @@ -424,7 +486,11 @@ func TestAutocommitDirectTarget(t *testing.T) { // TestAutocommitDirectRangeTarget: no instant-commit. func TestAutocommitDirectRangeTarget(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() session := &vtgatepb.Session{ TargetString: "TestExecutor[-]@primary", diff --git a/go/vt/vtgate/bench_test.go b/go/vt/vtgate/bench_test.go deleted file mode 100644 index 42f96735252..00000000000 --- a/go/vt/vtgate/bench_test.go +++ /dev/null @@ -1,110 +0,0 @@ -/* -Copyright 2019 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 vtgate - -import ( - "bytes" - "context" - "fmt" - "testing" - - topodatapb "vitess.io/vitess/go/vt/proto/topodata" - vtgatepb "vitess.io/vitess/go/vt/proto/vtgate" -) - -// Benchmark run on 6/27/17, with optimized byte-level operations -// using bytes2.Buffer: -// BenchmarkWithNormalizer-4 300 5694660 ns/op - -// Benchmark run on 6/25/17, prior to improvements: -// BenchmarkWithNormalizer-4 100 10629161 ns/op -// BenchmarkWithoutNormalizer-4 200000 8584 ns/op - -var benchQuery string - -func init() { - // benchQuerySize is the approximate size of the query. - benchQuerySize := 1000000 - - // Size of value is 1/10 size of query. Then we add - // 10 such values to the where clause. - baseval := &bytes.Buffer{} - for i := 0; i < benchQuerySize/100; i++ { - // Add an escape character: This will force the upcoming - // tokenizer improvement to still create a copy of the string. - // Then we can see if avoiding the copy will be worth it. - baseval.WriteString("\\'123456789") - } - - buf := &bytes.Buffer{} - buf.WriteString("select a from t1 where v = 1") - for i := 0; i < 10; i++ { - fmt.Fprintf(buf, " and v%d = '%d%s'", i, i, baseval.String()) - } - benchQuery = buf.String() - // fmt.Printf("len: %d\n", len(benchQuery)) -} - -func BenchmarkWithNormalizer(b *testing.B) { - createSandbox(KsTestUnsharded) - hcVTGateTest.Reset() - _ = hcVTGateTest.AddTestTablet("aa", "1.1.1.1", 1001, KsTestUnsharded, "0", topodatapb.TabletType_PRIMARY, true, 1, nil) - saved := rpcVTGate.executor.normalize - rpcVTGate.executor.normalize = true - defer func() { rpcVTGate.executor.normalize = saved }() - - for i := 0; i < b.N; i++ { - _, _, err := rpcVTGate.Execute( - context.Background(), - nil, - &vtgatepb.Session{ - TargetString: "@primary", - Options: executeOptions, - }, - benchQuery, - nil, - ) - if err != nil { - panic(err) - } - } -} - -func BenchmarkWithoutNormalizer(b *testing.B) { - createSandbox(KsTestUnsharded) - hcVTGateTest.Reset() - _ = hcVTGateTest.AddTestTablet("aa", "1.1.1.1", 1001, KsTestUnsharded, "0", topodatapb.TabletType_PRIMARY, true, 1, nil) - saved := rpcVTGate.executor.normalize - rpcVTGate.executor.normalize = false - defer func() { rpcVTGate.executor.normalize = saved }() - - for i := 0; i < b.N; i++ { - _, _, err := rpcVTGate.Execute( - context.Background(), - nil, - &vtgatepb.Session{ - TargetString: "@primary", - Options: executeOptions, - }, - benchQuery, - nil, - ) - if err != nil { - panic(err) - } - } -} diff --git a/go/vt/vtgate/evalengine/integration/comparison_test.go b/go/vt/vtgate/evalengine/integration/comparison_test.go index 9677a39bad0..a28c7f40a96 100644 --- a/go/vt/vtgate/evalengine/integration/comparison_test.go +++ b/go/vt/vtgate/evalengine/integration/comparison_test.go @@ -29,14 +29,14 @@ import ( "github.com/spf13/pflag" - "vitess.io/vitess/go/mysql/format" - "vitess.io/vitess/go/vt/callerid" - "vitess.io/vitess/go/vt/proto/vtrpc" - "vitess.io/vitess/go/mysql" "vitess.io/vitess/go/mysql/collations" + "vitess.io/vitess/go/mysql/format" "vitess.io/vitess/go/sqltypes" + "vitess.io/vitess/go/test/utils" + "vitess.io/vitess/go/vt/callerid" querypb "vitess.io/vitess/go/vt/proto/query" + "vitess.io/vitess/go/vt/proto/vtrpc" "vitess.io/vitess/go/vt/servenv" "vitess.io/vitess/go/vt/vtgate/evalengine" "vitess.io/vitess/go/vt/vtgate/evalengine/testcases" @@ -250,6 +250,7 @@ func initTimezoneData(t *testing.T, conn *mysql.Conn) { } func TestMySQL(t *testing.T) { + defer utils.EnsureNoLeaks(t) var conn = mysqlconn(t) defer conn.Close() diff --git a/go/vt/vtgate/executor.go b/go/vt/vtgate/executor.go index 0adf111a371..013d3e68e16 100644 --- a/go/vt/vtgate/executor.go +++ b/go/vt/vtgate/executor.go @@ -32,6 +32,7 @@ import ( "time" "github.com/spf13/pflag" + "vitess.io/vitess/go/streamlog" "vitess.io/vitess/go/acl" "vitess.io/vitess/go/cache" @@ -118,6 +119,9 @@ type Executor struct { // truncateErrorLen truncates errors sent to client if they are above this value // (0 means do not truncate). truncateErrorLen int + + // queryLogger is passed in for logging from this vtgate executor. + queryLogger *streamlog.StreamLogger[*logstats.LogStats] } var executorOnce sync.Once @@ -138,6 +142,7 @@ func NewExecutor( schemaTracker SchemaInfo, noScatter bool, pv plancontext.PlannerVersion, + queryLogger *streamlog.StreamLogger[*logstats.LogStats], ) *Executor { e := &Executor{ serv: serv, @@ -152,6 +157,7 @@ func NewExecutor( schemaTracker: schemaTracker, allowScatter: !noScatter, pv: pv, + queryLogger: queryLogger, } vschemaacl.Init() @@ -215,7 +221,7 @@ func (e *Executor) Execute(ctx context.Context, mysqlCtx vtgateservice.MySQLConn } logStats.SaveEndTime() - QueryLogger.Send(logStats) + e.queryLogger.Send(logStats) err = vterrors.TruncateError(err, truncateErrorLen) return result, err } @@ -349,7 +355,7 @@ func (e *Executor) StreamExecute( } logStats.SaveEndTime() - QueryLogger.Send(logStats) + e.queryLogger.Send(logStats) return vterrors.TruncateError(err, truncateErrorLen) } @@ -1267,7 +1273,7 @@ func (e *Executor) Prepare(ctx context.Context, method string, safeSession *Safe // it was a no-op record (i.e. didn't issue any queries) if !(logStats.StmtType == "ROLLBACK" && logStats.ShardQueries == 0) { logStats.SaveEndTime() - QueryLogger.Send(logStats) + e.queryLogger.Send(logStats) } return fld, vterrors.TruncateError(err, truncateErrorLen) } @@ -1506,3 +1512,13 @@ func (e *Executor) planPrepareStmt(ctx context.Context, vcursor *vcursorImpl, qu } return plan, stmt, nil } + +func (e *Executor) Close() { + e.scatterConn.Close() + topo, err := e.serv.GetTopoServer() + if err != nil { + panic(err) + } + topo.Close() + e.plans.Close() +} diff --git a/go/vt/vtgate/executor_ddl_test.go b/go/vt/vtgate/executor_ddl_test.go index 7948f1b6208..189e3ebfaea 100644 --- a/go/vt/vtgate/executor_ddl_test.go +++ b/go/vt/vtgate/executor_ddl_test.go @@ -17,16 +17,22 @@ limitations under the License. package vtgate import ( + "context" "fmt" "testing" + "vitess.io/vitess/go/test/utils" vtgatepb "vitess.io/vitess/go/vt/proto/vtgate" "github.com/stretchr/testify/require" ) func TestDDLFlags(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded}) defer func() { enableOnlineDDL = true diff --git a/go/vt/vtgate/executor_dml_test.go b/go/vt/vtgate/executor_dml_test.go index 4ef0a0e5bfc..b376349826b 100644 --- a/go/vt/vtgate/executor_dml_test.go +++ b/go/vt/vtgate/executor_dml_test.go @@ -38,13 +38,17 @@ import ( ) func TestUpdateEqual(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) // Update by primary vindex. - _, err := executorExec(executor, "update user set a=2 where id = 1", nil) + _, err := executorExec(ctx, executor, "update user set a=2 where id = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "update `user` set a = 2 where id = 1", @@ -52,10 +56,10 @@ func TestUpdateEqual(t *testing.T) { }} assertQueries(t, sbc1, wantQueries) assertQueries(t, sbc2, nil) - testQueryLog(t, logChan, "TestExecute", "UPDATE", "update `user` set a = 2 where id = 1", 1) + testQueryLog(t, executor, logChan, "TestExecute", "UPDATE", "update `user` set a = 2 where id = 1", 1) sbc1.Queries = nil - _, err = executorExec(executor, "update user set a=2 where id = 3", nil) + _, err = executorExec(ctx, executor, "update user set a=2 where id = 3", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "update `user` set a = 2 where id = 3", @@ -68,7 +72,7 @@ func TestUpdateEqual(t *testing.T) { sbc1.Queries = nil sbc2.Queries = nil sbclookup.SetResults([]*sqltypes.Result{{}}) - _, err = executorExec(executor, "update music set a=2 where id = 2", nil) + _, err = executorExec(ctx, executor, "update music set a=2 where id = 2", nil) require.NoError(t, err) vars, err := sqltypes.BuildBindVariable([]any{sqltypes.NewInt64(2)}) require.NoError(t, err) @@ -92,7 +96,7 @@ func TestUpdateEqual(t *testing.T) { ), }) - _, err = executorExec(executor, "update user2 set `name`='myname', lastname='mylastname' where id = 1", nil) + _, err = executorExec(ctx, executor, "update user2 set `name`='myname', lastname='mylastname' where id = 1", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{ { @@ -130,10 +134,14 @@ func TestUpdateEqual(t *testing.T) { } func TestUpdateFromSubQuery(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() executor.pv = querypb.ExecuteOptions_Gen4 - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) fields := []*querypb.Field{ {Name: "count(*)", Type: sqltypes.Int64}, @@ -146,7 +154,7 @@ func TestUpdateFromSubQuery(t *testing.T) { }}) // Update by primary vindex, but first execute subquery - _, err := executorExec(executor, "update user set a=(select count(*) from user where id = 3) where id = 1", nil) + _, err := executorExec(ctx, executor, "update user set a=(select count(*) from user where id = 3) where id = 1", nil) require.NoError(t, err) wantQueriesSbc1 := []*querypb.BoundQuery{{ Sql: "update `user` set a = :__sq1 where id = 1", @@ -160,10 +168,13 @@ func TestUpdateFromSubQuery(t *testing.T) { }} assertQueries(t, sbc1, wantQueriesSbc1) assertQueries(t, sbc2, wantQueriesSbc2) - testQueryLog(t, logChan, "TestExecute", "UPDATE", "update `user` set a = (select count(*) from `user` where id = 3) where id = 1", 2) + testQueryLog(t, executor, logChan, "TestExecute", "UPDATE", "update `user` set a = (select count(*) from `user` where id = 3) where id = 1", 2) } func TestUpdateEqualWithNoVerifyAndWriteOnlyLookupUniqueVindexes(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() res := []*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields( "id|wo_lu_col|erl_lu_col|srl_lu_col|nrl_lu_col|nv_lu_col|lu_col|t2_lu_vdx", @@ -171,9 +182,10 @@ func TestUpdateEqualWithNoVerifyAndWriteOnlyLookupUniqueVindexes(t *testing.T) { ), "1|2|2|2|2|2|1|0", )} - executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(executorVSchema, res) + executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, res) + defer executor.Close() - _, err := executorExec(executor, "update t2_lookup set lu_col = 5 where wo_lu_col = 2", nil) + _, err := executorExec(ctx, executor, "update t2_lookup set lu_col = 5 where wo_lu_col = 2", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{ { @@ -215,6 +227,9 @@ func TestUpdateEqualWithNoVerifyAndWriteOnlyLookupUniqueVindexes(t *testing.T) { } func TestUpdateInTransactionLookupDefaultReadLock(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() res := []*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields( "id|wo_lu_col|erl_lu_col|srl_lu_col|nrl_lu_col|nv_lu_col|lu_col|t2_lu_vdx", @@ -222,10 +237,11 @@ func TestUpdateInTransactionLookupDefaultReadLock(t *testing.T) { ), "1|2|2|2|2|2|1|0", )} - executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(executorVSchema, res) + executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, res) + defer executor.Close() safeSession := NewSafeSession(&vtgatepb.Session{InTransaction: true}) - _, err := executorExecSession( + _, err := executorExecSession(ctx, executor, "update t2_lookup set lu_col = 5 where nv_lu_col = 2", nil, @@ -277,6 +293,9 @@ func TestUpdateInTransactionLookupDefaultReadLock(t *testing.T) { } func TestUpdateInTransactionLookupExclusiveReadLock(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() res := []*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields( "id|wo_lu_col|erl_lu_col|srl_lu_col|nrl_lu_col|nv_lu_col|lu_col|t2_lu_vdx", @@ -284,10 +303,11 @@ func TestUpdateInTransactionLookupExclusiveReadLock(t *testing.T) { ), "1|2|2|2|2|2|1|0", )} - executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(executorVSchema, res) + executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, res) + defer executor.Close() safeSession := NewSafeSession(&vtgatepb.Session{InTransaction: true}) - _, err := executorExecSession( + _, err := executorExecSession(ctx, executor, "update t2_lookup set lu_col = 5 where erl_lu_col = 2", nil, @@ -339,6 +359,9 @@ func TestUpdateInTransactionLookupExclusiveReadLock(t *testing.T) { } func TestUpdateInTransactionLookupSharedReadLock(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() res := []*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields( "id|wo_lu_col|erl_lu_col|srl_lu_col|nrl_lu_col|nv_lu_col|lu_col|t2_lu_vdx", @@ -346,10 +369,11 @@ func TestUpdateInTransactionLookupSharedReadLock(t *testing.T) { ), "1|2|2|2|2|2|1|0", )} - executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(executorVSchema, res) + executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, res) + defer executor.Close() safeSession := NewSafeSession(&vtgatepb.Session{InTransaction: true}) - _, err := executorExecSession( + _, err := executorExecSession(ctx, executor, "update t2_lookup set lu_col = 5 where srl_lu_col = 2", nil, @@ -401,6 +425,9 @@ func TestUpdateInTransactionLookupSharedReadLock(t *testing.T) { } func TestUpdateInTransactionLookupNoReadLock(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() res := []*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields( "id|wo_lu_col|erl_lu_col|srl_lu_col|nrl_lu_col|nv_lu_col|lu_col|t2_lu_vdx", @@ -408,10 +435,11 @@ func TestUpdateInTransactionLookupNoReadLock(t *testing.T) { ), "1|2|2|2|2|2|1|0", )} - executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(executorVSchema, res) + executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, res) + defer executor.Close() safeSession := NewSafeSession(&vtgatepb.Session{InTransaction: true}) - _, err := executorExecSession( + _, err := executorExecSession(ctx, executor, "update t2_lookup set lu_col = 5 where nrl_lu_col = 2", nil, @@ -522,7 +550,11 @@ func TestUpdateMultiOwned(t *testing.T) { } } ` - executor, sbc1, sbc2, sbclookup := createCustomExecutor(vschema) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbclookup := createCustomExecutor(ctx, vschema) + defer executor.Close() sbc1.SetResults([]*sqltypes.Result{ sqltypes.MakeTestResult( @@ -530,7 +562,7 @@ func TestUpdateMultiOwned(t *testing.T) { "1|10|20|30|40|50|60|0|0", ), }) - _, err := executorExec(executor, "update user set a=1, b=2, f=4, e=3 where id=1", nil) + _, err := executorExec(ctx, executor, "update user set a=1, b=2, f=4, e=3 where id=1", nil) if err != nil { t.Fatal(err) } @@ -578,9 +610,13 @@ func TestUpdateMultiOwned(t *testing.T) { } func TestUpdateComments(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() - _, err := executorExec(executor, "update user set a=2 where id = 1 /* trailing */", nil) + _, err := executorExec(ctx, executor, "update user set a=2 where id = 1 /* trailing */", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "update `user` set a = 2 where id = 1 /* trailing */", @@ -591,10 +627,14 @@ func TestUpdateComments(t *testing.T) { } func TestUpdateNormalize(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true - _, err := executorExec(executor, "/* leading */ update user set a=2 where id = 1 /* trailing */", nil) + _, err := executorExec(ctx, executor, "/* leading */ update user set a=2 where id = 1 /* trailing */", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "/* leading */ update `user` set a = :a /* INT64 */ where id = :id /* INT64 */ /* trailing */", @@ -609,7 +649,7 @@ func TestUpdateNormalize(t *testing.T) { // Force the query to go to the "wrong" shard and ensure that normalization still happens primarySession.TargetString = "TestExecutor/40-60" - _, err = executorExec(executor, "/* leading */ update user set a=2 where id = 1 /* trailing */", nil) + _, err = executorExec(ctx, executor, "/* leading */ update user set a=2 where id = 1 /* trailing */", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "/* leading */ update `user` set a = :a /* INT64 */ where id = :id /* INT64 */ /* trailing */", @@ -625,7 +665,11 @@ func TestUpdateNormalize(t *testing.T) { } func TestDeleteEqual(t *testing.T) { - executor, sbc, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() sbc.SetResults([]*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -639,7 +683,7 @@ func TestDeleteEqual(t *testing.T) { sqltypes.NewVarChar("myname"), }}, }}) - _, err := executorExec(executor, "delete from user where id = 1", nil) + _, err := executorExec(ctx, executor, "delete from user where id = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select Id, `name` from `user` where id = 1 for update", @@ -662,7 +706,7 @@ func TestDeleteEqual(t *testing.T) { sbc.Queries = nil sbclookup.Queries = nil sbc.SetResults([]*sqltypes.Result{{}}) - _, err = executorExec(executor, "delete from user where id = 1", nil) + _, err = executorExec(ctx, executor, "delete from user where id = 1", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "select Id, `name` from `user` where id = 1 for update", @@ -677,7 +721,7 @@ func TestDeleteEqual(t *testing.T) { sbc.Queries = nil sbclookup.Queries = nil sbclookup.SetResults([]*sqltypes.Result{{}}) - _, err = executorExec(executor, "delete from music where id = 1", nil) + _, err = executorExec(ctx, executor, "delete from music where id = 1", nil) require.NoError(t, err) vars, err := sqltypes.BuildBindVariable([]any{sqltypes.NewInt64(1)}) require.NoError(t, err) @@ -693,7 +737,7 @@ func TestDeleteEqual(t *testing.T) { sbc.Queries = nil sbclookup.Queries = nil sbclookup.SetResults([]*sqltypes.Result{{}}) - _, err = executorExec(executor, "delete from user_extra where user_id = 1", nil) + _, err = executorExec(ctx, executor, "delete from user_extra where user_id = 1", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "delete from user_extra where user_id = 1", @@ -709,7 +753,7 @@ func TestDeleteEqual(t *testing.T) { "1|1|foo", ), }) - _, err = executorExec(executor, "delete from user2 where id = 1", nil) + _, err = executorExec(ctx, executor, "delete from user2 where id = 1", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{ { @@ -738,8 +782,12 @@ func TestDeleteEqual(t *testing.T) { } func TestUpdateScatter(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() - _, err := executorExec(executor, "update user_extra set col = 2", nil) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() + _, err := executorExec(ctx, executor, "update user_extra set col = 2", nil) require.NoError(t, err) // Queries get annotatted. wantQueries := []*querypb.BoundQuery{{ @@ -751,8 +799,12 @@ func TestUpdateScatter(t *testing.T) { } func TestDeleteScatter(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() - _, err := executorExec(executor, "delete from user_extra", nil) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() + _, err := executorExec(ctx, executor, "delete from user_extra", nil) require.NoError(t, err) // Queries get annotatted. wantQueries := []*querypb.BoundQuery{{ @@ -764,7 +816,11 @@ func TestDeleteScatter(t *testing.T) { } func TestUpdateEqualWithMultipleLookupVindex(t *testing.T) { - executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(executorVSchema, nil) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, nil) + defer executor.Close() sbcLookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("lu_col|keyspace_id", "int64|varbinary"), @@ -779,7 +835,7 @@ func TestUpdateEqualWithMultipleLookupVindex(t *testing.T) { "1|2|2|2|2|2|1|0", )}) - _, err := executorExec(executor, "update t2_lookup set lu_col = 5 where wo_lu_col = 2 and lu_col = 1", nil) + _, err := executorExec(ctx, executor, "update t2_lookup set lu_col = 5 where wo_lu_col = 2 and lu_col = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{ { @@ -817,7 +873,11 @@ func TestUpdateEqualWithMultipleLookupVindex(t *testing.T) { } func TestUpdateUseHigherCostVindexIfBackfilling(t *testing.T) { - executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(executorVSchema, nil) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, nil) + defer executor.Close() sbcLookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("lu_col|keyspace_id", "int64|varbinary"), @@ -834,7 +894,7 @@ func TestUpdateUseHigherCostVindexIfBackfilling(t *testing.T) { "1|2|2|2|2|2|2|0", )}) - _, err := executorExec(executor, "update t2_lookup set lu_col = 5 where wo_lu_col = 2 and lu_col in (1, 2)", nil) + _, err := executorExec(ctx, executor, "update t2_lookup set lu_col = 5 where wo_lu_col = 2 and lu_col in (1, 2)", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{ { @@ -892,9 +952,13 @@ func TestDeleteEqualWithNoVerifyAndWriteOnlyLookupUniqueVindex(t *testing.T) { ), "1|1|1|1|1|1|1", )} - executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(executorVSchema, res) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, res) + defer executor.Close() - _, err := executorExec(executor, "delete from t2_lookup where wo_lu_col = 1", nil) + _, err := executorExec(ctx, executor, "delete from t2_lookup where wo_lu_col = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{ { @@ -963,7 +1027,11 @@ func TestDeleteEqualWithNoVerifyAndWriteOnlyLookupUniqueVindex(t *testing.T) { } func TestDeleteEqualWithMultipleLookupVindex(t *testing.T) { - executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(executorVSchema, nil) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, nil) + defer executor.Close() sbcLookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("lu_col|keyspace_id", "int64|varbinary"), @@ -978,7 +1046,7 @@ func TestDeleteEqualWithMultipleLookupVindex(t *testing.T) { "1|1|1|1|1|1|1", )}) - _, err := executorExec(executor, "delete from t2_lookup where wo_lu_col = 1 and lu_col = 1", nil) + _, err := executorExec(ctx, executor, "delete from t2_lookup where wo_lu_col = 1 and lu_col = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{ { @@ -1041,7 +1109,11 @@ func TestDeleteEqualWithMultipleLookupVindex(t *testing.T) { } func TestDeleteUseHigherCostVindexIfBackfilling(t *testing.T) { - executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(executorVSchema, nil) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, nil) + defer executor.Close() sbcLookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("lu_col|keyspace_id", "int64|varbinary"), @@ -1058,7 +1130,7 @@ func TestDeleteUseHigherCostVindexIfBackfilling(t *testing.T) { "1|1|1|1|1|1|2", )}) - _, err := executorExec(executor, "delete from t2_lookup where wo_lu_col = 1 and lu_col in (1, 2)", nil) + _, err := executorExec(ctx, executor, "delete from t2_lookup where wo_lu_col = 1 and lu_col in (1, 2)", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{ { @@ -1158,8 +1230,12 @@ func TestDeleteUseHigherCostVindexIfBackfilling(t *testing.T) { } func TestDeleteByDestination(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() - _, err := executorExec(executor, "delete from `TestExecutor[-]`.user_extra limit 10", nil) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() + _, err := executorExec(ctx, executor, "delete from `TestExecutor[-]`.user_extra limit 10", nil) require.NoError(t, err) // Queries get annotatted. wantQueries := []*querypb.BoundQuery{{ @@ -1171,7 +1247,11 @@ func TestDeleteByDestination(t *testing.T) { } func TestDeleteComments(t *testing.T) { - executor, sbc, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() sbc.SetResults([]*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -1185,7 +1265,7 @@ func TestDeleteComments(t *testing.T) { sqltypes.NewVarChar("myname"), }}, }}) - _, err := executorExec(executor, "delete from user where id = 1 /* trailing */", nil) + _, err := executorExec(ctx, executor, "delete from user where id = 1 /* trailing */", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select Id, `name` from `user` where id = 1 for update /* trailing */", @@ -1207,12 +1287,16 @@ func TestDeleteComments(t *testing.T) { } func TestInsertSharded(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) - _, err := executorExec(executor, "insert into user(id, v, name) values (1, 2, 'myname')", nil) + _, err := executorExec(ctx, executor, "insert into user(id, v, name) values (1, 2, 'myname')", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into `user`(id, v, `name`) values (:_Id_0, 2, :_name_0)", @@ -1233,13 +1317,13 @@ func TestInsertSharded(t *testing.T) { }} assertQueries(t, sbclookup, wantQueries) - testQueryLog(t, logChan, "MarkSavepoint", "SAVEPOINT", "savepoint x", 0) - testQueryLog(t, logChan, "VindexCreate", "INSERT", "insert into name_user_map(`name`, user_id) values (:name_0, :user_id_0)", 1) - testQueryLog(t, logChan, "TestExecute", "INSERT", "insert into `user`(id, v, `name`) values (1, 2, 'myname')", 1) + testQueryLog(t, executor, logChan, "MarkSavepoint", "SAVEPOINT", "savepoint x", 0) + testQueryLog(t, executor, logChan, "VindexCreate", "INSERT", "insert into name_user_map(`name`, user_id) values (:name_0, :user_id_0)", 1) + testQueryLog(t, executor, logChan, "TestExecute", "INSERT", "insert into `user`(id, v, `name`) values (1, 2, 'myname')", 1) sbc1.Queries = nil sbclookup.Queries = nil - _, err = executorExec(executor, "insert into user(id, v, name) values (3, 2, 'myname2')", nil) + _, err = executorExec(ctx, executor, "insert into user(id, v, name) values (3, 2, 'myname2')", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "insert into `user`(id, v, `name`) values (:_Id_0, 2, :_name_0)", @@ -1259,12 +1343,12 @@ func TestInsertSharded(t *testing.T) { }, }} assertQueries(t, sbclookup, wantQueries) - testQueryLog(t, logChan, "MarkSavepoint", "SAVEPOINT", "savepoint x", 2) - testQueryLog(t, logChan, "VindexCreate", "INSERT", "insert into name_user_map(`name`, user_id) values (:name_0, :user_id_0)", 1) - testQueryLog(t, logChan, "TestExecute", "INSERT", "insert into `user`(id, v, `name`) values (3, 2, 'myname2')", 1) + testQueryLog(t, executor, logChan, "MarkSavepoint", "SAVEPOINT", "savepoint x", 2) + testQueryLog(t, executor, logChan, "VindexCreate", "INSERT", "insert into name_user_map(`name`, user_id) values (:name_0, :user_id_0)", 1) + testQueryLog(t, executor, logChan, "TestExecute", "INSERT", "insert into `user`(id, v, `name`) values (3, 2, 'myname2')", 1) sbc1.Queries = nil - _, err = executorExec(executor, "insert into user2(id, name, lastname) values (2, 'myname', 'mylastname')", nil) + _, err = executorExec(ctx, executor, "insert into user2(id, name, lastname) values (2, 'myname', 'mylastname')", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "insert into user2(id, `name`, lastname) values (:_id_0, :_name_0, :_lastname_0)", @@ -1275,16 +1359,16 @@ func TestInsertSharded(t *testing.T) { }, }} assertQueries(t, sbc1, wantQueries) - testQueryLog(t, logChan, "MarkSavepoint", "SAVEPOINT", "savepoint x", 3) - testQueryLog(t, logChan, "VindexCreate", "INSERT", "insert into name_lastname_keyspace_id_map(`name`, lastname, keyspace_id) values (:name_0, :lastname_0, :keyspace_id_0)", 1) - testQueryLog(t, logChan, "TestExecute", "INSERT", "insert into user2(id, `name`, lastname) values (2, 'myname', 'mylastname')", 1) + testQueryLog(t, executor, logChan, "MarkSavepoint", "SAVEPOINT", "savepoint x", 3) + testQueryLog(t, executor, logChan, "VindexCreate", "INSERT", "insert into name_lastname_keyspace_id_map(`name`, lastname, keyspace_id) values (:name_0, :lastname_0, :keyspace_id_0)", 1) + testQueryLog(t, executor, logChan, "TestExecute", "INSERT", "insert into user2(id, `name`, lastname) values (2, 'myname', 'mylastname')", 1) // insert with binary values executor.normalize = true sbc1.Queries = nil sbc2.Queries = nil sbclookup.Queries = nil - _, err = executorExec(executor, "insert into user(id, v, name) values (1, 2, _binary 'myname')", nil) + _, err = executorExec(ctx, executor, "insert into user(id, v, name) values (1, 2, _binary 'myname')", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "insert into `user`(id, v, `name`) values (:_Id_0, :vtg2 /* INT64 */, :_name_0)", @@ -1308,16 +1392,20 @@ func TestInsertSharded(t *testing.T) { }} assertQueries(t, sbclookup, wantQueries) - testQueryLog(t, logChan, "MarkSavepoint", "SAVEPOINT", "savepoint x", 3) - testQueryLog(t, logChan, "VindexCreate", "INSERT", "insert into name_user_map(`name`, user_id) values (:name_0, :user_id_0)", 1) - testQueryLog(t, logChan, "TestExecute", "INSERT", "insert into `user`(id, v, `name`) values (:vtg1 /* INT64 */, :vtg2 /* INT64 */, _binary :vtg3 /* VARCHAR */)", 1) + testQueryLog(t, executor, logChan, "MarkSavepoint", "SAVEPOINT", "savepoint x", 3) + testQueryLog(t, executor, logChan, "VindexCreate", "INSERT", "insert into name_user_map(`name`, user_id) values (:name_0, :user_id_0)", 1) + testQueryLog(t, executor, logChan, "TestExecute", "INSERT", "insert into `user`(id, v, `name`) values (:vtg1 /* INT64 */, :vtg2 /* INT64 */, _binary :vtg3 /* VARCHAR */)", 1) } func TestInsertShardedKeyrange(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() // If a unique vindex returns a keyrange, we fail the insert - _, err := executorExec(executor, "insert into keyrange_table(krcol_unique, krcol) values(1, 1)", nil) + _, err := executorExec(ctx, executor, "insert into keyrange_table(krcol_unique, krcol) values(1, 1)", nil) require.EqualError(t, err, "could not map [INT64(1)] to a unique keyspace id: DestinationKeyRange(-10)") } @@ -1381,9 +1469,13 @@ func TestInsertShardedAutocommitLookup(t *testing.T) { } } ` - executor, sbc1, sbc2, sbclookup := createCustomExecutor(vschema) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbclookup := createCustomExecutor(ctx, vschema) + defer executor.Close() - _, err := executorExecSession(executor, "insert into user(id, v, name, music) values (1, 2, 'myname', 'star')", nil, &vtgatepb.Session{}) + _, err := executorExecSession(ctx, executor, "insert into user(id, v, name, music) values (1, 2, 'myname', 'star')", nil, &vtgatepb.Session{}) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into `user`(id, v, `name`, music) values (:_Id_0, 2, :_name_0, :_music_0)", @@ -1414,7 +1506,11 @@ func TestInsertShardedAutocommitLookup(t *testing.T) { } func TestInsertShardedIgnore(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() // Build the sequence of responses for sbclookup. This should // match the sequence of queries we validate below. @@ -1443,7 +1539,7 @@ func TestInsertShardedIgnore(t *testing.T) { // Fifth row: first shard. // Sixth row: second shard (because 3 hash maps to 40-60). query := "insert ignore into insert_ignore_test(pv, owned, verify) values (1, 1, 1), (2, 2, 2), (3, 3, 1), (4, 4, 4), (5, 5, 1), (6, 6, 3)" - _, err := executorExec(executor, query, nil) + _, err := executorExec(ctx, executor, query, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert ignore into insert_ignore_test(pv, owned, verify) values (:_pv_0, :_owned_0, :_verify_0),(:_pv_4, :_owned_4, :_verify_4)", @@ -1545,7 +1641,7 @@ func TestInsertShardedIgnore(t *testing.T) { {}, }) query = "insert ignore into insert_ignore_test(pv, owned, verify) values (1, 1, 1)" - qr, err := executorExec(executor, query, nil) + qr, err := executorExec(ctx, executor, query, nil) require.NoError(t, err) if !qr.Equal(&sqltypes.Result{}) { t.Errorf("qr: %v, want empty result", qr) @@ -1566,13 +1662,17 @@ func TestInsertShardedIgnore(t *testing.T) { func TestInsertOnDupKey(t *testing.T) { // This test just sanity checks that the statement is getting passed through // correctly. The full set of use cases are covered by TestInsertShardedIgnore. - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() sbclookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("b|a", "int64|varbinary"), "1|1", )}) query := "insert into insert_ignore_test(pv, owned, verify) values (1, 1, 1) on duplicate key update col = 2" - _, err := executorExec(executor, query, nil) + _, err := executorExec(ctx, executor, query, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into insert_ignore_test(pv, owned, verify) values (:_pv_0, :_owned_0, :_verify_0) on duplicate key update col = 2", @@ -1608,7 +1708,11 @@ func TestInsertOnDupKey(t *testing.T) { } func TestAutocommitFail(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() query := "insert into user (id) values (1)" sbc1.MustFailCodes[vtrpcpb.Code_ALREADY_EXISTS] = 1 @@ -1617,7 +1721,7 @@ func TestAutocommitFail(t *testing.T) { defer func() { primarySession.Autocommit = false }() - _, err := executorExec(executor, query, nil) + _, err := executorExec(ctx, executor, query, nil) require.Error(t, err) // make sure we have closed and rolled back any transactions started @@ -1625,9 +1729,13 @@ func TestAutocommitFail(t *testing.T) { } func TestInsertComments(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() - _, err := executorExec(executor, "insert into user(id, v, name) values (1, 2, 'myname') /* trailing */", nil) + _, err := executorExec(ctx, executor, "insert into user(id, v, name) values (1, 2, 'myname') /* trailing */", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into `user`(id, v, `name`) values (:_Id_0, 2, :_name_0) /* trailing */", @@ -1650,7 +1758,11 @@ func TestInsertComments(t *testing.T) { } func TestInsertGeneratorSharded(t *testing.T) { - executor, sbc, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() sbclookup.SetResults([]*sqltypes.Result{{ Rows: [][]sqltypes.Value{{ @@ -1659,7 +1771,7 @@ func TestInsertGeneratorSharded(t *testing.T) { RowsAffected: 1, InsertID: 1, }}) - result, err := executorExec(executor, "insert into user(v, `name`) values (2, 'myname')", nil) + result, err := executorExec(ctx, executor, "insert into user(v, `name`) values (2, 'myname')", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into `user`(v, `name`, id) values (2, :_name_0, :_Id_0)", @@ -1689,7 +1801,11 @@ func TestInsertGeneratorSharded(t *testing.T) { } func TestInsertAutoincSharded(t *testing.T) { - router, sbc, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + router, sbc, _, _ := createExecutorEnv(ctx) + defer router.Close() // Fake a mysql auto-inc response. wantResult := &sqltypes.Result{ @@ -1700,7 +1816,7 @@ func TestInsertAutoincSharded(t *testing.T) { InsertID: 2, } sbc.SetResults([]*sqltypes.Result{wantResult}) - result, err := executorExec(router, "insert into user_extra(user_id) values (2)", nil) + result, err := executorExec(ctx, router, "insert into user_extra(user_id) values (2)", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into user_extra(user_id) values (:_user_id_0)", @@ -1716,8 +1832,12 @@ func TestInsertAutoincSharded(t *testing.T) { } func TestInsertGeneratorUnsharded(t *testing.T) { - executor, _, _, sbclookup := createExecutorEnv() - result, err := executorExec(executor, "insert into main1(id, name) values (null, 'myname')", nil) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() + result, err := executorExec(ctx, executor, "insert into main1(id, name) values (null, 'myname')", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select next :n /* INT64 */ values from user_seq", @@ -1737,10 +1857,14 @@ func TestInsertGeneratorUnsharded(t *testing.T) { } func TestInsertAutoincUnsharded(t *testing.T) { - router, _, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + router, _, _, sbclookup := createExecutorEnv(ctx) + defer router.Close() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := router.queryLogger.Subscribe("Test") + defer router.queryLogger.Unsubscribe(logChan) // Fake a mysql auto-inc response. query := "insert into `simple`(val) values ('val')" @@ -1753,7 +1877,7 @@ func TestInsertAutoincUnsharded(t *testing.T) { } sbclookup.SetResults([]*sqltypes.Result{wantResult}) - result, err := executorExec(router, query, nil) + result, err := executorExec(ctx, router, query, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: query, @@ -1762,13 +1886,17 @@ func TestInsertAutoincUnsharded(t *testing.T) { assertQueries(t, sbclookup, wantQueries) assert.Equal(t, result, wantResult) - testQueryLog(t, logChan, "TestExecute", "INSERT", "insert into `simple`(val) values ('val')", 1) + testQueryLog(t, router, logChan, "TestExecute", "INSERT", "insert into `simple`(val) values ('val')", 1) } func TestInsertLookupOwned(t *testing.T) { - executor, sbc, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() - _, err := executorExec(executor, "insert into music(user_id, id) values (2, 3)", nil) + _, err := executorExec(ctx, executor, "insert into music(user_id, id) values (2, 3)", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into music(user_id, id) values (:_user_id_0, :_id_0)", @@ -1790,7 +1918,11 @@ func TestInsertLookupOwned(t *testing.T) { } func TestInsertLookupOwnedGenerator(t *testing.T) { - executor, sbc, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() sbclookup.SetResults([]*sqltypes.Result{{ Rows: [][]sqltypes.Value{{ @@ -1799,7 +1931,7 @@ func TestInsertLookupOwnedGenerator(t *testing.T) { RowsAffected: 1, InsertID: 1, }}) - result, err := executorExec(executor, "insert into music(user_id) values (2)", nil) + result, err := executorExec(ctx, executor, "insert into music(user_id) values (2)", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into music(user_id, id) values (:_user_id_0, :_id_0)", @@ -1829,9 +1961,13 @@ func TestInsertLookupOwnedGenerator(t *testing.T) { } func TestInsertLookupUnowned(t *testing.T) { - executor, sbc, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() - _, err := executorExec(executor, "insert into music_extra(user_id, music_id) values (2, 3)", nil) + _, err := executorExec(ctx, executor, "insert into music_extra(user_id, music_id) values (2, 3)", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into music_extra(user_id, music_id) values (:_user_id_0, :_music_id_0)", @@ -1852,12 +1988,16 @@ func TestInsertLookupUnowned(t *testing.T) { } func TestInsertLookupUnownedUnsupplied(t *testing.T) { - executor, sbc, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() sbclookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("b|a", "int64|varbinary"), "3|1", )}) - _, err := executorExec(executor, "insert into music_extra_reversed(music_id) values (3)", nil) + _, err := executorExec(ctx, executor, "insert into music_extra_reversed(music_id) values (3)", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into music_extra_reversed(music_id, user_id) values (:_music_id_0, :_user_id_0)", @@ -1881,7 +2021,11 @@ func TestInsertLookupUnownedUnsupplied(t *testing.T) { // If a statement gets broken up into two, and the first one fails, // then an error should be returned normally. func TestInsertPartialFail1(t *testing.T) { - executor, _, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() // Make the first DML fail, there should be no rollback. sbclookup.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 1 @@ -1901,7 +2045,11 @@ func TestInsertPartialFail1(t *testing.T) { // after successful execution of the first, then the transaction must // be rolled back due to partial execution. func TestInsertPartialFail2(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() // Make the second DML fail, it should result in a rollback. sbc1.MustFailExecute[sqlparser.StmtInsert] = 1 @@ -1941,9 +2089,13 @@ func TestInsertPartialFail2(t *testing.T) { } func TestMultiInsertSharded(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() - _, err := executorExec(executor, "insert into user(id, v, name) values (1, 1, 'myname1'),(3, 3, 'myname3')", nil) + _, err := executorExec(ctx, executor, "insert into user(id, v, name) values (1, 1, 'myname1'),(3, 3, 'myname3')", nil) require.NoError(t, err) wantQueries1 := []*querypb.BoundQuery{{ Sql: "insert into `user`(id, v, `name`) values (:_Id_0, 1, :_name_0)", @@ -1985,7 +2137,7 @@ func TestMultiInsertSharded(t *testing.T) { sbc1.Queries = nil sbclookup.Queries = nil sbc2.Queries = nil - _, err = executorExec(executor, "insert into user(id, v, name) values (1, 1, 'myname1'),(2, 2, 'myname2')", nil) + _, err = executorExec(ctx, executor, "insert into user(id, v, name) values (1, 1, 'myname1'),(2, 2, 'myname2')", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into `user`(id, v, `name`) values (:_Id_0, 1, :_name_0),(:_Id_1, 2, :_name_1)", @@ -2016,7 +2168,7 @@ func TestMultiInsertSharded(t *testing.T) { sbc1.Queries = nil sbclookup.Queries = nil sbc2.Queries = nil - _, err = executorExec(executor, "insert into user2(id, `name`, lastname) values (2, 'myname', 'mylastname'), (3, 'myname2', 'mylastname2')", nil) + _, err = executorExec(ctx, executor, "insert into user2(id, `name`, lastname) values (2, 'myname', 'mylastname'), (3, 'myname2', 'mylastname2')", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "insert into user2(id, `name`, lastname) values (:_id_0, :_name_0, :_lastname_0)", @@ -2045,7 +2197,11 @@ func TestMultiInsertSharded(t *testing.T) { } func TestMultiInsertGenerator(t *testing.T) { - executor, sbc, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() sbclookup.SetResults([]*sqltypes.Result{{ Rows: [][]sqltypes.Value{{ @@ -2054,7 +2210,7 @@ func TestMultiInsertGenerator(t *testing.T) { RowsAffected: 1, InsertID: 1, }}) - result, err := executorExec(executor, "insert into music(user_id, `name`) values (:u, 'myname1'),(:u, 'myname2')", map[string]*querypb.BindVariable{"u": sqltypes.Int64BindVariable(2)}) + result, err := executorExec(ctx, executor, "insert into music(user_id, `name`) values (:u, 'myname1'),(:u, 'myname2')", map[string]*querypb.BindVariable{"u": sqltypes.Int64BindVariable(2)}) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into music(user_id, `name`, id) values (:_user_id_0, 'myname1', :_id_0),(:_user_id_1, 'myname2', :_id_1)", @@ -2090,7 +2246,11 @@ func TestMultiInsertGenerator(t *testing.T) { } func TestMultiInsertGeneratorSparse(t *testing.T) { - executor, sbc, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() sbclookup.SetResults([]*sqltypes.Result{{ Rows: [][]sqltypes.Value{{ @@ -2099,7 +2259,7 @@ func TestMultiInsertGeneratorSparse(t *testing.T) { RowsAffected: 1, InsertID: 1, }}) - result, err := executorExec(executor, "insert into music(id, user_id, name) values (NULL, :u, 'myname1'),(2, :u, 'myname2'), (NULL, :u, 'myname3')", map[string]*querypb.BindVariable{"u": sqltypes.Int64BindVariable(2)}) + result, err := executorExec(ctx, executor, "insert into music(id, user_id, name) values (NULL, :u, 'myname1'),(2, :u, 'myname2'), (NULL, :u, 'myname3')", map[string]*querypb.BindVariable{"u": sqltypes.Int64BindVariable(2)}) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into music(id, user_id, `name`) values (:_id_0, :_user_id_0, 'myname1'),(:_id_1, :_user_id_1, 'myname2'),(:_id_2, :_user_id_2, 'myname3')", @@ -2164,10 +2324,14 @@ func TestInsertBadAutoInc(t *testing.T) { } } ` - executor, _, _, _ := createCustomExecutor(vschema) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createCustomExecutor(ctx, vschema) + defer executor.Close() // If auto inc table cannot be found, the table should not be added to vschema. - _, err := executorExec(executor, "insert into bad_auto(v, name) values (1, 'myname')", nil) + _, err := executorExec(ctx, executor, "insert into bad_auto(v, name) values (1, 'myname')", nil) want := "table bad_auto not found" if err == nil || err.Error() != want { t.Errorf("bad auto inc err: %v, want %v", err, want) @@ -2234,10 +2398,14 @@ func TestKeyDestRangeQuery(t *testing.T) { for _, tc := range tests { t.Run(tc.targetString+" - "+tc.inputQuery, func(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() primarySession.TargetString = tc.targetString - _, err := executorExec(executor, tc.inputQuery, nil) + _, err := executorExec(ctx, executor, tc.inputQuery, nil) require.NoError(t, err) if tc.expectedSbc1Query == "" { @@ -2255,9 +2423,13 @@ func TestKeyDestRangeQuery(t *testing.T) { } // it does not work for inserts - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() primarySession.TargetString = "TestExecutor[-]" - _, err := executorExec(executor, insertInput, nil) + _, err := executorExec(ctx, executor, insertInput, nil) require.EqualError(t, err, "VT03023: INSERT not supported when targeting a key range: TestExecutor[-]") @@ -2275,12 +2447,16 @@ func assertQueriesContain(t *testing.T, sql, sbcName string, sbc *sandboxconn.Sa // Prepared statement tests func TestUpdateEqualWithPrepare(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) - _, err := executorPrepare(executor, "update music set a = :a0 where id = :id0", map[string]*querypb.BindVariable{ + _, err := executorPrepare(ctx, executor, "update music set a = :a0 where id = :id0", map[string]*querypb.BindVariable{ "a0": sqltypes.Int64BindVariable(3), "id0": sqltypes.Int64BindVariable(2), }) @@ -2293,12 +2469,16 @@ func TestUpdateEqualWithPrepare(t *testing.T) { assertQueries(t, sbc1, nil) } func TestInsertShardedWithPrepare(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) - _, err := executorPrepare(executor, "insert into user(id, v, name) values (:_Id0, 2, ':_name_0')", map[string]*querypb.BindVariable{ + _, err := executorPrepare(ctx, executor, "insert into user(id, v, name) values (:_Id0, 2, ':_name_0')", map[string]*querypb.BindVariable{ "_Id0": sqltypes.Int64BindVariable(1), "_name_0": sqltypes.BytesBindVariable([]byte("myname")), "__seq0": sqltypes.Int64BindVariable(1), @@ -2314,8 +2494,12 @@ func TestInsertShardedWithPrepare(t *testing.T) { } func TestDeleteEqualWithPrepare(t *testing.T) { - executor, sbc, _, sbclookup := createExecutorEnv() - _, err := executorPrepare(executor, "delete from user where id = :id0", map[string]*querypb.BindVariable{ + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() + _, err := executorPrepare(ctx, executor, "delete from user where id = :id0", map[string]*querypb.BindVariable{ "id0": sqltypes.Int64BindVariable(1), }) require.NoError(t, err) @@ -2328,12 +2512,16 @@ func TestDeleteEqualWithPrepare(t *testing.T) { } func TestUpdateLastInsertID(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true sql := "update user set a = last_insert_id() where id = 1" primarySession.LastInsertId = 43 - _, err := executorExec(executor, sql, map[string]*querypb.BindVariable{}) + _, err := executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{}) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "update `user` set a = :__lastInsertId where id = :id /* INT64 */", @@ -2346,12 +2534,16 @@ func TestUpdateLastInsertID(t *testing.T) { } func TestUpdateReference(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) - _, err := executorExec(executor, "update zip_detail set status = 'CLOSED' where id = 1", nil) + _, err := executorExec(ctx, executor, "update zip_detail set status = 'CLOSED' where id = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "update zip_detail set `status` = 'CLOSED' where id = 1", @@ -2361,11 +2553,11 @@ func TestUpdateReference(t *testing.T) { assertQueries(t, sbc2, nil) assertQueries(t, sbclookup, wantQueries) - testQueryLog(t, logChan, "TestExecute", "UPDATE", "update zip_detail set `status` = 'CLOSED' where id = 1", 1) + testQueryLog(t, executor, logChan, "TestExecute", "UPDATE", "update zip_detail set `status` = 'CLOSED' where id = 1", 1) sbclookup.Queries = nil - _, err = executorExec(executor, "update TestUnsharded.zip_detail set status = 'CLOSED' where id = 1", nil) + _, err = executorExec(ctx, executor, "update TestUnsharded.zip_detail set status = 'CLOSED' where id = 1", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "update zip_detail set `status` = 'CLOSED' where id = 1", @@ -2375,22 +2567,26 @@ func TestUpdateReference(t *testing.T) { assertQueries(t, sbc2, nil) assertQueries(t, sbclookup, wantQueries) - testQueryLog(t, logChan, "TestExecute", "UPDATE", + testQueryLog(t, executor, logChan, "TestExecute", "UPDATE", "update TestUnsharded.zip_detail set `status` = 'CLOSED' where id = 1", 1) sbclookup.Queries = nil - _, err = executorExec(executor, "update TestExecutor.zip_detail set status = 'CLOSED' where id = 1", nil) + _, err = executorExec(ctx, executor, "update TestExecutor.zip_detail set status = 'CLOSED' where id = 1", nil) require.NoError(t, err) // Gen4 planner can redirect the query to correct source for update when reference table is involved. } func TestDeleteLookupOwnedEqual(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() sbc1.SetResults([]*sqltypes.Result{ sqltypes.MakeTestResult(sqltypes.MakeTestFields("uniq_col|keyspace_id", "int64|varbinary"), "1|N±\u0090ɢú\u0016\u009C"), }) - _, err := executorExec(executor, "delete from t1 where unq_col = 1", nil) + _, err := executorExec(ctx, executor, "delete from t1 where unq_col = 1", nil) require.NoError(t, err) tupleBindVar, _ := sqltypes.BuildBindVariable([]int64{1}) sbc1wantQueries := []*querypb.BoundQuery{{ @@ -2412,12 +2608,16 @@ func TestDeleteLookupOwnedEqual(t *testing.T) { } func TestDeleteReference(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) - _, err := executorExec(executor, "delete from zip_detail where id = 1", nil) + _, err := executorExec(ctx, executor, "delete from zip_detail where id = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "delete from zip_detail where id = 1", @@ -2427,11 +2627,11 @@ func TestDeleteReference(t *testing.T) { assertQueries(t, sbc2, nil) assertQueries(t, sbclookup, wantQueries) - testQueryLog(t, logChan, "TestExecute", "DELETE", "delete from zip_detail where id = 1", 1) + testQueryLog(t, executor, logChan, "TestExecute", "DELETE", "delete from zip_detail where id = 1", 1) sbclookup.Queries = nil - _, err = executorExec(executor, "delete from zip_detail where id = 1", nil) + _, err = executorExec(ctx, executor, "delete from zip_detail where id = 1", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "delete from zip_detail where id = 1", @@ -2441,21 +2641,24 @@ func TestDeleteReference(t *testing.T) { assertQueries(t, sbc2, nil) assertQueries(t, sbclookup, wantQueries) - testQueryLog(t, logChan, "TestExecute", "DELETE", "delete from zip_detail where id = 1", 1) + testQueryLog(t, executor, logChan, "TestExecute", "DELETE", "delete from zip_detail where id = 1", 1) sbclookup.Queries = nil - _, err = executorExec(executor, "delete from TestExecutor.zip_detail where id = 1", nil) + _, err = executorExec(ctx, executor, "delete from TestExecutor.zip_detail where id = 1", nil) require.NoError(t, err) // Gen4 planner can redirect the query to correct source for update when reference table is involved. } func TestReservedConnDML(t *testing.T) { - executor, _, _, sbc := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, sbc := createExecutorEnv(ctx) + defer executor.Close() - logChan := QueryLogger.Subscribe("TestReservedConnDML") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("TestReservedConnDML") + defer executor.queryLogger.Unsubscribe(logChan) - ctx := context.Background() session := NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true}) _, err := executor.Execute(ctx, nil, "TestReservedConnDML", session, "use "+KsTestUnsharded, nil) @@ -2503,12 +2706,15 @@ func TestReservedConnDML(t *testing.T) { func TestStreamingDML(t *testing.T) { method := "TestStreamingDML" - executor, _, _, sbc := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, sbc := createExecutorEnv(ctx) + defer executor.Close() - logChan := QueryLogger.Subscribe(method) - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe(method) + defer executor.queryLogger.Unsubscribe(logChan) - ctx := context.Background() session := NewAutocommitSession(&vtgatepb.Session{}) tcases := []struct { @@ -2588,16 +2794,20 @@ func TestStreamingDML(t *testing.T) { } func TestPartialVindexInsertQueryFailure(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) session := NewAutocommitSession(&vtgatepb.Session{}) require.True(t, session.GetAutocommit()) require.False(t, session.InTransaction()) - _, err := executorExecSession(executor, "begin", nil, session.Session) + _, err := executorExecSession(ctx, executor, "begin", nil, session.Session) require.NoError(t, err) require.True(t, session.GetAutocommit()) require.True(t, session.InTransaction()) @@ -2622,7 +2832,7 @@ func TestPartialVindexInsertQueryFailure(t *testing.T) { BindVariables: map[string]*querypb.BindVariable{}, }} - _, err = executorExecSession(executor, "insert into t1(id, unq_col) values (1, 1), (2, 3)", nil, session.Session) + _, err = executorExecSession(ctx, executor, "insert into t1(id, unq_col) values (1, 1), (2, 3)", nil, session.Session) require.Error(t, err) require.Contains(t, err.Error(), "reverted partial DML execution failure") require.True(t, session.GetAutocommit()) @@ -2634,17 +2844,21 @@ func TestPartialVindexInsertQueryFailure(t *testing.T) { wantQ[1].Sql = "insert into t1_lkp_idx(unq_col, keyspace_id) values (:_unq_col_1, :keyspace_id_1)" assertQueriesWithSavepoint(t, sbc2, wantQ) - testQueryLog(t, logChan, "TestExecute", "BEGIN", "begin", 0) - testQueryLog(t, logChan, "MarkSavepoint", "SAVEPOINT", "savepoint x", 0) - testQueryLog(t, logChan, "VindexCreate", "SAVEPOINT_ROLLBACK", "rollback to x", 0) - testQueryLog(t, logChan, "TestExecute", "INSERT", "insert into t1(id, unq_col) values (1, 1), (2, 3)", 0) + testQueryLog(t, executor, logChan, "TestExecute", "BEGIN", "begin", 0) + testQueryLog(t, executor, logChan, "MarkSavepoint", "SAVEPOINT", "savepoint x", 0) + testQueryLog(t, executor, logChan, "VindexCreate", "SAVEPOINT_ROLLBACK", "rollback to x", 0) + testQueryLog(t, executor, logChan, "TestExecute", "INSERT", "insert into t1(id, unq_col) values (1, 1), (2, 3)", 0) } func TestPartialVindexInsertQueryFailureAutoCommit(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) session := NewAutocommitSession(&vtgatepb.Session{}) require.True(t, session.GetAutocommit()) @@ -2664,7 +2878,7 @@ func TestPartialVindexInsertQueryFailureAutoCommit(t *testing.T) { }, }} - _, err := executorExecSession(executor, "insert into t1(id, unq_col) values (1, 1), (2, 3)", nil, session.Session) + _, err := executorExecSession(ctx, executor, "insert into t1(id, unq_col) values (1, 1), (2, 3)", nil, session.Session) require.Error(t, err) assert.Contains(t, err.Error(), "transaction rolled back to reverse changes of partial DML execution") assert.True(t, session.GetAutocommit()) @@ -2676,8 +2890,8 @@ func TestPartialVindexInsertQueryFailureAutoCommit(t *testing.T) { wantQ[0].Sql = "insert into t1_lkp_idx(unq_col, keyspace_id) values (:_unq_col_1, :keyspace_id_1)" assertQueriesWithSavepoint(t, sbc2, wantQ) - testQueryLog(t, logChan, "VindexCreate", "INSERT", "insert into t1_lkp_idx(unq_col, keyspace_id) values (:unq_col_0, :keyspace_id_0), (:unq_col_1, :keyspace_id_1)", 2) - testQueryLog(t, logChan, "TestExecute", "INSERT", "insert into t1(id, unq_col) values (1, 1), (2, 3)", 0) + testQueryLog(t, executor, logChan, "VindexCreate", "INSERT", "insert into t1_lkp_idx(unq_col, keyspace_id) values (:unq_col_0, :keyspace_id_0), (:unq_col_1, :keyspace_id_1)", 2) + testQueryLog(t, executor, logChan, "TestExecute", "INSERT", "insert into t1(id, unq_col) values (1, 1), (2, 3)", 0) } // TestMultiInternalSavepoint shows that the internal savepoint created for rolling back any partial dml changes on a failure is not removed from the savepoint list. @@ -2685,14 +2899,18 @@ func TestPartialVindexInsertQueryFailureAutoCommit(t *testing.T) { // The change for it cannot be done as the executor level and will be made at the VTGate entry point. // Test TestMultiInternalSavepointVtGate shows that it fixes the behaviour. func TestMultiInternalSavepoint(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() session := NewAutocommitSession(&vtgatepb.Session{}) - _, err := executorExecSession(executor, "begin", nil, session.Session) + _, err := executorExecSession(ctx, executor, "begin", nil, session.Session) require.NoError(t, err) // this query goes to multiple shards so internal savepoint will be created. - _, err = executorExecSession(executor, "insert into user_extra(user_id) values (1), (4)", nil, session.Session) + _, err = executorExecSession(ctx, executor, "insert into user_extra(user_id) values (1), (4)", nil, session.Session) require.NoError(t, err) wantQ := []*querypb.BoundQuery{{ @@ -2709,7 +2927,7 @@ func TestMultiInternalSavepoint(t *testing.T) { require.Len(t, sbc2.Queries, 0) sbc1.Queries = nil - _, err = executorExecSession(executor, "insert into user_extra(user_id) values (3), (6)", nil, session.Session) + _, err = executorExecSession(ctx, executor, "insert into user_extra(user_id) values (3), (6)", nil, session.Session) require.NoError(t, err) wantQ = []*querypb.BoundQuery{{ Sql: "savepoint x", @@ -2733,10 +2951,14 @@ func TestMultiInternalSavepoint(t *testing.T) { } func TestInsertSelectFromDual(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() - logChan := QueryLogger.Subscribe("TestInsertSelect") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("TestInsertSelect") + defer executor.queryLogger.Unsubscribe(logChan) session := NewAutocommitSession(&vtgatepb.Session{}) @@ -2780,18 +3002,22 @@ func TestInsertSelectFromDual(t *testing.T) { assertQueries(t, sbc2, nil) assertQueries(t, sbclookup, wantlkpQueries) - testQueryLog(t, logChan, "TestInsertSelect", "SET", wQuery, 0) - testQueryLog(t, logChan, "VindexCreate", "INSERT", "insert into name_user_map(`name`, user_id) values (:name_0, :user_id_0)", 1) - testQueryLog(t, logChan, "TestInsertSelect", "INSERT", "insert into `user`(id, v, `name`) select 1, 2, 'myname' from dual", 2) + testQueryLog(t, executor, logChan, "TestInsertSelect", "SET", wQuery, 0) + testQueryLog(t, executor, logChan, "VindexCreate", "INSERT", "insert into name_user_map(`name`, user_id) values (:name_0, :user_id_0)", 1) + testQueryLog(t, executor, logChan, "TestInsertSelect", "INSERT", "insert into `user`(id, v, `name`) select 1, 2, 'myname' from dual", 2) }) } } func TestInsertSelectFromTable(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() - logChan := QueryLogger.Subscribe("TestInsertSelect") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("TestInsertSelect") + defer executor.queryLogger.Unsubscribe(logChan) session := NewAutocommitSession(&vtgatepb.Session{}) @@ -2842,19 +3068,23 @@ func TestInsertSelectFromTable(t *testing.T) { assertQueries(t, sbc2, wantQueries[:1]) // select scatter query went scatter. assertQueries(t, sbclookup, wantlkpQueries) - testQueryLog(t, logChan, "TestInsertSelect", "SET", wQuery, 0) - testQueryLog(t, logChan, "VindexCreate", "INSERT", "insert into name_user_map(`name`, user_id) values (:name_0, :user_id_0), (:name_1, :user_id_1), (:name_2, :user_id_2), (:name_3, :user_id_3), (:name_4, :user_id_4), (:name_5, :user_id_5), (:name_6, :user_id_6), (:name_7, :user_id_7)", 1) - testQueryLog(t, logChan, "TestInsertSelect", "INSERT", "insert into `user`(id, `name`) select c1, c2 from music", 9) // 8 from select and 1 from insert. + testQueryLog(t, executor, logChan, "TestInsertSelect", "SET", wQuery, 0) + testQueryLog(t, executor, logChan, "VindexCreate", "INSERT", "insert into name_user_map(`name`, user_id) values (:name_0, :user_id_0), (:name_1, :user_id_1), (:name_2, :user_id_2), (:name_3, :user_id_3), (:name_4, :user_id_4), (:name_5, :user_id_5), (:name_6, :user_id_6), (:name_7, :user_id_7)", 1) + testQueryLog(t, executor, logChan, "TestInsertSelect", "INSERT", "insert into `user`(id, `name`) select c1, c2 from music", 9) // 8 from select and 1 from insert. } } func TestInsertReference(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) - _, err := executorExec(executor, "insert into zip_detail(id, status) values (1, 'CLOSED')", nil) + _, err := executorExec(ctx, executor, "insert into zip_detail(id, status) values (1, 'CLOSED')", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into zip_detail(id, `status`) values (1, 'CLOSED')", @@ -2864,11 +3094,11 @@ func TestInsertReference(t *testing.T) { assertQueries(t, sbc2, nil) assertQueries(t, sbclookup, wantQueries) - testQueryLog(t, logChan, "TestExecute", "INSERT", "insert into zip_detail(id, `status`) values (1, 'CLOSED')", 1) + testQueryLog(t, executor, logChan, "TestExecute", "INSERT", "insert into zip_detail(id, `status`) values (1, 'CLOSED')", 1) sbclookup.Queries = nil - _, err = executorExec(executor, "insert into TestUnsharded.zip_detail(id, status) values (1, 'CLOSED')", nil) + _, err = executorExec(ctx, executor, "insert into TestUnsharded.zip_detail(id, status) values (1, 'CLOSED')", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "insert into zip_detail(id, `status`) values (1, 'CLOSED')", @@ -2878,11 +3108,11 @@ func TestInsertReference(t *testing.T) { assertQueries(t, sbc2, nil) assertQueries(t, sbclookup, wantQueries) - testQueryLog(t, logChan, "TestExecute", "INSERT", + testQueryLog(t, executor, logChan, "TestExecute", "INSERT", "insert into TestUnsharded.zip_detail(id, `status`) values (1, 'CLOSED')", 1) sbclookup.Queries = nil - _, err = executorExec(executor, "insert into TestExecutor.zip_detail(id, status) values (1, 'CLOSED')", nil) + _, err = executorExec(ctx, executor, "insert into TestExecutor.zip_detail(id, status) values (1, 'CLOSED')", nil) require.NoError(t, err) // Gen4 planner can redirect the query to correct source for update when reference table is involved. } diff --git a/go/vt/vtgate/executor_framework_test.go b/go/vt/vtgate/executor_framework_test.go index 01e3854a21d..81c7ecde423 100644 --- a/go/vt/vtgate/executor_framework_test.go +++ b/go/vt/vtgate/executor_framework_test.go @@ -125,13 +125,13 @@ func init() { vindexes.Register("keyrange_lookuper_unique", newKeyRangeLookuperUnique) } -func createExecutorEnv() (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn.SandboxConn) { +func createExecutorEnv(ctx context.Context) (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn.SandboxConn) { cell := "aa" hc := discovery.NewFakeHealthCheck(make(chan *discovery.TabletHealth)) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema - serv := newSandboxForCells([]string{cell}) - serv.topoServer.CreateKeyspace(context.Background(), "TestExecutor", &topodatapb.Keyspace{SidecarDbName: sidecar.DefaultName}) + serv := newSandboxForCells(ctx, []string{cell}) + serv.topoServer.CreateKeyspace(ctx, "TestExecutor", &topodatapb.Keyspace{SidecarDbName: sidecar.DefaultName}) // Force a new cache to use for lookups of the sidecar database identifier // in use by each keyspace -- as we want to use a different load function // than the one already created by the vtgate as it uses a different topo. @@ -148,7 +148,7 @@ func createExecutorEnv() (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn if !created { log.Fatal("Failed to [re]create a sidecar database identifier cache!") } - resolver := newTestResolver(hc, serv, cell) + resolver := newTestResolver(ctx, hc, serv, cell) sbc1 = hc.AddTestTablet(cell, "-20", 1, "TestExecutor", "-20", topodatapb.TabletType_PRIMARY, true, 1, nil) sbc2 = hc.AddTestTablet(cell, "40-60", 1, "TestExecutor", "40-60", topodatapb.TabletType_PRIMARY, true, 1, nil) // Create these connections so scatter queries don't fail. @@ -170,7 +170,8 @@ func createExecutorEnv() (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn bad.VSchema = badVSchema getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - executor = NewExecutor(context.Background(), serv, cell, resolver, false, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4) + queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize) + executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) key.AnyShardPicker = DestinationAnyShardPickerFirstShard{} // create a new session each time so that ShardSessions don't get re-used across tests @@ -180,13 +181,13 @@ func createExecutorEnv() (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn return executor, sbc1, sbc2, sbclookup } -func createCustomExecutor(vschema string) (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn.SandboxConn) { +func createCustomExecutor(ctx context.Context, vschema string) (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn.SandboxConn) { cell := "aa" hc := discovery.NewFakeHealthCheck(nil) s := createSandbox(KsTestSharded) s.VSchema = vschema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) sbc1 = hc.AddTestTablet(cell, "-20", 1, "TestExecutor", "-20", topodatapb.TabletType_PRIMARY, true, 1, nil) sbc2 = hc.AddTestTablet(cell, "40-60", 1, "TestExecutor", "40-60", topodatapb.TabletType_PRIMARY, true, 1, nil) @@ -194,7 +195,8 @@ func createCustomExecutor(vschema string) (executor *Executor, sbc1, sbc2, sbclo sbclookup = hc.AddTestTablet(cell, "0", 1, KsTestUnsharded, "0", topodatapb.TabletType_PRIMARY, true, 1, nil) getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - executor = NewExecutor(context.Background(), serv, cell, resolver, false, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4) + queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize) + executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) // create a new session each time so that ShardSessions don't get re-used across tests primarySession = &vtgatepb.Session{ TargetString: "@primary", @@ -202,13 +204,13 @@ func createCustomExecutor(vschema string) (executor *Executor, sbc1, sbc2, sbclo return executor, sbc1, sbc2, sbclookup } -func createCustomExecutorSetValues(vschema string, values []*sqltypes.Result) (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn.SandboxConn) { +func createCustomExecutorSetValues(ctx context.Context, vschema string, values []*sqltypes.Result) (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn.SandboxConn) { cell := "aa" hc := discovery.NewFakeHealthCheck(nil) s := createSandbox(KsTestSharded) s.VSchema = vschema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} sbcs := []*sandboxconn.SandboxConn{} for _, shard := range shards { @@ -223,7 +225,8 @@ func createCustomExecutorSetValues(vschema string, values []*sqltypes.Result) (e sbclookup = hc.AddTestTablet(cell, "0", 1, KsTestUnsharded, "0", topodatapb.TabletType_PRIMARY, true, 1, nil) getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - executor = NewExecutor(context.Background(), serv, cell, resolver, false, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4) + queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize) + executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) // create a new session each time so that ShardSessions don't get re-used across tests primarySession = &vtgatepb.Session{ TargetString: "@primary", @@ -231,9 +234,9 @@ func createCustomExecutorSetValues(vschema string, values []*sqltypes.Result) (e return executor, sbcs[0], sbcs[1], sbclookup } -func executorExecSession(executor *Executor, sql string, bv map[string]*querypb.BindVariable, session *vtgatepb.Session) (*sqltypes.Result, error) { +func executorExecSession(ctx context.Context, executor *Executor, sql string, bv map[string]*querypb.BindVariable, session *vtgatepb.Session) (*sqltypes.Result, error) { return executor.Execute( - context.Background(), + ctx, nil, "TestExecute", NewSafeSession(session), @@ -241,23 +244,23 @@ func executorExecSession(executor *Executor, sql string, bv map[string]*querypb. bv) } -func executorExec(executor *Executor, sql string, bv map[string]*querypb.BindVariable) (*sqltypes.Result, error) { - return executorExecSession(executor, sql, bv, primarySession) +func executorExec(ctx context.Context, executor *Executor, sql string, bv map[string]*querypb.BindVariable) (*sqltypes.Result, error) { + return executorExecSession(ctx, executor, sql, bv, primarySession) } -func executorPrepare(executor *Executor, sql string, bv map[string]*querypb.BindVariable) ([]*querypb.Field, error) { +func executorPrepare(ctx context.Context, executor *Executor, sql string, bv map[string]*querypb.BindVariable) ([]*querypb.Field, error) { return executor.Prepare( - context.Background(), + ctx, "TestExecute", NewSafeSession(primarySession), sql, bv) } -func executorStream(executor *Executor, sql string) (qr *sqltypes.Result, err error) { +func executorStream(ctx context.Context, executor *Executor, sql string) (qr *sqltypes.Result, err error) { results := make(chan *sqltypes.Result, 100) err = executor.StreamExecute( - context.Background(), + ctx, nil, "TestExecuteStream", NewSafeSession(nil), @@ -362,14 +365,14 @@ func getQueryLog(logChan chan *logstats.LogStats) *logstats.LogStats { // is a repeat query. var testPlannedQueries = map[string]bool{} -func testQueryLog(t *testing.T, logChan chan *logstats.LogStats, method, stmtType, sql string, shardQueries int) *logstats.LogStats { +func testQueryLog(t *testing.T, executor *Executor, logChan chan *logstats.LogStats, method, stmtType, sql string, shardQueries int) *logstats.LogStats { t.Helper() logStats := getQueryLog(logChan) require.NotNil(t, logStats) var log bytes.Buffer - streamlog.GetFormatter(QueryLogger)(&log, nil, logStats) + streamlog.GetFormatter(executor.queryLogger)(&log, nil, logStats) fields := strings.Split(log.String(), "\t") // fields[0] is the method @@ -425,8 +428,8 @@ func testQueryLog(t *testing.T, logChan chan *logstats.LogStats, method, stmtTyp return logStats } -func newTestResolver(hc discovery.HealthCheck, serv srvtopo.Server, cell string) *Resolver { - sc := newTestScatterConn(hc, serv, cell) +func newTestResolver(ctx context.Context, hc discovery.HealthCheck, serv srvtopo.Server, cell string) *Resolver { + sc := newTestScatterConn(ctx, hc, serv, cell) srvResolver := srvtopo.NewResolver(serv, sc.gateway, cell) return NewResolver(srvResolver, serv, cell, sc) } diff --git a/go/vt/vtgate/executor_scatter_stats_test.go b/go/vt/vtgate/executor_scatter_stats_test.go index 59eba522bb8..5c655767284 100644 --- a/go/vt/vtgate/executor_scatter_stats_test.go +++ b/go/vt/vtgate/executor_scatter_stats_test.go @@ -22,15 +22,21 @@ import ( "testing" "github.com/stretchr/testify/require" + "vitess.io/vitess/go/test/utils" vtgatepb "vitess.io/vitess/go/vt/proto/vtgate" ) func TestScatterStatsWithNoScatterQuery(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() + session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) - _, err := executor.Execute(context.Background(), nil, "TestExecutorResultsExceeded", session, "select * from main1", nil) + _, err := executor.Execute(ctx, nil, "TestExecutorResultsExceeded", session, "select * from main1", nil) require.NoError(t, err) result, err := executor.gatherScatterStats() @@ -39,10 +45,14 @@ func TestScatterStatsWithNoScatterQuery(t *testing.T) { } func TestScatterStatsWithSingleScatterQuery(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) - _, err := executor.Execute(context.Background(), nil, "TestExecutorResultsExceeded", session, "select * from user", nil) + _, err := executor.Execute(ctx, nil, "TestExecutorResultsExceeded", session, "select * from user", nil) require.NoError(t, err) result, err := executor.gatherScatterStats() @@ -51,20 +61,24 @@ func TestScatterStatsWithSingleScatterQuery(t *testing.T) { } func TestScatterStatsHttpWriting(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) - _, err := executor.Execute(context.Background(), nil, "TestExecutorResultsExceeded", session, "select * from user", nil) + _, err := executor.Execute(ctx, nil, "TestExecutorResultsExceeded", session, "select * from user", nil) require.NoError(t, err) - _, err = executor.Execute(context.Background(), nil, "TestExecutorResultsExceeded", session, "select * from user where Id = 15", nil) + _, err = executor.Execute(ctx, nil, "TestExecutorResultsExceeded", session, "select * from user where Id = 15", nil) require.NoError(t, err) - _, err = executor.Execute(context.Background(), nil, "TestExecutorResultsExceeded", session, "select * from user where Id > 15", nil) + _, err = executor.Execute(ctx, nil, "TestExecutorResultsExceeded", session, "select * from user where Id > 15", nil) require.NoError(t, err) query4 := "select * from user as u1 join user as u2 on u1.Id = u2.Id" - _, err = executor.Execute(context.Background(), nil, "TestExecutorResultsExceeded", session, query4, nil) + _, err = executor.Execute(ctx, nil, "TestExecutorResultsExceeded", session, query4, nil) require.NoError(t, err) executor.plans.Wait() diff --git a/go/vt/vtgate/executor_select_test.go b/go/vt/vtgate/executor_select_test.go index a19e9d3e5b4..4fe2e6a697b 100644 --- a/go/vt/vtgate/executor_select_test.go +++ b/go/vt/vtgate/executor_select_test.go @@ -28,6 +28,8 @@ import ( _flag "vitess.io/vitess/go/internal/flag" "vitess.io/vitess/go/mysql/collations" + "vitess.io/vitess/go/streamlog" + "vitess.io/vitess/go/vt/vtgate/logstats" "vitess.io/vitess/go/vt/sqlparser" @@ -51,7 +53,11 @@ import ( ) func TestSelectNext(t *testing.T) { - executor, _, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() query := "select next :n values from user_seq" bv := map[string]*querypb.BindVariable{"n": sqltypes.Int64BindVariable(2)} @@ -105,7 +111,11 @@ func TestSelectNext(t *testing.T) { } func TestSelectDBA(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() query := "select * from INFORMATION_SCHEMA.foo" _, err := executor.Execute(context.Background(), nil, "TestSelectDBA", @@ -159,7 +169,11 @@ func TestSelectDBA(t *testing.T) { } func TestSystemVariablesMySQLBelow80(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true sqlparser.SetParserVersion("57000") @@ -195,7 +209,11 @@ func TestSystemVariablesMySQLBelow80(t *testing.T) { } func TestSystemVariablesWithSetVarDisabled(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true sqlparser.SetParserVersion("80000") @@ -233,7 +251,11 @@ func TestSystemVariablesWithSetVarDisabled(t *testing.T) { } func TestSetSystemVariablesTx(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true sqlparser.SetParserVersion("80001") @@ -281,7 +303,11 @@ func TestSetSystemVariablesTx(t *testing.T) { } func TestSetSystemVariables(t *testing.T) { - executor, _, _, lookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, lookup := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true sqlparser.SetParserVersion("80001") @@ -394,7 +420,11 @@ func TestSetSystemVariables(t *testing.T) { } func TestSetSystemVariablesWithReservedConnection(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true session := NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true, SystemVariables: map[string]string{}}) @@ -437,7 +467,11 @@ func TestSetSystemVariablesWithReservedConnection(t *testing.T) { } func TestCreateTableValidTimestamp(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor", SystemVariables: map[string]string{"sql_mode": "ALLOW_INVALID_DATES"}}) @@ -456,7 +490,11 @@ func TestCreateTableValidTimestamp(t *testing.T) { } func TestGen4SelectDBA(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 @@ -531,9 +569,13 @@ func TestGen4SelectDBA(t *testing.T) { } func TestUnsharded(t *testing.T) { - executor, _, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() - _, err := executorExec(executor, "select id from music_user_map where id = 1", nil) + _, err := executorExec(ctx, executor, "select id from music_user_map where id = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select id from music_user_map where id = 1", @@ -543,9 +585,13 @@ func TestUnsharded(t *testing.T) { } func TestUnshardedComments(t *testing.T) { - executor, _, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() - _, err := executorExec(executor, "/* leading */ select id from music_user_map where id = 1 /* trailing */", nil) + _, err := executorExec(ctx, executor, "/* leading */ select id from music_user_map where id = 1 /* trailing */", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "/* leading */ select id from music_user_map where id = 1 /* trailing */", @@ -553,7 +599,7 @@ func TestUnshardedComments(t *testing.T) { }} utils.MustMatch(t, wantQueries, sbclookup.Queries) - _, err = executorExec(executor, "update music_user_map set id = 1 /* trailing */", nil) + _, err = executorExec(ctx, executor, "update music_user_map set id = 1 /* trailing */", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "/* leading */ select id from music_user_map where id = 1 /* trailing */", @@ -565,7 +611,7 @@ func TestUnshardedComments(t *testing.T) { assertQueries(t, sbclookup, wantQueries) sbclookup.Queries = nil - _, err = executorExec(executor, "delete from music_user_map /* trailing */", nil) + _, err = executorExec(ctx, executor, "delete from music_user_map /* trailing */", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "delete from music_user_map /* trailing */", @@ -574,7 +620,7 @@ func TestUnshardedComments(t *testing.T) { assertQueries(t, sbclookup, wantQueries) sbclookup.Queries = nil - _, err = executorExec(executor, "insert into music_user_map values (1) /* trailing */", nil) + _, err = executorExec(ctx, executor, "insert into music_user_map values (1) /* trailing */", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "insert into music_user_map values (1) /* trailing */", @@ -584,23 +630,31 @@ func TestUnshardedComments(t *testing.T) { } func TestStreamUnsharded(t *testing.T) { - executor, _, _, _ := createExecutorEnv() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) sql := "select id from music_user_map where id = 1" - result, err := executorStream(executor, sql) + result, err := executorStream(ctx, executor, sql) require.NoError(t, err) wantResult := sandboxconn.StreamRowResult if !result.Equal(wantResult) { diff := cmp.Diff(wantResult, result) t.Errorf("result: %+v, want %+v\ndiff: %s", result, wantResult, diff) } - testQueryLog(t, logChan, "TestExecuteStream", "SELECT", sql, 1) + testQueryLog(t, executor, logChan, "TestExecuteStream", "SELECT", sql, 1) } func TestStreamBuffering(t *testing.T) { - executor, _, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() // This test is similar to TestStreamUnsharded except that it returns a Result > 10 bytes, // such that the splitting of the Result into multiple Result responses gets tested. @@ -652,7 +706,11 @@ func TestStreamBuffering(t *testing.T) { } func TestStreamLimitOffset(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() // This test is similar to TestStreamUnsharded except that it returns a Result > 10 bytes, // such that the splitting of the Result into multiple Result responses gets tested. @@ -727,14 +785,18 @@ func TestStreamLimitOffset(t *testing.T) { } func TestSelectLastInsertId(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() primarySession.LastInsertId = 52 executor.normalize = true - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) sql := "select last_insert_id()" - result, err := executorExec(executor, sql, map[string]*querypb.BindVariable{}) + result, err := executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{}) wantResult := &sqltypes.Result{ Fields: []*querypb.Field{ {Name: "last_insert_id()", Type: sqltypes.Uint64, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NOT_NULL_FLAG | querypb.MySqlFlag_NUM_FLAG | querypb.MySqlFlag_UNSIGNED_FLAG)}, @@ -748,21 +810,25 @@ func TestSelectLastInsertId(t *testing.T) { } func TestSelectSystemVariables(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() primarySession.ReadAfterWrite = &vtgatepb.ReadAfterWrite{ ReadAfterWriteGtid: "a fine gtid", ReadAfterWriteTimeout: 13, SessionTrackGtids: true, } executor.normalize = true - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) sql := "select @@autocommit, @@client_found_rows, @@skip_query_plan_cache, @@enable_system_settings, " + "@@sql_select_limit, @@transaction_mode, @@workload, @@read_after_write_gtid, " + "@@read_after_write_timeout, @@session_track_gtids, @@ddl_strategy, @@migration_context, @@socket, @@query_timeout" - result, err := executorExec(executor, sql, map[string]*querypb.BindVariable{}) + result, err := executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{}) wantResult := &sqltypes.Result{ Fields: []*querypb.Field{ {Name: "@@autocommit", Type: sqltypes.Int64, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NOT_NULL_FLAG | querypb.MySqlFlag_NUM_FLAG)}, @@ -804,10 +870,14 @@ func TestSelectSystemVariables(t *testing.T) { } func TestSelectInitializedVitessAwareVariable(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) primarySession.Autocommit = true primarySession.EnableSystemSettings = true @@ -821,7 +891,7 @@ func TestSelectInitializedVitessAwareVariable(t *testing.T) { sql := "select @@autocommit, @@enable_system_settings, @@query_timeout" - result, err := executorExec(executor, sql, nil) + result, err := executorExec(ctx, executor, sql, nil) wantResult := &sqltypes.Result{ Fields: []*querypb.Field{ {Name: "@@autocommit", Type: sqltypes.Int64, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NOT_NULL_FLAG | querypb.MySqlFlag_NUM_FLAG)}, @@ -839,13 +909,17 @@ func TestSelectInitializedVitessAwareVariable(t *testing.T) { } func TestSelectUserDefinedVariable(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) sql := "select @foo" - result, err := executorExec(executor, sql, map[string]*querypb.BindVariable{}) + result, err := executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{}) require.NoError(t, err) wantResult := &sqltypes.Result{ Fields: []*querypb.Field{ @@ -858,7 +932,7 @@ func TestSelectUserDefinedVariable(t *testing.T) { utils.MustMatch(t, wantResult, result, "Mismatch") primarySession = &vtgatepb.Session{UserDefinedVariables: createMap([]string{"foo"}, []any{"bar"})} - result, err = executorExec(executor, sql, map[string]*querypb.BindVariable{}) + result, err = executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{}) require.NoError(t, err) wantResult = &sqltypes.Result{ Fields: []*querypb.Field{ @@ -872,17 +946,21 @@ func TestSelectUserDefinedVariable(t *testing.T) { } func TestFoundRows(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) // run this extra query so we can assert on the number of rows found - _, err := executorExec(executor, "select 42", map[string]*querypb.BindVariable{}) + _, err := executorExec(ctx, executor, "select 42", map[string]*querypb.BindVariable{}) require.NoError(t, err) sql := "select found_rows()" - result, err := executorExec(executor, sql, map[string]*querypb.BindVariable{}) + result, err := executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{}) wantResult := &sqltypes.Result{ Fields: []*querypb.Field{ {Name: "found_rows()", Type: sqltypes.Int64, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NOT_NULL_FLAG | querypb.MySqlFlag_NUM_FLAG)}, @@ -896,23 +974,27 @@ func TestFoundRows(t *testing.T) { } func TestRowCount(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) - _, err := executorExec(executor, "select 42", map[string]*querypb.BindVariable{}) + _, err := executorExec(ctx, executor, "select 42", map[string]*querypb.BindVariable{}) require.NoError(t, err) - testRowCount(t, executor, -1) + testRowCount(t, ctx, executor, -1) - _, err = executorExec(executor, "delete from user where id in (42, 24)", map[string]*querypb.BindVariable{}) + _, err = executorExec(ctx, executor, "delete from user where id in (42, 24)", map[string]*querypb.BindVariable{}) require.NoError(t, err) - testRowCount(t, executor, 2) + testRowCount(t, ctx, executor, 2) } -func testRowCount(t *testing.T, executor *Executor, wantRowCount int64) { +func testRowCount(t *testing.T, ctx context.Context, executor *Executor, wantRowCount int64) { t.Helper() - result, err := executorExec(executor, "select row_count()", map[string]*querypb.BindVariable{}) + result, err := executorExec(ctx, executor, "select row_count()", map[string]*querypb.BindVariable{}) wantResult := &sqltypes.Result{ Fields: []*querypb.Field{ {Name: "row_count()", Type: sqltypes.Int64, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NOT_NULL_FLAG | querypb.MySqlFlag_NUM_FLAG)}, @@ -926,7 +1008,11 @@ func testRowCount(t *testing.T, executor *Executor, wantRowCount int64) { } func TestSelectLastInsertIdInUnion(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true primarySession.LastInsertId = 52 @@ -942,7 +1028,7 @@ func TestSelectLastInsertIdInUnion(t *testing.T) { sbc1.SetResults(result1) sql := "select last_insert_id() as id union select last_insert_id() as id" - got, err := executorExec(executor, sql, map[string]*querypb.BindVariable{}) + got, err := executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{}) require.NoError(t, err) wantResult := &sqltypes.Result{ Fields: []*querypb.Field{ @@ -956,13 +1042,17 @@ func TestSelectLastInsertIdInUnion(t *testing.T) { } func TestSelectLastInsertIdInWhere(t *testing.T) { - executor, _, _, lookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, lookup := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) sql := "select id from music_user_map where id = last_insert_id()" - _, err := executorExec(executor, sql, map[string]*querypb.BindVariable{}) + _, err := executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{}) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select id from music_user_map where id = :__lastInsertId", @@ -973,7 +1063,11 @@ func TestSelectLastInsertIdInWhere(t *testing.T) { } func TestLastInsertIDInVirtualTable(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true result1 := []*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -987,7 +1081,7 @@ func TestLastInsertIDInVirtualTable(t *testing.T) { }}, }} sbc1.SetResults(result1) - _, err := executorExec(executor, "select * from (select last_insert_id()) as t", nil) + _, err := executorExec(ctx, executor, "select * from (select last_insert_id()) as t", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select t.`last_insert_id()` from (select :__lastInsertId as `last_insert_id()` from dual) as t", @@ -998,14 +1092,18 @@ func TestLastInsertIDInVirtualTable(t *testing.T) { } func TestLastInsertIDInSubQueryExpression(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true primarySession.LastInsertId = 12345 defer func() { // clean up global state primarySession.LastInsertId = 0 }() - rs, err := executorExec(executor, "select (select last_insert_id()) as x", nil) + rs, err := executorExec(ctx, executor, "select (select last_insert_id()) as x", nil) require.NoError(t, err) wantResult := &sqltypes.Result{ Fields: []*querypb.Field{ @@ -1023,7 +1121,11 @@ func TestLastInsertIDInSubQueryExpression(t *testing.T) { } func TestSelectDatabase(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true sql := "select database()" newSession := proto.Clone(primarySession).(*vtgatepb.Session) @@ -1050,9 +1152,13 @@ func TestSelectDatabase(t *testing.T) { } func TestSelectBindvars(t *testing.T) { - executor, sbc1, sbc2, lookup := createExecutorEnv() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, lookup := createExecutorEnv(ctx) + defer executor.Close() + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) lookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("b|a", "varbinary|varbinary"), @@ -1063,7 +1169,7 @@ func TestSelectBindvars(t *testing.T) { )}) sql := "select id from `user` where id = :id" - _, err := executorExec(executor, sql, map[string]*querypb.BindVariable{ + _, err := executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{ "id": sqltypes.Int64BindVariable(1), }) require.NoError(t, err) @@ -1074,11 +1180,11 @@ func TestSelectBindvars(t *testing.T) { utils.MustMatch(t, sbc1.Queries, wantQueries) assert.Empty(t, sbc2.Queries) sbc1.Queries = nil - testQueryLog(t, logChan, "TestExecute", "SELECT", sql, 1) + testQueryLog(t, executor, logChan, "TestExecute", "SELECT", sql, 1) // Test with StringBindVariable sql = "select id from `user` where `name` in (:name1, :name2)" - _, err = executorExec(executor, sql, map[string]*querypb.BindVariable{ + _, err = executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{ "name1": sqltypes.StringBindVariable("foo1"), "name2": sqltypes.StringBindVariable("foo2"), }) @@ -1093,11 +1199,11 @@ func TestSelectBindvars(t *testing.T) { }} utils.MustMatch(t, wantQueries, sbc1.Queries) sbc1.Queries = nil - testQueryLog(t, logChan, "TestExecute", "SELECT", "select id from `user` where `name` in (:name1, :name2)", 3) + testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select id from `user` where `name` in (:name1, :name2)", 3) // Test with BytesBindVariable sql = "select id from `user` where `name` in (:name1, :name2)" - _, err = executorExec(executor, sql, map[string]*querypb.BindVariable{ + _, err = executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{ "name1": sqltypes.BytesBindVariable([]byte("foo1")), "name2": sqltypes.BytesBindVariable([]byte("foo2")), }) @@ -1111,7 +1217,7 @@ func TestSelectBindvars(t *testing.T) { }, }} utils.MustMatch(t, wantQueries, sbc1.Queries) - testQueryLog(t, logChan, "TestExecute", "SELECT", sql, 3) + testQueryLog(t, executor, logChan, "TestExecute", "SELECT", sql, 3) // Test no match in the lookup vindex sbc1.Queries = nil @@ -1126,7 +1232,7 @@ func TestSelectBindvars(t *testing.T) { }}) sql = "select id from user where name = :name" - _, err = executorExec(executor, sql, map[string]*querypb.BindVariable{ + _, err = executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{ "name": sqltypes.StringBindVariable("nonexistent"), }) require.NoError(t, err) @@ -1150,13 +1256,17 @@ func TestSelectBindvars(t *testing.T) { }} utils.MustMatch(t, wantLookupQueries, lookup.Queries) - testQueryLog(t, logChan, "TestExecute", "SELECT", "select id from `user` where `name` = :name", 2) + testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select id from `user` where `name` = :name", 2) } func TestSelectEqual(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() - _, err := executorExec(executor, "select id from user where id = 1", nil) + _, err := executorExec(ctx, executor, "select id from user where id = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select id from `user` where id = 1", @@ -1168,7 +1278,7 @@ func TestSelectEqual(t *testing.T) { } sbc1.Queries = nil - _, err = executorExec(executor, "select id from user where id = 3", nil) + _, err = executorExec(ctx, executor, "select id from user where id = 3", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "select id from `user` where id = 3", @@ -1183,7 +1293,7 @@ func TestSelectEqual(t *testing.T) { } sbc2.Queries = nil - _, err = executorExec(executor, "select id from user where id = '3'", nil) + _, err = executorExec(ctx, executor, "select id from user where id = '3'", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "select id from `user` where id = '3'", @@ -1202,7 +1312,7 @@ func TestSelectEqual(t *testing.T) { sqltypes.MakeTestFields("b|a", "varbinary|varbinary"), "foo|1", )}) - _, err = executorExec(executor, "select id from user where name = 'foo'", nil) + _, err = executorExec(ctx, executor, "select id from user where name = 'foo'", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "select id from `user` where `name` = 'foo'", @@ -1221,10 +1331,14 @@ func TestSelectEqual(t *testing.T) { } func TestSelectINFromOR(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.pv = querypb.ExecuteOptions_Gen4 - _, err := executorExec(executor, "select 1 from user where id = 1 and name = 'apa' or id = 2 and name = 'toto'", nil) + _, err := executorExec(ctx, executor, "select 1 from user where id = 1 and name = 'apa' or id = 2 and name = 'toto'", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select 1 from `user` where id = 1 and `name` = 'apa' or id = 2 and `name` = 'toto'", @@ -1236,9 +1350,13 @@ func TestSelectINFromOR(t *testing.T) { } func TestSelectDual(t *testing.T) { - executor, sbc1, _, lookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, _, lookup := createExecutorEnv(ctx) + defer executor.Close() - _, err := executorExec(executor, "select @@aa.bb from dual", nil) + _, err := executorExec(ctx, executor, "select @@aa.bb from dual", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select @@`aa.bb` from dual", @@ -1246,15 +1364,19 @@ func TestSelectDual(t *testing.T) { }} utils.MustMatch(t, wantQueries, sbc1.Queries) - _, err = executorExec(executor, "select @@aa.bb from TestUnsharded.dual", nil) + _, err = executorExec(ctx, executor, "select @@aa.bb from TestUnsharded.dual", nil) require.NoError(t, err) utils.MustMatch(t, wantQueries, lookup.Queries) } func TestSelectComments(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() - _, err := executorExec(executor, "/* leading */ select id from user where id = 1 /* trailing */", nil) + _, err := executorExec(ctx, executor, "/* leading */ select id from user where id = 1 /* trailing */", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "/* leading */ select id from `user` where id = 1 /* trailing */", @@ -1268,10 +1390,14 @@ func TestSelectComments(t *testing.T) { } func TestSelectNormalize(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true - _, err := executorExec(executor, "/* leading */ select id from user where id = 1 /* trailing */", nil) + _, err := executorExec(ctx, executor, "/* leading */ select id from user where id = 1 /* trailing */", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "/* leading */ select id from `user` where id = :id /* INT64 */ /* trailing */", @@ -1287,7 +1413,7 @@ func TestSelectNormalize(t *testing.T) { // Force the query to go to the "wrong" shard and ensure that normalization still happens primarySession.TargetString = "TestExecutor/40-60" - _, err = executorExec(executor, "/* leading */ select id from user where id = 1 /* trailing */", nil) + _, err = executorExec(ctx, executor, "/* leading */ select id from user where id = 1 /* trailing */", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "/* leading */ select id from `user` where id = :id /* INT64 */ /* trailing */", @@ -1302,9 +1428,13 @@ func TestSelectNormalize(t *testing.T) { } func TestSelectCaseSensitivity(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() - _, err := executorExec(executor, "select Id from user where iD = 1", nil) + _, err := executorExec(ctx, executor, "select Id from user where iD = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select Id from `user` where iD = 1", @@ -1318,10 +1448,14 @@ func TestSelectCaseSensitivity(t *testing.T) { } func TestStreamSelectEqual(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() sql := "select id from user where id = 1" - result, err := executorStream(executor, sql) + result, err := executorStream(ctx, executor, sql) require.NoError(t, err) wantResult := sandboxconn.StreamRowResult if !result.Equal(wantResult) { @@ -1330,9 +1464,13 @@ func TestStreamSelectEqual(t *testing.T) { } func TestSelectKeyRange(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() - _, err := executorExec(executor, "select krcol_unique, krcol from keyrange_table where krcol = 1", nil) + _, err := executorExec(ctx, executor, "select krcol_unique, krcol from keyrange_table where krcol = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select krcol_unique, krcol from keyrange_table where krcol = 1", @@ -1346,9 +1484,13 @@ func TestSelectKeyRange(t *testing.T) { } func TestSelectKeyRangeUnique(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() - _, err := executorExec(executor, "select krcol_unique, krcol from keyrange_table where krcol_unique = 1", nil) + _, err := executorExec(ctx, executor, "select krcol_unique, krcol from keyrange_table where krcol_unique = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select krcol_unique, krcol from keyrange_table where krcol_unique = 1", @@ -1362,10 +1504,14 @@ func TestSelectKeyRangeUnique(t *testing.T) { } func TestSelectIN(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() // Constant in IN clause is just a number, not a bind variable. - _, err := executorExec(executor, "select id from user where id in (1)", nil) + _, err := executorExec(ctx, executor, "select id from user where id in (1)", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select id from `user` where id in (1)", @@ -1380,7 +1526,7 @@ func TestSelectIN(t *testing.T) { // They result in two different queries on two shards. sbc1.Queries = nil sbc2.Queries = nil - _, err = executorExec(executor, "select id from user where id in (1, 3)", nil) + _, err = executorExec(ctx, executor, "select id from user where id in (1, 3)", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "select id from `user` where id in ::__vals", @@ -1401,7 +1547,7 @@ func TestSelectIN(t *testing.T) { // This is using []any for the bind variable list. sbc1.Queries = nil sbc2.Queries = nil - _, err = executorExec(executor, "select id from user where id in ::vals", map[string]*querypb.BindVariable{ + _, err = executorExec(ctx, executor, "select id from user where id in ::vals", map[string]*querypb.BindVariable{ "vals": sqltypes.TestBindVariable([]any{int64(1), int64(3)}), }) require.NoError(t, err) @@ -1429,7 +1575,7 @@ func TestSelectIN(t *testing.T) { sqltypes.MakeTestFields("b|a", "varbinary|varbinary"), "foo|1", )}) - _, err = executorExec(executor, "select id from user where name = 'foo'", nil) + _, err = executorExec(ctx, executor, "select id from user where name = 'foo'", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "select id from `user` where `name` = 'foo'", @@ -1448,10 +1594,14 @@ func TestSelectIN(t *testing.T) { } func TestStreamSelectIN(t *testing.T) { - executor, _, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() sql := "select id from user where id in (1)" - result, err := executorStream(executor, sql) + result, err := executorStream(ctx, executor, sql) require.NoError(t, err) wantResult := sandboxconn.StreamRowResult if !result.Equal(wantResult) { @@ -1459,7 +1609,7 @@ func TestStreamSelectIN(t *testing.T) { } sql = "select id from user where id in (1, 3)" - result, err = executorStream(executor, sql) + result, err = executorStream(ctx, executor, sql) require.NoError(t, err) wantResult = &sqltypes.Result{ Fields: sandboxconn.StreamRowResult.Fields, @@ -1474,7 +1624,7 @@ func TestStreamSelectIN(t *testing.T) { } sql = "select id from user where name = 'foo'" - result, err = executorStream(executor, sql) + result, err = executorStream(ctx, executor, sql) require.NoError(t, err) wantResult = sandboxconn.StreamRowResult if !result.Equal(wantResult) { @@ -1492,31 +1642,37 @@ func TestStreamSelectIN(t *testing.T) { utils.MustMatch(t, wantQueries, sbclookup.Queries) } -func createExecutor(serv *sandboxTopo, cell string, resolver *Resolver) *Executor { - return NewExecutor(context.Background(), serv, cell, resolver, false, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4) +func createExecutor(ctx context.Context, serv *sandboxTopo, cell string, resolver *Resolver) *Executor { + queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize) + return NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) } func TestSelectScatter(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} var conns []*sandboxconn.SandboxConn for _, shard := range shards { sbc := hc.AddTestTablet(cell, shard, 1, "TestExecutor", shard, topodatapb.TabletType_PRIMARY, true, 1, nil) conns = append(conns, sbc) } - executor := createExecutor(serv, cell, resolver) - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + executor := createExecutor(ctx, serv, cell, resolver) + defer executor.Close() + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) sql := "select id from `user`" - _, err := executorExec(executor, sql, nil) + _, err := executorExec(ctx, executor, sql, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select id from `user`", @@ -1525,10 +1681,14 @@ func TestSelectScatter(t *testing.T) { for _, conn := range conns { utils.MustMatch(t, wantQueries, conn.Queries) } - testQueryLog(t, logChan, "TestExecute", "SELECT", sql, 8) + testQueryLog(t, executor, logChan, "TestExecute", "SELECT", sql, 8) } func TestSelectScatterPartial(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Special setup: Don't use createExecutorEnv. primarySession = &vtgatepb.Session{ TargetString: "@primary", @@ -1538,8 +1698,8 @@ func TestSelectScatterPartial(t *testing.T) { s := createSandbox(KsTestSharded) s.VSchema = executorVSchema getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} var conns []*sandboxconn.SandboxConn for _, shard := range shards { @@ -1547,13 +1707,14 @@ func TestSelectScatterPartial(t *testing.T) { conns = append(conns, sbc) } - executor := createExecutor(serv, cell, resolver) - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + executor := createExecutor(ctx, serv, cell, resolver) + defer executor.Close() + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) // Fail 1 of N without the directive fails the whole operation conns[2].MustFailCodes[vtrpcpb.Code_RESOURCE_EXHAUSTED] = 1000 - results, err := executorExec(executor, "select id from `user`", nil) + results, err := executorExec(ctx, executor, "select id from `user`", nil) wantErr := "TestExecutor.40-60.primary" if err == nil || !strings.Contains(err.Error(), wantErr) { t.Errorf("want error %v, got %v", wantErr, err) @@ -1564,15 +1725,15 @@ func TestSelectScatterPartial(t *testing.T) { if results != nil { t.Errorf("want nil results, got %v", results) } - testQueryLog(t, logChan, "TestExecute", "SELECT", "select id from `user`", 8) + testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select id from `user`", 8) // Fail 1 of N with the directive succeeds with 7 rows - results, err = executorExec(executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user", nil) + results, err = executorExec(ctx, executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user", nil) require.NoError(t, err) if results == nil || len(results.Rows) != 7 { t.Errorf("want 7 results, got %v", results) } - testQueryLog(t, logChan, "TestExecute", "SELECT", "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from `user`", 8) + testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from `user`", 8) // When all shards fail, the execution should also fail conns[0].MustFailCodes[vtrpcpb.Code_RESOURCE_EXHAUSTED] = 1000 @@ -1583,23 +1744,27 @@ func TestSelectScatterPartial(t *testing.T) { conns[6].MustFailCodes[vtrpcpb.Code_RESOURCE_EXHAUSTED] = 1000 conns[7].MustFailCodes[vtrpcpb.Code_RESOURCE_EXHAUSTED] = 1000 - _, err = executorExec(executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user", nil) + _, err = executorExec(ctx, executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user", nil) require.Error(t, err) - testQueryLog(t, logChan, "TestExecute", "SELECT", "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from `user`", 8) + testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from `user`", 8) - _, err = executorExec(executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user order by id", nil) + _, err = executorExec(ctx, executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user order by id", nil) require.Error(t, err) } func TestSelectScatterPartialOLAP(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} var conns []*sandboxconn.SandboxConn for _, shard := range shards { @@ -1607,23 +1772,24 @@ func TestSelectScatterPartialOLAP(t *testing.T) { conns = append(conns, sbc) } - executor := createExecutor(serv, cell, resolver) - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + executor := createExecutor(ctx, serv, cell, resolver) + defer executor.Close() + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) // Fail 1 of N without the directive fails the whole operation conns[2].MustFailCodes[vtrpcpb.Code_RESOURCE_EXHAUSTED] = 1000 - results, err := executorStream(executor, "select id from `user`") + results, err := executorStream(ctx, executor, "select id from `user`") assert.EqualError(t, err, "target: TestExecutor.40-60.primary: RESOURCE_EXHAUSTED error") assert.Equal(t, vtrpcpb.Code_RESOURCE_EXHAUSTED, vterrors.Code(err)) assert.Nil(t, results) - testQueryLog(t, logChan, "TestExecuteStream", "SELECT", "select id from `user`", 8) + testQueryLog(t, executor, logChan, "TestExecuteStream", "SELECT", "select id from `user`", 8) // Fail 1 of N with the directive succeeds with 7 rows - results, err = executorStream(executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user") + results, err = executorStream(ctx, executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user") require.NoError(t, err) assert.EqualValues(t, 7, len(results.Rows)) - testQueryLog(t, logChan, "TestExecuteStream", "SELECT", "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from `user`", 8) + testQueryLog(t, executor, logChan, "TestExecuteStream", "SELECT", "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from `user`", 8) // If all shards fail, the operation should also fail conns[0].MustFailCodes[vtrpcpb.Code_RESOURCE_EXHAUSTED] = 1000 @@ -1634,23 +1800,27 @@ func TestSelectScatterPartialOLAP(t *testing.T) { conns[6].MustFailCodes[vtrpcpb.Code_RESOURCE_EXHAUSTED] = 1000 conns[7].MustFailCodes[vtrpcpb.Code_RESOURCE_EXHAUSTED] = 1000 - _, err = executorStream(executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user") + _, err = executorStream(ctx, executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user") require.Error(t, err) - testQueryLog(t, logChan, "TestExecuteStream", "SELECT", "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from `user`", 8) + testQueryLog(t, executor, logChan, "TestExecuteStream", "SELECT", "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from `user`", 8) - _, err = executorStream(executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user order by id") + _, err = executorStream(ctx, executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user order by id") require.Error(t, err) } func TestSelectScatterPartialOLAP2(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} var conns []*sandboxconn.SandboxConn for _, shard := range shards { @@ -1658,9 +1828,10 @@ func TestSelectScatterPartialOLAP2(t *testing.T) { conns = append(conns, sbc) } - executor := createExecutor(serv, cell, resolver) - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + executor := createExecutor(ctx, serv, cell, resolver) + defer executor.Close() + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) // Fail 1 of N without the directive fails the whole operation tablet0 := conns[2].Tablet() @@ -1672,49 +1843,54 @@ func TestSelectScatterPartialOLAP2(t *testing.T) { sbc0Th := ths[0] sbc0Th.Serving = false - results, err := executorStream(executor, "select id from `user`") + results, err := executorStream(ctx, executor, "select id from `user`") require.Error(t, err) assert.Contains(t, err.Error(), `no healthy tablet available for 'keyspace:"TestExecutor" shard:"40-60"`) assert.Equal(t, vtrpcpb.Code_UNAVAILABLE, vterrors.Code(err)) assert.Nil(t, results) - testQueryLog(t, logChan, "TestExecuteStream", "SELECT", "select id from `user`", 8) + testQueryLog(t, executor, logChan, "TestExecuteStream", "SELECT", "select id from `user`", 8) // Fail 1 of N with the directive succeeds with 7 rows - results, err = executorStream(executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user") + results, err = executorStream(ctx, executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user") require.NoError(t, err) assert.EqualValues(t, 7, len(results.Rows)) - testQueryLog(t, logChan, "TestExecuteStream", "SELECT", "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from `user`", 8) + testQueryLog(t, executor, logChan, "TestExecuteStream", "SELECT", "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from `user`", 8) // order by - results, err = executorStream(executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user order by id") + results, err = executorStream(ctx, executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user order by id") require.NoError(t, err) assert.EqualValues(t, 7, len(results.Rows)) - testQueryLog(t, logChan, "TestExecuteStream", "SELECT", "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from `user` order by id asc", 8) + testQueryLog(t, executor, logChan, "TestExecuteStream", "SELECT", "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from `user` order by id asc", 8) // order by and limit - results, err = executorStream(executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user order by id limit 5") + results, err = executorStream(ctx, executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user order by id limit 5") require.NoError(t, err) assert.EqualValues(t, 5, len(results.Rows)) - testQueryLog(t, logChan, "TestExecuteStream", "SELECT", "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from `user` order by id asc limit 5", 8) + testQueryLog(t, executor, logChan, "TestExecuteStream", "SELECT", "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from `user` order by id asc limit 5", 8) } func TestStreamSelectScatter(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} for _, shard := range shards { _ = hc.AddTestTablet(cell, shard, 1, "TestExecutor", shard, topodatapb.TabletType_PRIMARY, true, 1, nil) } - executor := createExecutor(serv, cell, resolver) + executor := createExecutor(ctx, serv, cell, resolver) + defer executor.Close() sql := "select id from `user`" - result, err := executorStream(executor, sql) + result, err := executorStream(ctx, executor, sql) require.NoError(t, err) wantResult := &sqltypes.Result{ Fields: sandboxconn.SingleRowResult.Fields, @@ -1734,14 +1910,18 @@ func TestStreamSelectScatter(t *testing.T) { // TestSelectScatterOrderBy will run an ORDER BY query that will scatter out to 8 shards and return the 8 rows (one per shard) sorted. func TestSelectScatterOrderBy(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} var conns []*sandboxconn.SandboxConn for i, shard := range shards { @@ -1764,10 +1944,11 @@ func TestSelectScatterOrderBy(t *testing.T) { }}) conns = append(conns, sbc) } - executor := createExecutor(serv, cell, resolver) + executor := createExecutor(ctx, serv, cell, resolver) + defer executor.Close() query := "select col1, col2 from user order by col2 desc" - gotResult, err := executorExec(executor, query, nil) + gotResult, err := executorExec(ctx, executor, query, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ @@ -1800,14 +1981,18 @@ func TestSelectScatterOrderBy(t *testing.T) { // TestSelectScatterOrderByVarChar will run an ORDER BY query that will scatter out to 8 shards and return the 8 rows (one per shard) sorted. func TestSelectScatterOrderByVarChar(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} var conns []*sandboxconn.SandboxConn for i, shard := range shards { @@ -1828,10 +2013,11 @@ func TestSelectScatterOrderByVarChar(t *testing.T) { }}) conns = append(conns, sbc) } - executor := createExecutor(serv, cell, resolver) + executor := createExecutor(ctx, serv, cell, resolver) + defer executor.Close() query := "select col1, textcol from user order by textcol desc" - gotResult, err := executorExec(executor, query, nil) + gotResult, err := executorExec(ctx, executor, query, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ @@ -1863,14 +2049,18 @@ func TestSelectScatterOrderByVarChar(t *testing.T) { } func TestStreamSelectScatterOrderBy(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} var conns []*sandboxconn.SandboxConn for i, shard := range shards { @@ -1890,10 +2080,11 @@ func TestStreamSelectScatterOrderBy(t *testing.T) { }}) conns = append(conns, sbc) } - executor := createExecutor(serv, cell, resolver) + executor := createExecutor(ctx, serv, cell, resolver) + defer executor.Close() query := "select id, col from user order by col desc" - gotResult, err := executorStream(executor, query) + gotResult, err := executorStream(ctx, executor, query) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ @@ -1921,14 +2112,18 @@ func TestStreamSelectScatterOrderBy(t *testing.T) { } func TestStreamSelectScatterOrderByVarChar(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} var conns []*sandboxconn.SandboxConn for i, shard := range shards { @@ -1946,10 +2141,11 @@ func TestStreamSelectScatterOrderByVarChar(t *testing.T) { }}) conns = append(conns, sbc) } - executor := createExecutor(serv, cell, resolver) + executor := createExecutor(ctx, serv, cell, resolver) + defer executor.Close() query := "select id, textcol from user order by textcol desc" - gotResult, err := executorStream(executor, query) + gotResult, err := executorStream(ctx, executor, query) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ @@ -1978,14 +2174,18 @@ func TestStreamSelectScatterOrderByVarChar(t *testing.T) { // TestSelectScatterAggregate will run an aggregate query that will scatter out to 8 shards and return 4 aggregated rows. func TestSelectScatterAggregate(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} var conns []*sandboxconn.SandboxConn for i, shard := range shards { @@ -2005,10 +2205,11 @@ func TestSelectScatterAggregate(t *testing.T) { }}) conns = append(conns, sbc) } - executor := createExecutor(serv, cell, resolver) + executor := createExecutor(ctx, serv, cell, resolver) + defer executor.Close() query := "select col, sum(foo) from user group by col" - gotResult, err := executorExec(executor, query, nil) + gotResult, err := executorExec(ctx, executor, query, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ @@ -2037,14 +2238,18 @@ func TestSelectScatterAggregate(t *testing.T) { } func TestStreamSelectScatterAggregate(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} var conns []*sandboxconn.SandboxConn for i, shard := range shards { @@ -2064,10 +2269,11 @@ func TestStreamSelectScatterAggregate(t *testing.T) { }}) conns = append(conns, sbc) } - executor := createExecutor(serv, cell, resolver) + executor := createExecutor(ctx, serv, cell, resolver) + defer executor.Close() query := "select col, sum(foo) from user group by col" - gotResult, err := executorStream(executor, query) + gotResult, err := executorStream(ctx, executor, query) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ @@ -2097,14 +2303,18 @@ func TestStreamSelectScatterAggregate(t *testing.T) { // TestSelectScatterLimit will run a limit query (ordered for consistency) against // a scatter route and verify that the limit primitive works as intended. func TestSelectScatterLimit(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} var conns []*sandboxconn.SandboxConn for i, shard := range shards { @@ -2124,10 +2334,11 @@ func TestSelectScatterLimit(t *testing.T) { }}) conns = append(conns, sbc) } - executor := createExecutor(serv, cell, resolver) + executor := createExecutor(ctx, serv, cell, resolver) + defer executor.Close() query := "select col1, col2 from user order by col2 desc limit 3" - gotResult, err := executorExec(executor, query, nil) + gotResult, err := executorExec(ctx, executor, query, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ @@ -2165,14 +2376,18 @@ func TestSelectScatterLimit(t *testing.T) { // TestStreamSelectScatterLimit will run a streaming limit query (ordered for consistency) against // a scatter route and verify that the limit primitive works as intended. func TestStreamSelectScatterLimit(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} var conns []*sandboxconn.SandboxConn for i, shard := range shards { @@ -2192,10 +2407,11 @@ func TestStreamSelectScatterLimit(t *testing.T) { }}) conns = append(conns, sbc) } - executor := createExecutor(serv, cell, resolver) + executor := createExecutor(ctx, serv, cell, resolver) + defer executor.Close() query := "select col1, col2 from user order by col2 desc limit 3" - gotResult, err := executorStream(executor, query) + gotResult, err := executorStream(ctx, executor, query) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ @@ -2232,12 +2448,17 @@ func TestStreamSelectScatterLimit(t *testing.T) { // TODO(sougou): stream and non-stream testing are very similar. // Could reuse code, func TestSimpleJoin(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) sql := "select u1.id, u2.id from user u1 join user u2 where u1.id = 1 and u2.id = 3" - result, err := executorExec(executor, sql, nil) + result, err := executorExec(ctx, executor, sql, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select u1.id from `user` as u1 where u1.id = 1", @@ -2265,16 +2486,21 @@ func TestSimpleJoin(t *testing.T) { t.Errorf("result: %+v, want %+v", result, wantResult) } - testQueryLog(t, logChan, "TestExecute", "SELECT", "select u1.id, u2.id from `user` as u1 join `user` as u2 where u1.id = 1 and u2.id = 3", 2) + testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select u1.id, u2.id from `user` as u1 join `user` as u2 where u1.id = 1 and u2.id = 3", 2) } func TestJoinComments(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) sql := "select u1.id, u2.id from user u1 join user u2 where u1.id = 1 and u2.id = 3 /* trailing */" - _, err := executorExec(executor, sql, nil) + _, err := executorExec(ctx, executor, sql, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select u1.id from `user` as u1 where u1.id = 1 /* trailing */", @@ -2287,16 +2513,21 @@ func TestJoinComments(t *testing.T) { }} utils.MustMatch(t, wantQueries, sbc2.Queries) - testQueryLog(t, logChan, "TestExecute", "SELECT", "select u1.id, u2.id from `user` as u1 join `user` as u2 where u1.id = 1 and u2.id = 3 /* trailing */", 2) + testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select u1.id, u2.id from `user` as u1 join `user` as u2 where u1.id = 1 and u2.id = 3 /* trailing */", 2) } func TestSimpleJoinStream(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) sql := "select u1.id, u2.id from user u1 join user u2 where u1.id = 1 and u2.id = 3" - result, err := executorStream(executor, sql) + result, err := executorStream(ctx, executor, sql) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select u1.id from `user` as u1 where u1.id = 1", @@ -2325,13 +2556,18 @@ func TestSimpleJoinStream(t *testing.T) { t.Errorf("result: %+v, want %+v", result, wantResult) } - testQueryLog(t, logChan, "TestExecuteStream", "SELECT", "select u1.id, u2.id from `user` as u1 join `user` as u2 where u1.id = 1 and u2.id = 3", 2) + testQueryLog(t, executor, logChan, "TestExecuteStream", "SELECT", "select u1.id, u2.id from `user` as u1 join `user` as u2 where u1.id = 1 and u2.id = 3", 2) } func TestVarJoin(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) result1 := []*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -2346,7 +2582,7 @@ func TestVarJoin(t *testing.T) { }} sbc1.SetResults(result1) sql := "select u1.id, u2.id from user u1 join user u2 on u2.id = u1.col where u1.id = 1" - _, err := executorExec(executor, sql, nil) + _, err := executorExec(ctx, executor, sql, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select u1.id, u1.col from `user` as u1 where u1.id = 1", @@ -2360,13 +2596,18 @@ func TestVarJoin(t *testing.T) { t.Errorf("sbc2.Queries: %s, want %s\n", got, want) } - testQueryLog(t, logChan, "TestExecute", "SELECT", "select u1.id, u2.id from `user` as u1 join `user` as u2 on u2.id = u1.col where u1.id = 1", 2) + testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select u1.id, u2.id from `user` as u1 join `user` as u2 on u2.id = u1.col where u1.id = 1", 2) } func TestVarJoinStream(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) result1 := []*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -2381,7 +2622,7 @@ func TestVarJoinStream(t *testing.T) { }} sbc1.SetResults(result1) sql := "select u1.id, u2.id from user u1 join user u2 on u2.id = u1.col where u1.id = 1" - _, err := executorStream(executor, sql) + _, err := executorStream(ctx, executor, sql) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select u1.id, u1.col from `user` as u1 where u1.id = 1", @@ -2395,13 +2636,18 @@ func TestVarJoinStream(t *testing.T) { t.Errorf("sbc2.Queries: %s, want %s\n", got, want) } - testQueryLog(t, logChan, "TestExecuteStream", "SELECT", "select u1.id, u2.id from `user` as u1 join `user` as u2 on u2.id = u1.col where u1.id = 1", 2) + testQueryLog(t, executor, logChan, "TestExecuteStream", "SELECT", "select u1.id, u2.id from `user` as u1 join `user` as u2 on u2.id = u1.col where u1.id = 1", 2) } func TestLeftJoin(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) result1 := []*sqltypes.Result{{ Fields: []*querypb.Field{ {Name: "id", Type: sqltypes.Int32, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NUM_FLAG)}, @@ -2421,7 +2667,7 @@ func TestLeftJoin(t *testing.T) { sbc1.SetResults(result1) sbc2.SetResults(emptyResult) sql := "select u1.id, u2.id from user u1 left join user u2 on u2.id = u1.col where u1.id = 1" - result, err := executorExec(executor, sql, nil) + result, err := executorExec(ctx, executor, sql, nil) require.NoError(t, err) wantResult := &sqltypes.Result{ Fields: []*querypb.Field{ @@ -2438,11 +2684,16 @@ func TestLeftJoin(t *testing.T) { if !result.Equal(wantResult) { t.Errorf("result: \n%+v, want \n%+v", result, wantResult) } - testQueryLog(t, logChan, "TestExecute", "SELECT", "select u1.id, u2.id from `user` as u1 left join `user` as u2 on u2.id = u1.col where u1.id = 1", 2) + testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select u1.id, u2.id from `user` as u1 left join `user` as u2 on u2.id = u1.col where u1.id = 1", 2) } func TestLeftJoinStream(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() result1 := []*sqltypes.Result{{ Fields: []*querypb.Field{ {Name: "id", Type: sqltypes.Int32, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NUM_FLAG)}, @@ -2461,7 +2712,7 @@ func TestLeftJoinStream(t *testing.T) { }} sbc1.SetResults(result1) sbc2.SetResults(emptyResult) - result, err := executorStream(executor, "select u1.id, u2.id from user u1 left join user u2 on u2.id = u1.col where u1.id = 1") + result, err := executorStream(ctx, executor, "select u1.id, u2.id from user u1 left join user u2 on u2.id = u1.col where u1.id = 1") require.NoError(t, err) wantResult := &sqltypes.Result{ Fields: []*querypb.Field{ @@ -2482,7 +2733,12 @@ func TestLeftJoinStream(t *testing.T) { } func TestEmptyJoin(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() // Empty result requires a field query for the second part of join, // which is sent to shard 0. sbc1.SetResults([]*sqltypes.Result{{ @@ -2496,7 +2752,7 @@ func TestEmptyJoin(t *testing.T) { {Name: "col", Type: sqltypes.Int32, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NUM_FLAG)}, }, }}) - result, err := executorExec(executor, "select u1.id, u2.id from user u1 join user u2 on u2.id = u1.col where u1.id = 1", nil) + result, err := executorExec(ctx, executor, "select u1.id, u2.id from user u1 join user u2 on u2.id = u1.col where u1.id = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select u1.id, u1.col from `user` as u1 where u1.id = 1", @@ -2520,7 +2776,12 @@ func TestEmptyJoin(t *testing.T) { } func TestEmptyJoinStream(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() // Empty result requires a field query for the second part of join, // which is sent to shard 0. sbc1.SetResults([]*sqltypes.Result{{ @@ -2532,7 +2793,7 @@ func TestEmptyJoinStream(t *testing.T) { {Name: "id", Type: sqltypes.Int32, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NUM_FLAG)}, }, }}) - result, err := executorStream(executor, "select u1.id, u2.id from user u1 join user u2 on u2.id = u1.col where u1.id = 1") + result, err := executorStream(ctx, executor, "select u1.id, u2.id from user u1 join user u2 on u2.id = u1.col where u1.id = 1") require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select u1.id, u1.col from `user` as u1 where u1.id = 1", @@ -2556,7 +2817,12 @@ func TestEmptyJoinStream(t *testing.T) { } func TestEmptyJoinRecursive(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() // Make sure it also works recursively. sbc1.SetResults([]*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -2572,7 +2838,7 @@ func TestEmptyJoinRecursive(t *testing.T) { {Name: "id", Type: sqltypes.Int32, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NUM_FLAG)}, }, }}) - result, err := executorExec(executor, "select u1.id, u2.id, u3.id from user u1 join (user u2 join user u3 on u3.id = u2.col) where u1.id = 1", nil) + result, err := executorExec(ctx, executor, "select u1.id, u2.id, u3.id from user u1 join (user u2 join user u3 on u3.id = u2.col) where u1.id = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select u1.id from `user` as u1 where u1.id = 1", @@ -2600,7 +2866,12 @@ func TestEmptyJoinRecursive(t *testing.T) { } func TestEmptyJoinRecursiveStream(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() // Make sure it also works recursively. sbc1.SetResults([]*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -2616,7 +2887,7 @@ func TestEmptyJoinRecursiveStream(t *testing.T) { {Name: "id", Type: sqltypes.Int32, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NUM_FLAG)}, }, }}) - result, err := executorStream(executor, "select u1.id, u2.id, u3.id from user u1 join (user u2 join user u3 on u3.id = u2.col) where u1.id = 1") + result, err := executorStream(ctx, executor, "select u1.id, u2.id, u3.id from user u1 join (user u2 join user u3 on u3.id = u2.col) where u1.id = 1") require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select u1.id from `user` as u1 where u1.id = 1", @@ -2644,7 +2915,12 @@ func TestEmptyJoinRecursiveStream(t *testing.T) { } func TestCrossShardSubquery(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() result1 := []*sqltypes.Result{{ Fields: []*querypb.Field{ {Name: "id", Type: sqltypes.Int32}, @@ -2657,7 +2933,7 @@ func TestCrossShardSubquery(t *testing.T) { }}, }} sbc1.SetResults(result1) - result, err := executorExec(executor, "select id1 from (select u1.id id1, u2.id from user u1 join user u2 on u2.id = u1.col where u1.id = 1) as t", nil) + result, err := executorExec(ctx, executor, "select id1 from (select u1.id id1, u2.id from user u1 join user u2 on u2.id = u1.col where u1.id = 1) as t", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select id1, t.`u1.col` from (select u1.id as id1, u1.col as `u1.col` from `user` as u1 where u1.id = 1) as t", @@ -2678,7 +2954,12 @@ func TestCrossShardSubquery(t *testing.T) { } func TestSubQueryAndQueryWithLimit(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() result1 := []*sqltypes.Result{{ Fields: []*querypb.Field{ {Name: "id", Type: sqltypes.Int32, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NUM_FLAG)}, @@ -2720,7 +3001,12 @@ func TestSubQueryAndQueryWithLimit(t *testing.T) { } func TestCrossShardSubqueryStream(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() result1 := []*sqltypes.Result{{ Fields: []*querypb.Field{ {Name: "id", Type: sqltypes.Int32, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NUM_FLAG)}, @@ -2733,7 +3019,7 @@ func TestCrossShardSubqueryStream(t *testing.T) { }}, }} sbc1.SetResults(result1) - result, err := executorStream(executor, "select id1 from (select u1.id id1, u2.id from user u1 join user u2 on u2.id = u1.col where u1.id = 1) as t") + result, err := executorStream(ctx, executor, "select id1 from (select u1.id id1, u2.id from user u1 join user u2 on u2.id = u1.col where u1.id = 1) as t") require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select id1, t.`u1.col` from (select u1.id as id1, u1.col as `u1.col` from `user` as u1 where u1.id = 1) as t", @@ -2760,7 +3046,12 @@ func TestCrossShardSubqueryStream(t *testing.T) { } func TestCrossShardSubqueryGetFields(t *testing.T) { - executor, sbc1, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() sbclookup.SetResults([]*sqltypes.Result{{ Fields: []*querypb.Field{ {Name: "col", Type: sqltypes.Int32, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NUM_FLAG)}, @@ -2773,7 +3064,7 @@ func TestCrossShardSubqueryGetFields(t *testing.T) { }, }} sbc1.SetResults(result1) - result, err := executorExec(executor, "select main1.col, t.id1 from main1 join (select u1.id id1, u2.id from user u1 join user u2 on u2.id = u1.col where u1.id = 1) as t", nil) + result, err := executorExec(ctx, executor, "select main1.col, t.id1 from main1 join (select u1.id id1, u2.id from user u1 join user u2 on u2.id = u1.col where u1.id = 1) as t", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select t.id1, t.`u1.col` from (select u1.id as id1, u1.col as `u1.col` from `user` as u1 where 1 != 1) as t where 1 != 1", @@ -2798,12 +3089,17 @@ func TestCrossShardSubqueryGetFields(t *testing.T) { } func TestSelectBindvarswithPrepare(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) sql := "select id from `user` where id = :id" - _, err := executorPrepare(executor, sql, map[string]*querypb.BindVariable{ + _, err := executorPrepare(ctx, executor, sql, map[string]*querypb.BindVariable{ "id": sqltypes.Int64BindVariable(1), }) require.NoError(t, err) @@ -2819,18 +3115,28 @@ func TestSelectBindvarswithPrepare(t *testing.T) { } func TestSelectDatabasePrepare(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) sql := "select database()" - _, err := executorPrepare(executor, sql, map[string]*querypb.BindVariable{}) + _, err := executorPrepare(ctx, executor, sql, map[string]*querypb.BindVariable{}) require.NoError(t, err) } func TestSelectWithUnionAll(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true sql := "select id from user where id in (1, 2, 3) union all select id from user where id in (1, 2, 3)" bv, _ := sqltypes.BuildBindVariable([]int64{1, 2, 3}) @@ -2866,7 +3172,7 @@ func TestSelectWithUnionAll(t *testing.T) { "vtg2": bv, }, }} - _, err := executorExec(executor, sql, map[string]*querypb.BindVariable{}) + _, err := executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{}) require.NoError(t, err) utils.MustMatch(t, sbc1WantQueries, sbc1.Queries, "sbc1") utils.MustMatch(t, sbc2WantQueries, sbc2.Queries, "sbc2") @@ -2875,14 +3181,19 @@ func TestSelectWithUnionAll(t *testing.T) { sbc1.Queries = nil sbc2.Queries = nil - _, err = executorStream(executor, sql) + _, err = executorStream(ctx, executor, sql) require.NoError(t, err) utils.MustMatch(t, sbc1WantQueries, sbc1.Queries, "sbc1") utils.MustMatch(t, sbc2WantQueries, sbc2.Queries, "sbc2") } func TestSelectLock(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(nil) session.Session.InTransaction = true session.ShardSessions = []*vtgatepb.Session_ShardSession{{ @@ -2941,6 +3252,10 @@ func TestSelectLock(t *testing.T) { } func TestLockReserve(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // no connection should be reserved for these queries. tcases := []string{ "select is_free_lock('lock name') from dual", @@ -2949,7 +3264,8 @@ func TestLockReserve(t *testing.T) { "select release_lock('lock name') from dual", } - executor, _, _, _ := createExecutorEnv() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() session := NewAutocommitSession(&vtgatepb.Session{}) for _, sql := range tcases { @@ -2968,7 +3284,12 @@ func TestLockReserve(t *testing.T) { } func TestSelectFromInformationSchema(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(nil) // check failure when trying to query two keyspaces @@ -2993,14 +3314,18 @@ func TestSelectFromInformationSchema(t *testing.T) { } func TestStreamOrderByLimitWithMultipleResults(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} count := 1 for _, shard := range shards { @@ -3012,13 +3337,15 @@ func TestStreamOrderByLimitWithMultipleResults(t *testing.T) { count++ } - executor := NewExecutor(context.Background(), serv, cell, resolver, true, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4) + queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize) + executor := NewExecutor(ctx, serv, cell, resolver, true, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) + defer executor.Close() // some sleep for all goroutines to start time.Sleep(100 * time.Millisecond) before := runtime.NumGoroutine() query := "select id, col from user order by id limit 2" - gotResult, err := executorStream(executor, query) + gotResult, err := executorStream(ctx, executor, query) require.NoError(t, err) wantResult := sqltypes.MakeTestResult(sqltypes.MakeTestFields("id|col", "int32|int32"), "1|1", "2|2") @@ -3029,14 +3356,18 @@ func TestStreamOrderByLimitWithMultipleResults(t *testing.T) { } func TestSelectScatterFails(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + sess := &vtgatepb.Session{} cell := "aa" hc := discovery.NewFakeHealthCheck(nil) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} for i, shard := range shards { @@ -3056,35 +3387,41 @@ func TestSelectScatterFails(t *testing.T) { }}) } - executor := createExecutor(serv, cell, resolver) + executor := createExecutor(ctx, serv, cell, resolver) + defer executor.Close() executor.allowScatter = false - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) - _, err := executorExecSession(executor, "select id from `user`", nil, sess) + _, err := executorExecSession(ctx, executor, "select id from `user`", nil, sess) require.Error(t, err) assert.Contains(t, err.Error(), "scatter") // Run the test again, to ensure it behaves the same for a cached query - _, err = executorExecSession(executor, "select id from `user`", nil, sess) + _, err = executorExecSession(ctx, executor, "select id from `user`", nil, sess) require.Error(t, err) assert.Contains(t, err.Error(), "scatter") - _, err = executorExecSession(executor, "select /*vt+ ALLOW_SCATTER */ id from user", nil, sess) + _, err = executorExecSession(ctx, executor, "select /*vt+ ALLOW_SCATTER */ id from user", nil, sess) require.NoError(t, err) - _, err = executorExecSession(executor, "begin", nil, sess) + _, err = executorExecSession(ctx, executor, "begin", nil, sess) require.NoError(t, err) - _, err = executorExecSession(executor, "commit", nil, sess) + _, err = executorExecSession(ctx, executor, "commit", nil, sess) require.NoError(t, err) - _, err = executorExecSession(executor, "savepoint a", nil, sess) + _, err = executorExecSession(ctx, executor, "savepoint a", nil, sess) require.NoError(t, err) } func TestGen4SelectStraightJoin(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) @@ -3112,7 +3449,12 @@ func TestGen4SelectStraightJoin(t *testing.T) { } func TestGen4MultiColumnVindexEqual(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 @@ -3151,7 +3493,12 @@ func TestGen4MultiColumnVindexEqual(t *testing.T) { } func TestGen4MultiColumnVindexIn(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 @@ -3190,7 +3537,12 @@ func TestGen4MultiColumnVindexIn(t *testing.T) { } func TestGen4MultiColMixedColComparision(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 @@ -3227,7 +3579,12 @@ func TestGen4MultiColMixedColComparision(t *testing.T) { } func TestGen4MultiColBestVindexSel(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 @@ -3273,7 +3630,12 @@ func TestGen4MultiColBestVindexSel(t *testing.T) { } func TestGen4MultiColMultiEqual(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 @@ -3296,11 +3658,16 @@ func TestGen4MultiColMultiEqual(t *testing.T) { } func TestGen4SelectUnqualifiedReferenceTable(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() executor.pv = querypb.ExecuteOptions_Gen4 query := "select * from zip_detail" - _, err := executorExec(executor, query, nil) + _, err := executorExec(ctx, executor, query, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{ { @@ -3314,11 +3681,16 @@ func TestGen4SelectUnqualifiedReferenceTable(t *testing.T) { } func TestGen4SelectQualifiedReferenceTable(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() executor.pv = querypb.ExecuteOptions_Gen4 query := fmt.Sprintf("select * from %s.zip_detail", KsTestSharded) - _, err := executorExec(executor, query, nil) + _, err := executorExec(ctx, executor, query, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{ { @@ -3332,11 +3704,16 @@ func TestGen4SelectQualifiedReferenceTable(t *testing.T) { } func TestGen4JoinUnqualifiedReferenceTable(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() executor.pv = querypb.ExecuteOptions_Gen4 query := "select * from user join zip_detail on user.zip_detail_id = zip_detail.id" - _, err := executorExec(executor, query, nil) + _, err := executorExec(ctx, executor, query, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{ { @@ -3352,7 +3729,7 @@ func TestGen4JoinUnqualifiedReferenceTable(t *testing.T) { sbc2.Queries = nil query = "select * from simple join zip_detail on simple.zip_detail_id = zip_detail.id" - _, err = executorExec(executor, query, nil) + _, err = executorExec(ctx, executor, query, nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{ { @@ -3366,11 +3743,16 @@ func TestGen4JoinUnqualifiedReferenceTable(t *testing.T) { } func TestGen4CrossShardJoinQualifiedReferenceTable(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() executor.pv = querypb.ExecuteOptions_Gen4 query := "select user.id from user join TestUnsharded.zip_detail on user.zip_detail_id = TestUnsharded.zip_detail.id" - _, err := executorExec(executor, query, nil) + _, err := executorExec(ctx, executor, query, nil) require.NoError(t, err) shardedWantQueries := []*querypb.BoundQuery{ @@ -3388,7 +3770,7 @@ func TestGen4CrossShardJoinQualifiedReferenceTable(t *testing.T) { sbc2.Queries = nil query = "select simple.id from simple join TestExecutor.zip_detail on simple.zip_detail_id = TestExecutor.zip_detail.id" - _, err = executorExec(executor, query, nil) + _, err = executorExec(ctx, executor, query, nil) require.NoError(t, err) unshardedWantQueries := []*querypb.BoundQuery{ { @@ -3402,23 +3784,27 @@ func TestGen4CrossShardJoinQualifiedReferenceTable(t *testing.T) { } func TestRegionRange(t *testing.T) { - // Special setup: Don't use createExecutorEnv. + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Special setup: Don't use createExecutorEnv. cell := "regioncell" ks := "TestExecutor" hc := discovery.NewFakeHealthCheck(nil) s := createSandbox(ks) s.ShardSpec = "-20-20a0-" s.VSchema = executorVSchema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-20a0", "20a0-"} var conns []*sandboxconn.SandboxConn for _, shard := range shards { sbc := hc.AddTestTablet(cell, shard, 1, ks, shard, topodatapb.TabletType_PRIMARY, true, 1, nil) conns = append(conns, sbc) } - executor := createExecutor(serv, cell, resolver) + executor := createExecutor(ctx, serv, cell, resolver) + defer executor.Close() executor.pv = querypb.ExecuteOptions_Gen4 tcases := []struct { @@ -3450,6 +3836,10 @@ func TestRegionRange(t *testing.T) { } func TestMultiCol(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Special setup: Don't use createLegacyExecutorEnv. cell := "multicol" ks := "TestMultiCol" @@ -3457,15 +3847,16 @@ func TestMultiCol(t *testing.T) { s := createSandbox(ks) s.ShardSpec = "-20-20a0-" s.VSchema = multiColVschema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-20a0", "20a0-"} var conns []*sandboxconn.SandboxConn for _, shard := range shards { sbc := hc.AddTestTablet(cell, shard, 1, ks, shard, topodatapb.TabletType_PRIMARY, true, 1, nil) conns = append(conns, sbc) } - executor := createExecutor(serv, cell, resolver) + executor := createExecutor(ctx, serv, cell, resolver) + defer executor.Close() executor.pv = querypb.ExecuteOptions_Gen4 tcases := []struct { @@ -3482,7 +3873,6 @@ func TestMultiCol(t *testing.T) { shards: []string{"20a0-"}, }} - ctx := context.Background() session := NewAutocommitSession(&vtgatepb.Session{}) for _, tcase := range tcases { @@ -3529,6 +3919,10 @@ var multiColVschema = ` ` func TestMultiColPartial(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Special setup: Don't use createLegacyExecutorEnv. cell := "multicol" ks := "TestMultiCol" @@ -3536,15 +3930,16 @@ func TestMultiColPartial(t *testing.T) { s := createSandbox(ks) s.ShardSpec = "-20-20a0c0-" s.VSchema = multiColVschema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-20a0c0", "20a0c0-"} var conns []*sandboxconn.SandboxConn for _, shard := range shards { sbc := hc.AddTestTablet(cell, shard, 1, ks, shard, topodatapb.TabletType_PRIMARY, true, 1, nil) conns = append(conns, sbc) } - executor := createExecutor(serv, cell, resolver) + executor := createExecutor(ctx, serv, cell, resolver) + defer executor.Close() executor.pv = querypb.ExecuteOptions_Gen4 tcases := []struct { @@ -3570,7 +3965,6 @@ func TestMultiColPartial(t *testing.T) { shards: []string{"20a0c0-"}, }} - ctx := context.Background() session := NewAutocommitSession(&vtgatepb.Session{}) for _, tcase := range tcases { @@ -3591,20 +3985,25 @@ func TestMultiColPartial(t *testing.T) { } func TestSelectAggregationNoData(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) createSandbox(KsTestSharded).VSchema = executorVSchema getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} var conns []*sandboxconn.SandboxConn for _, shard := range shards { sbc := hc.AddTestTablet(cell, shard, 1, KsTestSharded, shard, topodatapb.TabletType_PRIMARY, true, 1, nil) conns = append(conns, sbc) } - executor := createExecutor(serv, cell, resolver) + executor := createExecutor(ctx, serv, cell, resolver) + defer executor.Close() executor.pv = querypb.ExecuteOptions_Gen4 tcases := []struct { @@ -3664,7 +4063,7 @@ func TestSelectAggregationNoData(t *testing.T) { sbc.SetResults([]*sqltypes.Result{tc.sandboxRes}) sbc.Queries = nil } - qr, err := executorExec(executor, tc.sql, nil) + qr, err := executorExec(ctx, executor, tc.sql, nil) require.NoError(t, err) assert.Equal(t, tc.expField, fmt.Sprintf("%v", qr.Fields)) assert.Equal(t, tc.expRow, fmt.Sprintf("%v", qr.Rows)) @@ -3675,20 +4074,25 @@ func TestSelectAggregationNoData(t *testing.T) { } func TestSelectAggregationData(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) createSandbox(KsTestSharded).VSchema = executorVSchema getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} var conns []*sandboxconn.SandboxConn for _, shard := range shards { sbc := hc.AddTestTablet(cell, shard, 1, KsTestSharded, shard, topodatapb.TabletType_PRIMARY, true, 1, nil) conns = append(conns, sbc) } - executor := createExecutor(serv, cell, resolver) + executor := createExecutor(ctx, serv, cell, resolver) + defer executor.Close() executor.pv = querypb.ExecuteOptions_Gen4 tcases := []struct { @@ -3804,7 +4208,7 @@ func TestSelectAggregationData(t *testing.T) { sbc.SetResults([]*sqltypes.Result{tc.sandboxRes}) sbc.Queries = nil } - qr, err := executorExec(executor, tc.sql, nil) + qr, err := executorExec(ctx, executor, tc.sql, nil) require.NoError(t, err) assert.Equal(t, tc.expField, fmt.Sprintf("%v", qr.Fields)) assert.Equal(t, tc.expRow, fmt.Sprintf("%v", qr.Rows)) @@ -3815,12 +4219,16 @@ func TestSelectAggregationData(t *testing.T) { } func TestSelectAggregationRandom(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + cell := "aa" hc := discovery.NewFakeHealthCheck(nil) createSandbox(KsTestSharded).VSchema = executorVSchema getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} var conns []*sandboxconn.SandboxConn for _, shard := range shards { @@ -3838,7 +4246,8 @@ func TestSelectAggregationRandom(t *testing.T) { "10|1", )}) - executor := createExecutor(serv, cell, resolver) + executor := createExecutor(ctx, serv, cell, resolver) + defer executor.Close() executor.pv = querypb.ExecuteOptions_Gen4 session := NewAutocommitSession(&vtgatepb.Session{}) @@ -3848,7 +4257,12 @@ func TestSelectAggregationRandom(t *testing.T) { } func TestSelectHexAndBit(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true session := NewAutocommitSession(&vtgatepb.Session{}) @@ -3864,7 +4278,12 @@ func TestSelectHexAndBit(t *testing.T) { // TestSelectCFC tests validates that cfc vindex plan gets cached and same plan is getting reused. // This also validates that cache_size is able to calculate the cfc vindex plan size. func TestSelectCFC(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true session := NewAutocommitSession(&vtgatepb.Session{}) @@ -3889,7 +4308,12 @@ func TestSelectCFC(t *testing.T) { } func TestSelectView(t *testing.T) { - executor, sbc, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc, _, _ := createExecutorEnv(ctx) + defer executor.Close() // add the view to local vschema err := executor.vschema.AddView(KsTestSharded, "user_details_view", "select user.id, user_extra.col from user join user_extra on user.id = user_extra.user_id") require.NoError(t, err) diff --git a/go/vt/vtgate/executor_set_test.go b/go/vt/vtgate/executor_set_test.go index 5810d70b768..6041275030c 100644 --- a/go/vt/vtgate/executor_set_test.go +++ b/go/vt/vtgate/executor_set_test.go @@ -41,7 +41,11 @@ import ( ) func TestExecutorSet(t *testing.T) { - executorEnv, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executorEnv, _, _, _ := createExecutorEnv(ctx) + defer executorEnv.Close() testcases := []struct { in string @@ -282,7 +286,11 @@ func TestExecutorSet(t *testing.T) { } func TestExecutorSetOp(t *testing.T) { - executor, _, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() sysVarSetEnabled = true returnResult := func(columnName, typ, value string) *sqltypes.Result { @@ -379,7 +387,11 @@ func TestExecutorSetOp(t *testing.T) { } func TestExecutorSetMetadata(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) set := "set @@vitess_metadata.app_keyspace_v1= '1'" @@ -391,7 +403,8 @@ func TestExecutorSetMetadata(t *testing.T) { vschemaacl.AuthorizedDDLUsers = "" }() - executor, _, _, _ = createExecutorEnv() + executor, _, _, _ = createExecutorEnv(ctx) + defer executor.Close() session = NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) set = "set @@vitess_metadata.app_keyspace_v1= '1'" @@ -436,7 +449,11 @@ func TestExecutorSetMetadata(t *testing.T) { } func TestPlanExecutorSetUDV(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() testcases := []struct { in string @@ -466,7 +483,11 @@ func TestPlanExecutorSetUDV(t *testing.T) { } func TestSetUDVFromTabletInput(t *testing.T) { - executor, sbc1, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, sbc1, _, _ := createExecutorEnv(ctx) + defer executor.Close() fields := sqltypes.MakeTestFields("some", "VARCHAR") sbc1.SetResults([]*sqltypes.Result{ @@ -480,7 +501,7 @@ func TestSetUDVFromTabletInput(t *testing.T) { defer func() { primarySession.TargetString = "" }() - _, err := executorExec(executor, "set @foo = concat('a','b','c')", nil) + _, err := executorExec(ctx, executor, "set @foo = concat('a','b','c')", nil) require.NoError(t, err) want := map[string]*querypb.BindVariable{"foo": sqltypes.StringBindVariable("abc")} @@ -500,7 +521,11 @@ func createMap(keys []string, values []any) map[string]*querypb.BindVariable { } func TestSetVar(t *testing.T) { - executor, _, _, sbc := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, sbc := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true oldVersion := sqlparser.GetParserVersion() @@ -544,7 +569,11 @@ func TestSetVar(t *testing.T) { } func TestSetVarShowVariables(t *testing.T) { - executor, _, _, sbc := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, sbc := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true oldVersion := sqlparser.GetParserVersion() @@ -573,7 +602,11 @@ func TestSetVarShowVariables(t *testing.T) { } func TestExecutorSetAndSelect(t *testing.T) { - e, _, _, sbc := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + e, _, _, sbc := createExecutorEnv(ctx) + defer e.Close() e.normalize = true testcases := []struct { diff --git a/go/vt/vtgate/executor_stream_test.go b/go/vt/vtgate/executor_stream_test.go index 0f2eb3677b7..fa142d28d7b 100644 --- a/go/vt/vtgate/executor_stream_test.go +++ b/go/vt/vtgate/executor_stream_test.go @@ -21,7 +21,10 @@ import ( "testing" "time" + "vitess.io/vitess/go/streamlog" + "vitess.io/vitess/go/test/utils" querypb "vitess.io/vitess/go/vt/proto/query" + "vitess.io/vitess/go/vt/vtgate/logstats" "vitess.io/vitess/go/cache" "vitess.io/vitess/go/vt/discovery" @@ -35,9 +38,13 @@ import ( ) func TestStreamSQLUnsharded(t *testing.T) { - executor, _, _, _ := createExecutorEnv() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) sql := "stream * from user_msgs" result, err := executorStreamMessages(executor, sql) @@ -49,18 +56,23 @@ func TestStreamSQLUnsharded(t *testing.T) { } func TestStreamSQLSharded(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() cell := "aa" hc := discovery.NewFakeHealthCheck(nil) s := createSandbox("TestExecutor") s.VSchema = executorVSchema getSandbox(KsTestUnsharded).VSchema = unshardedVSchema - serv := newSandboxForCells([]string{cell}) - resolver := newTestResolver(hc, serv, cell) + serv := newSandboxForCells(ctx, []string{cell}) + resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} for _, shard := range shards { _ = hc.AddTestTablet(cell, shard, 1, "TestExecutor", shard, topodatapb.TabletType_PRIMARY, true, 1, nil) } - executor := NewExecutor(context.Background(), serv, cell, resolver, false, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4) + queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize) + executor := NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) + defer executor.Close() sql := "stream * from sharded_user_msgs" result, err := executorStreamMessages(executor, sql) diff --git a/go/vt/vtgate/executor_test.go b/go/vt/vtgate/executor_test.go index ce247cc69c8..3250783261b 100644 --- a/go/vt/vtgate/executor_test.go +++ b/go/vt/vtgate/executor_test.go @@ -58,11 +58,16 @@ import ( ) func TestExecutorResultsExceeded(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + save := warnMemoryRows warnMemoryRows = 3 defer func() { warnMemoryRows = save }() - executor, _, _, sbclookup := createExecutorEnv() + executor, _, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) initial := warnings.Counts()["ResultsExceeded"] @@ -81,11 +86,16 @@ func TestExecutorResultsExceeded(t *testing.T) { } func TestExecutorMaxMemoryRowsExceeded(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + save := maxMemoryRows maxMemoryRows = 3 defer func() { maxMemoryRows = save }() - executor, _, _, sbclookup := createExecutorEnv() + executor, _, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) result := sqltypes.MakeTestResult(sqltypes.MakeTestFields("col", "int64"), "1", "2", "3", "4") fn := func(r *sqltypes.Result) error { @@ -118,11 +128,16 @@ func TestExecutorMaxMemoryRowsExceeded(t *testing.T) { } func TestExecutorTransactionsNoAutoCommit(t *testing.T) { - executor, _, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", SessionUUID: "suuid"}) - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) // begin. _, err := executor.Execute(ctx, nil, "TestExecute", session, "begin", nil) @@ -130,14 +145,14 @@ func TestExecutorTransactionsNoAutoCommit(t *testing.T) { wantSession := &vtgatepb.Session{InTransaction: true, TargetString: "@primary", SessionUUID: "suuid"} utils.MustMatch(t, wantSession, session.Session, "session") assert.EqualValues(t, 0, sbclookup.CommitCount.Load(), "commit count") - logStats := testQueryLog(t, logChan, "TestExecute", "BEGIN", "begin", 0) + logStats := testQueryLog(t, executor, logChan, "TestExecute", "BEGIN", "begin", 0) assert.EqualValues(t, 0, logStats.CommitTime, "logstats: expected zero CommitTime") assert.EqualValues(t, "suuid", logStats.SessionUUID, "logstats: expected non-empty SessionUUID") // commit. _, err = executor.Execute(ctx, nil, "TestExecute", session, "select id from main1", nil) require.NoError(t, err) - logStats = testQueryLog(t, logChan, "TestExecute", "SELECT", "select id from main1", 1) + logStats = testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select id from main1", 1) assert.EqualValues(t, 0, logStats.CommitTime, "logstats: expected zero CommitTime") assert.EqualValues(t, "suuid", logStats.SessionUUID, "logstats: expected non-empty SessionUUID") @@ -152,7 +167,7 @@ func TestExecutorTransactionsNoAutoCommit(t *testing.T) { if commitCount := sbclookup.CommitCount.Load(); commitCount != 1 { t.Errorf("want 1, got %d", commitCount) } - logStats = testQueryLog(t, logChan, "TestExecute", "COMMIT", "commit", 1) + logStats = testQueryLog(t, executor, logChan, "TestExecute", "COMMIT", "commit", 1) if logStats.CommitTime == 0 { t.Errorf("logstats: expected non-zero CommitTime") } @@ -168,9 +183,9 @@ func TestExecutorTransactionsNoAutoCommit(t *testing.T) { wantSession = &vtgatepb.Session{TargetString: "@primary", SessionUUID: "suuid"} utils.MustMatch(t, wantSession, session.Session, "session") assert.EqualValues(t, 1, sbclookup.RollbackCount.Load(), "rollback count") - _ = testQueryLog(t, logChan, "TestExecute", "BEGIN", "begin", 0) - _ = testQueryLog(t, logChan, "TestExecute", "SELECT", "select id from main1", 1) - logStats = testQueryLog(t, logChan, "TestExecute", "ROLLBACK", "rollback", 1) + _ = testQueryLog(t, executor, logChan, "TestExecute", "BEGIN", "begin", 0) + _ = testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select id from main1", 1) + logStats = testQueryLog(t, executor, logChan, "TestExecute", "ROLLBACK", "rollback", 1) if logStats.CommitTime == 0 { t.Errorf("logstats: expected non-zero CommitTime") } @@ -191,7 +206,12 @@ func TestExecutorTransactionsNoAutoCommit(t *testing.T) { } func TestDirectTargetRewrites(t *testing.T) { - executor, _, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true session := &vtgatepb.Session{ @@ -210,11 +230,16 @@ func TestDirectTargetRewrites(t *testing.T) { } func TestExecutorTransactionsAutoCommit(t *testing.T) { - executor, _, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true, SessionUUID: "suuid"}) - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) // begin. _, err := executor.Execute(ctx, nil, "TestExecute", session, "begin", nil) @@ -224,7 +249,7 @@ func TestExecutorTransactionsAutoCommit(t *testing.T) { if commitCount := sbclookup.CommitCount.Load(); commitCount != 0 { t.Errorf("want 0, got %d", commitCount) } - logStats := testQueryLog(t, logChan, "TestExecute", "BEGIN", "begin", 0) + logStats := testQueryLog(t, executor, logChan, "TestExecute", "BEGIN", "begin", 0) assert.EqualValues(t, "suuid", logStats.SessionUUID, "logstats: expected non-empty SessionUUID") // commit. @@ -236,10 +261,10 @@ func TestExecutorTransactionsAutoCommit(t *testing.T) { utils.MustMatch(t, wantSession, session.Session, "session") assert.EqualValues(t, 1, sbclookup.CommitCount.Load()) - logStats = testQueryLog(t, logChan, "TestExecute", "SELECT", "select id from main1", 1) + logStats = testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select id from main1", 1) assert.EqualValues(t, 0, logStats.CommitTime) assert.EqualValues(t, "suuid", logStats.SessionUUID, "logstats: expected non-empty SessionUUID") - logStats = testQueryLog(t, logChan, "TestExecute", "COMMIT", "commit", 1) + logStats = testQueryLog(t, executor, logChan, "TestExecute", "COMMIT", "commit", 1) assert.NotEqual(t, 0, logStats.CommitTime) assert.EqualValues(t, "suuid", logStats.SessionUUID, "logstats: expected non-empty SessionUUID") @@ -255,14 +280,19 @@ func TestExecutorTransactionsAutoCommit(t *testing.T) { if rollbackCount := sbclookup.RollbackCount.Load(); rollbackCount != 1 { t.Errorf("want 1, got %d", rollbackCount) } - _ = testQueryLog(t, logChan, "TestExecute", "BEGIN", "begin", 0) - _ = testQueryLog(t, logChan, "TestExecute", "SELECT", "select id from main1", 1) - logStats = testQueryLog(t, logChan, "TestExecute", "ROLLBACK", "rollback", 1) + _ = testQueryLog(t, executor, logChan, "TestExecute", "BEGIN", "begin", 0) + _ = testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select id from main1", 1) + logStats = testQueryLog(t, executor, logChan, "TestExecute", "ROLLBACK", "rollback", 1) assert.EqualValues(t, "suuid", logStats.SessionUUID, "logstats: expected non-empty SessionUUID") } func TestExecutorTransactionsAutoCommitStreaming(t *testing.T) { - executor, _, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() oltpOptions := &querypb.ExecuteOptions{Workload: querypb.ExecuteOptions_OLTP} session := NewSafeSession(&vtgatepb.Session{ TargetString: "@primary", @@ -271,8 +301,8 @@ func TestExecutorTransactionsAutoCommitStreaming(t *testing.T) { SessionUUID: "suuid", }) - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) var results []*sqltypes.Result @@ -295,7 +325,7 @@ func TestExecutorTransactionsAutoCommitStreaming(t *testing.T) { } utils.MustMatch(t, wantSession, session.Session, "session") assert.Zero(t, sbclookup.CommitCount.Load()) - logStats := testQueryLog(t, logChan, "TestExecute", "BEGIN", "begin", 0) + logStats := testQueryLog(t, executor, logChan, "TestExecute", "BEGIN", "begin", 0) assert.EqualValues(t, "suuid", logStats.SessionUUID, "logstats: expected non-empty SessionUUID") // commit. @@ -307,10 +337,10 @@ func TestExecutorTransactionsAutoCommitStreaming(t *testing.T) { utils.MustMatch(t, wantSession, session.Session, "session") assert.EqualValues(t, 1, sbclookup.CommitCount.Load()) - logStats = testQueryLog(t, logChan, "TestExecute", "SELECT", "select id from main1", 1) + logStats = testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select id from main1", 1) assert.EqualValues(t, 0, logStats.CommitTime) assert.EqualValues(t, "suuid", logStats.SessionUUID, "logstats: expected non-empty SessionUUID") - logStats = testQueryLog(t, logChan, "TestExecute", "COMMIT", "commit", 1) + logStats = testQueryLog(t, executor, logChan, "TestExecute", "COMMIT", "commit", 1) assert.NotEqual(t, 0, logStats.CommitTime) assert.EqualValues(t, "suuid", logStats.SessionUUID, "logstats: expected non-empty SessionUUID") @@ -327,12 +357,17 @@ func TestExecutorTransactionsAutoCommitStreaming(t *testing.T) { } func TestExecutorDeleteMetadata(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + vschemaacl.AuthorizedDDLUsers = "%" defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, _, _, _ := createExecutorEnv() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) set := "set @@vitess_metadata.app_v1= '1'" @@ -359,11 +394,16 @@ func TestExecutorDeleteMetadata(t *testing.T) { } func TestExecutorAutocommit(t *testing.T) { - executor, _, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) // autocommit = 0 startCount := sbclookup.CommitCount.Load() @@ -374,7 +414,7 @@ func TestExecutorAutocommit(t *testing.T) { testSession.ShardSessions = nil utils.MustMatch(t, wantSession, testSession, "session does not match for autocommit=0") - logStats := testQueryLog(t, logChan, "TestExecute", "SELECT", "select id from main1", 1) + logStats := testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select id from main1", 1) if logStats.CommitTime != 0 { t.Errorf("logstats: expected zero CommitTime") } @@ -385,7 +425,7 @@ func TestExecutorAutocommit(t *testing.T) { // autocommit = 1 _, err = executor.Execute(ctx, nil, "TestExecute", session, "set autocommit=1", nil) require.NoError(t, err) - _ = testQueryLog(t, logChan, "TestExecute", "SET", "set @@autocommit = 1", 0) + _ = testQueryLog(t, executor, logChan, "TestExecute", "SET", "set @@autocommit = 1", 0) // Setting autocommit=1 commits existing transaction. if got, want := sbclookup.CommitCount.Load(), startCount+1; got != want { @@ -397,7 +437,7 @@ func TestExecutorAutocommit(t *testing.T) { wantSession = &vtgatepb.Session{Autocommit: true, TargetString: "@primary", FoundRows: 0, RowCount: 1} utils.MustMatch(t, wantSession, session.Session, "session does not match for autocommit=1") - logStats = testQueryLog(t, logChan, "TestExecute", "UPDATE", "update main1 set id = 1", 1) + logStats = testQueryLog(t, executor, logChan, "TestExecute", "UPDATE", "update main1 set id = 1", 1) assert.NotZero(t, logStats.CommitTime, "logstats: expected non-zero CommitTime") assert.NotEqual(t, uint64(0), logStats.RowsAffected, "logstats: expected non-zero RowsAffected") @@ -406,7 +446,7 @@ func TestExecutorAutocommit(t *testing.T) { startCount = sbclookup.CommitCount.Load() _, err = executor.Execute(ctx, nil, "TestExecute", session, "begin", nil) require.NoError(t, err) - _ = testQueryLog(t, logChan, "TestExecute", "BEGIN", "begin", 0) + _ = testQueryLog(t, executor, logChan, "TestExecute", "BEGIN", "begin", 0) _, err = executor.Execute(ctx, nil, "TestExecute", session, "update main1 set id=1", nil) require.NoError(t, err) @@ -418,7 +458,7 @@ func TestExecutorAutocommit(t *testing.T) { t.Errorf("Commit count: %d, want %d", got, want) } - logStats = testQueryLog(t, logChan, "TestExecute", "UPDATE", "update main1 set id = 1", 1) + logStats = testQueryLog(t, executor, logChan, "TestExecute", "UPDATE", "update main1 set id = 1", 1) if logStats.CommitTime != 0 { t.Errorf("logstats: expected zero CommitTime") } @@ -435,7 +475,7 @@ func TestExecutorAutocommit(t *testing.T) { if got, want := sbclookup.CommitCount.Load(), startCount+1; got != want { t.Errorf("Commit count: %d, want %d", got, want) } - _ = testQueryLog(t, logChan, "TestExecute", "COMMIT", "commit", 1) + _ = testQueryLog(t, executor, logChan, "TestExecute", "COMMIT", "commit", 1) // transition autocommit from 0 to 1 in the middle of a transaction. startCount = sbclookup.CommitCount.Load() @@ -459,7 +499,12 @@ func TestExecutorAutocommit(t *testing.T) { } func TestExecutorShowColumns(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: ""}) queries := []string{ @@ -510,7 +555,12 @@ func assertMatchesNoOrder(t *testing.T, expected, got string) { } func TestExecutorShow(t *testing.T) { - executor, _, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) for _, query := range []string{"show vitess_keyspaces", "show keyspaces"} { @@ -1068,7 +1118,12 @@ func TestExecutorShow(t *testing.T) { } func TestExecutorShowTargeted(t *testing.T) { - executor, _, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor/40-60"}) queries := []string{ @@ -1094,7 +1149,12 @@ func TestExecutorShowTargeted(t *testing.T) { } func TestExecutorUse(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(&vtgatepb.Session{Autocommit: true, TargetString: "@primary"}) stmts := []string{ @@ -1125,7 +1185,12 @@ func TestExecutorUse(t *testing.T) { } func TestExecutorComment(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() stmts := []string{ "/*! SET autocommit=1*/", @@ -1145,7 +1210,12 @@ func TestExecutorComment(t *testing.T) { } func TestExecutorOther(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() type cnts struct { Sbc1Cnt int64 @@ -1237,10 +1307,14 @@ func TestExecutorOther(t *testing.T) { } func TestExecutorDDL(t *testing.T) { - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) type cnts struct { Sbc1Cnt int64 @@ -1328,7 +1402,7 @@ func TestExecutorDDL(t *testing.T) { t.Errorf("stmt: %s\ntc: %+v\n-want,+got:\n%s", stmt, tc, diff) } - testQueryLog(t, logChan, "TestExecute", stmtType, stmt, tc.shardQueryCnt) + testQueryLog(t, executor, logChan, "TestExecute", stmtType, stmt, tc.shardQueryCnt) } } @@ -1352,16 +1426,21 @@ func TestExecutorDDL(t *testing.T) { _, err := executor.Execute(ctx, nil, "TestExecute", NewSafeSession(&vtgatepb.Session{TargetString: ""}), stmt.input, nil) if stmt.hasErr { require.EqualError(t, err, errNoKeyspace.Error(), "expect query to fail") - testQueryLog(t, logChan, "TestExecute", "", stmt.input, 0) + testQueryLog(t, executor, logChan, "TestExecute", "", stmt.input, 0) } else { require.NoError(t, err) - testQueryLog(t, logChan, "TestExecute", "DDL", stmt.input, 8) + testQueryLog(t, executor, logChan, "TestExecute", "DDL", stmt.input, 8) } } } func TestExecutorDDLFk(t *testing.T) { - executor, _, _, sbc := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, sbc := createExecutorEnv(ctx) + defer executor.Close() mName := "TestExecutorDDLFk" stmts := []string{ @@ -1388,11 +1467,16 @@ func TestExecutorDDLFk(t *testing.T) { } func TestExecutorAlterVSchemaKeyspace(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + vschemaacl.AuthorizedDDLUsers = "%" defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, _, _, _ := createExecutorEnv() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) vschemaUpdates := make(chan *vschemapb.SrvVSchema, 2) @@ -1416,11 +1500,16 @@ func TestExecutorAlterVSchemaKeyspace(t *testing.T) { } func TestExecutorCreateVindexDDL(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + vschemaacl.AuthorizedDDLUsers = "%" defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() ks := "TestExecutor" vschemaUpdates := make(chan *vschemapb.SrvVSchema, 4) @@ -1486,11 +1575,16 @@ func TestExecutorCreateVindexDDL(t *testing.T) { } func TestExecutorAddDropVschemaTableDDL(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + vschemaacl.AuthorizedDDLUsers = "%" defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() ks := KsTestUnsharded vschemaUpdates := make(chan *vschemapb.SrvVSchema, 4) @@ -1538,7 +1632,12 @@ func TestExecutorAddDropVschemaTableDDL(t *testing.T) { } func TestExecutorVindexDDLACL(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() ks := "TestExecutor" session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) @@ -1583,7 +1682,12 @@ func TestExecutorVindexDDLACL(t *testing.T) { } func TestExecutorUnrecognized(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() _, err := executor.Execute(ctx, nil, "TestExecute", NewSafeSession(&vtgatepb.Session{}), "invalid statement", nil) require.Error(t, err, "unrecognized statement: invalid statement'") } @@ -1591,7 +1695,12 @@ func TestExecutorUnrecognized(t *testing.T) { // TestVSchemaStats makes sure the building and displaying of the // VSchemaStats works. func TestVSchemaStats(t *testing.T) { - r, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + r, _, _, _ := createExecutorEnv(ctx) + defer r.Close() stats := r.VSchemaStats() @@ -1614,18 +1723,23 @@ func TestVSchemaStats(t *testing.T) { var pv = querypb.ExecuteOptions_Gen4 func TestGetPlanUnnormalized(t *testing.T) { - r, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + r, _, _, _ := createExecutorEnv(ctx) + defer r.Close() emptyvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) unshardedvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) query1 := "select * from music_user_map where id = 1" - plan1, logStats1 := getPlanCached(t, r, emptyvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) + plan1, logStats1 := getPlanCached(t, ctx, r, emptyvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) wantSQL := query1 + " /* comment */" if logStats1.SQL != wantSQL { t.Errorf("logstats sql want \"%s\" got \"%s\"", wantSQL, logStats1.SQL) } - plan2, logStats2 := getPlanCached(t, r, emptyvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) + plan2, logStats2 := getPlanCached(t, ctx, r, emptyvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) if plan1 != plan2 { t.Errorf("getPlan(query1): plans must be equal: %p %p", plan1, plan2) } @@ -1633,14 +1747,14 @@ func TestGetPlanUnnormalized(t *testing.T) { if logStats2.SQL != wantSQL { t.Errorf("logstats sql want \"%s\" got \"%s\"", wantSQL, logStats2.SQL) } - plan3, logStats3 := getPlanCached(t, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) + plan3, logStats3 := getPlanCached(t, ctx, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) if plan1 == plan3 { t.Errorf("getPlan(query1, ks): plans must not be equal: %p %p", plan1, plan3) } if logStats3.SQL != wantSQL { t.Errorf("logstats sql want \"%s\" got \"%s\"", wantSQL, logStats3.SQL) } - plan4, logStats4 := getPlanCached(t, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) + plan4, logStats4 := getPlanCached(t, ctx, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) if plan3 != plan4 { t.Errorf("getPlan(query1, ks): plans must be equal: %p %p", plan3, plan4) } @@ -1685,7 +1799,7 @@ func assertCacheContains(t *testing.T, e *Executor, vc *vcursorImpl, sql string) return plan } -func getPlanCached(t *testing.T, e *Executor, vcursor *vcursorImpl, sql string, comments sqlparser.MarginComments, bindVars map[string]*querypb.BindVariable, skipQueryPlanCache bool) (*engine.Plan, *logstats.LogStats) { +func getPlanCached(t *testing.T, ctx context.Context, e *Executor, vcursor *vcursorImpl, sql string, comments sqlparser.MarginComments, bindVars map[string]*querypb.BindVariable, skipQueryPlanCache bool) (*engine.Plan, *logstats.LogStats) { logStats := logstats.NewLogStats(ctx, "Test", "", "", nil) vcursor.safeSession = &SafeSession{ Session: &vtgatepb.Session{ @@ -1703,11 +1817,16 @@ func getPlanCached(t *testing.T, e *Executor, vcursor *vcursorImpl, sql string, } func TestGetPlanCacheUnnormalized(t *testing.T) { - r, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + r, _, _, _ := createExecutorEnv(ctx) + defer r.Close() emptyvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) query1 := "select * from music_user_map where id = 1" - _, logStats1 := getPlanCached(t, r, emptyvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, true) + _, logStats1 := getPlanCached(t, ctx, r, emptyvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, true) assertCacheSize(t, r.plans, 0) wantSQL := query1 + " /* comment */" @@ -1715,7 +1834,7 @@ func TestGetPlanCacheUnnormalized(t *testing.T) { t.Errorf("logstats sql want \"%s\" got \"%s\"", wantSQL, logStats1.SQL) } - _, logStats2 := getPlanCached(t, r, emptyvc, query1, makeComments(" /* comment 2 */"), map[string]*querypb.BindVariable{}, false) + _, logStats2 := getPlanCached(t, ctx, r, emptyvc, query1, makeComments(" /* comment 2 */"), map[string]*querypb.BindVariable{}, false) assertCacheSize(t, r.plans, 1) wantSQL = query1 + " /* comment 2 */" @@ -1724,69 +1843,81 @@ func TestGetPlanCacheUnnormalized(t *testing.T) { } // Skip cache using directive - r, _, _, _ = createExecutorEnv() + r, _, _, _ = createExecutorEnv(ctx) + defer r.Close() unshardedvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) query1 = "insert /*vt+ SKIP_QUERY_PLAN_CACHE=1 */ into user(id) values (1), (2)" - getPlanCached(t, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) + getPlanCached(t, ctx, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) assertCacheSize(t, r.plans, 0) query1 = "insert into user(id) values (1), (2)" - getPlanCached(t, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) + getPlanCached(t, ctx, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) assertCacheSize(t, r.plans, 1) // the target string will be resolved and become part of the plan cache key, which adds a new entry ksIDVc1, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "[deadbeef]"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) - getPlanCached(t, r, ksIDVc1, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) + getPlanCached(t, ctx, r, ksIDVc1, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) assertCacheSize(t, r.plans, 2) // the target string will be resolved and become part of the plan cache key, as it's an unsharded ks, it will be the same entry as above ksIDVc2, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "[beefdead]"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) - getPlanCached(t, r, ksIDVc2, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) + getPlanCached(t, ctx, r, ksIDVc2, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) assertCacheSize(t, r.plans, 2) } func TestGetPlanCacheNormalized(t *testing.T) { - r, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + r, _, _, _ := createExecutorEnv(ctx) + defer r.Close() r.normalize = true emptyvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) query1 := "select * from music_user_map where id = 1" - _, logStats1 := getPlanCached(t, r, emptyvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, true /* skipQueryPlanCache */) + _, logStats1 := getPlanCached(t, ctx, r, emptyvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, true /* skipQueryPlanCache */) assertCacheSize(t, r.plans, 0) wantSQL := "select * from music_user_map where id = :id /* INT64 */ /* comment */" assert.Equal(t, wantSQL, logStats1.SQL) - _, logStats2 := getPlanCached(t, r, emptyvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false /* skipQueryPlanCache */) + _, logStats2 := getPlanCached(t, ctx, r, emptyvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false /* skipQueryPlanCache */) assertCacheSize(t, r.plans, 1) assert.Equal(t, wantSQL, logStats2.SQL) // Skip cache using directive - r, _, _, _ = createExecutorEnv() + r, _, _, _ = createExecutorEnv(ctx) + defer r.Close() r.normalize = true unshardedvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) query1 = "insert /*vt+ SKIP_QUERY_PLAN_CACHE=1 */ into user(id) values (1), (2)" - getPlanCached(t, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) + getPlanCached(t, ctx, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) assertCacheSize(t, r.plans, 0) query1 = "insert into user(id) values (1), (2)" - getPlanCached(t, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) + getPlanCached(t, ctx, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) assertCacheSize(t, r.plans, 1) // the target string will be resolved and become part of the plan cache key, which adds a new entry ksIDVc1, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "[deadbeef]"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) - getPlanCached(t, r, ksIDVc1, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) + getPlanCached(t, ctx, r, ksIDVc1, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) assertCacheSize(t, r.plans, 2) // the target string will be resolved and become part of the plan cache key, as it's an unsharded ks, it will be the same entry as above ksIDVc2, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "[beefdead]"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) - getPlanCached(t, r, ksIDVc2, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) + getPlanCached(t, ctx, r, ksIDVc2, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) assertCacheSize(t, r.plans, 2) } func TestGetPlanNormalized(t *testing.T) { - r, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + r, _, _, _ := createExecutorEnv(ctx) + defer r.Close() r.normalize = true emptyvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) unshardedvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) @@ -1795,8 +1926,8 @@ func TestGetPlanNormalized(t *testing.T) { query2 := "select * from music_user_map where id = 2" normalized := "select * from music_user_map where id = :id /* INT64 */" - plan1, logStats1 := getPlanCached(t, r, emptyvc, query1, makeComments(" /* comment 1 */"), map[string]*querypb.BindVariable{}, false) - plan2, logStats2 := getPlanCached(t, r, emptyvc, query1, makeComments(" /* comment 2 */"), map[string]*querypb.BindVariable{}, false) + plan1, logStats1 := getPlanCached(t, ctx, r, emptyvc, query1, makeComments(" /* comment 1 */"), map[string]*querypb.BindVariable{}, false) + plan2, logStats2 := getPlanCached(t, ctx, r, emptyvc, query1, makeComments(" /* comment 2 */"), map[string]*querypb.BindVariable{}, false) assert.Equal(t, plan1, plan2) assertCacheContains(t, r, emptyvc, normalized) @@ -1806,18 +1937,18 @@ func TestGetPlanNormalized(t *testing.T) { wantSQL = normalized + " /* comment 2 */" assert.Equal(t, wantSQL, logStats2.SQL) - plan3, logStats3 := getPlanCached(t, r, emptyvc, query2, makeComments(" /* comment 3 */"), map[string]*querypb.BindVariable{}, false) + plan3, logStats3 := getPlanCached(t, ctx, r, emptyvc, query2, makeComments(" /* comment 3 */"), map[string]*querypb.BindVariable{}, false) assert.Equal(t, plan1, plan3) wantSQL = normalized + " /* comment 3 */" assert.Equal(t, wantSQL, logStats3.SQL) var logStats5 *logstats.LogStats - plan3, logStats5 = getPlanCached(t, r, unshardedvc, query1, makeComments(" /* comment 5 */"), map[string]*querypb.BindVariable{}, false) + plan3, logStats5 = getPlanCached(t, ctx, r, unshardedvc, query1, makeComments(" /* comment 5 */"), map[string]*querypb.BindVariable{}, false) assert.Equal(t, plan1, plan3) wantSQL = normalized + " /* comment 5 */" assert.Equal(t, wantSQL, logStats5.SQL) - plan4, _ := getPlanCached(t, r, unshardedvc, query1, makeComments(" /* comment 6 */"), map[string]*querypb.BindVariable{}, false) + plan4, _ := getPlanCached(t, ctx, r, unshardedvc, query1, makeComments(" /* comment 6 */"), map[string]*querypb.BindVariable{}, false) assert.Equal(t, plan1, plan4) assertCacheContains(t, r, emptyvc, normalized) assertCacheContains(t, r, unshardedvc, normalized) @@ -1842,7 +1973,12 @@ func TestGetPlanPriority(t *testing.T) { testCase := aTestCase t.Run(testCase.name, func(t *testing.T) { - r, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + r, _, _, _ := createExecutorEnv(ctx) + defer r.Close() r.normalize = true logStats := logstats.NewLogStats(ctx, "Test", "", "", nil) vCursor, err := newVCursorImpl(session, makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) @@ -1866,11 +2002,16 @@ func TestGetPlanPriority(t *testing.T) { } func TestPassthroughDDL(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() primarySession.TargetString = "TestExecutor" alterDDL := "/* leading */ alter table passthrough_ddl add columne col bigint default 123 /* trailing */" - _, err := executorExec(executor, alterDDL, nil) + _, err := executorExec(ctx, executor, alterDDL, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: alterDDL, @@ -1889,7 +2030,7 @@ func TestPassthroughDDL(t *testing.T) { primarySession.TargetString = "TestExecutor/40-60" executor.normalize = true - _, err = executorExec(executor, alterDDL, nil) + _, err = executorExec(ctx, executor, alterDDL, nil) require.NoError(t, err) require.Nil(t, sbc1.Queries) if !reflect.DeepEqual(sbc2.Queries, wantQueries) { @@ -1902,7 +2043,7 @@ func TestPassthroughDDL(t *testing.T) { primarySession.TargetString = "TestExecutor[-]" executor.normalize = true - _, err = executorExec(executor, alterDDL, nil) + _, err = executorExec(ctx, executor, alterDDL, nil) require.NoError(t, err) if !reflect.DeepEqual(sbc1.Queries, wantQueries) { t.Errorf("sbc2.Queries: %+v, want %+v\n", sbc1.Queries, wantQueries) @@ -1915,7 +2056,12 @@ func TestPassthroughDDL(t *testing.T) { } func TestParseEmptyTargetSingleKeyspace(t *testing.T) { - r, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + r, _, _, _ := createExecutorEnv(ctx) + defer r.Close() altVSchema := &vindexes.VSchema{ Keyspaces: map[string]*vindexes.KeyspaceSchema{ KsTestUnsharded: r.vschema.Keyspaces[KsTestUnsharded], @@ -1937,7 +2083,12 @@ func TestParseEmptyTargetSingleKeyspace(t *testing.T) { } func TestParseEmptyTargetMultiKeyspace(t *testing.T) { - r, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + r, _, _, _ := createExecutorEnv(ctx) + defer r.Close() altVSchema := &vindexes.VSchema{ Keyspaces: map[string]*vindexes.KeyspaceSchema{ KsTestUnsharded: r.vschema.Keyspaces[KsTestUnsharded], @@ -1960,7 +2111,12 @@ func TestParseEmptyTargetMultiKeyspace(t *testing.T) { } func TestParseTargetSingleKeyspace(t *testing.T) { - r, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + r, _, _, _ := createExecutorEnv(ctx) + defer r.Close() altVSchema := &vindexes.VSchema{ Keyspaces: map[string]*vindexes.KeyspaceSchema{ KsTestUnsharded: r.vschema.Keyspaces[KsTestUnsharded], @@ -1982,10 +2138,15 @@ func TestParseTargetSingleKeyspace(t *testing.T) { } func TestDebugVSchema(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + resp := httptest.NewRecorder() req, _ := http.NewRequest("GET", "/debug/vschema", nil) - executor, _, _, _ := createExecutorEnv() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.ServeHTTP(resp, req) v := make(map[string]any) if err := json.Unmarshal(resp.Body.Bytes(), &v); err != nil { @@ -2000,6 +2161,10 @@ func TestDebugVSchema(t *testing.T) { } func TestExecutorMaxPayloadSizeExceeded(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + saveMax := maxPayloadSize saveWarn := warnPayloadSize maxPayloadSize = 10 @@ -2009,7 +2174,8 @@ func TestExecutorMaxPayloadSizeExceeded(t *testing.T) { warnPayloadSize = saveWarn }() - executor, _, _, _ := createExecutorEnv() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) warningCount := warnings.Counts()["WarnPayloadSizeExceeded"] testMaxPayloadSizeExceeded := []string{ @@ -2046,7 +2212,12 @@ func TestExecutorMaxPayloadSizeExceeded(t *testing.T) { } func TestOlapSelectDatabase(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true session := &vtgatepb.Session{Autocommit: true} @@ -2063,7 +2234,12 @@ func TestOlapSelectDatabase(t *testing.T) { } func TestExecutorClearsWarnings(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(&vtgatepb.Session{ Warnings: []*querypb.QueryWarning{{Code: 234, Message: "oh noes"}}, }) @@ -2074,10 +2250,17 @@ func TestExecutorClearsWarnings(t *testing.T) { // TestServingKeyspaces tests that the dual queries are routed to the correct keyspaces from the list of serving keyspaces. func TestServingKeyspaces(t *testing.T) { - buffer.EnableBuffering() - defer buffer.DisableBuffering() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() - executor, sbc1, _, sbclookup := createExecutorEnv() + buffer.SetBufferingModeInTestingEnv(true) + defer func() { + buffer.SetBufferingModeInTestingEnv(false) + }() + + executor, sbc1, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() executor.pv = querypb.ExecuteOptions_Gen4 gw, ok := executor.resolver.resolver.GetGateway().(*TabletGateway) require.True(t, ok) @@ -2119,7 +2302,12 @@ func TestServingKeyspaces(t *testing.T) { } func TestExecutorOtherRead(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() type cnts struct { Sbc1Cnt int64 @@ -2194,27 +2382,37 @@ func TestExecutorOtherRead(t *testing.T) { } func TestExecutorVExplain(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) bindVars := map[string]*querypb.BindVariable{} - result, err := executorExec(executor, "vexplain plan select * from user", bindVars) + result, err := executorExec(ctx, executor, "vexplain plan select * from user", bindVars) require.NoError(t, err) require.Equal(t, `[[VARCHAR("{\n\t\"OperatorType\": \"Route\",\n\t\"Variant\": \"Scatter\",\n\t\"Keyspace\": {\n\t\t\"Name\": \"TestExecutor\",\n\t\t\"Sharded\": true\n\t},\n\t\"FieldQuery\": \"select * from `+"`user`"+` where 1 != 1\",\n\t\"Query\": \"select * from `+"`user`"+`\",\n\t\"Table\": \"`+"`user`"+`\"\n}")]]`, fmt.Sprintf("%v", result.Rows)) - result, err = executorExec(executor, "vexplain plan select 42", bindVars) + result, err = executorExec(ctx, executor, "vexplain plan select 42", bindVars) require.NoError(t, err) expected := `[[VARCHAR("{\n\t\"OperatorType\": \"Projection\",\n\t\"Expressions\": [\n\t\t\"INT64(42) as 42\"\n\t],\n\t\"Inputs\": [\n\t\t{\n\t\t\t\"OperatorType\": \"SingleRow\"\n\t\t}\n\t]\n}")]]` require.Equal(t, expected, fmt.Sprintf("%v", result.Rows)) } func TestExecutorOtherAdmin(t *testing.T) { - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() type cnts struct { Sbc1Cnt int64 @@ -2288,9 +2486,14 @@ func TestExecutorOtherAdmin(t *testing.T) { } func TestExecutorSavepointInTx(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() - logChan := QueryLogger.Subscribe("TestExecutorSavepoint") - defer QueryLogger.Unsubscribe(logChan) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() + logChan := executor.queryLogger.Subscribe("TestExecutorSavepoint") + defer executor.queryLogger.Unsubscribe(logChan) session := NewSafeSession(&vtgatepb.Session{Autocommit: false, TargetString: "@primary"}) _, err := exec(executor, session, "savepoint a") @@ -2358,21 +2561,26 @@ func TestExecutorSavepointInTx(t *testing.T) { }} utils.MustMatch(t, sbc1WantQueries, sbc1.Queries, "") utils.MustMatch(t, sbc2WantQueries, sbc2.Queries, "") - testQueryLog(t, logChan, "TestExecute", "SAVEPOINT", "savepoint a", 0) - testQueryLog(t, logChan, "TestExecute", "SAVEPOINT_ROLLBACK", "rollback to a", 0) - testQueryLog(t, logChan, "TestExecute", "RELEASE", "release savepoint a", 0) - testQueryLog(t, logChan, "TestExecute", "SELECT", "select id from `user` where id = 1", 1) - testQueryLog(t, logChan, "TestExecute", "SAVEPOINT", "savepoint b", 1) - testQueryLog(t, logChan, "TestExecute", "SAVEPOINT_ROLLBACK", "rollback to b", 1) - testQueryLog(t, logChan, "TestExecute", "RELEASE", "release savepoint b", 1) - testQueryLog(t, logChan, "TestExecute", "SELECT", "select id from `user` where id = 3", 1) - testQueryLog(t, logChan, "TestExecute", "ROLLBACK", "rollback", 2) + testQueryLog(t, executor, logChan, "TestExecute", "SAVEPOINT", "savepoint a", 0) + testQueryLog(t, executor, logChan, "TestExecute", "SAVEPOINT_ROLLBACK", "rollback to a", 0) + testQueryLog(t, executor, logChan, "TestExecute", "RELEASE", "release savepoint a", 0) + testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select id from `user` where id = 1", 1) + testQueryLog(t, executor, logChan, "TestExecute", "SAVEPOINT", "savepoint b", 1) + testQueryLog(t, executor, logChan, "TestExecute", "SAVEPOINT_ROLLBACK", "rollback to b", 1) + testQueryLog(t, executor, logChan, "TestExecute", "RELEASE", "release savepoint b", 1) + testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select id from `user` where id = 3", 1) + testQueryLog(t, executor, logChan, "TestExecute", "ROLLBACK", "rollback", 2) } func TestExecutorSavepointInTxWithReservedConn(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() - logChan := QueryLogger.Subscribe("TestExecutorSavepoint") - defer QueryLogger.Unsubscribe(logChan) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() + logChan := executor.queryLogger.Subscribe("TestExecutorSavepoint") + defer executor.queryLogger.Unsubscribe(logChan) session := NewSafeSession(&vtgatepb.Session{Autocommit: true, TargetString: "TestExecutor", EnableSystemSettings: true}) sbc1.SetResults([]*sqltypes.Result{ @@ -2425,20 +2633,25 @@ func TestExecutorSavepointInTxWithReservedConn(t *testing.T) { utils.MustMatch(t, sbc1WantQueries, sbc1.Queries, "") utils.MustMatch(t, sbc2WantQueries, sbc2.Queries, "") - testQueryLog(t, logChan, "TestExecute", "SET", "set @@sql_mode = ''", 1) - testQueryLog(t, logChan, "TestExecute", "BEGIN", "begin", 0) - testQueryLog(t, logChan, "TestExecute", "SAVEPOINT", "savepoint a", 0) - testQueryLog(t, logChan, "TestExecute", "SELECT", "select id from `user` where id = 1", 1) - testQueryLog(t, logChan, "TestExecute", "SAVEPOINT", "savepoint b", 1) - testQueryLog(t, logChan, "TestExecute", "RELEASE", "release savepoint a", 1) - testQueryLog(t, logChan, "TestExecute", "SELECT", "select id from `user` where id = 3", 1) - testQueryLog(t, logChan, "TestExecute", "COMMIT", "commit", 2) + testQueryLog(t, executor, logChan, "TestExecute", "SET", "set @@sql_mode = ''", 1) + testQueryLog(t, executor, logChan, "TestExecute", "BEGIN", "begin", 0) + testQueryLog(t, executor, logChan, "TestExecute", "SAVEPOINT", "savepoint a", 0) + testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select id from `user` where id = 1", 1) + testQueryLog(t, executor, logChan, "TestExecute", "SAVEPOINT", "savepoint b", 1) + testQueryLog(t, executor, logChan, "TestExecute", "RELEASE", "release savepoint a", 1) + testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select id from `user` where id = 3", 1) + testQueryLog(t, executor, logChan, "TestExecute", "COMMIT", "commit", 2) } func TestExecutorSavepointWithoutTx(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() - logChan := QueryLogger.Subscribe("TestExecutorSavepoint") - defer QueryLogger.Unsubscribe(logChan) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() + logChan := executor.queryLogger.Subscribe("TestExecutorSavepoint") + defer executor.queryLogger.Unsubscribe(logChan) session := NewSafeSession(&vtgatepb.Session{Autocommit: true, TargetString: "@primary", InTransaction: false}) _, err := exec(executor, session, "savepoint a") @@ -2468,18 +2681,23 @@ func TestExecutorSavepointWithoutTx(t *testing.T) { }} utils.MustMatch(t, sbc1WantQueries, sbc1.Queries, "") utils.MustMatch(t, sbc2WantQueries, sbc2.Queries, "") - testQueryLog(t, logChan, "TestExecute", "SAVEPOINT", "savepoint a", 0) - testQueryLog(t, logChan, "TestExecute", "SAVEPOINT_ROLLBACK", "rollback to a", 0) - testQueryLog(t, logChan, "TestExecute", "RELEASE", "release savepoint a", 0) - testQueryLog(t, logChan, "TestExecute", "SELECT", "select id from `user` where id = 1", 1) - testQueryLog(t, logChan, "TestExecute", "SAVEPOINT", "savepoint b", 0) - testQueryLog(t, logChan, "TestExecute", "SAVEPOINT_ROLLBACK", "rollback to b", 0) - testQueryLog(t, logChan, "TestExecute", "RELEASE", "release savepoint b", 0) - testQueryLog(t, logChan, "TestExecute", "SELECT", "select id from `user` where id = 3", 1) + testQueryLog(t, executor, logChan, "TestExecute", "SAVEPOINT", "savepoint a", 0) + testQueryLog(t, executor, logChan, "TestExecute", "SAVEPOINT_ROLLBACK", "rollback to a", 0) + testQueryLog(t, executor, logChan, "TestExecute", "RELEASE", "release savepoint a", 0) + testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select id from `user` where id = 1", 1) + testQueryLog(t, executor, logChan, "TestExecute", "SAVEPOINT", "savepoint b", 0) + testQueryLog(t, executor, logChan, "TestExecute", "SAVEPOINT_ROLLBACK", "rollback to b", 0) + testQueryLog(t, executor, logChan, "TestExecute", "RELEASE", "release savepoint b", 0) + testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select id from `user` where id = 3", 1) } func TestExecutorCallProc(t *testing.T) { - executor, sbc1, sbc2, sbcUnsharded := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, sbcUnsharded := createExecutorEnv(ctx) + defer executor.Close() type cnts struct { Sbc1Cnt int64 @@ -2544,11 +2762,15 @@ func TestExecutorCallProc(t *testing.T) { } func TestExecutorTempTable(t *testing.T) { - executor, _, _, sbcUnsharded := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, sbcUnsharded := createExecutorEnv(ctx) + defer executor.Close() executor.warnShardedOnly = true creatQuery := "create temporary table temp_t(id bigint primary key)" session := NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded}) - ctx := context.Background() _, err := executor.Execute(ctx, nil, "TestExecutorTempTable", session, creatQuery, nil) require.NoError(t, err) assert.EqualValues(t, 1, sbcUnsharded.ExecCount.Load()) @@ -2563,10 +2785,14 @@ func TestExecutorTempTable(t *testing.T) { } func TestExecutorShowVitessMigrations(t *testing.T) { - executor, sbc1, sbc2, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, sbc1, sbc2, _ := createExecutorEnv(ctx) + defer executor.Close() showQuery := "show vitess_migrations" session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) - ctx := context.Background() _, err := executor.Execute(ctx, nil, "", session, showQuery, nil) require.NoError(t, err) assert.Contains(t, sbc1.StringQueries(), "show vitess_migrations") @@ -2574,16 +2800,25 @@ func TestExecutorShowVitessMigrations(t *testing.T) { } func TestExecutorDescHash(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() showQuery := "desc hash_index" session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) - ctx := context.Background() _, err := executor.Execute(ctx, nil, "", session, showQuery, nil) require.NoError(t, err) } func TestExecutorVExplainQueries(t *testing.T) { - executor, _, _, sbclookup := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, sbclookup := createExecutorEnv(ctx) + defer executor.Close() session := NewAutocommitSession(&vtgatepb.Session{}) sbclookup.SetResults([]*sqltypes.Result{ @@ -2608,7 +2843,12 @@ func TestExecutorVExplainQueries(t *testing.T) { } func TestExecutorStartTxnStmt(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() session := NewAutocommitSession(&vtgatepb.Session{}) tcases := []struct { @@ -2653,7 +2893,12 @@ func TestExecutorStartTxnStmt(t *testing.T) { } func TestExecutorPrepareExecute(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() executor.normalize = true session := NewAutocommitSession(&vtgatepb.Session{}) @@ -2692,11 +2937,16 @@ func TestExecutorPrepareExecute(t *testing.T) { } func TestExecutorTruncateErrors(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + save := truncateErrorLen truncateErrorLen = 32 defer func() { truncateErrorLen = save }() - executor, _, _, _ := createExecutorEnv() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(&vtgatepb.Session{}) fn := func(r *sqltypes.Result) error { return nil @@ -2713,7 +2963,12 @@ func TestExecutorTruncateErrors(t *testing.T) { } func TestExecutorFlushStmt(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() tcs := []struct { targetStr string @@ -2770,7 +3025,12 @@ func TestExecutorFlushStmt(t *testing.T) { // TestExecutorKillStmt tests the kill statements on executor. func TestExecutorKillStmt(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() tcs := []struct { errStr string diff --git a/go/vt/vtgate/executor_vschema_ddl_test.go b/go/vt/vtgate/executor_vschema_ddl_test.go index a2a4cc34e31..73b4864cde0 100644 --- a/go/vt/vtgate/executor_vschema_ddl_test.go +++ b/go/vt/vtgate/executor_vschema_ddl_test.go @@ -95,7 +95,6 @@ func waitForVschemaTables(t *testing.T, ks string, tables []string, executor *Ex return nil } -// nolint func waitForColVindexes(t *testing.T, ks, table string, names []string, executor *Executor) *vschemapb.SrvVSchema { t.Helper() @@ -131,15 +130,19 @@ func waitForColVindexes(t *testing.T, ks, table string, names []string, executor } func TestPlanExecutorAlterVSchemaKeyspace(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() vschemaacl.AuthorizedDDLUsers = "%" defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, _, _, _ := createExecutorEnv() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) vschemaUpdates := make(chan *vschemapb.SrvVSchema, 2) - executor.serv.WatchSrvVSchema(context.Background(), "aa", func(vschema *vschemapb.SrvVSchema, err error) bool { + executor.serv.WatchSrvVSchema(ctx, "aa", func(vschema *vschemapb.SrvVSchema, err error) bool { vschemaUpdates <- vschema return true }) @@ -151,7 +154,7 @@ func TestPlanExecutorAlterVSchemaKeyspace(t *testing.T) { } stmt := "alter vschema create vindex TestExecutor.test_vindex using hash" - _, err := executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err := executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) _, vindex := waitForVindex(t, "TestExecutor", "test_vindex", vschemaUpdates, executor) @@ -159,15 +162,19 @@ func TestPlanExecutorAlterVSchemaKeyspace(t *testing.T) { } func TestPlanExecutorCreateVindexDDL(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() vschemaacl.AuthorizedDDLUsers = "%" defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, _, _, _ := createExecutorEnv() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() ks := "TestExecutor" vschemaUpdates := make(chan *vschemapb.SrvVSchema, 4) - executor.serv.WatchSrvVSchema(context.Background(), "aa", func(vschema *vschemapb.SrvVSchema, err error) bool { + executor.serv.WatchSrvVSchema(ctx, "aa", func(vschema *vschemapb.SrvVSchema, err error) bool { vschemaUpdates <- vschema return true }) @@ -180,7 +187,7 @@ func TestPlanExecutorCreateVindexDDL(t *testing.T) { session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) stmt := "alter vschema create vindex test_vindex using hash" - _, err := executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err := executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) _, vindex := waitForVindex(t, ks, "test_vindex", vschemaUpdates, executor) @@ -188,7 +195,7 @@ func TestPlanExecutorCreateVindexDDL(t *testing.T) { t.Errorf("updated vschema did not contain test_vindex") } - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) wantErr := "vindex test_vindex already exists in keyspace TestExecutor" if err == nil || err.Error() != wantErr { t.Errorf("create duplicate vindex: %v, want %s", err, wantErr) @@ -201,15 +208,19 @@ func TestPlanExecutorCreateVindexDDL(t *testing.T) { } func TestPlanExecutorDropVindexDDL(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() vschemaacl.AuthorizedDDLUsers = "%" defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, _, _, _ := createExecutorEnv() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() ks := "TestExecutor" vschemaUpdates := make(chan *vschemapb.SrvVSchema, 4) - executor.serv.WatchSrvVSchema(context.Background(), "aa", func(vschema *vschemapb.SrvVSchema, err error) bool { + executor.serv.WatchSrvVSchema(ctx, "aa", func(vschema *vschemapb.SrvVSchema, err error) bool { vschemaUpdates <- vschema return true }) @@ -222,14 +233,14 @@ func TestPlanExecutorDropVindexDDL(t *testing.T) { session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) stmt := "alter vschema drop vindex test_vindex" - _, err := executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err := executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) wantErr := "vindex test_vindex does not exists in keyspace TestExecutor" if err == nil || err.Error() != wantErr { t.Errorf("want error %v got %v", wantErr, err) } stmt = "alter vschema drop vindex TestExecutor.test_vindex" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) wantErr = "vindex test_vindex does not exists in keyspace TestExecutor" if err == nil || err.Error() != wantErr { t.Errorf("want error %v got %v", wantErr, err) @@ -237,7 +248,7 @@ func TestPlanExecutorDropVindexDDL(t *testing.T) { // add one vindex that has never been used by the tables stmt = "alter vschema create vindex test_vindex using hash" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) _, vindex := waitForVindex(t, ks, "test_vindex", vschemaUpdates, executor) @@ -247,7 +258,7 @@ func TestPlanExecutorDropVindexDDL(t *testing.T) { // drop an existing vindex that has never been used by the tables stmt = "alter vschema drop vindex TestExecutor.test_vindex" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) vschema = <-vschemaUpdates _, ok = vschema.Keyspaces[ks].Vindexes["test_vindex"] @@ -257,7 +268,7 @@ func TestPlanExecutorDropVindexDDL(t *testing.T) { // drop an existing vindex that is used by at least one table stmt = "alter vschema drop vindex TestExecutor.keyspace_id" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) wantErr = "can not drop vindex cause keyspace_id still defined on table ksid_table" if err == nil || err.Error() != wantErr { t.Errorf("drop vindex still defined: %v, want %s", err, wantErr) @@ -270,15 +281,19 @@ func TestPlanExecutorDropVindexDDL(t *testing.T) { } func TestPlanExecutorAddDropVschemaTableDDL(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() vschemaacl.AuthorizedDDLUsers = "%" defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, sbc1, sbc2, sbclookup := createExecutorEnv() + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() ks := KsTestUnsharded vschemaUpdates := make(chan *vschemapb.SrvVSchema, 4) - executor.serv.WatchSrvVSchema(context.Background(), "aa", func(vschema *vschemapb.SrvVSchema, err error) bool { + executor.serv.WatchSrvVSchema(ctx, "aa", func(vschema *vschemapb.SrvVSchema, err error) bool { vschemaUpdates <- vschema return true }) @@ -296,19 +311,19 @@ func TestPlanExecutorAddDropVschemaTableDDL(t *testing.T) { session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) stmt := "alter vschema add table test_table" - _, err := executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err := executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) _ = waitForVschemaTables(t, ks, append([]string{"test_table"}, vschemaTables...), executor) stmt = "alter vschema add table test_table2" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) _ = waitForVschemaTables(t, ks, append([]string{"test_table", "test_table2"}, vschemaTables...), executor) // Should fail adding a table on a sharded keyspace session = NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) stmt = "alter vschema add table test_table" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) wantErr := "add vschema table: unsupported on sharded keyspace TestExecutor" if err == nil || err.Error() != wantErr { t.Errorf("want error %v got %v", wantErr, err) @@ -327,11 +342,15 @@ func TestPlanExecutorAddDropVschemaTableDDL(t *testing.T) { } func TestExecutorAddSequenceDDL(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() vschemaacl.AuthorizedDDLUsers = "%" defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, _, _, _ := createExecutorEnv() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() ks := KsTestUnsharded vschema := executor.vm.GetCurrentSrvVschema() @@ -343,7 +362,7 @@ func TestExecutorAddSequenceDDL(t *testing.T) { session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) stmt := "alter vschema add sequence test_seq" - _, err := executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err := executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) _ = waitForVschemaTables(t, ks, append(vschemaTables, []string{"test_seq"}...), executor) vschema = executor.vm.GetCurrentSrvVschema() @@ -358,7 +377,7 @@ func TestExecutorAddSequenceDDL(t *testing.T) { session = NewSafeSession(&vtgatepb.Session{TargetString: ksSharded}) stmt = "alter vschema add sequence sequence_table" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) wantErr := "add sequence table: unsupported on sharded keyspace TestExecutor" if err == nil || err.Error() != wantErr { @@ -367,13 +386,13 @@ func TestExecutorAddSequenceDDL(t *testing.T) { // Should be able to add autoincrement to table in sharded keyspace stmt = "alter vschema on test_table add vindex hash_index (id)" - if _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil); err != nil { + if _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil); err != nil { t.Error(err) } time.Sleep(10 * time.Millisecond) stmt = "alter vschema on test_table add auto_increment id using `db-name`.`test_seq`" - if _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil); err != nil { + if _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil); err != nil { t.Error(err) } time.Sleep(10 * time.Millisecond) @@ -387,15 +406,19 @@ func TestExecutorAddSequenceDDL(t *testing.T) { } func TestExecutorAddDropVindexDDL(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() vschemaacl.AuthorizedDDLUsers = "%" defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, sbc1, sbc2, sbclookup := createExecutorEnv() // nolint + executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) + defer executor.Close() ks := "TestExecutor" session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) vschemaUpdates := make(chan *vschemapb.SrvVSchema, 4) - executor.serv.WatchSrvVSchema(context.Background(), "aa", func(vschema *vschemapb.SrvVSchema, err error) bool { + executor.serv.WatchSrvVSchema(ctx, "aa", func(vschema *vschemapb.SrvVSchema, err error) bool { vschemaUpdates <- vschema return true }) @@ -406,14 +429,14 @@ func TestExecutorAddDropVindexDDL(t *testing.T) { // Create a new vindex implicitly with the statement stmt := "alter vschema on test add vindex test_hash (id) using hash " - _, err := executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err := executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) _, vindex := waitForVindex(t, ks, "test_hash", vschemaUpdates, executor) require.Equal(t, "hash", vindex.Type) _ = waitForColVindexes(t, ks, "test", []string{"test_hash"}, executor) - qr, err := executor.Execute(context.Background(), nil, "TestExecute", session, "show vschema vindexes on TestExecutor.test", nil) + qr, err := executor.Execute(ctx, nil, "TestExecute", session, "show vschema vindexes on TestExecutor.test", nil) require.NoError(t, err) wantqr := &sqltypes.Result{ Fields: buildVarCharFields("Columns", "Name", "Type", "Params", "Owner"), @@ -425,17 +448,17 @@ func TestExecutorAddDropVindexDDL(t *testing.T) { // Drop it stmt = "alter vschema on test drop vindex test_hash" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) _, _ = waitForVindex(t, ks, "test_hash", vschemaUpdates, executor) _ = waitForColVindexes(t, ks, "test", []string{}, executor) - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, "show vschema vindexes on TestExecutor.test", nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, "show vschema vindexes on TestExecutor.test", nil) require.EqualError(t, err, "VT05005: table 'test' does not exist in keyspace 'TestExecutor'") // add it again using the same syntax stmt = "alter vschema on test add vindex test_hash (id) using hash " - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) _, vindex = waitForVindex(t, ks, "test_hash", vschemaUpdates, executor) @@ -443,7 +466,7 @@ func TestExecutorAddDropVindexDDL(t *testing.T) { _ = waitForColVindexes(t, ks, "test", []string{"test_hash"}, executor) - qr, err = executor.Execute(context.Background(), nil, "TestExecute", session, "show vschema vindexes on TestExecutor.test", nil) + qr, err = executor.Execute(ctx, nil, "TestExecute", session, "show vschema vindexes on TestExecutor.test", nil) require.NoError(t, err) wantqr = &sqltypes.Result{ Fields: buildVarCharFields("Columns", "Name", "Type", "Params", "Owner"), @@ -456,7 +479,7 @@ func TestExecutorAddDropVindexDDL(t *testing.T) { // add another stmt = "alter vschema on test add vindex test_lookup (c1,c2) using lookup with owner=`test`, from=`c1,c2`, table=test_lookup, to=keyspace_id" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) vschema, vindex = waitForVindex(t, ks, "test_lookup", vschemaUpdates, executor) @@ -473,7 +496,7 @@ func TestExecutorAddDropVindexDDL(t *testing.T) { t.Fatalf("table test not defined in vschema") } - qr, err = executor.Execute(context.Background(), nil, "TestExecute", session, "show vschema vindexes on TestExecutor.test", nil) + qr, err = executor.Execute(ctx, nil, "TestExecute", session, "show vschema vindexes on TestExecutor.test", nil) require.NoError(t, err) wantqr = &sqltypes.Result{ Fields: buildVarCharFields("Columns", "Name", "Type", "Params", "Owner"), @@ -485,7 +508,7 @@ func TestExecutorAddDropVindexDDL(t *testing.T) { utils.MustMatch(t, wantqr, qr) stmt = "alter vschema on test add vindex test_hash_id2 (id2) using hash" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) vschema, vindex = waitForVindex(t, ks, "test_hash_id2", vschemaUpdates, executor) @@ -502,7 +525,7 @@ func TestExecutorAddDropVindexDDL(t *testing.T) { t.Fatalf("table test not defined in vschema") } - qr, err = executor.Execute(context.Background(), nil, "TestExecute", session, "show vschema vindexes on TestExecutor.test", nil) + qr, err = executor.Execute(ctx, nil, "TestExecute", session, "show vschema vindexes on TestExecutor.test", nil) require.NoError(t, err) wantqr = &sqltypes.Result{ Fields: buildVarCharFields("Columns", "Name", "Type", "Params", "Owner"), @@ -516,13 +539,13 @@ func TestExecutorAddDropVindexDDL(t *testing.T) { // drop one stmt = "alter vschema on test drop vindex test_lookup" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) // wait for up to 50ms for it to disappear deadline := time.Now().Add(50 * time.Millisecond) for { - qr, err = executor.Execute(context.Background(), nil, "TestExecute", session, "show vschema vindexes on TestExecutor.test", nil) + qr, err = executor.Execute(ctx, nil, "TestExecute", session, "show vschema vindexes on TestExecutor.test", nil) require.NoError(t, err) wantqr = &sqltypes.Result{ Fields: buildVarCharFields("Columns", "Name", "Type", "Params", "Owner"), @@ -543,7 +566,7 @@ func TestExecutorAddDropVindexDDL(t *testing.T) { // use the newly created vindex on a new table stmt = "alter vschema on test2 add vindex test_hash (id)" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) vschema, vindex = waitForVindex(t, ks, "test_hash", vschemaUpdates, executor) @@ -556,7 +579,7 @@ func TestExecutorAddDropVindexDDL(t *testing.T) { // create an identical vindex definition on a different table stmt = "alter vschema on test2 add vindex test_lookup (c1,c2) using lookup with owner=`test`, from=`c1,c2`, table=test_lookup, to=keyspace_id" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) vschema, vindex = waitForVindex(t, ks, "test_lookup", vschemaUpdates, executor) @@ -567,7 +590,7 @@ func TestExecutorAddDropVindexDDL(t *testing.T) { require.Len(t, table.ColumnVindexes, 2) require.Equal(t, "test_lookup", table.ColumnVindexes[1].Name) - qr, err = executor.Execute(context.Background(), nil, "TestExecute", session, "show vschema vindexes on TestExecutor.test2", nil) + qr, err = executor.Execute(ctx, nil, "TestExecute", session, "show vschema vindexes on TestExecutor.test2", nil) require.NoError(t, err) wantqr = &sqltypes.Result{ Fields: buildVarCharFields("Columns", "Name", "Type", "Params", "Owner"), @@ -580,7 +603,7 @@ func TestExecutorAddDropVindexDDL(t *testing.T) { // now make sure we can create another vindex that references a table with dashes (i.e. escaping is necessary) stmt = "alter vschema on test2 add vindex test_lookup_fqn(c1,c2) using consistent_lookup_unique with owner=`test`, from=`c1,c2`, table=`test-keyspace`.`lookup-fqn`, to=`keyspace_id`" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) _, vindex = waitForVindex(t, ks, "test_lookup_fqn", vschemaUpdates, executor) @@ -591,35 +614,35 @@ func TestExecutorAddDropVindexDDL(t *testing.T) { require.Equal(t, "keyspace_id", vindex.Params["to"]) stmt = "alter vschema on test2 add vindex nonexistent (c1,c2)" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.EqualError(t, err, "vindex nonexistent does not exist in keyspace TestExecutor") stmt = "alter vschema on test2 add vindex test_hash (c1,c2) using lookup" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.EqualError(t, err, "vindex test_hash defined with type hash not lookup") stmt = "alter vschema on test2 add vindex test_lookup (c1,c2) using lookup with owner=xyz" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.EqualError(t, err, "vindex test_lookup defined with owner test not xyz") stmt = "alter vschema on test2 add vindex test_lookup (c1,c2) using lookup with owner=`test`, foo=bar" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.EqualError(t, err, "vindex test_lookup defined with different parameters") stmt = "alter vschema on nonexistent drop vindex test_lookup" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.EqualError(t, err, "table TestExecutor.nonexistent not defined in vschema") stmt = "alter vschema on nonexistent drop vindex test_lookup" - _, err = executor.Execute(context.Background(), nil, "TestExecute", NewSafeSession(&vtgatepb.Session{TargetString: "InvalidKeyspace"}), stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", NewSafeSession(&vtgatepb.Session{TargetString: "InvalidKeyspace"}), stmt, nil) require.EqualError(t, err, "VT05003: unknown database 'InvalidKeyspace' in vschema") stmt = "alter vschema on nowhere.nohow drop vindex test_lookup" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.EqualError(t, err, "VT05003: unknown database 'nowhere' in vschema") stmt = "alter vschema on test drop vindex test_lookup" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, stmt, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.EqualError(t, err, "vindex test_lookup not defined in table TestExecutor.test") // no queries should have gone to any tablets @@ -634,12 +657,16 @@ func TestExecutorAddDropVindexDDL(t *testing.T) { func TestPlanExecutorVindexDDLACL(t *testing.T) { // t.Skip("not yet planned") - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() ks := "TestExecutor" session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) - ctxRedUser := callerid.NewContext(context.Background(), &vtrpcpb.CallerID{}, &querypb.VTGateCallerID{Username: "redUser"}) - ctxBlueUser := callerid.NewContext(context.Background(), &vtrpcpb.CallerID{}, &querypb.VTGateCallerID{Username: "blueUser"}) + ctxRedUser := callerid.NewContext(ctx, &vtrpcpb.CallerID{}, &querypb.VTGateCallerID{Username: "redUser"}) + ctxBlueUser := callerid.NewContext(ctx, &vtrpcpb.CallerID{}, &querypb.VTGateCallerID{Username: "blueUser"}) // test that by default no users can perform the operation stmt := "alter vschema create vindex test_hash using hash" diff --git a/go/vt/vtgate/executor_vstream_test.go b/go/vt/vtgate/executor_vstream_test.go index b2176b7f38b..e7e514e8ef3 100644 --- a/go/vt/vtgate/executor_vstream_test.go +++ b/go/vt/vtgate/executor_vstream_test.go @@ -21,6 +21,7 @@ import ( "testing" "time" + "vitess.io/vitess/go/test/utils" "vitess.io/vitess/go/vt/vtgate/engine" querypb "vitess.io/vitess/go/vt/proto/query" @@ -38,9 +39,12 @@ import ( // TestVStreamSQLUnsharded tests the experimental 'vstream * from' vtgate olap query func TestVStreamSQLUnsharded(t *testing.T) { t.Skip("this test is failing due to races") // FIXME - executor, _, _, sbcLookup := createExecutorEnv() - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, sbcLookup := createExecutorEnv(ctx) + logChan := executor.queryLogger.Subscribe("Test") + defer executor.queryLogger.Unsubscribe(logChan) send1 := []*binlogdatapb.VEvent{ {Type: binlogdatapb.VEventType_GTID, Gtid: "gtid01"}, {Type: binlogdatapb.VEventType_FIELD, FieldEvent: &binlogdatapb.FieldEvent{TableName: "t1", Fields: []*querypb.Field{ @@ -76,8 +80,6 @@ func TestVStreamSQLUnsharded(t *testing.T) { sql := "vstream * from t1" results := make(chan *sqltypes.Result, 20) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() go func() { err := executor.StreamExecute(ctx, nil, "TestExecuteStream", NewAutocommitSession(&vtgatepb.Session{TargetString: KsTestUnsharded}), sql, nil, func(qr *sqltypes.Result) error { results <- qr diff --git a/go/vt/vtgate/legacy_scatter_conn_test.go b/go/vt/vtgate/legacy_scatter_conn_test.go index ee5e9a2f5f6..ede4e46303c 100644 --- a/go/vt/vtgate/legacy_scatter_conn_test.go +++ b/go/vt/vtgate/legacy_scatter_conn_test.go @@ -44,10 +44,13 @@ import ( // This file uses the sandbox_test framework. func TestLegacyExecuteFailOnAutocommit(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() createSandbox("TestExecuteFailOnAutocommit") hc := discovery.NewFakeHealthCheck(nil) - sc := newTestScatterConn(hc, newSandboxForCells([]string{"aa"}), "aa") + sc := newTestScatterConn(ctx, hc, newSandboxForCells(ctx, []string{"aa"}), "aa") sbc0 := hc.AddTestTablet("aa", "0", 1, "TestExecuteFailOnAutocommit", "0", topodatapb.TabletType_PRIMARY, true, 1, nil) sbc1 := hc.AddTestTablet("aa", "1", 1, "TestExecuteFailOnAutocommit", "1", topodatapb.TabletType_PRIMARY, true, 1, nil) @@ -108,7 +111,10 @@ func TestLegacyExecuteFailOnAutocommit(t *testing.T) { func TestScatterConnExecuteMulti(t *testing.T) { testScatterConnGeneric(t, "TestScatterConnExecuteMultiShard", func(sc *ScatterConn, shards []string) (*sqltypes.Result, error) { - res := srvtopo.NewResolver(newSandboxForCells([]string{"aa"}), sc.gateway, "aa") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + res := srvtopo.NewResolver(newSandboxForCells(ctx, []string{"aa"}), sc.gateway, "aa") rss, err := res.ResolveDestination(ctx, "TestScatterConnExecuteMultiShard", topodatapb.TabletType_REPLICA, key.DestinationShards(shards)) if err != nil { return nil, err @@ -129,7 +135,10 @@ func TestScatterConnExecuteMulti(t *testing.T) { func TestScatterConnStreamExecuteMulti(t *testing.T) { testScatterConnGeneric(t, "TestScatterConnStreamExecuteMulti", func(sc *ScatterConn, shards []string) (*sqltypes.Result, error) { - res := srvtopo.NewResolver(newSandboxForCells([]string{"aa"}), sc.gateway, "aa") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + res := srvtopo.NewResolver(newSandboxForCells(ctx, []string{"aa"}), sc.gateway, "aa") rss, err := res.ResolveDestination(ctx, "TestScatterConnStreamExecuteMulti", topodatapb.TabletType_REPLICA, key.DestinationShards(shards)) if err != nil { return nil, err @@ -156,11 +165,15 @@ func verifyScatterConnError(t *testing.T, err error, wantErr string, wantCode vt } func testScatterConnGeneric(t *testing.T, name string, f func(sc *ScatterConn, shards []string) (*sqltypes.Result, error)) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + hc := discovery.NewFakeHealthCheck(nil) // no shard s := createSandbox(name) - sc := newTestScatterConn(hc, newSandboxForCells([]string{"aa"}), "aa") + sc := newTestScatterConn(ctx, hc, newSandboxForCells(ctx, []string{"aa"}), "aa") qr, err := f(sc, nil) require.NoError(t, err) if qr.RowsAffected != 0 { @@ -169,7 +182,7 @@ func testScatterConnGeneric(t *testing.T, name string, f func(sc *ScatterConn, s // single shard s.Reset() - sc = newTestScatterConn(hc, newSandboxForCells([]string{"aa"}), "aa") + sc = newTestScatterConn(ctx, hc, newSandboxForCells(ctx, []string{"aa"}), "aa") sbc := hc.AddTestTablet("aa", "0", 1, name, "0", topodatapb.TabletType_REPLICA, true, 1, nil) sbc.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 1 _, err = f(sc, []string{"0"}) @@ -186,7 +199,7 @@ func testScatterConnGeneric(t *testing.T, name string, f func(sc *ScatterConn, s // two shards s.Reset() hc.Reset() - sc = newTestScatterConn(hc, newSandboxForCells([]string{"aa"}), "aa") + sc = newTestScatterConn(ctx, hc, newSandboxForCells(ctx, []string{"aa"}), "aa") sbc0 := hc.AddTestTablet("aa", "0", 1, name, "0", topodatapb.TabletType_REPLICA, true, 1, nil) sbc1 := hc.AddTestTablet("aa", "1", 1, name, "1", topodatapb.TabletType_REPLICA, true, 1, nil) sbc0.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 1 @@ -206,7 +219,7 @@ func testScatterConnGeneric(t *testing.T, name string, f func(sc *ScatterConn, s // two shards with different errors s.Reset() hc.Reset() - sc = newTestScatterConn(hc, newSandboxForCells([]string{"aa"}), "aa") + sc = newTestScatterConn(ctx, hc, newSandboxForCells(ctx, []string{"aa"}), "aa") sbc0 = hc.AddTestTablet("aa", "0", 1, name, "0", topodatapb.TabletType_REPLICA, true, 1, nil) sbc1 = hc.AddTestTablet("aa", "1", 1, name, "1", topodatapb.TabletType_REPLICA, true, 1, nil) sbc0.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 1 @@ -227,7 +240,7 @@ func testScatterConnGeneric(t *testing.T, name string, f func(sc *ScatterConn, s // duplicate shards s.Reset() hc.Reset() - sc = newTestScatterConn(hc, newSandboxForCells([]string{"aa"}), "aa") + sc = newTestScatterConn(ctx, hc, newSandboxForCells(ctx, []string{"aa"}), "aa") sbc = hc.AddTestTablet("aa", "0", 1, name, "0", topodatapb.TabletType_REPLICA, true, 1, nil) _, _ = f(sc, []string{"0", "0"}) // Ensure that we executed only once. @@ -238,7 +251,7 @@ func testScatterConnGeneric(t *testing.T, name string, f func(sc *ScatterConn, s // no errors s.Reset() hc.Reset() - sc = newTestScatterConn(hc, newSandboxForCells([]string{"aa"}), "aa") + sc = newTestScatterConn(ctx, hc, newSandboxForCells(ctx, []string{"aa"}), "aa") sbc0 = hc.AddTestTablet("aa", "0", 1, name, "0", topodatapb.TabletType_REPLICA, true, 1, nil) sbc1 = hc.AddTestTablet("aa", "1", 1, name, "1", topodatapb.TabletType_REPLICA, true, 1, nil) qr, err = f(sc, []string{"0", "1"}) @@ -260,17 +273,21 @@ func testScatterConnGeneric(t *testing.T, name string, f func(sc *ScatterConn, s } func TestMaxMemoryRows(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + save := maxMemoryRows maxMemoryRows = 3 defer func() { maxMemoryRows = save }() createSandbox("TestMaxMemoryRows") hc := discovery.NewFakeHealthCheck(nil) - sc := newTestScatterConn(hc, newSandboxForCells([]string{"aa"}), "aa") + sc := newTestScatterConn(ctx, hc, newSandboxForCells(ctx, []string{"aa"}), "aa") sbc0 := hc.AddTestTablet("aa", "0", 1, "TestMaxMemoryRows", "0", topodatapb.TabletType_REPLICA, true, 1, nil) sbc1 := hc.AddTestTablet("aa", "1", 1, "TestMaxMemoryRows", "1", topodatapb.TabletType_REPLICA, true, 1, nil) - res := srvtopo.NewResolver(newSandboxForCells([]string{"aa"}), sc.gateway, "aa") + res := srvtopo.NewResolver(newSandboxForCells(ctx, []string{"aa"}), sc.gateway, "aa") rss, _, err := res.ResolveDestinations(ctx, "TestMaxMemoryRows", topodatapb.TabletType_REPLICA, nil, []key.Destination{key.DestinationShard("0"), key.DestinationShard("1")}) require.NoError(t, err) @@ -315,12 +332,15 @@ func TestMaxMemoryRows(t *testing.T) { } func TestLegaceHealthCheckFailsOnReservedConnections(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() keyspace := "keyspace" createSandbox(keyspace) hc := discovery.NewFakeHealthCheck(nil) - sc := newTestScatterConn(hc, newSandboxForCells([]string{"aa"}), "aa") + sc := newTestScatterConn(ctx, hc, newSandboxForCells(ctx, []string{"aa"}), "aa") - res := srvtopo.NewResolver(newSandboxForCells([]string{"aa"}), sc.gateway, "aa") + res := srvtopo.NewResolver(newSandboxForCells(ctx, []string{"aa"}), sc.gateway, "aa") session := NewSafeSession(&vtgatepb.Session{InTransaction: false, InReservedConn: true}) destinations := []key.Destination{key.DestinationShard("0")} @@ -340,12 +360,12 @@ func TestLegaceHealthCheckFailsOnReservedConnections(t *testing.T) { require.Error(t, vterrors.Aggregate(errs)) } -func executeOnShards(t *testing.T, res *srvtopo.Resolver, keyspace string, sc *ScatterConn, session *SafeSession, destinations []key.Destination) { +func executeOnShards(t *testing.T, ctx context.Context, res *srvtopo.Resolver, keyspace string, sc *ScatterConn, session *SafeSession, destinations []key.Destination) { t.Helper() - require.Empty(t, executeOnShardsReturnsErr(t, res, keyspace, sc, session, destinations)) + require.Empty(t, executeOnShardsReturnsErr(t, ctx, res, keyspace, sc, session, destinations)) } -func executeOnShardsReturnsErr(t *testing.T, res *srvtopo.Resolver, keyspace string, sc *ScatterConn, session *SafeSession, destinations []key.Destination) error { +func executeOnShardsReturnsErr(t *testing.T, ctx context.Context, res *srvtopo.Resolver, keyspace string, sc *ScatterConn, session *SafeSession, destinations []key.Destination) error { t.Helper() rss, _, err := res.ResolveDestinations(ctx, keyspace, topodatapb.TabletType_REPLICA, nil, destinations) require.NoError(t, err) @@ -364,9 +384,12 @@ func executeOnShardsReturnsErr(t *testing.T, res *srvtopo.Resolver, keyspace str } func TestMultiExecs(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() createSandbox("TestMultiExecs") hc := discovery.NewFakeHealthCheck(nil) - sc := newTestScatterConn(hc, newSandboxForCells([]string{"aa"}), "aa") + sc := newTestScatterConn(ctx, hc, newSandboxForCells(ctx, []string{"aa"}), "aa") sbc0 := hc.AddTestTablet("aa", "0", 1, "TestMultiExecs", "0", topodatapb.TabletType_REPLICA, true, 1, nil) sbc1 := hc.AddTestTablet("aa", "1", 1, "TestMultiExecs", "1", topodatapb.TabletType_REPLICA, true, 1, nil) @@ -460,15 +483,18 @@ func TestMultiExecs(t *testing.T) { } func TestScatterConnSingleDB(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() createSandbox("TestScatterConnSingleDB") hc := discovery.NewFakeHealthCheck(nil) hc.Reset() - sc := newTestScatterConn(hc, newSandboxForCells([]string{"aa"}), "aa") + sc := newTestScatterConn(ctx, hc, newSandboxForCells(ctx, []string{"aa"}), "aa") hc.AddTestTablet("aa", "0", 1, "TestScatterConnSingleDB", "0", topodatapb.TabletType_PRIMARY, true, 1, nil) hc.AddTestTablet("aa", "1", 1, "TestScatterConnSingleDB", "1", topodatapb.TabletType_PRIMARY, true, 1, nil) - res := srvtopo.NewResolver(newSandboxForCells([]string{"aa"}), sc.gateway, "aa") + res := srvtopo.NewResolver(newSandboxForCells(ctx, []string{"aa"}), sc.gateway, "aa") rss0, err := res.ResolveDestination(ctx, "TestScatterConnSingleDB", topodatapb.TabletType_PRIMARY, key.DestinationShard("0")) require.NoError(t, err) rss1, err := res.ResolveDestination(ctx, "TestScatterConnSingleDB", topodatapb.TabletType_PRIMARY, key.DestinationShard("1")) @@ -553,10 +579,13 @@ func TestAppendResult(t *testing.T) { } func TestReservePrequeries(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() keyspace := "keyspace" createSandbox(keyspace) hc := discovery.NewFakeHealthCheck(nil) - sc := newTestScatterConn(hc, newSandboxForCells([]string{"aa"}), "aa") + sc := newTestScatterConn(ctx, hc, newSandboxForCells(ctx, []string{"aa"}), "aa") sbc0 := hc.AddTestTablet("aa", "0", 1, keyspace, "0", topodatapb.TabletType_REPLICA, true, 1, nil) sbc1 := hc.AddTestTablet("aa", "1", 1, keyspace, "1", topodatapb.TabletType_REPLICA, true, 1, nil) @@ -564,7 +593,7 @@ func TestReservePrequeries(t *testing.T) { sbc0.SetResults([]*sqltypes.Result{{}}) sbc1.SetResults([]*sqltypes.Result{{}}) - res := srvtopo.NewResolver(newSandboxForCells([]string{"aa"}), sc.gateway, "aa") + res := srvtopo.NewResolver(newSandboxForCells(ctx, []string{"aa"}), sc.gateway, "aa") session := NewSafeSession(&vtgatepb.Session{ InTransaction: false, @@ -576,11 +605,11 @@ func TestReservePrequeries(t *testing.T) { }) destinations := []key.Destination{key.DestinationShard("0")} - executeOnShards(t, res, keyspace, sc, session, destinations) + executeOnShards(t, ctx, res, keyspace, sc, session, destinations) assert.Equal(t, 1+1, len(sbc0.StringQueries())) } -func newTestScatterConn(hc discovery.HealthCheck, serv srvtopo.Server, cell string) *ScatterConn { +func newTestScatterConn(ctx context.Context, hc discovery.HealthCheck, serv srvtopo.Server, cell string) *ScatterConn { // The topo.Server is used to start watching the cells described // in '-cells_to_watch' command line parameter, which is // empty by default. So it's unused in this test, set to nil. @@ -588,5 +617,3 @@ func newTestScatterConn(hc discovery.HealthCheck, serv srvtopo.Server, cell stri tc := NewTxConn(gw, vtgatepb.TransactionMode_TWOPC) return NewScatterConn("", tc, gw) } - -var ctx = context.Background() diff --git a/go/vt/vtgate/mysql_protocol_test.go b/go/vt/vtgate/mysql_protocol_test.go deleted file mode 100644 index 75afa04f5fe..00000000000 --- a/go/vt/vtgate/mysql_protocol_test.go +++ /dev/null @@ -1,177 +0,0 @@ -/* -Copyright 2019 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 vtgate - -import ( - "context" - "net" - "strconv" - "testing" - - "vitess.io/vitess/go/test/utils" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - - "google.golang.org/protobuf/proto" - - "vitess.io/vitess/go/mysql" - "vitess.io/vitess/go/vt/vttablet/sandboxconn" - - querypb "vitess.io/vitess/go/vt/proto/query" - topodatapb "vitess.io/vitess/go/vt/proto/topodata" -) - -func TestMySQLProtocolExecute(t *testing.T) { - createSandbox(KsTestUnsharded) - hcVTGateTest.Reset() - sbc := hcVTGateTest.AddTestTablet("aa", "1.1.1.1", 1001, KsTestUnsharded, "0", topodatapb.TabletType_PRIMARY, true, 1, nil) - - c, err := mysqlConnect(&mysql.ConnParams{}) - if err != nil { - t.Fatal(err) - } - defer c.Close() - - qr, err := c.ExecuteFetch("select id from t1", 10, true /* wantfields */) - require.NoError(t, err) - utils.MustMatch(t, sandboxconn.SingleRowResult, qr, "mismatch in rows") - - options := &querypb.ExecuteOptions{ - IncludedFields: querypb.ExecuteOptions_ALL, - Workload: querypb.ExecuteOptions_OLTP, - } - if !proto.Equal(sbc.Options[0], options) { - t.Errorf("got ExecuteOptions \n%+v, want \n%+v", sbc.Options[0], options) - } -} - -func TestMySQLProtocolStreamExecute(t *testing.T) { - createSandbox(KsTestUnsharded) - hcVTGateTest.Reset() - sbc := hcVTGateTest.AddTestTablet("aa", "1.1.1.1", 1001, KsTestUnsharded, "0", topodatapb.TabletType_PRIMARY, true, 1, nil) - - c, err := mysqlConnect(&mysql.ConnParams{}) - if err != nil { - t.Fatal(err) - } - defer c.Close() - - _, err = c.ExecuteFetch("set workload='olap'", 1, true /* wantfields */) - require.NoError(t, err) - - qr, err := c.ExecuteFetch("select id from t1", 10, true /* wantfields */) - require.NoError(t, err) - utils.MustMatch(t, sandboxconn.SingleRowResult, qr, "mismatch in rows") - - options := &querypb.ExecuteOptions{ - IncludedFields: querypb.ExecuteOptions_ALL, - Workload: querypb.ExecuteOptions_OLAP, - } - if !proto.Equal(sbc.Options[0], options) { - t.Errorf("got ExecuteOptions \n%+v, want \n%+v", sbc.Options[0], options) - } -} - -func TestMySQLProtocolExecuteUseStatement(t *testing.T) { - createSandbox(KsTestUnsharded) - hcVTGateTest.Reset() - hcVTGateTest.AddTestTablet("aa", "1.1.1.1", 1001, KsTestUnsharded, "0", topodatapb.TabletType_PRIMARY, true, 1, nil) - - c, err := mysqlConnect(&mysql.ConnParams{DbName: "@primary"}) - if err != nil { - t.Fatal(err) - } - defer c.Close() - - qr, err := c.ExecuteFetch("select id from t1", 10, true /* wantfields */) - require.NoError(t, err) - utils.MustMatch(t, sandboxconn.SingleRowResult, qr) - - qr, err = c.ExecuteFetch("show vitess_target", 1, false) - require.NoError(t, err) - assert.Equal(t, "VARCHAR(\"@primary\")", qr.Rows[0][0].String()) - - _, err = c.ExecuteFetch("use TestUnsharded", 0, false) - require.NoError(t, err) - - qr, err = c.ExecuteFetch("select id from t1", 10, true /* wantfields */) - require.NoError(t, err) - utils.MustMatch(t, sandboxconn.SingleRowResult, qr) - - // No such keyspace this will fail - _, err = c.ExecuteFetch("use InvalidKeyspace", 0, false) - require.Error(t, err) - assert.Contains(t, err.Error(), "VT05003: unknown database 'InvalidKeyspace' in vschema (errno 1049) (sqlstate 42000)") - - // That doesn't reset the vitess_target - qr, err = c.ExecuteFetch("show vitess_target", 1, false) - require.NoError(t, err) - assert.Equal(t, "VARCHAR(\"TestUnsharded\")", qr.Rows[0][0].String()) - - _, err = c.ExecuteFetch("use @replica", 0, false) - require.NoError(t, err) - - // No replica tablets, this should also fail - _, err = c.ExecuteFetch("select id from t1", 10, true /* wantfields */) - require.Error(t, err) - assert.Contains(t, err.Error(), `no healthy tablet available for 'keyspace:"TestUnsharded" shard:"0" tablet_type:REPLICA`) -} - -func TestMysqlProtocolInvalidDB(t *testing.T) { - _, err := mysqlConnect(&mysql.ConnParams{DbName: "invalidDB"}) - require.EqualError(t, err, "VT05003: unknown database 'invalidDB' in vschema (errno 1049) (sqlstate 42000)") -} - -func TestMySQLProtocolClientFoundRows(t *testing.T) { - createSandbox(KsTestUnsharded) - hcVTGateTest.Reset() - sbc := hcVTGateTest.AddTestTablet("aa", "1.1.1.1", 1001, KsTestUnsharded, "0", topodatapb.TabletType_PRIMARY, true, 1, nil) - - c, err := mysqlConnect(&mysql.ConnParams{Flags: mysql.CapabilityClientFoundRows}) - if err != nil { - t.Fatal(err) - } - defer c.Close() - - qr, err := c.ExecuteFetch("select id from t1", 10, true /* wantfields */) - require.NoError(t, err) - utils.MustMatch(t, sandboxconn.SingleRowResult, qr) - - options := &querypb.ExecuteOptions{ - IncludedFields: querypb.ExecuteOptions_ALL, - ClientFoundRows: true, - Workload: querypb.ExecuteOptions_OLTP, - } - - if !proto.Equal(sbc.Options[0], options) { - t.Errorf("got ExecuteOptions \n%+v, want \n%+v", sbc.Options[0], options) - } -} - -// mysqlConnect fills the host & port into params and connects -// to the mysql protocol port. -func mysqlConnect(params *mysql.ConnParams) (*mysql.Conn, error) { - host, port, err := net.SplitHostPort(mysqlListener.Addr().String()) - if err != nil { - return nil, err - } - portnum, _ := strconv.Atoi(port) - params.Host = host - params.Port = portnum - return mysql.Connect(context.Background(), params) -} diff --git a/go/vt/vtgate/planbuilder/plan_test.go b/go/vt/vtgate/planbuilder/plan_test.go index c548f1c0ca9..b139d2573cc 100644 --- a/go/vt/vtgate/planbuilder/plan_test.go +++ b/go/vt/vtgate/planbuilder/plan_test.go @@ -55,6 +55,7 @@ func makeTestOutput(t *testing.T) string { } func TestPlan(t *testing.T) { + defer utils.EnsureNoLeaks(t) vschemaWrapper := &vschemawrapper.VSchemaWrapper{ V: loadSchema(t, "vschemas/schema.json", true), TabletType_: topodatapb.TabletType_PRIMARY, diff --git a/go/vt/vtgate/plugin_mysql_server.go b/go/vt/vtgate/plugin_mysql_server.go index 70f16ac3eea..436ed23755f 100644 --- a/go/vt/vtgate/plugin_mysql_server.go +++ b/go/vt/vtgate/plugin_mysql_server.go @@ -443,7 +443,7 @@ var sigChan chan os.Signal var vtgateHandle *vtgateHandler // initTLSConfig inits tls config for the given mysql listener -func initTLSConfig(mysqlListener *mysql.Listener, mysqlSslCert, mysqlSslKey, mysqlSslCa, mysqlSslCrl, mysqlSslServerCA string, mysqlServerRequireSecureTransport bool, mysqlMinTLSVersion uint16) error { +func initTLSConfig(ctx context.Context, mysqlListener *mysql.Listener, mysqlSslCert, mysqlSslKey, mysqlSslCa, mysqlSslCrl, mysqlSslServerCA string, mysqlServerRequireSecureTransport bool, mysqlMinTLSVersion uint16) error { serverConfig, err := vttls.ServerConfig(mysqlSslCert, mysqlSslKey, mysqlSslCa, mysqlSslCrl, mysqlSslServerCA, mysqlMinTLSVersion) if err != nil { log.Exitf("grpcutils.TLSServerConfig failed: %v", err) @@ -454,13 +454,18 @@ func initTLSConfig(mysqlListener *mysql.Listener, mysqlSslCert, mysqlSslKey, mys sigChan = make(chan os.Signal, 1) signal.Notify(sigChan, syscall.SIGHUP) go func() { - for range sigChan { - serverConfig, err := vttls.ServerConfig(mysqlSslCert, mysqlSslKey, mysqlSslCa, mysqlSslCrl, mysqlSslServerCA, mysqlMinTLSVersion) - if err != nil { - log.Errorf("grpcutils.TLSServerConfig failed: %v", err) - } else { - log.Info("grpcutils.TLSServerConfig updated") - mysqlListener.TLSConfig.Store(serverConfig) + for { + select { + case <-ctx.Done(): + return + case <-sigChan: + serverConfig, err := vttls.ServerConfig(mysqlSslCert, mysqlSslKey, mysqlSslCa, mysqlSslCrl, mysqlSslServerCA, mysqlMinTLSVersion) + if err != nil { + log.Errorf("grpcutils.TLSServerConfig failed: %v", err) + } else { + log.Info("grpcutils.TLSServerConfig updated") + mysqlListener.TLSConfig.Store(serverConfig) + } } } }() @@ -524,7 +529,7 @@ func initMySQLProtocol() { log.Exitf("mysql.NewListener failed: %v", err) } - _ = initTLSConfig(mysqlListener, mysqlSslCert, mysqlSslKey, mysqlSslCa, mysqlSslCrl, mysqlSslServerCA, mysqlServerRequireSecureTransport, tlsVersion) + _ = initTLSConfig(context.Background(), mysqlListener, mysqlSslCert, mysqlSslKey, mysqlSslCa, mysqlSslCrl, mysqlSslServerCA, mysqlServerRequireSecureTransport, tlsVersion) } mysqlListener.AllowClearTextWithoutTLS.Store(mysqlAllowClearTextWithoutTLS) // Check for the connection threshold diff --git a/go/vt/vtgate/plugin_mysql_server_test.go b/go/vt/vtgate/plugin_mysql_server_test.go index ed8d5b8c2fa..b068662a7e2 100644 --- a/go/vt/vtgate/plugin_mysql_server_test.go +++ b/go/vt/vtgate/plugin_mysql_server_test.go @@ -29,6 +29,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "vitess.io/vitess/go/test/utils" "vitess.io/vitess/go/mysql/replication" @@ -271,6 +272,10 @@ func TestInitTLSConfigWithServerCA(t *testing.T) { func testInitTLSConfig(t *testing.T, serverCA bool) { // Create the certs. + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + root := t.TempDir() tlstest.CreateCA(root) tlstest.CreateCRL(root, tlstest.CA) @@ -282,7 +287,7 @@ func testInitTLSConfig(t *testing.T, serverCA bool) { } listener := &mysql.Listener{} - if err := initTLSConfig(listener, path.Join(root, "server-cert.pem"), path.Join(root, "server-key.pem"), path.Join(root, "ca-cert.pem"), path.Join(root, "ca-crl.pem"), serverCACert, true, tls.VersionTLS12); err != nil { + if err := initTLSConfig(ctx, listener, path.Join(root, "server-cert.pem"), path.Join(root, "server-key.pem"), path.Join(root, "ca-cert.pem"), path.Join(root, "ca-crl.pem"), serverCACert, true, tls.VersionTLS12); err != nil { t.Fatalf("init tls config failure due to: +%v", err) } @@ -301,7 +306,11 @@ func testInitTLSConfig(t *testing.T, serverCA bool) { // TestKillMethods test the mysql plugin for kill method calls. func TestKillMethods(t *testing.T) { - executor, _, _, _ := createExecutorEnv() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() vh := newVtgateHandler(&VTGate{executor: executor}) // connection does not exist @@ -319,21 +328,21 @@ func TestKillMethods(t *testing.T) { // connection exists // updating context. - ctx, cancel := context.WithCancel(context.Background()) - mysqlConn.UpdateCancelCtx(cancel) + cancelCtx, cancelFunc := context.WithCancel(context.Background()) + mysqlConn.UpdateCancelCtx(cancelFunc) // kill query err = vh.KillQuery(1) assert.NoError(t, err) - require.EqualError(t, ctx.Err(), "context canceled") + require.EqualError(t, cancelCtx.Err(), "context canceled") // updating context. - ctx, cancel = context.WithCancel(context.Background()) - mysqlConn.UpdateCancelCtx(cancel) + cancelCtx, cancelFunc = context.WithCancel(context.Background()) + mysqlConn.UpdateCancelCtx(cancelFunc) // kill connection err = vh.KillConnection(context.Background(), 1) assert.NoError(t, err) - require.EqualError(t, ctx.Err(), "context canceled") + require.EqualError(t, cancelCtx.Err(), "context canceled") require.True(t, mysqlConn.IsMarkedForClose()) } diff --git a/go/vt/vtgate/querylog.go b/go/vt/vtgate/querylog.go index c501c5af2a4..e8a3ae9ac2b 100644 --- a/go/vt/vtgate/querylog.go +++ b/go/vt/vtgate/querylog.go @@ -18,7 +18,6 @@ package vtgate import ( "net/http" - "sync" "vitess.io/vitess/go/streamlog" "vitess.io/vitess/go/vt/servenv" @@ -34,25 +33,15 @@ var ( // QueryzHandler is the debug UI path for exposing query plan stats QueryzHandler = "/debug/queryz" - - // QueryLogger enables streaming logging of queries - QueryLogger *streamlog.StreamLogger[*logstats.LogStats] - queryLoggerMu sync.Mutex ) -func SetQueryLogger(logger *streamlog.StreamLogger[*logstats.LogStats]) { - queryLoggerMu.Lock() - defer queryLoggerMu.Unlock() - QueryLogger = logger -} - -func initQueryLogger(vtg *VTGate) error { - SetQueryLogger(streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize)) - QueryLogger.ServeLogs(QueryLogHandler, streamlog.GetFormatter(QueryLogger)) +func initQueryLogger(vtg *VTGate) (*streamlog.StreamLogger[*logstats.LogStats], error) { + queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize) + queryLogger.ServeLogs(QueryLogHandler, streamlog.GetFormatter(queryLogger)) servenv.HTTPHandleFunc(QueryLogzHandler, func(w http.ResponseWriter, r *http.Request) { - ch := QueryLogger.Subscribe("querylogz") - defer QueryLogger.Unsubscribe(ch) + ch := queryLogger.Subscribe("querylogz") + defer queryLogger.Unsubscribe(ch) querylogzHandler(ch, w, r) }) @@ -61,11 +50,11 @@ func initQueryLogger(vtg *VTGate) error { }) if queryLogToFile != "" { - _, err := QueryLogger.LogToFile(queryLogToFile, streamlog.GetFormatter(QueryLogger)) + _, err := queryLogger.LogToFile(queryLogToFile, streamlog.GetFormatter(queryLogger)) if err != nil { - return err + return nil, err } } - return nil + return queryLogger, nil } diff --git a/go/vt/vtgate/queryz_test.go b/go/vt/vtgate/queryz_test.go index cb03bafcbab..1f140887ad4 100644 --- a/go/vt/vtgate/queryz_test.go +++ b/go/vt/vtgate/queryz_test.go @@ -17,6 +17,7 @@ limitations under the License. package vtgate import ( + "context" "io" "net/http" "net/http/httptest" @@ -26,6 +27,7 @@ import ( "time" "github.com/stretchr/testify/require" + "vitess.io/vitess/go/test/utils" "vitess.io/vitess/go/sqltypes" "vitess.io/vitess/go/vt/vtgate/engine" @@ -34,14 +36,19 @@ import ( ) func TestQueryzHandler(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + resp := httptest.NewRecorder() req, _ := http.NewRequest("GET", "/schemaz", nil) - executor, _, _, _ := createExecutorEnv() + executor, _, _, _ := createExecutorEnv(ctx) + defer executor.Close() // single shard query sql := "select id from user where id = 1" - _, err := executorExec(executor, sql, nil) + _, err := executorExec(ctx, executor, sql, nil) require.NoError(t, err) executor.plans.Wait() plan1 := assertCacheContains(t, executor, nil, "select id from `user` where id = 1") @@ -49,14 +56,14 @@ func TestQueryzHandler(t *testing.T) { // scatter sql = "select id from user" - _, err = executorExec(executor, sql, nil) + _, err = executorExec(ctx, executor, sql, nil) require.NoError(t, err) executor.plans.Wait() plan2 := assertCacheContains(t, executor, nil, "select id from `user`") plan2.ExecTime = uint64(1 * time.Second) sql = "insert into user (id, name) values (:id, :name)" - _, err = executorExec(executor, sql, map[string]*querypb.BindVariable{ + _, err = executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{ "id": sqltypes.Uint64BindVariable(1), "name": sqltypes.BytesBindVariable([]byte("myname")), }) @@ -69,7 +76,7 @@ func TestQueryzHandler(t *testing.T) { // same query again should add query counts to existing plans sql = "insert into user (id, name) values (:id, :name)" - _, err = executorExec(executor, sql, map[string]*querypb.BindVariable{ + _, err = executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{ "id": sqltypes.Uint64BindVariable(1), "name": sqltypes.BytesBindVariable([]byte("myname")), }) diff --git a/go/vt/vtgate/sandbox_test.go b/go/vt/vtgate/sandbox_test.go index 91ae9589591..22c124020a1 100644 --- a/go/vt/vtgate/sandbox_test.go +++ b/go/vt/vtgate/sandbox_test.go @@ -220,9 +220,9 @@ type sandboxTopo struct { // the given cells. // // when this version is used, WatchSrvVSchema can properly simulate watches -func newSandboxForCells(cells []string) *sandboxTopo { +func newSandboxForCells(ctx context.Context, cells []string) *sandboxTopo { return &sandboxTopo{ - topoServer: memorytopo.NewServer(cells...), + topoServer: memorytopo.NewServer(ctx, cells...), } } @@ -314,6 +314,9 @@ func (sct *sandboxTopo) WatchSrvVSchema(ctx context.Context, cell string, callba if !callback(current.Value, nil) { panic("sandboxTopo callback returned false") } + if updateChan == nil { + panic("sandboxTopo updateChan is nil") + } currentHash := GetSrvVSchemaHash(current.Value) go func() { for { diff --git a/go/vt/vtgate/scatter_conn_test.go b/go/vt/vtgate/scatter_conn_test.go index 33b6798b187..a441227c6c2 100644 --- a/go/vt/vtgate/scatter_conn_test.go +++ b/go/vt/vtgate/scatter_conn_test.go @@ -17,6 +17,7 @@ limitations under the License. package vtgate import ( + "context" "testing" "vitess.io/vitess/go/mysql/sqlerror" @@ -42,10 +43,13 @@ import ( // This file uses the sandbox_test framework. func TestExecuteFailOnAutocommit(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() createSandbox("TestExecuteFailOnAutocommit") hc := discovery.NewFakeHealthCheck(nil) - sc := newTestScatterConn(hc, newSandboxForCells([]string{"aa"}), "aa") + sc := newTestScatterConn(ctx, hc, newSandboxForCells(ctx, []string{"aa"}), "aa") sbc0 := hc.AddTestTablet("aa", "0", 1, "TestExecuteFailOnAutocommit", "0", topodatapb.TabletType_PRIMARY, true, 1, nil) sbc1 := hc.AddTestTablet("aa", "1", 1, "TestExecuteFailOnAutocommit", "1", topodatapb.TabletType_PRIMARY, true, 1, nil) @@ -105,10 +109,14 @@ func TestExecuteFailOnAutocommit(t *testing.T) { } func TestReservedOnMultiReplica(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + keyspace := "keyspace" createSandbox(keyspace) hc := discovery.NewFakeHealthCheck(nil) - sc := newTestScatterConn(hc, newSandboxForCells([]string{"aa"}), "aa") + sc := newTestScatterConn(ctx, hc, newSandboxForCells(ctx, []string{"aa"}), "aa") sbc0_1 := hc.AddTestTablet("aa", "0", 1, keyspace, "0", topodatapb.TabletType_REPLICA, true, 1, nil) sbc0_2 := hc.AddTestTablet("aa", "2", 1, keyspace, "0", topodatapb.TabletType_REPLICA, true, 1, nil) // sbc1 := hc.AddTestTablet("aa", "1", 1, keyspace, "1", topodatapb.TabletType_REPLICA, true, 1, nil) @@ -117,18 +125,22 @@ func TestReservedOnMultiReplica(t *testing.T) { sbc0_1.SetResults([]*sqltypes.Result{{}}) sbc0_2.SetResults([]*sqltypes.Result{{}}) - res := srvtopo.NewResolver(newSandboxForCells([]string{"aa"}), sc.gateway, "aa") + res := srvtopo.NewResolver(newSandboxForCells(ctx, []string{"aa"}), sc.gateway, "aa") session := NewSafeSession(&vtgatepb.Session{InTransaction: false, InReservedConn: true}) destinations := []key.Destination{key.DestinationShard("0")} for i := 0; i < 10; i++ { - executeOnShards(t, res, keyspace, sc, session, destinations) + executeOnShards(t, ctx, res, keyspace, sc, session, destinations) assert.EqualValues(t, 1, sbc0_1.ReserveCount.Load()+sbc0_2.ReserveCount.Load(), "sbc0 reserve count") assert.EqualValues(t, 0, sbc0_1.BeginCount.Load()+sbc0_2.BeginCount.Load(), "sbc0 begin count") } } func TestReservedBeginTableDriven(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + type testAction struct { transaction, reserved bool shards []string @@ -253,7 +265,7 @@ func TestReservedBeginTableDriven(t *testing.T) { keyspace := "keyspace" createSandbox(keyspace) hc := discovery.NewFakeHealthCheck(nil) - sc := newTestScatterConn(hc, newSandboxForCells([]string{"aa"}), "aa") + sc := newTestScatterConn(ctx, hc, newSandboxForCells(ctx, []string{"aa"}), "aa") sbc0 := hc.AddTestTablet("aa", "0", 1, keyspace, "0", topodatapb.TabletType_REPLICA, true, 1, nil) sbc1 := hc.AddTestTablet("aa", "1", 1, keyspace, "1", topodatapb.TabletType_REPLICA, true, 1, nil) @@ -261,7 +273,7 @@ func TestReservedBeginTableDriven(t *testing.T) { sbc0.SetResults([]*sqltypes.Result{{}}) sbc1.SetResults([]*sqltypes.Result{{}}) - res := srvtopo.NewResolver(newSandboxForCells([]string{"aa"}), sc.gateway, "aa") + res := srvtopo.NewResolver(newSandboxForCells(ctx, []string{"aa"}), sc.gateway, "aa") t.Run(test.name, func(t *testing.T) { session := NewSafeSession(&vtgatepb.Session{}) @@ -272,7 +284,7 @@ func TestReservedBeginTableDriven(t *testing.T) { for _, shard := range action.shards { destinations = append(destinations, key.DestinationShard(shard)) } - executeOnShards(t, res, keyspace, sc, session, destinations) + executeOnShards(t, ctx, res, keyspace, sc, session, destinations) assert.EqualValues(t, action.sbc0Reserve, sbc0.ReserveCount.Load(), "sbc0 reserve count") assert.EqualValues(t, action.sbc0Begin, sbc0.BeginCount.Load(), "sbc0 begin count") assert.EqualValues(t, action.sbc1Reserve, sbc1.ReserveCount.Load(), "sbc1 reserve count") @@ -287,23 +299,27 @@ func TestReservedBeginTableDriven(t *testing.T) { } func TestReservedConnFail(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + keyspace := "keyspace" createSandbox(keyspace) hc := discovery.NewFakeHealthCheck(nil) - sc := newTestScatterConn(hc, newSandboxForCells([]string{"aa"}), "aa") + sc := newTestScatterConn(ctx, hc, newSandboxForCells(ctx, []string{"aa"}), "aa") sbc0 := hc.AddTestTablet("aa", "0", 1, keyspace, "0", topodatapb.TabletType_REPLICA, true, 1, nil) _ = hc.AddTestTablet("aa", "1", 1, keyspace, "1", topodatapb.TabletType_REPLICA, true, 1, nil) - res := srvtopo.NewResolver(newSandboxForCells([]string{"aa"}), sc.gateway, "aa") + res := srvtopo.NewResolver(newSandboxForCells(ctx, []string{"aa"}), sc.gateway, "aa") session := NewSafeSession(&vtgatepb.Session{InTransaction: false, InReservedConn: true}) destinations := []key.Destination{key.DestinationShard("0")} - executeOnShards(t, res, keyspace, sc, session, destinations) + executeOnShards(t, ctx, res, keyspace, sc, session, destinations) assert.Equal(t, 1, len(session.ShardSessions)) oldRId := session.Session.ShardSessions[0].ReservedId sbc0.EphemeralShardErr = sqlerror.NewSQLError(sqlerror.CRServerGone, sqlerror.SSUnknownSQLState, "lost connection") - _ = executeOnShardsReturnsErr(t, res, keyspace, sc, session, destinations) + _ = executeOnShardsReturnsErr(t, ctx, res, keyspace, sc, session, destinations) assert.Equal(t, 3, len(sbc0.Queries), "1 for the successful run, one for the failed attempt, and one for the retry") require.Equal(t, 1, len(session.ShardSessions)) assert.NotEqual(t, oldRId, session.Session.ShardSessions[0].ReservedId, "should have recreated a reserved connection since the last connection was lost") @@ -311,7 +327,7 @@ func TestReservedConnFail(t *testing.T) { sbc0.Queries = nil sbc0.EphemeralShardErr = sqlerror.NewSQLError(sqlerror.ERQueryInterrupted, sqlerror.SSUnknownSQLState, "transaction 123 not found") - _ = executeOnShardsReturnsErr(t, res, keyspace, sc, session, destinations) + _ = executeOnShardsReturnsErr(t, ctx, res, keyspace, sc, session, destinations) assert.Equal(t, 2, len(sbc0.Queries), "one for the failed attempt, and one for the retry") require.Equal(t, 1, len(session.ShardSessions)) assert.NotEqual(t, oldRId, session.Session.ShardSessions[0].ReservedId, "should have recreated a reserved connection since the last connection was lost") @@ -319,7 +335,7 @@ func TestReservedConnFail(t *testing.T) { sbc0.Queries = nil sbc0.EphemeralShardErr = sqlerror.NewSQLError(sqlerror.ERQueryInterrupted, sqlerror.SSUnknownSQLState, "transaction 123 ended at 2020-01-20") - _ = executeOnShardsReturnsErr(t, res, keyspace, sc, session, destinations) + _ = executeOnShardsReturnsErr(t, ctx, res, keyspace, sc, session, destinations) assert.Equal(t, 2, len(sbc0.Queries), "one for the failed attempt, and one for the retry") require.Equal(t, 1, len(session.ShardSessions)) assert.NotEqual(t, oldRId, session.Session.ShardSessions[0].ReservedId, "should have recreated a reserved connection since the last connection was lost") @@ -327,7 +343,7 @@ func TestReservedConnFail(t *testing.T) { sbc0.Queries = nil sbc0.EphemeralShardErr = sqlerror.NewSQLError(sqlerror.ERQueryInterrupted, sqlerror.SSUnknownSQLState, "transaction 123 in use: for tx killer rollback") - _ = executeOnShardsReturnsErr(t, res, keyspace, sc, session, destinations) + _ = executeOnShardsReturnsErr(t, ctx, res, keyspace, sc, session, destinations) assert.Equal(t, 2, len(sbc0.Queries), "one for the failed attempt, and one for the retry") require.Equal(t, 1, len(session.ShardSessions)) assert.NotEqual(t, oldRId, session.Session.ShardSessions[0].ReservedId, "should have recreated a reserved connection since the last connection was lost") @@ -335,7 +351,7 @@ func TestReservedConnFail(t *testing.T) { sbc0.Queries = nil sbc0.EphemeralShardErr = vterrors.New(vtrpcpb.Code_CLUSTER_EVENT, "operation not allowed in state NOT_SERVING during query: query1") - _ = executeOnShardsReturnsErr(t, res, keyspace, sc, session, destinations) + _ = executeOnShardsReturnsErr(t, ctx, res, keyspace, sc, session, destinations) assert.Equal(t, 2, len(sbc0.Queries), "one for the failed attempt, and one for the retry") require.Equal(t, 1, len(session.ShardSessions)) assert.NotEqual(t, oldRId, session.Session.ShardSessions[0].ReservedId, "should have recreated a reserved connection since the last connection was lost") @@ -343,7 +359,7 @@ func TestReservedConnFail(t *testing.T) { sbc0.Queries = nil sbc0.EphemeralShardErr = vterrors.New(vtrpcpb.Code_FAILED_PRECONDITION, "invalid tablet type: REPLICA, want: PRIMARY") - _ = executeOnShardsReturnsErr(t, res, keyspace, sc, session, destinations) + _ = executeOnShardsReturnsErr(t, ctx, res, keyspace, sc, session, destinations) assert.Equal(t, 2, len(sbc0.Queries), "one for the failed attempt, and one for the retry") require.Equal(t, 1, len(session.ShardSessions)) assert.NotEqual(t, oldRId, session.Session.ShardSessions[0].ReservedId, "should have recreated a reserved connection since the last connection was lost") @@ -364,7 +380,7 @@ func TestReservedConnFail(t *testing.T) { sbc0.Queries = nil sbc0.ExecCount.Store(0) - _ = executeOnShardsReturnsErr(t, res, keyspace, sc, session, destinations) + _ = executeOnShardsReturnsErr(t, ctx, res, keyspace, sc, session, destinations) assert.EqualValues(t, 1, sbc0.ExecCount.Load(), "first attempt should be made on original tablet") assert.EqualValues(t, 0, len(sbc0.Queries), "no query should be executed on it") assert.Equal(t, 1, len(sbc0Rep.Queries), "this attempt on new healthy tablet should pass") @@ -394,7 +410,7 @@ func TestReservedConnFail(t *testing.T) { sbc0Rep.Queries = nil sbc0Rep.ExecCount.Store(0) - _ = executeOnShardsReturnsErr(t, res, keyspace, sc, session, destinations) + _ = executeOnShardsReturnsErr(t, ctx, res, keyspace, sc, session, destinations) assert.EqualValues(t, 1, sbc0Rep.ExecCount.Load(), "first attempt should be made on the changed tablet type") assert.EqualValues(t, 0, len(sbc0Rep.Queries), "no query should be executed on it") assert.Equal(t, 1, len(sbc0.Queries), "this attempt should pass as it is on new healthy tablet and matches the target") diff --git a/go/vt/vtgate/tabletgateway.go b/go/vt/vtgate/tabletgateway.go index 7b833928e2e..f7caf6819e2 100644 --- a/go/vt/vtgate/tabletgateway.go +++ b/go/vt/vtgate/tabletgateway.go @@ -283,12 +283,12 @@ func (gw *TabletGateway) withRetry(ctx context.Context, target *querypb.Target, // if we have a keyspace event watcher, check if the reason why our primary is not available is that it's currently being resharded // or if a reparent operation is in progress. if kev := gw.kev; kev != nil { - if kev.TargetIsBeingResharded(target) { + if kev.TargetIsBeingResharded(ctx, target) { log.V(2).Infof("current keyspace is being resharded, retrying: %s: %s", target.Keyspace, debug.Stack()) err = vterrors.Errorf(vtrpcpb.Code_CLUSTER_EVENT, buffer.ClusterEventReshardingInProgress) continue } - primary, notServing := kev.PrimaryIsNotServing(target) + primary, notServing := kev.PrimaryIsNotServing(ctx, target) if notServing { err = vterrors.Errorf(vtrpcpb.Code_CLUSTER_EVENT, buffer.ClusterEventReparentInProgress) continue diff --git a/go/vt/vtgate/tabletgateway_flaky_test.go b/go/vt/vtgate/tabletgateway_flaky_test.go index 16ffb5febfb..e98dc65a0ce 100644 --- a/go/vt/vtgate/tabletgateway_flaky_test.go +++ b/go/vt/vtgate/tabletgateway_flaky_test.go @@ -22,6 +22,7 @@ import ( "time" "github.com/stretchr/testify/require" + "vitess.io/vitess/go/test/utils" "vitess.io/vitess/go/mysql/collations" @@ -35,6 +36,10 @@ import ( // TestGatewayBufferingWhenPrimarySwitchesServingState is used to test that the buffering mechanism buffers the queries when a primary goes to a non serving state and // stops buffering when the primary is healthy again func TestGatewayBufferingWhenPrimarySwitchesServingState(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + buffer.SetBufferingModeInTestingEnv(true) defer func() { buffer.SetBufferingModeInTestingEnv(false) @@ -55,7 +60,8 @@ func TestGatewayBufferingWhenPrimarySwitchesServingState(t *testing.T) { // create a new fake health check. We want to check the buffering code which uses Subscribe, so we must also pass a channel hc := discovery.NewFakeHealthCheck(make(chan *discovery.TabletHealth)) // create a new tablet gateway - tg := NewTabletGateway(context.Background(), hc, ts, "cell") + tg := NewTabletGateway(ctx, hc, ts, "cell") + defer tg.Close(ctx) // add a primary tabelt which is serving sbc := hc.AddTestTablet("cell", host, port, keyspace, shard, tabletType, true, 10, nil) @@ -75,7 +81,7 @@ func TestGatewayBufferingWhenPrimarySwitchesServingState(t *testing.T) { sbc.SetResults([]*sqltypes.Result{sqlResult1}) // run a query that we indeed get the result added to the sandbox connection back - res, err := tg.Execute(context.Background(), target, "query", nil, 0, 0, nil) + res, err := tg.Execute(ctx, target, "query", nil, 0, 0, nil) require.NoError(t, err) require.Equal(t, res, sqlResult1) @@ -93,7 +99,7 @@ func TestGatewayBufferingWhenPrimarySwitchesServingState(t *testing.T) { // execute the query in a go routine since it should be buffered, and check that it eventually succeed queryChan := make(chan struct{}) go func() { - res, err = tg.Execute(context.Background(), target, "query", nil, 0, 0, nil) + res, err = tg.Execute(ctx, target, "query", nil, 0, 0, nil) queryChan <- struct{}{} }() @@ -117,6 +123,10 @@ func TestGatewayBufferingWhenPrimarySwitchesServingState(t *testing.T) { // TestGatewayBufferingWhileReparenting is used to test that the buffering mechanism buffers the queries when a PRS happens // the healthchecks that happen during a PRS are simulated in this test func TestGatewayBufferingWhileReparenting(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + buffer.SetBufferingModeInTestingEnv(true) defer func() { buffer.SetBufferingModeInTestingEnv(false) @@ -139,7 +149,8 @@ func TestGatewayBufferingWhileReparenting(t *testing.T) { // create a new fake health check. We want to check the buffering code which uses Subscribe, so we must also pass a channel hc := discovery.NewFakeHealthCheck(make(chan *discovery.TabletHealth)) // create a new tablet gateway - tg := NewTabletGateway(context.Background(), hc, ts, "cell") + tg := NewTabletGateway(ctx, hc, ts, "cell") + defer tg.Close(ctx) // add a primary tabelt which is serving sbc := hc.AddTestTablet("cell", host, port, keyspace, shard, tabletType, true, 10, nil) @@ -162,7 +173,7 @@ func TestGatewayBufferingWhileReparenting(t *testing.T) { // run a query that we indeed get the result added to the sandbox connection back // this also checks that the query reaches the primary tablet and not the replica - res, err := tg.Execute(context.Background(), target, "query", nil, 0, 0, nil) + res, err := tg.Execute(ctx, target, "query", nil, 0, 0, nil) require.NoError(t, err) require.Equal(t, res, sqlResult1) @@ -191,7 +202,7 @@ func TestGatewayBufferingWhileReparenting(t *testing.T) { hc.Broadcast(primaryTablet) require.Len(t, tg.hc.GetHealthyTabletStats(target), 0, "GetHealthyTabletStats has tablets even though it shouldn't") - _, isNotServing := tg.kev.PrimaryIsNotServing(target) + _, isNotServing := tg.kev.PrimaryIsNotServing(ctx, target) require.True(t, isNotServing) // add a result to the sandbox connection of the new primary @@ -200,7 +211,7 @@ func TestGatewayBufferingWhileReparenting(t *testing.T) { // execute the query in a go routine since it should be buffered, and check that it eventually succeed queryChan := make(chan struct{}) go func() { - res, err = tg.Execute(context.Background(), target, "query", nil, 0, 0, nil) + res, err = tg.Execute(ctx, target, "query", nil, 0, 0, nil) queryChan <- struct{}{} }() @@ -222,7 +233,7 @@ outer: case <-timeout: require.Fail(t, "timed out - could not verify the new primary") case <-time.After(10 * time.Millisecond): - newPrimary, notServing := tg.kev.PrimaryIsNotServing(target) + newPrimary, notServing := tg.kev.PrimaryIsNotServing(ctx, target) if newPrimary != nil && newPrimary.Uid == 1 && !notServing { break outer } @@ -245,6 +256,10 @@ outer: // This is inconsistent and we want to fail properly. This scenario used to panic since no error and no results were // returned. func TestInconsistentStateDetectedBuffering(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + buffer.SetBufferingModeInTestingEnv(true) defer func() { buffer.SetBufferingModeInTestingEnv(false) @@ -265,7 +280,8 @@ func TestInconsistentStateDetectedBuffering(t *testing.T) { // create a new fake health check. We want to check the buffering code which uses Subscribe, so we must also pass a channel hc := discovery.NewFakeHealthCheck(make(chan *discovery.TabletHealth)) // create a new tablet gateway - tg := NewTabletGateway(context.Background(), hc, ts, "cell") + tg := NewTabletGateway(ctx, hc, ts, "cell") + defer tg.Close(ctx) tg.retryCount = 0 @@ -304,7 +320,7 @@ func TestInconsistentStateDetectedBuffering(t *testing.T) { var err error queryChan := make(chan struct{}) go func() { - res, err = tg.Execute(context.Background(), target, "query", nil, 0, 0, nil) + res, err = tg.Execute(ctx, target, "query", nil, 0, 0, nil) queryChan <- struct{}{} }() diff --git a/go/vt/vtgate/tabletgateway_test.go b/go/vt/vtgate/tabletgateway_test.go index 643ec7f9d26..986c4cd3dfd 100644 --- a/go/vt/vtgate/tabletgateway_test.go +++ b/go/vt/vtgate/tabletgateway_test.go @@ -24,6 +24,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "vitess.io/vitess/go/test/utils" "vitess.io/vitess/go/sqltypes" "vitess.io/vitess/go/vt/discovery" @@ -83,9 +84,14 @@ func TestTabletGatewayBeginExecute(t *testing.T) { } func TestTabletGatewayShuffleTablets(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + hc := discovery.NewFakeHealthCheck(nil) ts := &fakeTopoServer{} - tg := NewTabletGateway(context.Background(), hc, ts, "local") + tg := NewTabletGateway(ctx, hc, ts, "local") + defer tg.Close(ctx) ts1 := &discovery.TabletHealth{ Tablet: topo.NewTablet(1, "cell1", "host1"), @@ -142,6 +148,10 @@ func TestTabletGatewayShuffleTablets(t *testing.T) { } func TestTabletGatewayReplicaTransactionError(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + keyspace := "ks" shard := "0" // transactions on REPLICA are not allowed from tabletgateway @@ -156,7 +166,8 @@ func TestTabletGatewayReplicaTransactionError(t *testing.T) { } hc := discovery.NewFakeHealthCheck(nil) ts := &fakeTopoServer{} - tg := NewTabletGateway(context.Background(), hc, ts, "cell") + tg := NewTabletGateway(ctx, hc, ts, "cell") + defer tg.Close(ctx) _ = hc.AddTestTablet("cell", host, port, keyspace, shard, tabletType, true, 10, nil) _, err := tg.Execute(context.Background(), target, "query", nil, 1, 0, nil) @@ -165,6 +176,10 @@ func TestTabletGatewayReplicaTransactionError(t *testing.T) { func testTabletGatewayGeneric(t *testing.T, f func(tg *TabletGateway, target *querypb.Target) error) { t.Helper() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + keyspace := "ks" shard := "0" tabletType := topodatapb.TabletType_REPLICA @@ -177,7 +192,8 @@ func testTabletGatewayGeneric(t *testing.T, f func(tg *TabletGateway, target *qu } hc := discovery.NewFakeHealthCheck(nil) ts := &fakeTopoServer{} - tg := NewTabletGateway(context.Background(), hc, ts, "cell") + tg := NewTabletGateway(ctx, hc, ts, "cell") + defer tg.Close(ctx) // no tablet want := []string{"target: ks.0.replica", `no healthy tablet available for 'keyspace:"ks" shard:"0" tablet_type:REPLICA`} @@ -231,6 +247,10 @@ func testTabletGatewayGeneric(t *testing.T, f func(tg *TabletGateway, target *qu func testTabletGatewayTransact(t *testing.T, f func(tg *TabletGateway, target *querypb.Target) error) { t.Helper() + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + keyspace := "ks" shard := "0" // test with PRIMARY because replica transactions don't use gateway's queryservice @@ -245,7 +265,8 @@ func testTabletGatewayTransact(t *testing.T, f func(tg *TabletGateway, target *q } hc := discovery.NewFakeHealthCheck(nil) ts := &fakeTopoServer{} - tg := NewTabletGateway(context.Background(), hc, ts, "cell") + tg := NewTabletGateway(ctx, hc, ts, "cell") + defer tg.Close(ctx) // retry error - no retry sc1 := hc.AddTestTablet("cell", host, port, keyspace, shard, tabletType, true, 10, nil) diff --git a/go/vt/vtgate/tx_conn_test.go b/go/vt/vtgate/tx_conn_test.go index f7dff51accd..bb531a74c0e 100644 --- a/go/vt/vtgate/tx_conn_test.go +++ b/go/vt/vtgate/tx_conn_test.go @@ -43,7 +43,11 @@ var queries = []*querypb.BoundQuery{{Sql: "query1"}} var twoQueries = []*querypb.BoundQuery{{Sql: "query1"}, {Sql: "query1"}} func TestTxConnBegin(t *testing.T) { - sc, sbc0, _, rss0, _, _ := newTestTxConnEnv(t, "TestTxConn") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, _, rss0, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") session := &vtgatepb.Session{} // begin @@ -63,7 +67,11 @@ func TestTxConnBegin(t *testing.T) { } func TestTxConnCommitFailure(t *testing.T) { - sc, sbc0, sbc1, rss0, rss1, rss01 := newTestTxConnEnv(t, "TestTxConn") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, rss0, rss1, rss01 := newTestTxConnEnv(t, ctx, "TestTxConn") sc.txConn.mode = vtgatepb.TransactionMode_MULTI // Sequence the executes to ensure commit order @@ -120,7 +128,11 @@ func TestTxConnCommitFailure(t *testing.T) { } func TestTxConnCommitSuccess(t *testing.T) { - sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, "TestTxConn") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TestTxConn") sc.txConn.mode = vtgatepb.TransactionMode_MULTI // Sequence the executes to ensure commit order @@ -171,7 +183,11 @@ func TestTxConnCommitSuccess(t *testing.T) { } func TestTxConnReservedCommitSuccess(t *testing.T) { - sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, "TestTxConn") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TestTxConn") sc.txConn.mode = vtgatepb.TransactionMode_MULTI // Sequence the executes to ensure commit order @@ -253,8 +269,12 @@ func TestTxConnReservedCommitSuccess(t *testing.T) { } func TestTxConnReservedOn2ShardTxOn1ShardAndCommit(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + keyspace := "TestTxConn" - sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, keyspace) + sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, ctx, keyspace) sc.txConn.mode = vtgatepb.TransactionMode_MULTI // Sequence the executes to ensure shard session order @@ -346,8 +366,12 @@ func TestTxConnReservedOn2ShardTxOn1ShardAndCommit(t *testing.T) { } func TestTxConnReservedOn2ShardTxOn1ShardAndRollback(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + keyspace := "TestTxConn" - sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, keyspace) + sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, ctx, keyspace) sc.txConn.mode = vtgatepb.TransactionMode_MULTI // Sequence the executes to ensure shard session order @@ -439,7 +463,11 @@ func TestTxConnReservedOn2ShardTxOn1ShardAndRollback(t *testing.T) { } func TestTxConnCommitOrderFailure1(t *testing.T) { - sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, "TestTxConn") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, ctx, "TestTxConn") sc.txConn.mode = vtgatepb.TransactionMode_MULTI queries := []*querypb.BoundQuery{{Sql: "query1"}} @@ -470,7 +498,11 @@ func TestTxConnCommitOrderFailure1(t *testing.T) { } func TestTxConnCommitOrderFailure2(t *testing.T) { - sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, "TestTxConn") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, ctx, "TestTxConn") sc.txConn.mode = vtgatepb.TransactionMode_MULTI queries := []*querypb.BoundQuery{{ @@ -502,7 +534,11 @@ func TestTxConnCommitOrderFailure2(t *testing.T) { } func TestTxConnCommitOrderFailure3(t *testing.T) { - sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, "TestTxConn") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, ctx, "TestTxConn") sc.txConn.mode = vtgatepb.TransactionMode_MULTI queries := []*querypb.BoundQuery{{ @@ -542,7 +578,11 @@ func TestTxConnCommitOrderFailure3(t *testing.T) { } func TestTxConnCommitOrderSuccess(t *testing.T) { - sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, "TestTxConn") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, ctx, "TestTxConn") sc.txConn.mode = vtgatepb.TransactionMode_MULTI queries := []*querypb.BoundQuery{{ @@ -638,7 +678,11 @@ func TestTxConnCommitOrderSuccess(t *testing.T) { } func TestTxConnReservedCommitOrderSuccess(t *testing.T) { - sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, "TestTxConn") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, ctx, "TestTxConn") sc.txConn.mode = vtgatepb.TransactionMode_MULTI queries := []*querypb.BoundQuery{{ @@ -779,7 +823,11 @@ func TestTxConnReservedCommitOrderSuccess(t *testing.T) { } func TestTxConnCommit2PC(t *testing.T) { - sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, "TestTxConnCommit2PC") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PC") session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false) @@ -795,7 +843,11 @@ func TestTxConnCommit2PC(t *testing.T) { } func TestTxConnCommit2PCOneParticipant(t *testing.T) { - sc, sbc0, _, rss0, _, _ := newTestTxConnEnv(t, "TestTxConnCommit2PCOneParticipant") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, _, rss0, _, _ := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PCOneParticipant") session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false) session.TransactionMode = vtgatepb.TransactionMode_TWOPC @@ -805,7 +857,11 @@ func TestTxConnCommit2PCOneParticipant(t *testing.T) { } func TestTxConnCommit2PCCreateTransactionFail(t *testing.T) { - sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, "TestTxConnCommit2PCCreateTransactionFail") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PCCreateTransactionFail") session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false) @@ -827,7 +883,11 @@ func TestTxConnCommit2PCCreateTransactionFail(t *testing.T) { } func TestTxConnCommit2PCPrepareFail(t *testing.T) { - sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, "TestTxConnCommit2PCPrepareFail") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PCPrepareFail") session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false) @@ -847,7 +907,11 @@ func TestTxConnCommit2PCPrepareFail(t *testing.T) { } func TestTxConnCommit2PCStartCommitFail(t *testing.T) { - sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, "TestTxConnCommit2PCStartCommitFail") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PCStartCommitFail") session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false) @@ -867,7 +931,11 @@ func TestTxConnCommit2PCStartCommitFail(t *testing.T) { } func TestTxConnCommit2PCCommitPreparedFail(t *testing.T) { - sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, "TestTxConnCommit2PCCommitPreparedFail") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PCCommitPreparedFail") session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false) @@ -887,7 +955,11 @@ func TestTxConnCommit2PCCommitPreparedFail(t *testing.T) { } func TestTxConnCommit2PCConcludeTransactionFail(t *testing.T) { - sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, "TestTxConnCommit2PCConcludeTransactionFail") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PCConcludeTransactionFail") session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false) @@ -907,7 +979,11 @@ func TestTxConnCommit2PCConcludeTransactionFail(t *testing.T) { } func TestTxConnRollback(t *testing.T) { - sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, "TxConnRollback") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TxConnRollback") session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false) @@ -921,7 +997,11 @@ func TestTxConnRollback(t *testing.T) { } func TestTxConnReservedRollback(t *testing.T) { - sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, "TxConnReservedRollback") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TxConnReservedRollback") session := NewSafeSession(&vtgatepb.Session{InTransaction: true, InReservedConn: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false) @@ -956,7 +1036,11 @@ func TestTxConnReservedRollback(t *testing.T) { } func TestTxConnReservedRollbackFailure(t *testing.T) { - sc, sbc0, sbc1, rss0, rss1, rss01 := newTestTxConnEnv(t, "TxConnReservedRollback") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, rss0, rss1, rss01 := newTestTxConnEnv(t, ctx, "TxConnReservedRollback") session := NewSafeSession(&vtgatepb.Session{InTransaction: true, InReservedConn: true}) sc.ExecuteMultiShard(ctx, nil, rss0, queries, session, false, false) @@ -985,7 +1069,11 @@ func TestTxConnReservedRollbackFailure(t *testing.T) { } func TestTxConnResolveOnPrepare(t *testing.T) { - sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, "TestTxConn") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") dtid := "TestTxConn:0:1234" sbc0.ReadTransactionResults = []*querypb.TransactionMetadata{{ @@ -1006,7 +1094,11 @@ func TestTxConnResolveOnPrepare(t *testing.T) { } func TestTxConnResolveOnRollback(t *testing.T) { - sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, "TestTxConn") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") dtid := "TestTxConn:0:1234" sbc0.ReadTransactionResults = []*querypb.TransactionMetadata{{ @@ -1027,7 +1119,11 @@ func TestTxConnResolveOnRollback(t *testing.T) { } func TestTxConnResolveOnCommit(t *testing.T) { - sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, "TestTxConn") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") dtid := "TestTxConn:0:1234" sbc0.ReadTransactionResults = []*querypb.TransactionMetadata{{ @@ -1048,7 +1144,11 @@ func TestTxConnResolveOnCommit(t *testing.T) { } func TestTxConnResolveInvalidDTID(t *testing.T) { - sc, _, _, _, _, _ := newTestTxConnEnv(t, "TestTxConn") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, _, _, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") err := sc.txConn.Resolve(ctx, "abcd") want := "invalid parts in dtid: abcd" @@ -1056,7 +1156,11 @@ func TestTxConnResolveInvalidDTID(t *testing.T) { } func TestTxConnResolveReadTransactionFail(t *testing.T) { - sc, sbc0, _, _, _, _ := newTestTxConnEnv(t, "TestTxConn") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, _, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") dtid := "TestTxConn:0:1234" sbc0.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 1 @@ -1067,7 +1171,11 @@ func TestTxConnResolveReadTransactionFail(t *testing.T) { } func TestTxConnResolveInternalError(t *testing.T) { - sc, sbc0, _, _, _, _ := newTestTxConnEnv(t, "TestTxConn") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, _, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") dtid := "TestTxConn:0:1234" sbc0.ReadTransactionResults = []*querypb.TransactionMetadata{{ @@ -1086,7 +1194,11 @@ func TestTxConnResolveInternalError(t *testing.T) { } func TestTxConnResolveSetRollbackFail(t *testing.T) { - sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, "TestTxConn") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") dtid := "TestTxConn:0:1234" sbc0.ReadTransactionResults = []*querypb.TransactionMetadata{{ @@ -1110,7 +1222,11 @@ func TestTxConnResolveSetRollbackFail(t *testing.T) { } func TestTxConnResolveRollbackPreparedFail(t *testing.T) { - sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, "TestTxConn") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") dtid := "TestTxConn:0:1234" sbc0.ReadTransactionResults = []*querypb.TransactionMetadata{{ @@ -1134,7 +1250,11 @@ func TestTxConnResolveRollbackPreparedFail(t *testing.T) { } func TestTxConnResolveCommitPreparedFail(t *testing.T) { - sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, "TestTxConn") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") dtid := "TestTxConn:0:1234" sbc0.ReadTransactionResults = []*querypb.TransactionMetadata{{ @@ -1158,7 +1278,11 @@ func TestTxConnResolveCommitPreparedFail(t *testing.T) { } func TestTxConnResolveConcludeTransactionFail(t *testing.T) { - sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, "TestTxConn") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") dtid := "TestTxConn:0:1234" sbc0.ReadTransactionResults = []*querypb.TransactionMetadata{{ @@ -1182,6 +1306,10 @@ func TestTxConnResolveConcludeTransactionFail(t *testing.T) { } func TestTxConnMultiGoSessions(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + txc := &TxConn{} input := []*vtgatepb.Session_ShardSession{{ @@ -1249,7 +1377,11 @@ func TestTxConnMultiGoTargets(t *testing.T) { } func TestTxConnAccessModeReset(t *testing.T) { - sc, _, _, _, _, _ := newTestTxConnEnv(t, "TestTxConn") + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sc, _, _, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") tcases := []struct { name string @@ -1290,14 +1422,14 @@ func TestTxConnAccessModeReset(t *testing.T) { } } -func newTestTxConnEnv(t *testing.T, name string) (sc *ScatterConn, sbc0, sbc1 *sandboxconn.SandboxConn, rss0, rss1, rss01 []*srvtopo.ResolvedShard) { +func newTestTxConnEnv(t *testing.T, ctx context.Context, name string) (sc *ScatterConn, sbc0, sbc1 *sandboxconn.SandboxConn, rss0, rss1, rss01 []*srvtopo.ResolvedShard) { t.Helper() createSandbox(name) hc := discovery.NewFakeHealthCheck(nil) - sc = newTestScatterConn(hc, newSandboxForCells([]string{"aa"}), "aa") + sc = newTestScatterConn(ctx, hc, newSandboxForCells(ctx, []string{"aa"}), "aa") sbc0 = hc.AddTestTablet("aa", "0", 1, name, "0", topodatapb.TabletType_PRIMARY, true, 1, nil) sbc1 = hc.AddTestTablet("aa", "1", 1, name, "1", topodatapb.TabletType_PRIMARY, true, 1, nil) - res := srvtopo.NewResolver(newSandboxForCells([]string{"aa"}), sc.gateway, "aa") + res := srvtopo.NewResolver(newSandboxForCells(ctx, []string{"aa"}), sc.gateway, "aa") var err error rss0, err = res.ResolveDestination(ctx, name, topodatapb.TabletType_PRIMARY, key.DestinationShard("0")) require.NoError(t, err) diff --git a/go/vt/vtgate/vstream_manager_test.go b/go/vt/vtgate/vstream_manager_test.go index 926af63e9ac..5e51f0640c2 100644 --- a/go/vt/vtgate/vstream_manager_test.go +++ b/go/vt/vtgate/vstream_manager_test.go @@ -25,6 +25,7 @@ import ( "testing" "time" + "vitess.io/vitess/go/test/utils" "vitess.io/vitess/go/vt/topo" vtgatepb "vitess.io/vitess/go/vt/proto/vtgate" @@ -94,7 +95,7 @@ func TestVStreamSkew(t *testing.T) { var sbc0, sbc1 *sandboxconn.SandboxConn if tcase.shard0idx != 0 { sbc0 = hc.AddTestTablet(cell, "1.1.1.1", 1001, ks, "-20", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "-20", sbc0.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "-20", sbc0.Tablet()) sbc0.VStreamCh = make(chan *binlogdatapb.VEvent) want += 2 * tcase.numEventsPerShard vgtid.ShardGtids = append(vgtid.ShardGtids, &binlogdatapb.ShardGtid{Keyspace: ks, Gtid: "pos", Shard: "-20"}) @@ -102,7 +103,7 @@ func TestVStreamSkew(t *testing.T) { } if tcase.shard1idx != 0 { sbc1 = hc.AddTestTablet(cell, "1.1.1.1", 1002, ks, "20-40", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "20-40", sbc1.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "20-40", sbc1.Tablet()) sbc1.VStreamCh = make(chan *binlogdatapb.VEvent) want += 2 * tcase.numEventsPerShard vgtid.ShardGtids = append(vgtid.ShardGtids, &binlogdatapb.ShardGtid{Keyspace: ks, Gtid: "pos", Shard: "20-40"}) @@ -136,7 +137,7 @@ func TestVStreamEvents(t *testing.T) { vsm := newTestVStreamManager(hc, st, cell) sbc0 := hc.AddTestTablet(cell, "1.1.1.1", 1001, ks, "-20", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "-20", sbc0.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "-20", sbc0.Tablet()) send1 := []*binlogdatapb.VEvent{ {Type: binlogdatapb.VEventType_GTID, Gtid: "gtid01"}, @@ -213,9 +214,9 @@ func TestVStreamChunks(t *testing.T) { st := getSandboxTopo(ctx, cell, ks, []string{"-20", "20-40"}) vsm := newTestVStreamManager(hc, st, cell) sbc0 := hc.AddTestTablet("aa", "1.1.1.1", 1001, ks, "-20", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "-20", sbc0.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "-20", sbc0.Tablet()) sbc1 := hc.AddTestTablet("aa", "1.1.1.1", 1002, ks, "20-40", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "20-40", sbc1.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "20-40", sbc1.Tablet()) for i := 0; i < 100; i++ { sbc0.AddVStreamEvents([]*binlogdatapb.VEvent{{Type: binlogdatapb.VEventType_DDL}}, nil) @@ -281,9 +282,9 @@ func TestVStreamMulti(t *testing.T) { st := getSandboxTopo(ctx, cell, ks, []string{"-20", "20-40"}) vsm := newTestVStreamManager(hc, st, "aa") sbc0 := hc.AddTestTablet(cell, "1.1.1.1", 1001, ks, "-20", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "-20", sbc0.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "-20", sbc0.Tablet()) sbc1 := hc.AddTestTablet(cell, "1.1.1.1", 1002, ks, "20-40", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "20-40", sbc1.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "20-40", sbc1.Tablet()) send0 := []*binlogdatapb.VEvent{ {Type: binlogdatapb.VEventType_GTID, Gtid: "gtid01"}, @@ -345,9 +346,9 @@ func TestVStreamsCreatedAndLagMetrics(t *testing.T) { vsm.vstreamsCreated.ResetAll() vsm.vstreamsLag.ResetAll() sbc0 := hc.AddTestTablet(cell, "1.1.1.1", 1001, ks, "-20", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "-20", sbc0.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "-20", sbc0.Tablet()) sbc1 := hc.AddTestTablet(cell, "1.1.1.1", 1002, ks, "20-40", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "20-40", sbc1.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "20-40", sbc1.Tablet()) send0 := []*binlogdatapb.VEvent{ {Type: binlogdatapb.VEventType_GTID, Gtid: "gtid01"}, @@ -398,7 +399,7 @@ func TestVStreamRetry(t *testing.T) { st := getSandboxTopo(ctx, cell, ks, []string{"-20"}) vsm := newTestVStreamManager(hc, st, "aa") sbc0 := hc.AddTestTablet(cell, "1.1.1.1", 1001, ks, "-20", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "-20", sbc0.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "-20", sbc0.Tablet()) commit := []*binlogdatapb.VEvent{ {Type: binlogdatapb.VEventType_COMMIT}, } @@ -438,7 +439,7 @@ func TestVStreamShouldNotSendSourceHeartbeats(t *testing.T) { st := getSandboxTopo(ctx, cell, ks, []string{"-20"}) vsm := newTestVStreamManager(hc, st, cell) sbc0 := hc.AddTestTablet(cell, "1.1.1.1", 1001, ks, "-20", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "-20", sbc0.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "-20", sbc0.Tablet()) send0 := []*binlogdatapb.VEvent{ {Type: binlogdatapb.VEventType_HEARTBEAT}, @@ -488,11 +489,11 @@ func TestVStreamJournalOneToMany(t *testing.T) { st := getSandboxTopo(ctx, cell, ks, []string{"-20", "-10", "10-20"}) vsm := newTestVStreamManager(hc, st, "aa") sbc0 := hc.AddTestTablet(cell, "1.1.1.1", 1001, ks, "-20", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "-20", sbc0.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "-20", sbc0.Tablet()) sbc1 := hc.AddTestTablet(cell, "1.1.1.1", 1002, ks, "-10", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "-10", sbc1.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "-10", sbc1.Tablet()) sbc2 := hc.AddTestTablet(cell, "1.1.1.1", 1003, ks, "10-20", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "10-20", sbc2.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "10-20", sbc2.Tablet()) send1 := []*binlogdatapb.VEvent{ {Type: binlogdatapb.VEventType_GTID, Gtid: "gtid01"}, @@ -601,11 +602,11 @@ func TestVStreamJournalManyToOne(t *testing.T) { st := getSandboxTopo(ctx, cell, ks, []string{"-20", "-10", "10-20"}) vsm := newTestVStreamManager(hc, st, cell) sbc0 := hc.AddTestTablet(cell, "1.1.1.1", 1001, ks, "-20", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "-20", sbc0.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "-20", sbc0.Tablet()) sbc1 := hc.AddTestTablet(cell, "1.1.1.1", 1002, ks, "-10", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "-10", sbc1.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "-10", sbc1.Tablet()) sbc2 := hc.AddTestTablet(cell, "1.1.1.1", 1003, ks, "10-20", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "10-20", sbc2.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "10-20", sbc2.Tablet()) send3 := []*binlogdatapb.VEvent{ {Type: binlogdatapb.VEventType_GTID, Gtid: "gtid03"}, @@ -718,7 +719,7 @@ func TestVStreamJournalNoMatch(t *testing.T) { st := getSandboxTopo(ctx, cell, ks, []string{"-20"}) vsm := newTestVStreamManager(hc, st, "aa") sbc0 := hc.AddTestTablet("aa", "1.1.1.1", 1001, ks, "-20", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "-20", sbc0.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "-20", sbc0.Tablet()) send1 := []*binlogdatapb.VEvent{ {Type: binlogdatapb.VEventType_GTID, Gtid: "gtid01"}, @@ -847,9 +848,9 @@ func TestVStreamJournalPartialMatch(t *testing.T) { st := getSandboxTopo(ctx, cell, ks, []string{"-20", "-10", "10-20"}) vsm := newTestVStreamManager(hc, st, "aa") sbc1 := hc.AddTestTablet("aa", "1.1.1.1", 1002, ks, "-10", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "-10", sbc1.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "-10", sbc1.Tablet()) sbc2 := hc.AddTestTablet("aa", "1.1.1.1", 1003, ks, "10-20", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "10-20", sbc2.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "10-20", sbc2.Tablet()) send := []*binlogdatapb.VEvent{ {Type: binlogdatapb.VEventType_JOURNAL, Journal: &binlogdatapb.Journal{ @@ -922,10 +923,13 @@ func TestVStreamJournalPartialMatch(t *testing.T) { } func TestResolveVStreamParams(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + name := "TestVStream" _ = createSandbox(name) hc := discovery.NewFakeHealthCheck(nil) - vsm := newTestVStreamManager(hc, newSandboxForCells([]string{"aa"}), "aa") + vsm := newTestVStreamManager(hc, newSandboxForCells(ctx, []string{"aa"}), "aa") testcases := []struct { input *binlogdatapb.VGtid output *binlogdatapb.VGtid @@ -1133,6 +1137,10 @@ func TestResolveVStreamParams(t *testing.T) { } func TestVStreamIdleHeartbeat(t *testing.T) { + defer utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + cell := "aa" ks := "TestVStream" _ = createSandbox(ks) @@ -1140,7 +1148,7 @@ func TestVStreamIdleHeartbeat(t *testing.T) { st := getSandboxTopo(ctx, cell, ks, []string{"-20"}) vsm := newTestVStreamManager(hc, st, cell) sbc0 := hc.AddTestTablet("aa", "1.1.1.1", 1001, ks, "-20", topodatapb.TabletType_PRIMARY, true, 1, nil) - addTabletToSandboxTopo(t, st, ks, "-20", sbc0.Tablet()) + addTabletToSandboxTopo(t, ctx, st, ks, "-20", sbc0.Tablet()) vgtid := &binlogdatapb.VGtid{ ShardGtids: []*binlogdatapb.ShardGtid{{ Keyspace: ks, @@ -1194,6 +1202,7 @@ func newTestVStreamManager(hc discovery.HealthCheck, serv srvtopo.Server, cell s } func startVStream(ctx context.Context, t *testing.T, vsm *vstreamManager, vgtid *binlogdatapb.VGtid, flags *vtgatepb.VStreamFlags) <-chan *binlogdatapb.VStreamResponse { + t.Helper() if flags == nil { flags = &vtgatepb.VStreamFlags{} } @@ -1250,7 +1259,7 @@ func getVEvents(keyspace, shard string, count, idx int64) []*binlogdatapb.VEvent } func getSandboxTopo(ctx context.Context, cell string, keyspace string, shards []string) *sandboxTopo { - st := newSandboxForCells([]string{cell}) + st := newSandboxForCells(ctx, []string{cell}) ts := st.topoServer ts.CreateCellInfo(ctx, cell, &topodatapb.CellInfo{}) ts.CreateKeyspace(ctx, keyspace, &topodatapb.Keyspace{}) @@ -1260,7 +1269,7 @@ func getSandboxTopo(ctx context.Context, cell string, keyspace string, shards [] return st } -func addTabletToSandboxTopo(t *testing.T, st *sandboxTopo, ks, shard string, tablet *topodatapb.Tablet) { +func addTabletToSandboxTopo(t *testing.T, ctx context.Context, st *sandboxTopo, ks, shard string, tablet *topodatapb.Tablet) { _, err := st.topoServer.UpdateShardFields(ctx, ks, shard, func(si *topo.ShardInfo) error { si.PrimaryAlias = tablet.Alias return nil diff --git a/go/vt/vtgate/vtgate.go b/go/vt/vtgate/vtgate.go index eef9c11a061..61c029f928a 100644 --- a/go/vt/vtgate/vtgate.go +++ b/go/vt/vtgate/vtgate.go @@ -301,6 +301,11 @@ func Init( LFU: queryPlanCacheLFU, } + queryLogger, err := initQueryLogger(rpcVTGate) + if err != nil { + log.Fatalf("error initializing query logger: %v", err) + } + executor := NewExecutor( ctx, serv, @@ -313,6 +318,7 @@ func Init( si, noScatter, pv, + queryLogger, ) // connect the schema tracker with the vschema manager @@ -370,10 +376,6 @@ func Init( }) rpcVTGate.registerDebugHealthHandler() rpcVTGate.registerDebugEnvHandler() - err = initQueryLogger(rpcVTGate) - if err != nil { - log.Fatalf("error initializing query logger: %v", err) - } initAPI(gw.hc) return rpcVTGate From 0ad34ceafe566ca5be59a2af12d0d85450df424c Mon Sep 17 00:00:00 2001 From: Dirkjan Bussink Date: Mon, 21 Aug 2023 16:01:26 +0200 Subject: [PATCH 02/17] Remove global primary session Signed-off-by: Dirkjan Bussink --- go/vt/vtgate/executor_dml_test.go | 280 ++++++++++++++------ go/vt/vtgate/executor_framework_test.go | 20 +- go/vt/vtgate/executor_select_test.go | 330 +++++++++++++++++------- go/vt/vtgate/executor_set_test.go | 43 ++- go/vt/vtgate/executor_stream_test.go | 4 +- go/vt/vtgate/executor_test.go | 23 +- go/vt/vtgate/queryz_test.go | 10 +- go/vt/vtgate/vtgate_test.go | 61 +---- 8 files changed, 493 insertions(+), 278 deletions(-) diff --git a/go/vt/vtgate/executor_dml_test.go b/go/vt/vtgate/executor_dml_test.go index b376349826b..6cbcf922c13 100644 --- a/go/vt/vtgate/executor_dml_test.go +++ b/go/vt/vtgate/executor_dml_test.go @@ -48,7 +48,10 @@ func TestUpdateEqual(t *testing.T) { defer executor.queryLogger.Unsubscribe(logChan) // Update by primary vindex. - _, err := executorExec(ctx, executor, "update user set a=2 where id = 1", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "update user set a=2 where id = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "update `user` set a = 2 where id = 1", @@ -59,7 +62,7 @@ func TestUpdateEqual(t *testing.T) { testQueryLog(t, executor, logChan, "TestExecute", "UPDATE", "update `user` set a = 2 where id = 1", 1) sbc1.Queries = nil - _, err = executorExec(ctx, executor, "update user set a=2 where id = 3", nil) + _, err = executorExec(ctx, executor, session, "update user set a=2 where id = 3", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "update `user` set a = 2 where id = 3", @@ -72,7 +75,7 @@ func TestUpdateEqual(t *testing.T) { sbc1.Queries = nil sbc2.Queries = nil sbclookup.SetResults([]*sqltypes.Result{{}}) - _, err = executorExec(ctx, executor, "update music set a=2 where id = 2", nil) + _, err = executorExec(ctx, executor, session, "update music set a=2 where id = 2", nil) require.NoError(t, err) vars, err := sqltypes.BuildBindVariable([]any{sqltypes.NewInt64(2)}) require.NoError(t, err) @@ -96,7 +99,7 @@ func TestUpdateEqual(t *testing.T) { ), }) - _, err = executorExec(ctx, executor, "update user2 set `name`='myname', lastname='mylastname' where id = 1", nil) + _, err = executorExec(ctx, executor, session, "update user2 set `name`='myname', lastname='mylastname' where id = 1", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{ { @@ -154,7 +157,10 @@ func TestUpdateFromSubQuery(t *testing.T) { }}) // Update by primary vindex, but first execute subquery - _, err := executorExec(ctx, executor, "update user set a=(select count(*) from user where id = 3) where id = 1", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "update user set a=(select count(*) from user where id = 3) where id = 1", nil) require.NoError(t, err) wantQueriesSbc1 := []*querypb.BoundQuery{{ Sql: "update `user` set a = :__sq1 where id = 1", @@ -185,7 +191,10 @@ func TestUpdateEqualWithNoVerifyAndWriteOnlyLookupUniqueVindexes(t *testing.T) { executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, res) defer executor.Close() - _, err := executorExec(ctx, executor, "update t2_lookup set lu_col = 5 where wo_lu_col = 2", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "update t2_lookup set lu_col = 5 where wo_lu_col = 2", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{ { @@ -562,7 +571,10 @@ func TestUpdateMultiOwned(t *testing.T) { "1|10|20|30|40|50|60|0|0", ), }) - _, err := executorExec(ctx, executor, "update user set a=1, b=2, f=4, e=3 where id=1", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "update user set a=1, b=2, f=4, e=3 where id=1", nil) if err != nil { t.Fatal(err) } @@ -616,7 +628,10 @@ func TestUpdateComments(t *testing.T) { executor, sbc1, sbc2, _ := createExecutorEnv(ctx) defer executor.Close() - _, err := executorExec(ctx, executor, "update user set a=2 where id = 1 /* trailing */", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "update user set a=2 where id = 1 /* trailing */", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "update `user` set a = 2 where id = 1 /* trailing */", @@ -634,7 +649,10 @@ func TestUpdateNormalize(t *testing.T) { defer executor.Close() executor.normalize = true - _, err := executorExec(ctx, executor, "/* leading */ update user set a=2 where id = 1 /* trailing */", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "/* leading */ update user set a=2 where id = 1 /* trailing */", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "/* leading */ update `user` set a = :a /* INT64 */ where id = :id /* INT64 */ /* trailing */", @@ -648,8 +666,8 @@ func TestUpdateNormalize(t *testing.T) { sbc1.Queries = nil // Force the query to go to the "wrong" shard and ensure that normalization still happens - primarySession.TargetString = "TestExecutor/40-60" - _, err = executorExec(ctx, executor, "/* leading */ update user set a=2 where id = 1 /* trailing */", nil) + session.TargetString = "TestExecutor/40-60" + _, err = executorExec(ctx, executor, session, "/* leading */ update user set a=2 where id = 1 /* trailing */", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "/* leading */ update `user` set a = :a /* INT64 */ where id = :id /* INT64 */ /* trailing */", @@ -661,7 +679,6 @@ func TestUpdateNormalize(t *testing.T) { assertQueries(t, sbc1, nil) assertQueries(t, sbc2, wantQueries) sbc2.Queries = nil - primarySession.TargetString = "" } func TestDeleteEqual(t *testing.T) { @@ -683,7 +700,10 @@ func TestDeleteEqual(t *testing.T) { sqltypes.NewVarChar("myname"), }}, }}) - _, err := executorExec(ctx, executor, "delete from user where id = 1", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "delete from user where id = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select Id, `name` from `user` where id = 1 for update", @@ -706,7 +726,7 @@ func TestDeleteEqual(t *testing.T) { sbc.Queries = nil sbclookup.Queries = nil sbc.SetResults([]*sqltypes.Result{{}}) - _, err = executorExec(ctx, executor, "delete from user where id = 1", nil) + _, err = executorExec(ctx, executor, session, "delete from user where id = 1", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "select Id, `name` from `user` where id = 1 for update", @@ -721,7 +741,7 @@ func TestDeleteEqual(t *testing.T) { sbc.Queries = nil sbclookup.Queries = nil sbclookup.SetResults([]*sqltypes.Result{{}}) - _, err = executorExec(ctx, executor, "delete from music where id = 1", nil) + _, err = executorExec(ctx, executor, session, "delete from music where id = 1", nil) require.NoError(t, err) vars, err := sqltypes.BuildBindVariable([]any{sqltypes.NewInt64(1)}) require.NoError(t, err) @@ -737,7 +757,7 @@ func TestDeleteEqual(t *testing.T) { sbc.Queries = nil sbclookup.Queries = nil sbclookup.SetResults([]*sqltypes.Result{{}}) - _, err = executorExec(ctx, executor, "delete from user_extra where user_id = 1", nil) + _, err = executorExec(ctx, executor, session, "delete from user_extra where user_id = 1", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "delete from user_extra where user_id = 1", @@ -753,7 +773,7 @@ func TestDeleteEqual(t *testing.T) { "1|1|foo", ), }) - _, err = executorExec(ctx, executor, "delete from user2 where id = 1", nil) + _, err = executorExec(ctx, executor, session, "delete from user2 where id = 1", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{ { @@ -787,7 +807,10 @@ func TestUpdateScatter(t *testing.T) { defer cancel() executor, sbc1, sbc2, _ := createExecutorEnv(ctx) defer executor.Close() - _, err := executorExec(ctx, executor, "update user_extra set col = 2", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "update user_extra set col = 2", nil) require.NoError(t, err) // Queries get annotatted. wantQueries := []*querypb.BoundQuery{{ @@ -804,7 +827,10 @@ func TestDeleteScatter(t *testing.T) { defer cancel() executor, sbc1, sbc2, _ := createExecutorEnv(ctx) defer executor.Close() - _, err := executorExec(ctx, executor, "delete from user_extra", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "delete from user_extra", nil) require.NoError(t, err) // Queries get annotatted. wantQueries := []*querypb.BoundQuery{{ @@ -835,7 +861,10 @@ func TestUpdateEqualWithMultipleLookupVindex(t *testing.T) { "1|2|2|2|2|2|1|0", )}) - _, err := executorExec(ctx, executor, "update t2_lookup set lu_col = 5 where wo_lu_col = 2 and lu_col = 1", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "update t2_lookup set lu_col = 5 where wo_lu_col = 2 and lu_col = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{ { @@ -894,7 +923,10 @@ func TestUpdateUseHigherCostVindexIfBackfilling(t *testing.T) { "1|2|2|2|2|2|2|0", )}) - _, err := executorExec(ctx, executor, "update t2_lookup set lu_col = 5 where wo_lu_col = 2 and lu_col in (1, 2)", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "update t2_lookup set lu_col = 5 where wo_lu_col = 2 and lu_col in (1, 2)", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{ { @@ -958,7 +990,10 @@ func TestDeleteEqualWithNoVerifyAndWriteOnlyLookupUniqueVindex(t *testing.T) { executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, res) defer executor.Close() - _, err := executorExec(ctx, executor, "delete from t2_lookup where wo_lu_col = 1", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "delete from t2_lookup where wo_lu_col = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{ { @@ -1046,7 +1081,10 @@ func TestDeleteEqualWithMultipleLookupVindex(t *testing.T) { "1|1|1|1|1|1|1", )}) - _, err := executorExec(ctx, executor, "delete from t2_lookup where wo_lu_col = 1 and lu_col = 1", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "delete from t2_lookup where wo_lu_col = 1 and lu_col = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{ { @@ -1130,7 +1168,10 @@ func TestDeleteUseHigherCostVindexIfBackfilling(t *testing.T) { "1|1|1|1|1|1|2", )}) - _, err := executorExec(ctx, executor, "delete from t2_lookup where wo_lu_col = 1 and lu_col in (1, 2)", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "delete from t2_lookup where wo_lu_col = 1 and lu_col in (1, 2)", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{ { @@ -1235,7 +1276,10 @@ func TestDeleteByDestination(t *testing.T) { defer cancel() executor, sbc1, sbc2, _ := createExecutorEnv(ctx) defer executor.Close() - _, err := executorExec(ctx, executor, "delete from `TestExecutor[-]`.user_extra limit 10", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "delete from `TestExecutor[-]`.user_extra limit 10", nil) require.NoError(t, err) // Queries get annotatted. wantQueries := []*querypb.BoundQuery{{ @@ -1265,7 +1309,10 @@ func TestDeleteComments(t *testing.T) { sqltypes.NewVarChar("myname"), }}, }}) - _, err := executorExec(ctx, executor, "delete from user where id = 1 /* trailing */", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "delete from user where id = 1 /* trailing */", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select Id, `name` from `user` where id = 1 for update /* trailing */", @@ -1296,7 +1343,10 @@ func TestInsertSharded(t *testing.T) { logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) - _, err := executorExec(ctx, executor, "insert into user(id, v, name) values (1, 2, 'myname')", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "insert into user(id, v, name) values (1, 2, 'myname')", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into `user`(id, v, `name`) values (:_Id_0, 2, :_name_0)", @@ -1323,7 +1373,7 @@ func TestInsertSharded(t *testing.T) { sbc1.Queries = nil sbclookup.Queries = nil - _, err = executorExec(ctx, executor, "insert into user(id, v, name) values (3, 2, 'myname2')", nil) + _, err = executorExec(ctx, executor, session, "insert into user(id, v, name) values (3, 2, 'myname2')", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "insert into `user`(id, v, `name`) values (:_Id_0, 2, :_name_0)", @@ -1348,7 +1398,7 @@ func TestInsertSharded(t *testing.T) { testQueryLog(t, executor, logChan, "TestExecute", "INSERT", "insert into `user`(id, v, `name`) values (3, 2, 'myname2')", 1) sbc1.Queries = nil - _, err = executorExec(ctx, executor, "insert into user2(id, name, lastname) values (2, 'myname', 'mylastname')", nil) + _, err = executorExec(ctx, executor, session, "insert into user2(id, name, lastname) values (2, 'myname', 'mylastname')", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "insert into user2(id, `name`, lastname) values (:_id_0, :_name_0, :_lastname_0)", @@ -1368,7 +1418,7 @@ func TestInsertSharded(t *testing.T) { sbc1.Queries = nil sbc2.Queries = nil sbclookup.Queries = nil - _, err = executorExec(ctx, executor, "insert into user(id, v, name) values (1, 2, _binary 'myname')", nil) + _, err = executorExec(ctx, executor, session, "insert into user(id, v, name) values (1, 2, _binary 'myname')", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "insert into `user`(id, v, `name`) values (:_Id_0, :vtg2 /* INT64 */, :_name_0)", @@ -1405,7 +1455,10 @@ func TestInsertShardedKeyrange(t *testing.T) { defer executor.Close() // If a unique vindex returns a keyrange, we fail the insert - _, err := executorExec(ctx, executor, "insert into keyrange_table(krcol_unique, krcol) values(1, 1)", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "insert into keyrange_table(krcol_unique, krcol) values(1, 1)", nil) require.EqualError(t, err, "could not map [INT64(1)] to a unique keyspace id: DestinationKeyRange(-10)") } @@ -1539,7 +1592,10 @@ func TestInsertShardedIgnore(t *testing.T) { // Fifth row: first shard. // Sixth row: second shard (because 3 hash maps to 40-60). query := "insert ignore into insert_ignore_test(pv, owned, verify) values (1, 1, 1), (2, 2, 2), (3, 3, 1), (4, 4, 4), (5, 5, 1), (6, 6, 3)" - _, err := executorExec(ctx, executor, query, nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, query, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert ignore into insert_ignore_test(pv, owned, verify) values (:_pv_0, :_owned_0, :_verify_0),(:_pv_4, :_owned_4, :_verify_4)", @@ -1641,7 +1697,7 @@ func TestInsertShardedIgnore(t *testing.T) { {}, }) query = "insert ignore into insert_ignore_test(pv, owned, verify) values (1, 1, 1)" - qr, err := executorExec(ctx, executor, query, nil) + qr, err := executorExec(ctx, executor, session, query, nil) require.NoError(t, err) if !qr.Equal(&sqltypes.Result{}) { t.Errorf("qr: %v, want empty result", qr) @@ -1672,7 +1728,10 @@ func TestInsertOnDupKey(t *testing.T) { "1|1", )}) query := "insert into insert_ignore_test(pv, owned, verify) values (1, 1, 1) on duplicate key update col = 2" - _, err := executorExec(ctx, executor, query, nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, query, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into insert_ignore_test(pv, owned, verify) values (:_pv_0, :_owned_0, :_verify_0) on duplicate key update col = 2", @@ -1716,16 +1775,16 @@ func TestAutocommitFail(t *testing.T) { query := "insert into user (id) values (1)" sbc1.MustFailCodes[vtrpcpb.Code_ALREADY_EXISTS] = 1 - primarySession.Reset() - primarySession.Autocommit = true - defer func() { - primarySession.Autocommit = false - }() - _, err := executorExec(ctx, executor, query, nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + Autocommit: true, + } + + _, err := executorExec(ctx, executor, session, query, nil) require.Error(t, err) // make sure we have closed and rolled back any transactions started - assert.False(t, primarySession.InTransaction, "left with tx open") + assert.False(t, session.InTransaction, "left with tx open") } func TestInsertComments(t *testing.T) { @@ -1735,7 +1794,10 @@ func TestInsertComments(t *testing.T) { executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) defer executor.Close() - _, err := executorExec(ctx, executor, "insert into user(id, v, name) values (1, 2, 'myname') /* trailing */", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "insert into user(id, v, name) values (1, 2, 'myname') /* trailing */", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into `user`(id, v, `name`) values (:_Id_0, 2, :_name_0) /* trailing */", @@ -1771,7 +1833,10 @@ func TestInsertGeneratorSharded(t *testing.T) { RowsAffected: 1, InsertID: 1, }}) - result, err := executorExec(ctx, executor, "insert into user(v, `name`) values (2, 'myname')", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + result, err := executorExec(ctx, executor, session, "insert into user(v, `name`) values (2, 'myname')", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into `user`(v, `name`, id) values (2, :_name_0, :_Id_0)", @@ -1816,7 +1881,10 @@ func TestInsertAutoincSharded(t *testing.T) { InsertID: 2, } sbc.SetResults([]*sqltypes.Result{wantResult}) - result, err := executorExec(ctx, router, "insert into user_extra(user_id) values (2)", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + result, err := executorExec(ctx, router, session, "insert into user_extra(user_id) values (2)", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into user_extra(user_id) values (:_user_id_0)", @@ -1828,7 +1896,7 @@ func TestInsertAutoincSharded(t *testing.T) { if !result.Equal(wantResult) { t.Errorf("result: %+v, want %+v", result, wantResult) } - assert.EqualValues(t, 2, primarySession.LastInsertId) + assert.EqualValues(t, 2, session.LastInsertId) } func TestInsertGeneratorUnsharded(t *testing.T) { @@ -1837,7 +1905,10 @@ func TestInsertGeneratorUnsharded(t *testing.T) { defer cancel() executor, _, _, sbclookup := createExecutorEnv(ctx) defer executor.Close() - result, err := executorExec(ctx, executor, "insert into main1(id, name) values (null, 'myname')", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + result, err := executorExec(ctx, executor, session, "insert into main1(id, name) values (null, 'myname')", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select next :n /* INT64 */ values from user_seq", @@ -1877,7 +1948,10 @@ func TestInsertAutoincUnsharded(t *testing.T) { } sbclookup.SetResults([]*sqltypes.Result{wantResult}) - result, err := executorExec(ctx, router, query, nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + result, err := executorExec(ctx, router, session, query, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: query, @@ -1896,7 +1970,10 @@ func TestInsertLookupOwned(t *testing.T) { executor, sbc, _, sbclookup := createExecutorEnv(ctx) defer executor.Close() - _, err := executorExec(ctx, executor, "insert into music(user_id, id) values (2, 3)", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "insert into music(user_id, id) values (2, 3)", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into music(user_id, id) values (:_user_id_0, :_id_0)", @@ -1931,7 +2008,10 @@ func TestInsertLookupOwnedGenerator(t *testing.T) { RowsAffected: 1, InsertID: 1, }}) - result, err := executorExec(ctx, executor, "insert into music(user_id) values (2)", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + result, err := executorExec(ctx, executor, session, "insert into music(user_id) values (2)", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into music(user_id, id) values (:_user_id_0, :_id_0)", @@ -1967,7 +2047,10 @@ func TestInsertLookupUnowned(t *testing.T) { executor, sbc, _, sbclookup := createExecutorEnv(ctx) defer executor.Close() - _, err := executorExec(ctx, executor, "insert into music_extra(user_id, music_id) values (2, 3)", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "insert into music_extra(user_id, music_id) values (2, 3)", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into music_extra(user_id, music_id) values (:_user_id_0, :_music_id_0)", @@ -1997,7 +2080,10 @@ func TestInsertLookupUnownedUnsupplied(t *testing.T) { sqltypes.MakeTestFields("b|a", "int64|varbinary"), "3|1", )}) - _, err := executorExec(ctx, executor, "insert into music_extra_reversed(music_id) values (3)", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "insert into music_extra_reversed(music_id) values (3)", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into music_extra_reversed(music_id, user_id) values (:_music_id_0, :_user_id_0)", @@ -2095,7 +2181,10 @@ func TestMultiInsertSharded(t *testing.T) { executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) defer executor.Close() - _, err := executorExec(ctx, executor, "insert into user(id, v, name) values (1, 1, 'myname1'),(3, 3, 'myname3')", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "insert into user(id, v, name) values (1, 1, 'myname1'),(3, 3, 'myname3')", nil) require.NoError(t, err) wantQueries1 := []*querypb.BoundQuery{{ Sql: "insert into `user`(id, v, `name`) values (:_Id_0, 1, :_name_0)", @@ -2137,7 +2226,7 @@ func TestMultiInsertSharded(t *testing.T) { sbc1.Queries = nil sbclookup.Queries = nil sbc2.Queries = nil - _, err = executorExec(ctx, executor, "insert into user(id, v, name) values (1, 1, 'myname1'),(2, 2, 'myname2')", nil) + _, err = executorExec(ctx, executor, session, "insert into user(id, v, name) values (1, 1, 'myname1'),(2, 2, 'myname2')", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into `user`(id, v, `name`) values (:_Id_0, 1, :_name_0),(:_Id_1, 2, :_name_1)", @@ -2168,7 +2257,7 @@ func TestMultiInsertSharded(t *testing.T) { sbc1.Queries = nil sbclookup.Queries = nil sbc2.Queries = nil - _, err = executorExec(ctx, executor, "insert into user2(id, `name`, lastname) values (2, 'myname', 'mylastname'), (3, 'myname2', 'mylastname2')", nil) + _, err = executorExec(ctx, executor, session, "insert into user2(id, `name`, lastname) values (2, 'myname', 'mylastname'), (3, 'myname2', 'mylastname2')", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "insert into user2(id, `name`, lastname) values (:_id_0, :_name_0, :_lastname_0)", @@ -2210,7 +2299,10 @@ func TestMultiInsertGenerator(t *testing.T) { RowsAffected: 1, InsertID: 1, }}) - result, err := executorExec(ctx, executor, "insert into music(user_id, `name`) values (:u, 'myname1'),(:u, 'myname2')", map[string]*querypb.BindVariable{"u": sqltypes.Int64BindVariable(2)}) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + result, err := executorExec(ctx, executor, session, "insert into music(user_id, `name`) values (:u, 'myname1'),(:u, 'myname2')", map[string]*querypb.BindVariable{"u": sqltypes.Int64BindVariable(2)}) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into music(user_id, `name`, id) values (:_user_id_0, 'myname1', :_id_0),(:_user_id_1, 'myname2', :_id_1)", @@ -2259,7 +2351,10 @@ func TestMultiInsertGeneratorSparse(t *testing.T) { RowsAffected: 1, InsertID: 1, }}) - result, err := executorExec(ctx, executor, "insert into music(id, user_id, name) values (NULL, :u, 'myname1'),(2, :u, 'myname2'), (NULL, :u, 'myname3')", map[string]*querypb.BindVariable{"u": sqltypes.Int64BindVariable(2)}) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + result, err := executorExec(ctx, executor, session, "insert into music(id, user_id, name) values (NULL, :u, 'myname1'),(2, :u, 'myname2'), (NULL, :u, 'myname3')", map[string]*querypb.BindVariable{"u": sqltypes.Int64BindVariable(2)}) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into music(id, user_id, `name`) values (:_id_0, :_user_id_0, 'myname1'),(:_id_1, :_user_id_1, 'myname2'),(:_id_2, :_user_id_2, 'myname3')", @@ -2331,7 +2426,10 @@ func TestInsertBadAutoInc(t *testing.T) { defer executor.Close() // If auto inc table cannot be found, the table should not be added to vschema. - _, err := executorExec(ctx, executor, "insert into bad_auto(v, name) values (1, 'myname')", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "insert into bad_auto(v, name) values (1, 'myname')", nil) want := "table bad_auto not found" if err == nil || err.Error() != want { t.Errorf("bad auto inc err: %v, want %v", err, want) @@ -2404,8 +2502,10 @@ func TestKeyDestRangeQuery(t *testing.T) { executor, sbc1, sbc2, _ := createExecutorEnv(ctx) defer executor.Close() - primarySession.TargetString = tc.targetString - _, err := executorExec(ctx, executor, tc.inputQuery, nil) + session := &vtgatepb.Session{ + TargetString: tc.targetString, + } + _, err := executorExec(ctx, executor, session, tc.inputQuery, nil) require.NoError(t, err) if tc.expectedSbc1Query == "" { @@ -2428,12 +2528,12 @@ func TestKeyDestRangeQuery(t *testing.T) { defer cancel() executor, _, _, _ := createExecutorEnv(ctx) defer executor.Close() - primarySession.TargetString = "TestExecutor[-]" - _, err := executorExec(ctx, executor, insertInput, nil) + session := &vtgatepb.Session{ + TargetString: "TestExecutor[-]", + } + _, err := executorExec(ctx, executor, session, insertInput, nil) require.EqualError(t, err, "VT03023: INSERT not supported when targeting a key range: TestExecutor[-]") - - primarySession.TargetString = "" } func assertQueriesContain(t *testing.T, sql, sbcName string, sbc *sandboxconn.SandboxConn) { @@ -2456,7 +2556,10 @@ func TestUpdateEqualWithPrepare(t *testing.T) { logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) - _, err := executorPrepare(ctx, executor, "update music set a = :a0 where id = :id0", map[string]*querypb.BindVariable{ + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorPrepare(ctx, executor, session, "update music set a = :a0 where id = :id0", map[string]*querypb.BindVariable{ "a0": sqltypes.Int64BindVariable(3), "id0": sqltypes.Int64BindVariable(2), }) @@ -2478,7 +2581,10 @@ func TestInsertShardedWithPrepare(t *testing.T) { logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) - _, err := executorPrepare(ctx, executor, "insert into user(id, v, name) values (:_Id0, 2, ':_name_0')", map[string]*querypb.BindVariable{ + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorPrepare(ctx, executor, session, "insert into user(id, v, name) values (:_Id0, 2, ':_name_0')", map[string]*querypb.BindVariable{ "_Id0": sqltypes.Int64BindVariable(1), "_name_0": sqltypes.BytesBindVariable([]byte("myname")), "__seq0": sqltypes.Int64BindVariable(1), @@ -2499,7 +2605,10 @@ func TestDeleteEqualWithPrepare(t *testing.T) { defer cancel() executor, sbc, _, sbclookup := createExecutorEnv(ctx) defer executor.Close() - _, err := executorPrepare(ctx, executor, "delete from user where id = :id0", map[string]*querypb.BindVariable{ + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorPrepare(ctx, executor, session, "delete from user where id = :id0", map[string]*querypb.BindVariable{ "id0": sqltypes.Int64BindVariable(1), }) require.NoError(t, err) @@ -2520,8 +2629,11 @@ func TestUpdateLastInsertID(t *testing.T) { executor.normalize = true sql := "update user set a = last_insert_id() where id = 1" - primarySession.LastInsertId = 43 - _, err := executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{}) + session := &vtgatepb.Session{ + TargetString: "@primary", + LastInsertId: 43, + } + _, err := executorExec(ctx, executor, session, sql, map[string]*querypb.BindVariable{}) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "update `user` set a = :__lastInsertId where id = :id /* INT64 */", @@ -2543,7 +2655,10 @@ func TestUpdateReference(t *testing.T) { logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) - _, err := executorExec(ctx, executor, "update zip_detail set status = 'CLOSED' where id = 1", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "update zip_detail set status = 'CLOSED' where id = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "update zip_detail set `status` = 'CLOSED' where id = 1", @@ -2557,7 +2672,7 @@ func TestUpdateReference(t *testing.T) { sbclookup.Queries = nil - _, err = executorExec(ctx, executor, "update TestUnsharded.zip_detail set status = 'CLOSED' where id = 1", nil) + _, err = executorExec(ctx, executor, session, "update TestUnsharded.zip_detail set status = 'CLOSED' where id = 1", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "update zip_detail set `status` = 'CLOSED' where id = 1", @@ -2572,7 +2687,7 @@ func TestUpdateReference(t *testing.T) { sbclookup.Queries = nil - _, err = executorExec(ctx, executor, "update TestExecutor.zip_detail set status = 'CLOSED' where id = 1", nil) + _, err = executorExec(ctx, executor, session, "update TestExecutor.zip_detail set status = 'CLOSED' where id = 1", nil) require.NoError(t, err) // Gen4 planner can redirect the query to correct source for update when reference table is involved. } @@ -2586,7 +2701,10 @@ func TestDeleteLookupOwnedEqual(t *testing.T) { sbc1.SetResults([]*sqltypes.Result{ sqltypes.MakeTestResult(sqltypes.MakeTestFields("uniq_col|keyspace_id", "int64|varbinary"), "1|N±\u0090ɢú\u0016\u009C"), }) - _, err := executorExec(ctx, executor, "delete from t1 where unq_col = 1", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "delete from t1 where unq_col = 1", nil) require.NoError(t, err) tupleBindVar, _ := sqltypes.BuildBindVariable([]int64{1}) sbc1wantQueries := []*querypb.BoundQuery{{ @@ -2617,7 +2735,10 @@ func TestDeleteReference(t *testing.T) { logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) - _, err := executorExec(ctx, executor, "delete from zip_detail where id = 1", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "delete from zip_detail where id = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "delete from zip_detail where id = 1", @@ -2631,7 +2752,7 @@ func TestDeleteReference(t *testing.T) { sbclookup.Queries = nil - _, err = executorExec(ctx, executor, "delete from zip_detail where id = 1", nil) + _, err = executorExec(ctx, executor, session, "delete from zip_detail where id = 1", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "delete from zip_detail where id = 1", @@ -2645,7 +2766,7 @@ func TestDeleteReference(t *testing.T) { sbclookup.Queries = nil - _, err = executorExec(ctx, executor, "delete from TestExecutor.zip_detail where id = 1", nil) + _, err = executorExec(ctx, executor, session, "delete from TestExecutor.zip_detail where id = 1", nil) require.NoError(t, err) // Gen4 planner can redirect the query to correct source for update when reference table is involved. } @@ -3084,7 +3205,10 @@ func TestInsertReference(t *testing.T) { logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) - _, err := executorExec(ctx, executor, "insert into zip_detail(id, status) values (1, 'CLOSED')", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "insert into zip_detail(id, status) values (1, 'CLOSED')", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "insert into zip_detail(id, `status`) values (1, 'CLOSED')", @@ -3098,7 +3222,7 @@ func TestInsertReference(t *testing.T) { sbclookup.Queries = nil - _, err = executorExec(ctx, executor, "insert into TestUnsharded.zip_detail(id, status) values (1, 'CLOSED')", nil) + _, err = executorExec(ctx, executor, session, "insert into TestUnsharded.zip_detail(id, status) values (1, 'CLOSED')", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "insert into zip_detail(id, `status`) values (1, 'CLOSED')", @@ -3113,6 +3237,6 @@ func TestInsertReference(t *testing.T) { sbclookup.Queries = nil - _, err = executorExec(ctx, executor, "insert into TestExecutor.zip_detail(id, status) values (1, 'CLOSED')", nil) + _, err = executorExec(ctx, executor, session, "insert into TestExecutor.zip_detail(id, status) values (1, 'CLOSED')", nil) require.NoError(t, err) // Gen4 planner can redirect the query to correct source for update when reference table is involved. } diff --git a/go/vt/vtgate/executor_framework_test.go b/go/vt/vtgate/executor_framework_test.go index 81c7ecde423..4c1fe85fe59 100644 --- a/go/vt/vtgate/executor_framework_test.go +++ b/go/vt/vtgate/executor_framework_test.go @@ -174,10 +174,6 @@ func createExecutorEnv(ctx context.Context) (executor *Executor, sbc1, sbc2, sbc executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) key.AnyShardPicker = DestinationAnyShardPickerFirstShard{} - // create a new session each time so that ShardSessions don't get re-used across tests - primarySession = &vtgatepb.Session{ - TargetString: "@primary", - } return executor, sbc1, sbc2, sbclookup } @@ -197,10 +193,6 @@ func createCustomExecutor(ctx context.Context, vschema string) (executor *Execut queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize) executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) - // create a new session each time so that ShardSessions don't get re-used across tests - primarySession = &vtgatepb.Session{ - TargetString: "@primary", - } return executor, sbc1, sbc2, sbclookup } @@ -227,10 +219,6 @@ func createCustomExecutorSetValues(ctx context.Context, vschema string, values [ queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize) executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) - // create a new session each time so that ShardSessions don't get re-used across tests - primarySession = &vtgatepb.Session{ - TargetString: "@primary", - } return executor, sbcs[0], sbcs[1], sbclookup } @@ -244,15 +232,15 @@ func executorExecSession(ctx context.Context, executor *Executor, sql string, bv bv) } -func executorExec(ctx context.Context, executor *Executor, sql string, bv map[string]*querypb.BindVariable) (*sqltypes.Result, error) { - return executorExecSession(ctx, executor, sql, bv, primarySession) +func executorExec(ctx context.Context, executor *Executor, session *vtgatepb.Session, sql string, bv map[string]*querypb.BindVariable) (*sqltypes.Result, error) { + return executorExecSession(ctx, executor, sql, bv, session) } -func executorPrepare(ctx context.Context, executor *Executor, sql string, bv map[string]*querypb.BindVariable) ([]*querypb.Field, error) { +func executorPrepare(ctx context.Context, executor *Executor, session *vtgatepb.Session, sql string, bv map[string]*querypb.BindVariable) ([]*querypb.Field, error) { return executor.Prepare( ctx, "TestExecute", - NewSafeSession(primarySession), + NewSafeSession(session), sql, bv) } diff --git a/go/vt/vtgate/executor_select_test.go b/go/vt/vtgate/executor_select_test.go index 4fe2e6a697b..6c095e6c739 100644 --- a/go/vt/vtgate/executor_select_test.go +++ b/go/vt/vtgate/executor_select_test.go @@ -36,8 +36,6 @@ import ( "github.com/google/go-cmp/cmp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - "google.golang.org/protobuf/proto" - "vitess.io/vitess/go/cache" "vitess.io/vitess/go/sqltypes" "vitess.io/vitess/go/test/utils" @@ -575,7 +573,10 @@ func TestUnsharded(t *testing.T) { executor, _, _, sbclookup := createExecutorEnv(ctx) defer executor.Close() - _, err := executorExec(ctx, executor, "select id from music_user_map where id = 1", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "select id from music_user_map where id = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select id from music_user_map where id = 1", @@ -591,7 +592,10 @@ func TestUnshardedComments(t *testing.T) { executor, _, _, sbclookup := createExecutorEnv(ctx) defer executor.Close() - _, err := executorExec(ctx, executor, "/* leading */ select id from music_user_map where id = 1 /* trailing */", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "/* leading */ select id from music_user_map where id = 1 /* trailing */", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "/* leading */ select id from music_user_map where id = 1 /* trailing */", @@ -599,7 +603,7 @@ func TestUnshardedComments(t *testing.T) { }} utils.MustMatch(t, wantQueries, sbclookup.Queries) - _, err = executorExec(ctx, executor, "update music_user_map set id = 1 /* trailing */", nil) + _, err = executorExec(ctx, executor, session, "update music_user_map set id = 1 /* trailing */", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "/* leading */ select id from music_user_map where id = 1 /* trailing */", @@ -611,7 +615,7 @@ func TestUnshardedComments(t *testing.T) { assertQueries(t, sbclookup, wantQueries) sbclookup.Queries = nil - _, err = executorExec(ctx, executor, "delete from music_user_map /* trailing */", nil) + _, err = executorExec(ctx, executor, session, "delete from music_user_map /* trailing */", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "delete from music_user_map /* trailing */", @@ -620,7 +624,7 @@ func TestUnshardedComments(t *testing.T) { assertQueries(t, sbclookup, wantQueries) sbclookup.Queries = nil - _, err = executorExec(ctx, executor, "insert into music_user_map values (1) /* trailing */", nil) + _, err = executorExec(ctx, executor, session, "insert into music_user_map values (1) /* trailing */", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "insert into music_user_map values (1) /* trailing */", @@ -673,11 +677,15 @@ func TestStreamBuffering(t *testing.T) { }}) var results []*sqltypes.Result + session := &vtgatepb.Session{ + TargetString: "@primary", + } + err := executor.StreamExecute( context.Background(), nil, "TestStreamBuffering", - NewSafeSession(primarySession), + NewSafeSession(session), "select id from music_user_map where id = 1", nil, func(qr *sqltypes.Result) error { @@ -745,11 +753,14 @@ func TestStreamLimitOffset(t *testing.T) { }}) results := make(chan *sqltypes.Result, 10) + session := &vtgatepb.Session{ + TargetString: "@primary", + } err := executor.StreamExecute( context.Background(), nil, "TestStreamLimitOffset", - NewSafeSession(primarySession), + NewSafeSession(session), "select id, textcol from user order by id limit 2 offset 2", nil, func(qr *sqltypes.Result) error { @@ -790,13 +801,16 @@ func TestSelectLastInsertId(t *testing.T) { defer cancel() executor, _, _, _ := createExecutorEnv(ctx) defer executor.Close() - primarySession.LastInsertId = 52 + session := &vtgatepb.Session{ + TargetString: "@primary", + LastInsertId: 52, + } executor.normalize = true logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) sql := "select last_insert_id()" - result, err := executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{}) + result, err := executorExec(ctx, executor, session, sql, map[string]*querypb.BindVariable{}) wantResult := &sqltypes.Result{ Fields: []*querypb.Field{ {Name: "last_insert_id()", Type: sqltypes.Uint64, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NOT_NULL_FLAG | querypb.MySqlFlag_NUM_FLAG | querypb.MySqlFlag_UNSIGNED_FLAG)}, @@ -815,10 +829,14 @@ func TestSelectSystemVariables(t *testing.T) { defer cancel() executor, _, _, _ := createExecutorEnv(ctx) defer executor.Close() - primarySession.ReadAfterWrite = &vtgatepb.ReadAfterWrite{ - ReadAfterWriteGtid: "a fine gtid", - ReadAfterWriteTimeout: 13, - SessionTrackGtids: true, + + session := &vtgatepb.Session{ + TargetString: "@primary", + ReadAfterWrite: &vtgatepb.ReadAfterWrite{ + ReadAfterWriteGtid: "a fine gtid", + ReadAfterWriteTimeout: 13, + SessionTrackGtids: true, + }, } executor.normalize = true logChan := executor.queryLogger.Subscribe("Test") @@ -828,7 +846,7 @@ func TestSelectSystemVariables(t *testing.T) { "@@sql_select_limit, @@transaction_mode, @@workload, @@read_after_write_gtid, " + "@@read_after_write_timeout, @@session_track_gtids, @@ddl_strategy, @@migration_context, @@socket, @@query_timeout" - result, err := executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{}) + result, err := executorExec(ctx, executor, session, sql, map[string]*querypb.BindVariable{}) wantResult := &sqltypes.Result{ Fields: []*querypb.Field{ {Name: "@@autocommit", Type: sqltypes.Int64, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NOT_NULL_FLAG | querypb.MySqlFlag_NUM_FLAG)}, @@ -879,19 +897,16 @@ func TestSelectInitializedVitessAwareVariable(t *testing.T) { logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) - primarySession.Autocommit = true - primarySession.EnableSystemSettings = true - primarySession.QueryTimeout = 75 - - defer func() { - primarySession.Autocommit = false - primarySession.EnableSystemSettings = false - primarySession.QueryTimeout = 0 - }() + session := &vtgatepb.Session{ + TargetString: "@primary", + Autocommit: true, + EnableSystemSettings: true, + QueryTimeout: 75, + } sql := "select @@autocommit, @@enable_system_settings, @@query_timeout" - result, err := executorExec(ctx, executor, sql, nil) + result, err := executorExec(ctx, executor, session, sql, nil) wantResult := &sqltypes.Result{ Fields: []*querypb.Field{ {Name: "@@autocommit", Type: sqltypes.Int64, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NOT_NULL_FLAG | querypb.MySqlFlag_NUM_FLAG)}, @@ -919,7 +934,10 @@ func TestSelectUserDefinedVariable(t *testing.T) { defer executor.queryLogger.Unsubscribe(logChan) sql := "select @foo" - result, err := executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{}) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + result, err := executorExec(ctx, executor, session, sql, map[string]*querypb.BindVariable{}) require.NoError(t, err) wantResult := &sqltypes.Result{ Fields: []*querypb.Field{ @@ -931,8 +949,8 @@ func TestSelectUserDefinedVariable(t *testing.T) { } utils.MustMatch(t, wantResult, result, "Mismatch") - primarySession = &vtgatepb.Session{UserDefinedVariables: createMap([]string{"foo"}, []any{"bar"})} - result, err = executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{}) + session = &vtgatepb.Session{UserDefinedVariables: createMap([]string{"foo"}, []any{"bar"})} + result, err = executorExec(ctx, executor, session, sql, map[string]*querypb.BindVariable{}) require.NoError(t, err) wantResult = &sqltypes.Result{ Fields: []*querypb.Field{ @@ -955,12 +973,15 @@ func TestFoundRows(t *testing.T) { logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) + session := &vtgatepb.Session{ + TargetString: "@primary", + } // run this extra query so we can assert on the number of rows found - _, err := executorExec(ctx, executor, "select 42", map[string]*querypb.BindVariable{}) + _, err := executorExec(ctx, executor, session, "select 42", map[string]*querypb.BindVariable{}) require.NoError(t, err) sql := "select found_rows()" - result, err := executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{}) + result, err := executorExec(ctx, executor, session, sql, map[string]*querypb.BindVariable{}) wantResult := &sqltypes.Result{ Fields: []*querypb.Field{ {Name: "found_rows()", Type: sqltypes.Int64, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NOT_NULL_FLAG | querypb.MySqlFlag_NUM_FLAG)}, @@ -983,18 +1004,21 @@ func TestRowCount(t *testing.T) { logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) - _, err := executorExec(ctx, executor, "select 42", map[string]*querypb.BindVariable{}) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "select 42", map[string]*querypb.BindVariable{}) require.NoError(t, err) - testRowCount(t, ctx, executor, -1) + testRowCount(t, ctx, executor, session, -1) - _, err = executorExec(ctx, executor, "delete from user where id in (42, 24)", map[string]*querypb.BindVariable{}) + _, err = executorExec(ctx, executor, session, "delete from user where id in (42, 24)", map[string]*querypb.BindVariable{}) require.NoError(t, err) - testRowCount(t, ctx, executor, 2) + testRowCount(t, ctx, executor, session, 2) } -func testRowCount(t *testing.T, ctx context.Context, executor *Executor, wantRowCount int64) { +func testRowCount(t *testing.T, ctx context.Context, executor *Executor, session *vtgatepb.Session, wantRowCount int64) { t.Helper() - result, err := executorExec(ctx, executor, "select row_count()", map[string]*querypb.BindVariable{}) + result, err := executorExec(ctx, executor, session, "select row_count()", map[string]*querypb.BindVariable{}) wantResult := &sqltypes.Result{ Fields: []*querypb.Field{ {Name: "row_count()", Type: sqltypes.Int64, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NOT_NULL_FLAG | querypb.MySqlFlag_NUM_FLAG)}, @@ -1014,7 +1038,11 @@ func TestSelectLastInsertIdInUnion(t *testing.T) { executor, sbc1, _, _ := createExecutorEnv(ctx) defer executor.Close() executor.normalize = true - primarySession.LastInsertId = 52 + + session := &vtgatepb.Session{ + TargetString: "@primary", + LastInsertId: 52, + } result1 := []*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -1028,7 +1056,7 @@ func TestSelectLastInsertIdInUnion(t *testing.T) { sbc1.SetResults(result1) sql := "select last_insert_id() as id union select last_insert_id() as id" - got, err := executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{}) + got, err := executorExec(ctx, executor, session, sql, map[string]*querypb.BindVariable{}) require.NoError(t, err) wantResult := &sqltypes.Result{ Fields: []*querypb.Field{ @@ -1052,7 +1080,10 @@ func TestSelectLastInsertIdInWhere(t *testing.T) { defer executor.queryLogger.Unsubscribe(logChan) sql := "select id from music_user_map where id = last_insert_id()" - _, err := executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{}) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, sql, map[string]*querypb.BindVariable{}) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select id from music_user_map where id = :__lastInsertId", @@ -1081,7 +1112,10 @@ func TestLastInsertIDInVirtualTable(t *testing.T) { }}, }} sbc1.SetResults(result1) - _, err := executorExec(ctx, executor, "select * from (select last_insert_id()) as t", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "select * from (select last_insert_id()) as t", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select t.`last_insert_id()` from (select :__lastInsertId as `last_insert_id()` from dual) as t", @@ -1098,12 +1132,11 @@ func TestLastInsertIDInSubQueryExpression(t *testing.T) { executor, sbc1, sbc2, _ := createExecutorEnv(ctx) defer executor.Close() executor.normalize = true - primarySession.LastInsertId = 12345 - defer func() { - // clean up global state - primarySession.LastInsertId = 0 - }() - rs, err := executorExec(ctx, executor, "select (select last_insert_id()) as x", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + LastInsertId: 12345, + } + rs, err := executorExec(ctx, executor, session, "select (select last_insert_id()) as x", nil) require.NoError(t, err) wantResult := &sqltypes.Result{ Fields: []*querypb.Field{ @@ -1128,7 +1161,9 @@ func TestSelectDatabase(t *testing.T) { defer executor.Close() executor.normalize = true sql := "select database()" - newSession := proto.Clone(primarySession).(*vtgatepb.Session) + newSession := &vtgatepb.Session{ + TargetString: "@primary", + } session := NewSafeSession(newSession) session.TargetString = "TestExecutor@primary" result, err := executor.Execute( @@ -1169,7 +1204,10 @@ func TestSelectBindvars(t *testing.T) { )}) sql := "select id from `user` where id = :id" - _, err := executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{ + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, sql, map[string]*querypb.BindVariable{ "id": sqltypes.Int64BindVariable(1), }) require.NoError(t, err) @@ -1184,7 +1222,7 @@ func TestSelectBindvars(t *testing.T) { // Test with StringBindVariable sql = "select id from `user` where `name` in (:name1, :name2)" - _, err = executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{ + _, err = executorExec(ctx, executor, session, sql, map[string]*querypb.BindVariable{ "name1": sqltypes.StringBindVariable("foo1"), "name2": sqltypes.StringBindVariable("foo2"), }) @@ -1203,7 +1241,7 @@ func TestSelectBindvars(t *testing.T) { // Test with BytesBindVariable sql = "select id from `user` where `name` in (:name1, :name2)" - _, err = executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{ + _, err = executorExec(ctx, executor, session, sql, map[string]*querypb.BindVariable{ "name1": sqltypes.BytesBindVariable([]byte("foo1")), "name2": sqltypes.BytesBindVariable([]byte("foo2")), }) @@ -1232,7 +1270,7 @@ func TestSelectBindvars(t *testing.T) { }}) sql = "select id from user where name = :name" - _, err = executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{ + _, err = executorExec(ctx, executor, session, sql, map[string]*querypb.BindVariable{ "name": sqltypes.StringBindVariable("nonexistent"), }) require.NoError(t, err) @@ -1266,7 +1304,10 @@ func TestSelectEqual(t *testing.T) { executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) defer executor.Close() - _, err := executorExec(ctx, executor, "select id from user where id = 1", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "select id from user where id = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select id from `user` where id = 1", @@ -1278,7 +1319,7 @@ func TestSelectEqual(t *testing.T) { } sbc1.Queries = nil - _, err = executorExec(ctx, executor, "select id from user where id = 3", nil) + _, err = executorExec(ctx, executor, session, "select id from user where id = 3", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "select id from `user` where id = 3", @@ -1293,7 +1334,7 @@ func TestSelectEqual(t *testing.T) { } sbc2.Queries = nil - _, err = executorExec(ctx, executor, "select id from user where id = '3'", nil) + _, err = executorExec(ctx, executor, session, "select id from user where id = '3'", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "select id from `user` where id = '3'", @@ -1312,7 +1353,7 @@ func TestSelectEqual(t *testing.T) { sqltypes.MakeTestFields("b|a", "varbinary|varbinary"), "foo|1", )}) - _, err = executorExec(ctx, executor, "select id from user where name = 'foo'", nil) + _, err = executorExec(ctx, executor, session, "select id from user where name = 'foo'", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "select id from `user` where `name` = 'foo'", @@ -1338,7 +1379,10 @@ func TestSelectINFromOR(t *testing.T) { defer executor.Close() executor.pv = querypb.ExecuteOptions_Gen4 - _, err := executorExec(ctx, executor, "select 1 from user where id = 1 and name = 'apa' or id = 2 and name = 'toto'", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "select 1 from user where id = 1 and name = 'apa' or id = 2 and name = 'toto'", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select 1 from `user` where id = 1 and `name` = 'apa' or id = 2 and `name` = 'toto'", @@ -1356,7 +1400,10 @@ func TestSelectDual(t *testing.T) { executor, sbc1, _, lookup := createExecutorEnv(ctx) defer executor.Close() - _, err := executorExec(ctx, executor, "select @@aa.bb from dual", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "select @@aa.bb from dual", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select @@`aa.bb` from dual", @@ -1364,7 +1411,7 @@ func TestSelectDual(t *testing.T) { }} utils.MustMatch(t, wantQueries, sbc1.Queries) - _, err = executorExec(ctx, executor, "select @@aa.bb from TestUnsharded.dual", nil) + _, err = executorExec(ctx, executor, session, "select @@aa.bb from TestUnsharded.dual", nil) require.NoError(t, err) utils.MustMatch(t, wantQueries, lookup.Queries) } @@ -1376,7 +1423,10 @@ func TestSelectComments(t *testing.T) { executor, sbc1, sbc2, _ := createExecutorEnv(ctx) defer executor.Close() - _, err := executorExec(ctx, executor, "/* leading */ select id from user where id = 1 /* trailing */", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "/* leading */ select id from user where id = 1 /* trailing */", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "/* leading */ select id from `user` where id = 1 /* trailing */", @@ -1397,7 +1447,10 @@ func TestSelectNormalize(t *testing.T) { defer executor.Close() executor.normalize = true - _, err := executorExec(ctx, executor, "/* leading */ select id from user where id = 1 /* trailing */", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "/* leading */ select id from user where id = 1 /* trailing */", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "/* leading */ select id from `user` where id = :id /* INT64 */ /* trailing */", @@ -1412,8 +1465,8 @@ func TestSelectNormalize(t *testing.T) { sbc1.Queries = nil // Force the query to go to the "wrong" shard and ensure that normalization still happens - primarySession.TargetString = "TestExecutor/40-60" - _, err = executorExec(ctx, executor, "/* leading */ select id from user where id = 1 /* trailing */", nil) + session.TargetString = "TestExecutor/40-60" + _, err = executorExec(ctx, executor, session, "/* leading */ select id from user where id = 1 /* trailing */", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "/* leading */ select id from `user` where id = :id /* INT64 */ /* trailing */", @@ -1424,7 +1477,6 @@ func TestSelectNormalize(t *testing.T) { require.Empty(t, sbc1.Queries) utils.MustMatch(t, wantQueries, sbc2.Queries, "sbc2.Queries") sbc2.Queries = nil - primarySession.TargetString = "" } func TestSelectCaseSensitivity(t *testing.T) { @@ -1434,7 +1486,10 @@ func TestSelectCaseSensitivity(t *testing.T) { executor, sbc1, sbc2, _ := createExecutorEnv(ctx) defer executor.Close() - _, err := executorExec(ctx, executor, "select Id from user where iD = 1", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "select Id from user where iD = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select Id from `user` where iD = 1", @@ -1470,7 +1525,10 @@ func TestSelectKeyRange(t *testing.T) { executor, sbc1, sbc2, _ := createExecutorEnv(ctx) defer executor.Close() - _, err := executorExec(ctx, executor, "select krcol_unique, krcol from keyrange_table where krcol = 1", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "select krcol_unique, krcol from keyrange_table where krcol = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select krcol_unique, krcol from keyrange_table where krcol = 1", @@ -1490,7 +1548,10 @@ func TestSelectKeyRangeUnique(t *testing.T) { executor, sbc1, sbc2, _ := createExecutorEnv(ctx) defer executor.Close() - _, err := executorExec(ctx, executor, "select krcol_unique, krcol from keyrange_table where krcol_unique = 1", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "select krcol_unique, krcol from keyrange_table where krcol_unique = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select krcol_unique, krcol from keyrange_table where krcol_unique = 1", @@ -1511,7 +1572,10 @@ func TestSelectIN(t *testing.T) { defer executor.Close() // Constant in IN clause is just a number, not a bind variable. - _, err := executorExec(ctx, executor, "select id from user where id in (1)", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, "select id from user where id in (1)", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select id from `user` where id in (1)", @@ -1526,7 +1590,7 @@ func TestSelectIN(t *testing.T) { // They result in two different queries on two shards. sbc1.Queries = nil sbc2.Queries = nil - _, err = executorExec(ctx, executor, "select id from user where id in (1, 3)", nil) + _, err = executorExec(ctx, executor, session, "select id from user where id in (1, 3)", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "select id from `user` where id in ::__vals", @@ -1547,7 +1611,7 @@ func TestSelectIN(t *testing.T) { // This is using []any for the bind variable list. sbc1.Queries = nil sbc2.Queries = nil - _, err = executorExec(ctx, executor, "select id from user where id in ::vals", map[string]*querypb.BindVariable{ + _, err = executorExec(ctx, executor, session, "select id from user where id in ::vals", map[string]*querypb.BindVariable{ "vals": sqltypes.TestBindVariable([]any{int64(1), int64(3)}), }) require.NoError(t, err) @@ -1575,7 +1639,7 @@ func TestSelectIN(t *testing.T) { sqltypes.MakeTestFields("b|a", "varbinary|varbinary"), "foo|1", )}) - _, err = executorExec(ctx, executor, "select id from user where name = 'foo'", nil) + _, err = executorExec(ctx, executor, session, "select id from user where name = 'foo'", nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{{ Sql: "select id from `user` where `name` = 'foo'", @@ -1672,7 +1736,10 @@ func TestSelectScatter(t *testing.T) { defer executor.queryLogger.Unsubscribe(logChan) sql := "select id from `user`" - _, err := executorExec(ctx, executor, sql, nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, sql, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select id from `user`", @@ -1690,7 +1757,7 @@ func TestSelectScatterPartial(t *testing.T) { defer cancel() // Special setup: Don't use createExecutorEnv. - primarySession = &vtgatepb.Session{ + session := &vtgatepb.Session{ TargetString: "@primary", } cell := "aa" @@ -1714,7 +1781,7 @@ func TestSelectScatterPartial(t *testing.T) { // Fail 1 of N without the directive fails the whole operation conns[2].MustFailCodes[vtrpcpb.Code_RESOURCE_EXHAUSTED] = 1000 - results, err := executorExec(ctx, executor, "select id from `user`", nil) + results, err := executorExec(ctx, executor, session, "select id from `user`", nil) wantErr := "TestExecutor.40-60.primary" if err == nil || !strings.Contains(err.Error(), wantErr) { t.Errorf("want error %v, got %v", wantErr, err) @@ -1728,7 +1795,7 @@ func TestSelectScatterPartial(t *testing.T) { testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select id from `user`", 8) // Fail 1 of N with the directive succeeds with 7 rows - results, err = executorExec(ctx, executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user", nil) + results, err = executorExec(ctx, executor, session, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user", nil) require.NoError(t, err) if results == nil || len(results.Rows) != 7 { t.Errorf("want 7 results, got %v", results) @@ -1744,11 +1811,11 @@ func TestSelectScatterPartial(t *testing.T) { conns[6].MustFailCodes[vtrpcpb.Code_RESOURCE_EXHAUSTED] = 1000 conns[7].MustFailCodes[vtrpcpb.Code_RESOURCE_EXHAUSTED] = 1000 - _, err = executorExec(ctx, executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user", nil) + _, err = executorExec(ctx, executor, session, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user", nil) require.Error(t, err) testQueryLog(t, executor, logChan, "TestExecute", "SELECT", "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from `user`", 8) - _, err = executorExec(ctx, executor, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user order by id", nil) + _, err = executorExec(ctx, executor, session, "select /*vt+ SCATTER_ERRORS_AS_WARNINGS=1 */ id from user order by id", nil) require.Error(t, err) } @@ -1948,7 +2015,10 @@ func TestSelectScatterOrderBy(t *testing.T) { defer executor.Close() query := "select col1, col2 from user order by col2 desc" - gotResult, err := executorExec(ctx, executor, query, nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + gotResult, err := executorExec(ctx, executor, session, query, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ @@ -2017,7 +2087,10 @@ func TestSelectScatterOrderByVarChar(t *testing.T) { defer executor.Close() query := "select col1, textcol from user order by textcol desc" - gotResult, err := executorExec(ctx, executor, query, nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + gotResult, err := executorExec(ctx, executor, session, query, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ @@ -2209,7 +2282,10 @@ func TestSelectScatterAggregate(t *testing.T) { defer executor.Close() query := "select col, sum(foo) from user group by col" - gotResult, err := executorExec(ctx, executor, query, nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + gotResult, err := executorExec(ctx, executor, session, query, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ @@ -2338,7 +2414,10 @@ func TestSelectScatterLimit(t *testing.T) { defer executor.Close() query := "select col1, col2 from user order by col2 desc limit 3" - gotResult, err := executorExec(ctx, executor, query, nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + gotResult, err := executorExec(ctx, executor, session, query, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ @@ -2458,7 +2537,10 @@ func TestSimpleJoin(t *testing.T) { defer executor.queryLogger.Unsubscribe(logChan) sql := "select u1.id, u2.id from user u1 join user u2 where u1.id = 1 and u2.id = 3" - result, err := executorExec(ctx, executor, sql, nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + result, err := executorExec(ctx, executor, session, sql, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select u1.id from `user` as u1 where u1.id = 1", @@ -2500,7 +2582,10 @@ func TestJoinComments(t *testing.T) { defer executor.queryLogger.Unsubscribe(logChan) sql := "select u1.id, u2.id from user u1 join user u2 where u1.id = 1 and u2.id = 3 /* trailing */" - _, err := executorExec(ctx, executor, sql, nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, sql, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select u1.id from `user` as u1 where u1.id = 1 /* trailing */", @@ -2582,7 +2667,10 @@ func TestVarJoin(t *testing.T) { }} sbc1.SetResults(result1) sql := "select u1.id, u2.id from user u1 join user u2 on u2.id = u1.col where u1.id = 1" - _, err := executorExec(ctx, executor, sql, nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, sql, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select u1.id, u1.col from `user` as u1 where u1.id = 1", @@ -2667,7 +2755,10 @@ func TestLeftJoin(t *testing.T) { sbc1.SetResults(result1) sbc2.SetResults(emptyResult) sql := "select u1.id, u2.id from user u1 left join user u2 on u2.id = u1.col where u1.id = 1" - result, err := executorExec(ctx, executor, sql, nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + result, err := executorExec(ctx, executor, session, sql, nil) require.NoError(t, err) wantResult := &sqltypes.Result{ Fields: []*querypb.Field{ @@ -2752,7 +2843,10 @@ func TestEmptyJoin(t *testing.T) { {Name: "col", Type: sqltypes.Int32, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NUM_FLAG)}, }, }}) - result, err := executorExec(ctx, executor, "select u1.id, u2.id from user u1 join user u2 on u2.id = u1.col where u1.id = 1", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + result, err := executorExec(ctx, executor, session, "select u1.id, u2.id from user u1 join user u2 on u2.id = u1.col where u1.id = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select u1.id, u1.col from `user` as u1 where u1.id = 1", @@ -2838,7 +2932,10 @@ func TestEmptyJoinRecursive(t *testing.T) { {Name: "id", Type: sqltypes.Int32, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NUM_FLAG)}, }, }}) - result, err := executorExec(ctx, executor, "select u1.id, u2.id, u3.id from user u1 join (user u2 join user u3 on u3.id = u2.col) where u1.id = 1", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + result, err := executorExec(ctx, executor, session, "select u1.id, u2.id, u3.id from user u1 join (user u2 join user u3 on u3.id = u2.col) where u1.id = 1", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select u1.id from `user` as u1 where u1.id = 1", @@ -2933,7 +3030,10 @@ func TestCrossShardSubquery(t *testing.T) { }}, }} sbc1.SetResults(result1) - result, err := executorExec(ctx, executor, "select id1 from (select u1.id id1, u2.id from user u1 join user u2 on u2.id = u1.col where u1.id = 1) as t", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + result, err := executorExec(ctx, executor, session, "select id1 from (select u1.id id1, u2.id from user u1 join user u2 on u2.id = u1.col where u1.id = 1) as t", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select id1, t.`u1.col` from (select u1.id as id1, u1.col as `u1.col` from `user` as u1 where u1.id = 1) as t", @@ -3064,7 +3164,10 @@ func TestCrossShardSubqueryGetFields(t *testing.T) { }, }} sbc1.SetResults(result1) - result, err := executorExec(ctx, executor, "select main1.col, t.id1 from main1 join (select u1.id id1, u2.id from user u1 join user u2 on u2.id = u1.col where u1.id = 1) as t", nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + result, err := executorExec(ctx, executor, session, "select main1.col, t.id1 from main1 join (select u1.id id1, u2.id from user u1 join user u2 on u2.id = u1.col where u1.id = 1) as t", nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: "select t.id1, t.`u1.col` from (select u1.id as id1, u1.col as `u1.col` from `user` as u1 where 1 != 1) as t where 1 != 1", @@ -3099,7 +3202,10 @@ func TestSelectBindvarswithPrepare(t *testing.T) { defer executor.queryLogger.Unsubscribe(logChan) sql := "select id from `user` where id = :id" - _, err := executorPrepare(ctx, executor, sql, map[string]*querypb.BindVariable{ + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorPrepare(ctx, executor, session, sql, map[string]*querypb.BindVariable{ "id": sqltypes.Int64BindVariable(1), }) require.NoError(t, err) @@ -3126,7 +3232,10 @@ func TestSelectDatabasePrepare(t *testing.T) { defer executor.queryLogger.Unsubscribe(logChan) sql := "select database()" - _, err := executorPrepare(ctx, executor, sql, map[string]*querypb.BindVariable{}) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorPrepare(ctx, executor, session, sql, map[string]*querypb.BindVariable{}) require.NoError(t, err) } @@ -3172,7 +3281,10 @@ func TestSelectWithUnionAll(t *testing.T) { "vtg2": bv, }, }} - _, err := executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{}) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, sql, map[string]*querypb.BindVariable{}) require.NoError(t, err) utils.MustMatch(t, sbc1WantQueries, sbc1.Queries, "sbc1") utils.MustMatch(t, sbc2WantQueries, sbc2.Queries, "sbc2") @@ -3667,7 +3779,10 @@ func TestGen4SelectUnqualifiedReferenceTable(t *testing.T) { executor.pv = querypb.ExecuteOptions_Gen4 query := "select * from zip_detail" - _, err := executorExec(ctx, executor, query, nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, query, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{ { @@ -3690,7 +3805,10 @@ func TestGen4SelectQualifiedReferenceTable(t *testing.T) { executor.pv = querypb.ExecuteOptions_Gen4 query := fmt.Sprintf("select * from %s.zip_detail", KsTestSharded) - _, err := executorExec(ctx, executor, query, nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, query, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{ { @@ -3713,7 +3831,10 @@ func TestGen4JoinUnqualifiedReferenceTable(t *testing.T) { executor.pv = querypb.ExecuteOptions_Gen4 query := "select * from user join zip_detail on user.zip_detail_id = zip_detail.id" - _, err := executorExec(ctx, executor, query, nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, query, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{ { @@ -3729,7 +3850,7 @@ func TestGen4JoinUnqualifiedReferenceTable(t *testing.T) { sbc2.Queries = nil query = "select * from simple join zip_detail on simple.zip_detail_id = zip_detail.id" - _, err = executorExec(ctx, executor, query, nil) + _, err = executorExec(ctx, executor, session, query, nil) require.NoError(t, err) wantQueries = []*querypb.BoundQuery{ { @@ -3752,7 +3873,10 @@ func TestGen4CrossShardJoinQualifiedReferenceTable(t *testing.T) { executor.pv = querypb.ExecuteOptions_Gen4 query := "select user.id from user join TestUnsharded.zip_detail on user.zip_detail_id = TestUnsharded.zip_detail.id" - _, err := executorExec(ctx, executor, query, nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, err := executorExec(ctx, executor, session, query, nil) require.NoError(t, err) shardedWantQueries := []*querypb.BoundQuery{ @@ -3770,7 +3894,7 @@ func TestGen4CrossShardJoinQualifiedReferenceTable(t *testing.T) { sbc2.Queries = nil query = "select simple.id from simple join TestExecutor.zip_detail on simple.zip_detail_id = TestExecutor.zip_detail.id" - _, err = executorExec(ctx, executor, query, nil) + _, err = executorExec(ctx, executor, session, query, nil) require.NoError(t, err) unshardedWantQueries := []*querypb.BoundQuery{ { @@ -4063,7 +4187,10 @@ func TestSelectAggregationNoData(t *testing.T) { sbc.SetResults([]*sqltypes.Result{tc.sandboxRes}) sbc.Queries = nil } - qr, err := executorExec(ctx, executor, tc.sql, nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + qr, err := executorExec(ctx, executor, session, tc.sql, nil) require.NoError(t, err) assert.Equal(t, tc.expField, fmt.Sprintf("%v", qr.Fields)) assert.Equal(t, tc.expRow, fmt.Sprintf("%v", qr.Rows)) @@ -4208,7 +4335,10 @@ func TestSelectAggregationData(t *testing.T) { sbc.SetResults([]*sqltypes.Result{tc.sandboxRes}) sbc.Queries = nil } - qr, err := executorExec(ctx, executor, tc.sql, nil) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + qr, err := executorExec(ctx, executor, session, tc.sql, nil) require.NoError(t, err) assert.Equal(t, tc.expField, fmt.Sprintf("%v", qr.Fields)) assert.Equal(t, tc.expRow, fmt.Sprintf("%v", qr.Rows)) diff --git a/go/vt/vtgate/executor_set_test.go b/go/vt/vtgate/executor_set_test.go index 6041275030c..17d5fa13038 100644 --- a/go/vt/vtgate/executor_set_test.go +++ b/go/vt/vtgate/executor_set_test.go @@ -274,7 +274,7 @@ func TestExecutorSet(t *testing.T) { for i, tcase := range testcases { t.Run(fmt.Sprintf("%d-%s", i, tcase.in), func(t *testing.T) { session := NewSafeSession(&vtgatepb.Session{Autocommit: true}) - _, err := executorEnv.Execute(context.Background(), nil, "TestExecute", session, tcase.in, nil) + _, err := executorEnv.Execute(ctx, nil, "TestExecute", session, tcase.in, nil) if tcase.err == "" { require.NoError(t, err) utils.MustMatch(t, tcase.out, session.Session, "new executor") @@ -374,11 +374,13 @@ func TestExecutorSetOp(t *testing.T) { }} for _, tcase := range testcases { t.Run(tcase.in, func(t *testing.T) { - session := NewAutocommitSession(primarySession) + session := NewAutocommitSession(&vtgatepb.Session{ + TargetString: "@primary", + }) session.TargetString = KsTestUnsharded session.EnableSystemSettings = !tcase.disallowResConn sbclookup.SetResults([]*sqltypes.Result{tcase.result}) - _, err := executor.Execute(context.Background(), nil, "TestExecute", session, tcase.in, nil) + _, err := executor.Execute(ctx, nil, "TestExecute", session, tcase.in, nil) require.NoError(t, err) utils.MustMatch(t, tcase.warning, session.Warnings, "") utils.MustMatch(t, tcase.sysVars, session.SystemVariables, "") @@ -395,7 +397,7 @@ func TestExecutorSetMetadata(t *testing.T) { session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) set := "set @@vitess_metadata.app_keyspace_v1= '1'" - _, err := executor.Execute(context.Background(), nil, "TestExecute", session, set, nil) + _, err := executor.Execute(ctx, nil, "TestExecute", session, set, nil) assert.Equalf(t, vtrpcpb.Code_PERMISSION_DENIED, vterrors.Code(err), "expected error %v, got error: %v", vtrpcpb.Code_PERMISSION_DENIED, err) vschemaacl.AuthorizedDDLUsers = "%" @@ -408,11 +410,11 @@ func TestExecutorSetMetadata(t *testing.T) { session = NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) set = "set @@vitess_metadata.app_keyspace_v1= '1'" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, set, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, set, nil) assert.NoError(t, err, "%s error: %v", set, err) show := `show vitess_metadata variables like 'app\\_keyspace\\_v_'` - result, err := executor.Execute(context.Background(), nil, "TestExecute", session, show, nil) + result, err := executor.Execute(ctx, nil, "TestExecute", session, show, nil) assert.NoError(t, err) want := "1" @@ -421,11 +423,11 @@ func TestExecutorSetMetadata(t *testing.T) { // Update metadata set = "set @@vitess_metadata.app_keyspace_v2='2'" - _, err = executor.Execute(context.Background(), nil, "TestExecute", session, set, nil) + _, err = executor.Execute(ctx, nil, "TestExecute", session, set, nil) assert.NoError(t, err, "%s error: %v", set, err) show = `show vitess_metadata variables like 'app\\_keyspace\\_v%'` - gotqr, err := executor.Execute(context.Background(), nil, "TestExecute", session, show, nil) + gotqr, err := executor.Execute(ctx, nil, "TestExecute", session, show, nil) assert.NoError(t, err) wantqr := &sqltypes.Result{ @@ -441,7 +443,7 @@ func TestExecutorSetMetadata(t *testing.T) { assert.ElementsMatch(t, wantqr.Rows, gotqr.Rows) show = "show vitess_metadata variables" - gotqr, err = executor.Execute(context.Background(), nil, "TestExecute", session, show, nil) + gotqr, err = executor.Execute(ctx, nil, "TestExecute", session, show, nil) require.NoError(t, err) assert.Equal(t, wantqr.Fields, gotqr.Fields) @@ -472,7 +474,7 @@ func TestPlanExecutorSetUDV(t *testing.T) { for _, tcase := range testcases { t.Run(tcase.in, func(t *testing.T) { session := NewSafeSession(&vtgatepb.Session{Autocommit: true}) - _, err := executor.Execute(context.Background(), nil, "TestExecute", session, tcase.in, nil) + _, err := executor.Execute(ctx, nil, "TestExecute", session, tcase.in, nil) if err != nil { require.EqualError(t, err, tcase.err) } else { @@ -497,15 +499,12 @@ func TestSetUDVFromTabletInput(t *testing.T) { ), }) - primarySession.TargetString = "TestExecutor" - defer func() { - primarySession.TargetString = "" - }() - _, err := executorExec(ctx, executor, "set @foo = concat('a','b','c')", nil) + session := &vtgatepb.Session{TargetString: "TestExecutor"} + _, err := executorExec(ctx, executor, session, "set @foo = concat('a','b','c')", nil) require.NoError(t, err) want := map[string]*querypb.BindVariable{"foo": sqltypes.StringBindVariable("abc")} - utils.MustMatch(t, want, primarySession.UserDefinedVariables, "") + utils.MustMatch(t, want, session.UserDefinedVariables, "") } func createMap(keys []string, values []any) map[string]*querypb.BindVariable { @@ -539,7 +538,7 @@ func TestSetVar(t *testing.T) { sqltypes.MakeTestFields("orig|new", "varchar|varchar"), "|only_full_group_by")}) - _, err := executor.Execute(context.Background(), nil, "TestSetVar", session, "set @@sql_mode = only_full_group_by", map[string]*querypb.BindVariable{}) + _, err := executor.Execute(ctx, nil, "TestSetVar", session, "set @@sql_mode = only_full_group_by", map[string]*querypb.BindVariable{}) require.NoError(t, err) tcases := []struct { @@ -561,7 +560,7 @@ func TestSetVar(t *testing.T) { // reset reserved conn need. session.SetReservedConn(false) - _, err = executor.Execute(context.Background(), nil, "TestSetVar", session, tc.sql, map[string]*querypb.BindVariable{}) + _, err = executor.Execute(ctx, nil, "TestSetVar", session, tc.sql, map[string]*querypb.BindVariable{}) require.NoError(t, err) assert.Equal(t, tc.rc, session.InReservedConn()) }) @@ -591,11 +590,11 @@ func TestSetVarShowVariables(t *testing.T) { sqltypes.MakeTestResult(sqltypes.MakeTestFields("Variable_name|Value", "varchar|varchar"), "sql_mode|ONLY_FULL_GROUP_BY,STRICT_TRANS_TABLES,NO_ZERO_IN_DATE")}) - _, err := executor.Execute(context.Background(), nil, "TestSetVar", session, "set @@sql_mode = only_full_group_by", map[string]*querypb.BindVariable{}) + _, err := executor.Execute(ctx, nil, "TestSetVar", session, "set @@sql_mode = only_full_group_by", map[string]*querypb.BindVariable{}) require.NoError(t, err) // this should return the updated value of sql_mode. - qr, err := executor.Execute(context.Background(), nil, "TestSetVar", session, "show variables like 'sql_mode'", map[string]*querypb.BindVariable{}) + qr, err := executor.Execute(ctx, nil, "TestSetVar", session, "show variables like 'sql_mode'", map[string]*querypb.BindVariable{}) require.NoError(t, err) assert.False(t, session.InReservedConn(), "reserved connection should not be used") assert.Equal(t, `[[VARCHAR("sql_mode") VARCHAR("only_full_group_by")]]`, fmt.Sprintf("%v", qr.Rows)) @@ -642,7 +641,7 @@ func TestExecutorSetAndSelect(t *testing.T) { sqltypes.MakeTestResult(nil)}) // third one for new set query setQ := fmt.Sprintf("set %s = '%s'", tcase.sysVar, tcase.val) - _, err := e.Execute(context.Background(), nil, "TestExecutorSetAndSelect", session, setQ, nil) + _, err := e.Execute(ctx, nil, "TestExecutorSetAndSelect", session, setQ, nil) require.NoError(t, err) } @@ -650,7 +649,7 @@ func TestExecutorSetAndSelect(t *testing.T) { // if the query reaches the shard, it will return REPEATABLE-READ isolation level. sbc.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult(sqltypes.MakeTestFields(tcase.sysVar, "varchar"), "REPEATABLE-READ")}) - qr, err := e.Execute(context.Background(), nil, "TestExecutorSetAndSelect", session, selectQ, nil) + qr, err := e.Execute(ctx, nil, "TestExecutorSetAndSelect", session, selectQ, nil) require.NoError(t, err) assert.Equal(t, tcase.exp, fmt.Sprintf("%v", qr.Rows)) }) diff --git a/go/vt/vtgate/executor_stream_test.go b/go/vt/vtgate/executor_stream_test.go index fa142d28d7b..065527b7ce2 100644 --- a/go/vt/vtgate/executor_stream_test.go +++ b/go/vt/vtgate/executor_stream_test.go @@ -24,6 +24,7 @@ import ( "vitess.io/vitess/go/streamlog" "vitess.io/vitess/go/test/utils" querypb "vitess.io/vitess/go/vt/proto/query" + vtgatepb "vitess.io/vitess/go/vt/proto/vtgate" "vitess.io/vitess/go/vt/vtgate/logstats" "vitess.io/vitess/go/cache" @@ -99,11 +100,12 @@ func executorStreamMessages(executor *Executor, sql string) (qr *sqltypes.Result results := make(chan *sqltypes.Result, 100) ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) defer cancel() + session := &vtgatepb.Session{TargetString: "@primary"} err = executor.StreamExecute( ctx, nil, "TestExecuteStream", - NewSafeSession(primarySession), + NewSafeSession(session), sql, nil, func(qr *sqltypes.Result) error { diff --git a/go/vt/vtgate/executor_test.go b/go/vt/vtgate/executor_test.go index 3250783261b..b2ce904e6e3 100644 --- a/go/vt/vtgate/executor_test.go +++ b/go/vt/vtgate/executor_test.go @@ -2008,10 +2008,12 @@ func TestPassthroughDDL(t *testing.T) { executor, sbc1, sbc2, _ := createExecutorEnv(ctx) defer executor.Close() - primarySession.TargetString = "TestExecutor" + session := &vtgatepb.Session{ + TargetString: "TestExecutor", + } alterDDL := "/* leading */ alter table passthrough_ddl add columne col bigint default 123 /* trailing */" - _, err := executorExec(ctx, executor, alterDDL, nil) + _, err := executorExec(ctx, executor, session, alterDDL, nil) require.NoError(t, err) wantQueries := []*querypb.BoundQuery{{ Sql: alterDDL, @@ -2027,23 +2029,22 @@ func TestPassthroughDDL(t *testing.T) { sbc2.Queries = nil // Force the query to go to only one shard. Normalization doesn't make any difference. - primarySession.TargetString = "TestExecutor/40-60" + session.TargetString = "TestExecutor/40-60" executor.normalize = true - _, err = executorExec(ctx, executor, alterDDL, nil) + _, err = executorExec(ctx, executor, session, alterDDL, nil) require.NoError(t, err) require.Nil(t, sbc1.Queries) if !reflect.DeepEqual(sbc2.Queries, wantQueries) { t.Errorf("sbc2.Queries: %+v, want %+v\n", sbc2.Queries, wantQueries) } sbc2.Queries = nil - primarySession.TargetString = "" // Use range query - primarySession.TargetString = "TestExecutor[-]" + session.TargetString = "TestExecutor[-]" executor.normalize = true - _, err = executorExec(ctx, executor, alterDDL, nil) + _, err = executorExec(ctx, executor, session, alterDDL, nil) require.NoError(t, err) if !reflect.DeepEqual(sbc1.Queries, wantQueries) { t.Errorf("sbc2.Queries: %+v, want %+v\n", sbc1.Queries, wantQueries) @@ -2052,7 +2053,6 @@ func TestPassthroughDDL(t *testing.T) { t.Errorf("sbc2.Queries: %+v, want %+v\n", sbc2.Queries, wantQueries) } sbc2.Queries = nil - primarySession.TargetString = "" } func TestParseEmptyTargetSingleKeyspace(t *testing.T) { @@ -2393,14 +2393,17 @@ func TestExecutorVExplain(t *testing.T) { defer executor.queryLogger.Unsubscribe(logChan) bindVars := map[string]*querypb.BindVariable{} - result, err := executorExec(ctx, executor, "vexplain plan select * from user", bindVars) + session := &vtgatepb.Session{ + TargetString: "@primary", + } + result, err := executorExec(ctx, executor, session, "vexplain plan select * from user", bindVars) require.NoError(t, err) require.Equal(t, `[[VARCHAR("{\n\t\"OperatorType\": \"Route\",\n\t\"Variant\": \"Scatter\",\n\t\"Keyspace\": {\n\t\t\"Name\": \"TestExecutor\",\n\t\t\"Sharded\": true\n\t},\n\t\"FieldQuery\": \"select * from `+"`user`"+` where 1 != 1\",\n\t\"Query\": \"select * from `+"`user`"+`\",\n\t\"Table\": \"`+"`user`"+`\"\n}")]]`, fmt.Sprintf("%v", result.Rows)) - result, err = executorExec(ctx, executor, "vexplain plan select 42", bindVars) + result, err = executorExec(ctx, executor, session, "vexplain plan select 42", bindVars) require.NoError(t, err) expected := `[[VARCHAR("{\n\t\"OperatorType\": \"Projection\",\n\t\"Expressions\": [\n\t\t\"INT64(42) as 42\"\n\t],\n\t\"Inputs\": [\n\t\t{\n\t\t\t\"OperatorType\": \"SingleRow\"\n\t\t}\n\t]\n}")]]` require.Equal(t, expected, fmt.Sprintf("%v", result.Rows)) diff --git a/go/vt/vtgate/queryz_test.go b/go/vt/vtgate/queryz_test.go index 1f140887ad4..1f46d077e3d 100644 --- a/go/vt/vtgate/queryz_test.go +++ b/go/vt/vtgate/queryz_test.go @@ -28,6 +28,7 @@ import ( "github.com/stretchr/testify/require" "vitess.io/vitess/go/test/utils" + vtgatepb "vitess.io/vitess/go/vt/proto/vtgate" "vitess.io/vitess/go/sqltypes" "vitess.io/vitess/go/vt/vtgate/engine" @@ -46,9 +47,10 @@ func TestQueryzHandler(t *testing.T) { executor, _, _, _ := createExecutorEnv(ctx) defer executor.Close() + session := &vtgatepb.Session{TargetString: "@primary"} // single shard query sql := "select id from user where id = 1" - _, err := executorExec(ctx, executor, sql, nil) + _, err := executorExec(ctx, executor, session, sql, nil) require.NoError(t, err) executor.plans.Wait() plan1 := assertCacheContains(t, executor, nil, "select id from `user` where id = 1") @@ -56,14 +58,14 @@ func TestQueryzHandler(t *testing.T) { // scatter sql = "select id from user" - _, err = executorExec(ctx, executor, sql, nil) + _, err = executorExec(ctx, executor, session, sql, nil) require.NoError(t, err) executor.plans.Wait() plan2 := assertCacheContains(t, executor, nil, "select id from `user`") plan2.ExecTime = uint64(1 * time.Second) sql = "insert into user (id, name) values (:id, :name)" - _, err = executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{ + _, err = executorExec(ctx, executor, session, sql, map[string]*querypb.BindVariable{ "id": sqltypes.Uint64BindVariable(1), "name": sqltypes.BytesBindVariable([]byte("myname")), }) @@ -76,7 +78,7 @@ func TestQueryzHandler(t *testing.T) { // same query again should add query counts to existing plans sql = "insert into user (id, name) values (:id, :name)" - _, err = executorExec(ctx, executor, sql, map[string]*querypb.BindVariable{ + _, err = executorExec(ctx, executor, session, sql, map[string]*querypb.BindVariable{ "id": sqltypes.Uint64BindVariable(1), "name": sqltypes.BytesBindVariable([]byte("myname")), }) diff --git a/go/vt/vtgate/vtgate_test.go b/go/vt/vtgate/vtgate_test.go index cb52ac31af9..7bd6fa3353e 100644 --- a/go/vt/vtgate/vtgate_test.go +++ b/go/vt/vtgate/vtgate_test.go @@ -47,39 +47,6 @@ var executeOptions = &querypb.ExecuteOptions{ IncludedFields: querypb.ExecuteOptions_TYPE_ONLY, } -var primarySession *vtgatepb.Session - -func init() { - createSandbox(KsTestUnsharded).VSchema = ` -{ - "sharded": false, - "tables": { - "t1": {} - } -} -` - createSandbox(KsTestBadVSchema).VSchema = ` -{ - "sharded": true, - "tables": { - "t2": { - "auto_increment": { - "column": "id", - "sequence": "id_seq" - } - } - } -} -` - hcVTGateTest = discovery.NewFakeHealthCheck(nil) - transactionMode = "MULTI" - Init(context.Background(), hcVTGateTest, newSandboxForCells([]string{"aa"}), "aa", nil, querypb.ExecuteOptions_Gen4) - - mysqlServerPort = 0 - mysqlAuthServerImpl = "none" - initMySQLProtocol() -} - func TestVTGateExecute(t *testing.T) { counts := rpcVTGate.timings.Timings.Counts() @@ -367,10 +334,14 @@ func testErrorPropagation(t *testing.T, sbcs []*sandboxconn.SandboxConn, before for _, sbc := range sbcs { before(sbc) } + session := &vtgatepb.Session{ + TargetString: "@primary", + } + _, _, err := rpcVTGate.Execute( context.Background(), nil, - primarySession, + session, "select id from t1", nil, ) @@ -393,7 +364,7 @@ func testErrorPropagation(t *testing.T, sbcs []*sandboxconn.SandboxConn, before _, err = rpcVTGate.StreamExecute( context.Background(), nil, - primarySession, + session, "select id from t1", nil, func(r *sqltypes.Result) error { @@ -420,10 +391,6 @@ func testErrorPropagation(t *testing.T, sbcs []*sandboxconn.SandboxConn, before func TestErrorPropagation(t *testing.T) { createSandbox(KsTestUnsharded) hcVTGateTest.Reset() - // create a new session each time so that ShardSessions don't get re-used across tests - primarySession = &vtgatepb.Session{ - TargetString: "@primary", - } sbcm := hcVTGateTest.AddTestTablet("aa", "1.1.1.1", 1001, KsTestUnsharded, "0", topodatapb.TabletType_PRIMARY, true, 1, nil) sbcrdonly := hcVTGateTest.AddTestTablet("aa", "1.1.1.2", 1001, KsTestUnsharded, "0", topodatapb.TabletType_RDONLY, true, 1, nil) @@ -651,8 +618,8 @@ func TestMultiInternalSavepointVtGate(t *testing.T) { sbc2 := hcVTGateTest.AddTestTablet("aa", "1.1.1.1", 2, KsTestSharded, "40-80", topodatapb.TabletType_PRIMARY, true, 1, nil) sbc3 := hcVTGateTest.AddTestTablet("aa", "1.1.1.1", 3, KsTestSharded, "80-", topodatapb.TabletType_PRIMARY, true, 1, nil) - logChan := QueryLogger.Subscribe("Test") - defer QueryLogger.Unsubscribe(logChan) + logChan := rpcVTGate.executor.queryLogger.Subscribe("Test") + defer rpcVTGate.executor.queryLogger.Unsubscribe(logChan) session := &vtgatepb.Session{Autocommit: true} require.True(t, session.GetAutocommit()) @@ -724,12 +691,12 @@ func TestMultiInternalSavepointVtGate(t *testing.T) { }} assertQueriesWithSavepoint(t, sbc2, wantQ) - testQueryLog(t, logChan, "Execute", "BEGIN", "begin", 0) - testQueryLog(t, logChan, "MarkSavepoint", "SAVEPOINT", "savepoint x", 0) - testQueryLog(t, logChan, "Execute", "INSERT", "insert into sp_tbl(user_id) values (:vtg1 /* INT64 */), (:vtg2 /* INT64 */)", 2) - testQueryLog(t, logChan, "MarkSavepoint", "SAVEPOINT", "savepoint y", 2) - testQueryLog(t, logChan, "Execute", "INSERT", "insert into sp_tbl(user_id) values (:vtg1 /* INT64 */), (:vtg2 /* INT64 */)", 2) - testQueryLog(t, logChan, "Execute", "INSERT", "insert into sp_tbl(user_id) values (:vtg1 /* INT64 */)", 1) + testQueryLog(t, rpcVTGate.executor, logChan, "Execute", "BEGIN", "begin", 0) + testQueryLog(t, rpcVTGate.executor, logChan, "MarkSavepoint", "SAVEPOINT", "savepoint x", 0) + testQueryLog(t, rpcVTGate.executor, logChan, "Execute", "INSERT", "insert into sp_tbl(user_id) values (:vtg1 /* INT64 */), (:vtg2 /* INT64 */)", 2) + testQueryLog(t, rpcVTGate.executor, logChan, "MarkSavepoint", "SAVEPOINT", "savepoint y", 2) + testQueryLog(t, rpcVTGate.executor, logChan, "Execute", "INSERT", "insert into sp_tbl(user_id) values (:vtg1 /* INT64 */), (:vtg2 /* INT64 */)", 2) + testQueryLog(t, rpcVTGate.executor, logChan, "Execute", "INSERT", "insert into sp_tbl(user_id) values (:vtg1 /* INT64 */)", 1) } func TestVSchemaVindexUnknownParams(t *testing.T) { From 47b0e3a6969d758bad629f70d688279efde57a23 Mon Sep 17 00:00:00 2001 From: Dirkjan Bussink Date: Tue, 22 Aug 2023 13:40:45 +0200 Subject: [PATCH 03/17] Pass in context for all topos Signed-off-by: Dirkjan Bussink --- go/cmd/vtcombo/main.go | 4 +- go/cmd/vtgate/vtgate.go | 2 +- go/cmd/vttablet/vttablet.go | 6 +- go/test/fuzzing/vtctl_fuzzer.go | 2 +- go/vt/discovery/healthcheck_test.go | 134 +++- go/vt/discovery/keyspace_events_test.go | 23 +- go/vt/discovery/replicationlag_test.go | 4 + go/vt/discovery/tablet_picker_test.go | 162 +++-- go/vt/discovery/topology_watcher.go | 3 +- go/vt/discovery/topology_watcher_test.go | 33 +- go/vt/mysqlctl/backup_blackbox_test.go | 44 +- go/vt/schemamanager/schemamanager_test.go | 5 +- go/vt/schemamanager/tablet_executor_test.go | 5 +- go/vt/srvtopo/discover_test.go | 7 +- .../srvtopo/keyspace_filtering_server_test.go | 41 +- go/vt/srvtopo/resilient_server.go | 7 +- go/vt/srvtopo/resilient_server_test.go | 73 +- go/vt/srvtopo/resolver_test.go | 14 +- go/vt/srvtopo/watch.go | 17 +- go/vt/srvtopo/watch_srvkeyspace.go | 14 +- go/vt/srvtopo/watch_srvvschema.go | 15 +- go/vt/throttler/demo/throttler_demo.go | 8 +- go/vt/topo/cell_info.go | 3 + go/vt/topo/consultopo/server_flaky_test.go | 12 +- go/vt/topo/etcd2topo/server_test.go | 4 +- go/vt/topo/faketopo/faketopo.go | 4 +- go/vt/topo/helpers/copy_test.go | 7 +- go/vt/topo/helpers/tee_topo_test.go | 11 +- go/vt/topo/memorytopo/memorytopo.go | 6 +- go/vt/topo/memorytopo/server_test.go | 7 +- go/vt/topo/test/directory.go | 7 +- go/vt/topo/test/election.go | 12 +- go/vt/topo/test/file.go | 7 +- go/vt/topo/test/keyspace.go | 6 +- go/vt/topo/test/lock.go | 6 +- go/vt/topo/test/replication.go | 6 +- go/vt/topo/test/serving.go | 11 +- go/vt/topo/test/shard.go | 6 +- go/vt/topo/test/tablet.go | 7 +- go/vt/topo/test/testing.go | 45 +- go/vt/topo/test/trylock.go | 3 +- go/vt/topo/test/vschema.go | 10 +- go/vt/topo/test/watch.go | 11 +- go/vt/topo/topotests/cell_info_test.go | 33 +- go/vt/topo/topotests/cells_aliases_test.go | 6 +- go/vt/topo/topotests/keyspace_test.go | 12 +- go/vt/topo/topotests/replication_test.go | 6 +- go/vt/topo/topotests/shard_watch_test.go | 18 +- go/vt/topo/topotests/srv_keyspace_test.go | 54 +- go/vt/topo/topotests/srv_vschema_test.go | 6 +- go/vt/topo/topotests/tablet_test.go | 6 +- go/vt/topo/topotests/wildcards_test.go | 18 +- go/vt/topo/zk2topo/server_test.go | 7 +- go/vt/topotools/routing_rules_test.go | 18 +- go/vt/topotools/shard_test.go | 28 +- go/vt/vtadmin/api_test.go | 74 +- go/vt/vtadmin/cluster/cluster_test.go | 57 +- go/vt/vtadmin/testutil/cluster.go | 4 +- go/vt/vtcombo/tablet_map.go | 2 +- go/vt/vtctl/endtoend/get_schema_test.go | 5 +- go/vt/vtctl/endtoend/onlineddl_show_test.go | 6 +- go/vt/vtctl/grpcvtctlclient/client_test.go | 9 +- go/vt/vtctl/grpcvtctldclient/client_test.go | 18 +- .../endtoend/init_shard_primary_test.go | 8 +- .../grpcvtctldserver/server_slow_test.go | 35 +- go/vt/vtctl/grpcvtctldserver/server_test.go | 643 ++++++++++-------- go/vt/vtctl/grpcvtctldserver/testutil/util.go | 2 + .../reparentutil/emergency_reparenter_test.go | 120 ++-- .../planned_reparenter_flaky_test.go | 114 ++-- go/vt/vtctl/schematools/reload_test.go | 6 +- go/vt/vtctl/vdiff2_test.go | 8 +- go/vt/vtctl/vdiff_env_test.go | 4 +- go/vt/vtctl/vtctl_env_test.go | 4 +- go/vt/vtctl/vtctl_test.go | 10 +- go/vt/vtctl/vtctlclienttest/client.go | 7 +- go/vt/vtctl/workflow/materializer_env_test.go | 4 +- go/vt/vtctl/workflow/materializer_test.go | 16 +- go/vt/vtctld/api_test.go | 6 +- go/vt/vtctld/explorer_test.go | 34 +- go/vt/vtctld/tablet_data_test.go | 11 +- go/vt/vtexplain/vtexplain_vtgate.go | 6 +- go/vt/vtexplain/vtexplain_vttablet.go | 4 +- go/vt/vtexplain/vtexplain_vttablet_test.go | 12 +- go/vt/vtgate/autocommit_test.go | 1 + go/vt/vtgate/executor.go | 1 + go/vt/vtgate/executor_scatter_stats_test.go | 1 + go/vt/vtgate/executor_select_test.go | 1 + go/vt/vtgate/planbuilder/plan_test.go | 8 +- go/vt/vtgate/plugin_mysql_server.go | 8 +- go/vt/vtgate/plugin_mysql_server_test.go | 1 + go/vt/vtgate/queryz_test.go | 1 + go/vt/vtgate/schema/tracker_test.go | 6 +- go/vt/vtgate/tabletgateway_flaky_test.go | 1 + go/vt/vtgate/tabletgateway_test.go | 1 + go/vt/vtgate/vtgate.go | 1 + .../logic/keyspace_shard_discovery_test.go | 32 +- go/vt/vtorc/logic/tablet_discovery_test.go | 9 +- go/vt/vtorc/logic/topology_recovery_test.go | 10 +- .../topocustomrule/topocustomrule_test.go | 6 +- .../endtoend/connkilling/main_test.go | 5 +- go/vt/vttablet/endtoend/framework/server.go | 10 +- go/vt/vttablet/endtoend/main_test.go | 5 +- .../endtoend/streamtimeout/main_test.go | 6 +- .../vttablet/tabletmanager/framework_test.go | 4 +- .../tabletmanager/rpc_vreplication_test.go | 20 +- .../vttablet/tabletmanager/shard_sync_test.go | 5 +- go/vt/vttablet/tabletmanager/tm_init_test.go | 49 +- go/vt/vttablet/tabletmanager/tm_state_test.go | 57 +- .../tabletmanager/vdiff/framework_test.go | 4 +- .../vreplication/framework_test.go | 16 +- go/vt/vttablet/tabletserver/bench_test.go | 15 +- .../tabletserver/health_streamer_test.go | 4 + .../tabletserver/query_executor_test.go | 6 +- .../tabletserver/state_manager_test.go | 4 + .../stateful_connection_pool_test.go | 18 +- go/vt/vttablet/tabletserver/tabletserver.go | 8 +- .../tabletserver/tabletserver_test.go | 262 +++++-- go/vt/vttablet/tabletserver/twopc_test.go | 13 +- go/vt/vttablet/tabletserver/tx_engine_test.go | 6 + .../vttablet/tabletserver/tx_executor_test.go | 118 ++-- go/vt/vttablet/tabletserver/tx_pool_test.go | 62 +- .../txthrottler/tx_throttler_test.go | 10 +- .../tabletserver/vstreamer/main_flaky_test.go | 9 +- .../tabletserver/vstreamer/testenv/testenv.go | 7 +- .../vstreamer/vstreamer_flaky_test.go | 18 +- go/vt/wrangler/external_cluster_test.go | 5 +- go/vt/wrangler/materializer_env_test.go | 4 +- go/vt/wrangler/materializer_test.go | 239 +++++-- go/vt/wrangler/resharder_env_test.go | 4 +- go/vt/wrangler/resharder_test.go | 69 +- go/vt/wrangler/tablet_test.go | 25 +- go/vt/wrangler/testlib/backup_test.go | 20 +- .../testlib/copy_schema_shard_test.go | 4 +- .../testlib/emergency_reparent_shard_test.go | 6 +- .../testlib/external_reparent_test.go | 30 +- go/vt/wrangler/testlib/find_tablet_test.go | 9 +- go/vt/wrangler/testlib/permissions_test.go | 5 +- .../testlib/planned_reparent_shard_test.go | 36 +- go/vt/wrangler/testlib/reparent_utils_test.go | 15 +- go/vt/wrangler/testlib/shard_test.go | 8 +- go/vt/wrangler/testlib/version_test.go | 5 +- go/vt/wrangler/testlib/vtctl_topo_test.go | 5 +- go/vt/wrangler/traffic_switcher_env_test.go | 6 +- go/vt/wrangler/vdiff_env_test.go | 4 +- go/vt/wrangler/vdiff_test.go | 20 +- go/vt/wrangler/vexec_test.go | 25 +- go/vt/wrangler/wrangler_env_test.go | 4 +- 147 files changed, 2242 insertions(+), 1390 deletions(-) diff --git a/go/cmd/vtcombo/main.go b/go/cmd/vtcombo/main.go index 1e2483428d7..02c8acbab7f 100644 --- a/go/cmd/vtcombo/main.go +++ b/go/cmd/vtcombo/main.go @@ -186,7 +186,7 @@ func main() { ts = topo.Open() } else { // Create topo server. We use a 'memorytopo' implementation. - ts = memorytopo.NewServer(tpb.Cells...) + ts = memorytopo.NewServer(context.Background(), tpb.Cells...) } // attempt to load any routing rules specified by tpb @@ -279,7 +279,7 @@ func main() { } // vtgate configuration and init - resilientServer = srvtopo.NewResilientServer(ts, "ResilientSrvTopoServer") + resilientServer = srvtopo.NewResilientServer(context.Background(), ts, "ResilientSrvTopoServer") tabletTypesToWait := []topodatapb.TabletType{ topodatapb.TabletType_PRIMARY, topodatapb.TabletType_REPLICA, diff --git a/go/cmd/vtgate/vtgate.go b/go/cmd/vtgate/vtgate.go index d98353bcd17..27d28aa35a8 100644 --- a/go/cmd/vtgate/vtgate.go +++ b/go/cmd/vtgate/vtgate.go @@ -130,7 +130,7 @@ func main() { ts := topo.Open() defer ts.Close() - resilientServer = srvtopo.NewResilientServer(ts, "ResilientSrvTopoServer") + resilientServer = srvtopo.NewResilientServer(context.Background(), ts, "ResilientSrvTopoServer") tabletTypes := make([]topodatapb.TabletType, 0, 1) if len(tabletTypesToWait) != 0 { diff --git a/go/cmd/vttablet/vttablet.go b/go/cmd/vttablet/vttablet.go index f879d2728e1..34f4ab77d64 100644 --- a/go/cmd/vttablet/vttablet.go +++ b/go/cmd/vttablet/vttablet.go @@ -95,7 +95,7 @@ func main() { config, mycnf := initConfig(tabletAlias) ts := topo.Open() - qsc := createTabletServer(config, ts, tabletAlias) + qsc := createTabletServer(context.Background(), config, ts, tabletAlias) mysqld := mysqlctl.NewMysqld(config.DB) servenv.OnClose(mysqld.Close) @@ -204,7 +204,7 @@ func extractOnlineDDL() error { return nil } -func createTabletServer(config *tabletenv.TabletConfig, ts *topo.Server, tabletAlias *topodatapb.TabletAlias) *tabletserver.TabletServer { +func createTabletServer(ctx context.Context, config *tabletenv.TabletConfig, ts *topo.Server, tabletAlias *topodatapb.TabletAlias) *tabletserver.TabletServer { if tableACLConfig != "" { // To override default simpleacl, other ACL plugins must set themselves to be default ACL factory tableacl.Register("simpleacl", &simpleacl.Factory{}) @@ -212,7 +212,7 @@ func createTabletServer(config *tabletenv.TabletConfig, ts *topo.Server, tabletA log.Exit("table acl config has to be specified with table-acl-config flag because enforce-tableacl-config is set.") } // creates and registers the query service - qsc := tabletserver.NewTabletServer("", config, ts, tabletAlias) + qsc := tabletserver.NewTabletServer(ctx, "", config, ts, tabletAlias) servenv.OnRun(func() { qsc.Register() addStatusParts(qsc) diff --git a/go/test/fuzzing/vtctl_fuzzer.go b/go/test/fuzzing/vtctl_fuzzer.go index aed11774cc8..82fdaa572de 100644 --- a/go/test/fuzzing/vtctl_fuzzer.go +++ b/go/test/fuzzing/vtctl_fuzzer.go @@ -189,6 +189,6 @@ func Fuzz(data []byte) int { } func createTopo(ctx context.Context) (*topo.Server, error) { - ts := memorytopo.NewServer("zone1", "zone2", "zone3") + ts := memorytopo.NewServer(ctx, "zone1", "zone2", "zone3") return ts, nil } diff --git a/go/vt/discovery/healthcheck_test.go b/go/vt/discovery/healthcheck_test.go index ea78b5c2647..2a5f5b18505 100644 --- a/go/vt/discovery/healthcheck_test.go +++ b/go/vt/discovery/healthcheck_test.go @@ -64,10 +64,14 @@ func init() { } func TestHealthCheck(t *testing.T) { + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() // reset error counters hcErrorCounters.ResetAll() - ts := memorytopo.NewServer("cell") - hc := createTestHc(ts) + ts := memorytopo.NewServer(ctx, "cell") + defer ts.Close() + hc := createTestHc(ctx, ts) // close healthcheck defer hc.Close() tablet := createTestTablet(0, "cell", "a") @@ -206,8 +210,13 @@ func TestHealthCheck(t *testing.T) { } func TestHealthCheckStreamError(t *testing.T) { - ts := memorytopo.NewServer("cell") - hc := createTestHc(ts) + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "cell") + defer ts.Close() + hc := createTestHc(ctx, ts) defer hc.Close() tablet := createTestTablet(0, "cell", "a") @@ -267,8 +276,13 @@ func TestHealthCheckStreamError(t *testing.T) { // TestHealthCheckErrorOnPrimary is the same as TestHealthCheckStreamError except for tablet type func TestHealthCheckErrorOnPrimary(t *testing.T) { - ts := memorytopo.NewServer("cell") - hc := createTestHc(ts) + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "cell") + defer ts.Close() + hc := createTestHc(ctx, ts) defer hc.Close() tablet := createTestTablet(0, "cell", "a") @@ -327,8 +341,13 @@ func TestHealthCheckErrorOnPrimary(t *testing.T) { } func TestHealthCheckErrorOnPrimaryAfterExternalReparent(t *testing.T) { - ts := memorytopo.NewServer("cell") - hc := createTestHc(ts) + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "cell") + defer ts.Close() + hc := createTestHc(ctx, ts) defer hc.Close() resultChan := hc.Subscribe() @@ -405,8 +424,13 @@ func TestHealthCheckErrorOnPrimaryAfterExternalReparent(t *testing.T) { } func TestHealthCheckVerifiesTabletAlias(t *testing.T) { - ts := memorytopo.NewServer("cell") - hc := createTestHc(ts) + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "cell") + defer ts.Close() + hc := createTestHc(ctx, ts) defer hc.Close() tablet := createTestTablet(0, "cell", "a") @@ -448,8 +472,14 @@ func TestHealthCheckVerifiesTabletAlias(t *testing.T) { // TestHealthCheckCloseWaitsForGoRoutines tests that Close() waits for all Go // routines to finish and the listener won't be called anymore. func TestHealthCheckCloseWaitsForGoRoutines(t *testing.T) { - ts := memorytopo.NewServer("cell") - hc := createTestHc(ts) + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "cell") + defer ts.Close() + hc := createTestHc(ctx, ts) + defer hc.Close() tablet := createTestTablet(0, "cell", "a") input := make(chan *querypb.StreamHealthResponse, 1) createFakeConn(tablet, input) @@ -508,10 +538,15 @@ func TestHealthCheckCloseWaitsForGoRoutines(t *testing.T) { } func TestHealthCheckTimeout(t *testing.T) { + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // reset counters hcErrorCounters.ResetAll() - ts := memorytopo.NewServer("cell") - hc := createTestHc(ts) + ts := memorytopo.NewServer(ctx, "cell") + defer ts.Close() + hc := createTestHc(ctx, ts) hc.healthCheckTimeout = 500 * time.Millisecond defer hc.Close() tablet := createTestTablet(0, "cell", "a") @@ -580,8 +615,13 @@ func TestHealthCheckTimeout(t *testing.T) { } func TestWaitForAllServingTablets(t *testing.T) { - ts := memorytopo.NewServer("cell") - hc := createTestHc(ts) + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "cell") + defer ts.Close() + hc := createTestHc(ctx, ts) defer hc.Close() tablet := createTestTablet(0, "cell", "a") tablet.Type = topodatapb.TabletType_REPLICA @@ -601,7 +641,7 @@ func TestWaitForAllServingTablets(t *testing.T) { // there will be a first result, get and discard it <-resultChan // empty - ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) + ctx, cancel = context.WithTimeout(ctx, 1*time.Second) defer cancel() err := hc.WaitForAllServingTablets(ctx, targets) @@ -672,8 +712,13 @@ func TestWaitForAllServingTablets(t *testing.T) { // TestRemoveTablet tests the behavior when a tablet goes away. func TestRemoveTablet(t *testing.T) { - ts := memorytopo.NewServer("cell") - hc := createTestHc(ts) + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "cell") + defer ts.Close() + hc := createTestHc(ctx, ts) defer hc.Close() tablet := createTestTablet(0, "cell", "a") tablet.Type = topodatapb.TabletType_REPLICA @@ -779,8 +824,13 @@ func TestRemoveTablet(t *testing.T) { // TestGetHealthyTablets tests the functionality of GetHealthyTabletStats. func TestGetHealthyTablets(t *testing.T) { - ts := memorytopo.NewServer("cell") - hc := createTestHc(ts) + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "cell") + defer ts.Close() + hc := createTestHc(ctx, ts) defer hc.Close() tablet := createTestTablet(0, "cell", "a") tablet.Type = topodatapb.TabletType_REPLICA @@ -965,8 +1015,13 @@ func TestGetHealthyTablets(t *testing.T) { } func TestPrimaryInOtherCell(t *testing.T) { - ts := memorytopo.NewServer("cell1", "cell2") - hc := NewHealthCheck(context.Background(), 1*time.Millisecond, time.Hour, ts, "cell1", "cell1, cell2") + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "cell1", "cell2") + defer ts.Close() + hc := NewHealthCheck(ctx, 1*time.Millisecond, time.Hour, ts, "cell1", "cell1, cell2") defer hc.Close() // add a tablet as primary in different cell @@ -1022,8 +1077,13 @@ func TestPrimaryInOtherCell(t *testing.T) { } func TestReplicaInOtherCell(t *testing.T) { - ts := memorytopo.NewServer("cell1", "cell2") - hc := NewHealthCheck(context.Background(), 1*time.Millisecond, time.Hour, ts, "cell1", "cell1, cell2") + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "cell1", "cell2") + defer ts.Close() + hc := NewHealthCheck(ctx, 1*time.Millisecond, time.Hour, ts, "cell1", "cell1, cell2") defer hc.Close() // add a tablet as replica @@ -1124,8 +1184,13 @@ func TestReplicaInOtherCell(t *testing.T) { } func TestCellAliases(t *testing.T) { - ts := memorytopo.NewServer("cell1", "cell2") - hc := NewHealthCheck(context.Background(), 1*time.Millisecond, time.Hour, ts, "cell1", "cell1, cell2") + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "cell1", "cell2") + defer ts.Close() + hc := NewHealthCheck(ctx, 1*time.Millisecond, time.Hour, ts, "cell1", "cell1, cell2") defer hc.Close() cellsAlias := &topodatapb.CellsAlias{ @@ -1184,8 +1249,13 @@ func TestCellAliases(t *testing.T) { } func TestHealthCheckChecksGrpcPort(t *testing.T) { - ts := memorytopo.NewServer("cell") - hc := createTestHc(ts) + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "cell") + defer ts.Close() + hc := createTestHc(ctx, ts) defer hc.Close() tablet := createTestTablet(0, "cell", "a") @@ -1204,6 +1274,7 @@ func TestHealthCheckChecksGrpcPort(t *testing.T) { } func TestTemplate(t *testing.T) { + utils.EnsureNoLeaks(t) TabletURLTemplateString = "http://{{.GetTabletHostPort}}" ParseTabletURLTemplateFromFlag() @@ -1231,6 +1302,7 @@ func TestTemplate(t *testing.T) { } func TestDebugURLFormatting(t *testing.T) { + utils.EnsureNoLeaks(t) TabletURLTemplateString = "https://{{.GetHostNameLevel 0}}.bastion.{{.Tablet.Alias.Cell}}.corp" ParseTabletURLTemplateFromFlag() @@ -1270,8 +1342,8 @@ func tabletDialer(tablet *topodatapb.Tablet, _ grpcclient.FailFast) (queryservic return nil, fmt.Errorf("tablet %v not found", key) } -func createTestHc(ts *topo.Server) *HealthCheckImpl { - return NewHealthCheck(context.Background(), 1*time.Millisecond, time.Hour, ts, "cell", "") +func createTestHc(ctx context.Context, ts *topo.Server) *HealthCheckImpl { + return NewHealthCheck(ctx, 1*time.Millisecond, time.Hour, ts, "cell", "") } type fakeConn struct { diff --git a/go/vt/discovery/keyspace_events_test.go b/go/vt/discovery/keyspace_events_test.go index eb2babdd44c..a64ef39558c 100644 --- a/go/vt/discovery/keyspace_events_test.go +++ b/go/vt/discovery/keyspace_events_test.go @@ -24,6 +24,7 @@ import ( "github.com/stretchr/testify/require" + "vitess.io/vitess/go/test/utils" "vitess.io/vitess/go/vt/topo" "vitess.io/vitess/go/vt/topo/faketopo" @@ -33,15 +34,18 @@ import ( ) func TestSrvKeyspaceWithNilNewKeyspace(t *testing.T) { + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() cell := "cell" keyspace := "testks" factory := faketopo.NewFakeTopoFactory() factory.AddCell(cell) - ts := faketopo.NewFakeTopoServer(factory) + ts := faketopo.NewFakeTopoServer(ctx, factory) ts2 := &fakeTopoServer{} - hc := NewHealthCheck(context.Background(), 1*time.Millisecond, time.Hour, ts, cell, "") + hc := NewHealthCheck(ctx, 1*time.Millisecond, time.Hour, ts, cell, "") defer hc.Close() - kew := NewKeyspaceEventWatcher(context.Background(), ts2, hc, cell) + kew := NewKeyspaceEventWatcher(ctx, ts2, hc, cell) kss := &keyspaceState{ kew: kew, keyspace: keyspace, @@ -76,15 +80,18 @@ func TestSrvKeyspaceWithNilNewKeyspace(t *testing.T) { // We should never consider both as a possible cause given the same // keyspace state. func TestKeyspaceEventTypes(t *testing.T) { + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() cell := "cell" keyspace := "testks" factory := faketopo.NewFakeTopoFactory() factory.AddCell(cell) - ts := faketopo.NewFakeTopoServer(factory) + ts := faketopo.NewFakeTopoServer(ctx, factory) ts2 := &fakeTopoServer{} - hc := NewHealthCheck(context.Background(), 1*time.Millisecond, time.Hour, ts, cell, "") + hc := NewHealthCheck(ctx, 1*time.Millisecond, time.Hour, ts, cell, "") defer hc.Close() - kew := NewKeyspaceEventWatcher(context.Background(), ts2, hc, cell) + kew := NewKeyspaceEventWatcher(ctx, ts2, hc, cell) type testCase struct { name string @@ -263,10 +270,10 @@ func TestKeyspaceEventTypes(t *testing.T) { require.NotNil(t, tc.kss.shards[tc.shardToCheck], "the specified shardToCheck of %q does not exist in the shardState", tc.shardToCheck) - resharding := kew.TargetIsBeingResharded(tc.kss.shards[tc.shardToCheck].target) + resharding := kew.TargetIsBeingResharded(ctx, tc.kss.shards[tc.shardToCheck].target) require.Equal(t, resharding, tc.expectResharding, "TargetIsBeingResharded should return %t", tc.expectResharding) - _, primaryDown := kew.PrimaryIsNotServing(tc.kss.shards[tc.shardToCheck].target) + _, primaryDown := kew.PrimaryIsNotServing(ctx, tc.kss.shards[tc.shardToCheck].target) require.Equal(t, primaryDown, tc.expectPrimaryNotServing, "PrimaryIsNotServing should return %t", tc.expectPrimaryNotServing) }) } diff --git a/go/vt/discovery/replicationlag_test.go b/go/vt/discovery/replicationlag_test.go index 9c047469fb2..4a96b68336a 100644 --- a/go/vt/discovery/replicationlag_test.go +++ b/go/vt/discovery/replicationlag_test.go @@ -45,6 +45,7 @@ func testSetMinNumTablets(newMin int) { } func TestFilterByReplicationLagUnhealthy(t *testing.T) { + utils.EnsureNoLeaks(t) // 1 healthy serving tablet, 1 not healthy ts1 := &TabletHealth{ Tablet: topo.NewTablet(1, "cell", "host1"), @@ -62,6 +63,7 @@ func TestFilterByReplicationLagUnhealthy(t *testing.T) { } func TestFilterByReplicationLag(t *testing.T) { + utils.EnsureNoLeaks(t) // Use simplified logic testSetLegacyReplicationLagAlgorithm(false) @@ -138,6 +140,7 @@ func TestFilterByReplicationLag(t *testing.T) { } func TestFilterByReplicationLagThreeTabletMin(t *testing.T) { + utils.EnsureNoLeaks(t) // Use at least 3 tablets if possible testSetMinNumTablets(3) // lags of (1s, 1s, 10m, 11m) - returns at least32 items where the slightly delayed ones that are returned are the 10m and 11m ones. @@ -194,6 +197,7 @@ func TestFilterByReplicationLagThreeTabletMin(t *testing.T) { } func TestFilterStatsByReplicationLagOneTabletMin(t *testing.T) { + utils.EnsureNoLeaks(t) // Use at least 1 tablets if possible testSetMinNumTablets(1) // lags of (1s, 100m) - return only healthy tablet if that is all that is available. diff --git a/go/vt/discovery/tablet_picker_test.go b/go/vt/discovery/tablet_picker_test.go index a6d0e8f4c42..add9ab59320 100644 --- a/go/vt/discovery/tablet_picker_test.go +++ b/go/vt/discovery/tablet_picker_test.go @@ -22,6 +22,8 @@ import ( "github.com/stretchr/testify/require" "google.golang.org/protobuf/proto" + "vitess.io/vitess/go/test/utils" + "vitess.io/vitess/go/vt/topo" "vitess.io/vitess/go/vt/topo/memorytopo" @@ -30,21 +32,25 @@ import ( ) func TestPickPrimary(t *testing.T) { - te := newPickerTestEnv(t, []string{"cell", "otherCell"}) - want := addTablet(te, 100, topodatapb.TabletType_PRIMARY, "cell", true, true) - defer deleteTablet(t, te, want) + utils.EnsureNoLeaks(t) ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) defer cancel() + + te := newPickerTestEnv(t, ctx, []string{"cell", "otherCell"}) + want := addTablet(ctx, te, 100, topodatapb.TabletType_PRIMARY, "cell", true, true) + defer deleteTablet(t, te, want) + ctx, cancel = context.WithTimeout(ctx, 200*time.Millisecond) + defer cancel() _, err := te.topoServ.UpdateShardFields(ctx, te.keyspace, te.shard, func(si *topo.ShardInfo) error { si.PrimaryAlias = want.Alias return nil }) require.NoError(t, err) - tp, err := NewTabletPicker(context.Background(), te.topoServ, []string{"otherCell"}, "cell", te.keyspace, te.shard, "primary", TabletPickerOptions{}) + tp, err := NewTabletPicker(ctx, te.topoServ, []string{"otherCell"}, "cell", te.keyspace, te.shard, "primary", TabletPickerOptions{}) require.NoError(t, err) - ctx2, cancel2 := context.WithTimeout(context.Background(), 200*time.Millisecond) + ctx2, cancel2 := context.WithTimeout(ctx, 200*time.Millisecond) defer cancel2() tablet, err := tp.PickForStreaming(ctx2) require.NoError(t, err) @@ -52,6 +58,7 @@ func TestPickPrimary(t *testing.T) { } func TestPickLocalPreferences(t *testing.T) { + utils.EnsureNoLeaks(t) type tablet struct { id uint32 typ topodatapb.TabletType @@ -263,20 +270,21 @@ func TestPickLocalPreferences(t *testing.T) { }, } - ctx := context.Background() for _, tcase := range tcases { t.Run(tcase.name, func(t *testing.T) { - te := newPickerTestEnv(t, tcase.envCells) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + te := newPickerTestEnv(t, ctx, tcase.envCells) var testTablets []*topodatapb.Tablet for _, tab := range tcase.tablets { - testTablets = append(testTablets, addTablet(te, int(tab.id), tab.typ, tab.cell, true, true)) + testTablets = append(testTablets, addTablet(ctx, te, int(tab.id), tab.typ, tab.cell, true, true)) } defer func() { for _, tab := range testTablets { deleteTablet(t, te, tab) } }() - tp, err := NewTabletPicker(context.Background(), te.topoServ, tcase.inCells, tcase.localCell, te.keyspace, te.shard, tcase.inTabletTypes, tcase.options) + tp, err := NewTabletPicker(ctx, te.topoServ, tcase.inCells, tcase.localCell, te.keyspace, te.shard, tcase.inTabletTypes, tcase.options) require.NoError(t, err) require.Equal(t, tp.localCellInfo.localCell, tcase.localCell) require.ElementsMatch(t, tp.cells, tcase.tpCells) @@ -297,26 +305,28 @@ func TestPickLocalPreferences(t *testing.T) { } func TestPickCellPreferenceLocalCell(t *testing.T) { + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + // test env puts all cells into an alias called "cella" - te := newPickerTestEnv(t, []string{"cell", "otherCell"}) - want1 := addTablet(te, 100, topodatapb.TabletType_REPLICA, "cell", true, true) + te := newPickerTestEnv(t, ctx, []string{"cell", "otherCell"}) + want1 := addTablet(ctx, te, 100, topodatapb.TabletType_REPLICA, "cell", true, true) defer deleteTablet(t, te, want1) // Local cell preference is default - tp, err := NewTabletPicker(context.Background(), te.topoServ, []string{"cella"}, "cell", te.keyspace, te.shard, "replica", TabletPickerOptions{}) + tp, err := NewTabletPicker(ctx, te.topoServ, []string{"cella"}, "cell", te.keyspace, te.shard, "replica", TabletPickerOptions{}) require.NoError(t, err) - ctx1, cancel1 := context.WithTimeout(context.Background(), 200*time.Millisecond) - defer cancel1() - tablet, err := tp.PickForStreaming(ctx1) + tablet, err := tp.PickForStreaming(ctx) require.NoError(t, err) assert.True(t, proto.Equal(want1, tablet), "Pick: %v, want %v", tablet, want1) // create a tablet in the other cell - want2 := addTablet(te, 101, topodatapb.TabletType_REPLICA, "otherCell", true, true) + want2 := addTablet(ctx, te, 101, topodatapb.TabletType_REPLICA, "otherCell", true, true) defer deleteTablet(t, te, want2) - ctx2, cancel2 := context.WithTimeout(context.Background(), 200*time.Millisecond) + ctx2, cancel2 := context.WithTimeout(ctx, 200*time.Millisecond) defer cancel2() // In 20 attempts, only tablet in "cell" will be picked because we give local cell priority by default @@ -336,49 +346,53 @@ func TestPickCellPreferenceLocalCell(t *testing.T) { } func TestPickCellPreferenceLocalAlias(t *testing.T) { + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + // test env puts all cells into an alias called "cella" - te := newPickerTestEnv(t, []string{"cell", "otherCell"}) - tp, err := NewTabletPicker(context.Background(), te.topoServ, []string{"cella"}, "cell", te.keyspace, te.shard, "replica", TabletPickerOptions{}) + te := newPickerTestEnv(t, ctx, []string{"cell", "otherCell"}) + tp, err := NewTabletPicker(ctx, te.topoServ, []string{"cella"}, "cell", te.keyspace, te.shard, "replica", TabletPickerOptions{}) require.NoError(t, err) // create a tablet in the other cell, it should be picked - want := addTablet(te, 101, topodatapb.TabletType_REPLICA, "otherCell", true, true) + want := addTablet(ctx, te, 101, topodatapb.TabletType_REPLICA, "otherCell", true, true) defer deleteTablet(t, te, want) - ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) - defer cancel() tablet, err := tp.PickForStreaming(ctx) require.NoError(t, err) assert.True(t, proto.Equal(want, tablet), "Pick: %v, want %v", tablet, want) } func TestPickUsingCellAliasOnlySpecified(t *testing.T) { + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + // test env puts all cells into an alias called "cella" - te := newPickerTestEnv(t, []string{"cell", "otherCell"}) - want1 := addTablet(te, 100, topodatapb.TabletType_REPLICA, "cell", true, true) + te := newPickerTestEnv(t, ctx, []string{"cell", "otherCell"}) + want1 := addTablet(ctx, te, 100, topodatapb.TabletType_REPLICA, "cell", true, true) defer deleteTablet(t, te, want1) - tp, err := NewTabletPicker(context.Background(), te.topoServ, []string{"cella"}, "cell", te.keyspace, te.shard, "replica", TabletPickerOptions{CellPreference: "OnlySpecified"}) + tp, err := NewTabletPicker(ctx, te.topoServ, []string{"cella"}, "cell", te.keyspace, te.shard, "replica", TabletPickerOptions{CellPreference: "OnlySpecified"}) require.NoError(t, err) - ctx1, cancel1 := context.WithTimeout(context.Background(), 200*time.Millisecond) - defer cancel1() - tablet, err := tp.PickForStreaming(ctx1) + tablet, err := tp.PickForStreaming(ctx) require.NoError(t, err) assert.True(t, proto.Equal(want1, tablet), "Pick: %v, want %v", tablet, want1) // create a tablet in the other cell, it should be picked deleteTablet(t, te, want1) - want2 := addTablet(te, 101, topodatapb.TabletType_REPLICA, "otherCell", true, true) + want2 := addTablet(ctx, te, 101, topodatapb.TabletType_REPLICA, "otherCell", true, true) defer deleteTablet(t, te, want2) - ctx2, cancel2 := context.WithTimeout(context.Background(), 200*time.Millisecond) + ctx2, cancel2 := context.WithTimeout(ctx, 200*time.Millisecond) defer cancel2() tablet, err = tp.PickForStreaming(ctx2) require.NoError(t, err) assert.True(t, proto.Equal(want2, tablet), "Pick: %v, want %v", tablet, want2) // addTablet again and test that both are picked at least once - want1 = addTablet(te, 100, topodatapb.TabletType_REPLICA, "cell", true, true) - ctx3, cancel3 := context.WithTimeout(context.Background(), 200*time.Millisecond) + want1 = addTablet(ctx, te, 100, topodatapb.TabletType_REPLICA, "cell", true, true) + ctx3, cancel3 := context.WithTimeout(ctx, 200*time.Millisecond) defer cancel3() // In 20 attempts each of the tablets should get picked at least once. @@ -399,8 +413,12 @@ func TestPickUsingCellAliasOnlySpecified(t *testing.T) { } func TestTabletAppearsDuringSleep(t *testing.T) { - te := newPickerTestEnv(t, []string{"cell"}) - tp, err := NewTabletPicker(context.Background(), te.topoServ, te.cells, "cell", te.keyspace, te.shard, "replica", TabletPickerOptions{}) + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + + te := newPickerTestEnv(t, ctx, []string{"cell"}) + tp, err := NewTabletPicker(ctx, te.topoServ, te.cells, "cell", te.keyspace, te.shard, "replica", TabletPickerOptions{}) require.NoError(t, err) delay := GetTabletPickerRetryDelay() @@ -412,14 +430,14 @@ func TestTabletAppearsDuringSleep(t *testing.T) { result := make(chan *topodatapb.Tablet) // start picker first, then add tablet go func() { - ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + ctx, cancel := context.WithTimeout(ctx, 20*time.Millisecond) defer cancel() tablet, err := tp.PickForStreaming(ctx) assert.NoError(t, err) result <- tablet }() - want := addTablet(te, 100, topodatapb.TabletType_REPLICA, "cell", true, true) + want := addTablet(ctx, te, 100, topodatapb.TabletType_REPLICA, "cell", true, true) defer deleteTablet(t, te, want) got := <-result require.NotNil(t, got, "Tablet should not be nil") @@ -427,11 +445,15 @@ func TestTabletAppearsDuringSleep(t *testing.T) { } func TestPickErrorLocalPreferenceDefault(t *testing.T) { - te := newPickerTestEnv(t, []string{"cell"}) - _, err := NewTabletPicker(context.Background(), te.topoServ, te.cells, "cell", te.keyspace, te.shard, "badtype", TabletPickerOptions{}) + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + te := newPickerTestEnv(t, ctx, []string{"cell"}) + _, err := NewTabletPicker(ctx, te.topoServ, te.cells, "cell", te.keyspace, te.shard, "badtype", TabletPickerOptions{}) assert.EqualError(t, err, "failed to parse list of tablet types: badtype") - tp, err := NewTabletPicker(context.Background(), te.topoServ, te.cells, "cell", te.keyspace, te.shard, "replica", TabletPickerOptions{}) + tp, err := NewTabletPicker(ctx, te.topoServ, te.cells, "cell", te.keyspace, te.shard, "replica", TabletPickerOptions{}) require.NoError(t, err) delay := GetTabletPickerRetryDelay() defer func() { @@ -439,25 +461,29 @@ func TestPickErrorLocalPreferenceDefault(t *testing.T) { }() SetTabletPickerRetryDelay(11 * time.Millisecond) - ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) - defer cancel() + timeoutCtx, timeoutCancel := context.WithTimeout(ctx, 20*time.Millisecond) + defer timeoutCancel() // no tablets - _, err = tp.PickForStreaming(ctx) + _, err = tp.PickForStreaming(timeoutCtx) require.EqualError(t, err, "context has expired") // no tablets of the correct type - defer deleteTablet(t, te, addTablet(te, 200, topodatapb.TabletType_RDONLY, "cell", true, true)) - ctx, cancel = context.WithTimeout(context.Background(), 20*time.Millisecond) - defer cancel() - _, err = tp.PickForStreaming(ctx) + defer deleteTablet(t, te, addTablet(ctx, te, 200, topodatapb.TabletType_RDONLY, "cell", true, true)) + timeoutCtx, timeoutCancel = context.WithTimeout(ctx, 20*time.Millisecond) + defer timeoutCancel() + _, err = tp.PickForStreaming(timeoutCtx) require.EqualError(t, err, "context has expired") // if local preference is selected, tp cells include's the local cell's alias require.Greater(t, globalTPStats.noTabletFoundError.Counts()["cell_cella.ks.0.replica"], int64(0)) } func TestPickErrorOnlySpecified(t *testing.T) { - te := newPickerTestEnv(t, []string{"cell"}) + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + te := newPickerTestEnv(t, ctx, []string{"cell"}) - tp, err := NewTabletPicker(context.Background(), te.topoServ, te.cells, "cell", te.keyspace, te.shard, "replica", TabletPickerOptions{CellPreference: "OnlySpecified"}) + tp, err := NewTabletPicker(ctx, te.topoServ, te.cells, "cell", te.keyspace, te.shard, "replica", TabletPickerOptions{CellPreference: "OnlySpecified"}) require.NoError(t, err) delay := GetTabletPickerRetryDelay() defer func() { @@ -465,16 +491,16 @@ func TestPickErrorOnlySpecified(t *testing.T) { }() SetTabletPickerRetryDelay(11 * time.Millisecond) - ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) - defer cancel() + timeoutCtx, timeoutCancel := context.WithTimeout(ctx, 20*time.Millisecond) + defer timeoutCancel() // no tablets - _, err = tp.PickForStreaming(ctx) + _, err = tp.PickForStreaming(timeoutCtx) require.EqualError(t, err, "context has expired") // no tablets of the correct type - defer deleteTablet(t, te, addTablet(te, 200, topodatapb.TabletType_RDONLY, "cell", true, true)) - ctx, cancel = context.WithTimeout(context.Background(), 20*time.Millisecond) - defer cancel() - _, err = tp.PickForStreaming(ctx) + defer deleteTablet(t, te, addTablet(ctx, te, 200, topodatapb.TabletType_RDONLY, "cell", true, true)) + timeoutCtx, timeoutCancel = context.WithTimeout(ctx, 20*time.Millisecond) + defer timeoutCancel() + _, err = tp.PickForStreaming(timeoutCtx) require.EqualError(t, err, "context has expired") require.Greater(t, globalTPStats.noTabletFoundError.Counts()["cell.ks.0.replica"], int64(0)) @@ -485,24 +511,28 @@ func TestPickErrorOnlySpecified(t *testing.T) { // type has no available healthy serving tablets that we select a healthy // serving tablet from the secondary/second type. func TestPickFallbackType(t *testing.T) { + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + cells := []string{"cell1", "cell2"} localCell := cells[0] tabletTypes := "replica,primary" options := TabletPickerOptions{ TabletOrder: "InOrder", } - te := newPickerTestEnv(t, cells) + te := newPickerTestEnv(t, ctx, cells) // This one should be selected even though it's the secondary type // as it is healthy and serving. - primaryTablet := addTablet(te, 100, topodatapb.TabletType_PRIMARY, localCell, true, true) + primaryTablet := addTablet(ctx, te, 100, topodatapb.TabletType_PRIMARY, localCell, true, true) defer deleteTablet(t, te, primaryTablet) // Replica tablet should not be selected as it is unhealthy. - replicaTablet := addTablet(te, 200, topodatapb.TabletType_REPLICA, localCell, false, false) + replicaTablet := addTablet(ctx, te, 200, topodatapb.TabletType_REPLICA, localCell, false, false) defer deleteTablet(t, te, replicaTablet) - ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + ctx, cancel = context.WithTimeout(ctx, 200*time.Millisecond) defer cancel() _, err := te.topoServ.UpdateShardFields(ctx, te.keyspace, te.shard, func(si *topo.ShardInfo) error { si.PrimaryAlias = primaryTablet.Alias @@ -510,9 +540,9 @@ func TestPickFallbackType(t *testing.T) { }) require.NoError(t, err) - tp, err := NewTabletPicker(context.Background(), te.topoServ, cells, localCell, te.keyspace, te.shard, tabletTypes, options) + tp, err := NewTabletPicker(ctx, te.topoServ, cells, localCell, te.keyspace, te.shard, tabletTypes, options) require.NoError(t, err) - ctx2, cancel2 := context.WithTimeout(context.Background(), 1*time.Second) + ctx2, cancel2 := context.WithTimeout(ctx, 1*time.Second) defer cancel2() tablet, err := tp.PickForStreaming(ctx2) require.NoError(t, err) @@ -528,15 +558,13 @@ type pickerTestEnv struct { topoServ *topo.Server } -func newPickerTestEnv(t *testing.T, cells []string) *pickerTestEnv { - ctx := context.Background() - +func newPickerTestEnv(t *testing.T, ctx context.Context, cells []string) *pickerTestEnv { te := &pickerTestEnv{ t: t, keyspace: "ks", shard: "0", cells: cells, - topoServ: memorytopo.NewServer(cells...), + topoServ: memorytopo.NewServer(ctx, cells...), } // create cell alias err := te.topoServ.CreateCellsAlias(ctx, "cella", &topodatapb.CellsAlias{ @@ -550,7 +578,7 @@ func newPickerTestEnv(t *testing.T, cells []string) *pickerTestEnv { return te } -func addTablet(te *pickerTestEnv, id int, tabletType topodatapb.TabletType, cell string, serving, healthy bool) *topodatapb.Tablet { +func addTablet(ctx context.Context, te *pickerTestEnv, id int, tabletType topodatapb.TabletType, cell string, serving, healthy bool) *topodatapb.Tablet { tablet := &topodatapb.Tablet{ Alias: &topodatapb.TabletAlias{ Cell: cell, @@ -564,7 +592,7 @@ func addTablet(te *pickerTestEnv, id int, tabletType topodatapb.TabletType, cell "test": int32(id), }, } - err := te.topoServ.CreateTablet(context.Background(), tablet) + err := te.topoServ.CreateTablet(ctx, tablet) require.NoError(te.t, err) shr := &querypb.StreamHealthResponse{ diff --git a/go/vt/discovery/topology_watcher.go b/go/vt/discovery/topology_watcher.go index 99d61e8c545..d1bd2d3acf8 100644 --- a/go/vt/discovery/topology_watcher.go +++ b/go/vt/discovery/topology_watcher.go @@ -18,6 +18,7 @@ package discovery import ( "bytes" + "context" "fmt" "hash/crc32" "sort" @@ -29,8 +30,6 @@ import ( "vitess.io/vitess/go/vt/key" - "context" - "vitess.io/vitess/go/stats" "vitess.io/vitess/go/trace" diff --git a/go/vt/discovery/topology_watcher_test.go b/go/vt/discovery/topology_watcher_test.go index dff8ba720c7..95c4e9e2e3c 100644 --- a/go/vt/discovery/topology_watcher_test.go +++ b/go/vt/discovery/topology_watcher_test.go @@ -26,6 +26,8 @@ import ( "github.com/stretchr/testify/require" "google.golang.org/protobuf/proto" + "vitess.io/vitess/go/test/utils" + "vitess.io/vitess/go/vt/logutil" topodatapb "vitess.io/vitess/go/vt/proto/topodata" "vitess.io/vitess/go/vt/topo" @@ -56,8 +58,14 @@ func checkChecksum(t *testing.T, tw *TopologyWatcher, want uint32) { } func TestStartAndCloseTopoWatcher(t *testing.T) { - ts := memorytopo.NewServer("aa") + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "aa") + defer ts.Close() fhc := NewFakeHealthCheck(nil) + defer fhc.Close() topologyWatcherOperations.ZeroAll() tw := NewCellTabletsWatcher(context.Background(), ts, fhc, nil, "aa", 100*time.Microsecond, true, 5) @@ -110,8 +118,14 @@ func TestCellTabletsWatcherNoRefreshKnown(t *testing.T) { } func checkWatcher(t *testing.T, refreshKnownTablets bool) { - ts := memorytopo.NewServer("aa") + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "aa") + defer ts.Close() fhc := NewFakeHealthCheck(nil) + defer fhc.Close() logger := logutil.NewMemoryLogger() topologyWatcherOperations.ZeroAll() counts := topologyWatcherOperations.Counts() @@ -429,9 +443,14 @@ var ( ) func TestFilterByKeyspace(t *testing.T) { + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + hc := NewFakeHealthCheck(nil) f := NewFilterByKeyspace(testKeyspacesToWatch) - ts := memorytopo.NewServer(testCell) + ts := memorytopo.NewServer(ctx, testCell) + defer ts.Close() tw := NewCellTabletsWatcher(context.Background(), ts, hc, f, testCell, 10*time.Minute, true, 5) for _, test := range testFilterByKeyspace { @@ -509,8 +528,14 @@ func TestFilterByKeyspace(t *testing.T) { // - does not continuosly call GetTablets for tablets that do not satisfy the filter // - does not add or remove these filtered out tablets from the its healtcheck func TestFilterByKeypsaceSkipsIgnoredTablets(t *testing.T) { - ts := memorytopo.NewServer("aa") + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "aa") + defer ts.Close() fhc := NewFakeHealthCheck(nil) + defer fhc.Close() topologyWatcherOperations.ZeroAll() counts := topologyWatcherOperations.Counts() f := NewFilterByKeyspace(testKeyspacesToWatch) diff --git a/go/vt/mysqlctl/backup_blackbox_test.go b/go/vt/mysqlctl/backup_blackbox_test.go index f7af0cdc306..5764eefc82b 100644 --- a/go/vt/mysqlctl/backup_blackbox_test.go +++ b/go/vt/mysqlctl/backup_blackbox_test.go @@ -29,6 +29,8 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "vitess.io/vitess/go/test/utils" + "vitess.io/vitess/go/mysql/replication" "vitess.io/vitess/go/sqltypes" @@ -78,6 +80,9 @@ func createBackupFiles(root string, fileCount int, ext string) error { } func TestExecuteBackup(t *testing.T) { + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() // Set up local backup directory backupRoot := "testdata/builtinbackup_test" filebackupstorage.FileBackupStorageRoot = backupRoot @@ -90,8 +95,6 @@ func TestExecuteBackup(t *testing.T) { require.NoError(t, createBackupFiles(path.Join(dataDir, "test2"), 2, "ibd")) defer os.RemoveAll(backupRoot) - ctx := context.Background() - needIt, err := needInnoDBRedoLogSubdir() require.NoError(t, err) if needIt { @@ -103,7 +106,7 @@ func TestExecuteBackup(t *testing.T) { // Set up topo keyspace, shard := "mykeyspace", "-80" - ts := memorytopo.NewServer("cell1") + ts := memorytopo.NewServer(ctx, "cell1") defer ts.Close() require.NoError(t, ts.CreateKeyspace(ctx, keyspace, &topodata.Keyspace{})) @@ -220,6 +223,10 @@ func TestExecuteBackup(t *testing.T) { } func TestExecuteBackupWithSafeUpgrade(t *testing.T) { + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Set up local backup directory backupRoot := "testdata/builtinbackup_test" filebackupstorage.FileBackupStorageRoot = backupRoot @@ -232,8 +239,6 @@ func TestExecuteBackupWithSafeUpgrade(t *testing.T) { require.NoError(t, createBackupFiles(path.Join(dataDir, "test2"), 2, "ibd")) defer os.RemoveAll(backupRoot) - ctx := context.Background() - needIt, err := needInnoDBRedoLogSubdir() require.NoError(t, err) if needIt { @@ -245,7 +250,7 @@ func TestExecuteBackupWithSafeUpgrade(t *testing.T) { // Set up topo keyspace, shard := "mykeyspace", "-80" - ts := memorytopo.NewServer("cell1") + ts := memorytopo.NewServer(ctx, "cell1") defer ts.Close() require.NoError(t, ts.CreateKeyspace(ctx, keyspace, &topodata.Keyspace{})) @@ -307,6 +312,10 @@ func TestExecuteBackupWithSafeUpgrade(t *testing.T) { // TestExecuteBackupWithCanceledContext tests the ability of the backup function to gracefully handle cases where errors // occur due to various reasons, such as context time cancel. The process should not panic in these situations. func TestExecuteBackupWithCanceledContext(t *testing.T) { + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Set up local backup directory id := fmt.Sprintf("%d", time.Now().UnixNano()) backupRoot := fmt.Sprintf("testdata/builtinbackup_test_%s", id) @@ -322,8 +331,8 @@ func TestExecuteBackupWithCanceledContext(t *testing.T) { defer os.RemoveAll(backupRoot) // Cancel the context deliberately - ctx, cancel := context.WithCancel(context.Background()) - cancel() + cancelledCtx, cancelCtx := context.WithCancel(context.Background()) + cancelCtx() needIt, err := needInnoDBRedoLogSubdir() require.NoError(t, err) if needIt { @@ -335,19 +344,19 @@ func TestExecuteBackupWithCanceledContext(t *testing.T) { // Set up topo keyspace, shard := "mykeyspace", "-80" - ts := memorytopo.NewServer("cell1") + ts := memorytopo.NewServer(ctx, "cell1") defer ts.Close() - require.NoError(t, ts.CreateKeyspace(ctx, keyspace, &topodata.Keyspace{})) - require.NoError(t, ts.CreateShard(ctx, keyspace, shard)) + require.NoError(t, ts.CreateKeyspace(cancelledCtx, keyspace, &topodata.Keyspace{})) + require.NoError(t, ts.CreateShard(cancelledCtx, keyspace, shard)) tablet := topo.NewTablet(100, "cell1", "mykeyspace-00-80-0100") tablet.Keyspace = keyspace tablet.Shard = shard - require.NoError(t, ts.CreateTablet(ctx, tablet)) + require.NoError(t, ts.CreateTablet(cancelledCtx, tablet)) - _, err = ts.UpdateShardFields(ctx, keyspace, shard, func(si *topo.ShardInfo) error { + _, err = ts.UpdateShardFields(cancelledCtx, keyspace, shard, func(si *topo.ShardInfo) error { si.PrimaryAlias = &topodata.TabletAlias{Uid: 100, Cell: "cell1"} now := time.Now() @@ -364,7 +373,7 @@ func TestExecuteBackupWithCanceledContext(t *testing.T) { mysqld := mysqlctl.NewFakeMysqlDaemon(fakesqldb.New(t)) mysqld.ExpectedExecuteSuperQueryList = []string{"STOP SLAVE", "START SLAVE"} - ok, err := be.ExecuteBackup(ctx, mysqlctl.BackupParams{ + ok, err := be.ExecuteBackup(cancelledCtx, mysqlctl.BackupParams{ Logger: logutil.NewConsoleLogger(), Mysqld: mysqld, Cnf: &mysqlctl.Mycnf{ @@ -389,6 +398,10 @@ func TestExecuteBackupWithCanceledContext(t *testing.T) { // TestExecuteRestoreWithCanceledContext tests the ability of the restore function to gracefully handle cases where errors // occur due to various reasons, such as context timed-out. The process should not panic in these situations. func TestExecuteRestoreWithTimedOutContext(t *testing.T) { + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Set up local backup directory id := fmt.Sprintf("%d", time.Now().UnixNano()) backupRoot := fmt.Sprintf("testdata/builtinbackup_test_%s", id) @@ -403,7 +416,6 @@ func TestExecuteRestoreWithTimedOutContext(t *testing.T) { require.NoError(t, createBackupFiles(path.Join(dataDir, "test2"), 2, "ibd")) defer os.RemoveAll(backupRoot) - ctx := context.Background() needIt, err := needInnoDBRedoLogSubdir() require.NoError(t, err) if needIt { @@ -415,7 +427,7 @@ func TestExecuteRestoreWithTimedOutContext(t *testing.T) { // Set up topo keyspace, shard := "mykeyspace", "-80" - ts := memorytopo.NewServer("cell1") + ts := memorytopo.NewServer(ctx, "cell1") defer ts.Close() require.NoError(t, ts.CreateKeyspace(ctx, keyspace, &topodata.Keyspace{})) diff --git a/go/vt/schemamanager/schemamanager_test.go b/go/vt/schemamanager/schemamanager_test.go index 6edb42bb5c4..154d985bba4 100644 --- a/go/vt/schemamanager/schemamanager_test.go +++ b/go/vt/schemamanager/schemamanager_test.go @@ -331,8 +331,9 @@ func (client *fakeTabletManagerClient) ExecuteFetchAsDba(ctx context.Context, ta // - 3 shards named '1', '2', '3'. // - A primary tablet for each shard. func newFakeTopo(t *testing.T) *topo.Server { - ts := memorytopo.NewServer("test_cell") - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "test_cell") if err := ts.CreateKeyspace(ctx, "test_keyspace", &topodatapb.Keyspace{}); err != nil { t.Fatalf("CreateKeyspace failed: %v", err) } diff --git a/go/vt/schemamanager/tablet_executor_test.go b/go/vt/schemamanager/tablet_executor_test.go index 498c1622da9..15022ecc527 100644 --- a/go/vt/schemamanager/tablet_executor_test.go +++ b/go/vt/schemamanager/tablet_executor_test.go @@ -55,8 +55,9 @@ func TestTabletExecutorOpen(t *testing.T) { } func TestTabletExecutorOpenWithEmptyPrimaryAlias(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("test_cell") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "test_cell") tablet := &topodatapb.Tablet{ Alias: &topodatapb.TabletAlias{ Cell: "test_cell", diff --git a/go/vt/srvtopo/discover_test.go b/go/vt/srvtopo/discover_test.go index c076ba0e7b7..ca4774a1b84 100644 --- a/go/vt/srvtopo/discover_test.go +++ b/go/vt/srvtopo/discover_test.go @@ -48,8 +48,9 @@ func (a TargetArray) Less(i, j int) bool { } func TestFindAllTargets(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("cell1", "cell2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2") srvTopoCacheRefresh = 0 srvTopoCacheTTL = 0 @@ -58,7 +59,7 @@ func TestFindAllTargets(t *testing.T) { srvTopoCacheTTL = 1 * time.Second }() - rs := NewResilientServer(ts, "TestFindAllKeyspaceShards") + rs := NewResilientServer(ctx, ts, "TestFindAllKeyspaceShards") // No keyspace / shards. ks, err := FindAllTargets(ctx, rs, "cell1", []topodatapb.TabletType{topodatapb.TabletType_PRIMARY}) diff --git a/go/vt/srvtopo/keyspace_filtering_server_test.go b/go/vt/srvtopo/keyspace_filtering_server_test.go index 83e1a18e062..bcd5681f3e8 100644 --- a/go/vt/srvtopo/keyspace_filtering_server_test.go +++ b/go/vt/srvtopo/keyspace_filtering_server_test.go @@ -48,10 +48,10 @@ var ( } ) -func newFiltering(filter []string) (*topo.Server, *srvtopotest.PassthroughSrvTopoServer, Server) { +func newFiltering(ctx context.Context, filter []string) (*topo.Server, *srvtopotest.PassthroughSrvTopoServer, Server) { testServer := srvtopotest.NewPassthroughSrvTopoServer() - testServer.TopoServer = memorytopo.NewServer(stockCell) + testServer.TopoServer = memorytopo.NewServer(ctx, stockCell) testServer.SrvKeyspaceNames = []string{"foo", "bar", "baz"} testServer.SrvKeyspace = &topodatapb.SrvKeyspace{} testServer.WatchedSrvVSchema = stockVSchema @@ -71,7 +71,9 @@ func TestFilteringServerHandlesNilUnderlying(t *testing.T) { } func TestFilteringServerReturnsUnderlyingServer(t *testing.T) { - _, _, f := newFiltering(nil) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, _, f := newFiltering(ctx, nil) got, gotErr := f.GetTopoServer() if gotErr != nil { t.Errorf("Got error getting topo.Server from FilteringServer") @@ -108,17 +110,23 @@ func doTestGetSrvKeyspaceNames( } func TestFilteringServerGetSrvKeyspameNamesFiltersEverythingOut(t *testing.T) { - _, _, f := newFiltering(nil) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, _, f := newFiltering(ctx, nil) doTestGetSrvKeyspaceNames(t, f, stockCell, []string{}, nil) } func TestFilteringServerGetSrvKeyspaceNamesFiltersKeyspaces(t *testing.T) { - _, _, f := newFiltering(stockFilters) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, _, f := newFiltering(ctx, stockFilters) doTestGetSrvKeyspaceNames(t, f, stockCell, stockFilters, nil) } func TestFilteringServerGetSrvKeyspaceNamesPassesThroughErrors(t *testing.T) { - _, mock, f := newFiltering(stockFilters) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, mock, f := newFiltering(ctx, stockFilters) wantErr := fmt.Errorf("some badcell error") mock.SrvKeyspaceNamesError = wantErr doTestGetSrvKeyspaceNames(t, f, "badcell", stockFilters, wantErr) @@ -140,28 +148,36 @@ func doTestGetSrvKeyspace( } func TestFilteringServerGetSrvKeyspaceReturnsSelectedKeyspaces(t *testing.T) { - _, mock, f := newFiltering(stockFilters) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, mock, f := newFiltering(ctx, stockFilters) mock.SrvKeyspace = stockKeyspaces["bar"] doTestGetSrvKeyspace(t, f, stockCell, "bar", stockKeyspaces["bar"], nil) } func TestFilteringServerGetSrvKeyspaceErrorPassthrough(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() wantErr := fmt.Errorf("some error") - _, mock, f := newFiltering(stockFilters) + _, mock, f := newFiltering(ctx, stockFilters) mock.SrvKeyspace = stockKeyspaces["bar"] mock.SrvKeyspaceError = wantErr doTestGetSrvKeyspace(t, f, "badcell", "bar", stockKeyspaces["bar"], wantErr) } func TestFilteringServerGetSrvKeyspaceFilters(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() wantErr := topo.NewError(topo.NoNode, "foo") - _, mock, f := newFiltering(stockFilters) + _, mock, f := newFiltering(ctx, stockFilters) mock.SrvKeyspaceError = wantErr doTestGetSrvKeyspace(t, f, stockCell, "foo", nil, wantErr) } func TestFilteringServerWatchSrvVSchemaFiltersPassthroughSrvVSchema(t *testing.T) { - _, mock, f := newFiltering(stockFilters) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, mock, f := newFiltering(ctx, stockFilters) allowed := map[string]bool{} for _, ks := range stockFilters { @@ -196,8 +212,11 @@ func TestFilteringServerWatchSrvVSchemaFiltersPassthroughSrvVSchema(t *testing.T } func TestFilteringServerWatchSrvVSchemaHandlesNilSchema(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + wantErr := fmt.Errorf("some err") - _, mock, f := newFiltering(stockFilters) + _, mock, f := newFiltering(ctx, stockFilters) mock.WatchedSrvVSchema = nil mock.WatchedSrvVSchemaError = wantErr diff --git a/go/vt/srvtopo/resilient_server.go b/go/vt/srvtopo/resilient_server.go index cac368dedd9..d1521952ab0 100644 --- a/go/vt/srvtopo/resilient_server.go +++ b/go/vt/srvtopo/resilient_server.go @@ -17,6 +17,7 @@ limitations under the License. package srvtopo import ( + "context" "time" "github.com/spf13/pflag" @@ -78,7 +79,7 @@ type ResilientServer struct { // NewResilientServer creates a new ResilientServer // based on the provided topo.Server. -func NewResilientServer(base *topo.Server, counterPrefix string) *ResilientServer { +func NewResilientServer(ctx context.Context, base *topo.Server, counterPrefix string) *ResilientServer { if srvTopoCacheRefresh > srvTopoCacheTTL { log.Fatalf("srv_topo_cache_refresh must be less than or equal to srv_topo_cache_ttl") } @@ -94,8 +95,8 @@ func NewResilientServer(base *topo.Server, counterPrefix string) *ResilientServe return &ResilientServer{ topoServer: base, counts: counts, - SrvKeyspaceWatcher: NewSrvKeyspaceWatcher(base, counts, srvTopoCacheRefresh, srvTopoCacheTTL), - SrvVSchemaWatcher: NewSrvVSchemaWatcher(base, counts, srvTopoCacheRefresh, srvTopoCacheTTL), + SrvKeyspaceWatcher: NewSrvKeyspaceWatcher(ctx, base, counts, srvTopoCacheRefresh, srvTopoCacheTTL), + SrvVSchemaWatcher: NewSrvVSchemaWatcher(ctx, base, counts, srvTopoCacheRefresh, srvTopoCacheTTL), SrvKeyspaceNamesQuery: NewSrvKeyspaceNamesQuery(base, counts, srvTopoCacheRefresh, srvTopoCacheTTL), } } diff --git a/go/vt/srvtopo/resilient_server_test.go b/go/vt/srvtopo/resilient_server_test.go index 95640028050..1dfe0340fdf 100644 --- a/go/vt/srvtopo/resilient_server_test.go +++ b/go/vt/srvtopo/resilient_server_test.go @@ -43,7 +43,10 @@ import ( // TestGetSrvKeyspace will test we properly return updated SrvKeyspace. func TestGetSrvKeyspace(t *testing.T) { - ts, factory := memorytopo.NewServerAndFactory("test_cell") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts, factory := memorytopo.NewServerAndFactory(ctx, "test_cell") + defer ts.Close() srvTopoCacheTTL = 200 * time.Millisecond srvTopoCacheRefresh = 80 * time.Millisecond defer func() { @@ -51,7 +54,7 @@ func TestGetSrvKeyspace(t *testing.T) { srvTopoCacheRefresh = 1 * time.Second }() - rs := NewResilientServer(ts, "TestGetSrvKeyspace") + rs := NewResilientServer(ctx, ts, "TestGetSrvKeyspace") // Ask for a not-yet-created keyspace _, err := rs.GetSrvKeyspace(context.Background(), "test_cell", "test_ks") @@ -359,17 +362,19 @@ func TestGetSrvKeyspace(t *testing.T) { // TestSrvKeyspaceCachedError will test we properly re-try to query // the topo server upon failure. func TestSrvKeyspaceCachedError(t *testing.T) { - ts := memorytopo.NewServer("test_cell") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "test_cell") + defer ts.Close() srvTopoCacheTTL = 100 * time.Millisecond srvTopoCacheRefresh = 40 * time.Millisecond defer func() { srvTopoCacheTTL = 1 * time.Second srvTopoCacheRefresh = 1 * time.Second }() - rs := NewResilientServer(ts, "TestSrvKeyspaceCachedErrors") + rs := NewResilientServer(ctx, ts, "TestSrvKeyspaceCachedErrors") // Ask for an unknown keyspace, should get an error. - ctx := context.Background() _, err := rs.GetSrvKeyspace(ctx, "test_cell", "unknown_ks") if err == nil { t.Fatalf("First GetSrvKeyspace didn't return an error") @@ -382,8 +387,6 @@ func TestSrvKeyspaceCachedError(t *testing.T) { time.Sleep(srvTopoCacheTTL + 10*time.Millisecond) // Ask again with a different context, should get an error and // save that context. - ctx, cancel := context.WithCancel(ctx) - defer cancel() _, err2 := rs.GetSrvKeyspace(ctx, "test_cell", "unknown_ks") if err2 == nil { t.Fatalf("Second GetSrvKeyspace didn't return an error") @@ -396,8 +399,11 @@ func TestSrvKeyspaceCachedError(t *testing.T) { // TestGetSrvKeyspaceCreated will test we properly get the initial // value if the SrvKeyspace already exists. func TestGetSrvKeyspaceCreated(t *testing.T) { - ts := memorytopo.NewServer("test_cell") - rs := NewResilientServer(ts, "TestGetSrvKeyspaceCreated") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "test_cell") + defer ts.Close() + rs := NewResilientServer(ctx, ts, "TestGetSrvKeyspaceCreated") // Set SrvKeyspace with value. want := &topodatapb.SrvKeyspace{} @@ -428,9 +434,11 @@ func TestGetSrvKeyspaceCreated(t *testing.T) { func TestWatchSrvVSchema(t *testing.T) { srvTopoCacheRefresh = 10 * time.Millisecond - ctx := context.Background() - ts := memorytopo.NewServer("test_cell") - rs := NewResilientServer(ts, "TestWatchSrvVSchema") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "test_cell") + defer ts.Close() + rs := NewResilientServer(ctx, ts, "TestWatchSrvVSchema") // mu protects watchValue and watchErr. mu := sync.Mutex{} @@ -512,14 +520,20 @@ func TestWatchSrvVSchema(t *testing.T) { } func TestGetSrvKeyspaceNames(t *testing.T) { - ts, factory := memorytopo.NewServerAndFactory("test_cell") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts, factory := memorytopo.NewServerAndFactory(ctx, "test_cell") + defer ts.Close() + + time.Sleep(1 * time.Second) + srvTopoCacheTTL = 100 * time.Millisecond srvTopoCacheRefresh = 40 * time.Millisecond defer func() { srvTopoCacheTTL = 1 * time.Second srvTopoCacheRefresh = 1 * time.Second }() - rs := NewResilientServer(ts, "TestGetSrvKeyspaceNames") + rs := NewResilientServer(ctx, ts, "TestGetSrvKeyspaceNames") // Set SrvKeyspace with value want := &topodatapb.SrvKeyspace{} @@ -529,7 +543,6 @@ func TestGetSrvKeyspaceNames(t *testing.T) { err = ts.UpdateSrvKeyspace(context.Background(), "test_cell", "test_ks2", want) require.NoError(t, err, "UpdateSrvKeyspace(test_cell, test_ks2, %s) failed", want) - ctx := context.Background() names, err := rs.GetSrvKeyspaceNames(ctx, "test_cell", false) if err != nil { t.Errorf("GetSrvKeyspaceNames unexpected error %v", err) @@ -644,8 +657,8 @@ func TestGetSrvKeyspaceNames(t *testing.T) { time.Sleep(srvTopoCacheTTL) - timeoutCtx, cancel := context.WithTimeout(context.Background(), srvTopoCacheRefresh*2) //nolint - defer cancel() + timeoutCtx, timeoutCancel := context.WithTimeout(context.Background(), srvTopoCacheRefresh*2) //nolint + defer timeoutCancel() _, err = rs.GetSrvKeyspaceNames(timeoutCtx, "test_cell", false) if err != context.DeadlineExceeded { t.Errorf("expected error '%v', got '%v'", context.DeadlineExceeded, err.Error()) @@ -666,7 +679,10 @@ func (w *watched) equals(other *watched) bool { } func TestSrvKeyspaceWatcher(t *testing.T) { - ts, factory := memorytopo.NewServerAndFactory("test_cell") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts, factory := memorytopo.NewServerAndFactory(ctx, "test_cell") + defer ts.Close() srvTopoCacheTTL = 100 * time.Millisecond srvTopoCacheRefresh = 40 * time.Millisecond defer func() { @@ -674,7 +690,7 @@ func TestSrvKeyspaceWatcher(t *testing.T) { srvTopoCacheRefresh = 1 * time.Second }() - rs := NewResilientServer(ts, "TestGetSrvKeyspaceWatcher") + rs := NewResilientServer(ctx, ts, "TestGetSrvKeyspaceWatcher") var wmu sync.Mutex var wseen []watched @@ -790,7 +806,10 @@ func TestSrvKeyspaceWatcher(t *testing.T) { } func TestSrvKeyspaceListener(t *testing.T) { - ts, _ := memorytopo.NewServerAndFactory("test_cell") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "test_cell") + defer ts.Close() srvTopoCacheTTL = 100 * time.Millisecond srvTopoCacheRefresh = 40 * time.Millisecond defer func() { @@ -798,16 +817,16 @@ func TestSrvKeyspaceListener(t *testing.T) { srvTopoCacheRefresh = 1 * time.Second }() - rs := NewResilientServer(ts, "TestGetSrvKeyspaceWatcher") + rs := NewResilientServer(ctx, ts, "TestGetSrvKeyspaceWatcher") - ctx, cancel := context.WithCancel(context.Background()) + cancelCtx, cancelFunc := context.WithCancel(context.Background()) var callbackCount atomic.Int32 // adding listener will perform callback. - rs.WatchSrvKeyspace(context.Background(), "test_cell", "test_ks", func(srvKs *topodatapb.SrvKeyspace, err error) bool { + rs.WatchSrvKeyspace(ctx, "test_cell", "test_ks", func(srvKs *topodatapb.SrvKeyspace, err error) bool { callbackCount.Add(1) select { - case <-ctx.Done(): + case <-cancelCtx.Done(): return false default: return true @@ -816,16 +835,16 @@ func TestSrvKeyspaceListener(t *testing.T) { // First update (callback - 2) want := &topodatapb.SrvKeyspace{} - err := ts.UpdateSrvKeyspace(context.Background(), "test_cell", "test_ks", want) + err := ts.UpdateSrvKeyspace(ctx, "test_cell", "test_ks", want) require.NoError(t, err) // Next callback to remove from listener - cancel() + cancelFunc() // multi updates thereafter for i := 0; i < 5; i++ { want = &topodatapb.SrvKeyspace{} - err = ts.UpdateSrvKeyspace(context.Background(), "test_cell", "test_ks", want) + err = ts.UpdateSrvKeyspace(ctx, "test_cell", "test_ks", want) require.NoError(t, err) time.Sleep(100 * time.Millisecond) } diff --git a/go/vt/srvtopo/resolver_test.go b/go/vt/srvtopo/resolver_test.go index 49b108fcffb..95e6dbe620c 100644 --- a/go/vt/srvtopo/resolver_test.go +++ b/go/vt/srvtopo/resolver_test.go @@ -17,12 +17,11 @@ limitations under the License. package srvtopo import ( + "context" "testing" "github.com/stretchr/testify/require" - "context" - "vitess.io/vitess/go/sqltypes" "vitess.io/vitess/go/vt/key" "vitess.io/vitess/go/vt/logutil" @@ -34,11 +33,10 @@ import ( topodatapb "vitess.io/vitess/go/vt/proto/topodata" ) -func initResolver(t *testing.T, name string) *Resolver { - ctx := context.Background() +func initResolver(t *testing.T, ctx context.Context, name string) *Resolver { cell := "cell1" - ts := memorytopo.NewServer(cell) - rs := NewResilientServer(ts, name) + ts := memorytopo.NewServer(ctx, cell) + rs := NewResilientServer(ctx, ts, name) // Create sharded keyspace and shards. if err := ts.CreateKeyspace(ctx, "sks", &topodatapb.Keyspace{}); err != nil { @@ -97,7 +95,9 @@ func initResolver(t *testing.T, name string) *Resolver { } func TestResolveDestinations(t *testing.T) { - resolver := initResolver(t, "TestResolveDestinations") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + resolver := initResolver(t, ctx, "TestResolveDestinations") id1 := &querypb.Value{ Type: sqltypes.VarChar, diff --git a/go/vt/srvtopo/watch.go b/go/vt/srvtopo/watch.go index 2d470327c4e..36d8fd428bd 100644 --- a/go/vt/srvtopo/watch.go +++ b/go/vt/srvtopo/watch.go @@ -33,6 +33,7 @@ const ( watchStateIdle watchState = iota watchStateStarting watchStateRunning + watchStateStopped ) type watchEntry struct { @@ -100,7 +101,10 @@ func (entry *watchEntry) addListener(ctx context.Context, callback func(any, err callback(v, err) } -func (entry *watchEntry) ensureWatchingLocked() { +func (entry *watchEntry) ensureWatchingLocked(ctx context.Context) { + if ctx.Err() != nil { + return + } switch entry.watchState { case watchStateRunning, watchStateStarting: case watchStateIdle: @@ -121,7 +125,7 @@ func (entry *watchEntry) currentValueLocked(ctx context.Context) (any, error) { return entry.value, entry.lastError } - entry.ensureWatchingLocked() + entry.ensureWatchingLocked(ctx) cacheValid := entry.value != nil && time.Since(entry.lastValueTime) < entry.rw.cacheTTL if cacheValid { @@ -146,12 +150,12 @@ func (entry *watchEntry) currentValueLocked(ctx context.Context) (any, error) { return nil, entry.lastError } -func (entry *watchEntry) update(value any, err error, init bool) { +func (entry *watchEntry) update(ctx context.Context, value any, err error, init bool) { entry.mutex.Lock() defer entry.mutex.Unlock() if err != nil { - entry.onErrorLocked(err, init) + entry.onErrorLocked(ctx, err, init) } else { entry.onValueLocked(value) } @@ -179,7 +183,7 @@ func (entry *watchEntry) onValueLocked(value any) { entry.lastErrorTime = time.Time{} } -func (entry *watchEntry) onErrorLocked(err error, init bool) { +func (entry *watchEntry) onErrorLocked(ctx context.Context, err error, init bool) { entry.rw.counts.Add(errorCategory, 1) entry.lastErrorTime = time.Now() @@ -217,12 +221,13 @@ func (entry *watchEntry) onErrorLocked(err error, init bool) { entry.watchState = watchStateIdle // only retry the watch if we haven't been explicitly interrupted + if len(entry.listeners) > 0 && !topo.IsErrType(err, topo.Interrupted) { go func() { time.Sleep(entry.rw.cacheRefreshInterval) entry.mutex.Lock() - entry.ensureWatchingLocked() + entry.ensureWatchingLocked(ctx) entry.mutex.Unlock() }() } diff --git a/go/vt/srvtopo/watch_srvkeyspace.go b/go/vt/srvtopo/watch_srvkeyspace.go index e47e810b615..cefe95c6951 100644 --- a/go/vt/srvtopo/watch_srvkeyspace.go +++ b/go/vt/srvtopo/watch_srvkeyspace.go @@ -37,25 +37,25 @@ func (k *srvKeyspaceKey) String() string { return k.cell + "." + k.keyspace } -func NewSrvKeyspaceWatcher(topoServer *topo.Server, counts *stats.CountersWithSingleLabel, cacheRefresh, cacheTTL time.Duration) *SrvKeyspaceWatcher { +func NewSrvKeyspaceWatcher(ctx context.Context, topoServer *topo.Server, counts *stats.CountersWithSingleLabel, cacheRefresh, cacheTTL time.Duration) *SrvKeyspaceWatcher { watch := func(entry *watchEntry) { key := entry.key.(*srvKeyspaceKey) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + requestCtx, requestCancel := context.WithCancel(context.Background()) + defer requestCancel() - current, changes, err := topoServer.WatchSrvKeyspace(ctx, key.cell, key.keyspace) + current, changes, err := topoServer.WatchSrvKeyspace(requestCtx, key.cell, key.keyspace) if err != nil { - entry.update(nil, err, true) + entry.update(ctx, nil, err, true) return } - entry.update(current.Value, current.Err, true) + entry.update(ctx, current.Value, current.Err, true) if current.Err != nil { return } for c := range changes { - entry.update(c.Value, c.Err, false) + entry.update(ctx, c.Value, c.Err, false) if c.Err != nil { return } diff --git a/go/vt/srvtopo/watch_srvvschema.go b/go/vt/srvtopo/watch_srvvschema.go index 251f5e55644..1b5536e623d 100644 --- a/go/vt/srvtopo/watch_srvvschema.go +++ b/go/vt/srvtopo/watch_srvvschema.go @@ -35,26 +35,25 @@ func (k cellName) String() string { return string(k) } -func NewSrvVSchemaWatcher(topoServer *topo.Server, counts *stats.CountersWithSingleLabel, cacheRefresh, cacheTTL time.Duration) *SrvVSchemaWatcher { +func NewSrvVSchemaWatcher(ctx context.Context, topoServer *topo.Server, counts *stats.CountersWithSingleLabel, cacheRefresh, cacheTTL time.Duration) *SrvVSchemaWatcher { watch := func(entry *watchEntry) { key := entry.key.(cellName) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + requestCtx, requestCancel := context.WithCancel(ctx) + defer requestCancel() - current, changes, err := topoServer.WatchSrvVSchema(ctx, key.String()) + current, changes, err := topoServer.WatchSrvVSchema(requestCtx, key.String()) if err != nil { - entry.update(nil, err, true) + entry.update(ctx, nil, err, true) return } - entry.update(current.Value, current.Err, true) + entry.update(ctx, current.Value, current.Err, true) if current.Err != nil { return } - defer cancel() for c := range changes { - entry.update(c.Value, c.Err, false) + entry.update(ctx, c.Value, c.Err, false) if c.Err != nil { return } diff --git a/go/vt/throttler/demo/throttler_demo.go b/go/vt/throttler/demo/throttler_demo.go index 615f7c4fc93..126b9098236 100644 --- a/go/vt/throttler/demo/throttler_demo.go +++ b/go/vt/throttler/demo/throttler_demo.go @@ -231,13 +231,13 @@ type client struct { healthcheckCh chan *discovery.TabletHealth } -func newClient(primary *primary, replica *replica, ts *topo.Server) *client { +func newClient(ctx context.Context, primary *primary, replica *replica, ts *topo.Server) *client { t, err := throttler.NewThrottler("client", "TPS", 1, throttler.MaxRateModuleDisabled, 5 /* seconds */) if err != nil { log.Fatal(err) } - healthCheck := discovery.NewHealthCheck(context.Background(), 5*time.Second, 1*time.Minute, ts, "cell1", "") + healthCheck := discovery.NewHealthCheck(ctx, 5*time.Second, 1*time.Minute, ts, "cell1", "") c := &client{ primary: primary, healthCheck: healthCheck, @@ -307,10 +307,10 @@ func main() { }) log.Infof("start rate set to: %v", rate) - ts := memorytopo.NewServer("cell1") + ts := memorytopo.NewServer(context.Background(), "cell1") replica := newReplica(lagUpdateInterval, replicaDegrationInterval, replicaDegrationDuration, ts) primary := &primary{replica: replica} - client := newClient(primary, replica, ts) + client := newClient(context.Background(), primary, replica, ts) client.run() time.Sleep(duration) diff --git a/go/vt/topo/cell_info.go b/go/vt/topo/cell_info.go index fd7a4a5249e..4a8112084cb 100644 --- a/go/vt/topo/cell_info.go +++ b/go/vt/topo/cell_info.go @@ -60,6 +60,9 @@ func (ts *Server) GetCellInfoNames(ctx context.Context) ([]string, error) { // GetCellInfo reads a CellInfo from the global Conn. func (ts *Server) GetCellInfo(ctx context.Context, cell string, strongRead bool) (*topodatapb.CellInfo, error) { conn := ts.globalCell + if ctx.Err() != nil { + return nil, ctx.Err() + } if !strongRead { conn = ts.globalReadOnlyCell } diff --git a/go/vt/topo/consultopo/server_flaky_test.go b/go/vt/topo/consultopo/server_flaky_test.go index 797ad4c955f..a987336dd01 100644 --- a/go/vt/topo/consultopo/server_flaky_test.go +++ b/go/vt/topo/consultopo/server_flaky_test.go @@ -144,7 +144,9 @@ func TestConsulTopo(t *testing.T) { // Run the TopoServerTestSuite tests. testIndex := 0 - test.TopoServerTestSuite(t, func() *topo.Server { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + test.TopoServerTestSuite(t, ctx, func() *topo.Server { // Each test will use its own sub-directories. testRoot := fmt.Sprintf("test-%v", testIndex) testIndex++ @@ -190,7 +192,9 @@ func TestConsulTopoWithChecks(t *testing.T) { // Run the TopoServerTestSuite tests. testIndex := 0 - test.TopoServerTestSuite(t, func() *topo.Server { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + test.TopoServerTestSuite(t, ctx, func() *topo.Server { // Each test will use its own sub-directories. testRoot := fmt.Sprintf("test-%v", testIndex) testIndex++ @@ -247,7 +251,9 @@ func TestConsulTopoWithAuth(t *testing.T) { t.Fatalf("couldn't write temp file: %v", err) } - test.TopoServerTestSuite(t, func() *topo.Server { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + test.TopoServerTestSuite(t, ctx, func() *topo.Server { // Each test will use its own sub-directories. testRoot := fmt.Sprintf("test-%v", testIndex) testIndex++ diff --git a/go/vt/topo/etcd2topo/server_test.go b/go/vt/topo/etcd2topo/server_test.go index 3bf0e29cfd7..732829ee78b 100644 --- a/go/vt/topo/etcd2topo/server_test.go +++ b/go/vt/topo/etcd2topo/server_test.go @@ -245,7 +245,9 @@ func TestEtcd2Topo(t *testing.T) { } // Run the TopoServerTestSuite tests. - test.TopoServerTestSuite(t, func() *topo.Server { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + test.TopoServerTestSuite(t, ctx, func() *topo.Server { return newServer() }, []string{}) diff --git a/go/vt/topo/faketopo/faketopo.go b/go/vt/topo/faketopo/faketopo.go index 9265ba699a3..8601d28f5b6 100644 --- a/go/vt/topo/faketopo/faketopo.go +++ b/go/vt/topo/faketopo/faketopo.go @@ -340,13 +340,13 @@ func (f *FakeConn) Close() { } // NewFakeTopoServer creates a new fake topo server -func NewFakeTopoServer(factory *FakeFactory) *topo.Server { +func NewFakeTopoServer(ctx context.Context, factory *FakeFactory) *topo.Server { ts, err := topo.NewWithFactory(factory, "" /*serverAddress*/, "" /*root*/) if err != nil { log.Exitf("topo.NewWithFactory() failed: %v", err) } for cell := range factory.cells { - if err := ts.CreateCellInfo(context.Background(), cell, &topodatapb.CellInfo{}); err != nil { + if err := ts.CreateCellInfo(ctx, cell, &topodatapb.CellInfo{}); err != nil { log.Exitf("ts.CreateCellInfo(%v) failed: %v", cell, err) } } diff --git a/go/vt/topo/helpers/copy_test.go b/go/vt/topo/helpers/copy_test.go index 73ecfabf66b..2086a2e6552 100644 --- a/go/vt/topo/helpers/copy_test.go +++ b/go/vt/topo/helpers/copy_test.go @@ -17,12 +17,11 @@ limitations under the License. package helpers import ( + "context" "testing" "github.com/stretchr/testify/require" - "context" - "vitess.io/vitess/go/vt/topo" "vitess.io/vitess/go/vt/topo/memorytopo" @@ -34,8 +33,8 @@ func createSetup(ctx context.Context, t *testing.T) (*topo.Server, *topo.Server) // Create a source and destination TS. They will have // different generations, so we test using the Version for // both works as expected. - fromTS := memorytopo.NewServer("test_cell") - toTS := memorytopo.NewServer("test_cell") + fromTS := memorytopo.NewServer(ctx, "test_cell") + toTS := memorytopo.NewServer(ctx, "test_cell") // create a keyspace and a couple tablets if err := fromTS.CreateKeyspace(ctx, "test_keyspace", &topodatapb.Keyspace{}); err != nil { diff --git a/go/vt/topo/helpers/tee_topo_test.go b/go/vt/topo/helpers/tee_topo_test.go index 519301eaafa..41e220e0a56 100644 --- a/go/vt/topo/helpers/tee_topo_test.go +++ b/go/vt/topo/helpers/tee_topo_test.go @@ -17,17 +17,22 @@ limitations under the License. package helpers import ( + "context" "testing" + "vitess.io/vitess/go/test/utils" "vitess.io/vitess/go/vt/topo" "vitess.io/vitess/go/vt/topo/memorytopo" "vitess.io/vitess/go/vt/topo/test" ) func TestTeeTopo(t *testing.T) { - test.TopoServerTestSuite(t, func() *topo.Server { - s1 := memorytopo.NewServer(test.LocalCellName) - s2 := memorytopo.NewServer(test.LocalCellName) + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + test.TopoServerTestSuite(t, ctx, func() *topo.Server { + s1 := memorytopo.NewServer(ctx, test.LocalCellName) + s2 := memorytopo.NewServer(ctx, test.LocalCellName) tee, err := NewTee(s1, s2, false) if err != nil { t.Fatalf("NewTee() failed: %v", err) diff --git a/go/vt/topo/memorytopo/memorytopo.go b/go/vt/topo/memorytopo/memorytopo.go index ea7c8112ceb..8619ad023fe 100644 --- a/go/vt/topo/memorytopo/memorytopo.go +++ b/go/vt/topo/memorytopo/memorytopo.go @@ -128,18 +128,16 @@ type Conn struct { } // dial returns immediately, unless the Conn points to the sentinel -// UnreachableServerAddr, in which case it will block until the context expires -// and return the context's error. +// UnreachableServerAddr, in which case it will block until the context expires. func (c *Conn) dial(ctx context.Context) error { if c.closed { return ErrConnectionClosed } if c.serverAddr == UnreachableServerAddr { <-ctx.Done() - return ctx.Err() } - return nil + return ctx.Err() } // Close is part of the topo.Conn interface. diff --git a/go/vt/topo/memorytopo/server_test.go b/go/vt/topo/memorytopo/server_test.go index 5bfa41c8a5e..c2d1cf6cfb5 100644 --- a/go/vt/topo/memorytopo/server_test.go +++ b/go/vt/topo/memorytopo/server_test.go @@ -17,6 +17,7 @@ limitations under the License. package memorytopo import ( + "context" "testing" "vitess.io/vitess/go/vt/topo" @@ -25,7 +26,9 @@ import ( func TestMemoryTopo(t *testing.T) { // Run the TopoServerTestSuite tests. - test.TopoServerTestSuite(t, func() *topo.Server { - return NewServer(test.LocalCellName) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + test.TopoServerTestSuite(t, ctx, func() *topo.Server { + return NewServer(ctx, test.LocalCellName) }, []string{"checkTryLock", "checkShardWithLock"}) } diff --git a/go/vt/topo/test/directory.go b/go/vt/topo/test/directory.go index b33404f8643..88001a56d87 100644 --- a/go/vt/topo/test/directory.go +++ b/go/vt/topo/test/directory.go @@ -17,18 +17,15 @@ limitations under the License. package test import ( + "context" "reflect" "testing" - "context" - "vitess.io/vitess/go/vt/topo" ) // checkDirectory tests the directory part of the topo.Conn API. -func checkDirectory(t *testing.T, ts *topo.Server) { - ctx := context.Background() - +func checkDirectory(t *testing.T, ctx context.Context, ts *topo.Server) { // global cell t.Logf("=== checkDirectoryInCell global") conn, err := ts.ConnForCell(ctx, topo.GlobalCell) diff --git a/go/vt/topo/test/election.go b/go/vt/topo/test/election.go index 594e6562eb2..377ea21ae09 100644 --- a/go/vt/topo/test/election.go +++ b/go/vt/topo/test/election.go @@ -46,8 +46,8 @@ func waitForLeaderID(t *testing.T, mp topo.LeaderParticipation, expected string) // checkElection runs the tests on the LeaderParticipation part of the // topo.Conn API. -func checkElection(t *testing.T, ts *topo.Server) { - conn, err := ts.ConnForCell(context.Background(), topo.GlobalCell) +func checkElection(t *testing.T, ctx context.Context, ts *topo.Server) { + conn, err := ts.ConnForCell(ctx, topo.GlobalCell) if err != nil { t.Fatalf("ConnForCell(global) failed: %v", err) } @@ -71,7 +71,7 @@ func checkElection(t *testing.T, ts *topo.Server) { // A lot of implementations use a toplevel directory for their elections. // Make sure it is marked as 'Ephemeral'. - entries, err := conn.ListDir(context.Background(), "/", true /*full*/) + entries, err := conn.ListDir(ctx, "/", true /*full*/) if err != nil { t.Fatalf("ListDir(/) failed: %v", err) } @@ -148,8 +148,8 @@ func checkElection(t *testing.T, ts *topo.Server) { } // checkWaitForNewLeader runs the WaitForLeadership test on the LeaderParticipation -func checkWaitForNewLeader(t *testing.T, ts *topo.Server) { - conn, err := ts.ConnForCell(context.Background(), topo.GlobalCell) +func checkWaitForNewLeader(t *testing.T, ctx context.Context, ts *topo.Server) { + conn, err := ts.ConnForCell(ctx, topo.GlobalCell) if err != nil { t.Fatalf("ConnForCell(global) failed: %v", err) } @@ -195,7 +195,7 @@ func checkWaitForNewLeader(t *testing.T, ts *topo.Server) { t.Fatalf("cannot create mp2: %v", err) } - leaders, err := mp2.WaitForNewLeader(context.Background()) + leaders, err := mp2.WaitForNewLeader(ctx) if topo.IsErrType(err, topo.NoImplementation) { t.Logf("%T does not support WaitForNewLeader()", mp2) return diff --git a/go/vt/topo/test/file.go b/go/vt/topo/test/file.go index 8e5858d17a3..70bb386fd80 100644 --- a/go/vt/topo/test/file.go +++ b/go/vt/topo/test/file.go @@ -26,9 +26,7 @@ import ( ) // checkFile tests the file part of the Conn API. -func checkFile(t *testing.T, ts *topo.Server) { - ctx := context.Background() - +func checkFile(t *testing.T, ctx context.Context, ts *topo.Server) { // global cell t.Logf("=== checkFileInCell global") conn, err := ts.ConnForCell(ctx, topo.GlobalCell) @@ -203,8 +201,7 @@ func checkFileInCell(t *testing.T, conn topo.Conn, hasCells bool) { } // checkList tests the file part of the Conn API. -func checkList(t *testing.T, ts *topo.Server) { - ctx := context.Background() +func checkList(t *testing.T, ctx context.Context, ts *topo.Server) { // global cell conn, err := ts.ConnForCell(ctx, LocalCellName) if err != nil { diff --git a/go/vt/topo/test/keyspace.go b/go/vt/topo/test/keyspace.go index c5b9af68009..0458e7fd2d7 100644 --- a/go/vt/topo/test/keyspace.go +++ b/go/vt/topo/test/keyspace.go @@ -17,9 +17,8 @@ limitations under the License. package test import ( - "testing" - "context" + "testing" "vitess.io/vitess/go/vt/topo" @@ -27,8 +26,7 @@ import ( ) // checkKeyspace tests the keyspace part of the API -func checkKeyspace(t *testing.T, ts *topo.Server) { - ctx := context.Background() +func checkKeyspace(t *testing.T, ctx context.Context, ts *topo.Server) { keyspaces, err := ts.GetKeyspaces(ctx) if err != nil { t.Errorf("GetKeyspaces(empty): %v", err) diff --git a/go/vt/topo/test/lock.go b/go/vt/topo/test/lock.go index 69cdeff2a55..dce51ed859d 100644 --- a/go/vt/topo/test/lock.go +++ b/go/vt/topo/test/lock.go @@ -17,12 +17,11 @@ limitations under the License. package test import ( + "context" "path" "testing" "time" - "context" - "vitess.io/vitess/go/vt/topo" topodatapb "vitess.io/vitess/go/vt/proto/topodata" @@ -35,8 +34,7 @@ var timeUntilLockIsTaken = 10 * time.Millisecond // checkLock checks we can lock / unlock as expected. It's using a keyspace // as the lock target. -func checkLock(t *testing.T, ts *topo.Server) { - ctx := context.Background() +func checkLock(t *testing.T, ctx context.Context, ts *topo.Server) { if err := ts.CreateKeyspace(ctx, "test_keyspace", &topodatapb.Keyspace{}); err != nil { t.Fatalf("CreateKeyspace: %v", err) } diff --git a/go/vt/topo/test/replication.go b/go/vt/topo/test/replication.go index e681749f68d..3080cb77145 100644 --- a/go/vt/topo/test/replication.go +++ b/go/vt/topo/test/replication.go @@ -17,20 +17,18 @@ limitations under the License. package test import ( + "context" "testing" "google.golang.org/protobuf/proto" - "context" - "vitess.io/vitess/go/vt/topo" topodatapb "vitess.io/vitess/go/vt/proto/topodata" ) // checkShardReplication tests ShardReplication objects -func checkShardReplication(t *testing.T, ts *topo.Server) { - ctx := context.Background() +func checkShardReplication(t *testing.T, ctx context.Context, ts *topo.Server) { if _, err := ts.GetShardReplication(ctx, LocalCellName, "test_keyspace", "-10"); !topo.IsErrType(err, topo.NoNode) { t.Errorf("GetShardReplication(not there): %v", err) } diff --git a/go/vt/topo/test/serving.go b/go/vt/topo/test/serving.go index dfeac442180..dd00f3da370 100644 --- a/go/vt/topo/test/serving.go +++ b/go/vt/topo/test/serving.go @@ -17,12 +17,11 @@ limitations under the License. package test import ( + "context" "testing" "google.golang.org/protobuf/proto" - "context" - "vitess.io/vitess/go/vt/topo" topodatapb "vitess.io/vitess/go/vt/proto/topodata" @@ -30,9 +29,7 @@ import ( ) // checkSrvKeyspace tests the SrvKeyspace methods (other than watch). -func checkSrvKeyspace(t *testing.T, ts *topo.Server) { - ctx := context.Background() - +func checkSrvKeyspace(t *testing.T, ctx context.Context, ts *topo.Server) { // Test GetSrvKeyspaceNames returns an empty list correctly. if names, err := ts.GetSrvKeyspaceNames(ctx, LocalCellName); err != nil || len(names) != 0 { t.Errorf("GetSrvKeyspace(not there): %v %v", names, err) @@ -91,9 +88,7 @@ func checkSrvKeyspace(t *testing.T, ts *topo.Server) { } // checkSrvVSchema tests the SrvVSchema methods (other than watch). -func checkSrvVSchema(t *testing.T, ts *topo.Server) { - ctx := context.Background() - +func checkSrvVSchema(t *testing.T, ctx context.Context, ts *topo.Server) { // check GetSrvVSchema returns topo.ErrNoNode if no SrvVSchema if _, err := ts.GetSrvVSchema(ctx, LocalCellName); !topo.IsErrType(err, topo.NoNode) { t.Errorf("GetSrvVSchema(not set): %v", err) diff --git a/go/vt/topo/test/shard.go b/go/vt/topo/test/shard.go index d285f382838..b5c92c4a3ec 100644 --- a/go/vt/topo/test/shard.go +++ b/go/vt/topo/test/shard.go @@ -31,8 +31,7 @@ import ( ) // checkShard verifies the Shard operations work correctly -func checkShard(t *testing.T, ts *topo.Server) { - ctx := context.Background() +func checkShard(t *testing.T, ctx context.Context, ts *topo.Server) { if err := ts.CreateKeyspace(ctx, "test_keyspace", &topodatapb.Keyspace{}); err != nil { t.Fatalf("CreateKeyspace: %v", err) } @@ -100,8 +99,7 @@ func checkShard(t *testing.T, ts *topo.Server) { // checkShardWithLock verifies that `TryLockShard` will keep failing with `NodeExists` error if there is // a lock already taken for given shard. Once we unlock that shard, then subsequent call to `TryLockShard` // should succeed. -func checkShardWithLock(t *testing.T, ts *topo.Server) { - ctx := context.Background() +func checkShardWithLock(t *testing.T, ctx context.Context, ts *topo.Server) { if err := ts.CreateKeyspace(ctx, "test_keyspace", &topodatapb.Keyspace{}); err != nil { t.Fatalf("CreateKeyspace: %v", err) } diff --git a/go/vt/topo/test/tablet.go b/go/vt/topo/test/tablet.go index 4a562b1e46a..63afc1abff0 100644 --- a/go/vt/topo/test/tablet.go +++ b/go/vt/topo/test/tablet.go @@ -17,21 +17,18 @@ limitations under the License. package test import ( + "context" "testing" "google.golang.org/protobuf/proto" - "context" - "vitess.io/vitess/go/vt/topo" topodatapb "vitess.io/vitess/go/vt/proto/topodata" ) // checkTablet verifies the topo server API is correct for managing tablets. -func checkTablet(t *testing.T, ts *topo.Server) { - ctx := context.Background() - +func checkTablet(t *testing.T, ctx context.Context, ts *topo.Server) { tablet := &topodatapb.Tablet{ Alias: &topodatapb.TabletAlias{ Cell: LocalCellName, diff --git a/go/vt/topo/test/testing.go b/go/vt/topo/test/testing.go index e8d014242ad..d189a7f4cf1 100644 --- a/go/vt/topo/test/testing.go +++ b/go/vt/topo/test/testing.go @@ -22,6 +22,7 @@ limitations under the License. package test import ( + "context" "testing" "vitess.io/vitess/go/vt/topo" @@ -40,7 +41,7 @@ func newKeyRange(value string) *topodatapb.KeyRange { return result } -func executeTestSuite(f func(*testing.T, *topo.Server), t *testing.T, ts *topo.Server, ignoreList []string, name string) { +func executeTestSuite(f func(*testing.T, context.Context, *topo.Server), t *testing.T, ctx context.Context, ts *topo.Server, ignoreList []string, name string) { // some test does not apply every where therefore we ignore them for _, n := range ignoreList { if n == name { @@ -48,7 +49,7 @@ func executeTestSuite(f func(*testing.T, *topo.Server), t *testing.T, ts *topo.S return } } - f(t, ts) + f(t, ctx, ts) } // TopoServerTestSuite runs the full topo.Server/Conn test suite. @@ -57,101 +58,101 @@ func executeTestSuite(f func(*testing.T, *topo.Server), t *testing.T, ts *topo.S // Not all tests are applicable for each Topo server, therefore we provide ignoreList in order to // avoid them for given Topo server tests. For example `TryLock` implementation is same as `Lock` for some Topo servers. // Hence, for these Topo servers we ignore executing TryLock Tests. -func TopoServerTestSuite(t *testing.T, factory func() *topo.Server, ignoreList []string) { +func TopoServerTestSuite(t *testing.T, ctx context.Context, factory func() *topo.Server, ignoreList []string) { var ts *topo.Server t.Log("=== checkKeyspace") ts = factory() - executeTestSuite(checkKeyspace, t, ts, ignoreList, "checkKeyspace") + executeTestSuite(checkKeyspace, t, ctx, ts, ignoreList, "checkKeyspace") ts.Close() t.Log("=== checkShard") ts = factory() - executeTestSuite(checkShard, t, ts, ignoreList, "checkShard") + executeTestSuite(checkShard, t, ctx, ts, ignoreList, "checkShard") ts.Close() t.Log("=== checkShardWithLock") ts = factory() - executeTestSuite(checkShardWithLock, t, ts, ignoreList, "checkShardWithLock") + executeTestSuite(checkShardWithLock, t, ctx, ts, ignoreList, "checkShardWithLock") ts.Close() t.Log("=== checkTablet") ts = factory() - executeTestSuite(checkTablet, t, ts, ignoreList, "checkTablet") + executeTestSuite(checkTablet, t, ctx, ts, ignoreList, "checkTablet") ts.Close() t.Log("=== checkShardReplication") ts = factory() - executeTestSuite(checkShardReplication, t, ts, ignoreList, "checkShardReplication") + executeTestSuite(checkShardReplication, t, ctx, ts, ignoreList, "checkShardReplication") ts.Close() t.Log("=== checkSrvKeyspace") ts = factory() - executeTestSuite(checkSrvKeyspace, t, ts, ignoreList, "checkSrvKeyspace") + executeTestSuite(checkSrvKeyspace, t, ctx, ts, ignoreList, "checkSrvKeyspace") ts.Close() t.Log("=== checkSrvVSchema") ts = factory() - executeTestSuite(checkSrvVSchema, t, ts, ignoreList, "checkSrvVSchema") + executeTestSuite(checkSrvVSchema, t, ctx, ts, ignoreList, "checkSrvVSchema") ts.Close() t.Log("=== checkLock") ts = factory() - executeTestSuite(checkLock, t, ts, ignoreList, "checkLock") + executeTestSuite(checkLock, t, ctx, ts, ignoreList, "checkLock") ts.Close() t.Log("=== checkTryLock") ts = factory() - executeTestSuite(checkTryLock, t, ts, ignoreList, "checkTryLock") + executeTestSuite(checkTryLock, t, ctx, ts, ignoreList, "checkTryLock") ts.Close() t.Log("=== checkVSchema") ts = factory() - executeTestSuite(checkVSchema, t, ts, ignoreList, "checkVSchema") + executeTestSuite(checkVSchema, t, ctx, ts, ignoreList, "checkVSchema") ts.Close() t.Log("=== checkRoutingRules") ts = factory() - executeTestSuite(checkRoutingRules, t, ts, ignoreList, "checkRoutingRules") + executeTestSuite(checkRoutingRules, t, ctx, ts, ignoreList, "checkRoutingRules") ts.Close() t.Log("=== checkElection") ts = factory() - executeTestSuite(checkElection, t, ts, ignoreList, "checkElection") + executeTestSuite(checkElection, t, ctx, ts, ignoreList, "checkElection") ts.Close() t.Log("=== checkWaitForNewLeader") ts = factory() - executeTestSuite(checkWaitForNewLeader, t, ts, ignoreList, "checkWaitForNewLeader") + executeTestSuite(checkWaitForNewLeader, t, ctx, ts, ignoreList, "checkWaitForNewLeader") ts.Close() t.Log("=== checkDirectory") ts = factory() - executeTestSuite(checkDirectory, t, ts, ignoreList, "checkDirectory") + executeTestSuite(checkDirectory, t, ctx, ts, ignoreList, "checkDirectory") ts.Close() t.Log("=== checkFile") ts = factory() - executeTestSuite(checkFile, t, ts, ignoreList, "checkFile") + executeTestSuite(checkFile, t, ctx, ts, ignoreList, "checkFile") ts.Close() t.Log("=== checkWatch") ts = factory() - executeTestSuite(checkWatch, t, ts, ignoreList, "checkWatch") + executeTestSuite(checkWatch, t, ctx, ts, ignoreList, "checkWatch") ts.Close() ts = factory() t.Log("=== checkWatchInterrupt") - executeTestSuite(checkWatchInterrupt, t, ts, ignoreList, "checkWatchInterrupt") + executeTestSuite(checkWatchInterrupt, t, ctx, ts, ignoreList, "checkWatchInterrupt") ts.Close() ts = factory() t.Log("=== checkList") - executeTestSuite(checkList, t, ts, ignoreList, "checkList") + executeTestSuite(checkList, t, ctx, ts, ignoreList, "checkList") ts.Close() ts = factory() t.Log("=== checkWatchRecursive") - executeTestSuite(checkWatchRecursive, t, ts, ignoreList, "checkWatchRecursive") + executeTestSuite(checkWatchRecursive, t, ctx, ts, ignoreList, "checkWatchRecursive") ts.Close() } diff --git a/go/vt/topo/test/trylock.go b/go/vt/topo/test/trylock.go index cace3cccc61..4519d1bcaab 100644 --- a/go/vt/topo/test/trylock.go +++ b/go/vt/topo/test/trylock.go @@ -31,8 +31,7 @@ import ( // checkTryLock checks if we can lock / unlock as expected. It's using a keyspace // as the lock target. -func checkTryLock(t *testing.T, ts *topo.Server) { - ctx := context.Background() +func checkTryLock(t *testing.T, ctx context.Context, ts *topo.Server) { if err := ts.CreateKeyspace(ctx, "test_keyspace", &topodatapb.Keyspace{}); err != nil { require.Fail(t, "CreateKeyspace fail", err.Error()) } diff --git a/go/vt/topo/test/vschema.go b/go/vt/topo/test/vschema.go index 0c2d58bdba7..5063addaefd 100644 --- a/go/vt/topo/test/vschema.go +++ b/go/vt/topo/test/vschema.go @@ -17,9 +17,8 @@ limitations under the License. package test import ( - "testing" - "context" + "testing" "github.com/stretchr/testify/require" "google.golang.org/protobuf/proto" @@ -31,8 +30,7 @@ import ( ) // checkVSchema runs the tests on the VSchema part of the API -func checkVSchema(t *testing.T, ts *topo.Server) { - ctx := context.Background() +func checkVSchema(t *testing.T, ctx context.Context, ts *topo.Server) { if err := ts.CreateKeyspace(ctx, "test_keyspace", &topodatapb.Keyspace{}); err != nil { t.Fatalf("CreateKeyspace: %v", err) } @@ -92,9 +90,7 @@ func checkVSchema(t *testing.T, ts *topo.Server) { } // checkRoutingRules runs the tests on the routing rules part of the API -func checkRoutingRules(t *testing.T, ts *topo.Server) { - ctx := context.Background() - +func checkRoutingRules(t *testing.T, ctx context.Context, ts *topo.Server) { if _, err := ts.GetRoutingRules(ctx); err != nil { t.Fatal(err) } diff --git a/go/vt/topo/test/watch.go b/go/vt/topo/test/watch.go index 08dec8cd56f..a4caaaf742d 100644 --- a/go/vt/topo/test/watch.go +++ b/go/vt/topo/test/watch.go @@ -114,8 +114,8 @@ func waitForInitialValueRecursive(t *testing.T, conn topo.Conn, srvKeyspace *top // checkWatch runs the tests on the Watch part of the Conn API. // We use a SrvKeyspace object. -func checkWatch(t *testing.T, ts *topo.Server) { - ctx, cancel := context.WithCancel(context.Background()) +func checkWatch(t *testing.T, ctx context.Context, ts *topo.Server) { + ctx, cancel := context.WithCancel(ctx) defer cancel() conn, err := ts.ConnForCell(ctx, LocalCellName) if err != nil { @@ -227,8 +227,7 @@ func checkWatch(t *testing.T, ts *topo.Server) { } // checkWatchInterrupt tests we can interrupt a watch. -func checkWatchInterrupt(t *testing.T, ts *topo.Server) { - ctx := context.Background() +func checkWatchInterrupt(t *testing.T, ctx context.Context, ts *topo.Server) { conn, err := ts.ConnForCell(ctx, LocalCellName) if err != nil { t.Fatalf("ConnForCell(test) failed: %v", err) @@ -298,8 +297,8 @@ func checkWatchInterrupt(t *testing.T, ts *topo.Server) { } // checkWatchRecursive tests we can setup a recursive watch -func checkWatchRecursive(t *testing.T, ts *topo.Server) { - ctx, cancel := context.WithCancel(context.Background()) +func checkWatchRecursive(t *testing.T, ctx context.Context, ts *topo.Server) { + ctx, cancel := context.WithCancel(ctx) defer cancel() conn, err := ts.ConnForCell(ctx, LocalCellName) if err != nil { diff --git a/go/vt/topo/topotests/cell_info_test.go b/go/vt/topo/topotests/cell_info_test.go index 89ad68043a1..becdbd8d14a 100644 --- a/go/vt/topo/topotests/cell_info_test.go +++ b/go/vt/topo/topotests/cell_info_test.go @@ -36,8 +36,10 @@ import ( func TestCellInfo(t *testing.T) { cell := "cell1" - ctx := context.Background() - ts := memorytopo.NewServer(cell) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cell) + defer ts.Close() // Check GetCellInfo returns what memorytopo created. ci, err := ts.GetCellInfo(ctx, cell, true /*strongRead*/) @@ -135,7 +137,8 @@ func TestCellInfo(t *testing.T) { } func TestExpandCells(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() var cells []string var err error var allCells = "cell1,cell2,cell3" @@ -162,10 +165,12 @@ func TestExpandCells(t *testing.T) { topoCells := strings.Split(cellsIn, ",") var ts *topo.Server if tCase.name == "bad" { - ts = memorytopo.NewServer() + ts = memorytopo.NewServer(ctx) } else { - ts = memorytopo.NewServer(topoCells...) + ts = memorytopo.NewServer(ctx, topoCells...) } + defer ts.Close() + cells, err = ts.ExpandCells(ctx, cellsIn) if tCase.errString != "" { require.Error(t, err) @@ -179,7 +184,7 @@ func TestExpandCells(t *testing.T) { t.Run("aliases", func(t *testing.T) { cells := []string{"cell1", "cell2", "cell3"} - ts := memorytopo.NewServer(cells...) + ts := memorytopo.NewServer(ctx, cells...) err := ts.CreateCellsAlias(ctx, "alias", &topodatapb.CellsAlias{Cells: cells}) require.NoError(t, err) @@ -228,8 +233,10 @@ func TestExpandCells(t *testing.T) { } func TestDeleteCellInfo(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("zone1", "unreachable") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1", "unreachable") + defer ts.Close() err := ts.UpdateCellInfoFields(ctx, "unreachable", func(ci *topodatapb.CellInfo) error { ci.ServerAddress = memorytopo.UnreachableServerAddr @@ -254,11 +261,11 @@ func TestDeleteCellInfo(t *testing.T) { }, } for _, tt := range tests { - func() { - ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) - defer cancel() + t.Run(fmt.Sprintf("force:%t", tt.force), func(t *testing.T) { + requestCtx, requestCancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer requestCancel() - err := ts.DeleteCellInfo(ctx, "unreachable", tt.force) + err := ts.DeleteCellInfo(requestCtx, "unreachable", tt.force) if tt.shouldErr { assert.Error(t, err, "force=%t", tt.force) } else { @@ -272,6 +279,6 @@ func TestDeleteCellInfo(t *testing.T) { } else { assert.True(t, topo.IsErrType(err, topo.NoNode), "expected cell %q to not exist", "unreachable") } - }() + }) } } diff --git a/go/vt/topo/topotests/cells_aliases_test.go b/go/vt/topo/topotests/cells_aliases_test.go index d124dcd8d47..7b8f0ebe3f5 100644 --- a/go/vt/topo/topotests/cells_aliases_test.go +++ b/go/vt/topo/topotests/cells_aliases_test.go @@ -33,8 +33,10 @@ func TestCellsAliases(t *testing.T) { // Create an alias cell := "cell1" - ctx := context.Background() - ts := memorytopo.NewServer(cell) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cell) + defer ts.Close() if err := ts.CreateCellsAlias(ctx, "alias", &topodatapb.CellsAlias{Cells: []string{"cell1", "cell2"}}); err != nil { t.Fatalf("CreateCellsAlias failed: %v", err) diff --git a/go/vt/topo/topotests/keyspace_test.go b/go/vt/topo/topotests/keyspace_test.go index b0b35c1421c..96eb9938353 100644 --- a/go/vt/topo/topotests/keyspace_test.go +++ b/go/vt/topo/topotests/keyspace_test.go @@ -31,8 +31,10 @@ import ( ) func TestCreateKeyspace(t *testing.T) { - ts := memorytopo.NewServer("zone1") - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") + defer ts.Close() t.Run("valid name", func(t *testing.T) { err := ts.CreateKeyspace(ctx, "ks", &topodatapb.Keyspace{}) @@ -46,8 +48,10 @@ func TestCreateKeyspace(t *testing.T) { } func TestGetKeyspace(t *testing.T) { - ts := memorytopo.NewServer("zone1") - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") + defer ts.Close() t.Run("valid name", func(t *testing.T) { // First, create the keyspace. diff --git a/go/vt/topo/topotests/replication_test.go b/go/vt/topo/topotests/replication_test.go index 0d6e148fa3d..d45aaf36551 100644 --- a/go/vt/topo/topotests/replication_test.go +++ b/go/vt/topo/topotests/replication_test.go @@ -37,8 +37,10 @@ func TestFixShardReplication(t *testing.T) { cell := "cell1" keyspace := "ks1" shard := "shard1" - ctx := context.Background() - ts := memorytopo.NewServer(cell) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cell) + defer ts.Close() // Create a tablet. alias := &topodatapb.TabletAlias{ diff --git a/go/vt/topo/topotests/shard_watch_test.go b/go/vt/topo/topotests/shard_watch_test.go index a8333251a2d..80b696c106d 100644 --- a/go/vt/topo/topotests/shard_watch_test.go +++ b/go/vt/topo/topotests/shard_watch_test.go @@ -56,8 +56,10 @@ func waitForInitialShard(t *testing.T, ts *topo.Server, keyspace, shard string) func TestWatchShardNoNode(t *testing.T) { keyspace := "ks1" shard := "0" - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") + defer ts.Close() // No Shard -> ErrNoNode _, _, err := ts.WatchShard(ctx, keyspace, shard) @@ -70,8 +72,10 @@ func TestWatchShard(t *testing.T) { cell := "cell1" keyspace := "ks1" shard := "0" - ctx := context.Background() - ts := memorytopo.NewServer(cell) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cell) + defer ts.Close() // Create keyspace if err := ts.CreateKeyspace(ctx, keyspace, &topodatapb.Keyspace{}); err != nil { @@ -205,8 +209,10 @@ func TestWatchShardCancel(t *testing.T) { cell := "cell1" keyspace := "ks1" shard := "0" - ctx := context.Background() - ts := memorytopo.NewServer(cell) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cell) + defer ts.Close() // No Shard -> ErrNoNode _, _, err := ts.WatchShard(ctx, keyspace, shard) diff --git a/go/vt/topo/topotests/srv_keyspace_test.go b/go/vt/topo/topotests/srv_keyspace_test.go index 8eeaf3f07ac..97e44e3a82a 100644 --- a/go/vt/topo/topotests/srv_keyspace_test.go +++ b/go/vt/topo/topotests/srv_keyspace_test.go @@ -62,8 +62,10 @@ func waitForInitialSrvKeyspace(t *testing.T, ts *topo.Server, cell, keyspace str func TestWatchSrvKeyspaceNoNode(t *testing.T) { cell := "cell1" keyspace := "ks1" - ctx := context.Background() - ts := memorytopo.NewServer(cell) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cell) + defer ts.Close() // No SrvKeyspace -> ErrNoNode _, _, err := ts.WatchSrvKeyspace(ctx, cell, keyspace) @@ -76,8 +78,10 @@ func TestWatchSrvKeyspace(t *testing.T) { cell := "cell1" keyspace := "ks1" - ctx := context.Background() - ts := memorytopo.NewServer(cell) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cell) + defer ts.Close() // Create initial value if err := ts.UpdateSrvKeyspace(ctx, cell, keyspace, &topodatapb.SrvKeyspace{}); err != nil { @@ -175,8 +179,10 @@ func TestWatchSrvKeyspace(t *testing.T) { func TestWatchSrvKeyspaceCancel(t *testing.T) { cell := "cell1" keyspace := "ks1" - ctx := context.Background() - ts := memorytopo.NewServer(cell) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cell) + defer ts.Close() // No SrvKeyspace -> ErrNoNode _, _, err := ts.WatchSrvKeyspace(ctx, cell, keyspace) @@ -223,8 +229,10 @@ func TestUpdateSrvKeyspacePartitions(t *testing.T) { cell := "cell1" cell2 := "cell2" keyspace := "ks1" - ctx := context.Background() - ts := memorytopo.NewServer(cell, cell2) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cell, cell2) + defer ts.Close() keyRange, err := key.ParseShardingSpec("-") if err != nil || len(keyRange) != 1 { @@ -464,8 +472,10 @@ func TestUpdateUpdateDisableQueryService(t *testing.T) { cell := "cell1" cell2 := "cell2" keyspace := "ks1" - ctx := context.Background() - ts := memorytopo.NewServer(cell, cell2) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cell, cell2) + defer ts.Close() leftKeyRange, err := key.ParseShardingSpec("-80") if err != nil || len(leftKeyRange) != 1 { @@ -656,8 +666,10 @@ func TestGetShardServingTypes(t *testing.T) { cell := "cell1" cell2 := "cell2" keyspace := "ks1" - ctx := context.Background() - ts := memorytopo.NewServer(cell, cell2) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cell, cell2) + defer ts.Close() leftKeyRange, err := key.ParseShardingSpec("-80") if err != nil || len(leftKeyRange) != 1 { @@ -764,8 +776,10 @@ func TestGetShardServingCells(t *testing.T) { cell := "cell1" cell2 := "cell2" keyspace := "ks1" - ctx := context.Background() - ts := memorytopo.NewServer(cell, cell2) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cell, cell2) + defer ts.Close() leftKeyRange, err := key.ParseShardingSpec("-80") if err != nil || len(leftKeyRange) != 1 { @@ -868,8 +882,10 @@ func TestMasterMigrateServedType(t *testing.T) { cell := "cell1" cell2 := "cell2" keyspace := "ks1" - ctx := context.Background() - ts := memorytopo.NewServer(cell, cell2) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cell, cell2) + defer ts.Close() initialKeyRange, err := key.ParseShardingSpec("-") if err != nil || len(initialKeyRange) != 1 { @@ -1152,8 +1168,10 @@ func TestValidateSrvKeyspace(t *testing.T) { cell := "cell1" cell2 := "cell2" keyspace := "ks1" - ctx := context.Background() - ts := memorytopo.NewServer(cell, cell2) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cell, cell2) + defer ts.Close() leftKeyRange, err := key.ParseShardingSpec("-80") if err != nil || len(leftKeyRange) != 1 { diff --git a/go/vt/topo/topotests/srv_vschema_test.go b/go/vt/topo/topotests/srv_vschema_test.go index 73745854ae1..85a2d65c4ec 100644 --- a/go/vt/topo/topotests/srv_vschema_test.go +++ b/go/vt/topo/topotests/srv_vschema_test.go @@ -28,7 +28,6 @@ import ( ) func TestRebuildVSchema(t *testing.T) { - ctx := context.Background() emptySrvVSchema := &vschemapb.SrvVSchema{ RoutingRules: &vschemapb.RoutingRules{}, ShardRoutingRules: &vschemapb.ShardRoutingRules{}, @@ -36,7 +35,10 @@ func TestRebuildVSchema(t *testing.T) { // Set up topology. cells := []string{"cell1", "cell2"} - ts := memorytopo.NewServer(cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cells...) + defer ts.Close() // Rebuild with no keyspace / no vschema if err := ts.RebuildSrvVSchema(ctx, cells); err != nil { diff --git a/go/vt/topo/topotests/tablet_test.go b/go/vt/topo/topotests/tablet_test.go index e59b4c6d060..96bcdba1ae5 100644 --- a/go/vt/topo/topotests/tablet_test.go +++ b/go/vt/topo/topotests/tablet_test.go @@ -34,8 +34,10 @@ func TestCreateTablet(t *testing.T) { cell := "cell1" keyspace := "ks1" shard := "shard1" - ctx := context.Background() - ts := memorytopo.NewServer(cell) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cell) + defer ts.Close() // Create a tablet. alias := &topodatapb.TabletAlias{ diff --git a/go/vt/topo/topotests/wildcards_test.go b/go/vt/topo/topotests/wildcards_test.go index a87992b28de..d373a91b686 100644 --- a/go/vt/topo/topotests/wildcards_test.go +++ b/go/vt/topo/topotests/wildcards_test.go @@ -49,10 +49,12 @@ func (l *topoLayout) initTopo(t *testing.T, ts *topo.Server) { } func validateKeyspaceWildcard(t *testing.T, l *topoLayout, param string, expected []string) { - ts := memorytopo.NewServer() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx) + defer ts.Close() l.initTopo(t, ts) - ctx := context.Background() r, err := ts.ResolveKeyspaceWildcard(ctx, param) if err != nil { if expected != nil { @@ -85,10 +87,12 @@ func TestResolveKeyspaceWildcard(t *testing.T) { } func validateShardWildcard(t *testing.T, l *topoLayout, param string, expected []topo.KeyspaceShard) { - ts := memorytopo.NewServer() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx) + defer ts.Close() l.initTopo(t, ts) - ctx := context.Background() r, err := ts.ResolveShardWildcard(ctx, param) if err != nil { if expected != nil { @@ -181,10 +185,12 @@ func TestResolveShardWildcard(t *testing.T) { } func validateWildcards(t *testing.T, l *topoLayout, param string, expected []string) { - ts := memorytopo.NewServer() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx) + defer ts.Close() l.initTopo(t, ts) - ctx := context.Background() r, err := ts.ResolveWildcards(ctx, topo.GlobalCell, []string{param}) if err != nil { if expected != nil { diff --git a/go/vt/topo/zk2topo/server_test.go b/go/vt/topo/zk2topo/server_test.go index ebbca9898ce..a35e84f2950 100644 --- a/go/vt/topo/zk2topo/server_test.go +++ b/go/vt/topo/zk2topo/server_test.go @@ -17,12 +17,11 @@ limitations under the License. package zk2topo import ( + "context" "fmt" "path" "testing" - "context" - "vitess.io/vitess/go/testfiles" "vitess.io/vitess/go/vt/topo" "vitess.io/vitess/go/vt/topo/test" @@ -38,7 +37,9 @@ func TestZk2Topo(t *testing.T) { // Run the test suite. testIndex := 0 - test.TopoServerTestSuite(t, func() *topo.Server { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + test.TopoServerTestSuite(t, ctx, func() *topo.Server { // Each test will use its own sub-directories. // The directories will be created when used the first time. testRoot := fmt.Sprintf("/test-%v", testIndex) diff --git a/go/vt/topotools/routing_rules_test.go b/go/vt/topotools/routing_rules_test.go index 6047bb441fe..0b4f265a77b 100644 --- a/go/vt/topotools/routing_rules_test.go +++ b/go/vt/topotools/routing_rules_test.go @@ -28,8 +28,10 @@ import ( ) func TestRoutingRulesRoundTrip(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") + defer ts.Close() rules := map[string][]string{ "t1": {"t2", "t3"}, @@ -46,8 +48,10 @@ func TestRoutingRulesRoundTrip(t *testing.T) { } func TestRoutingRulesErrors(t *testing.T) { - ctx := context.Background() - ts, factory := memorytopo.NewServerAndFactory("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts, factory := memorytopo.NewServerAndFactory(ctx, "zone1") + defer ts.Close() factory.SetError(errors.New("topo failure for testing")) t.Run("GetRoutingRules error", func(t *testing.T) { @@ -68,8 +72,10 @@ func TestRoutingRulesErrors(t *testing.T) { } func TestShardRoutingRulesRoundTrip(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") + defer ts.Close() srr := map[string]string{ "ks1.shard1": "ks2", diff --git a/go/vt/topotools/shard_test.go b/go/vt/topotools/shard_test.go index 8904c984715..f38ba315034 100644 --- a/go/vt/topotools/shard_test.go +++ b/go/vt/topotools/shard_test.go @@ -17,24 +17,23 @@ limitations under the License. package topotools import ( + "context" "fmt" "math/rand" "sync" "testing" "time" - "context" - topodatapb "vitess.io/vitess/go/vt/proto/topodata" "vitess.io/vitess/go/vt/topo/memorytopo" ) // TestCreateShard tests a few cases for topo.CreateShard func TestCreateShard(t *testing.T) { - ctx := context.Background() - - // Set up topology. - ts := memorytopo.NewServer("test_cell") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "test_cell") + defer ts.Close() keyspace := "test_keyspace" shard := "0" @@ -60,10 +59,10 @@ func TestCreateShard(t *testing.T) { // TODO(sougou): we should eventually disallow multiple shards // for unsharded keyspaces. func TestCreateShardMultiUnsharded(t *testing.T) { - ctx := context.Background() - - // Set up topology. - ts := memorytopo.NewServer("test_cell") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "test_cell") + defer ts.Close() // create keyspace keyspace := "test_keyspace" @@ -102,11 +101,10 @@ func TestCreateShardMultiUnsharded(t *testing.T) { // for a long time in parallel, making sure the locking and everything // works correctly. func TestGetOrCreateShard(t *testing.T) { - ctx := context.Background() - - // Set up topology. - cell := "test_cell" - ts := memorytopo.NewServer(cell) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "test_cell") + defer ts.Close() // and do massive parallel GetOrCreateShard keyspace := "test_keyspace" diff --git a/go/vt/vtadmin/api_test.go b/go/vt/vtadmin/api_test.go index 2470e2995d7..3bd01ac5821 100644 --- a/go/vt/vtadmin/api_test.go +++ b/go/vt/vtadmin/api_test.go @@ -1050,19 +1050,19 @@ func TestGetKeyspace(t *testing.T) { }, } - ctx := context.Background() - for _, tt := range tests { tt := tt t.Run(tt.name, func(t *testing.T) { t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() topos := make([]*topo.Server, len(tt.clusterShards)) vtctlds := make([]vtctlservicepb.VtctldServer, len(tt.clusterShards)) for i, shards := range tt.clusterShards { - ts := memorytopo.NewServer("cell1") + ts := memorytopo.NewServer(ctx, "cell1") testutil.AddShards(ctx, t, ts, shards...) topos[i] = ts vtctlds[i] = testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { @@ -1282,20 +1282,20 @@ func TestGetKeyspaces(t *testing.T) { }, } - ctx := context.Background() - for _, tt := range tests { tt := tt t.Run(tt.name, func(t *testing.T) { t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() // Note that these test cases were written prior to the existence of // WithTestServers, so they are all written with the assumption that // there are exactly 2 clusters. topos := []*topo.Server{ - memorytopo.NewServer("c0_cell1"), - memorytopo.NewServer("c1_cell1"), + memorytopo.NewServer(ctx, "c0_cell1"), + memorytopo.NewServer(ctx, "c1_cell1"), } for cdx, cks := range tt.clusterKeyspaces { @@ -1346,7 +1346,8 @@ func TestGetKeyspaces(t *testing.T) { } func TestGetSchema(t *testing.T) { - t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() tests := []struct { name string @@ -1361,7 +1362,7 @@ func TestGetSchema(t *testing.T) { { name: "success", clusterID: 1, - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ GetSchemaResults: map[string]struct { Schema *tabletmanagerdatapb.SchemaDefinition @@ -1416,7 +1417,7 @@ func TestGetSchema(t *testing.T) { { name: "cluster not found", clusterID: 1, // results in clusterId == "c1" - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tablets: nil, req: &vtadminpb.GetSchemaRequest{ ClusterId: "c2", @@ -1429,7 +1430,7 @@ func TestGetSchema(t *testing.T) { { name: "tablet not found for keyspace", clusterID: 1, - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tablets: []*vtadminpb.Tablet{ { Cluster: &vtadminpb.Cluster{ @@ -1457,7 +1458,7 @@ func TestGetSchema(t *testing.T) { { name: "no serving tablet found for keyspace", clusterID: 1, - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tablets: []*vtadminpb.Tablet{ { Cluster: &vtadminpb.Cluster{ @@ -1485,7 +1486,7 @@ func TestGetSchema(t *testing.T) { { name: "error in GetSchema call", clusterID: 1, - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ GetSchemaResults: map[string]struct { Schema *tabletmanagerdatapb.SchemaDefinition @@ -1535,13 +1536,13 @@ func TestGetSchema(t *testing.T) { }, } - ctx := context.Background() - for _, tt := range tests { tt := tt t.Run(tt.name, func(t *testing.T) { t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, tt.ts, tt.tmc, func(ts *topo.Server) vtctlservicepb.VtctldServer { return grpcvtctldserver.NewVtctldServer(ts) @@ -1580,8 +1581,6 @@ func TestGetSchema(t *testing.T) { } t.Run("size aggregation", func(t *testing.T) { - t.Parallel() - c1pb := &vtadminpb.Cluster{ Id: "c1", Name: "cluster1", @@ -2176,8 +2175,6 @@ func TestGetSchemas(t *testing.T) { }, } - ctx := context.Background() - for _, tt := range tests { // Note that these test cases were written prior to the existence of // WithTestServers, so they are all written with the assumption that @@ -2186,10 +2183,12 @@ func TestGetSchemas(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() topos := []*topo.Server{ - memorytopo.NewServer("c0_cell1"), - memorytopo.NewServer("c1_cell1"), + memorytopo.NewServer(ctx, "c0_cell1"), + memorytopo.NewServer(ctx, "c1_cell1"), } tmc := testutil.TabletManagerClient{ @@ -2469,7 +2468,7 @@ func TestGetSchemas(t *testing.T) { api := NewAPI([]*cluster.Cluster{c1, c2}, Options{}) defer api.Close() - resp, err := api.GetSchemas(ctx, &vtadminpb.GetSchemasRequest{ + resp, err := api.GetSchemas(context.Background(), &vtadminpb.GetSchemasRequest{ TableSizeOptions: &vtadminpb.GetSchemaTableSizeOptions{ AggregateSizes: true, }, @@ -2627,17 +2626,18 @@ func TestGetSrvKeyspace(t *testing.T) { }, } - ctx := context.Background() - for _, tt := range tests { tt := tt t.Run(tt.name, func(t *testing.T) { t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + tmc := testutil.TabletManagerClient{} - toposerver := memorytopo.NewServer(tt.cells...) + toposerver := memorytopo.NewServer(ctx, tt.cells...) vtctldserver := testutil.NewVtctldServerWithTabletManagerClient(t, toposerver, &tmc, func(ts *topo.Server) vtctlservicepb.VtctldServer { return grpcvtctldserver.NewVtctldServer(ts) @@ -2787,17 +2787,17 @@ func TestGetSrvKeyspaces(t *testing.T) { }, } - ctx := context.Background() - for _, tt := range tests { tt := tt t.Run(tt.name, func(t *testing.T) { t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() tmc := testutil.TabletManagerClient{} - toposerver := memorytopo.NewServer(tt.cells...) + toposerver := memorytopo.NewServer(ctx, tt.cells...) for _, ks := range tt.keyspaces { testutil.AddKeyspace(ctx, t, toposerver, ks) @@ -2956,17 +2956,17 @@ func TestGetSrvVSchema(t *testing.T) { }, } - ctx := context.Background() - for _, tt := range tests { tt := tt t.Run(tt.name, func(t *testing.T) { t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() tmc := testutil.TabletManagerClient{} - toposerver := memorytopo.NewServer(tt.cells...) + toposerver := memorytopo.NewServer(ctx, tt.cells...) vtctldserver := testutil.NewVtctldServerWithTabletManagerClient(t, toposerver, &tmc, func(ts *topo.Server) vtctlservicepb.VtctldServer { return grpcvtctldserver.NewVtctldServer(ts) @@ -3250,17 +3250,17 @@ func TestGetSrvVSchemas(t *testing.T) { }, } - ctx := context.Background() - for _, tt := range tests { tt := tt t.Run(tt.name, func(t *testing.T) { t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() tmc := testutil.TabletManagerClient{} - toposerver := memorytopo.NewServer(tt.cells...) + toposerver := memorytopo.NewServer(ctx, tt.cells...) vtctldserver := testutil.NewVtctldServerWithTabletManagerClient(t, toposerver, &tmc, func(ts *topo.Server) vtctlservicepb.VtctldServer { return grpcvtctldserver.NewVtctldServer(ts) @@ -5099,13 +5099,13 @@ func TestVTExplain(t *testing.T) { }, } - ctx := context.Background() - for _, tt := range tests { tt := tt t.Run(tt.name, func(t *testing.T) { - toposerver := memorytopo.NewServer("c0_cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + toposerver := memorytopo.NewServer(ctx, "c0_cell1") tmc := testutil.TabletManagerClient{ GetSchemaResults: map[string]struct { diff --git a/go/vt/vtadmin/cluster/cluster_test.go b/go/vt/vtadmin/cluster/cluster_test.go index cc37d65c487..8111329acd5 100644 --- a/go/vt/vtadmin/cluster/cluster_test.go +++ b/go/vt/vtadmin/cluster/cluster_test.go @@ -50,9 +50,11 @@ import ( ) func TestCreateKeyspace(t *testing.T) { + utils.EnsureNoLeaks(t) t.Parallel() - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() tests := []struct { name string cfg testutil.TestClusterConfig @@ -158,11 +160,9 @@ func TestCreateKeyspace(t *testing.T) { } for _, tt := range tests { - tt := tt t.Run(tt.name, func(t *testing.T) { - t.Parallel() - cluster := testutil.BuildCluster(t, tt.cfg) + defer cluster.Close() resp, err := cluster.CreateKeyspace(ctx, tt.req) if tt.shouldErr { @@ -177,6 +177,7 @@ func TestCreateKeyspace(t *testing.T) { } func TestCreateShard(t *testing.T) { + utils.EnsureNoLeaks(t) t.Parallel() type test struct { @@ -186,11 +187,12 @@ func TestCreateShard(t *testing.T) { shouldErr bool assertion func(t *testing.T, tt *test) } - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() tests := []*test{ { name: "ok", - tc: testutil.BuildIntegrationTestCluster(t, &vtadminpb.Cluster{ + tc: testutil.BuildIntegrationTestCluster(t, ctx, &vtadminpb.Cluster{ Id: "local", Name: "local", }, "zone1"), @@ -211,7 +213,7 @@ func TestCreateShard(t *testing.T) { }, { name: "nil request", - tc: testutil.BuildIntegrationTestCluster(t, &vtadminpb.Cluster{ + tc: testutil.BuildIntegrationTestCluster(t, ctx, &vtadminpb.Cluster{ Id: "local", Name: "local", }, "zone1"), @@ -220,7 +222,7 @@ func TestCreateShard(t *testing.T) { }, { name: "no keyspace in request", - tc: testutil.BuildIntegrationTestCluster(t, &vtadminpb.Cluster{ + tc: testutil.BuildIntegrationTestCluster(t, ctx, &vtadminpb.Cluster{ Id: "local", Name: "local", }, "zone1"), @@ -232,7 +234,7 @@ func TestCreateShard(t *testing.T) { }, { name: "no shard name in request", - tc: testutil.BuildIntegrationTestCluster(t, &vtadminpb.Cluster{ + tc: testutil.BuildIntegrationTestCluster(t, ctx, &vtadminpb.Cluster{ Id: "local", Name: "local", }, "zone1"), @@ -244,7 +246,7 @@ func TestCreateShard(t *testing.T) { }, { name: "vtctld.CreateShard fails", - tc: testutil.BuildIntegrationTestCluster(t, &vtadminpb.Cluster{ + tc: testutil.BuildIntegrationTestCluster(t, ctx, &vtadminpb.Cluster{ Id: "local", Name: "local", }, "zone1"), @@ -277,9 +279,10 @@ func TestCreateShard(t *testing.T) { } func TestDeleteKeyspace(t *testing.T) { - t.Parallel() + utils.EnsureNoLeaks(t) - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() tests := []struct { name string cfg testutil.TestClusterConfig @@ -344,11 +347,9 @@ func TestDeleteKeyspace(t *testing.T) { } for _, tt := range tests { - tt := tt t.Run(tt.name, func(t *testing.T) { - t.Parallel() - cluster := testutil.BuildCluster(t, tt.cfg) + defer cluster.Close() resp, err := cluster.DeleteKeyspace(ctx, tt.req) if tt.shouldErr { @@ -363,6 +364,7 @@ func TestDeleteKeyspace(t *testing.T) { } func TestDeleteShards(t *testing.T) { + utils.EnsureNoLeaks(t) t.Parallel() type test struct { @@ -373,16 +375,16 @@ func TestDeleteShards(t *testing.T) { shouldErr bool assertion func(t *testing.T, tt *test) } - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() tests := []*test{ { name: "ok", - tc: testutil.BuildIntegrationTestCluster(t, &vtadminpb.Cluster{ + tc: testutil.BuildIntegrationTestCluster(t, ctx, &vtadminpb.Cluster{ Id: "local", Name: "local", }, "zone1"), setup: func(t *testing.T, tt *test) { - ctx := context.Background() shards := []string{"-80", "80-"} for _, shard := range shards { _, err := tt.tc.Cluster.CreateShard(ctx, &vtctldatapb.CreateShardRequest{ @@ -419,7 +421,7 @@ func TestDeleteShards(t *testing.T) { }, { name: "nil request", - tc: testutil.BuildIntegrationTestCluster(t, &vtadminpb.Cluster{ + tc: testutil.BuildIntegrationTestCluster(t, ctx, &vtadminpb.Cluster{ Id: "local", Name: "local", }, "zone1"), @@ -428,7 +430,7 @@ func TestDeleteShards(t *testing.T) { }, { name: "vtctld.DeleteShards fails", - tc: testutil.BuildIntegrationTestCluster(t, &vtadminpb.Cluster{ + tc: testutil.BuildIntegrationTestCluster(t, ctx, &vtadminpb.Cluster{ Id: "local", Name: "local", }, "zone1"), @@ -476,7 +478,6 @@ func TestDeleteShards(t *testing.T) { } for _, tt := range tests { - tt := tt t.Run(tt.name, func(t *testing.T) { if tt.setup != nil { func() { @@ -611,6 +612,7 @@ func TestFindTablet(t *testing.T) { }, Tablets: tt.tablets, }) + defer cluster.Close() tablet, err := cluster.FindTablet(ctx, tt.filter) if tt.expectedError != nil { @@ -822,6 +824,7 @@ func TestFindTablets(t *testing.T) { }, Tablets: tt.tablets, }) + defer cluster.Close() tablets, err := cluster.FindTablets(ctx, tt.filter, tt.n) assert.NoError(t, err) @@ -1203,6 +1206,7 @@ func TestFindWorkflows(t *testing.T) { t.Parallel() c := testutil.BuildCluster(t, tt.cfg) + defer c.Close() workflows, err := c.FindWorkflows(ctx, tt.keyspaces, tt.opts) if tt.shouldErr { assert.Error(t, err) @@ -1483,6 +1487,7 @@ func TestGetCellInfos(t *testing.T) { Cluster: cpb, VtctldClient: tt.vtctld, }) + defer c.Close() cellInfos, err := c.GetCellInfos(context.Background(), tt.req) if tt.shouldErr { assert.Error(t, err) @@ -1565,6 +1570,7 @@ func TestGetCellsAliases(t *testing.T) { Cluster: cpb, VtctldClient: tt.vtctld, }) + defer c.Close() cellsAliases, err := c.GetCellsAliases(context.Background()) if tt.shouldErr { assert.Error(t, err) @@ -1717,6 +1723,7 @@ func TestGetSchema(t *testing.T) { Tablets: []*vtadminpb.Tablet{tt.tablet}, DBConfig: testutil.Dbcfg{}, }) + defer c.Close() schema, err := c.GetSchema(ctx, "testkeyspace", cluster.GetSchemaOptions{ BaseRequest: tt.req, @@ -1768,6 +1775,7 @@ func TestGetSchema(t *testing.T) { }, VtctldClient: vtctld, }) + defer c.Close() _, _ = c.GetSchema(ctx, "testkeyspace", cluster.GetSchemaOptions{ BaseRequest: req, @@ -2690,6 +2698,7 @@ func TestGetSchema(t *testing.T) { } c := testutil.BuildCluster(t, tt.cfg) + defer c.Close() schema, err := c.GetSchema(ctx, tt.keyspace, tt.opts) if tt.shouldErr { assert.Error(t, err) @@ -2945,6 +2954,7 @@ func TestGetShardReplicationPositions(t *testing.T) { t.Parallel() c := testutil.BuildCluster(t, tt.cfg) + defer c.Close() resp, err := c.GetShardReplicationPositions(ctx, tt.req) if tt.shouldErr { @@ -3034,6 +3044,7 @@ func TestGetVSchema(t *testing.T) { t.Parallel() cluster := testutil.BuildCluster(t, tt.cfg) + defer cluster.Close() vschema, err := cluster.GetVSchema(ctx, tt.keyspace) if tt.shouldErr { @@ -3192,6 +3203,7 @@ func TestGetWorkflow(t *testing.T) { t.Parallel() c := testutil.BuildCluster(t, tt.cfg) + defer c.Close() workflow, err := c.GetWorkflow(ctx, tt.keyspace, tt.workflow, tt.opts) if tt.shouldErr { assert.Error(t, err) @@ -3358,6 +3370,7 @@ func TestGetWorkflows(t *testing.T) { t.Parallel() c := testutil.BuildCluster(t, tt.cfg) + defer c.Close() workflows, err := c.GetWorkflows(ctx, tt.keyspaces, tt.opts) if tt.shouldErr { assert.Error(t, err) @@ -3435,6 +3448,7 @@ func TestSetWritable(t *testing.T) { t.Parallel() c := testutil.BuildCluster(t, tt.cfg) + defer c.Close() err := c.SetWritable(ctx, tt.req) tt.assertion(t, err, tt.assertionMsgExtra...) }) @@ -3581,6 +3595,7 @@ func TestToggleTabletReplication(t *testing.T) { t.Parallel() c := testutil.BuildCluster(t, tt.cfg) + defer c.Close() err := c.ToggleTabletReplication(ctx, tt.tablet, bool(tt.state)) tt.assertion(t, err, tt.assertionMsgExtra...) }) diff --git a/go/vt/vtadmin/testutil/cluster.go b/go/vt/vtadmin/testutil/cluster.go index 2f825093e1f..9141d6b0c22 100644 --- a/go/vt/vtadmin/testutil/cluster.go +++ b/go/vt/vtadmin/testutil/cluster.go @@ -164,10 +164,10 @@ type IntegrationTestCluster struct { // // (TODO|@ajm188): Unify this with the BuildCluster API. Also this does not // support any cluster methods that involve vtgate/vitessdriver queries. -func BuildIntegrationTestCluster(t testing.TB, c *vtadminpb.Cluster, cells ...string) *IntegrationTestCluster { +func BuildIntegrationTestCluster(t testing.TB, ctx context.Context, c *vtadminpb.Cluster, cells ...string) *IntegrationTestCluster { t.Helper() - ts, factory := memorytopo.NewServerAndFactory(cells...) + ts, factory := memorytopo.NewServerAndFactory(ctx, cells...) vtctld := grpcvtctldtestutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return grpcvtctldserver.NewVtctldServer(ts) }) diff --git a/go/vt/vtcombo/tablet_map.go b/go/vt/vtcombo/tablet_map.go index 6f3e8511e07..4c5cff2023c 100644 --- a/go/vt/vtcombo/tablet_map.go +++ b/go/vt/vtcombo/tablet_map.go @@ -89,7 +89,7 @@ func CreateTablet( } log.Infof("Creating %v tablet %v for %v/%v", tabletType, topoproto.TabletAliasString(alias), keyspace, shard) - controller := tabletserver.NewServer(topoproto.TabletAliasString(alias), ts, alias) + controller := tabletserver.NewServer(ctx, topoproto.TabletAliasString(alias), ts, alias) initTabletType := tabletType if tabletType == topodatapb.TabletType_PRIMARY { initTabletType = topodatapb.TabletType_REPLICA diff --git a/go/vt/vtctl/endtoend/get_schema_test.go b/go/vt/vtctl/endtoend/get_schema_test.go index fe25080c7ae..2373fb6e3a5 100644 --- a/go/vt/vtctl/endtoend/get_schema_test.go +++ b/go/vt/vtctl/endtoend/get_schema_test.go @@ -26,9 +26,10 @@ import ( ) func TestGetSchema(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() - topo := memorytopo.NewServer("zone1", "zone2", "zone3") + topo := memorytopo.NewServer(ctx, "zone1", "zone2", "zone3") tablet := &topodatapb.Tablet{ Alias: &topodatapb.TabletAlias{ diff --git a/go/vt/vtctl/endtoend/onlineddl_show_test.go b/go/vt/vtctl/endtoend/onlineddl_show_test.go index 0b982bc7545..fe795af752d 100644 --- a/go/vt/vtctl/endtoend/onlineddl_show_test.go +++ b/go/vt/vtctl/endtoend/onlineddl_show_test.go @@ -94,9 +94,11 @@ func TestShowOnlineDDL_Cancel(t *testing.T) { func onlineDDLTest(t *testing.T, args []string, expectedQuery string) { t.Helper() - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() - fakeTopo := memorytopo.NewServer("zone1", "zone2", "zone3") + fakeTopo := memorytopo.NewServer(ctx, "zone1", "zone2", "zone3") + defer fakeTopo.Close() tablet := &topodatapb.Tablet{ Alias: &topodatapb.TabletAlias{ diff --git a/go/vt/vtctl/grpcvtctlclient/client_test.go b/go/vt/vtctl/grpcvtctlclient/client_test.go index a50a79ecdce..50e1968533e 100644 --- a/go/vt/vtctl/grpcvtctlclient/client_test.go +++ b/go/vt/vtctl/grpcvtctlclient/client_test.go @@ -17,6 +17,7 @@ limitations under the License. package grpcvtctlclient import ( + "context" "fmt" "io" "net" @@ -38,7 +39,9 @@ import ( // the test here creates a fake server implementation, a fake client // implementation, and runs the test suite against the setup. func TestVtctlServer(t *testing.T) { - ts := vtctlclienttest.CreateTopoServer(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := vtctlclienttest.CreateTopoServer(t, ctx) // Listen on a random port listener, err := net.Listen("tcp", "127.0.0.1:0") @@ -65,7 +68,9 @@ func TestVtctlServer(t *testing.T) { // the test here creates a fake server implementation, a fake client with auth // implementation, and runs the test suite against the setup. func TestVtctlAuthClient(t *testing.T) { - ts := vtctlclienttest.CreateTopoServer(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := vtctlclienttest.CreateTopoServer(t, ctx) // Listen on a random port listener, err := net.Listen("tcp", "127.0.0.1:0") diff --git a/go/vt/vtctl/grpcvtctldclient/client_test.go b/go/vt/vtctl/grpcvtctldclient/client_test.go index 1b42f8e5270..93c95ffa607 100644 --- a/go/vt/vtctl/grpcvtctldclient/client_test.go +++ b/go/vt/vtctl/grpcvtctldclient/client_test.go @@ -36,8 +36,10 @@ import ( ) func TestFindAllShardsInKeyspace(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") + defer ts.Close() vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return grpcvtctldserver.NewVtctldServer(ts) }) @@ -80,9 +82,11 @@ func TestFindAllShardsInKeyspace(t *testing.T) { } func TestGetKeyspace(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() - ts := memorytopo.NewServer("cell1") + ts := memorytopo.NewServer(ctx, "cell1") + defer ts.Close() vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return grpcvtctldserver.NewVtctldServer(ts) }) @@ -107,9 +111,11 @@ func TestGetKeyspace(t *testing.T) { } func TestGetKeyspaces(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() - ts := memorytopo.NewServer("cell1") + ts := memorytopo.NewServer(ctx, "cell1") + defer ts.Close() vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return grpcvtctldserver.NewVtctldServer(ts) }) diff --git a/go/vt/vtctl/grpcvtctldserver/endtoend/init_shard_primary_test.go b/go/vt/vtctl/grpcvtctldserver/endtoend/init_shard_primary_test.go index fb1f1f5b40e..a80fb89638a 100644 --- a/go/vt/vtctl/grpcvtctldserver/endtoend/init_shard_primary_test.go +++ b/go/vt/vtctl/grpcvtctldserver/endtoend/init_shard_primary_test.go @@ -41,7 +41,9 @@ import ( ) func TestInitShardPrimary(t *testing.T) { - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") tmc := tmclient.NewTabletManagerClient() wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmc) @@ -101,7 +103,9 @@ func TestInitShardPrimary(t *testing.T) { } func TestInitShardPrimaryNoFormerPrimary(t *testing.T) { - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") tmc := tmclient.NewTabletManagerClient() wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmc) diff --git a/go/vt/vtctl/grpcvtctldserver/server_slow_test.go b/go/vt/vtctl/grpcvtctldserver/server_slow_test.go index 9e38d936c42..3100855e370 100644 --- a/go/vt/vtctl/grpcvtctldserver/server_slow_test.go +++ b/go/vt/vtctl/grpcvtctldserver/server_slow_test.go @@ -45,7 +45,6 @@ func TestEmergencyReparentShardSlow(t *testing.T) { tests := []struct { name string - ts *topo.Server tmc tmclient.TabletManagerClient tablets []*topodatapb.Tablet @@ -64,7 +63,6 @@ func TestEmergencyReparentShardSlow(t *testing.T) { // concurrently, so the total time is only around 30 seconds, but // that's still a long time for a unit test! name: "nil WaitReplicasTimeout and request takes 29 seconds is ok", - ts: memorytopo.NewServer("zone1"), tablets: []*topodatapb.Tablet{ { Alias: &topodatapb.TabletAlias{ @@ -182,7 +180,6 @@ func TestEmergencyReparentShardSlow(t *testing.T) { }, { name: "nil WaitReplicasTimeout and request takes 31 seconds is error", - ts: memorytopo.NewServer("zone1"), tablets: []*topodatapb.Tablet{ { Alias: &topodatapb.TabletAlias{ @@ -292,8 +289,6 @@ func TestEmergencyReparentShardSlow(t *testing.T) { }, } - ctx := context.Background() - for _, tt := range tests { tt := tt @@ -304,13 +299,17 @@ func TestEmergencyReparentShardSlow(t *testing.T) { t.Skip("tt.EmergencyReparentShardRequest = nil implies test not ready to run") } - testutil.AddTablets(ctx, t, tt.ts, &testutil.AddTabletOptions{ + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") + + testutil.AddTablets(ctx, t, ts, &testutil.AddTabletOptions{ AlsoSetShardPrimary: true, ForceSetShardPrimary: true, SkipShardCreation: false, }, tt.tablets...) - vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, tt.ts, tt.tmc, func(ts *topo.Server) vtctlservicepb.VtctldServer { + vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, tt.tmc, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) resp, err := vtctld.EmergencyReparentShard(ctx, tt.req) @@ -345,7 +344,6 @@ func TestPlannedReparentShardSlow(t *testing.T) { tests := []struct { name string - ts *topo.Server tmc tmclient.TabletManagerClient tablets []*topodatapb.Tablet @@ -359,7 +357,6 @@ func TestPlannedReparentShardSlow(t *testing.T) { // nil WaitReplicasTimeout in the request results in a default 30 // second WaitReplicasTimeout. name: "nil WaitReplicasTimeout and request takes 29 seconds is ok", - ts: memorytopo.NewServer("zone1"), tablets: []*topodatapb.Tablet{ { Alias: &topodatapb.TabletAlias{ @@ -477,7 +474,6 @@ func TestPlannedReparentShardSlow(t *testing.T) { }, { name: "nil WaitReplicasTimeout and request takes 31 seconds is error", - ts: memorytopo.NewServer("zone1"), tablets: []*topodatapb.Tablet{ { Alias: &topodatapb.TabletAlias{ @@ -595,21 +591,23 @@ func TestPlannedReparentShardSlow(t *testing.T) { }, } - ctx := context.Background() - for _, tt := range tests { tt := tt t.Run(tt.name, func(t *testing.T) { t.Parallel() - testutil.AddTablets(ctx, t, tt.ts, &testutil.AddTabletOptions{ + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "zone1") + testutil.AddTablets(ctx, t, ts, &testutil.AddTabletOptions{ AlsoSetShardPrimary: true, ForceSetShardPrimary: true, SkipShardCreation: false, }, tt.tablets...) - vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, tt.ts, tt.tmc, func(ts *topo.Server) vtctlservicepb.VtctldServer { + vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, tt.tmc, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) resp, err := vtctld.PlannedReparentShard(ctx, tt.req) @@ -642,8 +640,10 @@ func TestPlannedReparentShardSlow(t *testing.T) { func TestSleepTablet(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "zone1") testutil.AddTablet(ctx, t, ts, &topodatapb.Tablet{ Alias: &topodatapb.TabletAlias{ Cell: "zone1", @@ -736,10 +736,7 @@ func TestSleepTablet(t *testing.T) { } for _, tt := range tests { - tt := tt t.Run(tt.name, func(t *testing.T) { - t.Parallel() - vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, &tt.tmc, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) diff --git a/go/vt/vtctl/grpcvtctldserver/server_test.go b/go/vt/vtctl/grpcvtctldserver/server_test.go index 1a55be69d94..b8cba1d64bd 100644 --- a/go/vt/vtctl/grpcvtctldserver/server_test.go +++ b/go/vt/vtctl/grpcvtctldserver/server_test.go @@ -97,7 +97,8 @@ func TestPanicHandler(t *testing.T) { func TestAddCellInfo(t *testing.T) { t.Parallel() - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() tests := []struct { name string ts *topo.Server @@ -105,7 +106,7 @@ func TestAddCellInfo(t *testing.T) { shouldErr bool }{ { - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), req: &vtctldatapb.AddCellInfoRequest{ Name: "zone2", CellInfo: &topodatapb.CellInfo{ @@ -116,7 +117,7 @@ func TestAddCellInfo(t *testing.T) { }, { name: "cell already exists", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), req: &vtctldatapb.AddCellInfoRequest{ Name: "zone1", CellInfo: &topodatapb.CellInfo{ @@ -128,7 +129,7 @@ func TestAddCellInfo(t *testing.T) { }, { name: "no cell root", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), req: &vtctldatapb.AddCellInfoRequest{ Name: "zone2", CellInfo: &topodatapb.CellInfo{ @@ -140,10 +141,7 @@ func TestAddCellInfo(t *testing.T) { } for _, tt := range tests { - tt := tt t.Run(tt.name, func(t *testing.T) { - t.Parallel() - vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, tt.ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -164,7 +162,8 @@ func TestAddCellInfo(t *testing.T) { func TestAddCellsAlias(t *testing.T) { t.Parallel() - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() tests := []struct { name string ts *topo.Server @@ -173,7 +172,7 @@ func TestAddCellsAlias(t *testing.T) { shouldErr bool }{ { - ts: memorytopo.NewServer("zone1", "zone2", "zone3"), + ts: memorytopo.NewServer(ctx, "zone1", "zone2", "zone3"), req: &vtctldatapb.AddCellsAliasRequest{ Name: "zone", Cells: []string{"zone1", "zone2", "zone3"}, @@ -181,7 +180,7 @@ func TestAddCellsAlias(t *testing.T) { }, { name: "alias exists", - ts: memorytopo.NewServer("zone1", "zone2", "zone3"), + ts: memorytopo.NewServer(ctx, "zone1", "zone2", "zone3"), setup: func(ts *topo.Server) error { return ts.CreateCellsAlias(ctx, "zone", &topodatapb.CellsAlias{ Cells: []string{"zone1", "zone2"}, @@ -195,7 +194,7 @@ func TestAddCellsAlias(t *testing.T) { }, { name: "alias overlaps", - ts: memorytopo.NewServer("zone1", "zone2", "zone3"), + ts: memorytopo.NewServer(ctx, "zone1", "zone2", "zone3"), setup: func(ts *topo.Server) error { return ts.CreateCellsAlias(context.Background(), "zone_a", &topodatapb.CellsAlias{ Cells: []string{"zone1", "zone3"}, @@ -210,10 +209,7 @@ func TestAddCellsAlias(t *testing.T) { } for _, tt := range tests { - tt := tt t.Run(tt.name, func(t *testing.T) { - t.Parallel() - if tt.setup != nil { err := tt.setup(tt.ts) require.NoError(t, err, "test setup failed") @@ -239,7 +235,8 @@ func TestAddCellsAlias(t *testing.T) { func TestApplyRoutingRules(t *testing.T) { t.Parallel() - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() tests := []struct { name string cells []string @@ -324,11 +321,8 @@ func TestApplyRoutingRules(t *testing.T) { } for _, tt := range tests { - tt := tt t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - ts, factory := memorytopo.NewServerAndFactory(tt.cells...) + ts, factory := memorytopo.NewServerAndFactory(ctx, tt.cells...) if tt.topoDown { factory.SetError(errors.New("topo down for testing")) } @@ -424,12 +418,10 @@ func TestApplyVSchema(t *testing.T) { } for _, tt := range tests { - tt := tt t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -511,7 +503,8 @@ func TestApplyVSchema(t *testing.T) { } func TestBackup(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() tests := []struct { name string ts *topo.Server @@ -523,7 +516,7 @@ func TestBackup(t *testing.T) { }{ { name: "ok", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ Backups: map[string]struct { Events []*logutilpb.Event @@ -561,7 +554,7 @@ func TestBackup(t *testing.T) { }, { name: "cannot backup primary", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ Backups: map[string]struct { Events []*logutilpb.Event @@ -596,7 +589,7 @@ func TestBackup(t *testing.T) { }, { name: "allow-primary", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ Backups: map[string]struct { Events []*logutilpb.Event @@ -632,7 +625,7 @@ func TestBackup(t *testing.T) { }, { name: "no tablet", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ Backups: map[string]struct { Events []*logutilpb.Event @@ -667,7 +660,7 @@ func TestBackup(t *testing.T) { }, { name: "midstream error", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ Backups: map[string]struct { Events []*logutilpb.Event @@ -705,7 +698,6 @@ func TestBackup(t *testing.T) { } for _, tt := range tests { - tt := tt t.Run(tt.name, func(t *testing.T) { if tt.tablet != nil { testutil.AddTablet(ctx, t, tt.ts, tt.tablet, nil) @@ -743,7 +735,8 @@ func TestBackup(t *testing.T) { } func TestBackupShard(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() tests := []struct { name string ts *topo.Server @@ -755,7 +748,7 @@ func TestBackupShard(t *testing.T) { }{ { name: "ok", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ Backups: map[string]struct { Events []*logutilpb.Event @@ -820,7 +813,7 @@ func TestBackupShard(t *testing.T) { }, { name: "cannot backup primary", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ Backups: map[string]struct { Events []*logutilpb.Event @@ -863,7 +856,7 @@ func TestBackupShard(t *testing.T) { }, { name: "allow-primary", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ Backups: map[string]struct { Events []*logutilpb.Event @@ -924,7 +917,7 @@ func TestBackupShard(t *testing.T) { }, { name: "incremental-from-pos", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ Backups: map[string]struct { Events []*logutilpb.Event @@ -990,7 +983,7 @@ func TestBackupShard(t *testing.T) { }, { name: "no available tablet", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ Backups: map[string]struct { Events []*logutilpb.Event @@ -1043,7 +1036,6 @@ func TestBackupShard(t *testing.T) { } for _, tt := range tests { - tt := tt t.Run(tt.name, func(t *testing.T) { testutil.AddTablets(ctx, t, tt.ts, &testutil.AddTabletOptions{ @@ -1294,8 +1286,9 @@ func TestChangeTabletType(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer(tt.cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, tt.cells...) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, &testutil.TabletManagerClient{ TopoServer: ts, }, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -1339,8 +1332,9 @@ func TestChangeTabletType(t *testing.T) { t.Run("tabletmanager failure", func(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, &testutil.TabletManagerClient{ TopoServer: nil, }, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -1817,8 +1811,9 @@ func TestCreateKeyspace(t *testing.T) { t.Skip("test not yet implemented") } - ctx := context.Background() - ts := memorytopo.NewServer(cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cells...) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -2094,8 +2089,9 @@ func TestCreateShard(t *testing.T) { t.Skip("focusing on other tests") } - ctx := context.Background() - ts, topofactory := memorytopo.NewServerAndFactory("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts, topofactory := memorytopo.NewServerAndFactory(ctx, "zone1") vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -2125,7 +2121,8 @@ func TestCreateShard(t *testing.T) { func TestDeleteCellInfo(t *testing.T) { t.Parallel() - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() tests := []struct { name string ts *topo.Server @@ -2133,14 +2130,14 @@ func TestDeleteCellInfo(t *testing.T) { shouldErr bool }{ { - ts: memorytopo.NewServer("zone1", "zone2"), + ts: memorytopo.NewServer(ctx, "zone1", "zone2"), req: &vtctldatapb.DeleteCellInfoRequest{ Name: "zone2", }, }, { name: "cell does not exist", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), req: &vtctldatapb.DeleteCellInfoRequest{ Name: "zone2", }, @@ -2149,10 +2146,7 @@ func TestDeleteCellInfo(t *testing.T) { } for _, tt := range tests { - tt := tt t.Run(tt.name, func(t *testing.T) { - t.Parallel() - vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, tt.ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -2172,7 +2166,8 @@ func TestDeleteCellInfo(t *testing.T) { func TestDeleteCellsAlias(t *testing.T) { t.Parallel() - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() tests := []struct { name string ts *topo.Server @@ -2181,7 +2176,7 @@ func TestDeleteCellsAlias(t *testing.T) { shouldErr bool }{ { - ts: memorytopo.NewServer("zone1", "zone2"), + ts: memorytopo.NewServer(ctx, "zone1", "zone2"), setup: func(ts *topo.Server) error { return ts.CreateCellsAlias(ctx, "zone", &topodatapb.CellsAlias{ Cells: []string{"zone1", "zone2"}, @@ -2193,7 +2188,7 @@ func TestDeleteCellsAlias(t *testing.T) { }, { name: "alias does not exist", - ts: memorytopo.NewServer("zone1", "zone2"), + ts: memorytopo.NewServer(ctx, "zone1", "zone2"), setup: func(ts *topo.Server) error { return ts.CreateCellsAlias(ctx, "zone_a", &topodatapb.CellsAlias{ Cells: []string{"zone1", "zone2"}, @@ -2207,10 +2202,7 @@ func TestDeleteCellsAlias(t *testing.T) { } for _, tt := range tests { - tt := tt t.Run(tt.name, func(t *testing.T) { - t.Parallel() - if tt.setup != nil { err := tt.setup(tt.ts) require.NoError(t, err, "test setup failed") @@ -2240,7 +2232,7 @@ func TestDeleteKeyspace(t *testing.T) { keyspaces []*vtctldatapb.Keyspace shards []*vtctldatapb.Shard srvKeyspaces map[string]map[string]*topodatapb.SrvKeyspace - before func(t *testing.T, ts *topo.Server, tt testcase) func() + before func(t *testing.T, ctx context.Context, ts *topo.Server, tt testcase) (context.Context, func()) topoErr error req *vtctldatapb.DeleteKeyspaceRequest expected *vtctldatapb.DeleteKeyspaceResponse @@ -2387,10 +2379,10 @@ func TestDeleteKeyspace(t *testing.T) { shards: nil, srvKeyspaces: nil, topoErr: nil, - before: func(t *testing.T, ts *topo.Server, tt testcase) func() { - _, unlock, err := ts.LockKeyspace(context.Background(), tt.req.Keyspace, "test.DeleteKeyspace") + before: func(t *testing.T, ctx context.Context, ts *topo.Server, tt testcase) (context.Context, func()) { + lctx, unlock, err := ts.LockKeyspace(ctx, tt.req.Keyspace, "test.DeleteKeyspace") require.NoError(t, err, "failed to lock keyspace %s before test", tt.req.Keyspace) - return func() { + return lctx, func() { unlock(&err) if !topo.IsErrType(err, topo.NoNode) { assert.NoError(t, err, "error while unlocking keyspace %s after test", tt.req.Keyspace) @@ -2418,10 +2410,10 @@ func TestDeleteKeyspace(t *testing.T) { shards: nil, srvKeyspaces: nil, topoErr: nil, - before: func(t *testing.T, ts *topo.Server, tt testcase) func() { - _, unlock, err := ts.LockKeyspace(context.Background(), tt.req.Keyspace, "test.DeleteKeyspace") + before: func(t *testing.T, ctx context.Context, ts *topo.Server, tt testcase) (context.Context, func()) { + lctx, unlock, err := ts.LockKeyspace(ctx, tt.req.Keyspace, "test.DeleteKeyspace") require.NoError(t, err, "failed to lock keyspace %s before test", tt.req.Keyspace) - return func() { + return lctx, func() { unlock(&err) if !topo.IsErrType(err, topo.NoNode) { assert.NoError(t, err, "error while unlocking keyspace %s after test", tt.req.Keyspace) @@ -2439,18 +2431,15 @@ func TestDeleteKeyspace(t *testing.T) { }, } - for _, tt := range tests { - tt := tt + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - t.Parallel() - cells := []string{"zone1", "zone2", "zone3"} - ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*50) - defer cancel() - - ts, topofactory := memorytopo.NewServerAndFactory(cells...) + ts, topofactory := memorytopo.NewServerAndFactory(ctx, cells...) + defer ts.Close() vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -2491,12 +2480,15 @@ func TestDeleteKeyspace(t *testing.T) { }() if tt.before != nil { - if after := tt.before(t, ts, tt); after != nil { + var after func() + if ctx, after = tt.before(t, ctx, ts, tt); after != nil { defer after() } } - resp, err := vtctld.DeleteKeyspace(ctx, tt.req) + requestCtx, requestCancel := context.WithTimeout(ctx, time.Millisecond*50) + defer requestCancel() + resp, err := vtctld.DeleteKeyspace(requestCtx, tt.req) if tt.shouldErr { assert.Error(t, err) @@ -2519,7 +2511,7 @@ func TestDeleteShards(t *testing.T) { replicationGraphs []*topo.ShardReplicationInfo srvKeyspaces map[string]map[string]*topodatapb.SrvKeyspace topoErr error - before func(t *testing.T, ts *topo.Server, tt testcase) func() + before func(t *testing.T, ctx context.Context, ts *topo.Server, tt testcase) (context.Context, func()) req *vtctldatapb.DeleteShardsRequest expected *vtctldatapb.DeleteShardsResponse expectedRemainingShards []*vtctldatapb.Shard @@ -2886,11 +2878,11 @@ func TestDeleteShards(t *testing.T) { }, tablets: nil, topoErr: nil, - before: func(t *testing.T, ts *topo.Server, tt testcase) func() { + before: func(t *testing.T, ctx context.Context, ts *topo.Server, tt testcase) (context.Context, func()) { shard := tt.req.Shards[0] - _, unlock, err := ts.LockShard(context.Background(), shard.Keyspace, shard.Name, "test.DeleteShard") + lctx, unlock, err := ts.LockShard(ctx, shard.Keyspace, shard.Name, "test.DeleteShard") require.NoError(t, err, "failed to lock shard %s/%s before test", shard.Keyspace, shard.Name) - return func() { + return lctx, func() { unlock(&err) if !topo.IsErrType(err, topo.NoNode) { assert.NoError(t, err, "error while unlocking shard %s/%s after test", shard.Keyspace, shard.Name) @@ -2924,11 +2916,11 @@ func TestDeleteShards(t *testing.T) { }, tablets: nil, topoErr: nil, - before: func(t *testing.T, ts *topo.Server, tt testcase) func() { + before: func(t *testing.T, ctx context.Context, ts *topo.Server, tt testcase) (context.Context, func()) { shard := tt.req.Shards[0] - _, unlock, err := ts.LockShard(context.Background(), shard.Keyspace, shard.Name, "test.DeleteShard") + lctx, unlock, err := ts.LockShard(ctx, shard.Keyspace, shard.Name, "test.DeleteShard") require.NoError(t, err, "failed to lock shard %s/%s before test", shard.Keyspace, shard.Name) - return func() { + return lctx, func() { unlock(&err) if !topo.IsErrType(err, topo.NoNode) { assert.NoError(t, err, "error while unlocking shard %s/%s after test", shard.Keyspace, shard.Name) @@ -2961,7 +2953,7 @@ func TestDeleteShards(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*50) defer cancel() - ts, topofactory := memorytopo.NewServerAndFactory(cells...) + ts, topofactory := memorytopo.NewServerAndFactory(ctx, cells...) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -3001,7 +2993,8 @@ func TestDeleteShards(t *testing.T) { } if tt.before != nil { - if after := tt.before(t, ts, tt); after != nil { + var after func() + if ctx, after = tt.before(t, ctx, ts, tt); after != nil { defer after() } } @@ -3022,7 +3015,6 @@ func TestDeleteShards(t *testing.T) { func TestDeleteSrvKeyspace(t *testing.T) { t.Parallel() - ctx := context.Background() tests := []struct { name string vschemas map[string]*vschemapb.SrvVSchema @@ -3097,7 +3089,9 @@ func TestDeleteSrvKeyspace(t *testing.T) { finalVSchemas[cell] = vschema } - ts := memorytopo.NewServer(cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cells...) for cell, vschema := range tt.vschemas { err := ts.UpdateSrvVSchema(ctx, cell, vschema) require.NoError(t, err, "failed to update SrvVSchema in cell = %v, vschema = %+v", cell, vschema) @@ -3561,8 +3555,9 @@ func TestDeleteTablets(t *testing.T) { t.Skip("focusing on other tests") } - ctx := context.Background() - ts, topofactory := memorytopo.NewServerAndFactory("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts, topofactory := memorytopo.NewServerAndFactory(ctx, "zone1") vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -3626,6 +3621,9 @@ func TestDeleteTablets(t *testing.T) { func TestEmergencyReparentShard(t *testing.T) { t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + tests := []struct { name string ts *topo.Server @@ -3639,7 +3637,7 @@ func TestEmergencyReparentShard(t *testing.T) { }{ { name: "successful reparent", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tablets: []*topodatapb.Tablet{ { Alias: &topodatapb.TabletAlias{ @@ -3759,7 +3757,7 @@ func TestEmergencyReparentShard(t *testing.T) { // the simplest way to trigger a failure is to attempt an ERS on a // shard that does not exist. name: "failed reparent", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tablets: nil, req: &vtctldatapb.EmergencyReparentShardRequest{ @@ -3781,14 +3779,8 @@ func TestEmergencyReparentShard(t *testing.T) { }, } - ctx := context.Background() - for _, tt := range tests { - tt := tt - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - testutil.AddTablets(ctx, t, tt.ts, &testutil.AddTabletOptions{ AlsoSetShardPrimary: true, ForceSetShardPrimary: true, @@ -3932,8 +3924,9 @@ func TestExecuteFetchAsApp(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") testutil.AddTablet(ctx, t, ts, tt.tablet, nil) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, tt.tmc, func(ts *topo.Server) vtctlservicepb.VtctldServer { @@ -4058,8 +4051,9 @@ func TestExecuteFetchAsDBA(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") testutil.AddTablet(ctx, t, ts, tt.tablet, nil) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, tt.tmc, func(ts *topo.Server) vtctlservicepb.VtctldServer { @@ -4080,6 +4074,9 @@ func TestExecuteFetchAsDBA(t *testing.T) { func TestExecuteHook(t *testing.T) { t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + tests := []struct { name string ts *topo.Server @@ -4090,7 +4087,7 @@ func TestExecuteHook(t *testing.T) { }{ { name: "ok", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ ExecuteHookResults: map[string]struct { Response *hk.HookResult @@ -4119,7 +4116,7 @@ func TestExecuteHook(t *testing.T) { }, { name: "nil hook request", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ ExecuteHookResults: map[string]struct { Response *hk.HookResult @@ -4149,7 +4146,7 @@ func TestExecuteHook(t *testing.T) { }, { name: "hook with slash", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ ExecuteHookResults: map[string]struct { Response *hk.HookResult @@ -4181,7 +4178,7 @@ func TestExecuteHook(t *testing.T) { }, { name: "no such tablet", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ ExecuteHookResults: map[string]struct { Response *hk.HookResult @@ -4211,7 +4208,7 @@ func TestExecuteHook(t *testing.T) { }, { name: "tablet hook failure", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ ExecuteHookResults: map[string]struct { Response *hk.HookResult @@ -4241,12 +4238,8 @@ func TestExecuteHook(t *testing.T) { }, } - ctx := context.Background() for _, tt := range tests { - tt := tt t.Run(tt.name, func(t *testing.T) { - t.Parallel() - testutil.AddTablets(ctx, t, tt.ts, nil, tt.tablets...) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, tt.ts, tt.tmc, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) @@ -4266,8 +4259,9 @@ func TestExecuteHook(t *testing.T) { func TestFindAllShardsInKeyspace(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -4307,8 +4301,9 @@ func TestFindAllShardsInKeyspace(t *testing.T) { } func TestGetBackups(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -4414,8 +4409,9 @@ func TestGetBackups(t *testing.T) { func TestGetKeyspace(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -4439,8 +4435,9 @@ func TestGetKeyspace(t *testing.T) { func TestGetCellInfoNames(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer("cell1", "cell2", "cell3") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2", "cell3") vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -4449,7 +4446,7 @@ func TestGetCellInfoNames(t *testing.T) { assert.NoError(t, err) assert.ElementsMatch(t, []string{"cell1", "cell2", "cell3"}, resp.Names) - ts = memorytopo.NewServer() + ts = memorytopo.NewServer(ctx) vtctld = testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -4458,7 +4455,7 @@ func TestGetCellInfoNames(t *testing.T) { assert.NoError(t, err) assert.Empty(t, resp.Names) - ts, topofactory := memorytopo.NewServerAndFactory("cell1") + ts, topofactory := memorytopo.NewServerAndFactory(ctx, "cell1") vtctld = testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -4471,8 +4468,9 @@ func TestGetCellInfoNames(t *testing.T) { func TestGetCellInfo(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -4498,8 +4496,9 @@ func TestGetCellInfo(t *testing.T) { func TestGetCellsAliases(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer("c11", "c12", "c13", "c21", "c22") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "c11", "c12", "c13", "c21", "c22") vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -4526,7 +4525,7 @@ func TestGetCellsAliases(t *testing.T) { assert.NoError(t, err) utils.MustMatch(t, expected, resp.Aliases) - ts, topofactory := memorytopo.NewServerAndFactory() + ts, topofactory := memorytopo.NewServerAndFactory(ctx) vtctld = testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -4600,8 +4599,9 @@ func TestGetFullStatus(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer(tt.cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, tt.cells...) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, &testutil.TabletManagerClient{ TopoServer: ts, FullStatusResult: &replicationdatapb.FullStatus{ @@ -4628,8 +4628,9 @@ func TestGetFullStatus(t *testing.T) { func TestGetKeyspaces(t *testing.T) { t.Parallel() - ctx := context.Background() - ts, topofactory := memorytopo.NewServerAndFactory("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts, topofactory := memorytopo.NewServerAndFactory(ctx, "cell1") vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -4669,7 +4670,6 @@ func TestGetKeyspaces(t *testing.T) { func TestGetPermissions(t *testing.T) { t.Parallel() - ctx := context.Background() var testGetPermissionsReply = &tabletmanagerdatapb.Permissions{ UserPermissions: []*tabletmanagerdatapb.UserPermission{ { @@ -4793,7 +4793,10 @@ func TestGetPermissions(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "zone1") testutil.AddTablets(ctx, t, ts, nil, tt.tablets...) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, &tt.tmc, func(ts *topo.Server) vtctlservicepb.VtctldServer { @@ -4816,7 +4819,6 @@ func TestGetPermissions(t *testing.T) { func TestGetRoutingRules(t *testing.T) { t.Parallel() - ctx := context.Background() tests := []struct { name string topoDown bool @@ -4860,7 +4862,10 @@ func TestGetRoutingRules(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ts, factory := memorytopo.NewServerAndFactory() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts, factory := memorytopo.NewServerAndFactory(ctx) if tt.rrIn != nil { err := ts.SaveRoutingRules(ctx, tt.rrIn) require.NoError(t, err, "could not save routing rules: %+v", tt.rrIn) @@ -4886,8 +4891,9 @@ func TestGetRoutingRules(t *testing.T) { } func TestGetSchema(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") tmc := testutil.TabletManagerClient{ GetSchemaResults: map[string]struct { Schema *tabletmanagerdatapb.SchemaDefinition @@ -5255,8 +5261,9 @@ func TestGetSchemaMigrations(t *testing.T) { cells := []string{"zone1", "zone2", "zone3"} - ctx := context.Background() - ts, factory := memorytopo.NewServerAndFactory(cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts, factory := memorytopo.NewServerAndFactory(ctx, cells...) testutil.AddTablets(ctx, t, ts, &testutil.AddTabletOptions{AlsoSetShardPrimary: true}, test.tablets...) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, tmc, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) @@ -5346,8 +5353,9 @@ func TestGetShard(t *testing.T) { cells := []string{"zone1", "zone2", "zone3"} - ctx := context.Background() - ts, topofactory := memorytopo.NewServerAndFactory(cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts, topofactory := memorytopo.NewServerAndFactory(ctx, cells...) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -5372,7 +5380,6 @@ func TestGetShard(t *testing.T) { func TestGetSrvKeyspaceNames(t *testing.T) { t.Parallel() - ctx := context.Background() tests := []struct { name string srvKeyspacesByCell map[string]map[string]*topodatapb.SrvKeyspace @@ -5471,7 +5478,9 @@ func TestGetSrvKeyspaceNames(t *testing.T) { cells = append(cells, cell) } - ts, factory := memorytopo.NewServerAndFactory(cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts, factory := memorytopo.NewServerAndFactory(ctx, cells...) for cell, srvKeyspaces := range tt.srvKeyspacesByCell { for ks, srvks := range srvKeyspaces { @@ -5625,8 +5634,6 @@ func TestGetSrvKeyspaces(t *testing.T) { }, } - ctx := context.Background() - for _, tt := range tests { tt := tt @@ -5637,7 +5644,10 @@ func TestGetSrvKeyspaces(t *testing.T) { t.SkipNow() } - ts, topofactory := memorytopo.NewServerAndFactory(tt.cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts, topofactory := memorytopo.NewServerAndFactory(ctx, tt.cells...) testutil.AddSrvKeyspaces(t, ts, tt.srvKeyspaces...) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { @@ -5664,8 +5674,9 @@ func TestGetSrvKeyspaces(t *testing.T) { func TestGetSrvVSchema(t *testing.T) { t.Parallel() - ctx := context.Background() - ts, topofactory := memorytopo.NewServerAndFactory("zone1", "zone2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts, topofactory := memorytopo.NewServerAndFactory(ctx, "zone1", "zone2") vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -5874,8 +5885,9 @@ func TestGetSrvVSchemas(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ctx := context.Background() - ts, topofactory := memorytopo.NewServerAndFactory("zone1", "zone2", "zone3") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts, topofactory := memorytopo.NewServerAndFactory(ctx, "zone1", "zone2", "zone3") vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -5933,8 +5945,9 @@ func TestGetSrvVSchemas(t *testing.T) { func TestGetTablet(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -6412,8 +6425,9 @@ func TestGetTablets(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer(tt.cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, tt.cells...) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -6435,8 +6449,9 @@ func TestGetTablets(t *testing.T) { func TestGetTopologyPath(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer("cell1", "cell2", "cell3") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2", "cell3") vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -6523,15 +6538,14 @@ func TestGetTopologyPath(t *testing.T) { func TestGetVSchema(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) t.Run("found", func(t *testing.T) { - t.Parallel() - err := ts.SaveVSchema(ctx, "testkeyspace", &vschemapb.Keyspace{ Sharded: true, Vindexes: map[string]*vschemapb.Vindex{ @@ -6573,8 +6587,9 @@ func TestGetVSchema(t *testing.T) { func TestPingTablet(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") testutil.AddTablet(ctx, t, ts, &topodatapb.Tablet{ Alias: &topodatapb.TabletAlias{ Cell: "zone1", @@ -6639,10 +6654,7 @@ func TestPingTablet(t *testing.T) { } for _, tt := range tests { - tt := tt t.Run(tt.name, func(t *testing.T) { - t.Parallel() - vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, &tt.tmc, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -6663,6 +6675,9 @@ func TestPingTablet(t *testing.T) { func TestPlannedReparentShard(t *testing.T) { t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + tests := []struct { name string ts *topo.Server @@ -6676,7 +6691,7 @@ func TestPlannedReparentShard(t *testing.T) { }{ { name: "successful reparent", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tablets: []*topodatapb.Tablet{ { Alias: &topodatapb.TabletAlias{ @@ -6796,7 +6811,7 @@ func TestPlannedReparentShard(t *testing.T) { // the simplest way to trigger a failure is to attempt an PRS on a // shard that does not exist. name: "failed reparent", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tablets: nil, req: &vtctldatapb.PlannedReparentShardRequest{ Keyspace: "testkeyspace", @@ -6817,14 +6832,8 @@ func TestPlannedReparentShard(t *testing.T) { }, } - ctx := context.Background() - for _, tt := range tests { - tt := tt - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - testutil.AddTablets(ctx, t, tt.ts, &testutil.AddTabletOptions{ AlsoSetShardPrimary: true, ForceSetShardPrimary: true, @@ -6867,8 +6876,10 @@ func TestRebuildKeyspaceGraph(t *testing.T) { t.Run("ok", func(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "zone1") testutil.AddKeyspace(ctx, t, ts, &vtctldatapb.Keyspace{ Name: "testkeyspace", }) @@ -6885,7 +6896,10 @@ func TestRebuildKeyspaceGraph(t *testing.T) { t.Run("no such keyspace", func(t *testing.T) { t.Parallel() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "zone1") vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -6899,8 +6913,10 @@ func TestRebuildKeyspaceGraph(t *testing.T) { t.Run("topo unavailable", func(t *testing.T) { t.Parallel() - ctx := context.Background() - ts, factory := memorytopo.NewServerAndFactory("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts, factory := memorytopo.NewServerAndFactory(ctx, "zone1") testutil.AddKeyspace(ctx, t, ts, &vtctldatapb.Keyspace{ Name: "testkeyspace", }) @@ -6918,8 +6934,10 @@ func TestRebuildKeyspaceGraph(t *testing.T) { t.Run("lock error", func(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "zone1") testutil.AddKeyspace(ctx, t, ts, &vtctldatapb.Keyspace{ Name: "testkeyspace", }) @@ -6927,15 +6945,13 @@ func TestRebuildKeyspaceGraph(t *testing.T) { return NewVtctldServer(ts) }) - _, unlock, lerr := ts.LockKeyspace(context.Background(), "testkeyspace", "test lock") + lctx, unlock, lerr := ts.LockKeyspace(context.Background(), "testkeyspace", "test lock") require.NoError(t, lerr, "could not lock keyspace for testing") defer unlock(&lerr) defer func() { require.NoError(t, lerr, "could not unlock testkeyspace after test") }() - ctx, cancel := context.WithTimeout(ctx, time.Millisecond*50) - defer cancel() - _, err := vtctld.RebuildKeyspaceGraph(ctx, &vtctldatapb.RebuildKeyspaceGraphRequest{ + _, err := vtctld.RebuildKeyspaceGraph(lctx, &vtctldatapb.RebuildKeyspaceGraphRequest{ Keyspace: "testkeyspace", }) assert.Error(t, err) @@ -6945,7 +6961,6 @@ func TestRebuildKeyspaceGraph(t *testing.T) { func TestRebuildVSchemaGraph(t *testing.T) { t.Parallel() - ctx := context.Background() req := &vtctldatapb.RebuildVSchemaGraphRequest{} tests := []struct { name string @@ -6967,7 +6982,10 @@ func TestRebuildVSchemaGraph(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ts, factory := memorytopo.NewServerAndFactory("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts, factory := memorytopo.NewServerAndFactory(ctx, "zone1") if tt.topoDown { factory.SetError(errors.New("topo down for testing")) } @@ -6989,7 +7007,9 @@ func TestRebuildVSchemaGraph(t *testing.T) { func TestRefreshState(t *testing.T) { t.Parallel() - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + tests := []struct { name string ts *topo.Server @@ -7000,7 +7020,7 @@ func TestRefreshState(t *testing.T) { }{ { name: "success", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tablet: &topodatapb.Tablet{ Alias: &topodatapb.TabletAlias{ Cell: "zone1", @@ -7017,13 +7037,13 @@ func TestRefreshState(t *testing.T) { }, { name: "tablet alias nil", - ts: memorytopo.NewServer(), + ts: memorytopo.NewServer(ctx), req: &vtctldatapb.RefreshStateRequest{}, shouldErr: true, }, { name: "tablet not found", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tablet: &topodatapb.Tablet{ Alias: &topodatapb.TabletAlias{ Cell: "zone1", @@ -7041,7 +7061,7 @@ func TestRefreshState(t *testing.T) { }, { name: "RefreshState failed", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tablet: &topodatapb.Tablet{ Alias: &topodatapb.TabletAlias{ Cell: "zone1", @@ -7060,11 +7080,7 @@ func TestRefreshState(t *testing.T) { } for _, tt := range tests { - tt := tt - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - var tmc testutil.TabletManagerClient if tt.tablet != nil { testutil.AddTablet(ctx, t, tt.ts, tt.tablet, nil) @@ -7090,7 +7106,8 @@ func TestRefreshState(t *testing.T) { func TestRefreshStateByShard(t *testing.T) { t.Parallel() - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() tests := []struct { name string ts *topo.Server @@ -7102,7 +7119,7 @@ func TestRefreshStateByShard(t *testing.T) { }{ { name: "success", - ts: memorytopo.NewServer("zone1", "zone2"), + ts: memorytopo.NewServer(ctx, "zone1", "zone2"), tablets: []*topodatapb.Tablet{ { Hostname: "zone1-100", @@ -7135,7 +7152,7 @@ func TestRefreshStateByShard(t *testing.T) { }, { name: "cell filtering", - ts: memorytopo.NewServer("zone1", "zone2"), + ts: memorytopo.NewServer(ctx, "zone1", "zone2"), tablets: []*topodatapb.Tablet{ { Hostname: "zone1-100", @@ -7173,7 +7190,7 @@ func TestRefreshStateByShard(t *testing.T) { }, { name: "partial result", - ts: memorytopo.NewServer("zone1", "zone2"), + ts: memorytopo.NewServer(ctx, "zone1", "zone2"), tablets: []*topodatapb.Tablet{ { Hostname: "zone1-100", @@ -7210,13 +7227,13 @@ func TestRefreshStateByShard(t *testing.T) { }, { name: "missing keyspace argument", - ts: memorytopo.NewServer(), + ts: memorytopo.NewServer(ctx), req: &vtctldatapb.RefreshStateByShardRequest{}, shouldErr: true, }, { name: "missing shard argument", - ts: memorytopo.NewServer(), + ts: memorytopo.NewServer(ctx), req: &vtctldatapb.RefreshStateByShardRequest{ Keyspace: "ks", }, @@ -7224,7 +7241,7 @@ func TestRefreshStateByShard(t *testing.T) { }, { name: "shard not found", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tablets: []*topodatapb.Tablet{ { Alias: &topodatapb.TabletAlias{ @@ -7245,11 +7262,7 @@ func TestRefreshStateByShard(t *testing.T) { } for _, tt := range tests { - tt := tt - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - require.Equal(t, len(tt.tablets), len(tt.refreshStateErrors), "Invalid test case: must have one refreshStateError for each tablet") tmc := &testutil.TabletManagerClient{ @@ -7356,11 +7369,13 @@ func TestReloadSchema(t *testing.T) { }, } - ctx := context.Background() for _, tt := range tests { tt := tt t.Run(tt.name, func(t *testing.T) { - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "zone1") testutil.AddTablets(ctx, t, ts, nil, tt.tablets...) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, &tt.tmc, func(ts *topo.Server) vtctlservicepb.VtctldServer { @@ -7448,13 +7463,15 @@ func TestReloadSchemaKeyspace(t *testing.T) { }, } - ctx := context.Background() for _, tt := range tests { tt := tt t.Run(tt.name, func(t *testing.T) { t.Parallel() - ts := memorytopo.NewServer("zone1", "zone2", "zone3") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "zone1", "zone2", "zone3") testutil.AddTablets(ctx, t, ts, &testutil.AddTabletOptions{ AlsoSetShardPrimary: true, }, tt.tablets...) @@ -7604,13 +7621,15 @@ func TestReloadSchemaShard(t *testing.T) { }, } - ctx := context.Background() for _, tt := range tests { tt := tt t.Run(tt.name, func(t *testing.T) { t.Parallel() - ts := memorytopo.NewServer("zone1", "zone2", "zone3") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "zone1", "zone2", "zone3") testutil.AddTablets(ctx, t, ts, &testutil.AddTabletOptions{ AlsoSetShardPrimary: true, }, tt.tablets...) @@ -7631,8 +7650,9 @@ func TestReloadSchemaShard(t *testing.T) { } func TestRemoveBackup(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -7821,8 +7841,9 @@ func TestRemoveKeyspaceCell(t *testing.T) { cells := []string{"zone1", "zone2", "zone3"} - ctx := context.Background() - ts, topofactory := memorytopo.NewServerAndFactory(cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts, topofactory := memorytopo.NewServerAndFactory(ctx, cells...) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -8109,8 +8130,9 @@ func TestRemoveShardCell(t *testing.T) { cells := []string{"zone1", "zone2", "zone3"} - ctx := context.Background() - ts, topofactory := memorytopo.NewServerAndFactory(cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts, topofactory := memorytopo.NewServerAndFactory(ctx, cells...) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -8718,8 +8740,9 @@ func TestReparentTablet(t *testing.T) { cells := []string{"zone1", "zone2", "zone3"} - ctx := context.Background() - ts, topofactory := memorytopo.NewServerAndFactory(cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts, topofactory := memorytopo.NewServerAndFactory(ctx, cells...) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, tt.tmc, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -8746,7 +8769,9 @@ func TestReparentTablet(t *testing.T) { } func TestRestoreFromBackup(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + tests := []struct { name string ts *topo.Server @@ -8758,7 +8783,7 @@ func TestRestoreFromBackup(t *testing.T) { }{ { name: "ok", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ RestoreFromBackupResults: map[string]struct { Events []*logutilpb.Event @@ -8807,7 +8832,7 @@ func TestRestoreFromBackup(t *testing.T) { }, { name: "no such tablet", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ Backups: map[string]struct { Events []*logutilpb.Event @@ -8845,7 +8870,6 @@ func TestRestoreFromBackup(t *testing.T) { } for _, tt := range tests { - tt := tt t.Run(tt.name, func(t *testing.T) { testutil.AddTablets(ctx, t, tt.ts, &testutil.AddTabletOptions{ @@ -9090,7 +9114,6 @@ func TestRetrySchemaMigration(t *testing.T) { func TestRunHealthCheck(t *testing.T) { t.Parallel() - ctx := context.Background() tests := []struct { name string tablets []*topodatapb.Tablet @@ -9173,7 +9196,10 @@ func TestRunHealthCheck(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "zone1") testutil.AddTablets(ctx, t, ts, nil, tt.tablets...) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, &tt.tmc, func(ts *topo.Server) vtctlservicepb.VtctldServer { @@ -9245,14 +9271,15 @@ func TestSetKeyspaceDurabilityPolicy(t *testing.T) { }, } - ctx := context.Background() - for _, tt := range tests { tt := tt t.Run(tt.name, func(t *testing.T) { t.Parallel() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "zone1") testutil.AddKeyspaces(ctx, t, ts, tt.keyspaces...) vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { @@ -9273,6 +9300,9 @@ func TestSetKeyspaceDurabilityPolicy(t *testing.T) { func TestSetShardIsPrimaryServing(t *testing.T) { t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + type testcase struct { name string ctx context.Context @@ -9288,8 +9318,8 @@ func TestSetShardIsPrimaryServing(t *testing.T) { { name: "ok", setup: func(t *testing.T, tt *testcase) { - tt.ctx = context.Background() - tt.ts = memorytopo.NewServer("zone1") + tt.ctx = ctx + tt.ts = memorytopo.NewServer(ctx, "zone1") testutil.AddShards(tt.ctx, t, tt.ts, &vtctldatapb.Shard{ Keyspace: "testkeyspace", Name: "-", @@ -9311,8 +9341,8 @@ func TestSetShardIsPrimaryServing(t *testing.T) { name: "lock error", setup: func(t *testing.T, tt *testcase) { var cancel func() - tt.ctx, cancel = context.WithTimeout(context.Background(), time.Millisecond*50) - tt.ts = memorytopo.NewServer("zone1") + tt.ctx, cancel = context.WithTimeout(ctx, time.Millisecond*50) + tt.ts = memorytopo.NewServer(ctx, "zone1") testutil.AddShards(tt.ctx, t, tt.ts, &vtctldatapb.Shard{ Keyspace: "testkeyspace", Name: "-", @@ -9339,11 +9369,7 @@ func TestSetShardIsPrimaryServing(t *testing.T) { } for _, tt := range tests { - tt := tt - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - if tt.setup != nil { tt.setup(t, tt) } @@ -9369,6 +9395,9 @@ func TestSetShardIsPrimaryServing(t *testing.T) { func TestSetShardTabletControl(t *testing.T) { t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + type testcase struct { name string ctx context.Context @@ -9384,8 +9413,8 @@ func TestSetShardTabletControl(t *testing.T) { { name: "ok", setup: func(t *testing.T, tt *testcase) { - tt.ctx = context.Background() - tt.ts = memorytopo.NewServer("zone1", "zone2", "zone3") + tt.ctx = ctx + tt.ts = memorytopo.NewServer(ctx, "zone1", "zone2", "zone3") testutil.AddShards(tt.ctx, t, tt.ts, &vtctldatapb.Shard{ Keyspace: "testkeyspace", @@ -9433,8 +9462,8 @@ func TestSetShardTabletControl(t *testing.T) { { name: "remove tabletcontrols", setup: func(t *testing.T, tt *testcase) { - tt.ctx = context.Background() - tt.ts = memorytopo.NewServer("zone1", "zone2", "zone3") + tt.ctx = ctx + tt.ts = memorytopo.NewServer(ctx, "zone1", "zone2", "zone3") testutil.AddShards(tt.ctx, t, tt.ts, &vtctldatapb.Shard{ Keyspace: "testkeyspace", @@ -9468,8 +9497,8 @@ func TestSetShardTabletControl(t *testing.T) { { name: "disable queryservice", setup: func(t *testing.T, tt *testcase) { - tt.ctx = context.Background() - tt.ts = memorytopo.NewServer("zone1", "zone2", "zone3") + tt.ctx = ctx + tt.ts = memorytopo.NewServer(ctx, "zone1", "zone2", "zone3") testutil.AddShards(tt.ctx, t, tt.ts, &vtctldatapb.Shard{ Keyspace: "testkeyspace", @@ -9563,8 +9592,8 @@ func TestSetShardTabletControl(t *testing.T) { name: "keyspace lock error", setup: func(t *testing.T, tt *testcase) { var cancel func() - tt.ctx, cancel = context.WithTimeout(context.Background(), time.Millisecond*50) - tt.ts = memorytopo.NewServer("zone1") + tt.ctx, cancel = context.WithTimeout(ctx, time.Millisecond*50) + tt.ts = memorytopo.NewServer(ctx, "zone1") testutil.AddShards(tt.ctx, t, tt.ts, &vtctldatapb.Shard{ Keyspace: "testkeyspace", Name: "-", @@ -9590,10 +9619,7 @@ func TestSetShardTabletControl(t *testing.T) { }, } for _, tt := range tests { - tt := tt t.Run(tt.name, func(t *testing.T) { - t.Parallel() - if tt.setup != nil { tt.setup(t, tt) } @@ -9793,13 +9819,15 @@ func TestSetWritable(t *testing.T) { }, } - ctx := context.Background() for _, tt := range tests { tt := tt t.Run(tt.name, func(t *testing.T) { t.Parallel() - ts := memorytopo.NewServer(tt.cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, tt.cells...) defer ts.Close() testutil.AddTablets(ctx, t, ts, nil, tt.tablets...) @@ -9821,8 +9849,9 @@ func TestSetWritable(t *testing.T) { func TestShardReplicationAdd(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -9875,6 +9904,9 @@ func TestShardReplicationAdd(t *testing.T) { func TestShardReplicationPositions(t *testing.T) { t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + tests := []struct { name string ts *topo.Server @@ -9887,7 +9919,7 @@ func TestShardReplicationPositions(t *testing.T) { }{ { name: "success", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tablets: []*topodatapb.Tablet{ { Alias: &topodatapb.TabletAlias{ @@ -9966,7 +9998,7 @@ func TestShardReplicationPositions(t *testing.T) { }, { name: "timeouts are nonfatal", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tablets: []*topodatapb.Tablet{ { Alias: &topodatapb.TabletAlias{ @@ -10048,7 +10080,7 @@ func TestShardReplicationPositions(t *testing.T) { }, { name: "other rpc errors are fatal", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tablets: []*topodatapb.Tablet{ { Alias: &topodatapb.TabletAlias{ @@ -10098,7 +10130,7 @@ func TestShardReplicationPositions(t *testing.T) { }, { name: "nonexistent shard", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), req: &vtctldatapb.ShardReplicationPositionsRequest{ Keyspace: "testkeyspace", Shard: "-", @@ -10109,13 +10141,7 @@ func TestShardReplicationPositions(t *testing.T) { } for _, tt := range tests { - tt := tt - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - ctx := context.Background() - testutil.AddTablets(ctx, t, tt.ts, &testutil.AddTabletOptions{ AlsoSetShardPrimary: true, SkipShardCreation: false, @@ -10125,14 +10151,14 @@ func TestShardReplicationPositions(t *testing.T) { return NewVtctldServer(ts) }) + requestCtx := ctx if tt.ctxTimeout > 0 { - _ctx, cancel := context.WithTimeout(ctx, tt.ctxTimeout) - defer cancel() - - ctx = _ctx + var requestCancel func() + requestCtx, requestCancel = context.WithTimeout(ctx, tt.ctxTimeout) + defer requestCancel() } - resp, err := vtctld.ShardReplicationPositions(ctx, tt.req) + resp, err := vtctld.ShardReplicationPositions(requestCtx, tt.req) if tt.shouldErr { assert.Error(t, err) @@ -10148,8 +10174,10 @@ func TestShardReplicationPositions(t *testing.T) { func TestShardReplicationRemove(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "zone1") vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -10307,8 +10335,9 @@ func TestSourceShardAdd(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -10441,8 +10470,9 @@ func TestSourceShardDelete(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer { return NewVtctldServer(ts) }) @@ -10621,13 +10651,15 @@ func TestStartReplication(t *testing.T) { }, } - ctx := context.Background() for _, tt := range tests { tt := tt t.Run(tt.name, func(t *testing.T) { t.Parallel() - ts := memorytopo.NewServer(tt.cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, tt.cells...) defer ts.Close() testutil.AddTablets(ctx, t, ts, &testutil.AddTabletOptions{ @@ -10758,13 +10790,15 @@ func TestStopReplication(t *testing.T) { }, } - ctx := context.Background() for _, tt := range tests { tt := tt t.Run(tt.name, func(t *testing.T) { t.Parallel() - ts := memorytopo.NewServer(tt.cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, tt.cells...) defer ts.Close() testutil.AddTablets(ctx, t, ts, nil, tt.tablets...) @@ -11149,8 +11183,9 @@ func TestTabletExternallyReparented(t *testing.T) { cells := []string{"zone1", "zone2", "zone3"} - ctx := context.Background() - ts, topofactory := memorytopo.NewServerAndFactory(cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts, topofactory := memorytopo.NewServerAndFactory(ctx, cells...) tmc := testutil.TabletManagerClient{ TopoServer: ts, } @@ -11203,7 +11238,6 @@ func TestTabletExternallyReparented(t *testing.T) { func TestUpdateCellInfo(t *testing.T) { t.Parallel() - ctx := context.Background() tests := []struct { name string cells map[string]*topodatapb.CellInfo @@ -11327,7 +11361,10 @@ func TestUpdateCellInfo(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ts, factory := memorytopo.NewServerAndFactory() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts, factory := memorytopo.NewServerAndFactory(ctx) for name, cell := range tt.cells { err := ts.CreateCellInfo(ctx, name, cell) require.NoError(t, err, "failed to create cell %s: %+v for test", name, cell) @@ -11355,7 +11392,6 @@ func TestUpdateCellInfo(t *testing.T) { func TestUpdateCellsAlias(t *testing.T) { t.Parallel() - ctx := context.Background() tests := []struct { name string cells []string @@ -11466,7 +11502,10 @@ func TestUpdateCellsAlias(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ts := memorytopo.NewServer(tt.cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, tt.cells...) for name, cells := range tt.aliases { for _, cell := range cells { // We use UpdateCellInfoFields rather than CreateCellInfo @@ -11503,8 +11542,9 @@ func TestUpdateCellsAlias(t *testing.T) { func TestValidate(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer("zone1", "zone2", "zone3") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1", "zone2", "zone3") tablets := []*topodatapb.Tablet{ { Keyspace: "ks1", @@ -11618,8 +11658,9 @@ func TestValidate(t *testing.T) { } func TestValidateSchemaKeyspace(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("zone1", "zone2", "zone3") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1", "zone2", "zone3") tmc := testutil.TabletManagerClient{ GetSchemaResults: map[string]struct { Schema *tabletmanagerdatapb.SchemaDefinition @@ -11845,8 +11886,9 @@ func TestValidateSchemaKeyspace(t *testing.T) { } func TestValidateVersionKeyspace(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("zone1", "zone2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1", "zone2") tmc := testutil.TabletManagerClient{ GetSchemaResults: map[string]struct { Schema *tabletmanagerdatapb.SchemaDefinition @@ -11964,8 +12006,9 @@ func TestValidateVersionKeyspace(t *testing.T) { func TestValidateVersionShard(t *testing.T) { t.Parallel() - ctx := context.Background() - ts := memorytopo.NewServer("zone1", "zone2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1", "zone2") tmc := testutil.TabletManagerClient{ GetSchemaResults: map[string]struct { Schema *tabletmanagerdatapb.SchemaDefinition @@ -12084,11 +12127,12 @@ func TestValidateShard(t *testing.T) { shouldErr bool } - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() tests := []*testcase{ { name: "ok", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: nil, setup: func(t *testing.T, tt *testcase) { tablets := []*topodatapb.Tablet{ @@ -12125,7 +12169,7 @@ func TestValidateShard(t *testing.T) { }, { name: "no shard", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: nil, setup: func(t *testing.T, tt *testcase) { tablets := []*topodatapb.Tablet{ @@ -12166,7 +12210,7 @@ func TestValidateShard(t *testing.T) { }, { name: "no primary in shard", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: nil, setup: func(t *testing.T, tt *testcase) { tablets := []*topodatapb.Tablet{ @@ -12193,7 +12237,7 @@ func TestValidateShard(t *testing.T) { }, { name: "two primaries in shard", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: nil, setup: func(t *testing.T, tt *testcase) { tablets := []*topodatapb.Tablet{ @@ -12236,7 +12280,7 @@ func TestValidateShard(t *testing.T) { }, { name: "ping_tablets/ok", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ GetReplicasResults: map[string]struct { Replicas []string @@ -12305,7 +12349,7 @@ func TestValidateShard(t *testing.T) { }, { name: "ping_tablets/GetReplicas failed", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ GetReplicasResults: map[string]struct { Replicas []string @@ -12376,7 +12420,7 @@ func TestValidateShard(t *testing.T) { }, { name: "ping_tablets/no replicas", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ GetReplicasResults: map[string]struct { Replicas []string @@ -12447,7 +12491,7 @@ func TestValidateShard(t *testing.T) { }, { name: "ping_tablets/orphaned replica", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ GetReplicasResults: map[string]struct { Replicas []string @@ -12521,7 +12565,7 @@ func TestValidateShard(t *testing.T) { }, { name: "ping_tablets/Ping failed", - ts: memorytopo.NewServer("zone1"), + ts: memorytopo.NewServer(ctx, "zone1"), tmc: &testutil.TabletManagerClient{ GetReplicasResults: map[string]struct { Replicas []string @@ -12593,10 +12637,7 @@ func TestValidateShard(t *testing.T) { } for _, tt := range tests { - tt := tt t.Run(tt.name, func(t *testing.T) { - t.Parallel() - if tt.setup != nil { tt.setup(t, tt) } diff --git a/go/vt/vtctl/grpcvtctldserver/testutil/util.go b/go/vt/vtctl/grpcvtctldserver/testutil/util.go index 2b18d0bce68..155f99fca7d 100644 --- a/go/vt/vtctl/grpcvtctldserver/testutil/util.go +++ b/go/vt/vtctl/grpcvtctldserver/testutil/util.go @@ -149,6 +149,7 @@ type AddTabletOptions struct { // shard to serving. If that shard record already has a serving primary, then // AddTablet will fail the test. func AddTablet(ctx context.Context, t *testing.T, ts *topo.Server, tablet *topodatapb.Tablet, opts *AddTabletOptions) { + t.Helper() tablet = proto.Clone(tablet).(*topodatapb.Tablet) if opts == nil { opts = &AddTabletOptions{} @@ -200,6 +201,7 @@ func AddTablet(ctx context.Context, t *testing.T, ts *topo.Server, tablet *topod // AddTablets adds a list of tablets to the topology. See AddTablet for more // details. func AddTablets(ctx context.Context, t *testing.T, ts *topo.Server, opts *AddTabletOptions, tablets ...*topodatapb.Tablet) { + t.Helper() for _, tablet := range tablets { AddTablet(ctx, t, ts, tablet, opts) } diff --git a/go/vt/vtctl/reparentutil/emergency_reparenter_test.go b/go/vt/vtctl/reparentutil/emergency_reparenter_test.go index 5d03b60affa..a4cf95700d5 100644 --- a/go/vt/vtctl/reparentutil/emergency_reparenter_test.go +++ b/go/vt/vtctl/reparentutil/emergency_reparenter_test.go @@ -118,7 +118,7 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { emergencyReparentOps EmergencyReparentOptions tmc *testutil.TabletManagerClient // setup - ts *topo.Server + cells []string keyspace string shard string unlockTopo bool @@ -242,7 +242,7 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { }, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), + cells: []string{"zone1"}, shouldErr: false, }, { @@ -382,7 +382,7 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { }, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), + cells: []string{"zone1"}, shouldErr: false, }, { @@ -503,7 +503,7 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { }, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), + cells: []string{"zone1"}, shouldErr: false, }, { @@ -622,7 +622,7 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { }, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), + cells: []string{"zone1"}, shouldErr: false, }, { @@ -634,7 +634,7 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { shards: nil, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), + cells: []string{"zone1"}, shouldErr: true, errShouldContain: "node doesn't exist: keyspaces/testkeyspace/shards/-/Shard", }, @@ -693,7 +693,7 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { }, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), + cells: []string{"zone1"}, shouldErr: true, errShouldContain: "failed to stop replication and build status maps", }, @@ -752,7 +752,7 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { }, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), + cells: []string{"zone1"}, shouldErr: true, errShouldContain: "lost topology lock, aborting", }, @@ -826,7 +826,7 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { }, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), + cells: []string{"zone1"}, shouldErr: true, errShouldContain: "encountered tablet zone1-0000000102 with no relay log position", }, @@ -844,7 +844,7 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { shouldErr: true, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), + cells: []string{"zone1"}, errShouldContain: "no valid candidates for emergency reparent", }, { @@ -939,7 +939,7 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { shouldErr: true, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), + cells: []string{"zone1"}, errShouldContain: "could not apply all relay logs within the provided waitReplicasTimeout", }, { @@ -1028,7 +1028,7 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { }, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), + cells: []string{"zone1"}, shouldErr: true, errShouldContain: "primary elect zone1-0000000200 has errant GTIDs", }, @@ -1041,7 +1041,7 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { }}, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), + cells: []string{"zone1"}, tmc: &testutil.TabletManagerClient{ PopulateReparentJournalResults: map[string]error{ "zone1-0000000102": nil, @@ -1275,7 +1275,7 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { }, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), + cells: []string{"zone1"}, shouldErr: true, errShouldContain: "failed to be upgraded to primary", }, @@ -1390,7 +1390,7 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { }, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), + cells: []string{"zone1"}, shouldErr: true, errShouldContain: "no valid candidates for emergency reparent", }, @@ -1510,7 +1510,7 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { }, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), + cells: []string{"zone1"}, shouldErr: true, errShouldContain: "proposed primary zone1-0000000102 has a must not promotion rule", }, @@ -1637,7 +1637,7 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { }, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1", "zone2"), + cells: []string{"zone1", "zone2"}, shouldErr: true, errShouldContain: "no valid candidates for emergency reparent", }, @@ -1770,7 +1770,7 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { }, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1", "zone2"), + cells: []string{"zone1", "zone2"}, shouldErr: true, errShouldContain: "proposed primary zone1-0000000102 is is a different cell as the previous primary", }, @@ -1893,7 +1893,7 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { }, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), + cells: []string{"zone1"}, shouldErr: true, errShouldContain: "proposed primary zone1-0000000102 will not be able to make forward progress on being promoted", }, @@ -1904,7 +1904,9 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { t.Run(tt.name, func(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + logger := logutil.NewMemoryLogger() ev := &events.Reparent{} @@ -1915,12 +1917,14 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { tt.tablets[i] = tablet } - testutil.AddShards(ctx, t, tt.ts, tt.shards...) - testutil.AddTablets(ctx, t, tt.ts, nil, tt.tablets...) - reparenttestutil.SetKeyspaceDurability(ctx, t, tt.ts, tt.keyspace, tt.durability) + ts := memorytopo.NewServer(ctx, tt.cells...) + defer ts.Close() + testutil.AddShards(ctx, t, ts, tt.shards...) + testutil.AddTablets(ctx, t, ts, nil, tt.tablets...) + reparenttestutil.SetKeyspaceDurability(ctx, t, ts, tt.keyspace, tt.durability) if !tt.unlockTopo { - lctx, unlock, lerr := tt.ts.LockShard(ctx, tt.keyspace, tt.shard, "test lock") + lctx, unlock, lerr := ts.LockShard(ctx, tt.keyspace, tt.shard, "test lock") require.NoError(t, lerr, "could not lock %s/%s for testing", tt.keyspace, tt.shard) defer func() { @@ -1931,7 +1935,7 @@ func TestEmergencyReparenter_reparentShardLocked(t *testing.T) { ctx = lctx // make the reparentShardLocked call use the lock ctx } - erp := NewEmergencyReparenter(tt.ts, tt.tmc, logger) + erp := NewEmergencyReparenter(ts, tt.tmc, logger) err := erp.reparentShardLocked(ctx, ev, tt.keyspace, tt.shard, tt.emergencyReparentOps) if tt.shouldErr { @@ -1954,7 +1958,6 @@ func TestEmergencyReparenter_promoteNewPrimary(t *testing.T) { tmc *testutil.TabletManagerClient unlockTopo bool newPrimaryTabletAlias string - ts *topo.Server keyspace string shard string tablets []*topodatapb.Tablet @@ -2047,7 +2050,6 @@ func TestEmergencyReparenter_promoteNewPrimary(t *testing.T) { }, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), shouldErr: false, }, { @@ -2093,7 +2095,6 @@ func TestEmergencyReparenter_promoteNewPrimary(t *testing.T) { statusMap: map[string]*replicationdatapb.StopReplicationStatus{}, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), shouldErr: true, errShouldContain: "primary position error", }, @@ -2143,7 +2144,6 @@ func TestEmergencyReparenter_promoteNewPrimary(t *testing.T) { statusMap: map[string]*replicationdatapb.StopReplicationStatus{}, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), shouldErr: true, errShouldContain: "failed to PopulateReparentJournal on primary", }, @@ -2207,7 +2207,6 @@ func TestEmergencyReparenter_promoteNewPrimary(t *testing.T) { statusMap: map[string]*replicationdatapb.StopReplicationStatus{}, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), shouldErr: true, errShouldContain: " replica(s) failed", }, @@ -2276,7 +2275,6 @@ func TestEmergencyReparenter_promoteNewPrimary(t *testing.T) { shouldErr: true, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), errShouldContain: "context deadline exceeded", }, { @@ -2338,7 +2336,6 @@ func TestEmergencyReparenter_promoteNewPrimary(t *testing.T) { statusMap: map[string]*replicationdatapb.StopReplicationStatus{}, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), shouldErr: false, }, { @@ -2425,7 +2422,6 @@ func TestEmergencyReparenter_promoteNewPrimary(t *testing.T) { }, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), shouldErr: false, }, } @@ -2437,7 +2433,8 @@ func TestEmergencyReparenter_promoteNewPrimary(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() logger := logutil.NewMemoryLogger() ev := &events.Reparent{ShardInfo: topo.ShardInfo{ Shard: &topodatapb.Shard{ @@ -2451,7 +2448,10 @@ func TestEmergencyReparenter_promoteNewPrimary(t *testing.T) { ev.ShardInfo.PrimaryAlias = nil } - testutil.AddShards(ctx, t, tt.ts, &vtctldatapb.Shard{ + ts := memorytopo.NewServer(ctx, "zone1") + defer ts.Close() + + testutil.AddShards(ctx, t, ts, &vtctldatapb.Shard{ Keyspace: tt.keyspace, Name: tt.shard, }) @@ -2462,7 +2462,7 @@ func TestEmergencyReparenter_promoteNewPrimary(t *testing.T) { lerr error ) - ctx, unlock, lerr = tt.ts.LockShard(ctx, tt.keyspace, tt.shard, "test lock") + ctx, unlock, lerr = ts.LockShard(ctx, tt.keyspace, tt.shard, "test lock") require.NoError(t, lerr, "could not lock %s/%s for test", tt.keyspace, tt.shard) defer func() { @@ -2474,7 +2474,7 @@ func TestEmergencyReparenter_promoteNewPrimary(t *testing.T) { tt.emergencyReparentOps.durability = durability - erp := NewEmergencyReparenter(tt.ts, tt.tmc, logger) + erp := NewEmergencyReparenter(ts, tt.tmc, logger) err := erp.promoteNewPrimary(ctx, ev, tabletInfo.Tablet, tt.emergencyReparentOps, tt.tabletMap, tt.statusMap) if tt.shouldErr { assert.Error(t, err) @@ -2846,11 +2846,12 @@ func TestEmergencyReparenterCounters(t *testing.T) { } keyspace := "testkeyspace" shard := "-" - ts := memorytopo.NewServer("zone1") - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() logger := logutil.NewMemoryLogger() + ts := memorytopo.NewServer(ctx, "zone1") testutil.AddShards(ctx, t, ts, shards...) testutil.AddTablets(ctx, t, ts, &testutil.AddTabletOptions{ AlsoSetShardPrimary: true, @@ -3156,7 +3157,6 @@ func TestEmergencyReparenter_reparentReplicas(t *testing.T) { tmc *testutil.TabletManagerClient unlockTopo bool newPrimaryTabletAlias string - ts *topo.Server keyspace string shard string tablets []*topodatapb.Tablet @@ -3240,7 +3240,6 @@ func TestEmergencyReparenter_reparentReplicas(t *testing.T) { }, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), shouldErr: false, }, { @@ -3278,7 +3277,6 @@ func TestEmergencyReparenter_reparentReplicas(t *testing.T) { statusMap: map[string]*replicationdatapb.StopReplicationStatus{}, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), shouldErr: true, errShouldContain: "primary position error", }, @@ -3320,7 +3318,6 @@ func TestEmergencyReparenter_reparentReplicas(t *testing.T) { statusMap: map[string]*replicationdatapb.StopReplicationStatus{}, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), shouldErr: true, errShouldContain: "failed to PopulateReparentJournal on primary", }, @@ -3376,7 +3373,6 @@ func TestEmergencyReparenter_reparentReplicas(t *testing.T) { statusMap: map[string]*replicationdatapb.StopReplicationStatus{}, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), shouldErr: true, errShouldContain: " replica(s) failed", }, @@ -3436,7 +3432,6 @@ func TestEmergencyReparenter_reparentReplicas(t *testing.T) { shouldErr: true, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), errShouldContain: "context deadline exceeded", }, { @@ -3489,7 +3484,6 @@ func TestEmergencyReparenter_reparentReplicas(t *testing.T) { statusMap: map[string]*replicationdatapb.StopReplicationStatus{}, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), shouldErr: false, }, { name: "single replica failing to SetReplicationSource does not fail the promotion", @@ -3532,7 +3526,6 @@ func TestEmergencyReparenter_reparentReplicas(t *testing.T) { statusMap: map[string]*replicationdatapb.StopReplicationStatus{}, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), shouldErr: false, }, } @@ -3544,11 +3537,15 @@ func TestEmergencyReparenter_reparentReplicas(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ctx := context.Background() logger := logutil.NewMemoryLogger() ev := &events.Reparent{} - testutil.AddShards(ctx, t, tt.ts, &vtctldatapb.Shard{ + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") + defer ts.Close() + + testutil.AddShards(ctx, t, ts, &vtctldatapb.Shard{ Keyspace: tt.keyspace, Name: tt.shard, }) @@ -3559,7 +3556,7 @@ func TestEmergencyReparenter_reparentReplicas(t *testing.T) { lerr error ) - ctx, unlock, lerr = tt.ts.LockShard(ctx, tt.keyspace, tt.shard, "test lock") + ctx, unlock, lerr = ts.LockShard(ctx, tt.keyspace, tt.shard, "test lock") require.NoError(t, lerr, "could not lock %s/%s for test", tt.keyspace, tt.shard) defer func() { @@ -3571,7 +3568,7 @@ func TestEmergencyReparenter_reparentReplicas(t *testing.T) { tt.emergencyReparentOps.durability = durability - erp := NewEmergencyReparenter(tt.ts, tt.tmc, logger) + erp := NewEmergencyReparenter(ts, tt.tmc, logger) _, err := erp.reparentReplicas(ctx, ev, tabletInfo.Tablet, tt.tabletMap, tt.statusMap, tt.emergencyReparentOps, false /* waitForAllReplicas */, true /* populateReparentJournal */) if tt.shouldErr { assert.Error(t, err) @@ -3593,7 +3590,6 @@ func TestEmergencyReparenter_promoteIntermediateSource(t *testing.T) { tmc *testutil.TabletManagerClient unlockTopo bool newSourceTabletAlias string - ts *topo.Server keyspace string shard string tablets []*topodatapb.Tablet @@ -3679,7 +3675,6 @@ func TestEmergencyReparenter_promoteIntermediateSource(t *testing.T) { }, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), shouldErr: false, result: []*topodatapb.Tablet{ { @@ -3797,7 +3792,6 @@ func TestEmergencyReparenter_promoteIntermediateSource(t *testing.T) { statusMap: map[string]*replicationdatapb.StopReplicationStatus{}, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), shouldErr: true, errShouldContain: " replica(s) failed", }, @@ -3851,7 +3845,6 @@ func TestEmergencyReparenter_promoteIntermediateSource(t *testing.T) { statusMap: map[string]*replicationdatapb.StopReplicationStatus{}, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), shouldErr: false, validCandidateTablets: []*topodatapb.Tablet{ { @@ -3953,7 +3946,6 @@ func TestEmergencyReparenter_promoteIntermediateSource(t *testing.T) { }, keyspace: "testkeyspace", shard: "-", - ts: memorytopo.NewServer("zone1"), shouldErr: false, result: []*topodatapb.Tablet{ { @@ -3993,11 +3985,15 @@ func TestEmergencyReparenter_promoteIntermediateSource(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() logger := logutil.NewMemoryLogger() ev := &events.Reparent{} - testutil.AddShards(ctx, t, tt.ts, &vtctldatapb.Shard{ + ts := memorytopo.NewServer(ctx, "zone1") + defer ts.Close() + + testutil.AddShards(ctx, t, ts, &vtctldatapb.Shard{ Keyspace: tt.keyspace, Name: tt.shard, }) @@ -4008,7 +4004,7 @@ func TestEmergencyReparenter_promoteIntermediateSource(t *testing.T) { lerr error ) - ctx, unlock, lerr = tt.ts.LockShard(ctx, tt.keyspace, tt.shard, "test lock") + ctx, unlock, lerr = ts.LockShard(ctx, tt.keyspace, tt.shard, "test lock") require.NoError(t, lerr, "could not lock %s/%s for test", tt.keyspace, tt.shard) defer func() { @@ -4020,7 +4016,7 @@ func TestEmergencyReparenter_promoteIntermediateSource(t *testing.T) { tt.emergencyReparentOps.durability = durability - erp := NewEmergencyReparenter(tt.ts, tt.tmc, logger) + erp := NewEmergencyReparenter(ts, tt.tmc, logger) res, err := erp.promoteIntermediateSource(ctx, ev, tabletInfo.Tablet, tt.tabletMap, tt.statusMap, tt.validCandidateTablets, tt.emergencyReparentOps) if tt.shouldErr { assert.Error(t, err) @@ -4298,10 +4294,12 @@ func TestParentContextCancelled(t *testing.T) { } keyspace := "testkeyspace" shard := "-" - ts := memorytopo.NewServer("zone1") ctx, cancel := context.WithCancel(context.Background()) defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") + defer ts.Close() + logger := logutil.NewMemoryLogger() ev := &events.Reparent{} diff --git a/go/vt/vtctl/reparentutil/planned_reparenter_flaky_test.go b/go/vt/vtctl/reparentutil/planned_reparenter_flaky_test.go index c37fca9a37c..270bf97f87e 100644 --- a/go/vt/vtctl/reparentutil/planned_reparenter_flaky_test.go +++ b/go/vt/vtctl/reparentutil/planned_reparenter_flaky_test.go @@ -78,7 +78,6 @@ func TestPlannedReparenter_ReparentShard(t *testing.T) { tests := []struct { name string - ts *topo.Server tmc tmclient.TabletManagerClient tablets []*topodatapb.Tablet lockShardBeforeTest bool @@ -92,7 +91,6 @@ func TestPlannedReparenter_ReparentShard(t *testing.T) { }{ { name: "success", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ PrimaryPositionResults: map[string]struct { Position string @@ -178,7 +176,6 @@ func TestPlannedReparenter_ReparentShard(t *testing.T) { }, { name: "success - new primary not provided", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ ReplicationStatusResults: map[string]struct { Position *replicationdatapb.Status @@ -304,7 +301,6 @@ func TestPlannedReparenter_ReparentShard(t *testing.T) { }, { name: "already locked shard", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ PrimaryPositionResults: map[string]struct { Position string @@ -395,7 +391,6 @@ func TestPlannedReparenter_ReparentShard(t *testing.T) { // fail the preflight checks. Other functions are unit-tested // thoroughly to cover all the cases. name: "reparent fails", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ // This is only needed to verify reachability, so empty results are fine. PrimaryStatusResults: map[string]struct { @@ -446,7 +441,6 @@ func TestPlannedReparenter_ReparentShard(t *testing.T) { }, } - ctx := context.Background() logger := logutil.NewMemoryLogger() for _, tt := range tests { @@ -455,15 +449,18 @@ func TestPlannedReparenter_ReparentShard(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ctx := ctx + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") + defer ts.Close() - testutil.AddTablets(ctx, t, tt.ts, &testutil.AddTabletOptions{ + testutil.AddTablets(ctx, t, ts, &testutil.AddTabletOptions{ AlsoSetShardPrimary: true, SkipShardCreation: false, }, tt.tablets...) if tt.lockShardBeforeTest { - lctx, unlock, err := tt.ts.LockShard(ctx, tt.keyspace, tt.shard, "locking for test") + lctx, unlock, err := ts.LockShard(ctx, tt.keyspace, tt.shard, "locking for test") require.NoError(t, err, "could not lock %s/%s for test case", tt.keyspace, tt.shard) defer func() { @@ -474,7 +471,7 @@ func TestPlannedReparenter_ReparentShard(t *testing.T) { ctx = lctx } - pr := NewPlannedReparenter(tt.ts, tt.tmc, logger) + pr := NewPlannedReparenter(ts, tt.tmc, logger) ev, err := pr.ReparentShard(ctx, tt.keyspace, tt.shard, tt.opts) if tt.shouldErr { assert.Error(t, err) @@ -562,7 +559,6 @@ func TestPlannedReparenter_preflightChecks(t *testing.T) { tests := []struct { name string - ts *topo.Server tmc tmclient.TabletManagerClient tablets []*topodatapb.Tablet @@ -953,7 +949,6 @@ func TestPlannedReparenter_preflightChecks(t *testing.T) { }, } - ctx := context.Background() logger := logutil.NewMemoryLogger() for _, tt := range tests { @@ -962,6 +957,11 @@ func TestPlannedReparenter_preflightChecks(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") + defer ts.Close() + defer func() { if tt.expectedEvent != nil { AssertReparentEventsEqualWithMessage(t, tt.expectedEvent, tt.ev, "expected preflightChecks to mutate the passed-in event") @@ -972,7 +972,7 @@ func TestPlannedReparenter_preflightChecks(t *testing.T) { } }() - pr := NewPlannedReparenter(tt.ts, tt.tmc, logger) + pr := NewPlannedReparenter(ts, tt.tmc, logger) if tt.opts.durability == nil { durability, err := GetDurabilityPolicy("none") require.NoError(t, err) @@ -997,7 +997,6 @@ func TestPlannedReparenter_performGracefulPromotion(t *testing.T) { tests := []struct { name string - ts *topo.Server tmc tmclient.TabletManagerClient unlockTopo bool ctxTimeout time.Duration @@ -1019,7 +1018,6 @@ func TestPlannedReparenter_performGracefulPromotion(t *testing.T) { }{ { name: "successful promotion", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ DemotePrimaryResults: map[string]struct { Status *replicationdatapb.PrimaryStatus @@ -1076,7 +1074,6 @@ func TestPlannedReparenter_performGracefulPromotion(t *testing.T) { }, { name: "cannot get snapshot of current primary", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ PrimaryPositionResults: map[string]struct { Position string @@ -1110,7 +1107,6 @@ func TestPlannedReparenter_performGracefulPromotion(t *testing.T) { }, { name: "primary-elect fails to catch up to current primary snapshot position", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ PrimaryPositionResults: map[string]struct { Position string @@ -1147,7 +1143,6 @@ func TestPlannedReparenter_performGracefulPromotion(t *testing.T) { }, { name: "primary-elect times out catching up to current primary snapshot position", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ PrimaryPositionResults: map[string]struct { Position string @@ -1189,7 +1184,6 @@ func TestPlannedReparenter_performGracefulPromotion(t *testing.T) { }, { name: "lost topology lock", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ PrimaryPositionResults: map[string]struct { Position string @@ -1227,7 +1221,6 @@ func TestPlannedReparenter_performGracefulPromotion(t *testing.T) { }, { name: "failed to demote current primary", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ DemotePrimaryResults: map[string]struct { Status *replicationdatapb.PrimaryStatus @@ -1272,7 +1265,6 @@ func TestPlannedReparenter_performGracefulPromotion(t *testing.T) { }, { name: "primary-elect fails to catch up to current primary demotion position", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ DemotePrimaryResults: map[string]struct { Status *replicationdatapb.PrimaryStatus @@ -1329,7 +1321,6 @@ func TestPlannedReparenter_performGracefulPromotion(t *testing.T) { }, { name: "primary-elect times out catching up to current primary demotion position", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ DemotePrimaryResults: map[string]struct { Status *replicationdatapb.PrimaryStatus @@ -1391,7 +1382,6 @@ func TestPlannedReparenter_performGracefulPromotion(t *testing.T) { }, { name: "demotion succeeds but parent context times out", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ DemotePrimaryResults: map[string]struct { Status *replicationdatapb.PrimaryStatus @@ -1448,7 +1438,6 @@ func TestPlannedReparenter_performGracefulPromotion(t *testing.T) { }, { name: "rollback fails", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ DemotePrimaryResults: map[string]struct { Status *replicationdatapb.PrimaryStatus @@ -1511,7 +1500,6 @@ func TestPlannedReparenter_performGracefulPromotion(t *testing.T) { }, { name: "rollback succeeds", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ DemotePrimaryResults: map[string]struct { Status *replicationdatapb.PrimaryStatus @@ -1574,7 +1562,6 @@ func TestPlannedReparenter_performGracefulPromotion(t *testing.T) { }, } - ctx := context.Background() logger := logutil.NewMemoryLogger() for _, tt := range tests { @@ -1583,15 +1570,18 @@ func TestPlannedReparenter_performGracefulPromotion(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ctx := ctx + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") + defer ts.Close() - testutil.AddShards(ctx, t, tt.ts, &vtctldatapb.Shard{ + testutil.AddShards(ctx, t, ts, &vtctldatapb.Shard{ Keyspace: tt.keyspace, Name: tt.shard, }) if !tt.unlockTopo { - lctx, unlock, err := tt.ts.LockShard(ctx, tt.keyspace, tt.shard, "test lock") + lctx, unlock, err := ts.LockShard(ctx, tt.keyspace, tt.shard, "test lock") require.NoError(t, err, "could not lock %s/%s for testing", tt.keyspace, tt.shard) defer func() { @@ -1602,7 +1592,7 @@ func TestPlannedReparenter_performGracefulPromotion(t *testing.T) { ctx = lctx } - pr := NewPlannedReparenter(tt.ts, tt.tmc, logger) + pr := NewPlannedReparenter(ts, tt.tmc, logger) if tt.ctxTimeout > 0 { _ctx, cancel := context.WithTimeout(ctx, tt.ctxTimeout) @@ -1646,7 +1636,6 @@ func TestPlannedReparenter_performInitialPromotion(t *testing.T) { tests := []struct { name string - ts *topo.Server tmc tmclient.TabletManagerClient ctxTimeout time.Duration @@ -1660,7 +1649,6 @@ func TestPlannedReparenter_performInitialPromotion(t *testing.T) { }{ { name: "successful promotion", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ InitPrimaryResults: map[string]struct { Result string @@ -1686,7 +1674,6 @@ func TestPlannedReparenter_performInitialPromotion(t *testing.T) { }, { name: "primary-elect fails to promote", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ InitPrimaryResults: map[string]struct { Result string @@ -1710,7 +1697,6 @@ func TestPlannedReparenter_performInitialPromotion(t *testing.T) { }, { name: "promotion succeeds but parent context times out", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ InitPrimaryPostDelays: map[string]time.Duration{ "zone1-0000000200": time.Millisecond * 100, // 10x the parent context timeout @@ -1738,7 +1724,6 @@ func TestPlannedReparenter_performInitialPromotion(t *testing.T) { }, } - ctx := context.Background() logger := logutil.NewMemoryLogger() for _, tt := range tests { @@ -1747,14 +1732,17 @@ func TestPlannedReparenter_performInitialPromotion(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ctx := ctx + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") + defer ts.Close() - testutil.AddShards(ctx, t, tt.ts, &vtctldatapb.Shard{ + testutil.AddShards(ctx, t, ts, &vtctldatapb.Shard{ Keyspace: tt.keyspace, Name: tt.shard, }) - pr := NewPlannedReparenter(tt.ts, tt.tmc, logger) + pr := NewPlannedReparenter(ts, tt.tmc, logger) if tt.ctxTimeout > 0 { _ctx, cancel := context.WithTimeout(ctx, tt.ctxTimeout) @@ -1943,7 +1931,6 @@ func TestPlannedReparenter_performPotentialPromotion(t *testing.T) { tests := []struct { name string - ts *topo.Server tmc tmclient.TabletManagerClient timeout time.Duration unlockTopo bool @@ -1957,7 +1944,6 @@ func TestPlannedReparenter_performPotentialPromotion(t *testing.T) { }{ { name: "success", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ DemotePrimaryResults: map[string]struct { Status *replicationdatapb.PrimaryStatus @@ -2022,7 +2008,6 @@ func TestPlannedReparenter_performPotentialPromotion(t *testing.T) { }, { name: "failed to DemotePrimary on a tablet", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ DemotePrimaryResults: map[string]struct { Status *replicationdatapb.PrimaryStatus @@ -2057,7 +2042,6 @@ func TestPlannedReparenter_performPotentialPromotion(t *testing.T) { }, { name: "timed out during DemotePrimary on a tablet", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ DemotePrimaryDelays: map[string]time.Duration{ "zone1-0000000100": time.Millisecond * 50, @@ -2098,7 +2082,6 @@ func TestPlannedReparenter_performPotentialPromotion(t *testing.T) { }, { name: "failed to DecodePosition on a tablet's demote position", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ DemotePrimaryResults: map[string]struct { Status *replicationdatapb.PrimaryStatus @@ -2135,7 +2118,6 @@ func TestPlannedReparenter_performPotentialPromotion(t *testing.T) { }, { name: "primary-elect not in tablet map", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{}, unlockTopo: false, keyspace: "testkeyspace", @@ -2151,7 +2133,6 @@ func TestPlannedReparenter_performPotentialPromotion(t *testing.T) { }, { name: "primary-elect not most at most advanced position", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ DemotePrimaryResults: map[string]struct { Status *replicationdatapb.PrimaryStatus @@ -2216,7 +2197,6 @@ func TestPlannedReparenter_performPotentialPromotion(t *testing.T) { }, { name: "lost topology lock", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ DemotePrimaryResults: map[string]struct { Status *replicationdatapb.PrimaryStatus @@ -2281,7 +2261,6 @@ func TestPlannedReparenter_performPotentialPromotion(t *testing.T) { }, } - ctx := context.Background() logger := logutil.NewMemoryLogger() for _, tt := range tests { @@ -2290,16 +2269,20 @@ func TestPlannedReparenter_performPotentialPromotion(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ctx := ctx + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") + defer ts.Close() + pr := NewPlannedReparenter(nil, tt.tmc, logger) - testutil.AddShards(ctx, t, tt.ts, &vtctldatapb.Shard{ + testutil.AddShards(ctx, t, ts, &vtctldatapb.Shard{ Keyspace: tt.keyspace, Name: tt.shard, }) if !tt.unlockTopo { - lctx, unlock, err := tt.ts.LockShard(ctx, tt.keyspace, tt.shard, "test lock") + lctx, unlock, err := ts.LockShard(ctx, tt.keyspace, tt.shard, "test lock") require.NoError(t, err, "could not lock %s/%s for testing", tt.keyspace, tt.shard) defer func() { @@ -2336,7 +2319,6 @@ func TestPlannedReparenter_reparentShardLocked(t *testing.T) { tests := []struct { name string - ts *topo.Server tmc tmclient.TabletManagerClient tablets []*topodatapb.Tablet unlockTopo bool @@ -2351,7 +2333,6 @@ func TestPlannedReparenter_reparentShardLocked(t *testing.T) { }{ { name: "success: current primary cannot be determined", // "Case (1)" - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ DemotePrimaryResults: map[string]struct { Status *replicationdatapb.PrimaryStatus @@ -2443,7 +2424,6 @@ func TestPlannedReparenter_reparentShardLocked(t *testing.T) { }, { name: "success: current primary is desired primary", // "Case (2)" - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ PrimaryPositionResults: map[string]struct { Position string @@ -2530,7 +2510,6 @@ func TestPlannedReparenter_reparentShardLocked(t *testing.T) { }, { name: "success: graceful promotion", // "Case (3)" - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ DemotePrimaryResults: map[string]struct { Status *replicationdatapb.PrimaryStatus @@ -2626,7 +2605,6 @@ func TestPlannedReparenter_reparentShardLocked(t *testing.T) { }, { name: "shard not found", - ts: memorytopo.NewServer("zone1"), tmc: nil, tablets: nil, unlockTopo: true, @@ -2641,7 +2619,6 @@ func TestPlannedReparenter_reparentShardLocked(t *testing.T) { }, { name: "shard initialization", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ PopulateReparentJournalResults: map[string]error{ "zone1-0000000200": nil, @@ -2722,7 +2699,6 @@ func TestPlannedReparenter_reparentShardLocked(t *testing.T) { }, { name: "shard initialization with no new primary provided", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ PopulateReparentJournalResults: map[string]error{ "zone1-0000000200": nil, @@ -2810,7 +2786,6 @@ func TestPlannedReparenter_reparentShardLocked(t *testing.T) { }, { name: "preflight checks determine PRS is no-op", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ // This is only needed to verify reachability, so empty results are fine. PrimaryStatusResults: map[string]struct { @@ -2871,7 +2846,6 @@ func TestPlannedReparenter_reparentShardLocked(t *testing.T) { }, { name: "promotion step fails", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ SetReadWriteResults: map[string]error{ "zone1-0000000100": assert.AnError, @@ -2943,7 +2917,6 @@ func TestPlannedReparenter_reparentShardLocked(t *testing.T) { }, { name: "lost topology lock", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ // This is only needed to verify reachability, so empty results are fine. PrimaryStatusResults: map[string]struct { @@ -3026,7 +2999,6 @@ func TestPlannedReparenter_reparentShardLocked(t *testing.T) { }, { name: "failed to reparent tablets", - ts: memorytopo.NewServer("zone1"), tmc: &testutil.TabletManagerClient{ PrimaryPositionResults: map[string]struct { Position string @@ -3110,7 +3082,6 @@ func TestPlannedReparenter_reparentShardLocked(t *testing.T) { }, } - ctx := context.Background() logger := logutil.NewMemoryLogger() for _, tt := range tests { @@ -3119,16 +3090,19 @@ func TestPlannedReparenter_reparentShardLocked(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ctx := ctx + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") + defer ts.Close() - testutil.AddTablets(ctx, t, tt.ts, &testutil.AddTabletOptions{ + testutil.AddTablets(ctx, t, ts, &testutil.AddTabletOptions{ AlsoSetShardPrimary: true, ForceSetShardPrimary: true, // Some of our test cases count on having multiple primaries, so let the last one "win". SkipShardCreation: false, }, tt.tablets...) if !tt.unlockTopo { - lctx, unlock, err := tt.ts.LockShard(ctx, tt.keyspace, tt.shard, "locking for testing") + lctx, unlock, err := ts.LockShard(ctx, tt.keyspace, tt.shard, "locking for testing") require.NoError(t, err, "could not lock %s/%s for testing", tt.keyspace, tt.shard) defer func() { @@ -3145,7 +3119,7 @@ func TestPlannedReparenter_reparentShardLocked(t *testing.T) { }() } - pr := NewPlannedReparenter(tt.ts, tt.tmc, logger) + pr := NewPlannedReparenter(ts, tt.tmc, logger) err := pr.reparentShardLocked(ctx, tt.ev, tt.keyspace, tt.shard, tt.opts) if tt.shouldErr { @@ -3800,7 +3774,6 @@ func AssertReparentEventsEqual(t *testing.T, expected *events.Reparent, actual * func TestPlannedReparenter_verifyAllTabletsReachable(t *testing.T) { tests := []struct { name string - ts *topo.Server tmc tmclient.TabletManagerClient tabletMap map[string]*topo.TabletInfo remoteOpTime time.Duration @@ -3957,8 +3930,13 @@ func TestPlannedReparenter_verifyAllTabletsReachable(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") + defer ts.Close() + pr := &PlannedReparenter{ - ts: tt.ts, + ts: ts, tmc: tt.tmc, } if tt.remoteOpTime != 0 { diff --git a/go/vt/vtctl/schematools/reload_test.go b/go/vt/vtctl/schematools/reload_test.go index 6fbc7f152be..4f00e300d13 100644 --- a/go/vt/vtctl/schematools/reload_test.go +++ b/go/vt/vtctl/schematools/reload_test.go @@ -86,7 +86,6 @@ func (tmc *reloadSchemaTMC) ReloadSchema(ctx context.Context, tablet *topodatapb func TestReloadShard(t *testing.T) { t.Parallel() - ctx := context.Background() tests := []struct { name string @@ -330,7 +329,10 @@ func TestReloadShard(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - ts := memorytopo.NewServer(tt.cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, tt.cells...) + defer ts.Close() testutil.AddTablets(ctx, t, ts, &testutil.AddTabletOptions{ AlsoSetShardPrimary: true, }, tt.tablets...) diff --git a/go/vt/vtctl/vdiff2_test.go b/go/vt/vtctl/vdiff2_test.go index cb19621be0c..1348cd06448 100644 --- a/go/vt/vtctl/vdiff2_test.go +++ b/go/vt/vtctl/vdiff2_test.go @@ -35,7 +35,9 @@ var ( ) func TestVDiff2Unsharded(t *testing.T) { - env := newTestVDiffEnv(t, []string{"0"}, []string{"0"}, "", nil) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestVDiffEnv(t, ctx, []string{"0"}, []string{"0"}, "", nil) defer env.close() UUID := uuid.New().String() @@ -275,7 +277,9 @@ func TestVDiff2Unsharded(t *testing.T) { } func TestVDiff2Sharded(t *testing.T) { - env := newTestVDiffEnv(t, []string{"-40", "40-"}, []string{"-80", "80-"}, "", map[string]string{ + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestVDiffEnv(t, ctx, []string{"-40", "40-"}, []string{"-80", "80-"}, "", map[string]string{ "-80": "MySQL56/0e45e704-7cb9-11ed-a1eb-0242ac120002:1-890", "80-": "MySQL56/1497ddb0-7cb9-11ed-a1eb-0242ac120002:1-891", }) diff --git a/go/vt/vtctl/vdiff_env_test.go b/go/vt/vtctl/vdiff_env_test.go index 5fb854284ae..955d2673d20 100644 --- a/go/vt/vtctl/vdiff_env_test.go +++ b/go/vt/vtctl/vdiff_env_test.go @@ -69,11 +69,11 @@ type testVDiffEnv struct { //---------------------------------------------- // testVDiffEnv -func newTestVDiffEnv(t testing.TB, sourceShards, targetShards []string, query string, positions map[string]string) *testVDiffEnv { +func newTestVDiffEnv(t testing.TB, ctx context.Context, sourceShards, targetShards []string, query string, positions map[string]string) *testVDiffEnv { env := &testVDiffEnv{ workflow: "vdiffTest", tablets: make(map[int]*testVDiffTablet), - topoServ: memorytopo.NewServer("cell"), + topoServ: memorytopo.NewServer(ctx, "cell"), cell: "cell", tabletType: topodatapb.TabletType_REPLICA, tmc: newTestVDiffTMClient(), diff --git a/go/vt/vtctl/vtctl_env_test.go b/go/vt/vtctl/vtctl_env_test.go index 570088b9d13..e502fbdf86a 100644 --- a/go/vt/vtctl/vtctl_env_test.go +++ b/go/vt/vtctl/vtctl_env_test.go @@ -68,12 +68,12 @@ func init() { //---------------------------------------------- // testVTCtlEnv -func newTestVTCtlEnv() *testVTCtlEnv { +func newTestVTCtlEnv(ctx context.Context) *testVTCtlEnv { tabletconntest.SetProtocol("go.vt.vtctl.vtctl_env_test", "VTCtlTest") cellName := "cell1" env := &testVTCtlEnv{ tablets: make(map[int]*testVTCtlTablet), - topoServ: memorytopo.NewServer(cellName), + topoServ: memorytopo.NewServer(ctx, cellName), cell: cellName, tabletType: topodatapb.TabletType_REPLICA, tmc: newTestVTCtlTMClient(), diff --git a/go/vt/vtctl/vtctl_test.go b/go/vt/vtctl/vtctl_test.go index 94937e06ce2..eb6a5f5941f 100644 --- a/go/vt/vtctl/vtctl_test.go +++ b/go/vt/vtctl/vtctl_test.go @@ -48,8 +48,9 @@ var ( func TestApplyVSchema(t *testing.T) { shard := "0" ks := "ks" - ctx := context.Background() - env := newTestVTCtlEnv() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestVTCtlEnv(ctx) defer env.close() _ = env.addTablet(100, ks, shard, &topodatapb.KeyRange{}, topodatapb.TabletType_PRIMARY) @@ -152,8 +153,9 @@ func TestMoveTables(t *testing.T) { wf := "testwf" ksWf := fmt.Sprintf("%s.%s", targetKs, wf) minTableSize := 16384 // a single 16KiB InnoDB page - ctx := context.Background() - env := newTestVTCtlEnv() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestVTCtlEnv(ctx) defer env.close() source := env.addTablet(100, sourceKs, shard, &topodatapb.KeyRange{}, topodatapb.TabletType_PRIMARY) target := env.addTablet(200, targetKs, shard, &topodatapb.KeyRange{}, topodatapb.TabletType_PRIMARY) diff --git a/go/vt/vtctl/vtctlclienttest/client.go b/go/vt/vtctl/vtctlclienttest/client.go index df192997cbb..0bb2739c1b8 100644 --- a/go/vt/vtctl/vtctlclienttest/client.go +++ b/go/vt/vtctl/vtctlclienttest/client.go @@ -25,13 +25,12 @@ package vtctlclienttest // zookeeper) won't be drawn into production binaries as well. import ( + "context" "io" "strings" "testing" "time" - "context" - "vitess.io/vitess/go/vt/logutil" "vitess.io/vitess/go/vt/topo" "vitess.io/vitess/go/vt/topo/memorytopo" @@ -51,8 +50,8 @@ func init() { } // CreateTopoServer returns the test topo server properly configured -func CreateTopoServer(t *testing.T) *topo.Server { - return memorytopo.NewServer("cell1") +func CreateTopoServer(t *testing.T, ctx context.Context) *topo.Server { + return memorytopo.NewServer(ctx, "cell1") } // TestSuite runs the test suite on the given topo server and client diff --git a/go/vt/vtctl/workflow/materializer_env_test.go b/go/vt/vtctl/workflow/materializer_env_test.go index 430cfb71053..73646d7da80 100644 --- a/go/vt/vtctl/workflow/materializer_env_test.go +++ b/go/vt/vtctl/workflow/materializer_env_test.go @@ -67,14 +67,14 @@ func TestMain(m *testing.M) { os.Exit(m.Run()) } -func newTestMaterializerEnv(t *testing.T, ms *vtctldatapb.MaterializeSettings, sources, targets []string) *testMaterializerEnv { +func newTestMaterializerEnv(t *testing.T, ctx context.Context, ms *vtctldatapb.MaterializeSettings, sources, targets []string) *testMaterializerEnv { t.Helper() env := &testMaterializerEnv{ ms: ms, sources: sources, targets: targets, tablets: make(map[int]*topodatapb.Tablet), - topoServ: memorytopo.NewServer("cell"), + topoServ: memorytopo.NewServer(ctx, "cell"), cell: "cell", tmc: newTestMaterializerTMClient(), } diff --git a/go/vt/vtctl/workflow/materializer_test.go b/go/vt/vtctl/workflow/materializer_test.go index 3631b32fe6e..3c48223d93e 100644 --- a/go/vt/vtctl/workflow/materializer_test.go +++ b/go/vt/vtctl/workflow/materializer_test.go @@ -208,8 +208,10 @@ func TestStripConstraints(t *testing.T) { } func TestAddTablesToVSchema(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") + defer ts.Close() srcks := "source" ws := &Server{ ts: ts, @@ -417,6 +419,8 @@ func TestAddTablesToVSchema(t *testing.T) { } func TestMigrateVSchema(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() ms := &vtctldatapb.MaterializeSettings{ Workflow: "workflow", Cell: "cell", @@ -427,7 +431,7 @@ func TestMigrateVSchema(t *testing.T) { SourceExpression: "select * from t1", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() env.tmc.expectVRQuery(100, mzCheckJournal, &sqltypes.Result{}) @@ -438,7 +442,6 @@ func TestMigrateVSchema(t *testing.T) { env.tmc.expectVRQuery(200, mzGetWorkflowStatusQuery, getWorkflowStatusRes) env.tmc.expectVRQuery(200, mzGetLatestCopyState, &sqltypes.Result{}) - ctx := context.Background() _, err := env.ws.MoveTablesCreate(ctx, &vtctldatapb.MoveTablesCreateRequest{ Workflow: ms.Workflow, Cells: []string{ms.Cell}, @@ -468,7 +471,6 @@ func TestMigrateVSchema(t *testing.T) { // - TestPlayerDDL tests that the vplayer correctly implements the ddl behavior // - We have a manual e2e test for the full behavior: TestVReplicationDDLHandling func TestMoveTablesDDLFlag(t *testing.T) { - ctx := context.Background() ms := &vtctldatapb.MaterializeSettings{ Workflow: "workflow", SourceKeyspace: "sourceks", @@ -481,7 +483,9 @@ func TestMoveTablesDDLFlag(t *testing.T) { for onDDLAction := range binlogdatapb.OnDDLAction_value { t.Run(fmt.Sprintf("OnDDL Flag:%v", onDDLAction), func(t *testing.T) { - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() // This is the default and go does not marshal defaults // for prototext fields so we use the default insert stmt. diff --git a/go/vt/vtctld/api_test.go b/go/vt/vtctld/api_test.go index b00de6b0d09..6443d89a56b 100644 --- a/go/vt/vtctld/api_test.go +++ b/go/vt/vtctld/api_test.go @@ -42,9 +42,11 @@ func compactJSON(in []byte) string { } func TestAPI(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() cells := []string{"cell1", "cell2"} - ts := memorytopo.NewServer(cells...) + ts := memorytopo.NewServer(ctx, cells...) + defer ts.Close() actionRepo := NewActionRepository(ts) server := testutils.HTTPTestServer() defer server.Close() diff --git a/go/vt/vtctld/explorer_test.go b/go/vt/vtctld/explorer_test.go index 62eb7c01642..95ce6c6c3d9 100644 --- a/go/vt/vtctld/explorer_test.go +++ b/go/vt/vtctld/explorer_test.go @@ -17,12 +17,11 @@ limitations under the License. package vtctld import ( + "context" "path" "reflect" "testing" - "context" - "vitess.io/vitess/go/vt/topo" "vitess.io/vitess/go/vt/topo/memorytopo" @@ -35,7 +34,10 @@ func TestHandlePathRoot(t *testing.T) { cells := []string{"cell1", "cell2", "cell3"} want := []string{topo.GlobalCell, "cell1", "cell2", "cell3"} - ts := memorytopo.NewServer(cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cells...) + defer ts.Close() ex := newBackendExplorer(ts) result := ex.HandlePath(input, nil) if got := result.Children; !reflect.DeepEqual(got, want) { @@ -52,8 +54,10 @@ func TestHandlePathKeyspace(t *testing.T) { KeyspaceType: topodatapb.KeyspaceType_SNAPSHOT, } - ctx := context.Background() - ts := memorytopo.NewServer(cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cells...) + defer ts.Close() if err := ts.CreateKeyspace(ctx, "test_keyspace", keyspace); err != nil { t.Fatalf("CreateKeyspace error: %v", err) } @@ -101,8 +105,11 @@ func TestHandlePathShard(t *testing.T) { keyspace := &topodatapb.Keyspace{} want := "is_primary_serving:true" - ctx := context.Background() - ts := memorytopo.NewServer(cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cells...) + defer ts.Close() + if err := ts.CreateKeyspace(ctx, "test_keyspace", keyspace); err != nil { t.Fatalf("CreateKeyspace error: %v", err) } @@ -140,8 +147,11 @@ func TestHandlePathTablet(t *testing.T) { } want := "alias:{cell:\"cell1\" uid:123} hostname:\"example.com\" port_map:{key:\"vt\" value:4321}" - ctx := context.Background() - ts := memorytopo.NewServer(cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cells...) + defer ts.Close() + if err := ts.CreateTablet(ctx, tablet); err != nil { t.Fatalf("CreateTablet error: %v", err) } @@ -164,7 +174,11 @@ func TestHandleBadPath(t *testing.T) { cells := []string{"cell1", "cell2", "cell3"} want := "Invalid cell: node doesn't exist: cells/foo/CellInfo" - ts := memorytopo.NewServer(cells...) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cells...) + defer ts.Close() + ex := newBackendExplorer(ts) result := ex.HandlePath(input, nil) if got := result.Error; !reflect.DeepEqual(got, want) { diff --git a/go/vt/vtctld/tablet_data_test.go b/go/vt/vtctld/tablet_data_test.go index a9d173f6224..d40c6647ef3 100644 --- a/go/vt/vtctld/tablet_data_test.go +++ b/go/vt/vtctld/tablet_data_test.go @@ -108,7 +108,10 @@ func (s *streamHealthTabletServer) BroadcastHealth() { } func TestTabletData(t *testing.T) { - ts := memorytopo.NewServer("cell1", "cell2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2") + defer ts.Close() wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) if err := ts.CreateKeyspace(context.Background(), "ks", &topodatapb.Keyspace{}); err != nil { @@ -137,9 +140,9 @@ func TestTabletData(t *testing.T) { }() // Start streaming and wait for the first result. - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - result, err := thc.Get(ctx, tablet1.Tablet.Alias) - cancel() + requestCtx, requestCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer requestCancel() + result, err := thc.Get(requestCtx, tablet1.Tablet.Alias) close(stop) if err != nil { diff --git a/go/vt/vtexplain/vtexplain_vtgate.go b/go/vt/vtexplain/vtexplain_vtgate.go index 310fba4060d..6c9d60466ab 100644 --- a/go/vt/vtexplain/vtexplain_vtgate.go +++ b/go/vt/vtexplain/vtexplain_vtgate.go @@ -57,7 +57,7 @@ func (vte *VTExplain) initVtgateExecutor(ctx context.Context, vSchemaStr, ksShar resolver := vte.newFakeResolver(opts, vte.explainTopo, vtexplainCell) - err := vte.buildTopology(opts, vSchemaStr, ksShardMapStr, opts.NumShards) + err := vte.buildTopology(ctx, opts, vSchemaStr, ksShardMapStr, opts.NumShards) if err != nil { return err } @@ -94,7 +94,7 @@ func (vte *VTExplain) newFakeResolver(opts *Options, serv srvtopo.Server, cell s return vtgate.NewResolver(srvResolver, serv, cell, sc) } -func (vte *VTExplain) buildTopology(opts *Options, vschemaStr string, ksShardMapStr string, numShardsPerKeyspace int) error { +func (vte *VTExplain) buildTopology(ctx context.Context, opts *Options, vschemaStr string, ksShardMapStr string, numShardsPerKeyspace int) error { vte.explainTopo.Lock.Lock() defer vte.explainTopo.Lock.Unlock() @@ -142,7 +142,7 @@ func (vte *VTExplain) buildTopology(opts *Options, vschemaStr string, ksShardMap 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(opts, t) + return vte.newTablet(ctx, opts, t) }) vte.explainTopo.TabletConns[hostname] = tablet.(*explainTablet) vte.explainTopo.KeyspaceShards[ks][shard.Name] = shard diff --git a/go/vt/vtexplain/vtexplain_vttablet.go b/go/vt/vtexplain/vtexplain_vttablet.go index 98da79b75d7..f902eca8b07 100644 --- a/go/vt/vtexplain/vtexplain_vttablet.go +++ b/go/vt/vtexplain/vtexplain_vttablet.go @@ -102,7 +102,7 @@ type explainTablet struct { var _ queryservice.QueryService = (*explainTablet)(nil) -func (vte *VTExplain) newTablet(opts *Options, t *topodatapb.Tablet) *explainTablet { +func (vte *VTExplain) newTablet(ctx context.Context, opts *Options, t *topodatapb.Tablet) *explainTablet { db := fakesqldb.New(nil) sidecardb.AddSchemaInitQueries(db, true) @@ -117,7 +117,7 @@ func (vte *VTExplain) newTablet(opts *Options, t *topodatapb.Tablet) *explainTab config.EnableTableGC = false // XXX much of this is cloned from the tabletserver tests - tsv := tabletserver.NewTabletServer(topoproto.TabletAliasString(t.Alias), config, memorytopo.NewServer(""), t.Alias) + tsv := tabletserver.NewTabletServer(ctx, topoproto.TabletAliasString(t.Alias), config, memorytopo.NewServer(ctx, ""), t.Alias) tablet := explainTablet{db: db, tsv: tsv, vte: vte} db.Handler = &tablet diff --git a/go/vt/vtexplain/vtexplain_vttablet_test.go b/go/vt/vtexplain/vtexplain_vttablet_test.go index 9f6159ad2f9..bea8a91e3fc 100644 --- a/go/vt/vtexplain/vtexplain_vttablet_test.go +++ b/go/vt/vtexplain/vtexplain_vttablet_test.go @@ -17,6 +17,7 @@ limitations under the License. package vtexplain import ( + "context" "encoding/json" "testing" @@ -67,7 +68,9 @@ create table t2 ( NumShards: 2, } - vte, err := Init(testVSchema, testSchema, "", opts) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + vte, err := Init(ctx, testVSchema, testSchema, "", opts) require.NoError(t, err) defer vte.Stop() @@ -119,12 +122,15 @@ create table test_partitioned ( if err != nil { t.Fatalf("parseSchema: %v", err) } - vte := initTest(ModeMulti, defaultTestOpts(), &testopts{}, t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + vte := initTest(ctx, ModeMulti, defaultTestOpts(), &testopts{}, t) tabletEnv, _ := newTabletEnvironment(ddls, defaultTestOpts()) vte.setGlobalTabletEnv(tabletEnv) - tablet := vte.newTablet(defaultTestOpts(), &topodatapb.Tablet{ + tablet := vte.newTablet(ctx, defaultTestOpts(), &topodatapb.Tablet{ Keyspace: "test_keyspace", Shard: "-80", Alias: &topodatapb.TabletAlias{}, diff --git a/go/vt/vtgate/autocommit_test.go b/go/vt/vtgate/autocommit_test.go index 965bda4326a..3bd61a90cd3 100644 --- a/go/vt/vtgate/autocommit_test.go +++ b/go/vt/vtgate/autocommit_test.go @@ -21,6 +21,7 @@ import ( "testing" "github.com/stretchr/testify/require" + "vitess.io/vitess/go/test/utils" "vitess.io/vitess/go/sqltypes" diff --git a/go/vt/vtgate/executor.go b/go/vt/vtgate/executor.go index 013d3e68e16..8c2f41888d9 100644 --- a/go/vt/vtgate/executor.go +++ b/go/vt/vtgate/executor.go @@ -32,6 +32,7 @@ import ( "time" "github.com/spf13/pflag" + "vitess.io/vitess/go/streamlog" "vitess.io/vitess/go/acl" diff --git a/go/vt/vtgate/executor_scatter_stats_test.go b/go/vt/vtgate/executor_scatter_stats_test.go index 5c655767284..c3123260b7b 100644 --- a/go/vt/vtgate/executor_scatter_stats_test.go +++ b/go/vt/vtgate/executor_scatter_stats_test.go @@ -22,6 +22,7 @@ import ( "testing" "github.com/stretchr/testify/require" + "vitess.io/vitess/go/test/utils" vtgatepb "vitess.io/vitess/go/vt/proto/vtgate" diff --git a/go/vt/vtgate/executor_select_test.go b/go/vt/vtgate/executor_select_test.go index 6c095e6c739..0c64f1aabd0 100644 --- a/go/vt/vtgate/executor_select_test.go +++ b/go/vt/vtgate/executor_select_test.go @@ -36,6 +36,7 @@ import ( "github.com/google/go-cmp/cmp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "vitess.io/vitess/go/cache" "vitess.io/vitess/go/sqltypes" "vitess.io/vitess/go/test/utils" diff --git a/go/vt/vtgate/planbuilder/plan_test.go b/go/vt/vtgate/planbuilder/plan_test.go index b139d2573cc..06defefdc94 100644 --- a/go/vt/vtgate/planbuilder/plan_test.go +++ b/go/vt/vtgate/planbuilder/plan_test.go @@ -349,6 +349,8 @@ func TestBypassPlanningKeyrangeTargetFromFile(t *testing.T) { } func TestWithDefaultKeyspaceFromFile(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() // We are testing this separately so we can set a default keyspace vschema := &vschemawrapper.VSchemaWrapper{ V: loadSchema(t, "vschemas/schema.json", true), @@ -358,9 +360,9 @@ func TestWithDefaultKeyspaceFromFile(t *testing.T) { }, TabletType_: topodatapb.TabletType_PRIMARY, } - ts := memorytopo.NewServer("cell1") - ts.CreateKeyspace(context.Background(), "main", &topodatapb.Keyspace{}) - ts.CreateKeyspace(context.Background(), "user", &topodatapb.Keyspace{}) + ts := memorytopo.NewServer(ctx, "cell1") + ts.CreateKeyspace(ctx, "main", &topodatapb.Keyspace{}) + ts.CreateKeyspace(ctx, "user", &topodatapb.Keyspace{}) // Create a cache to use for lookups of the sidecar database identifier // in use by each keyspace. _, created := sidecardb.NewIdentifierCache(func(ctx context.Context, keyspace string) (string, error) { diff --git a/go/vt/vtgate/plugin_mysql_server.go b/go/vt/vtgate/plugin_mysql_server.go index 436ed23755f..6de6f12dba7 100644 --- a/go/vt/vtgate/plugin_mysql_server.go +++ b/go/vt/vtgate/plugin_mysql_server.go @@ -474,14 +474,14 @@ func initTLSConfig(ctx context.Context, mysqlListener *mysql.Listener, mysqlSslC // initiMySQLProtocol starts the mysql protocol. // It should be called only once in a process. -func initMySQLProtocol() { +func initMySQLProtocol(vtgate *VTGate) { // Flag is not set, just return. if mysqlServerPort < 0 && mysqlServerSocketPath == "" { return } // If no VTGate was created, just return. - if rpcVTGate == nil { + if vtgate == nil { return } @@ -506,7 +506,7 @@ func initMySQLProtocol() { // Create a Listener. var err error - vtgateHandle = newVtgateHandler(rpcVTGate) + vtgateHandle = newVtgateHandler(vtgate) if mysqlServerPort >= 0 { mysqlListener, err = mysql.NewListener( mysqlTCPVersion, @@ -680,8 +680,6 @@ func mysqlSocketPath() string { func init() { servenv.OnParseFor("vtgate", registerPluginFlags) servenv.OnParseFor("vtcombo", registerPluginFlags) - - servenv.OnRun(initMySQLProtocol) servenv.OnTermSync(shutdownMysqlProtocolAndDrain) servenv.OnClose(rollbackAtShutdown) } diff --git a/go/vt/vtgate/plugin_mysql_server_test.go b/go/vt/vtgate/plugin_mysql_server_test.go index b068662a7e2..3487d0253c8 100644 --- a/go/vt/vtgate/plugin_mysql_server_test.go +++ b/go/vt/vtgate/plugin_mysql_server_test.go @@ -29,6 +29,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "vitess.io/vitess/go/test/utils" "vitess.io/vitess/go/mysql/replication" diff --git a/go/vt/vtgate/queryz_test.go b/go/vt/vtgate/queryz_test.go index 1f46d077e3d..2de9c1bfe42 100644 --- a/go/vt/vtgate/queryz_test.go +++ b/go/vt/vtgate/queryz_test.go @@ -27,6 +27,7 @@ import ( "time" "github.com/stretchr/testify/require" + "vitess.io/vitess/go/test/utils" vtgatepb "vitess.io/vitess/go/vt/proto/vtgate" diff --git a/go/vt/vtgate/schema/tracker_test.go b/go/vt/vtgate/schema/tracker_test.go index 68b8bb206ce..4f514fec101 100644 --- a/go/vt/vtgate/schema/tracker_test.go +++ b/go/vt/vtgate/schema/tracker_test.go @@ -45,8 +45,10 @@ var ( func TestMain(m *testing.M) { exitCode := func() int { - ts := memorytopo.NewServer(cell) - ts.CreateKeyspace(context.Background(), keyspace, &topodatapb.Keyspace{}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, cell) + ts.CreateKeyspace(ctx, keyspace, &topodatapb.Keyspace{}) _, created := sidecardb.NewIdentifierCache(func(ctx context.Context, keyspace string) (string, error) { ki, err := ts.GetKeyspace(ctx, keyspace) if err != nil { diff --git a/go/vt/vtgate/tabletgateway_flaky_test.go b/go/vt/vtgate/tabletgateway_flaky_test.go index e98dc65a0ce..df96ed0ab1d 100644 --- a/go/vt/vtgate/tabletgateway_flaky_test.go +++ b/go/vt/vtgate/tabletgateway_flaky_test.go @@ -22,6 +22,7 @@ import ( "time" "github.com/stretchr/testify/require" + "vitess.io/vitess/go/test/utils" "vitess.io/vitess/go/mysql/collations" diff --git a/go/vt/vtgate/tabletgateway_test.go b/go/vt/vtgate/tabletgateway_test.go index 986c4cd3dfd..9a41029f7c8 100644 --- a/go/vt/vtgate/tabletgateway_test.go +++ b/go/vt/vtgate/tabletgateway_test.go @@ -24,6 +24,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "vitess.io/vitess/go/test/utils" "vitess.io/vitess/go/sqltypes" diff --git a/go/vt/vtgate/vtgate.go b/go/vt/vtgate/vtgate.go index 61c029f928a..b3ed103896e 100644 --- a/go/vt/vtgate/vtgate.go +++ b/go/vt/vtgate/vtgate.go @@ -368,6 +368,7 @@ func Init( if st != nil && enableSchemaChangeSignal { st.Start() } + initMySQLProtocol(rpcVTGate) }) servenv.OnTerm(func() { if st != nil && enableSchemaChangeSignal { diff --git a/go/vt/vtorc/logic/keyspace_shard_discovery_test.go b/go/vt/vtorc/logic/keyspace_shard_discovery_test.go index da83f050e3e..2911b3d29c2 100644 --- a/go/vt/vtorc/logic/keyspace_shard_discovery_test.go +++ b/go/vt/vtorc/logic/keyspace_shard_discovery_test.go @@ -66,18 +66,20 @@ func TestRefreshAllKeyspaces(t *testing.T) { db.ClearVTOrcDatabase() }() - ts = memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts = memorytopo.NewServer(ctx, "zone1") keyspaceNames := []string{"ks1", "ks2", "ks3", "ks4"} keyspaces := []*topodatapb.Keyspace{keyspaceDurabilityNone, keyspaceDurabilitySemiSync, keyspaceSnapshot, keyspaceDurabilityTest} // Create 4 keyspaces for i, keyspace := range keyspaces { - err := ts.CreateKeyspace(context.Background(), keyspaceNames[i], keyspace) + err := ts.CreateKeyspace(ctx, keyspaceNames[i], keyspace) require.NoError(t, err) for idx, shardName := range []string{"-80", "80-"} { - err = ts.CreateShard(context.Background(), keyspaceNames[i], shardName) + err = ts.CreateShard(ctx, keyspaceNames[i], shardName) require.NoError(t, err) - _, err = ts.UpdateShardFields(context.Background(), keyspaceNames[i], shardName, func(si *topo.ShardInfo) error { + _, err = ts.UpdateShardFields(ctx, keyspaceNames[i], shardName, func(si *topo.ShardInfo) error { si.PrimaryAlias = &topodatapb.TabletAlias{ Cell: fmt.Sprintf("zone_%v", keyspaceNames[i]), Uid: uint32(100 + idx), @@ -105,7 +107,7 @@ func TestRefreshAllKeyspaces(t *testing.T) { // Set clusters to watch to watch all keyspaces clustersToWatch = nil // Change the durability policy of ks1 - reparenttestutil.SetKeyspaceDurability(context.Background(), t, ts, "ks1", "semi_sync") + reparenttestutil.SetKeyspaceDurability(ctx, t, ts, "ks1", "semi_sync") RefreshAllKeyspacesAndShards() // Verify that all the keyspaces are correctly reloaded @@ -135,14 +137,12 @@ func TestRefreshKeyspace(t *testing.T) { name string keyspaceName string keyspace *topodatapb.Keyspace - ts *topo.Server keyspaceWanted *topodatapb.Keyspace err string }{ { name: "Success with keyspaceType and durability", keyspaceName: "ks1", - ts: memorytopo.NewServer("zone1"), keyspace: &topodatapb.Keyspace{ KeyspaceType: topodatapb.KeyspaceType_NORMAL, DurabilityPolicy: "semi_sync", @@ -152,7 +152,6 @@ func TestRefreshKeyspace(t *testing.T) { }, { name: "Success with keyspaceType and no durability", keyspaceName: "ks2", - ts: memorytopo.NewServer("zone1"), keyspace: &topodatapb.Keyspace{ KeyspaceType: topodatapb.KeyspaceType_NORMAL, }, @@ -161,7 +160,6 @@ func TestRefreshKeyspace(t *testing.T) { }, { name: "Success with snapshot keyspaceType", keyspaceName: "ks3", - ts: memorytopo.NewServer("zone1"), keyspace: &topodatapb.Keyspace{ KeyspaceType: topodatapb.KeyspaceType_SNAPSHOT, }, @@ -170,7 +168,6 @@ func TestRefreshKeyspace(t *testing.T) { }, { name: "Success with fields that are not stored", keyspaceName: "ks4", - ts: memorytopo.NewServer("zone1"), keyspace: &topodatapb.Keyspace{ KeyspaceType: topodatapb.KeyspaceType_NORMAL, DurabilityPolicy: "none", @@ -184,7 +181,6 @@ func TestRefreshKeyspace(t *testing.T) { }, { name: "No keyspace found", keyspaceName: "ks5", - ts: memorytopo.NewServer("zone1"), keyspace: nil, keyspaceWanted: nil, err: "node doesn't exist: keyspaces/ks5/Keyspace", @@ -196,7 +192,10 @@ func TestRefreshKeyspace(t *testing.T) { tt.keyspaceWanted = tt.keyspace } - ts = tt.ts + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts = memorytopo.NewServer(ctx, "zone1") if tt.keyspace != nil { err := ts.CreateKeyspace(context.Background(), tt.keyspaceName, tt.keyspace) require.NoError(t, err) @@ -243,7 +242,6 @@ func TestRefreshShard(t *testing.T) { keyspaceName string shardName string shard *topodatapb.Shard - ts *topo.Server primaryAliasWanted string err string }{ @@ -251,7 +249,6 @@ func TestRefreshShard(t *testing.T) { name: "Success with primaryAlias", keyspaceName: "ks1", shardName: "0", - ts: memorytopo.NewServer("zone1"), shard: &topodatapb.Shard{ PrimaryAlias: &topodatapb.TabletAlias{ Cell: "zone1", @@ -264,7 +261,6 @@ func TestRefreshShard(t *testing.T) { name: "Success with empty primaryAlias", keyspaceName: "ks1", shardName: "-80", - ts: memorytopo.NewServer("zone1"), shard: &topodatapb.Shard{}, primaryAliasWanted: "", err: "", @@ -272,13 +268,15 @@ func TestRefreshShard(t *testing.T) { name: "No shard found", keyspaceName: "ks2", shardName: "-", - ts: memorytopo.NewServer("zone1"), err: "node doesn't exist: keyspaces/ks2/shards/-/Shard", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - ts = tt.ts + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts = memorytopo.NewServer(ctx, "zone1") if tt.shard != nil { _, err := ts.GetOrCreateShard(context.Background(), tt.keyspaceName, tt.shardName) require.NoError(t, err) diff --git a/go/vt/vtorc/logic/tablet_discovery_test.go b/go/vt/vtorc/logic/tablet_discovery_test.go index a08c2a48da1..0e8ac72fabf 100644 --- a/go/vt/vtorc/logic/tablet_discovery_test.go +++ b/go/vt/vtorc/logic/tablet_discovery_test.go @@ -111,7 +111,10 @@ func TestRefreshTabletsInKeyspaceShard(t *testing.T) { }() // Create a memory topo-server and create the keyspace and shard records - ts = memorytopo.NewServer(cell1) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts = memorytopo.NewServer(ctx, cell1) _, err := ts.GetOrCreateShard(context.Background(), keyspace, shard) require.NoError(t, err) @@ -243,7 +246,9 @@ func TestShardPrimary(t *testing.T) { }() // Create a memory topo-server and create the keyspace and shard records - ts = memorytopo.NewServer(cell1) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts = memorytopo.NewServer(ctx, cell1) _, err := ts.GetOrCreateShard(context.Background(), keyspace, shard) require.NoError(t, err) diff --git a/go/vt/vtorc/logic/topology_recovery_test.go b/go/vt/vtorc/logic/topology_recovery_test.go index e588b284423..e5a69db1a09 100644 --- a/go/vt/vtorc/logic/topology_recovery_test.go +++ b/go/vt/vtorc/logic/topology_recovery_test.go @@ -125,7 +125,10 @@ func TestElectNewPrimaryPanic(t *testing.T) { analysisEntry := &inst.ReplicationAnalysis{ AnalyzedInstanceAlias: topoproto.TabletAliasString(tablet.Alias), } - ts = memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts = memorytopo.NewServer(ctx, "zone1") recoveryAttempted, _, err := electNewPrimary(context.Background(), analysisEntry) require.True(t, recoveryAttempted) require.Error(t, err) @@ -177,7 +180,10 @@ func TestDifferentAnalysescHaveDifferentCooldowns(t *testing.T) { AnalyzedInstanceAlias: topoproto.TabletAliasString(replica.Alias), Analysis: inst.DeadPrimary, } - ts = memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts = memorytopo.NewServer(ctx, "zone1") _, err = AttemptRecoveryRegistration(&replicaAnalysisEntry, false, true) require.Nil(t, err) diff --git a/go/vt/vttablet/customrule/topocustomrule/topocustomrule_test.go b/go/vt/vttablet/customrule/topocustomrule/topocustomrule_test.go index 2f7090e5450..e1d5c3fd2cd 100644 --- a/go/vt/vttablet/customrule/topocustomrule/topocustomrule_test.go +++ b/go/vt/vttablet/customrule/topocustomrule/topocustomrule_test.go @@ -78,11 +78,13 @@ func TestUpdate(t *testing.T) { cell := "cell1" filePath := "/keyspaces/ks1/configs/CustomRules" - ts := memorytopo.NewServer(cell) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, cell) qsc := tabletservermock.NewController() qsc.TS = ts sleepDuringTopoFailure = time.Millisecond - ctx := context.Background() cr, err := newTopoCustomRule(qsc, cell, filePath) if err != nil { diff --git a/go/vt/vttablet/endtoend/connkilling/main_test.go b/go/vt/vttablet/endtoend/connkilling/main_test.go index 98c4b9671d1..e7486c397eb 100644 --- a/go/vt/vttablet/endtoend/connkilling/main_test.go +++ b/go/vt/vttablet/endtoend/connkilling/main_test.go @@ -17,6 +17,7 @@ limitations under the License. package connkilling import ( + "context" "errors" "flag" "fmt" @@ -81,7 +82,9 @@ func TestMain(m *testing.M) { connAppDebugParams = cluster.MySQLAppDebugConnParams() config := tabletenv.NewDefaultConfig() _ = config.Oltp.TxTimeoutSeconds.Set("3s") - err := framework.StartCustomServer(connParams, connAppDebugParams, cluster.DbName(), config) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + err := framework.StartCustomServer(ctx, connParams, connAppDebugParams, cluster.DbName(), config) if err != nil { fmt.Fprintf(os.Stderr, "%v", err) return 1 diff --git a/go/vt/vttablet/endtoend/framework/server.go b/go/vt/vttablet/endtoend/framework/server.go index 169055faba3..5992ba494da 100644 --- a/go/vt/vttablet/endtoend/framework/server.go +++ b/go/vt/vttablet/endtoend/framework/server.go @@ -58,7 +58,7 @@ var ( // StartCustomServer starts the server and initializes // all the global variables. This function should only be called // once at the beginning of the test. -func StartCustomServer(connParams, connAppDebugParams mysql.ConnParams, dbName string, config *tabletenv.TabletConfig) error { +func StartCustomServer(ctx context.Context, connParams, connAppDebugParams mysql.ConnParams, dbName string, config *tabletenv.TabletConfig) error { // Setup a fake vtgate server. protocol := "resolveTest" vtgateconn.SetVTGateProtocol(protocol) @@ -75,9 +75,9 @@ func StartCustomServer(connParams, connAppDebugParams mysql.ConnParams, dbName s Shard: "0", TabletType: topodatapb.TabletType_PRIMARY, } - TopoServer = memorytopo.NewServer("") + TopoServer = memorytopo.NewServer(ctx, "") - Server = tabletserver.NewTabletServer("", config, TopoServer, &topodatapb.TabletAlias{}) + Server = tabletserver.NewTabletServer(ctx, "", config, TopoServer, &topodatapb.TabletAlias{}) Server.Register() err := Server.StartService(Target, dbcfgs, nil /* mysqld */) if err != nil { @@ -110,7 +110,7 @@ func StartCustomServer(connParams, connAppDebugParams mysql.ConnParams, dbName s // StartServer starts the server and initializes // all the global variables. This function should only be called // once at the beginning of the test. -func StartServer(connParams, connAppDebugParams mysql.ConnParams, dbName string) error { +func StartServer(ctx context.Context, connParams, connAppDebugParams mysql.ConnParams, dbName string) error { config := tabletenv.NewDefaultConfig() config.StrictTableACL = true config.TwoPCEnable = true @@ -126,7 +126,7 @@ func StartServer(connParams, connAppDebugParams mysql.ConnParams, dbName string) config.EnableViews = true gotBytes, _ := yaml2.Marshal(config) log.Infof("Config:\n%s", gotBytes) - return StartCustomServer(connParams, connAppDebugParams, dbName, config) + return StartCustomServer(ctx, connParams, connAppDebugParams, dbName, config) } // StopServer must be called once all the tests are done. diff --git a/go/vt/vttablet/endtoend/main_test.go b/go/vt/vttablet/endtoend/main_test.go index d4a6bd70899..e3f112e2b7f 100644 --- a/go/vt/vttablet/endtoend/main_test.go +++ b/go/vt/vttablet/endtoend/main_test.go @@ -17,6 +17,7 @@ limitations under the License. package endtoend import ( + "context" "errors" "flag" "fmt" @@ -84,7 +85,9 @@ func TestMain(m *testing.M) { connParams = cluster.MySQLConnParams() connAppDebugParams = cluster.MySQLAppDebugConnParams() - err = framework.StartServer(connParams, connAppDebugParams, cluster.DbName()) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + err = framework.StartServer(ctx, connParams, connAppDebugParams, cluster.DbName()) if err != nil { fmt.Fprintf(os.Stderr, "%v", err) return 1 diff --git a/go/vt/vttablet/endtoend/streamtimeout/main_test.go b/go/vt/vttablet/endtoend/streamtimeout/main_test.go index 406a70fe3a7..e676616b294 100644 --- a/go/vt/vttablet/endtoend/streamtimeout/main_test.go +++ b/go/vt/vttablet/endtoend/streamtimeout/main_test.go @@ -20,6 +20,7 @@ All tests in this package come with toxiproxy in front of the MySQL server package streamtimeout import ( + "context" "flag" "fmt" "os" @@ -86,7 +87,10 @@ func TestMain(m *testing.M) { _ = config.SchemaReloadIntervalSeconds.Set("2100ms") config.SchemaChangeReloadTimeout = 10 * time.Second config.SignalWhenSchemaChange = true - err = framework.StartCustomServer(connParams, connAppDebugParams, cluster.DbName(), config) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + err = framework.StartCustomServer(ctx, connParams, connAppDebugParams, cluster.DbName(), config) if err != nil { fmt.Fprintf(os.Stderr, "%v", err) return 1 diff --git a/go/vt/vttablet/tabletmanager/framework_test.go b/go/vt/vttablet/tabletmanager/framework_test.go index 642c699e284..0fbba42f31c 100644 --- a/go/vt/vttablet/tabletmanager/framework_test.go +++ b/go/vt/vttablet/tabletmanager/framework_test.go @@ -81,7 +81,7 @@ type testEnv struct { protoName string } -func newTestEnv(t *testing.T, sourceKeyspace string, sourceShards []string) *testEnv { +func newTestEnv(t *testing.T, ctx context.Context, sourceKeyspace string, sourceShards []string) *testEnv { tenv := &testEnv{ ctx: context.Background(), tmc: newFakeTMClient(), @@ -91,7 +91,7 @@ func newTestEnv(t *testing.T, sourceKeyspace string, sourceShards []string) *tes } tenv.mu.Lock() defer tenv.mu.Unlock() - tenv.ts = memorytopo.NewServer(tenv.cells...) + tenv.ts = memorytopo.NewServer(ctx, tenv.cells...) tenv.tmc.sourceKeyspace = sourceKeyspace tenv.tmc.sourceShards = sourceShards tenv.tmc.schema = defaultSchema diff --git a/go/vt/vttablet/tabletmanager/rpc_vreplication_test.go b/go/vt/vttablet/tabletmanager/rpc_vreplication_test.go index 942ef201dbe..bc4efedd85a 100644 --- a/go/vt/vttablet/tabletmanager/rpc_vreplication_test.go +++ b/go/vt/vttablet/tabletmanager/rpc_vreplication_test.go @@ -89,14 +89,15 @@ var ( // from a VtctldServer MoveTablesCreate request to ensure // that the VReplication stream(s) are created correctly. func TestCreateVReplicationWorkflow(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() sourceKs := "sourceks" sourceTabletUID := 200 targetKs := "targetks" targetTabletUID := 300 shard := "0" wf := "testwf" - tenv := newTestEnv(t, sourceKs, []string{shard}) + tenv := newTestEnv(t, ctx, sourceKs, []string{shard}) defer tenv.close() sourceTablet := tenv.addTablet(t, sourceTabletUID, sourceKs, shard) @@ -202,7 +203,8 @@ func TestCreateVReplicationWorkflow(t *testing.T) { // stream(s) are created correctly. Followed by ensuring that // SwitchTraffic and ReverseTraffic work as expected. func TestMoveTables(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() sourceKs := "sourceks" sourceTabletUID := 200 targetKs := "targetks" @@ -217,7 +219,7 @@ func TestMoveTables(t *testing.T) { topodatapb.TabletType_RDONLY, } - tenv := newTestEnv(t, sourceKs, []string{sourceShard}) + tenv := newTestEnv(t, ctx, sourceKs, []string{sourceShard}) defer tenv.close() sourceTablet := tenv.addTablet(t, sourceTabletUID, sourceKs, sourceShard) @@ -434,7 +436,8 @@ func TestMoveTables(t *testing.T) { } func TestUpdateVReplicationWorkflow(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() cells := []string{"zone1"} tabletTypes := []string{"replica"} workflow := "testwf" @@ -442,7 +445,7 @@ func TestUpdateVReplicationWorkflow(t *testing.T) { vreplID := 1 tabletUID := 100 - tenv := newTestEnv(t, keyspace, []string{shard}) + tenv := newTestEnv(t, ctx, keyspace, []string{shard}) defer tenv.close() tablet := tenv.addTablet(t, tabletUID, keyspace, shard) @@ -578,7 +581,8 @@ func TestUpdateVReplicationWorkflow(t *testing.T) { // fails -- specifically after the point where we have created // the workflow streams. func TestFailedMoveTablesCreateCleanup(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() sourceKs := "sourceks" sourceTabletUID := 200 shard := "0" @@ -589,7 +593,7 @@ func TestFailedMoveTablesCreateCleanup(t *testing.T) { invalidTimeZone := "NOPE" bls := fmt.Sprintf("keyspace:\"%s\" shard:\"%s\" filter:{rules:{match:\"%s\" filter:\"select * from %s\"}}", sourceKs, shard, table, table) - tenv := newTestEnv(t, sourceKs, []string{shard}) + tenv := newTestEnv(t, ctx, sourceKs, []string{shard}) defer tenv.close() ws := workflow.NewServer(tenv.ts, tenv.tmc) diff --git a/go/vt/vttablet/tabletmanager/shard_sync_test.go b/go/vt/vttablet/tabletmanager/shard_sync_test.go index 83a7cede2e1..24078efa977 100644 --- a/go/vt/vttablet/tabletmanager/shard_sync_test.go +++ b/go/vt/vttablet/tabletmanager/shard_sync_test.go @@ -44,8 +44,9 @@ const ( ) func TestShardSync(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") statsTabletTypeCount.ResetAll() tm := newTestTM(t, ts, 100, keyspace, shard) defer tm.Stop() diff --git a/go/vt/vttablet/tabletmanager/tm_init_test.go b/go/vt/vttablet/tabletmanager/tm_init_test.go index ccaedf0c272..d77251e20ab 100644 --- a/go/vt/vttablet/tabletmanager/tm_init_test.go +++ b/go/vt/vttablet/tabletmanager/tm_init_test.go @@ -164,10 +164,11 @@ func TestStartCreateKeyspaceShard(t *testing.T) { defer func(saved time.Duration) { rebuildKeyspaceRetryInterval = saved }(rebuildKeyspaceRetryInterval) rebuildKeyspaceRetryInterval = 10 * time.Millisecond - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() statsTabletTypeCount.ResetAll() cell := "cell1" - ts := memorytopo.NewServer(cell) + ts := memorytopo.NewServer(ctx, cell) tm := newTestTM(t, ts, 1, "ks", "0") defer tm.Stop() @@ -178,7 +179,7 @@ func TestStartCreateKeyspaceShard(t *testing.T) { _, err := ts.GetShard(ctx, "ks", "0") require.NoError(t, err) - ensureSrvKeyspace(t, ts, cell, "ks") + ensureSrvKeyspace(t, ctx, ts, cell, "ks") srvVSchema, err := ts.GetSrvVSchema(context.Background(), cell) require.NoError(t, err) @@ -192,7 +193,7 @@ func TestStartCreateKeyspaceShard(t *testing.T) { defer tm.Stop() _, err = ts.GetShard(ctx, "ks1", "0") require.NoError(t, err) - ensureSrvKeyspace(t, ts, cell, "ks1") + ensureSrvKeyspace(t, ctx, ts, cell, "ks1") srvVSchema, err = ts.GetSrvVSchema(context.Background(), cell) require.NoError(t, err) assert.Equal(t, wantVSchema, srvVSchema.Keyspaces["ks1"]) @@ -241,16 +242,17 @@ func TestStartCreateKeyspaceShard(t *testing.T) { tm2 := newTestTM(t, ts, 6, "ks4", "80-") defer tm2.Stop() // Now that we've started the tablet for the other shard, srvKeyspace will succeed. - ensureSrvKeyspace(t, ts, cell, "ks4") + ensureSrvKeyspace(t, ctx, ts, cell, "ks4") } func TestCheckPrimaryShip(t *testing.T) { defer func(saved time.Duration) { rebuildKeyspaceRetryInterval = saved }(rebuildKeyspaceRetryInterval) rebuildKeyspaceRetryInterval = 10 * time.Millisecond - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() cell := "cell1" - ts := memorytopo.NewServer(cell) + ts := memorytopo.NewServer(ctx, cell) alias := &topodatapb.TabletAlias{ Cell: "cell1", Uid: 1, @@ -260,7 +262,7 @@ func TestCheckPrimaryShip(t *testing.T) { // This will create the respective topology records. tm := newTestTM(t, ts, 1, "ks", "0") tablet := tm.Tablet() - ensureSrvKeyspace(t, ts, cell, "ks") + ensureSrvKeyspace(t, ctx, ts, cell, "ks") ti, err := ts.GetTablet(ctx, alias) require.NoError(t, err) assert.Equal(t, topodatapb.TabletType_REPLICA, ti.Type) @@ -388,9 +390,10 @@ func TestCheckPrimaryShip(t *testing.T) { } func TestStartCheckMysql(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() cell := "cell1" - ts := memorytopo.NewServer(cell) + ts := memorytopo.NewServer(ctx, cell) tablet := newTestTablet(t, 1, "ks", "0") cp := mysql.ConnParams{ Host: "foo", @@ -418,9 +421,10 @@ func TestStartFindMysqlPort(t *testing.T) { defer func(saved time.Duration) { mysqlPortRetryInterval = saved }(mysqlPortRetryInterval) mysqlPortRetryInterval = 50 * time.Millisecond - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() cell := "cell1" - ts := memorytopo.NewServer(cell) + ts := memorytopo.NewServer(ctx, cell) tablet := newTestTablet(t, 1, "ks", "0") fmd := newTestMysqlDaemon(t, -1) tm := &TabletManager{ @@ -460,9 +464,10 @@ func TestStartFindMysqlPort(t *testing.T) { // Init tablet fixes replication data when safe func TestStartFixesReplicationData(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() cell := "cell1" - ts := memorytopo.NewServer(cell, "cell2") + ts := memorytopo.NewServer(ctx, cell, "cell2") tm := newTestTM(t, ts, 1, "ks", "0") defer tm.Stop() tabletAlias := tm.tabletAlias @@ -493,8 +498,9 @@ func TestStartFixesReplicationData(t *testing.T) { // to be created due to a NodeExists error. During this particular error we were not doing // the sanity checks that the provided tablet was the same in the topo. func TestStartDoesNotUpdateReplicationDataForTabletInWrongShard(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("cell1", "cell2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2") tm := newTestTM(t, ts, 1, "ks", "0") tm.Stop() @@ -516,9 +522,10 @@ func TestCheckTabletTypeResets(t *testing.T) { defer func(saved time.Duration) { rebuildKeyspaceRetryInterval = saved }(rebuildKeyspaceRetryInterval) rebuildKeyspaceRetryInterval = 10 * time.Millisecond - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() cell := "cell1" - ts := memorytopo.NewServer(cell) + ts := memorytopo.NewServer(ctx, cell) alias := &topodatapb.TabletAlias{ Cell: "cell1", Uid: 1, @@ -528,7 +535,7 @@ func TestCheckTabletTypeResets(t *testing.T) { // This will create the respective topology records. tm := newTestTM(t, ts, 1, "ks", "0") tablet := tm.Tablet() - ensureSrvKeyspace(t, ts, cell, "ks") + ensureSrvKeyspace(t, ctx, ts, cell, "ks") ti, err := ts.GetTablet(ctx, alias) require.NoError(t, err) assert.Equal(t, topodatapb.TabletType_REPLICA, ti.Type) @@ -711,11 +718,11 @@ func newTestTablet(t *testing.T, uid int, keyspace, shard string) *topodatapb.Ta } } -func ensureSrvKeyspace(t *testing.T, ts *topo.Server, cell, keyspace string) { +func ensureSrvKeyspace(t *testing.T, ctx context.Context, ts *topo.Server, cell, keyspace string) { t.Helper() found := false for i := 0; i < 10; i++ { - _, err := ts.GetSrvKeyspace(context.Background(), cell, "ks") + _, err := ts.GetSrvKeyspace(ctx, cell, keyspace) if err == nil { found = true break diff --git a/go/vt/vttablet/tabletmanager/tm_state_test.go b/go/vt/vttablet/tabletmanager/tm_state_test.go index 537580d4853..8bd98edefff 100644 --- a/go/vt/vttablet/tabletmanager/tm_state_test.go +++ b/go/vt/vttablet/tabletmanager/tm_state_test.go @@ -42,7 +42,10 @@ import ( ) func TestStateOpenClose(t *testing.T) { - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "cell1") tm := newTestTM(t, ts, 1, "ks", "0") // Re-Open should be a no-op @@ -63,8 +66,9 @@ func TestStateOpenClose(t *testing.T) { } func TestStateRefreshFromTopo(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") tm := newTestTM(t, ts, 1, "ks", "0") defer tm.Stop() @@ -73,8 +77,9 @@ func TestStateRefreshFromTopo(t *testing.T) { } func TestStateResharding(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") tm := newTestTM(t, ts, 1, "ks", "0") defer tm.Stop() @@ -100,8 +105,9 @@ func TestStateResharding(t *testing.T) { } func TestStateDenyList(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") tm := newTestTM(t, ts, 1, "ks", "0") defer tm.Stop() @@ -131,8 +137,9 @@ func TestStateDenyList(t *testing.T) { } func TestStateTabletControls(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") tm := newTestTM(t, ts, 1, "ks", "0") defer tm.Stop() @@ -159,8 +166,9 @@ func TestStateTabletControls(t *testing.T) { } func TestStateIsShardServingisInSrvKeyspace(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") tm := newTestTM(t, ts, 1, "ks", "0") defer tm.Stop() @@ -330,8 +338,9 @@ func TestStateIsShardServingisInSrvKeyspace(t *testing.T) { } func TestStateNonServing(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") tm := newTestTM(t, ts, 1, "ks", "0") defer tm.Stop() @@ -346,8 +355,9 @@ func TestStateNonServing(t *testing.T) { } func TestStateChangeTabletType(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") statsTabletTypeCount.ResetAll() tm := newTestTM(t, ts, 2, "ks", "0") defer tm.Stop() @@ -387,8 +397,9 @@ func TestStateChangeTabletType(t *testing.T) { the new table type */ func TestStateChangeTabletTypeWithFailure(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") statsTabletTypeCount.ResetAll() // create TM with replica and put a hook to return error during SetServingType tm := newTestTM(t, ts, 2, "ks", "0") @@ -472,7 +483,7 @@ func TestChangeTypeErrorWhileWritingToTopo(t *testing.T) { factory := faketopo.NewFakeTopoFactory() // add cell1 to the factory. This returns a fake connection which we will use to set the get and update errors as we require. fakeConn := factory.AddCell("cell1") - ts := faketopo.NewFakeTopoServer(factory) + ts := faketopo.NewFakeTopoServer(context.TODO(), factory) statsTabletTypeCount.ResetAll() tm := newTestTM(t, ts, 2, "ks", "0") defer tm.Stop() @@ -519,8 +530,9 @@ func TestPublishStateNew(t *testing.T) { // we can't do using memorytopo, but we do test the retry // code path. - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") tm := newTestTM(t, ts, 42, "ks", "0") ttablet, err := tm.TopoServer.GetTablet(ctx, tm.tabletAlias) require.NoError(t, err) @@ -565,8 +577,9 @@ func TestPublishStateNew(t *testing.T) { } func TestPublishDeleted(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") tm := newTestTM(t, ts, 2, "ks", "0") defer tm.Stop() diff --git a/go/vt/vttablet/tabletmanager/vdiff/framework_test.go b/go/vt/vttablet/tabletmanager/vdiff/framework_test.go index 3e6236fa244..9f69e9ed86d 100644 --- a/go/vt/vttablet/tabletmanager/vdiff/framework_test.go +++ b/go/vt/vttablet/tabletmanager/vdiff/framework_test.go @@ -174,7 +174,9 @@ func init() { func TestMain(m *testing.M) { exitCode := func() int { var err error - tstenv, err = testenv.Init() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + tstenv, err = testenv.Init(ctx) if err != nil { fmt.Fprintf(os.Stderr, "%v", err) return 1 diff --git a/go/vt/vttablet/tabletmanager/vreplication/framework_test.go b/go/vt/vttablet/tabletmanager/vreplication/framework_test.go index 8484c71576a..cbbc564199c 100644 --- a/go/vt/vttablet/tabletmanager/vreplication/framework_test.go +++ b/go/vt/vttablet/tabletmanager/vreplication/framework_test.go @@ -127,9 +127,9 @@ func cleanup() { envMu.Unlock() } -func setup() (func(), int) { +func setup(ctx context.Context) (func(), int) { var err error - env, err = testenv.Init() + env, err = testenv.Init(ctx) if err != nil { fmt.Fprintf(os.Stderr, "%v", err) return nil, 1 @@ -145,12 +145,12 @@ func setup() (func(), int) { streamerEngine.InitDBConfig(env.KeyspaceName, env.ShardName) streamerEngine.Open() - if err := env.Mysqld.ExecuteSuperQuery(context.Background(), fmt.Sprintf("create database %s", vrepldb)); err != nil { + if err := env.Mysqld.ExecuteSuperQuery(ctx, fmt.Sprintf("create database %s", vrepldb)); err != nil { fmt.Fprintf(os.Stderr, "%v", err) return nil, 1 } - if err := env.Mysqld.ExecuteSuperQuery(context.Background(), "set @@global.innodb_lock_wait_timeout=1"); err != nil { + if err := env.Mysqld.ExecuteSuperQuery(ctx, "set @@global.innodb_lock_wait_timeout=1"); err != nil { fmt.Fprintf(os.Stderr, "%v", err) return nil, 1 } @@ -159,7 +159,7 @@ func setup() (func(), int) { "extb": env.Dbcfgs, } playerEngine = NewTestEngine(env.TopoServ, env.Cells[0], env.Mysqld, realDBClientFactory, realDBClientFactory, vrepldb, externalConfig) - playerEngine.Open(context.Background()) + playerEngine.Open(ctx) return cleanup, 0 } @@ -174,11 +174,13 @@ func TestMain(m *testing.M) { binlogplayer.SetProtocol("vreplication_test_framework", "test") _flag.ParseFlagsForTest() exitCode := func() int { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() if err := utils.SetBinlogRowImageMode("full", tempDir); err != nil { panic(err) } defer utils.SetBinlogRowImageMode("", tempDir) - cancel, ret := setup() + cancel, ret := setup(ctx) if ret > 0 { return ret } @@ -193,7 +195,7 @@ func TestMain(m *testing.M) { panic(err) } defer utils.SetBinlogRowImageMode("", tempDir) - cancel, ret = setup() + cancel, ret = setup(ctx) if ret > 0 { return ret } diff --git a/go/vt/vttablet/tabletserver/bench_test.go b/go/vt/vttablet/tabletserver/bench_test.go index 270e98019ec..fd2d86c2812 100644 --- a/go/vt/vttablet/tabletserver/bench_test.go +++ b/go/vt/vttablet/tabletserver/bench_test.go @@ -18,11 +18,10 @@ package tabletserver import ( "bytes" + "context" "fmt" "testing" - "context" - "vitess.io/vitess/go/sqltypes" querypb "vitess.io/vitess/go/vt/proto/query" @@ -55,7 +54,9 @@ func init() { } func BenchmarkExecuteVarBinary(b *testing.B) { - db, tsv := setupTabletServerTest(nil, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(b, ctx, "") defer db.Close() defer tsv.StopService() @@ -70,14 +71,16 @@ func BenchmarkExecuteVarBinary(b *testing.B) { target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY} db.SetAllowAll(true) for i := 0; i < b.N; i++ { - if _, err := tsv.Execute(context.Background(), &target, benchQuery, bv, 0, 0, nil); err != nil { + if _, err := tsv.Execute(ctx, &target, benchQuery, bv, 0, 0, nil); err != nil { panic(err) } } } func BenchmarkExecuteExpression(b *testing.B) { - db, tsv := setupTabletServerTest(nil, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(b, ctx, "") defer db.Close() defer tsv.StopService() @@ -95,7 +98,7 @@ func BenchmarkExecuteExpression(b *testing.B) { target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY} db.SetAllowAll(true) for i := 0; i < b.N; i++ { - if _, err := tsv.Execute(context.Background(), &target, benchQuery, bv, 0, 0, nil); err != nil { + if _, err := tsv.Execute(ctx, &target, benchQuery, bv, 0, 0, nil); err != nil { panic(err) } } diff --git a/go/vt/vttablet/tabletserver/health_streamer_test.go b/go/vt/vttablet/tabletserver/health_streamer_test.go index 98982b31c25..b2fbb2db1ea 100644 --- a/go/vt/vttablet/tabletserver/health_streamer_test.go +++ b/go/vt/vttablet/tabletserver/health_streamer_test.go @@ -210,6 +210,8 @@ func TestReloadSchema(t *testing.T) { for _, testcase := range testcases { t.Run(testcase.name, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() db := fakesqldb.New(t) defer db.Close() config := newConfig(db) @@ -325,6 +327,8 @@ func TestReloadSchema(t *testing.T) { // TestReloadView tests that the health streamer tracks view changes correctly func TestReloadView(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() db := fakesqldb.New(t) defer db.Close() config := newConfig(db) diff --git a/go/vt/vttablet/tabletserver/query_executor_test.go b/go/vt/vttablet/tabletserver/query_executor_test.go index 8a341791169..cf0e87a6e27 100644 --- a/go/vt/vttablet/tabletserver/query_executor_test.go +++ b/go/vt/vttablet/tabletserver/query_executor_test.go @@ -788,6 +788,8 @@ func TestQueryExecutorPlanNextval(t *testing.T) { } func TestQueryExecutorMessageStreamACL(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() aclName := fmt.Sprintf("simpleacl-test-%d", rand.Int63()) tableacl.Register(aclName, &simpleacl.Factory{}) tableacl.SetDefaultACL(aclName) @@ -817,7 +819,7 @@ func TestQueryExecutorMessageStreamACL(t *testing.T) { callerID := &querypb.VTGateCallerID{ Username: "u1", } - ctx := callerid.NewContext(context.Background(), nil, callerID) + ctx = callerid.NewContext(ctx, nil, callerID) qre := &QueryExecutor{ ctx: ctx, query: "stream from msg", @@ -1484,7 +1486,7 @@ func newTestTabletServer(ctx context.Context, flags executorFlags, db *fakesqldb } dbconfigs := newDBConfigs(db) config.DB = dbconfigs - tsv := NewTabletServer("TabletServerTest", config, memorytopo.NewServer(""), &topodatapb.TabletAlias{}) + tsv := NewTabletServer(ctx, "TabletServerTest", config, memorytopo.NewServer(ctx, ""), &topodatapb.TabletAlias{}) target := &querypb.Target{TabletType: topodatapb.TabletType_PRIMARY} err := tsv.StartService(target, dbconfigs, nil /* mysqld */) if config.TwoPCEnable { diff --git a/go/vt/vttablet/tabletserver/state_manager_test.go b/go/vt/vttablet/tabletserver/state_manager_test.go index 179411f212e..24c500ed087 100644 --- a/go/vt/vttablet/tabletserver/state_manager_test.go +++ b/go/vt/vttablet/tabletserver/state_manager_test.go @@ -517,6 +517,8 @@ func TestStateManagerCheckMySQL(t *testing.T) { } func TestStateManagerValidations(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() sm := newTestStateManager(t) target := &querypb.Target{TabletType: topodatapb.TabletType_PRIMARY} sm.target = proto.Clone(target).(*querypb.Target) @@ -579,6 +581,8 @@ func TestStateManagerValidations(t *testing.T) { } func TestStateManagerWaitForRequests(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() sm := newTestStateManager(t) defer sm.StopService() target := &querypb.Target{TabletType: topodatapb.TabletType_PRIMARY} diff --git a/go/vt/vttablet/tabletserver/stateful_connection_pool_test.go b/go/vt/vttablet/tabletserver/stateful_connection_pool_test.go index 79e70dc4ffe..b9ea4dfc185 100644 --- a/go/vt/vttablet/tabletserver/stateful_connection_pool_test.go +++ b/go/vt/vttablet/tabletserver/stateful_connection_pool_test.go @@ -29,9 +29,9 @@ import ( "vitess.io/vitess/go/vt/vttablet/tabletserver/tx" ) -var ctx = context.Background() - func TestActivePoolClientRowsFound(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() db := fakesqldb.New(t) defer db.Close() db.AddQuery("begin", &sqltypes.Result{}) @@ -58,6 +58,8 @@ func TestActivePoolClientRowsFound(t *testing.T) { } func TestActivePoolForAllTxProps(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() db := fakesqldb.New(t) defer db.Close() pool := newActivePool() @@ -84,6 +86,8 @@ func TestActivePoolForAllTxProps(t *testing.T) { } func TestStatefulPoolShutdownNonTx(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() db := fakesqldb.New(t) defer db.Close() pool := newActivePool() @@ -122,6 +126,8 @@ func TestStatefulPoolShutdownNonTx(t *testing.T) { } func TestStatefulPoolShutdownAll(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() db := fakesqldb.New(t) defer db.Close() pool := newActivePool() @@ -157,7 +163,7 @@ func TestActivePoolGetConnNonExistentTransaction(t *testing.T) { } func TestExecWithAbortedCtx(t *testing.T) { - ctx, cancel := context.WithCancel(ctx) + ctx, cancel := context.WithCancel(context.Background()) db := fakesqldb.New(t) defer db.Close() pool := newActivePool() @@ -170,6 +176,8 @@ func TestExecWithAbortedCtx(t *testing.T) { } func TestExecWithDbconnClosed(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() db := fakesqldb.New(t) defer db.Close() pool := newActivePool() @@ -183,6 +191,8 @@ func TestExecWithDbconnClosed(t *testing.T) { } func TestExecWithDbconnClosedHavingTx(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() db := fakesqldb.New(t) defer db.Close() pool := newActivePool() @@ -197,6 +207,8 @@ func TestExecWithDbconnClosedHavingTx(t *testing.T) { } func TestFailOnConnectionRegistering(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() db := fakesqldb.New(t) defer db.Close() pool := newActivePool() diff --git a/go/vt/vttablet/tabletserver/tabletserver.go b/go/vt/vttablet/tabletserver/tabletserver.go index deb13cbfdb9..d106dfd1596 100644 --- a/go/vt/vttablet/tabletserver/tabletserver.go +++ b/go/vt/vttablet/tabletserver/tabletserver.go @@ -140,8 +140,8 @@ var _ queryservice.QueryService = (*TabletServer)(nil) var RegisterFunctions []func(Controller) // NewServer creates a new TabletServer based on the command line flags. -func NewServer(name string, topoServer *topo.Server, alias *topodatapb.TabletAlias) *TabletServer { - return NewTabletServer(name, tabletenv.NewCurrentConfig(), topoServer, alias) +func NewServer(ctx context.Context, name string, topoServer *topo.Server, alias *topodatapb.TabletAlias) *TabletServer { + return NewTabletServer(ctx, name, tabletenv.NewCurrentConfig(), topoServer, alias) } var ( @@ -151,7 +151,7 @@ var ( // NewTabletServer creates an instance of TabletServer. Only the first // instance of TabletServer will expose its state variables. -func NewTabletServer(name string, config *tabletenv.TabletConfig, topoServer *topo.Server, alias *topodatapb.TabletAlias) *TabletServer { +func NewTabletServer(ctx context.Context, name string, config *tabletenv.TabletConfig, topoServer *topo.Server, alias *topodatapb.TabletAlias) *TabletServer { exporter := servenv.NewExporter(name, "Tablet") tsv := &TabletServer{ exporter: exporter, @@ -165,7 +165,7 @@ func NewTabletServer(name string, config *tabletenv.TabletConfig, topoServer *to } tsv.QueryTimeout.Store(config.Oltp.QueryTimeoutSeconds.Get().Nanoseconds()) - tsOnce.Do(func() { srvTopoServer = srvtopo.NewResilientServer(topoServer, "TabletSrvTopo") }) + tsOnce.Do(func() { srvTopoServer = srvtopo.NewResilientServer(ctx, topoServer, "TabletSrvTopo") }) tabletTypeFunc := func() topodatapb.TabletType { if tsv.sm == nil { diff --git a/go/vt/vttablet/tabletserver/tabletserver_test.go b/go/vt/vttablet/tabletserver/tabletserver_test.go index 0781904ae2a..f8fa06f5277 100644 --- a/go/vt/vttablet/tabletserver/tabletserver_test.go +++ b/go/vt/vttablet/tabletserver/tabletserver_test.go @@ -60,7 +60,9 @@ import ( ) func TestTabletServerHealthz(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") defer tsv.StopService() defer db.Close() @@ -87,7 +89,9 @@ func TestTabletServerHealthz(t *testing.T) { } func TestTabletServerHealthzNotConnected(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") defer tsv.StopService() defer db.Close() @@ -116,7 +120,9 @@ func TestTabletServerHealthzNotConnected(t *testing.T) { } func TestBeginOnReplica(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") defer tsv.StopService() defer db.Close() @@ -143,8 +149,10 @@ func TestBeginOnReplica(t *testing.T) { } func TestTabletServerPrimaryToReplica(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() // Reuse code from tx_executor_test. - _, tsv, db := newTestTxExecutor(t) + _, tsv, db := newTestTxExecutor(t, ctx) defer tsv.StopService() defer db.Close() target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY} @@ -182,8 +190,10 @@ func TestTabletServerPrimaryToReplica(t *testing.T) { } func TestTabletServerRedoLogIsKeptBetweenRestarts(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() // Reuse code from tx_executor_test. - _, tsv, db := newTestTxExecutor(t) + _, tsv, db := newTestTxExecutor(t, ctx) defer tsv.StopService() defer db.Close() tsv.SetServingType(topodatapb.TabletType_REPLICA, time.Time{}, true, "") @@ -267,7 +277,9 @@ func TestTabletServerRedoLogIsKeptBetweenRestarts(t *testing.T) { } func TestTabletServerCreateTransaction(t *testing.T) { - _, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, tsv, db := newTestTxExecutor(t, ctx) defer tsv.StopService() defer db.Close() target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY} @@ -282,43 +294,49 @@ func TestTabletServerCreateTransaction(t *testing.T) { } func TestTabletServerStartCommit(t *testing.T) { - _, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, tsv, db := newTestTxExecutor(t, ctx) defer tsv.StopService() defer db.Close() target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY} commitTransition := fmt.Sprintf("update _vt.dt_state set state = %d where dtid = 'aa' and state = %d", int(querypb.TransactionState_COMMIT), int(querypb.TransactionState_PREPARE)) db.AddQuery(commitTransition, &sqltypes.Result{RowsAffected: 1}) - txid := newTxForPrep(tsv) + txid := newTxForPrep(ctx, tsv) err := tsv.StartCommit(ctx, &target, txid, "aa") require.NoError(t, err) db.AddQuery(commitTransition, &sqltypes.Result{}) - txid = newTxForPrep(tsv) + txid = newTxForPrep(ctx, tsv) err = tsv.StartCommit(ctx, &target, txid, "aa") assert.EqualError(t, err, "could not transition to COMMIT: aa", "Prepare err") } func TestTabletserverSetRollback(t *testing.T) { - _, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, tsv, db := newTestTxExecutor(t, ctx) defer tsv.StopService() defer db.Close() target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY} rollbackTransition := fmt.Sprintf("update _vt.dt_state set state = %d where dtid = 'aa' and state = %d", int(querypb.TransactionState_ROLLBACK), int(querypb.TransactionState_PREPARE)) db.AddQuery(rollbackTransition, &sqltypes.Result{RowsAffected: 1}) - txid := newTxForPrep(tsv) + txid := newTxForPrep(ctx, tsv) err := tsv.SetRollback(ctx, &target, "aa", txid) require.NoError(t, err) db.AddQuery(rollbackTransition, &sqltypes.Result{}) - txid = newTxForPrep(tsv) + txid = newTxForPrep(ctx, tsv) err = tsv.SetRollback(ctx, &target, "aa", txid) assert.EqualError(t, err, "could not transition to ROLLBACK: aa", "Prepare err") } func TestTabletServerReadTransaction(t *testing.T) { - _, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, tsv, db := newTestTxExecutor(t, ctx) defer tsv.StopService() defer db.Close() target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY} @@ -411,7 +429,9 @@ func TestTabletServerReadTransaction(t *testing.T) { } func TestTabletServerConcludeTransaction(t *testing.T) { - _, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, tsv, db := newTestTxExecutor(t, ctx) defer tsv.StopService() defer db.Close() target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY} @@ -423,14 +443,16 @@ func TestTabletServerConcludeTransaction(t *testing.T) { } func TestTabletServerBeginFail(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() config := tabletenv.NewDefaultConfig() config.TxPool.Size = 1 - db, tsv := setupTabletServerTestCustom(t, config, "") + db, tsv := setupTabletServerTestCustom(t, ctx, config, "") defer tsv.StopService() defer db.Close() target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY} - ctx, cancel := context.WithTimeout(context.Background(), 1*time.Nanosecond) + ctx, cancel = context.WithTimeout(context.Background(), 1*time.Nanosecond) defer cancel() tsv.Begin(ctx, &target, nil) _, err := tsv.Begin(ctx, &target, nil) @@ -438,7 +460,9 @@ func TestTabletServerBeginFail(t *testing.T) { } func TestTabletServerCommitTransaction(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") defer tsv.StopService() defer db.Close() @@ -463,7 +487,9 @@ func TestTabletServerCommitTransaction(t *testing.T) { } func TestTabletServerCommiRollbacktFail(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") defer tsv.StopService() defer db.Close() @@ -476,7 +502,9 @@ func TestTabletServerCommiRollbacktFail(t *testing.T) { } func TestTabletServerRollback(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") defer tsv.StopService() defer db.Close() @@ -504,8 +532,10 @@ func TestTabletServerRollback(t *testing.T) { } func TestTabletServerPrepare(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() // Reuse code from tx_executor_test. - _, tsv, db := newTestTxExecutor(t) + _, tsv, db := newTestTxExecutor(t, ctx) defer tsv.StopService() defer db.Close() target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY} @@ -519,8 +549,10 @@ func TestTabletServerPrepare(t *testing.T) { } func TestTabletServerCommitPrepared(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() // Reuse code from tx_executor_test. - _, tsv, db := newTestTxExecutor(t) + _, tsv, db := newTestTxExecutor(t, ctx) defer tsv.StopService() defer db.Close() target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY} @@ -566,7 +598,9 @@ func TestSmallerTimeout(t *testing.T) { } func TestTabletServerReserveConnection(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") tsv.config.EnableSettingsPool = false defer tsv.StopService() defer db.Close() @@ -589,7 +623,9 @@ func TestTabletServerReserveConnection(t *testing.T) { } func TestTabletServerExecNonExistentConnection(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") defer tsv.StopService() defer db.Close() @@ -603,7 +639,9 @@ func TestTabletServerExecNonExistentConnection(t *testing.T) { } func TestTabletServerReleaseNonExistentConnection(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") defer tsv.StopService() defer db.Close() @@ -616,7 +654,9 @@ func TestTabletServerReleaseNonExistentConnection(t *testing.T) { } func TestMakeSureToCloseDbConnWhenBeginQueryFails(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") defer tsv.StopService() defer db.Close() @@ -630,7 +670,9 @@ func TestMakeSureToCloseDbConnWhenBeginQueryFails(t *testing.T) { } func TestTabletServerReserveAndBeginCommit(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") tsv.config.EnableSettingsPool = false defer tsv.StopService() defer db.Close() @@ -684,8 +726,10 @@ func TestTabletServerReserveAndBeginCommit(t *testing.T) { } func TestTabletServerRollbackPrepared(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() // Reuse code from tx_executor_test. - _, tsv, db := newTestTxExecutor(t) + _, tsv, db := newTestTxExecutor(t, ctx) defer tsv.StopService() defer db.Close() target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY} @@ -700,7 +744,9 @@ func TestTabletServerRollbackPrepared(t *testing.T) { } func TestTabletServerStreamExecute(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") defer tsv.StopService() defer db.Close() @@ -724,7 +770,9 @@ func TestTabletServerStreamExecute(t *testing.T) { } func TestTabletServerStreamExecuteComments(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") defer tsv.StopService() defer db.Close() @@ -762,7 +810,9 @@ func TestTabletServerStreamExecuteComments(t *testing.T) { } func TestTabletServerBeginStreamExecute(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") defer tsv.StopService() defer db.Close() @@ -790,7 +840,9 @@ func TestTabletServerBeginStreamExecute(t *testing.T) { } func TestTabletServerBeginStreamExecuteWithError(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") defer tsv.StopService() defer db.Close() @@ -817,6 +869,8 @@ func TestTabletServerBeginStreamExecuteWithError(t *testing.T) { } func TestSerializeTransactionsSameRow(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() // This test runs three transaction in parallel: // tx1 | tx2 | tx3 // However, tx1 and tx2 have the same WHERE clause (i.e. target the same row) @@ -829,7 +883,7 @@ func TestSerializeTransactionsSameRow(t *testing.T) { config.HotRowProtection.MaxConcurrency = 1 // Reduce the txpool to 2 because we should never consume more than two slots. config.TxPool.Size = 2 - db, tsv := setupTabletServerTestCustom(t, config, "") + db, tsv := setupTabletServerTestCustom(t, ctx, config, "") defer tsv.StopService() defer db.Close() @@ -930,11 +984,13 @@ func TestSerializeTransactionsSameRow(t *testing.T) { } func TestDMLQueryWithoutWhereClause(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() config := tabletenv.NewDefaultConfig() config.HotRowProtection.Mode = tabletenv.Enable config.HotRowProtection.MaxConcurrency = 1 config.TxPool.Size = 2 - db, tsv := setupTabletServerTestCustom(t, config, "") + db, tsv := setupTabletServerTestCustom(t, ctx, config, "") defer tsv.StopService() defer db.Close() @@ -950,6 +1006,8 @@ func TestDMLQueryWithoutWhereClause(t *testing.T) { } func TestSerializeTransactionsSameRow_ConcurrentTransactions(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() // This test runs three transaction in parallel: // tx1 | tx2 | tx3 // Out of these three, two can run in parallel because we increased the @@ -960,7 +1018,7 @@ func TestSerializeTransactionsSameRow_ConcurrentTransactions(t *testing.T) { config.HotRowProtection.MaxConcurrency = 2 // Reduce the txpool to 2 because we should never consume more than two slots. config.TxPool.Size = 2 - db, tsv := setupTabletServerTestCustom(t, config, "") + db, tsv := setupTabletServerTestCustom(t, ctx, config, "") defer tsv.StopService() defer db.Close() @@ -1085,6 +1143,8 @@ func waitForTxSerializationPendingQueries(tsv *TabletServer, key string, i int) } func TestSerializeTransactionsSameRow_TooManyPendingRequests(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() // This test is similar to TestSerializeTransactionsSameRow, but tests only // that there must not be too many pending BeginExecute() requests which are // serialized. @@ -1094,7 +1154,7 @@ func TestSerializeTransactionsSameRow_TooManyPendingRequests(t *testing.T) { config.HotRowProtection.Mode = tabletenv.Enable config.HotRowProtection.MaxQueueSize = 1 config.HotRowProtection.MaxConcurrency = 1 - db, tsv := setupTabletServerTestCustom(t, config, "") + db, tsv := setupTabletServerTestCustom(t, ctx, config, "") defer tsv.StopService() defer db.Close() @@ -1166,6 +1226,8 @@ func TestSerializeTransactionsSameRow_TooManyPendingRequests(t *testing.T) { } func TestSerializeTransactionsSameRow_RequestCanceled(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() // This test is similar to TestSerializeTransactionsSameRow, but tests only // that a queued request unblocks itself when its context is done. // @@ -1175,7 +1237,7 @@ func TestSerializeTransactionsSameRow_RequestCanceled(t *testing.T) { config := tabletenv.NewDefaultConfig() config.HotRowProtection.Mode = tabletenv.Enable config.HotRowProtection.MaxConcurrency = 1 - db, tsv := setupTabletServerTestCustom(t, config, "") + db, tsv := setupTabletServerTestCustom(t, ctx, config, "") defer tsv.StopService() defer db.Close() @@ -1283,7 +1345,9 @@ func TestSerializeTransactionsSameRow_RequestCanceled(t *testing.T) { } func TestMessageStream(t *testing.T) { - _, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY} @@ -1309,7 +1373,9 @@ func TestMessageStream(t *testing.T) { } func TestCheckMySQLGauge(t *testing.T) { - _, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() @@ -1337,7 +1403,9 @@ func TestCheckMySQLGauge(t *testing.T) { } func TestMessageAck(t *testing.T) { - _, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY} @@ -1366,7 +1434,9 @@ func TestMessageAck(t *testing.T) { } func TestRescheduleMessages(t *testing.T) { - _, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY} @@ -1390,7 +1460,9 @@ func TestRescheduleMessages(t *testing.T) { } func TestPurgeMessages(t *testing.T) { - _, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY} @@ -1415,9 +1487,11 @@ func TestPurgeMessages(t *testing.T) { } func TestHandleExecUnknownError(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() logStats := tabletenv.NewLogStats(ctx, "TestHandleExecError") config := tabletenv.NewDefaultConfig() - tsv := NewTabletServer("TabletServerTest", config, memorytopo.NewServer(""), &topodatapb.TabletAlias{}) + tsv := NewTabletServer(ctx, "TabletServerTest", config, memorytopo.NewServer(ctx, ""), &topodatapb.TabletAlias{}) defer tsv.handlePanicAndSendLogStats("select * from test_table", nil, logStats) panic("unknown exec error") } @@ -1427,10 +1501,12 @@ func TestHandleExecUnknownError(t *testing.T) { // truncate the error text in logs, but will not truncate the error text in the // error value. func TestHandlePanicAndSendLogStatsMessageTruncation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() tl := newTestLogger() defer tl.Close() logStats := tabletenv.NewLogStats(ctx, "TestHandlePanicAndSendLogStatsMessageTruncation") - db, tsv := setupTabletServerTest(t, "") + db, tsv := setupTabletServerTest(t, ctx, "") defer tsv.StopService() defer db.Close() @@ -1577,8 +1653,10 @@ func (tl *testLogger) getLogs() []string { } func TestHandleExecTabletError(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() config := tabletenv.NewDefaultConfig() - tsv := NewTabletServer("TabletServerTest", config, memorytopo.NewServer(""), &topodatapb.TabletAlias{}) + tsv := NewTabletServer(ctx, "TabletServerTest", config, memorytopo.NewServer(ctx, ""), &topodatapb.TabletAlias{}) tl := newTestLogger() defer tl.Close() err := tsv.convertAndLogError( @@ -1598,10 +1676,12 @@ func TestHandleExecTabletError(t *testing.T) { } func TestTerseErrors(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() config := tabletenv.NewDefaultConfig() config.TerseErrors = true config.SanitizeLogMessages = false - tsv := NewTabletServer("TabletServerTest", config, memorytopo.NewServer(""), &topodatapb.TabletAlias{}) + tsv := NewTabletServer(ctx, "TabletServerTest", config, memorytopo.NewServer(ctx, ""), &topodatapb.TabletAlias{}) tl := newTestLogger() defer tl.Close() @@ -1630,10 +1710,12 @@ func TestTerseErrors(t *testing.T) { } func TestSanitizeLogMessages(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() config := tabletenv.NewDefaultConfig() config.TerseErrors = false config.SanitizeLogMessages = true - tsv := NewTabletServer("TabletServerTest", config, memorytopo.NewServer(""), &topodatapb.TabletAlias{}) + tsv := NewTabletServer(ctx, "TabletServerTest", config, memorytopo.NewServer(ctx, ""), &topodatapb.TabletAlias{}) tl := newTestLogger() defer tl.Close() @@ -1662,9 +1744,11 @@ func TestSanitizeLogMessages(t *testing.T) { } func TestTerseErrorsNonSQLError(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() config := tabletenv.NewDefaultConfig() config.TerseErrors = true - tsv := NewTabletServer("TabletServerTest", config, memorytopo.NewServer(""), &topodatapb.TabletAlias{}) + tsv := NewTabletServer(ctx, "TabletServerTest", config, memorytopo.NewServer(ctx, ""), &topodatapb.TabletAlias{}) tl := newTestLogger() defer tl.Close() err := tsv.convertAndLogError( @@ -1684,10 +1768,12 @@ func TestTerseErrorsNonSQLError(t *testing.T) { } func TestSanitizeLogMessagesNonSQLError(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() config := tabletenv.NewDefaultConfig() config.TerseErrors = false config.SanitizeLogMessages = true - tsv := NewTabletServer("TabletServerTest", config, memorytopo.NewServer(""), &topodatapb.TabletAlias{}) + tsv := NewTabletServer(ctx, "TabletServerTest", config, memorytopo.NewServer(ctx, ""), &topodatapb.TabletAlias{}) tl := newTestLogger() defer tl.Close() err := tsv.convertAndLogError( @@ -1707,10 +1793,12 @@ func TestSanitizeLogMessagesNonSQLError(t *testing.T) { } func TestSanitizeMessagesBindVars(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() config := tabletenv.NewDefaultConfig() config.TerseErrors = true config.SanitizeLogMessages = true - tsv := NewTabletServer("TabletServerTest", config, memorytopo.NewServer(""), &topodatapb.TabletAlias{}) + tsv := NewTabletServer(ctx, "TabletServerTest", config, memorytopo.NewServer(ctx, ""), &topodatapb.TabletAlias{}) tl := newTestLogger() defer tl.Close() @@ -1736,10 +1824,12 @@ func TestSanitizeMessagesBindVars(t *testing.T) { } func TestSanitizeMessagesNoBindVars(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() config := tabletenv.NewDefaultConfig() config.TerseErrors = true config.SanitizeLogMessages = true - tsv := NewTabletServer("TabletServerTest", config, memorytopo.NewServer(""), &topodatapb.TabletAlias{}) + tsv := NewTabletServer(ctx, "TabletServerTest", config, memorytopo.NewServer(ctx, ""), &topodatapb.TabletAlias{}) tl := newTestLogger() defer tl.Close() err := tsv.convertAndLogError(ctx, "", nil, vterrors.Errorf(vtrpcpb.Code_DEADLINE_EXCEEDED, "sensitive message"), nil) @@ -1753,9 +1843,11 @@ func TestSanitizeMessagesNoBindVars(t *testing.T) { } func TestTruncateErrorLen(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() config := tabletenv.NewDefaultConfig() config.TruncateErrorLen = 32 - tsv := NewTabletServer("TabletServerTest", config, memorytopo.NewServer(""), &topodatapb.TabletAlias{}) + tsv := NewTabletServer(ctx, "TabletServerTest", config, memorytopo.NewServer(ctx, ""), &topodatapb.TabletAlias{}) tl := newTestLogger() defer tl.Close() err := tsv.convertAndLogError( @@ -1775,11 +1867,13 @@ func TestTruncateErrorLen(t *testing.T) { } func TestTruncateMessages(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() config := tabletenv.NewDefaultConfig() config.TerseErrors = false // Sanitize the log messages, which means that the bind vars are omitted config.SanitizeLogMessages = true - tsv := NewTabletServer("TabletServerTest", config, memorytopo.NewServer(""), &topodatapb.TabletAlias{}) + tsv := NewTabletServer(ctx, "TabletServerTest", config, memorytopo.NewServer(ctx, ""), &topodatapb.TabletAlias{}) tl := newTestLogger() defer tl.Close() @@ -1831,9 +1925,11 @@ func TestTruncateMessages(t *testing.T) { } func TestTerseErrorsIgnoreFailoverInProgress(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() config := tabletenv.NewDefaultConfig() config.TerseErrors = true - tsv := NewTabletServer("TabletServerTest", config, memorytopo.NewServer(""), &topodatapb.TabletAlias{}) + tsv := NewTabletServer(ctx, "TabletServerTest", config, memorytopo.NewServer(ctx, ""), &topodatapb.TabletAlias{}) tl := newTestLogger() defer tl.Close() err := tsv.convertAndLogError(ctx, "select * from test_table where id = :a", @@ -1871,9 +1967,11 @@ var aclJSON2 = `{ }` func TestACLHUP(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() tableacl.Register("simpleacl", &simpleacl.Factory{}) config := tabletenv.NewDefaultConfig() - tsv := NewTabletServer("TabletServerTest", config, memorytopo.NewServer(""), &topodatapb.TabletAlias{}) + tsv := NewTabletServer(ctx, "TabletServerTest", config, memorytopo.NewServer(ctx, ""), &topodatapb.TabletAlias{}) f, err := os.CreateTemp("", "tableacl") require.NoError(t, err) @@ -1916,7 +2014,9 @@ func TestACLHUP(t *testing.T) { } func TestConfigChanges(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") defer tsv.StopService() defer db.Close() @@ -1981,7 +2081,9 @@ func TestConfigChanges(t *testing.T) { } func TestReserveBeginExecute(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") tsv.config.EnableSettingsPool = false defer tsv.StopService() defer db.Close() @@ -2006,7 +2108,9 @@ func TestReserveBeginExecute(t *testing.T) { } func TestReserveExecute_WithoutTx(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") tsv.config.EnableSettingsPool = false defer tsv.StopService() defer db.Close() @@ -2029,7 +2133,9 @@ func TestReserveExecute_WithoutTx(t *testing.T) { } func TestReserveExecute_WithTx(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") tsv.config.EnableSettingsPool = false defer tsv.StopService() defer db.Close() @@ -2089,7 +2195,9 @@ func TestRelease(t *testing.T) { name += " reserve" } t.Run(name, func(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") tsv.config.EnableSettingsPool = false defer tsv.StopService() defer db.Close() @@ -2132,7 +2240,9 @@ func TestRelease(t *testing.T) { } func TestReserveStats(t *testing.T) { - db, tsv := setupTabletServerTest(t, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "") tsv.config.EnableSettingsPool = false defer tsv.StopService() defer db.Close() @@ -2142,7 +2252,7 @@ func TestReserveStats(t *testing.T) { callerID := &querypb.VTGateCallerID{ Username: "test", } - ctx := callerid.NewContext(context.Background(), nil, callerID) + ctx = callerid.NewContext(ctx, nil, callerID) // Starts reserved connection and transaction rbeState, _, err := tsv.ReserveBeginExecute(ctx, &target, nil, nil, "select 42", nil, &querypb.ExecuteOptions{}) @@ -2187,7 +2297,9 @@ func TestReserveStats(t *testing.T) { } func TestDatabaseNameReplaceByKeyspaceNameExecuteMethod(t *testing.T) { - db, tsv := setupTabletServerTest(t, "keyspaceName") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "keyspaceName") setDBName(db, tsv, "databaseInMysql") defer tsv.StopService() defer db.Close() @@ -2223,7 +2335,9 @@ func TestDatabaseNameReplaceByKeyspaceNameExecuteMethod(t *testing.T) { } func TestDatabaseNameReplaceByKeyspaceNameStreamExecuteMethod(t *testing.T) { - db, tsv := setupTabletServerTest(t, "keyspaceName") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "keyspaceName") setDBName(db, tsv, "databaseInMysql") defer tsv.StopService() defer db.Close() @@ -2260,7 +2374,9 @@ func TestDatabaseNameReplaceByKeyspaceNameStreamExecuteMethod(t *testing.T) { } func TestDatabaseNameReplaceByKeyspaceNameBeginExecuteMethod(t *testing.T) { - db, tsv := setupTabletServerTest(t, "keyspaceName") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "keyspaceName") setDBName(db, tsv, "databaseInMysql") defer tsv.StopService() defer db.Close() @@ -2299,7 +2415,9 @@ func setDBName(db *fakesqldb.DB, tsv *TabletServer, s string) { } func TestDatabaseNameReplaceByKeyspaceNameReserveExecuteMethod(t *testing.T) { - db, tsv := setupTabletServerTest(t, "keyspaceName") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "keyspaceName") tsv.config.EnableSettingsPool = false setDBName(db, tsv, "databaseInMysql") defer tsv.StopService() @@ -2334,7 +2452,9 @@ func TestDatabaseNameReplaceByKeyspaceNameReserveExecuteMethod(t *testing.T) { } func TestDatabaseNameReplaceByKeyspaceNameReserveBeginExecuteMethod(t *testing.T) { - db, tsv := setupTabletServerTest(t, "keyspaceName") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, tsv := setupTabletServerTest(t, ctx, "keyspaceName") setDBName(db, tsv, "databaseInMysql") defer tsv.StopService() defer db.Close() @@ -2367,15 +2487,15 @@ func TestDatabaseNameReplaceByKeyspaceNameReserveBeginExecuteMethod(t *testing.T require.NoError(t, err) } -func setupTabletServerTest(t *testing.T, keyspaceName string) (*fakesqldb.DB, *TabletServer) { +func setupTabletServerTest(t testing.TB, ctx context.Context, keyspaceName string) (*fakesqldb.DB, *TabletServer) { config := tabletenv.NewDefaultConfig() - return setupTabletServerTestCustom(t, config, keyspaceName) + return setupTabletServerTestCustom(t, ctx, config, keyspaceName) } -func setupTabletServerTestCustom(t *testing.T, config *tabletenv.TabletConfig, keyspaceName string) (*fakesqldb.DB, *TabletServer) { +func setupTabletServerTestCustom(t testing.TB, ctx context.Context, config *tabletenv.TabletConfig, keyspaceName string) (*fakesqldb.DB, *TabletServer) { db := setupFakeDB(t) sidecardb.AddSchemaInitQueries(db, true) - tsv := NewTabletServer("TabletServerTest", config, memorytopo.NewServer(""), &topodatapb.TabletAlias{}) + tsv := NewTabletServer(ctx, "TabletServerTest", config, memorytopo.NewServer(ctx, ""), &topodatapb.TabletAlias{}) require.Equal(t, StateNotConnected, tsv.sm.State()) dbcfgs := newDBConfigs(db) target := &querypb.Target{ @@ -2387,7 +2507,7 @@ func setupTabletServerTestCustom(t *testing.T, config *tabletenv.TabletConfig, k return db, tsv } -func setupFakeDB(t *testing.T) *fakesqldb.DB { +func setupFakeDB(t testing.TB) *fakesqldb.DB { db := fakesqldb.New(t) addTabletServerSupportedQueries(db) db.AddQueryPattern(baseShowTablesPattern, &sqltypes.Result{ diff --git a/go/vt/vttablet/tabletserver/twopc_test.go b/go/vt/vttablet/tabletserver/twopc_test.go index fe2acd983c8..2ef5e05b7c7 100644 --- a/go/vt/vttablet/tabletserver/twopc_test.go +++ b/go/vt/vttablet/tabletserver/twopc_test.go @@ -17,6 +17,7 @@ limitations under the License. package tabletserver import ( + "context" "encoding/json" "reflect" "testing" @@ -24,20 +25,19 @@ import ( "vitess.io/vitess/go/vt/vttablet/tabletserver/tx" - "context" - "vitess.io/vitess/go/sqltypes" querypb "vitess.io/vitess/go/vt/proto/query" ) func TestReadAllRedo(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() // Reuse code from tx_executor_test. - _, tsv, db := newTestTxExecutor(t) + _, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() tpc := tsv.te.twoPC - ctx := context.Background() conn, err := tsv.qe.conns.Get(ctx, nil) if err != nil { @@ -237,11 +237,12 @@ func TestReadAllRedo(t *testing.T) { } func TestReadAllTransactions(t *testing.T) { - _, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() tpc := tsv.te.twoPC - ctx := context.Background() conn, err := tsv.qe.conns.Get(ctx, nil) if err != nil { diff --git a/go/vt/vttablet/tabletserver/tx_engine_test.go b/go/vt/vttablet/tabletserver/tx_engine_test.go index 21ba81d84b6..6a86d044bcd 100644 --- a/go/vt/vttablet/tabletserver/tx_engine_test.go +++ b/go/vt/vttablet/tabletserver/tx_engine_test.go @@ -144,6 +144,8 @@ func TestTxEngineClose(t *testing.T) { } func TestTxEngineBegin(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() db := setUpQueryExecutorTest(t) defer db.Close() db.AddQueryPattern(".*", &sqltypes.Result{}) @@ -188,6 +190,8 @@ func TestTxEngineBegin(t *testing.T) { } func TestTxEngineRenewFails(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() db := setUpQueryExecutorTest(t) defer db.Close() db.AddQueryPattern(".*", &sqltypes.Result{}) @@ -556,6 +560,8 @@ func startTx(te *TxEngine, writeTransaction bool) error { } func TestTxEngineFailReserve(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() db := setUpQueryExecutorTest(t) defer db.Close() db.AddQueryPattern(".*", &sqltypes.Result{}) diff --git a/go/vt/vttablet/tabletserver/tx_executor_test.go b/go/vt/vttablet/tabletserver/tx_executor_test.go index edeb5ea5e1d..2651eb2a6cc 100644 --- a/go/vt/vttablet/tabletserver/tx_executor_test.go +++ b/go/vt/vttablet/tabletserver/tx_executor_test.go @@ -40,7 +40,9 @@ import ( ) func TestTxExecutorEmptyPrepare(t *testing.T) { - txe, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + txe, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() txid := newTransaction(tsv, nil) @@ -51,10 +53,12 @@ func TestTxExecutorEmptyPrepare(t *testing.T) { } func TestTxExecutorPrepare(t *testing.T) { - txe, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + txe, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() - txid := newTxForPrep(tsv) + txid := newTxForPrep(ctx, tsv) err := txe.Prepare(txid, "aa") require.NoError(t, err) err = txe.RollbackPrepared("aa", 1) @@ -68,7 +72,9 @@ func TestTxExecutorPrepare(t *testing.T) { } func TestTxExecutorPrepareNotInTx(t *testing.T) { - txe, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + txe, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() err := txe.Prepare(0, "aa") @@ -76,11 +82,13 @@ func TestTxExecutorPrepareNotInTx(t *testing.T) { } func TestTxExecutorPreparePoolFail(t *testing.T) { - txe, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + txe, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() - txid1 := newTxForPrep(tsv) - txid2 := newTxForPrep(tsv) + txid1 := newTxForPrep(ctx, tsv) + txid2 := newTxForPrep(ctx, tsv) err := txe.Prepare(txid1, "aa") require.NoError(t, err) defer txe.RollbackPrepared("aa", 0) @@ -90,10 +98,12 @@ func TestTxExecutorPreparePoolFail(t *testing.T) { } func TestTxExecutorPrepareRedoBeginFail(t *testing.T) { - txe, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + txe, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() - txid := newTxForPrep(tsv) + txid := newTxForPrep(ctx, tsv) db.AddRejectedQuery("begin", errors.New("begin fail")) err := txe.Prepare(txid, "aa") defer txe.RollbackPrepared("aa", 0) @@ -102,10 +112,12 @@ func TestTxExecutorPrepareRedoBeginFail(t *testing.T) { } func TestTxExecutorPrepareRedoFail(t *testing.T) { - txe, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + txe, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() - txid := newTxForPrep(tsv) + txid := newTxForPrep(ctx, tsv) err := txe.Prepare(txid, "bb") defer txe.RollbackPrepared("bb", 0) require.Error(t, err) @@ -113,10 +125,12 @@ func TestTxExecutorPrepareRedoFail(t *testing.T) { } func TestTxExecutorPrepareRedoCommitFail(t *testing.T) { - txe, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + txe, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() - txid := newTxForPrep(tsv) + txid := newTxForPrep(ctx, tsv) db.AddRejectedQuery("commit", errors.New("commit fail")) err := txe.Prepare(txid, "aa") defer txe.RollbackPrepared("aa", 0) @@ -125,10 +139,12 @@ func TestTxExecutorPrepareRedoCommitFail(t *testing.T) { } func TestTxExecutorCommit(t *testing.T) { - txe, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + txe, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() - txid := newTxForPrep(tsv) + txid := newTxForPrep(ctx, tsv) err := txe.Prepare(txid, "aa") require.NoError(t, err) err = txe.CommitPrepared("aa") @@ -139,10 +155,12 @@ func TestTxExecutorCommit(t *testing.T) { } func TestTxExecutorCommitRedoFail(t *testing.T) { - txe, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + txe, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() - txid := newTxForPrep(tsv) + txid := newTxForPrep(ctx, tsv) // Allow all additions to redo logs to succeed db.AddQueryPattern("insert into _vt\\.redo_state.*", &sqltypes.Result{}) err := txe.Prepare(txid, "bb") @@ -159,10 +177,12 @@ func TestTxExecutorCommitRedoFail(t *testing.T) { } func TestTxExecutorCommitRedoCommitFail(t *testing.T) { - txe, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + txe, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() - txid := newTxForPrep(tsv) + txid := newTxForPrep(ctx, tsv) err := txe.Prepare(txid, "aa") require.NoError(t, err) defer txe.RollbackPrepared("aa", 0) @@ -173,10 +193,12 @@ func TestTxExecutorCommitRedoCommitFail(t *testing.T) { } func TestTxExecutorRollbackBeginFail(t *testing.T) { - txe, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + txe, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() - txid := newTxForPrep(tsv) + txid := newTxForPrep(ctx, tsv) err := txe.Prepare(txid, "aa") require.NoError(t, err) db.AddRejectedQuery("begin", errors.New("begin fail")) @@ -186,10 +208,12 @@ func TestTxExecutorRollbackBeginFail(t *testing.T) { } func TestTxExecutorRollbackRedoFail(t *testing.T) { - txe, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + txe, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() - txid := newTxForPrep(tsv) + txid := newTxForPrep(ctx, tsv) // Allow all additions to redo logs to succeed db.AddQueryPattern("insert into _vt\\.redo_state.*", &sqltypes.Result{}) err := txe.Prepare(txid, "bb") @@ -200,7 +224,9 @@ func TestTxExecutorRollbackRedoFail(t *testing.T) { } func TestExecutorCreateTransaction(t *testing.T) { - txe, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + txe, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() @@ -214,43 +240,49 @@ func TestExecutorCreateTransaction(t *testing.T) { } func TestExecutorStartCommit(t *testing.T) { - txe, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + txe, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() commitTransition := fmt.Sprintf("update _vt.dt_state set state = %d where dtid = 'aa' and state = %d", int(querypb.TransactionState_COMMIT), int(querypb.TransactionState_PREPARE)) db.AddQuery(commitTransition, &sqltypes.Result{RowsAffected: 1}) - txid := newTxForPrep(tsv) + txid := newTxForPrep(ctx, tsv) err := txe.StartCommit(txid, "aa") require.NoError(t, err) db.AddQuery(commitTransition, &sqltypes.Result{}) - txid = newTxForPrep(tsv) + txid = newTxForPrep(ctx, tsv) err = txe.StartCommit(txid, "aa") require.Error(t, err) require.Contains(t, err.Error(), "could not transition to COMMIT: aa") } func TestExecutorSetRollback(t *testing.T) { - txe, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + txe, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() rollbackTransition := fmt.Sprintf("update _vt.dt_state set state = %d where dtid = 'aa' and state = %d", int(querypb.TransactionState_ROLLBACK), int(querypb.TransactionState_PREPARE)) db.AddQuery(rollbackTransition, &sqltypes.Result{RowsAffected: 1}) - txid := newTxForPrep(tsv) + txid := newTxForPrep(ctx, tsv) err := txe.SetRollback("aa", txid) require.NoError(t, err) db.AddQuery(rollbackTransition, &sqltypes.Result{}) - txid = newTxForPrep(tsv) + txid = newTxForPrep(ctx, tsv) err = txe.SetRollback("aa", txid) require.Error(t, err) require.Contains(t, err.Error(), "could not transition to ROLLBACK: aa") } func TestExecutorConcludeTransaction(t *testing.T) { - txe, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + txe, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() @@ -261,7 +293,9 @@ func TestExecutorConcludeTransaction(t *testing.T) { } func TestExecutorReadTransaction(t *testing.T) { - txe, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + txe, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() @@ -361,7 +395,9 @@ func TestExecutorReadTransaction(t *testing.T) { } func TestExecutorReadAllTransactions(t *testing.T) { - txe, tsv, db := newTestTxExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + txe, tsv, db := newTestTxExecutor(t, ctx) defer db.Close() defer tsv.StopService() @@ -410,6 +446,8 @@ func (conn *FakeVTGateConn) ResolveTransaction(ctx context.Context, dtid string) } func TestExecutorResolveTransaction(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() protocol := "resolveTest" oldValue := vtgateconn.GetVTGateProtocol() vtgateconn.SetVTGateProtocol(protocol) @@ -421,7 +459,7 @@ func TestExecutorResolveTransaction(t *testing.T) { FakeVTGateConn: fakerpcvtgateconn.FakeVTGateConn{}, }, nil }) - _, tsv, db := newShortAgeExecutor(t) + _, tsv, db := newShortAgeExecutor(t, ctx) defer db.Close() defer tsv.StopService() want := "aa" @@ -444,7 +482,9 @@ func TestExecutorResolveTransaction(t *testing.T) { } func TestNoTwopc(t *testing.T) { - txe, tsv, db := newNoTwopcExecutor(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + txe, tsv, db := newNoTwopcExecutor(t, ctx) defer db.Close() defer tsv.StopService() @@ -493,7 +533,7 @@ func TestNoTwopc(t *testing.T) { } } -func newTestTxExecutor(t *testing.T) (txe *TxExecutor, tsv *TabletServer, db *fakesqldb.DB) { +func newTestTxExecutor(t *testing.T, ctx context.Context) (txe *TxExecutor, tsv *TabletServer, db *fakesqldb.DB) { db = setUpQueryExecutorTest(t) logStats := tabletenv.NewLogStats(ctx, "TestTxExecutor") tsv = newTestTabletServer(ctx, smallTxPool, db) @@ -510,7 +550,7 @@ func newTestTxExecutor(t *testing.T) (txe *TxExecutor, tsv *TabletServer, db *fa } // newShortAgeExecutor is same as newTestTxExecutor, but shorter transaction abandon age. -func newShortAgeExecutor(t *testing.T) (txe *TxExecutor, tsv *TabletServer, db *fakesqldb.DB) { +func newShortAgeExecutor(t *testing.T, ctx context.Context) (txe *TxExecutor, tsv *TabletServer, db *fakesqldb.DB) { db = setUpQueryExecutorTest(t) logStats := tabletenv.NewLogStats(ctx, "TestTxExecutor") tsv = newTestTabletServer(ctx, smallTxPool|shortTwopcAge, db) @@ -527,7 +567,7 @@ func newShortAgeExecutor(t *testing.T) (txe *TxExecutor, tsv *TabletServer, db * } // newNoTwopcExecutor is same as newTestTxExecutor, but 2pc disabled. -func newNoTwopcExecutor(t *testing.T) (txe *TxExecutor, tsv *TabletServer, db *fakesqldb.DB) { +func newNoTwopcExecutor(t *testing.T, ctx context.Context) (txe *TxExecutor, tsv *TabletServer, db *fakesqldb.DB) { db = setUpQueryExecutorTest(t) logStats := tabletenv.NewLogStats(ctx, "TestTxExecutor") tsv = newTestTabletServer(ctx, noTwopc, db) @@ -539,7 +579,7 @@ func newNoTwopcExecutor(t *testing.T) (txe *TxExecutor, tsv *TabletServer, db *f } // newTxForPrep creates a non-empty transaction. -func newTxForPrep(tsv *TabletServer) int64 { +func newTxForPrep(ctx context.Context, tsv *TabletServer) int64 { txid := newTransaction(tsv, nil) target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY} _, err := tsv.Execute(ctx, &target, "update test_table set name = 2 where pk = 1", nil, txid, 0, nil) diff --git a/go/vt/vttablet/tabletserver/tx_pool_test.go b/go/vt/vttablet/tabletserver/tx_pool_test.go index 46b0552a0a3..3515310c481 100644 --- a/go/vt/vttablet/tabletserver/tx_pool_test.go +++ b/go/vt/vttablet/tabletserver/tx_pool_test.go @@ -39,6 +39,8 @@ import ( ) func TestTxPoolExecuteCommit(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() db, txPool, _, closer := setup(t) defer closer() @@ -73,6 +75,9 @@ func TestTxPoolExecuteCommit(t *testing.T) { } func TestTxPoolExecuteRollback(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, txPool, _, closer := setup(t) defer closer() @@ -91,6 +96,9 @@ func TestTxPoolExecuteRollback(t *testing.T) { } func TestTxPoolExecuteRollbackOnClosedConn(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, txPool, _, closer := setup(t) defer closer() @@ -108,6 +116,9 @@ func TestTxPoolExecuteRollbackOnClosedConn(t *testing.T) { } func TestTxPoolRollbackNonBusy(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, txPool, _, closer := setup(t) defer closer() @@ -135,6 +146,9 @@ func TestTxPoolRollbackNonBusy(t *testing.T) { } func TestTxPoolTransactionIsolation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, txPool, _, closer := setup(t) defer closer() @@ -146,6 +160,9 @@ func TestTxPoolTransactionIsolation(t *testing.T) { } func TestTxPoolAutocommit(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, txPool, _, closer := setup(t) defer closer() @@ -173,7 +190,10 @@ func TestTxPoolAutocommit(t *testing.T) { // db connection. DBConn.Exec() is going to reconnect and retry automatically // due to this connection error and the BEGIN will succeed. func TestTxPoolBeginWithPoolConnectionError_Errno2006_Transient(t *testing.T) { - db, txPool := primeTxPoolWithConnection(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + db, txPool := primeTxPoolWithConnection(t, ctx) defer db.Close() defer txPool.Close() @@ -189,7 +209,7 @@ func TestTxPoolBeginWithPoolConnectionError_Errno2006_Transient(t *testing.T) { // primeTxPoolWithConnection is a helper function. It reconstructs the // scenario where future transactions are going to reuse an open db connection. -func primeTxPoolWithConnection(t *testing.T) (*fakesqldb.DB, *TxPool) { +func primeTxPoolWithConnection(t *testing.T, ctx context.Context) (*fakesqldb.DB, *TxPool) { t.Helper() db := fakesqldb.New(t) txPool, _ := newTxPool() @@ -209,6 +229,9 @@ func primeTxPoolWithConnection(t *testing.T) (*fakesqldb.DB, *TxPool) { } func TestTxPoolBeginWithError(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, txPool, limiter, closer := setup(t) defer closer() db.AddRejectedQuery("begin", errRejected) @@ -244,6 +267,9 @@ func TestTxPoolBeginWithError(t *testing.T) { } func TestTxPoolBeginWithPreQueryError(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, txPool, _, closer := setup(t) defer closer() db.AddRejectedQuery("pre_query", errRejected) @@ -257,7 +283,7 @@ func TestTxPoolCancelledContextError(t *testing.T) { // given db, txPool, _, closer := setup(t) defer closer() - ctx, cancel := context.WithCancel(ctx) + ctx, cancel := context.WithCancel(context.Background()) cancel() // when @@ -271,6 +297,9 @@ func TestTxPoolCancelledContextError(t *testing.T) { } func TestTxPoolWaitTimeoutError(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newEnv("TabletServerTest") env.Config().TxPool.Size = 1 env.Config().TxPool.MaxWaiters = 0 @@ -297,6 +326,9 @@ func TestTxPoolWaitTimeoutError(t *testing.T) { } func TestTxPoolRollbackFailIsPassedThrough(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + sql := "alter table test_table add test_column int" db, txPool, _, closer := setup(t) defer closer() @@ -317,6 +349,9 @@ func TestTxPoolRollbackFailIsPassedThrough(t *testing.T) { } func TestTxPoolGetConnRecentlyRemovedTransaction(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + db, txPool, _, _ := setup(t) defer db.Close() conn1, _, _, _ := txPool.Begin(ctx, &querypb.ExecuteOptions{}, false, 0, nil, nil) @@ -384,6 +419,9 @@ func TestTxPoolCloseKillsStrayTransactions(t *testing.T) { } func TestTxTimeoutKillsTransactions(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newEnv("TabletServerTest") env.Config().TxPool.Size = 1 env.Config().TxPool.MaxWaiters = 0 @@ -430,6 +468,9 @@ func TestTxTimeoutKillsTransactions(t *testing.T) { } func TestTxTimeoutDoesNotKillShortLivedTransactions(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newEnv("TabletServerTest") env.Config().TxPool.Size = 1 env.Config().TxPool.MaxWaiters = 0 @@ -460,6 +501,9 @@ func TestTxTimeoutDoesNotKillShortLivedTransactions(t *testing.T) { } func TestTxTimeoutKillsOlapTransactions(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newEnv("TabletServerTest") env.Config().TxPool.Size = 1 env.Config().TxPool.MaxWaiters = 0 @@ -495,6 +539,9 @@ func TestTxTimeoutKillsOlapTransactions(t *testing.T) { } func TestTxTimeoutNotEnforcedForZeroLengthTimeouts(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newEnv("TabletServerTest") env.Config().TxPool.Size = 2 env.Config().TxPool.MaxWaiters = 0 @@ -535,6 +582,9 @@ func TestTxTimeoutNotEnforcedForZeroLengthTimeouts(t *testing.T) { } func TestTxTimeoutReservedConn(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newEnv("TabletServerTest") env.Config().TxPool.Size = 1 env.Config().TxPool.MaxWaiters = 0 @@ -575,6 +625,9 @@ func TestTxTimeoutReservedConn(t *testing.T) { } func TestTxTimeoutReusedReservedConn(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newEnv("TabletServerTest") env.Config().TxPool.Size = 1 env.Config().TxPool.MaxWaiters = 0 @@ -728,6 +781,9 @@ func TestTxPoolBeginStatements(t *testing.T) { for _, tc := range testCases { t.Run(fmt.Sprintf("%v:%v:readOnly:%v", tc.txIsolationLevel, tc.txAccessModes, tc.readOnly), func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + options := &querypb.ExecuteOptions{ TransactionIsolation: tc.txIsolationLevel, TransactionAccessMode: tc.txAccessModes, diff --git a/go/vt/vttablet/tabletserver/txthrottler/tx_throttler_test.go b/go/vt/vttablet/tabletserver/txthrottler/tx_throttler_test.go index b0d15a5b059..ea57d37ad8e 100644 --- a/go/vt/vttablet/tabletserver/txthrottler/tx_throttler_test.go +++ b/go/vt/vttablet/tabletserver/txthrottler/tx_throttler_test.go @@ -57,11 +57,13 @@ func TestDisabledThrottler(t *testing.T) { } func TestEnabledThrottler(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() mockCtrl := gomock.NewController(t) defer mockCtrl.Finish() defer resetTxThrottlerFactories() - ts := memorytopo.NewServer("cell1", "cell2") + ts := memorytopo.NewServer(ctx, "cell1", "cell2") mockHealthCheck := NewMockHealthCheck(mockCtrl) hcCall1 := mockHealthCheck.EXPECT().Subscribe() @@ -164,13 +166,15 @@ func TestEnabledThrottler(t *testing.T) { } func TestFetchKnownCells(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() { - ts := memorytopo.NewServer("cell1", "cell2") + ts := memorytopo.NewServer(ctx, "cell1", "cell2") cells := fetchKnownCells(context.Background(), ts, &querypb.Target{Cell: "cell1"}) assert.Equal(t, []string{"cell1", "cell2"}, cells) } { - ts := memorytopo.NewServer() + ts := memorytopo.NewServer(ctx) cells := fetchKnownCells(context.Background(), ts, &querypb.Target{Cell: "cell1"}) assert.Equal(t, []string{"cell1"}, cells) } diff --git a/go/vt/vttablet/tabletserver/vstreamer/main_flaky_test.go b/go/vt/vttablet/tabletserver/vstreamer/main_flaky_test.go index a62489e2fe0..e2cdec9b3f1 100644 --- a/go/vt/vttablet/tabletserver/vstreamer/main_flaky_test.go +++ b/go/vt/vttablet/tabletserver/vstreamer/main_flaky_test.go @@ -17,6 +17,7 @@ limitations under the License. package vstreamer import ( + "context" "fmt" "os" "testing" @@ -44,7 +45,9 @@ func TestMain(m *testing.M) { exitCode := func() int { var err error - env, err = testenv.Init() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env, err = testenv.Init(ctx) if err != nil { fmt.Fprintf(os.Stderr, "%v", err) return 1 @@ -63,7 +66,7 @@ func TestMain(m *testing.M) { os.Exit(exitCode) } -func newEngine(t *testing.T, binlogRowImage string) { +func newEngine(t *testing.T, ctx context.Context, binlogRowImage string) { if engine != nil { engine.Close() } @@ -71,7 +74,7 @@ func newEngine(t *testing.T, binlogRowImage string) { env.Close() } var err error - env, err = testenv.Init() + env, err = testenv.Init(ctx) require.NoError(t, err) setBinlogRowImage(t, binlogRowImage) diff --git a/go/vt/vttablet/tabletserver/vstreamer/testenv/testenv.go b/go/vt/vttablet/tabletserver/vstreamer/testenv/testenv.go index edec8d25d47..c40e180110f 100644 --- a/go/vt/vttablet/tabletserver/vstreamer/testenv/testenv.go +++ b/go/vt/vttablet/tabletserver/vstreamer/testenv/testenv.go @@ -61,22 +61,21 @@ type Env struct { } // Init initializes an Env. -func Init() (*Env, error) { +func Init(ctx context.Context) (*Env, error) { te := &Env{ KeyspaceName: "vttest", ShardName: "0", Cells: []string{"cell1"}, } - ctx := context.Background() - te.TopoServ = memorytopo.NewServer(te.Cells...) + te.TopoServ = memorytopo.NewServer(ctx, te.Cells...) if err := te.TopoServ.CreateKeyspace(ctx, te.KeyspaceName, &topodatapb.Keyspace{}); err != nil { return nil, err } if err := te.TopoServ.CreateShard(ctx, te.KeyspaceName, te.ShardName); err != nil { panic(err) } - te.SrvTopo = srvtopo.NewResilientServer(te.TopoServ, "TestTopo") + te.SrvTopo = srvtopo.NewResilientServer(ctx, te.TopoServ, "TestTopo") cfg := vttest.Config{ Topology: &vttestpb.VTTestTopology{ diff --git a/go/vt/vttablet/tabletserver/vstreamer/vstreamer_flaky_test.go b/go/vt/vttablet/tabletserver/vstreamer/vstreamer_flaky_test.go index 0c9eda07752..70cbf1896c8 100644 --- a/go/vt/vttablet/tabletserver/vstreamer/vstreamer_flaky_test.go +++ b/go/vt/vttablet/tabletserver/vstreamer/vstreamer_flaky_test.go @@ -86,8 +86,13 @@ func (tfe *TestFieldEvent) String() string { // TestPlayerNoBlob sets up a new environment with mysql running with binlog_row_image as noblob. It confirms that // the VEvents created are correct: that they don't contain the missing columns and that the DataColumns bitmap is sent func TestNoBlob(t *testing.T) { - newEngine(t, "noblob") - defer newEngine(t, "full") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + oldEngine := engine + newEngine(t, ctx, "noblob") + defer func() { + engine = oldEngine + }() execStatements(t, []string{ "create table t1(id int, blb blob, val varchar(4), primary key(id))", "create table t2(id int, txt text, val varchar(4), unique key(id, val))", @@ -1928,8 +1933,13 @@ func TestMinimalMode(t *testing.T) { t.Skip() } - newEngine(t, "minimal") - defer newEngine(t, "full") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + oldEngine := engine + newEngine(t, ctx, "minimal") + defer func() { + engine = oldEngine + }() err := engine.Stream(context.Background(), "current", nil, nil, throttlerapp.VStreamerName, func(evs []*binlogdatapb.VEvent) error { return nil }) require.Error(t, err, "minimal binlog_row_image is not supported by Vitess VReplication") } diff --git a/go/vt/wrangler/external_cluster_test.go b/go/vt/wrangler/external_cluster_test.go index 3592c4b6895..3c878411b6b 100644 --- a/go/vt/wrangler/external_cluster_test.go +++ b/go/vt/wrangler/external_cluster_test.go @@ -14,8 +14,9 @@ import ( ) func TestVitessCluster(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") tmc := newTestWranglerTMClient() wr := New(logutil.NewConsoleLogger(), ts, tmc) name, topoType, topoServer, topoRoot := "c1", "x", "y", "z" diff --git a/go/vt/wrangler/materializer_env_test.go b/go/vt/wrangler/materializer_env_test.go index 48cca1a0bb8..6c236a038bf 100644 --- a/go/vt/wrangler/materializer_env_test.go +++ b/go/vt/wrangler/materializer_env_test.go @@ -118,14 +118,14 @@ func TestMain(m *testing.M) { os.Exit(testMainWrapper(m)) } -func newTestMaterializerEnv(t *testing.T, ms *vtctldatapb.MaterializeSettings, sources, targets []string) *testMaterializerEnv { +func newTestMaterializerEnv(t *testing.T, ctx context.Context, ms *vtctldatapb.MaterializeSettings, sources, targets []string) *testMaterializerEnv { t.Helper() env := &testMaterializerEnv{ ms: ms, sources: sources, targets: targets, tablets: make(map[int]*topodatapb.Tablet), - topoServ: memorytopo.NewServer("cell"), + topoServ: memorytopo.NewServer(ctx, "cell"), cell: "cell", tmc: newTestMaterializerTMClient(), } diff --git a/go/vt/wrangler/materializer_test.go b/go/vt/wrangler/materializer_test.go index 0001058a7b3..0350d8f00fd 100644 --- a/go/vt/wrangler/materializer_test.go +++ b/go/vt/wrangler/materializer_test.go @@ -59,7 +59,9 @@ func TestMigrateTables(t *testing.T) { SourceExpression: "select * from t1", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() env.tmc.expectVRQuery(100, mzCheckJournal, &sqltypes.Result{}) @@ -68,7 +70,6 @@ func TestMigrateTables(t *testing.T) { env.tmc.expectVRQuery(200, mzSelectIDQuery, &sqltypes.Result{}) env.tmc.expectVRQuery(200, mzUpdateQuery, &sqltypes.Result{}) - ctx := context.Background() err := env.wr.MoveTables(ctx, "workflow", "sourceks", "targetks", "t1", "", "", false, "", true, false, "", false, false, "", defaultOnDDL, nil) require.NoError(t, err) vschema, err := env.wr.ts.GetSrvVSchema(ctx, env.cell) @@ -100,7 +101,10 @@ func TestMissingTables(t *testing.T) { SourceExpression: "select * from t3", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() env.tmc.expectVRQuery(100, mzCheckJournal, &sqltypes.Result{}) @@ -109,7 +113,6 @@ func TestMissingTables(t *testing.T) { env.tmc.expectVRQuery(200, mzSelectIDQuery, &sqltypes.Result{}) env.tmc.expectVRQuery(200, mzUpdateQuery, &sqltypes.Result{}) - ctx := context.Background() err := env.wr.MoveTables(ctx, "workflow", "sourceks", "targetks", "t1,tyt", "", "", false, "", true, false, "", false, false, "", defaultOnDDL, nil) require.EqualError(t, err, "table(s) not found in source keyspace sourceks: tyt") err = env.wr.MoveTables(ctx, "workflow", "sourceks", "targetks", "t1,tyt,t2,txt", "", "", false, "", true, false, "", false, false, "", defaultOnDDL, nil) @@ -135,10 +138,9 @@ func TestMoveTablesAllAndExclude(t *testing.T) { }}, } - ctx := context.Background() var err error - var targetTables = func(env *testMaterializerEnv) []string { + var targetTables = func(ctx context.Context, env *testMaterializerEnv) []string { vschema, err := env.wr.ts.GetSrvVSchema(ctx, env.cell) require.NoError(t, err) var targetTables []string @@ -163,7 +165,10 @@ func TestMoveTablesAllAndExclude(t *testing.T) { } for _, tcase := range testCases { t.Run("", func(t *testing.T) { - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() env.tmc.expectVRQuery(100, mzCheckJournal, &sqltypes.Result{}) env.tmc.expectVRQuery(200, mzSelectFrozenQuery, &sqltypes.Result{}) @@ -172,7 +177,7 @@ func TestMoveTablesAllAndExclude(t *testing.T) { env.tmc.expectVRQuery(200, mzUpdateQuery, &sqltypes.Result{}) err = env.wr.MoveTables(ctx, "workflow", "sourceks", "targetks", "", "", "", tcase.allTables, tcase.excludeTables, true, false, "", false, false, "", defaultOnDDL, nil) require.NoError(t, err) - require.EqualValues(t, tcase.want, targetTables(env)) + require.EqualValues(t, tcase.want, targetTables(ctx, env)) }) } @@ -190,10 +195,12 @@ func TestMoveTablesStopFlags(t *testing.T) { }}, } - ctx := context.Background() var err error t.Run("StopStartedAndStopAfterCopyFlags", func(t *testing.T) { - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() env.tmc.expectVRQuery(100, mzCheckJournal, &sqltypes.Result{}) env.tmc.expectVRQuery(200, mzSelectFrozenQuery, &sqltypes.Result{}) @@ -220,7 +227,10 @@ func TestMigrateVSchema(t *testing.T) { SourceExpression: "select * from t1", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() env.tmc.expectVRQuery(100, mzCheckJournal, &sqltypes.Result{}) @@ -229,7 +239,6 @@ func TestMigrateVSchema(t *testing.T) { env.tmc.expectVRQuery(200, mzSelectIDQuery, &sqltypes.Result{}) env.tmc.expectVRQuery(200, mzUpdateQuery, &sqltypes.Result{}) - ctx := context.Background() err := env.wr.MoveTables(ctx, "workflow", "sourceks", "targetks", `{"t1":{}}`, "", "", false, "", true, false, "", false, false, "", defaultOnDDL, nil) require.NoError(t, err) vschema, err := env.wr.ts.GetSrvVSchema(ctx, env.cell) @@ -251,7 +260,10 @@ func TestCreateLookupVindexFull(t *testing.T) { SourceKeyspace: "sourceks", TargetKeyspace: "targetks", } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() specs := &vschemapb.Keyspace{ @@ -323,7 +335,6 @@ func TestCreateLookupVindexFull(t *testing.T) { env.tmc.expectVRQuery(200, insertPrefix, &sqltypes.Result{}) env.tmc.expectVRQuery(200, "update _vt.vreplication set state='Running' where db_name='vt_targetks' and workflow='lkp_vdx'", &sqltypes.Result{}) - ctx := context.Background() err := env.wr.CreateLookupVindex(ctx, ms.SourceKeyspace, specs, "cell", "PRIMARY", false) require.NoError(t, err) @@ -375,7 +386,10 @@ func TestCreateLookupVindexCreateDDL(t *testing.T) { SourceKeyspace: "sourceks", TargetKeyspace: "targetks", } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() vs := &vschemapb.Keyspace{ Sharded: true, @@ -592,7 +606,10 @@ func TestCreateLookupVindexSourceVSchema(t *testing.T) { SourceKeyspace: "sourceks", TargetKeyspace: "targetks", } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() specs := &vschemapb.Keyspace{ @@ -828,7 +845,10 @@ func TestCreateLookupVindexTargetVSchema(t *testing.T) { SourceKeyspace: "sourceks", TargetKeyspace: "targetks", } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() sourcevs := &vschemapb.Keyspace{ Sharded: true, @@ -1065,7 +1085,10 @@ func TestCreateLookupVindexSameKeyspace(t *testing.T) { SourceKeyspace: "ks", TargetKeyspace: "ks", } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() specs := &vschemapb.Keyspace{ @@ -1174,7 +1197,10 @@ func TestCreateCustomizedVindex(t *testing.T) { SourceKeyspace: "ks", TargetKeyspace: "ks", } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() specs := &vschemapb.Keyspace{ @@ -1284,7 +1310,10 @@ func TestStopAfterCopyFlag(t *testing.T) { SourceKeyspace: "ks", TargetKeyspace: "ks", } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() specs := &vschemapb.Keyspace{ Vindexes: map[string]*vschemapb.Vindex{ @@ -1356,7 +1385,10 @@ func TestStopAfterCopyFlag(t *testing.T) { } func TestCreateLookupVindexFailures(t *testing.T) { - topoServ := memorytopo.NewServer("cell") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + topoServ := memorytopo.NewServer(ctx, "cell") wr := New(logutil.NewConsoleLogger(), topoServ, nil) unique := map[string]*vschemapb.Vindex{ @@ -1631,7 +1663,10 @@ func TestExternalizeVindex(t *testing.T) { SourceKeyspace: "sourceks", TargetKeyspace: "targetks", } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"-80", "80-"}) defer env.close() sourceVSchema := &vschemapb.Keyspace{ @@ -1779,7 +1814,10 @@ func TestMaterializerOneToOne(t *testing.T) { topodatapb.TabletType_RDONLY, }), } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() env.tmc.expectVRQuery(200, mzSelectFrozenQuery, &sqltypes.Result{}) @@ -1800,7 +1838,7 @@ func TestMaterializerOneToOne(t *testing.T) { ) env.tmc.expectVRQuery(200, mzUpdateQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.NoError(t, err) env.tmc.verifyQueries(t) } @@ -1820,7 +1858,10 @@ func TestMaterializerManyToOne(t *testing.T) { CreateDdl: "t2ddl", }}, } - env := newTestMaterializerEnv(t, ms, []string{"-80", "80-"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"-80", "80-"}, []string{"0"}) defer env.close() env.tmc.expectVRQuery(200, mzSelectFrozenQuery, &sqltypes.Result{}) @@ -1835,7 +1876,7 @@ func TestMaterializerManyToOne(t *testing.T) { ) env.tmc.expectVRQuery(200, mzUpdateQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.NoError(t, err) env.tmc.verifyQueries(t) } @@ -1851,7 +1892,10 @@ func TestMaterializerOneToMany(t *testing.T) { CreateDdl: "t1ddl", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"-80", "80-"}) defer env.close() vs := &vschemapb.Keyspace{ @@ -1892,7 +1936,7 @@ func TestMaterializerOneToMany(t *testing.T) { env.tmc.expectVRQuery(200, mzUpdateQuery, &sqltypes.Result{}) env.tmc.expectVRQuery(210, mzUpdateQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.NoError(t, err) env.tmc.verifyQueries(t) } @@ -1908,7 +1952,10 @@ func TestMaterializerManyToMany(t *testing.T) { CreateDdl: "t1ddl", }}, } - env := newTestMaterializerEnv(t, ms, []string{"-40", "40-"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"-40", "40-"}, []string{"-80", "80-"}) defer env.close() vs := &vschemapb.Keyspace{ @@ -1950,7 +1997,7 @@ func TestMaterializerManyToMany(t *testing.T) { ) env.tmc.expectVRQuery(200, mzUpdateQuery, &sqltypes.Result{}) env.tmc.expectVRQuery(210, mzUpdateQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.NoError(t, err) env.tmc.verifyQueries(t) } @@ -1966,7 +2013,10 @@ func TestMaterializerMulticolumnVindex(t *testing.T) { CreateDdl: "t1ddl", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"-80", "80-"}) defer env.close() vs := &vschemapb.Keyspace{ @@ -2010,7 +2060,7 @@ func TestMaterializerMulticolumnVindex(t *testing.T) { env.tmc.expectVRQuery(200, mzUpdateQuery, &sqltypes.Result{}) env.tmc.expectVRQuery(210, mzUpdateQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.NoError(t, err) env.tmc.verifyQueries(t) } @@ -2030,7 +2080,10 @@ func TestMaterializerDeploySchema(t *testing.T) { CreateDdl: "t2ddl", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() delete(env.tmc.schema, "targetks.t2") @@ -2046,7 +2099,7 @@ func TestMaterializerDeploySchema(t *testing.T) { ) env.tmc.expectVRQuery(200, mzUpdateQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.NoError(t, err) env.tmc.verifyQueries(t) require.Equal(t, env.tmc.getSchemaRequestCount(100), 1) @@ -2068,7 +2121,10 @@ func TestMaterializerCopySchema(t *testing.T) { CreateDdl: "t2ddl", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() delete(env.tmc.schema, "targetks.t1") @@ -2084,7 +2140,7 @@ func TestMaterializerCopySchema(t *testing.T) { ) env.tmc.expectVRQuery(200, mzUpdateQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.NoError(t, err) env.tmc.verifyQueries(t) require.Equal(t, env.tmc.getSchemaRequestCount(100), 1) @@ -2103,7 +2159,10 @@ func TestMaterializerExplicitColumns(t *testing.T) { CreateDdl: "t1ddl", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"-80", "80-"}) defer env.close() vs := &vschemapb.Keyspace{ @@ -2147,7 +2206,7 @@ func TestMaterializerExplicitColumns(t *testing.T) { env.tmc.expectVRQuery(200, mzUpdateQuery, &sqltypes.Result{}) env.tmc.expectVRQuery(210, mzUpdateQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.NoError(t, err) env.tmc.verifyQueries(t) } @@ -2163,7 +2222,10 @@ func TestMaterializerRenamedColumns(t *testing.T) { CreateDdl: "t1ddl", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"-80", "80-"}) defer env.close() vs := &vschemapb.Keyspace{ @@ -2207,7 +2269,7 @@ func TestMaterializerRenamedColumns(t *testing.T) { env.tmc.expectVRQuery(200, mzUpdateQuery, &sqltypes.Result{}) env.tmc.expectVRQuery(210, mzUpdateQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.NoError(t, err) env.tmc.verifyQueries(t) } @@ -2228,14 +2290,17 @@ func TestMaterializerStopAfterCopy(t *testing.T) { CreateDdl: "t2ddl", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() env.tmc.expectVRQuery(200, mzSelectFrozenQuery, &sqltypes.Result{}) env.tmc.expectVRQuery(200, insertPrefix+`.*stop_after_copy:true`, &sqltypes.Result{}) env.tmc.expectVRQuery(200, mzUpdateQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.NoError(t, err) env.tmc.verifyQueries(t) } @@ -2251,7 +2316,10 @@ func TestMaterializerNoTargetVSchema(t *testing.T) { CreateDdl: "t1ddl", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"-80", "80-"}) defer env.close() vs := &vschemapb.Keyspace{ @@ -2263,7 +2331,7 @@ func TestMaterializerNoTargetVSchema(t *testing.T) { } env.tmc.expectVRQuery(200, mzSelectFrozenQuery, &sqltypes.Result{}) env.tmc.expectVRQuery(210, mzSelectFrozenQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.EqualError(t, err, "table t1 not found in vschema for keyspace targetks") } @@ -2278,13 +2346,16 @@ func TestMaterializerNoDDL(t *testing.T) { CreateDdl: "", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() delete(env.tmc.schema, "targetks.t1") env.tmc.expectVRQuery(200, mzSelectFrozenQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.EqualError(t, err, "target table t1 does not exist and there is no create ddl defined") require.Equal(t, env.tmc.getSchemaRequestCount(100), 0) require.Equal(t, env.tmc.getSchemaRequestCount(200), 1) @@ -2305,13 +2376,16 @@ func TestMaterializerNoSourcePrimary(t *testing.T) { sources := []string{"0"} targets := []string{"0"} + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Copied from newTestMaterializerEnv env := &testMaterializerEnv{ ms: ms, sources: sources, targets: targets, tablets: make(map[int]*topodatapb.Tablet), - topoServ: memorytopo.NewServer("cell"), + topoServ: memorytopo.NewServer(ctx, "cell"), cell: "cell", tmc: newTestMaterializerTMClient(), } @@ -2334,7 +2408,7 @@ func TestMaterializerNoSourcePrimary(t *testing.T) { env.expectValidation() env.tmc.expectVRQuery(200, mzSelectFrozenQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.EqualError(t, err, "source shard must have a primary for copying schema: 0") } @@ -2349,13 +2423,16 @@ func TestMaterializerTableMismatchNonCopy(t *testing.T) { CreateDdl: "", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() delete(env.tmc.schema, "targetks.t1") env.tmc.expectVRQuery(200, mzSelectFrozenQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.EqualError(t, err, "target table t1 does not exist and there is no create ddl defined") } @@ -2370,13 +2447,16 @@ func TestMaterializerTableMismatchCopy(t *testing.T) { CreateDdl: "copy", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() delete(env.tmc.schema, "targetks.t1") env.tmc.expectVRQuery(200, mzSelectFrozenQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.EqualError(t, err, "source and target table names must match for copying schema: t2 vs t1") } @@ -2391,14 +2471,17 @@ func TestMaterializerNoSourceTable(t *testing.T) { CreateDdl: "copy", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() delete(env.tmc.schema, "targetks.t1") delete(env.tmc.schema, "sourceks.t1") env.tmc.expectVRQuery(200, mzSelectFrozenQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.EqualError(t, err, "source table t1 does not exist") } @@ -2413,11 +2496,14 @@ func TestMaterializerSyntaxError(t *testing.T) { CreateDdl: "t1ddl", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() env.tmc.expectVRQuery(200, mzSelectFrozenQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.EqualError(t, err, "syntax error at position 4 near 'bad'") } @@ -2432,11 +2518,14 @@ func TestMaterializerNotASelect(t *testing.T) { CreateDdl: "t1ddl", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() env.tmc.expectVRQuery(200, mzSelectFrozenQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.EqualError(t, err, "unrecognized statement: update t1 set val=1") } @@ -2451,7 +2540,10 @@ func TestMaterializerNoGoodVindex(t *testing.T) { CreateDdl: "t1ddl", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"-80", "80-"}) defer env.close() vs := &vschemapb.Keyspace{ @@ -2482,7 +2574,7 @@ func TestMaterializerNoGoodVindex(t *testing.T) { env.tmc.expectVRQuery(200, mzSelectFrozenQuery, &sqltypes.Result{}) env.tmc.expectVRQuery(210, mzSelectFrozenQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.EqualError(t, err, "could not find a vindex to compute keyspace id for table t1") } @@ -2497,7 +2589,9 @@ func TestMaterializerComplexVindexExpression(t *testing.T) { CreateDdl: "t1ddl", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"-80", "80-"}) defer env.close() vs := &vschemapb.Keyspace{ @@ -2523,7 +2617,7 @@ func TestMaterializerComplexVindexExpression(t *testing.T) { env.tmc.expectVRQuery(200, mzSelectFrozenQuery, &sqltypes.Result{}) env.tmc.expectVRQuery(210, mzSelectFrozenQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.EqualError(t, err, "vindex column cannot be a complex expression: a + b as c1") } @@ -2538,7 +2632,9 @@ func TestMaterializerNoVindexInExpression(t *testing.T) { CreateDdl: "t1ddl", }}, } - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"-80", "80-"}) defer env.close() vs := &vschemapb.Keyspace{ @@ -2564,7 +2660,7 @@ func TestMaterializerNoVindexInExpression(t *testing.T) { env.tmc.expectVRQuery(200, mzSelectFrozenQuery, &sqltypes.Result{}) env.tmc.expectVRQuery(210, mzSelectFrozenQuery, &sqltypes.Result{}) - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.EqualError(t, err, "could not find vindex column c1") } @@ -2778,7 +2874,9 @@ func TestMaterializerManyToManySomeUnreachable(t *testing.T) { for _, tcase := range testcases { t.Run("", func(t *testing.T) { - env := newTestMaterializerEnv(t, ms, tcase.sourceShards, tcase.targetShards) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestMaterializerEnv(t, ctx, ms, tcase.sourceShards, tcase.targetShards) if err := env.topoServ.SaveVSchema(context.Background(), "targetks", vs); err != nil { t.Fatal(err) } @@ -2799,7 +2897,7 @@ func TestMaterializerManyToManySomeUnreachable(t *testing.T) { ) env.tmc.expectVRQuery(tabletID, mzUpdateQuery, &sqltypes.Result{}) } - err := env.wr.Materialize(context.Background(), ms) + err := env.wr.Materialize(ctx, ms) require.NoError(t, err) env.tmc.verifyQueries(t) }) @@ -2825,7 +2923,9 @@ func TestMoveTablesDDLFlag(t *testing.T) { for onDDLAction := range binlogdatapb.OnDDLAction_value { t.Run(fmt.Sprintf("OnDDL Flag:%v", onDDLAction), func(t *testing.T) { - env := newTestMaterializerEnv(t, ms, []string{"0"}, []string{"0"}) + ctx, cancel := context.WithCancel(ctx) + defer cancel() + env := newTestMaterializerEnv(t, ctx, ms, []string{"0"}, []string{"0"}) defer env.close() env.tmc.expectVRQuery(100, mzCheckJournal, &sqltypes.Result{}) @@ -2849,8 +2949,9 @@ func TestMoveTablesDDLFlag(t *testing.T) { } func TestAddTablesToVSchema(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") srcks := "source" wr := &Wrangler{ logger: logutil.NewMemoryLogger(), diff --git a/go/vt/wrangler/resharder_env_test.go b/go/vt/wrangler/resharder_env_test.go index 2add7e6446f..ee39c7e5eaa 100644 --- a/go/vt/wrangler/resharder_env_test.go +++ b/go/vt/wrangler/resharder_env_test.go @@ -82,14 +82,14 @@ func initTopo(t *testing.T, topo *topo.Server, keyspace string, sources, targets topo.ValidateSrvKeyspace(ctx, keyspace, strings.Join(cells, ",")) } -func newTestResharderEnv(t *testing.T, sources, targets []string) *testResharderEnv { +func newTestResharderEnv(t *testing.T, ctx context.Context, sources, targets []string) *testResharderEnv { env := &testResharderEnv{ keyspace: "ks", workflow: "resharderTest", sources: sources, targets: targets, tablets: make(map[int]*topodatapb.Tablet), - topoServ: memorytopo.NewServer("cell"), + topoServ: memorytopo.NewServer(ctx, "cell"), cell: "cell", tmc: newTestResharderTMClient(), } diff --git a/go/vt/wrangler/resharder_test.go b/go/vt/wrangler/resharder_test.go index 9fdb5889a9e..40d31d36e1f 100644 --- a/go/vt/wrangler/resharder_test.go +++ b/go/vt/wrangler/resharder_test.go @@ -17,14 +17,13 @@ limitations under the License. package wrangler import ( + "context" "fmt" "strings" "testing" "github.com/stretchr/testify/require" - "context" - "github.com/stretchr/testify/assert" "vitess.io/vitess/go/sqltypes" @@ -39,7 +38,9 @@ const insertPrefix = `/insert into _vt.vreplication\(workflow, source, pos, max_ const eol = "$" func TestResharderOneToMany(t *testing.T) { - env := newTestResharderEnv(t, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestResharderEnv(t, ctx, []string{"0"}, []string{"-80", "80-"}) defer env.close() schm := &tabletmanagerdatapb.SchemaDefinition{ @@ -73,7 +74,7 @@ func TestResharderOneToMany(t *testing.T) { testCases = append(testCases, newTestCase("", "replica,rdonly")) for _, tc := range testCases { - env := newTestResharderEnv(t, []string{"0"}, []string{"-80", "80-"}) + env := newTestResharderEnv(t, ctx, []string{"0"}, []string{"-80", "80-"}) schm := &tabletmanagerdatapb.SchemaDefinition{ TableDefinitions: []*tabletmanagerdatapb.TableDefinition{{ @@ -115,7 +116,9 @@ func TestResharderOneToMany(t *testing.T) { } func TestResharderManyToOne(t *testing.T) { - env := newTestResharderEnv(t, []string{"-80", "80-"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestResharderEnv(t, ctx, []string{"-80", "80-"}, []string{"0"}) defer env.close() schm := &tabletmanagerdatapb.SchemaDefinition{ @@ -148,7 +151,9 @@ func TestResharderManyToOne(t *testing.T) { } func TestResharderManyToMany(t *testing.T) { - env := newTestResharderEnv(t, []string{"-40", "40-"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestResharderEnv(t, ctx, []string{"-40", "40-"}, []string{"-80", "80-"}) defer env.close() schm := &tabletmanagerdatapb.SchemaDefinition{ @@ -191,7 +196,9 @@ func TestResharderManyToMany(t *testing.T) { // TestResharderOneRefTable tests the case where there's one ref table, but no stream for it. // This means that the table is being updated manually. func TestResharderOneRefTable(t *testing.T) { - env := newTestResharderEnv(t, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestResharderEnv(t, ctx, []string{"0"}, []string{"-80", "80-"}) defer env.close() schm := &tabletmanagerdatapb.SchemaDefinition{ @@ -243,7 +250,9 @@ func TestResharderOneRefTable(t *testing.T) { // TestReshardStopFlags tests the flags -stop_started and -stop_after_copy func TestReshardStopFlags(t *testing.T) { - env := newTestResharderEnv(t, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestResharderEnv(t, ctx, []string{"0"}, []string{"-80", "80-"}) defer env.close() schm := &tabletmanagerdatapb.SchemaDefinition{ @@ -294,7 +303,9 @@ func TestReshardStopFlags(t *testing.T) { // TestResharderOneRefStream tests the case where there's one ref table and an associated stream. func TestResharderOneRefStream(t *testing.T) { - env := newTestResharderEnv(t, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestResharderEnv(t, ctx, []string{"0"}, []string{"-80", "80-"}) defer env.close() schm := &tabletmanagerdatapb.SchemaDefinition{ @@ -362,7 +373,9 @@ func TestResharderOneRefStream(t *testing.T) { // TestResharderNoRefStream tests the case where there's a stream, but it's not a reference. func TestResharderNoRefStream(t *testing.T) { - env := newTestResharderEnv(t, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestResharderEnv(t, ctx, []string{"0"}, []string{"-80", "80-"}) defer env.close() schm := &tabletmanagerdatapb.SchemaDefinition{ @@ -438,7 +451,9 @@ func TestResharderNoRefStream(t *testing.T) { } func TestResharderCopySchema(t *testing.T) { - env := newTestResharderEnv(t, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestResharderEnv(t, ctx, []string{"0"}, []string{"-80", "80-"}) defer env.close() schm := &tabletmanagerdatapb.SchemaDefinition{ @@ -478,7 +493,9 @@ func TestResharderCopySchema(t *testing.T) { } func TestResharderDupWorkflow(t *testing.T) { - env := newTestResharderEnv(t, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestResharderEnv(t, ctx, []string{"0"}, []string{"-80", "80-"}) defer env.close() schm := &tabletmanagerdatapb.SchemaDefinition{ @@ -509,7 +526,9 @@ func TestResharderDupWorkflow(t *testing.T) { } func TestResharderServingState(t *testing.T) { - env := newTestResharderEnv(t, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestResharderEnv(t, ctx, []string{"0"}, []string{"-80", "80-"}) defer env.close() schm := &tabletmanagerdatapb.SchemaDefinition{ @@ -551,7 +570,9 @@ func TestResharderServingState(t *testing.T) { } func TestResharderTargetAlreadyResharding(t *testing.T) { - env := newTestResharderEnv(t, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestResharderEnv(t, ctx, []string{"0"}, []string{"-80", "80-"}) defer env.close() schm := &tabletmanagerdatapb.SchemaDefinition{ @@ -584,7 +605,9 @@ func TestResharderTargetAlreadyResharding(t *testing.T) { } func TestResharderUnnamedStream(t *testing.T) { - env := newTestResharderEnv(t, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestResharderEnv(t, ctx, []string{"0"}, []string{"-80", "80-"}) defer env.close() schm := &tabletmanagerdatapb.SchemaDefinition{ @@ -632,7 +655,9 @@ func TestResharderUnnamedStream(t *testing.T) { } func TestResharderMismatchedRefStreams(t *testing.T) { - env := newTestResharderEnv(t, []string{"-80", "80-"}, []string{"0"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestResharderEnv(t, ctx, []string{"-80", "80-"}, []string{"0"}) defer env.close() schm := &tabletmanagerdatapb.SchemaDefinition{ @@ -699,7 +724,9 @@ func TestResharderMismatchedRefStreams(t *testing.T) { } func TestResharderTableNotInVSchema(t *testing.T) { - env := newTestResharderEnv(t, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestResharderEnv(t, ctx, []string{"0"}, []string{"-80", "80-"}) defer env.close() schm := &tabletmanagerdatapb.SchemaDefinition{ @@ -736,7 +763,9 @@ func TestResharderTableNotInVSchema(t *testing.T) { } func TestResharderMixedTablesOrder1(t *testing.T) { - env := newTestResharderEnv(t, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestResharderEnv(t, ctx, []string{"0"}, []string{"-80", "80-"}) defer env.close() schm := &tabletmanagerdatapb.SchemaDefinition{ @@ -803,7 +832,9 @@ func TestResharderMixedTablesOrder1(t *testing.T) { } func TestResharderMixedTablesOrder2(t *testing.T) { - env := newTestResharderEnv(t, []string{"0"}, []string{"-80", "80-"}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestResharderEnv(t, ctx, []string{"0"}, []string{"-80", "80-"}) defer env.close() schm := &tabletmanagerdatapb.SchemaDefinition{ diff --git a/go/vt/wrangler/tablet_test.go b/go/vt/wrangler/tablet_test.go index 8aa02864dd3..1350b6b574c 100644 --- a/go/vt/wrangler/tablet_test.go +++ b/go/vt/wrangler/tablet_test.go @@ -31,8 +31,11 @@ import ( // shard name to lower case when it's a keyrange, and populates // KeyRange properly. func TestInitTabletShardConversion(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + cell := "cell1" - ts := memorytopo.NewServer(cell) + ts := memorytopo.NewServer(ctx, cell) wr := New(logutil.NewConsoleLogger(), ts, nil) tablet := &topodatapb.Tablet{ @@ -62,8 +65,11 @@ func TestInitTabletShardConversion(t *testing.T) { // TestDeleteTabletBasic tests delete of non-primary tablet func TestDeleteTabletBasic(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + cell := "cell1" - ts := memorytopo.NewServer(cell) + ts := memorytopo.NewServer(ctx, cell) wr := New(logutil.NewConsoleLogger(), ts, nil) tablet := &topodatapb.Tablet{ @@ -91,8 +97,11 @@ func TestDeleteTabletBasic(t *testing.T) { // TestDeleteTabletTruePrimary tests that you can delete a true primary tablet // only if allowPrimary is set to true func TestDeleteTabletTruePrimary(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + cell := "cell1" - ts := memorytopo.NewServer(cell) + ts := memorytopo.NewServer(ctx, cell) wr := New(logutil.NewConsoleLogger(), ts, nil) tablet := &topodatapb.Tablet{ @@ -135,8 +144,11 @@ func TestDeleteTabletTruePrimary(t *testing.T) { // TestDeleteTabletFalsePrimary tests that you can delete a false primary tablet // with allowPrimary set to false func TestDeleteTabletFalsePrimary(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + cell := "cell1" - ts := memorytopo.NewServer(cell) + ts := memorytopo.NewServer(ctx, cell) wr := New(logutil.NewConsoleLogger(), ts, nil) tablet1 := &topodatapb.Tablet{ @@ -184,8 +196,11 @@ func TestDeleteTabletFalsePrimary(t *testing.T) { // TestDeleteTabletShardNonExisting tests that you can delete a true primary // tablet if a shard does not exists anymore. func TestDeleteTabletShardNonExisting(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + cell := "cell1" - ts := memorytopo.NewServer(cell) + ts := memorytopo.NewServer(ctx, cell) wr := New(logutil.NewConsoleLogger(), ts, nil) tablet := &topodatapb.Tablet{ diff --git a/go/vt/wrangler/testlib/backup_test.go b/go/vt/wrangler/testlib/backup_test.go index 2af9984261f..2cad68d71b3 100644 --- a/go/vt/wrangler/testlib/backup_test.go +++ b/go/vt/wrangler/testlib/backup_test.go @@ -85,10 +85,11 @@ func testBackupRestore(t *testing.T, cDetails *compressionDetails) error { discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) // Initialize our environment - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() db := fakesqldb.New(t) defer db.Close() - ts := memorytopo.NewServer("cell1", "cell2") + ts := memorytopo.NewServer(ctx, "cell1", "cell2") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) vp := NewVtctlPipe(t, ts) defer vp.Close() @@ -336,10 +337,11 @@ func TestBackupRestoreLagged(t *testing.T) { discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) // Initialize our environment - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() db := fakesqldb.New(t) defer db.Close() - ts := memorytopo.NewServer("cell1", "cell2") + ts := memorytopo.NewServer(ctx, "cell1", "cell2") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) vp := NewVtctlPipe(t, ts) defer vp.Close() @@ -554,10 +556,11 @@ func TestRestoreUnreachablePrimary(t *testing.T) { discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) // Initialize our environment - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() db := fakesqldb.New(t) defer db.Close() - ts := memorytopo.NewServer("cell1") + ts := memorytopo.NewServer(ctx, "cell1") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) vp := NewVtctlPipe(t, ts) defer vp.Close() @@ -728,10 +731,11 @@ func TestDisableActiveReparents(t *testing.T) { discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) // Initialize our environment - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() db := fakesqldb.New(t) defer db.Close() - ts := memorytopo.NewServer("cell1", "cell2") + ts := memorytopo.NewServer(ctx, "cell1", "cell2") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) vp := NewVtctlPipe(t, ts) defer vp.Close() diff --git a/go/vt/wrangler/testlib/copy_schema_shard_test.go b/go/vt/wrangler/testlib/copy_schema_shard_test.go index af80203ea9e..866ec2fe931 100644 --- a/go/vt/wrangler/testlib/copy_schema_shard_test.go +++ b/go/vt/wrangler/testlib/copy_schema_shard_test.go @@ -53,7 +53,9 @@ func copySchema(t *testing.T, useShardAsSource bool) { }() discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) - ts := memorytopo.NewServer("cell1", "cell2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) vp := NewVtctlPipe(t, ts) defer vp.Close() diff --git a/go/vt/wrangler/testlib/emergency_reparent_shard_test.go b/go/vt/wrangler/testlib/emergency_reparent_shard_test.go index 3ef9513e1ed..99cc1839186 100644 --- a/go/vt/wrangler/testlib/emergency_reparent_shard_test.go +++ b/go/vt/wrangler/testlib/emergency_reparent_shard_test.go @@ -47,7 +47,9 @@ func TestEmergencyReparentShard(t *testing.T) { }() discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) - ts := memorytopo.NewServer("cell1", "cell2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) vp := NewVtctlPipe(t, ts) defer vp.Close() @@ -201,7 +203,7 @@ func TestEmergencyReparentShardPrimaryElectNotBest(t *testing.T) { }() discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) - ts := memorytopo.NewServer("cell1", "cell2") + ts := memorytopo.NewServer(ctx, "cell1", "cell2") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) // Create a primary, a couple good replicas diff --git a/go/vt/wrangler/testlib/external_reparent_test.go b/go/vt/wrangler/testlib/external_reparent_test.go index c105d1dfd94..c0152de3cf3 100644 --- a/go/vt/wrangler/testlib/external_reparent_test.go +++ b/go/vt/wrangler/testlib/external_reparent_test.go @@ -47,8 +47,9 @@ func TestTabletExternallyReparentedBasic(t *testing.T) { }() discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithTimeout(context.Background(), time.Second*30) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) vp := NewVtctlPipe(t, ts) defer vp.Close() @@ -139,8 +140,9 @@ func TestTabletExternallyReparentedToReplica(t *testing.T) { }() discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithTimeout(context.Background(), time.Second*30) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) // Create an old primary, a new primary, two good replicas, one bad replica @@ -221,8 +223,9 @@ func TestTabletExternallyReparentedWithDifferentMysqlPort(t *testing.T) { }() discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithTimeout(context.Background(), time.Second*30) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) // Create an old primary, a new primary, two good replicas, one bad replica @@ -313,8 +316,9 @@ func TestTabletExternallyReparentedContinueOnUnexpectedPrimary(t *testing.T) { }() discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithTimeout(context.Background(), time.Second*30) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) // Create an old primary, a new primary, two good replicas, one bad replica @@ -398,8 +402,9 @@ func TestTabletExternallyReparentedRerun(t *testing.T) { }() discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithTimeout(context.Background(), time.Second*30) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) // Create an old primary, a new primary, and a good replica. @@ -501,8 +506,9 @@ func TestRPCTabletExternallyReparentedDemotesPrimaryToConfiguredTabletType(t *te flag.Set("disable_active_reparents", "true") defer flag.Set("disable_active_reparents", "false") - ctx := context.Background() - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithTimeout(context.Background(), time.Second*30) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) // Create an old primary and a new primary diff --git a/go/vt/wrangler/testlib/find_tablet_test.go b/go/vt/wrangler/testlib/find_tablet_test.go index 783e6258751..5b6f26f7056 100644 --- a/go/vt/wrangler/testlib/find_tablet_test.go +++ b/go/vt/wrangler/testlib/find_tablet_test.go @@ -17,9 +17,9 @@ limitations under the License. package testlib import ( - "testing" - "context" + "testing" + "time" "vitess.io/vitess/go/vt/logutil" "vitess.io/vitess/go/vt/topo" @@ -33,8 +33,9 @@ import ( ) func TestFindTablet(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("cell1", "cell2") + ctx, cancel := context.WithTimeout(context.Background(), time.Second*30) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) // Create an old primary, two good replicas diff --git a/go/vt/wrangler/testlib/permissions_test.go b/go/vt/wrangler/testlib/permissions_test.go index 9d2a950c8e3..4a0e71512f3 100644 --- a/go/vt/wrangler/testlib/permissions_test.go +++ b/go/vt/wrangler/testlib/permissions_test.go @@ -44,8 +44,9 @@ func TestPermissions(t *testing.T) { discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) // Initialize our environment - ctx := context.Background() - ts := memorytopo.NewServer("cell1", "cell2") + ctx, cancel := context.WithTimeout(context.Background(), time.Second*30) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) vp := NewVtctlPipe(t, ts) defer vp.Close() diff --git a/go/vt/wrangler/testlib/planned_reparent_shard_test.go b/go/vt/wrangler/testlib/planned_reparent_shard_test.go index e123727b4fb..0125e69cac0 100644 --- a/go/vt/wrangler/testlib/planned_reparent_shard_test.go +++ b/go/vt/wrangler/testlib/planned_reparent_shard_test.go @@ -48,7 +48,9 @@ func TestPlannedReparentShardNoPrimaryProvided(t *testing.T) { }() discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) - ts := memorytopo.NewServer("cell1", "cell2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) vp := NewVtctlPipe(t, ts) defer vp.Close() @@ -162,7 +164,9 @@ func TestPlannedReparentShardNoError(t *testing.T) { }() discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) - ts := memorytopo.NewServer("cell1", "cell2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) vp := NewVtctlPipe(t, ts) defer vp.Close() @@ -296,7 +300,9 @@ func TestPlannedReparentInitialization(t *testing.T) { }() discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) - ts := memorytopo.NewServer("cell1", "cell2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) vp := NewVtctlPipe(t, ts) defer vp.Close() @@ -380,7 +386,9 @@ func TestPlannedReparentShardWaitForPositionFail(t *testing.T) { }() discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) - ts := memorytopo.NewServer("cell1", "cell2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) vp := NewVtctlPipe(t, ts) defer vp.Close() @@ -486,7 +494,9 @@ func TestPlannedReparentShardWaitForPositionTimeout(t *testing.T) { }() discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) - ts := memorytopo.NewServer("cell1", "cell2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) vp := NewVtctlPipe(t, ts) defer vp.Close() @@ -590,7 +600,9 @@ func TestPlannedReparentShardRelayLogError(t *testing.T) { }() discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) vp := NewVtctlPipe(t, ts) defer vp.Close() @@ -668,7 +680,9 @@ func TestPlannedReparentShardRelayLogErrorStartReplication(t *testing.T) { }() discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) - ts := memorytopo.NewServer("cell1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) vp := NewVtctlPipe(t, ts) defer vp.Close() @@ -751,7 +765,9 @@ func TestPlannedReparentShardPromoteReplicaFail(t *testing.T) { }() discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) - ts := memorytopo.NewServer("cell1", "cell2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) vp := NewVtctlPipe(t, ts) defer vp.Close() @@ -889,7 +905,9 @@ func TestPlannedReparentShardSamePrimary(t *testing.T) { }() discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) - ts := memorytopo.NewServer("cell1", "cell2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) vp := NewVtctlPipe(t, ts) defer vp.Close() diff --git a/go/vt/wrangler/testlib/reparent_utils_test.go b/go/vt/wrangler/testlib/reparent_utils_test.go index d667c7ca9df..0d1d84e89f5 100644 --- a/go/vt/wrangler/testlib/reparent_utils_test.go +++ b/go/vt/wrangler/testlib/reparent_utils_test.go @@ -48,8 +48,9 @@ func TestShardReplicationStatuses(t *testing.T) { }() discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) - ctx := context.Background() - ts := memorytopo.NewServer("cell1", "cell2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) // create shard and tablets @@ -131,8 +132,9 @@ func TestReparentTablet(t *testing.T) { }() discovery.SetTabletPickerRetryDelay(5 * time.Millisecond) - ctx := context.Background() - ts := memorytopo.NewServer("cell1", "cell2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) // create shard and tablets @@ -187,8 +189,9 @@ func TestReparentTablet(t *testing.T) { // TestSetReplicationSource tests that SetReplicationSource works as intended under various circumstances. func TestSetReplicationSource(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("cell1", "cell2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) // create shard and tablets diff --git a/go/vt/wrangler/testlib/shard_test.go b/go/vt/wrangler/testlib/shard_test.go index 1b00ad82769..a0b1b0a3562 100644 --- a/go/vt/wrangler/testlib/shard_test.go +++ b/go/vt/wrangler/testlib/shard_test.go @@ -17,11 +17,10 @@ limitations under the License. package testlib import ( + "context" "strings" "testing" - "context" - "vitess.io/vitess/go/vt/logutil" "vitess.io/vitess/go/vt/topo" "vitess.io/vitess/go/vt/topo/memorytopo" @@ -33,8 +32,9 @@ import ( ) func TestDeleteShardCleanup(t *testing.T) { - ctx := context.Background() - ts := memorytopo.NewServer("cell1", "cell2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) vp := NewVtctlPipe(t, ts) defer vp.Close() diff --git a/go/vt/wrangler/testlib/version_test.go b/go/vt/wrangler/testlib/version_test.go index c54a0811948..634f48e4d7e 100644 --- a/go/vt/wrangler/testlib/version_test.go +++ b/go/vt/wrangler/testlib/version_test.go @@ -17,6 +17,7 @@ limitations under the License. package testlib import ( + "context" "encoding/json" "fmt" "net/http" @@ -70,7 +71,9 @@ func TestVersion(t *testing.T) { wrangler.ResetDebugVarsGetVersion() // Initialize our environment - ts := memorytopo.NewServer("cell1", "cell2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "cell1", "cell2") wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient()) vp := NewVtctlPipe(t, ts) defer vp.Close() diff --git a/go/vt/wrangler/testlib/vtctl_topo_test.go b/go/vt/wrangler/testlib/vtctl_topo_test.go index bd3e4134d62..a13535f4111 100644 --- a/go/vt/wrangler/testlib/vtctl_topo_test.go +++ b/go/vt/wrangler/testlib/vtctl_topo_test.go @@ -52,7 +52,10 @@ func testVtctlTopoCommand(t *testing.T, vp *VtctlPipe, args []string, want strin // TestVtctlTopoCommands tests all vtctl commands from the // "Topo" group. func TestVtctlTopoCommands(t *testing.T) { - ts := memorytopo.NewServer("cell1", "cell2") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ts := memorytopo.NewServer(ctx, "cell1", "cell2") if err := ts.CreateKeyspace(context.Background(), "ks1", &topodatapb.Keyspace{KeyspaceType: topodatapb.KeyspaceType_NORMAL}); err != nil { t.Fatalf("CreateKeyspace() failed: %v", err) } diff --git a/go/vt/wrangler/traffic_switcher_env_test.go b/go/vt/wrangler/traffic_switcher_env_test.go index 9445e849d2b..fb6d2c9fc58 100644 --- a/go/vt/wrangler/traffic_switcher_env_test.go +++ b/go/vt/wrangler/traffic_switcher_env_test.go @@ -119,7 +119,7 @@ func newTestTableMigrater(ctx context.Context, t *testing.T) *testMigraterEnv { // The test will Sprintf a from clause and where clause as needed. func newTestTableMigraterCustom(ctx context.Context, t *testing.T, sourceShards, targetShards []string, fmtQuery string) *testMigraterEnv { tme := &testMigraterEnv{} - tme.ts = memorytopo.NewServer("cell1", "cell2") + tme.ts = memorytopo.NewServer(ctx, "cell1", "cell2") tme.wr = New(logutil.NewConsoleLogger(), tme.ts, tmclient.NewTabletManagerClient()) tme.wr.sem = semaphore.NewWeighted(1) tme.sourceShards = sourceShards @@ -383,7 +383,7 @@ func newTestTableMigraterCustom(ctx context.Context, t *testing.T, sourceShards, func newTestTablePartialMigrater(ctx context.Context, t *testing.T, shards, shardsToMove []string, fmtQuery string) *testMigraterEnv { require.Greater(t, len(shards), 1, "shard by shard migrations can only be done on sharded keyspaces") tme := &testMigraterEnv{} - tme.ts = memorytopo.NewServer("cell1", "cell2") + tme.ts = memorytopo.NewServer(ctx, "cell1", "cell2") tme.wr = New(logutil.NewConsoleLogger(), tme.ts, tmclient.NewTabletManagerClient()) tme.wr.sem = semaphore.NewWeighted(1) tme.sourceShards = shards @@ -539,7 +539,7 @@ func newTestTablePartialMigrater(ctx context.Context, t *testing.T, shards, shar func newTestShardMigrater(ctx context.Context, t *testing.T, sourceShards, targetShards []string) *testShardMigraterEnv { tme := &testShardMigraterEnv{} - tme.ts = memorytopo.NewServer("cell1", "cell2") + tme.ts = memorytopo.NewServer(ctx, "cell1", "cell2") tme.wr = New(logutil.NewConsoleLogger(), tme.ts, tmclient.NewTabletManagerClient()) tme.sourceShards = sourceShards tme.targetShards = targetShards diff --git a/go/vt/wrangler/vdiff_env_test.go b/go/vt/wrangler/vdiff_env_test.go index 1fe06456aed..01f3a3a0f9e 100644 --- a/go/vt/wrangler/vdiff_env_test.go +++ b/go/vt/wrangler/vdiff_env_test.go @@ -69,11 +69,11 @@ type testVDiffEnv struct { //---------------------------------------------- // testVDiffEnv -func newTestVDiffEnv(t testing.TB, sourceShards, targetShards []string, query string, positions map[string]string) *testVDiffEnv { +func newTestVDiffEnv(t testing.TB, ctx context.Context, sourceShards, targetShards []string, query string, positions map[string]string) *testVDiffEnv { env := &testVDiffEnv{ workflow: "vdiffTest", tablets: make(map[int]*testVDiffTablet), - topoServ: memorytopo.NewServer("cell"), + topoServ: memorytopo.NewServer(ctx, "cell"), cell: "cell", tabletType: topodatapb.TabletType_REPLICA, tmc: newTestVDiffTMClient(), diff --git a/go/vt/wrangler/vdiff_test.go b/go/vt/wrangler/vdiff_test.go index b9ca2d0dbf2..df38bb5c752 100644 --- a/go/vt/wrangler/vdiff_test.go +++ b/go/vt/wrangler/vdiff_test.go @@ -493,7 +493,9 @@ func TestVDiffPlanFailure(t *testing.T) { } func TestVDiffUnsharded(t *testing.T) { - env := newTestVDiffEnv(t, []string{"0"}, []string{"0"}, "", nil) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestVDiffEnv(t, ctx, []string{"0"}, []string{"0"}, "", nil) defer env.close() schm := &tabletmanagerdatapb.SchemaDefinition{ @@ -771,7 +773,9 @@ func TestVDiffUnsharded(t *testing.T) { func TestVDiffSharded(t *testing.T) { // Also test that highest position ""MariaDB/5-456-892" will be used // if lower positions are found. - env := newTestVDiffEnv(t, []string{"-40", "40-"}, []string{"-80", "80-"}, "", map[string]string{ + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestVDiffEnv(t, ctx, []string{"-40", "40-"}, []string{"-80", "80-"}, "", map[string]string{ "-40-80": "MariaDB/5-456-890", "40-80-": "MariaDB/5-456-891", }) @@ -844,7 +848,9 @@ func TestVDiffSharded(t *testing.T) { } func TestVDiffAggregates(t *testing.T) { - env := newTestVDiffEnv(t, []string{"-40", "40-"}, []string{"-80", "80-"}, "select c1, count(*) c2, sum(c3) c3 from t group by c1", nil) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestVDiffEnv(t, ctx, []string{"-40", "40-"}, []string{"-80", "80-"}, "select c1, count(*) c2, sum(c3) c3 from t group by c1", nil) defer env.close() schm := &tabletmanagerdatapb.SchemaDefinition{ @@ -912,7 +918,9 @@ func TestVDiffAggregates(t *testing.T) { } func TestVDiffDefaults(t *testing.T) { - env := newTestVDiffEnv(t, []string{"0"}, []string{"0"}, "", nil) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestVDiffEnv(t, ctx, []string{"0"}, []string{"0"}, "", nil) defer env.close() schm := &tabletmanagerdatapb.SchemaDefinition{ @@ -966,7 +974,9 @@ func TestVDiffDefaults(t *testing.T) { } func TestVDiffReplicationWait(t *testing.T) { - env := newTestVDiffEnv(t, []string{"0"}, []string{"0"}, "", nil) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + env := newTestVDiffEnv(t, ctx, []string{"0"}, []string{"0"}, "", nil) defer env.close() schm := &tabletmanagerdatapb.SchemaDefinition{ diff --git a/go/vt/wrangler/vexec_test.go b/go/vt/wrangler/vexec_test.go index 031db6c7a08..ead2be6a56f 100644 --- a/go/vt/wrangler/vexec_test.go +++ b/go/vt/wrangler/vexec_test.go @@ -36,11 +36,12 @@ import ( ) func TestVExec(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() workflow := "wrWorkflow" keyspace := "target" query := "update _vt.vreplication set state = 'Running'" - env := newWranglerTestEnv(t, []string{"0"}, []string{"-80", "80-"}, "", nil, time.Now().Unix()) + env := newWranglerTestEnv(t, ctx, []string{"0"}, []string{"-80", "80-"}, "", nil, time.Now().Unix()) defer env.close() var logger = logutil.NewMemoryLogger() wr := New(logger, env.topoServ, env.tmc) @@ -181,10 +182,11 @@ func TestWorkflowStatusUpdate(t *testing.T) { } func TestWorkflowListStreams(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() workflow := "wrWorkflow" keyspace := "target" - env := newWranglerTestEnv(t, []string{"0"}, []string{"-80", "80-"}, "", nil, 1234) + env := newWranglerTestEnv(t, ctx, []string{"0"}, []string{"-80", "80-"}, "", nil, 1234) defer env.close() logger := logutil.NewMemoryLogger() wr := New(logger, env.topoServ, env.tmc) @@ -356,10 +358,11 @@ will be run on the following streams in keyspace target for workflow wrWorkflow: } func TestWorkflowListAll(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() keyspace := "target" workflow := "wrWorkflow" - env := newWranglerTestEnv(t, []string{"0"}, []string{"-80", "80-"}, "", nil, 0) + env := newWranglerTestEnv(t, ctx, []string{"0"}, []string{"-80", "80-"}, "", nil, 0) defer env.close() logger := logutil.NewMemoryLogger() wr := New(logger, env.topoServ, env.tmc) @@ -375,11 +378,12 @@ func TestWorkflowListAll(t *testing.T) { } func TestVExecValidations(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() workflow := "wf" keyspace := "ks" query := "" - env := newWranglerTestEnv(t, []string{"0"}, []string{"-80", "80-"}, "", nil, 0) + env := newWranglerTestEnv(t, ctx, []string{"0"}, []string{"-80", "80-"}, "", nil, 0) defer env.close() wr := New(logutil.NewConsoleLogger(), env.topoServ, env.tmc) @@ -461,10 +465,11 @@ func TestVExecValidations(t *testing.T) { // tabletmanager and the behavior is tested // there. func TestWorkflowUpdate(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() workflow := "wrWorkflow" keyspace := "target" - env := newWranglerTestEnv(t, []string{"0"}, []string{"-80", "80-"}, "", nil, 1234) + env := newWranglerTestEnv(t, ctx, []string{"0"}, []string{"-80", "80-"}, "", nil, 1234) defer env.close() logger := logutil.NewMemoryLogger() wr := New(logger, env.topoServ, env.tmc) diff --git a/go/vt/wrangler/wrangler_env_test.go b/go/vt/wrangler/wrangler_env_test.go index 23585fb1651..4dd5e342c35 100644 --- a/go/vt/wrangler/wrangler_env_test.go +++ b/go/vt/wrangler/wrangler_env_test.go @@ -60,10 +60,10 @@ type testWranglerEnv struct { //---------------------------------------------- // testWranglerEnv -func newWranglerTestEnv(t testing.TB, sourceShards, targetShards []string, query string, positions map[string]string, timeUpdated int64) *testWranglerEnv { +func newWranglerTestEnv(t testing.TB, ctx context.Context, sourceShards, targetShards []string, query string, positions map[string]string, timeUpdated int64) *testWranglerEnv { env := &testWranglerEnv{ workflow: "wrWorkflow", - topoServ: memorytopo.NewServer("zone1"), + topoServ: memorytopo.NewServer(ctx, "zone1"), cell: "zone1", tabletType: topodatapb.TabletType_REPLICA, tmc: newTestWranglerTMClient(), From f34e71fc7684d5e315e9b8779ec40cc4220752b8 Mon Sep 17 00:00:00 2001 From: Dirkjan Bussink Date: Tue, 22 Aug 2023 22:20:46 +0200 Subject: [PATCH 04/17] vtgate: Remove globals This removes a number of globals from vtgate so that it can be tested with the new leak checking enabled. Before the globals would leak things and globals have problematic side effects anyway. Signed-off-by: Dirkjan Bussink --- go/cache/ristretto/cache.go | 20 +- go/test/utils/noleak.go | 1 + go/vt/logutil/throttled.go | 2 +- go/vt/topo/helpers/tee_topo_test.go | 2 +- go/vt/vtexplain/vtexplain_vtgate.go | 3 +- go/vt/vtgate/bench_test.go | 106 ++++++++ go/vt/vtgate/executor.go | 4 +- go/vt/vtgate/executor_framework_test.go | 23 +- go/vt/vtgate/executor_select_test.go | 69 +++-- go/vt/vtgate/executor_stream_test.go | 6 +- go/vt/vtgate/executor_test.go | 2 +- go/vt/vtgate/plugin_mysql_server.go | 117 ++++----- go/vt/vtgate/plugin_mysql_server_test.go | 11 +- go/vt/vtgate/querylog.go | 5 +- go/vt/vtgate/queryz.go | 6 +- go/vt/vtgate/queryz_test.go | 2 +- go/vt/vtgate/sandbox_test.go | 4 + go/vt/vtgate/vtgate.go | 91 ++++--- go/vt/vtgate/vtgate_test.go | 311 +++++++++++++---------- 19 files changed, 494 insertions(+), 291 deletions(-) create mode 100644 go/vt/vtgate/bench_test.go diff --git a/go/cache/ristretto/cache.go b/go/cache/ristretto/cache.go index b745d6dc991..aa6aa2c2870 100644 --- a/go/cache/ristretto/cache.go +++ b/go/cache/ristretto/cache.go @@ -75,7 +75,7 @@ type Cache struct { // stop is used to stop the processItems goroutine. stop chan struct{} // indicates whether cache is closed. - isClosed bool + isClosed atomic.Bool // cost calculates cost from a value. cost func(value any) int64 // ignoreInternalCost dictates whether to ignore the cost of internally storing @@ -211,7 +211,7 @@ func NewCache(config *Config) (*Cache, error) { // Wait blocks until all the current cache operations have been processed in the background func (c *Cache) Wait() { - if c == nil || c.isClosed { + if c == nil || c.isClosed.Load() { return } wg := &sync.WaitGroup{} @@ -224,7 +224,7 @@ func (c *Cache) Wait() { // value was found or not. The value can be nil and the boolean can be true at // the same time. func (c *Cache) Get(key string) (any, bool) { - if c == nil || c.isClosed { + if c == nil || c.isClosed.Load() { return nil, false } keyHash, conflictHash := c.keyToHash(key) @@ -253,7 +253,7 @@ func (c *Cache) Set(key string, value any) bool { // cost. The built-in Cost function will not be called to evaluate the object's cost // and instead the given value will be used. func (c *Cache) SetWithCost(key string, value any, cost int64) bool { - if c == nil || c.isClosed { + if c == nil || c.isClosed.Load() { return false } @@ -289,7 +289,7 @@ func (c *Cache) SetWithCost(key string, value any, cost int64) bool { // Delete deletes the key-value item from the cache if it exists. func (c *Cache) Delete(key string) { - if c == nil || c.isClosed { + if c == nil || c.isClosed.Load() { return } keyHash, conflictHash := c.keyToHash(key) @@ -309,7 +309,11 @@ func (c *Cache) Delete(key string) { // Close stops all goroutines and closes all channels. func (c *Cache) Close() { - if c == nil || c.isClosed { + if c == nil { + return + } + wasClosed := c.isClosed.Swap(true) + if wasClosed { return } c.Clear() @@ -319,14 +323,14 @@ func (c *Cache) Close() { close(c.stop) close(c.setBuf) c.policy.Close() - c.isClosed = true + c.isClosed.Store(true) } // Clear empties the hashmap and zeroes all policy counters. Note that this is // not an atomic operation (but that shouldn't be a problem as it's assumed that // Set/Get calls won't be occurring until after this). func (c *Cache) Clear() { - if c == nil || c.isClosed { + if c == nil || c.isClosed.Load() { return } // Block until processItems goroutine is returned. diff --git a/go/test/utils/noleak.go b/go/test/utils/noleak.go index b24fb585579..0f0258a5d9b 100644 --- a/go/test/utils/noleak.go +++ b/go/test/utils/noleak.go @@ -27,6 +27,7 @@ func ensureNoGoroutines(t testing.TB) { goleak.IgnoreTopFunction("vitess.io/vitess/go/vt/vtgate.resetAggregators"), goleak.IgnoreTopFunction("vitess.io/vitess/go/vt/vtgate.processQueryInfo"), goleak.IgnoreTopFunction("github.com/patrickmn/go-cache.(*janitor).Run"), + goleak.IgnoreTopFunction("vitess.io/vitess/go/vt/logutil.(*ThrottledLogger).log.func1"), } var err error diff --git a/go/vt/logutil/throttled.go b/go/vt/logutil/throttled.go index 917798626bb..4ee11912e71 100644 --- a/go/vt/logutil/throttled.go +++ b/go/vt/logutil/throttled.go @@ -69,7 +69,7 @@ func (tl *ThrottledLogger) log(logF logFunc, format string, v ...any) { // to log and reset skippedCount if tl.skippedCount == 0 { go func(d time.Duration) { - time.Sleep(d) + <-time.After(d) tl.mu.Lock() defer tl.mu.Unlock() // Because of the go func(), we lose the stack trace, diff --git a/go/vt/topo/helpers/tee_topo_test.go b/go/vt/topo/helpers/tee_topo_test.go index 41e220e0a56..b1f1d379a5a 100644 --- a/go/vt/topo/helpers/tee_topo_test.go +++ b/go/vt/topo/helpers/tee_topo_test.go @@ -27,7 +27,7 @@ import ( ) func TestTeeTopo(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() test.TopoServerTestSuite(t, ctx, func() *topo.Server { diff --git a/go/vt/vtexplain/vtexplain_vtgate.go b/go/vt/vtexplain/vtexplain_vtgate.go index 6c9d60466ab..d7b3c3accf7 100644 --- a/go/vt/vtexplain/vtexplain_vtgate.go +++ b/go/vt/vtexplain/vtexplain_vtgate.go @@ -74,7 +74,8 @@ func (vte *VTExplain) initVtgateExecutor(ctx context.Context, vSchemaStr, ksShar streamSize := 10 var schemaTracker vtgate.SchemaInfo // no schema tracker for these tests queryLogBufferSize := 10 - vte.vtgateExecutor = vtgate.NewExecutor(ctx, vte.explainTopo, vtexplainCell, resolver, opts.Normalize, false, streamSize, cache.DefaultConfig, schemaTracker, false, opts.PlannerVersion, streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize)) + plans := cache.NewDefaultCacheImpl(cache.DefaultConfig) + vte.vtgateExecutor = vtgate.NewExecutor(ctx, vte.explainTopo, vtexplainCell, resolver, opts.Normalize, false, streamSize, plans, schemaTracker, false, opts.PlannerVersion, streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize)) return nil } diff --git a/go/vt/vtgate/bench_test.go b/go/vt/vtgate/bench_test.go new file mode 100644 index 00000000000..423407270b8 --- /dev/null +++ b/go/vt/vtgate/bench_test.go @@ -0,0 +1,106 @@ +/* +Copyright 2019 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 vtgate + +import ( + "bytes" + "context" + "fmt" + "testing" + + vtgatepb "vitess.io/vitess/go/vt/proto/vtgate" +) + +// Benchmark run on 6/27/17, with optimized byte-level operations +// using bytes2.Buffer: +// BenchmarkWithNormalizer-4 300 5694660 ns/op + +// Benchmark run on 6/25/17, prior to improvements: +// BenchmarkWithNormalizer-4 100 10629161 ns/op +// BenchmarkWithoutNormalizer-4 200000 8584 ns/op + +var benchQuery string + +func init() { + // benchQuerySize is the approximate size of the query. + benchQuerySize := 1000000 + + // Size of value is 1/10 size of query. Then we add + // 10 such values to the where clause. + baseval := &bytes.Buffer{} + for i := 0; i < benchQuerySize/100; i++ { + // Add an escape character: This will force the upcoming + // tokenizer improvement to still create a copy of the string. + // Then we can see if avoiding the copy will be worth it. + baseval.WriteString("\\'123456789") + } + + buf := &bytes.Buffer{} + buf.WriteString("select a from t1 where v = 1") + for i := 0; i < 10; i++ { + fmt.Fprintf(buf, " and v%d = '%d%s'", i, i, baseval.String()) + } + benchQuery = buf.String() + // fmt.Printf("len: %d\n", len(benchQuery)) +} + +func BenchmarkWithNormalizer(b *testing.B) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + vtgateInst, _ := createVtgateEnv(ctx) + defer vtgateInst.Close() + + for i := 0; i < b.N; i++ { + _, _, err := vtgateInst.Execute( + ctx, + nil, + &vtgatepb.Session{ + TargetString: "@primary", + Options: executeOptions, + }, + benchQuery, + nil, + ) + if err != nil { + panic(err) + } + } +} + +func BenchmarkWithoutNormalizer(b *testing.B) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + vtgateInst, _ := createVtgateEnv(ctx) + vtgateInst.executor.normalize = false + defer vtgateInst.Close() + + for i := 0; i < b.N; i++ { + _, _, err := vtgateInst.Execute( + ctx, + nil, + &vtgatepb.Session{ + TargetString: "@primary", + Options: executeOptions, + }, + benchQuery, + nil, + ) + if err != nil { + panic(err) + } + } +} diff --git a/go/vt/vtgate/executor.go b/go/vt/vtgate/executor.go index 8c2f41888d9..971967c0427 100644 --- a/go/vt/vtgate/executor.go +++ b/go/vt/vtgate/executor.go @@ -139,7 +139,7 @@ func NewExecutor( resolver *Resolver, normalize, warnOnShardedOnly bool, streamSize int, - cacheCfg *cache.Config, + plans cache.Cache, schemaTracker SchemaInfo, noScatter bool, pv plancontext.PlannerVersion, @@ -151,7 +151,7 @@ func NewExecutor( resolver: resolver, scatterConn: resolver.scatterConn, txConn: resolver.scatterConn.txConn, - plans: cache.NewDefaultCacheImpl(cacheCfg), + plans: plans, normalize: normalize, warnShardedOnly: warnOnShardedOnly, streamSize: streamSize, diff --git a/go/vt/vtgate/executor_framework_test.go b/go/vt/vtgate/executor_framework_test.go index 4c1fe85fe59..d3b896558e3 100644 --- a/go/vt/vtgate/executor_framework_test.go +++ b/go/vt/vtgate/executor_framework_test.go @@ -131,7 +131,7 @@ func createExecutorEnv(ctx context.Context) (executor *Executor, sbc1, sbc2, sbc s := createSandbox(KsTestSharded) s.VSchema = executorVSchema serv := newSandboxForCells(ctx, []string{cell}) - serv.topoServer.CreateKeyspace(ctx, "TestExecutor", &topodatapb.Keyspace{SidecarDbName: sidecar.DefaultName}) + serv.topoServer.CreateKeyspace(ctx, KsTestSharded, &topodatapb.Keyspace{SidecarDbName: sidecar.DefaultName}) // Force a new cache to use for lookups of the sidecar database identifier // in use by each keyspace -- as we want to use a different load function // than the one already created by the vtgate as it uses a different topo. @@ -160,7 +160,7 @@ func createExecutorEnv(ctx context.Context) (executor *Executor, sbc1, sbc2, sbc _ = hc.AddTestTablet(cell, "e0-", 1, "TestExecutor", "e0-", topodatapb.TabletType_PRIMARY, true, 1, nil) // Below is needed so that SendAnyWherePlan doesn't fail - createSandbox(KsTestUnsharded) + sb := createSandbox(KsTestUnsharded) sbclookup = hc.AddTestTablet(cell, "0", 1, KsTestUnsharded, "0", topodatapb.TabletType_PRIMARY, true, 1, nil) _ = hc.AddTestTablet(cell, "2", 3, KsTestUnsharded, "0", topodatapb.TabletType_REPLICA, true, 1, nil) @@ -169,9 +169,10 @@ func createExecutorEnv(ctx context.Context) (executor *Executor, sbc1, sbc2, sbc bad := createSandbox("TestXBadSharding") bad.VSchema = badVSchema - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + sb.VSchema = unshardedVSchema queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize) - executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) + plans := cache.NewDefaultCacheImpl(cache.DefaultConfig) + executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) key.AnyShardPicker = DestinationAnyShardPickerFirstShard{} return executor, sbc1, sbc2, sbclookup @@ -187,12 +188,13 @@ func createCustomExecutor(ctx context.Context, vschema string) (executor *Execut sbc1 = hc.AddTestTablet(cell, "-20", 1, "TestExecutor", "-20", topodatapb.TabletType_PRIMARY, true, 1, nil) sbc2 = hc.AddTestTablet(cell, "40-60", 1, "TestExecutor", "40-60", topodatapb.TabletType_PRIMARY, true, 1, nil) - createSandbox(KsTestUnsharded) + sb := createSandbox(KsTestUnsharded) sbclookup = hc.AddTestTablet(cell, "0", 1, KsTestUnsharded, "0", topodatapb.TabletType_PRIMARY, true, 1, nil) - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + sb.VSchema = unshardedVSchema queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize) - executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) + plans := cache.NewDefaultCacheImpl(cache.DefaultConfig) + executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) return executor, sbc1, sbc2, sbclookup } @@ -213,12 +215,13 @@ func createCustomExecutorSetValues(ctx context.Context, vschema string, values [ sbcs = append(sbcs, sbc) } - createSandbox(KsTestUnsharded) + sb := createSandbox(KsTestUnsharded) sbclookup = hc.AddTestTablet(cell, "0", 1, KsTestUnsharded, "0", topodatapb.TabletType_PRIMARY, true, 1, nil) - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + sb.VSchema = unshardedVSchema queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize) - executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) + plans := cache.NewDefaultCacheImpl(cache.DefaultConfig) + executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) return executor, sbcs[0], sbcs[1], sbclookup } diff --git a/go/vt/vtgate/executor_select_test.go b/go/vt/vtgate/executor_select_test.go index 0c64f1aabd0..23202637f04 100644 --- a/go/vt/vtgate/executor_select_test.go +++ b/go/vt/vtgate/executor_select_test.go @@ -1709,7 +1709,8 @@ func TestStreamSelectIN(t *testing.T) { func createExecutor(ctx context.Context, serv *sandboxTopo, cell string, resolver *Resolver) *Executor { queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize) - return NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) + plans := cache.NewDefaultCacheImpl(cache.DefaultConfig) + return NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) } func TestSelectScatter(t *testing.T) { @@ -1720,9 +1721,10 @@ func TestSelectScatter(t *testing.T) { // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) + u := createSandbox(KsTestUnsharded) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + u.VSchema = unshardedVSchema serv := newSandboxForCells(ctx, []string{cell}) resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} @@ -1763,9 +1765,10 @@ func TestSelectScatterPartial(t *testing.T) { } cell := "aa" hc := discovery.NewFakeHealthCheck(nil) + u := createSandbox(KsTestUnsharded) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + u.VSchema = unshardedVSchema serv := newSandboxForCells(ctx, []string{cell}) resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} @@ -1828,9 +1831,10 @@ func TestSelectScatterPartialOLAP(t *testing.T) { // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) + u := createSandbox(KsTestUnsharded) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + u.VSchema = unshardedVSchema serv := newSandboxForCells(ctx, []string{cell}) resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} @@ -1884,9 +1888,10 @@ func TestSelectScatterPartialOLAP2(t *testing.T) { // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) + u := createSandbox(KsTestUnsharded) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + u.VSchema = unshardedVSchema serv := newSandboxForCells(ctx, []string{cell}) resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} @@ -1945,9 +1950,10 @@ func TestStreamSelectScatter(t *testing.T) { // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) + u := createSandbox(KsTestUnsharded) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + u.VSchema = unshardedVSchema serv := newSandboxForCells(ctx, []string{cell}) resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} @@ -1985,9 +1991,10 @@ func TestSelectScatterOrderBy(t *testing.T) { // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) + u := createSandbox(KsTestUnsharded) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + u.VSchema = unshardedVSchema serv := newSandboxForCells(ctx, []string{cell}) resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} @@ -2059,9 +2066,10 @@ func TestSelectScatterOrderByVarChar(t *testing.T) { // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) + u := createSandbox(KsTestUnsharded) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + u.VSchema = unshardedVSchema serv := newSandboxForCells(ctx, []string{cell}) resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} @@ -2130,9 +2138,10 @@ func TestStreamSelectScatterOrderBy(t *testing.T) { // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) + u := createSandbox(KsTestUnsharded) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + u.VSchema = unshardedVSchema serv := newSandboxForCells(ctx, []string{cell}) resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} @@ -2193,9 +2202,10 @@ func TestStreamSelectScatterOrderByVarChar(t *testing.T) { // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) + u := createSandbox(KsTestUnsharded) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + u.VSchema = unshardedVSchema serv := newSandboxForCells(ctx, []string{cell}) resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} @@ -2255,9 +2265,10 @@ func TestSelectScatterAggregate(t *testing.T) { // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) + u := createSandbox(KsTestUnsharded) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + u.VSchema = unshardedVSchema serv := newSandboxForCells(ctx, []string{cell}) resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} @@ -2322,9 +2333,10 @@ func TestStreamSelectScatterAggregate(t *testing.T) { // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) + u := createSandbox(KsTestUnsharded) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + u.VSchema = unshardedVSchema serv := newSandboxForCells(ctx, []string{cell}) resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} @@ -2387,9 +2399,10 @@ func TestSelectScatterLimit(t *testing.T) { // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) + u := createSandbox(KsTestUnsharded) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + u.VSchema = unshardedVSchema serv := newSandboxForCells(ctx, []string{cell}) resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} @@ -2463,9 +2476,10 @@ func TestStreamSelectScatterLimit(t *testing.T) { // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) + u := createSandbox(KsTestUnsharded) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + u.VSchema = unshardedVSchema serv := newSandboxForCells(ctx, []string{cell}) resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} @@ -3434,9 +3448,10 @@ func TestStreamOrderByLimitWithMultipleResults(t *testing.T) { // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) + u := createSandbox(KsTestUnsharded) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + u.VSchema = unshardedVSchema serv := newSandboxForCells(ctx, []string{cell}) resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} @@ -3451,7 +3466,8 @@ func TestStreamOrderByLimitWithMultipleResults(t *testing.T) { } queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize) - executor := NewExecutor(ctx, serv, cell, resolver, true, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) + plans := cache.NewDefaultCacheImpl(cache.DefaultConfig) + executor := NewExecutor(ctx, serv, cell, resolver, true, false, testBufferSize, plans, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) defer executor.Close() // some sleep for all goroutines to start time.Sleep(100 * time.Millisecond) @@ -3476,9 +3492,10 @@ func TestSelectScatterFails(t *testing.T) { sess := &vtgatepb.Session{} cell := "aa" hc := discovery.NewFakeHealthCheck(nil) + u := createSandbox(KsTestUnsharded) s := createSandbox(KsTestSharded) s.VSchema = executorVSchema - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + u.VSchema = unshardedVSchema serv := newSandboxForCells(ctx, []string{cell}) resolver := newTestResolver(ctx, hc, serv, cell) @@ -4117,8 +4134,10 @@ func TestSelectAggregationNoData(t *testing.T) { // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) - createSandbox(KsTestSharded).VSchema = executorVSchema - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + u := createSandbox(KsTestUnsharded) + s := createSandbox(KsTestSharded) + s.VSchema = executorVSchema + u.VSchema = unshardedVSchema serv := newSandboxForCells(ctx, []string{cell}) resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} @@ -4209,8 +4228,10 @@ func TestSelectAggregationData(t *testing.T) { // Special setup: Don't use createExecutorEnv. cell := "aa" hc := discovery.NewFakeHealthCheck(nil) - createSandbox(KsTestSharded).VSchema = executorVSchema - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + u := createSandbox(KsTestUnsharded) + s := createSandbox(KsTestSharded) + s.VSchema = executorVSchema + u.VSchema = unshardedVSchema serv := newSandboxForCells(ctx, []string{cell}) resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} @@ -4356,8 +4377,10 @@ func TestSelectAggregationRandom(t *testing.T) { cell := "aa" hc := discovery.NewFakeHealthCheck(nil) - createSandbox(KsTestSharded).VSchema = executorVSchema - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + u := createSandbox(KsTestUnsharded) + s := createSandbox(KsTestSharded) + s.VSchema = executorVSchema + u.VSchema = unshardedVSchema serv := newSandboxForCells(ctx, []string{cell}) resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} diff --git a/go/vt/vtgate/executor_stream_test.go b/go/vt/vtgate/executor_stream_test.go index 065527b7ce2..9b041d095e1 100644 --- a/go/vt/vtgate/executor_stream_test.go +++ b/go/vt/vtgate/executor_stream_test.go @@ -62,9 +62,10 @@ func TestStreamSQLSharded(t *testing.T) { defer cancel() cell := "aa" hc := discovery.NewFakeHealthCheck(nil) + u := createSandbox(KsTestUnsharded) s := createSandbox("TestExecutor") s.VSchema = executorVSchema - getSandbox(KsTestUnsharded).VSchema = unshardedVSchema + u.VSchema = unshardedVSchema serv := newSandboxForCells(ctx, []string{cell}) resolver := newTestResolver(ctx, hc, serv, cell) shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"} @@ -72,7 +73,8 @@ func TestStreamSQLSharded(t *testing.T) { _ = hc.AddTestTablet(cell, shard, 1, "TestExecutor", shard, topodatapb.TabletType_PRIMARY, true, 1, nil) } queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize) - executor := NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, cache.DefaultConfig, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) + plans := cache.NewDefaultCacheImpl(cache.DefaultConfig) + executor := NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) defer executor.Close() sql := "stream * from sharded_user_msgs" diff --git a/go/vt/vtgate/executor_test.go b/go/vt/vtgate/executor_test.go index b2ce904e6e3..c3aa300bcda 100644 --- a/go/vt/vtgate/executor_test.go +++ b/go/vt/vtgate/executor_test.go @@ -1048,7 +1048,7 @@ func TestExecutorShow(t *testing.T) { utils.MustMatch(t, wantqr, qr, query) // Make sure it still works when one of the keyspaces is in a bad state - getSandbox("TestExecutor").SrvKeyspaceMustFail++ + getSandbox(KsTestSharded).SrvKeyspaceMustFail++ query = "show vitess_shards" qr, err = executor.Execute(ctx, nil, "TestExecute", session, query, nil) require.NoError(t, err) diff --git a/go/vt/vtgate/plugin_mysql_server.go b/go/vt/vtgate/plugin_mysql_server.go index 6de6f12dba7..bfbb7b105f8 100644 --- a/go/vt/vtgate/plugin_mysql_server.go +++ b/go/vt/vtgate/plugin_mysql_server.go @@ -74,8 +74,6 @@ var ( mysqlDefaultWorkloadName = "OLTP" mysqlDefaultWorkload int32 - - busyConnections int32 ) func registerPluginFlags(fs *pflag.FlagSet) { @@ -110,6 +108,8 @@ type vtgateHandler struct { vtg *VTGate connections map[uint32]*mysql.Conn + + busyConnections atomic.Int32 } func newVtgateHandler(vtg *VTGate) *vtgateHandler { @@ -135,7 +135,7 @@ func (vh *vtgateHandler) ComResetConnection(c *mysql.Conn) { ctx := context.Background() session := vh.session(c) if session.InTransaction { - defer atomic.AddInt32(&busyConnections, -1) + defer vh.busyConnections.Add(-1) } err := vh.vtg.CloseSession(ctx, session) if err != nil { @@ -161,7 +161,7 @@ func (vh *vtgateHandler) ConnectionClosed(c *mysql.Conn) { } session := vh.session(c) if session.InTransaction { - defer atomic.AddInt32(&busyConnections, -1) + defer vh.busyConnections.Add(-1) } _ = vh.vtg.CloseSession(ctx, session) } @@ -231,11 +231,11 @@ func (vh *vtgateHandler) ComQuery(c *mysql.Conn, query string, callback func(*sq session := vh.session(c) if !session.InTransaction { - atomic.AddInt32(&busyConnections, 1) + vh.busyConnections.Add(1) } defer func() { if !session.InTransaction { - atomic.AddInt32(&busyConnections, -1) + vh.busyConnections.Add(-1) } }() @@ -296,11 +296,11 @@ func (vh *vtgateHandler) ComPrepare(c *mysql.Conn, query string, bindVars map[st session := vh.session(c) if !session.InTransaction { - atomic.AddInt32(&busyConnections, 1) + vh.busyConnections.Add(1) } defer func() { if !session.InTransaction { - atomic.AddInt32(&busyConnections, -1) + vh.busyConnections.Add(-1) } }() @@ -337,11 +337,11 @@ func (vh *vtgateHandler) ComStmtExecute(c *mysql.Conn, prepare *mysql.PrepareDat session := vh.session(c) if !session.InTransaction { - atomic.AddInt32(&busyConnections, 1) + vh.busyConnections.Add(1) } defer func() { if !session.InTransaction { - atomic.AddInt32(&busyConnections, -1) + vh.busyConnections.Add(-1) } }() @@ -437,34 +437,36 @@ func (vh *vtgateHandler) session(c *mysql.Conn) *vtgatepb.Session { return session } -var mysqlListener *mysql.Listener -var mysqlUnixListener *mysql.Listener -var sigChan chan os.Signal -var vtgateHandle *vtgateHandler +type mysqlServer struct { + tcpListener *mysql.Listener + unixListener *mysql.Listener + sigChan chan os.Signal + vtgateHandle *vtgateHandler +} // initTLSConfig inits tls config for the given mysql listener -func initTLSConfig(ctx context.Context, mysqlListener *mysql.Listener, mysqlSslCert, mysqlSslKey, mysqlSslCa, mysqlSslCrl, mysqlSslServerCA string, mysqlServerRequireSecureTransport bool, mysqlMinTLSVersion uint16) error { +func initTLSConfig(ctx context.Context, srv *mysqlServer, mysqlSslCert, mysqlSslKey, mysqlSslCa, mysqlSslCrl, mysqlSslServerCA string, mysqlServerRequireSecureTransport bool, mysqlMinTLSVersion uint16) error { serverConfig, err := vttls.ServerConfig(mysqlSslCert, mysqlSslKey, mysqlSslCa, mysqlSslCrl, mysqlSslServerCA, mysqlMinTLSVersion) if err != nil { log.Exitf("grpcutils.TLSServerConfig failed: %v", err) return err } - mysqlListener.TLSConfig.Store(serverConfig) - mysqlListener.RequireSecureTransport = mysqlServerRequireSecureTransport - sigChan = make(chan os.Signal, 1) - signal.Notify(sigChan, syscall.SIGHUP) + srv.tcpListener.TLSConfig.Store(serverConfig) + srv.tcpListener.RequireSecureTransport = mysqlServerRequireSecureTransport + srv.sigChan = make(chan os.Signal, 1) + signal.Notify(srv.sigChan, syscall.SIGHUP) go func() { for { select { case <-ctx.Done(): return - case <-sigChan: + case <-srv.sigChan: serverConfig, err := vttls.ServerConfig(mysqlSslCert, mysqlSslKey, mysqlSslCa, mysqlSslCrl, mysqlSslServerCA, mysqlMinTLSVersion) if err != nil { log.Errorf("grpcutils.TLSServerConfig failed: %v", err) } else { log.Info("grpcutils.TLSServerConfig updated") - mysqlListener.TLSConfig.Store(serverConfig) + srv.tcpListener.TLSConfig.Store(serverConfig) } } } @@ -474,15 +476,15 @@ func initTLSConfig(ctx context.Context, mysqlListener *mysql.Listener, mysqlSslC // initiMySQLProtocol starts the mysql protocol. // It should be called only once in a process. -func initMySQLProtocol(vtgate *VTGate) { +func initMySQLProtocol(vtgate *VTGate) *mysqlServer { // Flag is not set, just return. if mysqlServerPort < 0 && mysqlServerSocketPath == "" { - return + return nil } // If no VTGate was created, just return. if vtgate == nil { - return + return nil } // Initialize registered AuthServer implementations (or other plugins) @@ -506,13 +508,14 @@ func initMySQLProtocol(vtgate *VTGate) { // Create a Listener. var err error - vtgateHandle = newVtgateHandler(vtgate) + srv := &mysqlServer{} + srv.vtgateHandle = newVtgateHandler(vtgate) if mysqlServerPort >= 0 { - mysqlListener, err = mysql.NewListener( + srv.tcpListener, err = mysql.NewListener( mysqlTCPVersion, net.JoinHostPort(mysqlServerBindAddress, fmt.Sprintf("%v", mysqlServerPort)), authServer, - vtgateHandle, + srv.vtgateHandle, mysqlConnReadTimeout, mysqlConnWriteTimeout, mysqlProxyProtocol, @@ -522,38 +525,39 @@ func initMySQLProtocol(vtgate *VTGate) { if err != nil { log.Exitf("mysql.NewListener failed: %v", err) } - mysqlListener.ServerVersion = servenv.MySQLServerVersion() + srv.tcpListener.ServerVersion = servenv.MySQLServerVersion() if mysqlSslCert != "" && mysqlSslKey != "" { tlsVersion, err := vttls.TLSVersionToNumber(mysqlTLSMinVersion) if err != nil { log.Exitf("mysql.NewListener failed: %v", err) } - _ = initTLSConfig(context.Background(), mysqlListener, mysqlSslCert, mysqlSslKey, mysqlSslCa, mysqlSslCrl, mysqlSslServerCA, mysqlServerRequireSecureTransport, tlsVersion) + _ = initTLSConfig(context.Background(), srv, mysqlSslCert, mysqlSslKey, mysqlSslCa, mysqlSslCrl, mysqlSslServerCA, mysqlServerRequireSecureTransport, tlsVersion) } - mysqlListener.AllowClearTextWithoutTLS.Store(mysqlAllowClearTextWithoutTLS) + srv.tcpListener.AllowClearTextWithoutTLS.Store(mysqlAllowClearTextWithoutTLS) // Check for the connection threshold if mysqlSlowConnectWarnThreshold != 0 { log.Infof("setting mysql slow connection threshold to %v", mysqlSlowConnectWarnThreshold) - mysqlListener.SlowConnectWarnThreshold.Store(mysqlSlowConnectWarnThreshold.Nanoseconds()) + srv.tcpListener.SlowConnectWarnThreshold.Store(mysqlSlowConnectWarnThreshold.Nanoseconds()) } // Start listening for tcp - go mysqlListener.Accept() + go srv.tcpListener.Accept() } if mysqlServerSocketPath != "" { // Let's create this unix socket with permissions to all users. In this way, // clients can connect to vtgate mysql server without being vtgate user oldMask := syscall.Umask(000) - mysqlUnixListener, err = newMysqlUnixSocket(mysqlServerSocketPath, authServer, vtgateHandle) + srv.unixListener, err = newMysqlUnixSocket(mysqlServerSocketPath, authServer, srv.vtgateHandle) _ = syscall.Umask(oldMask) if err != nil { log.Exitf("mysql.NewListener failed: %v", err) - return + return nil } // Listen for unix socket - go mysqlUnixListener.Accept() + go srv.unixListener.Accept() } + return srv } // newMysqlUnixSocket creates a new unix socket mysql listener. If a socket file already exists, attempts @@ -608,37 +612,38 @@ func newMysqlUnixSocket(address string, authServer mysql.AuthServer, handler mys } } -func shutdownMysqlProtocolAndDrain() { - if mysqlListener != nil { - mysqlListener.Close() - mysqlListener = nil +func (srv *mysqlServer) shutdownMysqlProtocolAndDrain() { + if srv.tcpListener != nil { + srv.tcpListener.Close() + srv.tcpListener = nil } - if mysqlUnixListener != nil { - mysqlUnixListener.Close() - mysqlUnixListener = nil + if srv.unixListener != nil { + srv.unixListener.Close() + srv.unixListener = nil } - if sigChan != nil { - signal.Stop(sigChan) + if srv.sigChan != nil { + signal.Stop(srv.sigChan) } - if atomic.LoadInt32(&busyConnections) > 0 { - log.Infof("Waiting for all client connections to be idle (%d active)...", atomic.LoadInt32(&busyConnections)) + if busy := srv.vtgateHandle.busyConnections.Load(); busy > 0 { + log.Infof("Waiting for all client connections to be idle (%d active)...", busy) start := time.Now() reported := start - for atomic.LoadInt32(&busyConnections) > 0 { + for busy > 0 { if time.Since(reported) > 2*time.Second { - log.Infof("Still waiting for client connections to be idle (%d active)...", atomic.LoadInt32(&busyConnections)) + log.Infof("Still waiting for client connections to be idle (%d active)...", busy) reported = time.Now() } time.Sleep(1 * time.Millisecond) + busy = srv.vtgateHandle.busyConnections.Load() } } } -func rollbackAtShutdown() { +func (srv *mysqlServer) rollbackAtShutdown() { defer log.Flush() - if vtgateHandle == nil { + if srv.vtgateHandle == nil { // we still haven't been able to initialise the vtgateHandler, so we don't need to rollback anything return } @@ -646,10 +651,10 @@ func rollbackAtShutdown() { // Close all open connections. If they're waiting for reads, this will cause // them to error out, which will automatically rollback open transactions. func() { - if vtgateHandle != nil { - vtgateHandle.mu.Lock() - defer vtgateHandle.mu.Unlock() - for id, c := range vtgateHandle.connections { + if srv.vtgateHandle != nil { + srv.vtgateHandle.mu.Lock() + defer srv.vtgateHandle.mu.Unlock() + for id, c := range srv.vtgateHandle.connections { if c != nil { log.Infof("Rolling back transactions associated with connection ID: %v", id) c.Close() @@ -661,7 +666,7 @@ func rollbackAtShutdown() { // If vtgate is instead busy executing a query, the number of open conns // will be non-zero. Give another second for those queries to finish. for i := 0; i < 100; i++ { - if vtgateHandle.numConnections() == 0 { + if srv.vtgateHandle.numConnections() == 0 { log.Infof("All connections have been rolled back.") return } @@ -680,8 +685,6 @@ func mysqlSocketPath() string { func init() { servenv.OnParseFor("vtgate", registerPluginFlags) servenv.OnParseFor("vtcombo", registerPluginFlags) - servenv.OnTermSync(shutdownMysqlProtocolAndDrain) - servenv.OnClose(rollbackAtShutdown) } var pluginInitializers []func() diff --git a/go/vt/vtgate/plugin_mysql_server_test.go b/go/vt/vtgate/plugin_mysql_server_test.go index 3487d0253c8..22c0d24e8f9 100644 --- a/go/vt/vtgate/plugin_mysql_server_test.go +++ b/go/vt/vtgate/plugin_mysql_server_test.go @@ -248,6 +248,7 @@ func newTestAuthServerStatic() *mysql.AuthServerStatic { func TestDefaultWorkloadEmpty(t *testing.T) { vh := &vtgateHandler{} + mysqlDefaultWorkload = int32(querypb.ExecuteOptions_OLTP) sess := vh.session(&mysql.Conn{}) if sess.Options.Workload != querypb.ExecuteOptions_OLTP { t.Fatalf("Expected default workload OLTP") @@ -287,20 +288,20 @@ func testInitTLSConfig(t *testing.T, serverCA bool) { serverCACert = path.Join(root, "ca-cert.pem") } - listener := &mysql.Listener{} - if err := initTLSConfig(ctx, listener, path.Join(root, "server-cert.pem"), path.Join(root, "server-key.pem"), path.Join(root, "ca-cert.pem"), path.Join(root, "ca-crl.pem"), serverCACert, true, tls.VersionTLS12); err != nil { + srv := &mysqlServer{tcpListener: &mysql.Listener{}} + if err := initTLSConfig(ctx, srv, path.Join(root, "server-cert.pem"), path.Join(root, "server-key.pem"), path.Join(root, "ca-cert.pem"), path.Join(root, "ca-crl.pem"), serverCACert, true, tls.VersionTLS12); err != nil { t.Fatalf("init tls config failure due to: +%v", err) } - serverConfig := listener.TLSConfig.Load() + serverConfig := srv.tcpListener.TLSConfig.Load() if serverConfig == nil { t.Fatalf("init tls config shouldn't create nil server config") } - sigChan <- syscall.SIGHUP + srv.sigChan <- syscall.SIGHUP time.Sleep(100 * time.Millisecond) // wait for signal handler - if listener.TLSConfig.Load() == serverConfig { + if srv.tcpListener.TLSConfig.Load() == serverConfig { t.Fatalf("init tls config should have been recreated after SIGHUP") } } diff --git a/go/vt/vtgate/querylog.go b/go/vt/vtgate/querylog.go index e8a3ae9ac2b..94c1b23263b 100644 --- a/go/vt/vtgate/querylog.go +++ b/go/vt/vtgate/querylog.go @@ -19,6 +19,7 @@ package vtgate import ( "net/http" + "vitess.io/vitess/go/cache" "vitess.io/vitess/go/streamlog" "vitess.io/vitess/go/vt/servenv" "vitess.io/vitess/go/vt/vtgate/logstats" @@ -35,7 +36,7 @@ var ( QueryzHandler = "/debug/queryz" ) -func initQueryLogger(vtg *VTGate) (*streamlog.StreamLogger[*logstats.LogStats], error) { +func initQueryLogger(plans cache.Cache) (*streamlog.StreamLogger[*logstats.LogStats], error) { queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize) queryLogger.ServeLogs(QueryLogHandler, streamlog.GetFormatter(queryLogger)) @@ -46,7 +47,7 @@ func initQueryLogger(vtg *VTGate) (*streamlog.StreamLogger[*logstats.LogStats], }) servenv.HTTPHandleFunc(QueryzHandler, func(w http.ResponseWriter, r *http.Request) { - queryzHandler(vtg.executor, w, r) + queryzHandler(plans, w, r) }) if queryLogToFile != "" { diff --git a/go/vt/vtgate/queryz.go b/go/vt/vtgate/queryz.go index 6f660127ad4..6c5d2b89dee 100644 --- a/go/vt/vtgate/queryz.go +++ b/go/vt/vtgate/queryz.go @@ -24,6 +24,8 @@ import ( "github.com/google/safehtml/template" + "vitess.io/vitess/go/cache" + "vitess.io/vitess/go/acl" "vitess.io/vitess/go/vt/log" "vitess.io/vitess/go/vt/logz" @@ -127,7 +129,7 @@ func (s *queryzSorter) Len() int { return len(s.rows) } func (s *queryzSorter) Swap(i, j int) { s.rows[i], s.rows[j] = s.rows[j], s.rows[i] } func (s *queryzSorter) Less(i, j int) bool { return s.less(s.rows[i], s.rows[j]) } -func queryzHandler(e *Executor, w http.ResponseWriter, r *http.Request) { +func queryzHandler(plans cache.Cache, w http.ResponseWriter, r *http.Request) { if err := acl.CheckAccessHTTP(r, acl.DEBUGGING); err != nil { acl.SendError(w, err) return @@ -143,7 +145,7 @@ func queryzHandler(e *Executor, w http.ResponseWriter, r *http.Request) { }, } - e.plans.ForEach(func(value any) bool { + plans.ForEach(func(value any) bool { plan := value.(*engine.Plan) Value := &queryzRow{ Query: logz.Wrappable(sqlparser.TruncateForUI(plan.Original)), diff --git a/go/vt/vtgate/queryz_test.go b/go/vt/vtgate/queryz_test.go index 2de9c1bfe42..1c8e9d9b190 100644 --- a/go/vt/vtgate/queryz_test.go +++ b/go/vt/vtgate/queryz_test.go @@ -89,7 +89,7 @@ func TestQueryzHandler(t *testing.T) { plan3.ExecTime = uint64(100 * time.Millisecond) plan4.ExecTime = uint64(200 * time.Millisecond) - queryzHandler(executor, resp, req) + queryzHandler(executor.plans, resp, req) body, _ := io.ReadAll(resp.Body) planPattern1 := []string{ ``, diff --git a/go/vt/vtgate/sandbox_test.go b/go/vt/vtgate/sandbox_test.go index 22c124020a1..1629e9a4faa 100644 --- a/go/vt/vtgate/sandbox_test.go +++ b/go/vt/vtgate/sandbox_test.go @@ -324,6 +324,10 @@ func (sct *sandboxTopo) WatchSrvVSchema(ctx context.Context, cell string, callba case <-ctx.Done(): return case update := <-updateChan: + // If the channel was closed, we're done. + if update == nil { + return + } newHash := GetSrvVSchemaHash(update.Value) if newHash == currentHash { // sometimes we get the same update multiple times. This results in the plan cache to be cleared diff --git a/go/vt/vtgate/vtgate.go b/go/vt/vtgate/vtgate.go index b3ed103896e..485759185de 100644 --- a/go/vt/vtgate/vtgate.go +++ b/go/vt/vtgate/vtgate.go @@ -178,8 +178,6 @@ func getTxMode() vtgatepb.TransactionMode { } var ( - rpcVTGate *VTGate - // vschemaCounters needs to be initialized before planner to // catch the initial load stats. vschemaCounters = stats.NewCountersWithSingleLabel("VtgateVSchemaCounts", "Vtgate vschema counts", "changes") @@ -193,6 +191,21 @@ var ( "Number of events that had to wait because the skew across shards was too high") vindexUnknownParams = stats.NewGauge("VindexUnknownParameters", "Number of parameterss unrecognized by Vindexes") + + timings = stats.NewMultiTimings( + "VtgateApi", + "VtgateApi timings", + []string{"Operation", "Keyspace", "DbType"}) + + rowsReturned = stats.NewCountersWithMultiLabels( + "VtgateApiRowsReturned", + "Rows returned through the VTgate API", + []string{"Operation", "Keyspace", "DbType"}) + + rowsAffected = stats.NewCountersWithMultiLabels( + "VtgateApiRowsAffected", + "Rows affected by a write (DML) operation through the VTgate API", + []string{"Operation", "Keyspace", "DbType"}) ) // VTGate is the rpc interface to vtgate. Only one instance @@ -234,10 +247,6 @@ func Init( tabletTypesToWait []topodatapb.TabletType, pv plancontext.PlannerVersion, ) *VTGate { - if rpcVTGate != nil { - log.Fatalf("VTGate already initialized") - } - // Build objects from low to high level. // Start with the gateway. If we can't reach the topology service, // we can't go on much further, so we log.Fatal out. @@ -301,7 +310,8 @@ func Init( LFU: queryPlanCacheLFU, } - queryLogger, err := initQueryLogger(rpcVTGate) + plans := cache.NewDefaultCacheImpl(cacheCfg) + queryLogger, err := initQueryLogger(plans) if err != nil { log.Fatalf("error initializing query logger: %v", err) } @@ -314,7 +324,7 @@ func Init( normalizeQueries, warnShardedOnly, streamBufferSize, - cacheCfg, + plans, si, noScatter, pv, @@ -328,33 +338,10 @@ func Init( // TODO: call serv.WatchSrvVSchema here - rpcVTGate = &VTGate{ - executor: executor, - resolver: resolver, - vsm: vsm, - txConn: tc, - gw: gw, - timings: stats.NewMultiTimings( - "VtgateApi", - "VtgateApi timings", - []string{"Operation", "Keyspace", "DbType"}), - rowsReturned: stats.NewCountersWithMultiLabels( - "VtgateApiRowsReturned", - "Rows returned through the VTgate API", - []string{"Operation", "Keyspace", "DbType"}), - rowsAffected: stats.NewCountersWithMultiLabels( - "VtgateApiRowsAffected", - "Rows affected by a write (DML) operation through the VTgate API", - []string{"Operation", "Keyspace", "DbType"}), - - logExecute: logutil.NewThrottledLogger("Execute", 5*time.Second), - logPrepare: logutil.NewThrottledLogger("Prepare", 5*time.Second), - logStreamExecute: logutil.NewThrottledLogger("StreamExecute", 5*time.Second), - } - - _ = stats.NewRates("QPSByOperation", stats.CounterForDimension(rpcVTGate.timings, "Operation"), 15, 1*time.Minute) - _ = stats.NewRates("QPSByKeyspace", stats.CounterForDimension(rpcVTGate.timings, "Keyspace"), 15, 1*time.Minute) - _ = stats.NewRates("QPSByDbType", stats.CounterForDimension(rpcVTGate.timings, "DbType"), 15*60/5, 5*time.Second) + vtgateInst := newVTGate(executor, resolver, vsm, tc, gw) + _ = stats.NewRates("QPSByOperation", stats.CounterForDimension(vtgateInst.timings, "Operation"), 15, 1*time.Minute) + _ = stats.NewRates("QPSByKeyspace", stats.CounterForDimension(vtgateInst.timings, "Keyspace"), 15, 1*time.Minute) + _ = stats.NewRates("QPSByDbType", stats.CounterForDimension(vtgateInst.timings, "DbType"), 15*60/5, 5*time.Second) _ = stats.NewRates("ErrorsByOperation", stats.CounterForDimension(errorCounts, "Operation"), 15, 1*time.Minute) _ = stats.NewRates("ErrorsByKeyspace", stats.CounterForDimension(errorCounts, "Keyspace"), 15, 1*time.Minute) @@ -363,23 +350,25 @@ func Init( servenv.OnRun(func() { for _, f := range RegisterVTGates { - f(rpcVTGate) + f(vtgateInst) } if st != nil && enableSchemaChangeSignal { st.Start() } - initMySQLProtocol(rpcVTGate) + srv := initMySQLProtocol(vtgateInst) + servenv.OnTermSync(srv.shutdownMysqlProtocolAndDrain) + servenv.OnClose(srv.rollbackAtShutdown) }) servenv.OnTerm(func() { if st != nil && enableSchemaChangeSignal { st.Stop() } }) - rpcVTGate.registerDebugHealthHandler() - rpcVTGate.registerDebugEnvHandler() + vtgateInst.registerDebugHealthHandler() + vtgateInst.registerDebugEnvHandler() initAPI(gw.hc) - return rpcVTGate + return vtgateInst } func addKeyspacesToTracker(ctx context.Context, srvResolver *srvtopo.Resolver, st *vtschema.Tracker, gw *TabletGateway) { @@ -676,3 +665,25 @@ func (vtg *VTGate) HandlePanic(err *error) { errorCounts.Add([]string{"Panic", "Unknown", "Unknown", vtrpcpb.Code_INTERNAL.String()}, 1) } } + +func (vtg *VTGate) Close() { + vtg.executor.Close() + vtg.gw.Close(context.Background()) +} + +func newVTGate(executor *Executor, resolver *Resolver, vsm *vstreamManager, tc *TxConn, gw *TabletGateway) *VTGate { + return &VTGate{ + executor: executor, + resolver: resolver, + vsm: vsm, + txConn: tc, + gw: gw, + timings: timings, + rowsReturned: rowsReturned, + rowsAffected: rowsAffected, + + logExecute: logutil.NewThrottledLogger("Execute", 5*time.Second), + logPrepare: logutil.NewThrottledLogger("Prepare", 5*time.Second), + logStreamExecute: logutil.NewThrottledLogger("StreamExecute", 5*time.Second), + } +} diff --git a/go/vt/vtgate/vtgate_test.go b/go/vt/vtgate/vtgate_test.go index 7bd6fa3353e..a2ed31af958 100644 --- a/go/vt/vtgate/vtgate_test.go +++ b/go/vt/vtgate/vtgate_test.go @@ -22,12 +22,11 @@ import ( "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "google.golang.org/protobuf/proto" "vitess.io/vitess/go/test/utils" - "github.com/stretchr/testify/require" - "vitess.io/vitess/go/sqltypes" "vitess.io/vitess/go/vt/discovery" "vitess.io/vitess/go/vt/vterrors" @@ -41,24 +40,24 @@ import ( // This file uses the sandbox_test framework. -var hcVTGateTest *discovery.FakeHealthCheck - var executeOptions = &querypb.ExecuteOptions{ IncludedFields: querypb.ExecuteOptions_TYPE_ONLY, } func TestVTGateExecute(t *testing.T) { - counts := rpcVTGate.timings.Timings.Counts() - - createSandbox(KsTestUnsharded) - hcVTGateTest.Reset() - sbc := hcVTGateTest.AddTestTablet("aa", "1.1.1.1", 1001, KsTestUnsharded, "0", topodatapb.TabletType_PRIMARY, true, 1, nil) - _, qr, err := rpcVTGate.Execute( - context.Background(), + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + vtg, sbc := createVtgateEnv(ctx) + defer vtg.Close() + counts := vtg.timings.Timings.Counts() + + _, qr, err := vtg.Execute( + ctx, nil, &vtgatepb.Session{ Autocommit: true, - TargetString: "@primary", + TargetString: KsTestUnsharded + "@primary", Options: executeOptions, }, "select id from t1", @@ -75,11 +74,11 @@ func TestVTGateExecute(t *testing.T) { t.Errorf("got ExecuteOptions \n%+v, want \n%+v", sbc.Options[0], executeOptions) } - newCounts := rpcVTGate.timings.Timings.Counts() + newCounts := vtg.timings.Timings.Counts() require.Contains(t, newCounts, "All") require.Equal(t, counts["All"]+1, newCounts["All"]) - require.Contains(t, newCounts, "Execute..primary") - require.Equal(t, counts["Execute..primary"]+1, newCounts["Execute..primary"]) + require.Contains(t, newCounts, "Execute.TestUnsharded.primary") + require.Equal(t, counts["Execute.TestUnsharded.primary"]+1, newCounts["Execute.TestUnsharded.primary"]) for k, v := range newCounts { if strings.HasPrefix(k, "Prepare") { @@ -89,17 +88,19 @@ func TestVTGateExecute(t *testing.T) { } func TestVTGateExecuteError(t *testing.T) { - counts := errorCounts.Counts() - - createSandbox(KsTestUnsharded) - hcVTGateTest.Reset() - hcVTGateTest.AddTestTablet("aa", "1.1.1.1", 1001, KsTestUnsharded, "0", topodatapb.TabletType_PRIMARY, true, 1, nil) - _, qr, err := rpcVTGate.Execute( - context.Background(), + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + vtg, _ := createVtgateEnv(ctx) + defer vtg.Close() + counts := vtg.timings.Timings.Counts() + + _, qr, err := vtg.Execute( + ctx, nil, &vtgatepb.Session{ Autocommit: true, - TargetString: "@primary", + TargetString: KsTestUnsharded + "@primary", Options: executeOptions, }, "bad select id from t1", @@ -109,8 +110,8 @@ func TestVTGateExecuteError(t *testing.T) { require.Nil(t, qr) newCounts := errorCounts.Counts() - require.Contains(t, newCounts, "Execute..primary.INVALID_ARGUMENT") - require.Equal(t, counts["Execute..primary.INVALID_ARGUMENT"]+1, newCounts["Execute..primary.INVALID_ARGUMENT"]) + require.Contains(t, newCounts, "Execute.TestUnsharded.primary.INVALID_ARGUMENT") + require.Equal(t, counts["Execute.TestUnsharded.primary.INVALID_ARGUMENT"]+1, newCounts["Execute.TestUnsharded.primary.INVALID_ARGUMENT"]) for k, v := range newCounts { if strings.HasPrefix(k, "Prepare") { @@ -120,16 +121,17 @@ func TestVTGateExecuteError(t *testing.T) { } func TestVTGatePrepare(t *testing.T) { - counts := rpcVTGate.timings.Timings.Counts() - - createSandbox(KsTestUnsharded) - hcVTGateTest.Reset() - sbc := hcVTGateTest.AddTestTablet("aa", "1.1.1.1", 1001, KsTestUnsharded, "0", topodatapb.TabletType_PRIMARY, true, 1, nil) - _, qr, err := rpcVTGate.Prepare( - context.Background(), + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + vtg, sbc := createVtgateEnv(ctx) + defer vtg.Close() + counts := vtg.timings.Timings.Counts() + _, qr, err := vtg.Prepare( + ctx, &vtgatepb.Session{ Autocommit: true, - TargetString: "@primary", + TargetString: KsTestUnsharded + "@primary", Options: executeOptions, }, "select id from t1", @@ -145,11 +147,11 @@ func TestVTGatePrepare(t *testing.T) { t.Errorf("got ExecuteOptions \n%+v, want \n%+v", sbc.Options[0], executeOptions) } - newCounts := rpcVTGate.timings.Timings.Counts() + newCounts := vtg.timings.Timings.Counts() require.Contains(t, newCounts, "All") require.Equal(t, counts["All"]+1, newCounts["All"]) - require.Contains(t, newCounts, "Prepare..primary") - require.Equal(t, counts["Prepare..primary"]+1, newCounts["Prepare..primary"]) + require.Contains(t, newCounts, "Prepare.TestUnsharded.primary") + require.Equal(t, counts["Prepare.TestUnsharded.primary"]+1, newCounts["Prepare.TestUnsharded.primary"]) for k, v := range newCounts { if strings.HasPrefix(k, "Execute") { @@ -159,16 +161,18 @@ func TestVTGatePrepare(t *testing.T) { } func TestVTGatePrepareError(t *testing.T) { + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + vtg, _ := createVtgateEnv(ctx) + defer vtg.Close() counts := errorCounts.Counts() - createSandbox(KsTestUnsharded) - hcVTGateTest.Reset() - hcVTGateTest.AddTestTablet("aa", "1.1.1.1", 1001, KsTestUnsharded, "0", topodatapb.TabletType_PRIMARY, true, 1, nil) - _, qr, err := rpcVTGate.Prepare( - context.Background(), + _, qr, err := vtg.Prepare( + ctx, &vtgatepb.Session{ Autocommit: true, - TargetString: "@primary", + TargetString: KsTestUnsharded + "@primary", Options: executeOptions, }, "bad select id from t1", @@ -178,8 +182,8 @@ func TestVTGatePrepareError(t *testing.T) { require.Nil(t, qr) newCounts := errorCounts.Counts() - require.Contains(t, newCounts, "Prepare..primary.INTERNAL") - require.Equal(t, counts["Prepare..primary.INTERNAL"]+1, newCounts["Prepare..primary.INTERNAL"]) + require.Contains(t, newCounts, "Prepare.TestUnsharded.primary.INTERNAL") + require.Equal(t, counts["Prepare.TestUnsharded.primary.INTERNAL"]+1, newCounts["Prepare.TestUnsharded.primary.INTERNAL"]) for k, v := range newCounts { if strings.HasPrefix(k, "Execute") { @@ -189,16 +193,18 @@ func TestVTGatePrepareError(t *testing.T) { } func TestVTGateExecuteWithKeyspaceShard(t *testing.T) { - createSandbox(KsTestUnsharded) - hcVTGateTest.Reset() - hcVTGateTest.AddTestTablet("aa", "1.1.1.1", 1001, KsTestUnsharded, "0", topodatapb.TabletType_PRIMARY, true, 1, nil) + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + vtg, _ := createVtgateEnv(ctx) + defer vtg.Close() // Valid keyspace. - _, qr, err := rpcVTGate.Execute( - context.Background(), + _, qr, err := vtg.Execute( + ctx, nil, &vtgatepb.Session{ - TargetString: KsTestUnsharded, + TargetString: KsTestSharded + ":-20@primary", }, "select id from none", nil, @@ -211,8 +217,8 @@ func TestVTGateExecuteWithKeyspaceShard(t *testing.T) { utils.MustMatch(t, &wantQr, qr) // Invalid keyspace. - _, _, err = rpcVTGate.Execute( - context.Background(), + _, _, err = vtg.Execute( + ctx, nil, &vtgatepb.Session{ TargetString: "invalid_keyspace", @@ -224,11 +230,11 @@ func TestVTGateExecuteWithKeyspaceShard(t *testing.T) { assert.EqualError(t, err, want) // Valid keyspace/shard. - _, qr, err = rpcVTGate.Execute( - context.Background(), + _, qr, err = vtg.Execute( + ctx, nil, &vtgatepb.Session{ - TargetString: KsTestUnsharded + ":0@primary", + TargetString: KsTestSharded + ":-20@primary", }, "select id from none", nil, @@ -239,31 +245,32 @@ func TestVTGateExecuteWithKeyspaceShard(t *testing.T) { utils.MustMatch(t, &wantQr, qr) // Invalid keyspace/shard. - _, _, err = rpcVTGate.Execute( - context.Background(), + _, _, err = vtg.Execute( + ctx, nil, &vtgatepb.Session{ - TargetString: KsTestUnsharded + ":noshard@primary", + TargetString: KsTestSharded + ":noshard@primary", }, "select id from none", nil, ) require.Error(t, err) - require.Contains(t, err.Error(), `no healthy tablet available for 'keyspace:"TestUnsharded" shard:"noshard" tablet_type:PRIMARY`) + require.Contains(t, err.Error(), `no healthy tablet available for 'keyspace:"TestExecutor" shard:"noshard" tablet_type:PRIMARY`) } func TestVTGateStreamExecute(t *testing.T) { - ks := KsTestUnsharded - shard := "0" - createSandbox(ks) - hcVTGateTest.Reset() - sbc := hcVTGateTest.AddTestTablet("aa", "1.1.1.1", 1001, ks, shard, topodatapb.TabletType_PRIMARY, true, 1, nil) + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + vtg, sbc := createVtgateEnv(ctx) + defer vtg.Close() + var qrs []*sqltypes.Result - _, err := rpcVTGate.StreamExecute( - context.Background(), + _, err := vtg.StreamExecute( + ctx, nil, &vtgatepb.Session{ - TargetString: "@primary", + TargetString: KsTestUnsharded + "@primary", Options: executeOptions, }, "select id from t1", @@ -286,10 +293,12 @@ func TestVTGateStreamExecute(t *testing.T) { } func TestVTGateBindVarError(t *testing.T) { - ks := KsTestUnsharded - createSandbox(ks) - hcVTGateTest.Reset() - ctx := context.Background() + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + vtg, _ := createVtgateEnv(ctx) + defer vtg.Close() + session := &vtgatepb.Session{} bindVars := map[string]*querypb.BindVariable{ "v": { @@ -305,19 +314,19 @@ func TestVTGateBindVarError(t *testing.T) { }{{ name: "Execute", f: func() error { - _, _, err := rpcVTGate.Execute(ctx, nil, session, "", bindVars) + _, _, err := vtg.Execute(ctx, nil, session, "", bindVars) return err }, }, { name: "ExecuteBatch", f: func() error { - _, _, err := rpcVTGate.ExecuteBatch(ctx, session, []string{""}, []map[string]*querypb.BindVariable{bindVars}) + _, _, err := vtg.ExecuteBatch(ctx, session, []string{""}, []map[string]*querypb.BindVariable{bindVars}) return err }, }, { name: "StreamExecute", f: func() error { - _, err := rpcVTGate.StreamExecute(ctx, nil, session, "", bindVars, func(_ *sqltypes.Result) error { return nil }) + _, err := vtg.StreamExecute(ctx, nil, session, "", bindVars, func(_ *sqltypes.Result) error { return nil }) return err }, }} @@ -328,18 +337,18 @@ func TestVTGateBindVarError(t *testing.T) { } } -func testErrorPropagation(t *testing.T, sbcs []*sandboxconn.SandboxConn, before func(sbc *sandboxconn.SandboxConn), after func(sbc *sandboxconn.SandboxConn), expected vtrpcpb.Code) { +func testErrorPropagation(t *testing.T, ctx context.Context, vtg *VTGate, sbcs []*sandboxconn.SandboxConn, before func(sbc *sandboxconn.SandboxConn), after func(sbc *sandboxconn.SandboxConn), expected vtrpcpb.Code) { // Execute for _, sbc := range sbcs { before(sbc) } session := &vtgatepb.Session{ - TargetString: "@primary", + TargetString: KsTestUnsharded + "@primary", } - _, _, err := rpcVTGate.Execute( - context.Background(), + _, _, err := vtg.Execute( + ctx, nil, session, "select id from t1", @@ -361,8 +370,8 @@ func testErrorPropagation(t *testing.T, sbcs []*sandboxconn.SandboxConn, before for _, sbc := range sbcs { before(sbc) } - _, err = rpcVTGate.StreamExecute( - context.Background(), + _, err = vtg.StreamExecute( + ctx, nil, session, "select id from t1", @@ -389,83 +398,83 @@ func testErrorPropagation(t *testing.T, sbcs []*sandboxconn.SandboxConn, before // tablet and a rdonly tablet because we don't control the routing of // Commit. func TestErrorPropagation(t *testing.T) { - createSandbox(KsTestUnsharded) - hcVTGateTest.Reset() + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + vtg, sbc := createVtgateEnv(ctx) + defer vtg.Close() - sbcm := hcVTGateTest.AddTestTablet("aa", "1.1.1.1", 1001, KsTestUnsharded, "0", topodatapb.TabletType_PRIMARY, true, 1, nil) - sbcrdonly := hcVTGateTest.AddTestTablet("aa", "1.1.1.2", 1001, KsTestUnsharded, "0", topodatapb.TabletType_RDONLY, true, 1, nil) sbcs := []*sandboxconn.SandboxConn{ - sbcm, - sbcrdonly, + sbc, } - testErrorPropagation(t, sbcs, func(sbc *sandboxconn.SandboxConn) { + testErrorPropagation(t, ctx, vtg, sbcs, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_CANCELED] = 20 }, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_CANCELED] = 0 }, vtrpcpb.Code_CANCELED) - testErrorPropagation(t, sbcs, func(sbc *sandboxconn.SandboxConn) { + testErrorPropagation(t, ctx, vtg, sbcs, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_UNKNOWN] = 20 }, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_UNKNOWN] = 0 }, vtrpcpb.Code_UNKNOWN) - testErrorPropagation(t, sbcs, func(sbc *sandboxconn.SandboxConn) { + testErrorPropagation(t, ctx, vtg, sbcs, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 20 }, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 0 }, vtrpcpb.Code_INVALID_ARGUMENT) - testErrorPropagation(t, sbcs, func(sbc *sandboxconn.SandboxConn) { + testErrorPropagation(t, ctx, vtg, sbcs, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_DEADLINE_EXCEEDED] = 20 }, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_DEADLINE_EXCEEDED] = 0 }, vtrpcpb.Code_DEADLINE_EXCEEDED) - testErrorPropagation(t, sbcs, func(sbc *sandboxconn.SandboxConn) { + testErrorPropagation(t, ctx, vtg, sbcs, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_ALREADY_EXISTS] = 20 }, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_ALREADY_EXISTS] = 0 }, vtrpcpb.Code_ALREADY_EXISTS) - testErrorPropagation(t, sbcs, func(sbc *sandboxconn.SandboxConn) { + testErrorPropagation(t, ctx, vtg, sbcs, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_PERMISSION_DENIED] = 20 }, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_PERMISSION_DENIED] = 0 }, vtrpcpb.Code_PERMISSION_DENIED) - testErrorPropagation(t, sbcs, func(sbc *sandboxconn.SandboxConn) { + testErrorPropagation(t, ctx, vtg, sbcs, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_RESOURCE_EXHAUSTED] = 20 }, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_RESOURCE_EXHAUSTED] = 0 }, vtrpcpb.Code_RESOURCE_EXHAUSTED) - testErrorPropagation(t, sbcs, func(sbc *sandboxconn.SandboxConn) { + testErrorPropagation(t, ctx, vtg, sbcs, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_FAILED_PRECONDITION] = 20 }, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_FAILED_PRECONDITION] = 0 }, vtrpcpb.Code_FAILED_PRECONDITION) - testErrorPropagation(t, sbcs, func(sbc *sandboxconn.SandboxConn) { + testErrorPropagation(t, ctx, vtg, sbcs, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_ABORTED] = 20 }, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_ABORTED] = 0 }, vtrpcpb.Code_ABORTED) - testErrorPropagation(t, sbcs, func(sbc *sandboxconn.SandboxConn) { + testErrorPropagation(t, ctx, vtg, sbcs, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_INTERNAL] = 20 }, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_INTERNAL] = 0 }, vtrpcpb.Code_INTERNAL) - testErrorPropagation(t, sbcs, func(sbc *sandboxconn.SandboxConn) { + testErrorPropagation(t, ctx, vtg, sbcs, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_UNAVAILABLE] = 20 }, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_UNAVAILABLE] = 0 }, vtrpcpb.Code_UNAVAILABLE) - testErrorPropagation(t, sbcs, func(sbc *sandboxconn.SandboxConn) { + testErrorPropagation(t, ctx, vtg, sbcs, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_UNAUTHENTICATED] = 20 }, func(sbc *sandboxconn.SandboxConn) { sbc.MustFailCodes[vtrpcpb.Code_UNAUTHENTICATED] = 0 @@ -475,18 +484,20 @@ func TestErrorPropagation(t *testing.T) { // This test makes sure that if we start a transaction and hit a critical // error, a rollback is issued. func TestErrorIssuesRollback(t *testing.T) { - createSandbox(KsTestUnsharded) - hcVTGateTest.Reset() - sbc := hcVTGateTest.AddTestTablet("aa", "1.1.1.1", 1001, KsTestUnsharded, "0", topodatapb.TabletType_PRIMARY, true, 1, nil) + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + vtg, sbc := createVtgateEnv(ctx) + defer vtg.Close() // Start a transaction, send one statement. // Simulate an error that should trigger a rollback: // vtrpcpb.Code_ABORTED case. - session, _, err := rpcVTGate.Execute(context.Background(), nil, &vtgatepb.Session{}, "begin", nil) + session, _, err := vtg.Execute(ctx, nil, &vtgatepb.Session{TargetString: KsTestUnsharded + "@primary"}, "begin", nil) if err != nil { t.Fatalf("cannot start a transaction: %v", err) } - session, _, err = rpcVTGate.Execute(context.Background(), nil, session, "select id from t1", nil) + session, _, err = vtg.Execute(ctx, nil, session, "select id from t1", nil) if err != nil { t.Fatalf("want nil, got %v", err) } @@ -494,7 +505,7 @@ func TestErrorIssuesRollback(t *testing.T) { t.Errorf("want 0, got %d", sbc.RollbackCount.Load()) } sbc.MustFailCodes[vtrpcpb.Code_ABORTED] = 20 - _, _, err = rpcVTGate.Execute(context.Background(), nil, session, "select id from t1", nil) + _, _, err = vtg.Execute(ctx, nil, session, "select id from t1", nil) if err == nil { t.Fatalf("want error but got nil") } @@ -507,11 +518,11 @@ func TestErrorIssuesRollback(t *testing.T) { // Start a transaction, send one statement. // Simulate an error that should trigger a rollback: // vtrpcpb.ErrorCode_RESOURCE_EXHAUSTED case. - session, _, err = rpcVTGate.Execute(context.Background(), nil, &vtgatepb.Session{}, "begin", nil) + session, _, err = vtg.Execute(ctx, nil, &vtgatepb.Session{TargetString: KsTestUnsharded + "@primary"}, "begin", nil) if err != nil { t.Fatalf("cannot start a transaction: %v", err) } - session, _, err = rpcVTGate.Execute(context.Background(), nil, session, "select id from t1", nil) + session, _, err = vtg.Execute(ctx, nil, session, "select id from t1", nil) if err != nil { t.Fatalf("want nil, got %v", err) } @@ -519,7 +530,7 @@ func TestErrorIssuesRollback(t *testing.T) { t.Errorf("want 0, got %d", sbc.RollbackCount.Load()) } sbc.MustFailCodes[vtrpcpb.Code_RESOURCE_EXHAUSTED] = 20 - _, _, err = rpcVTGate.Execute(context.Background(), nil, session, "select id from t1", nil) + _, _, err = vtg.Execute(ctx, nil, session, "select id from t1", nil) if err == nil { t.Fatalf("want error but got nil") } @@ -532,11 +543,11 @@ func TestErrorIssuesRollback(t *testing.T) { // Start a transaction, send one statement. // Simulate an error that should *not* trigger a rollback: // vtrpcpb.Code_ALREADY_EXISTS case. - session, _, err = rpcVTGate.Execute(context.Background(), nil, &vtgatepb.Session{}, "begin", nil) + session, _, err = vtg.Execute(ctx, nil, &vtgatepb.Session{TargetString: KsTestUnsharded + "@primary"}, "begin", nil) if err != nil { t.Fatalf("cannot start a transaction: %v", err) } - session, _, err = rpcVTGate.Execute(context.Background(), nil, session, "select id from t1", nil) + session, _, err = vtg.Execute(ctx, nil, session, "select id from t1", nil) if err != nil { t.Fatalf("want nil, got %v", err) } @@ -544,7 +555,7 @@ func TestErrorIssuesRollback(t *testing.T) { t.Errorf("want 0, got %d", sbc.RollbackCount.Load()) } sbc.MustFailCodes[vtrpcpb.Code_ALREADY_EXISTS] = 20 - _, _, err = rpcVTGate.Execute(context.Background(), nil, session, "select id from t1", nil) + _, _, err = vtg.Execute(ctx, nil, session, "select id from t1", nil) if err == nil { t.Fatalf("want error but got nil") } @@ -607,32 +618,40 @@ var shardedVSchemaUnknownParams = ` ` func TestMultiInternalSavepointVtGate(t *testing.T) { - s := createSandbox(KsTestSharded) + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + vtg, _ := createVtgateEnv(ctx) + defer vtg.Close() + + const customKeyspace = "CustomSharding" + s := createSandbox(customKeyspace) s.ShardSpec = "-40-80-" s.VSchema = shardedVSchema srvSchema := getSandboxSrvVSchema() - rpcVTGate.executor.vm.VSchemaUpdate(srvSchema, nil) - hcVTGateTest.Reset() + vtg.executor.vm.VSchemaUpdate(srvSchema, nil) + + hc := vtg.resolver.scatterConn.gateway.hc.(*discovery.FakeHealthCheck) - sbc1 := hcVTGateTest.AddTestTablet("aa", "1.1.1.1", 1, KsTestSharded, "-40", topodatapb.TabletType_PRIMARY, true, 1, nil) - sbc2 := hcVTGateTest.AddTestTablet("aa", "1.1.1.1", 2, KsTestSharded, "40-80", topodatapb.TabletType_PRIMARY, true, 1, nil) - sbc3 := hcVTGateTest.AddTestTablet("aa", "1.1.1.1", 3, KsTestSharded, "80-", topodatapb.TabletType_PRIMARY, true, 1, nil) + sbc1 := hc.AddTestTablet("aa", "-40", 1, customKeyspace, "-40", topodatapb.TabletType_PRIMARY, true, 1, nil) + sbc2 := hc.AddTestTablet("aa", "40-80", 1, customKeyspace, "40-80", topodatapb.TabletType_PRIMARY, true, 1, nil) + sbc3 := hc.AddTestTablet("aa", "80-", 1, customKeyspace, "80-", topodatapb.TabletType_PRIMARY, true, 1, nil) - logChan := rpcVTGate.executor.queryLogger.Subscribe("Test") - defer rpcVTGate.executor.queryLogger.Unsubscribe(logChan) + logChan := vtg.executor.queryLogger.Subscribe("Test") + defer vtg.executor.queryLogger.Unsubscribe(logChan) - session := &vtgatepb.Session{Autocommit: true} + session := &vtgatepb.Session{Autocommit: true, TargetString: customKeyspace + "@primary"} require.True(t, session.GetAutocommit()) require.False(t, session.InTransaction) var err error - session, _, err = rpcVTGate.Execute(context.Background(), nil, session, "begin", nil) + session, _, err = vtg.Execute(ctx, nil, session, "begin", nil) require.NoError(t, err) require.True(t, session.GetAutocommit()) require.True(t, session.InTransaction) // this query goes to multiple shards so internal savepoint will be created. - session, _, err = rpcVTGate.Execute(context.Background(), nil, session, "insert into sp_tbl(user_id) values (1), (3)", nil) + session, _, err = vtg.Execute(ctx, nil, session, "insert into sp_tbl(user_id) values (1), (3)", nil) require.NoError(t, err) require.True(t, session.GetAutocommit()) require.True(t, session.InTransaction) @@ -649,6 +668,7 @@ func TestMultiInternalSavepointVtGate(t *testing.T) { "vtg2": sqltypes.Int64BindVariable(3), }, }} + assertQueriesWithSavepoint(t, sbc1, wantQ) wantQ[1].Sql = "insert into sp_tbl(user_id) values (:_user_id_1)" assertQueriesWithSavepoint(t, sbc2, wantQ) @@ -659,7 +679,7 @@ func TestMultiInternalSavepointVtGate(t *testing.T) { sbc2.Queries = nil // multi shard so new savepoint will be created. - session, _, err = rpcVTGate.Execute(context.Background(), nil, session, "insert into sp_tbl(user_id) values (2), (4)", nil) + session, _, err = vtg.Execute(ctx, nil, session, "insert into sp_tbl(user_id) values (2), (4)", nil) require.NoError(t, err) wantQ = []*querypb.BoundQuery{{ Sql: "savepoint x", @@ -680,7 +700,7 @@ func TestMultiInternalSavepointVtGate(t *testing.T) { sbc3.Queries = nil // single shard so no savepoint will be created and neither any old savepoint will be executed - _, _, err = rpcVTGate.Execute(context.Background(), nil, session, "insert into sp_tbl(user_id) values (5)", nil) + _, _, err = vtg.Execute(ctx, nil, session, "insert into sp_tbl(user_id) values (5)", nil) require.NoError(t, err) wantQ = []*querypb.BoundQuery{{ Sql: "insert into sp_tbl(user_id) values (:_user_id_0)", @@ -691,37 +711,58 @@ func TestMultiInternalSavepointVtGate(t *testing.T) { }} assertQueriesWithSavepoint(t, sbc2, wantQ) - testQueryLog(t, rpcVTGate.executor, logChan, "Execute", "BEGIN", "begin", 0) - testQueryLog(t, rpcVTGate.executor, logChan, "MarkSavepoint", "SAVEPOINT", "savepoint x", 0) - testQueryLog(t, rpcVTGate.executor, logChan, "Execute", "INSERT", "insert into sp_tbl(user_id) values (:vtg1 /* INT64 */), (:vtg2 /* INT64 */)", 2) - testQueryLog(t, rpcVTGate.executor, logChan, "MarkSavepoint", "SAVEPOINT", "savepoint y", 2) - testQueryLog(t, rpcVTGate.executor, logChan, "Execute", "INSERT", "insert into sp_tbl(user_id) values (:vtg1 /* INT64 */), (:vtg2 /* INT64 */)", 2) - testQueryLog(t, rpcVTGate.executor, logChan, "Execute", "INSERT", "insert into sp_tbl(user_id) values (:vtg1 /* INT64 */)", 1) + testQueryLog(t, vtg.executor, logChan, "Execute", "BEGIN", "begin", 0) + testQueryLog(t, vtg.executor, logChan, "MarkSavepoint", "SAVEPOINT", "savepoint x", 0) + testQueryLog(t, vtg.executor, logChan, "Execute", "INSERT", "insert into sp_tbl(user_id) values (:vtg1 /* INT64 */), (:vtg2 /* INT64 */)", 2) + testQueryLog(t, vtg.executor, logChan, "MarkSavepoint", "SAVEPOINT", "savepoint y", 2) + testQueryLog(t, vtg.executor, logChan, "Execute", "INSERT", "insert into sp_tbl(user_id) values (:vtg1 /* INT64 */), (:vtg2 /* INT64 */)", 2) + testQueryLog(t, vtg.executor, logChan, "Execute", "INSERT", "insert into sp_tbl(user_id) values (:vtg1 /* INT64 */)", 1) } func TestVSchemaVindexUnknownParams(t *testing.T) { - s := createSandbox(KsTestSharded) + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + vtg, _ := createVtgateEnv(ctx) + defer vtg.Close() + + const customKeyspace = "CustomSharding" + s := createSandbox(customKeyspace) s.ShardSpec = "-40-80-" - s.VSchema = shardedVSchema srvSchema := getSandboxSrvVSchema() - rpcVTGate.executor.vm.VSchemaUpdate(srvSchema, nil) - hcVTGateTest.Reset() + vtg.executor.vm.VSchemaUpdate(srvSchema, nil) + + hc := vtg.resolver.scatterConn.gateway.hc.(*discovery.FakeHealthCheck) + _ = hc.AddTestTablet("aa", "-40", 1, customKeyspace, "-40", topodatapb.TabletType_PRIMARY, true, 1, nil) + _ = hc.AddTestTablet("aa", "40-80", 1, customKeyspace, "40-80", topodatapb.TabletType_PRIMARY, true, 1, nil) + _ = hc.AddTestTablet("aa", "80-", 1, customKeyspace, "80-", topodatapb.TabletType_PRIMARY, true, 1, nil) unknownParams := vindexUnknownParams.Get() require.Equal(t, int64(0), unknownParams) s.VSchema = shardedVSchemaUnknownParams srvSchema = getSandboxSrvVSchema() - rpcVTGate.executor.vm.VSchemaUpdate(srvSchema, nil) + vtg.executor.vm.VSchemaUpdate(srvSchema, nil) unknownParams = vindexUnknownParams.Get() require.Equal(t, int64(3), unknownParams) s.VSchema = shardedVSchema srvSchema = getSandboxSrvVSchema() - rpcVTGate.executor.vm.VSchemaUpdate(srvSchema, nil) + vtg.executor.vm.VSchemaUpdate(srvSchema, nil) unknownParams = vindexUnknownParams.Get() require.Equal(t, int64(0), unknownParams) } + +func createVtgateEnv(ctx context.Context) (*VTGate, *sandboxconn.SandboxConn) { + cell := "aa" + sb := createSandbox(KsTestSharded) + sb.ShardSpec = "-" + executor, _, _, sbc := createExecutorEnv(ctx) + executor.normalize = normalizeQueries + + vsm := newVStreamManager(executor.resolver.resolver, executor.serv, cell) + return newVTGate(executor, executor.resolver, vsm, nil, executor.scatterConn.gateway), sbc +} From 7798273437c5f6f505e1c9c543da94756a0cce00 Mon Sep 17 00:00:00 2001 From: Dirkjan Bussink Date: Wed, 23 Aug 2023 13:16:11 +0200 Subject: [PATCH 05/17] Fix additional tests and cleanup Signed-off-by: Dirkjan Bussink --- go/stats/rates_test.go | 3 + go/test/fuzzing/tablet_manager_fuzzer.go | 1 + go/test/utils/noleak.go | 1 + go/timer/timer.go | 18 ++--- .../binlog/binlogplayer/binlog_player_test.go | 32 ++++++-- go/vt/mysqlctl/backup_blackbox_test.go | 45 ++++++++---- go/vt/mysqlctl/backup_test.go | 16 +++- go/vt/vtadmin/cluster/cluster_test.go | 2 +- .../endtoend/init_shard_primary_test.go | 4 + go/vt/vtexplain/vtexplain.go | 5 ++ go/vt/vtexplain/vtexplain_test.go | 34 +++------ go/vt/vtexplain/vtexplain_vttablet_test.go | 1 + go/vt/vttablet/tabletserver/connpool/pool.go | 1 + .../vttablet/tabletserver/health_streamer.go | 4 + go/vt/vttablet/tabletserver/query_engine.go | 2 +- .../vttablet/tabletserver/tabletenv/stats.go | 4 + go/vt/vttablet/tabletserver/tabletserver.go | 5 +- .../vttablet/tabletserver/throttle/client.go | 2 +- .../tabletserver/throttle/throttler.go | 73 ++++++++++--------- .../tabletserver/vstreamer/engine_test.go | 1 + .../vstreamer/vstreamer_flaky_test.go | 8 ++ 21 files changed, 169 insertions(+), 93 deletions(-) diff --git a/go/stats/rates_test.go b/go/stats/rates_test.go index 0518cc73b1b..a25a055020a 100644 --- a/go/stats/rates_test.go +++ b/go/stats/rates_test.go @@ -44,6 +44,7 @@ func TestRates(t *testing.T) { clearStats() c := NewCountersWithSingleLabel("rcounter1", "rcounter help", "label") r := NewRates("rates1", c, 3, -1*time.Second) + defer r.Stop() r.snapshot() now = now.Add(epsilon) c.Add("tag1", 0) @@ -92,6 +93,7 @@ func TestRatesConsistency(t *testing.T) { clearStats() c := NewCountersWithSingleLabel("rcounter4", "rcounter4 help", "label") r := NewRates("rates4", c, 100, -1*time.Second) + defer r.Stop() r.snapshot() now = now.Add(epsilon) @@ -133,6 +135,7 @@ func TestRatesHook(t *testing.T) { }) v := NewRates("rates2", c, 2, 10*time.Second) + defer v.Stop() if gotname != "rates2" { t.Errorf("want rates2, got %s", gotname) } diff --git a/go/test/fuzzing/tablet_manager_fuzzer.go b/go/test/fuzzing/tablet_manager_fuzzer.go index 316cf75fb82..4c61afa64bc 100644 --- a/go/test/fuzzing/tablet_manager_fuzzer.go +++ b/go/test/fuzzing/tablet_manager_fuzzer.go @@ -41,6 +41,7 @@ func FuzzTabletManagerExecuteFetchAsDba(data []byte) int { ctx := context.Background() cp := mysql.ConnParams{} db := fakesqldb.New(t) + defer db.Close() db.AddQueryPattern(".*", &sqltypes.Result{}) daemon := mysqlctl.NewFakeMysqlDaemon(db) diff --git a/go/test/utils/noleak.go b/go/test/utils/noleak.go index 0f0258a5d9b..adfceb90b9a 100644 --- a/go/test/utils/noleak.go +++ b/go/test/utils/noleak.go @@ -28,6 +28,7 @@ func ensureNoGoroutines(t testing.TB) { goleak.IgnoreTopFunction("vitess.io/vitess/go/vt/vtgate.processQueryInfo"), goleak.IgnoreTopFunction("github.com/patrickmn/go-cache.(*janitor).Run"), goleak.IgnoreTopFunction("vitess.io/vitess/go/vt/logutil.(*ThrottledLogger).log.func1"), + goleak.IgnoreTopFunction("vitess.io/vitess/go/vt/vttablet/tabletserver/throttle.initThrottleTicker.func1.1"), } var err error diff --git a/go/timer/timer.go b/go/timer/timer.go index 5a28820274f..5407190ef55 100644 --- a/go/timer/timer.go +++ b/go/timer/timer.go @@ -59,7 +59,7 @@ type Timer struct { // state management mu sync.Mutex - running bool + running atomic.Bool // msg is used for out-of-band messages msg chan typeAction @@ -78,10 +78,10 @@ func NewTimer(interval time.Duration) *Timer { func (tm *Timer) Start(keephouse func()) { tm.mu.Lock() defer tm.mu.Unlock() - if tm.running { + isRunning := tm.running.Swap(true) + if isRunning { return } - tm.running = true go tm.run(keephouse) } @@ -118,7 +118,7 @@ func (tm *Timer) SetInterval(ns time.Duration) { tm.interval.Store(ns.Nanoseconds()) tm.mu.Lock() defer tm.mu.Unlock() - if tm.running { + if tm.running.Load() { tm.msg <- timerReset } } @@ -128,7 +128,7 @@ func (tm *Timer) SetInterval(ns time.Duration) { func (tm *Timer) Trigger() { tm.mu.Lock() defer tm.mu.Unlock() - if tm.running { + if tm.running.Load() { tm.msg <- timerTrigger } } @@ -146,9 +146,9 @@ func (tm *Timer) TriggerAfter(duration time.Duration) { func (tm *Timer) Stop() { tm.mu.Lock() defer tm.mu.Unlock() - if tm.running { + isRunning := tm.running.Swap(false) + if isRunning { tm.msg <- timerStop - tm.running = false } } @@ -158,7 +158,5 @@ func (tm *Timer) Interval() time.Duration { } func (tm *Timer) Running() bool { - tm.mu.Lock() - defer tm.mu.Unlock() - return tm.running + return tm.running.Load() } diff --git a/go/vt/binlog/binlogplayer/binlog_player_test.go b/go/vt/binlog/binlogplayer/binlog_player_test.go index db091139331..148c4fb386b 100644 --- a/go/vt/binlog/binlogplayer/binlog_player_test.go +++ b/go/vt/binlog/binlogplayer/binlog_player_test.go @@ -86,7 +86,9 @@ func TestNewBinlogPlayerKeyRange(t *testing.T) { } wantKeyRange := &topodatapb.KeyRange{End: []byte{0x80}} - blp := NewBinlogPlayerKeyRange(dbClient, wantTablet, wantKeyRange, 1, NewStats()) + stats := NewStats() + defer stats.Stop() + blp := NewBinlogPlayerKeyRange(dbClient, wantTablet, wantKeyRange, 1, stats) errfunc := applyEvents(blp) dbClient.Wait() @@ -117,7 +119,9 @@ func TestNewBinlogPlayerTables(t *testing.T) { } wantTables := []string{"a", "b"} - blp := NewBinlogPlayerTables(dbClient, wantTablet, wantTables, 1, NewStats()) + stats := NewStats() + defer stats.Stop() + blp := NewBinlogPlayerTables(dbClient, wantTablet, wantTables, 1, stats) errfunc := applyEvents(blp) dbClient.Wait() @@ -138,7 +142,9 @@ func TestApplyEventsFail(t *testing.T) { _ = newFakeBinlogClient() - blp := NewBinlogPlayerTables(dbClient, nil, []string{"a"}, 1, NewStats()) + stats := NewStats() + defer stats.Stop() + blp := NewBinlogPlayerTables(dbClient, nil, []string{"a"}, 1, stats) errfunc := applyEvents(blp) dbClient.Wait() @@ -188,7 +194,9 @@ func TestStopPosEqual(t *testing.T) { _ = newFakeBinlogClient() - blp := NewBinlogPlayerTables(dbClient, nil, []string{"a"}, 1, NewStats()) + stats := NewStats() + defer stats.Stop() + blp := NewBinlogPlayerTables(dbClient, nil, []string{"a"}, 1, stats) errfunc := applyEvents(blp) dbClient.Wait() @@ -225,7 +233,9 @@ func TestStopPosLess(t *testing.T) { _ = newFakeBinlogClient() - blp := NewBinlogPlayerTables(dbClient, nil, []string{"a"}, 1, NewStats()) + stats := NewStats() + defer stats.Stop() + blp := NewBinlogPlayerTables(dbClient, nil, []string{"a"}, 1, stats) errfunc := applyEvents(blp) dbClient.Wait() @@ -266,7 +276,9 @@ func TestStopPosGreater(t *testing.T) { _ = newFakeBinlogClient() - blp := NewBinlogPlayerTables(dbClient, nil, []string{"a"}, 1, NewStats()) + stats := NewStats() + defer stats.Stop() + blp := NewBinlogPlayerTables(dbClient, nil, []string{"a"}, 1, stats) errfunc := applyEvents(blp) dbClient.Wait() @@ -307,7 +319,9 @@ func TestContextCancel(t *testing.T) { _ = newFakeBinlogClient() - blp := NewBinlogPlayerTables(dbClient, nil, []string{"a"}, 1, NewStats()) + stats := NewStats() + defer stats.Stop() + blp := NewBinlogPlayerTables(dbClient, nil, []string{"a"}, 1, stats) errfunc := applyEvents(blp) dbClient.Wait() @@ -335,7 +349,9 @@ func TestRetryOnDeadlock(t *testing.T) { dbClient.ExpectRequestRE("update _vt.vreplication set pos='MariaDB/0-1-1235', time_updated=.*", testDMLResponse, nil) dbClient.ExpectRequest("commit", nil, nil) - blp := NewBinlogPlayerTables(dbClient, nil, []string{"a"}, 1, NewStats()) + stats := NewStats() + defer stats.Stop() + blp := NewBinlogPlayerTables(dbClient, nil, []string{"a"}, 1, stats) blp.deadlockRetry = 10 * time.Millisecond errfunc := applyEvents(blp) diff --git a/go/vt/mysqlctl/backup_blackbox_test.go b/go/vt/mysqlctl/backup_blackbox_test.go index 5764eefc82b..bf88f4b1791 100644 --- a/go/vt/mysqlctl/backup_blackbox_test.go +++ b/go/vt/mysqlctl/backup_blackbox_test.go @@ -138,7 +138,10 @@ func TestExecuteBackup(t *testing.T) { // Spin up a fake daemon to be used in backups. It needs to be allowed to receive: // "STOP SLAVE", "START SLAVE", in that order. - mysqld := mysqlctl.NewFakeMysqlDaemon(fakesqldb.New(t)) + fakedb := fakesqldb.New(t) + defer fakedb.Close() + mysqld := mysqlctl.NewFakeMysqlDaemon(fakedb) + defer mysqld.Close() mysqld.ExpectedExecuteSuperQueryList = []string{"STOP SLAVE", "START SLAVE"} // mysqld.ShutdownTime = time.Minute @@ -283,7 +286,10 @@ func TestExecuteBackupWithSafeUpgrade(t *testing.T) { // Spin up a fake daemon to be used in backups. It needs to be allowed to receive: // "STOP SLAVE", "START SLAVE", in that order. // It also needs to be allowed to receive the query to disable the innodb_fast_shutdown flag. - mysqld := mysqlctl.NewFakeMysqlDaemon(fakesqldb.New(t)) + fakedb := fakesqldb.New(t) + defer fakedb.Close() + mysqld := mysqlctl.NewFakeMysqlDaemon(fakedb) + defer mysqld.Close() mysqld.ExpectedExecuteSuperQueryList = []string{"STOP SLAVE", "START SLAVE"} mysqld.FetchSuperQueryMap = map[string]*sqltypes.Result{ "SET GLOBAL innodb_fast_shutdown=0": {}, @@ -330,9 +336,6 @@ func TestExecuteBackupWithCanceledContext(t *testing.T) { require.NoError(t, createBackupFiles(path.Join(dataDir, "test2"), 2, "ibd")) defer os.RemoveAll(backupRoot) - // Cancel the context deliberately - cancelledCtx, cancelCtx := context.WithCancel(context.Background()) - cancelCtx() needIt, err := needInnoDBRedoLogSubdir() require.NoError(t, err) if needIt { @@ -347,16 +350,16 @@ func TestExecuteBackupWithCanceledContext(t *testing.T) { ts := memorytopo.NewServer(ctx, "cell1") defer ts.Close() - require.NoError(t, ts.CreateKeyspace(cancelledCtx, keyspace, &topodata.Keyspace{})) - require.NoError(t, ts.CreateShard(cancelledCtx, keyspace, shard)) + require.NoError(t, ts.CreateKeyspace(ctx, keyspace, &topodata.Keyspace{})) + require.NoError(t, ts.CreateShard(ctx, keyspace, shard)) tablet := topo.NewTablet(100, "cell1", "mykeyspace-00-80-0100") tablet.Keyspace = keyspace tablet.Shard = shard - require.NoError(t, ts.CreateTablet(cancelledCtx, tablet)) + require.NoError(t, ts.CreateTablet(ctx, tablet)) - _, err = ts.UpdateShardFields(cancelledCtx, keyspace, shard, func(si *topo.ShardInfo) error { + _, err = ts.UpdateShardFields(ctx, keyspace, shard, func(si *topo.ShardInfo) error { si.PrimaryAlias = &topodata.TabletAlias{Uid: 100, Cell: "cell1"} now := time.Now() @@ -370,9 +373,16 @@ func TestExecuteBackupWithCanceledContext(t *testing.T) { bh := filebackupstorage.NewBackupHandle(nil, "", "", false) // Spin up a fake daemon to be used in backups. It needs to be allowed to receive: // "STOP SLAVE", "START SLAVE", in that order. - mysqld := mysqlctl.NewFakeMysqlDaemon(fakesqldb.New(t)) + fakedb := fakesqldb.New(t) + defer fakedb.Close() + mysqld := mysqlctl.NewFakeMysqlDaemon(fakedb) + defer mysqld.Close() mysqld.ExpectedExecuteSuperQueryList = []string{"STOP SLAVE", "START SLAVE"} + // Cancel the context deliberately + cancelledCtx, cancelCtx := context.WithCancel(context.Background()) + cancelCtx() + ok, err := be.ExecuteBackup(cancelledCtx, mysqlctl.BackupParams{ Logger: logutil.NewConsoleLogger(), Mysqld: mysqld, @@ -453,7 +463,10 @@ func TestExecuteRestoreWithTimedOutContext(t *testing.T) { bh := filebackupstorage.NewBackupHandle(nil, "", "", false) // Spin up a fake daemon to be used in backups. It needs to be allowed to receive: // "STOP SLAVE", "START SLAVE", in that order. - mysqld := mysqlctl.NewFakeMysqlDaemon(fakesqldb.New(t)) + fakedb := fakesqldb.New(t) + defer fakedb.Close() + mysqld := mysqlctl.NewFakeMysqlDaemon(fakedb) + defer mysqld.Close() mysqld.ExpectedExecuteSuperQueryList = []string{"STOP SLAVE", "START SLAVE"} ok, err := be.ExecuteBackup(ctx, mysqlctl.BackupParams{ @@ -477,7 +490,10 @@ func TestExecuteRestoreWithTimedOutContext(t *testing.T) { // Now try to restore the above backup. bh = filebackupstorage.NewBackupHandle(nil, "", "", true) - mysqld = mysqlctl.NewFakeMysqlDaemon(fakesqldb.New(t)) + fakedb = fakesqldb.New(t) + defer fakedb.Close() + mysqld = mysqlctl.NewFakeMysqlDaemon(fakedb) + defer mysqld.Close() mysqld.ExpectedExecuteSuperQueryList = []string{"STOP SLAVE", "START SLAVE"} fakeStats := backupstats.NewFakeStats() @@ -551,7 +567,10 @@ func TestExecuteRestoreWithTimedOutContext(t *testing.T) { require.Equal(t, 4, sourceReadStats) // Restore using timed-out context - mysqld = mysqlctl.NewFakeMysqlDaemon(fakesqldb.New(t)) + fakedb = fakesqldb.New(t) + defer fakedb.Close() + mysqld = mysqlctl.NewFakeMysqlDaemon(fakedb) + defer mysqld.Close() mysqld.ExpectedExecuteSuperQueryList = []string{"STOP SLAVE", "START SLAVE"} restoreParams.Mysqld = mysqld timedOutCtx, cancel := context.WithTimeout(ctx, 1*time.Second) diff --git a/go/vt/mysqlctl/backup_test.go b/go/vt/mysqlctl/backup_test.go index 75f54ebe66b..878adefb9a6 100644 --- a/go/vt/mysqlctl/backup_test.go +++ b/go/vt/mysqlctl/backup_test.go @@ -31,6 +31,8 @@ import ( "github.com/stretchr/testify/require" + "vitess.io/vitess/go/test/utils" + "vitess.io/vitess/go/mysql/replication" "vitess.io/vitess/go/mysql" @@ -43,6 +45,7 @@ import ( // TestBackupExecutesBackupWithScopedParams tests that Backup passes // a Scope()-ed stats to backupengine ExecuteBackup. func TestBackupExecutesBackupWithScopedParams(t *testing.T) { + utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() @@ -65,6 +68,7 @@ func TestBackupExecutesBackupWithScopedParams(t *testing.T) { // TestBackupNoStats tests that if BackupParams.Stats is nil, then Backup will // pass non-nil Stats to sub-components. func TestBackupNoStats(t *testing.T) { + utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() @@ -80,6 +84,7 @@ func TestBackupNoStats(t *testing.T) { // TestBackupParameterizesBackupStorageWithScopedStats tests that Backup passes // a Scope()-ed stats to BackupStorage.WithParams. func TestBackupParameterizesBackupStorageWithScopedStats(t *testing.T) { + utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() @@ -100,6 +105,7 @@ func TestBackupParameterizesBackupStorageWithScopedStats(t *testing.T) { // TestBackupEmitsStats tests that Backup emits stats. func TestBackupEmitsStats(t *testing.T) { + utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() @@ -117,6 +123,7 @@ func TestBackupEmitsStats(t *testing.T) { // backupstorage.Params to backupstorage, but only if it responds to // backupstorage.WithParams. func TestBackupTriesToParameterizeBackupStorage(t *testing.T) { + utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() @@ -322,6 +329,7 @@ func TestFindFilesToBackupWithRedoLog(t *testing.T) { // TestRestoreEmitsStats tests that Restore emits stats. func TestRestoreEmitsStats(t *testing.T) { + utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() @@ -339,6 +347,7 @@ func TestRestoreEmitsStats(t *testing.T) { // TestRestoreExecutesRestoreWithScopedParams tests that Restore passes // a Scope()-ed stats to backupengine ExecuteRestore. func TestRestoreExecutesRestoreWithScopedParams(t *testing.T) { + utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() @@ -362,6 +371,7 @@ func TestRestoreExecutesRestoreWithScopedParams(t *testing.T) { // TestRestoreNoStats tests that if RestoreParams.Stats is nil, then Restore will // pass non-nil Stats to sub-components. func TestRestoreNoStats(t *testing.T) { + utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() @@ -378,6 +388,7 @@ func TestRestoreNoStats(t *testing.T) { // TestRestoreParameterizesBackupStorageWithScopedStats tests that Restore passes // a Scope()-ed stats to BackupStorage.WithParams. func TestRestoreParameterizesBackupStorageWithScopedStats(t *testing.T) { + utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() @@ -401,6 +412,7 @@ func TestRestoreParameterizesBackupStorageWithScopedStats(t *testing.T) { // backupstorage.Params to backupstorage, but only if it responds to // backupstorage.WithParams. func TestRestoreTriesToParameterizeBackupStorage(t *testing.T) { + utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() @@ -512,6 +524,7 @@ func TestRestoreManifestMySQLVersionValidation(t *testing.T) { for _, tc := range testCases { t.Run(fmt.Sprintf("%s->%s upgradeSafe=%t", tc.fromVersion, tc.toVersion, tc.upgradeSafe), func(t *testing.T) { + utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() env.mysqld.Version = tc.toVersion @@ -575,7 +588,6 @@ func createFakeBackupRestoreEnv(t *testing.T) (*fakeBackupRestoreEnv, func()) { sqldb.SetNeverFail(true) mysqld := NewFakeMysqlDaemon(sqldb) require.Nil(t, mysqld.Shutdown(ctx, nil, false)) - defer mysqld.Close() dirName, err := os.MkdirTemp("", "vt_backup_test") require.Nil(t, err) @@ -660,6 +672,8 @@ func createFakeBackupRestoreEnv(t *testing.T) (*fakeBackupRestoreEnv, func()) { delete(backupstorage.BackupStorageMap, "fake") backupstorage.BackupStorageImplementation = previousBackupStorageImplementation + mysqld.Close() + sqldb.Close() } return &fakeBackupRestoreEnv{ diff --git a/go/vt/vtadmin/cluster/cluster_test.go b/go/vt/vtadmin/cluster/cluster_test.go index 8111329acd5..570fbaf0cb4 100644 --- a/go/vt/vtadmin/cluster/cluster_test.go +++ b/go/vt/vtadmin/cluster/cluster_test.go @@ -50,7 +50,7 @@ import ( ) func TestCreateKeyspace(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) t.Parallel() ctx, cancel := context.WithCancel(context.Background()) diff --git a/go/vt/vtctl/grpcvtctldserver/endtoend/init_shard_primary_test.go b/go/vt/vtctl/grpcvtctldserver/endtoend/init_shard_primary_test.go index a80fb89638a..8fbfaad8410 100644 --- a/go/vt/vtctl/grpcvtctldserver/endtoend/init_shard_primary_test.go +++ b/go/vt/vtctl/grpcvtctldserver/endtoend/init_shard_primary_test.go @@ -21,6 +21,7 @@ import ( "fmt" "testing" + "vitess.io/vitess/go/test/utils" "vitess.io/vitess/go/vt/mysqlctl" "github.com/stretchr/testify/assert" @@ -41,6 +42,7 @@ import ( ) func TestInitShardPrimary(t *testing.T) { + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() ts := memorytopo.NewServer(ctx, "cell1") @@ -48,6 +50,7 @@ func TestInitShardPrimary(t *testing.T) { wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmc) primaryDb := fakesqldb.New(t) + defer primaryDb.Close() primaryDb.AddQuery("create database if not exists `vt_test_keyspace`", &sqltypes.Result{InsertID: 0, RowsAffected: 0}) tablet1 := testlib.NewFakeTablet(t, wr, "cell1", 0, topodatapb.TabletType_PRIMARY, primaryDb) @@ -110,6 +113,7 @@ func TestInitShardPrimaryNoFormerPrimary(t *testing.T) { wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmc) primaryDb := fakesqldb.New(t) + defer primaryDb.Close() primaryDb.AddQuery("create database if not exists `vt_test_keyspace`", &sqltypes.Result{InsertID: 0, RowsAffected: 0}) tablet1 := testlib.NewFakeTablet(t, wr, "cell1", 0, topodatapb.TabletType_REPLICA, primaryDb) diff --git a/go/vt/vtexplain/vtexplain.go b/go/vt/vtexplain/vtexplain.go index 4b3c4c3bf47..55e76606e08 100644 --- a/go/vt/vtexplain/vtexplain.go +++ b/go/vt/vtexplain/vtexplain.go @@ -211,10 +211,15 @@ func Init(ctx context.Context, vSchemaStr, sqlSchema, ksShardMapStr string, opts // Stop and cleans up fake execution environment func (vte *VTExplain) Stop() { + if vte.vtgateExecutor != nil { + vte.vtgateExecutor.Close() + } + // Cleanup all created fake dbs. if vte.explainTopo != nil { for _, conn := range vte.explainTopo.TabletConns { conn.tsv.StopService() + conn.tsv.Close(context.Background()) } for _, conn := range vte.explainTopo.TabletConns { conn.db.Close() diff --git a/go/vt/vtexplain/vtexplain_test.go b/go/vt/vtexplain/vtexplain_test.go index 29b2cfdf415..85b41675e19 100644 --- a/go/vt/vtexplain/vtexplain_test.go +++ b/go/vt/vtexplain/vtexplain_test.go @@ -28,12 +28,12 @@ import ( "github.com/google/go-cmp/cmp" "github.com/stretchr/testify/require" + "vitess.io/vitess/go/test/utils" + "vitess.io/vitess/go/vt/key" "vitess.io/vitess/go/vt/proto/topodata" "vitess.io/vitess/go/vt/topo" "vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv/tabletenvtest" - - querypb "vitess.io/vitess/go/vt/proto/query" ) func defaultTestOpts() *Options { @@ -86,10 +86,12 @@ func testExplain(testcase string, opts *Options, t *testing.T) { func runTestCase(testcase, mode string, opts *Options, topts *testopts, t *testing.T) { t.Run(testcase, func(t *testing.T) { + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() vte := initTest(ctx, mode, opts, topts, t) + defer vte.Stop() sqlFile := fmt.Sprintf("testdata/%s-queries.sql", testcase) sql, err := os.ReadFile(sqlFile) @@ -145,28 +147,6 @@ func TestExplain(t *testing.T) { } tests := []test{ {"unsharded", defaultTestOpts()}, - {"selectsharded", defaultTestOpts()}, - {"insertsharded", defaultTestOpts()}, - {"updatesharded", defaultTestOpts()}, - {"deletesharded", defaultTestOpts()}, - {"comments", defaultTestOpts()}, - {"options", &Options{ - ReplicationMode: "STATEMENT", - NumShards: 4, - Normalize: false, - }}, - {"target", &Options{ - ReplicationMode: "ROW", - NumShards: 4, - Normalize: false, - Target: "ks_sharded/40-80", - }}, - {"gen4", &Options{ - ReplicationMode: "ROW", - NumShards: 4, - Normalize: true, - PlannerVersion: querypb.ExecuteOptions_Gen4, - }}, } for _, tst := range tests { @@ -175,10 +155,13 @@ func TestExplain(t *testing.T) { } func TestErrors(t *testing.T) { + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) defer cancel() vte := initTest(ctx, ModeMulti, defaultTestOpts(), &testopts{}, t) + defer vte.Stop() tests := []struct { SQL string @@ -215,10 +198,13 @@ func TestErrors(t *testing.T) { } func TestJSONOutput(t *testing.T) { + utils.EnsureNoLeaks(t) + ctx, cancel := context.WithCancel(context.Background()) defer cancel() vte := initTest(ctx, ModeMulti, defaultTestOpts(), &testopts{}, t) + defer vte.Stop() sql := "select 1 from user where id = 1" explains, err := vte.Run(sql) require.NoError(t, err, "vtexplain error") diff --git a/go/vt/vtexplain/vtexplain_vttablet_test.go b/go/vt/vtexplain/vtexplain_vttablet_test.go index bea8a91e3fc..614ad186224 100644 --- a/go/vt/vtexplain/vtexplain_vttablet_test.go +++ b/go/vt/vtexplain/vtexplain_vttablet_test.go @@ -126,6 +126,7 @@ create table test_partitioned ( defer cancel() vte := initTest(ctx, ModeMulti, defaultTestOpts(), &testopts{}, t) + defer vte.Stop() tabletEnv, _ := newTabletEnvironment(ddls, defaultTestOpts()) vte.setGlobalTabletEnv(tabletEnv) diff --git a/go/vt/vttablet/tabletserver/connpool/pool.go b/go/vt/vttablet/tabletserver/connpool/pool.go index 6cb9adf7387..a67de9f2a1a 100644 --- a/go/vt/vttablet/tabletserver/connpool/pool.go +++ b/go/vt/vttablet/tabletserver/connpool/pool.go @@ -170,6 +170,7 @@ func (cp *Pool) Close() { log.Infof("connpool - acquiring lock") cp.mu.Lock() log.Infof("connpool - acquired lock") + cp.connections.Close() cp.connections = nil cp.mu.Unlock() log.Infof("connpool - closing dbaPool") diff --git a/go/vt/vttablet/tabletserver/health_streamer.go b/go/vt/vttablet/tabletserver/health_streamer.go index 1170e4fbc39..fcd8644585b 100644 --- a/go/vt/vttablet/tabletserver/health_streamer.go +++ b/go/vt/vttablet/tabletserver/health_streamer.go @@ -156,6 +156,10 @@ func (hs *healthStreamer) Close() { hs.cancel() hs.cancel = nil } + if hs.conns != nil { + hs.conns.Close() + hs.conns = nil + } } func (hs *healthStreamer) Stream(ctx context.Context, callback func(*querypb.StreamHealthResponse) error) error { diff --git a/go/vt/vttablet/tabletserver/query_engine.go b/go/vt/vttablet/tabletserver/query_engine.go index a1185d16fee..7fc8b1efe30 100644 --- a/go/vt/vttablet/tabletserver/query_engine.go +++ b/go/vt/vttablet/tabletserver/query_engine.go @@ -314,7 +314,7 @@ func (qe *QueryEngine) Close() { } // Close in reverse order of Open. qe.se.UnregisterNotifier("qe") - qe.plans.Clear() + qe.plans.Close() qe.tables = make(map[string]*schema.Table) qe.streamConns.Close() qe.conns.Close() diff --git a/go/vt/vttablet/tabletserver/tabletenv/stats.go b/go/vt/vttablet/tabletserver/tabletenv/stats.go index 17b7101402d..1ad93532719 100644 --- a/go/vt/vttablet/tabletserver/tabletenv/stats.go +++ b/go/vt/vttablet/tabletserver/tabletenv/stats.go @@ -102,3 +102,7 @@ func NewStats(exporter *servenv.Exporter) *Stats { stats.QPSRates = exporter.NewRates("QPS", stats.QueryTimings, 15*60/5, 5*time.Second) return stats } + +func (st *Stats) Stop() { + st.QPSRates.Stop() +} diff --git a/go/vt/vttablet/tabletserver/tabletserver.go b/go/vt/vttablet/tabletserver/tabletserver.go index d106dfd1596..6a154fcea94 100644 --- a/go/vt/vttablet/tabletserver/tabletserver.go +++ b/go/vt/vttablet/tabletserver/tabletserver.go @@ -1721,8 +1721,11 @@ func (tsv *TabletServer) HandlePanic(err *error) { } } -// Close is a no-op. +// Close shuts down any remaining go routines func (tsv *TabletServer) Close(ctx context.Context) error { + tsv.sm.closeAll() + tsv.qe.Close() + tsv.stats.Stop() return nil } diff --git a/go/vt/vttablet/tabletserver/throttle/client.go b/go/vt/vttablet/tabletserver/throttle/client.go index 30d6c79afdf..69bf4bac0cb 100644 --- a/go/vt/vttablet/tabletserver/throttle/client.go +++ b/go/vt/vttablet/tabletserver/throttle/client.go @@ -68,7 +68,7 @@ func NewProductionClient(throttler *Throttler, appName throttlerapp.Name, checkT } } -// NewBackgroundClient creates a client suitable for background jobs, which have low priority over productio ntraffic, +// NewBackgroundClient creates a client suitable for background jobs, which have low priority over production traffic, // e.g. migration, table pruning, vreplication func NewBackgroundClient(throttler *Throttler, appName throttlerapp.Name, checkType ThrottleCheckType) *Client { initThrottleTicker() diff --git a/go/vt/vttablet/tabletserver/throttle/throttler.go b/go/vt/vttablet/tabletserver/throttle/throttler.go index c703bce8c51..260f024b72b 100644 --- a/go/vt/vttablet/tabletserver/throttle/throttler.go +++ b/go/vt/vttablet/tabletserver/throttle/throttler.go @@ -122,9 +122,9 @@ type Throttler struct { cell string check *ThrottlerCheck - isEnabled int64 - isLeader int64 - isOpen int64 + isEnabled atomic.Bool + isLeader atomic.Bool + isOpen atomic.Bool env tabletenv.Env pool *connpool.Pool @@ -162,6 +162,7 @@ type Throttler struct { initMutex sync.Mutex enableMutex sync.Mutex + cancelOpenContext context.CancelFunc cancelEnableContext context.CancelFunc throttledAppsMutex sync.Mutex watchSrvKeyspaceOnce sync.Once @@ -190,9 +191,6 @@ type ThrottlerStatus struct { // NewThrottler creates a Throttler func NewThrottler(env tabletenv.Env, srvTopoServer srvtopo.Server, ts *topo.Server, cell string, heartbeatWriter heartbeat.HeartbeatWriter, tabletTypeFunc func() topodatapb.TabletType) *Throttler { throttler := &Throttler{ - isLeader: 0, - isOpen: 0, - cell: cell, env: env, tabletTypeFunc: tabletTypeFunc, @@ -359,11 +357,11 @@ func (throttler *Throttler) applyThrottlerConfig(ctx context.Context, throttlerC } func (throttler *Throttler) IsEnabled() bool { - return atomic.LoadInt64(&throttler.isEnabled) > 0 + return throttler.isEnabled.Load() } func (throttler *Throttler) IsOpen() bool { - return atomic.LoadInt64(&throttler.isOpen) > 0 + return throttler.isOpen.Load() } // CheckIsOpen checks if this throttler is ready to serve. If not, it @@ -386,12 +384,12 @@ func (throttler *Throttler) Enable(ctx context.Context) bool { throttler.enableMutex.Lock() defer throttler.enableMutex.Unlock() - if throttler.IsEnabled() { + isEnabled := throttler.isEnabled.Swap(true) + if isEnabled { log.Infof("Throttler: already enabled") return false } log.Infof("Throttler: enabling") - atomic.StoreInt64(&throttler.isEnabled, 1) ctx, throttler.cancelEnableContext = context.WithCancel(ctx) throttler.check.SelfChecks(ctx) @@ -409,14 +407,13 @@ func (throttler *Throttler) Disable(ctx context.Context) bool { throttler.enableMutex.Lock() defer throttler.enableMutex.Unlock() - if !throttler.IsEnabled() { + isEnabled := throttler.isEnabled.Swap(false) + if !isEnabled { log.Infof("Throttler: already disabled") return false } log.Infof("Throttler: disabling") // _ = throttler.updateConfig(ctx, false, throttler.MetricsThreshold.Get()) // TODO(shlomi) - atomic.StoreInt64(&throttler.isEnabled, 0) - throttler.aggregatedMetrics.Flush() throttler.recentApps.Flush() throttler.nonLowPriorityAppRequestsThrottled.Flush() @@ -431,20 +428,22 @@ func (throttler *Throttler) Open() error { log.Infof("Throttler: started execution of Open. Acquiring initMutex lock") throttler.initMutex.Lock() defer throttler.initMutex.Unlock() - if throttler.IsOpen() { + + isOpen := throttler.isOpen.Swap(true) + if isOpen { // already open log.Infof("Throttler: throttler is already open") return nil } log.Infof("Throttler: opening") - ctx := context.Background() + var ctx context.Context + ctx, throttler.cancelOpenContext = context.WithCancel(context.Background()) // The query needs to be dynamically built because the sidecar database name // is not known when the TabletServer is created, which in turn creates the // Throttler. throttler.metricsQuery.Store(sqlparser.BuildParsedQuery(defaultReplicationLagQuery, sidecar.GetIdentifier()).Query) // default throttler.initConfig() throttler.pool.Open(throttler.env.Config().DB.AppWithDB(), throttler.env.Config().DB.DbaWithDB(), throttler.env.Config().DB.AppDebugWithDB()) - atomic.StoreInt64(&throttler.isOpen, 1) throttler.ThrottleApp("always-throttled-app", time.Now().Add(time.Hour*24*365*10), DefaultThrottleRatio, false) @@ -454,7 +453,7 @@ func (throttler *Throttler) Open() error { // opening of all other components. We thus read the throttler config in the background. // However, we want to handle a situation where the read errors out. // So we kick a loop that keeps retrying reading the config, for as long as this throttler is open. - retryReadAndApplyThrottlerConfig := func() { + retryReadAndApplyThrottlerConfig := func(ctx context.Context) { retryInterval := 10 * time.Second retryTicker := time.NewTicker(retryInterval) defer retryTicker.Stop() @@ -465,7 +464,9 @@ func (throttler *Throttler) Open() error { return } - throttlerConfig, err := throttler.readThrottlerConfig(ctx) + requestCtx, requestCancel := context.WithTimeout(ctx, 5*time.Second) + defer requestCancel() + throttlerConfig, err := throttler.readThrottlerConfig(requestCtx) if err == nil { log.Errorf("Throttler.retryReadAndApplyThrottlerConfig(): success reading throttler config: %+v", throttlerConfig) // It's possible that during a retry-sleep, the throttler is closed and opened again, leading @@ -482,10 +483,16 @@ func (throttler *Throttler) Open() error { return } log.Errorf("Throttler.retryReadAndApplyThrottlerConfig(): error reading throttler config. Will retry in %v. Err=%+v", retryInterval, err) - <-retryTicker.C + select { + case <-ctx.Done(): + // Throttler is not open so no need to keep retrying. + log.Errorf("Throttler.retryReadAndApplyThrottlerConfig(): throttler no longer seems to be open, exiting") + return + case <-retryTicker.C: + } } } - go retryReadAndApplyThrottlerConfig() + go retryReadAndApplyThrottlerConfig(ctx) return nil } @@ -496,17 +503,18 @@ func (throttler *Throttler) Close() { throttler.initMutex.Lock() log.Infof("Throttler: acquired initMutex lock") defer throttler.initMutex.Unlock() - if !throttler.IsOpen() { + isOpen := throttler.isOpen.Swap(false) + if !isOpen { log.Infof("Throttler: throttler is not open") return } ctx := context.Background() throttler.Disable(ctx) - atomic.StoreInt64(&throttler.isLeader, 0) + throttler.isLeader.Store(false) log.Infof("Throttler: closing pool") throttler.pool.Close() - atomic.StoreInt64(&throttler.isOpen, 0) + throttler.cancelOpenContext() log.Infof("Throttler: finished execution of Close") } @@ -619,24 +627,23 @@ func (throttler *Throttler) Operate(ctx context.Context) { defer throttler.initMutex.Unlock() // sparse - shouldBeLeader := int64(0) + shouldBeLeader := false if throttler.IsOpen() { if throttler.tabletTypeFunc() == topodatapb.TabletType_PRIMARY { - shouldBeLeader = 1 + shouldBeLeader = true } } + isLeader := throttler.isLeader.Swap(shouldBeLeader) transitionedIntoLeader := false - if shouldBeLeader > throttler.isLeader { + if shouldBeLeader && !isLeader { log.Infof("Throttler: transition into leadership") transitionedIntoLeader = true } - if shouldBeLeader < throttler.isLeader { + if !shouldBeLeader && isLeader { log.Infof("Throttler: transition out of leadership") } - atomic.StoreInt64(&throttler.isLeader, shouldBeLeader) - if transitionedIntoLeader { // transitioned into leadership, let's speed up the next 'refresh' and 'collect' ticks go mysqlRefreshTicker.TickNow() @@ -847,7 +854,7 @@ func (throttler *Throttler) refreshMySQLInventory(ctx context.Context) error { throttler.mysqlClusterProbesChan <- clusterProbes return } - if atomic.LoadInt64(&throttler.isLeader) == 0 { + if !throttler.isLeader.Load() { // not the leader (primary tablet)? Then no more work for us. return } @@ -1167,9 +1174,9 @@ func (throttler *Throttler) Status() *ThrottlerStatus { Keyspace: throttler.keyspace, Shard: throttler.shard, - IsLeader: (atomic.LoadInt64(&throttler.isLeader) > 0), - IsOpen: (atomic.LoadInt64(&throttler.isOpen) > 0), - IsEnabled: (atomic.LoadInt64(&throttler.isEnabled) > 0), + IsLeader: throttler.isLeader.Load(), + IsOpen: throttler.isOpen.Load(), + IsEnabled: throttler.isEnabled.Load(), IsDormant: throttler.isDormant(), Query: throttler.GetMetricsQuery(), diff --git a/go/vt/vttablet/tabletserver/vstreamer/engine_test.go b/go/vt/vttablet/tabletserver/vstreamer/engine_test.go index 825d251365a..7f9a8cdff99 100644 --- a/go/vt/vttablet/tabletserver/vstreamer/engine_test.go +++ b/go/vt/vttablet/tabletserver/vstreamer/engine_test.go @@ -174,6 +174,7 @@ func TestVStreamerWaitForMySQL(t *testing.T) { tableName := "test" expectedWaits := int64(0) testDB := fakesqldb.New(t) + defer testDB.Close() hostres := sqltypes.MakeTestResult(sqltypes.MakeTestFields( "hostname|port", "varchar|int64"), diff --git a/go/vt/vttablet/tabletserver/vstreamer/vstreamer_flaky_test.go b/go/vt/vttablet/tabletserver/vstreamer/vstreamer_flaky_test.go index 70cbf1896c8..fec0c10ded7 100644 --- a/go/vt/vttablet/tabletserver/vstreamer/vstreamer_flaky_test.go +++ b/go/vt/vttablet/tabletserver/vstreamer/vstreamer_flaky_test.go @@ -89,9 +89,13 @@ func TestNoBlob(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) defer cancel() oldEngine := engine + engine = nil + oldEnv := env + env = nil newEngine(t, ctx, "noblob") defer func() { engine = oldEngine + env = oldEnv }() execStatements(t, []string{ "create table t1(id int, blb blob, val varchar(4), primary key(id))", @@ -1936,9 +1940,13 @@ func TestMinimalMode(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) defer cancel() oldEngine := engine + engine = nil + oldEnv := env + env = nil newEngine(t, ctx, "minimal") defer func() { engine = oldEngine + env = oldEnv }() err := engine.Stream(context.Background(), "current", nil, nil, throttlerapp.VStreamerName, func(evs []*binlogdatapb.VEvent) error { return nil }) require.Error(t, err, "minimal binlog_row_image is not supported by Vitess VReplication") From 48b36dc760cbfb4ce8c9624e646ee126ba9d4320 Mon Sep 17 00:00:00 2001 From: Dirkjan Bussink Date: Wed, 23 Aug 2023 13:33:13 +0200 Subject: [PATCH 06/17] Plug more leaks Signed-off-by: Dirkjan Bussink --- go/test/endtoend/vtgate/queries/misc/misc_test.go | 2 ++ go/vt/vtadmin/api_test.go | 1 + go/vt/vtadmin/cluster/cluster.go | 3 ++- go/vt/vtadmin/cluster/cluster_test.go | 11 ++++++----- .../endtoend/init_shard_primary_test.go | 1 + go/vt/vttablet/tabletmanager/tm_init.go | 4 ++++ 6 files changed, 16 insertions(+), 6 deletions(-) diff --git a/go/test/endtoend/vtgate/queries/misc/misc_test.go b/go/test/endtoend/vtgate/queries/misc/misc_test.go index c52ab5687a7..fe9699b8267 100644 --- a/go/test/endtoend/vtgate/queries/misc/misc_test.go +++ b/go/test/endtoend/vtgate/queries/misc/misc_test.go @@ -211,10 +211,12 @@ func TestHighNumberOfParams(t *testing.T) { // connect to the vitess cluster db, err := sql.Open("mysql", fmt.Sprintf("@tcp(%s:%v)/%s", vtParams.Host, vtParams.Port, vtParams.DbName)) require.NoError(t, err) + defer db.Close() // run the query r, err := db.Query(fmt.Sprintf("SELECT /*vt+ QUERY_TIMEOUT_MS=10000 */ id1 FROM t1 WHERE id1 in (%s) ORDER BY id1 ASC", strings.Join(params, ", ")), vals...) require.NoError(t, err) + defer r.Close() // check the results we got, we should get 5 rows with each: 0, 1, 2, 3, 4 // count is the row number we are currently visiting, also correspond to the diff --git a/go/vt/vtadmin/api_test.go b/go/vt/vtadmin/api_test.go index 3bd01ac5821..a93fca0bb7e 100644 --- a/go/vt/vtadmin/api_test.go +++ b/go/vt/vtadmin/api_test.go @@ -5202,6 +5202,7 @@ func TestServeHTTP(t *testing.T) { }, }, }.Cluster(context.Background()) + defer testCluster.Close() tests := []struct { name string diff --git a/go/vt/vtadmin/cluster/cluster.go b/go/vt/vtadmin/cluster/cluster.go index 1341ce2a3a5..a0ddc62f5ad 100644 --- a/go/vt/vtadmin/cluster/cluster.go +++ b/go/vt/vtadmin/cluster/cluster.go @@ -231,12 +231,13 @@ func (c *Cluster) Close() error { rec.RecordError(closer.Close()) }(closer) } + wg.Wait() if rec.HasErrors() { return fmt.Errorf("failed to cleanly close cluster (id=%s): %w", c.ID, rec.Error()) } - return nil + return c.schemaCache.Close() } // ToProto returns a value-copy protobuf equivalent of the cluster. diff --git a/go/vt/vtadmin/cluster/cluster_test.go b/go/vt/vtadmin/cluster/cluster_test.go index 570fbaf0cb4..3d87399c8c9 100644 --- a/go/vt/vtadmin/cluster/cluster_test.go +++ b/go/vt/vtadmin/cluster/cluster_test.go @@ -50,11 +50,8 @@ import ( ) func TestCreateKeyspace(t *testing.T) { - defer utils.EnsureNoLeaks(t) - t.Parallel() + utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() tests := []struct { name string cfg testutil.TestClusterConfig @@ -161,6 +158,9 @@ func TestCreateKeyspace(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + cluster := testutil.BuildCluster(t, tt.cfg) defer cluster.Close() @@ -178,7 +178,6 @@ func TestCreateKeyspace(t *testing.T) { func TestCreateShard(t *testing.T) { utils.EnsureNoLeaks(t) - t.Parallel() type test struct { name string @@ -261,6 +260,7 @@ func TestCreateShard(t *testing.T) { for _, tt := range tests { tt := tt t.Run(tt.name, func(t *testing.T) { + defer tt.tc.Cluster.Close() _, err := tt.tc.Cluster.CreateShard(ctx, tt.req) if tt.shouldErr { assert.Error(t, err) @@ -486,6 +486,7 @@ func TestDeleteShards(t *testing.T) { }() } + defer tt.tc.Cluster.Close() _, err := tt.tc.Cluster.DeleteShards(ctx, tt.req) if tt.shouldErr { assert.Error(t, err) diff --git a/go/vt/vtctl/grpcvtctldserver/endtoend/init_shard_primary_test.go b/go/vt/vtctl/grpcvtctldserver/endtoend/init_shard_primary_test.go index 8fbfaad8410..8de5eceaeff 100644 --- a/go/vt/vtctl/grpcvtctldserver/endtoend/init_shard_primary_test.go +++ b/go/vt/vtctl/grpcvtctldserver/endtoend/init_shard_primary_test.go @@ -47,6 +47,7 @@ func TestInitShardPrimary(t *testing.T) { defer cancel() ts := memorytopo.NewServer(ctx, "cell1") tmc := tmclient.NewTabletManagerClient() + defer tmc.Close() wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmc) primaryDb := fakesqldb.New(t) diff --git a/go/vt/vttablet/tabletmanager/tm_init.go b/go/vt/vttablet/tabletmanager/tm_init.go index da4e142c0a5..314973e744b 100644 --- a/go/vt/vttablet/tabletmanager/tm_init.go +++ b/go/vt/vttablet/tabletmanager/tm_init.go @@ -452,6 +452,10 @@ func (tm *TabletManager) Stop() { tm.stopShardSync() tm.stopRebuildKeyspace() + if tm.QueryServiceControl != nil { + tm.QueryServiceControl.Stats().Stop() + } + if tm.UpdateStream != nil { tm.UpdateStream.Disable() } From 478a4dfb10921196612093d5cc3263452d900183 Mon Sep 17 00:00:00 2001 From: Dirkjan Bussink Date: Wed, 23 Aug 2023 21:37:21 +0200 Subject: [PATCH 07/17] Cleanup tests Signed-off-by: Dirkjan Bussink --- go/vt/vtgate/autocommit_test.go | 118 ++--- go/vt/vtgate/bench_test.go | 14 +- go/vt/vtgate/executor_ddl_test.go | 9 +- go/vt/vtgate/executor_dml_test.go | 412 +++++---------- go/vt/vtgate/executor_framework_test.go | 32 +- go/vt/vtgate/executor_scatter_stats_test.go | 24 +- go/vt/vtgate/executor_select_test.go | 530 ++++++-------------- go/vt/vtgate/executor_set_test.go | 61 +-- go/vt/vtgate/executor_stream_test.go | 7 +- go/vt/vtgate/executor_test.go | 393 +++++---------- go/vt/vtgate/executor_vschema_ddl_test.go | 50 +- go/vt/vtgate/executor_vstream_test.go | 8 +- go/vt/vtgate/plugin_mysql_server_test.go | 7 +- go/vt/vtgate/queryz_test.go | 10 +- go/vt/vtgate/vtgate.go | 5 - go/vt/vtgate/vtgate_test.go | 88 ++-- 16 files changed, 519 insertions(+), 1249 deletions(-) diff --git a/go/vt/vtgate/autocommit_test.go b/go/vt/vtgate/autocommit_test.go index 3bd61a90cd3..5e0c267259d 100644 --- a/go/vt/vtgate/autocommit_test.go +++ b/go/vt/vtgate/autocommit_test.go @@ -22,8 +22,6 @@ import ( "github.com/stretchr/testify/require" - "vitess.io/vitess/go/test/utils" - "vitess.io/vitess/go/sqltypes" querypb "vitess.io/vitess/go/vt/proto/query" @@ -37,11 +35,8 @@ import ( // TestAutocommitUpdateSharded: instant-commit. func TestAutocommitUpdateSharded(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) + defer closer() _, err := autocommitExec(executor, "update user set a=2 where id = 1") require.NoError(t, err) @@ -58,11 +53,8 @@ func TestAutocommitUpdateSharded(t *testing.T) { // TestAutocommitUpdateLookup: transaction: select before update. func TestAutocommitUpdateLookup(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, sbclookup, _, closer := createExecutorEnv(t) + defer closer() sbclookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("b|a", "int64|varbinary"), "2|1", @@ -91,11 +83,8 @@ func TestAutocommitUpdateLookup(t *testing.T) { // TestAutocommitUpdateVindexChange: transaction: select & update before final update. func TestAutocommitUpdateVindexChange(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc, _, sbclookup, _, closer := createExecutorEnv(t) + defer closer() sbc.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("id|name|lastname|name_lastname_keyspace_id_map", "int64|int32|varchar|int64"), "1|1|foo|0", @@ -134,11 +123,8 @@ func TestAutocommitUpdateVindexChange(t *testing.T) { // TestAutocommitDeleteSharded: instant-commit. func TestAutocommitDeleteSharded(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) + defer closer() _, err := autocommitExec(executor, "delete from user_extra where user_id = 1") require.NoError(t, err) @@ -155,11 +141,8 @@ func TestAutocommitDeleteSharded(t *testing.T) { // TestAutocommitDeleteLookup: transaction: select before update. func TestAutocommitDeleteLookup(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, sbclookup, _, closer := createExecutorEnv(t) + defer closer() sbc1.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("id|name|lastname", "int64|int32|varchar"), "1|1|foo", @@ -201,11 +184,8 @@ func TestAutocommitDeleteLookup(t *testing.T) { // TestAutocommitDeleteIn: instant-commit. func TestAutocommitDeleteIn(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) + defer closer() _, err := autocommitExec(executor, "delete from user_extra where user_id in (1, 2)") require.NoError(t, err) @@ -222,11 +202,8 @@ func TestAutocommitDeleteIn(t *testing.T) { // TestAutocommitDeleteMultiShard: instant-commit. func TestAutocommitDeleteMultiShard(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) + defer closer() _, err := autocommitExec(executor, "delete from user_extra where user_id = user_id + 1") require.NoError(t, err) @@ -246,11 +223,8 @@ func TestAutocommitDeleteMultiShard(t *testing.T) { // TestAutocommitDeleteMultiShardAutoCommit: instant-commit. func TestAutocommitDeleteMultiShardAutoCommit(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) + defer closer() _, err := autocommitExec(executor, "delete /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ from user_extra where user_id = user_id + 1") require.NoError(t, err) @@ -270,11 +244,8 @@ func TestAutocommitDeleteMultiShardAutoCommit(t *testing.T) { // TestAutocommitInsertSharded: instant-commit. func TestAutocommitInsertSharded(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) + defer closer() _, err := autocommitExec(executor, "insert into user_extra(user_id, v) values (1, 2)") require.NoError(t, err) @@ -293,11 +264,8 @@ func TestAutocommitInsertSharded(t *testing.T) { // TestAutocommitInsertLookup: transaction: select before update. func TestAutocommitInsertLookup(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, sbclookup, _, closer := createExecutorEnv(t) + defer closer() _, err := autocommitExec(executor, "insert into user(id, v, name) values (1, 2, 'myname')") require.NoError(t, err) @@ -324,11 +292,7 @@ func TestAutocommitInsertLookup(t *testing.T) { // TestAutocommitInsertShardAutoCommit: instant-commit. func TestAutocommitInsertMultishardAutoCommit(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) _, err := autocommitExec(executor, "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ into user_extra(user_id, v) values (1, 2), (3, 4)") require.NoError(t, err) @@ -351,8 +315,9 @@ func TestAutocommitInsertMultishardAutoCommit(t *testing.T) { }}) testCommitCount(t, "sbc2", sbc2, 0) - executor, sbc1, sbc2, _ = createExecutorEnv(ctx) - defer executor.Close() + closer() + executor, sbc1, sbc2, _, _, closer = createExecutorEnv(t) + defer closer() // Make the first shard fail - the second completes anyway sbc1.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 1 _, err = autocommitExec(executor, "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ into user_extra(user_id, v) values (1, 2), (3, 4)") @@ -373,11 +338,8 @@ func TestAutocommitInsertMultishardAutoCommit(t *testing.T) { } func TestAutocommitInsertMultishard(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) + defer closer() _, err := autocommitExec(executor, "insert into user_extra(user_id, v) values (1, 2), (3, 4)") require.NoError(t, err) @@ -403,11 +365,8 @@ func TestAutocommitInsertMultishard(t *testing.T) { // TestAutocommitInsertAutoinc: instant-commit: sequence fetch is not transactional. func TestAutocommitInsertAutoinc(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, sbclookup, _, closer := createExecutorEnv(t) + defer closer() _, err := autocommitExec(executor, "insert into main1(id, name) values (null, 'myname')") require.NoError(t, err) @@ -426,11 +385,8 @@ func TestAutocommitInsertAutoinc(t *testing.T) { // TestAutocommitTransactionStarted: no instant-commit. func TestAutocommitTransactionStarted(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, _, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "@primary", @@ -462,11 +418,8 @@ func TestAutocommitTransactionStarted(t *testing.T) { // TestAutocommitDirectTarget: instant-commit. func TestAutocommitDirectTarget(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, sbclookup, _, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "TestUnsharded/0@primary", @@ -487,11 +440,8 @@ func TestAutocommitDirectTarget(t *testing.T) { // TestAutocommitDirectRangeTarget: no instant-commit. func TestAutocommitDirectRangeTarget(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, _, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "TestExecutor[-]@primary", diff --git a/go/vt/vtgate/bench_test.go b/go/vt/vtgate/bench_test.go index 423407270b8..d021fb9f7b5 100644 --- a/go/vt/vtgate/bench_test.go +++ b/go/vt/vtgate/bench_test.go @@ -18,7 +18,6 @@ package vtgate import ( "bytes" - "context" "fmt" "testing" @@ -59,10 +58,8 @@ func init() { } func BenchmarkWithNormalizer(b *testing.B) { - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - vtgateInst, _ := createVtgateEnv(ctx) - defer vtgateInst.Close() + vtgateInst, _, ctx, closer := createVtgateEnv(b) + defer closer() for i := 0; i < b.N; i++ { _, _, err := vtgateInst.Execute( @@ -82,11 +79,10 @@ func BenchmarkWithNormalizer(b *testing.B) { } func BenchmarkWithoutNormalizer(b *testing.B) { - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - vtgateInst, _ := createVtgateEnv(ctx) + vtgateInst, _, ctx, closer := createVtgateEnv(b) + defer closer() + vtgateInst.executor.normalize = false - defer vtgateInst.Close() for i := 0; i < b.N; i++ { _, _, err := vtgateInst.Execute( diff --git a/go/vt/vtgate/executor_ddl_test.go b/go/vt/vtgate/executor_ddl_test.go index 189e3ebfaea..f5d4d58d068 100644 --- a/go/vt/vtgate/executor_ddl_test.go +++ b/go/vt/vtgate/executor_ddl_test.go @@ -17,22 +17,17 @@ limitations under the License. package vtgate import ( - "context" "fmt" "testing" - "vitess.io/vitess/go/test/utils" vtgatepb "vitess.io/vitess/go/vt/proto/vtgate" "github.com/stretchr/testify/require" ) func TestDDLFlags(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() session := NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded}) defer func() { enableOnlineDDL = true diff --git a/go/vt/vtgate/executor_dml_test.go b/go/vt/vtgate/executor_dml_test.go index 6cbcf922c13..6975413333b 100644 --- a/go/vt/vtgate/executor_dml_test.go +++ b/go/vt/vtgate/executor_dml_test.go @@ -38,11 +38,8 @@ import ( ) func TestUpdateEqual(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -137,11 +134,9 @@ func TestUpdateEqual(t *testing.T) { } func TestUpdateFromSubQuery(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() + executor.pv = querypb.ExecuteOptions_Gen4 logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -178,9 +173,6 @@ func TestUpdateFromSubQuery(t *testing.T) { } func TestUpdateEqualWithNoVerifyAndWriteOnlyLookupUniqueVindexes(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() res := []*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields( "id|wo_lu_col|erl_lu_col|srl_lu_col|nrl_lu_col|nv_lu_col|lu_col|t2_lu_vdx", @@ -188,8 +180,8 @@ func TestUpdateEqualWithNoVerifyAndWriteOnlyLookupUniqueVindexes(t *testing.T) { ), "1|2|2|2|2|2|1|0", )} - executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, res) - defer executor.Close() + executor, sbc1, sbc2, sbcLookup, ctx, closer := createCustomExecutorSetValues(t, executorVSchema, res) + defer closer() session := &vtgatepb.Session{ TargetString: "@primary", @@ -236,9 +228,6 @@ func TestUpdateEqualWithNoVerifyAndWriteOnlyLookupUniqueVindexes(t *testing.T) { } func TestUpdateInTransactionLookupDefaultReadLock(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() res := []*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields( "id|wo_lu_col|erl_lu_col|srl_lu_col|nrl_lu_col|nv_lu_col|lu_col|t2_lu_vdx", @@ -246,8 +235,8 @@ func TestUpdateInTransactionLookupDefaultReadLock(t *testing.T) { ), "1|2|2|2|2|2|1|0", )} - executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, res) - defer executor.Close() + executor, sbc1, sbc2, sbcLookup, ctx, closer := createCustomExecutorSetValues(t, executorVSchema, res) + defer closer() safeSession := NewSafeSession(&vtgatepb.Session{InTransaction: true}) _, err := executorExecSession(ctx, @@ -302,9 +291,6 @@ func TestUpdateInTransactionLookupDefaultReadLock(t *testing.T) { } func TestUpdateInTransactionLookupExclusiveReadLock(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() res := []*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields( "id|wo_lu_col|erl_lu_col|srl_lu_col|nrl_lu_col|nv_lu_col|lu_col|t2_lu_vdx", @@ -312,8 +298,8 @@ func TestUpdateInTransactionLookupExclusiveReadLock(t *testing.T) { ), "1|2|2|2|2|2|1|0", )} - executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, res) - defer executor.Close() + executor, sbc1, sbc2, sbcLookup, ctx, closer := createCustomExecutorSetValues(t, executorVSchema, res) + defer closer() safeSession := NewSafeSession(&vtgatepb.Session{InTransaction: true}) _, err := executorExecSession(ctx, @@ -368,9 +354,6 @@ func TestUpdateInTransactionLookupExclusiveReadLock(t *testing.T) { } func TestUpdateInTransactionLookupSharedReadLock(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() res := []*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields( "id|wo_lu_col|erl_lu_col|srl_lu_col|nrl_lu_col|nv_lu_col|lu_col|t2_lu_vdx", @@ -378,8 +361,8 @@ func TestUpdateInTransactionLookupSharedReadLock(t *testing.T) { ), "1|2|2|2|2|2|1|0", )} - executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, res) - defer executor.Close() + executor, sbc1, sbc2, sbcLookup, ctx, closer := createCustomExecutorSetValues(t, executorVSchema, res) + defer closer() safeSession := NewSafeSession(&vtgatepb.Session{InTransaction: true}) _, err := executorExecSession(ctx, @@ -434,9 +417,6 @@ func TestUpdateInTransactionLookupSharedReadLock(t *testing.T) { } func TestUpdateInTransactionLookupNoReadLock(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() res := []*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields( "id|wo_lu_col|erl_lu_col|srl_lu_col|nrl_lu_col|nv_lu_col|lu_col|t2_lu_vdx", @@ -444,8 +424,8 @@ func TestUpdateInTransactionLookupNoReadLock(t *testing.T) { ), "1|2|2|2|2|2|1|0", )} - executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, res) - defer executor.Close() + executor, sbc1, sbc2, sbcLookup, ctx, closer := createCustomExecutorSetValues(t, executorVSchema, res) + defer closer() safeSession := NewSafeSession(&vtgatepb.Session{InTransaction: true}) _, err := executorExecSession(ctx, @@ -559,11 +539,8 @@ func TestUpdateMultiOwned(t *testing.T) { } } ` - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbclookup := createCustomExecutor(ctx, vschema) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createCustomExecutor(t, vschema) + defer closer() sbc1.SetResults([]*sqltypes.Result{ sqltypes.MakeTestResult( @@ -622,11 +599,8 @@ func TestUpdateMultiOwned(t *testing.T) { } func TestUpdateComments(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "@primary", @@ -642,11 +616,8 @@ func TestUpdateComments(t *testing.T) { } func TestUpdateNormalize(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() executor.normalize = true session := &vtgatepb.Session{ @@ -682,11 +653,8 @@ func TestUpdateNormalize(t *testing.T) { } func TestDeleteEqual(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() sbc.SetResults([]*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -802,11 +770,9 @@ func TestDeleteEqual(t *testing.T) { } func TestUpdateScatter(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() + session := &vtgatepb.Session{ TargetString: "@primary", } @@ -822,11 +788,9 @@ func TestUpdateScatter(t *testing.T) { } func TestDeleteScatter(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() + session := &vtgatepb.Session{ TargetString: "@primary", } @@ -842,11 +806,8 @@ func TestDeleteScatter(t *testing.T) { } func TestUpdateEqualWithMultipleLookupVindex(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, nil) - defer executor.Close() + executor, sbc1, sbc2, sbcLookup, ctx, closer := createExecutorEnv(t) + defer closer() sbcLookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("lu_col|keyspace_id", "int64|varbinary"), @@ -902,11 +863,8 @@ func TestUpdateEqualWithMultipleLookupVindex(t *testing.T) { } func TestUpdateUseHigherCostVindexIfBackfilling(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, nil) - defer executor.Close() + executor, sbc1, sbc2, sbcLookup, ctx, closer := createExecutorEnv(t) + defer closer() sbcLookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("lu_col|keyspace_id", "int64|varbinary"), @@ -984,11 +942,8 @@ func TestDeleteEqualWithNoVerifyAndWriteOnlyLookupUniqueVindex(t *testing.T) { ), "1|1|1|1|1|1|1", )} - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, res) - defer executor.Close() + executor, sbc1, sbc2, sbcLookup, ctx, closer := createCustomExecutorSetValues(t, executorVSchema, res) + defer closer() session := &vtgatepb.Session{ TargetString: "@primary", @@ -1062,11 +1017,8 @@ func TestDeleteEqualWithNoVerifyAndWriteOnlyLookupUniqueVindex(t *testing.T) { } func TestDeleteEqualWithMultipleLookupVindex(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, nil) - defer executor.Close() + executor, sbc1, sbc2, sbcLookup, ctx, closer := createCustomExecutorSetValues(t, executorVSchema, nil) + defer closer() sbcLookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("lu_col|keyspace_id", "int64|varbinary"), @@ -1147,11 +1099,8 @@ func TestDeleteEqualWithMultipleLookupVindex(t *testing.T) { } func TestDeleteUseHigherCostVindexIfBackfilling(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbcLookup := createCustomExecutorSetValues(ctx, executorVSchema, nil) - defer executor.Close() + executor, sbc1, sbc2, sbcLookup, ctx, closer := createCustomExecutorSetValues(t, executorVSchema, nil) + defer closer() sbcLookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("lu_col|keyspace_id", "int64|varbinary"), @@ -1271,11 +1220,9 @@ func TestDeleteUseHigherCostVindexIfBackfilling(t *testing.T) { } func TestDeleteByDestination(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() + session := &vtgatepb.Session{ TargetString: "@primary", } @@ -1291,11 +1238,8 @@ func TestDeleteByDestination(t *testing.T) { } func TestDeleteComments(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() sbc.SetResults([]*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -1334,11 +1278,8 @@ func TestDeleteComments(t *testing.T) { } func TestInsertSharded(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -1448,11 +1389,8 @@ func TestInsertSharded(t *testing.T) { } func TestInsertShardedKeyrange(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() // If a unique vindex returns a keyrange, we fail the insert session := &vtgatepb.Session{ @@ -1522,11 +1460,8 @@ func TestInsertShardedAutocommitLookup(t *testing.T) { } } ` - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbclookup := createCustomExecutor(ctx, vschema) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createCustomExecutor(t, vschema) + defer closer() _, err := executorExecSession(ctx, executor, "insert into user(id, v, name, music) values (1, 2, 'myname', 'star')", nil, &vtgatepb.Session{}) require.NoError(t, err) @@ -1559,11 +1494,8 @@ func TestInsertShardedAutocommitLookup(t *testing.T) { } func TestInsertShardedIgnore(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() // Build the sequence of responses for sbclookup. This should // match the sequence of queries we validate below. @@ -1718,11 +1650,8 @@ func TestInsertShardedIgnore(t *testing.T) { func TestInsertOnDupKey(t *testing.T) { // This test just sanity checks that the statement is getting passed through // correctly. The full set of use cases are covered by TestInsertShardedIgnore. - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() sbclookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("b|a", "int64|varbinary"), "1|1", @@ -1767,11 +1696,8 @@ func TestInsertOnDupKey(t *testing.T) { } func TestAutocommitFail(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, ctx, closer := createExecutorEnv(t) + defer closer() query := "insert into user (id) values (1)" sbc1.MustFailCodes[vtrpcpb.Code_ALREADY_EXISTS] = 1 @@ -1788,11 +1714,8 @@ func TestAutocommitFail(t *testing.T) { } func TestInsertComments(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "@primary", @@ -1820,11 +1743,8 @@ func TestInsertComments(t *testing.T) { } func TestInsertGeneratorSharded(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() sbclookup.SetResults([]*sqltypes.Result{{ Rows: [][]sqltypes.Value{{ @@ -1866,11 +1786,8 @@ func TestInsertGeneratorSharded(t *testing.T) { } func TestInsertAutoincSharded(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - router, sbc, _, _ := createExecutorEnv(ctx) - defer router.Close() + router, sbc, _, _, ctx, closer := createExecutorEnv(t) + defer closer() // Fake a mysql auto-inc response. wantResult := &sqltypes.Result{ @@ -1900,11 +1817,8 @@ func TestInsertAutoincSharded(t *testing.T) { } func TestInsertGeneratorUnsharded(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "@primary", } @@ -1928,11 +1842,8 @@ func TestInsertGeneratorUnsharded(t *testing.T) { } func TestInsertAutoincUnsharded(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - router, _, _, sbclookup := createExecutorEnv(ctx) - defer router.Close() + router, _, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() logChan := router.queryLogger.Subscribe("Test") defer router.queryLogger.Unsubscribe(logChan) @@ -1964,11 +1875,8 @@ func TestInsertAutoincUnsharded(t *testing.T) { } func TestInsertLookupOwned(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "@primary", @@ -1995,11 +1903,8 @@ func TestInsertLookupOwned(t *testing.T) { } func TestInsertLookupOwnedGenerator(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() sbclookup.SetResults([]*sqltypes.Result{{ Rows: [][]sqltypes.Value{{ @@ -2041,11 +1946,8 @@ func TestInsertLookupOwnedGenerator(t *testing.T) { } func TestInsertLookupUnowned(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "@primary", @@ -2071,11 +1973,8 @@ func TestInsertLookupUnowned(t *testing.T) { } func TestInsertLookupUnownedUnsupplied(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() sbclookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("b|a", "int64|varbinary"), "3|1", @@ -2107,11 +2006,8 @@ func TestInsertLookupUnownedUnsupplied(t *testing.T) { // If a statement gets broken up into two, and the first one fails, // then an error should be returned normally. func TestInsertPartialFail1(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, sbclookup, _, closer := createExecutorEnv(t) + defer closer() // Make the first DML fail, there should be no rollback. sbclookup.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 1 @@ -2131,11 +2027,8 @@ func TestInsertPartialFail1(t *testing.T) { // after successful execution of the first, then the transaction must // be rolled back due to partial execution. func TestInsertPartialFail2(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, _, closer := createExecutorEnv(t) + defer closer() // Make the second DML fail, it should result in a rollback. sbc1.MustFailExecute[sqlparser.StmtInsert] = 1 @@ -2175,11 +2068,8 @@ func TestInsertPartialFail2(t *testing.T) { } func TestMultiInsertSharded(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "@primary", @@ -2286,11 +2176,8 @@ func TestMultiInsertSharded(t *testing.T) { } func TestMultiInsertGenerator(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() sbclookup.SetResults([]*sqltypes.Result{{ Rows: [][]sqltypes.Value{{ @@ -2338,11 +2225,8 @@ func TestMultiInsertGenerator(t *testing.T) { } func TestMultiInsertGeneratorSparse(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() sbclookup.SetResults([]*sqltypes.Result{{ Rows: [][]sqltypes.Value{{ @@ -2419,11 +2303,8 @@ func TestInsertBadAutoInc(t *testing.T) { } } ` - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createCustomExecutor(ctx, vschema) - defer executor.Close() + executor, _, _, _, ctx, closer := createCustomExecutor(t, vschema) + defer closer() // If auto inc table cannot be found, the table should not be added to vschema. session := &vtgatepb.Session{ @@ -2496,11 +2377,8 @@ func TestKeyDestRangeQuery(t *testing.T) { for _, tc := range tests { t.Run(tc.targetString+" - "+tc.inputQuery, func(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: tc.targetString, @@ -2523,11 +2401,8 @@ func TestKeyDestRangeQuery(t *testing.T) { } // it does not work for inserts - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "TestExecutor[-]", } @@ -2547,11 +2422,8 @@ func assertQueriesContain(t *testing.T, sql, sbcName string, sbc *sandboxconn.Sa // Prepared statement tests func TestUpdateEqualWithPrepare(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2572,11 +2444,8 @@ func TestUpdateEqualWithPrepare(t *testing.T) { assertQueries(t, sbc1, nil) } func TestInsertShardedWithPrepare(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2600,11 +2469,9 @@ func TestInsertShardedWithPrepare(t *testing.T) { } func TestDeleteEqualWithPrepare(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() + session := &vtgatepb.Session{ TargetString: "@primary", } @@ -2621,11 +2488,9 @@ func TestDeleteEqualWithPrepare(t *testing.T) { } func TestUpdateLastInsertID(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, ctx, closer := createExecutorEnv(t) + defer closer() + executor.normalize = true sql := "update user set a = last_insert_id() where id = 1" @@ -2646,11 +2511,8 @@ func TestUpdateLastInsertID(t *testing.T) { } func TestUpdateReference(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2692,11 +2554,8 @@ func TestUpdateReference(t *testing.T) { } func TestDeleteLookupOwnedEqual(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() sbc1.SetResults([]*sqltypes.Result{ sqltypes.MakeTestResult(sqltypes.MakeTestFields("uniq_col|keyspace_id", "int64|varbinary"), "1|N±\u0090ɢú\u0016\u009C"), @@ -2726,11 +2585,8 @@ func TestDeleteLookupOwnedEqual(t *testing.T) { } func TestDeleteReference(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2771,11 +2627,8 @@ func TestDeleteReference(t *testing.T) { } func TestReservedConnDML(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, sbc := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, sbc, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("TestReservedConnDML") defer executor.queryLogger.Unsubscribe(logChan) @@ -2827,11 +2680,8 @@ func TestReservedConnDML(t *testing.T) { func TestStreamingDML(t *testing.T) { method := "TestStreamingDML" - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, sbc := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, sbc, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe(method) defer executor.queryLogger.Unsubscribe(logChan) @@ -2915,11 +2765,8 @@ func TestStreamingDML(t *testing.T) { } func TestPartialVindexInsertQueryFailure(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2972,11 +2819,8 @@ func TestPartialVindexInsertQueryFailure(t *testing.T) { } func TestPartialVindexInsertQueryFailureAutoCommit(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -3020,11 +2864,8 @@ func TestPartialVindexInsertQueryFailureAutoCommit(t *testing.T) { // The change for it cannot be done as the executor level and will be made at the VTGate entry point. // Test TestMultiInternalSavepointVtGate shows that it fixes the behaviour. func TestMultiInternalSavepoint(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() session := NewAutocommitSession(&vtgatepb.Session{}) _, err := executorExecSession(ctx, executor, "begin", nil, session.Session) @@ -3072,11 +2913,8 @@ func TestMultiInternalSavepoint(t *testing.T) { } func TestInsertSelectFromDual(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, _, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("TestInsertSelect") defer executor.queryLogger.Unsubscribe(logChan) @@ -3131,11 +2969,8 @@ func TestInsertSelectFromDual(t *testing.T) { } func TestInsertSelectFromTable(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, _, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("TestInsertSelect") defer executor.queryLogger.Unsubscribe(logChan) @@ -3196,11 +3031,8 @@ func TestInsertSelectFromTable(t *testing.T) { } func TestInsertReference(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) diff --git a/go/vt/vtgate/executor_framework_test.go b/go/vt/vtgate/executor_framework_test.go index d3b896558e3..d3bbad9752e 100644 --- a/go/vt/vtgate/executor_framework_test.go +++ b/go/vt/vtgate/executor_framework_test.go @@ -28,6 +28,8 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "vitess.io/vitess/go/test/utils" + "vitess.io/vitess/go/cache" "vitess.io/vitess/go/constants/sidecar" "vitess.io/vitess/go/sqltypes" @@ -125,7 +127,9 @@ func init() { vindexes.Register("keyrange_lookuper_unique", newKeyRangeLookuperUnique) } -func createExecutorEnv(ctx context.Context) (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn.SandboxConn) { +func createExecutorEnv(t testing.TB) (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn.SandboxConn, ctx context.Context, closer func()) { + var cancel context.CancelFunc + ctx, cancel = context.WithCancel(context.Background()) cell := "aa" hc := discovery.NewFakeHealthCheck(make(chan *discovery.TabletHealth)) s := createSandbox(KsTestSharded) @@ -175,10 +179,16 @@ func createExecutorEnv(ctx context.Context) (executor *Executor, sbc1, sbc2, sbc executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) key.AnyShardPicker = DestinationAnyShardPickerFirstShard{} - return executor, sbc1, sbc2, sbclookup + return executor, sbc1, sbc2, sbclookup, ctx, func() { + executor.Close() + cancel() + utils.EnsureNoLeaks(t) + } } -func createCustomExecutor(ctx context.Context, vschema string) (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn.SandboxConn) { +func createCustomExecutor(t testing.TB, vschema string) (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn.SandboxConn, ctx context.Context, closer func()) { + var cancel context.CancelFunc + ctx, cancel = context.WithCancel(context.Background()) cell := "aa" hc := discovery.NewFakeHealthCheck(nil) s := createSandbox(KsTestSharded) @@ -195,10 +205,16 @@ func createCustomExecutor(ctx context.Context, vschema string) (executor *Execut queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize) plans := cache.NewDefaultCacheImpl(cache.DefaultConfig) executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) - return executor, sbc1, sbc2, sbclookup + return executor, sbc1, sbc2, sbclookup, ctx, func() { + defer utils.EnsureNoLeaks(t) + defer cancel() + executor.Close() + } } -func createCustomExecutorSetValues(ctx context.Context, vschema string, values []*sqltypes.Result) (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn.SandboxConn) { +func createCustomExecutorSetValues(t testing.TB, vschema string, values []*sqltypes.Result) (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn.SandboxConn, ctx context.Context, closer func()) { + var cancel context.CancelFunc + ctx, cancel = context.WithCancel(context.Background()) cell := "aa" hc := discovery.NewFakeHealthCheck(nil) s := createSandbox(KsTestSharded) @@ -222,7 +238,11 @@ func createCustomExecutorSetValues(ctx context.Context, vschema string, values [ queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize) plans := cache.NewDefaultCacheImpl(cache.DefaultConfig) executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) - return executor, sbcs[0], sbcs[1], sbclookup + return executor, sbcs[0], sbcs[1], sbclookup, ctx, func() { + executor.Close() + cancel() + utils.EnsureNoLeaks(t) + } } func executorExecSession(ctx context.Context, executor *Executor, sql string, bv map[string]*querypb.BindVariable, session *vtgatepb.Session) (*sqltypes.Result, error) { diff --git a/go/vt/vtgate/executor_scatter_stats_test.go b/go/vt/vtgate/executor_scatter_stats_test.go index c3123260b7b..64348ac8a2e 100644 --- a/go/vt/vtgate/executor_scatter_stats_test.go +++ b/go/vt/vtgate/executor_scatter_stats_test.go @@ -17,23 +17,17 @@ limitations under the License. package vtgate import ( - "context" "net/http/httptest" "testing" "github.com/stretchr/testify/require" - "vitess.io/vitess/go/test/utils" - vtgatepb "vitess.io/vitess/go/vt/proto/vtgate" ) func TestScatterStatsWithNoScatterQuery(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) @@ -46,11 +40,8 @@ func TestScatterStatsWithNoScatterQuery(t *testing.T) { } func TestScatterStatsWithSingleScatterQuery(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) _, err := executor.Execute(ctx, nil, "TestExecutorResultsExceeded", session, "select * from user", nil) @@ -62,11 +53,8 @@ func TestScatterStatsWithSingleScatterQuery(t *testing.T) { } func TestScatterStatsHttpWriting(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) _, err := executor.Execute(ctx, nil, "TestExecutorResultsExceeded", session, "select * from user", nil) diff --git a/go/vt/vtgate/executor_select_test.go b/go/vt/vtgate/executor_select_test.go index 23202637f04..8b83853fcf6 100644 --- a/go/vt/vtgate/executor_select_test.go +++ b/go/vt/vtgate/executor_select_test.go @@ -52,11 +52,8 @@ import ( ) func TestSelectNext(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, sbclookup, _, closer := createExecutorEnv(t) + defer closer() query := "select next :n values from user_seq" bv := map[string]*querypb.BindVariable{"n": sqltypes.Int64BindVariable(2)} @@ -110,11 +107,8 @@ func TestSelectNext(t *testing.T) { } func TestSelectDBA(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, _, closer := createExecutorEnv(t) + defer closer() query := "select * from INFORMATION_SCHEMA.foo" _, err := executor.Execute(context.Background(), nil, "TestSelectDBA", @@ -168,11 +162,8 @@ func TestSelectDBA(t *testing.T) { } func TestSystemVariablesMySQLBelow80(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, _, closer := createExecutorEnv(t) + defer closer() executor.normalize = true sqlparser.SetParserVersion("57000") @@ -208,11 +199,8 @@ func TestSystemVariablesMySQLBelow80(t *testing.T) { } func TestSystemVariablesWithSetVarDisabled(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, _, closer := createExecutorEnv(t) + defer closer() executor.normalize = true sqlparser.SetParserVersion("80000") @@ -250,11 +238,8 @@ func TestSystemVariablesWithSetVarDisabled(t *testing.T) { } func TestSetSystemVariablesTx(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, _, closer := createExecutorEnv(t) + defer closer() executor.normalize = true sqlparser.SetParserVersion("80001") @@ -302,11 +287,8 @@ func TestSetSystemVariablesTx(t *testing.T) { } func TestSetSystemVariables(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, lookup := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, lookup, _, closer := createExecutorEnv(t) + defer closer() executor.normalize = true sqlparser.SetParserVersion("80001") @@ -419,11 +401,8 @@ func TestSetSystemVariables(t *testing.T) { } func TestSetSystemVariablesWithReservedConnection(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, _, closer := createExecutorEnv(t) + defer closer() executor.normalize = true session := NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true, SystemVariables: map[string]string{}}) @@ -466,11 +445,8 @@ func TestSetSystemVariablesWithReservedConnection(t *testing.T) { } func TestCreateTableValidTimestamp(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, _, closer := createExecutorEnv(t) + defer closer() executor.normalize = true session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor", SystemVariables: map[string]string{"sql_mode": "ALLOW_INVALID_DATES"}}) @@ -489,11 +465,8 @@ func TestCreateTableValidTimestamp(t *testing.T) { } func TestGen4SelectDBA(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, _, closer := createExecutorEnv(t) + defer closer() executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 @@ -568,11 +541,8 @@ func TestGen4SelectDBA(t *testing.T) { } func TestUnsharded(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "@primary", @@ -587,11 +557,8 @@ func TestUnsharded(t *testing.T) { } func TestUnshardedComments(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "@primary", @@ -635,11 +602,8 @@ func TestUnshardedComments(t *testing.T) { } func TestStreamUnsharded(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -655,11 +619,8 @@ func TestStreamUnsharded(t *testing.T) { } func TestStreamBuffering(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, sbclookup, _, closer := createExecutorEnv(t) + defer closer() // This test is similar to TestStreamUnsharded except that it returns a Result > 10 bytes, // such that the splitting of the Result into multiple Result responses gets tested. @@ -715,11 +676,8 @@ func TestStreamBuffering(t *testing.T) { } func TestStreamLimitOffset(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) + defer closer() // This test is similar to TestStreamUnsharded except that it returns a Result > 10 bytes, // such that the splitting of the Result into multiple Result responses gets tested. @@ -797,11 +755,8 @@ func TestStreamLimitOffset(t *testing.T) { } func TestSelectLastInsertId(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "@primary", LastInsertId: 52, @@ -825,11 +780,8 @@ func TestSelectLastInsertId(t *testing.T) { } func TestSelectSystemVariables(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "@primary", @@ -889,11 +841,8 @@ func TestSelectSystemVariables(t *testing.T) { } func TestSelectInitializedVitessAwareVariable(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() executor.normalize = true logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -925,11 +874,8 @@ func TestSelectInitializedVitessAwareVariable(t *testing.T) { } func TestSelectUserDefinedVariable(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() executor.normalize = true logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -965,11 +911,8 @@ func TestSelectUserDefinedVariable(t *testing.T) { } func TestFoundRows(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() executor.normalize = true logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -996,11 +939,8 @@ func TestFoundRows(t *testing.T) { } func TestRowCount(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() executor.normalize = true logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -1033,11 +973,8 @@ func testRowCount(t *testing.T, ctx context.Context, executor *Executor, session } func TestSelectLastInsertIdInUnion(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, ctx, closer := createExecutorEnv(t) + defer closer() executor.normalize = true session := &vtgatepb.Session{ @@ -1071,11 +1008,8 @@ func TestSelectLastInsertIdInUnion(t *testing.T) { } func TestSelectLastInsertIdInWhere(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, lookup := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, lookup, ctx, closer := createExecutorEnv(t) + defer closer() executor.normalize = true logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -1095,11 +1029,8 @@ func TestSelectLastInsertIdInWhere(t *testing.T) { } func TestLastInsertIDInVirtualTable(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, ctx, closer := createExecutorEnv(t) + defer closer() executor.normalize = true result1 := []*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -1127,11 +1058,8 @@ func TestLastInsertIDInVirtualTable(t *testing.T) { } func TestLastInsertIDInSubQueryExpression(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() executor.normalize = true session := &vtgatepb.Session{ TargetString: "@primary", @@ -1155,11 +1083,8 @@ func TestLastInsertIDInSubQueryExpression(t *testing.T) { } func TestSelectDatabase(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, _, closer := createExecutorEnv(t) + defer closer() executor.normalize = true sql := "select database()" newSession := &vtgatepb.Session{ @@ -1188,11 +1113,8 @@ func TestSelectDatabase(t *testing.T) { } func TestSelectBindvars(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, lookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, lookup, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -1299,11 +1221,8 @@ func TestSelectBindvars(t *testing.T) { } func TestSelectEqual(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "@primary", @@ -1373,11 +1292,8 @@ func TestSelectEqual(t *testing.T) { } func TestSelectINFromOR(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, ctx, closer := createExecutorEnv(t) + defer closer() executor.pv = querypb.ExecuteOptions_Gen4 session := &vtgatepb.Session{ @@ -1395,11 +1311,8 @@ func TestSelectINFromOR(t *testing.T) { } func TestSelectDual(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, _, lookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, lookup, ctx, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "@primary", @@ -1418,11 +1331,8 @@ func TestSelectDual(t *testing.T) { } func TestSelectComments(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "@primary", @@ -1441,11 +1351,8 @@ func TestSelectComments(t *testing.T) { } func TestSelectNormalize(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() executor.normalize = true session := &vtgatepb.Session{ @@ -1481,11 +1388,8 @@ func TestSelectNormalize(t *testing.T) { } func TestSelectCaseSensitivity(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "@primary", @@ -1504,11 +1408,8 @@ func TestSelectCaseSensitivity(t *testing.T) { } func TestStreamSelectEqual(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() sql := "select id from user where id = 1" result, err := executorStream(ctx, executor, sql) @@ -1520,11 +1421,8 @@ func TestStreamSelectEqual(t *testing.T) { } func TestSelectKeyRange(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "@primary", @@ -1543,11 +1441,8 @@ func TestSelectKeyRange(t *testing.T) { } func TestSelectKeyRangeUnique(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "@primary", @@ -1566,11 +1461,8 @@ func TestSelectKeyRangeUnique(t *testing.T) { } func TestSelectIN(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() // Constant in IN clause is just a number, not a bind variable. session := &vtgatepb.Session{ @@ -1659,11 +1551,8 @@ func TestSelectIN(t *testing.T) { } func TestStreamSelectIN(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() sql := "select id from user where id in (1)" result, err := executorStream(ctx, executor, sql) @@ -2542,12 +2431,8 @@ func TestStreamSelectScatterLimit(t *testing.T) { // TODO(sougou): stream and non-stream testing are very similar. // Could reuse code, func TestSimpleJoin(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2587,12 +2472,8 @@ func TestSimpleJoin(t *testing.T) { } func TestJoinComments(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2617,12 +2498,8 @@ func TestJoinComments(t *testing.T) { } func TestSimpleJoinStream(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2660,12 +2537,8 @@ func TestSimpleJoinStream(t *testing.T) { } func TestVarJoin(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2703,12 +2576,8 @@ func TestVarJoin(t *testing.T) { } func TestVarJoinStream(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2743,12 +2612,8 @@ func TestVarJoinStream(t *testing.T) { } func TestLeftJoin(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) result1 := []*sqltypes.Result{{ @@ -2794,12 +2659,8 @@ func TestLeftJoin(t *testing.T) { } func TestLeftJoinStream(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() result1 := []*sqltypes.Result{{ Fields: []*querypb.Field{ {Name: "id", Type: sqltypes.Int32, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NUM_FLAG)}, @@ -2839,12 +2700,8 @@ func TestLeftJoinStream(t *testing.T) { } func TestEmptyJoin(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, ctx, closer := createExecutorEnv(t) + defer closer() // Empty result requires a field query for the second part of join, // which is sent to shard 0. sbc1.SetResults([]*sqltypes.Result{{ @@ -2885,12 +2742,8 @@ func TestEmptyJoin(t *testing.T) { } func TestEmptyJoinStream(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, ctx, closer := createExecutorEnv(t) + defer closer() // Empty result requires a field query for the second part of join, // which is sent to shard 0. sbc1.SetResults([]*sqltypes.Result{{ @@ -2926,12 +2779,8 @@ func TestEmptyJoinStream(t *testing.T) { } func TestEmptyJoinRecursive(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, ctx, closer := createExecutorEnv(t) + defer closer() // Make sure it also works recursively. sbc1.SetResults([]*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -2978,12 +2827,8 @@ func TestEmptyJoinRecursive(t *testing.T) { } func TestEmptyJoinRecursiveStream(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, ctx, closer := createExecutorEnv(t) + defer closer() // Make sure it also works recursively. sbc1.SetResults([]*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -3027,12 +2872,8 @@ func TestEmptyJoinRecursiveStream(t *testing.T) { } func TestCrossShardSubquery(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() result1 := []*sqltypes.Result{{ Fields: []*querypb.Field{ {Name: "id", Type: sqltypes.Int32}, @@ -3069,12 +2910,8 @@ func TestCrossShardSubquery(t *testing.T) { } func TestSubQueryAndQueryWithLimit(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) + defer closer() result1 := []*sqltypes.Result{{ Fields: []*querypb.Field{ {Name: "id", Type: sqltypes.Int32, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NUM_FLAG)}, @@ -3116,12 +2953,8 @@ func TestSubQueryAndQueryWithLimit(t *testing.T) { } func TestCrossShardSubqueryStream(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() result1 := []*sqltypes.Result{{ Fields: []*querypb.Field{ {Name: "id", Type: sqltypes.Int32, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NUM_FLAG)}, @@ -3161,12 +2994,8 @@ func TestCrossShardSubqueryStream(t *testing.T) { } func TestCrossShardSubqueryGetFields(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() sbclookup.SetResults([]*sqltypes.Result{{ Fields: []*querypb.Field{ {Name: "col", Type: sqltypes.Int32, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NUM_FLAG)}, @@ -3207,12 +3036,8 @@ func TestCrossShardSubqueryGetFields(t *testing.T) { } func TestSelectBindvarswithPrepare(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -3236,12 +3061,8 @@ func TestSelectBindvarswithPrepare(t *testing.T) { } func TestSelectDatabasePrepare(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() executor.normalize = true logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -3255,12 +3076,8 @@ func TestSelectDatabasePrepare(t *testing.T) { } func TestSelectWithUnionAll(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() executor.normalize = true sql := "select id from user where id in (1, 2, 3) union all select id from user where id in (1, 2, 3)" bv, _ := sqltypes.BuildBindVariable([]int64{1, 2, 3}) @@ -3315,12 +3132,8 @@ func TestSelectWithUnionAll(t *testing.T) { } func TestSelectLock(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, _, closer := createExecutorEnv(t) + defer closer() session := NewSafeSession(nil) session.Session.InTransaction = true session.ShardSessions = []*vtgatepb.Session_ShardSession{{ @@ -3379,9 +3192,8 @@ func TestSelectLock(t *testing.T) { } func TestLockReserve(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, _, _, closer := createExecutorEnv(t) + defer closer() // no connection should be reserved for these queries. tcases := []string{ @@ -3391,8 +3203,6 @@ func TestLockReserve(t *testing.T) { "select release_lock('lock name') from dual", } - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() session := NewAutocommitSession(&vtgatepb.Session{}) for _, sql := range tcases { @@ -3411,12 +3221,8 @@ func TestLockReserve(t *testing.T) { } func TestSelectFromInformationSchema(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, _, closer := createExecutorEnv(t) + defer closer() session := NewSafeSession(nil) // check failure when trying to query two keyspaces @@ -3546,12 +3352,8 @@ func TestSelectScatterFails(t *testing.T) { } func TestGen4SelectStraightJoin(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, _, closer := createExecutorEnv(t) + defer closer() executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) @@ -3579,12 +3381,8 @@ func TestGen4SelectStraightJoin(t *testing.T) { } func TestGen4MultiColumnVindexEqual(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) + defer closer() executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 @@ -3623,12 +3421,8 @@ func TestGen4MultiColumnVindexEqual(t *testing.T) { } func TestGen4MultiColumnVindexIn(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) + defer closer() executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 @@ -3667,12 +3461,8 @@ func TestGen4MultiColumnVindexIn(t *testing.T) { } func TestGen4MultiColMixedColComparision(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) + defer closer() executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 @@ -3709,12 +3499,8 @@ func TestGen4MultiColMixedColComparision(t *testing.T) { } func TestGen4MultiColBestVindexSel(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) + defer closer() executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 @@ -3760,12 +3546,8 @@ func TestGen4MultiColBestVindexSel(t *testing.T) { } func TestGen4MultiColMultiEqual(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) + defer closer() executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 @@ -3788,12 +3570,8 @@ func TestGen4MultiColMultiEqual(t *testing.T) { } func TestGen4SelectUnqualifiedReferenceTable(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() executor.pv = querypb.ExecuteOptions_Gen4 query := "select * from zip_detail" @@ -3814,12 +3592,8 @@ func TestGen4SelectUnqualifiedReferenceTable(t *testing.T) { } func TestGen4SelectQualifiedReferenceTable(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() executor.pv = querypb.ExecuteOptions_Gen4 query := fmt.Sprintf("select * from %s.zip_detail", KsTestSharded) @@ -3840,12 +3614,8 @@ func TestGen4SelectQualifiedReferenceTable(t *testing.T) { } func TestGen4JoinUnqualifiedReferenceTable(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() executor.pv = querypb.ExecuteOptions_Gen4 query := "select * from user join zip_detail on user.zip_detail_id = zip_detail.id" @@ -3882,12 +3652,8 @@ func TestGen4JoinUnqualifiedReferenceTable(t *testing.T) { } func TestGen4CrossShardJoinQualifiedReferenceTable(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() executor.pv = querypb.ExecuteOptions_Gen4 query := "select user.id from user join TestUnsharded.zip_detail on user.zip_detail_id = TestUnsharded.zip_detail.id" @@ -4411,12 +4177,8 @@ func TestSelectAggregationRandom(t *testing.T) { } func TestSelectHexAndBit(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, _, closer := createExecutorEnv(t) + defer closer() executor.normalize = true session := NewAutocommitSession(&vtgatepb.Session{}) @@ -4432,12 +4194,8 @@ func TestSelectHexAndBit(t *testing.T) { // TestSelectCFC tests validates that cfc vindex plan gets cached and same plan is getting reused. // This also validates that cache_size is able to calculate the cfc vindex plan size. func TestSelectCFC(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, _, closer := createExecutorEnv(t) + defer closer() executor.normalize = true session := NewAutocommitSession(&vtgatepb.Session{}) @@ -4462,12 +4220,8 @@ func TestSelectCFC(t *testing.T) { } func TestSelectView(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc, _, _, _, closer := createExecutorEnv(t) + defer closer() // add the view to local vschema err := executor.vschema.AddView(KsTestSharded, "user_details_view", "select user.id, user_extra.col from user join user_extra on user.id = user_extra.user_id") require.NoError(t, err) diff --git a/go/vt/vtgate/executor_set_test.go b/go/vt/vtgate/executor_set_test.go index 17d5fa13038..900593608fc 100644 --- a/go/vt/vtgate/executor_set_test.go +++ b/go/vt/vtgate/executor_set_test.go @@ -17,7 +17,6 @@ limitations under the License. package vtgate import ( - "context" "fmt" "testing" @@ -41,11 +40,8 @@ import ( ) func TestExecutorSet(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executorEnv, _, _, _ := createExecutorEnv(ctx) - defer executorEnv.Close() + executorEnv, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() testcases := []struct { in string @@ -286,11 +282,8 @@ func TestExecutorSet(t *testing.T) { } func TestExecutorSetOp(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() sysVarSetEnabled = true returnResult := func(columnName, typ, value string) *sqltypes.Result { @@ -389,11 +382,7 @@ func TestExecutorSetOp(t *testing.T) { } func TestExecutorSetMetadata(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) set := "set @@vitess_metadata.app_keyspace_v1= '1'" @@ -405,8 +394,9 @@ func TestExecutorSetMetadata(t *testing.T) { vschemaacl.AuthorizedDDLUsers = "" }() - executor, _, _, _ = createExecutorEnv(ctx) - defer executor.Close() + closer() + executor, _, _, _, ctx, closer = createExecutorEnv(t) + defer closer() session = NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) set = "set @@vitess_metadata.app_keyspace_v1= '1'" @@ -451,11 +441,8 @@ func TestExecutorSetMetadata(t *testing.T) { } func TestPlanExecutorSetUDV(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() testcases := []struct { in string @@ -485,11 +472,8 @@ func TestPlanExecutorSetUDV(t *testing.T) { } func TestSetUDVFromTabletInput(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, sbc1, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, _, ctx, closer := createExecutorEnv(t) + defer closer() fields := sqltypes.MakeTestFields("some", "VARCHAR") sbc1.SetResults([]*sqltypes.Result{ @@ -520,11 +504,8 @@ func createMap(keys []string, values []any) map[string]*querypb.BindVariable { } func TestSetVar(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, sbc := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, sbc, ctx, closer := createExecutorEnv(t) + defer closer() executor.normalize = true oldVersion := sqlparser.GetParserVersion() @@ -568,11 +549,8 @@ func TestSetVar(t *testing.T) { } func TestSetVarShowVariables(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, sbc := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, sbc, ctx, closer := createExecutorEnv(t) + defer closer() executor.normalize = true oldVersion := sqlparser.GetParserVersion() @@ -601,11 +579,8 @@ func TestSetVarShowVariables(t *testing.T) { } func TestExecutorSetAndSelect(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - e, _, _, sbc := createExecutorEnv(ctx) - defer e.Close() + e, _, _, sbc, ctx, closer := createExecutorEnv(t) + defer closer() e.normalize = true testcases := []struct { diff --git a/go/vt/vtgate/executor_stream_test.go b/go/vt/vtgate/executor_stream_test.go index 9b041d095e1..4f423696982 100644 --- a/go/vt/vtgate/executor_stream_test.go +++ b/go/vt/vtgate/executor_stream_test.go @@ -39,11 +39,8 @@ import ( ) func TestStreamSQLUnsharded(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, _, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) diff --git a/go/vt/vtgate/executor_test.go b/go/vt/vtgate/executor_test.go index c3aa300bcda..1dea9ef635b 100644 --- a/go/vt/vtgate/executor_test.go +++ b/go/vt/vtgate/executor_test.go @@ -58,16 +58,13 @@ import ( ) func TestExecutorResultsExceeded(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() save := warnMemoryRows warnMemoryRows = 3 defer func() { warnMemoryRows = save }() - executor, _, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) initial := warnings.Counts()["ResultsExceeded"] @@ -86,16 +83,13 @@ func TestExecutorResultsExceeded(t *testing.T) { } func TestExecutorMaxMemoryRowsExceeded(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() save := maxMemoryRows maxMemoryRows = 3 defer func() { maxMemoryRows = save }() - executor, _, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) result := sqltypes.MakeTestResult(sqltypes.MakeTestFields("col", "int64"), "1", "2", "3", "4") fn := func(r *sqltypes.Result) error { @@ -128,12 +122,9 @@ func TestExecutorMaxMemoryRowsExceeded(t *testing.T) { } func TestExecutorTransactionsNoAutoCommit(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() - executor, _, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", SessionUUID: "suuid"}) logChan := executor.queryLogger.Subscribe("Test") @@ -206,12 +197,9 @@ func TestExecutorTransactionsNoAutoCommit(t *testing.T) { } func TestDirectTargetRewrites(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() - executor, _, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() executor.normalize = true session := &vtgatepb.Session{ @@ -230,12 +218,9 @@ func TestDirectTargetRewrites(t *testing.T) { } func TestExecutorTransactionsAutoCommit(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() - executor, _, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true, SessionUUID: "suuid"}) logChan := executor.queryLogger.Subscribe("Test") @@ -287,12 +272,9 @@ func TestExecutorTransactionsAutoCommit(t *testing.T) { } func TestExecutorTransactionsAutoCommitStreaming(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() - executor, _, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() oltpOptions := &querypb.ExecuteOptions{Workload: querypb.ExecuteOptions_OLTP} session := NewSafeSession(&vtgatepb.Session{ TargetString: "@primary", @@ -357,17 +339,13 @@ func TestExecutorTransactionsAutoCommitStreaming(t *testing.T) { } func TestExecutorDeleteMetadata(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - vschemaacl.AuthorizedDDLUsers = "%" defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) set := "set @@vitess_metadata.app_v1= '1'" @@ -394,12 +372,9 @@ func TestExecutorDeleteMetadata(t *testing.T) { } func TestExecutorAutocommit(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() - executor, _, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) logChan := executor.queryLogger.Subscribe("Test") @@ -499,12 +474,9 @@ func TestExecutorAutocommit(t *testing.T) { } func TestExecutorShowColumns(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: ""}) queries := []string{ @@ -555,12 +527,9 @@ func assertMatchesNoOrder(t *testing.T, expected, got string) { } func TestExecutorShow(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() - executor, _, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) for _, query := range []string{"show vitess_keyspaces", "show keyspaces"} { @@ -1118,12 +1087,9 @@ func TestExecutorShow(t *testing.T) { } func TestExecutorShowTargeted(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() - executor, _, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor/40-60"}) queries := []string{ @@ -1149,12 +1115,9 @@ func TestExecutorShowTargeted(t *testing.T) { } func TestExecutorUse(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() session := NewSafeSession(&vtgatepb.Session{Autocommit: true, TargetString: "@primary"}) stmts := []string{ @@ -1185,12 +1148,8 @@ func TestExecutorUse(t *testing.T) { } func TestExecutorComment(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() stmts := []string{ "/*! SET autocommit=1*/", @@ -1210,12 +1169,8 @@ func TestExecutorComment(t *testing.T) { } func TestExecutorOther(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() type cnts struct { Sbc1Cnt int64 @@ -1307,12 +1262,9 @@ func TestExecutorOther(t *testing.T) { } func TestExecutorDDL(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -1435,12 +1387,8 @@ func TestExecutorDDL(t *testing.T) { } func TestExecutorDDLFk(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, _, _, sbc := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, sbc, ctx, closer := createExecutorEnv(t) + defer closer() mName := "TestExecutorDDLFk" stmts := []string{ @@ -1467,16 +1415,13 @@ func TestExecutorDDLFk(t *testing.T) { } func TestExecutorAlterVSchemaKeyspace(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - vschemaacl.AuthorizedDDLUsers = "%" defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) vschemaUpdates := make(chan *vschemapb.SrvVSchema, 2) @@ -1500,16 +1445,12 @@ func TestExecutorAlterVSchemaKeyspace(t *testing.T) { } func TestExecutorCreateVindexDDL(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - vschemaacl.AuthorizedDDLUsers = "%" defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() ks := "TestExecutor" vschemaUpdates := make(chan *vschemapb.SrvVSchema, 4) @@ -1575,16 +1516,12 @@ func TestExecutorCreateVindexDDL(t *testing.T) { } func TestExecutorAddDropVschemaTableDDL(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - vschemaacl.AuthorizedDDLUsers = "%" defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() ks := KsTestUnsharded vschemaUpdates := make(chan *vschemapb.SrvVSchema, 4) @@ -1632,12 +1569,9 @@ func TestExecutorAddDropVschemaTableDDL(t *testing.T) { } func TestExecutorVindexDDLACL(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() ks := "TestExecutor" session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) @@ -1682,12 +1616,9 @@ func TestExecutorVindexDDLACL(t *testing.T) { } func TestExecutorUnrecognized(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() _, err := executor.Execute(ctx, nil, "TestExecute", NewSafeSession(&vtgatepb.Session{}), "invalid statement", nil) require.Error(t, err, "unrecognized statement: invalid statement'") } @@ -1695,12 +1626,8 @@ func TestExecutorUnrecognized(t *testing.T) { // TestVSchemaStats makes sure the building and displaying of the // VSchemaStats works. func TestVSchemaStats(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - r, _, _, _ := createExecutorEnv(ctx) - defer r.Close() + r, _, _, _, _, closer := createExecutorEnv(t) + defer closer() stats := r.VSchemaStats() @@ -1723,12 +1650,9 @@ func TestVSchemaStats(t *testing.T) { var pv = querypb.ExecuteOptions_Gen4 func TestGetPlanUnnormalized(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + r, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() - r, _, _, _ := createExecutorEnv(ctx) - defer r.Close() emptyvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) unshardedvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) @@ -1817,12 +1741,8 @@ func getPlanCached(t *testing.T, ctx context.Context, e *Executor, vcursor *vcur } func TestGetPlanCacheUnnormalized(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + r, _, _, _, ctx, closer := createExecutorEnv(t) - r, _, _, _ := createExecutorEnv(ctx) - defer r.Close() emptyvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) query1 := "select * from music_user_map where id = 1" @@ -1842,9 +1762,11 @@ func TestGetPlanCacheUnnormalized(t *testing.T) { t.Errorf("logstats sql want \"%s\" got \"%s\"", wantSQL, logStats2.SQL) } + closer() // Skip cache using directive - r, _, _, _ = createExecutorEnv(ctx) - defer r.Close() + r, _, _, _, ctx, closer = createExecutorEnv(t) + defer closer() + unshardedvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) query1 = "insert /*vt+ SKIP_QUERY_PLAN_CACHE=1 */ into user(id) values (1), (2)" @@ -1867,12 +1789,7 @@ func TestGetPlanCacheUnnormalized(t *testing.T) { } func TestGetPlanCacheNormalized(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - r, _, _, _ := createExecutorEnv(ctx) - defer r.Close() + r, _, _, _, ctx, closer := createExecutorEnv(t) r.normalize = true emptyvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) @@ -1886,9 +1803,10 @@ func TestGetPlanCacheNormalized(t *testing.T) { assertCacheSize(t, r.plans, 1) assert.Equal(t, wantSQL, logStats2.SQL) + closer() // Skip cache using directive - r, _, _, _ = createExecutorEnv(ctx) - defer r.Close() + r, _, _, _, ctx, closer = createExecutorEnv(t) + defer closer() r.normalize = true unshardedvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) @@ -1912,12 +1830,9 @@ func TestGetPlanCacheNormalized(t *testing.T) { } func TestGetPlanNormalized(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + r, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() - r, _, _, _ := createExecutorEnv(ctx) - defer r.Close() r.normalize = true emptyvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) unshardedvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) @@ -1973,12 +1888,9 @@ func TestGetPlanPriority(t *testing.T) { testCase := aTestCase t.Run(testCase.name, func(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + r, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() - r, _, _, _ := createExecutorEnv(ctx) - defer r.Close() r.normalize = true logStats := logstats.NewLogStats(ctx, "Test", "", "", nil) vCursor, err := newVCursorImpl(session, makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) @@ -2002,12 +1914,8 @@ func TestGetPlanPriority(t *testing.T) { } func TestPassthroughDDL(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() session := &vtgatepb.Session{ TargetString: "TestExecutor", } @@ -2056,12 +1964,9 @@ func TestPassthroughDDL(t *testing.T) { } func TestParseEmptyTargetSingleKeyspace(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + r, _, _, _, _, closer := createExecutorEnv(t) + defer closer() - r, _, _, _ := createExecutorEnv(ctx) - defer r.Close() altVSchema := &vindexes.VSchema{ Keyspaces: map[string]*vindexes.KeyspaceSchema{ KsTestUnsharded: r.vschema.Keyspaces[KsTestUnsharded], @@ -2083,12 +1988,9 @@ func TestParseEmptyTargetSingleKeyspace(t *testing.T) { } func TestParseEmptyTargetMultiKeyspace(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + r, _, _, _, _, closer := createExecutorEnv(t) + defer closer() - r, _, _, _ := createExecutorEnv(ctx) - defer r.Close() altVSchema := &vindexes.VSchema{ Keyspaces: map[string]*vindexes.KeyspaceSchema{ KsTestUnsharded: r.vschema.Keyspaces[KsTestUnsharded], @@ -2111,12 +2013,9 @@ func TestParseEmptyTargetMultiKeyspace(t *testing.T) { } func TestParseTargetSingleKeyspace(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + r, _, _, _, _, closer := createExecutorEnv(t) + defer closer() - r, _, _, _ := createExecutorEnv(ctx) - defer r.Close() altVSchema := &vindexes.VSchema{ Keyspaces: map[string]*vindexes.KeyspaceSchema{ KsTestUnsharded: r.vschema.Keyspaces[KsTestUnsharded], @@ -2138,15 +2037,12 @@ func TestParseTargetSingleKeyspace(t *testing.T) { } func TestDebugVSchema(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, _, _, closer := createExecutorEnv(t) + defer closer() resp := httptest.NewRecorder() req, _ := http.NewRequest("GET", "/debug/vschema", nil) - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() executor.ServeHTTP(resp, req) v := make(map[string]any) if err := json.Unmarshal(resp.Body.Bytes(), &v); err != nil { @@ -2161,10 +2057,6 @@ func TestDebugVSchema(t *testing.T) { } func TestExecutorMaxPayloadSizeExceeded(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - saveMax := maxPayloadSize saveWarn := warnPayloadSize maxPayloadSize = 10 @@ -2174,8 +2066,9 @@ func TestExecutorMaxPayloadSizeExceeded(t *testing.T) { warnPayloadSize = saveWarn }() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, _, closer := createExecutorEnv(t) + defer closer() + session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) warningCount := warnings.Counts()["WarnPayloadSizeExceeded"] testMaxPayloadSizeExceeded := []string{ @@ -2212,12 +2105,9 @@ func TestExecutorMaxPayloadSizeExceeded(t *testing.T) { } func TestOlapSelectDatabase(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, _, _, closer := createExecutorEnv(t) + defer closer() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() executor.normalize = true session := &vtgatepb.Session{Autocommit: true} @@ -2234,12 +2124,9 @@ func TestOlapSelectDatabase(t *testing.T) { } func TestExecutorClearsWarnings(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, _, _, closer := createExecutorEnv(t) + defer closer() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() session := NewSafeSession(&vtgatepb.Session{ Warnings: []*querypb.QueryWarning{{Code: 234, Message: "oh noes"}}, }) @@ -2250,17 +2137,14 @@ func TestExecutorClearsWarnings(t *testing.T) { // TestServingKeyspaces tests that the dual queries are routed to the correct keyspaces from the list of serving keyspaces. func TestServingKeyspaces(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - buffer.SetBufferingModeInTestingEnv(true) defer func() { buffer.SetBufferingModeInTestingEnv(false) }() - executor, sbc1, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() + executor.pv = querypb.ExecuteOptions_Gen4 gw, ok := executor.resolver.resolver.GetGateway().(*TabletGateway) require.True(t, ok) @@ -2302,12 +2186,8 @@ func TestServingKeyspaces(t *testing.T) { } func TestExecutorOtherRead(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, _, closer := createExecutorEnv(t) + defer closer() type cnts struct { Sbc1Cnt int64 @@ -2382,12 +2262,9 @@ func TestExecutorOtherRead(t *testing.T) { } func TestExecutorVExplain(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() executor.normalize = true logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2410,12 +2287,8 @@ func TestExecutorVExplain(t *testing.T) { } func TestExecutorOtherAdmin(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, _, closer := createExecutorEnv(t) + defer closer() type cnts struct { Sbc1Cnt int64 @@ -2489,12 +2362,9 @@ func TestExecutorOtherAdmin(t *testing.T) { } func TestExecutorSavepointInTx(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) + defer closer() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() logChan := executor.queryLogger.Subscribe("TestExecutorSavepoint") defer executor.queryLogger.Unsubscribe(logChan) @@ -2576,12 +2446,9 @@ func TestExecutorSavepointInTx(t *testing.T) { } func TestExecutorSavepointInTxWithReservedConn(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) + defer closer() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() logChan := executor.queryLogger.Subscribe("TestExecutorSavepoint") defer executor.queryLogger.Unsubscribe(logChan) @@ -2647,12 +2514,9 @@ func TestExecutorSavepointInTxWithReservedConn(t *testing.T) { } func TestExecutorSavepointWithoutTx(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) + defer closer() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() logChan := executor.queryLogger.Subscribe("TestExecutorSavepoint") defer executor.queryLogger.Unsubscribe(logChan) @@ -2695,12 +2559,8 @@ func TestExecutorSavepointWithoutTx(t *testing.T) { } func TestExecutorCallProc(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, sbc1, sbc2, sbcUnsharded := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbcUnsharded, _, closer := createExecutorEnv(t) + defer closer() type cnts struct { Sbc1Cnt int64 @@ -2765,12 +2625,9 @@ func TestExecutorCallProc(t *testing.T) { } func TestExecutorTempTable(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, sbcUnsharded, ctx, closer := createExecutorEnv(t) + defer closer() - executor, _, _, sbcUnsharded := createExecutorEnv(ctx) - defer executor.Close() executor.warnShardedOnly = true creatQuery := "create temporary table temp_t(id bigint primary key)" session := NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded}) @@ -2788,12 +2645,9 @@ func TestExecutorTempTable(t *testing.T) { } func TestExecutorShowVitessMigrations(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) + defer closer() - executor, sbc1, sbc2, _ := createExecutorEnv(ctx) - defer executor.Close() showQuery := "show vitess_migrations" session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) _, err := executor.Execute(ctx, nil, "", session, showQuery, nil) @@ -2803,12 +2657,9 @@ func TestExecutorShowVitessMigrations(t *testing.T) { } func TestExecutorDescHash(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() showQuery := "desc hash_index" session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) _, err := executor.Execute(ctx, nil, "", session, showQuery, nil) @@ -2816,12 +2667,9 @@ func TestExecutorDescHash(t *testing.T) { } func TestExecutorVExplainQueries(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() - executor, _, _, sbclookup := createExecutorEnv(ctx) - defer executor.Close() session := NewAutocommitSession(&vtgatepb.Session{}) sbclookup.SetResults([]*sqltypes.Result{ @@ -2846,12 +2694,9 @@ func TestExecutorVExplainQueries(t *testing.T) { } func TestExecutorStartTxnStmt(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() session := NewAutocommitSession(&vtgatepb.Session{}) tcases := []struct { @@ -2896,12 +2741,9 @@ func TestExecutorStartTxnStmt(t *testing.T) { } func TestExecutorPrepareExecute(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, _, _, closer := createExecutorEnv(t) + defer closer() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() executor.normalize = true session := NewAutocommitSession(&vtgatepb.Session{}) @@ -2940,16 +2782,13 @@ func TestExecutorPrepareExecute(t *testing.T) { } func TestExecutorTruncateErrors(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() save := truncateErrorLen truncateErrorLen = 32 defer func() { truncateErrorLen = save }() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() session := NewSafeSession(&vtgatepb.Session{}) fn := func(r *sqltypes.Result) error { return nil @@ -2966,12 +2805,8 @@ func TestExecutorTruncateErrors(t *testing.T) { } func TestExecutorFlushStmt(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, _, closer := createExecutorEnv(t) + defer closer() tcs := []struct { targetStr string @@ -3028,12 +2863,8 @@ func TestExecutorFlushStmt(t *testing.T) { // TestExecutorKillStmt tests the kill statements on executor. func TestExecutorKillStmt(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, _, closer := createExecutorEnv(t) + defer closer() tcs := []struct { errStr string diff --git a/go/vt/vtgate/executor_vschema_ddl_test.go b/go/vt/vtgate/executor_vschema_ddl_test.go index 73b4864cde0..456fbb4b035 100644 --- a/go/vt/vtgate/executor_vschema_ddl_test.go +++ b/go/vt/vtgate/executor_vschema_ddl_test.go @@ -17,7 +17,6 @@ limitations under the License. package vtgate import ( - "context" "reflect" "sort" "testing" @@ -130,15 +129,12 @@ func waitForColVindexes(t *testing.T, ks, table string, names []string, executor } func TestPlanExecutorAlterVSchemaKeyspace(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() vschemaacl.AuthorizedDDLUsers = "%" defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) vschemaUpdates := make(chan *vschemapb.SrvVSchema, 2) @@ -162,15 +158,12 @@ func TestPlanExecutorAlterVSchemaKeyspace(t *testing.T) { } func TestPlanExecutorCreateVindexDDL(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() vschemaacl.AuthorizedDDLUsers = "%" defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() ks := "TestExecutor" vschemaUpdates := make(chan *vschemapb.SrvVSchema, 4) @@ -208,15 +201,12 @@ func TestPlanExecutorCreateVindexDDL(t *testing.T) { } func TestPlanExecutorDropVindexDDL(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() vschemaacl.AuthorizedDDLUsers = "%" defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() ks := "TestExecutor" vschemaUpdates := make(chan *vschemapb.SrvVSchema, 4) @@ -281,15 +271,12 @@ func TestPlanExecutorDropVindexDDL(t *testing.T) { } func TestPlanExecutorAddDropVschemaTableDDL(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() vschemaacl.AuthorizedDDLUsers = "%" defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() ks := KsTestUnsharded vschemaUpdates := make(chan *vschemapb.SrvVSchema, 4) @@ -342,15 +329,12 @@ func TestPlanExecutorAddDropVschemaTableDDL(t *testing.T) { } func TestExecutorAddSequenceDDL(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() vschemaacl.AuthorizedDDLUsers = "%" defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() ks := KsTestUnsharded vschema := executor.vm.GetCurrentSrvVschema() @@ -406,15 +390,12 @@ func TestExecutorAddSequenceDDL(t *testing.T) { } func TestExecutorAddDropVindexDDL(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() vschemaacl.AuthorizedDDLUsers = "%" defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, sbc1, sbc2, sbclookup := createExecutorEnv(ctx) - defer executor.Close() + executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) + defer closer() ks := "TestExecutor" session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) vschemaUpdates := make(chan *vschemapb.SrvVSchema, 4) @@ -657,11 +638,8 @@ func TestExecutorAddDropVindexDDL(t *testing.T) { func TestPlanExecutorVindexDDLACL(t *testing.T) { // t.Skip("not yet planned") - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() ks := "TestExecutor" session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) diff --git a/go/vt/vtgate/executor_vstream_test.go b/go/vt/vtgate/executor_vstream_test.go index e7e514e8ef3..b146efb27c5 100644 --- a/go/vt/vtgate/executor_vstream_test.go +++ b/go/vt/vtgate/executor_vstream_test.go @@ -17,11 +17,9 @@ limitations under the License. package vtgate import ( - "context" "testing" "time" - "vitess.io/vitess/go/test/utils" "vitess.io/vitess/go/vt/vtgate/engine" querypb "vitess.io/vitess/go/vt/proto/query" @@ -39,10 +37,8 @@ import ( // TestVStreamSQLUnsharded tests the experimental 'vstream * from' vtgate olap query func TestVStreamSQLUnsharded(t *testing.T) { t.Skip("this test is failing due to races") // FIXME - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, sbcLookup := createExecutorEnv(ctx) + executor, _, _, sbcLookup, ctx, closer := createExecutorEnv(t) + defer closer() logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) send1 := []*binlogdatapb.VEvent{ diff --git a/go/vt/vtgate/plugin_mysql_server_test.go b/go/vt/vtgate/plugin_mysql_server_test.go index 22c0d24e8f9..a63872d4791 100644 --- a/go/vt/vtgate/plugin_mysql_server_test.go +++ b/go/vt/vtgate/plugin_mysql_server_test.go @@ -308,11 +308,8 @@ func testInitTLSConfig(t *testing.T, serverCA bool) { // TestKillMethods test the mysql plugin for kill method calls. func TestKillMethods(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() + executor, _, _, _, _, closer := createExecutorEnv(t) + defer closer() vh := newVtgateHandler(&VTGate{executor: executor}) // connection does not exist diff --git a/go/vt/vtgate/queryz_test.go b/go/vt/vtgate/queryz_test.go index 1c8e9d9b190..ff864393899 100644 --- a/go/vt/vtgate/queryz_test.go +++ b/go/vt/vtgate/queryz_test.go @@ -17,7 +17,6 @@ limitations under the License. package vtgate import ( - "context" "io" "net/http" "net/http/httptest" @@ -28,7 +27,6 @@ import ( "github.com/stretchr/testify/require" - "vitess.io/vitess/go/test/utils" vtgatepb "vitess.io/vitess/go/vt/proto/vtgate" "vitess.io/vitess/go/sqltypes" @@ -38,16 +36,12 @@ import ( ) func TestQueryzHandler(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + executor, _, _, _, ctx, closer := createExecutorEnv(t) + defer closer() resp := httptest.NewRecorder() req, _ := http.NewRequest("GET", "/schemaz", nil) - executor, _, _, _ := createExecutorEnv(ctx) - defer executor.Close() - session := &vtgatepb.Session{TargetString: "@primary"} // single shard query sql := "select id from user where id = 1" diff --git a/go/vt/vtgate/vtgate.go b/go/vt/vtgate/vtgate.go index 485759185de..eee0442fda7 100644 --- a/go/vt/vtgate/vtgate.go +++ b/go/vt/vtgate/vtgate.go @@ -666,11 +666,6 @@ func (vtg *VTGate) HandlePanic(err *error) { } } -func (vtg *VTGate) Close() { - vtg.executor.Close() - vtg.gw.Close(context.Background()) -} - func newVTGate(executor *Executor, resolver *Resolver, vsm *vstreamManager, tc *TxConn, gw *TabletGateway) *VTGate { return &VTGate{ executor: executor, diff --git a/go/vt/vtgate/vtgate_test.go b/go/vt/vtgate/vtgate_test.go index a2ed31af958..244bfbf88b8 100644 --- a/go/vt/vtgate/vtgate_test.go +++ b/go/vt/vtgate/vtgate_test.go @@ -45,11 +45,8 @@ var executeOptions = &querypb.ExecuteOptions{ } func TestVTGateExecute(t *testing.T) { - utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - vtg, sbc := createVtgateEnv(ctx) - defer vtg.Close() + vtg, sbc, ctx, closer := createVtgateEnv(t) + defer closer() counts := vtg.timings.Timings.Counts() _, qr, err := vtg.Execute( @@ -88,11 +85,9 @@ func TestVTGateExecute(t *testing.T) { } func TestVTGateExecuteError(t *testing.T) { - utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - vtg, _ := createVtgateEnv(ctx) - defer vtg.Close() + vtg, _, ctx, closer := createVtgateEnv(t) + defer closer() + counts := vtg.timings.Timings.Counts() _, qr, err := vtg.Execute( @@ -121,11 +116,9 @@ func TestVTGateExecuteError(t *testing.T) { } func TestVTGatePrepare(t *testing.T) { - utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - vtg, sbc := createVtgateEnv(ctx) - defer vtg.Close() + vtg, sbc, ctx, closer := createVtgateEnv(t) + defer closer() + counts := vtg.timings.Timings.Counts() _, qr, err := vtg.Prepare( ctx, @@ -161,11 +154,9 @@ func TestVTGatePrepare(t *testing.T) { } func TestVTGatePrepareError(t *testing.T) { - utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - vtg, _ := createVtgateEnv(ctx) - defer vtg.Close() + vtg, _, ctx, closer := createVtgateEnv(t) + defer closer() + counts := errorCounts.Counts() _, qr, err := vtg.Prepare( @@ -193,11 +184,8 @@ func TestVTGatePrepareError(t *testing.T) { } func TestVTGateExecuteWithKeyspaceShard(t *testing.T) { - utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - vtg, _ := createVtgateEnv(ctx) - defer vtg.Close() + vtg, _, ctx, closer := createVtgateEnv(t) + defer closer() // Valid keyspace. _, qr, err := vtg.Execute( @@ -259,11 +247,8 @@ func TestVTGateExecuteWithKeyspaceShard(t *testing.T) { } func TestVTGateStreamExecute(t *testing.T) { - utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - vtg, sbc := createVtgateEnv(ctx) - defer vtg.Close() + vtg, sbc, ctx, closer := createVtgateEnv(t) + defer closer() var qrs []*sqltypes.Result _, err := vtg.StreamExecute( @@ -293,11 +278,8 @@ func TestVTGateStreamExecute(t *testing.T) { } func TestVTGateBindVarError(t *testing.T) { - utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - vtg, _ := createVtgateEnv(ctx) - defer vtg.Close() + vtg, _, ctx, closer := createVtgateEnv(t) + defer closer() session := &vtgatepb.Session{} bindVars := map[string]*querypb.BindVariable{ @@ -398,11 +380,8 @@ func testErrorPropagation(t *testing.T, ctx context.Context, vtg *VTGate, sbcs [ // tablet and a rdonly tablet because we don't control the routing of // Commit. func TestErrorPropagation(t *testing.T) { - utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - vtg, sbc := createVtgateEnv(ctx) - defer vtg.Close() + vtg, sbc, ctx, closer := createVtgateEnv(t) + defer closer() sbcs := []*sandboxconn.SandboxConn{ sbc, @@ -484,11 +463,8 @@ func TestErrorPropagation(t *testing.T) { // This test makes sure that if we start a transaction and hit a critical // error, a rollback is issued. func TestErrorIssuesRollback(t *testing.T) { - utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - vtg, sbc := createVtgateEnv(ctx) - defer vtg.Close() + vtg, sbc, ctx, closer := createVtgateEnv(t) + defer closer() // Start a transaction, send one statement. // Simulate an error that should trigger a rollback: @@ -618,11 +594,8 @@ var shardedVSchemaUnknownParams = ` ` func TestMultiInternalSavepointVtGate(t *testing.T) { - utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - vtg, _ := createVtgateEnv(ctx) - defer vtg.Close() + vtg, _, ctx, closer := createVtgateEnv(t) + defer closer() const customKeyspace = "CustomSharding" s := createSandbox(customKeyspace) @@ -720,11 +693,8 @@ func TestMultiInternalSavepointVtGate(t *testing.T) { } func TestVSchemaVindexUnknownParams(t *testing.T) { - utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - vtg, _ := createVtgateEnv(ctx) - defer vtg.Close() + vtg, _, _, closer := createVtgateEnv(t) + defer closer() const customKeyspace = "CustomSharding" s := createSandbox(customKeyspace) @@ -756,13 +726,15 @@ func TestVSchemaVindexUnknownParams(t *testing.T) { require.Equal(t, int64(0), unknownParams) } -func createVtgateEnv(ctx context.Context) (*VTGate, *sandboxconn.SandboxConn) { +func createVtgateEnv(t testing.TB) (*VTGate, *sandboxconn.SandboxConn, context.Context, func()) { cell := "aa" sb := createSandbox(KsTestSharded) sb.ShardSpec = "-" - executor, _, _, sbc := createExecutorEnv(ctx) + executor, _, _, sbc, ctx, closer := createExecutorEnv(t) executor.normalize = normalizeQueries vsm := newVStreamManager(executor.resolver.resolver, executor.serv, cell) - return newVTGate(executor, executor.resolver, vsm, nil, executor.scatterConn.gateway), sbc + vtg := newVTGate(executor, executor.resolver, vsm, nil, executor.scatterConn.gateway) + + return vtg, sbc, ctx, closer } From 864a52842368840b355459e7ea8f0bcbb4be343a Mon Sep 17 00:00:00 2001 From: Rohit Nayak Date: Wed, 23 Aug 2023 21:35:55 +0200 Subject: [PATCH 08/17] Plug some tabletmanager/vreplication unit test leaks: stats and rowstreamer Signed-off-by: Rohit Nayak --- go/vt/vttablet/tabletmanager/vreplication/stats_test.go | 3 ++- .../tabletmanager/vreplication/vplayer_flaky_test.go | 1 + .../vttablet/tabletmanager/vreplication/vreplicator_test.go | 2 ++ go/vt/vttablet/tabletserver/vstreamer/rowstreamer.go | 5 ++++- 4 files changed, 9 insertions(+), 2 deletions(-) diff --git a/go/vt/vttablet/tabletmanager/vreplication/stats_test.go b/go/vt/vttablet/tabletmanager/vreplication/stats_test.go index 1dc771f1fc3..d5b5eacbdf2 100644 --- a/go/vt/vttablet/tabletmanager/vreplication/stats_test.go +++ b/go/vt/vttablet/tabletmanager/vreplication/stats_test.go @@ -81,6 +81,7 @@ func TestStatusHtml(t *testing.T) { } blpStats := binlogplayer.NewStats() + defer blpStats.Stop() blpStats.SetLastPosition(pos) blpStats.ReplicationLagSeconds.Store(2) blpStats.History.Add(&binlogplayer.StatsHistoryRecord{Time: time.Now(), Message: "Test Message1"}) @@ -130,7 +131,7 @@ func TestStatusHtml(t *testing.T) { func TestVReplicationStats(t *testing.T) { blpStats := binlogplayer.NewStats() - + defer blpStats.Stop() testStats := &vrStats{} testStats.isOpen = true testStats.controllers = map[int32]*controller{ diff --git a/go/vt/vttablet/tabletmanager/vreplication/vplayer_flaky_test.go b/go/vt/vttablet/tabletmanager/vreplication/vplayer_flaky_test.go index 369080d6312..2aa5697df63 100644 --- a/go/vt/vttablet/tabletmanager/vreplication/vplayer_flaky_test.go +++ b/go/vt/vttablet/tabletmanager/vreplication/vplayer_flaky_test.go @@ -202,6 +202,7 @@ func TestHeartbeatFrequencyFlag(t *testing.T) { }() stats := binlogplayer.NewStats() + defer stats.Stop() vp := &vplayer{vr: &vreplicator{dbClient: newVDBClient(realDBClientFactory(), stats), stats: stats}} type testcount struct { diff --git a/go/vt/vttablet/tabletmanager/vreplication/vreplicator_test.go b/go/vt/vttablet/tabletmanager/vreplication/vreplicator_test.go index 66591bbcb81..9821c4144a8 100644 --- a/go/vt/vttablet/tabletmanager/vreplication/vreplicator_test.go +++ b/go/vt/vttablet/tabletmanager/vreplication/vreplicator_test.go @@ -207,6 +207,7 @@ func TestDeferSecondaryKeys(t *testing.T) { id := int32(1) vsclient := newTabletConnector(tablet) stats := binlogplayer.NewStats() + stats.Stop() dbClient := playerEngine.dbClientFactoryFiltered() err := dbClient.Connect() require.NoError(t, err) @@ -538,6 +539,7 @@ func TestCancelledDeferSecondaryKeys(t *testing.T) { id := int32(1) vsclient := newTabletConnector(tablet) stats := binlogplayer.NewStats() + stats.Stop() dbaconn := playerEngine.dbClientFactoryDba() err = dbaconn.Connect() require.NoError(t, err) diff --git a/go/vt/vttablet/tabletserver/vstreamer/rowstreamer.go b/go/vt/vttablet/tabletserver/vstreamer/rowstreamer.go index 7fcbbb65735..86b50947cdf 100644 --- a/go/vt/vttablet/tabletserver/vstreamer/rowstreamer.go +++ b/go/vt/vttablet/tabletserver/vstreamer/rowstreamer.go @@ -331,7 +331,10 @@ func (rs *rowStreamer) streamQuery(conn *snapshotConn, send func(*binlogdatapb.V heartbeatTicker := time.NewTicker(rowStreamertHeartbeatInterval) defer heartbeatTicker.Stop() go func() { - for range heartbeatTicker.C { + select { + case <-rs.ctx.Done(): + return + case <-heartbeatTicker.C: safeSend(&binlogdatapb.VStreamRowsResponse{Heartbeat: true}) } }() From cfe7a28501f63e054ab08bc1816e2b4831584438 Mon Sep 17 00:00:00 2001 From: Rohit Nayak Date: Wed, 23 Aug 2023 23:18:08 +0200 Subject: [PATCH 09/17] Added utility GetLeaks() to log leaks esp. in TestMain() and related refactor Signed-off-by: Rohit Nayak --- go/test/utils/noleak.go | 39 ++++++++++++++++++++++++++++++--------- 1 file changed, 30 insertions(+), 9 deletions(-) diff --git a/go/test/utils/noleak.go b/go/test/utils/noleak.go index adfceb90b9a..683a76483ae 100644 --- a/go/test/utils/noleak.go +++ b/go/test/utils/noleak.go @@ -9,17 +9,37 @@ import ( "time" "go.uber.org/goleak" + + "vitess.io/vitess/go/vt/log" ) +// EnsureNoLeaks checks for goroutine and socket leaks and fails the test if any are found. func EnsureNoLeaks(t testing.TB) { if t.Failed() { return } - ensureNoGoroutines(t) - ensureNoOpenSockets(t) + if err := ensureNoLeaks(); err != nil { + t.Fatal(err) + } +} + +// GetLeaks checks for goroutine and socket leaks and returns an error if any are found. +// One use case is in TestMain()s to ensure that all tests are cleaned up. +func GetLeaks() error { + return ensureNoLeaks() +} + +func ensureNoLeaks() error { + if err := ensureNoGoroutines(); err != nil { + return err + } + if err := ensureNoOpenSockets(); err != nil { + return err + } + return nil } -func ensureNoGoroutines(t testing.TB) { +func ensureNoGoroutines() error { var ignored = []goleak.Option{ goleak.IgnoreTopFunction("github.com/golang/glog.(*fileSink).flushDaemon"), goleak.IgnoreTopFunction("github.com/golang/glog.(*loggingT).flushDaemon"), @@ -35,23 +55,24 @@ func ensureNoGoroutines(t testing.TB) { for i := 0; i < 5; i++ { err = goleak.Find(ignored...) if err == nil { - return + return nil } time.Sleep(100 * time.Millisecond) } - t.Fatal(err) + return err } -func ensureNoOpenSockets(t testing.TB) { +func ensureNoOpenSockets() error { cmd := exec.Command("lsof", "-a", "-p", strconv.Itoa(os.Getpid()), "-i", "-P", "-V") cmd.Stderr = nil lsof, err := cmd.Output() if err == nil { - t.Errorf("found open sockets:\n%s", lsof) + log.Errorf("found open sockets:\n%s", lsof) } else { if strings.Contains(string(lsof), "no Internet files located") { - return + return nil } - t.Errorf("failed to run `lsof`: %v (%q)", err, lsof) + log.Errorf("failed to run `lsof`: %v (%q)", err, lsof) } + return err } From be9e89fcd55019a8224e3cfc7df9b27edd8dd5cd Mon Sep 17 00:00:00 2001 From: Rohit Nayak Date: Thu, 24 Aug 2023 10:07:31 +0200 Subject: [PATCH 10/17] Duh, defer stats.Stop() Signed-off-by: Rohit Nayak --- go/vt/vttablet/tabletmanager/vreplication/vreplicator_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/go/vt/vttablet/tabletmanager/vreplication/vreplicator_test.go b/go/vt/vttablet/tabletmanager/vreplication/vreplicator_test.go index 9821c4144a8..58f38246af3 100644 --- a/go/vt/vttablet/tabletmanager/vreplication/vreplicator_test.go +++ b/go/vt/vttablet/tabletmanager/vreplication/vreplicator_test.go @@ -207,7 +207,7 @@ func TestDeferSecondaryKeys(t *testing.T) { id := int32(1) vsclient := newTabletConnector(tablet) stats := binlogplayer.NewStats() - stats.Stop() + defer stats.Stop() dbClient := playerEngine.dbClientFactoryFiltered() err := dbClient.Connect() require.NoError(t, err) @@ -539,7 +539,7 @@ func TestCancelledDeferSecondaryKeys(t *testing.T) { id := int32(1) vsclient := newTabletConnector(tablet) stats := binlogplayer.NewStats() - stats.Stop() + defer stats.Stop() dbaconn := playerEngine.dbClientFactoryDba() err = dbaconn.Connect() require.NoError(t, err) From b81debc763a6a92902a40b1e5243839e3ab8b75f Mon Sep 17 00:00:00 2001 From: Rohit Nayak Date: Thu, 24 Aug 2023 18:33:40 +0200 Subject: [PATCH 11/17] License Signed-off-by: Rohit Nayak --- go/test/utils/noleak.go | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/go/test/utils/noleak.go b/go/test/utils/noleak.go index 683a76483ae..ee591654bfa 100644 --- a/go/test/utils/noleak.go +++ b/go/test/utils/noleak.go @@ -1,3 +1,19 @@ +/* +Copyright 2023 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 utils import ( From 649a7a46086dadf2595e787f161247803675032e Mon Sep 17 00:00:00 2001 From: Rohit Nayak Date: Thu, 24 Aug 2023 20:47:22 +0200 Subject: [PATCH 12/17] Change plain utils.EnsureNoLeaks(0 to defers Signed-off-by: Rohit Nayak --- go/vt/discovery/healthcheck_test.go | 32 ++++++++++++------------ go/vt/discovery/keyspace_events_test.go | 2 +- go/vt/discovery/replicationlag_test.go | 8 +++--- go/vt/discovery/tablet_picker_test.go | 18 ++++++------- go/vt/discovery/topology_watcher_test.go | 8 +++--- go/vt/mysqlctl/backup_blackbox_test.go | 8 +++--- go/vt/mysqlctl/backup_test.go | 22 ++++++++-------- go/vt/vtadmin/cluster/cluster_test.go | 8 +++--- go/vt/vtexplain/vtexplain_test.go | 4 +-- go/vt/vtgate/executor_framework_test.go | 4 +-- 10 files changed, 57 insertions(+), 57 deletions(-) diff --git a/go/vt/discovery/healthcheck_test.go b/go/vt/discovery/healthcheck_test.go index 2a5f5b18505..19780081bd2 100644 --- a/go/vt/discovery/healthcheck_test.go +++ b/go/vt/discovery/healthcheck_test.go @@ -64,7 +64,7 @@ func init() { } func TestHealthCheck(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() // reset error counters @@ -210,7 +210,7 @@ func TestHealthCheck(t *testing.T) { } func TestHealthCheckStreamError(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -276,7 +276,7 @@ func TestHealthCheckStreamError(t *testing.T) { // TestHealthCheckErrorOnPrimary is the same as TestHealthCheckStreamError except for tablet type func TestHealthCheckErrorOnPrimary(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -341,7 +341,7 @@ func TestHealthCheckErrorOnPrimary(t *testing.T) { } func TestHealthCheckErrorOnPrimaryAfterExternalReparent(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -424,7 +424,7 @@ func TestHealthCheckErrorOnPrimaryAfterExternalReparent(t *testing.T) { } func TestHealthCheckVerifiesTabletAlias(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -472,7 +472,7 @@ func TestHealthCheckVerifiesTabletAlias(t *testing.T) { // TestHealthCheckCloseWaitsForGoRoutines tests that Close() waits for all Go // routines to finish and the listener won't be called anymore. func TestHealthCheckCloseWaitsForGoRoutines(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -538,7 +538,7 @@ func TestHealthCheckCloseWaitsForGoRoutines(t *testing.T) { } func TestHealthCheckTimeout(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -615,7 +615,7 @@ func TestHealthCheckTimeout(t *testing.T) { } func TestWaitForAllServingTablets(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -712,7 +712,7 @@ func TestWaitForAllServingTablets(t *testing.T) { // TestRemoveTablet tests the behavior when a tablet goes away. func TestRemoveTablet(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -824,7 +824,7 @@ func TestRemoveTablet(t *testing.T) { // TestGetHealthyTablets tests the functionality of GetHealthyTabletStats. func TestGetHealthyTablets(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -1015,7 +1015,7 @@ func TestGetHealthyTablets(t *testing.T) { } func TestPrimaryInOtherCell(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -1077,7 +1077,7 @@ func TestPrimaryInOtherCell(t *testing.T) { } func TestReplicaInOtherCell(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -1184,7 +1184,7 @@ func TestReplicaInOtherCell(t *testing.T) { } func TestCellAliases(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -1249,7 +1249,7 @@ func TestCellAliases(t *testing.T) { } func TestHealthCheckChecksGrpcPort(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -1274,7 +1274,7 @@ func TestHealthCheckChecksGrpcPort(t *testing.T) { } func TestTemplate(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) TabletURLTemplateString = "http://{{.GetTabletHostPort}}" ParseTabletURLTemplateFromFlag() @@ -1302,7 +1302,7 @@ func TestTemplate(t *testing.T) { } func TestDebugURLFormatting(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) TabletURLTemplateString = "https://{{.GetHostNameLevel 0}}.bastion.{{.Tablet.Alias.Cell}}.corp" ParseTabletURLTemplateFromFlag() diff --git a/go/vt/discovery/keyspace_events_test.go b/go/vt/discovery/keyspace_events_test.go index a64ef39558c..09aedd35061 100644 --- a/go/vt/discovery/keyspace_events_test.go +++ b/go/vt/discovery/keyspace_events_test.go @@ -34,7 +34,7 @@ import ( ) func TestSrvKeyspaceWithNilNewKeyspace(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() cell := "cell" diff --git a/go/vt/discovery/replicationlag_test.go b/go/vt/discovery/replicationlag_test.go index 4a96b68336a..5cef05a3f4b 100644 --- a/go/vt/discovery/replicationlag_test.go +++ b/go/vt/discovery/replicationlag_test.go @@ -45,7 +45,7 @@ func testSetMinNumTablets(newMin int) { } func TestFilterByReplicationLagUnhealthy(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) // 1 healthy serving tablet, 1 not healthy ts1 := &TabletHealth{ Tablet: topo.NewTablet(1, "cell", "host1"), @@ -63,7 +63,7 @@ func TestFilterByReplicationLagUnhealthy(t *testing.T) { } func TestFilterByReplicationLag(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) // Use simplified logic testSetLegacyReplicationLagAlgorithm(false) @@ -140,7 +140,7 @@ func TestFilterByReplicationLag(t *testing.T) { } func TestFilterByReplicationLagThreeTabletMin(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) // Use at least 3 tablets if possible testSetMinNumTablets(3) // lags of (1s, 1s, 10m, 11m) - returns at least32 items where the slightly delayed ones that are returned are the 10m and 11m ones. @@ -197,7 +197,7 @@ func TestFilterByReplicationLagThreeTabletMin(t *testing.T) { } func TestFilterStatsByReplicationLagOneTabletMin(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) // Use at least 1 tablets if possible testSetMinNumTablets(1) // lags of (1s, 100m) - return only healthy tablet if that is all that is available. diff --git a/go/vt/discovery/tablet_picker_test.go b/go/vt/discovery/tablet_picker_test.go index add9ab59320..5bfdb144024 100644 --- a/go/vt/discovery/tablet_picker_test.go +++ b/go/vt/discovery/tablet_picker_test.go @@ -32,7 +32,7 @@ import ( ) func TestPickPrimary(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) defer cancel() @@ -58,7 +58,7 @@ func TestPickPrimary(t *testing.T) { } func TestPickLocalPreferences(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) type tablet struct { id uint32 typ topodatapb.TabletType @@ -305,7 +305,7 @@ func TestPickLocalPreferences(t *testing.T) { } func TestPickCellPreferenceLocalCell(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) defer cancel() @@ -346,7 +346,7 @@ func TestPickCellPreferenceLocalCell(t *testing.T) { } func TestPickCellPreferenceLocalAlias(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) defer cancel() @@ -364,7 +364,7 @@ func TestPickCellPreferenceLocalAlias(t *testing.T) { } func TestPickUsingCellAliasOnlySpecified(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) defer cancel() @@ -413,7 +413,7 @@ func TestPickUsingCellAliasOnlySpecified(t *testing.T) { } func TestTabletAppearsDuringSleep(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) defer cancel() @@ -445,7 +445,7 @@ func TestTabletAppearsDuringSleep(t *testing.T) { } func TestPickErrorLocalPreferenceDefault(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -477,7 +477,7 @@ func TestPickErrorLocalPreferenceDefault(t *testing.T) { } func TestPickErrorOnlySpecified(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -511,7 +511,7 @@ func TestPickErrorOnlySpecified(t *testing.T) { // type has no available healthy serving tablets that we select a healthy // serving tablet from the secondary/second type. func TestPickFallbackType(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() diff --git a/go/vt/discovery/topology_watcher_test.go b/go/vt/discovery/topology_watcher_test.go index 95c4e9e2e3c..6ca8b907aac 100644 --- a/go/vt/discovery/topology_watcher_test.go +++ b/go/vt/discovery/topology_watcher_test.go @@ -58,7 +58,7 @@ func checkChecksum(t *testing.T, tw *TopologyWatcher, want uint32) { } func TestStartAndCloseTopoWatcher(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -118,7 +118,7 @@ func TestCellTabletsWatcherNoRefreshKnown(t *testing.T) { } func checkWatcher(t *testing.T, refreshKnownTablets bool) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -443,7 +443,7 @@ var ( ) func TestFilterByKeyspace(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -528,7 +528,7 @@ func TestFilterByKeyspace(t *testing.T) { // - does not continuosly call GetTablets for tablets that do not satisfy the filter // - does not add or remove these filtered out tablets from the its healtcheck func TestFilterByKeypsaceSkipsIgnoredTablets(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() diff --git a/go/vt/mysqlctl/backup_blackbox_test.go b/go/vt/mysqlctl/backup_blackbox_test.go index bf88f4b1791..cd6fcdac535 100644 --- a/go/vt/mysqlctl/backup_blackbox_test.go +++ b/go/vt/mysqlctl/backup_blackbox_test.go @@ -80,7 +80,7 @@ func createBackupFiles(root string, fileCount int, ext string) error { } func TestExecuteBackup(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() // Set up local backup directory @@ -226,7 +226,7 @@ func TestExecuteBackup(t *testing.T) { } func TestExecuteBackupWithSafeUpgrade(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -318,7 +318,7 @@ func TestExecuteBackupWithSafeUpgrade(t *testing.T) { // TestExecuteBackupWithCanceledContext tests the ability of the backup function to gracefully handle cases where errors // occur due to various reasons, such as context time cancel. The process should not panic in these situations. func TestExecuteBackupWithCanceledContext(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -408,7 +408,7 @@ func TestExecuteBackupWithCanceledContext(t *testing.T) { // TestExecuteRestoreWithCanceledContext tests the ability of the restore function to gracefully handle cases where errors // occur due to various reasons, such as context timed-out. The process should not panic in these situations. func TestExecuteRestoreWithTimedOutContext(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() diff --git a/go/vt/mysqlctl/backup_test.go b/go/vt/mysqlctl/backup_test.go index 878adefb9a6..c16dec64e50 100644 --- a/go/vt/mysqlctl/backup_test.go +++ b/go/vt/mysqlctl/backup_test.go @@ -45,7 +45,7 @@ import ( // TestBackupExecutesBackupWithScopedParams tests that Backup passes // a Scope()-ed stats to backupengine ExecuteBackup. func TestBackupExecutesBackupWithScopedParams(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() @@ -68,7 +68,7 @@ func TestBackupExecutesBackupWithScopedParams(t *testing.T) { // TestBackupNoStats tests that if BackupParams.Stats is nil, then Backup will // pass non-nil Stats to sub-components. func TestBackupNoStats(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() @@ -84,7 +84,7 @@ func TestBackupNoStats(t *testing.T) { // TestBackupParameterizesBackupStorageWithScopedStats tests that Backup passes // a Scope()-ed stats to BackupStorage.WithParams. func TestBackupParameterizesBackupStorageWithScopedStats(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() @@ -105,7 +105,7 @@ func TestBackupParameterizesBackupStorageWithScopedStats(t *testing.T) { // TestBackupEmitsStats tests that Backup emits stats. func TestBackupEmitsStats(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() @@ -123,7 +123,7 @@ func TestBackupEmitsStats(t *testing.T) { // backupstorage.Params to backupstorage, but only if it responds to // backupstorage.WithParams. func TestBackupTriesToParameterizeBackupStorage(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() @@ -329,7 +329,7 @@ func TestFindFilesToBackupWithRedoLog(t *testing.T) { // TestRestoreEmitsStats tests that Restore emits stats. func TestRestoreEmitsStats(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() @@ -347,7 +347,7 @@ func TestRestoreEmitsStats(t *testing.T) { // TestRestoreExecutesRestoreWithScopedParams tests that Restore passes // a Scope()-ed stats to backupengine ExecuteRestore. func TestRestoreExecutesRestoreWithScopedParams(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() @@ -371,7 +371,7 @@ func TestRestoreExecutesRestoreWithScopedParams(t *testing.T) { // TestRestoreNoStats tests that if RestoreParams.Stats is nil, then Restore will // pass non-nil Stats to sub-components. func TestRestoreNoStats(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() @@ -388,7 +388,7 @@ func TestRestoreNoStats(t *testing.T) { // TestRestoreParameterizesBackupStorageWithScopedStats tests that Restore passes // a Scope()-ed stats to BackupStorage.WithParams. func TestRestoreParameterizesBackupStorageWithScopedStats(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() @@ -412,7 +412,7 @@ func TestRestoreParameterizesBackupStorageWithScopedStats(t *testing.T) { // backupstorage.Params to backupstorage, but only if it responds to // backupstorage.WithParams. func TestRestoreTriesToParameterizeBackupStorage(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() @@ -524,7 +524,7 @@ func TestRestoreManifestMySQLVersionValidation(t *testing.T) { for _, tc := range testCases { t.Run(fmt.Sprintf("%s->%s upgradeSafe=%t", tc.fromVersion, tc.toVersion, tc.upgradeSafe), func(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) env, closer := createFakeBackupRestoreEnv(t) defer closer() env.mysqld.Version = tc.toVersion diff --git a/go/vt/vtadmin/cluster/cluster_test.go b/go/vt/vtadmin/cluster/cluster_test.go index 3d87399c8c9..c1a7f4a47b1 100644 --- a/go/vt/vtadmin/cluster/cluster_test.go +++ b/go/vt/vtadmin/cluster/cluster_test.go @@ -50,7 +50,7 @@ import ( ) func TestCreateKeyspace(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) tests := []struct { name string @@ -177,7 +177,7 @@ func TestCreateKeyspace(t *testing.T) { } func TestCreateShard(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) type test struct { name string @@ -279,7 +279,7 @@ func TestCreateShard(t *testing.T) { } func TestDeleteKeyspace(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -364,7 +364,7 @@ func TestDeleteKeyspace(t *testing.T) { } func TestDeleteShards(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) t.Parallel() type test struct { diff --git a/go/vt/vtexplain/vtexplain_test.go b/go/vt/vtexplain/vtexplain_test.go index 85b41675e19..e764ce6f3e8 100644 --- a/go/vt/vtexplain/vtexplain_test.go +++ b/go/vt/vtexplain/vtexplain_test.go @@ -155,7 +155,7 @@ func TestExplain(t *testing.T) { } func TestErrors(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -198,7 +198,7 @@ func TestErrors(t *testing.T) { } func TestJSONOutput(t *testing.T) { - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() diff --git a/go/vt/vtgate/executor_framework_test.go b/go/vt/vtgate/executor_framework_test.go index d3bbad9752e..997d6da87aa 100644 --- a/go/vt/vtgate/executor_framework_test.go +++ b/go/vt/vtgate/executor_framework_test.go @@ -182,7 +182,7 @@ func createExecutorEnv(t testing.TB) (executor *Executor, sbc1, sbc2, sbclookup return executor, sbc1, sbc2, sbclookup, ctx, func() { executor.Close() cancel() - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) } } @@ -241,7 +241,7 @@ func createCustomExecutorSetValues(t testing.TB, vschema string, values []*sqlty return executor, sbcs[0], sbcs[1], sbclookup, ctx, func() { executor.Close() cancel() - utils.EnsureNoLeaks(t) + defer utils.EnsureNoLeaks(t) } } From 959cff2aeadc5dfd082bbcd4992e288f485aa518 Mon Sep 17 00:00:00 2001 From: Rohit Nayak Date: Thu, 24 Aug 2023 21:41:59 +0200 Subject: [PATCH 13/17] Found testing.tRunner.func1 for TestDeleteShards, for now adding to acceptable goleaks Signed-off-by: Rohit Nayak --- go/test/utils/noleak.go | 1 + 1 file changed, 1 insertion(+) diff --git a/go/test/utils/noleak.go b/go/test/utils/noleak.go index ee591654bfa..7a27a815cd9 100644 --- a/go/test/utils/noleak.go +++ b/go/test/utils/noleak.go @@ -65,6 +65,7 @@ func ensureNoGoroutines() error { goleak.IgnoreTopFunction("github.com/patrickmn/go-cache.(*janitor).Run"), goleak.IgnoreTopFunction("vitess.io/vitess/go/vt/logutil.(*ThrottledLogger).log.func1"), goleak.IgnoreTopFunction("vitess.io/vitess/go/vt/vttablet/tabletserver/throttle.initThrottleTicker.func1.1"), + goleak.IgnoreTopFunction("testing.tRunner.func1"), } var err error From d52a492119955d569fe761eedeebb6be3684e082 Mon Sep 17 00:00:00 2001 From: Vicent Marti Date: Fri, 25 Aug 2023 12:22:23 +0200 Subject: [PATCH 14/17] vttest: use testing.T.Cleanup Signed-off-by: Vicent Marti --- go/test/utils/noleak.go | 19 ++ go/vt/discovery/healthcheck_test.go | 57 +--- go/vt/discovery/keyspace_events_test.go | 4 +- go/vt/discovery/tablet_picker_test.go | 29 +- go/vt/discovery/topology_watcher_test.go | 16 +- go/vt/mysqlctl/backup_blackbox_test.go | 17 +- go/vt/mysqlctl/backup_test.go | 59 ++-- go/vt/topo/helpers/tee_topo_test.go | 5 +- go/vt/vtadmin/cluster/cluster_test.go | 14 +- .../endtoend/init_shard_primary_test.go | 4 +- go/vt/vtexplain/vtexplain_test.go | 14 +- go/vt/vtgate/autocommit_test.go | 131 ++++---- go/vt/vtgate/bench_test.go | 6 +- go/vt/vtgate/executor_ddl_test.go | 3 +- go/vt/vtgate/executor_dml_test.go | 174 ++++------ go/vt/vtgate/executor_framework_test.go | 33 +- go/vt/vtgate/executor_scatter_stats_test.go | 9 +- go/vt/vtgate/executor_select_test.go | 297 ++++++------------ go/vt/vtgate/executor_set_test.go | 124 ++++---- go/vt/vtgate/executor_stream_test.go | 7 +- go/vt/vtgate/executor_test.go | 297 ++++++++---------- go/vt/vtgate/executor_vschema_ddl_test.go | 21 +- go/vt/vtgate/executor_vstream_test.go | 3 +- go/vt/vtgate/legacy_scatter_conn_test.go | 36 +-- go/vt/vtgate/plugin_mysql_server_test.go | 7 +- go/vt/vtgate/queryz_test.go | 3 +- go/vt/vtgate/scatter_conn_test.go | 17 +- go/vt/vtgate/tabletgateway_flaky_test.go | 13 +- go/vt/vtgate/tabletgateway_test.go | 16 +- go/vt/vtgate/tx_conn_test.go | 132 ++------ go/vt/vtgate/vstream_manager_test.go | 4 +- go/vt/vtgate/vtgate_test.go | 39 +-- 32 files changed, 584 insertions(+), 1026 deletions(-) diff --git a/go/test/utils/noleak.go b/go/test/utils/noleak.go index 7a27a815cd9..a492ab670f5 100644 --- a/go/test/utils/noleak.go +++ b/go/test/utils/noleak.go @@ -17,6 +17,7 @@ limitations under the License. package utils import ( + "context" "os" "os/exec" "strconv" @@ -29,6 +30,24 @@ import ( "vitess.io/vitess/go/vt/log" ) +func LeakCheckContext(t testing.TB) context.Context { + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(func() { + cancel() + EnsureNoLeaks(t) + }) + return ctx +} + +func LeakCheckContextTimeout(t testing.TB, timeout time.Duration) context.Context { + ctx, cancel := context.WithTimeout(context.Background(), timeout) + t.Cleanup(func() { + cancel() + EnsureNoLeaks(t) + }) + return ctx +} + // EnsureNoLeaks checks for goroutine and socket leaks and fails the test if any are found. func EnsureNoLeaks(t testing.TB) { if t.Failed() { diff --git a/go/vt/discovery/healthcheck_test.go b/go/vt/discovery/healthcheck_test.go index 19780081bd2..5fadc57eb2e 100644 --- a/go/vt/discovery/healthcheck_test.go +++ b/go/vt/discovery/healthcheck_test.go @@ -64,9 +64,7 @@ func init() { } func TestHealthCheck(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // reset error counters hcErrorCounters.ResetAll() ts := memorytopo.NewServer(ctx, "cell") @@ -210,9 +208,7 @@ func TestHealthCheck(t *testing.T) { } func TestHealthCheckStreamError(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) ts := memorytopo.NewServer(ctx, "cell") defer ts.Close() @@ -276,9 +272,7 @@ func TestHealthCheckStreamError(t *testing.T) { // TestHealthCheckErrorOnPrimary is the same as TestHealthCheckStreamError except for tablet type func TestHealthCheckErrorOnPrimary(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) ts := memorytopo.NewServer(ctx, "cell") defer ts.Close() @@ -341,9 +335,7 @@ func TestHealthCheckErrorOnPrimary(t *testing.T) { } func TestHealthCheckErrorOnPrimaryAfterExternalReparent(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) ts := memorytopo.NewServer(ctx, "cell") defer ts.Close() @@ -424,9 +416,7 @@ func TestHealthCheckErrorOnPrimaryAfterExternalReparent(t *testing.T) { } func TestHealthCheckVerifiesTabletAlias(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) ts := memorytopo.NewServer(ctx, "cell") defer ts.Close() @@ -472,9 +462,7 @@ func TestHealthCheckVerifiesTabletAlias(t *testing.T) { // TestHealthCheckCloseWaitsForGoRoutines tests that Close() waits for all Go // routines to finish and the listener won't be called anymore. func TestHealthCheckCloseWaitsForGoRoutines(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) ts := memorytopo.NewServer(ctx, "cell") defer ts.Close() @@ -538,9 +526,7 @@ func TestHealthCheckCloseWaitsForGoRoutines(t *testing.T) { } func TestHealthCheckTimeout(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // reset counters hcErrorCounters.ResetAll() @@ -615,9 +601,7 @@ func TestHealthCheckTimeout(t *testing.T) { } func TestWaitForAllServingTablets(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) ts := memorytopo.NewServer(ctx, "cell") defer ts.Close() @@ -641,6 +625,7 @@ func TestWaitForAllServingTablets(t *testing.T) { // there will be a first result, get and discard it <-resultChan // empty + var cancel context.CancelFunc ctx, cancel = context.WithTimeout(ctx, 1*time.Second) defer cancel() @@ -712,9 +697,7 @@ func TestWaitForAllServingTablets(t *testing.T) { // TestRemoveTablet tests the behavior when a tablet goes away. func TestRemoveTablet(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) ts := memorytopo.NewServer(ctx, "cell") defer ts.Close() @@ -824,9 +807,7 @@ func TestRemoveTablet(t *testing.T) { // TestGetHealthyTablets tests the functionality of GetHealthyTabletStats. func TestGetHealthyTablets(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) ts := memorytopo.NewServer(ctx, "cell") defer ts.Close() @@ -1015,9 +996,7 @@ func TestGetHealthyTablets(t *testing.T) { } func TestPrimaryInOtherCell(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) ts := memorytopo.NewServer(ctx, "cell1", "cell2") defer ts.Close() @@ -1077,9 +1056,7 @@ func TestPrimaryInOtherCell(t *testing.T) { } func TestReplicaInOtherCell(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) ts := memorytopo.NewServer(ctx, "cell1", "cell2") defer ts.Close() @@ -1184,9 +1161,7 @@ func TestReplicaInOtherCell(t *testing.T) { } func TestCellAliases(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) ts := memorytopo.NewServer(ctx, "cell1", "cell2") defer ts.Close() @@ -1249,9 +1224,7 @@ func TestCellAliases(t *testing.T) { } func TestHealthCheckChecksGrpcPort(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) ts := memorytopo.NewServer(ctx, "cell") defer ts.Close() diff --git a/go/vt/discovery/keyspace_events_test.go b/go/vt/discovery/keyspace_events_test.go index 09aedd35061..43af4bf49de 100644 --- a/go/vt/discovery/keyspace_events_test.go +++ b/go/vt/discovery/keyspace_events_test.go @@ -34,9 +34,7 @@ import ( ) func TestSrvKeyspaceWithNilNewKeyspace(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) cell := "cell" keyspace := "testks" factory := faketopo.NewFakeTopoFactory() diff --git a/go/vt/discovery/tablet_picker_test.go b/go/vt/discovery/tablet_picker_test.go index 5bfdb144024..ae94cb295d1 100644 --- a/go/vt/discovery/tablet_picker_test.go +++ b/go/vt/discovery/tablet_picker_test.go @@ -305,9 +305,7 @@ func TestPickLocalPreferences(t *testing.T) { } func TestPickCellPreferenceLocalCell(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) - defer cancel() + ctx := utils.LeakCheckContext(t) // test env puts all cells into an alias called "cella" te := newPickerTestEnv(t, ctx, []string{"cell", "otherCell"}) @@ -346,9 +344,7 @@ func TestPickCellPreferenceLocalCell(t *testing.T) { } func TestPickCellPreferenceLocalAlias(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) - defer cancel() + ctx := utils.LeakCheckContext(t) // test env puts all cells into an alias called "cella" te := newPickerTestEnv(t, ctx, []string{"cell", "otherCell"}) @@ -364,9 +360,7 @@ func TestPickCellPreferenceLocalAlias(t *testing.T) { } func TestPickUsingCellAliasOnlySpecified(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) - defer cancel() + ctx := utils.LeakCheckContextTimeout(t, 200*time.Millisecond) // test env puts all cells into an alias called "cella" te := newPickerTestEnv(t, ctx, []string{"cell", "otherCell"}) @@ -413,9 +407,7 @@ func TestPickUsingCellAliasOnlySpecified(t *testing.T) { } func TestTabletAppearsDuringSleep(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) - defer cancel() + ctx := utils.LeakCheckContextTimeout(t, 200*time.Millisecond) te := newPickerTestEnv(t, ctx, []string{"cell"}) tp, err := NewTabletPicker(ctx, te.topoServ, te.cells, "cell", te.keyspace, te.shard, "replica", TabletPickerOptions{}) @@ -445,9 +437,7 @@ func TestTabletAppearsDuringSleep(t *testing.T) { } func TestPickErrorLocalPreferenceDefault(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) te := newPickerTestEnv(t, ctx, []string{"cell"}) _, err := NewTabletPicker(ctx, te.topoServ, te.cells, "cell", te.keyspace, te.shard, "badtype", TabletPickerOptions{}) @@ -477,9 +467,7 @@ func TestPickErrorLocalPreferenceDefault(t *testing.T) { } func TestPickErrorOnlySpecified(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) te := newPickerTestEnv(t, ctx, []string{"cell"}) @@ -511,9 +499,7 @@ func TestPickErrorOnlySpecified(t *testing.T) { // type has no available healthy serving tablets that we select a healthy // serving tablet from the secondary/second type. func TestPickFallbackType(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) cells := []string{"cell1", "cell2"} localCell := cells[0] @@ -532,6 +518,7 @@ func TestPickFallbackType(t *testing.T) { replicaTablet := addTablet(ctx, te, 200, topodatapb.TabletType_REPLICA, localCell, false, false) defer deleteTablet(t, te, replicaTablet) + var cancel context.CancelFunc ctx, cancel = context.WithTimeout(ctx, 200*time.Millisecond) defer cancel() _, err := te.topoServ.UpdateShardFields(ctx, te.keyspace, te.shard, func(si *topo.ShardInfo) error { diff --git a/go/vt/discovery/topology_watcher_test.go b/go/vt/discovery/topology_watcher_test.go index 6ca8b907aac..f0dfe2da833 100644 --- a/go/vt/discovery/topology_watcher_test.go +++ b/go/vt/discovery/topology_watcher_test.go @@ -58,9 +58,7 @@ func checkChecksum(t *testing.T, tw *TopologyWatcher, want uint32) { } func TestStartAndCloseTopoWatcher(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) ts := memorytopo.NewServer(ctx, "aa") defer ts.Close() @@ -118,9 +116,7 @@ func TestCellTabletsWatcherNoRefreshKnown(t *testing.T) { } func checkWatcher(t *testing.T, refreshKnownTablets bool) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) ts := memorytopo.NewServer(ctx, "aa") defer ts.Close() @@ -443,9 +439,7 @@ var ( ) func TestFilterByKeyspace(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) hc := NewFakeHealthCheck(nil) f := NewFilterByKeyspace(testKeyspacesToWatch) @@ -528,9 +522,7 @@ func TestFilterByKeyspace(t *testing.T) { // - does not continuosly call GetTablets for tablets that do not satisfy the filter // - does not add or remove these filtered out tablets from the its healtcheck func TestFilterByKeypsaceSkipsIgnoredTablets(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) ts := memorytopo.NewServer(ctx, "aa") defer ts.Close() diff --git a/go/vt/mysqlctl/backup_blackbox_test.go b/go/vt/mysqlctl/backup_blackbox_test.go index cd6fcdac535..62b58f2a5c8 100644 --- a/go/vt/mysqlctl/backup_blackbox_test.go +++ b/go/vt/mysqlctl/backup_blackbox_test.go @@ -80,9 +80,8 @@ func createBackupFiles(root string, fileCount int, ext string) error { } func TestExecuteBackup(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) + // Set up local backup directory backupRoot := "testdata/builtinbackup_test" filebackupstorage.FileBackupStorageRoot = backupRoot @@ -226,9 +225,7 @@ func TestExecuteBackup(t *testing.T) { } func TestExecuteBackupWithSafeUpgrade(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Set up local backup directory backupRoot := "testdata/builtinbackup_test" @@ -318,9 +315,7 @@ func TestExecuteBackupWithSafeUpgrade(t *testing.T) { // TestExecuteBackupWithCanceledContext tests the ability of the backup function to gracefully handle cases where errors // occur due to various reasons, such as context time cancel. The process should not panic in these situations. func TestExecuteBackupWithCanceledContext(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Set up local backup directory id := fmt.Sprintf("%d", time.Now().UnixNano()) @@ -408,9 +403,7 @@ func TestExecuteBackupWithCanceledContext(t *testing.T) { // TestExecuteRestoreWithCanceledContext tests the ability of the restore function to gracefully handle cases where errors // occur due to various reasons, such as context timed-out. The process should not panic in these situations. func TestExecuteRestoreWithTimedOutContext(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Set up local backup directory id := fmt.Sprintf("%d", time.Now().UnixNano()) diff --git a/go/vt/mysqlctl/backup_test.go b/go/vt/mysqlctl/backup_test.go index c16dec64e50..5a135c26a30 100644 --- a/go/vt/mysqlctl/backup_test.go +++ b/go/vt/mysqlctl/backup_test.go @@ -45,9 +45,7 @@ import ( // TestBackupExecutesBackupWithScopedParams tests that Backup passes // a Scope()-ed stats to backupengine ExecuteBackup. func TestBackupExecutesBackupWithScopedParams(t *testing.T) { - defer utils.EnsureNoLeaks(t) - env, closer := createFakeBackupRestoreEnv(t) - defer closer() + env := createFakeBackupRestoreEnv(t) require.Nil(t, Backup(env.ctx, env.backupParams), env.logger.Events) @@ -68,10 +66,7 @@ func TestBackupExecutesBackupWithScopedParams(t *testing.T) { // TestBackupNoStats tests that if BackupParams.Stats is nil, then Backup will // pass non-nil Stats to sub-components. func TestBackupNoStats(t *testing.T) { - defer utils.EnsureNoLeaks(t) - env, closer := createFakeBackupRestoreEnv(t) - defer closer() - + env := createFakeBackupRestoreEnv(t) env.setStats(nil) require.Nil(t, Backup(env.ctx, env.backupParams), env.logger.Events) @@ -84,9 +79,7 @@ func TestBackupNoStats(t *testing.T) { // TestBackupParameterizesBackupStorageWithScopedStats tests that Backup passes // a Scope()-ed stats to BackupStorage.WithParams. func TestBackupParameterizesBackupStorageWithScopedStats(t *testing.T) { - defer utils.EnsureNoLeaks(t) - env, closer := createFakeBackupRestoreEnv(t) - defer closer() + env := createFakeBackupRestoreEnv(t) require.Nil(t, Backup(env.ctx, env.backupParams), env.logger.Events) @@ -105,9 +98,7 @@ func TestBackupParameterizesBackupStorageWithScopedStats(t *testing.T) { // TestBackupEmitsStats tests that Backup emits stats. func TestBackupEmitsStats(t *testing.T) { - defer utils.EnsureNoLeaks(t) - env, closer := createFakeBackupRestoreEnv(t) - defer closer() + env := createFakeBackupRestoreEnv(t) // Force ExecuteBackup to take time so we can test stats emission. env.backupEngine.ExecuteBackupDuration = 1001 * time.Millisecond @@ -123,9 +114,7 @@ func TestBackupEmitsStats(t *testing.T) { // backupstorage.Params to backupstorage, but only if it responds to // backupstorage.WithParams. func TestBackupTriesToParameterizeBackupStorage(t *testing.T) { - defer utils.EnsureNoLeaks(t) - env, closer := createFakeBackupRestoreEnv(t) - defer closer() + env := createFakeBackupRestoreEnv(t) require.Nil(t, Backup(env.ctx, env.backupParams), env.logger.Events) @@ -329,9 +318,7 @@ func TestFindFilesToBackupWithRedoLog(t *testing.T) { // TestRestoreEmitsStats tests that Restore emits stats. func TestRestoreEmitsStats(t *testing.T) { - defer utils.EnsureNoLeaks(t) - env, closer := createFakeBackupRestoreEnv(t) - defer closer() + env := createFakeBackupRestoreEnv(t) // Force ExecuteRestore to take time so we can test stats emission. env.backupEngine.ExecuteRestoreDuration = 1001 * time.Millisecond @@ -347,9 +334,7 @@ func TestRestoreEmitsStats(t *testing.T) { // TestRestoreExecutesRestoreWithScopedParams tests that Restore passes // a Scope()-ed stats to backupengine ExecuteRestore. func TestRestoreExecutesRestoreWithScopedParams(t *testing.T) { - defer utils.EnsureNoLeaks(t) - env, closer := createFakeBackupRestoreEnv(t) - defer closer() + env := createFakeBackupRestoreEnv(t) _, err := Restore(env.ctx, env.restoreParams) require.Nil(t, err, env.logger.Events) @@ -371,10 +356,7 @@ func TestRestoreExecutesRestoreWithScopedParams(t *testing.T) { // TestRestoreNoStats tests that if RestoreParams.Stats is nil, then Restore will // pass non-nil Stats to sub-components. func TestRestoreNoStats(t *testing.T) { - defer utils.EnsureNoLeaks(t) - env, closer := createFakeBackupRestoreEnv(t) - defer closer() - + env := createFakeBackupRestoreEnv(t) env.setStats(nil) _, err := Restore(env.ctx, env.restoreParams) @@ -388,9 +370,7 @@ func TestRestoreNoStats(t *testing.T) { // TestRestoreParameterizesBackupStorageWithScopedStats tests that Restore passes // a Scope()-ed stats to BackupStorage.WithParams. func TestRestoreParameterizesBackupStorageWithScopedStats(t *testing.T) { - defer utils.EnsureNoLeaks(t) - env, closer := createFakeBackupRestoreEnv(t) - defer closer() + env := createFakeBackupRestoreEnv(t) _, err := Restore(env.ctx, env.restoreParams) require.Nil(t, err, env.logger.Events) @@ -412,9 +392,7 @@ func TestRestoreParameterizesBackupStorageWithScopedStats(t *testing.T) { // backupstorage.Params to backupstorage, but only if it responds to // backupstorage.WithParams. func TestRestoreTriesToParameterizeBackupStorage(t *testing.T) { - defer utils.EnsureNoLeaks(t) - env, closer := createFakeBackupRestoreEnv(t) - defer closer() + env := createFakeBackupRestoreEnv(t) _, err := Restore(env.ctx, env.restoreParams) require.Nil(t, err, env.logger.Events) @@ -524,9 +502,7 @@ func TestRestoreManifestMySQLVersionValidation(t *testing.T) { for _, tc := range testCases { t.Run(fmt.Sprintf("%s->%s upgradeSafe=%t", tc.fromVersion, tc.toVersion, tc.upgradeSafe), func(t *testing.T) { - defer utils.EnsureNoLeaks(t) - env, closer := createFakeBackupRestoreEnv(t) - defer closer() + env := createFakeBackupRestoreEnv(t) env.mysqld.Version = tc.toVersion manifest := BackupManifest{ @@ -580,7 +556,7 @@ type fakeBackupRestoreEnv struct { stats *backupstats.FakeStats } -func createFakeBackupRestoreEnv(t *testing.T) (*fakeBackupRestoreEnv, func()) { +func createFakeBackupRestoreEnv(t *testing.T) *fakeBackupRestoreEnv { ctx := context.Background() logger := logutil.NewMemoryLogger() @@ -663,7 +639,12 @@ func createFakeBackupRestoreEnv(t *testing.T) (*fakeBackupRestoreEnv, func()) { previousBackupStorageImplementation := backupstorage.BackupStorageImplementation backupstorage.BackupStorageImplementation = "fake" - closer := func() { + // all restore integration tests must be leak checked + t.Cleanup(func() { + utils.EnsureNoLeaks(t) + }) + + t.Cleanup(func() { backupstats.DeprecatedBackupDurationS.Reset() backupstats.DeprecatedRestoreDurationS.Reset() @@ -674,7 +655,7 @@ func createFakeBackupRestoreEnv(t *testing.T) (*fakeBackupRestoreEnv, func()) { backupstorage.BackupStorageImplementation = previousBackupStorageImplementation mysqld.Close() sqldb.Close() - } + }) return &fakeBackupRestoreEnv{ backupEngine: &testBackupEngine, @@ -685,7 +666,7 @@ func createFakeBackupRestoreEnv(t *testing.T) (*fakeBackupRestoreEnv, func()) { mysqld: mysqld, restoreParams: restoreParams, stats: stats, - }, closer + } } func (fbe *fakeBackupRestoreEnv) setStats(stats *backupstats.FakeStats) { diff --git a/go/vt/topo/helpers/tee_topo_test.go b/go/vt/topo/helpers/tee_topo_test.go index b1f1d379a5a..8a4c5690846 100644 --- a/go/vt/topo/helpers/tee_topo_test.go +++ b/go/vt/topo/helpers/tee_topo_test.go @@ -17,7 +17,6 @@ limitations under the License. package helpers import ( - "context" "testing" "vitess.io/vitess/go/test/utils" @@ -27,9 +26,7 @@ import ( ) func TestTeeTopo(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) test.TopoServerTestSuite(t, ctx, func() *topo.Server { s1 := memorytopo.NewServer(ctx, test.LocalCellName) s2 := memorytopo.NewServer(ctx, test.LocalCellName) diff --git a/go/vt/vtadmin/cluster/cluster_test.go b/go/vt/vtadmin/cluster/cluster_test.go index c1a7f4a47b1..c3b86a7526d 100644 --- a/go/vt/vtadmin/cluster/cluster_test.go +++ b/go/vt/vtadmin/cluster/cluster_test.go @@ -177,7 +177,7 @@ func TestCreateKeyspace(t *testing.T) { } func TestCreateShard(t *testing.T) { - defer utils.EnsureNoLeaks(t) + ctx := utils.LeakCheckContext(t) type test struct { name string @@ -186,8 +186,7 @@ func TestCreateShard(t *testing.T) { shouldErr bool assertion func(t *testing.T, tt *test) } - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + tests := []*test{ { name: "ok", @@ -279,10 +278,8 @@ func TestCreateShard(t *testing.T) { } func TestDeleteKeyspace(t *testing.T) { - defer utils.EnsureNoLeaks(t) + ctx := utils.LeakCheckContext(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() tests := []struct { name string cfg testutil.TestClusterConfig @@ -364,8 +361,8 @@ func TestDeleteKeyspace(t *testing.T) { } func TestDeleteShards(t *testing.T) { - defer utils.EnsureNoLeaks(t) t.Parallel() + ctx := utils.LeakCheckContext(t) type test struct { name string @@ -375,8 +372,7 @@ func TestDeleteShards(t *testing.T) { shouldErr bool assertion func(t *testing.T, tt *test) } - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + tests := []*test{ { name: "ok", diff --git a/go/vt/vtctl/grpcvtctldserver/endtoend/init_shard_primary_test.go b/go/vt/vtctl/grpcvtctldserver/endtoend/init_shard_primary_test.go index 8de5eceaeff..f5f7847b499 100644 --- a/go/vt/vtctl/grpcvtctldserver/endtoend/init_shard_primary_test.go +++ b/go/vt/vtctl/grpcvtctldserver/endtoend/init_shard_primary_test.go @@ -42,9 +42,7 @@ import ( ) func TestInitShardPrimary(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) ts := memorytopo.NewServer(ctx, "cell1") tmc := tmclient.NewTabletManagerClient() defer tmc.Close() diff --git a/go/vt/vtexplain/vtexplain_test.go b/go/vt/vtexplain/vtexplain_test.go index e764ce6f3e8..54f1efbc522 100644 --- a/go/vt/vtexplain/vtexplain_test.go +++ b/go/vt/vtexplain/vtexplain_test.go @@ -86,9 +86,7 @@ func testExplain(testcase string, opts *Options, t *testing.T) { func runTestCase(testcase, mode string, opts *Options, topts *testopts, t *testing.T) { t.Run(testcase, func(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) vte := initTest(ctx, mode, opts, topts, t) defer vte.Stop() @@ -155,10 +153,7 @@ func TestExplain(t *testing.T) { } func TestErrors(t *testing.T) { - defer utils.EnsureNoLeaks(t) - - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) vte := initTest(ctx, ModeMulti, defaultTestOpts(), &testopts{}, t) defer vte.Stop() @@ -198,10 +193,7 @@ func TestErrors(t *testing.T) { } func TestJSONOutput(t *testing.T) { - defer utils.EnsureNoLeaks(t) - - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) vte := initTest(ctx, ModeMulti, defaultTestOpts(), &testopts{}, t) defer vte.Stop() diff --git a/go/vt/vtgate/autocommit_test.go b/go/vt/vtgate/autocommit_test.go index 5e0c267259d..0d55bbf2875 100644 --- a/go/vt/vtgate/autocommit_test.go +++ b/go/vt/vtgate/autocommit_test.go @@ -35,8 +35,7 @@ import ( // TestAutocommitUpdateSharded: instant-commit. func TestAutocommitUpdateSharded(t *testing.T) { - executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, _ := createExecutorEnv(t) _, err := autocommitExec(executor, "update user set a=2 where id = 1") require.NoError(t, err) @@ -53,8 +52,7 @@ func TestAutocommitUpdateSharded(t *testing.T) { // TestAutocommitUpdateLookup: transaction: select before update. func TestAutocommitUpdateLookup(t *testing.T) { - executor, sbc1, _, sbclookup, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, sbclookup, _ := createExecutorEnv(t) sbclookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("b|a", "int64|varbinary"), "2|1", @@ -83,8 +81,7 @@ func TestAutocommitUpdateLookup(t *testing.T) { // TestAutocommitUpdateVindexChange: transaction: select & update before final update. func TestAutocommitUpdateVindexChange(t *testing.T) { - executor, sbc, _, sbclookup, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc, _, sbclookup, _ := createExecutorEnv(t) sbc.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("id|name|lastname|name_lastname_keyspace_id_map", "int64|int32|varchar|int64"), "1|1|foo|0", @@ -123,8 +120,7 @@ func TestAutocommitUpdateVindexChange(t *testing.T) { // TestAutocommitDeleteSharded: instant-commit. func TestAutocommitDeleteSharded(t *testing.T) { - executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, _ := createExecutorEnv(t) _, err := autocommitExec(executor, "delete from user_extra where user_id = 1") require.NoError(t, err) @@ -141,8 +137,7 @@ func TestAutocommitDeleteSharded(t *testing.T) { // TestAutocommitDeleteLookup: transaction: select before update. func TestAutocommitDeleteLookup(t *testing.T) { - executor, sbc1, _, sbclookup, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, sbclookup, _ := createExecutorEnv(t) sbc1.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("id|name|lastname", "int64|int32|varchar"), "1|1|foo", @@ -184,8 +179,7 @@ func TestAutocommitDeleteLookup(t *testing.T) { // TestAutocommitDeleteIn: instant-commit. func TestAutocommitDeleteIn(t *testing.T) { - executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, _ := createExecutorEnv(t) _, err := autocommitExec(executor, "delete from user_extra where user_id in (1, 2)") require.NoError(t, err) @@ -202,8 +196,7 @@ func TestAutocommitDeleteIn(t *testing.T) { // TestAutocommitDeleteMultiShard: instant-commit. func TestAutocommitDeleteMultiShard(t *testing.T) { - executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, _ := createExecutorEnv(t) _, err := autocommitExec(executor, "delete from user_extra where user_id = user_id + 1") require.NoError(t, err) @@ -223,8 +216,7 @@ func TestAutocommitDeleteMultiShard(t *testing.T) { // TestAutocommitDeleteMultiShardAutoCommit: instant-commit. func TestAutocommitDeleteMultiShardAutoCommit(t *testing.T) { - executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, _ := createExecutorEnv(t) _, err := autocommitExec(executor, "delete /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ from user_extra where user_id = user_id + 1") require.NoError(t, err) @@ -244,8 +236,7 @@ func TestAutocommitDeleteMultiShardAutoCommit(t *testing.T) { // TestAutocommitInsertSharded: instant-commit. func TestAutocommitInsertSharded(t *testing.T) { - executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, _ := createExecutorEnv(t) _, err := autocommitExec(executor, "insert into user_extra(user_id, v) values (1, 2)") require.NoError(t, err) @@ -264,8 +255,7 @@ func TestAutocommitInsertSharded(t *testing.T) { // TestAutocommitInsertLookup: transaction: select before update. func TestAutocommitInsertLookup(t *testing.T) { - executor, sbc1, _, sbclookup, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, sbclookup, _ := createExecutorEnv(t) _, err := autocommitExec(executor, "insert into user(id, v, name) values (1, 2, 'myname')") require.NoError(t, err) @@ -292,54 +282,55 @@ func TestAutocommitInsertLookup(t *testing.T) { // TestAutocommitInsertShardAutoCommit: instant-commit. func TestAutocommitInsertMultishardAutoCommit(t *testing.T) { - executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) - - _, err := autocommitExec(executor, "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ into user_extra(user_id, v) values (1, 2), (3, 4)") - require.NoError(t, err) - - assertQueries(t, sbc1, []*querypb.BoundQuery{{ - Sql: "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ into user_extra(user_id, v) values (:_user_id_0, 2)", - BindVariables: map[string]*querypb.BindVariable{ - "_user_id_0": sqltypes.Int64BindVariable(1), - "_user_id_1": sqltypes.Int64BindVariable(3), - }, - }}) - testCommitCount(t, "sbc1", sbc1, 0) - - assertQueries(t, sbc2, []*querypb.BoundQuery{{ - Sql: "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ into user_extra(user_id, v) values (:_user_id_1, 4)", - BindVariables: map[string]*querypb.BindVariable{ - "_user_id_0": sqltypes.Int64BindVariable(1), - "_user_id_1": sqltypes.Int64BindVariable(3), - }, - }}) - testCommitCount(t, "sbc2", sbc2, 0) - - closer() - executor, sbc1, sbc2, _, _, closer = createExecutorEnv(t) - defer closer() - // Make the first shard fail - the second completes anyway - sbc1.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 1 - _, err = autocommitExec(executor, "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ into user_extra(user_id, v) values (1, 2), (3, 4)") - require.Error(t, err) - require.Contains(t, err.Error(), "INVALID_ARGUMENT", "expected invalid argument error") - - testCommitCount(t, "sbc1", sbc1, 0) - - assertQueries(t, sbc2, []*querypb.BoundQuery{{ - Sql: "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ into user_extra(user_id, v) values (:_user_id_1, 4)", - BindVariables: map[string]*querypb.BindVariable{ - "_user_id_0": sqltypes.Int64BindVariable(1), - "_user_id_1": sqltypes.Int64BindVariable(3), - }, - }}) - testCommitCount(t, "sbc2", sbc2, 0) + t.Run("1", func(t *testing.T) { + executor, sbc1, sbc2, _, _ := createExecutorEnv(t) + + _, err := autocommitExec(executor, "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ into user_extra(user_id, v) values (1, 2), (3, 4)") + require.NoError(t, err) + + assertQueries(t, sbc1, []*querypb.BoundQuery{{ + Sql: "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ into user_extra(user_id, v) values (:_user_id_0, 2)", + BindVariables: map[string]*querypb.BindVariable{ + "_user_id_0": sqltypes.Int64BindVariable(1), + "_user_id_1": sqltypes.Int64BindVariable(3), + }, + }}) + testCommitCount(t, "sbc1", sbc1, 0) + + assertQueries(t, sbc2, []*querypb.BoundQuery{{ + Sql: "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ into user_extra(user_id, v) values (:_user_id_1, 4)", + BindVariables: map[string]*querypb.BindVariable{ + "_user_id_0": sqltypes.Int64BindVariable(1), + "_user_id_1": sqltypes.Int64BindVariable(3), + }, + }}) + testCommitCount(t, "sbc2", sbc2, 0) + }) + t.Run("2", func(t *testing.T) { + executor, sbc1, sbc2, _, _ := createExecutorEnv(t) + + // Make the first shard fail - the second completes anyway + sbc1.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 1 + _, err := autocommitExec(executor, "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ into user_extra(user_id, v) values (1, 2), (3, 4)") + require.Error(t, err) + require.Contains(t, err.Error(), "INVALID_ARGUMENT", "expected invalid argument error") + + testCommitCount(t, "sbc1", sbc1, 0) + + assertQueries(t, sbc2, []*querypb.BoundQuery{{ + Sql: "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ into user_extra(user_id, v) values (:_user_id_1, 4)", + BindVariables: map[string]*querypb.BindVariable{ + "_user_id_0": sqltypes.Int64BindVariable(1), + "_user_id_1": sqltypes.Int64BindVariable(3), + }, + }}) + testCommitCount(t, "sbc2", sbc2, 0) + }) } func TestAutocommitInsertMultishard(t *testing.T) { - executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, _ := createExecutorEnv(t) _, err := autocommitExec(executor, "insert into user_extra(user_id, v) values (1, 2), (3, 4)") require.NoError(t, err) @@ -365,8 +356,7 @@ func TestAutocommitInsertMultishard(t *testing.T) { // TestAutocommitInsertAutoinc: instant-commit: sequence fetch is not transactional. func TestAutocommitInsertAutoinc(t *testing.T) { - executor, _, _, sbclookup, _, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbclookup, _ := createExecutorEnv(t) _, err := autocommitExec(executor, "insert into main1(id, name) values (null, 'myname')") require.NoError(t, err) @@ -385,8 +375,7 @@ func TestAutocommitInsertAutoinc(t *testing.T) { // TestAutocommitTransactionStarted: no instant-commit. func TestAutocommitTransactionStarted(t *testing.T) { - executor, sbc1, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, _ := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", @@ -418,8 +407,7 @@ func TestAutocommitTransactionStarted(t *testing.T) { // TestAutocommitDirectTarget: instant-commit. func TestAutocommitDirectTarget(t *testing.T) { - executor, _, _, sbclookup, _, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbclookup, _ := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "TestUnsharded/0@primary", @@ -440,8 +428,7 @@ func TestAutocommitDirectTarget(t *testing.T) { // TestAutocommitDirectRangeTarget: no instant-commit. func TestAutocommitDirectRangeTarget(t *testing.T) { - executor, sbc1, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, _ := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "TestExecutor[-]@primary", diff --git a/go/vt/vtgate/bench_test.go b/go/vt/vtgate/bench_test.go index d021fb9f7b5..5c64c7e3473 100644 --- a/go/vt/vtgate/bench_test.go +++ b/go/vt/vtgate/bench_test.go @@ -58,8 +58,7 @@ func init() { } func BenchmarkWithNormalizer(b *testing.B) { - vtgateInst, _, ctx, closer := createVtgateEnv(b) - defer closer() + vtgateInst, _, ctx := createVtgateEnv(b) for i := 0; i < b.N; i++ { _, _, err := vtgateInst.Execute( @@ -79,8 +78,7 @@ func BenchmarkWithNormalizer(b *testing.B) { } func BenchmarkWithoutNormalizer(b *testing.B) { - vtgateInst, _, ctx, closer := createVtgateEnv(b) - defer closer() + vtgateInst, _, ctx := createVtgateEnv(b) vtgateInst.executor.normalize = false diff --git a/go/vt/vtgate/executor_ddl_test.go b/go/vt/vtgate/executor_ddl_test.go index f5d4d58d068..b2502ab247a 100644 --- a/go/vt/vtgate/executor_ddl_test.go +++ b/go/vt/vtgate/executor_ddl_test.go @@ -26,8 +26,7 @@ import ( ) func TestDDLFlags(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) session := NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded}) defer func() { enableOnlineDDL = true diff --git a/go/vt/vtgate/executor_dml_test.go b/go/vt/vtgate/executor_dml_test.go index 6975413333b..59fbe314346 100644 --- a/go/vt/vtgate/executor_dml_test.go +++ b/go/vt/vtgate/executor_dml_test.go @@ -38,8 +38,7 @@ import ( ) func TestUpdateEqual(t *testing.T) { - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -134,8 +133,7 @@ func TestUpdateEqual(t *testing.T) { } func TestUpdateFromSubQuery(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) executor.pv = querypb.ExecuteOptions_Gen4 logChan := executor.queryLogger.Subscribe("Test") @@ -180,8 +178,7 @@ func TestUpdateEqualWithNoVerifyAndWriteOnlyLookupUniqueVindexes(t *testing.T) { ), "1|2|2|2|2|2|1|0", )} - executor, sbc1, sbc2, sbcLookup, ctx, closer := createCustomExecutorSetValues(t, executorVSchema, res) - defer closer() + executor, sbc1, sbc2, sbcLookup, ctx := createCustomExecutorSetValues(t, executorVSchema, res) session := &vtgatepb.Session{ TargetString: "@primary", @@ -235,8 +232,7 @@ func TestUpdateInTransactionLookupDefaultReadLock(t *testing.T) { ), "1|2|2|2|2|2|1|0", )} - executor, sbc1, sbc2, sbcLookup, ctx, closer := createCustomExecutorSetValues(t, executorVSchema, res) - defer closer() + executor, sbc1, sbc2, sbcLookup, ctx := createCustomExecutorSetValues(t, executorVSchema, res) safeSession := NewSafeSession(&vtgatepb.Session{InTransaction: true}) _, err := executorExecSession(ctx, @@ -298,8 +294,7 @@ func TestUpdateInTransactionLookupExclusiveReadLock(t *testing.T) { ), "1|2|2|2|2|2|1|0", )} - executor, sbc1, sbc2, sbcLookup, ctx, closer := createCustomExecutorSetValues(t, executorVSchema, res) - defer closer() + executor, sbc1, sbc2, sbcLookup, ctx := createCustomExecutorSetValues(t, executorVSchema, res) safeSession := NewSafeSession(&vtgatepb.Session{InTransaction: true}) _, err := executorExecSession(ctx, @@ -361,8 +356,7 @@ func TestUpdateInTransactionLookupSharedReadLock(t *testing.T) { ), "1|2|2|2|2|2|1|0", )} - executor, sbc1, sbc2, sbcLookup, ctx, closer := createCustomExecutorSetValues(t, executorVSchema, res) - defer closer() + executor, sbc1, sbc2, sbcLookup, ctx := createCustomExecutorSetValues(t, executorVSchema, res) safeSession := NewSafeSession(&vtgatepb.Session{InTransaction: true}) _, err := executorExecSession(ctx, @@ -424,8 +418,7 @@ func TestUpdateInTransactionLookupNoReadLock(t *testing.T) { ), "1|2|2|2|2|2|1|0", )} - executor, sbc1, sbc2, sbcLookup, ctx, closer := createCustomExecutorSetValues(t, executorVSchema, res) - defer closer() + executor, sbc1, sbc2, sbcLookup, ctx := createCustomExecutorSetValues(t, executorVSchema, res) safeSession := NewSafeSession(&vtgatepb.Session{InTransaction: true}) _, err := executorExecSession(ctx, @@ -539,8 +532,7 @@ func TestUpdateMultiOwned(t *testing.T) { } } ` - executor, sbc1, sbc2, sbclookup, ctx, closer := createCustomExecutor(t, vschema) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createCustomExecutor(t, vschema) sbc1.SetResults([]*sqltypes.Result{ sqltypes.MakeTestResult( @@ -599,8 +591,7 @@ func TestUpdateMultiOwned(t *testing.T) { } func TestUpdateComments(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", @@ -616,8 +607,7 @@ func TestUpdateComments(t *testing.T) { } func TestUpdateNormalize(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) executor.normalize = true session := &vtgatepb.Session{ @@ -653,8 +643,7 @@ func TestUpdateNormalize(t *testing.T) { } func TestDeleteEqual(t *testing.T) { - executor, sbc, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc, _, sbclookup, ctx := createExecutorEnv(t) sbc.SetResults([]*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -770,8 +759,7 @@ func TestDeleteEqual(t *testing.T) { } func TestUpdateScatter(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", @@ -788,8 +776,7 @@ func TestUpdateScatter(t *testing.T) { } func TestDeleteScatter(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", @@ -806,8 +793,7 @@ func TestDeleteScatter(t *testing.T) { } func TestUpdateEqualWithMultipleLookupVindex(t *testing.T) { - executor, sbc1, sbc2, sbcLookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbcLookup, ctx := createExecutorEnv(t) sbcLookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("lu_col|keyspace_id", "int64|varbinary"), @@ -863,8 +849,7 @@ func TestUpdateEqualWithMultipleLookupVindex(t *testing.T) { } func TestUpdateUseHigherCostVindexIfBackfilling(t *testing.T) { - executor, sbc1, sbc2, sbcLookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbcLookup, ctx := createExecutorEnv(t) sbcLookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("lu_col|keyspace_id", "int64|varbinary"), @@ -942,8 +927,7 @@ func TestDeleteEqualWithNoVerifyAndWriteOnlyLookupUniqueVindex(t *testing.T) { ), "1|1|1|1|1|1|1", )} - executor, sbc1, sbc2, sbcLookup, ctx, closer := createCustomExecutorSetValues(t, executorVSchema, res) - defer closer() + executor, sbc1, sbc2, sbcLookup, ctx := createCustomExecutorSetValues(t, executorVSchema, res) session := &vtgatepb.Session{ TargetString: "@primary", @@ -1017,8 +1001,7 @@ func TestDeleteEqualWithNoVerifyAndWriteOnlyLookupUniqueVindex(t *testing.T) { } func TestDeleteEqualWithMultipleLookupVindex(t *testing.T) { - executor, sbc1, sbc2, sbcLookup, ctx, closer := createCustomExecutorSetValues(t, executorVSchema, nil) - defer closer() + executor, sbc1, sbc2, sbcLookup, ctx := createCustomExecutorSetValues(t, executorVSchema, nil) sbcLookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("lu_col|keyspace_id", "int64|varbinary"), @@ -1099,8 +1082,7 @@ func TestDeleteEqualWithMultipleLookupVindex(t *testing.T) { } func TestDeleteUseHigherCostVindexIfBackfilling(t *testing.T) { - executor, sbc1, sbc2, sbcLookup, ctx, closer := createCustomExecutorSetValues(t, executorVSchema, nil) - defer closer() + executor, sbc1, sbc2, sbcLookup, ctx := createCustomExecutorSetValues(t, executorVSchema, nil) sbcLookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("lu_col|keyspace_id", "int64|varbinary"), @@ -1220,8 +1202,7 @@ func TestDeleteUseHigherCostVindexIfBackfilling(t *testing.T) { } func TestDeleteByDestination(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", @@ -1238,8 +1219,7 @@ func TestDeleteByDestination(t *testing.T) { } func TestDeleteComments(t *testing.T) { - executor, sbc, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc, _, sbclookup, ctx := createExecutorEnv(t) sbc.SetResults([]*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -1278,8 +1258,7 @@ func TestDeleteComments(t *testing.T) { } func TestInsertSharded(t *testing.T) { - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -1389,8 +1368,7 @@ func TestInsertSharded(t *testing.T) { } func TestInsertShardedKeyrange(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) // If a unique vindex returns a keyrange, we fail the insert session := &vtgatepb.Session{ @@ -1460,8 +1438,7 @@ func TestInsertShardedAutocommitLookup(t *testing.T) { } } ` - executor, sbc1, sbc2, sbclookup, ctx, closer := createCustomExecutor(t, vschema) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createCustomExecutor(t, vschema) _, err := executorExecSession(ctx, executor, "insert into user(id, v, name, music) values (1, 2, 'myname', 'star')", nil, &vtgatepb.Session{}) require.NoError(t, err) @@ -1494,8 +1471,7 @@ func TestInsertShardedAutocommitLookup(t *testing.T) { } func TestInsertShardedIgnore(t *testing.T) { - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) // Build the sequence of responses for sbclookup. This should // match the sequence of queries we validate below. @@ -1650,8 +1626,7 @@ func TestInsertShardedIgnore(t *testing.T) { func TestInsertOnDupKey(t *testing.T) { // This test just sanity checks that the statement is getting passed through // correctly. The full set of use cases are covered by TestInsertShardedIgnore. - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) sbclookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("b|a", "int64|varbinary"), "1|1", @@ -1696,8 +1671,7 @@ func TestInsertOnDupKey(t *testing.T) { } func TestAutocommitFail(t *testing.T) { - executor, sbc1, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, ctx := createExecutorEnv(t) query := "insert into user (id) values (1)" sbc1.MustFailCodes[vtrpcpb.Code_ALREADY_EXISTS] = 1 @@ -1714,8 +1688,7 @@ func TestAutocommitFail(t *testing.T) { } func TestInsertComments(t *testing.T) { - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", @@ -1743,8 +1716,7 @@ func TestInsertComments(t *testing.T) { } func TestInsertGeneratorSharded(t *testing.T) { - executor, sbc, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc, _, sbclookup, ctx := createExecutorEnv(t) sbclookup.SetResults([]*sqltypes.Result{{ Rows: [][]sqltypes.Value{{ @@ -1786,8 +1758,7 @@ func TestInsertGeneratorSharded(t *testing.T) { } func TestInsertAutoincSharded(t *testing.T) { - router, sbc, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + router, sbc, _, _, ctx := createExecutorEnv(t) // Fake a mysql auto-inc response. wantResult := &sqltypes.Result{ @@ -1817,8 +1788,7 @@ func TestInsertAutoincSharded(t *testing.T) { } func TestInsertGeneratorUnsharded(t *testing.T) { - executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbclookup, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", } @@ -1842,8 +1812,7 @@ func TestInsertGeneratorUnsharded(t *testing.T) { } func TestInsertAutoincUnsharded(t *testing.T) { - router, _, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + router, _, _, sbclookup, ctx := createExecutorEnv(t) logChan := router.queryLogger.Subscribe("Test") defer router.queryLogger.Unsubscribe(logChan) @@ -1875,8 +1844,7 @@ func TestInsertAutoincUnsharded(t *testing.T) { } func TestInsertLookupOwned(t *testing.T) { - executor, sbc, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc, _, sbclookup, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", @@ -1903,8 +1871,7 @@ func TestInsertLookupOwned(t *testing.T) { } func TestInsertLookupOwnedGenerator(t *testing.T) { - executor, sbc, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc, _, sbclookup, ctx := createExecutorEnv(t) sbclookup.SetResults([]*sqltypes.Result{{ Rows: [][]sqltypes.Value{{ @@ -1946,8 +1913,7 @@ func TestInsertLookupOwnedGenerator(t *testing.T) { } func TestInsertLookupUnowned(t *testing.T) { - executor, sbc, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc, _, sbclookup, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", @@ -1973,8 +1939,7 @@ func TestInsertLookupUnowned(t *testing.T) { } func TestInsertLookupUnownedUnsupplied(t *testing.T) { - executor, sbc, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc, _, sbclookup, ctx := createExecutorEnv(t) sbclookup.SetResults([]*sqltypes.Result{sqltypes.MakeTestResult( sqltypes.MakeTestFields("b|a", "int64|varbinary"), "3|1", @@ -2006,8 +1971,7 @@ func TestInsertLookupUnownedUnsupplied(t *testing.T) { // If a statement gets broken up into two, and the first one fails, // then an error should be returned normally. func TestInsertPartialFail1(t *testing.T) { - executor, _, _, sbclookup, _, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbclookup, _ := createExecutorEnv(t) // Make the first DML fail, there should be no rollback. sbclookup.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 1 @@ -2027,8 +1991,7 @@ func TestInsertPartialFail1(t *testing.T) { // after successful execution of the first, then the transaction must // be rolled back due to partial execution. func TestInsertPartialFail2(t *testing.T) { - executor, sbc1, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, _ := createExecutorEnv(t) // Make the second DML fail, it should result in a rollback. sbc1.MustFailExecute[sqlparser.StmtInsert] = 1 @@ -2068,8 +2031,7 @@ func TestInsertPartialFail2(t *testing.T) { } func TestMultiInsertSharded(t *testing.T) { - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", @@ -2176,8 +2138,7 @@ func TestMultiInsertSharded(t *testing.T) { } func TestMultiInsertGenerator(t *testing.T) { - executor, sbc, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc, _, sbclookup, ctx := createExecutorEnv(t) sbclookup.SetResults([]*sqltypes.Result{{ Rows: [][]sqltypes.Value{{ @@ -2225,8 +2186,7 @@ func TestMultiInsertGenerator(t *testing.T) { } func TestMultiInsertGeneratorSparse(t *testing.T) { - executor, sbc, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc, _, sbclookup, ctx := createExecutorEnv(t) sbclookup.SetResults([]*sqltypes.Result{{ Rows: [][]sqltypes.Value{{ @@ -2303,8 +2263,7 @@ func TestInsertBadAutoInc(t *testing.T) { } } ` - executor, _, _, _, ctx, closer := createCustomExecutor(t, vschema) - defer closer() + executor, _, _, _, ctx := createCustomExecutor(t, vschema) // If auto inc table cannot be found, the table should not be added to vschema. session := &vtgatepb.Session{ @@ -2377,8 +2336,7 @@ func TestKeyDestRangeQuery(t *testing.T) { for _, tc := range tests { t.Run(tc.targetString+" - "+tc.inputQuery, func(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: tc.targetString, @@ -2401,8 +2359,7 @@ func TestKeyDestRangeQuery(t *testing.T) { } // it does not work for inserts - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "TestExecutor[-]", } @@ -2422,8 +2379,7 @@ func assertQueriesContain(t *testing.T, sql, sbcName string, sbc *sandboxconn.Sa // Prepared statement tests func TestUpdateEqualWithPrepare(t *testing.T) { - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2444,8 +2400,7 @@ func TestUpdateEqualWithPrepare(t *testing.T) { assertQueries(t, sbc1, nil) } func TestInsertShardedWithPrepare(t *testing.T) { - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2469,8 +2424,7 @@ func TestInsertShardedWithPrepare(t *testing.T) { } func TestDeleteEqualWithPrepare(t *testing.T) { - executor, sbc, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc, _, sbclookup, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", @@ -2488,8 +2442,7 @@ func TestDeleteEqualWithPrepare(t *testing.T) { } func TestUpdateLastInsertID(t *testing.T) { - executor, sbc1, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, ctx := createExecutorEnv(t) executor.normalize = true @@ -2511,8 +2464,7 @@ func TestUpdateLastInsertID(t *testing.T) { } func TestUpdateReference(t *testing.T) { - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2554,8 +2506,7 @@ func TestUpdateReference(t *testing.T) { } func TestDeleteLookupOwnedEqual(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) sbc1.SetResults([]*sqltypes.Result{ sqltypes.MakeTestResult(sqltypes.MakeTestFields("uniq_col|keyspace_id", "int64|varbinary"), "1|N±\u0090ɢú\u0016\u009C"), @@ -2585,8 +2536,7 @@ func TestDeleteLookupOwnedEqual(t *testing.T) { } func TestDeleteReference(t *testing.T) { - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2627,8 +2577,7 @@ func TestDeleteReference(t *testing.T) { } func TestReservedConnDML(t *testing.T) { - executor, _, _, sbc, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbc, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("TestReservedConnDML") defer executor.queryLogger.Unsubscribe(logChan) @@ -2680,8 +2629,7 @@ func TestReservedConnDML(t *testing.T) { func TestStreamingDML(t *testing.T) { method := "TestStreamingDML" - executor, _, _, sbc, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbc, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe(method) defer executor.queryLogger.Unsubscribe(logChan) @@ -2765,8 +2713,7 @@ func TestStreamingDML(t *testing.T) { } func TestPartialVindexInsertQueryFailure(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2819,8 +2766,7 @@ func TestPartialVindexInsertQueryFailure(t *testing.T) { } func TestPartialVindexInsertQueryFailureAutoCommit(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2864,8 +2810,7 @@ func TestPartialVindexInsertQueryFailureAutoCommit(t *testing.T) { // The change for it cannot be done as the executor level and will be made at the VTGate entry point. // Test TestMultiInternalSavepointVtGate shows that it fixes the behaviour. func TestMultiInternalSavepoint(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) session := NewAutocommitSession(&vtgatepb.Session{}) _, err := executorExecSession(ctx, executor, "begin", nil, session.Session) @@ -2913,8 +2858,7 @@ func TestMultiInternalSavepoint(t *testing.T) { } func TestInsertSelectFromDual(t *testing.T) { - executor, sbc1, sbc2, sbclookup, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, _ := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("TestInsertSelect") defer executor.queryLogger.Unsubscribe(logChan) @@ -2969,8 +2913,7 @@ func TestInsertSelectFromDual(t *testing.T) { } func TestInsertSelectFromTable(t *testing.T) { - executor, sbc1, sbc2, sbclookup, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, _ := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("TestInsertSelect") defer executor.queryLogger.Unsubscribe(logChan) @@ -3031,8 +2974,7 @@ func TestInsertSelectFromTable(t *testing.T) { } func TestInsertReference(t *testing.T) { - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) diff --git a/go/vt/vtgate/executor_framework_test.go b/go/vt/vtgate/executor_framework_test.go index 997d6da87aa..4917ce2864a 100644 --- a/go/vt/vtgate/executor_framework_test.go +++ b/go/vt/vtgate/executor_framework_test.go @@ -127,7 +127,7 @@ func init() { vindexes.Register("keyrange_lookuper_unique", newKeyRangeLookuperUnique) } -func createExecutorEnv(t testing.TB) (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn.SandboxConn, ctx context.Context, closer func()) { +func createExecutorEnv(t testing.TB) (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn.SandboxConn, ctx context.Context) { var cancel context.CancelFunc ctx, cancel = context.WithCancel(context.Background()) cell := "aa" @@ -179,14 +179,17 @@ func createExecutorEnv(t testing.TB) (executor *Executor, sbc1, sbc2, sbclookup executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) key.AnyShardPicker = DestinationAnyShardPickerFirstShard{} - return executor, sbc1, sbc2, sbclookup, ctx, func() { + + t.Cleanup(func() { + defer utils.EnsureNoLeaks(t) executor.Close() cancel() - defer utils.EnsureNoLeaks(t) - } + }) + + return executor, sbc1, sbc2, sbclookup, ctx } -func createCustomExecutor(t testing.TB, vschema string) (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn.SandboxConn, ctx context.Context, closer func()) { +func createCustomExecutor(t testing.TB, vschema string) (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn.SandboxConn, ctx context.Context) { var cancel context.CancelFunc ctx, cancel = context.WithCancel(context.Background()) cell := "aa" @@ -205,14 +208,17 @@ func createCustomExecutor(t testing.TB, vschema string) (executor *Executor, sbc queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize) plans := cache.NewDefaultCacheImpl(cache.DefaultConfig) executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) - return executor, sbc1, sbc2, sbclookup, ctx, func() { + + t.Cleanup(func() { defer utils.EnsureNoLeaks(t) - defer cancel() executor.Close() - } + cancel() + }) + + return executor, sbc1, sbc2, sbclookup, ctx } -func createCustomExecutorSetValues(t testing.TB, vschema string, values []*sqltypes.Result) (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn.SandboxConn, ctx context.Context, closer func()) { +func createCustomExecutorSetValues(t testing.TB, vschema string, values []*sqltypes.Result) (executor *Executor, sbc1, sbc2, sbclookup *sandboxconn.SandboxConn, ctx context.Context) { var cancel context.CancelFunc ctx, cancel = context.WithCancel(context.Background()) cell := "aa" @@ -238,11 +244,14 @@ func createCustomExecutorSetValues(t testing.TB, vschema string, values []*sqlty queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize) plans := cache.NewDefaultCacheImpl(cache.DefaultConfig) executor = NewExecutor(ctx, serv, cell, resolver, false, false, testBufferSize, plans, nil, false, querypb.ExecuteOptions_Gen4, queryLogger) - return executor, sbcs[0], sbcs[1], sbclookup, ctx, func() { + + t.Cleanup(func() { + defer utils.EnsureNoLeaks(t) executor.Close() cancel() - defer utils.EnsureNoLeaks(t) - } + }) + + return executor, sbcs[0], sbcs[1], sbclookup, ctx } func executorExecSession(ctx context.Context, executor *Executor, sql string, bv map[string]*querypb.BindVariable, session *vtgatepb.Session) (*sqltypes.Result, error) { diff --git a/go/vt/vtgate/executor_scatter_stats_test.go b/go/vt/vtgate/executor_scatter_stats_test.go index 64348ac8a2e..9a9e54517e7 100644 --- a/go/vt/vtgate/executor_scatter_stats_test.go +++ b/go/vt/vtgate/executor_scatter_stats_test.go @@ -26,8 +26,7 @@ import ( ) func TestScatterStatsWithNoScatterQuery(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) @@ -40,8 +39,7 @@ func TestScatterStatsWithNoScatterQuery(t *testing.T) { } func TestScatterStatsWithSingleScatterQuery(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) _, err := executor.Execute(ctx, nil, "TestExecutorResultsExceeded", session, "select * from user", nil) @@ -53,8 +51,7 @@ func TestScatterStatsWithSingleScatterQuery(t *testing.T) { } func TestScatterStatsHttpWriting(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) _, err := executor.Execute(ctx, nil, "TestExecutorResultsExceeded", session, "select * from user", nil) diff --git a/go/vt/vtgate/executor_select_test.go b/go/vt/vtgate/executor_select_test.go index 8b83853fcf6..b8d2625ec04 100644 --- a/go/vt/vtgate/executor_select_test.go +++ b/go/vt/vtgate/executor_select_test.go @@ -52,8 +52,7 @@ import ( ) func TestSelectNext(t *testing.T) { - executor, _, _, sbclookup, _, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbclookup, _ := createExecutorEnv(t) query := "select next :n values from user_seq" bv := map[string]*querypb.BindVariable{"n": sqltypes.Int64BindVariable(2)} @@ -107,8 +106,7 @@ func TestSelectNext(t *testing.T) { } func TestSelectDBA(t *testing.T) { - executor, sbc1, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, _ := createExecutorEnv(t) query := "select * from INFORMATION_SCHEMA.foo" _, err := executor.Execute(context.Background(), nil, "TestSelectDBA", @@ -162,8 +160,7 @@ func TestSelectDBA(t *testing.T) { } func TestSystemVariablesMySQLBelow80(t *testing.T) { - executor, sbc1, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, _ := createExecutorEnv(t) executor.normalize = true sqlparser.SetParserVersion("57000") @@ -199,8 +196,7 @@ func TestSystemVariablesMySQLBelow80(t *testing.T) { } func TestSystemVariablesWithSetVarDisabled(t *testing.T) { - executor, sbc1, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, _ := createExecutorEnv(t) executor.normalize = true sqlparser.SetParserVersion("80000") @@ -238,8 +234,7 @@ func TestSystemVariablesWithSetVarDisabled(t *testing.T) { } func TestSetSystemVariablesTx(t *testing.T) { - executor, sbc1, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, _ := createExecutorEnv(t) executor.normalize = true sqlparser.SetParserVersion("80001") @@ -287,8 +282,7 @@ func TestSetSystemVariablesTx(t *testing.T) { } func TestSetSystemVariables(t *testing.T) { - executor, _, _, lookup, _, closer := createExecutorEnv(t) - defer closer() + executor, _, _, lookup, _ := createExecutorEnv(t) executor.normalize = true sqlparser.SetParserVersion("80001") @@ -401,8 +395,7 @@ func TestSetSystemVariables(t *testing.T) { } func TestSetSystemVariablesWithReservedConnection(t *testing.T) { - executor, sbc1, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, _ := createExecutorEnv(t) executor.normalize = true session := NewAutocommitSession(&vtgatepb.Session{EnableSystemSettings: true, SystemVariables: map[string]string{}}) @@ -445,8 +438,7 @@ func TestSetSystemVariablesWithReservedConnection(t *testing.T) { } func TestCreateTableValidTimestamp(t *testing.T) { - executor, sbc1, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, _ := createExecutorEnv(t) executor.normalize = true session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor", SystemVariables: map[string]string{"sql_mode": "ALLOW_INVALID_DATES"}}) @@ -465,8 +457,7 @@ func TestCreateTableValidTimestamp(t *testing.T) { } func TestGen4SelectDBA(t *testing.T) { - executor, sbc1, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, _ := createExecutorEnv(t) executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 @@ -541,8 +532,7 @@ func TestGen4SelectDBA(t *testing.T) { } func TestUnsharded(t *testing.T) { - executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbclookup, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", @@ -557,8 +547,7 @@ func TestUnsharded(t *testing.T) { } func TestUnshardedComments(t *testing.T) { - executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbclookup, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", @@ -602,8 +591,7 @@ func TestUnshardedComments(t *testing.T) { } func TestStreamUnsharded(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -619,8 +607,7 @@ func TestStreamUnsharded(t *testing.T) { } func TestStreamBuffering(t *testing.T) { - executor, _, _, sbclookup, _, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbclookup, _ := createExecutorEnv(t) // This test is similar to TestStreamUnsharded except that it returns a Result > 10 bytes, // such that the splitting of the Result into multiple Result responses gets tested. @@ -676,8 +663,7 @@ func TestStreamBuffering(t *testing.T) { } func TestStreamLimitOffset(t *testing.T) { - executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, _ := createExecutorEnv(t) // This test is similar to TestStreamUnsharded except that it returns a Result > 10 bytes, // such that the splitting of the Result into multiple Result responses gets tested. @@ -755,8 +741,7 @@ func TestStreamLimitOffset(t *testing.T) { } func TestSelectLastInsertId(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", LastInsertId: 52, @@ -780,8 +765,7 @@ func TestSelectLastInsertId(t *testing.T) { } func TestSelectSystemVariables(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", @@ -841,8 +825,7 @@ func TestSelectSystemVariables(t *testing.T) { } func TestSelectInitializedVitessAwareVariable(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) executor.normalize = true logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -874,8 +857,7 @@ func TestSelectInitializedVitessAwareVariable(t *testing.T) { } func TestSelectUserDefinedVariable(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) executor.normalize = true logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -911,8 +893,7 @@ func TestSelectUserDefinedVariable(t *testing.T) { } func TestFoundRows(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) executor.normalize = true logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -939,8 +920,7 @@ func TestFoundRows(t *testing.T) { } func TestRowCount(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) executor.normalize = true logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -973,8 +953,7 @@ func testRowCount(t *testing.T, ctx context.Context, executor *Executor, session } func TestSelectLastInsertIdInUnion(t *testing.T) { - executor, sbc1, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, ctx := createExecutorEnv(t) executor.normalize = true session := &vtgatepb.Session{ @@ -1008,8 +987,7 @@ func TestSelectLastInsertIdInUnion(t *testing.T) { } func TestSelectLastInsertIdInWhere(t *testing.T) { - executor, _, _, lookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, lookup, ctx := createExecutorEnv(t) executor.normalize = true logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -1029,8 +1007,7 @@ func TestSelectLastInsertIdInWhere(t *testing.T) { } func TestLastInsertIDInVirtualTable(t *testing.T) { - executor, sbc1, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, ctx := createExecutorEnv(t) executor.normalize = true result1 := []*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -1058,8 +1035,7 @@ func TestLastInsertIDInVirtualTable(t *testing.T) { } func TestLastInsertIDInSubQueryExpression(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) executor.normalize = true session := &vtgatepb.Session{ TargetString: "@primary", @@ -1083,8 +1059,7 @@ func TestLastInsertIDInSubQueryExpression(t *testing.T) { } func TestSelectDatabase(t *testing.T) { - executor, _, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, _ := createExecutorEnv(t) executor.normalize = true sql := "select database()" newSession := &vtgatepb.Session{ @@ -1113,8 +1088,7 @@ func TestSelectDatabase(t *testing.T) { } func TestSelectBindvars(t *testing.T) { - executor, sbc1, sbc2, lookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, lookup, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -1221,8 +1195,7 @@ func TestSelectBindvars(t *testing.T) { } func TestSelectEqual(t *testing.T) { - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", @@ -1292,8 +1265,7 @@ func TestSelectEqual(t *testing.T) { } func TestSelectINFromOR(t *testing.T) { - executor, sbc1, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, ctx := createExecutorEnv(t) executor.pv = querypb.ExecuteOptions_Gen4 session := &vtgatepb.Session{ @@ -1311,8 +1283,7 @@ func TestSelectINFromOR(t *testing.T) { } func TestSelectDual(t *testing.T) { - executor, sbc1, _, lookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, lookup, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", @@ -1331,8 +1302,7 @@ func TestSelectDual(t *testing.T) { } func TestSelectComments(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", @@ -1351,8 +1321,7 @@ func TestSelectComments(t *testing.T) { } func TestSelectNormalize(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) executor.normalize = true session := &vtgatepb.Session{ @@ -1388,8 +1357,7 @@ func TestSelectNormalize(t *testing.T) { } func TestSelectCaseSensitivity(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", @@ -1408,8 +1376,7 @@ func TestSelectCaseSensitivity(t *testing.T) { } func TestStreamSelectEqual(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) sql := "select id from user where id = 1" result, err := executorStream(ctx, executor, sql) @@ -1421,8 +1388,7 @@ func TestStreamSelectEqual(t *testing.T) { } func TestSelectKeyRange(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", @@ -1441,8 +1407,7 @@ func TestSelectKeyRange(t *testing.T) { } func TestSelectKeyRangeUnique(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "@primary", @@ -1461,8 +1426,7 @@ func TestSelectKeyRangeUnique(t *testing.T) { } func TestSelectIN(t *testing.T) { - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) // Constant in IN clause is just a number, not a bind variable. session := &vtgatepb.Session{ @@ -1551,8 +1515,7 @@ func TestSelectIN(t *testing.T) { } func TestStreamSelectIN(t *testing.T) { - executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbclookup, ctx := createExecutorEnv(t) sql := "select id from user where id in (1)" result, err := executorStream(ctx, executor, sql) @@ -1603,9 +1566,7 @@ func createExecutor(ctx context.Context, serv *sandboxTopo, cell string, resolve } func TestSelectScatter(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Special setup: Don't use createExecutorEnv. cell := "aa" @@ -1644,9 +1605,7 @@ func TestSelectScatter(t *testing.T) { } func TestSelectScatterPartial(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Special setup: Don't use createExecutorEnv. session := &vtgatepb.Session{ @@ -1713,9 +1672,7 @@ func TestSelectScatterPartial(t *testing.T) { } func TestSelectScatterPartialOLAP(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Special setup: Don't use createExecutorEnv. cell := "aa" @@ -1770,9 +1727,7 @@ func TestSelectScatterPartialOLAP(t *testing.T) { } func TestSelectScatterPartialOLAP2(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Special setup: Don't use createExecutorEnv. cell := "aa" @@ -1832,9 +1787,7 @@ func TestSelectScatterPartialOLAP2(t *testing.T) { } func TestStreamSelectScatter(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Special setup: Don't use createExecutorEnv. cell := "aa" @@ -1873,9 +1826,7 @@ func TestStreamSelectScatter(t *testing.T) { // TestSelectScatterOrderBy will run an ORDER BY query that will scatter out to 8 shards and return the 8 rows (one per shard) sorted. func TestSelectScatterOrderBy(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Special setup: Don't use createExecutorEnv. cell := "aa" @@ -1948,9 +1899,7 @@ func TestSelectScatterOrderBy(t *testing.T) { // TestSelectScatterOrderByVarChar will run an ORDER BY query that will scatter out to 8 shards and return the 8 rows (one per shard) sorted. func TestSelectScatterOrderByVarChar(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Special setup: Don't use createExecutorEnv. cell := "aa" @@ -2020,9 +1969,7 @@ func TestSelectScatterOrderByVarChar(t *testing.T) { } func TestStreamSelectScatterOrderBy(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Special setup: Don't use createExecutorEnv. cell := "aa" @@ -2084,9 +2031,7 @@ func TestStreamSelectScatterOrderBy(t *testing.T) { } func TestStreamSelectScatterOrderByVarChar(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Special setup: Don't use createExecutorEnv. cell := "aa" @@ -2147,9 +2092,7 @@ func TestStreamSelectScatterOrderByVarChar(t *testing.T) { // TestSelectScatterAggregate will run an aggregate query that will scatter out to 8 shards and return 4 aggregated rows. func TestSelectScatterAggregate(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Special setup: Don't use createExecutorEnv. cell := "aa" @@ -2215,9 +2158,7 @@ func TestSelectScatterAggregate(t *testing.T) { } func TestStreamSelectScatterAggregate(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Special setup: Don't use createExecutorEnv. cell := "aa" @@ -2281,9 +2222,7 @@ func TestStreamSelectScatterAggregate(t *testing.T) { // TestSelectScatterLimit will run a limit query (ordered for consistency) against // a scatter route and verify that the limit primitive works as intended. func TestSelectScatterLimit(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Special setup: Don't use createExecutorEnv. cell := "aa" @@ -2358,9 +2297,7 @@ func TestSelectScatterLimit(t *testing.T) { // TestStreamSelectScatterLimit will run a streaming limit query (ordered for consistency) against // a scatter route and verify that the limit primitive works as intended. func TestStreamSelectScatterLimit(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Special setup: Don't use createExecutorEnv. cell := "aa" @@ -2431,8 +2368,7 @@ func TestStreamSelectScatterLimit(t *testing.T) { // TODO(sougou): stream and non-stream testing are very similar. // Could reuse code, func TestSimpleJoin(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2472,8 +2408,7 @@ func TestSimpleJoin(t *testing.T) { } func TestJoinComments(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2498,8 +2433,7 @@ func TestJoinComments(t *testing.T) { } func TestSimpleJoinStream(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2537,8 +2471,7 @@ func TestSimpleJoinStream(t *testing.T) { } func TestVarJoin(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2576,8 +2509,7 @@ func TestVarJoin(t *testing.T) { } func TestVarJoinStream(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -2612,8 +2544,7 @@ func TestVarJoinStream(t *testing.T) { } func TestLeftJoin(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) result1 := []*sqltypes.Result{{ @@ -2659,8 +2590,7 @@ func TestLeftJoin(t *testing.T) { } func TestLeftJoinStream(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) result1 := []*sqltypes.Result{{ Fields: []*querypb.Field{ {Name: "id", Type: sqltypes.Int32, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NUM_FLAG)}, @@ -2700,8 +2630,7 @@ func TestLeftJoinStream(t *testing.T) { } func TestEmptyJoin(t *testing.T) { - executor, sbc1, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, ctx := createExecutorEnv(t) // Empty result requires a field query for the second part of join, // which is sent to shard 0. sbc1.SetResults([]*sqltypes.Result{{ @@ -2742,8 +2671,7 @@ func TestEmptyJoin(t *testing.T) { } func TestEmptyJoinStream(t *testing.T) { - executor, sbc1, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, ctx := createExecutorEnv(t) // Empty result requires a field query for the second part of join, // which is sent to shard 0. sbc1.SetResults([]*sqltypes.Result{{ @@ -2779,8 +2707,7 @@ func TestEmptyJoinStream(t *testing.T) { } func TestEmptyJoinRecursive(t *testing.T) { - executor, sbc1, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, ctx := createExecutorEnv(t) // Make sure it also works recursively. sbc1.SetResults([]*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -2827,8 +2754,7 @@ func TestEmptyJoinRecursive(t *testing.T) { } func TestEmptyJoinRecursiveStream(t *testing.T) { - executor, sbc1, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, ctx := createExecutorEnv(t) // Make sure it also works recursively. sbc1.SetResults([]*sqltypes.Result{{ Fields: []*querypb.Field{ @@ -2872,8 +2798,7 @@ func TestEmptyJoinRecursiveStream(t *testing.T) { } func TestCrossShardSubquery(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) result1 := []*sqltypes.Result{{ Fields: []*querypb.Field{ {Name: "id", Type: sqltypes.Int32}, @@ -2910,8 +2835,7 @@ func TestCrossShardSubquery(t *testing.T) { } func TestSubQueryAndQueryWithLimit(t *testing.T) { - executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, _ := createExecutorEnv(t) result1 := []*sqltypes.Result{{ Fields: []*querypb.Field{ {Name: "id", Type: sqltypes.Int32, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NUM_FLAG)}, @@ -2953,8 +2877,7 @@ func TestSubQueryAndQueryWithLimit(t *testing.T) { } func TestCrossShardSubqueryStream(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) result1 := []*sqltypes.Result{{ Fields: []*querypb.Field{ {Name: "id", Type: sqltypes.Int32, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NUM_FLAG)}, @@ -2994,8 +2917,7 @@ func TestCrossShardSubqueryStream(t *testing.T) { } func TestCrossShardSubqueryGetFields(t *testing.T) { - executor, sbc1, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, sbclookup, ctx := createExecutorEnv(t) sbclookup.SetResults([]*sqltypes.Result{{ Fields: []*querypb.Field{ {Name: "col", Type: sqltypes.Int32, Charset: collations.CollationBinaryID, Flags: uint32(querypb.MySqlFlag_NUM_FLAG)}, @@ -3036,8 +2958,7 @@ func TestCrossShardSubqueryGetFields(t *testing.T) { } func TestSelectBindvarswithPrepare(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -3061,8 +2982,7 @@ func TestSelectBindvarswithPrepare(t *testing.T) { } func TestSelectDatabasePrepare(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) executor.normalize = true logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -3076,8 +2996,7 @@ func TestSelectDatabasePrepare(t *testing.T) { } func TestSelectWithUnionAll(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) executor.normalize = true sql := "select id from user where id in (1, 2, 3) union all select id from user where id in (1, 2, 3)" bv, _ := sqltypes.BuildBindVariable([]int64{1, 2, 3}) @@ -3132,8 +3051,7 @@ func TestSelectWithUnionAll(t *testing.T) { } func TestSelectLock(t *testing.T) { - executor, sbc1, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, _ := createExecutorEnv(t) session := NewSafeSession(nil) session.Session.InTransaction = true session.ShardSessions = []*vtgatepb.Session_ShardSession{{ @@ -3192,8 +3110,7 @@ func TestSelectLock(t *testing.T) { } func TestLockReserve(t *testing.T) { - executor, _, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, _ := createExecutorEnv(t) // no connection should be reserved for these queries. tcases := []string{ @@ -3221,8 +3138,7 @@ func TestLockReserve(t *testing.T) { } func TestSelectFromInformationSchema(t *testing.T) { - executor, sbc1, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, _ := createExecutorEnv(t) session := NewSafeSession(nil) // check failure when trying to query two keyspaces @@ -3247,9 +3163,7 @@ func TestSelectFromInformationSchema(t *testing.T) { } func TestStreamOrderByLimitWithMultipleResults(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Special setup: Don't use createExecutorEnv. cell := "aa" @@ -3291,9 +3205,7 @@ func TestStreamOrderByLimitWithMultipleResults(t *testing.T) { } func TestSelectScatterFails(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sess := &vtgatepb.Session{} cell := "aa" @@ -3352,8 +3264,7 @@ func TestSelectScatterFails(t *testing.T) { } func TestGen4SelectStraightJoin(t *testing.T) { - executor, sbc1, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, _ := createExecutorEnv(t) executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) @@ -3381,8 +3292,7 @@ func TestGen4SelectStraightJoin(t *testing.T) { } func TestGen4MultiColumnVindexEqual(t *testing.T) { - executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, _ := createExecutorEnv(t) executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 @@ -3421,8 +3331,7 @@ func TestGen4MultiColumnVindexEqual(t *testing.T) { } func TestGen4MultiColumnVindexIn(t *testing.T) { - executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, _ := createExecutorEnv(t) executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 @@ -3461,8 +3370,7 @@ func TestGen4MultiColumnVindexIn(t *testing.T) { } func TestGen4MultiColMixedColComparision(t *testing.T) { - executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, _ := createExecutorEnv(t) executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 @@ -3499,8 +3407,7 @@ func TestGen4MultiColMixedColComparision(t *testing.T) { } func TestGen4MultiColBestVindexSel(t *testing.T) { - executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, _ := createExecutorEnv(t) executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 @@ -3546,8 +3453,7 @@ func TestGen4MultiColBestVindexSel(t *testing.T) { } func TestGen4MultiColMultiEqual(t *testing.T) { - executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, _ := createExecutorEnv(t) executor.normalize = true executor.pv = querypb.ExecuteOptions_Gen4 @@ -3570,8 +3476,7 @@ func TestGen4MultiColMultiEqual(t *testing.T) { } func TestGen4SelectUnqualifiedReferenceTable(t *testing.T) { - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) executor.pv = querypb.ExecuteOptions_Gen4 query := "select * from zip_detail" @@ -3592,8 +3497,7 @@ func TestGen4SelectUnqualifiedReferenceTable(t *testing.T) { } func TestGen4SelectQualifiedReferenceTable(t *testing.T) { - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) executor.pv = querypb.ExecuteOptions_Gen4 query := fmt.Sprintf("select * from %s.zip_detail", KsTestSharded) @@ -3614,8 +3518,7 @@ func TestGen4SelectQualifiedReferenceTable(t *testing.T) { } func TestGen4JoinUnqualifiedReferenceTable(t *testing.T) { - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) executor.pv = querypb.ExecuteOptions_Gen4 query := "select * from user join zip_detail on user.zip_detail_id = zip_detail.id" @@ -3652,8 +3555,7 @@ func TestGen4JoinUnqualifiedReferenceTable(t *testing.T) { } func TestGen4CrossShardJoinQualifiedReferenceTable(t *testing.T) { - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) executor.pv = querypb.ExecuteOptions_Gen4 query := "select user.id from user join TestUnsharded.zip_detail on user.zip_detail_id = TestUnsharded.zip_detail.id" @@ -3692,9 +3594,7 @@ func TestGen4CrossShardJoinQualifiedReferenceTable(t *testing.T) { } func TestRegionRange(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Special setup: Don't use createExecutorEnv. cell := "regioncell" @@ -3744,9 +3644,7 @@ func TestRegionRange(t *testing.T) { } func TestMultiCol(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Special setup: Don't use createLegacyExecutorEnv. cell := "multicol" @@ -3827,9 +3725,7 @@ var multiColVschema = ` ` func TestMultiColPartial(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Special setup: Don't use createLegacyExecutorEnv. cell := "multicol" @@ -3893,9 +3789,7 @@ func TestMultiColPartial(t *testing.T) { } func TestSelectAggregationNoData(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Special setup: Don't use createExecutorEnv. cell := "aa" @@ -3987,9 +3881,7 @@ func TestSelectAggregationNoData(t *testing.T) { } func TestSelectAggregationData(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) // Special setup: Don't use createExecutorEnv. cell := "aa" @@ -4137,9 +4029,7 @@ func TestSelectAggregationData(t *testing.T) { } func TestSelectAggregationRandom(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) cell := "aa" hc := discovery.NewFakeHealthCheck(nil) @@ -4177,8 +4067,7 @@ func TestSelectAggregationRandom(t *testing.T) { } func TestSelectHexAndBit(t *testing.T) { - executor, _, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, _ := createExecutorEnv(t) executor.normalize = true session := NewAutocommitSession(&vtgatepb.Session{}) @@ -4194,8 +4083,7 @@ func TestSelectHexAndBit(t *testing.T) { // TestSelectCFC tests validates that cfc vindex plan gets cached and same plan is getting reused. // This also validates that cache_size is able to calculate the cfc vindex plan size. func TestSelectCFC(t *testing.T) { - executor, _, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, _ := createExecutorEnv(t) executor.normalize = true session := NewAutocommitSession(&vtgatepb.Session{}) @@ -4220,8 +4108,7 @@ func TestSelectCFC(t *testing.T) { } func TestSelectView(t *testing.T) { - executor, sbc, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc, _, _, _ := createExecutorEnv(t) // add the view to local vschema err := executor.vschema.AddView(KsTestSharded, "user_details_view", "select user.id, user_extra.col from user join user_extra on user.id = user_extra.user_id") require.NoError(t, err) diff --git a/go/vt/vtgate/executor_set_test.go b/go/vt/vtgate/executor_set_test.go index 900593608fc..e71a41eeb7f 100644 --- a/go/vt/vtgate/executor_set_test.go +++ b/go/vt/vtgate/executor_set_test.go @@ -40,8 +40,7 @@ import ( ) func TestExecutorSet(t *testing.T) { - executorEnv, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executorEnv, _, _, _, ctx := createExecutorEnv(t) testcases := []struct { in string @@ -282,8 +281,7 @@ func TestExecutorSet(t *testing.T) { } func TestExecutorSetOp(t *testing.T) { - executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbclookup, ctx := createExecutorEnv(t) sysVarSetEnabled = true returnResult := func(columnName, typ, value string) *sqltypes.Result { @@ -382,67 +380,69 @@ func TestExecutorSetOp(t *testing.T) { } func TestExecutorSetMetadata(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) - set := "set @@vitess_metadata.app_keyspace_v1= '1'" - _, err := executor.Execute(ctx, nil, "TestExecute", session, set, nil) - assert.Equalf(t, vtrpcpb.Code_PERMISSION_DENIED, vterrors.Code(err), "expected error %v, got error: %v", vtrpcpb.Code_PERMISSION_DENIED, err) + t.Run("Session 1", func(t *testing.T) { + executor, _, _, _, ctx := createExecutorEnv(t) + session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) - vschemaacl.AuthorizedDDLUsers = "%" - defer func() { - vschemaacl.AuthorizedDDLUsers = "" - }() + set := "set @@vitess_metadata.app_keyspace_v1= '1'" + _, err := executor.Execute(ctx, nil, "TestExecute", session, set, nil) + assert.Equalf(t, vtrpcpb.Code_PERMISSION_DENIED, vterrors.Code(err), "expected error %v, got error: %v", vtrpcpb.Code_PERMISSION_DENIED, err) + }) - closer() - executor, _, _, _, ctx, closer = createExecutorEnv(t) - defer closer() - session = NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) - - set = "set @@vitess_metadata.app_keyspace_v1= '1'" - _, err = executor.Execute(ctx, nil, "TestExecute", session, set, nil) - assert.NoError(t, err, "%s error: %v", set, err) - - show := `show vitess_metadata variables like 'app\\_keyspace\\_v_'` - result, err := executor.Execute(ctx, nil, "TestExecute", session, show, nil) - assert.NoError(t, err) - - want := "1" - got := result.Rows[0][1].ToString() - assert.Equalf(t, want, got, "want migrations %s, result %s", want, got) - - // Update metadata - set = "set @@vitess_metadata.app_keyspace_v2='2'" - _, err = executor.Execute(ctx, nil, "TestExecute", session, set, nil) - assert.NoError(t, err, "%s error: %v", set, err) - - show = `show vitess_metadata variables like 'app\\_keyspace\\_v%'` - gotqr, err := executor.Execute(ctx, nil, "TestExecute", session, show, nil) - assert.NoError(t, err) - - wantqr := &sqltypes.Result{ - Fields: buildVarCharFields("Key", "Value"), - Rows: [][]sqltypes.Value{ - buildVarCharRow("app_keyspace_v1", "1"), - buildVarCharRow("app_keyspace_v2", "2"), - }, - RowsAffected: 2, - } + t.Run("Session 2", func(t *testing.T) { + vschemaacl.AuthorizedDDLUsers = "%" + defer func() { + vschemaacl.AuthorizedDDLUsers = "" + }() + + executor, _, _, _, ctx := createExecutorEnv(t) + session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) + + set := "set @@vitess_metadata.app_keyspace_v1= '1'" + _, err := executor.Execute(ctx, nil, "TestExecute", session, set, nil) + assert.NoError(t, err, "%s error: %v", set, err) + + show := `show vitess_metadata variables like 'app\\_keyspace\\_v_'` + result, err := executor.Execute(ctx, nil, "TestExecute", session, show, nil) + assert.NoError(t, err) + + want := "1" + got := result.Rows[0][1].ToString() + assert.Equalf(t, want, got, "want migrations %s, result %s", want, got) + + // Update metadata + set = "set @@vitess_metadata.app_keyspace_v2='2'" + _, err = executor.Execute(ctx, nil, "TestExecute", session, set, nil) + assert.NoError(t, err, "%s error: %v", set, err) + + show = `show vitess_metadata variables like 'app\\_keyspace\\_v%'` + gotqr, err := executor.Execute(ctx, nil, "TestExecute", session, show, nil) + assert.NoError(t, err) + + wantqr := &sqltypes.Result{ + Fields: buildVarCharFields("Key", "Value"), + Rows: [][]sqltypes.Value{ + buildVarCharRow("app_keyspace_v1", "1"), + buildVarCharRow("app_keyspace_v2", "2"), + }, + RowsAffected: 2, + } - assert.Equal(t, wantqr.Fields, gotqr.Fields) - assert.ElementsMatch(t, wantqr.Rows, gotqr.Rows) + assert.Equal(t, wantqr.Fields, gotqr.Fields) + assert.ElementsMatch(t, wantqr.Rows, gotqr.Rows) - show = "show vitess_metadata variables" - gotqr, err = executor.Execute(ctx, nil, "TestExecute", session, show, nil) - require.NoError(t, err) + show = "show vitess_metadata variables" + gotqr, err = executor.Execute(ctx, nil, "TestExecute", session, show, nil) + require.NoError(t, err) - assert.Equal(t, wantqr.Fields, gotqr.Fields) - assert.ElementsMatch(t, wantqr.Rows, gotqr.Rows) + assert.Equal(t, wantqr.Fields, gotqr.Fields) + assert.ElementsMatch(t, wantqr.Rows, gotqr.Rows) + }) } func TestPlanExecutorSetUDV(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) testcases := []struct { in string @@ -472,8 +472,7 @@ func TestPlanExecutorSetUDV(t *testing.T) { } func TestSetUDVFromTabletInput(t *testing.T) { - executor, sbc1, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, _, ctx := createExecutorEnv(t) fields := sqltypes.MakeTestFields("some", "VARCHAR") sbc1.SetResults([]*sqltypes.Result{ @@ -504,8 +503,7 @@ func createMap(keys []string, values []any) map[string]*querypb.BindVariable { } func TestSetVar(t *testing.T) { - executor, _, _, sbc, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbc, ctx := createExecutorEnv(t) executor.normalize = true oldVersion := sqlparser.GetParserVersion() @@ -549,8 +547,7 @@ func TestSetVar(t *testing.T) { } func TestSetVarShowVariables(t *testing.T) { - executor, _, _, sbc, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbc, ctx := createExecutorEnv(t) executor.normalize = true oldVersion := sqlparser.GetParserVersion() @@ -579,8 +576,7 @@ func TestSetVarShowVariables(t *testing.T) { } func TestExecutorSetAndSelect(t *testing.T) { - e, _, _, sbc, ctx, closer := createExecutorEnv(t) - defer closer() + e, _, _, sbc, ctx := createExecutorEnv(t) e.normalize = true testcases := []struct { diff --git a/go/vt/vtgate/executor_stream_test.go b/go/vt/vtgate/executor_stream_test.go index 4f423696982..6f6dcd9f6b4 100644 --- a/go/vt/vtgate/executor_stream_test.go +++ b/go/vt/vtgate/executor_stream_test.go @@ -39,8 +39,7 @@ import ( ) func TestStreamSQLUnsharded(t *testing.T) { - executor, _, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, _ := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -54,9 +53,7 @@ func TestStreamSQLUnsharded(t *testing.T) { } func TestStreamSQLSharded(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) cell := "aa" hc := discovery.NewFakeHealthCheck(nil) u := createSandbox(KsTestUnsharded) diff --git a/go/vt/vtgate/executor_test.go b/go/vt/vtgate/executor_test.go index 1dea9ef635b..99a40b2aef4 100644 --- a/go/vt/vtgate/executor_test.go +++ b/go/vt/vtgate/executor_test.go @@ -58,8 +58,7 @@ import ( ) func TestExecutorResultsExceeded(t *testing.T) { - executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbclookup, ctx := createExecutorEnv(t) save := warnMemoryRows warnMemoryRows = 3 @@ -83,8 +82,7 @@ func TestExecutorResultsExceeded(t *testing.T) { } func TestExecutorMaxMemoryRowsExceeded(t *testing.T) { - executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbclookup, ctx := createExecutorEnv(t) save := maxMemoryRows maxMemoryRows = 3 @@ -122,8 +120,7 @@ func TestExecutorMaxMemoryRowsExceeded(t *testing.T) { } func TestExecutorTransactionsNoAutoCommit(t *testing.T) { - executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbclookup, ctx := createExecutorEnv(t) session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", SessionUUID: "suuid"}) @@ -197,8 +194,7 @@ func TestExecutorTransactionsNoAutoCommit(t *testing.T) { } func TestDirectTargetRewrites(t *testing.T) { - executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbclookup, ctx := createExecutorEnv(t) executor.normalize = true @@ -218,8 +214,7 @@ func TestDirectTargetRewrites(t *testing.T) { } func TestExecutorTransactionsAutoCommit(t *testing.T) { - executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbclookup, ctx := createExecutorEnv(t) session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true, SessionUUID: "suuid"}) @@ -272,8 +267,7 @@ func TestExecutorTransactionsAutoCommit(t *testing.T) { } func TestExecutorTransactionsAutoCommitStreaming(t *testing.T) { - executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbclookup, ctx := createExecutorEnv(t) oltpOptions := &querypb.ExecuteOptions{Workload: querypb.ExecuteOptions_OLTP} session := NewSafeSession(&vtgatepb.Session{ @@ -344,8 +338,7 @@ func TestExecutorDeleteMetadata(t *testing.T) { vschemaacl.AuthorizedDDLUsers = "" }() - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) set := "set @@vitess_metadata.app_v1= '1'" @@ -372,8 +365,7 @@ func TestExecutorDeleteMetadata(t *testing.T) { } func TestExecutorAutocommit(t *testing.T) { - executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbclookup, ctx := createExecutorEnv(t) session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) @@ -474,8 +466,7 @@ func TestExecutorAutocommit(t *testing.T) { } func TestExecutorShowColumns(t *testing.T) { - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) session := NewSafeSession(&vtgatepb.Session{TargetString: ""}) @@ -527,8 +518,7 @@ func assertMatchesNoOrder(t *testing.T, expected, got string) { } func TestExecutorShow(t *testing.T) { - executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbclookup, ctx := createExecutorEnv(t) session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) @@ -1087,8 +1077,7 @@ func TestExecutorShow(t *testing.T) { } func TestExecutorShowTargeted(t *testing.T) { - executor, _, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, sbc2, _, ctx := createExecutorEnv(t) session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor/40-60"}) @@ -1115,8 +1104,7 @@ func TestExecutorShowTargeted(t *testing.T) { } func TestExecutorUse(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) session := NewSafeSession(&vtgatepb.Session{Autocommit: true, TargetString: "@primary"}) @@ -1148,8 +1136,7 @@ func TestExecutorUse(t *testing.T) { } func TestExecutorComment(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) stmts := []string{ "/*! SET autocommit=1*/", @@ -1169,8 +1156,7 @@ func TestExecutorComment(t *testing.T) { } func TestExecutorOther(t *testing.T) { - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) type cnts struct { Sbc1Cnt int64 @@ -1262,8 +1248,7 @@ func TestExecutorOther(t *testing.T) { } func TestExecutorDDL(t *testing.T) { - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) @@ -1387,8 +1372,7 @@ func TestExecutorDDL(t *testing.T) { } func TestExecutorDDLFk(t *testing.T) { - executor, _, _, sbc, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbc, ctx := createExecutorEnv(t) mName := "TestExecutorDDLFk" stmts := []string{ @@ -1420,8 +1404,7 @@ func TestExecutorAlterVSchemaKeyspace(t *testing.T) { vschemaacl.AuthorizedDDLUsers = "" }() - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) vschemaUpdates := make(chan *vschemapb.SrvVSchema, 2) @@ -1449,8 +1432,7 @@ func TestExecutorCreateVindexDDL(t *testing.T) { defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) ks := "TestExecutor" vschemaUpdates := make(chan *vschemapb.SrvVSchema, 4) @@ -1520,8 +1502,7 @@ func TestExecutorAddDropVschemaTableDDL(t *testing.T) { defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) ks := KsTestUnsharded vschemaUpdates := make(chan *vschemapb.SrvVSchema, 4) @@ -1569,8 +1550,7 @@ func TestExecutorAddDropVschemaTableDDL(t *testing.T) { } func TestExecutorVindexDDLACL(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) ks := "TestExecutor" session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) @@ -1616,8 +1596,7 @@ func TestExecutorVindexDDLACL(t *testing.T) { } func TestExecutorUnrecognized(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) _, err := executor.Execute(ctx, nil, "TestExecute", NewSafeSession(&vtgatepb.Session{}), "invalid statement", nil) require.Error(t, err, "unrecognized statement: invalid statement'") @@ -1626,9 +1605,7 @@ func TestExecutorUnrecognized(t *testing.T) { // TestVSchemaStats makes sure the building and displaying of the // VSchemaStats works. func TestVSchemaStats(t *testing.T) { - r, _, _, _, _, closer := createExecutorEnv(t) - defer closer() - + r, _, _, _, _ := createExecutorEnv(t) stats := r.VSchemaStats() templ := template.New("") @@ -1650,8 +1627,7 @@ func TestVSchemaStats(t *testing.T) { var pv = querypb.ExecuteOptions_Gen4 func TestGetPlanUnnormalized(t *testing.T) { - r, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + r, _, _, _, ctx := createExecutorEnv(t) emptyvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) unshardedvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) @@ -1741,97 +1717,100 @@ func getPlanCached(t *testing.T, ctx context.Context, e *Executor, vcursor *vcur } func TestGetPlanCacheUnnormalized(t *testing.T) { - r, _, _, _, ctx, closer := createExecutorEnv(t) + t.Run("Cache", func(t *testing.T) { + r, _, _, _, ctx := createExecutorEnv(t) - emptyvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) - query1 := "select * from music_user_map where id = 1" + emptyvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) + query1 := "select * from music_user_map where id = 1" - _, logStats1 := getPlanCached(t, ctx, r, emptyvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, true) - assertCacheSize(t, r.plans, 0) + _, logStats1 := getPlanCached(t, ctx, r, emptyvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, true) + assertCacheSize(t, r.plans, 0) - wantSQL := query1 + " /* comment */" - if logStats1.SQL != wantSQL { - t.Errorf("logstats sql want \"%s\" got \"%s\"", wantSQL, logStats1.SQL) - } + wantSQL := query1 + " /* comment */" + if logStats1.SQL != wantSQL { + t.Errorf("logstats sql want \"%s\" got \"%s\"", wantSQL, logStats1.SQL) + } - _, logStats2 := getPlanCached(t, ctx, r, emptyvc, query1, makeComments(" /* comment 2 */"), map[string]*querypb.BindVariable{}, false) - assertCacheSize(t, r.plans, 1) + _, logStats2 := getPlanCached(t, ctx, r, emptyvc, query1, makeComments(" /* comment 2 */"), map[string]*querypb.BindVariable{}, false) + assertCacheSize(t, r.plans, 1) - wantSQL = query1 + " /* comment 2 */" - if logStats2.SQL != wantSQL { - t.Errorf("logstats sql want \"%s\" got \"%s\"", wantSQL, logStats2.SQL) - } + wantSQL = query1 + " /* comment 2 */" + if logStats2.SQL != wantSQL { + t.Errorf("logstats sql want \"%s\" got \"%s\"", wantSQL, logStats2.SQL) + } + }) - closer() - // Skip cache using directive - r, _, _, _, ctx, closer = createExecutorEnv(t) - defer closer() + t.Run("Skip Cache", func(t *testing.T) { + // Skip cache using directive + r, _, _, _, ctx := createExecutorEnv(t) - unshardedvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) + unshardedvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) - query1 = "insert /*vt+ SKIP_QUERY_PLAN_CACHE=1 */ into user(id) values (1), (2)" - getPlanCached(t, ctx, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) - assertCacheSize(t, r.plans, 0) + query1 := "insert /*vt+ SKIP_QUERY_PLAN_CACHE=1 */ into user(id) values (1), (2)" + getPlanCached(t, ctx, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) + assertCacheSize(t, r.plans, 0) - query1 = "insert into user(id) values (1), (2)" - getPlanCached(t, ctx, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) - assertCacheSize(t, r.plans, 1) + query1 = "insert into user(id) values (1), (2)" + getPlanCached(t, ctx, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) + assertCacheSize(t, r.plans, 1) - // the target string will be resolved and become part of the plan cache key, which adds a new entry - ksIDVc1, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "[deadbeef]"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) - getPlanCached(t, ctx, r, ksIDVc1, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) - assertCacheSize(t, r.plans, 2) + // the target string will be resolved and become part of the plan cache key, which adds a new entry + ksIDVc1, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "[deadbeef]"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) + getPlanCached(t, ctx, r, ksIDVc1, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) + assertCacheSize(t, r.plans, 2) - // the target string will be resolved and become part of the plan cache key, as it's an unsharded ks, it will be the same entry as above - ksIDVc2, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "[beefdead]"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) - getPlanCached(t, ctx, r, ksIDVc2, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) - assertCacheSize(t, r.plans, 2) + // the target string will be resolved and become part of the plan cache key, as it's an unsharded ks, it will be the same entry as above + ksIDVc2, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "[beefdead]"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) + getPlanCached(t, ctx, r, ksIDVc2, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) + assertCacheSize(t, r.plans, 2) + }) } func TestGetPlanCacheNormalized(t *testing.T) { - r, _, _, _, ctx, closer := createExecutorEnv(t) - r.normalize = true - emptyvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) - - query1 := "select * from music_user_map where id = 1" - _, logStats1 := getPlanCached(t, ctx, r, emptyvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, true /* skipQueryPlanCache */) - assertCacheSize(t, r.plans, 0) - wantSQL := "select * from music_user_map where id = :id /* INT64 */ /* comment */" - assert.Equal(t, wantSQL, logStats1.SQL) - - _, logStats2 := getPlanCached(t, ctx, r, emptyvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false /* skipQueryPlanCache */) - assertCacheSize(t, r.plans, 1) - assert.Equal(t, wantSQL, logStats2.SQL) - - closer() - // Skip cache using directive - r, _, _, _, ctx, closer = createExecutorEnv(t) - defer closer() - r.normalize = true - unshardedvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) - - query1 = "insert /*vt+ SKIP_QUERY_PLAN_CACHE=1 */ into user(id) values (1), (2)" - getPlanCached(t, ctx, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) - assertCacheSize(t, r.plans, 0) - - query1 = "insert into user(id) values (1), (2)" - getPlanCached(t, ctx, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) - assertCacheSize(t, r.plans, 1) - - // the target string will be resolved and become part of the plan cache key, which adds a new entry - ksIDVc1, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "[deadbeef]"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) - getPlanCached(t, ctx, r, ksIDVc1, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) - assertCacheSize(t, r.plans, 2) + t.Run("Cache", func(t *testing.T) { + r, _, _, _, ctx := createExecutorEnv(t) + r.normalize = true + emptyvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) + + query1 := "select * from music_user_map where id = 1" + _, logStats1 := getPlanCached(t, ctx, r, emptyvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, true /* skipQueryPlanCache */) + assertCacheSize(t, r.plans, 0) + wantSQL := "select * from music_user_map where id = :id /* INT64 */ /* comment */" + assert.Equal(t, wantSQL, logStats1.SQL) + + _, logStats2 := getPlanCached(t, ctx, r, emptyvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false /* skipQueryPlanCache */) + assertCacheSize(t, r.plans, 1) + assert.Equal(t, wantSQL, logStats2.SQL) + }) - // the target string will be resolved and become part of the plan cache key, as it's an unsharded ks, it will be the same entry as above - ksIDVc2, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "[beefdead]"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) - getPlanCached(t, ctx, r, ksIDVc2, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) - assertCacheSize(t, r.plans, 2) + t.Run("Skip Cache", func(t *testing.T) { + // Skip cache using directive + r, _, _, _, ctx := createExecutorEnv(t) + r.normalize = true + unshardedvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) + + query1 := "insert /*vt+ SKIP_QUERY_PLAN_CACHE=1 */ into user(id) values (1), (2)" + getPlanCached(t, ctx, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) + assertCacheSize(t, r.plans, 0) + + query1 = "insert into user(id) values (1), (2)" + getPlanCached(t, ctx, r, unshardedvc, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) + assertCacheSize(t, r.plans, 1) + + // the target string will be resolved and become part of the plan cache key, which adds a new entry + ksIDVc1, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "[deadbeef]"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) + getPlanCached(t, ctx, r, ksIDVc1, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) + assertCacheSize(t, r.plans, 2) + + // the target string will be resolved and become part of the plan cache key, as it's an unsharded ks, it will be the same entry as above + ksIDVc2, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: KsTestUnsharded + "[beefdead]"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) + getPlanCached(t, ctx, r, ksIDVc2, query1, makeComments(" /* comment */"), map[string]*querypb.BindVariable{}, false) + assertCacheSize(t, r.plans, 2) + }) } func TestGetPlanNormalized(t *testing.T) { - r, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + r, _, _, _, ctx := createExecutorEnv(t) r.normalize = true emptyvc, _ := newVCursorImpl(NewSafeSession(&vtgatepb.Session{TargetString: "@unknown"}), makeComments(""), r, nil, r.vm, r.VSchema(), r.resolver.resolver, nil, false, pv) @@ -1888,8 +1867,7 @@ func TestGetPlanPriority(t *testing.T) { testCase := aTestCase t.Run(testCase.name, func(t *testing.T) { - r, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + r, _, _, _, ctx := createExecutorEnv(t) r.normalize = true logStats := logstats.NewLogStats(ctx, "Test", "", "", nil) @@ -1914,8 +1892,7 @@ func TestGetPlanPriority(t *testing.T) { } func TestPassthroughDDL(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) session := &vtgatepb.Session{ TargetString: "TestExecutor", } @@ -1964,8 +1941,7 @@ func TestPassthroughDDL(t *testing.T) { } func TestParseEmptyTargetSingleKeyspace(t *testing.T) { - r, _, _, _, _, closer := createExecutorEnv(t) - defer closer() + r, _, _, _, _ := createExecutorEnv(t) altVSchema := &vindexes.VSchema{ Keyspaces: map[string]*vindexes.KeyspaceSchema{ @@ -1988,8 +1964,7 @@ func TestParseEmptyTargetSingleKeyspace(t *testing.T) { } func TestParseEmptyTargetMultiKeyspace(t *testing.T) { - r, _, _, _, _, closer := createExecutorEnv(t) - defer closer() + r, _, _, _, _ := createExecutorEnv(t) altVSchema := &vindexes.VSchema{ Keyspaces: map[string]*vindexes.KeyspaceSchema{ @@ -2013,8 +1988,7 @@ func TestParseEmptyTargetMultiKeyspace(t *testing.T) { } func TestParseTargetSingleKeyspace(t *testing.T) { - r, _, _, _, _, closer := createExecutorEnv(t) - defer closer() + r, _, _, _, _ := createExecutorEnv(t) altVSchema := &vindexes.VSchema{ Keyspaces: map[string]*vindexes.KeyspaceSchema{ @@ -2037,8 +2011,7 @@ func TestParseTargetSingleKeyspace(t *testing.T) { } func TestDebugVSchema(t *testing.T) { - executor, _, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, _ := createExecutorEnv(t) resp := httptest.NewRecorder() req, _ := http.NewRequest("GET", "/debug/vschema", nil) @@ -2066,8 +2039,7 @@ func TestExecutorMaxPayloadSizeExceeded(t *testing.T) { warnPayloadSize = saveWarn }() - executor, _, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, _ := createExecutorEnv(t) session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) warningCount := warnings.Counts()["WarnPayloadSizeExceeded"] @@ -2105,8 +2077,7 @@ func TestExecutorMaxPayloadSizeExceeded(t *testing.T) { } func TestOlapSelectDatabase(t *testing.T) { - executor, _, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, _ := createExecutorEnv(t) executor.normalize = true @@ -2124,8 +2095,7 @@ func TestOlapSelectDatabase(t *testing.T) { } func TestExecutorClearsWarnings(t *testing.T) { - executor, _, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, _ := createExecutorEnv(t) session := NewSafeSession(&vtgatepb.Session{ Warnings: []*querypb.QueryWarning{{Code: 234, Message: "oh noes"}}, @@ -2142,8 +2112,7 @@ func TestServingKeyspaces(t *testing.T) { buffer.SetBufferingModeInTestingEnv(false) }() - executor, sbc1, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, _, sbclookup, ctx := createExecutorEnv(t) executor.pv = querypb.ExecuteOptions_Gen4 gw, ok := executor.resolver.resolver.GetGateway().(*TabletGateway) @@ -2186,8 +2155,7 @@ func TestServingKeyspaces(t *testing.T) { } func TestExecutorOtherRead(t *testing.T) { - executor, sbc1, sbc2, sbclookup, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, _ := createExecutorEnv(t) type cnts struct { Sbc1Cnt int64 @@ -2262,8 +2230,7 @@ func TestExecutorOtherRead(t *testing.T) { } func TestExecutorVExplain(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) executor.normalize = true logChan := executor.queryLogger.Subscribe("Test") @@ -2287,8 +2254,7 @@ func TestExecutorVExplain(t *testing.T) { } func TestExecutorOtherAdmin(t *testing.T) { - executor, sbc1, sbc2, sbclookup, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, _ := createExecutorEnv(t) type cnts struct { Sbc1Cnt int64 @@ -2362,8 +2328,7 @@ func TestExecutorOtherAdmin(t *testing.T) { } func TestExecutorSavepointInTx(t *testing.T) { - executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, _ := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("TestExecutorSavepoint") defer executor.queryLogger.Unsubscribe(logChan) @@ -2446,8 +2411,7 @@ func TestExecutorSavepointInTx(t *testing.T) { } func TestExecutorSavepointInTxWithReservedConn(t *testing.T) { - executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, _ := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("TestExecutorSavepoint") defer executor.queryLogger.Unsubscribe(logChan) @@ -2514,8 +2478,7 @@ func TestExecutorSavepointInTxWithReservedConn(t *testing.T) { } func TestExecutorSavepointWithoutTx(t *testing.T) { - executor, sbc1, sbc2, _, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, _ := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("TestExecutorSavepoint") defer executor.queryLogger.Unsubscribe(logChan) @@ -2559,8 +2522,7 @@ func TestExecutorSavepointWithoutTx(t *testing.T) { } func TestExecutorCallProc(t *testing.T) { - executor, sbc1, sbc2, sbcUnsharded, _, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbcUnsharded, _ := createExecutorEnv(t) type cnts struct { Sbc1Cnt int64 @@ -2625,8 +2587,7 @@ func TestExecutorCallProc(t *testing.T) { } func TestExecutorTempTable(t *testing.T) { - executor, _, _, sbcUnsharded, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbcUnsharded, ctx := createExecutorEnv(t) executor.warnShardedOnly = true creatQuery := "create temporary table temp_t(id bigint primary key)" @@ -2645,8 +2606,7 @@ func TestExecutorTempTable(t *testing.T) { } func TestExecutorShowVitessMigrations(t *testing.T) { - executor, sbc1, sbc2, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, _, ctx := createExecutorEnv(t) showQuery := "show vitess_migrations" session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) @@ -2657,8 +2617,7 @@ func TestExecutorShowVitessMigrations(t *testing.T) { } func TestExecutorDescHash(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) showQuery := "desc hash_index" session := NewSafeSession(&vtgatepb.Session{TargetString: "TestExecutor"}) @@ -2667,8 +2626,7 @@ func TestExecutorDescHash(t *testing.T) { } func TestExecutorVExplainQueries(t *testing.T) { - executor, _, _, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbclookup, ctx := createExecutorEnv(t) session := NewAutocommitSession(&vtgatepb.Session{}) @@ -2694,8 +2652,7 @@ func TestExecutorVExplainQueries(t *testing.T) { } func TestExecutorStartTxnStmt(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) session := NewAutocommitSession(&vtgatepb.Session{}) @@ -2741,8 +2698,7 @@ func TestExecutorStartTxnStmt(t *testing.T) { } func TestExecutorPrepareExecute(t *testing.T) { - executor, _, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, _ := createExecutorEnv(t) executor.normalize = true session := NewAutocommitSession(&vtgatepb.Session{}) @@ -2782,8 +2738,7 @@ func TestExecutorPrepareExecute(t *testing.T) { } func TestExecutorTruncateErrors(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) save := truncateErrorLen truncateErrorLen = 32 @@ -2805,8 +2760,7 @@ func TestExecutorTruncateErrors(t *testing.T) { } func TestExecutorFlushStmt(t *testing.T) { - executor, _, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, _ := createExecutorEnv(t) tcs := []struct { targetStr string @@ -2863,8 +2817,7 @@ func TestExecutorFlushStmt(t *testing.T) { // TestExecutorKillStmt tests the kill statements on executor. func TestExecutorKillStmt(t *testing.T) { - executor, _, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, _ := createExecutorEnv(t) tcs := []struct { errStr string diff --git a/go/vt/vtgate/executor_vschema_ddl_test.go b/go/vt/vtgate/executor_vschema_ddl_test.go index 456fbb4b035..a7708f3e255 100644 --- a/go/vt/vtgate/executor_vschema_ddl_test.go +++ b/go/vt/vtgate/executor_vschema_ddl_test.go @@ -133,8 +133,7 @@ func TestPlanExecutorAlterVSchemaKeyspace(t *testing.T) { defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary", Autocommit: true}) vschemaUpdates := make(chan *vschemapb.SrvVSchema, 2) @@ -162,8 +161,7 @@ func TestPlanExecutorCreateVindexDDL(t *testing.T) { defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) ks := "TestExecutor" vschemaUpdates := make(chan *vschemapb.SrvVSchema, 4) @@ -205,8 +203,7 @@ func TestPlanExecutorDropVindexDDL(t *testing.T) { defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) ks := "TestExecutor" vschemaUpdates := make(chan *vschemapb.SrvVSchema, 4) @@ -275,8 +272,7 @@ func TestPlanExecutorAddDropVschemaTableDDL(t *testing.T) { defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) ks := KsTestUnsharded vschemaUpdates := make(chan *vschemapb.SrvVSchema, 4) @@ -333,8 +329,7 @@ func TestExecutorAddSequenceDDL(t *testing.T) { defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) ks := KsTestUnsharded vschema := executor.vm.GetCurrentSrvVschema() @@ -394,8 +389,7 @@ func TestExecutorAddDropVindexDDL(t *testing.T) { defer func() { vschemaacl.AuthorizedDDLUsers = "" }() - executor, sbc1, sbc2, sbclookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, sbc1, sbc2, sbclookup, ctx := createExecutorEnv(t) ks := "TestExecutor" session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) vschemaUpdates := make(chan *vschemapb.SrvVSchema, 4) @@ -638,8 +632,7 @@ func TestExecutorAddDropVindexDDL(t *testing.T) { func TestPlanExecutorVindexDDLACL(t *testing.T) { // t.Skip("not yet planned") - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) ks := "TestExecutor" session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) diff --git a/go/vt/vtgate/executor_vstream_test.go b/go/vt/vtgate/executor_vstream_test.go index b146efb27c5..5466e9e8f3f 100644 --- a/go/vt/vtgate/executor_vstream_test.go +++ b/go/vt/vtgate/executor_vstream_test.go @@ -37,8 +37,7 @@ import ( // TestVStreamSQLUnsharded tests the experimental 'vstream * from' vtgate olap query func TestVStreamSQLUnsharded(t *testing.T) { t.Skip("this test is failing due to races") // FIXME - executor, _, _, sbcLookup, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, sbcLookup, ctx := createExecutorEnv(t) logChan := executor.queryLogger.Subscribe("Test") defer executor.queryLogger.Unsubscribe(logChan) send1 := []*binlogdatapb.VEvent{ diff --git a/go/vt/vtgate/legacy_scatter_conn_test.go b/go/vt/vtgate/legacy_scatter_conn_test.go index ede4e46303c..c4dcd3746fa 100644 --- a/go/vt/vtgate/legacy_scatter_conn_test.go +++ b/go/vt/vtgate/legacy_scatter_conn_test.go @@ -44,9 +44,7 @@ import ( // This file uses the sandbox_test framework. func TestLegacyExecuteFailOnAutocommit(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) createSandbox("TestExecuteFailOnAutocommit") hc := discovery.NewFakeHealthCheck(nil) @@ -111,9 +109,7 @@ func TestLegacyExecuteFailOnAutocommit(t *testing.T) { func TestScatterConnExecuteMulti(t *testing.T) { testScatterConnGeneric(t, "TestScatterConnExecuteMultiShard", func(sc *ScatterConn, shards []string) (*sqltypes.Result, error) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) res := srvtopo.NewResolver(newSandboxForCells(ctx, []string{"aa"}), sc.gateway, "aa") rss, err := res.ResolveDestination(ctx, "TestScatterConnExecuteMultiShard", topodatapb.TabletType_REPLICA, key.DestinationShards(shards)) if err != nil { @@ -135,9 +131,7 @@ func TestScatterConnExecuteMulti(t *testing.T) { func TestScatterConnStreamExecuteMulti(t *testing.T) { testScatterConnGeneric(t, "TestScatterConnStreamExecuteMulti", func(sc *ScatterConn, shards []string) (*sqltypes.Result, error) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) res := srvtopo.NewResolver(newSandboxForCells(ctx, []string{"aa"}), sc.gateway, "aa") rss, err := res.ResolveDestination(ctx, "TestScatterConnStreamExecuteMulti", topodatapb.TabletType_REPLICA, key.DestinationShards(shards)) if err != nil { @@ -165,9 +159,7 @@ func verifyScatterConnError(t *testing.T, err error, wantErr string, wantCode vt } func testScatterConnGeneric(t *testing.T, name string, f func(sc *ScatterConn, shards []string) (*sqltypes.Result, error)) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) hc := discovery.NewFakeHealthCheck(nil) @@ -273,9 +265,7 @@ func testScatterConnGeneric(t *testing.T, name string, f func(sc *ScatterConn, s } func TestMaxMemoryRows(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) save := maxMemoryRows maxMemoryRows = 3 @@ -332,9 +322,7 @@ func TestMaxMemoryRows(t *testing.T) { } func TestLegaceHealthCheckFailsOnReservedConnections(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) keyspace := "keyspace" createSandbox(keyspace) hc := discovery.NewFakeHealthCheck(nil) @@ -384,9 +372,7 @@ func executeOnShardsReturnsErr(t *testing.T, ctx context.Context, res *srvtopo.R } func TestMultiExecs(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) createSandbox("TestMultiExecs") hc := discovery.NewFakeHealthCheck(nil) sc := newTestScatterConn(ctx, hc, newSandboxForCells(ctx, []string{"aa"}), "aa") @@ -483,9 +469,7 @@ func TestMultiExecs(t *testing.T) { } func TestScatterConnSingleDB(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) createSandbox("TestScatterConnSingleDB") hc := discovery.NewFakeHealthCheck(nil) @@ -579,9 +563,7 @@ func TestAppendResult(t *testing.T) { } func TestReservePrequeries(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) keyspace := "keyspace" createSandbox(keyspace) hc := discovery.NewFakeHealthCheck(nil) diff --git a/go/vt/vtgate/plugin_mysql_server_test.go b/go/vt/vtgate/plugin_mysql_server_test.go index a63872d4791..1b161dfb171 100644 --- a/go/vt/vtgate/plugin_mysql_server_test.go +++ b/go/vt/vtgate/plugin_mysql_server_test.go @@ -274,9 +274,7 @@ func TestInitTLSConfigWithServerCA(t *testing.T) { func testInitTLSConfig(t *testing.T, serverCA bool) { // Create the certs. - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) root := t.TempDir() tlstest.CreateCA(root) @@ -308,8 +306,7 @@ func testInitTLSConfig(t *testing.T, serverCA bool) { // TestKillMethods test the mysql plugin for kill method calls. func TestKillMethods(t *testing.T) { - executor, _, _, _, _, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, _ := createExecutorEnv(t) vh := newVtgateHandler(&VTGate{executor: executor}) // connection does not exist diff --git a/go/vt/vtgate/queryz_test.go b/go/vt/vtgate/queryz_test.go index ff864393899..83fd064df7d 100644 --- a/go/vt/vtgate/queryz_test.go +++ b/go/vt/vtgate/queryz_test.go @@ -36,8 +36,7 @@ import ( ) func TestQueryzHandler(t *testing.T) { - executor, _, _, _, ctx, closer := createExecutorEnv(t) - defer closer() + executor, _, _, _, ctx := createExecutorEnv(t) resp := httptest.NewRecorder() req, _ := http.NewRequest("GET", "/schemaz", nil) diff --git a/go/vt/vtgate/scatter_conn_test.go b/go/vt/vtgate/scatter_conn_test.go index a441227c6c2..6e57c10bbbd 100644 --- a/go/vt/vtgate/scatter_conn_test.go +++ b/go/vt/vtgate/scatter_conn_test.go @@ -17,7 +17,6 @@ limitations under the License. package vtgate import ( - "context" "testing" "vitess.io/vitess/go/mysql/sqlerror" @@ -43,9 +42,7 @@ import ( // This file uses the sandbox_test framework. func TestExecuteFailOnAutocommit(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) createSandbox("TestExecuteFailOnAutocommit") hc := discovery.NewFakeHealthCheck(nil) @@ -109,9 +106,7 @@ func TestExecuteFailOnAutocommit(t *testing.T) { } func TestReservedOnMultiReplica(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) keyspace := "keyspace" createSandbox(keyspace) @@ -137,9 +132,7 @@ func TestReservedOnMultiReplica(t *testing.T) { } func TestReservedBeginTableDriven(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) type testAction struct { transaction, reserved bool @@ -299,9 +292,7 @@ func TestReservedBeginTableDriven(t *testing.T) { } func TestReservedConnFail(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) keyspace := "keyspace" createSandbox(keyspace) diff --git a/go/vt/vtgate/tabletgateway_flaky_test.go b/go/vt/vtgate/tabletgateway_flaky_test.go index df96ed0ab1d..f625b5599cd 100644 --- a/go/vt/vtgate/tabletgateway_flaky_test.go +++ b/go/vt/vtgate/tabletgateway_flaky_test.go @@ -17,7 +17,6 @@ limitations under the License. package vtgate import ( - "context" "testing" "time" @@ -37,9 +36,7 @@ import ( // TestGatewayBufferingWhenPrimarySwitchesServingState is used to test that the buffering mechanism buffers the queries when a primary goes to a non serving state and // stops buffering when the primary is healthy again func TestGatewayBufferingWhenPrimarySwitchesServingState(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) buffer.SetBufferingModeInTestingEnv(true) defer func() { @@ -124,9 +121,7 @@ func TestGatewayBufferingWhenPrimarySwitchesServingState(t *testing.T) { // TestGatewayBufferingWhileReparenting is used to test that the buffering mechanism buffers the queries when a PRS happens // the healthchecks that happen during a PRS are simulated in this test func TestGatewayBufferingWhileReparenting(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) buffer.SetBufferingModeInTestingEnv(true) defer func() { @@ -257,9 +252,7 @@ outer: // This is inconsistent and we want to fail properly. This scenario used to panic since no error and no results were // returned. func TestInconsistentStateDetectedBuffering(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) buffer.SetBufferingModeInTestingEnv(true) defer func() { diff --git a/go/vt/vtgate/tabletgateway_test.go b/go/vt/vtgate/tabletgateway_test.go index 9a41029f7c8..e88f3cbf32b 100644 --- a/go/vt/vtgate/tabletgateway_test.go +++ b/go/vt/vtgate/tabletgateway_test.go @@ -85,9 +85,7 @@ func TestTabletGatewayBeginExecute(t *testing.T) { } func TestTabletGatewayShuffleTablets(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) hc := discovery.NewFakeHealthCheck(nil) ts := &fakeTopoServer{} @@ -149,9 +147,7 @@ func TestTabletGatewayShuffleTablets(t *testing.T) { } func TestTabletGatewayReplicaTransactionError(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) keyspace := "ks" shard := "0" @@ -177,9 +173,7 @@ func TestTabletGatewayReplicaTransactionError(t *testing.T) { func testTabletGatewayGeneric(t *testing.T, f func(tg *TabletGateway, target *querypb.Target) error) { t.Helper() - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) keyspace := "ks" shard := "0" @@ -248,9 +242,7 @@ func testTabletGatewayGeneric(t *testing.T, f func(tg *TabletGateway, target *qu func testTabletGatewayTransact(t *testing.T, f func(tg *TabletGateway, target *querypb.Target) error) { t.Helper() - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) keyspace := "ks" shard := "0" diff --git a/go/vt/vtgate/tx_conn_test.go b/go/vt/vtgate/tx_conn_test.go index bb531a74c0e..3fc141c64ac 100644 --- a/go/vt/vtgate/tx_conn_test.go +++ b/go/vt/vtgate/tx_conn_test.go @@ -43,9 +43,7 @@ var queries = []*querypb.BoundQuery{{Sql: "query1"}} var twoQueries = []*querypb.BoundQuery{{Sql: "query1"}, {Sql: "query1"}} func TestTxConnBegin(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, _, rss0, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") session := &vtgatepb.Session{} @@ -67,9 +65,7 @@ func TestTxConnBegin(t *testing.T) { } func TestTxConnCommitFailure(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, rss0, rss1, rss01 := newTestTxConnEnv(t, ctx, "TestTxConn") sc.txConn.mode = vtgatepb.TransactionMode_MULTI @@ -128,9 +124,7 @@ func TestTxConnCommitFailure(t *testing.T) { } func TestTxConnCommitSuccess(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TestTxConn") sc.txConn.mode = vtgatepb.TransactionMode_MULTI @@ -183,9 +177,7 @@ func TestTxConnCommitSuccess(t *testing.T) { } func TestTxConnReservedCommitSuccess(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TestTxConn") sc.txConn.mode = vtgatepb.TransactionMode_MULTI @@ -269,9 +261,7 @@ func TestTxConnReservedCommitSuccess(t *testing.T) { } func TestTxConnReservedOn2ShardTxOn1ShardAndCommit(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) keyspace := "TestTxConn" sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, ctx, keyspace) @@ -366,9 +356,7 @@ func TestTxConnReservedOn2ShardTxOn1ShardAndCommit(t *testing.T) { } func TestTxConnReservedOn2ShardTxOn1ShardAndRollback(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) keyspace := "TestTxConn" sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, ctx, keyspace) @@ -463,9 +451,7 @@ func TestTxConnReservedOn2ShardTxOn1ShardAndRollback(t *testing.T) { } func TestTxConnCommitOrderFailure1(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, ctx, "TestTxConn") sc.txConn.mode = vtgatepb.TransactionMode_MULTI @@ -498,9 +484,7 @@ func TestTxConnCommitOrderFailure1(t *testing.T) { } func TestTxConnCommitOrderFailure2(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, ctx, "TestTxConn") sc.txConn.mode = vtgatepb.TransactionMode_MULTI @@ -534,9 +518,7 @@ func TestTxConnCommitOrderFailure2(t *testing.T) { } func TestTxConnCommitOrderFailure3(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, ctx, "TestTxConn") sc.txConn.mode = vtgatepb.TransactionMode_MULTI @@ -578,9 +560,7 @@ func TestTxConnCommitOrderFailure3(t *testing.T) { } func TestTxConnCommitOrderSuccess(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, ctx, "TestTxConn") sc.txConn.mode = vtgatepb.TransactionMode_MULTI @@ -678,9 +658,7 @@ func TestTxConnCommitOrderSuccess(t *testing.T) { } func TestTxConnReservedCommitOrderSuccess(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, ctx, "TestTxConn") sc.txConn.mode = vtgatepb.TransactionMode_MULTI @@ -823,9 +801,7 @@ func TestTxConnReservedCommitOrderSuccess(t *testing.T) { } func TestTxConnCommit2PC(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PC") @@ -843,9 +819,7 @@ func TestTxConnCommit2PC(t *testing.T) { } func TestTxConnCommit2PCOneParticipant(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, _, rss0, _, _ := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PCOneParticipant") session := NewSafeSession(&vtgatepb.Session{InTransaction: true}) @@ -857,9 +831,7 @@ func TestTxConnCommit2PCOneParticipant(t *testing.T) { } func TestTxConnCommit2PCCreateTransactionFail(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, rss0, rss1, _ := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PCCreateTransactionFail") @@ -883,9 +855,7 @@ func TestTxConnCommit2PCCreateTransactionFail(t *testing.T) { } func TestTxConnCommit2PCPrepareFail(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PCPrepareFail") @@ -907,9 +877,7 @@ func TestTxConnCommit2PCPrepareFail(t *testing.T) { } func TestTxConnCommit2PCStartCommitFail(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PCStartCommitFail") @@ -931,9 +899,7 @@ func TestTxConnCommit2PCStartCommitFail(t *testing.T) { } func TestTxConnCommit2PCCommitPreparedFail(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PCCommitPreparedFail") @@ -955,9 +921,7 @@ func TestTxConnCommit2PCCommitPreparedFail(t *testing.T) { } func TestTxConnCommit2PCConcludeTransactionFail(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TestTxConnCommit2PCConcludeTransactionFail") @@ -979,9 +943,7 @@ func TestTxConnCommit2PCConcludeTransactionFail(t *testing.T) { } func TestTxConnRollback(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TxConnRollback") @@ -997,9 +959,7 @@ func TestTxConnRollback(t *testing.T) { } func TestTxConnReservedRollback(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, rss0, _, rss01 := newTestTxConnEnv(t, ctx, "TxConnReservedRollback") @@ -1036,9 +996,7 @@ func TestTxConnReservedRollback(t *testing.T) { } func TestTxConnReservedRollbackFailure(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, rss0, rss1, rss01 := newTestTxConnEnv(t, ctx, "TxConnReservedRollback") @@ -1069,9 +1027,7 @@ func TestTxConnReservedRollbackFailure(t *testing.T) { } func TestTxConnResolveOnPrepare(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") @@ -1094,9 +1050,7 @@ func TestTxConnResolveOnPrepare(t *testing.T) { } func TestTxConnResolveOnRollback(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") @@ -1119,9 +1073,7 @@ func TestTxConnResolveOnRollback(t *testing.T) { } func TestTxConnResolveOnCommit(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") @@ -1144,9 +1096,7 @@ func TestTxConnResolveOnCommit(t *testing.T) { } func TestTxConnResolveInvalidDTID(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, _, _, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") @@ -1156,9 +1106,7 @@ func TestTxConnResolveInvalidDTID(t *testing.T) { } func TestTxConnResolveReadTransactionFail(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, _, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") @@ -1171,9 +1119,7 @@ func TestTxConnResolveReadTransactionFail(t *testing.T) { } func TestTxConnResolveInternalError(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, _, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") @@ -1194,9 +1140,7 @@ func TestTxConnResolveInternalError(t *testing.T) { } func TestTxConnResolveSetRollbackFail(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") @@ -1222,9 +1166,7 @@ func TestTxConnResolveSetRollbackFail(t *testing.T) { } func TestTxConnResolveRollbackPreparedFail(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") @@ -1250,9 +1192,7 @@ func TestTxConnResolveRollbackPreparedFail(t *testing.T) { } func TestTxConnResolveCommitPreparedFail(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") @@ -1278,9 +1218,7 @@ func TestTxConnResolveCommitPreparedFail(t *testing.T) { } func TestTxConnResolveConcludeTransactionFail(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, sbc0, sbc1, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") @@ -1306,9 +1244,7 @@ func TestTxConnResolveConcludeTransactionFail(t *testing.T) { } func TestTxConnMultiGoSessions(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) txc := &TxConn{} @@ -1377,9 +1313,7 @@ func TestTxConnMultiGoTargets(t *testing.T) { } func TestTxConnAccessModeReset(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) sc, _, _, _, _, _ := newTestTxConnEnv(t, ctx, "TestTxConn") diff --git a/go/vt/vtgate/vstream_manager_test.go b/go/vt/vtgate/vstream_manager_test.go index 5e51f0640c2..13aae235d44 100644 --- a/go/vt/vtgate/vstream_manager_test.go +++ b/go/vt/vtgate/vstream_manager_test.go @@ -1137,9 +1137,7 @@ func TestResolveVStreamParams(t *testing.T) { } func TestVStreamIdleHeartbeat(t *testing.T) { - defer utils.EnsureNoLeaks(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := utils.LeakCheckContext(t) cell := "aa" ks := "TestVStream" diff --git a/go/vt/vtgate/vtgate_test.go b/go/vt/vtgate/vtgate_test.go index 244bfbf88b8..6f21158d7bb 100644 --- a/go/vt/vtgate/vtgate_test.go +++ b/go/vt/vtgate/vtgate_test.go @@ -45,8 +45,7 @@ var executeOptions = &querypb.ExecuteOptions{ } func TestVTGateExecute(t *testing.T) { - vtg, sbc, ctx, closer := createVtgateEnv(t) - defer closer() + vtg, sbc, ctx := createVtgateEnv(t) counts := vtg.timings.Timings.Counts() _, qr, err := vtg.Execute( @@ -85,8 +84,7 @@ func TestVTGateExecute(t *testing.T) { } func TestVTGateExecuteError(t *testing.T) { - vtg, _, ctx, closer := createVtgateEnv(t) - defer closer() + vtg, _, ctx := createVtgateEnv(t) counts := vtg.timings.Timings.Counts() @@ -116,8 +114,7 @@ func TestVTGateExecuteError(t *testing.T) { } func TestVTGatePrepare(t *testing.T) { - vtg, sbc, ctx, closer := createVtgateEnv(t) - defer closer() + vtg, sbc, ctx := createVtgateEnv(t) counts := vtg.timings.Timings.Counts() _, qr, err := vtg.Prepare( @@ -154,8 +151,7 @@ func TestVTGatePrepare(t *testing.T) { } func TestVTGatePrepareError(t *testing.T) { - vtg, _, ctx, closer := createVtgateEnv(t) - defer closer() + vtg, _, ctx := createVtgateEnv(t) counts := errorCounts.Counts() @@ -184,8 +180,7 @@ func TestVTGatePrepareError(t *testing.T) { } func TestVTGateExecuteWithKeyspaceShard(t *testing.T) { - vtg, _, ctx, closer := createVtgateEnv(t) - defer closer() + vtg, _, ctx := createVtgateEnv(t) // Valid keyspace. _, qr, err := vtg.Execute( @@ -247,8 +242,7 @@ func TestVTGateExecuteWithKeyspaceShard(t *testing.T) { } func TestVTGateStreamExecute(t *testing.T) { - vtg, sbc, ctx, closer := createVtgateEnv(t) - defer closer() + vtg, sbc, ctx := createVtgateEnv(t) var qrs []*sqltypes.Result _, err := vtg.StreamExecute( @@ -278,8 +272,7 @@ func TestVTGateStreamExecute(t *testing.T) { } func TestVTGateBindVarError(t *testing.T) { - vtg, _, ctx, closer := createVtgateEnv(t) - defer closer() + vtg, _, ctx := createVtgateEnv(t) session := &vtgatepb.Session{} bindVars := map[string]*querypb.BindVariable{ @@ -380,8 +373,7 @@ func testErrorPropagation(t *testing.T, ctx context.Context, vtg *VTGate, sbcs [ // tablet and a rdonly tablet because we don't control the routing of // Commit. func TestErrorPropagation(t *testing.T) { - vtg, sbc, ctx, closer := createVtgateEnv(t) - defer closer() + vtg, sbc, ctx := createVtgateEnv(t) sbcs := []*sandboxconn.SandboxConn{ sbc, @@ -463,8 +455,7 @@ func TestErrorPropagation(t *testing.T) { // This test makes sure that if we start a transaction and hit a critical // error, a rollback is issued. func TestErrorIssuesRollback(t *testing.T) { - vtg, sbc, ctx, closer := createVtgateEnv(t) - defer closer() + vtg, sbc, ctx := createVtgateEnv(t) // Start a transaction, send one statement. // Simulate an error that should trigger a rollback: @@ -594,8 +585,7 @@ var shardedVSchemaUnknownParams = ` ` func TestMultiInternalSavepointVtGate(t *testing.T) { - vtg, _, ctx, closer := createVtgateEnv(t) - defer closer() + vtg, _, ctx := createVtgateEnv(t) const customKeyspace = "CustomSharding" s := createSandbox(customKeyspace) @@ -693,8 +683,7 @@ func TestMultiInternalSavepointVtGate(t *testing.T) { } func TestVSchemaVindexUnknownParams(t *testing.T) { - vtg, _, _, closer := createVtgateEnv(t) - defer closer() + vtg, _, _ := createVtgateEnv(t) const customKeyspace = "CustomSharding" s := createSandbox(customKeyspace) @@ -726,15 +715,15 @@ func TestVSchemaVindexUnknownParams(t *testing.T) { require.Equal(t, int64(0), unknownParams) } -func createVtgateEnv(t testing.TB) (*VTGate, *sandboxconn.SandboxConn, context.Context, func()) { +func createVtgateEnv(t testing.TB) (*VTGate, *sandboxconn.SandboxConn, context.Context) { cell := "aa" sb := createSandbox(KsTestSharded) sb.ShardSpec = "-" - executor, _, _, sbc, ctx, closer := createExecutorEnv(t) + executor, _, _, sbc, ctx := createExecutorEnv(t) executor.normalize = normalizeQueries vsm := newVStreamManager(executor.resolver.resolver, executor.serv, cell) vtg := newVTGate(executor, executor.resolver, vsm, nil, executor.scatterConn.gateway) - return vtg, sbc, ctx, closer + return vtg, sbc, ctx } From fbcaba4968a9ea09c04014c713713654cbdad877 Mon Sep 17 00:00:00 2001 From: Vicent Marti Date: Fri, 25 Aug 2023 12:58:38 +0200 Subject: [PATCH 15/17] utils: documentat the two new noleak APIs Signed-off-by: Vicent Marti --- go/test/utils/noleak.go | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/go/test/utils/noleak.go b/go/test/utils/noleak.go index a492ab670f5..697702b8861 100644 --- a/go/test/utils/noleak.go +++ b/go/test/utils/noleak.go @@ -30,6 +30,9 @@ import ( "vitess.io/vitess/go/vt/log" ) +// LeakCheckContext returns a Context that will be automatically cancelled at the end +// of this test. If the test has finished successfully, it will be checked for goroutine +// leaks after context cancellation. func LeakCheckContext(t testing.TB) context.Context { ctx, cancel := context.WithCancel(context.Background()) t.Cleanup(func() { @@ -39,6 +42,8 @@ func LeakCheckContext(t testing.TB) context.Context { return ctx } +// LeakCheckContextTimeout behaves like LeakCheckContext but the returned Context will +// be cancelled after `timeout`, or after the test finishes, whichever happens first. func LeakCheckContextTimeout(t testing.TB, timeout time.Duration) context.Context { ctx, cancel := context.WithTimeout(context.Background(), timeout) t.Cleanup(func() { From 488a15598f3e960f558fb45e60429ff42e7c5294 Mon Sep 17 00:00:00 2001 From: Dirkjan Bussink Date: Mon, 28 Aug 2023 10:41:51 +0200 Subject: [PATCH 16/17] Fix rebased tests Signed-off-by: Dirkjan Bussink --- go/vt/vtctl/grpcvtctldserver/server_test.go | 26 ++++++++++----------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/go/vt/vtctl/grpcvtctldserver/server_test.go b/go/vt/vtctl/grpcvtctldserver/server_test.go index b8cba1d64bd..e61f3eba607 100644 --- a/go/vt/vtctl/grpcvtctldserver/server_test.go +++ b/go/vt/vtctl/grpcvtctldserver/server_test.go @@ -1548,11 +1548,10 @@ func TestCleanupSchemaMigration(t *testing.T) { } for _, test := range tests { - test := test t.Run(test.name, func(t *testing.T) { - - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") testutil.AddTablets(ctx, t, ts, &testutil.AddTabletOptions{ AlsoSetShardPrimary: true, @@ -9085,11 +9084,10 @@ func TestRetrySchemaMigration(t *testing.T) { } for _, test := range tests { - test := test t.Run(test.name, func(t *testing.T) { - - ctx := context.Background() - ts := memorytopo.NewServer("zone1") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ts := memorytopo.NewServer(ctx, "zone1") testutil.AddTablets(ctx, t, ts, &testutil.AddTabletOptions{ AlsoSetShardPrimary: true, @@ -9307,7 +9305,7 @@ func TestSetShardIsPrimaryServing(t *testing.T) { name string ctx context.Context ts *topo.Server - setup func(*testing.T, *testcase) + setup func(*testing.T, *testcase) context.Context teardown func(*testing.T, *testcase) req *vtctldatapb.SetShardIsPrimaryServingRequest expected *vtctldatapb.SetShardIsPrimaryServingResponse @@ -9317,7 +9315,7 @@ func TestSetShardIsPrimaryServing(t *testing.T) { tests := []*testcase{ { name: "ok", - setup: func(t *testing.T, tt *testcase) { + setup: func(t *testing.T, tt *testcase) context.Context { tt.ctx = ctx tt.ts = memorytopo.NewServer(ctx, "zone1") testutil.AddShards(tt.ctx, t, tt.ts, &vtctldatapb.Shard{ @@ -9325,6 +9323,7 @@ func TestSetShardIsPrimaryServing(t *testing.T) { Name: "-", Shard: &topodatapb.Shard{}, }) + return tt.ctx }, req: &vtctldatapb.SetShardIsPrimaryServingRequest{ Keyspace: "testkeyspace", @@ -9339,7 +9338,7 @@ func TestSetShardIsPrimaryServing(t *testing.T) { }, { name: "lock error", - setup: func(t *testing.T, tt *testcase) { + setup: func(t *testing.T, tt *testcase) context.Context { var cancel func() tt.ctx, cancel = context.WithTimeout(ctx, time.Millisecond*50) tt.ts = memorytopo.NewServer(ctx, "zone1") @@ -9349,7 +9348,7 @@ func TestSetShardIsPrimaryServing(t *testing.T) { Shard: &topodatapb.Shard{}, }) - _, unlock, err := tt.ts.LockKeyspace(tt.ctx, "testkeyspace", "test lock") + lctx, unlock, err := tt.ts.LockKeyspace(tt.ctx, "testkeyspace", "test lock") require.NoError(t, err) tt.teardown = func(t *testing.T, tt *testcase) { var err error @@ -9357,6 +9356,7 @@ func TestSetShardIsPrimaryServing(t *testing.T) { assert.NoError(t, err) cancel() } + return lctx }, req: &vtctldatapb.SetShardIsPrimaryServingRequest{ Keyspace: "testkeyspace", @@ -9371,7 +9371,7 @@ func TestSetShardIsPrimaryServing(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { if tt.setup != nil { - tt.setup(t, tt) + tt.ctx = tt.setup(t, tt) } if tt.teardown != nil { defer tt.teardown(t, tt) From 27ab83edba5888e2bff809dfdd2869a092c5771b Mon Sep 17 00:00:00 2001 From: Dirkjan Bussink Date: Mon, 28 Aug 2023 12:30:16 +0200 Subject: [PATCH 17/17] Allow additional throttler for leak check Signed-off-by: Dirkjan Bussink --- go/test/utils/noleak.go | 1 + 1 file changed, 1 insertion(+) diff --git a/go/test/utils/noleak.go b/go/test/utils/noleak.go index 697702b8861..21207b6b989 100644 --- a/go/test/utils/noleak.go +++ b/go/test/utils/noleak.go @@ -89,6 +89,7 @@ func ensureNoGoroutines() error { goleak.IgnoreTopFunction("github.com/patrickmn/go-cache.(*janitor).Run"), goleak.IgnoreTopFunction("vitess.io/vitess/go/vt/logutil.(*ThrottledLogger).log.func1"), goleak.IgnoreTopFunction("vitess.io/vitess/go/vt/vttablet/tabletserver/throttle.initThrottleTicker.func1.1"), + goleak.IgnoreTopFunction("vitess.io/vitess/go/vt/vttablet/tabletserver/throttle.NewBackgroundClient.initThrottleTicker.func1.1"), goleak.IgnoreTopFunction("testing.tRunner.func1"), }