github.com/lingyao2333/mo-zero@v1.4.1/zrpc/internal/rpcserver.go (about)

     1  package internal
     2  
     3  import (
     4  	"net"
     5  
     6  	"github.com/lingyao2333/mo-zero/core/proc"
     7  	"github.com/lingyao2333/mo-zero/core/stat"
     8  	"github.com/lingyao2333/mo-zero/zrpc/internal/serverinterceptors"
     9  	"google.golang.org/grpc"
    10  	"google.golang.org/grpc/health/grpc_health_v1"
    11  )
    12  
    13  type (
    14  	// ServerOption defines the method to customize a rpcServerOptions.
    15  	ServerOption func(options *rpcServerOptions)
    16  
    17  	rpcServerOptions struct {
    18  		metrics *stat.Metrics
    19  		health  bool
    20  	}
    21  
    22  	rpcServer struct {
    23  		name string
    24  		*baseRpcServer
    25  	}
    26  )
    27  
    28  // NewRpcServer returns a Server.
    29  func NewRpcServer(address string, opts ...ServerOption) Server {
    30  	var options rpcServerOptions
    31  	for _, opt := range opts {
    32  		opt(&options)
    33  	}
    34  	if options.metrics == nil {
    35  		options.metrics = stat.NewMetrics(address)
    36  	}
    37  
    38  	return &rpcServer{
    39  		baseRpcServer: newBaseRpcServer(address, &options),
    40  	}
    41  }
    42  
    43  func (s *rpcServer) SetName(name string) {
    44  	s.name = name
    45  	s.baseRpcServer.SetName(name)
    46  }
    47  
    48  func (s *rpcServer) Start(register RegisterFn) error {
    49  	lis, err := net.Listen("tcp", s.address)
    50  	if err != nil {
    51  		return err
    52  	}
    53  
    54  	unaryInterceptors := []grpc.UnaryServerInterceptor{
    55  		serverinterceptors.UnaryTracingInterceptor,
    56  		serverinterceptors.UnaryCrashInterceptor,
    57  		serverinterceptors.UnaryStatInterceptor(s.metrics),
    58  		serverinterceptors.UnaryPrometheusInterceptor,
    59  		serverinterceptors.UnaryBreakerInterceptor,
    60  	}
    61  	unaryInterceptors = append(unaryInterceptors, s.unaryInterceptors...)
    62  	streamInterceptors := []grpc.StreamServerInterceptor{
    63  		serverinterceptors.StreamTracingInterceptor,
    64  		serverinterceptors.StreamCrashInterceptor,
    65  		serverinterceptors.StreamBreakerInterceptor,
    66  	}
    67  	streamInterceptors = append(streamInterceptors, s.streamInterceptors...)
    68  	options := append(s.options, WithUnaryServerInterceptors(unaryInterceptors...),
    69  		WithStreamServerInterceptors(streamInterceptors...))
    70  	server := grpc.NewServer(options...)
    71  	register(server)
    72  
    73  	// register the health check service
    74  	if s.health != nil {
    75  		grpc_health_v1.RegisterHealthServer(server, s.health)
    76  		s.health.Resume()
    77  	}
    78  
    79  	// we need to make sure all others are wrapped up,
    80  	// so we do graceful stop at shutdown phase instead of wrap up phase
    81  	waitForCalled := proc.AddWrapUpListener(func() {
    82  		if s.health != nil {
    83  			s.health.Shutdown()
    84  		}
    85  		server.GracefulStop()
    86  	})
    87  	defer waitForCalled()
    88  
    89  	return server.Serve(lis)
    90  }
    91  
    92  // WithMetrics returns a func that sets metrics to a Server.
    93  func WithMetrics(metrics *stat.Metrics) ServerOption {
    94  	return func(options *rpcServerOptions) {
    95  		options.metrics = metrics
    96  	}
    97  }
    98  
    99  // WithRpcHealth returns a func that sets rpc health switch to a Server.
   100  func WithRpcHealth(health bool) ServerOption {
   101  	return func(options *rpcServerOptions) {
   102  		options.health = health
   103  	}
   104  }