Sort constraints by SpreadByField

jd-opensource · Aug 14, 2024 · 7eab13a · 7eab13a
1 parent f287669
commit 7eab13a
Show file tree

Hide file tree

Showing 2 changed files with 106 additions and 2 deletions.
diff --git a/pkg/scheduler/core/spreadconstraint/group_builder.go b/pkg/scheduler/core/spreadconstraint/group_builder.go
@@ -37,7 +37,8 @@ func init() {
 func (builder *groupBuilder) Create(ctx SelectionCtx) (Selection, error) {
 	root := &groupRoot{}
 	root.Name = ""
-	root.DisableConstraint = len(ctx.Placement.SpreadConstraints) == 0 || disableSpreadConstraint(ctx.Placement)
+	constraints := ctx.Placement.SpreadConstraints
+	root.DisableConstraint = len(constraints) == 0 || disableSpreadConstraint(ctx.Placement)
 	root.Replicas = ctx.Spec.Replicas
 	if root.DisableConstraint {
 		root.Replicas = InvalidReplicas
@@ -52,7 +53,8 @@ func (builder *groupBuilder) Create(ctx SelectionCtx) (Selection, error) {
 		}
 		root.Clusters = createClustersWithReplicas(ctx.ClusterScores, ctx.Spec, replicasFunc)
 		sortClusters(root.Clusters)
-		born(&root.groupNode, ctx.Placement.SpreadConstraints, 0)
+
+		born(&root.groupNode, sortConstraints(constraints), 0)
 	}
 
 	return root, nil
@@ -105,6 +107,35 @@ func disableAvailableResource(placement *policyv1alpha1.Placement) bool {
 	return false
 }
 
+// sortConstraints sorts the given slice of SpreadConstraints.
+func sortConstraints(constraints []policyv1alpha1.SpreadConstraint) []policyv1alpha1.SpreadConstraint {
+	var sorts = make([]policyv1alpha1.SpreadConstraint, len(constraints))
+	copy(sorts, constraints)
+	sort.Slice(sorts, func(i, j int) bool {
+		if sorts[i].SpreadByField == policyv1alpha1.SpreadByFieldCluster {
+			return false
+		} else if sorts[j].SpreadByField == policyv1alpha1.SpreadByFieldCluster {
+			return true
+		} else if sorts[i].SpreadByField != "" && sorts[j].SpreadByField != "" {
+			if sorts[i].SpreadByField == policyv1alpha1.SpreadByFieldProvider {
+				return true
+			} else if sorts[j].SpreadByField == policyv1alpha1.SpreadByFieldProvider {
+				return false
+			} else if sorts[i].SpreadByField == policyv1alpha1.SpreadByFieldRegion {
+				return true
+			} else if sorts[j].SpreadByField == policyv1alpha1.SpreadByFieldRegion {
+				return false
+			} else if sorts[i].SpreadByField == policyv1alpha1.SpreadByFieldZone {
+				return true
+			} else if sorts[j].SpreadByField == policyv1alpha1.SpreadByFieldZone {
+				return false
+			}
+		}
+		return i < j
+	})
+	return sorts
+}
+
 // born recursively groups clusters based on the provided constraints.
 // It updates the parent group with its child groups and their respective clusters.
 // If the current group is at the last constraint, it marks it as a leaf node.

diff --git a/pkg/scheduler/core/spreadconstraint/group_builder_test.go b/pkg/scheduler/core/spreadconstraint/group_builder_test.go
@@ -347,6 +347,79 @@ func Test_GroupClusters(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "test SpreadConstraints is zone/region/provider",
+			ctx: SelectionCtx{
+				ClusterScores: generateClusterScore(),
+				Placement: &policyv1alpha1.Placement{
+					SpreadConstraints: []policyv1alpha1.SpreadConstraint{
+						{
+							SpreadByField: policyv1alpha1.SpreadByFieldZone,
+							MaxGroups:     1,
+							MinGroups:     1,
+						},
+						{
+							SpreadByField: policyv1alpha1.SpreadByFieldRegion,
+							MaxGroups:     1,
+							MinGroups:     1,
+						},
+						{
+							SpreadByField: policyv1alpha1.SpreadByFieldProvider,
+							MaxGroups:     1,
+							MinGroups:     1,
+						},
+					},
+				},
+				Spec:         &workv1alpha2.ResourceBindingSpec{},
+				ReplicasFunc: replicasFunc,
+			},
+			want: want{
+				name:     "",
+				clusters: []string{"member4", "member2", "member3", "member1"},
+				groups: []want{
+					{
+						name:     "P2",
+						clusters: []string{"member4", "member3"},
+						groups: []want{
+							{
+								name:     "R2",
+								clusters: []string{"member4", "member3"},
+								groups: []want{
+									{
+										name:     "Z4",
+										clusters: []string{"member4"},
+									},
+									{
+										name:     "Z3",
+										clusters: []string{"member3"},
+									},
+								},
+							},
+						},
+					},
+					{
+						name:     "P1",
+						clusters: []string{"member2", "member1"},
+						groups: []want{
+							{
+								name:     "R1",
+								clusters: []string{"member2", "member1"},
+								groups: []want{
+									{
+										name:     "Z2",
+										clusters: []string{"member2"},
+									},
+									{
+										name:     "Z1",
+										clusters: []string{"member1"},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
 		{
 			name: "test SpreadConstraints is label unit/cell",
 			ctx: SelectionCtx{