From 14ddfe679fc580da01192307bcb83ef1365d02f8 Mon Sep 17 00:00:00 2001
From: Jiquan Long <jiquan.long@zilliz.com>
Date: Mon, 24 Oct 2022 15:57:29 +0800
Subject: [PATCH] Fix failed to show partitions (#20007)

Signed-off-by: longjiquan <jiquan.long@zilliz.com>

Signed-off-by: longjiquan <jiquan.long@zilliz.com>
---
 internal/rootcoord/meta_table.go      | 41 +++++++++++++++++++++++----
 internal/rootcoord/meta_table_test.go | 40 ++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 6 deletions(-)

diff --git a/internal/rootcoord/meta_table.go b/internal/rootcoord/meta_table.go
index 79c077875b..7d1961d4cb 100644
--- a/internal/rootcoord/meta_table.go
+++ b/internal/rootcoord/meta_table.go
@@ -205,6 +205,35 @@ func (mt *MetaTable) ChangeCollectionState(ctx context.Context, collectionID Uni
 	return nil
 }
 
+func (mt *MetaTable) removeIfNameMatchedInternal(collectionID UniqueID, name string) {
+	id, ok := mt.collName2ID[name]
+	if ok && id == collectionID {
+		delete(mt.collName2ID, name)
+	}
+}
+
+func (mt *MetaTable) removeIfAliasMatchedInternal(collectionID UniqueID, alias string) {
+	id, ok := mt.collAlias2ID[alias]
+	if ok && id == collectionID {
+		delete(mt.collAlias2ID, alias)
+	}
+}
+
+func (mt *MetaTable) removeIfMatchedInternal(collectionID UniqueID, name string) {
+	mt.removeIfNameMatchedInternal(collectionID, name)
+	mt.removeIfAliasMatchedInternal(collectionID, name)
+}
+
+func (mt *MetaTable) removeAllNamesIfMatchedInternal(collectionID UniqueID, names []string) {
+	for _, name := range names {
+		mt.removeIfMatchedInternal(collectionID, name)
+	}
+}
+
+func (mt *MetaTable) removeCollectionByIDInternal(collectionID UniqueID) {
+	delete(mt.collID2Meta, collectionID)
+}
+
 func (mt *MetaTable) RemoveCollection(ctx context.Context, collectionID UniqueID, ts Timestamp) error {
 	mt.ddLock.Lock()
 	defer mt.ddLock.Unlock()
@@ -218,18 +247,18 @@ func (mt *MetaTable) RemoveCollection(ctx context.Context, collectionID UniqueID
 		return err
 	}
 
+	allNames := common.CloneStringList(aliases)
+
 	var name string
 	coll, ok := mt.collID2Meta[collectionID]
 	if ok && coll != nil {
 		name = coll.Name
-		delete(mt.collName2ID, name)
+		allNames = append(allNames, name)
 	}
 
-	for _, alias := range aliases {
-		delete(mt.collAlias2ID, alias)
-	}
-
-	delete(mt.collID2Meta, collectionID)
+	// We cannot delete the name directly, since newly collection with same name may be created.
+	mt.removeAllNamesIfMatchedInternal(collectionID, allNames)
+	mt.removeCollectionByIDInternal(collectionID)
 
 	log.Info("remove collection", zap.String("name", name), zap.Int64("id", collectionID), zap.Strings("aliases", aliases))
 	return nil
diff --git a/internal/rootcoord/meta_table_test.go b/internal/rootcoord/meta_table_test.go
index 2e6d890d85..483e14595d 100644
--- a/internal/rootcoord/meta_table_test.go
+++ b/internal/rootcoord/meta_table_test.go
@@ -980,3 +980,43 @@ func TestMetaTable_getLatestCollectionByIDInternal(t *testing.T) {
 		assert.Equal(t, 1, len(coll.Partitions))
 	})
 }
+
+func TestMetaTable_RemoveCollection(t *testing.T) {
+	t.Run("catalog error", func(t *testing.T) {
+		catalog := mocks.NewRootCoordCatalog(t)
+		catalog.On("DropCollection",
+			mock.Anything, // context.Context
+			mock.Anything, // model.Collection
+			mock.AnythingOfType("uint64"),
+		).Return(errors.New("error mock DropCollection"))
+		meta := &MetaTable{catalog: catalog}
+		ctx := context.Background()
+		err := meta.RemoveCollection(ctx, 100, 9999)
+		assert.Error(t, err)
+	})
+
+	t.Run("normal case", func(t *testing.T) {
+		catalog := mocks.NewRootCoordCatalog(t)
+		catalog.On("DropCollection",
+			mock.Anything, // context.Context
+			mock.Anything, // model.Collection
+			mock.AnythingOfType("uint64"),
+		).Return(nil)
+		meta := &MetaTable{
+			catalog: catalog,
+			collAlias2ID: map[string]typeutil.UniqueID{
+				"alias1": 100,
+				"alias2": 100,
+			},
+			collID2Meta: map[typeutil.UniqueID]*model.Collection{
+				100: {Name: "collection"},
+			},
+			collName2ID: map[string]typeutil.UniqueID{
+				"collection": 100,
+			},
+		}
+		ctx := context.Background()
+		err := meta.RemoveCollection(ctx, 100, 9999)
+		assert.NoError(t, err)
+	})
+}