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:
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user