From 2cb42d1f919a4eaec895b33b3a18ddb12c430f57 Mon Sep 17 00:00:00 2001 From: Tim Vaillancourt Date: Sat, 24 Feb 2024 00:52:47 +0100 Subject: [PATCH] Filter by keyspace earlier in `tabletgateway`s `WaitForTablets(...)` Signed-off-by: Tim Vaillancourt --- go/vt/discovery/healthcheck.go | 22 ------------- go/vt/discovery/healthcheck_test.go | 7 ----- go/vt/srvtopo/discover.go | 19 +++++++++++- go/vt/srvtopo/discover_test.go | 48 ++++++++++++++++++++++++++--- go/vt/vtgate/tabletgateway.go | 2 +- 5 files changed, 62 insertions(+), 36 deletions(-) diff --git a/go/vt/discovery/healthcheck.go b/go/vt/discovery/healthcheck.go index 5d6a5e32662..6b30a794893 100644 --- a/go/vt/discovery/healthcheck.go +++ b/go/vt/discovery/healthcheck.go @@ -750,30 +750,8 @@ func (hc *HealthCheckImpl) WaitForAllServingTablets(ctx context.Context, targets return hc.waitForTablets(ctx, targets, true) } -// FilterTargetsByKeyspaces only returns the targets that are part of the provided keyspaces -func FilterTargetsByKeyspaces(keyspaces []string, targets []*query.Target) []*query.Target { - filteredTargets := make([]*query.Target, 0) - - // Keep them all if there are no keyspaces to watch - if len(KeyspacesToWatch) == 0 { - return append(filteredTargets, targets...) - } - - // Let's remove from the target shards that are not in the keyspaceToWatch list. - for _, target := range targets { - for _, keyspaceToWatch := range keyspaces { - if target.Keyspace == keyspaceToWatch { - filteredTargets = append(filteredTargets, target) - } - } - } - return filteredTargets -} - // waitForTablets is the internal method that polls for tablets. func (hc *HealthCheckImpl) waitForTablets(ctx context.Context, targets []*query.Target, requireServing bool) error { - targets = FilterTargetsByKeyspaces(KeyspacesToWatch, targets) - for { // We nil targets as we find them. allPresent := true diff --git a/go/vt/discovery/healthcheck_test.go b/go/vt/discovery/healthcheck_test.go index 9563d9bfdc5..adc415cc497 100644 --- a/go/vt/discovery/healthcheck_test.go +++ b/go/vt/discovery/healthcheck_test.go @@ -686,13 +686,6 @@ func TestWaitForAllServingTablets(t *testing.T) { TabletType: tablet.Type, }, } - - KeyspacesToWatch = []string{tablet.Keyspace} - - err = hc.WaitForAllServingTablets(ctx, targets) - assert.Nil(t, err, "error should be nil. Keyspace with no tablets is filtered") - - KeyspacesToWatch = []string{} } // TestRemoveTablet tests the behavior when a tablet goes away. diff --git a/go/vt/srvtopo/discover.go b/go/vt/srvtopo/discover.go index 91aaea9daf6..7c7e4c58023 100644 --- a/go/vt/srvtopo/discover.go +++ b/go/vt/srvtopo/discover.go @@ -29,10 +29,24 @@ import ( topodatapb "vitess.io/vitess/go/vt/proto/topodata" ) +// shouldFindKeyspace returns true if a keyspace should be fetched. +// true is also returned if no keyspaces are defined. +func shouldFindKeyspace(keyspaces []string, ksName string) bool { + if len(keyspaces) == 0 { + return true + } + for _, keyspace := range keyspaces { + if ksName == keyspace { + return true + } + } + return false +} + // FindAllTargets goes through all serving shards in the topology // for the provided tablet types. It returns one Target object per // keyspace / shard / matching TabletType. -func FindAllTargets(ctx context.Context, ts Server, cell string, tabletTypes []topodatapb.TabletType) ([]*querypb.Target, error) { +func FindAllTargets(ctx context.Context, ts Server, cell string, keyspaces []string, tabletTypes []topodatapb.TabletType) ([]*querypb.Target, error) { ksNames, err := ts.GetSrvKeyspaceNames(ctx, cell, true) if err != nil { return nil, err @@ -43,6 +57,9 @@ func FindAllTargets(ctx context.Context, ts Server, cell string, tabletTypes []t var mu sync.Mutex var errRecorder concurrency.AllErrorRecorder for _, ksName := range ksNames { + if !shouldFindKeyspace(keyspaces, ksName) { + continue + } wg.Add(1) go func(keyspace string) { defer wg.Done() diff --git a/go/vt/srvtopo/discover_test.go b/go/vt/srvtopo/discover_test.go index 81279b7d61a..fbfde18e2e3 100644 --- a/go/vt/srvtopo/discover_test.go +++ b/go/vt/srvtopo/discover_test.go @@ -64,7 +64,7 @@ func TestFindAllTargets(t *testing.T) { rs := NewResilientServer(ctx, ts, counts) // No keyspace / shards. - ks, err := FindAllTargets(ctx, rs, "cell1", []topodatapb.TabletType{topodatapb.TabletType_PRIMARY}) + ks, err := FindAllTargets(ctx, rs, "cell1", []string{"test_keyspace"}, []topodatapb.TabletType{topodatapb.TabletType_PRIMARY}) if err != nil { t.Errorf("unexpected error: %v", err) } @@ -89,7 +89,23 @@ func TestFindAllTargets(t *testing.T) { } // Get it. - ks, err = FindAllTargets(ctx, rs, "cell1", []topodatapb.TabletType{topodatapb.TabletType_PRIMARY}) + ks, err = FindAllTargets(ctx, rs, "cell1", []string{"test_keyspace"}, []topodatapb.TabletType{topodatapb.TabletType_PRIMARY}) + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if !reflect.DeepEqual(ks, []*querypb.Target{ + { + Cell: "cell1", + Keyspace: "test_keyspace", + Shard: "test_shard0", + TabletType: topodatapb.TabletType_PRIMARY, + }, + }) { + t.Errorf("got wrong value: %v", ks) + } + + // Get any keyspace. + ks, err = FindAllTargets(ctx, rs, "cell1", []string{}, []topodatapb.TabletType{topodatapb.TabletType_PRIMARY}) if err != nil { t.Errorf("unexpected error: %v", err) } @@ -129,7 +145,7 @@ func TestFindAllTargets(t *testing.T) { } // Get it for all types. - ks, err = FindAllTargets(ctx, rs, "cell1", []topodatapb.TabletType{topodatapb.TabletType_PRIMARY, topodatapb.TabletType_REPLICA}) + ks, err = FindAllTargets(ctx, rs, "cell1", []string{"test_keyspace", "test_keyspace2"}, []topodatapb.TabletType{topodatapb.TabletType_PRIMARY, topodatapb.TabletType_REPLICA}) if err != nil { t.Errorf("unexpected error: %v", err) } @@ -157,8 +173,30 @@ func TestFindAllTargets(t *testing.T) { t.Errorf("got wrong value: %v", ks) } - // Only get the REPLICA targets. - ks, err = FindAllTargets(ctx, rs, "cell1", []topodatapb.TabletType{topodatapb.TabletType_REPLICA}) + // Only get 1 keyspaces only for all types. + ks, err = FindAllTargets(ctx, rs, "cell1", []string{"test_keyspace2"}, []topodatapb.TabletType{topodatapb.TabletType_PRIMARY, topodatapb.TabletType_REPLICA}) + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if !reflect.DeepEqual(ks, []*querypb.Target{ + { + Cell: "cell1", + Keyspace: "test_keyspace2", + Shard: "test_shard1", + TabletType: topodatapb.TabletType_PRIMARY, + }, + { + Cell: "cell1", + Keyspace: "test_keyspace2", + Shard: "test_shard2", + TabletType: topodatapb.TabletType_REPLICA, + }, + }) { + t.Errorf("got wrong value: %v", ks) + } + + // Only get the REPLICA targets for any keyspace. + ks, err = FindAllTargets(ctx, rs, "cell1", []string{}, []topodatapb.TabletType{topodatapb.TabletType_REPLICA}) if err != nil { t.Errorf("unexpected error: %v", err) } diff --git a/go/vt/vtgate/tabletgateway.go b/go/vt/vtgate/tabletgateway.go index b289cf3b74e..49b9dac1128 100644 --- a/go/vt/vtgate/tabletgateway.go +++ b/go/vt/vtgate/tabletgateway.go @@ -191,7 +191,7 @@ func (gw *TabletGateway) WaitForTablets(ctx context.Context, tabletTypesToWait [ } // Finds the targets to look for. - targets, err := srvtopo.FindAllTargets(ctx, gw.srvTopoServer, gw.localCell, tabletTypesToWait) + targets, err := srvtopo.FindAllTargets(ctx, gw.srvTopoServer, gw.localCell, discovery.KeyspacesToWatch, tabletTypesToWait) if err != nil { return err }