From 3626f49025d93833be466373941e2d0d40c39619 Mon Sep 17 00:00:00 2001
From: congqixia <congqi.xia@zilliz.com>
Date: Fri, 5 Jan 2024 10:08:47 +0800
Subject: [PATCH] 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>
---
 .../balance/score_based_balancer.go           | 39 ++++++++++---------
 1 file changed, 21 insertions(+), 18 deletions(-)

diff --git a/internal/querycoordv2/balance/score_based_balancer.go b/internal/querycoordv2/balance/score_based_balancer.go
index dd74ad2d07..d9305a1ecd 100644
--- a/internal/querycoordv2/balance/score_based_balancer.go
+++ b/internal/querycoordv2/balance/score_based_balancer.go
@@ -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
 }