xds/cdsbalancer: correctly remove the unwanted cds watchers (#8428)

Author: eshitachandwani

xds/internal/balancer/cdsbalancer/aggregate_cluster_test.go

@@ -866,3 +866,53 @@ func (s) TestAggregatedCluster_CycleWithLeafNode(t *testing.T) {
 		t.Fatalf("EmptyCall() failed: %v", err)
 	}
 }
+
+// Tests the scenario where the cluster tree changes, and verifies that the
+// watchers for the cds balancer are updated accordingly. That is the cluster
+// removed from the tree no longer has a watcher and the new cluster added has a
+// new watcher.
+func (s) TestWatchers(t *testing.T) {
+	mgmtServer, nodeID, _, _, _, cdsResourceRequestedCh, _ := setupWithManagementServer(t)
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	const (
+		clusterA = clusterName
+		clusterB = clusterName + "-B"
+		clusterC = clusterName + "-C"
+		clusterD = clusterName + "-D"
+	)
+
+	// Initial CDS resources: A -> B, C
+	initialResources := e2e.UpdateOptions{
+		NodeID: nodeID,
+		Clusters: []*v3clusterpb.Cluster{
+			makeAggregateClusterResource(clusterA, []string{clusterB, clusterC}),
+		},
+		SkipValidation: true,
+	}
+	if err := mgmtServer.Update(ctx, initialResources); err != nil {
+		t.Fatalf("Update failed: %v", err)
+	}
+	wantNames := []string{clusterA, clusterB, clusterC}
+	if err := waitForResourceNames(ctx, cdsResourceRequestedCh, wantNames); err != nil {
+		t.Fatal(err)
+	}
+
+	// Update the CDS resources to remove cluster C and add cluster D.
+	updatedResources := e2e.UpdateOptions{
+		NodeID: nodeID,
+		Clusters: []*v3clusterpb.Cluster{
+			makeAggregateClusterResource(clusterA, []string{clusterB, clusterD}),
+		},
+		SkipValidation: true,
+	}
+	if err := mgmtServer.Update(ctx, updatedResources); err != nil {
+		t.Fatalf("Update failed: %v", err)
+	}
+	wantNames = []string{clusterA, clusterB, clusterD}
+	if err := waitForResourceNames(ctx, cdsResourceRequestedCh, wantNames); err != nil {
+		t.Fatal(err)
+	}
+}

xds/internal/balancer/cdsbalancer/cdsbalancer.go

@@ -530,13 +530,11 @@ func (b *cdsBalancer) onClusterUpdate(name string, update xdsresource.ClusterUpd
 	}
 	// We no longer need the clusters that we did not see in this iteration of
 	// generateDMsForCluster().
-	for cluster := range clustersSeen {
-		state, ok := b.watchers[cluster]
-		if ok {
-			continue
+	for cluster, state := range b.watchers {
+		if !clustersSeen[cluster] {
+			state.cancelWatch()
+			delete(b.watchers, cluster)
 		}
-		state.cancelWatch()
-		delete(b.watchers, cluster)
 	}
 }
 

xds/internal/balancer/cdsbalancer/cdsbalancer_test.go

@@ -26,6 +26,7 @@ import (
 	"time"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/google/uuid"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/balancer"
@@ -85,7 +86,9 @@ func waitForResourceNames(ctx context.Context, resourceNamesCh chan []string, wa
 		select {
 		case <-ctx.Done():
 		case gotNames := <-resourceNamesCh:
-			if cmp.Equal(gotNames, wantNames) {
+			// Sort both slices before comparing them, as the order of clusters
+			// does not matter.
+			if cmp.Equal(gotNames, wantNames, cmpopts.SortSlices(func(a, b string) bool { return a < b })) {
 				return nil
 			}
 		}