From 22851eac4701921ae520e4040f420f47def2dd1c 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
---
 k8sutils/statefulset.go | 94 ++++++++++++++++++-----------------------
 1 file changed, 42 insertions(+), 52 deletions(-)

diff --git a/k8sutils/statefulset.go b/k8sutils/statefulset.go
index 280ef3c82..5ddf53a28 100644
--- a/k8sutils/statefulset.go
+++ b/k8sutils/statefulset.go
@@ -336,8 +336,7 @@ func generateContainerDef(name string, containerParams containerParameters, enab
 	return containerDefinition
 }
 
-func GenerateTLSEnvironmentVariables(tlsconfig *redisv1beta1.TLSConfig) []corev1.EnvVar {
-	var envVars []corev1.EnvVar
+func GenerateTLSEnvironmentVariables(enabledMetric bool, tlsconfig *redisv1beta1.TLSConfig) []corev1.EnvVar {
 	root := "/tls/"
 
 	// get and set Defaults
@@ -355,23 +354,21 @@ func GenerateTLSEnvironmentVariables(tlsconfig *redisv1beta1.TLSConfig) []corev1
 		tlsCertKey = tlsconfig.KeyFile
 	}
 
-	envVars = append(envVars, corev1.EnvVar{
-		Name:  "TLS_MODE",
-		Value: "true",
-	})
-	envVars = append(envVars, corev1.EnvVar{
-		Name:  "REDIS_TLS_CA_KEY",
-		Value: path.Join(root, caCert),
-	})
-	envVars = append(envVars, corev1.EnvVar{
-		Name:  "REDIS_TLS_CERT",
-		Value: path.Join(root, tlsCert),
-	})
-	envVars = append(envVars, corev1.EnvVar{
-		Name:  "REDIS_TLS_CERT_KEY",
-		Value: path.Join(root, tlsCertKey),
-	})
-	return envVars
+	if enabledMetric {
+		return []corev1.EnvVar{
+			{Name: "REDIS_EXPORTER_SKIP_TLS_VERIFICATION", Value: "true"},
+			{Name: "REDIS_EXPORTER_TLS_CA_CERT_FILE", Value: path.Join(root, caCert)},
+			{Name: "REDIS_EXPORTER_TLS_CLIENT_CERT_FILE", Value: path.Join(root, tlsCert)},
+			{Name: "REDIS_EXPORTER_TLS_CLIENT_KEY_FILE", Value: path.Join(root, tlsCertKey)},
+		}
+	} else {
+		return []corev1.EnvVar{
+			{Name: "TLS_MODE", Value: "true"},
+			{Name: "REDIS_TLS_CA_KEY", Value: path.Join(root, caCert)},
+			{Name: "REDIS_TLS_CERT", Value: path.Join(root, tlsCert)},
+			{Name: "REDIS_TLS_CERT_KEY", Value: path.Join(root, tlsCertKey)},
+		}
+	}
 }
 
 // enableRedisMonitoring will add Redis Exporter as sidecar container
@@ -447,8 +444,8 @@ func getProbeInfo(probe *redisv1beta1.Probe) *corev1.Probe {
 		ProbeHandler: corev1.ProbeHandler{
 			Exec: &corev1.ExecAction{
 				Command: []string{
-					"bash",
-					"/usr/bin/healthcheck.sh",
+					"redis-cli", "-s", "redis.sock",
+					"ping",
 				},
 			},
 		},
@@ -457,43 +454,36 @@ func getProbeInfo(probe *redisv1beta1.Probe) *corev1.Probe {
 
 // getEnvironmentVariables returns all the required Environment Variables
 func getEnvironmentVariables(role string, enabledMetric bool, enabledPassword *bool, secretName *string, secretKey *string, persistenceEnabled *bool, exporterEnvVar *[]corev1.EnvVar, tlsConfig *redisv1beta1.TLSConfig) []corev1.EnvVar {
-	envVars := []corev1.EnvVar{
-		{Name: "SERVER_MODE", Value: role},
-		{Name: "SETUP_MODE", Value: role},
-	}
+	var envVars []corev1.EnvVar
 
-	redisHost := "redis://localhost:6379"
 	if tlsConfig != nil {
-		redisHost = "rediss://localhost:6379"
-		envVars = append(envVars, GenerateTLSEnvironmentVariables(tlsConfig)...)
-		if enabledMetric {
-			envVars = append(envVars, corev1.EnvVar{
-				Name:  "REDIS_EXPORTER_TLS_CLIENT_KEY_FILE",
-				Value: "/tls/tls.key",
-			})
-			envVars = append(envVars, corev1.EnvVar{
-				Name:  "REDIS_EXPORTER_TLS_CLIENT_CERT_FILE",
-				Value: "/tls/tls.crt",
-			})
-			envVars = append(envVars, corev1.EnvVar{
-				Name:  "REDIS_EXPORTER_TLS_CA_CERT_FILE",
-				Value: "/tls/ca.crt",
-			})
-			envVars = append(envVars, corev1.EnvVar{
-				Name:  "REDIS_EXPORTER_SKIP_TLS_VERIFICATION",
-				Value: "true",
-			})
-		}
+		envVars = append(envVars, GenerateTLSEnvironmentVariables(enabledMetric, tlsConfig)...)
 	}
 
-	envVars = append(envVars, corev1.EnvVar{
-		Name:  "REDIS_ADDR",
-		Value: redisHost,
-	})
+	RedisPassword := "REDISCLI_AUTH"
+	if enabledMetric {
+		envVars = append(envVars, corev1.EnvVar{
+			Name: "REDIS_ADDR",
+			Value: func(role string, tls *redisv1beta1.TLSConfig) string {
+				prefix := "redis://localhost:"
+				if tls != nil {
+					prefix = "rediss://localhost:"
+				}
+				if role == "sentinel" {
+					return prefix + strconv.Itoa(sentinelPort)
+				} else {
+					return prefix + strconv.Itoa(redisPort)
+				}
+			}(role, tlsConfig),
+		})
+		RedisPassword = "REDIS_PASSWORD"
+	} else {
+		envVars = append(envVars, corev1.EnvVar{Name: "SERVER_MODE", Value: role})
+	}
 
 	if enabledPassword != nil && *enabledPassword {
 		envVars = append(envVars, corev1.EnvVar{
-			Name: "REDIS_PASSWORD",
+			Name: RedisPassword,
 			ValueFrom: &corev1.EnvVarSource{
 				SecretKeyRef: &corev1.SecretKeySelector{
 					LocalObjectReference: corev1.LocalObjectReference{
@@ -508,7 +498,7 @@ func getEnvironmentVariables(role string, enabledMetric bool, enabledPassword *b
 		envVars = append(envVars, corev1.EnvVar{Name: "PERSISTENCE_ENABLED", Value: "true"})
 	}
 
-	if exporterEnvVar != nil {
+	if enabledMetric && exporterEnvVar != nil {
 		envVars = append(envVars, *exporterEnvVar...)
 	}
 	sort.SliceStable(envVars, func(i, j int) bool {