Filter by keyspace earlier in tabletgateways WaitForTablets(...)

Signed-off-by: Tim Vaillancourt <[email protected]>

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

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.

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()

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)
 	}

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
 	}