-
Notifications
You must be signed in to change notification settings - Fork 127
/
affinity.go
65 lines (56 loc) · 2.19 KB
/
affinity.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
package affinity
import (
"github.com/pkg/errors"
"golang.org/x/exp/slices"
v1 "k8s.io/api/core/v1"
)
func AddNodeAntiAffinity(affinity *v1.Affinity, labelName string, labelValue string) error {
if affinity == nil {
return errors.Errorf("failed to add not anti affinity, as provided affinity is nil")
}
ensureAffinityHasNodeSelectorTerms(affinity)
addAvoidNodeAffinityToNodeSelectorTerms(affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms, labelName, labelValue)
return nil
}
func ensureAffinityHasNodeSelectorTerms(affinity *v1.Affinity) {
if affinity.NodeAffinity == nil {
affinity.NodeAffinity = &v1.NodeAffinity{}
}
nodeAffinity := affinity.NodeAffinity
if nodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution == nil {
nodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = &v1.NodeSelector{}
}
ns := nodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution
if len(ns.NodeSelectorTerms) == 0 {
ns.NodeSelectorTerms = []v1.NodeSelectorTerm{{}}
}
}
func addAvoidNodeAffinityToNodeSelectorTerms(terms []v1.NodeSelectorTerm, labelName string, labelValue string) {
// The node NotIn needs to be added to all NodeSelectorTerms, because if any single NodeSelectorTerm matches
// the pod is considered to match.
for i := range terms {
addAvoidNodeAffinityToNodeSelectorTerm(&terms[i], labelName, labelValue)
}
}
func addAvoidNodeAffinityToNodeSelectorTerm(term *v1.NodeSelectorTerm, labelName string, labelValue string) {
mexp := findMatchExpression(term.MatchExpressions, labelName, v1.NodeSelectorOpNotIn)
if mexp == nil {
term.MatchExpressions = append(term.MatchExpressions, v1.NodeSelectorRequirement{
Key: labelName,
Operator: v1.NodeSelectorOpNotIn,
Values: []string{},
})
mexp = &term.MatchExpressions[len(term.MatchExpressions)-1]
}
if !slices.Contains(mexp.Values, labelValue) {
mexp.Values = append(mexp.Values, labelValue)
}
}
func findMatchExpression(matchExpressions []v1.NodeSelectorRequirement, key string, operator v1.NodeSelectorOperator) *v1.NodeSelectorRequirement {
for i, me := range matchExpressions {
if me.Key == key && me.Operator == operator {
return &matchExpressions[i]
}
}
return nil
}