From 32ab95a21c5efc5e3ccb2696951a4324e3262cdd Mon Sep 17 00:00:00 2001
From: Mostafa Moradian <mostafa@gatewayd.io>
Date: Tue, 24 Oct 2023 19:37:52 +0200
Subject: [PATCH] Add healthz endpoint to HTTP server

---
 api/http_server.go | 38 ++++++++++++++++++++++++++++----------
 cmd/run.go         | 19 +++++++++++--------
 2 files changed, 39 insertions(+), 18 deletions(-)

diff --git a/api/http_server.go b/api/http_server.go
index 3f3c23fa..882a083b 100644
--- a/api/http_server.go
+++ b/api/http_server.go
@@ -2,6 +2,7 @@ package api
 
 import (
 	"context"
+	"encoding/json"
 	"io/fs"
 	"net/http"
 
@@ -12,6 +13,10 @@ import (
 	"google.golang.org/grpc/credentials/insecure"
 )
 
+type Healthz struct {
+	Status string `json:"status"`
+}
+
 // StartHTTPAPI starts the HTTP API.
 func StartHTTPAPI(options *Options) {
 	ctx := context.Background()
@@ -30,25 +35,38 @@ func StartHTTPAPI(options *Options) {
 
 	mux := http.NewServeMux()
 	mux.Handle("/", rmux)
-	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
-		w.WriteHeader(http.StatusOK)
+	mux.HandleFunc("/healthz", func(writer http.ResponseWriter, r *http.Request) {
+		if liveness(options.Servers) {
+			writer.Header().Set("Content-Type", "application/json")
+			writer.WriteHeader(http.StatusOK)
+			if err := json.NewEncoder(writer).Encode(Healthz{Status: "SERVING"}); err != nil {
+				options.Logger.Err(err).Msg("failed to serve healthcheck")
+				writer.WriteHeader(http.StatusInternalServerError)
+			}
+		} else {
+			writer.Header().Set("Content-Type", "application/json")
+			writer.WriteHeader(http.StatusServiceUnavailable)
+			if err := json.NewEncoder(writer).Encode(Healthz{Status: "NOT_SERVING"}); err != nil {
+				options.Logger.Err(err).Msg("failed to serve healthcheck")
+			}
+		}
 	})
 
-	mux.HandleFunc("/version", func(w http.ResponseWriter, r *http.Request) {
-		w.WriteHeader(http.StatusOK)
-		if _, err := w.Write([]byte(config.Version)); err != nil {
+	mux.HandleFunc("/version", func(writer http.ResponseWriter, r *http.Request) {
+		writer.WriteHeader(http.StatusOK)
+		if _, err := writer.Write([]byte(config.Version)); err != nil {
 			options.Logger.Err(err).Msg("failed to serve version")
-			w.WriteHeader(http.StatusInternalServerError)
+			writer.WriteHeader(http.StatusInternalServerError)
 		}
 	})
 
 	if IsSwaggerEmbedded() {
-		mux.HandleFunc("/swagger.json", func(w http.ResponseWriter, r *http.Request) {
-			w.WriteHeader(http.StatusOK)
+		mux.HandleFunc("/swagger.json", func(writer http.ResponseWriter, r *http.Request) {
+			writer.WriteHeader(http.StatusOK)
 			data, _ := swaggerUI.ReadFile("v1/api.swagger.json")
-			if _, err := w.Write(data); err != nil {
+			if _, err := writer.Write(data); err != nil {
 				options.Logger.Err(err).Msg("failed to serve swagger.json")
-				w.WriteHeader(http.StatusInternalServerError)
+				writer.WriteHeader(http.StatusInternalServerError)
 			}
 		})
 
diff --git a/cmd/run.go b/cmd/run.go
index a865f91e..3a0944e9 100644
--- a/cmd/run.go
+++ b/cmd/run.go
@@ -696,16 +696,19 @@ var runCmd = &cobra.Command{
 				GRPCNetwork: conf.Global.API.GRPCNetwork,
 				GRPCAddress: conf.Global.API.GRPCAddress,
 				HTTPAddress: conf.Global.API.HTTPAddress,
+				Servers:     servers,
 			}
 
-			go api.StartGRPCAPI(&api.API{
-				Options:        &apiOptions,
-				Config:         conf,
-				PluginRegistry: pluginRegistry,
-				Pools:          pools,
-				Proxies:        proxies,
-				Servers:        servers,
-			})
+			go api.StartGRPCAPI(
+				&api.API{
+					Options:        &apiOptions,
+					Config:         conf,
+					PluginRegistry: pluginRegistry,
+					Pools:          pools,
+					Proxies:        proxies,
+					Servers:        servers,
+				},
+				&api.HealthChecker{Servers: servers})
 			logger.Info().Str("address", apiOptions.HTTPAddress).Msg("Started the HTTP API")
 
 			go api.StartHTTPAPI(&apiOptions)