Explorar o código

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

* feat: rpc add health check function configuration optional

* update config field name
MarkJoyMa %!s(int64=2) %!d(string=hai) anos
pai
achega
040c9e0954
Modificáronse 4 ficheiros con 23 adicións e 4 borrados
  1. 2 0
      zrpc/config.go
  2. 15 3
      zrpc/internal/rpcserver.go
  3. 5 1
      zrpc/internal/server.go
  4. 1 0
      zrpc/server.go

+ 2 - 0
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.

+ 15 - 3
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
+	}
+}

+ 5 - 1
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,

+ 1 - 0
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() {