From 2990f2cd0e5724289cae6e2860eff98f832b5e46 Mon Sep 17 00:00:00 2001
From: Evan Baker <rbtr@users.noreply.github.com>
Date: Thu, 7 Dec 2023 19:59:27 +0000
Subject: [PATCH] feat: enable readyz/healthz endpoints

Signed-off-by: Evan Baker <rbtr@users.noreply.github.com>
---
 cns/healthserver/server.go |  6 +++---
 cns/service/main.go        | 31 +++++++++++++++++++++----------
 2 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/cns/healthserver/server.go b/cns/healthserver/server.go
index f2bf6416df..e9b75df11e 100644
--- a/cns/healthserver/server.go
+++ b/cns/healthserver/server.go
@@ -6,14 +6,14 @@ import (
 	"github.com/labstack/echo/v4"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"go.uber.org/zap"
-	"sigs.k8s.io/controller-runtime/pkg/healthz"
 	"sigs.k8s.io/controller-runtime/pkg/metrics"
 )
 
-func Start(log *zap.Logger, addr string) {
+func Start(log *zap.Logger, addr string, readyz, healthz http.Handler) {
 	e := echo.New()
 	e.HideBanner = true
-	e.GET("/healthz", echo.WrapHandler(http.StripPrefix("/healthz", &healthz.Handler{})))
+	e.GET("/healthz", echo.WrapHandler(http.StripPrefix("/healthz", healthz)))
+	e.GET("/readyz", echo.WrapHandler(http.StripPrefix("/readyz", readyz)))
 	e.GET("/metrics", echo.WrapHandler(promhttp.HandlerFor(metrics.Registry, promhttp.HandlerOpts{
 		ErrorHandling: promhttp.HTTPErrorOnError,
 	})))
diff --git a/cns/service/main.go b/cns/service/main.go
index d024876776..d80b90b286 100644
--- a/cns/service/main.go
+++ b/cns/service/main.go
@@ -60,6 +60,7 @@ import (
 	"github.com/avast/retry-go/v4"
 	"github.com/pkg/errors"
 	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
 	corev1 "k8s.io/api/core/v1"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -579,9 +580,24 @@ func main() {
 		}
 	}
 
-	// start the health server
-	z, _ := zap.NewProduction()
-	go healthserver.Start(z, cnsconfig.MetricsBindAddress)
+	// configure zap logger
+	zconfig := zap.NewProductionConfig()
+	zconfig.EncoderConfig.EncodeTime = zapcore.ISO8601TimeEncoder
+	z, _ := zconfig.Build()
+
+	// start the healthz/readyz/metrics server
+	readyCh := make(chan interface{})
+	readyChecker := healthz.CheckHandler{
+		Checker: healthz.Checker(func(*http.Request) error {
+			select {
+			default:
+				return errors.New("not ready")
+			case <-readyCh:
+			}
+			return nil
+		}),
+	}
+	go healthserver.Start(z, cnsconfig.MetricsBindAddress, &healthz.Handler{}, readyChecker)
 
 	nmaConfig, err := nmagent.NewConfig(cnsconfig.WireserverIP)
 	if err != nil {
@@ -953,6 +969,8 @@ func main() {
 		}
 	}
 
+	// mark the service as "ready"
+	close(readyCh)
 	// block until process exiting
 	<-rootCtx.Done()
 
@@ -1353,13 +1371,6 @@ func InitializeCRDState(ctx context.Context, httpRestService cns.HTTPService, cn
 		httpRestService.AttachSWIFTv2Middleware(&swiftV2Middleware)
 	}
 
-	// adding some routes to the root service mux
-	mux := httpRestServiceImplementation.Listener.GetMux()
-	mux.Handle("/readyz", http.StripPrefix("/readyz", &healthz.Handler{}))
-	if cnsconfig.EnablePprof {
-		httpRestServiceImplementation.RegisterPProfEndpoints()
-	}
-
 	// start the pool Monitor before the Reconciler, since it needs to be ready to receive an
 	// NodeNetworkConfig update by the time the Reconciler tries to send it.
 	go func() {