fix: make sure balance candidate is alway pushed back (#29702)

See also #29699

Querycoord panicked when tried to pop from an empty heap. We assume the
heap shall not be empty, but in some branch, the candidate is never
pushed back.

This PR put pop & push in a closure and adds a defer call to push item
back.

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
pull/29686/head
congqixia 2024-01-05 10:08:47 +08:00 committed by GitHub
parent dc6a6a50fa
commit 3626f49025
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 21 additions and 18 deletions

View File

@ -77,27 +77,30 @@ func (b *ScoreBasedBalancer) AssignSegment(collectionID int64, segments []*meta.
plans := make([]SegmentAssignPlan, 0, len(segments))
for _, s := range segments {
// for each segment, pick the node with the least score
targetNode := queue.pop().(*nodeItem)
priorityChange := b.calculateSegmentScore(s)
func(s *meta.Segment) {
// for each segment, pick the node with the least score
targetNode := queue.pop().(*nodeItem)
// make sure candidate is always push back
defer queue.push(targetNode)
priorityChange := b.calculateSegmentScore(s)
sourceNode := nodeItemsMap[s.Node]
// if segment's node exist, which means this segment comes from balancer. we should consider the benefit
// if the segment reassignment doesn't got enough benefit, we should skip this reassignment
if sourceNode != nil && !b.hasEnoughBenefit(sourceNode, targetNode, priorityChange) {
continue
}
sourceNode := nodeItemsMap[s.Node]
// if segment's node exist, which means this segment comes from balancer. we should consider the benefit
// if the segment reassignment doesn't got enough benefit, we should skip this reassignment
if sourceNode != nil && !b.hasEnoughBenefit(sourceNode, targetNode, priorityChange) {
return
}
plan := SegmentAssignPlan{
From: -1,
To: targetNode.nodeID,
Segment: s,
}
plans = append(plans, plan)
plan := SegmentAssignPlan{
From: -1,
To: targetNode.nodeID,
Segment: s,
}
plans = append(plans, plan)
// update the targetNode's score
targetNode.setPriority(targetNode.getPriority() + priorityChange)
queue.push(targetNode)
// update the targetNode's score
targetNode.setPriority(targetNode.getPriority() + priorityChange)
}(s)
}
return plans
}