feat: rpc add health check function configuration optional (#2288)

* feat: rpc add health check function configuration optional

* update config field name
This commit is contained in:
MarkJoyMa
2022-08-23 13:44:21 +08:00
committed by GitHub
parent 1c85d39add
commit 040c9e0954
4 changed files with 23 additions and 4 deletions

View File

@@ -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
}
}