From b6b74b43a5e99e73eb7e3f1efb0c25b881a4a941 Mon Sep 17 00:00:00 2001
From: bigsheeper <yihao.dai@zilliz.com>
Date: Fri, 17 Mar 2023 11:27:13 +0800
Subject: [PATCH] Fix load partition getting stuck

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
---
 client/client_grpc_collection.go | 10 +++-----
 client/client_grpc_partition.go  | 42 ++++----------------------------
 2 files changed, 9 insertions(+), 43 deletions(-)

diff --git a/client/client_grpc_collection.go b/client/client_grpc_collection.go
index 30be56e7f..b71b6e56e 100644
--- a/client/client_grpc_collection.go
+++ b/client/client_grpc_collection.go
@@ -372,16 +372,14 @@ func (c *GrpcClient) LoadCollection(ctx context.Context, collName string, async
 				return errors.New("context deadline exceeded")
 			default:
 			}
-
-			coll, err := c.ShowCollection(ctx, collName)
+			progress, err := c.GetLoadingProgress(ctx, collName, nil)
 			if err != nil {
 				return err
 			}
-			if coll.Loaded {
-				break
+			if progress == 100 {
+				return nil
 			}
-
-			time.Sleep(200 * time.Millisecond) // TODO change to configuration
+			time.Sleep(time.Millisecond * 500)
 		}
 	}
 	return nil
diff --git a/client/client_grpc_partition.go b/client/client_grpc_partition.go
index 87612331e..fa1d1d1ba 100644
--- a/client/client_grpc_partition.go
+++ b/client/client_grpc_partition.go
@@ -14,9 +14,8 @@ package client
 import (
 	"context"
 	"fmt"
-	"time"
-
 	"github.com/cockroachdb/errors"
+	"time"
 
 	common "github.com/milvus-io/milvus-proto/go-api/commonpb"
 	server "github.com/milvus-io/milvus-proto/go-api/milvuspb"
@@ -149,23 +148,6 @@ func (c *GrpcClient) LoadPartitions(ctx context.Context, collName string, partit
 			return err
 		}
 	}
-	partitions, err := c.ShowPartitions(ctx, collName)
-	if err != nil {
-		return err
-	}
-	m := make(map[string]int64)
-	for _, partition := range partitions {
-		m[partition.Name] = partition.ID
-	}
-	// load partitions ids
-	ids := make(map[int64]struct{})
-	for _, partitionName := range partitionNames {
-		id, has := m[partitionName]
-		if !has {
-			return fmt.Errorf("Collection %s does not has partitions %s", collName, partitionName)
-		}
-		ids[id] = struct{}{}
-	}
 
 	req := &server.LoadPartitionsRequest{
 		DbName:         "", // reserved
@@ -187,28 +169,14 @@ func (c *GrpcClient) LoadPartitions(ctx context.Context, collName string, partit
 				return errors.New("context deadline exceeded")
 			default:
 			}
-			partitions, err := c.ShowPartitions(ctx, collName)
+			progress, err := c.GetLoadingProgress(ctx, collName, partitionNames)
 			if err != nil {
 				return err
 			}
-			foundLoading := false
-			loaded := 0
-			for _, partition := range partitions {
-				if _, has := ids[partition.ID]; !has {
-					continue
-				}
-				if !partition.Loaded {
-					//Not loaded
-					foundLoading = true
-					break
-				}
-				loaded++
-			}
-			if foundLoading || loaded < len(partitionNames) {
-				time.Sleep(time.Millisecond * 100)
-				continue
+			if progress == 100 {
+				return nil
 			}
-			break
+			time.Sleep(time.Millisecond * 500)
 		}
 	}