Add WithHealthzEndpoint as ServeMuxOption to register a `/healthz…

…`endpoint (#2319) * Add HealthCheck endpoint option to ServeMux * Use outboundMarshaler for HealthCheck reponses * Rename WithHealthCheckEnabled to WithHealthzEndpoint * Add docs for WithHealthzEndpoint * Cleanup health_check.md and imports * Fix health_check.md header * Fix health_check.md header whitespace * Replace panic with error return in test grpc server Co-authored-by: Johan Brandhorst-Satzkorn <[email protected]> * Update runtime/BUILD.bazel * Use HealthClient for WithHealthzEndpoint option Signed-off-by: Tobias Brumhard <[email protected]> * Regenerate bazel files Signed-off-by: Tobias Brumhard <[email protected]> Co-authored-by: Johan Brandhorst-Satzkorn <[email protected]>
grpc-ecosystem · Mar 2, 2022 · db9c664 · db9c664
1 parent 8ca093b
commit db9c664
Show file tree

Hide file tree

Showing 4 changed files with 167 additions and 0 deletions.
diff --git a/docs/docs/operations/health_check.md b/docs/docs/operations/health_check.md
@@ -36,3 +36,21 @@ func (s *serviceServer) Watch(in *health.HealthCheckRequest, _ health.Health_Wat
 ```
 
 3. You can test the functionality with [GRPC health probe](https://github.com/grpc-ecosystem/grpc-health-probe).
+
+## Adding `/healthz` endpoint to runtime.ServeMux
+
+To automatically register a `/healthz` endpoint in your `ServeMux` you can use
+the `ServeMuxOption` `WithHealthzEndpoint`
+which takes in a connection to your registered gRPC server.
+
+This endpoint will forward a request to the `Check` method described above to really check the health of the
+whole system, not only the gateway itself. If your server doesn't implement the health checking protocol each request
+to `/healthz` will result in the following:
+
+```json
+{"code":12,"message":"unknown service grpc.health.v1.Health","details":[]}
+```
+
+If you've implemented multiple services in your server you can target specific services with the `?service=<service>`
+query parameter. This will then be added to the `health.HealthCheckRequest` in the `Service` property. With that you can
+write your own logic to handle that in the health checking methods.
diff --git a/runtime/BUILD.bazel b/runtime/BUILD.bazel
@@ -30,6 +30,7 @@ go_library(
         "@io_bazel_rules_go//proto/wkt:field_mask_go_proto",
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//grpclog",
+        "@org_golang_google_grpc//health/grpc_health_v1",
         "@org_golang_google_grpc//metadata",
         "@org_golang_google_grpc//status",
         "@org_golang_google_protobuf//encoding/protojson",
@@ -71,7 +72,9 @@ go_test(
         "@go_googleapis//google/rpc:errdetails_go_proto",
         "@go_googleapis//google/rpc:status_go_proto",
         "@io_bazel_rules_go//proto/wkt:field_mask_go_proto",
+        "@org_golang_google_grpc//:go_default_library",
         "@org_golang_google_grpc//codes",
+        "@org_golang_google_grpc//health/grpc_health_v1",
         "@org_golang_google_grpc//metadata",
         "@org_golang_google_grpc//status",
         "@org_golang_google_protobuf//encoding/protojson",

diff --git a/runtime/mux.go b/runtime/mux.go
@@ -12,6 +12,7 @@ import (
 	"github.com/grpc-ecosystem/grpc-gateway/v2/internal/httprule"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/grpclog"
+	"google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"
@@ -204,6 +205,48 @@ func WithDisablePathLengthFallback() ServeMuxOption {
 	}
 }
 
+// WithHealthzEndpoint returns a ServeMuxOption that will add a /healthz endpoint to the created ServeMux.
+// When called the handler will forward the request to the upstream grpc service health check (defined in the
+// gRPC Health Checking Protocol).
+// See here https://grpc-ecosystem.github.io/grpc-gateway/docs/operations/health_check/ for more information on how
+// to setup the protocol in the grpc server.
+// If you define a service as query parameter, this will also be forwarded as service in the HealthCheckRequest.
+func WithHealthzEndpoint(healthCheckClient grpc_health_v1.HealthClient) ServeMuxOption {
+	return func(s *ServeMux) {
+		// error can be ignored since pattern is definitely valid
+		_ = s.HandlePath(
+			http.MethodGet, "/healthz", func(w http.ResponseWriter, r *http.Request, _ map[string]string,
+			) {
+				_, outboundMarshaler := MarshalerForRequest(s, r)
+
+				serviceQueryParam := r.URL.Query().Get("service")
+
+				resp, err := healthCheckClient.Check(r.Context(), &grpc_health_v1.HealthCheckRequest{
+					Service: serviceQueryParam,
+				})
+				if err != nil {
+					s.errorHandler(r.Context(), s, outboundMarshaler, w, r, err)
+					return
+				}
+
+				if resp.GetStatus() != grpc_health_v1.HealthCheckResponse_SERVING {
+					var err error
+					switch resp.GetStatus() {
+					case grpc_health_v1.HealthCheckResponse_NOT_SERVING, grpc_health_v1.HealthCheckResponse_UNKNOWN:
+						err = status.Error(codes.Unavailable, resp.String())
+					case grpc_health_v1.HealthCheckResponse_SERVICE_UNKNOWN:
+						err = status.Error(codes.NotFound, resp.String())
+					}
+
+					s.errorHandler(r.Context(), s, outboundMarshaler, w, r, err)
+					return
+				}
+
+				_ = outboundMarshaler.NewEncoder(w).Encode(resp)
+			})
+	}
+}
+
 // NewServeMux returns a new ServeMux whose internal mapping is empty.
 func NewServeMux(opts ...ServeMuxOption) *ServeMux {
 	serveMux := &ServeMux{

diff --git a/runtime/mux_test.go b/runtime/mux_test.go
@@ -2,15 +2,21 @@ package runtime_test
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
 	"strconv"
+	"strings"
 	"testing"
 
 	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
 	"github.com/grpc-ecosystem/grpc-gateway/v2/utilities"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/health/grpc_health_v1"
+	"google.golang.org/grpc/status"
 )
 
 func TestMuxServeHTTP(t *testing.T) {
@@ -598,3 +604,100 @@ func TestServeMux_HandlePath(t *testing.T) {
 		})
 	}
 }
+
+var healthCheckTests = []struct {
+	name           string
+	code           codes.Code
+	status         grpc_health_v1.HealthCheckResponse_ServingStatus
+	httpStatusCode int
+}{
+	{
+		"Test grpc error code",
+		codes.NotFound,
+		grpc_health_v1.HealthCheckResponse_UNKNOWN,
+		http.StatusNotFound,
+	},
+	{
+		"Test HealthCheckResponse_SERVING",
+		codes.OK,
+		grpc_health_v1.HealthCheckResponse_SERVING,
+		http.StatusOK,
+	},
+	{
+		"Test HealthCheckResponse_NOT_SERVING",
+		codes.OK,
+		grpc_health_v1.HealthCheckResponse_NOT_SERVING,
+		http.StatusServiceUnavailable,
+	},
+	{
+		"Test HealthCheckResponse_UNKNOWN",
+		codes.OK,
+		grpc_health_v1.HealthCheckResponse_UNKNOWN,
+		http.StatusServiceUnavailable,
+	},
+	{
+		"Test HealthCheckResponse_SERVICE_UNKNOWN",
+		codes.OK,
+		grpc_health_v1.HealthCheckResponse_SERVICE_UNKNOWN,
+		http.StatusNotFound,
+	},
+}
+
+func TestWithHealthzEndpoint_codes(t *testing.T) {
+	for _, tt := range healthCheckTests {
+		t.Run(tt.name, func(t *testing.T) {
+			mux := runtime.NewServeMux(runtime.WithHealthzEndpoint(&dummyHealthCheckClient{status: tt.status, code: tt.code}))
+
+			r := httptest.NewRequest(http.MethodGet, "/healthz", nil)
+			rr := httptest.NewRecorder()
+
+			mux.ServeHTTP(rr, r)
+
+			if rr.Code != tt.httpStatusCode {
+				t.Errorf(
+					"result http status code for grpc code %q and status %q should be %d, got %d",
+					tt.code, tt.status, tt.httpStatusCode, rr.Code,
+				)
+			}
+		})
+	}
+}
+
+func TestWithHealthzEndpoint_serviceParam(t *testing.T) {
+	service := "test"
+
+	// trigger error to output service in body
+	dummyClient := dummyHealthCheckClient{status: grpc_health_v1.HealthCheckResponse_UNKNOWN, code: codes.Unknown}
+	mux := runtime.NewServeMux(runtime.WithHealthzEndpoint(&dummyClient))
+
+	r := httptest.NewRequest(http.MethodGet, "/healthz?service="+service, nil)
+	rr := httptest.NewRecorder()
+
+	mux.ServeHTTP(rr, r)
+
+	if !strings.Contains(rr.Body.String(), service) {
+		t.Errorf(
+			"service query parameter should be translated to HealthCheckRequest: expected %s to contain %s",
+			rr.Body.String(), service,
+		)
+	}
+}
+
+var _ grpc_health_v1.HealthClient = (*dummyHealthCheckClient)(nil)
+
+type dummyHealthCheckClient struct {
+	status grpc_health_v1.HealthCheckResponse_ServingStatus
+	code   codes.Code
+}
+
+func (g *dummyHealthCheckClient) Check(ctx context.Context, r *grpc_health_v1.HealthCheckRequest, opts ...grpc.CallOption) (*grpc_health_v1.HealthCheckResponse, error) {
+	if g.code != codes.OK {
+		return nil, status.Error(g.code, r.GetService())
+	}
+
+	return &grpc_health_v1.HealthCheckResponse{Status: g.status}, nil
+}
+
+func (g *dummyHealthCheckClient) Watch(ctx context.Context, r *grpc_health_v1.HealthCheckRequest, opts ...grpc.CallOption) (grpc_health_v1.Health_WatchClient, error) {
+	return nil, status.Error(codes.Unimplemented, "unimplemented")
+}