github.com/shuguocloud/go-zero@v1.3.0/zrpc/internal/rpcserver.go (about)

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