fix: Fix collection leak in querynode

internal/querycoordv2/job/job_release.go

@@ -89,14 +89,14 @@ func (job *ReleaseCollectionJob) Execute() error {
 		return errors.Wrap(err, msg)
 	}
 
+	job.targetObserver.ReleaseCollection(req.GetCollectionID())
+	waitCollectionReleased(job.dist, job.checkerController, req.GetCollectionID())
+	cleanQueryNodesCollectionMetaCache(job.ctx, job.meta, job.cluster, req.GetCollectionID())
 	err = job.meta.ReplicaManager.RemoveCollection(req.GetCollectionID())
 	if err != nil {
 		msg := "failed to remove replicas"
 		log.Warn(msg, zap.Error(err))
 	}
-
-	job.targetObserver.ReleaseCollection(req.GetCollectionID())
-	waitCollectionReleased(job.dist, job.checkerController, req.GetCollectionID())
 	metrics.QueryCoordNumCollections.WithLabelValues().Dec()
 	metrics.QueryCoordNumPartitions.WithLabelValues().Sub(float64(len(toRelease)))
 	metrics.QueryCoordReleaseCount.WithLabelValues(metrics.TotalLabel).Inc()

internal/querycoordv2/job/job_test.go

@@ -45,11 +45,6 @@ import (
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
 )
 
-const (
-	defaultVecFieldID = 1
-	defaultIndexID    = 1
-)
-
 type JobSuite struct {
 	suite.Suite
 
@@ -141,7 +136,10 @@ func (suite *JobSuite) SetupSuite() {
 		Return(merr.Success(), nil)
 	suite.cluster.EXPECT().
 		ReleasePartitions(mock.Anything, mock.Anything, mock.Anything).
-		Return(merr.Success(), nil).Maybe()
+		Return(merr.Success(), nil)
+	suite.cluster.EXPECT().
+		ReleaseCollection(mock.Anything, mock.Anything, mock.Anything).
+		Return(merr.Success(), nil)
 
 	suite.proxyManager = proxyutil.NewMockProxyClientManager(suite.T())
 	suite.proxyManager.EXPECT().InvalidateCollectionMetaCache(mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
@@ -849,7 +847,6 @@ func (suite *JobSuite) TestReleaseCollection() {
 			suite.cluster,
 			suite.targetMgr,
 			suite.targetObserver,
-
 			suite.checkerController,
 		)
 		suite.scheduler.Add(job)

internal/querycoordv2/job/utils.go

@@ -147,3 +147,27 @@
 		}
 	}
 }
+
+func cleanQueryNodesCollectionMetaCache(ctx context.Context,
+	meta *meta.Meta,
+	cluster session.Cluster,
+	collection int64,
+) {
+	log := log.Ctx(ctx).With(zap.Int64("collection", collection))
+	replicas := meta.ReplicaManager.GetByCollection(collection)
+	releaseReq := &querypb.ReleaseCollectionRequest{
+		Base: &commonpb.MsgBase{
+			MsgType: commonpb.MsgType_ReleaseCollection,
+		},
+		CollectionID: collection,
+	}
+	for _, replica := range replicas {
+		for _, node := range replica.GetNodes() {
+			status, err := cluster.ReleaseCollection(ctx, node, releaseReq)
+			err = merr.CheckRPCCall(status, err)
+			if err != nil {
+				log.Warn("failed to release collection", zap.Int64("node", node), zap.Error(err))
+			}
+		}
+	}
+}

internal/querycoordv2/services.go

@@ -45,16 +45,6 @@ import (
 	"github.com/milvus-io/milvus/pkg/util/typeutil"
 )
 
-var (
-// ErrRemoveNodeFromRGFailed      = errors.New("failed to remove node from resource group")
-// ErrTransferNodeFailed          = errors.New("failed to transfer node between resource group")
-// ErrTransferReplicaFailed       = errors.New("failed to transfer replica between resource group")
-// ErrListResourceGroupsFailed    = errors.New("failed to list resource group")
-// ErrDescribeResourceGroupFailed = errors.New("failed to describe resource group")
-// ErrLoadUseWrongRG              = errors.New("load operation should use collection's resource group")
-// ErrLoadWithDefaultRG           = errors.New("load operation can't use default resource group and other resource group together")
-)
-
 func (s *Server) ShowCollections(ctx context.Context, req *querypb.ShowCollectionsRequest) (*querypb.ShowCollectionsResponse, error) {
 	log.Ctx(ctx).Info("show collections request received", zap.Int64s("collections", req.GetCollectionIDs()))
 

internal/querycoordv2/services_test.go

@@ -1007,6 +1007,8 @@ func (suite *ServiceSuite) TestReleaseCollection() {
 
 	suite.cluster.EXPECT().ReleasePartitions(mock.Anything, mock.Anything, mock.Anything).
 		Return(merr.Success(), nil)
+	suite.cluster.EXPECT().ReleaseCollection(mock.Anything, mock.Anything, mock.Anything).
+		Return(merr.Success(), nil)
 
 	// Test release all collections
 	for _, collection := range suite.collections {

internal/querycoordv2/session/cluster.go

@@ -48,6 +48,7 @@
 	ReleaseSegments(ctx context.Context, nodeID int64, req *querypb.ReleaseSegmentsRequest) (*commonpb.Status, error)
 	LoadPartitions(ctx context.Context, nodeID int64, req *querypb.LoadPartitionsRequest) (*commonpb.Status, error)
 	ReleasePartitions(ctx context.Context, nodeID int64, req *querypb.ReleasePartitionsRequest) (*commonpb.Status, error)
+	ReleaseCollection(ctx context.Context, nodeID int64, req *querypb.ReleaseCollectionRequest) (*commonpb.Status, error)
 	GetDataDistribution(ctx context.Context, nodeID int64, req *querypb.GetDataDistributionRequest) (*querypb.GetDataDistributionResponse, error)
 	GetMetrics(ctx context.Context, nodeID int64, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error)
 	SyncDistribution(ctx context.Context, nodeID int64, req *querypb.SyncDistributionRequest) (*commonpb.Status, error)
@@ -197,6 +198,20 @@
 	return status, err
 }
 
+func (c *QueryCluster) ReleaseCollection(ctx context.Context, nodeID int64, req *querypb.ReleaseCollectionRequest) (*commonpb.Status, error) {
+	var status *commonpb.Status
+	var err error
+	err1 := c.send(ctx, nodeID, func(cli types.QueryNodeClient) {
+		req := proto.Clone(req).(*querypb.ReleaseCollectionRequest)
+		req.Base.TargetID = nodeID
+		status, err = cli.ReleaseCollection(ctx, req)
+	})
+	if err1 != nil {
+		return nil, err1
+	}
+	return status, err
+}
+
 func (c *QueryCluster) GetDataDistribution(ctx context.Context, nodeID int64, req *querypb.GetDataDistributionRequest) (*querypb.GetDataDistributionResponse, error) {
 	var resp *querypb.GetDataDistributionResponse
 	var err error

internal/querycoordv2/session/cluster_test.go

@@ -140,6 +140,10 @@ func (suite *ClusterTestSuite) createDefaultMockServer() querypb.QueryNodeServer
 		mock.Anything,
 		mock.AnythingOfType("*querypb.ReleasePartitionsRequest"),
 	).Maybe().Return(succStatus, nil)
+	svr.EXPECT().ReleaseCollection(
+		mock.Anything,
+		mock.AnythingOfType("*querypb.ReleaseCollectionRequest"),
+	).Maybe().Return(succStatus, nil)
 	svr.EXPECT().GetDataDistribution(
 		mock.Anything,
 		mock.AnythingOfType("*querypb.GetDataDistributionRequest"),
@@ -193,6 +197,10 @@ func (suite *ClusterTestSuite) createFailedMockServer() querypb.QueryNodeServer
 		mock.Anything,
 		mock.AnythingOfType("*querypb.ReleasePartitionsRequest"),
 	).Maybe().Return(failStatus, nil)
+	svr.EXPECT().ReleaseCollection(
+		mock.Anything,
+		mock.AnythingOfType("*querypb.ReleaseCollectionRequest"),
+	).Maybe().Return(failStatus, nil)
 	svr.EXPECT().GetDataDistribution(
 		mock.Anything,
 		mock.AnythingOfType("*querypb.GetDataDistributionRequest"),
@@ -318,6 +326,22 @@ func (suite *ClusterTestSuite) TestLoadAndReleasePartitions() {
 	suite.Equal("unexpected error", status.GetReason())
 }
 
+func (suite *ClusterTestSuite) TestReleaseCollection() {
+	ctx := context.TODO()
+	status, err := suite.cluster.ReleaseCollection(ctx, 0, &querypb.ReleaseCollectionRequest{
+		Base: &commonpb.MsgBase{},
+	})
+	suite.NoError(err)
+	merr.Ok(status)
+
+	status, err = suite.cluster.ReleaseCollection(ctx, 1, &querypb.ReleaseCollectionRequest{
+		Base: &commonpb.MsgBase{},
+	})
+	suite.NoError(err)
+	suite.Equal(commonpb.ErrorCode_UnexpectedError, status.GetErrorCode())
+	suite.Equal("unexpected error", status.GetReason())
+}
+
 func (suite *ClusterTestSuite) TestGetDataDistribution() {
 	ctx := context.TODO()
 	resp, err := suite.cluster.GetDataDistribution(ctx, 0, &querypb.GetDataDistributionRequest{

internal/querynodev2/delegator/delegator_data_test.go

@@ -103,7 +103,7 @@ func (s *DelegatorDataSuite) TearDownSuite() {
 }
 
 func (s *DelegatorDataSuite) genNormalCollection() {
-	s.manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
+	s.manager.Collection.Put(s.collectionID, &schemapb.CollectionSchema{
 		Name: "TestCollection",
 		Fields: []*schemapb.FieldSchema{
 			{
@@ -158,7 +158,7 @@ func (s *DelegatorDataSuite) genNormalCollection() {
 }
 
 func (s *DelegatorDataSuite) genCollectionWithFunction() {
-	s.manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
+	s.manager.Collection.Put(s.collectionID, &schemapb.CollectionSchema{
 		Name: "TestCollection",
 		Fields: []*schemapb.FieldSchema{
 			{

internal/querynodev2/delegator/delegator_test.go

@@ -104,7 +104,7 @@ func (s *DelegatorSuite) SetupTest() {
 	}, nil)
 
 	// init schema
-	s.manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
+	s.manager.Collection.Put(s.collectionID, &schemapb.CollectionSchema{
 		Name: "TestCollection",
 		Fields: []*schemapb.FieldSchema{
 			{
@@ -181,7 +181,7 @@ func (s *DelegatorSuite) TearDownTest() {
 func (s *DelegatorSuite) TestCreateDelegatorWithFunction() {
 	s.Run("init function failed", func() {
 		manager := segments.NewManager()
-		manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
+		manager.Collection.Put(s.collectionID, &schemapb.CollectionSchema{
 			Name: "TestCollection",
 			Fields: []*schemapb.FieldSchema{
 				{
@@ -214,7 +214,7 @@ func (s *DelegatorSuite) TestCreateDelegatorWithFunction() {
 
 	s.Run("init function failed", func() {
 		manager := segments.NewManager()
-		manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
+		manager.Collection.Put(s.collectionID, &schemapb.CollectionSchema{
 			Name: "TestCollection",
 			Fields: []*schemapb.FieldSchema{
 				{

internal/querynodev2/delegator/delta_forward_test.go

@@ -91,7 +91,7 @@ func (s *StreamingForwardSuite) SetupTest() {
 	}, nil)
 
 	// init schema
-	s.manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
+	s.manager.Collection.Put(s.collectionID, &schemapb.CollectionSchema{
 		Name: "TestCollection",
 		Fields: []*schemapb.FieldSchema{
 			{

internal/querynodev2/local_worker_test.go

@@ -102,7 +102,7 @@ func (suite *LocalWorkerTestSuite) BeforeTest(suiteName, testName string) {
 		LoadType:     querypb.LoadType_LoadCollection,
 		CollectionID: suite.collectionID,
 	}
-	suite.node.manager.Collection.PutOrRef(suite.collectionID, collection.Schema(), suite.indexMeta, loadMata)
+	suite.node.manager.Collection.Put(suite.collectionID, collection.Schema(), suite.indexMeta, loadMata)
 	suite.worker = NewLocalWorker(suite.node)
 }