Update the customizations to operate on individual pods, for more flexibility later

pkg/scheduler/framework/plugins/defaultpreemption/default_preemption.go

@@ -46,15 +46,15 @@ import (
 // Name of the plugin used in the plugin registry and configurations.
 const Name = names.DefaultPreemption
 
-// EligiblePodsFunc is a function which may be assigned to the DefaultPreemption plugin.
-// This function selects pods from the provided nodeInfo which are eligible to be preempted
-// in order to fit the provided preemptor.
-type EligiblePodsFunc func(nodeInfo *framework.NodeInfo, preemptor *v1.Pod) []*framework.PodInfo
+// EligiblePodFunc is a function which may be assigned to the DefaultPreemption plugin.
+// This function returns whether a given victim pod on the provided nodeInfo is eligible
+// to be preempted by the provided preemptor.
+type EligiblePodFunc func(nodeInfo *framework.NodeInfo, victim *framework.PodInfo, preemptor *v1.Pod) bool
 
-// OrderPodsFunc is a function which may be assigned to the DefaultPreemption plugin.
-// This function orders the provided eligible pods in-place in descending order of highest
-// to lowest priority, where pods at the start of the slice are less likely to be preempted.
-type OrderPodsFunc func(eligible []*framework.PodInfo)
+// MoreImportantPodFunc is a function which may be assigned to the DefaultPreemption plugin.
+// This function returns true if the priority of the first pod is higher than the second pod.
+// If the two pods are of equal priority then this will return false.
+type MoreImportantPodFunc func(pod1, pod2 *v1.Pod) bool
 
 // DefaultPreemption is a PostFilter plugin implements the preemption logic.
 type DefaultPreemption struct {
@@ -65,13 +65,14 @@ type DefaultPreemption struct {
 	pdbLister policylisters.PodDisruptionBudgetLister
 	Evaluator *preemption.Evaluator
 
-	// EligiblePods returns victim pods which are allowed to be preempted by the provided preemptor.
-	// The default behavior is to allow any pods of lower priority to be preempted by any pods of higher priority.
-	EligiblePods EligiblePodsFunc
+	// EligiblePod returns whether a victim pod is allowed to be preempted by a preemptor pod.
+	// The default behavior is to allow any pods of lower priority to be preempted by any pods
+	// of higher priority.
+	EligiblePod EligiblePodFunc
 
-	// OrderPods sorts eligible victims in-place in descending order of highest to lowest priority.
-	// Pods at the start of the slice are less likely to be preempted.
-	OrderPods OrderPodsFunc
+	// MoreImportantPod is used to sort eligible victims in-place in descending order of highest to
+	// lowest importance. Pods with higher importance are less likely to be preempted.
+	MoreImportantPod MoreImportantPodFunc
 }
 
 var _ framework.PostFilterPlugin = &DefaultPreemption{}
@@ -110,21 +111,12 @@ func NewDefaultPreemption(dpArgs runtime.Object, fh framework.Handle, fts featur
 	pl.Evaluator = preemption.NewEvaluator(Name, fh, &pl, fts.EnableAsyncPreemption)
 
 	// Default behavior: Any pods of lower priority may be preempted by any pods of higher priority.
-	pl.EligiblePods = func(nodeInfo *framework.NodeInfo, preemptor *v1.Pod) []*framework.PodInfo {
-		var eligible []*framework.PodInfo
-		podPriority := corev1helpers.PodPriority(preemptor)
-		for _, pi := range nodeInfo.Pods {
-			if corev1helpers.PodPriority(pi.Pod) < podPriority {
-				eligible = append(eligible, pi)
-			}
-		}
-		return eligible
+	pl.EligiblePod = func(nodeInfo *framework.NodeInfo, victim *framework.PodInfo, preemptor *v1.Pod) bool {
+		return corev1helpers.PodPriority(victim.Pod) < corev1helpers.PodPriority(preemptor)
 	}
 
 	// Default behavior: Sort by descending priority, then by descending runtime duration as secondary ordering.
-	pl.OrderPods = func(eligible []*framework.PodInfo) {
-		sort.Slice(eligible, func(i, j int) bool { return util.MoreImportantPod(eligible[i].Pod, eligible[j].Pod) })
-	}
+	pl.MoreImportantPod = util.MoreImportantPod
 
 	return &pl, nil
 }
@@ -205,6 +197,7 @@ func (pl *DefaultPreemption) SelectVictimsOnNode(
 	nodeInfo *framework.NodeInfo,
 	pdbs []*policy.PodDisruptionBudget) ([]*v1.Pod, int, *framework.Status) {
 	logger := klog.FromContext(ctx)
+	var potentialVictims []*framework.PodInfo
 	removePod := func(rpi *framework.PodInfo) error {
 		if err := nodeInfo.RemovePod(logger, rpi.Pod); err != nil {
 			return err
@@ -224,11 +217,13 @@ func (pl *DefaultPreemption) SelectVictimsOnNode(
 		return nil
 	}
 	// As the first step, remove all pods eligible for preemption from the node and
-	// check if the given pod could be scheduled without them present.
-	potentialVictims := pl.EligiblePods(nodeInfo, pod)
-	for _, pi := range potentialVictims {
-		if err := removePod(pi); err != nil {
-			return nil, 0, framework.AsStatus(err)
+	// check if the given pod can be scheduled without them present.
+	for _, pi := range nodeInfo.Pods {
+		if pl.EligiblePod(nodeInfo, pi, pod) {
+			potentialVictims = append(potentialVictims, pi)
+			if err := removePod(pi); err != nil {
+				return nil, 0, framework.AsStatus(err)
+			}
 		}
 	}
 
@@ -250,7 +245,7 @@ func (pl *DefaultPreemption) SelectVictimsOnNode(
 	numViolatingVictim := 0
 	// Sort potentialVictims by descending importance, which ensures reprieve of
 	// higher importance pods first.
-	pl.OrderPods(potentialVictims)
+	sort.Slice(potentialVictims, func(i, j int) bool { return pl.MoreImportantPod(potentialVictims[i].Pod, potentialVictims[j].Pod) })
 	// Try to reprieve as many pods as possible. We first try to reprieve the PDB
 	// violating victims and then other non-violating ones. In both cases, we start
 	// from the highest priority victims.
@@ -286,7 +281,7 @@ func (pl *DefaultPreemption) SelectVictimsOnNode(
 
 	// Sort victims after reprieving pods to keep the pods in the victims sorted in order of priority from high to low.
 	if len(violatingVictims) != 0 && len(nonViolatingVictims) != 0 {
-		pl.OrderPods(victims)
+		sort.Slice(victims, func(i, j int) bool { return pl.MoreImportantPod(victims[i].Pod, victims[j].Pod) })
 	}
 	var victimPods []*v1.Pod
 	for _, pi := range victims {
@@ -318,8 +313,8 @@ func (pl *DefaultPreemption) PodEligibleToPreemptOthers(_ context.Context, pod *
 		}
 
 		if nodeInfo, _ := nodeInfos.Get(nomNodeName); nodeInfo != nil {
-			for _, p := range pl.EligiblePods(nodeInfo, pod) {
-				if podTerminatingByPreemption(p.Pod) {
+			for _, p := range nodeInfo.Pods {
+				if pl.EligiblePod(nodeInfo, p, pod) && podTerminatingByPreemption(p.Pod) {
 					// There is a terminating pod on the nominated node.
 					return false, "not eligible due to a terminating pod on the nominated node."
 				}

pkg/scheduler/framework/plugins/defaultpreemption/default_preemption_test.go

@@ -1460,57 +1460,42 @@ func TestSelectBestCandidate(t *testing.T) {
 }
 
 func TestCustomSelection(t *testing.T) {
-	labelIsEligible := func(key, val string) EligiblePodsFunc {
-		return func(nodeInfo *framework.NodeInfo, preemptor *v1.Pod) []*framework.PodInfo {
-			eligible := []*framework.PodInfo{}
-			for _, p := range nodeInfo.Pods {
-				pval, ok := p.Pod.Labels[key]
-				if !ok {
-					continue
-				}
-				if pval == val {
-					eligible = append(eligible, p)
-				}
+	podLabelIsEligible := func(key, val string) EligiblePodFunc {
+		return func(nodeInfo *framework.NodeInfo, victim *framework.PodInfo, preemptor *v1.Pod) bool {
+			pval, ok := victim.Pod.Labels[key]
+			if !ok {
+				return false
 			}
-			return eligible
+			return pval == val
 		}
 	}
-	priorityBelowThresholdCannotPreempt := func(minPreempting int32) EligiblePodsFunc {
-		return func(nodeInfo *framework.NodeInfo, preemptor *v1.Pod) []*framework.PodInfo {
-			priority := corev1helpers.PodPriority(preemptor)
-			if priority >= minPreempting {
-				// For this example we allow all pods to be preempted, regardless of their priority
-				return nodeInfo.Pods
-			}
-			return []*framework.PodInfo{}
+	nodeNameIsEligible := func(name string) EligiblePodFunc {
+		return func(nodeInfo *framework.NodeInfo, victim *framework.PodInfo, preemptor *v1.Pod) bool {
+			return nodeInfo.Node().Name == name
 		}
 	}
-	priorityAboveThresholdCannotBePreempted := func(maxPreemptible int32) EligiblePodsFunc {
-		return func(nodeInfo *framework.NodeInfo, _ *v1.Pod) []*framework.PodInfo {
-			eligible := []*framework.PodInfo{}
-			for _, p := range nodeInfo.Pods {
-				if corev1helpers.PodPriority(p.Pod) <= maxPreemptible {
-					// For this example we allow all pods to be preempted, regardless of preemptor priority
-					eligible = append(eligible, p)
-				}
-			}
-			return eligible
+	priorityBelowThresholdCannotPreempt := func(minPreempting int32) EligiblePodFunc {
+		return func(nodeInfo *framework.NodeInfo, victim *framework.PodInfo, preemptor *v1.Pod) bool {
+			return corev1helpers.PodPriority(preemptor) >= minPreempting
+		}
+	}
+	priorityAboveThresholdCannotBePreempted := func(maxPreemptible int32) EligiblePodFunc {
+		return func(nodeInfo *framework.NodeInfo, victim *framework.PodInfo, preemptor *v1.Pod) bool {
+			return corev1helpers.PodPriority(victim.Pod) <= maxPreemptible
 		}
 	}
 
-	orderByOldestStart := func(eligible []*framework.PodInfo) {
-		sort.Slice(eligible, func(i, j int) bool {
-			return util.GetPodStartTime(eligible[i].Pod).Before(util.GetPodStartTime(eligible[j].Pod))
-		})
+	orderByOldestStart := func(pod1, pod2 *v1.Pod) bool {
+		return util.GetPodStartTime(pod1).Before(util.GetPodStartTime(pod2))
 	}
-	orderByPodName := func(eligible []*framework.PodInfo) {
-		sort.Slice(eligible, func(i, j int) bool { return eligible[i].Pod.Name < eligible[j].Pod.Name })
+	orderByPodName := func(pod1, pod2 *v1.Pod) bool {
+		return pod1.Name < pod2.Name
 	}
 
 	tests := []struct {
 		name         string
-		eligiblePods EligiblePodsFunc
-		orderPods    OrderPodsFunc
+		eligiblePods EligiblePodFunc
+		orderPods    MoreImportantPodFunc
 		nodeNames    []string
 		pod          *v1.Pod
 		pods         []*v1.Pod
@@ -1518,7 +1503,7 @@ func TestCustomSelection(t *testing.T) {
 	}{
 		{
 			name:         "only pods with specified label are preemptable: high priority",
-			eligiblePods: labelIsEligible("preemptible", "yes"),
+			eligiblePods: podLabelIsEligible("preemptible", "yes"),
 			nodeNames:    []string{"node1", "node2", "node3"},
 			pod:          st.MakePod().Name("p1").UID("p1").Priority(highPriority).Req(largeRes).Obj(),
 			pods: []*v1.Pod{
@@ -1526,12 +1511,12 @@ func TestCustomSelection(t *testing.T) {
 				st.MakePod().Name("v2").UID("v2").Label("preemptible", "yes").Node("node2").Priority(midPriority).Req(largeRes).StartTime(epochTime).Obj(),
 				st.MakePod().Name("v3").UID("v3").Node("node3").Priority(lowPriority).Req(largeRes).StartTime(epochTime).Obj(),
 			},
-			// must have "preemptible"="yes" label
+			// must have "preemptible"="yes" label, others are not preemptible despite high priority
 			expected: map[string][]string{"node2": {"v2"}},
 		},
 		{
 			name:         "only pods with specified label are preemptable: mid priority",
-			eligiblePods: labelIsEligible("preemptible", "yes"),
+			eligiblePods: podLabelIsEligible("preemptible", "yes"),
 			nodeNames:    []string{"node1", "node2", "node3"},
 			pod:          st.MakePod().Name("p2").UID("p2").Priority(midPriority).Req(largeRes).Obj(),
 			pods: []*v1.Pod{
@@ -1539,12 +1524,12 @@ func TestCustomSelection(t *testing.T) {
 				st.MakePod().Name("v2").UID("v2").Label("preemptible", "yes").Node("node2").Priority(midPriority).Req(largeRes).StartTime(epochTime).Obj(),
 				st.MakePod().Name("v3").UID("v3").Node("node3").Priority(lowPriority).Req(largeRes).StartTime(epochTime).Obj(),
 			},
-			// midPriority can preempt equal priority
+			// midPriority can preempt equal priority since we only look for matching pod label
 			expected: map[string][]string{"node2": {"v2"}},
 		},
 		{
 			name:         "only pods with specified label are preemptable: low priority",
-			eligiblePods: labelIsEligible("preemptible", "yes"),
+			eligiblePods: podLabelIsEligible("preemptible", "yes"),
 			nodeNames:    []string{"node1", "node2", "node3"},
 			pod:          st.MakePod().Name("p3").UID("p3").Priority(lowPriority).Req(largeRes).Obj(),
 			pods: []*v1.Pod{
@@ -1552,9 +1537,22 @@ func TestCustomSelection(t *testing.T) {
 				st.MakePod().Name("v2").UID("v2").Label("preemptible", "yes").Node("node2").Priority(midPriority).Req(largeRes).StartTime(epochTime).Obj(),
 				st.MakePod().Name("v3").UID("v3").Node("node3").Priority(lowPriority).Req(largeRes).StartTime(epochTime).Obj(),
 			},
-			// lowPriority can preempt higher priority
+			// lowPriority can preempt higher priority since we only look for matching pod label
 			expected: map[string][]string{"node2": {"v2"}},
 		},
+		{
+			name:         "only pods on specified node name are preemptable",
+			eligiblePods: nodeNameIsEligible("node1"),
+			nodeNames:    []string{"node1", "node2", "node3"},
+			pod:          st.MakePod().Name("p3").UID("p3").Priority(lowPriority).Req(largeRes).Obj(),
+			pods: []*v1.Pod{
+				st.MakePod().Name("v1").UID("v1").Node("node1").Priority(midPriority).Req(largeRes).StartTime(epochTime).Obj(),
+				st.MakePod().Name("v2").UID("v2").Node("node2").Priority(midPriority).Req(largeRes).StartTime(epochTime).Obj(),
+				st.MakePod().Name("v3").UID("v3").Node("node3").Priority(lowPriority).Req(largeRes).StartTime(epochTime).Obj(),
+			},
+			// lowPriority can preempt higher priority since we only look for matching node name
+			expected: map[string][]string{"node1": {"v1"}},
+		},
 		{
 			name:         "only pods at or above specified priority can preempted: high priority",
 			eligiblePods: priorityBelowThresholdCannotPreempt(highPriority),
@@ -1701,10 +1699,10 @@ func TestCustomSelection(t *testing.T) {
 			}
 			// Override selection logic
 			if tt.eligiblePods != nil {
-				pl.EligiblePods = tt.eligiblePods
+				pl.EligiblePod = tt.eligiblePods
 			}
 			if tt.orderPods != nil {
-				pl.OrderPods = tt.orderPods
+				pl.MoreImportantPod = tt.orderPods
 			}
 			offset, numCandidates := pl.GetOffsetAndNumCandidates(int32(len(nodeInfos)))
 			candidates, _, _ := pl.Evaluator.DryRunPreemption(ctx, state, tt.pod, nodeInfos, nil, offset, numCandidates)