Merge pull request #353 from gatewayd-io/add-healthz-endpoints

Add healthz endpoints
gatewayd-io · Oct 24, 2023 · 98354ae · 98354ae
2 parents 0cb307a + 32ab95a
commit 98354ae
Show file tree

Hide file tree

Showing 6 changed files with 96 additions and 19 deletions.
diff --git a/api/api.go b/api/api.go
@@ -24,7 +24,9 @@ type Options struct {
 	GRPCNetwork string
 	GRPCAddress string
 	HTTPAddress string
+	Servers     map[string]*network.Server
 }
+
 type API struct {
 	v1.GatewayDAdminAPIServiceServer
 

diff --git a/api/grpc_server.go b/api/grpc_server.go
@@ -5,11 +5,12 @@ import (
 
 	v1 "github.com/gatewayd-io/gatewayd/api/v1"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/grpc/reflection"
 )
 
 // StartGRPCAPI starts the gRPC API.
-func StartGRPCAPI(api *API) {
+func StartGRPCAPI(api *API, healthchecker *HealthChecker) {
 	listener, err := net.Listen(api.Options.GRPCNetwork, api.Options.GRPCAddress)
 	if err != nil {
 		api.Options.Logger.Err(err).Msg("failed to start gRPC API")
@@ -18,6 +19,7 @@ func StartGRPCAPI(api *API) {
 	grpcServer := grpc.NewServer()
 	reflection.Register(grpcServer)
 	v1.RegisterGatewayDAdminAPIServiceServer(grpcServer, api)
+	grpc_health_v1.RegisterHealthServer(grpcServer, healthchecker)
 	if err := grpcServer.Serve(listener); err != nil {
 		api.Options.Logger.Err(err).Msg("failed to start gRPC API")
 	}

diff --git a/api/healthcheck.go b/api/healthcheck.go
@@ -0,0 +1,38 @@
+package api
+
+import (
+	"context"
+
+	"github.com/gatewayd-io/gatewayd/network"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/health/grpc_health_v1"
+	"google.golang.org/grpc/status"
+)
+
+type HealthChecker struct {
+	grpc_health_v1.UnimplementedHealthServer
+
+	Servers map[string]*network.Server
+}
+
+func (h *HealthChecker) Check(
+	context.Context, *grpc_health_v1.HealthCheckRequest,
+) (*grpc_health_v1.HealthCheckResponse, error) {
+	// Check if all servers are running
+	if liveness(h.Servers) {
+		return &grpc_health_v1.HealthCheckResponse{
+			Status: grpc_health_v1.HealthCheckResponse_SERVING,
+		}, nil
+	}
+
+	return &grpc_health_v1.HealthCheckResponse{
+		Status: grpc_health_v1.HealthCheckResponse_NOT_SERVING,
+	}, nil
+}
+
+func (h *HealthChecker) Watch(
+	*grpc_health_v1.HealthCheckRequest,
+	grpc_health_v1.Health_WatchServer,
+) error {
+	return status.Error(codes.Unimplemented, "not implemented") //nolint:wrapcheck
+}
diff --git 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/api/utils.go b/api/utils.go
@@ -0,0 +1,14 @@
+package api
+
+import (
+	"github.com/gatewayd-io/gatewayd/network"
+)
+
+func liveness(servers map[string]*network.Server) bool {
+	for _, v := range servers {
+		if !v.IsRunning() {
+			return false
+		}
+	}
+	return true
+}
diff --git 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)