refactor: Remove balance constraints between channel and segment tasks (milvus-io#42177)

issue: milvus-io#42176

Remove the mutual exclusion constraints between channel and segment
balance tasks to allow them to run concurrently.

Changes include:
- Remove permitBalanceChannel() and permitBalanceSegment() methods from
RoundRobinBalancer
- Update ChannelLevelScoreBalancer, MultiTargetBalancer,
RowCountBasedBalancer, and ScoreBasedBalancer to remove constraint
checks
- Allow segment balance tasks to proceed even when channel balance tasks
are running
- Update test cases to reflect new behavior where balance tasks no
longer block each other

This change improves the efficiency of load balancing by removing
unnecessary coordination overhead between different types of balance
operations.

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

Author: weiliu1031

internal/querycoordv2/balance/balance.go

@@ -155,14 +155,6 @@ func (b *RoundRobinBalancer) BalanceReplica(ctx context.Context, replica *meta.R
 	return nil, nil
 }
 
-func (b *RoundRobinBalancer) permitBalanceChannel(collectionID int64) bool {
-	return b.scheduler.GetSegmentTaskNum(task.WithCollectionID2TaskFilter(collectionID), task.WithTaskTypeFilter(task.TaskTypeMove)) == 0
-}
-
-func (b *RoundRobinBalancer) permitBalanceSegment(collectionID int64) bool {
-	return b.scheduler.GetChannelTaskNum(task.WithCollectionID2TaskFilter(collectionID), task.WithTaskTypeFilter(task.TaskTypeMove)) == 0
-}
-
 func (b *RoundRobinBalancer) getNodes(nodes []int64) []*session.NodeInfo {
 	ret := make([]*session.NodeInfo, 0, len(nodes))
 	for _, n := range nodes {

internal/querycoordv2/balance/channel_level_score_balancer.go

@@ -135,19 +135,19 @@ func (b *ChannelLevelScoreBalancer) BalanceReplica(ctx context.Context, replica
 				zap.Any("available nodes", rwNodes),
 			)
 			// handle stopped nodes here, have to assign segments on stopping nodes to nodes with the smallest score
-			if b.permitBalanceChannel(replica.GetCollectionID()) && !streamingutil.IsStreamingServiceEnabled() {
+			if !streamingutil.IsStreamingServiceEnabled() {
 				channelPlans = append(channelPlans, b.genStoppingChannelPlan(ctx, replica, channelName, rwNodes, roNodes)...)
 			}
 
-			if len(channelPlans) == 0 && b.permitBalanceSegment(replica.GetCollectionID()) {
+			if len(channelPlans) == 0 {
 				segmentPlans = append(segmentPlans, b.genStoppingSegmentPlan(ctx, replica, channelName, rwNodes, roNodes)...)
 			}
 		} else {
-			if paramtable.Get().QueryCoordCfg.AutoBalanceChannel.GetAsBool() && b.permitBalanceChannel(replica.GetCollectionID()) && !streamingutil.IsStreamingServiceEnabled() {
+			if paramtable.Get().QueryCoordCfg.AutoBalanceChannel.GetAsBool() && !streamingutil.IsStreamingServiceEnabled() {
 				channelPlans = append(channelPlans, b.genChannelPlan(ctx, replica, channelName, rwNodes)...)
 			}
 
-			if len(channelPlans) == 0 && b.permitBalanceSegment(replica.GetCollectionID()) {
+			if len(channelPlans) == 0 {
 				segmentPlans = append(segmentPlans, b.genSegmentPlan(ctx, br, replica, channelName, rwNodes)...)
 			}
 		}

internal/querycoordv2/balance/multi_target_balance.go

@@ -503,7 +503,7 @@ func (b *MultiTargetBalancer) balanceChannels(ctx context.Context, br *balanceRe
 		rwNodes, roNodes = replica.GetRWNodes(), replica.GetRONodes()
 	}
 
-	if len(rwNodes) == 0 || !b.permitBalanceChannel(replica.GetCollectionID()) {
+	if len(rwNodes) == 0 {
 		return nil
 	}
 
@@ -525,7 +525,7 @@ func (b *MultiTargetBalancer) balanceSegments(ctx context.Context, replica *meta
 	rwNodes := replica.GetRWNodes()
 	roNodes := replica.GetRONodes()
 
-	if len(rwNodes) == 0 || !b.permitBalanceSegment(replica.GetCollectionID()) {
+	if len(rwNodes) == 0 {
 		return nil
 	}
 	// print current distribution before generating plans

internal/querycoordv2/balance/rowcount_based_balancer.go

@@ -213,7 +213,7 @@ func (b *RowCountBasedBalancer) balanceChannels(ctx context.Context, br *balance
 		rwNodes, roNodes = replica.GetRWNodes(), replica.GetRONodes()
 	}
 
-	if len(rwNodes) == 0 || !b.permitBalanceChannel(replica.GetCollectionID()) {
+	if len(rwNodes) == 0 {
 		return nil
 	}
 	if len(roNodes) != 0 {
@@ -234,7 +234,7 @@ func (b *RowCountBasedBalancer) balanceSegments(ctx context.Context, replica *me
 	rwNodes := replica.GetRWNodes()
 	roNodes := replica.GetRONodes()
 
-	if len(rwNodes) == 0 || !b.permitBalanceSegment(replica.GetCollectionID()) {
+	if len(rwNodes) == 0 {
 		return nil
 	}
 	// print current distribution before generating plans

internal/querycoordv2/balance/score_based_balancer.go

@@ -485,7 +485,7 @@ func (b *ScoreBasedBalancer) balanceChannels(ctx context.Context, br *balanceRep
 		rwNodes, roNodes = replica.GetRWNodes(), replica.GetRONodes()
 	}
 
-	if len(rwNodes) == 0 || !b.permitBalanceChannel(replica.GetCollectionID()) {
+	if len(rwNodes) == 0 {
 		return nil
 	}
 
@@ -515,9 +515,6 @@ func (b *ScoreBasedBalancer) balanceSegments(ctx context.Context, br *balanceRep
 		br.AddRecord(StrRecord("no rwNodes to balance"))
 		return nil
 	}
-	if !b.permitBalanceSegment(replica.GetCollectionID()) {
-		return nil
-	}
 	// print current distribution before generating plans
 	if len(roNodes) != 0 {
 		if !stoppingBalance {

internal/querycoordv2/balance/score_based_balancer_test.go

@@ -1203,14 +1203,14 @@ func (suite *ScoreBasedBalancerTestSuite) TestBalanceSegmentAndChannel() {
 	segmentPlans, _ := suite.getCollectionBalancePlans(balancer, collectionID)
 	suite.Equal(len(segmentPlans), 2)
 
-	// mock balance channel is executing, expect to generate 0 balance segment task
+	// mock balance channel is executing, expect to generate 2 balance segment task, balance segment won't be blocked by channel balance
 	suite.mockScheduler.ExpectedCalls = nil
 	suite.mockScheduler.EXPECT().GetChannelTaskNum(mock.Anything, mock.Anything).Return(1).Maybe()
 	suite.mockScheduler.EXPECT().GetSegmentTaskNum(mock.Anything, mock.Anything).Return(0).Maybe()
 	suite.mockScheduler.EXPECT().GetSegmentTaskDelta(mock.Anything, mock.Anything).Return(0).Maybe()
 	suite.mockScheduler.EXPECT().GetChannelTaskDelta(mock.Anything, mock.Anything).Return(0).Maybe()
 	segmentPlans, _ = suite.getCollectionBalancePlans(balancer, collectionID)
-	suite.Equal(len(segmentPlans), 0)
+	suite.Equal(len(segmentPlans), 2)
 
 	// set unbalance channel distribution
 	balancer.dist.ChannelDistManager.Update(1, []*meta.DmChannel{
@@ -1228,14 +1228,14 @@ func (suite *ScoreBasedBalancerTestSuite) TestBalanceSegmentAndChannel() {
 	_, channelPlans := suite.getCollectionBalancePlans(balancer, collectionID)
 	suite.Equal(len(channelPlans), 2)
 
-	// mock balance channel is executing, expect to generate 0 balance segment task
+	// mock balance channel is executing, expect to generate 2 balance segment task, balance segment won't be blocked by channel balance
 	suite.mockScheduler.ExpectedCalls = nil
 	suite.mockScheduler.EXPECT().GetChannelTaskNum(mock.Anything, mock.Anything).Return(0).Maybe()
 	suite.mockScheduler.EXPECT().GetSegmentTaskNum(mock.Anything, mock.Anything).Return(1).Maybe()
 	suite.mockScheduler.EXPECT().GetSegmentTaskDelta(mock.Anything, mock.Anything).Return(0).Maybe()
 	suite.mockScheduler.EXPECT().GetChannelTaskDelta(mock.Anything, mock.Anything).Return(0).Maybe()
 	_, channelPlans = suite.getCollectionBalancePlans(balancer, collectionID)
-	suite.Equal(len(channelPlans), 0)
+	suite.Equal(len(channelPlans), 2)
 }
 
 func (suite *ScoreBasedBalancerTestSuite) TestBalanceChannelOnMultiCollections() {