|
|
|
@ -16,6 +16,7 @@ type (
|
|
|
|
|
|
|
|
|
|
rpcServerOptions struct {
|
|
|
|
|
metrics *stat.Metrics
|
|
|
|
|
health bool
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
rpcServer struct {
|
|
|
|
@ -74,13 +75,17 @@ func (s *rpcServer) Start(register RegisterFn) error {
|
|
|
|
|
register(server)
|
|
|
|
|
|
|
|
|
|
// register the health check service
|
|
|
|
|
grpc_health_v1.RegisterHealthServer(server, s.health)
|
|
|
|
|
s.health.Resume()
|
|
|
|
|
if s.health != nil {
|
|
|
|
|
grpc_health_v1.RegisterHealthServer(server, s.health)
|
|
|
|
|
s.health.Resume()
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// we need to make sure all others are wrapped up,
|
|
|
|
|
// so we do graceful stop at shutdown phase instead of wrap up phase
|
|
|
|
|
waitForCalled := proc.AddWrapUpListener(func() {
|
|
|
|
|
s.health.Shutdown()
|
|
|
|
|
if s.health != nil {
|
|
|
|
|
s.health.Shutdown()
|
|
|
|
|
}
|
|
|
|
|
server.GracefulStop()
|
|
|
|
|
})
|
|
|
|
|
defer waitForCalled()
|
|
|
|
@ -94,3 +99,10 @@ func WithMetrics(metrics *stat.Metrics) ServerOption {
|
|
|
|
|
options.metrics = metrics
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// WithRpcHealth returns a func that sets rpc health switch to a Server.
|
|
|
|
|
func WithRpcHealth(health bool) ServerOption {
|
|
|
|
|
return func(options *rpcServerOptions) {
|
|
|
|
|
options.health = health
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|