Merge pull request #30 from cwdsuzhou/tensile-kube

Only recover NodeAffinity required for Affinity
virtual-kubelet · Aug 6, 2020 · 6da539d · 6da539d
2 parents 4f0f69b + 538a909
commit 6da539d
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 11 deletions.
diff --git a/pkg/util/conversions.go b/pkg/util/conversions.go
@@ -46,17 +46,7 @@ func TrimPod(pod *corev1.Pod, ignoreLabels []string) *corev1.Pod {
 	}
 	podCopy.Labels[VirtualPodLabel] = "true"
 	cns := ConvertAnnotations(pod.Annotations)
-	// remove selector
-	if cns != nil {
-		podCopy.Spec.NodeSelector = cns.NodeSelector
-		podCopy.Spec.Affinity = cns.Affinity
-		podCopy.Spec.Tolerations = cns.Tolerations
-	} else {
-		podCopy.Spec.NodeSelector = nil
-		podCopy.Spec.Affinity = nil
-		podCopy.Spec.Tolerations = nil
-	}
-
+	recoverSelectors(podCopy, cns)
 	podCopy.Spec.Containers = trimContainers(pod.Spec.Containers)
 	podCopy.Spec.InitContainers = trimContainers(pod.Spec.InitContainers)
 	podCopy.Spec.Volumes = vols
@@ -145,6 +135,46 @@ func RecoverLabels(labels map[string]string, annotations map[string]string) {
 	}
 }
 
+// recoverSelectors recover some affinity, tolerations and nodeSelector from
+// ClusterSelector
+func recoverSelectors(pod *corev1.Pod, cns *ClustersNodeSelection) {
+	if cns != nil {
+		pod.Spec.NodeSelector = cns.NodeSelector
+		pod.Spec.Tolerations = cns.Tolerations
+		if pod.Spec.Affinity == nil {
+			pod.Spec.Affinity = cns.Affinity
+		} else {
+			if cns.Affinity != nil && cns.Affinity.NodeAffinity != nil {
+				if pod.Spec.Affinity.NodeAffinity != nil {
+					pod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = cns.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution
+				} else {
+					pod.Spec.Affinity.NodeAffinity = cns.Affinity.NodeAffinity
+				}
+			} else {
+				pod.Spec.Affinity.NodeAffinity = nil
+			}
+		}
+	} else {
+		pod.Spec.NodeSelector = nil
+		pod.Spec.Tolerations = nil
+		if pod.Spec.Affinity != nil && pod.Spec.Affinity.NodeAffinity != nil {
+			pod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = nil
+		}
+	}
+	if pod.Spec.Affinity != nil {
+		if pod.Spec.Affinity.NodeAffinity != nil {
+			if pod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution == nil &&
+				pod.Spec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution == nil {
+				pod.Spec.Affinity.NodeAffinity = nil
+			}
+		}
+		if pod.Spec.Affinity.NodeAffinity == nil && pod.Spec.Affinity.PodAffinity == nil &&
+			pod.Spec.Affinity.PodAntiAffinity == nil {
+			pod.Spec.Affinity = nil
+		}
+	}
+}
+
 // trimLabels removes label from labels according to ignoreLabels
 func trimLabels(labels map[string]string, ignoreLabels []string) map[string]string {
 	if ignoreLabels == nil {

diff --git a/pkg/util/conversions_test.go b/pkg/util/conversions_test.go
@@ -89,42 +89,50 @@ func TestTrimPod(t *testing.T) {
 	basePod2.Labels = map[string]string{"test": "test"}
 
 	cases := []struct {
+		name      string
 		pod       *corev1.Pod
 		desire    *corev1.Pod
 		trimLabel []string
 	}{
 		{
+			name:      "base test",
 			pod:       basePod,
 			desire:    desired,
 			trimLabel: nil,
 		},
 		{
+			name:      "base test with annotation clusterSelector",
 			pod:       basePod1,
 			desire:    desired1,
 			trimLabel: nil,
 		},
 		{
+			name:      "base test with label",
 			pod:       basePod2,
 			desire:    desired2,
 			trimLabel: []string{"test"},
 		},
 		{
+			name:      "base test toleration",
 			pod:       testbase.PodForTestWithOtherTolerations(),
 			desire:    desired,
 			trimLabel: nil,
 		},
 		{
+			name:      "base test node selector",
 			pod:       testbase.PodForTestWithNodeSelector(),
 			desire:    desired,
 			trimLabel: nil,
 		},
 		{
+			name:      "base affinity",
 			pod:       testbase.PodForTestWithAffinity(),
 			desire:    desired,
 			trimLabel: nil,
 		},
 	}
 	for _, d := range cases {
+		t.Log(d.name)
 		new := TrimPod(d.pod, d.trimLabel)
 		if new.String() != d.desire.String() {
 			t.Fatalf("Desired:\n %v\n, get:\n %v", d.desire, new)