zeromicro/go-zero
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add middlewares config for zrpc (#2766)

Browse Source 
* feat: add middlewares config for zrpc

* chore: add tests

* chore: improve codecov

* chore: improve codecov
This commit is contained in:
Kevin Wan
2023-01-08 19:34:05 +08:00
committed by GitHub
parent ade6f9ee46
commit 26c541b9cb
13 changed files with 272 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
zrpc/internal/rpcserver.go
View File

@@ -26,12 +26,13 @@ type (
rpcServer struct {
*baseRpcServer
name string
middlewares ServerMiddlewaresConf
healthManager health.Probe
}
)
// NewRpcServer returns a Server.
func NewRpcServer(addr string, opts ...ServerOption) Server {
func NewRpcServer(addr string, middlewares ServerMiddlewaresConf, opts ...ServerOption) Server {
var options rpcServerOptions
for _, opt := range opts {
opt(&options)
@@ -42,6 +43,7 @@ func NewRpcServer(addr string, opts ...ServerOption) Server {
return &rpcServer{
baseRpcServer: newBaseRpcServer(addr, &options),
middlewares: middlewares,
healthManager: health.NewHealthManager(fmt.Sprintf("%s-%s", probeNamePrefix, addr)),
}
}
@@ -57,19 +59,9 @@ func (s *rpcServer) Start(register RegisterFn) error {
return err
}
unaryInterceptors := []grpc.UnaryServerInterceptor{
serverinterceptors.UnaryTracingInterceptor,
serverinterceptors.UnaryCrashInterceptor,
serverinterceptors.UnaryStatInterceptor(s.metrics),
serverinterceptors.UnaryPrometheusInterceptor,
serverinterceptors.UnaryBreakerInterceptor,
}
unaryInterceptors := s.buildUnaryInterceptors()
unaryInterceptors = append(unaryInterceptors, s.unaryInterceptors...)
streamInterceptors := []grpc.StreamServerInterceptor{
serverinterceptors.StreamTracingInterceptor,
serverinterceptors.StreamCrashInterceptor,
serverinterceptors.StreamBreakerInterceptor,
}
streamInterceptors := s.buildStreamInterceptors()
streamInterceptors = append(streamInterceptors, s.streamInterceptors...)
options := append(s.options, WithUnaryServerInterceptors(unaryInterceptors...),
WithStreamServerInterceptors(streamInterceptors...))
@@ -97,6 +89,44 @@ func (s *rpcServer) Start(register RegisterFn) error {
return server.Serve(lis)
}
func (s *rpcServer) buildStreamInterceptors() []grpc.StreamServerInterceptor {
var interceptors []grpc.StreamServerInterceptor
if s.middlewares.Trace {
interceptors = append(interceptors, serverinterceptors.StreamTracingInterceptor)
}
if s.middlewares.Recover {
interceptors = append(interceptors, serverinterceptors.StreamRecoverInterceptor)
}
if s.middlewares.Breaker {
interceptors = append(interceptors, serverinterceptors.StreamBreakerInterceptor)
}
return interceptors
}
func (s *rpcServer) buildUnaryInterceptors() []grpc.UnaryServerInterceptor {
var interceptors []grpc.UnaryServerInterceptor
if s.middlewares.Trace {
interceptors = append(interceptors, serverinterceptors.UnaryTracingInterceptor)
}
if s.middlewares.Recover {
interceptors = append(interceptors, serverinterceptors.UnaryRecoverInterceptor)
}
if s.middlewares.Stat {
interceptors = append(interceptors, serverinterceptors.UnaryStatInterceptor(s.metrics))
}
if s.middlewares.Prometheus {
interceptors = append(interceptors, serverinterceptors.UnaryPrometheusInterceptor)
}
if s.middlewares.Breaker {
interceptors = append(interceptors, serverinterceptors.UnaryBreakerInterceptor)
}
return interceptors
}
// WithMetrics returns a func that sets metrics to a Server.
func WithMetrics(metrics *stat.Metrics) ServerOption {
return func(options *rpcServerOptions) {