Changes

pkg/common/scheduler/statefulset/scheduler.go

@@ -260,18 +260,13 @@ func (s *StatefulSetScheduler) removeFromExistingReplicas(state *state, logger *
 
 	if state.schedulerPolicy == EVENSPREAD_BYNODE {
 		placementsByDomain = s.getPlacementsByNodeKey(state, placements)
-
-		for nodeName := range state.nodeToZoneMap {
-			domainNames = append(domainNames, nodeName)
-		}
-		sort.Strings(domainNames) //for ordered accessing of map
 	} else {
 		placementsByDomain = s.getPlacementsByZoneKey(state, placements)
-		for zoneName := range placementsByDomain {
-			domainNames = append(domainNames, zoneName)
-		}
-		sort.Strings(domainNames) //for ordered accessing of map
 	}
+	for domainName := range placementsByDomain {
+		domainNames = append(domainNames, domainName)
+	}
+	sort.Strings(domainNames) //for ordered accessing of map
 
 	for i := 0; i < len(domainNames); i++ { //iterate through each domain
 		var totalInDomain int32
@@ -389,18 +384,13 @@ func (s *StatefulSetScheduler) addToExistingReplicas(state *state, logger *zap.S
 
 	if state.schedulerPolicy == EVENSPREAD_BYNODE {
 		placementsByDomain = s.getPlacementsByNodeKey(state, placements)
-
-		for nodeName := range state.nodeToZoneMap {
-			domainNames = append(domainNames, nodeName)
-		}
-		sort.Strings(domainNames) //for ordered accessing of map
 	} else {
 		placementsByDomain = s.getPlacementsByZoneKey(state, placements)
-		for zoneName := range placementsByDomain {
-			domainNames = append(domainNames, zoneName)
-		}
-		sort.Strings(domainNames) //for ordered accessing of map
 	}
+	for domainName := range placementsByDomain {
+		domainNames = append(domainNames, domainName)
+	}
+	sort.Strings(domainNames) //for ordered accessing of map
 
 	for i := 0; i < len(domainNames); i++ { //iterate through each domain
 		var totalInDomain int32

pkg/common/scheduler/statefulset/scheduler_test.go

@@ -491,7 +491,6 @@ func TestStatefulsetScheduler(t *testing.T) {
 					for j := int32(0); j < numNodes/numZones; j++ {
 						nodeName := "node" + fmt.Sprint((j*((numNodes/numZones)+1))+i)
 						zoneName := "zone" + fmt.Sprint(i)
-						t.Logf("cluster info %s %s", nodeName, zoneName)
 						node, err := kubeclient.Get(ctx).CoreV1().Nodes().Create(ctx, makeNode(nodeName, zoneName), metav1.CreateOptions{})
 						if err != nil {
 							t.Fatal("unexpected error", err)