diff --git a/zrpc/config.go b/zrpc/config.go index 2a9595a1..36dd0074 100644 --- a/zrpc/config.go +++ b/zrpc/config.go @@ -19,6 +19,8 @@ type ( // setting 0 means no timeout Timeout int64 `json:",default=2000"` CpuThreshold int64 `json:",default=900,range=[0:1000]"` + // grpc health check switch + Health bool `json:",default=true"` } // A RpcClientConf is a rpc client config. diff --git a/zrpc/internal/rpcserver.go b/zrpc/internal/rpcserver.go index 98ef126f..36408776 100644 --- a/zrpc/internal/rpcserver.go +++ b/zrpc/internal/rpcserver.go @@ -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 + } +} diff --git a/zrpc/internal/server.go b/zrpc/internal/server.go index 5a5d5655..ad405b34 100644 --- a/zrpc/internal/server.go +++ b/zrpc/internal/server.go @@ -35,9 +35,13 @@ type ( ) func newBaseRpcServer(address string, rpcServerOpts *rpcServerOptions) *baseRpcServer { + var h *health.Server + if rpcServerOpts.health { + h = health.NewServer() + } return &baseRpcServer{ address: address, - health: health.NewServer(), + health: h, metrics: rpcServerOpts.metrics, options: []grpc.ServerOption{grpc.KeepaliveParams(keepalive.ServerParameters{ MaxConnectionIdle: defaultConnectionIdleDuration, diff --git a/zrpc/server.go b/zrpc/server.go index c7386184..6e6b1cfc 100644 --- a/zrpc/server.go +++ b/zrpc/server.go @@ -40,6 +40,7 @@ func NewServer(c RpcServerConf, register internal.RegisterFn) (*RpcServer, error metrics := stat.NewMetrics(c.ListenOn) serverOptions := []internal.ServerOption{ internal.WithMetrics(metrics), + internal.WithRpcHealth(c.Health), } if c.HasEtcd() {