Add TCP probe to target port if no readiness probe is found in the API spec (#2379)

Author: RobertLucian

cmd/proxy/main.go

@@ -19,6 +19,7 @@ package main
 import (
 	"context"
 	"flag"
+	"fmt"
 	"net"
 	"net/http"
 	"os"
@@ -48,6 +49,7 @@ func main() {
 		userContainerPort int
 		maxConcurrency    int
 		maxQueueLength    int
+		hasTCPProbe       bool
 		clusterConfigPath string
 	)
 
@@ -56,6 +58,7 @@ func main() {
 	flag.IntVar(&userContainerPort, "user-port", 8080, "port where the proxy will redirect to the traffic to")
 	flag.IntVar(&maxConcurrency, "max-concurrency", 0, "max concurrency allowed for user container")
 	flag.IntVar(&maxQueueLength, "max-queue-length", 0, "max request queue length for user container")
+	flag.BoolVar(&hasTCPProbe, "has-tcp-probe", false, "tcp probe to the user-provided container port")
 	flag.StringVar(&clusterConfigPath, "cluster-config", "", "cluster config path")
 	flag.Parse()
 
@@ -142,7 +145,7 @@ func main() {
 
 	adminHandler := http.NewServeMux()
 	adminHandler.Handle("/metrics", promStats)
-	adminHandler.Handle("/healthz", readinessTCPHandler(userContainerPort, log))
+	adminHandler.Handle("/healthz", readinessTCPHandler(userContainerPort, hasTCPProbe, log))
 
 	servers := map[string]*http.Server{
 		"proxy": {
@@ -201,19 +204,22 @@ func exit(log *zap.SugaredLogger, err error, wrapStrs ...string) {
 	os.Exit(1)
 }
 
-func readinessTCPHandler(port int, logger *zap.SugaredLogger) http.HandlerFunc {
+func readinessTCPHandler(port int, enableTCPProbe bool, logger *zap.SugaredLogger) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
-		timeout := time.Duration(1) * time.Second
-		address := net.JoinHostPort("localhost", strconv.FormatInt(int64(port), 10))
-
-		conn, err := net.DialTimeout("tcp", address, timeout)
-		if err != nil {
-			logger.Warn(errors.Wrap(err, "TCP probe to user-provided container port failed"))
-			w.WriteHeader(http.StatusInternalServerError)
-			_, _ = w.Write([]byte("unhealthy"))
-			return
+		if enableTCPProbe {
+			ctx := r.Context()
+			address := net.JoinHostPort("localhost", fmt.Sprintf("%d", port))
+
+			var d net.Dialer
+			conn, err := d.DialContext(ctx, "tcp", address)
+			if err != nil {
+				logger.Warn(errors.Wrap(err, "TCP probe to user-provided container port failed"))
+				w.WriteHeader(http.StatusInternalServerError)
+				_, _ = w.Write([]byte("unhealthy"))
+				return
+			}
+			_ = conn.Close()
 		}
-		_ = conn.Close()
 
 		w.WriteHeader(http.StatusOK)
 		_, _ = w.Write([]byte("healthy"))

pkg/workloads/helpers.go

@@ -83,7 +83,6 @@ func GetReadinessProbesFromContainers(containers []*userconfig.Container) map[st
 		if container == nil {
 			continue
 		}
-
 		if container.ReadinessProbe != nil {
 			probes[container.Name] = *GetProbeSpec(container.ReadinessProbe)
 		}
@@ -92,6 +91,21 @@ func GetReadinessProbesFromContainers(containers []*userconfig.Container) map[st
 	return probes
 }
 
+func HasReadinessProbesTargetingPort(containers []*userconfig.Container, targetPort int32) bool {
+	for _, container := range containers {
+		if container == nil || container.ReadinessProbe == nil {
+			continue
+		}
+
+		probe := container.ReadinessProbe
+		if (probe.TCPSocket != nil && probe.TCPSocket.Port == targetPort) ||
+			probe.HTTPGet != nil && probe.HTTPGet.Port == targetPort {
+			return true
+		}
+	}
+	return false
+}
+
 func BaseClusterEnvVars() []kcore.EnvFromSource {
 	envVars := []kcore.EnvFromSource{
 		{

pkg/workloads/k8s.go

@@ -172,6 +172,8 @@ func batchDequeuerProxyContainer(api spec.API, jobID, queueURL string) (kcore.Co
 }
 
 func realtimeProxyContainer(api spec.API) (kcore.Container, kcore.Volume) {
+	proxyHasTCPProbe := !HasReadinessProbesTargetingPort(api.Pod.Containers, *api.Pod.Port)
+
 	return kcore.Container{
 		Name:            ProxyContainerName,
 		Image:           config.ClusterConfig.ImageProxy,
@@ -189,6 +191,8 @@ func realtimeProxyContainer(api spec.API) (kcore.Container, kcore.Volume) {
 			s.Int32(int32(api.Pod.MaxConcurrency)),
 			"--max-queue-length",
 			s.Int32(int32(api.Pod.MaxQueueLength)),
+			"--has-tcp-probe",
+			s.Bool(proxyHasTCPProbe),
 		},
 		Ports: []kcore.ContainerPort{
 			{Name: consts.AdminPortName, ContainerPort: consts.AdminPortInt32},
@@ -213,10 +217,10 @@ func realtimeProxyContainer(api spec.API) (kcore.Container, kcore.Volume) {
 				},
 			},
 			InitialDelaySeconds: 1,
-			TimeoutSeconds:      1,
+			TimeoutSeconds:      3,
 			PeriodSeconds:       10,
 			SuccessThreshold:    1,
-			FailureThreshold:    1,
+			FailureThreshold:    3,
 		},
 	}, ClusterConfigVolume()
 }