fix: failed to release collection with more than 128 partitions (milv…

…us-io#28446) milvus-io#28343 Signed-off-by: wayblink <[email protected]>
bigsheeper · Nov 19, 2023 · 7f46f4c · 7f46f4c
1 parent a1c505d
commit 7f46f4c
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 8 deletions.
diff --git a/internal/metastore/kv/querycoord/kv_catalog.go b/internal/metastore/kv/querycoord/kv_catalog.go
@@ -21,6 +21,8 @@ const (
 	CollectionMetaPrefixV1   = "queryCoord-collectionMeta"
 	ReplicaMetaPrefixV1      = "queryCoord-ReplicaMeta"
 	ResourceGroupPrefix      = "queryCoord-ResourceGroup"
+
+	MetaOpsBatchSize = 128
 )
 
 type Catalog struct {
@@ -195,27 +197,42 @@ func (s Catalog) ReleaseCollection(collection int64) error {
 	if err != nil {
 		return err
 	}
-	partitions := make([]*querypb.PartitionLoadInfo, 0)
+	partitionIDs := make([]int64, 0)
 	for _, v := range values {
 		info := querypb.PartitionLoadInfo{}
 		if err = proto.Unmarshal([]byte(v), &info); err != nil {
 			return err
 		}
-		partitions = append(partitions, &info)
+		partitionIDs = append(partitionIDs, (&info).GetPartitionID())
 	}
 	// remove collection and obtained partitions
-	keys := lo.Map(partitions, func(partition *querypb.PartitionLoadInfo, _ int) string {
-		return EncodePartitionLoadInfoKey(collection, partition.GetPartitionID())
-	})
-	k := EncodeCollectionLoadInfoKey(collection)
-	keys = append(keys, k)
-	return s.cli.MultiRemove(keys)
+	collectionKey := EncodeCollectionLoadInfoKey(collection)
+	err = s.cli.Remove(collectionKey)
+	if err != nil {
+		return err
+	}
+	return s.ReleasePartition(collection, partitionIDs...)
 }
 
 func (s Catalog) ReleasePartition(collection int64, partitions ...int64) error {
 	keys := lo.Map(partitions, func(partition int64, _ int) string {
 		return EncodePartitionLoadInfoKey(collection, partition)
 	})
+	if len(partitions) >= MetaOpsBatchSize {
+		index := 0
+		for index < len(partitions) {
+			endIndex := index + MetaOpsBatchSize
+			if endIndex > len(partitions) {
+				endIndex = len(partitions)
+			}
+			err := s.cli.MultiRemove(keys[index:endIndex])
+			if err != nil {
+				return err
+			}
+			index = endIndex
+		}
+		return nil
+	}
 	return s.cli.MultiRemove(keys)
 }
 

diff --git a/internal/metastore/kv/querycoord/kv_catalog_test.go b/internal/metastore/kv/querycoord/kv_catalog_test.go
@@ -121,6 +121,23 @@ func (suite *CatalogTestSuite) TestPartition() {
 	suite.Len(partitions, 1)
 }
 
+func (suite *CatalogTestSuite) TestReleaseManyPartitions() {
+	partitionIds := make([]int64, 0)
+	for i := 1; i <= 150; i++ {
+		suite.catalog.SavePartition(&querypb.PartitionLoadInfo{
+			CollectionID: 1,
+			PartitionID:  int64(i),
+		})
+		partitionIds = append(partitionIds, int64(i))
+	}
+
+	err := suite.catalog.ReleasePartition(1, partitionIds...)
+	suite.NoError(err)
+	partitions, err := suite.catalog.GetPartitions()
+	suite.NoError(err)
+	suite.Len(partitions, 0)
+}
+
 func (suite *CatalogTestSuite) TestReplica() {
 	suite.catalog.SaveReplica(&querypb.Replica{
 		CollectionID: 1,