Fix dml stream leakage in Proxy (#19450)

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

internal/common/error.go

@@ -16,7 +16,11 @@
 
 package common
 
-import "fmt"
+import (
+	"fmt"
+
+	"github.com/milvus-io/milvus/api/commonpb"
+)
 
 type IgnorableError struct {
 	msg string
@@ -55,3 +59,47 @@ type KeyNotExistError struct {
 func (k *KeyNotExistError) Error() string {
 	return fmt.Sprintf("there is no value on key = %s", k.key)
 }
+
+type statusError struct {
+	commonpb.Status
+}
+
+func (e *statusError) Error() string {
+	return fmt.Sprintf("code: %s, reason: %s", e.GetErrorCode().String(), e.GetReason())
+}
+
+func NewStatusError(code commonpb.ErrorCode, reason string) *statusError {
+	return &statusError{Status: commonpb.Status{ErrorCode: code, Reason: reason}}
+}
+
+func IsStatusError(e error) bool {
+	_, ok := e.(*statusError)
+	return ok
+}
+
+var (
+	// static variable, save temporary memory.
+	collectionNotExistCodes = []commonpb.ErrorCode{
+		commonpb.ErrorCode_UnexpectedError, // TODO: remove this after SDK remove this dependency.
+		commonpb.ErrorCode_CollectionNotExists,
+	}
+)
+
+func NewCollectionNotExistError(msg string) *statusError {
+	return NewStatusError(commonpb.ErrorCode_CollectionNotExists, msg)
+}
+
+func IsCollectionNotExistError(e error) bool {
+	statusError, ok := e.(*statusError)
+	if !ok {
+		return false
+	}
+	// cycle import: common -> funcutil -> types -> sessionutil -> common
+	// return funcutil.SliceContain(collectionNotExistCodes, statusError.GetErrorCode())
+	for _, code := range collectionNotExistCodes {
+		if code == statusError.GetErrorCode() {
+			return true
+		}
+	}
+	return false
+}

internal/common/error_test.go

@@ -21,6 +21,8 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/milvus-io/milvus/api/commonpb"
+
 	"github.com/stretchr/testify/assert"
 )
 
@@ -36,3 +38,26 @@ func TestNotExistError(t *testing.T) {
 	assert.Equal(t, false, IsKeyNotExistError(err))
 	assert.Equal(t, true, IsKeyNotExistError(NewKeyNotExistError("foo")))
 }
+
+func TestStatusError_Error(t *testing.T) {
+	err := NewCollectionNotExistError("collection not exist")
+	fmt.Println("test status error: ", err.Error())
+}
+
+func TestIsStatusError(t *testing.T) {
+	err := NewCollectionNotExistError("collection not exist")
+	assert.True(t, IsStatusError(err))
+	assert.False(t, IsStatusError(errors.New("not status error")))
+	assert.False(t, IsStatusError(nil))
+}
+
+func Test_IsCollectionNotExistError(t *testing.T) {
+	assert.False(t, IsCollectionNotExistError(nil))
+	assert.False(t, IsCollectionNotExistError(errors.New("not status error")))
+	for _, code := range collectionNotExistCodes {
+		err := NewStatusError(code, "collection not exist")
+		assert.True(t, IsCollectionNotExistError(err))
+	}
+	assert.True(t, IsCollectionNotExistError(NewCollectionNotExistError("collection not exist")))
+	assert.False(t, IsCollectionNotExistError(NewStatusError(commonpb.ErrorCode_BuildIndexError, "")))
+}

internal/metastore/catalog.go

@@ -10,6 +10,7 @@ import (
 	"github.com/milvus-io/milvus/internal/util/typeutil"
 )
 
+//go:generate mockery --name=RootCoordCatalog
 type RootCoordCatalog interface {
 	CreateCollection(ctx context.Context, collectionInfo *model.Collection, ts typeutil.Timestamp) error
 	GetCollectionByID(ctx context.Context, collectionID typeutil.UniqueID, ts typeutil.Timestamp) (*model.Collection, error)

internal/metastore/kv/rootcoord/kv_catalog.go

@@ -137,7 +137,7 @@ func (kc *Catalog) loadCollection(ctx context.Context, collectionID typeutil.Uni
 	collKey := buildCollectionKey(collectionID)
 	collVal, err := kc.Snapshot.Load(collKey, ts)
 	if err != nil {
-		return nil, fmt.Errorf("can't find collection: %d", collectionID)
+		return nil, common.NewCollectionNotExistError(fmt.Sprintf("collection not found: %d", collectionID))
 	}
 
 	collMeta := &pb.CollectionInfo{}
@@ -487,7 +487,7 @@ func (kc *Catalog) GetCollectionByName(ctx context.Context, collectionName strin
 		}
 	}
 
-	return nil, fmt.Errorf("can't find collection: %s, at timestamp = %d", collectionName, ts)
+	return nil, common.NewCollectionNotExistError(fmt.Sprintf("can't find collection: %s, at timestamp = %d", collectionName, ts))
 }
 
 func (kc *Catalog) ListCollections(ctx context.Context, ts typeutil.Timestamp) (map[string]*model.Collection, error) {

internal/metastore/mocks/RootCoordCatalog.go

+// Code generated by mockery v2.14.0. DO NOT EDIT.
+
+package mocks
+
+import (
+	context "context"
+
+	milvuspb "github.com/milvus-io/milvus/api/milvuspb"
+	metastore "github.com/milvus-io/milvus/internal/metastore"
+
+	mock "github.com/stretchr/testify/mock"
+
+	model "github.com/milvus-io/milvus/internal/metastore/model"
+)
+
+// RootCoordCatalog is an autogenerated mock type for the RootCoordCatalog type
+type RootCoordCatalog struct {
+	mock.Mock
+}
+
+// AlterAlias provides a mock function with given fields: ctx, alias, ts
+func (_m *RootCoordCatalog) AlterAlias(ctx context.Context, alias *model.Alias, ts uint64) error {
+	ret := _m.Called(ctx, alias, ts)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, *model.Alias, uint64) error); ok {
+		r0 = rf(ctx, alias, ts)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// AlterCollection provides a mock function with given fields: ctx, oldColl, newColl, alterType, ts
+func (_m *RootCoordCatalog) AlterCollection(ctx context.Context, oldColl *model.Collection, newColl *model.Collection, alterType metastore.AlterType, ts uint64) error {
+	ret := _m.Called(ctx, oldColl, newColl, alterType, ts)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, *model.Collection, *model.Collection, metastore.AlterType, uint64) error); ok {
+		r0 = rf(ctx, oldColl, newColl, alterType, ts)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// AlterCredential provides a mock function with given fields: ctx, credential
+func (_m *RootCoordCatalog) AlterCredential(ctx context.Context, credential *model.Credential) error {
+	ret := _m.Called(ctx, credential)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, *model.Credential) error); ok {
+		r0 = rf(ctx, credential)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// AlterGrant provides a mock function with given fields: ctx, tenant, entity, operateType
+func (_m *RootCoordCatalog) AlterGrant(ctx context.Context, tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error {
+	ret := _m.Called(ctx, tenant, entity, operateType)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, string, *milvuspb.GrantEntity, milvuspb.OperatePrivilegeType) error); ok {
+		r0 = rf(ctx, tenant, entity, operateType)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// AlterPartition provides a mock function with given fields: ctx, oldPart, newPart, alterType, ts
+func (_m *RootCoordCatalog) AlterPartition(ctx context.Context, oldPart *model.Partition, newPart *model.Partition, alterType metastore.AlterType, ts uint64) error {
+	ret := _m.Called(ctx, oldPart, newPart, alterType, ts)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, *model.Partition, *model.Partition, metastore.AlterType, uint64) error); ok {
+		r0 = rf(ctx, oldPart, newPart, alterType, ts)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// AlterUserRole provides a mock function with given fields: ctx, tenant, userEntity, roleEntity, operateType
+func (_m *RootCoordCatalog) AlterUserRole(ctx context.Context, tenant string, userEntity *milvuspb.UserEntity, roleEntity *milvuspb.RoleEntity, operateType milvuspb.OperateUserRoleType) error {
+	ret := _m.Called(ctx, tenant, userEntity, roleEntity, operateType)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, string, *milvuspb.UserEntity, *milvuspb.RoleEntity, milvuspb.OperateUserRoleType) error); ok {
+		r0 = rf(ctx, tenant, userEntity, roleEntity, operateType)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// Close provides a mock function with given fields:
+func (_m *RootCoordCatalog) Close() {
+	_m.Called()
+}
+
+// CollectionExists provides a mock function with given fields: ctx, collectionID, ts
+func (_m *RootCoordCatalog) CollectionExists(ctx context.Context, collectionID int64, ts uint64) bool {
+	ret := _m.Called(ctx, collectionID, ts)
+
+	var r0 bool
+	if rf, ok := ret.Get(0).(func(context.Context, int64, uint64) bool); ok {
+		r0 = rf(ctx, collectionID, ts)
+	} else {
+		r0 = ret.Get(0).(bool)
+	}
+
+	return r0
+}
+
+// CreateAlias provides a mock function with given fields: ctx, alias, ts
+func (_m *RootCoordCatalog) CreateAlias(ctx context.Context, alias *model.Alias, ts uint64) error {
+	ret := _m.Called(ctx, alias, ts)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, *model.Alias, uint64) error); ok {
+		r0 = rf(ctx, alias, ts)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// CreateCollection provides a mock function with given fields: ctx, collectionInfo, ts
+func (_m *RootCoordCatalog) CreateCollection(ctx context.Context, collectionInfo *model.Collection, ts uint64) error {
+	ret := _m.Called(ctx, collectionInfo, ts)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, *model.Collection, uint64) error); ok {
+		r0 = rf(ctx, collectionInfo, ts)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// CreateCredential provides a mock function with given fields: ctx, credential
+func (_m *RootCoordCatalog) CreateCredential(ctx context.Context, credential *model.Credential) error {
+	ret := _m.Called(ctx, credential)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, *model.Credential) error); ok {
+		r0 = rf(ctx, credential)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// CreatePartition provides a mock function with given fields: ctx, partition, ts
+func (_m *RootCoordCatalog) CreatePartition(ctx context.Context, partition *model.Partition, ts uint64) error {
+	ret := _m.Called(ctx, partition, ts)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, *model.Partition, uint64) error); ok {
+		r0 = rf(ctx, partition, ts)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// CreateRole provides a mock function with given fields: ctx, tenant, entity
+func (_m *RootCoordCatalog) CreateRole(ctx context.Context, tenant string, entity *milvuspb.RoleEntity) error {
+	ret := _m.Called(ctx, tenant, entity)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, string, *milvuspb.RoleEntity) error); ok {
+		r0 = rf(ctx, tenant, entity)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// DeleteGrant provides a mock function with given fields: ctx, tenant, role
+func (_m *RootCoordCatalog) DeleteGrant(ctx context.Context, tenant string, role *milvuspb.RoleEntity) error {
+	ret := _m.Called(ctx, tenant, role)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, string, *milvuspb.RoleEntity) error); ok {
+		r0 = rf(ctx, tenant, role)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// DropAlias provides a mock function with given fields: ctx, alias, ts
+func (_m *RootCoordCatalog) DropAlias(ctx context.Context, alias string, ts uint64) error {
+	ret := _m.Called(ctx, alias, ts)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, string, uint64) error); ok {
+		r0 = rf(ctx, alias, ts)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// DropCollection provides a mock function with given fields: ctx, collectionInfo, ts
+func (_m *RootCoordCatalog) DropCollection(ctx context.Context, collectionInfo *model.Collection, ts uint64) error {
+	ret := _m.Called(ctx, collectionInfo, ts)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, *model.Collection, uint64) error); ok {
+		r0 = rf(ctx, collectionInfo, ts)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// DropCredential provides a mock function with given fields: ctx, username
+func (_m *RootCoordCatalog) DropCredential(ctx context.Context, username string) error {
+	ret := _m.Called(ctx, username)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, string) error); ok {
+		r0 = rf(ctx, username)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// DropPartition provides a mock function with given fields: ctx, collectionID, partitionID, ts
+func (_m *RootCoordCatalog) DropPartition(ctx context.Context, collectionID int64, partitionID int64, ts uint64) error {
+	ret := _m.Called(ctx, collectionID, partitionID, ts)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, int64, int64, uint64) error); ok {
+		r0 = rf(ctx, collectionID, partitionID, ts)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// DropRole provides a mock function with given fields: ctx, tenant, roleName
+func (_m *RootCoordCatalog) DropRole(ctx context.Context, tenant string, roleName string) error {
+	ret := _m.Called(ctx, tenant, roleName)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, string, string) error); ok {
+		r0 = rf(ctx, tenant, roleName)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// GetCollectionByID provides a mock function with given fields: ctx, collectionID, ts
+func (_m *RootCoordCatalog) GetCollectionByID(ctx context.Context, collectionID int64, ts uint64) (*model.Collection, error) {
+	ret := _m.Called(ctx, collectionID, ts)
+
+	var r0 *model.Collection
+	if rf, ok := ret.Get(0).(func(context.Context, int64, uint64) *model.Collection); ok {
+		r0 = rf(ctx, collectionID, ts)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*model.Collection)
+		}
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func(context.Context, int64, uint64) error); ok {
+		r1 = rf(ctx, collectionID, ts)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// GetCollectionByName provides a mock function with given fields: ctx, collectionName, ts
+func (_m *RootCoordCatalog) GetCollectionByName(ctx context.Context, collectionName string, ts uint64) (*model.Collection, error) {
+	ret := _m.Called(ctx, collectionName, ts)
+
+	var r0 *model.Collection
+	if rf, ok := ret.Get(0).(func(context.Context, string, uint64) *model.Collection); ok {
+		r0 = rf(ctx, collectionName, ts)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*model.Collection)
+		}
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func(context.Context, string, uint64) error); ok {
+		r1 = rf(ctx, collectionName, ts)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// GetCredential provides a mock function with given fields: ctx, username
+func (_m *RootCoordCatalog) GetCredential(ctx context.Context, username string) (*model.Credential, error) {
+	ret := _m.Called(ctx, username)
+
+	var r0 *model.Credential
+	if rf, ok := ret.Get(0).(func(context.Context, string) *model.Credential); ok {
+		r0 = rf(ctx, username)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*model.Credential)
+		}
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = rf(ctx, username)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// ListAliases provides a mock function with given fields: ctx, ts
+func (_m *RootCoordCatalog) ListAliases(ctx context.Context, ts uint64) ([]*model.Alias, error) {
+	ret := _m.Called(ctx, ts)
+
+	var r0 []*model.Alias
+	if rf, ok := ret.Get(0).(func(context.Context, uint64) []*model.Alias); ok {
+		r0 = rf(ctx, ts)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]*model.Alias)
+		}
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func(context.Context, uint64) error); ok {
+		r1 = rf(ctx, ts)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// ListCollections provides a mock function with given fields: ctx, ts
+func (_m *RootCoordCatalog) ListCollections(ctx context.Context, ts uint64) (map[string]*model.Collection, error) {
+	ret := _m.Called(ctx, ts)
+
+	var r0 map[string]*model.Collection
+	if rf, ok := ret.Get(0).(func(context.Context, uint64) map[string]*model.Collection); ok {
+		r0 = rf(ctx, ts)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(map[string]*model.Collection)
+		}
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func(context.Context, uint64) error); ok {
+		r1 = rf(ctx, ts)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// ListCredentials provides a mock function with given fields: ctx
+func (_m *RootCoordCatalog) ListCredentials(ctx context.Context) ([]string, error) {
+	ret := _m.Called(ctx)
+
+	var r0 []string
+	if rf, ok := ret.Get(0).(func(context.Context) []string); ok {
+		r0 = rf(ctx)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]string)
+		}
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func(context.Context) error); ok {
+		r1 = rf(ctx)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// ListGrant provides a mock function with given fields: ctx, tenant, entity
+func (_m *RootCoordCatalog) ListGrant(ctx context.Context, tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error) {
+	ret := _m.Called(ctx, tenant, entity)
+
+	var r0 []*milvuspb.GrantEntity
+	if rf, ok := ret.Get(0).(func(context.Context, string, *milvuspb.GrantEntity) []*milvuspb.GrantEntity); ok {
+		r0 = rf(ctx, tenant, entity)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]*milvuspb.GrantEntity)
+		}
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func(context.Context, string, *milvuspb.GrantEntity) error); ok {
+		r1 = rf(ctx, tenant, entity)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// ListPolicy provides a mock function with given fields: ctx, tenant
+func (_m *RootCoordCatalog) ListPolicy(ctx context.Context, tenant string) ([]string, error) {
+	ret := _m.Called(ctx, tenant)
+
+	var r0 []string
+	if rf, ok := ret.Get(0).(func(context.Context, string) []string); ok {
+		r0 = rf(ctx, tenant)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]string)
+		}
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = rf(ctx, tenant)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// ListRole provides a mock function with given fields: ctx, tenant, entity, includeUserInfo
+func (_m *RootCoordCatalog) ListRole(ctx context.Context, tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
+	ret := _m.Called(ctx, tenant, entity, includeUserInfo)
+
+	var r0 []*milvuspb.RoleResult
+	if rf, ok := ret.Get(0).(func(context.Context, string, *milvuspb.RoleEntity, bool) []*milvuspb.RoleResult); ok {
+		r0 = rf(ctx, tenant, entity, includeUserInfo)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]*milvuspb.RoleResult)
+		}
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func(context.Context, string, *milvuspb.RoleEntity, bool) error); ok {
+		r1 = rf(ctx, tenant, entity, includeUserInfo)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// ListUser provides a mock function with given fields: ctx, tenant, entity, includeRoleInfo
+func (_m *RootCoordCatalog) ListUser(ctx context.Context, tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error) {
+	ret := _m.Called(ctx, tenant, entity, includeRoleInfo)
+
+	var r0 []*milvuspb.UserResult
+	if rf, ok := ret.Get(0).(func(context.Context, string, *milvuspb.UserEntity, bool) []*milvuspb.UserResult); ok {
+		r0 = rf(ctx, tenant, entity, includeRoleInfo)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]*milvuspb.UserResult)
+		}
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func(context.Context, string, *milvuspb.UserEntity, bool) error); ok {
+		r1 = rf(ctx, tenant, entity, includeRoleInfo)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// ListUserRole provides a mock function with given fields: ctx, tenant
+func (_m *RootCoordCatalog) ListUserRole(ctx context.Context, tenant string) ([]string, error) {
+	ret := _m.Called(ctx, tenant)
+
+	var r0 []string
+	if rf, ok := ret.Get(0).(func(context.Context, string) []string); ok {
+		r0 = rf(ctx, tenant)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]string)
+		}
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = rf(ctx, tenant)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+type mockConstructorTestingTNewRootCoordCatalog interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewRootCoordCatalog creates a new instance of RootCoordCatalog. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewRootCoordCatalog(t mockConstructorTestingTNewRootCoordCatalog) *RootCoordCatalog {
+	mock := &RootCoordCatalog{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

internal/proto/proxy.proto

@@ -23,6 +23,9 @@ service Proxy {
 }
 
 message InvalidateCollMetaCacheRequest {
+  // MsgType:
+  //  DropCollection    ->  {meta cache, dml channels}
+  //  Other             ->  {meta cache}
   common.MsgBase base = 1;
   string db_name = 2;
   string collection_name = 3;

internal/proto/proxypb/proxy.pb.go

@@ -28,6 +28,9 @@ var _ = math.Inf
 const _ = proto.ProtoPackageIsVersion3 // please upgrade the proto package
 
 type InvalidateCollMetaCacheRequest struct {
+	// MsgType:
+	//  DropCollection    ->  {meta cache, dml channels}
+	//  Other             ->  {meta cache}
 	Base                 *commonpb.MsgBase `protobuf:"bytes,1,opt,name=base,proto3" json:"base,omitempty"`
 	DbName               string            `protobuf:"bytes,2,opt,name=db_name,json=dbName,proto3" json:"db_name,omitempty"`
 	CollectionName       string            `protobuf:"bytes,3,opt,name=collection_name,json=collectionName,proto3" json:"collection_name,omitempty"`

internal/proxy/channels_mgr.go

@@ -284,8 +284,10 @@ func (mgr *singleTypeChannelsMgr) removeStream(collectionID UniqueID) error {
 	defer mgr.mu.Unlock()
 	if info, ok := mgr.infos[collectionID]; ok {
 		decPChanMetrics(info.channelInfos.pchans)
+		info.stream.Close()
 		delete(mgr.infos, collectionID)
 	}
+	log.Info("dml stream removed", zap.Int64("collection_id", collectionID))
 	return nil
 }
 
@@ -294,9 +296,11 @@ func (mgr *singleTypeChannelsMgr) removeAllStream() error {
 	mgr.mu.Lock()
 	defer mgr.mu.Unlock()
 	for _, info := range mgr.infos {
+		info.stream.Close()
 		decPChanMetrics(info.channelInfos.pchans)
 	}
 	mgr.infos = make(map[UniqueID]streamInfos)
+	log.Info("all dml stream removed")
 	return nil
 }
 

internal/proxy/error_test.go

@@ -20,7 +20,6 @@ import (
 	"testing"
 
 	"github.com/milvus-io/milvus/api/schemapb"
-
 	"github.com/milvus-io/milvus/internal/log"
 	"go.uber.org/zap"
 )

internal/proxy/impl.go

@@ -113,6 +113,10 @@ func (node *Proxy) InvalidateCollectionMetaCache(ctx context.Context, request *p
 			globalMetaCache.RemoveCollectionsByID(ctx, collectionID)
 		}
 	}
+	if request.GetBase().GetMsgType() == commonpb.MsgType_DropCollection {
+		// no need to handle error, since this Proxy may not create dml stream for the collection.
+		_ = node.chMgr.removeDMLStream(request.GetCollectionID())
+	}
 	logutil.Logger(ctx).Info("complete to invalidate collection meta cache",
 		zap.String("role", typeutil.ProxyRole),
 		zap.String("db", request.DbName),

internal/proxy/impl_test.go

+package proxy
+
+import (
+	"context"
+	"testing"
+
+	"github.com/milvus-io/milvus/api/commonpb"
+	"github.com/milvus-io/milvus/internal/log"
+	"github.com/milvus-io/milvus/internal/proto/proxypb"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestProxy_InvalidateCollectionMetaCache_remove_stream(t *testing.T) {
+	cache := globalMetaCache
+	globalMetaCache = nil
+	defer func() { globalMetaCache = cache }()
+
+	chMgr := newMockChannelsMgr()
+	chMgr.removeDMLStreamFuncType = func(collectionID UniqueID) error {
+		log.Debug("TestProxy_InvalidateCollectionMetaCache_remove_stream, remove dml stream")
+		return nil
+	}
+
+	node := &Proxy{chMgr: chMgr}
+
+	ctx := context.Background()
+	req := &proxypb.InvalidateCollMetaCacheRequest{
+		Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_DropCollection},
+	}
+
+	status, err := node.InvalidateCollectionMetaCache(ctx, req)
+	assert.NoError(t, err)
+	assert.Equal(t, commonpb.ErrorCode_Success, status.GetErrorCode())
+}

internal/proxy/meta_cache.go

@@ -406,7 +406,7 @@ func (m *MetaCache) describeCollection(ctx context.Context, collectionName strin
 		return nil, err
 	}
 	if coll.Status.ErrorCode != commonpb.ErrorCode_Success {
-		return nil, errors.New(coll.Status.Reason)
+		return nil, common.NewStatusError(coll.GetStatus().GetErrorCode(), coll.GetStatus().GetReason())
 	}
 	resp := &milvuspb.DescribeCollectionResponse{
 		Status: coll.Status,

internal/proxy/proxy.go

@@ -412,6 +412,10 @@ func (node *Proxy) Stop() error {
 		node.shardMgr.Close()
 	}
 
+	if node.chMgr != nil {
+		_ = node.chMgr.removeAllDMLStream()
+	}
+
 	// https://github.com/milvus-io/milvus/issues/12282
 	node.UpdateStateCode(internalpb.StateCode_Abnormal)
 

internal/proxy/task.go

@@ -312,25 +312,17 @@ func (dct *dropCollectionTask) PreExecute(ctx context.Context) error {
 }
 
 func (dct *dropCollectionTask) Execute(ctx context.Context) error {
-	collID, err := globalMetaCache.GetCollectionID(ctx, dct.CollectionName)
-	if err != nil {
+	var err error
+	dct.result, err = dct.rootCoord.DropCollection(ctx, dct.DropCollectionRequest)
+	if common.IsCollectionNotExistError(err) {
 		// make dropping collection idempotent.
 		dct.result = &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}
 		return nil
 	}
-
-	dct.result, err = dct.rootCoord.DropCollection(ctx, dct.DropCollectionRequest)
-	if err != nil {
-		return err
-	}
-
-	_ = dct.chMgr.removeDMLStream(collID)
-	globalMetaCache.RemoveCollection(ctx, dct.CollectionName)
-	return nil
+	return err
 }
 
 func (dct *dropCollectionTask) PostExecute(ctx context.Context) error {
-	globalMetaCache.RemoveCollection(ctx, dct.CollectionName)
 	return nil
 }
 

internal/proxy/task_test.go

@@ -27,6 +27,9 @@ import (
 	"testing"
 	"time"
 
+	"github.com/milvus-io/milvus/internal/mocks"
+	"github.com/stretchr/testify/mock"
+
 	"github.com/milvus-io/milvus/internal/allocator"
 	"github.com/milvus-io/milvus/internal/mq/msgstream"
 
@@ -747,79 +750,6 @@ func TestCreateCollectionTask(t *testing.T) {
 	})
 }
 
-func TestDropCollectionTask(t *testing.T) {
-	Params.InitOnce()
-
-	prefix := "TestDropCollectionTask"
-	dbName := ""
-	collectionName := prefix + funcutil.GenRandomStr()
-	ctx := context.Background()
-
-	task := &dropCollectionTask{
-		Condition: NewTaskCondition(ctx),
-		DropCollectionRequest: &milvuspb.DropCollectionRequest{
-			Base: &commonpb.MsgBase{
-				MsgType:   commonpb.MsgType_DropCollection,
-				MsgID:     100,
-				Timestamp: 100,
-			},
-			DbName:         dbName,
-			CollectionName: collectionName,
-		},
-		ctx:    ctx,
-		result: nil,
-	}
-
-	task.SetID(100)
-	assert.Equal(t, UniqueID(100), task.ID())
-	assert.Equal(t, DropCollectionTaskName, task.Name())
-	assert.Equal(t, commonpb.MsgType_DropCollection, task.Type())
-	task.SetTs(100)
-	assert.Equal(t, Timestamp(100), task.BeginTs())
-	assert.Equal(t, Timestamp(100), task.EndTs())
-
-	err := task.PreExecute(ctx)
-	assert.NoError(t, err)
-	assert.Equal(t, Params.ProxyCfg.GetNodeID(), task.GetBase().GetSourceID())
-
-	task.CollectionName = "#0xc0de"
-	err = task.PreExecute(ctx)
-	assert.Error(t, err)
-	task.CollectionName = collectionName
-
-	cache := newMockCache()
-	chMgr := newMockChannelsMgr()
-	rc := newMockRootCoord()
-
-	globalMetaCache = cache
-	task.chMgr = chMgr
-	task.rootCoord = rc
-
-	cache.setGetIDFunc(func(ctx context.Context, collectionName string) (typeutil.UniqueID, error) {
-		return 0, errors.New("mock")
-	})
-	err = task.Execute(ctx)
-	assert.NoError(t, err)
-	cache.setGetIDFunc(func(ctx context.Context, collectionName string) (typeutil.UniqueID, error) {
-		return 0, nil
-	})
-
-	rc.DropCollectionFunc = func(ctx context.Context, request *milvuspb.DropCollectionRequest) (*commonpb.Status, error) {
-		return nil, errors.New("mock")
-	}
-	err = task.Execute(ctx)
-	assert.Error(t, err)
-	rc.DropCollectionFunc = func(ctx context.Context, request *milvuspb.DropCollectionRequest) (*commonpb.Status, error) {
-		return &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil
-	}
-
-	// normal case
-	err = task.Execute(ctx)
-	assert.NoError(t, err)
-	err = task.PostExecute(ctx)
-	assert.NoError(t, err)
-}
-
 func TestHasCollectionTask(t *testing.T) {
 	Params.InitOnce()
 	rc := NewRootCoordMock()
@@ -2192,3 +2122,54 @@ func Test_createIndexTask_PreExecute(t *testing.T) {
 		assert.Error(t, cit.PreExecute(context.Background()))
 	})
 }
+
+func Test_dropCollectionTask_PreExecute(t *testing.T) {
+	Params.InitOnce()
+	dct := &dropCollectionTask{DropCollectionRequest: &milvuspb.DropCollectionRequest{
+		Base:           &commonpb.MsgBase{},
+		CollectionName: "0xffff", // invalid
+	}}
+	ctx := context.Background()
+	err := dct.PreExecute(ctx)
+	assert.Error(t, err)
+	dct.DropCollectionRequest.CollectionName = "valid"
+	err = dct.PreExecute(ctx)
+	assert.NoError(t, err)
+}
+
+func Test_dropCollectionTask_Execute(t *testing.T) {
+	mockRC := mocks.NewRootCoord(t)
+	mockRC.On("DropCollection",
+		mock.Anything, // context.Context
+		mock.Anything, // *milvuspb.DropCollectionRequest
+	).Return(&commonpb.Status{}, func(ctx context.Context, request *milvuspb.DropCollectionRequest) error {
+		switch request.GetCollectionName() {
+		case "c1":
+			return errors.New("error mock DropCollection")
+		case "c2":
+			return common.NewStatusError(commonpb.ErrorCode_CollectionNotExists, "collection not exist")
+		default:
+			return nil
+		}
+	})
+
+	ctx := context.Background()
+
+	dct := &dropCollectionTask{rootCoord: mockRC, DropCollectionRequest: &milvuspb.DropCollectionRequest{CollectionName: "normal"}}
+	err := dct.Execute(ctx)
+	assert.NoError(t, err)
+
+	dct.DropCollectionRequest.CollectionName = "c1"
+	err = dct.Execute(ctx)
+	assert.Error(t, err)
+
+	dct.DropCollectionRequest.CollectionName = "c2"
+	err = dct.Execute(ctx)
+	assert.NoError(t, err)
+	assert.Equal(t, commonpb.ErrorCode_Success, dct.result.GetErrorCode())
+}
+
+func Test_dropCollectionTask_PostExecute(t *testing.T) {
+	dct := &dropCollectionTask{}
+	assert.NoError(t, dct.PostExecute(context.Background()))
+}

internal/rootcoord/drop_collection_task.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"fmt"
 
+	"github.com/milvus-io/milvus/internal/common"
+
 	"github.com/milvus-io/milvus/internal/log"
 	"go.uber.org/zap"
 
@@ -41,12 +43,16 @@ func (t *dropCollectionTask) Execute(ctx context.Context) error {
 	// dropping collection with `ts1` but a collection exists in catalog with newer ts which is bigger than `ts1`.
 	// fortunately, if ddls are promised to execute in sequence, then everything is OK. The `ts1` will always be latest.
 	collMeta, err := t.core.meta.GetCollectionByName(ctx, t.Req.GetCollectionName(), typeutil.MaxTimestamp)
-	if err != nil {
+	if common.IsCollectionNotExistError(err) {
 		// make dropping collection idempotent.
 		log.Warn("drop non-existent collection", zap.String("collection", t.Req.GetCollectionName()))
 		return nil
 	}
 
+	if err != nil {
+		return err
+	}
+
 	// meta cache of all aliases should also be cleaned.
 	aliases := t.core.meta.ListAliasesByID(collMeta.CollectionID)
 
@@ -65,6 +71,7 @@ func (t *dropCollectionTask) Execute(ctx context.Context) error {
 		collectionNames: append(aliases, collMeta.Name),
 		collectionID:    collMeta.CollectionID,
 		ts:              ts,
+		opts:            []expireCacheOpt{expireCacheWithDropFlag()},
 	})
 
 	redoTask.AddAsyncStep(&releaseCollectionStep{

internal/rootcoord/drop_collection_task_test.go

@@ -5,6 +5,8 @@ import (
 	"errors"
 	"testing"
 
+	"github.com/milvus-io/milvus/internal/common"
+
 	"github.com/milvus-io/milvus/api/commonpb"
 	"github.com/milvus-io/milvus/api/milvuspb"
 	"github.com/milvus-io/milvus/internal/metastore/model"
@@ -66,7 +68,18 @@ func Test_dropCollectionTask_Prepare(t *testing.T) {
 func Test_dropCollectionTask_Execute(t *testing.T) {
 	t.Run("drop non-existent collection", func(t *testing.T) {
 		collectionName := funcutil.GenRandomStr()
-		core := newTestCore(withInvalidMeta())
+		meta := mockrootcoord.NewIMetaTable(t)
+		meta.On("GetCollectionByName",
+			mock.Anything, // context.Context.
+			mock.AnythingOfType("string"),
+			mock.AnythingOfType("uint64"),
+		).Return(nil, func(ctx context.Context, name string, ts Timestamp) error {
+			if collectionName == name {
+				return common.NewCollectionNotExistError("collection not exist")
+			}
+			return errors.New("error mock GetCollectionByName")
+		})
+		core := newTestCore(withMeta(meta))
 		task := &dropCollectionTask{
 			baseTask: baseTask{core: core},
 			Req: &milvuspb.DropCollectionRequest{
@@ -76,6 +89,9 @@ func Test_dropCollectionTask_Execute(t *testing.T) {
 		}
 		err := task.Execute(context.Background())
 		assert.NoError(t, err)
+		task.Req.CollectionName = collectionName + "_test"
+		err = task.Execute(context.Background())
+		assert.Error(t, err)
 	})
 
 	t.Run("failed to expire cache", func(t *testing.T) {

internal/rootcoord/expire_cache.go

+package rootcoord
+
+import (
+	"context"
+
+	"github.com/milvus-io/milvus/api/commonpb"
+	"github.com/milvus-io/milvus/internal/proto/proxypb"
+	"github.com/milvus-io/milvus/internal/util/typeutil"
+)
+
+type expireCacheConfig struct {
+	withDropFlag bool
+}
+
+func (c expireCacheConfig) apply(req *proxypb.InvalidateCollMetaCacheRequest) {
+	if !c.withDropFlag {
+		return
+	}
+	if req.GetBase() == nil {
+		req.Base = &commonpb.MsgBase{}
+	}
+	req.Base.MsgType = commonpb.MsgType_DropCollection
+}
+
+func defaultExpireCacheConfig() expireCacheConfig {
+	return expireCacheConfig{withDropFlag: false}
+}
+
+type expireCacheOpt func(c *expireCacheConfig)
+
+func expireCacheWithDropFlag() expireCacheOpt {
+	return func(c *expireCacheConfig) {
+		c.withDropFlag = true
+	}
+}
+
+// ExpireMetaCache will call invalidate collection meta cache
+func (c *Core) ExpireMetaCache(ctx context.Context, collNames []string, collectionID UniqueID, ts typeutil.Timestamp, opts ...expireCacheOpt) error {
+	// if collectionID is specified, invalidate all the collection meta cache with the specified collectionID and return
+	if collectionID != InvalidCollectionID {
+		req := proxypb.InvalidateCollMetaCacheRequest{
+			Base: &commonpb.MsgBase{
+				Timestamp: ts,
+				SourceID:  c.session.ServerID,
+			},
+			CollectionID: collectionID,
+		}
+		return c.proxyClientManager.InvalidateCollectionMetaCache(ctx, &req, opts...)
+	}
+
+	// if only collNames are specified, invalidate the collection meta cache with the specified collectionName
+	for _, collName := range collNames {
+		req := proxypb.InvalidateCollMetaCacheRequest{
+			Base: &commonpb.MsgBase{
+				MsgType:   0, //TODO, msg type
+				MsgID:     0, //TODO, msg id
+				Timestamp: ts,
+				SourceID:  c.session.ServerID,
+			},
+			CollectionName: collName,
+		}
+		err := c.proxyClientManager.InvalidateCollectionMetaCache(ctx, &req, opts...)
+		if err != nil {
+			// TODO: try to expire all or directly return err?
+			return err
+		}
+	}
+	return nil
+}

internal/rootcoord/expire_cache_test.go

+package rootcoord
+
+import (
+	"testing"
+
+	"github.com/milvus-io/milvus/api/commonpb"
+
+	"github.com/milvus-io/milvus/internal/proto/proxypb"
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_expireCacheConfig_apply(t *testing.T) {
+	c := defaultExpireCacheConfig()
+	req := &proxypb.InvalidateCollMetaCacheRequest{}
+	c.apply(req)
+	assert.Nil(t, req.GetBase())
+	opt := expireCacheWithDropFlag()
+	opt(&c)
+	c.apply(req)
+	assert.Equal(t, commonpb.MsgType_DropCollection, req.GetBase().GetMsgType())
+}

internal/rootcoord/meta_table.go

@@ -248,7 +248,7 @@ func (mt *MetaTable) getCollectionByIDInternal(ctx context.Context, collectionID
 	var err error
 
 	coll, ok := mt.collID2Meta[collectionID]
-	if !ok || !coll.Available() || coll.CreateTime > ts {
+	if !ok || coll == nil || !coll.Available() || coll.CreateTime > ts {
 		// travel meta information from catalog.
 		ctx1 := contextutil.WithTenantID(ctx, Params.CommonCfg.ClusterName)
 		coll, err = mt.catalog.GetCollectionByID(ctx1, collectionID, ts)
@@ -257,9 +257,9 @@ func (mt *MetaTable) getCollectionByIDInternal(ctx context.Context, collectionID
 		}
 	}
 
-	if !coll.Available() {
+	if coll == nil || !coll.Available() {
 		// use coll.Name to match error message of regression. TODO: remove this after error code is ready.
-		return nil, fmt.Errorf("can't find collection: %s", coll.Name)
+		return nil, common.NewCollectionNotExistError(fmt.Sprintf("can't find collection: %s", coll.Name))
 	}
 
 	clone := coll.Clone()
@@ -296,7 +296,14 @@ func (mt *MetaTable) GetCollectionByName(ctx context.Context, collectionName str
 		return nil, err
 	}
 	if !coll.Available() {
-		return nil, fmt.Errorf("can't find collection: %s", collectionName)
+		return nil, common.NewCollectionNotExistError(fmt.Sprintf("can't find collection: %s", collectionName))
+	}
+	partitions := coll.Partitions
+	coll.Partitions = nil
+	for _, partition := range partitions {
+		if partition.Available() {
+			coll.Partitions = append(coll.Partitions, partition.Clone())
+		}
 	}
 	return coll, nil
 }

internal/rootcoord/meta_table_test.go

@@ -17,6 +17,7 @@
 package rootcoord
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -25,6 +26,13 @@ import (
 	"testing"
 	"time"
 
+	pb "github.com/milvus-io/milvus/internal/proto/etcdpb"
+
+	"github.com/milvus-io/milvus/internal/metastore/mocks"
+	"github.com/stretchr/testify/mock"
+
+	"github.com/milvus-io/milvus/internal/metastore/model"
+
 	"github.com/milvus-io/milvus/internal/util/funcutil"
 
 	"github.com/milvus-io/milvus/api/commonpb"
@@ -676,3 +684,160 @@ func TestRbacListUserRole(t *testing.T) {
 	assert.Nil(t, err)
 	assert.Equal(t, 4, len(userRoles))
 }
+
+func TestMetaTable_getCollectionByIDInternal(t *testing.T) {
+	t.Run("failed to get from catalog", func(t *testing.T) {
+		catalog := mocks.NewRootCoordCatalog(t)
+		catalog.On("GetCollectionByID",
+			mock.Anything, // context.Context
+			mock.AnythingOfType("int64"),
+			mock.AnythingOfType("uint64"),
+		).Return(nil, errors.New("error mock GetCollectionByID"))
+		meta := &MetaTable{
+			catalog:     catalog,
+			collID2Meta: map[typeutil.UniqueID]*model.Collection{},
+		}
+		ctx := context.Background()
+		_, err := meta.getCollectionByIDInternal(ctx, 100, 101)
+		assert.Error(t, err)
+	})
+
+	t.Run("collection not available", func(t *testing.T) {
+		catalog := mocks.NewRootCoordCatalog(t)
+		catalog.On("GetCollectionByID",
+			mock.Anything, // context.Context
+			mock.AnythingOfType("int64"),
+			mock.AnythingOfType("uint64"),
+		).Return(&model.Collection{State: pb.CollectionState_CollectionDropped}, nil)
+		meta := &MetaTable{
+			catalog:     catalog,
+			collID2Meta: map[typeutil.UniqueID]*model.Collection{},
+		}
+		ctx := context.Background()
+		_, err := meta.getCollectionByIDInternal(ctx, 100, 101)
+		assert.Error(t, err)
+		assert.True(t, common.IsCollectionNotExistError(err))
+	})
+
+	t.Run("normal case, filter unavailable partitions", func(t *testing.T) {
+		Params.InitOnce()
+		meta := &MetaTable{
+			collID2Meta: map[typeutil.UniqueID]*model.Collection{
+				100: {
+					State:      pb.CollectionState_CollectionCreated,
+					CreateTime: 99,
+					Partitions: []*model.Partition{
+						{PartitionID: 11, PartitionName: Params.CommonCfg.DefaultPartitionName, State: pb.PartitionState_PartitionCreated},
+						{PartitionID: 22, PartitionName: "dropped", State: pb.PartitionState_PartitionDropped},
+					},
+				},
+			},
+		}
+		ctx := context.Background()
+		coll, err := meta.getCollectionByIDInternal(ctx, 100, 101)
+		assert.NoError(t, err)
+		assert.Equal(t, 1, len(coll.Partitions))
+		assert.Equal(t, UniqueID(11), coll.Partitions[0].PartitionID)
+		assert.Equal(t, Params.CommonCfg.DefaultPartitionName, coll.Partitions[0].PartitionName)
+	})
+}
+
+func TestMetaTable_GetCollectionByName(t *testing.T) {
+	t.Run("get by alias", func(t *testing.T) {
+		meta := &MetaTable{
+			collAlias2ID: map[string]typeutil.UniqueID{
+				"alias": 100,
+			},
+			collID2Meta: map[typeutil.UniqueID]*model.Collection{
+				100: {
+					State:      pb.CollectionState_CollectionCreated,
+					CreateTime: 99,
+					Partitions: []*model.Partition{
+						{PartitionID: 11, PartitionName: Params.CommonCfg.DefaultPartitionName, State: pb.PartitionState_PartitionCreated},
+						{PartitionID: 22, PartitionName: "dropped", State: pb.PartitionState_PartitionDropped},
+					},
+				},
+			},
+		}
+		ctx := context.Background()
+		coll, err := meta.GetCollectionByName(ctx, "alias", 101)
+		assert.NoError(t, err)
+		assert.Equal(t, 1, len(coll.Partitions))
+		assert.Equal(t, UniqueID(11), coll.Partitions[0].PartitionID)
+		assert.Equal(t, Params.CommonCfg.DefaultPartitionName, coll.Partitions[0].PartitionName)
+	})
+
+	t.Run("get by name", func(t *testing.T) {
+		meta := &MetaTable{
+			collName2ID: map[string]typeutil.UniqueID{
+				"name": 100,
+			},
+			collID2Meta: map[typeutil.UniqueID]*model.Collection{
+				100: {
+					State:      pb.CollectionState_CollectionCreated,
+					CreateTime: 99,
+					Partitions: []*model.Partition{
+						{PartitionID: 11, PartitionName: Params.CommonCfg.DefaultPartitionName, State: pb.PartitionState_PartitionCreated},
+						{PartitionID: 22, PartitionName: "dropped", State: pb.PartitionState_PartitionDropped},
+					},
+				},
+			},
+		}
+		ctx := context.Background()
+		coll, err := meta.GetCollectionByName(ctx, "name", 101)
+		assert.NoError(t, err)
+		assert.Equal(t, 1, len(coll.Partitions))
+		assert.Equal(t, UniqueID(11), coll.Partitions[0].PartitionID)
+		assert.Equal(t, Params.CommonCfg.DefaultPartitionName, coll.Partitions[0].PartitionName)
+	})
+
+	t.Run("failed to get from catalog", func(t *testing.T) {
+		catalog := mocks.NewRootCoordCatalog(t)
+		catalog.On("GetCollectionByName",
+			mock.Anything, // context.Context
+			mock.AnythingOfType("string"),
+			mock.AnythingOfType("uint64"),
+		).Return(nil, errors.New("error mock GetCollectionByName"))
+		meta := &MetaTable{catalog: catalog}
+		ctx := context.Background()
+		_, err := meta.GetCollectionByName(ctx, "name", 101)
+		assert.Error(t, err)
+	})
+
+	t.Run("collection not available", func(t *testing.T) {
+		catalog := mocks.NewRootCoordCatalog(t)
+		catalog.On("GetCollectionByName",
+			mock.Anything, // context.Context
+			mock.AnythingOfType("string"),
+			mock.AnythingOfType("uint64"),
+		).Return(&model.Collection{State: pb.CollectionState_CollectionDropped}, nil)
+		meta := &MetaTable{catalog: catalog}
+		ctx := context.Background()
+		_, err := meta.GetCollectionByName(ctx, "name", 101)
+		assert.Error(t, err)
+		assert.True(t, common.IsCollectionNotExistError(err))
+	})
+
+	t.Run("normal case, filter unavailable partitions", func(t *testing.T) {
+		catalog := mocks.NewRootCoordCatalog(t)
+		catalog.On("GetCollectionByName",
+			mock.Anything, // context.Context
+			mock.AnythingOfType("string"),
+			mock.AnythingOfType("uint64"),
+		).Return(&model.Collection{
+			State:      pb.CollectionState_CollectionCreated,
+			CreateTime: 99,
+			Partitions: []*model.Partition{
+				{PartitionID: 11, PartitionName: Params.CommonCfg.DefaultPartitionName, State: pb.PartitionState_PartitionCreated},
+				{PartitionID: 22, PartitionName: "dropped", State: pb.PartitionState_PartitionDropped},
+			},
+		}, nil)
+		meta := &MetaTable{catalog: catalog}
+		ctx := context.Background()
+		coll, err := meta.GetCollectionByName(ctx, "name", 101)
+		assert.NoError(t, err)
+		assert.Equal(t, 1, len(coll.Partitions))
+		assert.Equal(t, UniqueID(11), coll.Partitions[0].PartitionID)
+		assert.Equal(t, Params.CommonCfg.DefaultPartitionName, coll.Partitions[0].PartitionName)
+	})
+}

internal/rootcoord/proxy_client_manager.go

@@ -117,10 +117,16 @@ func (p *proxyClientManager) DelProxyClient(s *sessionutil.Session) {
 	log.Debug("remove proxy client", zap.String("proxy address", s.Address), zap.Int64("proxy id", s.ServerID))
 }
 
-func (p *proxyClientManager) InvalidateCollectionMetaCache(ctx context.Context, request *proxypb.InvalidateCollMetaCacheRequest) error {
+func (p *proxyClientManager) InvalidateCollectionMetaCache(ctx context.Context, request *proxypb.InvalidateCollMetaCacheRequest, opts ...expireCacheOpt) error {
 	p.lock.Lock()
 	defer p.lock.Unlock()
 
+	c := defaultExpireCacheConfig()
+	for _, opt := range opts {
+		opt(&c)
+	}
+	c.apply(request)
+
 	if len(p.proxyClient) == 0 {
 		log.Warn("proxy client is empty, InvalidateCollectionMetaCache will not send to any client")
 		return nil

internal/rootcoord/root_coord.go

@@ -256,40 +256,6 @@ func (c *Core) SetQueryCoord(s types.QueryCoord) error {
 	return nil
 }
 
-// ExpireMetaCache will call invalidate collection meta cache
-func (c *Core) ExpireMetaCache(ctx context.Context, collNames []string, collectionID UniqueID, ts typeutil.Timestamp) error {
-	// if collectionID is specified, invalidate all the collection meta cache with the specified collectionID and return
-	if collectionID != InvalidCollectionID {
-		req := proxypb.InvalidateCollMetaCacheRequest{
-			Base: &commonpb.MsgBase{
-				Timestamp: ts,
-				SourceID:  c.session.ServerID,
-			},
-			CollectionID: collectionID,
-		}
-		return c.proxyClientManager.InvalidateCollectionMetaCache(ctx, &req)
-	}
-
-	// if only collNames are specified, invalidate the collection meta cache with the specified collectionName
-	for _, collName := range collNames {
-		req := proxypb.InvalidateCollMetaCacheRequest{
-			Base: &commonpb.MsgBase{
-				MsgType:   0, //TODO, msg type
-				MsgID:     0, //TODO, msg id
-				Timestamp: ts,
-				SourceID:  c.session.ServerID,
-			},
-			CollectionName: collName,
-		}
-		err := c.proxyClientManager.InvalidateCollectionMetaCache(ctx, &req)
-		if err != nil {
-			// TODO: try to expire all or directly return err?
-			return err
-		}
-	}
-	return nil
-}
-
 // Register register rootcoord at etcd
 func (c *Core) Register() error {
 	c.session.Register()
@@ -828,6 +794,7 @@ func (c *Core) HasCollection(ctx context.Context, in *milvuspb.HasCollectionRequ
 	log.Info("received request to has collection")
 
 	_, err := c.meta.GetCollectionByName(ctx, in.GetCollectionName(), ts)
+	// TODO: what if err != nil && common.IsCollectionNotExistError == false, should we consider this RPC as failure?
 	has := err == nil
 
 	metrics.RootCoordDDLReqCounter.WithLabelValues("HasCollection", metrics.SuccessLabel).Inc()
@@ -894,6 +861,8 @@ func (c *Core) DescribeCollection(ctx context.Context, in *milvuspb.DescribeColl
 
 	coll, err := c.describeCollection(ctx, in)
 	if err != nil {
+		// TODO: check whether err indicates the collection not exist.
+
 		log.Error("failed to describe collection", zap.Error(err))
 		metrics.RootCoordDDLReqCounter.WithLabelValues("DescribeCollection", metrics.FailLabel).Inc()
 
@@ -1077,6 +1046,7 @@ func (c *Core) HasPartition(ctx context.Context, in *milvuspb.HasPartitionReques
 
 	coll, err := c.meta.GetCollectionByName(ctx, in.GetCollectionName(), ts)
 	if err != nil {
+		// TODO: check if err indicates collection not exist.
 		log.Error("failed to has partition", zap.Error(err))
 		metrics.RootCoordDDLReqCounter.WithLabelValues("HasPartition", metrics.FailLabel).Inc()
 		// TODO: use commonpb.ErrorCode_CollectionNotExists. SDK use commonpb.ErrorCode_UnexpectedError now.
@@ -1127,6 +1097,7 @@ func (c *Core) ShowPartitions(ctx context.Context, in *milvuspb.ShowPartitionsRe
 	}
 
 	if err != nil {
+		// TODO: check if err indicates collection not exist.
 		log.Error("failed to show partitions", zap.Error(err))
 		metrics.RootCoordDDLReqCounter.WithLabelValues("ShowPartitions", metrics.FailLabel).Inc()
 		// TODO: use commonpb.ErrorCode_CollectionNotExists. SDK use commonpb.ErrorCode_UnexpectedError now.

internal/rootcoord/step.go

@@ -149,10 +149,11 @@ type expireCacheStep struct {
 	collectionNames []string
 	collectionID    UniqueID
 	ts              Timestamp
+	opts            []expireCacheOpt
 }
 
 func (s *expireCacheStep) Execute(ctx context.Context) ([]nestedStep, error) {
-	err := s.core.ExpireMetaCache(ctx, s.collectionNames, s.collectionID, s.ts)
+	err := s.core.ExpireMetaCache(ctx, s.collectionNames, s.collectionID, s.ts, s.opts...)
 	return nil, err
 }