kubernetes/test/e2e/framework/pod/node_selection.go
Michelle Au fb9f02b5e1 Don't set NodeName directly in Pods so that it still goes through the scheduler
Change-Id: I244b6aac0289a13339f3ac228c4ad9ecf8c07b42
2020-02-11 19:17:41 -08:00

106 lines
4.0 KiB
Go

/*
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 pod
import (
v1 "k8s.io/api/core/v1"
)
// NodeSelection specifies where to run a pod, using a combination of fixed node name,
// node selector and/or affinity.
type NodeSelection struct {
Name string
Selector map[string]string
Affinity *v1.Affinity
}
// setNodeAffinityRequirement sets affinity with specified operator to nodeName to nodeSelection
func setNodeAffinityRequirement(nodeSelection *NodeSelection, operator v1.NodeSelectorOperator, nodeName string) {
// Add node-anti-affinity.
if nodeSelection.Affinity == nil {
nodeSelection.Affinity = &v1.Affinity{}
}
if nodeSelection.Affinity.NodeAffinity == nil {
nodeSelection.Affinity.NodeAffinity = &v1.NodeAffinity{}
}
if nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution == nil {
nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = &v1.NodeSelector{}
}
nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms = append(nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms,
v1.NodeSelectorTerm{
MatchFields: []v1.NodeSelectorRequirement{
{Key: "metadata.name", Operator: operator, Values: []string{nodeName}},
},
})
}
// SetNodeAffinityTopologyRequirement sets node affinity to a specified topology
func SetNodeAffinityTopologyRequirement(nodeSelection *NodeSelection, topology map[string]string) {
if nodeSelection.Affinity == nil {
nodeSelection.Affinity = &v1.Affinity{}
}
if nodeSelection.Affinity.NodeAffinity == nil {
nodeSelection.Affinity.NodeAffinity = &v1.NodeAffinity{}
}
if nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution == nil {
nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = &v1.NodeSelector{}
}
for k, v := range topology {
nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms = append(nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms,
v1.NodeSelectorTerm{
MatchExpressions: []v1.NodeSelectorRequirement{
{Key: k, Operator: v1.NodeSelectorOpIn, Values: []string{v}},
},
})
}
}
// SetAffinity sets affinity to nodeName to nodeSelection
func SetAffinity(nodeSelection *NodeSelection, nodeName string) {
setNodeAffinityRequirement(nodeSelection, v1.NodeSelectorOpIn, nodeName)
}
// SetAntiAffinity sets anti-affinity to nodeName to nodeSelection
func SetAntiAffinity(nodeSelection *NodeSelection, nodeName string) {
setNodeAffinityRequirement(nodeSelection, v1.NodeSelectorOpNotIn, nodeName)
}
// SetNodeAffinity modifies the given pod object with
// NodeAffinity to the given node name.
func SetNodeAffinity(pod *v1.Pod, nodeName string) {
nodeSelection := &NodeSelection{}
SetAffinity(nodeSelection, nodeName)
pod.Spec.Affinity = nodeSelection.Affinity
}
// SetNodeSelection modifies the given pod object with
// the specified NodeSelection
func SetNodeSelection(pod *v1.Pod, nodeSelection NodeSelection) {
pod.Spec.NodeSelector = nodeSelection.Selector
pod.Spec.Affinity = nodeSelection.Affinity
// pod.Spec.NodeName should not be set directly because
// it will bypass the scheduler, potentially causing
// kubelet to Fail the pod immediately if it's out of
// resources. Instead, we want the pod to remain
// pending in the scheduler until the node has resources
// freed up.
if nodeSelection.Name != "" {
SetNodeAffinity(pod, nodeSelection.Name)
}
}