Update the region.

Signed-off-by: corvofeng <corvofeng@gmail.com>
corvofeng · May 19, 2024 · efe8d53 · efe8d53
1 parent 3a07309
commit efe8d53
Showing 1 changed file with 35 additions and 22 deletions.
diff --git a/lib/cloud_provider/km_aws/eks.go b/lib/cloud_provider/km_aws/eks.go
@@ -164,29 +164,42 @@ func (c *AWSProvider) ListRegions() ([]string, error) {
 	return regions, nil
 }
 
-func (c *AWSProvider) ListClusters(regions []string, setProgress func(int)) ([]*cluster.CPCluster, error) {
-	var wg sync.WaitGroup
-	totalRegions := len(regions)
-	completedRegions := 0
+type ClusterGetter interface {
+	GetClusters(ch chan<- *cluster.CPCluster)
+}
+
+func getEKSClients(regions []string) []ClusterGetter {
+	clients := make([]ClusterGetter, 0, len(regions))
 
-	clustersChan := make(chan *cluster.CPCluster)
 	for _, region := range regions {
-		wg.Add(1)
-		regionCh := make(chan *cluster.CPCluster)
-		go func(r string) {
-			k, err := NewEKS(r)
-			if err != nil {
-				log.WithFields(log.Fields{
-					"region": r,
-					"error":  err.Error(),
-				}).Error("Failed to create EKS client")
-				return
-			}
+		log.WithFields(log.Fields{
+			"region": region,
+		}).Debug("Initialize client")
+		eks, err := NewEKS(region)
+		if err != nil {
 			log.WithFields(log.Fields{
-				"region": r,
-			}).Debug("Start Query clusters")
-			k.GetClusters(regionCh)
-		}(region)
+				"region": region,
+				"error":  err.Error(),
+			}).Error("Failed to create AWS SDK session")
+			continue
+		}
+
+		clients = append(clients, ClusterGetter(eks))
+	}
+	return clients
+}
+
+func (c *AWSProvider) ListClusters(regions []string, setProgress func(int)) ([]*cluster.CPCluster, error) {
+	clients := getEKSClients(regions)
+	totalRegionClients := len(clients)
+	completedRegionClients := 0
+	clustersChan := make(chan *cluster.CPCluster)
+
+	var wg sync.WaitGroup
+	wg.Add(totalRegionClients)
+	for _, c := range clients {
+		regionCh := make(chan *cluster.CPCluster)
+		go c.GetClusters(regionCh)
 
 		go func(clusterCh chan *cluster.CPCluster) {
 			defer wg.Done()
@@ -205,8 +218,8 @@ func (c *AWSProvider) ListClusters(regions []string, setProgress func(int)) ([]*
 	var allClusters []*cluster.CPCluster
 	for cluster := range clustersChan {
 		allClusters = append(allClusters, cluster)
-		completedRegions++
-		progress := (completedRegions * 100) / totalRegions
+		completedRegionClients++
+		progress := (completedRegionClients * 100) / totalRegionClients
 		if setProgress != nil {
 			setProgress(progress)
 		}