Add grpc retry (#1160)

* 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

zrpc/internal/serverinterceptors/retryinterceptor.go

@@ -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)
+	}
+}

zrpc/internal/serverinterceptors/retryinterceptor_test.go

@@ -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)
+	})
+
+}