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

Add grpc retry (#1160)

Browse Source 
* Add grpc retry

* Update grpc retry

* Add tests

* Fix a bug

* Add api && some tests

* Add comment

* Add double check

* Add server retry quota

* Update optimize code

* Fix bug

* Update optimize code

* Update optimize code

* Fix bug
This commit is contained in:
chenquan
2021-10-27 19:46:07 +08:00
committed by GitHub
parent 496a2f341e
commit 462ddbb145
17 changed files with 544 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
zrpc/client.go
View File

@@ -14,6 +14,8 @@ var (
WithDialOption = internal.WithDialOption
// WithTimeout is an alias of internal.WithTimeout.
WithTimeout = internal.WithTimeout
// WithRetry is an alias of internal.WithRetry.
WithRetry = internal.WithRetry
// WithUnaryClientInterceptor is an alias of internal.WithUnaryClientInterceptor.
WithUnaryClientInterceptor = internal.WithUnaryClientInterceptor
)
@@ -52,6 +54,9 @@ func NewClient(c RpcClientConf, options ...ClientOption) (Client, error) {
if c.Timeout > 0 {
opts = append(opts, WithTimeout(time.Duration(c.Timeout)*time.Millisecond))
}
if c.Retry {
opts = append(opts, WithRetry())
}
opts = append(opts, options...)
var target string

2
zrpc/config.go
View File

@@ -18,6 +18,7 @@ type (
// setting 0 means no timeout
Timeout int64 `json:",default=2000"`
CpuThreshold int64 `json:",default=900,range=[0:1000]"`
MaxRetries int `json:",range=[0:]"`
}
// A RpcClientConf is a rpc client config.
@@ -27,6 +28,7 @@ type (
Target string `json:",optional"`
App string `json:",optional"`
Token string `json:",optional"`
Retry bool `json:",optional"` // grpc auto retry
Timeout int64 `json:",default=2000"`
}
)

9
zrpc/internal/client.go
View File

@@ -31,6 +31,7 @@ type (
// A ClientOptions is a client options.
ClientOptions struct {
Timeout time.Duration
Retry bool
DialOptions []grpc.DialOption
}
@@ -72,6 +73,7 @@ func (c *client) buildDialOptions(opts ...ClientOption) []grpc.DialOption {
clientinterceptors.PrometheusInterceptor,
clientinterceptors.BreakerInterceptor,
clientinterceptors.TimeoutInterceptor(cliOpts.Timeout),
clientinterceptors.RetryInterceptor(cliOpts.Retry),
),
WithStreamClientInterceptors(
clientinterceptors.StreamTracingInterceptor,
@@ -117,6 +119,13 @@ func WithTimeout(timeout time.Duration) ClientOption {
}
}
// WithRetry returns a func to customize a ClientOptions with auto retry.
func WithRetry() ClientOption {
return func(options *ClientOptions) {
options.Retry = true
}
}
// WithUnaryClientInterceptor returns a func to customize a ClientOptions with given interceptor.
func WithUnaryClientInterceptor(interceptor grpc.UnaryClientInterceptor) ClientOption {
return func(options *ClientOptions) {

19
zrpc/internal/clientinterceptors/retryinterceptor.go Normal file
View File

@@ -0,0 +1,19 @@
package clientinterceptors
import (
"context"
"github.com/tal-tech/go-zero/core/retry"
"google.golang.org/grpc"
)
// RetryInterceptor retry interceptor
func RetryInterceptor(enable bool) grpc.UnaryClientInterceptor {
return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
if !enable {
return invoker(ctx, method, req, reply, cc, opts...)
}
return retry.Do(ctx, func(ctx context.Context, callOpts ...grpc.CallOption) error {
return invoker(ctx, method, req, reply, cc, callOpts...)
}, opts...)
}
}

27
zrpc/internal/clientinterceptors/retryinterceptor_test.go Normal file
View File

@@ -0,0 +1,27 @@
package clientinterceptors
import (
"context"
"github.com/stretchr/testify/assert"
"github.com/tal-tech/go-zero/core/retry"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"testing"
)
func TestRetryInterceptor_WithMax(t *testing.T) {
n := 4
for i := 0; i < n; i++ {
count := 0
cc := new(grpc.ClientConn)
err := RetryInterceptor(true)(context.Background(), "/1", nil, nil, cc,
func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, opts ...grpc.CallOption) error {
count++
return status.Error(codes.ResourceExhausted, "ResourceExhausted")
}, retry.WithMax(i))
assert.Error(t, err)
assert.Equal(t, i+1, count)
}
}

13
zrpc/internal/rpcserver.go
View File

@@ -14,7 +14,8 @@ type (
ServerOption func(options *rpcServerOptions)
rpcServerOptions struct {
metrics *stat.Metrics
metrics *stat.Metrics
MaxRetries int
}
rpcServer struct {
@@ -38,7 +39,7 @@ func NewRpcServer(address string, opts ...ServerOption) Server {
}
return &rpcServer{
baseRpcServer: newBaseRpcServer(address, options.metrics),
baseRpcServer: newBaseRpcServer(address, &options),
}
}
@@ -55,6 +56,7 @@ func (s *rpcServer) Start(register RegisterFn) error {
unaryInterceptors := []grpc.UnaryServerInterceptor{
serverinterceptors.UnaryTracingInterceptor,
serverinterceptors.RetryInterceptor(s.maxRetries),
serverinterceptors.UnaryCrashInterceptor,
serverinterceptors.UnaryStatInterceptor(s.metrics),
serverinterceptors.UnaryPrometheusInterceptor,
@@ -87,3 +89,10 @@ func WithMetrics(metrics *stat.Metrics) ServerOption {
options.metrics = metrics
}
}
// WithMaxRetries returns a func that sets a max retries to a Server.
func WithMaxRetries(maxRetries int) ServerOption {
return func(options *rpcServerOptions) {
options.MaxRetries = maxRetries
}
}

8
zrpc/internal/server.go
View File

@@ -21,16 +21,18 @@ type (
baseRpcServer struct {
address string
metrics *stat.Metrics
maxRetries int
options []grpc.ServerOption
streamInterceptors []grpc.StreamServerInterceptor
unaryInterceptors []grpc.UnaryServerInterceptor
}
)
func newBaseRpcServer(address string, metrics *stat.Metrics) *baseRpcServer {
func newBaseRpcServer(address string, rpcServerOpts *rpcServerOptions) *baseRpcServer {
return &baseRpcServer{
address: address,
metrics: metrics,
address: address,
metrics: rpcServerOpts.metrics,
maxRetries: rpcServerOpts.MaxRetries,
}
}

6
zrpc/internal/server_test.go
View File

@@ -11,7 +11,7 @@ import (
func TestBaseRpcServer_AddOptions(t *testing.T) {
metrics := stat.NewMetrics("foo")
server := newBaseRpcServer("foo", metrics)
server := newBaseRpcServer("foo", &rpcServerOptions{metrics: metrics})
server.SetName("bar")
var opt grpc.EmptyServerOption
server.AddOptions(opt)
@@ -20,7 +20,7 @@ func TestBaseRpcServer_AddOptions(t *testing.T) {
func TestBaseRpcServer_AddStreamInterceptors(t *testing.T) {
metrics := stat.NewMetrics("foo")
server := newBaseRpcServer("foo", metrics)
server := newBaseRpcServer("foo", &rpcServerOptions{metrics: metrics})
server.SetName("bar")
var vals []int
f := func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
@@ -36,7 +36,7 @@ func TestBaseRpcServer_AddStreamInterceptors(t *testing.T) {
func TestBaseRpcServer_AddUnaryInterceptors(t *testing.T) {
metrics := stat.NewMetrics("foo")
server := newBaseRpcServer("foo", metrics)
server := newBaseRpcServer("foo", &rpcServerOptions{metrics: metrics})
server.SetName("bar")
var vals []int
f := func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (

33
zrpc/internal/serverinterceptors/retryinterceptor.go Normal file
View File

@@ -0,0 +1,33 @@
package serverinterceptors
import (
"context"
"github.com/tal-tech/go-zero/core/logx"
"github.com/tal-tech/go-zero/core/retry"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/status"
"strconv"
)
func RetryInterceptor(maxAttempt int) grpc.UnaryServerInterceptor {
return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
var md metadata.MD
requestMd, ok := metadata.FromIncomingContext(ctx)
if ok {
md = requestMd.Copy()
attemptMd := md.Get(retry.AttemptMetadataKey)
if len(attemptMd) != 0 && attemptMd[0] != "" {
if attempt, err := strconv.Atoi(attemptMd[0]); err == nil {
if attempt > maxAttempt {
logx.WithContext(ctx).Errorf("retries exceeded:%d, max retries:%d", attempt, maxAttempt)
return nil, status.Error(codes.FailedPrecondition, "Retries exceeded")
}
}
}
}
return handler(ctx, req)
}
}

40
zrpc/internal/serverinterceptors/retryinterceptor_test.go Normal file
View File

@@ -0,0 +1,40 @@
package serverinterceptors
import (
"context"
"github.com/stretchr/testify/assert"
"github.com/tal-tech/go-zero/core/retry"
"google.golang.org/grpc/metadata"
"testing"
)
func TestRetryInterceptor(t *testing.T) {
t.Run("retries exceeded", func(t *testing.T) {
interceptor := RetryInterceptor(2)
ctx := metadata.NewIncomingContext(context.Background(), metadata.New(map[string]string{retry.AttemptMetadataKey: "3"}))
resp, err := interceptor(ctx, nil, nil, func(ctx context.Context, req interface{}) (interface{}, error) {
return nil, nil
})
assert.Error(t, err)
assert.Nil(t, resp)
})
t.Run("reasonable retries", func(t *testing.T) {
interceptor := RetryInterceptor(2)
ctx := metadata.NewIncomingContext(context.Background(), metadata.New(map[string]string{retry.AttemptMetadataKey: "2"}))
resp, err := interceptor(ctx, nil, nil, func(ctx context.Context, req interface{}) (interface{}, error) {
return nil, nil
})
assert.NoError(t, err)
assert.Nil(t, resp)
})
t.Run("no retries", func(t *testing.T) {
interceptor := RetryInterceptor(0)
resp, err := interceptor(context.Background(), nil, nil, func(ctx context.Context, req interface{}) (interface{}, error) {
return nil, nil
})
assert.NoError(t, err)
assert.Nil(t, resp)
})
}

6
zrpc/server.go
View File

@@ -38,13 +38,15 @@ func NewServer(c RpcServerConf, register internal.RegisterFn) (*RpcServer, error
var server internal.Server
metrics := stat.NewMetrics(c.ListenOn)
serverOptions := []internal.ServerOption{internal.WithMetrics(metrics), internal.WithMaxRetries(c.MaxRetries)}
if c.HasEtcd() {
server, err = internal.NewRpcPubServer(c.Etcd.Hosts, c.Etcd.Key, c.ListenOn, internal.WithMetrics(metrics))
server, err = internal.NewRpcPubServer(c.Etcd.Hosts, c.Etcd.Key, c.ListenOn, serverOptions...)
if err != nil {
return nil, err
}
} else {
server = internal.NewRpcServer(c.ListenOn, internal.WithMetrics(metrics))
server = internal.NewRpcServer(c.ListenOn, serverOptions...)
}
server.SetName(c.Name)