From 635106564b54059ee2e349789941be3f648dbe86 Mon Sep 17 00:00:00 2001
From: muicoder <muicoder@gmail.com>
Date: Tue, 21 Feb 2023 15:01:29 +0800
Subject: [PATCH] Probe use built-in, discarded healthcheck.sh

Signed-off-by: muicoder <muicoder@gmail.com>
https://github.com/redis/redis/blob/unstable/TLS.md
The probe uses env mode to load the TLS certificate.
---
 k8sutils/statefulset.go      | 41 ++++++++++++++++++++++++++++++------
 k8sutils/statefulset_test.go | 16 ++++++++++++++
 2 files changed, 50 insertions(+), 7 deletions(-)

diff --git a/k8sutils/statefulset.go b/k8sutils/statefulset.go
index 67f394a6e..69323d972 100644
--- a/k8sutils/statefulset.go
+++ b/k8sutils/statefulset.go
@@ -363,8 +363,8 @@ func generateContainerDef(name string, containerParams containerParameters, clus
 				containerParams.Port,
 				clusterVersion,
 			),
-			ReadinessProbe: getProbeInfo(containerParams.ReadinessProbe),
-			LivenessProbe:  getProbeInfo(containerParams.LivenessProbe),
+			ReadinessProbe: getProbeInfo(containerParams, "R"),
+			LivenessProbe:  getProbeInfo(containerParams, "L"),
 			VolumeMounts:   getVolumeMount(name, containerParams.PersistenceEnabled, clusterMode, nodeConfVolume, externalConfig, mountpath, containerParams.TLSConfig, containerParams.ACLConfig),
 		},
 	}
@@ -591,7 +591,33 @@ func getVolumeMount(name string, persistenceEnabled *bool, clusterMode bool, nod
 }
 
 // getProbeInfo generate probe for Redis StatefulSet
-func getProbeInfo(probe *commonapi.Probe) *corev1.Probe {
+func getProbeInfo(params containerParameters, probeType string) *corev1.Probe {
+	probePort := redisPort
+	if params.Role == "sentinel" {
+		probePort = sentinelPort
+	}
+
+	probeCommand := []string{
+		"redis-cli", "-p", strconv.Itoa(probePort),
+	}
+
+	if params.TLSConfig != nil {
+		probeCommand = append(probeCommand, "--tls")
+		probeCommand = append(probeCommand, "--cacert", "$(REDIS_TLS_CA_KEY)")
+		probeCommand = append(probeCommand, "--cert", "$(REDIS_TLS_CERT)", "--key", "$(REDIS_TLS_CERT_KEY)")
+	}
+	probeCommand = append(probeCommand, "ping")
+
+	var probe *commonapi.Probe
+	switch probeType {
+	case "R":
+		probe = params.ReadinessProbe
+	case "L":
+		probe = params.LivenessProbe
+	default:
+		probe = params.LivenessProbe
+	}
+
 	return &corev1.Probe{
 		InitialDelaySeconds: probe.InitialDelaySeconds,
 		PeriodSeconds:       probe.PeriodSeconds,
@@ -600,10 +626,7 @@ func getProbeInfo(probe *commonapi.Probe) *corev1.Probe {
 		SuccessThreshold:    probe.SuccessThreshold,
 		ProbeHandler: corev1.ProbeHandler{
 			Exec: &corev1.ExecAction{
-				Command: []string{
-					"bash",
-					"/usr/bin/healthcheck.sh",
-				},
+				Command: probeCommand,
 			},
 		},
 	}
@@ -670,6 +693,10 @@ func getEnvironmentVariables(role string, enabledPassword *bool, secretName *str
 				},
 			},
 		})
+		envVars = append(envVars, corev1.EnvVar{
+			Name:      "REDISCLI_AUTH",
+			ValueFrom: envVars[len(envVars)-1].ValueFrom,
+		})
 	}
 	if persistenceEnabled != nil && *persistenceEnabled {
 		envVars = append(envVars, corev1.EnvVar{Name: "PERSISTENCE_ENABLED", Value: "true"})
diff --git a/k8sutils/statefulset_test.go b/k8sutils/statefulset_test.go
index 879c3242e..6b8e77caa 100644
--- a/k8sutils/statefulset_test.go
+++ b/k8sutils/statefulset_test.go
@@ -275,6 +275,14 @@ func TestGetEnvironmentVariables(t *testing.T) {
 						Key: "test-key",
 					},
 				}},
+				{Name: "REDISCLI_AUTH", ValueFrom: &corev1.EnvVarSource{
+					SecretKeyRef: &corev1.SecretKeySelector{
+						LocalObjectReference: corev1.LocalObjectReference{
+							Name: "test-secret",
+						},
+						Key: "test-key",
+					},
+				}},
 				{Name: "SERVER_MODE", Value: "sentinel"},
 				{Name: "SETUP_MODE", Value: "sentinel"},
 				{Name: "TEST_ENV", Value: "test-value"},
@@ -340,6 +348,14 @@ func TestGetEnvironmentVariables(t *testing.T) {
 						Key: "test-key",
 					},
 				}},
+				{Name: "REDISCLI_AUTH", ValueFrom: &corev1.EnvVarSource{
+					SecretKeyRef: &corev1.SecretKeySelector{
+						LocalObjectReference: corev1.LocalObjectReference{
+							Name: "test-secret",
+						},
+						Key: "test-key",
+					},
+				}},
 				{Name: "SERVER_MODE", Value: "cluster"},
 				{Name: "SETUP_MODE", Value: "cluster"},
 				{Name: "TEST_ENV", Value: "test-value"},