Merge pull request #77760 from Huang-Wei/eps-pred-meta

Even Pods Spread - 2. Calculating Predicates Metadata

pkg/scheduler/algorithm/predicates/metadata.go

@@ -19,6 +19,7 @@ package predicates
 import (
 	"context"
 	"fmt"
+	"math"
 	"sync"
 
 	"k8s.io/klog"
@@ -66,6 +67,17 @@ type topologyPairsMaps struct {
 	podToTopologyPairs map[string]topologyPairSet
 }
 
+// topologyPairsPodSpreadMap combines topologyKeyToMinPodsMap and topologyPairsMaps
+// to represent:
+// (1) minimum number of pods matched on the spread constraints.
+// (2) how existing pods match incoming pod on its spread constraints.
+type topologyPairsPodSpreadMap struct {
+	// This map is keyed with a topology key, and valued with minimum number
+	// of pods matched on that topology domain.
+	topologyKeyToMinPodsMap map[string]int32
+	*topologyPairsMaps
+}
+
 // NOTE: When new fields are added/removed or logic is changed, please make sure that
 // RemovePod, AddPod, and ShallowCopy functions are updated to work with the new changes.
 type predicateMetadata struct {
@@ -91,6 +103,9 @@ type predicateMetadata struct {
 	// which should be accounted only by the extenders. This set is synthesized
 	// from scheduler extender configuration and does not change per pod.
 	ignoredExtendedResources sets.String
+	// Similar to the map for pod (anti-)affinity, but imposes additional min matches info
+	// to describe minimum match number on each topology spread constraint.
+	topologyPairsPodSpreadMap *topologyPairsPodSpreadMap
 }
 
 // Ensure that predicateMetadata implements algorithm.PredicateMetadata.
@@ -137,17 +152,24 @@ func (pfactory *PredicateMetadataFactory) GetMetadata(pod *v1.Pod, nodeNameToInf
 	if pod == nil {
 		return nil
 	}
+	// existingPodSpreadConstraintsMap represents how existing pods match "pod"
+	// on its spread constraints
+	existingPodSpreadConstraintsMap, err := getTPMapMatchingSpreadConstraints(pod, nodeNameToInfoMap)
+	if err != nil {
+		klog.Errorf("Error calculating spreadConstraintsMap: %v", err)
+		return nil
+	}
 	// existingPodAntiAffinityMap will be used later for efficient check on existing pods' anti-affinity
 	existingPodAntiAffinityMap, err := getTPMapMatchingExistingAntiAffinity(pod, nodeNameToInfoMap)
 	if err != nil {
-		klog.Errorf("[predicate meta data generation] error finding pods whose affinity terms are matched: %v", err)
+		klog.Errorf("Error calculating existingPodAntiAffinityMap: %v", err)
 		return nil
 	}
 	// incomingPodAffinityMap will be used later for efficient check on incoming pod's affinity
 	// incomingPodAntiAffinityMap will be used later for efficient check on incoming pod's anti-affinity
 	incomingPodAffinityMap, incomingPodAntiAffinityMap, err := getTPMapMatchingIncomingAffinityAntiAffinity(pod, nodeNameToInfoMap)
 	if err != nil {
-		klog.Errorf("[predicate meta data generation] error finding pods that match affinity terms: %v", err)
+		klog.Errorf("Error calculating incomingPod(Anti)AffinityMap: %v", err)
 		return nil
 	}
 	predicateMetadata := &predicateMetadata{
@@ -158,6 +180,7 @@ func (pfactory *PredicateMetadataFactory) GetMetadata(pod *v1.Pod, nodeNameToInf
 		topologyPairsPotentialAffinityPods:     incomingPodAffinityMap,
 		topologyPairsPotentialAntiAffinityPods: incomingPodAntiAffinityMap,
 		topologyPairsAntiAffinityPodsMap:       existingPodAntiAffinityMap,
+		topologyPairsPodSpreadMap:              existingPodSpreadConstraintsMap,
 	}
 	for predicateName, precomputeFunc := range predicateMetadataProducers {
 		klog.V(10).Infof("Precompute: %v", predicateName)
@@ -166,44 +189,204 @@ func (pfactory *PredicateMetadataFactory) GetMetadata(pod *v1.Pod, nodeNameToInf
 	return predicateMetadata
 }
 
+func getTPMapMatchingSpreadConstraints(pod *v1.Pod, nodeInfoMap map[string]*schedulernodeinfo.NodeInfo) (*topologyPairsPodSpreadMap, error) {
+	// We have feature gating in APIServer to strip the spec
+	// so don't need to re-check feature gate, just check length of constraints.
+	constraints := getHardTopologySpreadConstraints(pod)
+	if len(constraints) == 0 {
+		return nil, nil
+	}
+
+	allNodeNames := make([]string, 0, len(nodeInfoMap))
+	for name := range nodeInfoMap {
+		allNodeNames = append(allNodeNames, name)
+	}
+
+	errCh := schedutil.NewErrorChannel()
+	var lock sync.Mutex
+
+	topologyPairsPodSpreadMap := &topologyPairsPodSpreadMap{
+		// topologyKeyToMinPodsMap will be initialized with proper size later.
+		topologyPairsMaps: newTopologyPairsMaps(),
+	}
+
+	appendTopologyPairsMaps := func(toAppend *topologyPairsMaps) {
+		lock.Lock()
+		topologyPairsPodSpreadMap.appendMaps(toAppend)
+		lock.Unlock()
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	processNode := func(i int) {
+		nodeInfo := nodeInfoMap[allNodeNames[i]]
+		node := nodeInfo.Node()
+		if node == nil {
+			klog.Errorf("node %q not found", allNodeNames[i])
+			return
+		}
+		// In accordance to design, if NodeAffinity or NodeSelector is defined,
+		// spreading is applied to nodes that pass those filters.
+		if !podMatchesNodeSelectorAndAffinityTerms(pod, node) {
+			return
+		}
+		// Ensure current node's labels contains all topologyKeys in 'constraints'.
+		for _, constraint := range constraints {
+			if _, ok := node.Labels[constraint.TopologyKey]; !ok {
+				return
+			}
+		}
+
+		nodeTopologyMaps := newTopologyPairsMaps()
+		// nodeInfo.Pods() can be empty; or all pods don't fit
+		for _, existingPod := range nodeInfo.Pods() {
+			if existingPod.Namespace != pod.Namespace {
+				continue
+			}
+			podLabelSet := labels.Set(existingPod.Labels)
+			for _, constraint := range constraints {
+				ok, err := podMatchesSpreadConstraint(podLabelSet, constraint)
+				if err != nil {
+					errCh.SendErrorWithCancel(err, cancel)
+					return
+				}
+				if ok {
+					// constraint.TopologyKey is already guaranteed to be present
+					pair := topologyPair{key: constraint.TopologyKey, value: node.Labels[constraint.TopologyKey]}
+					nodeTopologyMaps.addTopologyPair(pair, existingPod)
+				}
+			}
+		}
+		// If needed, append topology pair without entry of pods.
+		// For example, on node-x, there is no pod matching spread constraints,
+		// but node-x should be also considered as a match (with match number 0)
+		// i.e. <node: node-x>: {}
+		for _, constraint := range constraints {
+			// constraint.TopologyKey is already guaranteed to be present
+			pair := topologyPair{
+				key:   constraint.TopologyKey,
+				value: node.Labels[constraint.TopologyKey],
+			}
+			// addTopologyPairWithoutPods is a non-op if other pods match this pair
+			nodeTopologyMaps.addTopologyPairWithoutPods(pair)
+		}
+
+		appendTopologyPairsMaps(nodeTopologyMaps)
+	}
+	workqueue.ParallelizeUntil(ctx, 16, len(allNodeNames), processNode)
+
+	if err := errCh.ReceiveError(); err != nil {
+		return nil, err
+	}
+
+	// calculate min match for each topology pair
+	topologyPairsPodSpreadMap.topologyKeyToMinPodsMap = make(map[string]int32, len(constraints))
+	for _, constraint := range constraints {
+		topologyPairsPodSpreadMap.topologyKeyToMinPodsMap[constraint.TopologyKey] = math.MaxInt32
+	}
+	for pair, podSet := range topologyPairsPodSpreadMap.topologyPairToPods {
+		// TODO(Huang-Wei): short circuit unvisited portions of <topologyKey: any value>
+		// if we already see 0 as min match of that topologyKey.
+		if l := int32(len(podSet)); l < topologyPairsPodSpreadMap.topologyKeyToMinPodsMap[pair.key] {
+			topologyPairsPodSpreadMap.topologyKeyToMinPodsMap[pair.key] = l
+		}
+	}
+	return topologyPairsPodSpreadMap, nil
+}
+
+func getHardTopologySpreadConstraints(pod *v1.Pod) (constraints []v1.TopologySpreadConstraint) {
+	if pod != nil {
+		for _, constraint := range pod.Spec.TopologySpreadConstraints {
+			if constraint.WhenUnsatisfiable == v1.DoNotSchedule {
+				constraints = append(constraints, constraint)
+			}
+		}
+	}
+	return
+}
+
+// some corner cases:
+// 1. podLabelSet = nil => returns (false, nil)
+// 2. constraint.LabelSelector = nil => returns (false, nil)
+func podMatchesSpreadConstraint(podLabelSet labels.Set, constraint v1.TopologySpreadConstraint) (bool, error) {
+	selector, err := metav1.LabelSelectorAsSelector(constraint.LabelSelector)
+	if err != nil {
+		return false, err
+	}
+	if !selector.Matches(podLabelSet) {
+		return false, nil
+	}
+	return true, nil
+}
+
 // returns a pointer to a new topologyPairsMaps
 func newTopologyPairsMaps() *topologyPairsMaps {
 	return &topologyPairsMaps{topologyPairToPods: make(map[topologyPair]podSet),
 		podToTopologyPairs: make(map[string]topologyPairSet)}
 }
 
-func (topologyPairsMaps *topologyPairsMaps) addTopologyPair(pair topologyPair, pod *v1.Pod) {
+func (m *topologyPairsMaps) addTopologyPair(pair topologyPair, pod *v1.Pod) {
 	podFullName := schedutil.GetPodFullName(pod)
-	if topologyPairsMaps.topologyPairToPods[pair] == nil {
+	m.addTopologyPairWithoutPods(pair)
-		topologyPairsMaps.topologyPairToPods[pair] = make(map[*v1.Pod]struct{})
+	m.topologyPairToPods[pair][pod] = struct{}{}
+	if m.podToTopologyPairs[podFullName] == nil {
+		m.podToTopologyPairs[podFullName] = make(map[topologyPair]struct{})
 	}
-	topologyPairsMaps.topologyPairToPods[pair][pod] = struct{}{}
+	m.podToTopologyPairs[podFullName][pair] = struct{}{}
-	if topologyPairsMaps.podToTopologyPairs[podFullName] == nil {
-		topologyPairsMaps.podToTopologyPairs[podFullName] = make(map[topologyPair]struct{})
-	}
-	topologyPairsMaps.podToTopologyPairs[podFullName][pair] = struct{}{}
 }
 
-func (topologyPairsMaps *topologyPairsMaps) removePod(deletedPod *v1.Pod) {
+// add a topology pair holder if needed
+func (m *topologyPairsMaps) addTopologyPairWithoutPods(pair topologyPair) {
+	if m.topologyPairToPods[pair] == nil {
+		m.topologyPairToPods[pair] = make(map[*v1.Pod]struct{})
+	}
+}
+
+func (m *topologyPairsMaps) removePod(deletedPod *v1.Pod) {
 	deletedPodFullName := schedutil.GetPodFullName(deletedPod)
-	for pair := range topologyPairsMaps.podToTopologyPairs[deletedPodFullName] {
+	for pair := range m.podToTopologyPairs[deletedPodFullName] {
-		delete(topologyPairsMaps.topologyPairToPods[pair], deletedPod)
+		delete(m.topologyPairToPods[pair], deletedPod)
-		if len(topologyPairsMaps.topologyPairToPods[pair]) == 0 {
+		if len(m.topologyPairToPods[pair]) == 0 {
-			delete(topologyPairsMaps.topologyPairToPods, pair)
+			delete(m.topologyPairToPods, pair)
 		}
 	}
-	delete(topologyPairsMaps.podToTopologyPairs, deletedPodFullName)
+	delete(m.podToTopologyPairs, deletedPodFullName)
 }
 
-func (topologyPairsMaps *topologyPairsMaps) appendMaps(toAppend *topologyPairsMaps) {
+func (m *topologyPairsMaps) appendMaps(toAppend *topologyPairsMaps) {
 	if toAppend == nil {
 		return
 	}
 	for pair := range toAppend.topologyPairToPods {
-		for pod := range toAppend.topologyPairToPods[pair] {
+		if podSet := toAppend.topologyPairToPods[pair]; len(podSet) == 0 {
-			topologyPairsMaps.addTopologyPair(pair, pod)
+			m.addTopologyPairWithoutPods(pair)
+		} else {
+			for pod := range podSet {
+				m.addTopologyPair(pair, pod)
 			}
 		}
+	}
+}
+
+func (m *topologyPairsMaps) clone() *topologyPairsMaps {
+	copy := newTopologyPairsMaps()
+	copy.appendMaps(m)
+	return copy
+}
+
+func (m *topologyPairsPodSpreadMap) clone() *topologyPairsPodSpreadMap {
+	// m could be nil when EvenPodsSpread feature is disabled
+	if m == nil {
+		return nil
+	}
+	copy := &topologyPairsPodSpreadMap{
+		topologyKeyToMinPodsMap: make(map[string]int32),
+		topologyPairsMaps:       m.topologyPairsMaps.clone(),
+	}
+	for key, minMatched := range m.topologyKeyToMinPodsMap {
+		copy.topologyKeyToMinPodsMap[key] = minMatched
+	}
+	return copy
 }
 
 // RemovePod changes predicateMetadata assuming that the given `deletedPod` is
@@ -301,12 +484,10 @@ func (meta *predicateMetadata) ShallowCopy() PredicateMetadata {
 		ignoredExtendedResources: meta.ignoredExtendedResources,
 	}
 	newPredMeta.podPorts = append([]*v1.ContainerPort(nil), meta.podPorts...)
-	newPredMeta.topologyPairsPotentialAffinityPods = newTopologyPairsMaps()
+	newPredMeta.topologyPairsPotentialAffinityPods = meta.topologyPairsPotentialAffinityPods.clone()
-	newPredMeta.topologyPairsPotentialAffinityPods.appendMaps(meta.topologyPairsPotentialAffinityPods)
+	newPredMeta.topologyPairsPotentialAntiAffinityPods = meta.topologyPairsPotentialAntiAffinityPods.clone()
-	newPredMeta.topologyPairsPotentialAntiAffinityPods = newTopologyPairsMaps()
+	newPredMeta.topologyPairsAntiAffinityPodsMap = meta.topologyPairsAntiAffinityPodsMap.clone()
-	newPredMeta.topologyPairsPotentialAntiAffinityPods.appendMaps(meta.topologyPairsPotentialAntiAffinityPods)
+	newPredMeta.topologyPairsPodSpreadMap = meta.topologyPairsPodSpreadMap.clone()
-	newPredMeta.topologyPairsAntiAffinityPodsMap = newTopologyPairsMaps()
-	newPredMeta.topologyPairsAntiAffinityPodsMap.appendMaps(meta.topologyPairsAntiAffinityPodsMap)
 	newPredMeta.serviceAffinityMatchingPodServices = append([]*v1.Service(nil),
 		meta.serviceAffinityMatchingPodServices...)
 	newPredMeta.serviceAffinityMatchingPodList = append([]*v1.Pod(nil),

pkg/scheduler/algorithm/predicates/metadata_test.go

@@ -24,8 +24,9 @@ import (
 
 	"k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/labels"
 	schedulernodeinfo "k8s.io/kubernetes/pkg/scheduler/nodeinfo"
-	schedulertesting "k8s.io/kubernetes/pkg/scheduler/testing"
+	st "k8s.io/kubernetes/pkg/scheduler/testing"
 )
 
 // sortablePods lets us to sort pods.
@@ -352,16 +353,16 @@ func TestPredicateMetadata_AddRemovePod(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			allPodLister := schedulertesting.FakePodLister(append(test.existingPods, test.addedPod))
+			allPodLister := st.FakePodLister(append(test.existingPods, test.addedPod))
 			// getMeta creates predicate meta data given the list of pods.
-			getMeta := func(lister schedulertesting.FakePodLister) (*predicateMetadata, map[string]*schedulernodeinfo.NodeInfo) {
+			getMeta := func(lister st.FakePodLister) (*predicateMetadata, map[string]*schedulernodeinfo.NodeInfo) {
 				nodeInfoMap := schedulernodeinfo.CreateNodeNameToInfoMap(lister, test.nodes)
 				// nodeList is a list of non-pointer nodes to feed to FakeNodeListInfo.
 				nodeList := []v1.Node{}
 				for _, n := range test.nodes {
 					nodeList = append(nodeList, *n)
 				}
-				_, precompute := NewServiceAffinityPredicate(lister, schedulertesting.FakeServiceLister(test.services), FakeNodeListInfo(nodeList), nil)
+				_, precompute := NewServiceAffinityPredicate(lister, st.FakeServiceLister(test.services), FakeNodeListInfo(nodeList), nil)
 				RegisterPredicateMetadataProducer("ServiceAffinityMetaProducer", precompute)
 				pmf := PredicateMetadataFactory{lister}
 				meta := pmf.GetMetadata(test.pendingPod, nodeInfoMap)
@@ -372,7 +373,7 @@ func TestPredicateMetadata_AddRemovePod(t *testing.T) {
 			// are given to the metadata producer.
 			allPodsMeta, _ := getMeta(allPodLister)
 			// existingPodsMeta1 is meta data produced for test.existingPods (without test.addedPod).
-			existingPodsMeta1, nodeInfoMap := getMeta(schedulertesting.FakePodLister(test.existingPods))
+			existingPodsMeta1, nodeInfoMap := getMeta(st.FakePodLister(test.existingPods))
 			// Add test.addedPod to existingPodsMeta1 and make sure meta is equal to allPodsMeta
 			nodeInfo := nodeInfoMap[test.addedPod.Spec.NodeName]
 			if err := existingPodsMeta1.AddPod(test.addedPod, nodeInfo); err != nil {
@@ -383,7 +384,7 @@ func TestPredicateMetadata_AddRemovePod(t *testing.T) {
 			}
 			// Remove the added pod and from existingPodsMeta1 an make sure it is equal
 			// to meta generated for existing pods.
-			existingPodsMeta2, _ := getMeta(schedulertesting.FakePodLister(test.existingPods))
+			existingPodsMeta2, _ := getMeta(st.FakePodLister(test.existingPods))
 			if err := existingPodsMeta1.RemovePod(test.addedPod); err != nil {
 				t.Errorf("error removing pod from meta: %v", err)
 			}
@@ -511,6 +512,39 @@ func TestPredicateMetadata_ShallowCopy(t *testing.T) {
 				},
 			},
 		},
+		topologyPairsPodSpreadMap: &topologyPairsPodSpreadMap{
+			topologyKeyToMinPodsMap: map[string]int32{"name": 1},
+			topologyPairsMaps: &topologyPairsMaps{
+				topologyPairToPods: map[topologyPair]podSet{
+					{key: "name", value: "nodeA"}: {
+						&v1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "p1", Labels: selector1},
+							Spec: v1.PodSpec{NodeName: "nodeA"},
+						}: struct{}{},
+					},
+					{key: "name", value: "nodeC"}: {
+						&v1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "p2"},
+							Spec: v1.PodSpec{
+								NodeName: "nodeC",
+							},
+						}: struct{}{},
+						&v1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "p6", Labels: selector1},
+							Spec: v1.PodSpec{NodeName: "nodeC"},
+						}: struct{}{},
+					},
+				},
+				podToTopologyPairs: map[string]topologyPairSet{
+					"p1_": {
+						topologyPair{key: "name", value: "nodeA"}: struct{}{},
+					},
+					"p2_": {
+						topologyPair{key: "name", value: "nodeC"}: struct{}{},
+					},
+					"p6_": {
+						topologyPair{key: "name", value: "nodeC"}: struct{}{},
+					},
+				},
+			},
+		},
 		serviceAffinityInUse: true,
 		serviceAffinityMatchingPodList: []*v1.Pod{
 			{ObjectMeta: metav1.ObjectMeta{Name: "pod1"}},
@@ -791,3 +825,434 @@ func TestGetTPMapMatchingIncomingAffinityAntiAffinity(t *testing.T) {
 		})
 	}
 }
+
+func TestPodMatchesSpreadConstraint(t *testing.T) {
+	tests := []struct {
+		name       string
+		podLabels  map[string]string
+		constraint v1.TopologySpreadConstraint
+		want       bool
+		wantErr    bool
+	}{
+		{
+			name:      "normal match",
+			podLabels: map[string]string{"foo": "", "bar": ""},
+			constraint: v1.TopologySpreadConstraint{
+				LabelSelector: &metav1.LabelSelector{
+					MatchExpressions: []metav1.LabelSelectorRequirement{
+						{
+							Key:      "foo",
+							Operator: metav1.LabelSelectorOpExists,
+						},
+					},
+				},
+			},
+			want: true,
+		},
+		{
+			name:      "normal mismatch",
+			podLabels: map[string]string{"foo": "", "baz": ""},
+			constraint: v1.TopologySpreadConstraint{
+				LabelSelector: &metav1.LabelSelector{
+					MatchExpressions: []metav1.LabelSelectorRequirement{
+						{
+							Key:      "foo",
+							Operator: metav1.LabelSelectorOpExists,
+						},
+						{
+							Key:      "bar",
+							Operator: metav1.LabelSelectorOpExists,
+						},
+					},
+				},
+			},
+			want: false,
+		},
+		{
+			name: "podLabels is nil",
+			constraint: v1.TopologySpreadConstraint{
+				LabelSelector: &metav1.LabelSelector{
+					MatchExpressions: []metav1.LabelSelectorRequirement{
+						{
+							Key:      "foo",
+							Operator: metav1.LabelSelectorOpExists,
+						},
+					},
+				},
+			},
+			want: false,
+		},
+		{
+			name: "constraint.LabelSelector is nil",
+			podLabels: map[string]string{
+				"foo": "",
+				"bar": "",
+			},
+			constraint: v1.TopologySpreadConstraint{
+				MaxSkew: 1,
+			},
+			want: false,
+		},
+		{
+			name: "both podLabels and constraint.LabelSelector are nil",
+			constraint: v1.TopologySpreadConstraint{
+				MaxSkew: 1,
+			},
+			want: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			podLabelSet := labels.Set(tt.podLabels)
+			got, err := podMatchesSpreadConstraint(podLabelSet, tt.constraint)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("podMatchesSpreadConstraint() error = %v, wantErr %v", err, tt.wantErr)
+			}
+			if got != tt.want {
+				t.Errorf("podMatchesSpreadConstraint() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestGetTPMapMatchingSpreadConstraints(t *testing.T) {
+	// we need to inject the exact pod pointers to want.topologyPairsMaps.topologyPairToPods
+	// otherwise, *pod (as key of a map) will always fail in reflect.DeepEqual()
+	tests := []struct {
+		name              string
+		pod               *v1.Pod
+		nodes             []*v1.Node
+		existingPods      []*v1.Pod
+		injectPodPointers map[topologyPair][]int
+		want              *topologyPairsPodSpreadMap
+	}{
+		{
+			name: "clean cluster with one spreadConstraint",
+			pod: st.MakePod().Name("p").Label("foo", "").SpreadConstraint(
+				1, "zone", hardSpread, st.MakeLabelSelector().Exists("foo").Obj(),
+			).Obj(),
+			nodes: []*v1.Node{
+				st.MakeNode().Name("node-a").Label("zone", "zone1").Label("node", "node-a").Obj(),
+				st.MakeNode().Name("node-b").Label("zone", "zone1").Label("node", "node-b").Obj(),
+				st.MakeNode().Name("node-x").Label("zone", "zone2").Label("node", "node-x").Obj(),
+				st.MakeNode().Name("node-y").Label("zone", "zone2").Label("node", "node-y").Obj(),
+			},
+			injectPodPointers: map[topologyPair][]int{
+				// denotes no existing pod is matched on this zone pair, but still needed to be
+				// calculated if incoming pod matches its own spread constraints
+				{key: "zone", value: "zone1"}: {},
+				{key: "zone", value: "zone2"}: {},
+			},
+			want: &topologyPairsPodSpreadMap{
+				topologyKeyToMinPodsMap: map[string]int32{"zone": 0},
+				topologyPairsMaps: &topologyPairsMaps{
+					podToTopologyPairs: make(map[string]topologyPairSet),
+				},
+			},
+		},
+		{
+			name: "normal case with one spreadConstraint",
+			pod: st.MakePod().Name("p").Label("foo", "").SpreadConstraint(
+				1, "zone", hardSpread, st.MakeLabelSelector().Exists("foo").Obj(),
+			).Obj(),
+			nodes: []*v1.Node{
+				st.MakeNode().Name("node-a").Label("zone", "zone1").Label("node", "node-a").Obj(),
+				st.MakeNode().Name("node-b").Label("zone", "zone1").Label("node", "node-b").Obj(),
+				st.MakeNode().Name("node-x").Label("zone", "zone2").Label("node", "node-x").Obj(),
+				st.MakeNode().Name("node-y").Label("zone", "zone2").Label("node", "node-y").Obj(),
+			},
+			existingPods: []*v1.Pod{
+				st.MakePod().Name("p-a1").Node("node-a").Label("foo", "").Obj(),
+				st.MakePod().Name("p-a2").Node("node-a").Label("foo", "").Obj(),
+				st.MakePod().Name("p-b1").Node("node-b").Label("foo", "").Obj(),
+				st.MakePod().Name("p-y1").Node("node-y").Label("foo", "").Obj(),
+				st.MakePod().Name("p-y2").Node("node-y").Label("foo", "").Obj(),
+			},
+			injectPodPointers: map[topologyPair][]int{
+				// denotes existingPods[0,1,2]
+				{key: "zone", value: "zone1"}: {0, 1, 2},
+				// denotes existingPods[3,4]
+				{key: "zone", value: "zone2"}: {3, 4},
+			},
+			want: &topologyPairsPodSpreadMap{
+				topologyKeyToMinPodsMap: map[string]int32{"zone": 2},
+				topologyPairsMaps: &topologyPairsMaps{
+					podToTopologyPairs: map[string]topologyPairSet{
+						"p-a1_": newPairSet("zone", "zone1"),
+						"p-a2_": newPairSet("zone", "zone1"),
+						"p-b1_": newPairSet("zone", "zone1"),
+						"p-y1_": newPairSet("zone", "zone2"),
+						"p-y2_": newPairSet("zone", "zone2"),
+					},
+				},
+			},
+		},
+		{
+			name: "namespace mismatch doesn't count",
+			pod: st.MakePod().Name("p").Label("foo", "").SpreadConstraint(
+				1, "zone", hardSpread, st.MakeLabelSelector().Exists("foo").Obj(),
+			).Obj(),
+			nodes: []*v1.Node{
+				st.MakeNode().Name("node-a").Label("zone", "zone1").Label("node", "node-a").Obj(),
+				st.MakeNode().Name("node-b").Label("zone", "zone1").Label("node", "node-b").Obj(),
+				st.MakeNode().Name("node-x").Label("zone", "zone2").Label("node", "node-x").Obj(),
+				st.MakeNode().Name("node-y").Label("zone", "zone2").Label("node", "node-y").Obj(),
+			},
+			existingPods: []*v1.Pod{
+				st.MakePod().Name("p-a1").Node("node-a").Label("foo", "").Obj(),
+				st.MakePod().Name("p-a2").Namespace("ns1").Node("node-a").Label("foo", "").Obj(),
+				st.MakePod().Name("p-b1").Node("node-b").Label("foo", "").Obj(),
+				st.MakePod().Name("p-y1").Namespace("ns2").Node("node-y").Label("foo", "").Obj(),
+				st.MakePod().Name("p-y2").Node("node-y").Label("foo", "").Obj(),
+			},
+			injectPodPointers: map[topologyPair][]int{
+				{key: "zone", value: "zone1"}: {0, 2},
+				{key: "zone", value: "zone2"}: {4},
+			},
+			want: &topologyPairsPodSpreadMap{
+				topologyKeyToMinPodsMap: map[string]int32{"zone": 1},
+				topologyPairsMaps: &topologyPairsMaps{
+					podToTopologyPairs: map[string]topologyPairSet{
+						"p-a1_": newPairSet("zone", "zone1"),
+						"p-b1_": newPairSet("zone", "zone1"),
+						"p-y2_": newPairSet("zone", "zone2"),
+					},
+				},
+			},
+		},
+		{
+			name: "normal case with two spreadConstraints",
+			pod: st.MakePod().Name("p").Label("foo", "").
+				SpreadConstraint(1, "zone", hardSpread, st.MakeLabelSelector().Exists("foo").Obj()).
+				SpreadConstraint(1, "node", hardSpread, st.MakeLabelSelector().Exists("foo").Obj()).
+				Obj(),
+			nodes: []*v1.Node{
+				st.MakeNode().Name("node-a").Label("zone", "zone1").Label("node", "node-a").Obj(),
+				st.MakeNode().Name("node-b").Label("zone", "zone1").Label("node", "node-b").Obj(),
+				st.MakeNode().Name("node-x").Label("zone", "zone2").Label("node", "node-x").Obj(),
+				st.MakeNode().Name("node-y").Label("zone", "zone2").Label("node", "node-y").Obj(),
+			},
+			existingPods: []*v1.Pod{
+				st.MakePod().Name("p-a1").Node("node-a").Label("foo", "").Obj(),
+				st.MakePod().Name("p-a2").Node("node-a").Label("foo", "").Obj(),
+				st.MakePod().Name("p-b1").Node("node-b").Label("foo", "").Obj(),
+				st.MakePod().Name("p-y1").Node("node-y").Label("foo", "").Obj(),
+				st.MakePod().Name("p-y2").Node("node-y").Label("foo", "").Obj(),
+				st.MakePod().Name("p-y3").Node("node-y").Label("foo", "").Obj(),
+				st.MakePod().Name("p-y4").Node("node-y").Label("foo", "").Obj(),
+			},
+			injectPodPointers: map[topologyPair][]int{
+				{key: "zone", value: "zone1"}:  {0, 1, 2},
+				{key: "zone", value: "zone2"}:  {3, 4, 5, 6},
+				{key: "node", value: "node-a"}: {0, 1},
+				{key: "node", value: "node-b"}: {2},
+				{key: "node", value: "node-x"}: {},
+				{key: "node", value: "node-y"}: {3, 4, 5, 6},
+			},
+			want: &topologyPairsPodSpreadMap{
+				topologyKeyToMinPodsMap: map[string]int32{"zone": 3, "node": 0},
+				topologyPairsMaps: &topologyPairsMaps{
+					podToTopologyPairs: map[string]topologyPairSet{
+						"p-a1_": newPairSet("zone", "zone1", "node", "node-a"),
+						"p-a2_": newPairSet("zone", "zone1", "node", "node-a"),
+						"p-b1_": newPairSet("zone", "zone1", "node", "node-b"),
+						"p-y1_": newPairSet("zone", "zone2", "node", "node-y"),
+						"p-y2_": newPairSet("zone", "zone2", "node", "node-y"),
+						"p-y3_": newPairSet("zone", "zone2", "node", "node-y"),
+						"p-y4_": newPairSet("zone", "zone2", "node", "node-y"),
+					},
+				},
+			},
+		},
+		{
+			name: "soft spreadConstraints should be bypassed",
+			pod: st.MakePod().Name("p").Label("foo", "").
+				SpreadConstraint(1, "zone", softSpread, st.MakeLabelSelector().Exists("foo").Obj()).
+				SpreadConstraint(1, "zone", hardSpread, st.MakeLabelSelector().Exists("foo").Obj()).
+				SpreadConstraint(1, "zone", softSpread, st.MakeLabelSelector().Exists("foo").Obj()).
+				SpreadConstraint(1, "node", hardSpread, st.MakeLabelSelector().Exists("foo").Obj()).
+				Obj(),
+			nodes: []*v1.Node{
+				st.MakeNode().Name("node-a").Label("zone", "zone1").Label("node", "node-a").Obj(),
+				st.MakeNode().Name("node-b").Label("zone", "zone1").Label("node", "node-b").Obj(),
+				st.MakeNode().Name("node-y").Label("zone", "zone2").Label("node", "node-y").Obj(),
+			},
+			existingPods: []*v1.Pod{
+				st.MakePod().Name("p-a1").Node("node-a").Label("foo", "").Obj(),
+				st.MakePod().Name("p-a2").Node("node-a").Label("foo", "").Obj(),
+				st.MakePod().Name("p-b1").Node("node-b").Label("foo", "").Obj(),
+				st.MakePod().Name("p-y1").Node("node-y").Label("foo", "").Obj(),
+				st.MakePod().Name("p-y2").Node("node-y").Label("foo", "").Obj(),
+				st.MakePod().Name("p-y3").Node("node-y").Label("foo", "").Obj(),
+				st.MakePod().Name("p-y4").Node("node-y").Label("foo", "").Obj(),
+			},
+			injectPodPointers: map[topologyPair][]int{
+				{key: "zone", value: "zone1"}:  {0, 1, 2},
+				{key: "zone", value: "zone2"}:  {3, 4, 5, 6},
+				{key: "node", value: "node-a"}: {0, 1},
+				{key: "node", value: "node-b"}: {2},
+				{key: "node", value: "node-y"}: {3, 4, 5, 6},
+			},
+			want: &topologyPairsPodSpreadMap{
+				topologyKeyToMinPodsMap: map[string]int32{"zone": 3, "node": 1},
+				topologyPairsMaps: &topologyPairsMaps{
+					podToTopologyPairs: map[string]topologyPairSet{
+						"p-a1_": newPairSet("zone", "zone1", "node", "node-a"),
+						"p-a2_": newPairSet("zone", "zone1", "node", "node-a"),
+						"p-b1_": newPairSet("zone", "zone1", "node", "node-b"),
+						"p-y1_": newPairSet("zone", "zone2", "node", "node-y"),
+						"p-y2_": newPairSet("zone", "zone2", "node", "node-y"),
+						"p-y3_": newPairSet("zone", "zone2", "node", "node-y"),
+						"p-y4_": newPairSet("zone", "zone2", "node", "node-y"),
+					},
+				},
+			},
+		},
+		{
+			name: "different labelSelectors - simple version",
+			pod: st.MakePod().Name("p").Label("foo", "").Label("bar", "").
+				SpreadConstraint(1, "zone", hardSpread, st.MakeLabelSelector().Exists("foo").Obj()).
+				SpreadConstraint(1, "node", hardSpread, st.MakeLabelSelector().Exists("bar").Obj()).
+				Obj(),
+			nodes: []*v1.Node{
+				st.MakeNode().Name("node-a").Label("zone", "zone1").Label("node", "node-a").Obj(),
+				st.MakeNode().Name("node-b").Label("zone", "zone1").Label("node", "node-b").Obj(),
+				st.MakeNode().Name("node-y").Label("zone", "zone2").Label("node", "node-y").Obj(),
+			},
+			existingPods: []*v1.Pod{
+				st.MakePod().Name("p-a").Node("node-a").Label("foo", "").Obj(),
+			},
+			injectPodPointers: map[topologyPair][]int{
+				{key: "zone", value: "zone1"}:  {0},
+				{key: "zone", value: "zone2"}:  {},
+				{key: "node", value: "node-a"}: {},
+				{key: "node", value: "node-b"}: {},
+				{key: "node", value: "node-y"}: {},
+			},
+			want: &topologyPairsPodSpreadMap{
+				topologyKeyToMinPodsMap: map[string]int32{"zone": 0, "node": 0},
+				topologyPairsMaps: &topologyPairsMaps{
+					podToTopologyPairs: map[string]topologyPairSet{
+						"p-a_": newPairSet("zone", "zone1"),
+					},
+				},
+			},
+		},
+		{
+			name: "different labelSelectors - complex version",
+			pod: st.MakePod().Name("p").Label("foo", "").Label("bar", "").
+				SpreadConstraint(1, "zone", hardSpread, st.MakeLabelSelector().Exists("foo").Obj()).
+				SpreadConstraint(1, "node", hardSpread, st.MakeLabelSelector().Exists("bar").Obj()).
+				Obj(),
+			nodes: []*v1.Node{
+				st.MakeNode().Name("node-a").Label("zone", "zone1").Label("node", "node-a").Obj(),
+				st.MakeNode().Name("node-b").Label("zone", "zone1").Label("node", "node-b").Obj(),
+				st.MakeNode().Name("node-y").Label("zone", "zone2").Label("node", "node-y").Obj(),
+			},
+			existingPods: []*v1.Pod{
+				st.MakePod().Name("p-a1").Node("node-a").Label("foo", "").Obj(),
+				st.MakePod().Name("p-a2").Node("node-a").Label("foo", "").Label("bar", "").Obj(),
+				st.MakePod().Name("p-b1").Node("node-b").Label("foo", "").Obj(),
+				st.MakePod().Name("p-y1").Node("node-y").Label("foo", "").Obj(),
+				st.MakePod().Name("p-y2").Node("node-y").Label("foo", "").Label("bar", "").Obj(),
+				st.MakePod().Name("p-y3").Node("node-y").Label("foo", "").Obj(),
+				st.MakePod().Name("p-y4").Node("node-y").Label("foo", "").Label("bar", "").Obj(),
+			},
+			injectPodPointers: map[topologyPair][]int{
+				{key: "zone", value: "zone1"}:  {0, 1, 2},
+				{key: "zone", value: "zone2"}:  {3, 4, 5, 6},
+				{key: "node", value: "node-a"}: {1},
+				{key: "node", value: "node-b"}: {},
+				{key: "node", value: "node-y"}: {4, 6},
+			},
+			want: &topologyPairsPodSpreadMap{
+				topologyKeyToMinPodsMap: map[string]int32{"zone": 3, "node": 0},
+				topologyPairsMaps: &topologyPairsMaps{
+					podToTopologyPairs: map[string]topologyPairSet{
+						"p-a1_": newPairSet("zone", "zone1"),
+						"p-a2_": newPairSet("zone", "zone1", "node", "node-a"),
+						"p-b1_": newPairSet("zone", "zone1"),
+						"p-y1_": newPairSet("zone", "zone2"),
+						"p-y2_": newPairSet("zone", "zone2", "node", "node-y"),
+						"p-y3_": newPairSet("zone", "zone2"),
+						"p-y4_": newPairSet("zone", "zone2", "node", "node-y"),
+					},
+				},
+			},
+		},
+		{
+			name: "two spreadConstraints, and with podAffinity",
+			pod: st.MakePod().Name("p").Label("foo", "").
+				NodeAffinityNotIn("node", []string{"node-x"}). // exclude node-x
+				SpreadConstraint(1, "zone", hardSpread, st.MakeLabelSelector().Exists("foo").Obj()).
+				SpreadConstraint(1, "node", hardSpread, st.MakeLabelSelector().Exists("foo").Obj()).
+				Obj(),
+			nodes: []*v1.Node{
+				st.MakeNode().Name("node-a").Label("zone", "zone1").Label("node", "node-a").Obj(),
+				st.MakeNode().Name("node-b").Label("zone", "zone1").Label("node", "node-b").Obj(),
+				st.MakeNode().Name("node-x").Label("zone", "zone2").Label("node", "node-x").Obj(),
+				st.MakeNode().Name("node-y").Label("zone", "zone2").Label("node", "node-y").Obj(),
+			},
+			existingPods: []*v1.Pod{
+				st.MakePod().Name("p-a1").Node("node-a").Label("foo", "").Obj(),
+				st.MakePod().Name("p-a2").Node("node-a").Label("foo", "").Obj(),
+				st.MakePod().Name("p-b1").Node("node-b").Label("foo", "").Obj(),
+				st.MakePod().Name("p-y1").Node("node-y").Label("foo", "").Obj(),
+				st.MakePod().Name("p-y2").Node("node-y").Label("foo", "").Obj(),
+				st.MakePod().Name("p-y3").Node("node-y").Label("foo", "").Obj(),
+				st.MakePod().Name("p-y4").Node("node-y").Label("foo", "").Obj(),
+			},
+			injectPodPointers: map[topologyPair][]int{
+				{key: "zone", value: "zone1"}:  {0, 1, 2},
+				{key: "zone", value: "zone2"}:  {3, 4, 5, 6},
+				{key: "node", value: "node-a"}: {0, 1},
+				{key: "node", value: "node-b"}: {2},
+				{key: "node", value: "node-y"}: {3, 4, 5, 6},
+			},
+			want: &topologyPairsPodSpreadMap{
+				topologyKeyToMinPodsMap: map[string]int32{"zone": 3, "node": 1},
+				topologyPairsMaps: &topologyPairsMaps{
+					podToTopologyPairs: map[string]topologyPairSet{
+						"p-a1_": newPairSet("zone", "zone1", "node", "node-a"),
+						"p-a2_": newPairSet("zone", "zone1", "node", "node-a"),
+						"p-b1_": newPairSet("zone", "zone1", "node", "node-b"),
+						"p-y1_": newPairSet("zone", "zone2", "node", "node-y"),
+						"p-y2_": newPairSet("zone", "zone2", "node", "node-y"),
+						"p-y3_": newPairSet("zone", "zone2", "node", "node-y"),
+						"p-y4_": newPairSet("zone", "zone2", "node", "node-y"),
+					},
+				},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tt.want.topologyPairToPods = make(map[topologyPair]podSet)
+			for pair, indexes := range tt.injectPodPointers {
+				pSet := make(podSet)
+				for _, i := range indexes {
+					pSet[tt.existingPods[i]] = struct{}{}
+				}
+				tt.want.topologyPairToPods[pair] = pSet
+			}
+			nodeInfoMap := schedulernodeinfo.CreateNodeNameToInfoMap(tt.existingPods, tt.nodes)
+			if got, _ := getTPMapMatchingSpreadConstraints(tt.pod, nodeInfoMap); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("getTPMapMatchingSpreadConstraints() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+var (
+	hardSpread = v1.DoNotSchedule
+	softSpread = v1.ScheduleAnyway
+)
+
+func newPairSet(kv ...string) topologyPairSet {
+	result := make(topologyPairSet)
+	for i := 0; i < len(kv); i += 2 {
+		pair := topologyPair{key: kv[i], value: kv[i+1]}
+		result[pair] = struct{}{}
+	}
+	return result
+}

pkg/scheduler/algorithm/predicates/utils.go

@@ -19,7 +19,7 @@ package predicates
 import (
 	"strings"
 
-	v1 "k8s.io/api/core/v1"
+	"k8s.io/api/core/v1"
 	storagev1beta1 "k8s.io/api/storage/v1beta1"
 	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/apimachinery/pkg/util/sets"

pkg/scheduler/testing/BUILD

@@ -4,7 +4,10 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library")
 
 go_library(
     name = "go_default_library",
-    srcs = ["fake_lister.go"],
+    srcs = [
+        "fake_lister.go",
+        "wrappers.go",
+    ],
     importpath = "k8s.io/kubernetes/pkg/scheduler/testing",
     deps = [
         "//pkg/scheduler/algorithm:go_default_library",

pkg/scheduler/testing/wrappers.go

@@ -0,0 +1,243 @@
+/*
+Copyright 2019 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package testing
+
+import (
+	"k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+// NodeSelectorWrapper wraps a NodeSelector inside.
+type NodeSelectorWrapper struct{ v1.NodeSelector }
+
+// MakeNodeSelector creates a NodeSelector wrapper.
+func MakeNodeSelector() *NodeSelectorWrapper {
+	return &NodeSelectorWrapper{v1.NodeSelector{}}
+}
+
+// In injects a matchExpression (with an operator IN) as a selectorTerm
+// to the inner nodeSelector.
+// NOTE: appended selecterTerms are ORed.
+func (s *NodeSelectorWrapper) In(key string, vals []string) *NodeSelectorWrapper {
+	expression := v1.NodeSelectorRequirement{
+		Key:      key,
+		Operator: v1.NodeSelectorOpIn,
+		Values:   vals,
+	}
+	selectorTerm := v1.NodeSelectorTerm{}
+	selectorTerm.MatchExpressions = append(selectorTerm.MatchExpressions, expression)
+	s.NodeSelectorTerms = append(s.NodeSelectorTerms, selectorTerm)
+	return s
+}
+
+// NotIn injects a matchExpression (with an operator NotIn) as a selectorTerm
+// to the inner nodeSelector.
+func (s *NodeSelectorWrapper) NotIn(key string, vals []string) *NodeSelectorWrapper {
+	expression := v1.NodeSelectorRequirement{
+		Key:      key,
+		Operator: v1.NodeSelectorOpNotIn,
+		Values:   vals,
+	}
+	selectorTerm := v1.NodeSelectorTerm{}
+	selectorTerm.MatchExpressions = append(selectorTerm.MatchExpressions, expression)
+	s.NodeSelectorTerms = append(s.NodeSelectorTerms, selectorTerm)
+	return s
+}
+
+// Obj returns the inner NodeSelector.
+func (s *NodeSelectorWrapper) Obj() *v1.NodeSelector {
+	return &s.NodeSelector
+}
+
+// LabelSelectorWrapper wraps a LabelSelector inside.
+type LabelSelectorWrapper struct{ metav1.LabelSelector }
+
+// MakeLabelSelector creates a LabelSelector wrapper.
+func MakeLabelSelector() *LabelSelectorWrapper {
+	return &LabelSelectorWrapper{metav1.LabelSelector{}}
+}
+
+// Label applies a {k,v} pair to the inner LabelSelector.
+func (s *LabelSelectorWrapper) Label(k, v string) *LabelSelectorWrapper {
+	if s.MatchLabels == nil {
+		s.MatchLabels = make(map[string]string)
+	}
+	s.MatchLabels[k] = v
+	return s
+}
+
+// In injects a matchExpression (with an operator In) to the inner labelSelector.
+func (s *LabelSelectorWrapper) In(key string, vals []string) *LabelSelectorWrapper {
+	expression := metav1.LabelSelectorRequirement{
+		Key:      key,
+		Operator: metav1.LabelSelectorOpIn,
+		Values:   vals,
+	}
+	s.MatchExpressions = append(s.MatchExpressions, expression)
+	return s
+}
+
+// NotIn injects a matchExpression (with an operator NotIn) to the inner labelSelector.
+func (s *LabelSelectorWrapper) NotIn(key string, vals []string) *LabelSelectorWrapper {
+	expression := metav1.LabelSelectorRequirement{
+		Key:      key,
+		Operator: metav1.LabelSelectorOpNotIn,
+		Values:   vals,
+	}
+	s.MatchExpressions = append(s.MatchExpressions, expression)
+	return s
+}
+
+// Exists injects a matchExpression (with an operator Exists) to the inner labelSelector.
+func (s *LabelSelectorWrapper) Exists(k string) *LabelSelectorWrapper {
+	expression := metav1.LabelSelectorRequirement{
+		Key:      k,
+		Operator: metav1.LabelSelectorOpExists,
+	}
+	s.MatchExpressions = append(s.MatchExpressions, expression)
+	return s
+}
+
+// NotExist injects a matchExpression (with an operator NotExist) to the inner labelSelector.
+func (s *LabelSelectorWrapper) NotExist(k string) *LabelSelectorWrapper {
+	expression := metav1.LabelSelectorRequirement{
+		Key:      k,
+		Operator: metav1.LabelSelectorOpDoesNotExist,
+	}
+	s.MatchExpressions = append(s.MatchExpressions, expression)
+	return s
+}
+
+// Obj returns the inner LabelSelector.
+func (s *LabelSelectorWrapper) Obj() *metav1.LabelSelector {
+	return &s.LabelSelector
+}
+
+// PodWrapper wraps a Pod inside.
+type PodWrapper struct{ v1.Pod }
+
+// MakePod creates a Pod wrapper.
+func MakePod() *PodWrapper {
+	return &PodWrapper{v1.Pod{}}
+}
+
+// Obj returns the inner Pod.
+func (p *PodWrapper) Obj() *v1.Pod {
+	return &p.Pod
+}
+
+// Name sets `s` as the name of the inner pod.
+func (p *PodWrapper) Name(s string) *PodWrapper {
+	p.SetName(s)
+	return p
+}
+
+// Namespace sets `s` as the namespace of the inner pod.
+func (p *PodWrapper) Namespace(s string) *PodWrapper {
+	p.SetNamespace(s)
+	return p
+}
+
+// Node sets `s` as the nodeName of the inner pod.
+func (p *PodWrapper) Node(s string) *PodWrapper {
+	p.Spec.NodeName = s
+	return p
+}
+
+// NodeSelector sets `m` as the nodeSelector of the inner pod.
+func (p *PodWrapper) NodeSelector(m map[string]string) *PodWrapper {
+	p.Spec.NodeSelector = m
+	return p
+}
+
+// NodeAffinityIn creates a HARD node affinity (with the operator In)
+// and injects into the innner pod.
+func (p *PodWrapper) NodeAffinityIn(key string, vals []string) *PodWrapper {
+	if p.Spec.Affinity == nil {
+		p.Spec.Affinity = &v1.Affinity{}
+	}
+	if p.Spec.Affinity.NodeAffinity == nil {
+		p.Spec.Affinity.NodeAffinity = &v1.NodeAffinity{}
+	}
+	nodeSelector := MakeNodeSelector().In(key, vals).Obj()
+	p.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = nodeSelector
+	return p
+}
+
+// NodeAffinityNotIn creates a HARD node affinity (with the operator NotIn)
+// and injects into the innner pod.
+func (p *PodWrapper) NodeAffinityNotIn(key string, vals []string) *PodWrapper {
+	if p.Spec.Affinity == nil {
+		p.Spec.Affinity = &v1.Affinity{}
+	}
+	if p.Spec.Affinity.NodeAffinity == nil {
+		p.Spec.Affinity.NodeAffinity = &v1.NodeAffinity{}
+	}
+	nodeSelector := MakeNodeSelector().NotIn(key, vals).Obj()
+	p.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = nodeSelector
+	return p
+}
+
+// SpreadConstraint constructs a TopologySpreadConstraint object and injects
+// into the inner pod.
+func (p *PodWrapper) SpreadConstraint(maxSkew int, tpKey string, mode v1.UnsatisfiableConstraintAction, selector *metav1.LabelSelector) *PodWrapper {
+	c := v1.TopologySpreadConstraint{
+		MaxSkew:           int32(maxSkew),
+		TopologyKey:       tpKey,
+		WhenUnsatisfiable: mode,
+		LabelSelector:     selector,
+	}
+	p.Spec.TopologySpreadConstraints = append(p.Spec.TopologySpreadConstraints, c)
+	return p
+}
+
+// Label sets a {k,v} pair to the inner pod.
+func (p *PodWrapper) Label(k, v string) *PodWrapper {
+	if p.Labels == nil {
+		p.Labels = make(map[string]string)
+	}
+	p.Labels[k] = v
+	return p
+}
+
+// NodeWrapper wraps a Node inside.
+type NodeWrapper struct{ v1.Node }
+
+// MakeNode creates a Node wrapper.
+func MakeNode() *NodeWrapper {
+	return &NodeWrapper{v1.Node{}}
+}
+
+// Obj returns the inner Node.
+func (n *NodeWrapper) Obj() *v1.Node {
+	return &n.Node
+}
+
+// Name sets `s` as the name of the inner pod.
+func (n *NodeWrapper) Name(s string) *NodeWrapper {
+	n.SetName(s)
+	return n
+}
+
+// Label applies a {k,v} label pair to the inner node.
+func (n *NodeWrapper) Label(k, v string) *NodeWrapper {
+	if n.Labels == nil {
+		n.Labels = make(map[string]string)
+	}
+	n.Labels[k] = v
+	return n
+}

pkg/scheduler/util/utils.go

@@ -21,7 +21,7 @@ import (
 
 	"time"
 
-	v1 "k8s.io/api/core/v1"
+	"k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/klog"
 	"k8s.io/kubernetes/pkg/apis/scheduling"