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

change grpc interceptor to chain interceptor (#200)

Browse Source 
* change grpc interceptor to chain interceptor

* change server rpc interceptors, del testing code
This commit is contained in:
SunJun
2020-11-11 17:15:22 +08:00
committed by GitHub
parent a5fb29a6f0
commit 0a2c2d1eca
4 changed files with 8 additions and 346 deletions
Download Patch File Download Diff File Expand all files Collapse all files

74
zrpc/internal/chainserverinterceptors.go
View File

@@ -1,81 +1,13 @@
package internal
import (
"context"
"google.golang.org/grpc"
)
func WithStreamServerInterceptors(interceptors ...grpc.StreamServerInterceptor) grpc.ServerOption {
return grpc.StreamInterceptor(chainStreamServerInterceptors(interceptors...))
return grpc.ChainStreamInterceptor(interceptors...)
}
func WithUnaryServerInterceptors(interceptors ...grpc.UnaryServerInterceptor) grpc.ServerOption {
return grpc.UnaryInterceptor(chainUnaryServerInterceptors(interceptors...))
}
func chainStreamServerInterceptors(interceptors ...grpc.StreamServerInterceptor) grpc.StreamServerInterceptor {
switch len(interceptors) {
case 0:
return func(srv interface{}, stream grpc.ServerStream, _ *grpc.StreamServerInfo,
handler grpc.StreamHandler) error {
return handler(srv, stream)
}
case 1:
return interceptors[0]
default:
last := len(interceptors) - 1
return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo,
handler grpc.StreamHandler) error {
var chainHandler grpc.StreamHandler
var current int
chainHandler = func(curSrv interface{}, curStream grpc.ServerStream) error {
if current == last {
return handler(curSrv, curStream)
}
current++
err := interceptors[current](curSrv, curStream, info, chainHandler)
current--
return err
}
return interceptors[0](srv, stream, info, chainHandler)
}
}
}
func chainUnaryServerInterceptors(interceptors ...grpc.UnaryServerInterceptor) grpc.UnaryServerInterceptor {
switch len(interceptors) {
case 0:
return func(ctx context.Context, req interface{}, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (
interface{}, error) {
return handler(ctx, req)
}
case 1:
return interceptors[0]
default:
last := len(interceptors) - 1
return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (
interface{}, error) {
var chainHandler grpc.UnaryHandler
var current int
chainHandler = func(curCtx context.Context, curReq interface{}) (interface{}, error) {
if current == last {
return handler(curCtx, curReq)
}
current++
resp, err := interceptors[current](curCtx, curReq, info, chainHandler)
current--
return resp, err
}
return interceptors[0](ctx, req, info, chainHandler)
}
}
}
return grpc.ChainUnaryInterceptor(interceptors...)
}