diff --git a/go/vt/discovery/fake_healthcheck.go b/go/vt/discovery/fake_healthcheck.go index 6ae7ee105c2..131f841acc8 100644 --- a/go/vt/discovery/fake_healthcheck.go +++ b/go/vt/discovery/fake_healthcheck.go @@ -192,7 +192,7 @@ func (fhc *FakeHealthCheck) Unsubscribe(c chan *TabletHealth) { } // GetLoadTabletsTrigger is not implemented. -func (fhc *FakeHealthCheck) GetLoadTabletsTrigger() chan struct{} { +func (fhc *FakeHealthCheck) GetLoadTabletsTrigger() chan topo.KeyspaceShard { return nil } diff --git a/go/vt/discovery/healthcheck.go b/go/vt/discovery/healthcheck.go index 5734749b167..3964df83248 100644 --- a/go/vt/discovery/healthcheck.go +++ b/go/vt/discovery/healthcheck.go @@ -247,7 +247,7 @@ type HealthCheck interface { Unsubscribe(c chan *TabletHealth) // GetLoadTabletsTrigger returns a channel that is used to inform when to load tablets. - GetLoadTabletsTrigger() chan struct{} + GetLoadTabletsTrigger() chan topo.KeyspaceShard } var _ HealthCheck = (*HealthCheckImpl)(nil) @@ -297,8 +297,8 @@ type HealthCheckImpl struct { subMu sync.Mutex // subscribers subscribers map[chan *TabletHealth]struct{} - // loadTablets trigger is used to immediately load a new primary tablet when the current one has been demoted - loadTabletsTrigger chan struct{} + // loadTabletsTrigger is used to immediately load information about tablets of a specific shard. + loadTabletsTrigger chan topo.KeyspaceShard } // NewVTGateHealthCheckFilters returns healthcheck filters for vtgate. @@ -363,7 +363,7 @@ func NewHealthCheck(ctx context.Context, retryDelay, healthCheckTimeout time.Dur healthy: make(map[KeyspaceShardTabletType][]*TabletHealth), subscribers: make(map[chan *TabletHealth]struct{}), cellAliases: make(map[string]string), - loadTabletsTrigger: make(chan struct{}, 1), + loadTabletsTrigger: make(chan topo.KeyspaceShard, 1024), } var topoWatchers []*TopologyWatcher cells := strings.Split(cellsToWatch, ",") @@ -535,18 +535,21 @@ func (hc *HealthCheckImpl) updateHealth(th *TabletHealth, prevTarget *query.Targ } // If the previous tablet type was primary, we need to check if the next new primary has already been assigned. - // If no new primary has been assigned, we will trigger a `loadTablets` call to immediately redirect traffic to the new primary. + // If no new primary has been assigned, we will trigger loading of tablets for this keyspace shard to immediately redirect traffic to the new primary. // // This is to avoid a situation where a newly primary tablet for a shard has just been started and the tableRefreshInterval has not yet passed, // causing an interruption where no primary is assigned to the shard. if prevTarget.TabletType == topodata.TabletType_PRIMARY { if primaries := hc.healthData[oldTargetKey]; len(primaries) == 0 { log.Infof("We will have no health data for the next new primary tablet after demoting the tablet: %v, so start loading tablets now", topotools.TabletIdent(th.Tablet)) - // We want to trigger a loadTablets call, but if the channel is not empty - // then a trigger is already scheduled, we don't need to trigger another one. - // This also prevents the code from deadlocking as described in https://github.com/vitessio/vitess/issues/16994. + // We want to trigger a call to load tablets for this keyspace-shard, + // but we want this to be non-blocking to prevent the code from deadlocking as described in https://github.com/vitessio/vitess/issues/16994. + // If the buffer is exhausted, then we'll just receive the update when all the tablets are loaded on the ticker. select { - case hc.loadTabletsTrigger <- struct{}{}: + case hc.loadTabletsTrigger <- topo.KeyspaceShard{ + Keyspace: prevTarget.Keyspace, + Shard: prevTarget.Shard, + }: default: } } @@ -662,7 +665,7 @@ func (hc *HealthCheckImpl) broadcast(th *TabletHealth) { } // GetLoadTabletsTrigger returns a channel that is used to inform when to load tablets. -func (hc *HealthCheckImpl) GetLoadTabletsTrigger() chan struct{} { +func (hc *HealthCheckImpl) GetLoadTabletsTrigger() chan topo.KeyspaceShard { return hc.loadTabletsTrigger } diff --git a/go/vt/discovery/healthcheck_test.go b/go/vt/discovery/healthcheck_test.go index 81f7d8b80b1..e76ad72ee37 100644 --- a/go/vt/discovery/healthcheck_test.go +++ b/go/vt/discovery/healthcheck_test.go @@ -1227,6 +1227,68 @@ func TestPrimaryInOtherCell(t *testing.T) { mustMatch(t, want, a[0], "Expecting healthy primary") } +// TestLoadTabletsTrigger tests that we send the correct information on the load tablets trigger. +func TestLoadTabletsTrigger(t *testing.T) { + ctx := utils.LeakCheckContext(t) + + // create a health check instance. + hc := NewHealthCheck(ctx, time.Hour, time.Hour, nil, "", "", nil) + defer hc.Close() + + ks := "keyspace" + shard := "shard" + // Add a tablet to the topology. + tablet1 := &topodatapb.Tablet{ + Alias: &topodatapb.TabletAlias{ + Cell: "zone-1", + Uid: 100, + }, + Type: topodatapb.TabletType_REPLICA, + Hostname: "host1", + PortMap: map[string]int32{ + "grpc": 123, + }, + Keyspace: ks, + Shard: shard, + } + + // We want to run updateHealth with arguments that always + // make it trigger load Tablets. + th := &TabletHealth{ + Tablet: tablet1, + Target: &querypb.Target{ + Keyspace: ks, + Shard: shard, + TabletType: topodatapb.TabletType_REPLICA, + }, + } + prevTarget := &querypb.Target{ + Keyspace: ks, + Shard: shard, + TabletType: topodatapb.TabletType_PRIMARY, + } + hc.AddTablet(tablet1) + + numTriggers := 10 + for i := 0; i < numTriggers; i++ { + // Since the previous target was a primary, and there are no other + // primary tablets for the given keyspace shard, we will see the healtcheck + // send on the loadTablets trigger. We just want to verify the information + // there is correct. + hc.updateHealth(th, prevTarget, false, false) + } + + ch := hc.GetLoadTabletsTrigger() + require.Len(t, ch, numTriggers) + for i := 0; i < numTriggers; i++ { + // Read from the channel and verify we indeed have the right values. + kss := <-ch + require.EqualValues(t, ks, kss.Keyspace) + require.EqualValues(t, shard, kss.Shard) + } + require.Len(t, ch, 0) +} + func TestReplicaInOtherCell(t *testing.T) { ctx := utils.LeakCheckContext(t) diff --git a/go/vt/discovery/topology_watcher.go b/go/vt/discovery/topology_watcher.go index d1e358e1aa5..30e6fc82c07 100644 --- a/go/vt/discovery/topology_watcher.go +++ b/go/vt/discovery/topology_watcher.go @@ -110,20 +110,39 @@ func (tw *TopologyWatcher) getTablets() ([]*topo.TabletInfo, error) { return tw.topoServer.GetTabletsByCell(tw.ctx, tw.cell, nil) } +func (tw *TopologyWatcher) getTabletsByShard(keyspace string, shard string) ([]*topo.TabletInfo, error) { + return tw.topoServer.GetTabletsByShardCell(tw.ctx, keyspace, shard, []string{tw.cell}) +} + // Start starts the topology watcher. func (tw *TopologyWatcher) Start() { tw.wg.Add(1) + // Goroutine to refresh the tablets list periodically. go func(t *TopologyWatcher) { defer t.wg.Done() ticker := time.NewTicker(t.refreshInterval) defer ticker.Stop() + t.loadTablets() for { - t.loadTablets() select { case <-t.ctx.Done(): return - case <-tw.healthcheck.GetLoadTabletsTrigger(): + case kss := <-t.healthcheck.GetLoadTabletsTrigger(): + t.loadTabletsForKeyspaceShard(kss.Keyspace, kss.Shard) case <-ticker.C: + // Since we are going to load all the tablets, + // we can clear out the entire list for reloading + // specific keyspace shards. + func() { + for { + select { + case <-t.healthcheck.GetLoadTabletsTrigger(): + default: + return + } + } + }() + t.loadTablets() } } }(tw) @@ -136,10 +155,23 @@ func (tw *TopologyWatcher) Stop() { tw.wg.Wait() } +func (tw *TopologyWatcher) loadTabletsForKeyspaceShard(keyspace string, shard string) { + if keyspace == "" || shard == "" { + log.Errorf("topologyWatcher: loadTabletsForKeyspaceShard: keyspace and shard are required") + return + } + tabletInfos, err := tw.getTabletsByShard(keyspace, shard) + if err != nil { + log.Errorf("error getting tablets for keyspace-shard: %v:%v: %v", keyspace, shard, err) + return + } + // Since we are only reading tablets for a keyspace shard, + // this is by default a partial result. + tw.storeTabletInfos(tabletInfos /* partialResults */, true) +} + func (tw *TopologyWatcher) loadTablets() { - newTablets := make(map[string]*tabletInfo) var partialResult bool - // First get the list of all tablets. tabletInfos, err := tw.getTablets() topologyWatcherOperations.Add(topologyWatcherOpListTablets, 1) @@ -155,6 +187,11 @@ func (tw *TopologyWatcher) loadTablets() { } } + tw.storeTabletInfos(tabletInfos, partialResult) +} + +func (tw *TopologyWatcher) storeTabletInfos(tabletInfos []*topo.TabletInfo, partialResult bool) { + newTablets := make(map[string]*tabletInfo) // Accumulate a list of all known alias strings to use later // when sorting. tabletAliasStrs := make([]string, 0, len(tabletInfos)) @@ -243,7 +280,6 @@ func (tw *TopologyWatcher) loadTablets() { } tw.topoChecksum = crc32.ChecksumIEEE(buf.Bytes()) tw.lastRefresh = time.Now() - } // RefreshLag returns the time since the last refresh. diff --git a/go/vt/discovery/topology_watcher_test.go b/go/vt/discovery/topology_watcher_test.go index 89a656c0982..73677a729f1 100644 --- a/go/vt/discovery/topology_watcher_test.go +++ b/go/vt/discovery/topology_watcher_test.go @@ -487,6 +487,71 @@ func TestFilterByKeyspace(t *testing.T) { } } +// TestLoadTablets tests that loadTablets works as intended for the given inputs. +func TestLoadTablets(t *testing.T) { + ctx := utils.LeakCheckContext(t) + + hc := NewFakeHealthCheck(nil) + f := TabletFilters{NewFilterByKeyspace(testKeyspacesToWatch)} + ts := memorytopo.NewServer(ctx, testCell) + defer ts.Close() + tw := NewTopologyWatcher(context.Background(), ts, hc, f, testCell, 10*time.Minute, true) + + // Add 2 tablets from 2 different tracked keyspaces to the topology. + tablet1 := &topodatapb.Tablet{ + Alias: &topodatapb.TabletAlias{ + Cell: testCell, + Uid: 0, + }, + Hostname: "host1", + PortMap: map[string]int32{ + "vt": 123, + }, + Keyspace: "ks1", + Shard: "shard", + } + tablet2 := &topodatapb.Tablet{ + Alias: &topodatapb.TabletAlias{ + Cell: testCell, + Uid: 10, + }, + Hostname: "host2", + PortMap: map[string]int32{ + "vt": 124, + }, + Keyspace: "ks4", + Shard: "shard", + } + for _, ks := range testKeyspacesToWatch { + _, err := ts.GetOrCreateShard(ctx, ks, "shard") + require.NoError(t, err) + } + require.NoError(t, ts.CreateTablet(ctx, tablet1)) + require.NoError(t, ts.CreateTablet(ctx, tablet2)) + + // Let's refresh the information for a different keyspace shard. We shouldn't + // reload either tablet's information. + tw.loadTabletsForKeyspaceShard("ks2", "shard") + key1 := TabletToMapKey(tablet1) + key2 := TabletToMapKey(tablet2) + allTablets := hc.GetAllTablets() + assert.NotContains(t, allTablets, key1) + assert.NotContains(t, allTablets, key2) + + // Now, if we reload the first tablet's shard, we should see this tablet + // but not the other. + tw.loadTabletsForKeyspaceShard("ks1", "shard") + allTablets = hc.GetAllTablets() + assert.Contains(t, allTablets, key1) + assert.NotContains(t, allTablets, key2) + + // Finally, if we load all the tablets, both the tablets should be visible. + tw.loadTablets() + allTablets = hc.GetAllTablets() + assert.Contains(t, allTablets, key1) + assert.Contains(t, allTablets, key2) +} + // TestFilterByKeyspaceSkipsIgnoredTablets confirms a bug fix for the case when a TopologyWatcher // has a FilterByKeyspace TabletFilter configured along with refreshKnownTablets turned off. We want // to ensure that the TopologyWatcher: diff --git a/go/vt/topo/shard.go b/go/vt/topo/shard.go index a610cac885a..dc49ad8578b 100644 --- a/go/vt/topo/shard.go +++ b/go/vt/topo/shard.go @@ -27,8 +27,6 @@ import ( "sync" "time" - "golang.org/x/sync/errgroup" - "vitess.io/vitess/go/constants/sidecar" "vitess.io/vitess/go/event" "vitess.io/vitess/go/protoutil" @@ -648,47 +646,20 @@ func (ts *Server) GetTabletsByShardCell(ctx context.Context, keyspace, shard str defer span.Finish() var err error - if len(cells) == 0 { - cells, err = ts.GetCellInfoNames(ctx) - if err != nil { - return nil, err - } - if len(cells) == 0 { // Nothing to do - return nil, nil - } + // if we get a partial result, we keep going. It most likely means + // a cell is out of commission. + aliases, err := ts.FindAllTabletAliasesInShardByCell(ctx, keyspace, shard, cells) + if err != nil && !IsErrType(err, PartialResult) { + return nil, err } - mu := sync.Mutex{} - eg, ctx := errgroup.WithContext(ctx) - - tablets := make([]*TabletInfo, 0, len(cells)) - var kss *KeyspaceShard - if keyspace != "" { - kss = &KeyspaceShard{ - Keyspace: keyspace, - Shard: shard, - } - } - options := &GetTabletsByCellOptions{ - KeyspaceShard: kss, - } - for _, cell := range cells { - eg.Go(func() error { - t, err := ts.GetTabletsByCell(ctx, cell, options) - if err != nil { - return vterrors.Wrapf(err, "GetTabletsByCell for %v failed.", cell) - } - mu.Lock() - defer mu.Unlock() - tablets = append(tablets, t...) - return nil - }) - } - if err := eg.Wait(); err != nil { - log.Warningf("GetTabletsByShardCell(%v,%v): got partial result: %v", keyspace, shard, err) - return tablets, NewError(PartialResult, shard) + // get the tablets for the cells we were able to reach, forward + // ErrPartialResult from FindAllTabletAliasesInShard + result, gerr := ts.GetTabletList(ctx, aliases, nil) + if gerr == nil && err != nil { + gerr = err } - return tablets, nil + return result, gerr } // GetTabletMapForShard returns the tablets for a shard. It can return diff --git a/go/vt/topo/shard_topo_test.go b/go/vt/topo/shard_topo_test.go new file mode 100644 index 00000000000..d9b987e132a --- /dev/null +++ b/go/vt/topo/shard_topo_test.go @@ -0,0 +1,98 @@ +/* +Copyright 2025 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 topo_test + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" + "golang.org/x/exp/maps" + + "vitess.io/vitess/go/vt/topo" + "vitess.io/vitess/go/vt/topo/memorytopo" +) + +// TestGetTabletsAndMapByShardCell tests GetTabletMapForShardByCell and GetTabletsByShardCell calls. +func TestGetTabletsAndMapByShardCell(t *testing.T) { + tests := []struct { + name string + keyspace string + shard string + cells []string + want map[string]*topo.TabletInfo + }{ + { + name: "no cells provided", + keyspace: kss[0], + shard: shards[1], + want: map[string]*topo.TabletInfo{ + "zone1-0000000002": { + Tablet: tablets[1], + }, + "zone2-0000000006": { + Tablet: tablets[5], + }, + }, + }, + { + name: "multiple cells", + keyspace: kss[0], + shard: shards[1], + cells: cells, + want: map[string]*topo.TabletInfo{ + "zone1-0000000002": { + Tablet: tablets[1], + }, + "zone2-0000000006": { + Tablet: tablets[5], + }, + }, + }, + { + name: "only one cell", + keyspace: kss[0], + shard: shards[1], + cells: []string{cells[0]}, + want: map[string]*topo.TabletInfo{ + "zone1-0000000002": { + Tablet: tablets[1], + }, + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + ts := memorytopo.NewServer(ctx, cells...) + defer ts.Close() + // This creates a tablet in each cell, keyspace, and shard, totalling 8 tablets. + setupFunc(t, ctx, ts) + + tabletMap, err := ts.GetTabletMapForShardByCell(ctx, tt.keyspace, tt.shard, tt.cells) + require.NoError(t, err) + checkTabletMapEqual(t, tt.want, tabletMap) + + tabletList, err := ts.GetTabletsByShardCell(ctx, tt.keyspace, tt.shard, tt.cells) + require.NoError(t, err) + checkTabletListEqual(t, maps.Values(tt.want), tabletList) + + }) + } +} diff --git a/go/vt/topo/tablet.go b/go/vt/topo/tablet.go index f2c9ab81d67..d830a25ae10 100644 --- a/go/vt/topo/tablet.go +++ b/go/vt/topo/tablet.go @@ -252,6 +252,13 @@ func (ts *Server) GetTabletsByCell(ctx context.Context, cellAlias string, opt *G // In the etcd case, it is possible that the response is too large. We also fall // back to fetching the tablets one by one in that case. if IsErrType(err, NoImplementation) || IsErrType(err, ResourceExhausted) { + // Getting all the tablets individually gets all the tablet records and filters + // them afterward. This is inefficient especially if we want to filter by + // keyspace and shard. So, we check for that case and use the ShardReplication + // to our advantage to reduce the number of topo calls. + if opt != nil && opt.KeyspaceShard != nil && opt.KeyspaceShard.Keyspace != "" && opt.KeyspaceShard.Shard != "" { + return ts.GetTabletsByShardCell(ctx, opt.KeyspaceShard.Keyspace, opt.KeyspaceShard.Shard, []string{cellAlias}) + } return ts.GetTabletsIndividuallyByCell(ctx, cellAlias, opt) } if IsErrType(err, NoNode) { @@ -509,6 +516,14 @@ func (ts *Server) GetTabletMap(ctx context.Context, tabletAliases []*topodatapb. returnErr = NewError(PartialResult, tabletAlias.GetCell()) } } else { + if opt != nil && opt.KeyspaceShard != nil { + if opt.KeyspaceShard.Keyspace != "" && opt.KeyspaceShard.Keyspace != tabletInfo.Keyspace { + return + } + if opt.KeyspaceShard.Shard != "" && opt.KeyspaceShard.Shard != tabletInfo.Shard { + return + } + } tabletMap[topoproto.TabletAliasString(tabletAlias)] = tabletInfo } }(tabletAlias) @@ -517,6 +532,56 @@ func (ts *Server) GetTabletMap(ctx context.Context, tabletAliases []*topodatapb. return tabletMap, returnErr } +// GetTabletList tries to read all the tablets in the provided list, +// and returns them in a list. +// If error is ErrPartialResult, the results in the list are +// incomplete, meaning some tablets couldn't be read. +func (ts *Server) GetTabletList(ctx context.Context, tabletAliases []*topodatapb.TabletAlias, opt *GetTabletsByCellOptions) ([]*TabletInfo, error) { + span, ctx := trace.NewSpan(ctx, "topo.GetTabletList") + span.Annotate("num_tablets", len(tabletAliases)) + defer span.Finish() + + var ( + mu sync.Mutex + wg sync.WaitGroup + tabletList = make([]*TabletInfo, 0) + returnErr error + ) + + for _, tabletAlias := range tabletAliases { + if tabletAlias == nil { + return nil, vterrors.Errorf(vtrpc.Code_INVALID_ARGUMENT, "nil tablet alias in list") + } + wg.Add(1) + go func(tabletAlias *topodatapb.TabletAlias) { + defer wg.Done() + tabletInfo, err := ts.GetTablet(ctx, tabletAlias) + mu.Lock() + defer mu.Unlock() + if err != nil { + log.Warningf("%v: %v", tabletAlias, err) + // There can be data races removing nodes - ignore them for now. + // We only need to set this on first error. + if returnErr == nil && !IsErrType(err, NoNode) { + returnErr = NewError(PartialResult, tabletAlias.GetCell()) + } + } else { + if opt != nil && opt.KeyspaceShard != nil { + if opt.KeyspaceShard.Keyspace != "" && opt.KeyspaceShard.Keyspace != tabletInfo.Keyspace { + return + } + if opt.KeyspaceShard.Shard != "" && opt.KeyspaceShard.Shard != tabletInfo.Shard { + return + } + } + tabletList = append(tabletList, tabletInfo) + } + }(tabletAlias) + } + wg.Wait() + return tabletList, returnErr +} + // InitTablet creates or updates a tablet. If no parent is specified // in the tablet, and the tablet has a replica type, we will find the // appropriate parent. If createShardAndKeyspace is true and the diff --git a/go/vt/topo/tablet_test.go b/go/vt/topo/tablet_test.go index 1c242e8778b..6b69e3e1f3b 100644 --- a/go/vt/topo/tablet_test.go +++ b/go/vt/topo/tablet_test.go @@ -22,16 +22,39 @@ import ( "errors" "slices" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "golang.org/x/exp/maps" "google.golang.org/protobuf/proto" topodatapb "vitess.io/vitess/go/vt/proto/topodata" "vitess.io/vitess/go/vt/topo" "vitess.io/vitess/go/vt/topo/memorytopo" + "vitess.io/vitess/go/vt/topo/topoproto" ) +var ( + cells = []string{"zone1", "zone2"} + kss = []string{"ks1", "ks2"} + shards = []string{"-80", "80-"} + tablets []*topodatapb.Tablet +) + +func init() { + uid := 1 + for _, cell := range cells { + for _, ks := range kss { + for _, shard := range shards { + tablet := getTablet(ks, shard, cell, int32(uid)) + tablets = append(tablets, tablet) + uid++ + } + } + } +} + // Test various cases of calls to GetTabletsByCell. // GetTabletsByCell first tries to get all the tablets using List. // If the response is too large, we will get an error, and fall back to one tablet at a time. @@ -369,6 +392,12 @@ func TestServerGetTabletsByCell(t *testing.T) { require.NoError(t, err) require.Len(t, out, len(tt.expectedTablets)) + // We also check that the results for getting tablets individually + // matches the output we get from listing them. + out2, err := ts.GetTabletsIndividuallyByCell(ctx, cell, tt.opt) + require.NoError(t, err) + require.ElementsMatch(t, out, out2) + slices.SortFunc(out, func(i, j *topo.TabletInfo) int { return cmp.Compare(i.Alias.Uid, j.Alias.Uid) }) @@ -377,10 +406,7 @@ func TestServerGetTabletsByCell(t *testing.T) { }) for i, tablet := range out { - expected := tt.expectedTablets[i] - require.Equal(t, expected.Alias.String(), tablet.Alias.String()) - require.Equal(t, expected.Keyspace, tablet.Keyspace) - require.Equal(t, expected.Shard, tablet.Shard) + checkTabletsEqual(t, tt.expectedTablets[i], tablet.Tablet) } }) } @@ -437,3 +463,263 @@ func TestServerGetTabletsByCellPartialResults(t *testing.T) { assert.True(t, proto.Equal(tablets[0].Tablet, out[0].Tablet), "Got: %v, want %v", tablets[0].Tablet, out[0].Tablet) assert.True(t, proto.Equal(tablets[2].Tablet, out[1].Tablet), "Got: %v, want %v", tablets[2].Tablet, out[1].Tablet) } + +func getTablet(keyspace string, shard string, cell string, uid int32) *topodatapb.Tablet { + return &topodatapb.Tablet{ + Alias: &topodatapb.TabletAlias{ + Cell: cell, + Uid: uint32(uid), + }, + Hostname: "host1", + PortMap: map[string]int32{ + "vt": uid, + }, + Keyspace: keyspace, + Shard: shard, + } +} + +func checkTabletsEqual(t *testing.T, expected, tablet *topodatapb.Tablet) { + t.Helper() + require.Equal(t, expected.Alias.String(), tablet.Alias.String()) + require.Equal(t, expected.Keyspace, tablet.Keyspace) + require.Equal(t, expected.Shard, tablet.Shard) +} + +func checkTabletMapEqual(t *testing.T, expected, tabletMap map[string]*topo.TabletInfo) { + t.Helper() + require.Len(t, tabletMap, len(expected)) + for key, tablet := range tabletMap { + expectedTablet, ok := expected[key] + require.True(t, ok, "unexpected tablet %v", key) + checkTabletsEqual(t, expectedTablet.Tablet, tablet.Tablet) + } +} + +func checkTabletListEqual(t *testing.T, expected, tabletMap []*topo.TabletInfo) { + t.Helper() + require.Len(t, tabletMap, len(expected)) + for _, tablet := range tabletMap { + found := false + for _, expectedTablet := range expected { + if topoproto.TabletAliasString(tablet.Alias) == topoproto.TabletAliasString(expectedTablet.Alias) { + checkTabletsEqual(t, expectedTablet.Tablet, tablet.Tablet) + found = true + break + } + } + require.True(t, found, "unexpected tablet %v", tablet) + } +} + +func setupFunc(t *testing.T, ctx context.Context, ts *topo.Server) { + for _, ks := range kss { + for _, shard := range shards { + _, err := ts.GetOrCreateShard(ctx, ks, shard) + require.NoError(t, err) + } + } + for _, tablet := range tablets { + require.NoError(t, ts.CreateTablet(ctx, tablet)) + } +} + +// TestServerGetTabletMapAndList tests the GetTabletMap and GetTabletList methods. +func TestServerGetTabletMapAndList(t *testing.T) { + tests := []struct { + name string + tabletAliases []*topodatapb.TabletAlias + opt *topo.GetTabletsByCellOptions + want map[string]*topo.TabletInfo + }{ + { + name: "single tablet without filtering - found", + tabletAliases: []*topodatapb.TabletAlias{ + { + Cell: cells[0], + Uid: 2, + }, + }, + opt: nil, + want: map[string]*topo.TabletInfo{ + "zone1-0000000002": { + Tablet: tablets[1], + }, + }, + }, + { + name: "single tablet without filtering - not found", + tabletAliases: []*topodatapb.TabletAlias{ + { + Cell: cells[0], + Uid: 2050, + }, + }, + opt: nil, + want: map[string]*topo.TabletInfo{}, + }, + { + name: "multiple tablets without filtering", + tabletAliases: []*topodatapb.TabletAlias{ + { + Cell: cells[0], + Uid: 2, + }, + { + Cell: cells[0], + Uid: 3, + }, + { + Cell: cells[0], + Uid: 4, + }, + { + Cell: cells[1], + Uid: 5, + }, + { + Cell: cells[1], + Uid: 205, + }, + }, + opt: nil, + want: map[string]*topo.TabletInfo{ + "zone1-0000000002": { + Tablet: tablets[1], + }, + "zone1-0000000003": { + Tablet: tablets[2], + }, + "zone1-0000000004": { + Tablet: tablets[3], + }, + "zone2-0000000005": { + Tablet: tablets[4], + }, + }, + }, + { + name: "multiple tablets with filtering", + tabletAliases: []*topodatapb.TabletAlias{ + { + Cell: cells[0], + Uid: 2, + }, + { + Cell: cells[0], + Uid: 3, + }, + { + Cell: cells[0], + Uid: 4, + }, + { + Cell: cells[1], + Uid: 5, + }, + { + Cell: cells[1], + Uid: 6, + }, + { + Cell: cells[1], + Uid: 205, + }, + }, + opt: &topo.GetTabletsByCellOptions{ + KeyspaceShard: &topo.KeyspaceShard{ + Keyspace: kss[0], + Shard: shards[1], + }, + }, + want: map[string]*topo.TabletInfo{ + "zone1-0000000002": { + Tablet: tablets[1], + }, + "zone2-0000000006": { + Tablet: tablets[5], + }, + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + ts := memorytopo.NewServer(ctx, cells...) + defer ts.Close() + // This creates a tablet in each cell, keyspace, and shard, totalling 8 tablets. + setupFunc(t, ctx, ts) + + tabletMap, err := ts.GetTabletMap(ctx, tt.tabletAliases, tt.opt) + require.NoError(t, err) + checkTabletMapEqual(t, tt.want, tabletMap) + + tabletList, err := ts.GetTabletList(ctx, tt.tabletAliases, tt.opt) + require.NoError(t, err) + checkTabletListEqual(t, maps.Values(tt.want), tabletList) + }) + } +} + +// TestGetTabletsIndividuallyByCell tests the GetTabletsIndividuallyByCell function. +func TestGetTabletsIndividuallyByCell(t *testing.T) { + tests := []struct { + name string + keyspace string + shard string + cell string + want []*topo.TabletInfo + }{ + { + name: "cell with filtering", + keyspace: kss[0], + shard: shards[1], + cell: cells[0], + want: []*topo.TabletInfo{ + { + Tablet: tablets[1], + }, + }, + }, + { + name: "cell without filtering", + cell: cells[0], + want: []*topo.TabletInfo{ + { + Tablet: tablets[0], + }, + { + Tablet: tablets[1], + }, + { + Tablet: tablets[2], + }, + { + Tablet: tablets[3], + }, + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + ts := memorytopo.NewServer(ctx, cells...) + defer ts.Close() + // This creates a tablet in each cell, keyspace, and shard, totalling 8 tablets. + setupFunc(t, ctx, ts) + + tabletList, err := ts.GetTabletsIndividuallyByCell(ctx, tt.cell, &topo.GetTabletsByCellOptions{KeyspaceShard: &topo.KeyspaceShard{Keyspace: tt.keyspace, Shard: tt.shard}}) + require.NoError(t, err) + checkTabletListEqual(t, tt.want, tabletList) + + if tt.keyspace != "" && tt.shard != "" { + // We can also check that this result matches what we get from GetTabletsByShardCell. + tl, err := ts.GetTabletsByShardCell(ctx, tt.keyspace, tt.shard, []string{tt.cell}) + require.NoError(t, err) + checkTabletListEqual(t, tabletList, tl) + } + }) + } +} diff --git a/go/vt/vttablet/tabletserver/txthrottler/mock_healthcheck_test.go b/go/vt/vttablet/tabletserver/txthrottler/mock_healthcheck_test.go index ecc6688fb9d..043c30a815c 100644 --- a/go/vt/vttablet/tabletserver/txthrottler/mock_healthcheck_test.go +++ b/go/vt/vttablet/tabletserver/txthrottler/mock_healthcheck_test.go @@ -1,5 +1,10 @@ // Code generated by MockGen. DO NOT EDIT. // Source: vitess.io/vitess/go/vt/discovery (interfaces: HealthCheck) +// +// Generated by this command: +// +// mockgen -destination mock_healthcheck_test.go -package txthrottler -mock_names HealthCheck=MockHealthCheck vitess.io/vitess/go/vt/discovery HealthCheck +// // Package txthrottler is a generated GoMock package. package txthrottler @@ -13,6 +18,7 @@ import ( discovery "vitess.io/vitess/go/vt/discovery" query "vitess.io/vitess/go/vt/proto/query" topodata "vitess.io/vitess/go/vt/proto/topodata" + topo "vitess.io/vitess/go/vt/topo" queryservice "vitess.io/vitess/go/vt/vttablet/queryservice" ) @@ -20,6 +26,7 @@ import ( type MockHealthCheck struct { ctrl *gomock.Controller recorder *MockHealthCheckMockRecorder + isgomock struct{} } // MockHealthCheckMockRecorder is the mock recorder for MockHealthCheck. @@ -40,15 +47,15 @@ func (m *MockHealthCheck) EXPECT() *MockHealthCheckMockRecorder { } // AddTablet mocks base method. -func (m *MockHealthCheck) AddTablet(arg0 *topodata.Tablet) { +func (m *MockHealthCheck) AddTablet(tablet *topodata.Tablet) { m.ctrl.T.Helper() - m.ctrl.Call(m, "AddTablet", arg0) + m.ctrl.Call(m, "AddTablet", tablet) } // AddTablet indicates an expected call of AddTablet. -func (mr *MockHealthCheckMockRecorder) AddTablet(arg0 interface{}) *gomock.Call { +func (mr *MockHealthCheckMockRecorder) AddTablet(tablet any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AddTablet", reflect.TypeOf((*MockHealthCheck)(nil).AddTablet), arg0) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AddTablet", reflect.TypeOf((*MockHealthCheck)(nil).AddTablet), tablet) } // CacheStatus mocks base method. @@ -59,14 +66,6 @@ func (m *MockHealthCheck) CacheStatus() discovery.TabletsCacheStatusList { return ret0 } -// HealthyStatus mocks base method. -func (m *MockHealthCheck) HealthyStatus() discovery.TabletsCacheStatusList { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "HealthyStatus") - ret0, _ := ret[0].(discovery.TabletsCacheStatusList) - return ret0 -} - // CacheStatus indicates an expected call of CacheStatus. func (mr *MockHealthCheckMockRecorder) CacheStatus() *gomock.Call { mr.mock.ctrl.T.Helper() @@ -102,24 +101,24 @@ func (mr *MockHealthCheckMockRecorder) Close() *gomock.Call { } // GetHealthyTabletStats mocks base method. -func (m *MockHealthCheck) GetHealthyTabletStats(arg0 *query.Target) []*discovery.TabletHealth { +func (m *MockHealthCheck) GetHealthyTabletStats(target *query.Target) []*discovery.TabletHealth { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetHealthyTabletStats", arg0) + ret := m.ctrl.Call(m, "GetHealthyTabletStats", target) ret0, _ := ret[0].([]*discovery.TabletHealth) return ret0 } // GetHealthyTabletStats indicates an expected call of GetHealthyTabletStats. -func (mr *MockHealthCheckMockRecorder) GetHealthyTabletStats(arg0 interface{}) *gomock.Call { +func (mr *MockHealthCheckMockRecorder) GetHealthyTabletStats(target any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetHealthyTabletStats", reflect.TypeOf((*MockHealthCheck)(nil).GetHealthyTabletStats), arg0) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetHealthyTabletStats", reflect.TypeOf((*MockHealthCheck)(nil).GetHealthyTabletStats), target) } // GetLoadTabletsTrigger mocks base method. -func (m *MockHealthCheck) GetLoadTabletsTrigger() chan struct{} { +func (m *MockHealthCheck) GetLoadTabletsTrigger() chan topo.KeyspaceShard { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "GetLoadTabletsTrigger") - ret0, _ := ret[0].(chan struct{}) + ret0, _ := ret[0].(chan topo.KeyspaceShard) return ret0 } @@ -130,33 +129,47 @@ func (mr *MockHealthCheckMockRecorder) GetLoadTabletsTrigger() *gomock.Call { } // GetTabletHealth mocks base method. -func (m *MockHealthCheck) GetTabletHealth(arg0 discovery.KeyspaceShardTabletType, arg1 *topodata.TabletAlias) (*discovery.TabletHealth, error) { +func (m *MockHealthCheck) GetTabletHealth(kst discovery.KeyspaceShardTabletType, alias *topodata.TabletAlias) (*discovery.TabletHealth, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetTabletHealth", arg0, arg1) + ret := m.ctrl.Call(m, "GetTabletHealth", kst, alias) ret0, _ := ret[0].(*discovery.TabletHealth) ret1, _ := ret[1].(error) return ret0, ret1 } // GetTabletHealth indicates an expected call of GetTabletHealth. -func (mr *MockHealthCheckMockRecorder) GetTabletHealth(arg0, arg1 interface{}) *gomock.Call { +func (mr *MockHealthCheckMockRecorder) GetTabletHealth(kst, alias any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTabletHealth", reflect.TypeOf((*MockHealthCheck)(nil).GetTabletHealth), arg0, arg1) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTabletHealth", reflect.TypeOf((*MockHealthCheck)(nil).GetTabletHealth), kst, alias) } // GetTabletHealthByAlias mocks base method. -func (m *MockHealthCheck) GetTabletHealthByAlias(arg0 *topodata.TabletAlias) (*discovery.TabletHealth, error) { +func (m *MockHealthCheck) GetTabletHealthByAlias(alias *topodata.TabletAlias) (*discovery.TabletHealth, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetTabletHealthByAlias", arg0) + ret := m.ctrl.Call(m, "GetTabletHealthByAlias", alias) ret0, _ := ret[0].(*discovery.TabletHealth) ret1, _ := ret[1].(error) return ret0, ret1 } // GetTabletHealthByAlias indicates an expected call of GetTabletHealthByAlias. -func (mr *MockHealthCheckMockRecorder) GetTabletHealthByAlias(arg0 interface{}) *gomock.Call { +func (mr *MockHealthCheckMockRecorder) GetTabletHealthByAlias(alias any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTabletHealthByAlias", reflect.TypeOf((*MockHealthCheck)(nil).GetTabletHealthByAlias), alias) +} + +// HealthyStatus mocks base method. +func (m *MockHealthCheck) HealthyStatus() discovery.TabletsCacheStatusList { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "HealthyStatus") + ret0, _ := ret[0].(discovery.TabletsCacheStatusList) + return ret0 +} + +// HealthyStatus indicates an expected call of HealthyStatus. +func (mr *MockHealthCheckMockRecorder) HealthyStatus() *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTabletHealthByAlias", reflect.TypeOf((*MockHealthCheck)(nil).GetTabletHealthByAlias), arg0) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HealthyStatus", reflect.TypeOf((*MockHealthCheck)(nil).HealthyStatus)) } // RegisterStats mocks base method. @@ -172,27 +185,27 @@ func (mr *MockHealthCheckMockRecorder) RegisterStats() *gomock.Call { } // RemoveTablet mocks base method. -func (m *MockHealthCheck) RemoveTablet(arg0 *topodata.Tablet) { +func (m *MockHealthCheck) RemoveTablet(tablet *topodata.Tablet) { m.ctrl.T.Helper() - m.ctrl.Call(m, "RemoveTablet", arg0) + m.ctrl.Call(m, "RemoveTablet", tablet) } // RemoveTablet indicates an expected call of RemoveTablet. -func (mr *MockHealthCheckMockRecorder) RemoveTablet(arg0 interface{}) *gomock.Call { +func (mr *MockHealthCheckMockRecorder) RemoveTablet(tablet any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemoveTablet", reflect.TypeOf((*MockHealthCheck)(nil).RemoveTablet), arg0) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemoveTablet", reflect.TypeOf((*MockHealthCheck)(nil).RemoveTablet), tablet) } // ReplaceTablet mocks base method. -func (m *MockHealthCheck) ReplaceTablet(arg0, arg1 *topodata.Tablet) { +func (m *MockHealthCheck) ReplaceTablet(old, new *topodata.Tablet) { m.ctrl.T.Helper() - m.ctrl.Call(m, "ReplaceTablet", arg0, arg1) + m.ctrl.Call(m, "ReplaceTablet", old, new) } // ReplaceTablet indicates an expected call of ReplaceTablet. -func (mr *MockHealthCheckMockRecorder) ReplaceTablet(arg0, arg1 interface{}) *gomock.Call { +func (mr *MockHealthCheckMockRecorder) ReplaceTablet(old, new any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReplaceTablet", reflect.TypeOf((*MockHealthCheck)(nil).ReplaceTablet), arg0, arg1) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReplaceTablet", reflect.TypeOf((*MockHealthCheck)(nil).ReplaceTablet), old, new) } // Subscribe mocks base method. @@ -210,42 +223,42 @@ func (mr *MockHealthCheckMockRecorder) Subscribe() *gomock.Call { } // TabletConnection mocks base method. -func (m *MockHealthCheck) TabletConnection(arg0 context.Context, arg1 *topodata.TabletAlias, arg2 *query.Target) (queryservice.QueryService, error) { +func (m *MockHealthCheck) TabletConnection(ctx context.Context, alias *topodata.TabletAlias, target *query.Target) (queryservice.QueryService, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "TabletConnection", arg0, arg1, arg2) + ret := m.ctrl.Call(m, "TabletConnection", ctx, alias, target) ret0, _ := ret[0].(queryservice.QueryService) ret1, _ := ret[1].(error) return ret0, ret1 } // TabletConnection indicates an expected call of TabletConnection. -func (mr *MockHealthCheckMockRecorder) TabletConnection(arg0, arg1 interface{}) *gomock.Call { +func (mr *MockHealthCheckMockRecorder) TabletConnection(ctx, alias, target any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "TabletConnection", reflect.TypeOf((*MockHealthCheck)(nil).TabletConnection), arg0, arg1) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "TabletConnection", reflect.TypeOf((*MockHealthCheck)(nil).TabletConnection), ctx, alias, target) } // Unsubscribe mocks base method. -func (m *MockHealthCheck) Unsubscribe(arg0 chan *discovery.TabletHealth) { +func (m *MockHealthCheck) Unsubscribe(c chan *discovery.TabletHealth) { m.ctrl.T.Helper() - m.ctrl.Call(m, "Unsubscribe", arg0) + m.ctrl.Call(m, "Unsubscribe", c) } // Unsubscribe indicates an expected call of Unsubscribe. -func (mr *MockHealthCheckMockRecorder) Unsubscribe(arg0 interface{}) *gomock.Call { +func (mr *MockHealthCheckMockRecorder) Unsubscribe(c any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Unsubscribe", reflect.TypeOf((*MockHealthCheck)(nil).Unsubscribe), arg0) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Unsubscribe", reflect.TypeOf((*MockHealthCheck)(nil).Unsubscribe), c) } // WaitForAllServingTablets mocks base method. -func (m *MockHealthCheck) WaitForAllServingTablets(arg0 context.Context, arg1 []*query.Target) error { +func (m *MockHealthCheck) WaitForAllServingTablets(ctx context.Context, targets []*query.Target) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "WaitForAllServingTablets", arg0, arg1) + ret := m.ctrl.Call(m, "WaitForAllServingTablets", ctx, targets) ret0, _ := ret[0].(error) return ret0 } // WaitForAllServingTablets indicates an expected call of WaitForAllServingTablets. -func (mr *MockHealthCheckMockRecorder) WaitForAllServingTablets(arg0, arg1 interface{}) *gomock.Call { +func (mr *MockHealthCheckMockRecorder) WaitForAllServingTablets(ctx, targets any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WaitForAllServingTablets", reflect.TypeOf((*MockHealthCheck)(nil).WaitForAllServingTablets), arg0, arg1) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WaitForAllServingTablets", reflect.TypeOf((*MockHealthCheck)(nil).WaitForAllServingTablets), ctx, targets) } diff --git a/go/vt/vttablet/tabletserver/txthrottler/mock_throttler_test.go b/go/vt/vttablet/tabletserver/txthrottler/mock_throttler_test.go index 327a37dc43f..00ac81b010b 100644 --- a/go/vt/vttablet/tabletserver/txthrottler/mock_throttler_test.go +++ b/go/vt/vttablet/tabletserver/txthrottler/mock_throttler_test.go @@ -1,5 +1,10 @@ // Code generated by MockGen. DO NOT EDIT. // Source: vitess.io/vitess/go/vt/throttler (interfaces: Throttler) +// +// Generated by this command: +// +// mockgen -destination mock_throttler_test.go -package txthrottler vitess.io/vitess/go/vt/throttler Throttler +// // Package txthrottler is a generated GoMock package. package txthrottler @@ -20,6 +25,7 @@ import ( type MockThrottler struct { ctrl *gomock.Controller recorder *MockThrottlerMockRecorder + isgomock struct{} } // MockThrottlerMockRecorder is the mock recorder for MockThrottler. @@ -80,17 +86,17 @@ func (mr *MockThrottlerMockRecorder) Log() *gomock.Call { } // MaxLag mocks base method. -func (m *MockThrottler) MaxLag(arg0 topodata.TabletType) uint32 { +func (m *MockThrottler) MaxLag(tabletType topodata.TabletType) uint32 { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "MaxLag", arg0) + ret := m.ctrl.Call(m, "MaxLag", tabletType) ret0, _ := ret[0].(uint32) return ret0 } // MaxLag indicates an expected call of MaxLag. -func (mr *MockThrottlerMockRecorder) MaxLag(arg0 interface{}) *gomock.Call { +func (mr *MockThrottlerMockRecorder) MaxLag(tabletType any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MaxLag", reflect.TypeOf((*MockThrottler)(nil).MaxLag), arg0) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MaxLag", reflect.TypeOf((*MockThrottler)(nil).MaxLag), tabletType) } // MaxRate mocks base method. @@ -108,15 +114,15 @@ func (mr *MockThrottlerMockRecorder) MaxRate() *gomock.Call { } // RecordReplicationLag mocks base method. -func (m *MockThrottler) RecordReplicationLag(arg0 time.Time, arg1 *discovery.TabletHealth) { +func (m *MockThrottler) RecordReplicationLag(time time.Time, th *discovery.TabletHealth) { m.ctrl.T.Helper() - m.ctrl.Call(m, "RecordReplicationLag", arg0, arg1) + m.ctrl.Call(m, "RecordReplicationLag", time, th) } // RecordReplicationLag indicates an expected call of RecordReplicationLag. -func (mr *MockThrottlerMockRecorder) RecordReplicationLag(arg0, arg1 interface{}) *gomock.Call { +func (mr *MockThrottlerMockRecorder) RecordReplicationLag(time, th any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RecordReplicationLag", reflect.TypeOf((*MockThrottler)(nil).RecordReplicationLag), arg0, arg1) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RecordReplicationLag", reflect.TypeOf((*MockThrottler)(nil).RecordReplicationLag), time, th) } // ResetConfiguration mocks base method. @@ -132,53 +138,53 @@ func (mr *MockThrottlerMockRecorder) ResetConfiguration() *gomock.Call { } // SetMaxRate mocks base method. -func (m *MockThrottler) SetMaxRate(arg0 int64) { +func (m *MockThrottler) SetMaxRate(rate int64) { m.ctrl.T.Helper() - m.ctrl.Call(m, "SetMaxRate", arg0) + m.ctrl.Call(m, "SetMaxRate", rate) } // SetMaxRate indicates an expected call of SetMaxRate. -func (mr *MockThrottlerMockRecorder) SetMaxRate(arg0 interface{}) *gomock.Call { +func (mr *MockThrottlerMockRecorder) SetMaxRate(rate any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetMaxRate", reflect.TypeOf((*MockThrottler)(nil).SetMaxRate), arg0) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetMaxRate", reflect.TypeOf((*MockThrottler)(nil).SetMaxRate), rate) } // ThreadFinished mocks base method. -func (m *MockThrottler) ThreadFinished(arg0 int) { +func (m *MockThrottler) ThreadFinished(threadID int) { m.ctrl.T.Helper() - m.ctrl.Call(m, "ThreadFinished", arg0) + m.ctrl.Call(m, "ThreadFinished", threadID) } // ThreadFinished indicates an expected call of ThreadFinished. -func (mr *MockThrottlerMockRecorder) ThreadFinished(arg0 interface{}) *gomock.Call { +func (mr *MockThrottlerMockRecorder) ThreadFinished(threadID any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ThreadFinished", reflect.TypeOf((*MockThrottler)(nil).ThreadFinished), arg0) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ThreadFinished", reflect.TypeOf((*MockThrottler)(nil).ThreadFinished), threadID) } // Throttle mocks base method. -func (m *MockThrottler) Throttle(arg0 int) time.Duration { +func (m *MockThrottler) Throttle(threadID int) time.Duration { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "Throttle", arg0) + ret := m.ctrl.Call(m, "Throttle", threadID) ret0, _ := ret[0].(time.Duration) return ret0 } // Throttle indicates an expected call of Throttle. -func (mr *MockThrottlerMockRecorder) Throttle(arg0 interface{}) *gomock.Call { +func (mr *MockThrottlerMockRecorder) Throttle(threadID any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Throttle", reflect.TypeOf((*MockThrottler)(nil).Throttle), arg0) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Throttle", reflect.TypeOf((*MockThrottler)(nil).Throttle), threadID) } // UpdateConfiguration mocks base method. -func (m *MockThrottler) UpdateConfiguration(arg0 *throttlerdata.Configuration, arg1 bool) error { +func (m *MockThrottler) UpdateConfiguration(configuration *throttlerdata.Configuration, copyZeroValues bool) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "UpdateConfiguration", arg0, arg1) + ret := m.ctrl.Call(m, "UpdateConfiguration", configuration, copyZeroValues) ret0, _ := ret[0].(error) return ret0 } // UpdateConfiguration indicates an expected call of UpdateConfiguration. -func (mr *MockThrottlerMockRecorder) UpdateConfiguration(arg0, arg1 interface{}) *gomock.Call { +func (mr *MockThrottlerMockRecorder) UpdateConfiguration(configuration, copyZeroValues any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateConfiguration", reflect.TypeOf((*MockThrottler)(nil).UpdateConfiguration), arg0, arg1) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateConfiguration", reflect.TypeOf((*MockThrottler)(nil).UpdateConfiguration), configuration, copyZeroValues) }