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

     1  package zrpc
     2  
     3  import (
     4  	"log"
     5  	"time"
     6  
     7  	"github.com/lingyao2333/mo-zero/core/load"
     8  	"github.com/lingyao2333/mo-zero/core/logx"
     9  	"github.com/lingyao2333/mo-zero/core/stat"
    10  	"github.com/lingyao2333/mo-zero/zrpc/internal"
    11  	"github.com/lingyao2333/mo-zero/zrpc/internal/auth"
    12  	"github.com/lingyao2333/mo-zero/zrpc/internal/serverinterceptors"
    13  	"google.golang.org/grpc"
    14  )
    15  
    16  // A RpcServer is a rpc server.
    17  type RpcServer struct {
    18  	server   internal.Server
    19  	register internal.RegisterFn
    20  }
    21  
    22  // MustNewServer returns a RpcSever, exits on any error.
    23  func MustNewServer(c RpcServerConf, register internal.RegisterFn) *RpcServer {
    24  	server, err := NewServer(c, register)
    25  	if err != nil {
    26  		log.Fatal(err)
    27  	}
    28  
    29  	return server
    30  }
    31  
    32  // NewServer returns a RpcServer.
    33  func NewServer(c RpcServerConf, register internal.RegisterFn) (*RpcServer, error) {
    34  	var err error
    35  	if err = c.Validate(); err != nil {
    36  		return nil, err
    37  	}
    38  
    39  	var server internal.Server
    40  	metrics := stat.NewMetrics(c.ListenOn)
    41  	serverOptions := []internal.ServerOption{
    42  		internal.WithMetrics(metrics),
    43  		internal.WithRpcHealth(c.Health),
    44  	}
    45  
    46  	if c.HasEtcd() {
    47  		server, err = internal.NewRpcPubServer(c.Etcd, c.ListenOn, serverOptions...)
    48  		if err != nil {
    49  			return nil, err
    50  		}
    51  	} else {
    52  		server = internal.NewRpcServer(c.ListenOn, serverOptions...)
    53  	}
    54  
    55  	server.SetName(c.Name)
    56  	if err = setupInterceptors(server, c, metrics); err != nil {
    57  		return nil, err
    58  	}
    59  
    60  	rpcServer := &RpcServer{
    61  		server:   server,
    62  		register: register,
    63  	}
    64  	if err = c.SetUp(); err != nil {
    65  		return nil, err
    66  	}
    67  
    68  	return rpcServer, nil
    69  }
    70  
    71  // AddOptions adds given options.
    72  func (rs *RpcServer) AddOptions(options ...grpc.ServerOption) {
    73  	rs.server.AddOptions(options...)
    74  }
    75  
    76  // AddStreamInterceptors adds given stream interceptors.
    77  func (rs *RpcServer) AddStreamInterceptors(interceptors ...grpc.StreamServerInterceptor) {
    78  	rs.server.AddStreamInterceptors(interceptors...)
    79  }
    80  
    81  // AddUnaryInterceptors adds given unary interceptors.
    82  func (rs *RpcServer) AddUnaryInterceptors(interceptors ...grpc.UnaryServerInterceptor) {
    83  	rs.server.AddUnaryInterceptors(interceptors...)
    84  }
    85  
    86  // Start starts the RpcServer.
    87  // Graceful shutdown is enabled by default.
    88  // Use proc.SetTimeToForceQuit to customize the graceful shutdown period.
    89  func (rs *RpcServer) Start() {
    90  	if err := rs.server.Start(rs.register); err != nil {
    91  		logx.Error(err)
    92  		panic(err)
    93  	}
    94  }
    95  
    96  // Stop stops the RpcServer.
    97  func (rs *RpcServer) Stop() {
    98  	logx.Close()
    99  }
   100  
   101  // DontLogContentForMethod disable logging content for given method.
   102  func DontLogContentForMethod(method string) {
   103  	serverinterceptors.DontLogContentForMethod(method)
   104  }
   105  
   106  // SetServerSlowThreshold sets the slow threshold on server side.
   107  func SetServerSlowThreshold(threshold time.Duration) {
   108  	serverinterceptors.SetSlowThreshold(threshold)
   109  }
   110  
   111  func setupInterceptors(server internal.Server, c RpcServerConf, metrics *stat.Metrics) error {
   112  	if c.CpuThreshold > 0 {
   113  		shedder := load.NewAdaptiveShedder(load.WithCpuThreshold(c.CpuThreshold))
   114  		server.AddUnaryInterceptors(serverinterceptors.UnarySheddingInterceptor(shedder, metrics))
   115  	}
   116  
   117  	if c.Timeout > 0 {
   118  		server.AddUnaryInterceptors(serverinterceptors.UnaryTimeoutInterceptor(
   119  			time.Duration(c.Timeout) * time.Millisecond))
   120  	}
   121  
   122  	if c.Auth {
   123  		authenticator, err := auth.NewAuthenticator(c.Redis.NewRedis(), c.Redis.Key, c.StrictControl)
   124  		if err != nil {
   125  			return err
   126  		}
   127  
   128  		server.AddStreamInterceptors(serverinterceptors.StreamAuthorizeInterceptor(authenticator))
   129  		server.AddUnaryInterceptors(serverinterceptors.UnaryAuthorizeInterceptor(authenticator))
   130  	}
   131  
   132  	return nil
   133  }