Support custom HTTP health check endpoint paths (#2587)

* feature/custom HTTP health check endpoint

* code: improve comment to function

Co-authored-by: Johan Brandhorst-Satzkorn <johan.brandhorst@gmail.com>

Co-authored-by: Johan Brandhorst-Satzkorn <johan.brandhorst@gmail.com>

docs/docs/operations/health_check.md

@@ -53,4 +53,8 @@ to `/healthz` will result in the following:
 
 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.
+write your own logic to handle that in the health checking methods.
+
+Analogously, to register an `{/endpoint/path}` endpoint in your `ServeMux` with a user-defined endpoint path, you can use
+the `ServeMuxOption` `WithHealthEndpointAt`, which accepts a connection to your registered gRPC server
+together with a custom `endpointPath string` parameter.

runtime/mux.go

@@ -205,24 +205,24 @@ func WithDisablePathLengthFallback() ServeMuxOption {
 	}
 }
 
-// WithHealthzEndpoint returns a ServeMuxOption that will add a /healthz endpoint to the created ServeMux.
+// WithHealthEndpointAt returns a ServeMuxOption that will add an endpoint to the created ServeMux at the path specified by endpointPath.
 // 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 {
+func WithHealthEndpointAt(healthCheckClient grpc_health_v1.HealthClient, endpointPath string) 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,
+			http.MethodGet, endpointPath, 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,
+					Service: r.URL.Query().Get("service"),
 				})
 				if err != nil {
 					s.errorHandler(r.Context(), s, outboundMarshaler, w, r, err)
@@ -247,6 +247,13 @@ func WithHealthzEndpoint(healthCheckClient grpc_health_v1.HealthClient) ServeMux
 	}
 }
 
+// WithHealthzEndpoint returns a ServeMuxOption that will add a /healthz endpoint to the created ServeMux.
+//
+// See WithHealthEndpointAt for the general implementation.
+func WithHealthzEndpoint(healthCheckClient grpc_health_v1.HealthClient) ServeMuxOption {
+	return WithHealthEndpointAt(healthCheckClient, "/healthz")
+}
+
 // NewServeMux returns a new ServeMux whose internal mapping is empty.
 func NewServeMux(opts ...ServeMuxOption) *ServeMux {
 	serveMux := &ServeMux{

runtime/mux_test.go

@@ -663,6 +663,42 @@ func TestWithHealthzEndpoint_codes(t *testing.T) {
 	}
 }
 
+func TestWithHealthEndpointAt_consistentWithHealthz(t *testing.T) {
+	const endpointPath = "/healthz"
+
+	r := httptest.NewRequest(http.MethodGet, endpointPath, nil)
+
+	for _, tt := range healthCheckTests {
+		tt := tt
+
+		t.Run(tt.name, func(t *testing.T) {
+			client := &dummyHealthCheckClient{
+				status: tt.status,
+				code:   tt.code,
+			}
+
+			w := httptest.NewRecorder()
+
+			runtime.NewServeMux(
+				runtime.WithHealthEndpointAt(client, endpointPath),
+			).ServeHTTP(w, r)
+
+			refW := httptest.NewRecorder()
+
+			runtime.NewServeMux(
+				runtime.WithHealthzEndpoint(client),
+			).ServeHTTP(refW, r)
+
+			if w.Code != refW.Code {
+				t.Errorf(
+					"result http status code for grpc code %q and status %q should be equal to %d, but got %d",
+					tt.code, tt.status, refW.Code, w.Code,
+				)
+			}
+		})
+	}
+}
+
 func TestWithHealthzEndpoint_serviceParam(t *testing.T) {
 	service := "test"