From cafe1ded92a481eb182cac130239e172d650b254 Mon Sep 17 00:00:00 2001
From: Federico Tedin <federicotedin@gmail.com>
Date: Wed, 27 Nov 2024 11:22:03 +0100
Subject: [PATCH] Add healthz check

---
 Dockerfile       |  2 +-
 cmd/main.go      |  2 +-
 pkg/grpc/grpc.go | 31 ++++++++++++++++++++++++++-----
 3 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index cf5a017..29070c8 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -24,5 +24,5 @@ FROM gcr.io/distroless/static-debian11
 
 COPY --from=builder /bin/quickpizza /bin
 
-EXPOSE 3333 3334
+EXPOSE 3333 3334 3335
 ENTRYPOINT [ "/bin/quickpizza" ]
diff --git a/cmd/main.go b/cmd/main.go
index 5ef66bc..6a39524 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -129,7 +129,7 @@ func main() {
 	}
 
 	if envServe("QUICKPIZZA_GRPC") {
-		grpcServer := qpgrpc.NewServer(":3334")
+		grpcServer := qpgrpc.NewServer(":3334", ":3335")
 		go func() {
 			err := grpcServer.ListenAndServe()
 			if err != nil {
diff --git a/pkg/grpc/grpc.go b/pkg/grpc/grpc.go
index 851064b..c5c09f4 100644
--- a/pkg/grpc/grpc.go
+++ b/pkg/grpc/grpc.go
@@ -6,6 +6,7 @@ import (
 	"log/slog"
 	"math/rand"
 	"net"
+	"net/http"
 
 	pb "github.com/grafana/quickpizza/pkg/grpc/quickpizza"
 	"google.golang.org/grpc"
@@ -16,8 +17,9 @@ type serverImplementation struct {
 }
 
 type Server struct {
-	grpcServer *grpc.Server
-	listen     string
+	grpcServer    *grpc.Server
+	listen        string
+	healthzListen string
 }
 
 func (s *serverImplementation) Status(_ context.Context, in *pb.StatusRequest) (*pb.StatusResponse, error) {
@@ -34,11 +36,28 @@ func (s *serverImplementation) EvaluatePizza(_ context.Context, in *pb.PizzaEval
 	}, nil
 }
 
-func NewServer(listen string) *Server {
+func NewServer(listen string, healthzListen string) *Server {
 	s := grpc.NewServer()
 	pb.RegisterGRPCServer(s, &serverImplementation{})
 
-	return &Server{grpcServer: s, listen: listen}
+	return &Server{grpcServer: s, listen: listen, healthzListen: healthzListen}
+}
+
+func (s *Server) listenHealthz() {
+	mux := http.NewServeMux()
+	mux.HandleFunc("/healthz", func(w http.ResponseWriter, req *http.Request) {
+		w.WriteHeader(http.StatusNoContent)
+	})
+
+	health := &http.Server{
+		Addr:    s.healthzListen,
+		Handler: mux,
+	}
+
+	slog.Info("Starting QuickPizza gRPC health check server", "listenAddress", s.healthzListen)
+	if err := health.ListenAndServe(); err != nil {
+		slog.Error("Error listening for gRPC health check server", "err", err)
+	}
 }
 
 func (s *Server) ListenAndServe() error {
@@ -47,6 +66,8 @@ func (s *Server) ListenAndServe() error {
 		return fmt.Errorf("failed to listen on port: %w", err)
 	}
 
-	slog.Info("Starting QuickPizza gRPC", "listenAddress", s.listen)
+	go s.listenHealthz()
+
+	slog.Info("Starting QuickPizza gRPC server", "listenAddress", s.listen)
 	return s.grpcServer.Serve(lis)
 }