enable balance channel (#23227)

Signed-off-by: Wei Liu <wei.liu@zilliz.com>
This commit is contained in:
wei liu 2023-04-07 19:06:28 +08:00 committed by GitHub
parent c855ea3171
commit 9f127dae47
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 97 additions and 6 deletions

View File

@ -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
}

View File

@ -211,19 +211,55 @@ func (suite *RowCountBasedBalancerTestSuite) TestBalance() {
},
},
{
name: "already balanced",
nodes: []int64{1, 2},
notExistedNodes: []int64{10},
segmentCnts: []int{1, 2},
states: []session.State{session.NodeStateNormal, session.NodeStateNormal},
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{
1: {{SegmentInfo: &datapb.SegmentInfo{ID: 1, CollectionID: 1, NumOfRows: 30}, Node: 1}},
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{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: 11}},
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{},
},