enable balance channel (#23227)

Signed-off-by: NWei Liu <wei.liu@zilliz.com>

internal/querycoordv2/balance/rowcount_based_balancer.go

@@ -268,6 +268,61 @@ func (b *RowCountBasedBalancer) getChannelPlan(replica *meta.Replica, onlineNode
 		}
 		channelPlans = append(channelPlans, plans...)
 	}
+
+	if len(channelPlans) == 0 && len(onlineNodes) > 1 {
+		// start to balance channels on all available nodes
+		channels := b.dist.ChannelDistManager.GetByCollection(replica.CollectionID)
+		channelsOnNode := lo.GroupBy(channels, func(channel *meta.DmChannel) int64 { return channel.Node })
+
+		nodes := replica.GetNodes()
+		getChannelNum := func(node int64) int {
+			if channelsOnNode[node] == nil {
+				return 0
+			}
+			return len(channelsOnNode[node])
+		}
+		sort.Slice(nodes, func(i, j int) bool { return getChannelNum(nodes[i]) < getChannelNum(nodes[j]) })
+
+		start := int64(0)
+		end := int64(len(nodes) - 1)
+		averageChannel := len(channels) / len(onlineNodes)
+		if averageChannel == 0 || getChannelNum(nodes[start]) >= getChannelNum(nodes[end]) {
+			return channelPlans
+		}
+
+		for start < end {
+			sourceNode := nodes[start]
+			targetNode := nodes[end]
+
+			// remove channel from end node
+			selectChannel := channelsOnNode[targetNode][0]
+			channelsOnNode[targetNode] = channelsOnNode[targetNode][1:]
+
+			// add channel to start node
+			if channelsOnNode[sourceNode] == nil {
+				channelsOnNode[sourceNode] = make([]*meta.DmChannel, 0)
+			}
+			channelsOnNode[sourceNode] = append(channelsOnNode[sourceNode], selectChannel)
+
+			// generate channel plan
+			plan := ChannelAssignPlan{
+				Channel:   selectChannel,
+				From:      targetNode,
+				To:        sourceNode,
+				ReplicaID: replica.ID,
+				Weight:    GetWeight(1),
+			}
+			channelPlans = append(channelPlans, plan)
+			for end > 0 && getChannelNum(nodes[end]) <= averageChannel {
+				end--
+			}
+
+			for start < end && getChannelNum(nodes[start]) >= averageChannel {
+				start++
+			}
+		}
+
+	}
 	return channelPlans
 }
 

internal/querycoordv2/balance/rowcount_based_balancer_test.go

@@ -210,20 +210,56 @@ func (suite *RowCountBasedBalancerTestSuite) TestBalance() {
 				{Channel: &meta.DmChannel{VchannelInfo: &datapb.VchannelInfo{CollectionID: 1, ChannelName: "v3"}, Node: 3}, From: 3, To: 1, ReplicaID: 1, Weight: weightHigh},
 			},
 		},
+		{
+			name:        "balance channel",
+			nodes:       []int64{2, 3},
+			segmentCnts: []int{2, 2},
+			states:      []session.State{session.NodeStateNormal, session.NodeStateNormal},
+			shouldMock:  true,
+			distributions: map[int64][]*meta.Segment{
+				2: {
+					{SegmentInfo: &datapb.SegmentInfo{ID: 2, CollectionID: 1, NumOfRows: 20}, Node: 2},
+					{SegmentInfo: &datapb.SegmentInfo{ID: 3, CollectionID: 1, NumOfRows: 30}, Node: 2},
+				},
+				3: {
+					{SegmentInfo: &datapb.SegmentInfo{ID: 4, CollectionID: 1, NumOfRows: 10}, Node: 3},
+					{SegmentInfo: &datapb.SegmentInfo{ID: 5, CollectionID: 1, NumOfRows: 10}, Node: 3},
+				},
+			},
+			distributionChannels: map[int64][]*meta.DmChannel{
+				2: {
+					{VchannelInfo: &datapb.VchannelInfo{CollectionID: 1, ChannelName: "v2"}, Node: 2},
+					{VchannelInfo: &datapb.VchannelInfo{CollectionID: 1, ChannelName: "v3"}, Node: 2},
+				},
+				3: {},
+			},
+			expectPlans: []SegmentAssignPlan{},
+			expectChannelPlans: []ChannelAssignPlan{
+				{Channel: &meta.DmChannel{VchannelInfo: &datapb.VchannelInfo{CollectionID: 1, ChannelName: "v2"}, Node: 2}, From: 2, To: 3, ReplicaID: 1, Weight: weightHigh},
+			},
+		},
 		{
 			name:            "already balanced",
-			nodes:           []int64{1, 2},
+			nodes:           []int64{11, 22},
 			notExistedNodes: []int64{10},
 			segmentCnts:     []int{1, 2},
 			states:          []session.State{session.NodeStateNormal, session.NodeStateNormal},
 			distributions: map[int64][]*meta.Segment{
-				1: {{SegmentInfo: &datapb.SegmentInfo{ID: 1, CollectionID: 1, NumOfRows: 30}, Node: 1}},
+				1: {{SegmentInfo: &datapb.SegmentInfo{ID: 1, CollectionID: 1, NumOfRows: 30}, Node: 11}},
 				2: {
-					{SegmentInfo: &datapb.SegmentInfo{ID: 2, CollectionID: 1, NumOfRows: 20}, Node: 2},
-					{SegmentInfo: &datapb.SegmentInfo{ID: 3, CollectionID: 1, NumOfRows: 30}, Node: 2},
+					{SegmentInfo: &datapb.SegmentInfo{ID: 2, CollectionID: 1, NumOfRows: 20}, Node: 22},
+					{SegmentInfo: &datapb.SegmentInfo{ID: 3, CollectionID: 1, NumOfRows: 30}, Node: 22},
 				},
 				10: {{SegmentInfo: &datapb.SegmentInfo{ID: 4, CollectionID: 1, NumOfRows: 30}, Node: 10}},
 			},
+			// distributionChannels: map[int64][]*meta.DmChannel{
+			// 	1: {
+			// 		{VchannelInfo: &datapb.VchannelInfo{CollectionID: 1, ChannelName: "v2"}, Node: 1},
+			// 	},
+			// 	2: {
+			// 		{VchannelInfo: &datapb.VchannelInfo{CollectionID: 1, ChannelName: "v3"}, Node: 2},
+			// 	},
+			// },
 			expectPlans:        []SegmentAssignPlan{},
 			expectChannelPlans: []ChannelAssignPlan{},
 		},