Add profile label key to node affinity (#1302) (#1324)

Co-authored-by: khewonc <[email protected]>
DataDog · Jul 26, 2024 · e8bf30c · e8bf30c
1 parent e4190d5
commit e8bf30c
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 2 deletions.
diff --git a/controllers/datadogagent_controller_v2_profiles_test.go b/controllers/datadogagent_controller_v2_profiles_test.go
@@ -169,6 +169,11 @@ var _ = Describe("V2 Controller - DatadogAgentProfile", func() {
 											Operator: v1.NodeSelectorOpIn,
 											Values:   []string{"3"},
 										},
+										{
+											Key:      agentprofile.ProfileLabelKey,
+											Operator: v1.NodeSelectorOpIn,
+											Values:   []string{profile.Name},
+										},
 									},
 								},
 							},
@@ -291,6 +296,11 @@ var _ = Describe("V2 Controller - DatadogAgentProfile", func() {
 											Operator: v1.NodeSelectorOpIn,
 											Values:   []string{"1", "2"},
 										},
+										{
+											Key:      agentprofile.ProfileLabelKey,
+											Operator: v1.NodeSelectorOpIn,
+											Values:   []string{profile.Name},
+										},
 									},
 								},
 							},
@@ -418,6 +428,11 @@ var _ = Describe("V2 Controller - DatadogAgentProfile", func() {
 											Operator: v1.NodeSelectorOpIn,
 											Values:   []string{"1"},
 										},
+										{
+											Key:      agentprofile.ProfileLabelKey,
+											Operator: v1.NodeSelectorOpIn,
+											Values:   []string{profile.Name},
+										},
 									},
 								},
 							},
@@ -577,6 +592,11 @@ var _ = Describe("V2 Controller - DatadogAgentProfile", func() {
 											Operator: v1.NodeSelectorOpIn,
 											Values:   []string{"1"},
 										},
+										{
+											Key:      agentprofile.ProfileLabelKey,
+											Operator: v1.NodeSelectorOpIn,
+											Values:   []string{profiles[0].Name},
+										},
 									},
 								},
 							},
@@ -609,6 +629,11 @@ var _ = Describe("V2 Controller - DatadogAgentProfile", func() {
 											Operator: v1.NodeSelectorOpIn,
 											Values:   []string{"2"},
 										},
+										{
+											Key:      agentprofile.ProfileLabelKey,
+											Operator: v1.NodeSelectorOpIn,
+											Values:   []string{profiles[1].Name},
+										},
 									},
 								},
 							},
@@ -762,6 +787,11 @@ var _ = Describe("V2 Controller - DatadogAgentProfile", func() {
 											Operator: v1.NodeSelectorOpIn,
 											Values:   []string{"1"},
 										},
+										{
+											Key:      agentprofile.ProfileLabelKey,
+											Operator: v1.NodeSelectorOpIn,
+											Values:   []string{profiles[0].Name},
+										},
 									},
 								},
 							},
@@ -903,6 +933,11 @@ var _ = Describe("V2 Controller - DatadogAgentProfile", func() {
 											Operator: v1.NodeSelectorOpIn,
 											Values:   []string{"1"},
 										},
+										{
+											Key:      agentprofile.ProfileLabelKey,
+											Operator: v1.NodeSelectorOpIn,
+											Values:   []string{profile.Name},
+										},
 									},
 								},
 							},
@@ -1027,6 +1062,11 @@ var _ = Describe("V2 Controller - DatadogAgentProfile", func() {
 											Operator: v1.NodeSelectorOpIn,
 											Values:   []string{"2"},
 										},
+										{
+											Key:      agentprofile.ProfileLabelKey,
+											Operator: v1.NodeSelectorOpIn,
+											Values:   []string{profile.Name},
+										},
 									},
 								},
 							},

diff --git a/pkg/agentprofile/agent_profile.go b/pkg/agentprofile/agent_profile.go
@@ -190,7 +190,7 @@ func affinityOverride(profile *datadoghqv1alpha1.DatadogAgentProfile) *v1.Affini
 		RequiredDuringSchedulingIgnoredDuringExecution: &v1.NodeSelector{
 			NodeSelectorTerms: []v1.NodeSelectorTerm{
 				{
-					MatchExpressions: profile.Spec.ProfileAffinity.ProfileNodeAffinity,
+					MatchExpressions: append(profile.Spec.ProfileAffinity.ProfileNodeAffinity, profileLabelKeyNSR(profile.Name)),
 				},
 			},
 		},
@@ -199,6 +199,17 @@ func affinityOverride(profile *datadoghqv1alpha1.DatadogAgentProfile) *v1.Affini
 	return affinity
 }
 
+// profileLabelKeyNSR returns the NodeSelectorRequirement for a profile to be
+// applied to nodes with the following label:
+// agent.datadoghq.com/datadogagentprofile:<profile-name>
+func profileLabelKeyNSR(profileName string) v1.NodeSelectorRequirement {
+	return v1.NodeSelectorRequirement{
+		Key:      ProfileLabelKey,
+		Operator: v1.NodeSelectorOpIn,
+		Values:   []string{profileName},
+	}
+}
+
 // affinityOverrideForDefaultProfile returns the affinity override that should
 // be applied to the default profile. The default profile should be applied to
 // all nodes that don't have the agent.datadoghq.com/datadogagentprofile label.

diff --git a/pkg/agentprofile/agent_profile_test.go b/pkg/agentprofile/agent_profile_test.go
@@ -320,6 +320,11 @@ func TestOverrideFromProfile(t *testing.T) {
 											Operator: v1.NodeSelectorOpIn,
 											Values:   []string{"linux"},
 										},
+										{
+											Key:      ProfileLabelKey,
+											Operator: v1.NodeSelectorOpIn,
+											Values:   []string{"linux"},
+										},
 									},
 								},
 							},
@@ -396,7 +401,7 @@ func TestDaemonSetName(t *testing.T) {
 	}
 }
 
-func TestPriorityClassNameOverride(t *testing.T) {
+func Test_priorityClassNameOverride(t *testing.T) {
 	tests := []struct {
 		name                  string
 		profile               v1alpha1.DatadogAgentProfile
@@ -474,6 +479,7 @@ func TestPriorityClassNameOverride(t *testing.T) {
 		})
 	}
 }
+
 func Test_labelsOverride(t *testing.T) {
 	tests := []struct {
 		name           string