From 2d71bda83afefc849e870fcd1c460db5d60793b4 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 | 22 ++++++++++++++++++++--
 cmd/run.go         | 19 +++++++++++--------
 2 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/api/http_server.go b/api/http_server.go
index 3f3c23fa..af5f54c2 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,8 +35,21 @@ 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(w http.ResponseWriter, r *http.Request) {
+		if liveness(options.Servers) {
+			w.Header().Set("Content-Type", "application/json")
+			w.WriteHeader(http.StatusOK)
+			if err := json.NewEncoder(w).Encode(Healthz{Status: "SERVING"}); err != nil {
+				options.Logger.Err(err).Msg("failed to serve healthcheck")
+				w.WriteHeader(http.StatusInternalServerError)
+			}
+		} else {
+			w.Header().Set("Content-Type", "application/json")
+			w.WriteHeader(http.StatusServiceUnavailable)
+			if err := json.NewEncoder(w).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) {
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)