fix: ReleasePartition cause delegator unserviceable.

issue: milvus-io#42098 milvus-io#42404
related to: #milvus-io#42009 milvus-io#41937

Implement new method to handle partition removal from next target without
directly modifying current target.

Changes include:
- Add RemovePartitionFromNextTarget method and deprecate RemovePartition
- Update target_observer to use new method for ReleasePartition operations
- Add unit tests and mock methods for new functionality

This ensures that all changes to next target will propagates to
delegator's query view.

Signed-off-by: Wei Liu <[email protected]>

Author: weiliu1031

internal/querycoordv2/meta/mock_target_manager.go

@@ -54,6 +54,7 @@ type TargetManagerInterface interface {
 	UpdateCollectionNextTarget(ctx context.Context, collectionID int64) error
 	RemoveCollection(ctx context.Context, collectionID int64)
 	RemovePartition(ctx context.Context, collectionID int64, partitionIDs ...int64)
+	RemovePartitionFromNextTarget(ctx context.Context, collectionID int64, partitionIDs ...int64)
 	GetGrowingSegmentsByCollection(ctx context.Context, collectionID int64, scope TargetScope) typeutil.UniqueSet
 	GetGrowingSegmentsByChannel(ctx context.Context, collectionID int64, channelName string, scope TargetScope) typeutil.UniqueSet
 	GetSealedSegmentsByCollection(ctx context.Context, collectionID int64, scope TargetScope) map[int64]*datapb.SegmentInfo
@@ -229,6 +230,7 @@ func (mgr *TargetManager) RemoveCollection(ctx context.Context, collectionID int
 
 // RemovePartition removes all segment in the given partition,
 // NOTE: this doesn't remove any channel even the given one is the only partition
+// Deprecated: use RemovePartitionFromNextTarget instead @weiliu1031
 func (mgr *TargetManager) RemovePartition(ctx context.Context, collectionID int64, partitionIDs ...int64) {
 	log := log.With(zap.Int64("collectionID", collectionID),
 		zap.Int64s("PartitionIDs", partitionIDs))
@@ -266,6 +268,32 @@ func (mgr *TargetManager) RemovePartition(ctx context.Context, collectionID int6
 	}
 }
 
+// remove partition from next target
+// NOTE: don't edit current target directly, it will be updated by target observer, which push the new next target as current target
+// need the full progress to update next target to current target, so the query view on delegator could be updated when current target is updated
+func (mgr *TargetManager) RemovePartitionFromNextTarget(ctx context.Context, collectionID int64, partitionIDs ...int64) {
+	log := log.With(zap.Int64("collectionID", collectionID),
+		zap.Int64s("PartitionIDs", partitionIDs))
+
+	partitionSet := typeutil.NewUniqueSet(partitionIDs...)
+
+	log.Info("remove partition from next target")
+	oleNextTarget := mgr.next.getCollectionTarget(collectionID)
+	if oleNextTarget != nil {
+		newTarget := mgr.removePartitionFromCollectionTarget(oleNextTarget, partitionSet)
+		if newTarget != nil {
+			mgr.next.updateCollectionTarget(collectionID, newTarget)
+			log.Info("finish to remove partition from next target for collection",
+				zap.Int64s("segments", newTarget.GetAllSegmentIDs()),
+				zap.Strings("channels", newTarget.GetAllDmChannelNames()))
+		} else {
+			log.Info("all partitions have been released, release the collection current target now")
+			mgr.current.removeCollectionTarget(collectionID)
+			mgr.next.removeCollectionTarget(collectionID)
+		}
+	}
+}
+
 func (mgr *TargetManager) removePartitionFromCollectionTarget(oldTarget *CollectionTarget, partitionSet typeutil.UniqueSet) *CollectionTarget {
 	segments := make(map[int64]*datapb.SegmentInfo)
 	for _, segment := range oldTarget.GetAllSegments() {

internal/querycoordv2/meta/target_manager.go

@@ -243,9 +243,6 @@ func (suite *TargetManagerSuite) TestUpdateNextTarget() {
 	suite.broker.EXPECT().GetRecoveryInfoV2(mock.Anything, collectionID).Return(nextTargetChannels, nextTargetSegments, nil)
 	err := suite.mgr.UpdateCollectionNextTarget(ctx, collectionID)
 	suite.NoError(err)
-
-	err = suite.mgr.UpdateCollectionNextTarget(ctx, collectionID)
-	suite.NoError(err)
 }
 
 func (suite *TargetManagerSuite) TestRemovePartition() {
@@ -263,6 +260,27 @@ func (suite *TargetManagerSuite) TestRemovePartition() {
 	suite.assertChannels([]string{}, suite.mgr.GetDmChannelsByCollection(ctx, collectionID, CurrentTarget))
 }
 
+func (suite *TargetManagerSuite) TestRemovePartitionFromNextTarget() {
+	ctx := suite.ctx
+	collectionID := int64(1000)
+	ret := suite.mgr.UpdateCollectionCurrentTarget(ctx, collectionID)
+	suite.True(ret)
+
+	err := suite.mgr.UpdateCollectionNextTarget(ctx, collectionID)
+	suite.NoError(err)
+
+	suite.assertSegments(suite.getAllSegment(collectionID, suite.partitions[collectionID]), suite.mgr.GetSealedSegmentsByCollection(ctx, collectionID, NextTarget))
+	suite.assertChannels(suite.channels[collectionID], suite.mgr.GetDmChannelsByCollection(ctx, collectionID, NextTarget))
+	suite.assertSegments(suite.getAllSegment(collectionID, suite.partitions[collectionID]), suite.mgr.GetSealedSegmentsByCollection(ctx, collectionID, CurrentTarget))
+	suite.assertChannels(suite.channels[collectionID], suite.mgr.GetDmChannelsByCollection(ctx, collectionID, CurrentTarget))
+
+	suite.mgr.RemovePartitionFromNextTarget(ctx, collectionID, 100)
+	suite.assertSegments([]int64{3, 4}, suite.mgr.GetSealedSegmentsByCollection(ctx, collectionID, NextTarget))
+	suite.assertChannels(suite.channels[collectionID], suite.mgr.GetDmChannelsByCollection(ctx, collectionID, NextTarget))
+	suite.assertSegments(suite.getAllSegment(collectionID, suite.partitions[collectionID]), suite.mgr.GetSealedSegmentsByCollection(ctx, collectionID, CurrentTarget))
+	suite.assertChannels(suite.channels[collectionID], suite.mgr.GetDmChannelsByCollection(ctx, collectionID, CurrentTarget))
+}
+
 func (suite *TargetManagerSuite) TestRemoveCollection() {
 	ctx := suite.ctx
 	collectionID := int64(1000)

internal/querycoordv2/meta/target_manager_test.go

@@ -225,7 +225,7 @@ func (ob *TargetObserver) schedule(ctx context.Context) {
 				ob.keylocks.Unlock(req.CollectionID)
 				req.Notifier <- nil
 			case ReleasePartition:
-				ob.targetMgr.RemovePartition(ctx, req.CollectionID, req.PartitionIDs...)
+				ob.targetMgr.RemovePartitionFromNextTarget(ctx, req.CollectionID, req.PartitionIDs...)
 				req.Notifier <- nil
 			}
 			log.Info("manually trigger update target done",