feat: added health handler

Author: Semior001

pkg/proxy/middleware/middleware.go

@@ -11,6 +11,8 @@ import (
 	"google.golang.org/grpc/peer"
 	"google.golang.org/grpc/status"
 	"github.com/Semior001/groxy/pkg/grpcx"
+	"google.golang.org/grpc/health"
+	healthpb "google.golang.org/grpc/health/grpc_health_v1"
 )
 
 // Middleware is a function that intercepts the execution of a gRPC handler.
@@ -64,8 +66,41 @@ func PassMetadata() Middleware {
 	}
 }
 
+// Health serves the health check requests.
+func Health(h *health.Server) Middleware {
+	return func(next grpc.StreamHandler) grpc.StreamHandler {
+		return func(srv any, stream grpc.ServerStream) error {
+			ctx := stream.Context()
+			mtd, ok := grpc.Method(ctx)
+			if !ok {
+				return next(srv, stream)
+			}
+
+			switch mtd {
+			case "/grpc.health.v1.Health/Check":
+				req := &healthpb.HealthCheckRequest{}
+				if err := stream.RecvMsg(req); err != nil {
+					return status.Error(codes.InvalidArgument, err.Error())
+				}
+
+				resp, err := h.Check(ctx, req)
+				if err != nil {
+					return err
+				}
+
+				return stream.SendMsg(resp)
+			case "/grpc.health.v1.Health/Watch":
+				// a dumb kludge to not write own WatchServer
+				return healthpb.Health_ServiceDesc.Streams[0].Handler(h, stream)
+			default:
+				return next(srv, stream)
+			}
+		}
+	}
+}
+
 // Recoverer is a middleware that recovers from panics, logs the panic and returns a gRPC error if possible.
-func Recoverer() Middleware {
+func Recoverer(responseMessage string) Middleware {
 	return func(next grpc.StreamHandler) grpc.StreamHandler {
 		return func(srv any, stream grpc.ServerStream) (err error) {
 			defer func() {
@@ -88,7 +123,7 @@ func Recoverer() Middleware {
 						slog.Any("panic", rvr),
 						slogx.Error(err))
 
-					err = status.Error(codes.ResourceExhausted, "{groxy} panic")
+					err = status.Error(codes.ResourceExhausted, responseMessage)
 				}
 			}()
 			return next(srv, stream)

pkg/proxy/middleware/middleware_test.go

@@ -13,6 +13,10 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
+	"google.golang.org/grpc/health"
+	healthpb "google.golang.org/grpc/health/grpc_health_v1"
+	"github.com/Semior001/groxy/pkg/grpcx/grpctest"
+	"google.golang.org/grpc/credentials/insecure"
 )
 
 func TestAppInfo(t *testing.T) {
@@ -38,7 +42,7 @@ func TestAppInfo(t *testing.T) {
 func TestRecoverer(t *testing.T) {
 	bts := bytes.NewBuffer(nil)
 	slog.SetDefault(slog.New(slog.NewTextHandler(bts, &slog.HandlerOptions{})))
-	mw := Recoverer()(func(_ any, _ grpc.ServerStream) error { panic("test") })
+	mw := Recoverer("{groxy} panic")(func(_ any, _ grpc.ServerStream) error { panic("test") })
 	var err error
 	require.NotPanics(t, func() {
 		err = mw(nil, &mocks.ServerStreamMock{
@@ -75,3 +79,60 @@ func TestChain(t *testing.T) {
 	require.NoError(t, h(nil, nil))
 	assert.Equal(t, []string{"mw1", "mw2", "mw3"}, calls)
 }
+
+func TestHealth(t *testing.T) {
+	prepare := func() (*health.Server, healthpb.HealthClient) {
+		h := health.NewServer()
+		h.SetServingStatus("", healthpb.HealthCheckResponse_SERVING)
+
+		srv := grpc.NewServer(grpc.UnknownServiceHandler(Health(h)(func(_ any, _ grpc.ServerStream) error {
+			return status.Error(codes.Internal, "must not be called")
+		})))
+
+		addr := grpctest.StartServer(t, srv)
+
+		conn, err := grpc.Dial(addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
+		require.NoError(t, err)
+
+		cl := healthpb.NewHealthClient(conn)
+
+		return h, cl
+	}
+
+	t.Run("unary", func(t *testing.T) {
+		h, cl := prepare()
+
+		resp, err := cl.Check(context.Background(), &healthpb.HealthCheckRequest{})
+		require.NoError(t, err)
+
+		assert.Equal(t, healthpb.HealthCheckResponse_SERVING, resp.Status)
+
+		h.SetServingStatus("", healthpb.HealthCheckResponse_NOT_SERVING)
+
+		resp, err = cl.Check(context.Background(), &healthpb.HealthCheckRequest{})
+		require.NoError(t, err)
+
+		assert.Equal(t, healthpb.HealthCheckResponse_NOT_SERVING, resp.Status)
+	})
+
+	t.Run("watch", func(t *testing.T) {
+		h, cl := prepare()
+
+		stream, err := cl.Watch(context.Background(), &healthpb.HealthCheckRequest{})
+		require.NoError(t, err)
+		defer stream.CloseSend()
+
+		resp, err := stream.Recv()
+		require.NoError(t, err)
+
+		assert.Equal(t, healthpb.HealthCheckResponse_SERVING, resp.Status)
+
+		h.SetServingStatus("", healthpb.HealthCheckResponse_NOT_SERVING)
+
+		resp, err = stream.Recv()
+		require.NoError(t, err)
+
+		assert.Equal(t, healthpb.HealthCheckResponse_NOT_SERVING, resp.Status)
+	})
+
+}

pkg/proxy/proxy.go

@@ -17,6 +17,8 @@ import (
 	"google.golang.org/grpc/status"
 	"github.com/Semior001/groxy/pkg/grpcx"
 	"context"
+	"google.golang.org/grpc/health"
+	healthpb "google.golang.org/grpc/health/grpc_health_v1"
 )
 
 //go:generate moq -out mocks/mocks.go --skip-ensure -pkg mocks . Matcher ServerStream
@@ -69,12 +71,16 @@ func (s *Server) Listen(addr string) (err error) {
 	slog.Info("starting gRPC server", slog.Any("addr", addr))
 	defer slog.Warn("gRPC server stopped", slogx.Error(err))
 
+	healthHandler := health.NewServer()
+	healthHandler.SetServingStatus("", healthpb.HealthCheckResponse_SERVING)
+
 	s.grpc = grpc.NewServer(append(s.serverOpts,
 		grpc.UnknownServiceHandler(middleware.Wrap(s.handle,
-			middleware.Recoverer(),
+			middleware.Recoverer("{groxy} panic"),
 			middleware.Maybe(s.signature, middleware.AppInfo("groxy", "Semior001", s.version)),
 			middleware.Log(s.debug, "/grpc.reflection."),
 			middleware.PassMetadata(),
+			middleware.Health(healthHandler),
 			middleware.Maybe(s.reflection, middleware.Chain(
 				middleware.Reflector{
 					Logger:        slog.Default().With(slog.String("subsystem", "reflection")),