Define Health Server for GRPC servers (#3712)

jaegertracing · May 30, 2022 · 22e9143 · 22e9143
1 parent b708823
commit 22e9143
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 6 deletions.
diff --git a/cmd/collector/app/server/grpc.go b/cmd/collector/app/server/grpc.go
@@ -22,6 +22,8 @@ import (
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/health"
+	"google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/grpc/keepalive"
 	"google.golang.org/grpc/reflection"
 
@@ -89,9 +91,16 @@ func StartGRPCServer(params *GRPCServerParams) (*grpc.Server, error) {
 }
 
 func serveGRPC(server *grpc.Server, listener net.Listener, params *GRPCServerParams) error {
+	healthServer := health.NewServer()
+
 	api_v2.RegisterCollectorServiceServer(server, params.Handler)
 	api_v2.RegisterSamplingManagerServer(server, sampling.NewGRPCHandler(params.SamplingStore))
 
+	healthServer.SetServingStatus("jaeger.api_v2.CollectorService", grpc_health_v1.HealthCheckResponse_SERVING)
+	healthServer.SetServingStatus("jaeger.api_v2.SamplingManager", grpc_health_v1.HealthCheckResponse_SERVING)
+
+	grpc_health_v1.RegisterHealthServer(server, healthServer)
+
 	params.Logger.Info("Starting jaeger-collector gRPC server", zap.String("grpc.host-port", params.HostPortActual))
 	go func() {
 		if err := server.Serve(listener); err != nil {

diff --git a/cmd/collector/app/server/grpc_test.go b/cmd/collector/app/server/grpc_test.go
@@ -130,6 +130,7 @@ func TestCollectorReflection(t *testing.T) {
 		ExpectedServices: []string{
 			"jaeger.api_v2.CollectorService",
 			"jaeger.api_v2.SamplingManager",
+			"grpc.health.v1.Health",
 		},
 	}.Execute(t)
 }
diff --git a/cmd/query/app/server.go b/cmd/query/app/server.go
@@ -29,6 +29,8 @@ import (
 	"go.uber.org/zap/zapcore"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/health"
+	"google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/grpc/reflection"
 
 	"github.com/jaegertracing/jaeger/cmd/query/app/apiv3"
@@ -63,7 +65,6 @@ type Server struct {
 
 // NewServer creates and initializes Server
 func NewServer(logger *zap.Logger, querySvc *querysvc.QueryService, metricsQuerySvc querysvc.MetricsQueryService, options *QueryOptions, tracer opentracing.Tracer) (*Server, error) {
-
 	_, httpPort, err := net.SplitHostPort(options.HTTPHostPort)
 	if err != nil {
 		return nil, err
@@ -129,9 +130,17 @@ func createGRPCServer(querySvc *querysvc.QueryService, metricsQuerySvc querysvc.
 		tracer:              tracer,
 		nowFn:               time.Now,
 	}
+	healthServer := health.NewServer()
+
 	api_v2.RegisterQueryServiceServer(server, handler)
 	metrics.RegisterMetricsQueryServiceServer(server, handler)
 	api_v3.RegisterQueryServiceServer(server, &apiv3.Handler{QueryService: querySvc})
+
+	healthServer.SetServingStatus("jaeger.api_v2.QueryService", grpc_health_v1.HealthCheckResponse_SERVING)
+	healthServer.SetServingStatus("jaeger.api_v2.metrics.MetricsQueryService", grpc_health_v1.HealthCheckResponse_SERVING)
+	healthServer.SetServingStatus("jaeger.api_v3.QueryService", grpc_health_v1.HealthCheckResponse_SERVING)
+
+	grpc_health_v1.RegisterHealthServer(server, healthServer)
 	return server, nil
 }
 

diff --git a/cmd/query/app/server_test.go b/cmd/query/app/server_test.go
@@ -323,7 +323,8 @@ func TestServerHTTPTLS(t *testing.T) {
 				HTTPHostPort:           ports.GetAddressFromCLIOptions(ports.QueryHTTP, ""),
 				TLSHTTP:                test.TLS,
 				TLSGRPC:                TLSGRPC,
-				BearerTokenPropagation: true}
+				BearerTokenPropagation: true,
+			}
 			flagsSvc := flags.NewService(ports.QueryAdminHTTP)
 			flagsSvc.Logger = zap.NewNop()
 
@@ -418,7 +419,6 @@ func TestServerHTTPTLS(t *testing.T) {
 			server.Close()
 			wg.Wait()
 			assert.Equal(t, healthcheck.Unavailable, flagsSvc.HC().Get())
-
 		})
 	}
 }
@@ -483,7 +483,8 @@ func TestServerGRPCTLS(t *testing.T) {
 				HTTPHostPort:           ports.GetAddressFromCLIOptions(ports.QueryHTTP, ""),
 				TLSHTTP:                TLSHTTP,
 				TLSGRPC:                test.TLS,
-				BearerTokenPropagation: true}
+				BearerTokenPropagation: true,
+			}
 			flagsSvc := flags.NewService(ports.QueryAdminHTTP)
 			flagsSvc.Logger = zap.NewNop()
 
@@ -544,8 +545,8 @@ func TestServerGRPCTLS(t *testing.T) {
 			assert.Equal(t, healthcheck.Unavailable, flagsSvc.HC().Get())
 		})
 	}
-
 }
+
 func TestServerBadHostPort(t *testing.T) {
 	_, err := NewServer(zap.NewNop(), &querysvc.QueryService{}, nil,
 		&QueryOptions{HTTPHostPort: "8080", GRPCHostPort: "127.0.0.1:8081", BearerTokenPropagation: true},
@@ -632,7 +633,6 @@ func TestServerSinglePort(t *testing.T) {
 
 		}
 		wg.Done()
-
 	}()
 
 	client := newGRPCClient(t, hostPort)
@@ -710,6 +710,7 @@ func TestServerHandlesPortZero(t *testing.T) {
 			"jaeger.api_v2.QueryService",
 			"jaeger.api_v3.QueryService",
 			"jaeger.api_v2.metrics.MetricsQueryService",
+			"grpc.health.v1.Health",
 		},
 	}.Execute(t)
 }