From aa1fc6675569dfd37281c2a7a579a504b0d74948 Mon Sep 17 00:00:00 2001 From: David Ortiz Date: Tue, 12 Dec 2023 09:38:16 +0100 Subject: [PATCH] [override/podtemplatespec] Merge affinities (#1004) --- .../datadogagent/override/podtemplatespec.go | 163 +++++++++++++++- .../override/podtemplatespec_test.go | 174 +++++++++++++++--- 2 files changed, 308 insertions(+), 29 deletions(-) diff --git a/controllers/datadogagent/override/podtemplatespec.go b/controllers/datadogagent/override/podtemplatespec.go index 9a7af7f5f..58d16966d 100644 --- a/controllers/datadogagent/override/podtemplatespec.go +++ b/controllers/datadogagent/override/podtemplatespec.go @@ -10,6 +10,9 @@ import ( "sort" "strings" + "github.com/go-logr/logr" + v1 "k8s.io/api/core/v1" + apicommon "github.com/DataDog/datadog-operator/apis/datadoghq/common" "github.com/DataDog/datadog-operator/apis/datadoghq/common/v1" "github.com/DataDog/datadog-operator/apis/datadoghq/v2alpha1" @@ -18,8 +21,6 @@ import ( "github.com/DataDog/datadog-operator/controllers/datadogagent/object/volume" "github.com/DataDog/datadog-operator/pkg/controller/utils/comparison" "github.com/DataDog/datadog-operator/pkg/defaulting" - - "github.com/go-logr/logr" ) // PodTemplateSpec use to override a corev1.PodTemplateSpec with a 2alpha1.DatadogAgentPodTemplateOverride. @@ -118,7 +119,7 @@ func PodTemplateSpec(logger logr.Logger, manager feature.PodTemplateManagers, ov } if override.Affinity != nil { - manager.PodTemplateSpec().Spec.Affinity = override.Affinity + manager.PodTemplateSpec().Spec.Affinity = mergeAffinities(manager.PodTemplateSpec().Spec.Affinity, override.Affinity) } for selectorKey, selectorVal := range override.NodeSelector { @@ -214,6 +215,162 @@ func overrideImage(currentImg string, overrideImg *common.AgentImageConfig) stri return apicommon.GetImage(&overrideImgCopy, ®istry) } +func mergeAffinities(affinity1 *v1.Affinity, affinity2 *v1.Affinity) *v1.Affinity { + if affinity1 == nil && affinity2 == nil { + return nil + } + + if affinity1 == nil { + return affinity2 + } + + if affinity2 == nil { + return affinity1 + } + + merged := &v1.Affinity{} + + merged.NodeAffinity = mergeNodeAffinities(affinity1.NodeAffinity, affinity2.NodeAffinity) + merged.PodAffinity = mergePodAffinities(affinity1.PodAffinity, affinity2.PodAffinity) + merged.PodAntiAffinity = mergePodAntiAffinities(affinity1.PodAntiAffinity, affinity2.PodAntiAffinity) + + return merged +} + +func mergeNodeAffinities(affinity1 *v1.NodeAffinity, affinity2 *v1.NodeAffinity) *v1.NodeAffinity { + if affinity1 == nil && affinity2 == nil { + return nil + } + + if affinity1 == nil { + return affinity2 + } + + if affinity2 == nil { + return affinity1 + } + + merged := &v1.NodeAffinity{} + + merged.RequiredDuringSchedulingIgnoredDuringExecution = mergeNodeSelectors( + affinity1.RequiredDuringSchedulingIgnoredDuringExecution, + affinity2.RequiredDuringSchedulingIgnoredDuringExecution, + ) + + merged.PreferredDuringSchedulingIgnoredDuringExecution = append( + merged.PreferredDuringSchedulingIgnoredDuringExecution, + affinity1.PreferredDuringSchedulingIgnoredDuringExecution..., + ) + merged.PreferredDuringSchedulingIgnoredDuringExecution = append( + merged.PreferredDuringSchedulingIgnoredDuringExecution, + affinity2.PreferredDuringSchedulingIgnoredDuringExecution..., + ) + + return merged +} + +func mergeNodeSelectors(selector1 *v1.NodeSelector, selector2 *v1.NodeSelector) *v1.NodeSelector { + if selector1 == nil && selector2 == nil { + return nil + } + + if selector1 == nil { + return selector2 + } + + if selector2 == nil { + return selector1 + } + + merged := &v1.NodeSelector{} + + // Note that the NodeSelectorTerms are ORed together. + for _, term1 := range selector1.NodeSelectorTerms { + for _, term2 := range selector2.NodeSelectorTerms { + mergedTerm := v1.NodeSelectorTerm{ + // These are ANDed together. + MatchExpressions: append(term1.MatchExpressions, term2.MatchExpressions...), + MatchFields: append(term1.MatchFields, term2.MatchFields...), + } + merged.NodeSelectorTerms = append(merged.NodeSelectorTerms, mergedTerm) + } + } + + return merged +} + +func mergePodAffinities(affinity1 *v1.PodAffinity, affinity2 *v1.PodAffinity) *v1.PodAffinity { + if affinity1 == nil && affinity2 == nil { + return nil + } + + if affinity1 == nil { + return affinity2 + } + + if affinity2 == nil { + return affinity1 + } + + merged := &v1.PodAffinity{} + + merged.RequiredDuringSchedulingIgnoredDuringExecution = append( + merged.RequiredDuringSchedulingIgnoredDuringExecution, + affinity1.RequiredDuringSchedulingIgnoredDuringExecution..., + ) + merged.RequiredDuringSchedulingIgnoredDuringExecution = append( + merged.RequiredDuringSchedulingIgnoredDuringExecution, + affinity2.RequiredDuringSchedulingIgnoredDuringExecution..., + ) + + merged.PreferredDuringSchedulingIgnoredDuringExecution = append( + merged.PreferredDuringSchedulingIgnoredDuringExecution, + affinity1.PreferredDuringSchedulingIgnoredDuringExecution..., + ) + merged.PreferredDuringSchedulingIgnoredDuringExecution = append( + merged.PreferredDuringSchedulingIgnoredDuringExecution, + affinity2.PreferredDuringSchedulingIgnoredDuringExecution..., + ) + + return merged +} + +func mergePodAntiAffinities(affinity1 *v1.PodAntiAffinity, affinity2 *v1.PodAntiAffinity) *v1.PodAntiAffinity { + if affinity1 == nil && affinity2 == nil { + return nil + } + + if affinity1 == nil { + return affinity2 + } + + if affinity2 == nil { + return affinity1 + } + + merged := &v1.PodAntiAffinity{} + + merged.RequiredDuringSchedulingIgnoredDuringExecution = append( + merged.RequiredDuringSchedulingIgnoredDuringExecution, + affinity1.RequiredDuringSchedulingIgnoredDuringExecution..., + ) + merged.RequiredDuringSchedulingIgnoredDuringExecution = append( + merged.RequiredDuringSchedulingIgnoredDuringExecution, + affinity2.RequiredDuringSchedulingIgnoredDuringExecution..., + ) + + merged.PreferredDuringSchedulingIgnoredDuringExecution = append( + merged.PreferredDuringSchedulingIgnoredDuringExecution, + affinity1.PreferredDuringSchedulingIgnoredDuringExecution..., + ) + merged.PreferredDuringSchedulingIgnoredDuringExecution = append( + merged.PreferredDuringSchedulingIgnoredDuringExecution, + affinity2.PreferredDuringSchedulingIgnoredDuringExecution..., + ) + + return merged +} + func sortKeys(keysMap map[v2alpha1.AgentConfigFileName]v2alpha1.CustomConfig) []v2alpha1.AgentConfigFileName { sortedKeys := make([]v2alpha1.AgentConfigFileName, 0, len(keysMap)) for key := range keysMap { diff --git a/controllers/datadogagent/override/podtemplatespec_test.go b/controllers/datadogagent/override/podtemplatespec_test.go index 72ae97d88..dc97f1af7 100644 --- a/controllers/datadogagent/override/podtemplatespec_test.go +++ b/controllers/datadogagent/override/podtemplatespec_test.go @@ -8,16 +8,16 @@ package override import ( "testing" + "github.com/stretchr/testify/assert" + v1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "sigs.k8s.io/controller-runtime/pkg/log/zap" + "github.com/DataDog/datadog-operator/apis/datadoghq/common" commonv1 "github.com/DataDog/datadog-operator/apis/datadoghq/common/v1" "github.com/DataDog/datadog-operator/apis/datadoghq/v2alpha1" apiutils "github.com/DataDog/datadog-operator/apis/utils" "github.com/DataDog/datadog-operator/controllers/datadogagent/feature/fake" - - "github.com/stretchr/testify/assert" - v1 "k8s.io/api/core/v1" - metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - "sigs.k8s.io/controller-runtime/pkg/log/zap" ) func TestPodTemplateSpec(t *testing.T) { @@ -663,18 +663,41 @@ func TestPodTemplateSpec(t *testing.T) { existingManager: func() *fake.PodTemplateManagers { manager := fake.NewPodTemplateManagers(t, v1.PodTemplateSpec{}) manager.PodTemplateSpec().Spec.Affinity = &v1.Affinity{ + NodeAffinity: &v1.NodeAffinity{ + RequiredDuringSchedulingIgnoredDuringExecution: &v1.NodeSelector{ + NodeSelectorTerms: []v1.NodeSelectorTerm{ + { + MatchExpressions: []v1.NodeSelectorRequirement{ + { + Key: "node-label-1", + Operator: v1.NodeSelectorOpExists, + }, + }, + }, + }, + }, + }, + PodAffinity: &v1.PodAffinity{ + RequiredDuringSchedulingIgnoredDuringExecution: []v1.PodAffinityTerm{ + { + LabelSelector: &metav1.LabelSelector{ + MatchLabels: map[string]string{ + "pod-label-1": "value-1", + }, + }, + TopologyKey: v1.LabelHostname, + }, + }, + }, PodAntiAffinity: &v1.PodAntiAffinity{ - PreferredDuringSchedulingIgnoredDuringExecution: []v1.WeightedPodAffinityTerm{ + RequiredDuringSchedulingIgnoredDuringExecution: []v1.PodAffinityTerm{ { - Weight: 50, - PodAffinityTerm: v1.PodAffinityTerm{ - LabelSelector: &metav1.LabelSelector{ - MatchLabels: map[string]string{ - "old-label": "123", - }, + LabelSelector: &metav1.LabelSelector{ + MatchLabels: map[string]string{ + "pod-label-2": "value-2", }, - TopologyKey: "kubernetes.io/hostname", }, + TopologyKey: v1.LabelHostname, }, }, }, @@ -683,29 +706,128 @@ func TestPodTemplateSpec(t *testing.T) { }, override: v2alpha1.DatadogAgentComponentOverride{ Affinity: &v1.Affinity{ + NodeAffinity: &v1.NodeAffinity{ + RequiredDuringSchedulingIgnoredDuringExecution: &v1.NodeSelector{ + NodeSelectorTerms: []v1.NodeSelectorTerm{ + { + MatchExpressions: []v1.NodeSelectorRequirement{ + { + Key: "node-label-2", + Operator: v1.NodeSelectorOpExists, + }, + }, + }, + { + MatchExpressions: []v1.NodeSelectorRequirement{ + { + Key: "node-label-3", + Operator: v1.NodeSelectorOpExists, + }, + }, + }, + }, + }, + }, + PodAffinity: &v1.PodAffinity{ + RequiredDuringSchedulingIgnoredDuringExecution: []v1.PodAffinityTerm{ + { + LabelSelector: &metav1.LabelSelector{ + MatchLabels: map[string]string{ + "pod-label-3": "value-3", + }, + }, + TopologyKey: v1.LabelHostname, + }, + }, + }, PodAntiAffinity: &v1.PodAntiAffinity{ - PreferredDuringSchedulingIgnoredDuringExecution: []v1.WeightedPodAffinityTerm{ + RequiredDuringSchedulingIgnoredDuringExecution: []v1.PodAffinityTerm{ { - Weight: 50, - PodAffinityTerm: v1.PodAffinityTerm{ - LabelSelector: &metav1.LabelSelector{ - MatchLabels: map[string]string{ - "new-label": "456", // Changed - }, + LabelSelector: &metav1.LabelSelector{ + MatchLabels: map[string]string{ + "pod-label-4": "value-4", }, - TopologyKey: "kubernetes.io/hostname", }, + TopologyKey: v1.LabelHostname, }, }, }, }, }, validateManager: func(t *testing.T, manager *fake.PodTemplateManagers) { - assert.Equal(t, - map[string]string{"new-label": "456"}, - manager.PodTemplateSpec().Spec.Affinity.PodAntiAffinity. - PreferredDuringSchedulingIgnoredDuringExecution[0]. - PodAffinityTerm.LabelSelector.MatchLabels) + expectedAffinity := &v1.Affinity{ + NodeAffinity: &v1.NodeAffinity{ + RequiredDuringSchedulingIgnoredDuringExecution: &v1.NodeSelector{ + NodeSelectorTerms: []v1.NodeSelectorTerm{ + { + MatchExpressions: []v1.NodeSelectorRequirement{ + { + Key: "node-label-1", + Operator: v1.NodeSelectorOpExists, + }, + { + Key: "node-label-2", + Operator: v1.NodeSelectorOpExists, + }, + }, + }, + { + MatchExpressions: []v1.NodeSelectorRequirement{ + { + Key: "node-label-1", + Operator: v1.NodeSelectorOpExists, + }, + { + Key: "node-label-3", + Operator: v1.NodeSelectorOpExists, + }, + }, + }, + }, + }, + }, + PodAffinity: &v1.PodAffinity{ + RequiredDuringSchedulingIgnoredDuringExecution: []v1.PodAffinityTerm{ + { + LabelSelector: &metav1.LabelSelector{ + MatchLabels: map[string]string{ + "pod-label-1": "value-1", + }, + }, + TopologyKey: v1.LabelHostname, + }, + { + LabelSelector: &metav1.LabelSelector{ + MatchLabels: map[string]string{ + "pod-label-3": "value-3", + }, + }, + TopologyKey: v1.LabelHostname, + }, + }, + }, + PodAntiAffinity: &v1.PodAntiAffinity{ + RequiredDuringSchedulingIgnoredDuringExecution: []v1.PodAffinityTerm{ + { + LabelSelector: &metav1.LabelSelector{ + MatchLabels: map[string]string{ + "pod-label-2": "value-2", + }, + }, + TopologyKey: v1.LabelHostname, + }, + { + LabelSelector: &metav1.LabelSelector{ + MatchLabels: map[string]string{ + "pod-label-4": "value-4", + }, + }, + TopologyKey: v1.LabelHostname, + }, + }, + }, + } + assert.Equal(t, expectedAffinity, manager.PodTemplateSpec().Spec.Affinity) }, }, {