Fix readinessprobe and livenessprobe for non-nil probehandler overrid…

…es (#1119) (#1124) * Fix readiness and liveness probes for non-nil probehandlers * check for probeHandler * small test update (cherry picked from commit 678c4ec)
DataDog · Mar 18, 2024 · 6f32056 · 6f32056
1 parent 0d24f72
commit 6f32056
Show file tree

Hide file tree

Showing 3 changed files with 100 additions and 4 deletions.
diff --git a/controllers/datadogagent/override/container.go b/controllers/datadogagent/override/container.go
@@ -221,8 +221,8 @@ func overrideAppArmorProfile(containerName commonv1.AgentContainerName, manager
 }
 
 func overrideReadinessProbe(readinessProbeOverride *corev1.Probe) *corev1.Probe {
-	// Add default httpGet.path and httpGet.port if not present in readinessProbe override
-	if readinessProbeOverride.HTTPGet == nil {
+	// Add default httpGet probeHandler if probeHandler is not configured in readinessProbe override
+	if !hasProbeHandler(readinessProbeOverride) {
 		readinessProbeOverride.HTTPGet = &corev1.HTTPGetAction{
 			Path: common.DefaultReadinessProbeHTTPPath,
 			Port: intstr.IntOrString{IntVal: common.DefaultAgentHealthPort}}
@@ -231,8 +231,8 @@ func overrideReadinessProbe(readinessProbeOverride *corev1.Probe) *corev1.Probe
 }
 
 func overrideLivenessProbe(livenessProbeOverride *corev1.Probe) *corev1.Probe {
-	// Add default httpGet.path and httpGet.port if not present in livenessProbe override
-	if livenessProbeOverride.HTTPGet == nil {
+	// Add default httpGet probeHandler if probeHandler is not configured in livenessProbe override
+	if !hasProbeHandler(livenessProbeOverride) {
 		livenessProbeOverride.HTTPGet = &corev1.HTTPGetAction{
 			Path: common.DefaultLivenessProbeHTTPPath,
 			Port: intstr.IntOrString{IntVal: common.DefaultAgentHealthPort}}

diff --git a/controllers/datadogagent/override/container_test.go b/controllers/datadogagent/override/container_test.go
@@ -414,6 +414,49 @@ func TestContainer(t *testing.T) {
 				})
 			},
 		},
+		{
+			name:          "override readiness probe with non-HTTPGet handler",
+			containerName: commonv1.CoreAgentContainerName,
+			existingManager: func() *fake.PodTemplateManagers {
+				return fake.NewPodTemplateManagers(t, corev1.PodTemplateSpec{
+					Spec: corev1.PodSpec{
+						Containers: []corev1.Container{*agentContainer},
+					},
+				})
+			},
+			override: v2alpha1.DatadogAgentGenericContainer{
+				ReadinessProbe: &corev1.Probe{
+					InitialDelaySeconds: 10,
+					TimeoutSeconds:      5,
+					PeriodSeconds:       30,
+					SuccessThreshold:    1,
+					FailureThreshold:    5,
+					ProbeHandler: corev1.ProbeHandler{
+						Exec: &corev1.ExecAction{
+							Command: []string{"echo", "foo", "bar"},
+						},
+					},
+				},
+			},
+			validateManager: func(t *testing.T, manager *fake.PodTemplateManagers, containerName string) {
+				assertContainerMatch(t, manager.PodTemplateSpec().Spec.Containers, containerName, func(container corev1.Container) bool {
+					return reflect.DeepEqual(
+						&corev1.Probe{
+							InitialDelaySeconds: 10,
+							TimeoutSeconds:      5,
+							PeriodSeconds:       30,
+							SuccessThreshold:    1,
+							FailureThreshold:    5,
+							ProbeHandler: corev1.ProbeHandler{
+								Exec: &corev1.ExecAction{
+									Command: []string{"echo", "foo", "bar"},
+								},
+							},
+						},
+						container.ReadinessProbe)
+				})
+			},
+		},
 		{
 			name:          "override readiness probe",
 			containerName: commonv1.CoreAgentContainerName,
@@ -504,6 +547,49 @@ func TestContainer(t *testing.T) {
 				})
 			},
 		},
+		{
+			name:          "override liveness probe with non-HTTPGet handler",
+			containerName: commonv1.CoreAgentContainerName,
+			existingManager: func() *fake.PodTemplateManagers {
+				return fake.NewPodTemplateManagers(t, corev1.PodTemplateSpec{
+					Spec: corev1.PodSpec{
+						Containers: []corev1.Container{*agentContainer},
+					},
+				})
+			},
+			override: v2alpha1.DatadogAgentGenericContainer{
+				LivenessProbe: &corev1.Probe{
+					InitialDelaySeconds: 10,
+					TimeoutSeconds:      5,
+					PeriodSeconds:       30,
+					SuccessThreshold:    1,
+					FailureThreshold:    5,
+					ProbeHandler: corev1.ProbeHandler{
+						Exec: &corev1.ExecAction{
+							Command: []string{"echo", "foo", "bar"},
+						},
+					},
+				},
+			},
+			validateManager: func(t *testing.T, manager *fake.PodTemplateManagers, containerName string) {
+				assertContainerMatch(t, manager.PodTemplateSpec().Spec.Containers, containerName, func(container corev1.Container) bool {
+					return reflect.DeepEqual(
+						&corev1.Probe{
+							InitialDelaySeconds: 10,
+							TimeoutSeconds:      5,
+							PeriodSeconds:       30,
+							SuccessThreshold:    1,
+							FailureThreshold:    5,
+							ProbeHandler: corev1.ProbeHandler{
+								Exec: &corev1.ExecAction{
+									Command: []string{"echo", "foo", "bar"},
+								},
+							},
+						},
+						container.LivenessProbe)
+				})
+			},
+		},
 		{
 			name:          "override liveness probe",
 			containerName: commonv1.CoreAgentContainerName,

diff --git a/controllers/datadogagent/override/utils.go b/controllers/datadogagent/override/utils.go
@@ -8,8 +8,18 @@ package override
 import (
 	"fmt"
 	"strings"
+
+	corev1 "k8s.io/api/core/v1"
 )
 
 func getDefaultConfigMapName(ddaName, fileName string) string {
 	return fmt.Sprintf("%s-%s-yaml", ddaName, strings.Split(fileName, ".")[0])
 }
+
+func hasProbeHandler(probe *corev1.Probe) bool {
+	handler := &probe.ProbeHandler
+	if handler.Exec != nil || handler.HTTPGet != nil || handler.TCPSocket != nil || handler.GRPC != nil {
+		return true
+	}
+	return false
+}