Fix a bug that mmap can enabled even the collection is loaded

Signed-off-by: yhmo <[email protected]>
yhmo · Mar 4, 2024 · 8b64dcf · 8b64dcf
1 parent fa73520
commit 8b64dcf
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 10 deletions.
diff --git a/internal/proxy/task.go b/internal/proxy/task.go
@@ -881,6 +881,18 @@ func (t *alterCollectionTask) PreExecute(ctx context.Context) error {
 	t.Base.MsgType = commonpb.MsgType_AlterCollection
 	t.Base.SourceID = paramtable.GetNodeID()
 
+	// The proto definition of AlterCollectionRequest accepts a collection name or a collection id.
+	// Usually, the client input a collection name, we use the collection name
+	// to get collection id to call isCollectionLoaded().
+	collName := t.GetCollectionName()
+	if collName != "" {
+		collectionID, err := globalMetaCache.GetCollectionID(ctx, t.GetDbName(), collName)
+		if err != nil {
+			return err
+		}
+		t.CollectionID = collectionID
+	}
+
 	if hasMmapProp(t.Properties...) || hasLazyLoadProp(t.Properties...) {
 		loaded, err := isCollectionLoaded(ctx, t.queryCoord, t.CollectionID)
 		if err != nil {

diff --git a/internal/proxy/task_test.go b/internal/proxy/task_test.go
@@ -3605,14 +3605,38 @@ func TestAlterCollectionCheckLoaded(t *testing.T) {
 		CollectionIDs:       []int64{1},
 		InMemoryPercentages: []int64{100},
 	}, nil)
-	task := &alterCollectionTask{
-		AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
-			Base:         &commonpb.MsgBase{},
-			CollectionID: 1,
-			Properties:   []*commonpb.KeyValuePair{{Key: common.MmapEnabledKey, Value: "true"}},
-		},
-		queryCoord: qc,
-	}
-	err := task.PreExecute(context.Background())
-	assert.Equal(t, merr.Code(merr.ErrCollectionLoaded), merr.Code(err))
+
+	t.Run("collection id not loaded", func(t *testing.T) {
+		task := &alterCollectionTask{
+			AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
+				Base:         &commonpb.MsgBase{},
+				CollectionID: 1,
+				Properties:   []*commonpb.KeyValuePair{{Key: common.MmapEnabledKey, Value: "true"}},
+			},
+			queryCoord: qc,
+		}
+		err := task.PreExecute(context.Background())
+		assert.Equal(t, merr.Code(merr.ErrCollectionLoaded), merr.Code(err))
+	})
+
+	t.Run("collection name not loaded", func(t *testing.T) {
+		cache := NewMockCache(t)
+		cache.On("GetCollectionID",
+			mock.Anything, // context.Context
+			mock.AnythingOfType("string"),
+			mock.AnythingOfType("string"),
+		).Return(UniqueID(1), nil)
+		globalMetaCache = cache
+
+		task := &alterCollectionTask{
+			AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
+				Base:           &commonpb.MsgBase{},
+				CollectionName: "dummy",
+				Properties:     []*commonpb.KeyValuePair{{Key: common.MmapEnabledKey, Value: "true"}},
+			},
+			queryCoord: qc,
+		}
+		err := task.PreExecute(context.Background())
+		assert.Equal(t, merr.Code(merr.ErrCollectionLoaded), merr.Code(err))
+	})
 }