...
1
16
17 package pod
18
19 import (
20 v1 "k8s.io/api/core/v1"
21 )
22
23
24
25 type NodeSelection struct {
26 Name string
27 Selector map[string]string
28 Affinity *v1.Affinity
29 }
30
31
32 func setNodeAffinityRequirement(nodeSelection *NodeSelection, operator v1.NodeSelectorOperator, nodeName string) {
33
34 if nodeSelection.Affinity == nil {
35 nodeSelection.Affinity = &v1.Affinity{}
36 }
37 if nodeSelection.Affinity.NodeAffinity == nil {
38 nodeSelection.Affinity.NodeAffinity = &v1.NodeAffinity{}
39 }
40 if nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution == nil {
41 nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = &v1.NodeSelector{}
42 }
43 nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms = append(nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms,
44 v1.NodeSelectorTerm{
45 MatchFields: []v1.NodeSelectorRequirement{
46 {Key: "metadata.name", Operator: operator, Values: []string{nodeName}},
47 },
48 })
49 }
50
51
52 func SetNodeAffinityTopologyRequirement(nodeSelection *NodeSelection, topology map[string]string) {
53 if nodeSelection.Affinity == nil {
54 nodeSelection.Affinity = &v1.Affinity{}
55 }
56 if nodeSelection.Affinity.NodeAffinity == nil {
57 nodeSelection.Affinity.NodeAffinity = &v1.NodeAffinity{}
58 }
59 if nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution == nil {
60 nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = &v1.NodeSelector{}
61 }
62 for k, v := range topology {
63 nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms = append(nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms,
64 v1.NodeSelectorTerm{
65 MatchExpressions: []v1.NodeSelectorRequirement{
66 {Key: k, Operator: v1.NodeSelectorOpIn, Values: []string{v}},
67 },
68 })
69
70 }
71 }
72
73
74 func SetAffinity(nodeSelection *NodeSelection, nodeName string) {
75 setNodeAffinityRequirement(nodeSelection, v1.NodeSelectorOpIn, nodeName)
76 }
77
78
79 func SetAntiAffinity(nodeSelection *NodeSelection, nodeName string) {
80 setNodeAffinityRequirement(nodeSelection, v1.NodeSelectorOpNotIn, nodeName)
81 }
82
83
84
85 func SetNodeAffinity(podSpec *v1.PodSpec, nodeName string) {
86 nodeSelection := &NodeSelection{}
87 SetAffinity(nodeSelection, nodeName)
88 podSpec.Affinity = nodeSelection.Affinity
89 }
90
91
92
93 func SetNodeSelection(podSpec *v1.PodSpec, nodeSelection NodeSelection) {
94 podSpec.NodeSelector = nodeSelection.Selector
95 podSpec.Affinity = nodeSelection.Affinity
96
97
98
99
100
101
102 if nodeSelection.Name != "" {
103 SetNodeAffinity(podSpec, nodeSelection.Name)
104 }
105 }
106
View as plain text