Repository: milvus-io/milvus-sdk-go
Pull request #675 (2024-03-18 07:32:16)

fix: Support retry on milvus errors and retry on ErrInconsistentRequery
1. Support retry on milvus errors.
2. Retry on ErrInconsistentRequery for search and hybrid search operations.

issue: https://github.com/milvus-io/milvus/issues/29656

related milvus pr: https://github.com/milvus-io/milvus/pull/31343



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

This pull-request has been approved by: *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/675#" title="Author self-approved">bigsheeper</a>*
To complete the [pull request process](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process), please assign **sunby** after the PR has been reviewed.
You can assign the PR to them by writing `/assign @sunby` in a comment when ready.

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).

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

- **[OWNERS](https://github.com/milvus-io/milvus-sdk-go/blob/master/OWNERS)**

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":["sunby"]} -->
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/675?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) Report
Attention: Patch coverage is `75.80645%` with `15 lines` in your changes are missing coverage. Please review.
> Project coverage is 71.87%. Comparing base [(`1a1a092`)](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/commit/1a1a0924c28098573302a68474c79e1d97d01fe0?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) to head [(`ecbd6ad`)](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/675?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io).

<details><summary>Additional details and impacted files</summary>


[![Impacted file tree graph](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/675/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/675?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io)

```diff
@@            Coverage Diff             @@
##           master     #675      +/-   ##
==========================================
- Coverage   71.88%   71.87%   -0.01%     
==========================================
  Files          49       50       +1     
  Lines        4564     4623      +59     
==========================================
+ Hits         3281     3323      +42     
- Misses       1169     1184      +15     
- Partials      114      116       +2     
```


| [Files](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/675?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) | Coverage Δ | |
|---|---|---|
| [merr/errors.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/675?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-bWVyci9lcnJvcnMuZ28=) | `55.55% <ø> (ø)` | |
| [client/data.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/675?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L2RhdGEuZ28=) | `64.40% <40.00%> (-0.38%)` | :arrow_down: |
| [client/util.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/675?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L3V0aWwuZ28=) | `82.69% <82.69%> (ø)` | |

... and [1 file with indirect coverage changes](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/675/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io)

</details>
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/675#" title="Author self-approved">bigsheeper</a>*, *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/675#pullrequestreview-1942567119" title="LGTM">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/master/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/data.go b/client/data.go
index 7192681..ad322bd 100644
--- a/client/data.go
+++ b/client/data.go
@@ -79,12 +79,16 @@ func (c *GrpcClient) HybridSearch(ctx context.Context, collName string, partitio
 		RankParams:       params,
 	}
 
-	result, err := c.Service.HybridSearch(ctx, req)
-
-	err = merr.CheckRPCCall(result, err)
+	r, err := RetryOnMilvusErrors(ctx, func() (interface{}, error) {
+		return c.Service.HybridSearch(ctx, req)
+	}, OnMerrCodes(merr.Code(merr.ErrInconsistentRequery)))
 	if err != nil {
 		return nil, err
 	}
+	result := r.(*milvuspb.SearchResults)
+	if err = handleRespStatus(result.GetStatus()); err != nil {
+		return nil, err
+	}
 
 	return c.handleSearchResult(schema, outputFields, nq, result)
 }
@@ -117,10 +121,13 @@ func (c *GrpcClient) Search(ctx context.Context, collName string, partitions []s
 		return nil, err
 	}
 
-	resp, err := c.Service.Search(ctx, req)
+	r, err := RetryOnMilvusErrors(ctx, func() (interface{}, error) {
+		return c.Service.Search(ctx, req)
+	}, OnMerrCodes(merr.Code(merr.ErrInconsistentRequery)))
 	if err != nil {
 		return nil, err
 	}
+	resp := r.(*milvuspb.SearchResults)
 	if err := handleRespStatus(resp.GetStatus()); err != nil {
 		return nil, err
 	}
diff --git a/client/util.go b/client/util.go
new file mode 100644
index 0000000..c70e204
--- /dev/null
+++ b/client/util.go
@@ -0,0 +1,141 @@
+// 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"
+
+	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
+)
+
+type config struct {
+	attempts     uint
+	sleep        time.Duration
+	maxSleepTime time.Duration
+	merrCodes    []int32
+}
+
+func newDefaultConfig() *config {
+	return &config{
+		attempts:     uint(10),
+		sleep:        200 * time.Millisecond,
+		maxSleepTime: 3 * time.Second,
+	}
+}
+
+// Option is used to config the retry function.
+type Option func(*config)
+
+// Attempts is used to config the max retry times.
+func Attempts(attempts uint) Option {
+	return func(c *config) {
+		c.attempts = attempts
+	}
+}
+
+// Sleep is used to config the initial interval time of each execution.
+func Sleep(sleep time.Duration) Option {
+	return func(c *config) {
+		c.sleep = sleep
+		// ensure max retry interval is always larger than retry interval
+		if c.sleep*2 > c.maxSleepTime {
+			c.maxSleepTime = 2 * c.sleep
+		}
+	}
+}
+
+// MaxSleepTime is used to config the max interval time of each execution.
+func MaxSleepTime(maxSleepTime time.Duration) Option {
+	return func(c *config) {
+		// ensure max retry interval is always larger than retry interval
+		if c.sleep*2 > maxSleepTime {
+			c.maxSleepTime = 2 * c.sleep
+		} else {
+			c.maxSleepTime = maxSleepTime
+		}
+	}
+}
+
+func OnMerrCodes(codes ...int32) Option {
+	return func(c *config) {
+		c.merrCodes = append(c.merrCodes, codes...)
+	}
+}
+
+func contains(codes []int32, target int32) bool {
+	for _, c := range codes {
+		if c == target {
+			return true
+		}
+	}
+	return false
+}
+
+func RetryOnMilvusErrors(ctx context.Context, fn func() (interface{}, error), opts ...Option) (interface{}, error) {
+	if ctx.Err() != nil {
+		return nil, ctx.Err()
+	}
+
+	c := newDefaultConfig()
+	for _, opt := range opts {
+		opt(c)
+	}
+
+	if len(c.merrCodes) == 0 {
+		return fn()
+	}
+
+	var lastResp interface{}
+	for i := uint(0); i < c.attempts; i++ {
+		resp, err := fn()
+		if err != nil {
+			return resp, err
+		}
+		var code int32
+		switch r := resp.(type) {
+		case interface{ GetStatus() *commonpb.Status }:
+			code = r.GetStatus().GetCode()
+		case interface{ GetCode() int32 }:
+			code = r.GetCode()
+		default:
+			return resp, nil
+		}
+		if code == 0 || !contains(c.merrCodes, code) {
+			return resp, nil
+		}
+
+		deadline, ok := ctx.Deadline()
+		if ok && time.Until(deadline) < c.sleep {
+			return resp, context.DeadlineExceeded
+		}
+
+		lastResp = resp
+
+		select {
+		case <-time.After(c.sleep):
+		case <-ctx.Done():
+			return lastResp, ctx.Err()
+		}
+
+		c.sleep *= 2
+		if c.sleep > c.maxSleepTime {
+			c.sleep = c.maxSleepTime
+		}
+	}
+	return lastResp, nil
+}
diff --git a/merr/errors.go b/merr/errors.go
index e7c86bf..c045029 100644
--- a/merr/errors.go
+++ b/merr/errors.go
@@ -144,6 +144,9 @@ var (
 
 	// import
 	ErrImportFailed = newMilvusError("importing data failed", 2100, false)
+
+	// Search/Query related
+	ErrInconsistentRequery = newMilvusError("inconsistent requery result", 2200, true)
 )
 
 type milvusError struct {


Test patch:
diff --git a/client/util_test.go b/client/util_test.go
new file mode 100644
index 0000000..92d4be3
--- /dev/null
+++ b/client/util_test.go
@@ -0,0 +1,180 @@
+// 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"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
+)
+
+func TestRetryOnMilvusErrors(t *testing.T) {
+	ctx := context.Background()
+
+	n := 0
+	testFn := func() (interface{}, error) {
+		if n < 3 {
+			n++
+			return &commonpb.Status{
+				Code: 0,
+			}, nil
+		}
+		return &commonpb.Status{
+			Reason: "mock err",
+			Code:   100,
+		}, nil
+	}
+
+	resp, err := RetryOnMilvusErrors(ctx, testFn, OnMerrCodes(100))
+	assert.Equal(t, int32(0), resp.(interface{ GetCode() int32 }).GetCode())
+	assert.NoError(t, err)
+	t.Log(resp)
+}
+
+func TestOnNoCode(t *testing.T) {
+	ctx := context.Background()
+
+	n := 0
+	testFn := func() (interface{}, error) {
+		if n < 3 {
+			n++
+			return &commonpb.Status{
+				Code: 0,
+			}, nil
+		}
+		return &commonpb.Status{
+			Reason: "mock err",
+			Code:   100,
+		}, nil
+	}
+
+	resp, err := RetryOnMilvusErrors(ctx, testFn)
+	assert.Equal(t, int32(0), resp.(interface{ GetCode() int32 }).GetCode())
+	assert.NoError(t, err)
+	t.Log(resp)
+}
+
+func TestReturnErr(t *testing.T) {
+	ctx := context.Background()
+
+	testFn := func() (interface{}, error) {
+		return nil, errors.New("mock err")
+	}
+
+	_, err := RetryOnMilvusErrors(ctx, testFn)
+	assert.Error(t, err)
+	t.Log(err)
+}
+
+func TestAttempts(t *testing.T) {
+	ctx := context.Background()
+
+	testFn := func() (interface{}, error) {
+		t.Log("executed")
+		return &commonpb.Status{
+			Reason: "mock err",
+			Code:   100,
+		}, nil
+	}
+
+	resp, err := RetryOnMilvusErrors(ctx, testFn, OnMerrCodes(100), Attempts(1))
+	assert.Equal(t, int32(100), resp.(interface{ GetCode() int32 }).GetCode())
+	assert.NoError(t, err)
+	t.Log(resp)
+}
+
+func TestMaxSleepTime(t *testing.T) {
+	ctx := context.Background()
+
+	testFn := func() (interface{}, error) {
+		t.Log("executed")
+		return &commonpb.Status{
+			Reason: "mock err",
+			Code:   100,
+		}, nil
+	}
+
+	resp, err := RetryOnMilvusErrors(ctx, testFn, OnMerrCodes(100), Attempts(3), MaxSleepTime(200*time.Millisecond))
+	assert.Equal(t, int32(100), resp.(interface{ GetCode() int32 }).GetCode())
+	assert.NoError(t, err)
+	t.Log(resp)
+}
+
+func TestSleep(t *testing.T) {
+	ctx := context.Background()
+
+	testFn := func() (interface{}, error) {
+		t.Log("executed")
+		return &commonpb.Status{
+			Reason: "mock err",
+			Code:   100,
+		}, nil
+	}
+
+	resp, err := RetryOnMilvusErrors(ctx, testFn, OnMerrCodes(100), Attempts(3), Sleep(500*time.Millisecond))
+	assert.Equal(t, int32(100), resp.(interface{ GetCode() int32 }).GetCode())
+	assert.NoError(t, err)
+	t.Log(resp)
+}
+
+func TestContextDeadline(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel()
+
+	testFn := func() (interface{}, error) {
+		t.Log("executed")
+		return &commonpb.Status{
+			Reason: "mock err",
+			Code:   100,
+		}, nil
+	}
+
+	resp, err := RetryOnMilvusErrors(ctx, testFn, OnMerrCodes(100))
+	assert.Equal(t, int32(100), resp.(interface{ GetCode() int32 }).GetCode())
+	assert.Error(t, err)
+	assert.ErrorIs(t, err, context.DeadlineExceeded)
+	t.Log(resp)
+}
+
+func TestContextCancel(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	testFn := func() (interface{}, error) {
+		t.Log("executed")
+		return &commonpb.Status{
+			Reason: "mock err",
+			Code:   100,
+		}, nil
+	}
+
+	go func() {
+		time.Sleep(100 * time.Millisecond)
+		cancel()
+	}()
+
+	resp, err := RetryOnMilvusErrors(ctx, testFn, OnMerrCodes(100))
+	assert.Equal(t, int32(100), resp.(interface{ GetCode() int32 }).GetCode())
+	assert.Error(t, err)
+	assert.ErrorIs(t, err, context.Canceled)
+	t.Log(resp)
+	t.Log(err)
+}
diff --git a/test/testcases/search_test.go b/test/testcases/search_test.go
index bbab1bb..1d5f034 100644
--- a/test/testcases/search_test.go
+++ b/test/testcases/search_test.go
@@ -1433,6 +1433,40 @@ func TestRangeSearchScannBinary(t *testing.T) {
 	}
 }
 
+func TestVectorOutputField(t *testing.T) {
+	ctx := createContext(t, time.Second*common.DefaultTimeout)
+	// connect
+	mc := createMilvusClient(ctx, t)
+
+	// create collection with data
+	collName, _ := createCollectionWithDataIndex(ctx, t, mc, false, true)
+
+	// load collection
+	errLoad := mc.LoadCollection(ctx, collName, false)
+	common.CheckErr(t, errLoad, true)
+
+	// search vector
+	for i := 0; i < 20; i++ {
+		sp, _ := entity.NewIndexHNSWSearchParam(74)
+		searchResult, errSearch := mc.Search(
+			ctx, collName,
+			[]string{common.DefaultPartition},
+			"",
+			[]string{common.DefaultFloatVecFieldName},
+			//[]entity.Vector{entity.FloatVector([]float32{0.1, 0.2})},
+			common.GenSearchVectors(common.DefaultNq, common.DefaultDim, entity.FieldTypeFloatVector),
+			common.DefaultFloatVecFieldName,
+			entity.L2,
+			common.DefaultTopK,
+			sp,
+		)
+		common.CheckErr(t, errSearch, true)
+		common.CheckOutputFields(t, searchResult[0].Fields, []string{common.DefaultFloatVecFieldName})
+		common.CheckSearchResult(t, searchResult, common.DefaultNq, common.DefaultTopK)
+		log.Printf("search %d done\n", i)
+	}
+}
+
 // TODO offset and limit
 // TODO consistency level
 // TODO WithGuaranteeTimestamp