Retry on rate limit (#320)

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>

client/client.go

@@ -18,16 +18,16 @@ import (
 	"math"
 	"time"
 
+	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
 	grpc_retry "github.com/grpc-ecosystem/go-grpc-middleware/retry"
+	"google.golang.org/grpc"
 	"google.golang.org/grpc/backoff"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/keepalive"
 
 	"github.com/milvus-io/milvus-sdk-go/v2/entity"
-
-	"google.golang.org/grpc"
 )
 
 // Client is the interface used to communicate with Milvus
@@ -197,12 +197,17 @@ func NewDefaultGrpcClient(ctx context.Context, addr string) (Client, error) {
 	defaultOpts := append(DefaultGrpcOpts,
 		grpc.WithTransportCredentials(insecure.NewCredentials()),
 		grpc.WithUnaryInterceptor(
-			grpc_retry.UnaryClientInterceptor(
-				grpc_retry.WithMax(6),
-				grpc_retry.WithBackoff(func(attempt uint) time.Duration {
-					return 60 * time.Millisecond * time.Duration(math.Pow(3, float64(attempt)))
+			grpc_middleware.ChainUnaryClient(
+				grpc_retry.UnaryClientInterceptor(
+					grpc_retry.WithMax(6),
+					grpc_retry.WithBackoff(func(attempt uint) time.Duration {
+						return 60 * time.Millisecond * time.Duration(math.Pow(3, float64(attempt)))
+					}),
+					grpc_retry.WithCodes(codes.Unavailable, codes.ResourceExhausted)),
+				RetryOnRateLimitInterceptor(10, func(ctx context.Context, attempt uint) time.Duration {
+					return 10 * time.Millisecond * time.Duration(math.Pow(3, float64(attempt)))
 				}),
-				grpc_retry.WithCodes(codes.Unavailable, codes.ResourceExhausted)),
+			),
 		),
 	)
 	err := c.connect(ctx, addr, defaultOpts...)

client/rate_limit_interceptor.go

@@ -0,0 +1,123 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package client
+
+import (
+	"context"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	grpc_retry "github.com/grpc-ecosystem/go-grpc-middleware/retry"
+	"github.com/milvus-io/milvus-sdk-go/v2/internal/proto/common"
+	"github.com/milvus-io/milvus-sdk-go/v2/internal/proto/server"
+)
+
+// ref: https://github.com/grpc-ecosystem/go-grpc-middleware
+
+type ctxKey int
+
+const (
+	RetryOnRateLimit ctxKey = iota
+)
+
+var MaxBackOff = 60 * time.Second
+
+// RetryOnRateLimitInterceptor returns a new retrying unary client interceptor.
+func RetryOnRateLimitInterceptor(maxRetry uint, backoffFunc grpc_retry.BackoffFuncContext) grpc.UnaryClientInterceptor {
+	return func(parentCtx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+		if maxRetry == 0 {
+			return invoker(parentCtx, method, req, reply, cc, opts...)
+		}
+		var lastErr error
+		for attempt := uint(0); attempt < maxRetry; attempt++ {
+			_, err := waitRetryBackoff(parentCtx, attempt, backoffFunc)
+			if err != nil {
+				return err
+			}
+			lastErr = invoker(parentCtx, method, req, reply, cc, opts...)
+			rspStatus := getResultStatus(reply)
+			if retryOnRateLimit(parentCtx) && rspStatus.GetErrorCode() == common.ErrorCode_RateLimit {
+				//log.Printf("rate limit retry attempt: %d, backoff for %v, reson: %v\n", attempt, backoff, rspStatus.GetReason())
+				continue
+			}
+			return lastErr
+		}
+		return lastErr
+	}
+}
+
+func retryOnRateLimit(ctx context.Context) bool {
+	retry, ok := ctx.Value(RetryOnRateLimit).(bool)
+	if !ok {
+		return true // default true
+	}
+	return retry
+}
+
+// getResultStatus returns status of response.
+func getResultStatus(reply interface{}) *common.Status {
+	switch r := reply.(type) {
+	case *common.Status:
+		return r
+	case *server.MutationResult:
+		return r.GetStatus()
+	case *server.BoolResponse:
+		return r.GetStatus()
+	case *server.SearchResults:
+		return r.GetStatus()
+	case *server.QueryResults:
+		return r.GetStatus()
+	case *server.FlushResponse:
+		return r.GetStatus()
+	default:
+		return nil
+	}
+}
+
+func contextErrToGrpcErr(err error) error {
+	switch err {
+	case context.DeadlineExceeded:
+		return status.Error(codes.DeadlineExceeded, err.Error())
+	case context.Canceled:
+		return status.Error(codes.Canceled, err.Error())
+	default:
+		return status.Error(codes.Unknown, err.Error())
+	}
+}
+
+func waitRetryBackoff(parentCtx context.Context, attempt uint, backoffFunc grpc_retry.BackoffFuncContext) (time.Duration, error) {
+	var waitTime time.Duration
+	if attempt > 0 {
+		waitTime = backoffFunc(parentCtx, attempt)
+	}
+	if waitTime > 0 {
+		if waitTime > MaxBackOff {
+			waitTime = MaxBackOff
+		}
+		timer := time.NewTimer(waitTime)
+		select {
+		case <-parentCtx.Done():
+			timer.Stop()
+			return waitTime, contextErrToGrpcErr(parentCtx.Err())
+		case <-timer.C:
+		}
+	}
+	return waitTime, nil
+}

client/rate_limit_interceptor_test.go

@@ -0,0 +1,65 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package client
+
+import (
+	"context"
+	"math"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"google.golang.org/grpc"
+
+	"github.com/milvus-io/milvus-sdk-go/v2/internal/proto/common"
+)
+
+var mockInvokerError error
+var mockInvokerReply interface{}
+var mockInvokeTimes = 0
+
+var mockInvoker grpc.UnaryInvoker = func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, opts ...grpc.CallOption) error {
+	mockInvokeTimes++
+	return mockInvokerError
+}
+
+func resetMockInvokeTimes() {
+	mockInvokeTimes = 0
+}
+
+func TestRateLimitInterceptor(t *testing.T) {
+	maxRetry := uint(3)
+	inter := RetryOnRateLimitInterceptor(maxRetry, func(ctx context.Context, attempt uint) time.Duration {
+		return 60 * time.Millisecond * time.Duration(math.Pow(2, float64(attempt)))
+	})
+
+	ctx := context.Background()
+
+	// with retry
+	mockInvokerReply = &common.Status{ErrorCode: common.ErrorCode_RateLimit}
+	resetMockInvokeTimes()
+	err := inter(ctx, "", nil, mockInvokerReply, nil, mockInvoker)
+	assert.NoError(t, err)
+	assert.Equal(t, maxRetry, uint(mockInvokeTimes))
+
+	// without retry
+	ctx1 := context.WithValue(ctx, RetryOnRateLimit, false)
+	resetMockInvokeTimes()
+	err = inter(ctx1, "", nil, mockInvokerReply, nil, mockInvoker)
+	assert.NoError(t, err)
+	assert.Equal(t, uint(1), uint(mockInvokeTimes))
+}