diff --git a/pkg/internal/affinityassistant/transformer.go b/pkg/internal/affinityassistant/transformer.go index ef75c5cc4ea..d83f40c1597 100644 --- 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 index c845c8d0359..deb0c248388 100644 --- 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)) + } + }) + } +}