Repository: milvus-io/milvus-sdk-go
Pull request #194 (2021-09-02 04:59:53)

Respect context deadline when doing sync wait
Respect deadline in sync wait

Fix: #192 

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

codecov[bot]: # [Codecov](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/194?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) Report
> :exclamation: No coverage uploaded for pull request base (`master@2c6699e`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#section-missing-base-commit).
> The diff coverage is `n/a`.

[![Impacted file tree graph](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/194/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://codecov.io/gh/milvus-io/milvus-sdk-go/pull/194?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     #194   +/-   ##
=========================================
  Coverage          ?   88.89%           
=========================================
  Files             ?       16           
  Lines             ?     2252           
  Branches          ?        0           
=========================================
  Hits              ?     2002           
  Misses            ?      165           
  Partials          ?       85           
```



sunby: /lgtm
/approve
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/194#" title="Author self-approved">congqixia</a>*, *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/194#issuecomment-911196849" title="Approved">sunby</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,sunby]

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_grpc_collection.go b/client/client_grpc_collection.go
index 9e18225..dfc87c7 100644
--- a/client/client_grpc_collection.go
+++ b/client/client_grpc_collection.go
@@ -261,8 +261,8 @@ func (c *grpcClient) HasCollection(ctx context.Context, collName string) (bool,
 	if err != nil {
 		return false, err
 	}
-	if resp.Status.ErrorCode != common.ErrorCode_Success {
-		return false, errors.New("request failed")
+	if err := handleRespStatus(resp.GetStatus()); err != nil {
+		return false, err
 	}
 	return resp.GetValue(), nil
 }
@@ -343,6 +343,12 @@ func (c *grpcClient) LoadCollection(ctx context.Context, collName string, async
 
 	if !async {
 		for {
+			select {
+			case <-ctx.Done():
+				return errors.New("context deadline exceeded")
+			default:
+			}
+
 			coll, err := c.ShowCollection(ctx, collName)
 			if err != nil {
 				return err
@@ -350,7 +356,8 @@ func (c *grpcClient) LoadCollection(ctx context.Context, collName string, async
 			if coll.Loaded {
 				break
 			}
-			time.Sleep(100 * time.Millisecond) // TODO change to configuration
+
+			time.Sleep(200 * time.Millisecond) // TODO change to configuration
 		}
 	}
 	return nil
diff --git a/client/client_grpc_data.go b/client/client_grpc_data.go
index 4426fb2..6532463 100644
--- a/client/client_grpc_data.go
+++ b/client/client_grpc_data.go
@@ -139,7 +139,6 @@ func (c *grpcClient) Flush(ctx context.Context, collName string, async bool) err
 	if err := handleRespStatus(resp.GetStatus()); err != nil {
 		return err
 	}
-	//TODO silverxia, add context done respect logic while waiting
 	if !async {
 		segmentIDs, has := resp.GetCollSegIDs()[collName]
 		if has {
@@ -165,7 +164,13 @@ func (c *grpcClient) Flush(ctx context.Context, collName string, async bool) err
 				return len(waitingSet) == flushed
 			}
 			for !flushed() {
-				time.Sleep(500 * time.Millisecond)
+				// respect context deadline/cancel
+				select {
+				case <-ctx.Done():
+					return errors.New("deadline exceeded")
+				default:
+				}
+				time.Sleep(200 * time.Millisecond)
 			}
 		}
 	}
diff --git a/client/client_grpc_partition.go b/client/client_grpc_partition.go
index f408420..4309abc 100644
--- a/client/client_grpc_partition.go
+++ b/client/client_grpc_partition.go
@@ -173,6 +173,11 @@ func (c *grpcClient) LoadPartitions(ctx context.Context, collName string, partit
 
 	if !async {
 		for {
+			select {
+			case <-ctx.Done():
+				return errors.New("context deadline exceeded")
+			default:
+			}
 			partitions, err := c.ShowPartitions(ctx, collName)
 			if err != nil {
 				return err


Test patch:
diff --git a/client/client_grpc_collection_test.go b/client/client_grpc_collection_test.go
index 0937e44..678a0dc 100644
--- a/client/client_grpc_collection_test.go
+++ b/client/client_grpc_collection_test.go
@@ -279,6 +279,7 @@ func TestGrpcClientDropCollection(t *testing.T) {
 func TestGrpcClientLoadCollection(t *testing.T) {
 	ctx := context.Background()
 	c := testClient(ctx, t)
+	mock.setInjection(mHasCollection, hasCollectionDefault)
 	// injection check collection name equals
 	mock.setInjection(mLoadCollection, func(_ context.Context, raw proto.Message) (proto.Message, error) {
 		req, ok := raw.(*server.LoadCollectionRequest)
@@ -293,47 +294,9 @@ func TestGrpcClientLoadCollection(t *testing.T) {
 	})
 	t.Run("Load collection sync", func(t *testing.T) {
 
-		loadTime := rand.Intn(900) + 100 // in milli seconds, 100~1000 milliseconds
+		loadTime := rand.Intn(500) + 500 // in milli seconds, 100~1000 milliseconds
 		passed := false                  //### flag variable
 		start := time.Now()
-		/*
-			mock.setInjection(mGetPersistentSegmentInfo, func(_ context.Context, raw proto.Message) (proto.Message, error) {
-				s, err := successStatus()
-				r := &server.GetPersistentSegmentInfoResponse{
-					Status: s,
-					Infos:  make([]*server.PersistentSegmentInfo, 0, segmentCount),
-				}
-				for i := 0; i < segmentCount; i++ {
-					r.Infos = append(r.Infos, &server.PersistentSegmentInfo{
-						SegmentID: int64(i),
-						NumRows:   int64(rowCounts),
-					})
-				}
-				r.Infos = append(r.Infos, &server.PersistentSegmentInfo{
-					SegmentID: int64(segmentCount),
-					NumRows:   0, // handcrafted empty segment
-				})
-				return r, err
-			})
-			mock.setInjection(mGetQuerySegmentInfo, func(_ context.Context, raw proto.Message) (proto.Message, error) {
-				s, err := successStatus()
-				r := &server.GetQuerySegmentInfoResponse{
-					Status: s,
-					Infos:  make([]*server.QuerySegmentInfo, 0, segmentCount),
-				}
-				rc := 0
-				if time.Since(start) > time.Duration(loadTime)*time.Millisecond {
-					rc = rowCounts // after load time, row counts set to full amount
-					ok = true
-				}
-				for i := 0; i < segmentCount; i++ {
-					r.Infos = append(r.Infos, &server.QuerySegmentInfo{
-						SegmentID: int64(i),
-						NumRows:   int64(rc),
-					})
-				}
-				return r, err
-			})*/
 
 		mock.setInjection(mShowCollections, func(_ context.Context, raw proto.Message) (proto.Message, error) {
 			req, ok := raw.(*server.ShowCollectionsRequest)
@@ -348,6 +311,7 @@ func TestGrpcClientLoadCollection(t *testing.T) {
 			r.CollectionIds = []int64{1}
 			var perc int64 = 0
 			if time.Since(start) > time.Duration(loadTime)*time.Millisecond {
+				t.Log("passed")
 				perc = 100
 				passed = true
 			}
@@ -357,11 +321,13 @@ func TestGrpcClientLoadCollection(t *testing.T) {
 		assert.Nil(t, c.LoadCollection(ctx, testCollectionName, false))
 		assert.True(t, passed)
 
+		start = time.Now()
+		passed = false
+		quickCtx, cancel := context.WithTimeout(ctx, 50*time.Millisecond)
+		defer cancel()
+		assert.NotNil(t, c.LoadCollection(quickCtx, testCollectionName, false))
+
 		// remove injection
-		/*
-			mock.delInjection(mGetPersistentSegmentInfo)
-			mock.delInjection(mGetQuerySegmentInfo)
-		*/
 		mock.delInjection(mShowCollections)
 	})
 }
diff --git a/client/client_grpc_data_test.go b/client/client_grpc_data_test.go
index ea8b16f..96a35c2 100644
--- a/client/client_grpc_data_test.go
+++ b/client/client_grpc_data_test.go
@@ -120,8 +120,8 @@ func TestGrpcClientFlush(t *testing.T) {
 		for i := 0; i < segCount; i++ {
 			segments = append(segments, rand.Int63())
 		}
-		// 500ms ~ 2s
-		flushTime := 500 + rand.Intn(1500)
+		// 510ms ~ 2s
+		flushTime := 510 + rand.Intn(1500)
 		start := time.Now()
 		flag := false
 		mock.setInjection(mFlush, func(_ context.Context, raw proto.Message) (proto.Message, error) {
@@ -172,6 +172,12 @@ func TestGrpcClientFlush(t *testing.T) {
 		})
 		assert.Nil(t, c.Flush(ctx, testCollectionName, false))
 		assert.True(t, flag)
+
+		start = time.Now()
+		flag = false
+		quickCtx, cancel := context.WithTimeout(ctx, 10*time.Millisecond)
+		defer cancel()
+		assert.NotNil(t, c.Flush(quickCtx, testCollectionName, false))
 	})
 }
 
diff --git a/client/client_grpc_partition_test.go b/client/client_grpc_partition_test.go
index 381d2f5..623b715 100644
--- a/client/client_grpc_partition_test.go
+++ b/client/client_grpc_partition_test.go
@@ -278,6 +278,33 @@ func TestGrpcClientLoadPartitions(t *testing.T) {
 			},
 			loadNames: []string{"_part1", "_part2"},
 		}
+
+		// paritions will not be loaded
+		mock.setInjection(mShowPartitions, func(ctx context.Context, raw proto.Message) (proto.Message, error) {
+			req, ok := raw.(*server.ShowPartitionsRequest)
+			resp := &server.ShowPartitionsResponse{}
+			if !ok {
+				s, err := badRequestStatus()
+				resp.Status = s
+				return resp, err
+			}
+			assert.Equal(t, testCollectionName, req.GetCollectionName())
+			resp.PartitionIDs = make([]int64, 0, len(tc.loadNames))
+			resp.PartitionNames = make([]string, 0, len(tc.loadNames))
+			var perc int64 = 0
+			for _, part := range tc.partitions {
+				resp.PartitionIDs = append(resp.PartitionIDs, part.ID)
+				resp.PartitionNames = append(resp.PartitionNames, part.Name)
+				resp.InMemoryPercentages = append(resp.InMemoryPercentages, perc)
+			}
+			s, err := successStatus()
+			resp.Status = s
+			return resp, err
+		})
+		quickCtx, cancel := context.WithTimeout(ctx, 10*time.Millisecond)
+		defer cancel()
+		assert.NotNil(t, c.LoadPartitions(quickCtx, tc.collName, tc.loadNames, false))
+
 		mock.setInjection(mShowPartitions, func(_ context.Context, raw proto.Message) (proto.Message, error) {
 			return &server.ShowPartitionsResponse{}, errors.New("always fail")
 		})