Repository: milvus-io/milvus-sdk-go
Pull request #522 (2023-07-21 03:29:50)

Make retry on ratelimit configurable
See also milvus-io/milvus#25693
/kind improvement

sre-ci-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/522#" title="Author self-approved">congqixia</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=milvus-io%2Fmilvus-sdk-go).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/milvus-io/milvus-sdk-go/blob/v2.2.x/OWNERS)~~ [congqixia]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/522?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) Report
> Merging [#522](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/522?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) (1deb6eb) into [v2.2.x](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/commit/6a10d5cbee8b03c9189db617ae1c7a7fdb04c0de?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) (6a10d5c) will **increase** coverage by `0.03%`.
> The diff coverage is `57.14%`.

[![Impacted file tree graph](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/522/graphs/tree.svg?width=650&height=150&src=pr&token=WJK43PV390&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io)](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/522?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io)

```diff
@@            Coverage Diff             @@
##           v2.2.x     #522      +/-   ##
==========================================
+ Coverage   79.99%   80.02%   +0.03%     
==========================================
  Files          35       36       +1     
  Lines        2844     2849       +5     
==========================================
+ Hits         2275     2280       +5     
+ Misses        458      457       -1     
- Partials      111      112       +1     
```


| [Impacted Files](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/522?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) | Coverage Δ | |
|---|---|---|
| [client/client.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/522?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L2NsaWVudC5nbw==) | `28.94% <ø> (-2.76%)` | :arrow_down: |
| [client/rate\_limit\_interceptor.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/522?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L3JhdGVfbGltaXRfaW50ZXJjZXB0b3IuZ28=) | `69.23% <33.33%> (ø)` | |
| [client/config.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/522?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L2NvbmZpZy5nbw==) | `70.37% <62.50%> (+3.02%)` | :arrow_up: |
| [client/results.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/522?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L3Jlc3VsdHMuZ28=) | `66.66% <66.66%> (ø)` | |

sre-ci-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/522#" title="Author self-approved">congqixia</a>*, *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/522#pullrequestreview-1540423881" title="LGTM">yah01</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=milvus-io%2Fmilvus-sdk-go).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/milvus-io/milvus-sdk-go/blob/v2.2.x/OWNERS)~~ [congqixia]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->

Patch:
diff --git a/client/client.go b/client/client.go
index 2cdfe30..9a74dce 100644
--- a/client/client.go
+++ b/client/client.go
@@ -204,31 +204,6 @@ type Client interface {
 	GetVersion(ctx context.Context) (string, error)
 }
 
-// SearchResult contains the result from Search api of client
-// IDs is the auto generated id values for the entities
-// Fields contains the data of `outputFieleds` specified or all columns if non
-// Scores is actually the distance between the vector current record contains and the search target vector
-type SearchResult struct {
-	ResultCount int           // the returning entry count
-	IDs         entity.Column // auto generated id, can be mapped to the columns from `Insert` API
-	Fields      ResultSet     // output field data
-	Scores      []float32     // distance to the target vector
-	Err         error         // search error if any
-}
-
-// ResultSet is an alias type for column slice.
-type ResultSet []entity.Column
-
-// GetColumn returns column with provided field name.
-func (rs ResultSet) GetColumn(fieldName string) entity.Column {
-	for _, column := range rs {
-		if column.Name() == fieldName {
-			return column
-		}
-	}
-	return nil
-}
-
 // Check if GrpcClient implement Client.
 var _ Client = &GrpcClient{}
 
diff --git a/client/config.go b/client/config.go
index 4386d67..241cc18 100644
--- a/client/config.go
+++ b/client/config.go
@@ -63,11 +63,18 @@ type Config struct {
 
 	parsedAddress *url.URL
 
+	RetryRateLimit *RetryRateLimitOption // option for retry on rate limit inteceptor
+
 	DisableConn bool
 
 	flags uint64 // internal flags
 }
 
+type RetryRateLimitOption struct {
+	MaxRetry   uint
+	MaxBackoff time.Duration
+}
+
 // Copy a new config, dialOption may shared with old config.
 func (c *Config) Copy() Config {
 	newConfig := Config{
@@ -147,9 +154,7 @@ func (c *Config) getDialOption() []grpc.DialOption {
 				return 60 * time.Millisecond * time.Duration(math.Pow(3, float64(attempt)))
 			}),
 			grpc_retry.WithCodes(codes.Unavailable, codes.ResourceExhausted)),
-			RetryOnRateLimitInterceptor(75, func(ctx context.Context, attempt uint) time.Duration {
-				return 10 * time.Millisecond * time.Duration(math.Pow(3, float64(attempt)))
-			}),
+			c.getRetryOnRateLimitInterceptor(),
 		))
 
 	options = append(options, grpc.WithChainUnaryInterceptor(
@@ -158,6 +163,23 @@ func (c *Config) getDialOption() []grpc.DialOption {
 	return options
 }
 
+func (c *Config) getRetryOnRateLimitInterceptor() grpc.UnaryClientInterceptor {
+	if c.RetryRateLimit == nil {
+		c.RetryRateLimit = c.defaultRetryRateLimitOption()
+	}
+
+	return RetryOnRateLimitInterceptor(c.RetryRateLimit.MaxRetry, c.RetryRateLimit.MaxBackoff, func(ctx context.Context, attempt uint) time.Duration {
+		return 10 * time.Millisecond * time.Duration(math.Pow(3, float64(attempt)))
+	})
+}
+
+func (c *Config) defaultRetryRateLimitOption() *RetryRateLimitOption {
+	return &RetryRateLimitOption{
+		MaxRetry:   75,
+		MaxBackoff: 3 * time.Second,
+	}
+}
+
 // addFlags set internal flags
 func (c *Config) addFlags(flags uint64) {
 	c.flags |= flags
diff --git a/client/rate_limit_interceptor.go b/client/rate_limit_interceptor.go
index c8d112e..ff728be 100644
--- a/client/rate_limit_interceptor.go
+++ b/client/rate_limit_interceptor.go
@@ -37,24 +37,21 @@ const (
 	RetryOnRateLimit ctxKey = iota
 )
 
-var MaxBackOff = 3 * time.Second
-
 // RetryOnRateLimitInterceptor returns a new retrying unary client interceptor.
-func RetryOnRateLimitInterceptor(maxRetry uint, backoffFunc grpc_retry.BackoffFuncContext) grpc.UnaryClientInterceptor {
+func RetryOnRateLimitInterceptor(maxRetry uint, maxBackoff time.Duration, 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)
+			_, err := waitRetryBackoff(parentCtx, attempt, maxBackoff, 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
@@ -102,14 +99,14 @@ func contextErrToGrpcErr(err error) error {
 	}
 }
 
-func waitRetryBackoff(parentCtx context.Context, attempt uint, backoffFunc grpc_retry.BackoffFuncContext) (time.Duration, error) {
+func waitRetryBackoff(parentCtx context.Context, attempt uint, maxBackoff time.Duration, 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
+		if waitTime > maxBackoff {
+			waitTime = maxBackoff
 		}
 		timer := time.NewTimer(waitTime)
 		select {
diff --git a/client/results.go b/client/results.go
new file mode 100644
index 0000000..2fde53c
--- /dev/null
+++ b/client/results.go
@@ -0,0 +1,28 @@
+package client
+
+import "github.com/milvus-io/milvus-sdk-go/v2/entity"
+
+// SearchResult contains the result from Search api of client
+// IDs is the auto generated id values for the entities
+// Fields contains the data of `outputFieleds` specified or all columns if non
+// Scores is actually the distance between the vector current record contains and the search target vector
+type SearchResult struct {
+	ResultCount int           // the returning entry count
+	IDs         entity.Column // auto generated id, can be mapped to the columns from `Insert` API
+	Fields      ResultSet     // output field data
+	Scores      []float32     // distance to the target vector
+	Err         error         // search error if any
+}
+
+// ResultSet is an alias type for column slice.
+type ResultSet []entity.Column
+
+// GetColumn returns column with provided field name.
+func (rs ResultSet) GetColumn(fieldName string) entity.Column {
+	for _, column := range rs {
+		if column.Name() == fieldName {
+			return column
+		}
+	}
+	return nil
+}


Test patch:
diff --git a/client/rate_limit_interceptor_test.go b/client/rate_limit_interceptor_test.go
index af5a949..784ee86 100644
--- a/client/rate_limit_interceptor_test.go
+++ b/client/rate_limit_interceptor_test.go
@@ -43,7 +43,8 @@ func resetMockInvokeTimes() {
 
 func TestRateLimitInterceptor(t *testing.T) {
 	maxRetry := uint(3)
-	inter := RetryOnRateLimitInterceptor(maxRetry, func(ctx context.Context, attempt uint) time.Duration {
+	maxBackoff := 3 * time.Second
+	inter := RetryOnRateLimitInterceptor(maxRetry, maxBackoff, func(ctx context.Context, attempt uint) time.Duration {
 		return 60 * time.Millisecond * time.Duration(math.Pow(2, float64(attempt)))
 	})