Merge affinity from podtempalte and affinity-assistant

No matter user provide affinity in podtemplate or not the affinity-assistant will overwrite the affinity in pod. However if user set affinity in podtemplate merge podtemplate's affinity with affinity-assistant's affinity is better way. So we inject the affinity of the affinity-assitant into the affinity of the podtempalte as the affinity of the pod to make sure that all the affinities work properly. Related Issue: #5241 Signed-off-by: yuzhipeng <yuzp1996@qq.com>
tektoncd · Aug 11, 2022 · e1bcf4b · e1bcf4b
1 parent 55f0441
commit e1bcf4b
Show file tree

Hide file tree

Showing 2 changed files with 170 additions and 20 deletions.
diff --git a/pkg/internal/affinityassistant/transformer.go b/pkg/internal/affinityassistant/transformer.go
@@ -28,32 +28,42 @@ import (
 // NewTransformer returns a pod.Transformer that will pod affinity if needed
 func NewTransformer(_ context.Context, annotations map[string]string) pod.Transformer {
 	return func(p *corev1.Pod) (*corev1.Pod, error) {
-		// Using node affinity on taskRuns sharing PVC workspace, with an Affinity Assistant
-		// is mutually exclusive with other affinity on taskRun pods. If other
-		// affinity is wanted, that should be added on the Affinity Assistant pod unless
-		// assistant is disabled. When Affinity Assistant is disabled, an affinityAssistantName is not set.
+		// Using node affinity on taskRuns sharing PVC workspace.  When Affinity Assistant
+		// is disabled, an affinityAssistantName is not set.
 		if affinityAssistantName := annotations[workspace.AnnotationAffinityAssistantName]; affinityAssistantName != "" {
-			p.Spec.Affinity = nodeAffinityUsingAffinityAssistant(affinityAssistantName)
+			if p.Spec.Affinity == nil {
+				p.Spec.Affinity = &corev1.Affinity{}
+			}
+			mergeAffinityWithAffinityAssistant(p.Spec.Affinity, affinityAssistantName)
 		}
 		return p, nil
 	}
 }
 
-// nodeAffinityUsingAffinityAssistant achieves Node Affinity for taskRun pods
-// sharing PVC workspace by setting PodAffinity so that taskRuns is
-// scheduled to the Node were the Affinity Assistant pod is scheduled.
-func nodeAffinityUsingAffinityAssistant(affinityAssistantName string) *corev1.Affinity {
-	return &corev1.Affinity{
-		PodAffinity: &corev1.PodAffinity{
-			RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{{
-				LabelSelector: &metav1.LabelSelector{
-					MatchLabels: map[string]string{
-						workspace.LabelInstance:  affinityAssistantName,
-						workspace.LabelComponent: workspace.ComponentNameAffinityAssistant,
-					},
-				},
-				TopologyKey: "kubernetes.io/hostname",
-			}},
+func mergeAffinityWithAffinityAssistant(affinity *corev1.Affinity, affinityAssistantName string) {
+	podAffinityTerm := podAffinityTermUsingAffinityAssistant(affinityAssistantName)
+
+	if affinity == nil {
+		affinity = &corev1.Affinity{}
+	}
+	if affinity.PodAffinity == nil {
+		affinity.PodAffinity = &corev1.PodAffinity{}
+	}
+
+	affinity.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution =
+		append(affinity.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution, *podAffinityTerm)
+}
+
+// podAffinityTermUsingAffinityAssistant achieves pod Affinity term for taskRun
+// pods so that taskRuns is scheduled to the Node were the Affinity Assistant pod
+// is scheduled.
+func podAffinityTermUsingAffinityAssistant(affinityAssistantName string) *corev1.PodAffinityTerm {
+	return &corev1.PodAffinityTerm{LabelSelector: &metav1.LabelSelector{
+		MatchLabels: map[string]string{
+			workspace.LabelInstance:  affinityAssistantName,
+			workspace.LabelComponent: workspace.ComponentNameAffinityAssistant,
 		},
+	},
+		TopologyKey: "kubernetes.io/hostname",
 	}
 }
diff --git a/pkg/internal/affinityassistant/transformer_test.go b/pkg/internal/affinityassistant/transformer_test.go
@@ -88,3 +88,143 @@ func TestNewTransformer(t *testing.T) {
 		})
 	}
 }
+
+func TestNewTransformerWithNodeAffinity(t *testing.T) {
+
+	nodeAffinity := &corev1.NodeAffinity{
+		RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
+			NodeSelectorTerms: []corev1.NodeSelectorTerm{{
+				MatchFields: []corev1.NodeSelectorRequirement{{
+					Key:      "kubernetes.io/hostname",
+					Operator: corev1.NodeSelectorOpNotIn,
+					Values:   []string{"192.0.0.1"},
+				}},
+			}},
+		},
+	}
+
+	podAffinityTerm := &corev1.PodAffinityTerm{
+		LabelSelector: &metav1.LabelSelector{
+			MatchLabels: map[string]string{
+				"from": "podtemplate",
+			},
+		},
+		TopologyKey: "kubernetes.io/hostname",
+	}
+
+	for _, tc := range []struct {
+		description string
+		annotations map[string]string
+		pod         *corev1.Pod
+		expected    *corev1.Affinity
+	}{{
+		description: "no affinity annotation and pod contains nodeAffinity",
+		annotations: map[string]string{
+			"foo": "bar",
+		},
+		pod: &corev1.Pod{
+			Spec: corev1.PodSpec{
+				Affinity: &corev1.Affinity{
+					NodeAffinity: nodeAffinity,
+				},
+			},
+		},
+		expected: &corev1.Affinity{
+			NodeAffinity: nodeAffinity,
+		},
+	}, {
+		description: "affinity annotation and pod contains nodeAffinity",
+		annotations: map[string]string{
+			"foo": "bar",
+			workspace.AnnotationAffinityAssistantName: "baz",
+		},
+		pod: &corev1.Pod{
+			Spec: corev1.PodSpec{
+				Affinity: &corev1.Affinity{
+					NodeAffinity: nodeAffinity,
+				},
+			},
+		},
+
+		expected: &corev1.Affinity{
+			NodeAffinity: nodeAffinity,
+			PodAffinity: &corev1.PodAffinity{
+				RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{{
+					LabelSelector: &metav1.LabelSelector{
+						MatchLabels: map[string]string{
+							workspace.LabelInstance:  "baz",
+							workspace.LabelComponent: workspace.ComponentNameAffinityAssistant,
+						},
+					},
+					TopologyKey: "kubernetes.io/hostname",
+				}},
+			},
+		},
+	}, {
+		description: "affinity annotation with a different name and pod contains podAffinity",
+		annotations: map[string]string{
+			workspace.AnnotationAffinityAssistantName: "helloworld",
+		},
+		pod: &corev1.Pod{
+			Spec: corev1.PodSpec{
+				Affinity: &corev1.Affinity{PodAffinity: &corev1.PodAffinity{
+					RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{*podAffinityTerm},
+				}},
+			},
+		},
+
+		expected: &corev1.Affinity{
+			PodAffinity: &corev1.PodAffinity{
+				RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{*podAffinityTerm, {
+					LabelSelector: &metav1.LabelSelector{
+						MatchLabels: map[string]string{
+							workspace.LabelInstance:  "helloworld",
+							workspace.LabelComponent: workspace.ComponentNameAffinityAssistant,
+						},
+					},
+					TopologyKey: "kubernetes.io/hostname",
+				}},
+			},
+		},
+	}, {
+		description: "affinity annotation with a different name and pod contains podAffinity and nodeAffinity",
+		annotations: map[string]string{
+			workspace.AnnotationAffinityAssistantName: "helloworld",
+		},
+		pod: &corev1.Pod{
+			Spec: corev1.PodSpec{
+				Affinity: &corev1.Affinity{PodAffinity: &corev1.PodAffinity{
+					RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{*podAffinityTerm},
+				}, NodeAffinity: nodeAffinity},
+			},
+		},
+
+		expected: &corev1.Affinity{
+			PodAffinity: &corev1.PodAffinity{
+				RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{*podAffinityTerm, {
+					LabelSelector: &metav1.LabelSelector{
+						MatchLabels: map[string]string{
+							workspace.LabelInstance:  "helloworld",
+							workspace.LabelComponent: workspace.ComponentNameAffinityAssistant,
+						},
+					},
+					TopologyKey: "kubernetes.io/hostname",
+				}},
+			},
+			NodeAffinity: nodeAffinity,
+		},
+	},
+	} {
+		t.Run(tc.description, func(t *testing.T) {
+			ctx := context.Background()
+			f := affinityassistant.NewTransformer(ctx, tc.annotations)
+			got, err := f(tc.pod)
+			if err != nil {
+				t.Fatalf("Transformer failed: %v", err)
+			}
+			if d := cmp.Diff(tc.expected, got.Spec.Affinity); d != "" {
+				t.Errorf("AffinityAssistant diff: %s", diff.PrintWantGot(d))
+			}
+		})
+	}
+}