From a6d9eb7f206f8e9055b508cd701705e2cfe61d8b Mon Sep 17 00:00:00 2001
From: congqixia <congqi.xia@zilliz.com>
Date: Sun, 18 Feb 2024 17:24:50 +0800
Subject: [PATCH] fix: Remove balance plan of which From, To nodes are same
 when merging (#30634)

See also #30627

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
---
 .../balance/multi_target_balance.go           | 34 ++++++++++++-------
 .../balance/multi_target_balancer_test.go     |  2 +-
 2 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/internal/querycoordv2/balance/multi_target_balance.go b/internal/querycoordv2/balance/multi_target_balance.go
index bfd196cb77..ba2a5141d2 100644
--- a/internal/querycoordv2/balance/multi_target_balance.go
+++ b/internal/querycoordv2/balance/multi_target_balance.go
@@ -16,6 +16,7 @@ import (
 	"github.com/milvus-io/milvus/internal/querycoordv2/task"
 	"github.com/milvus-io/milvus/pkg/log"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
+	"github.com/milvus-io/milvus/pkg/util/typeutil"
 )
 
 func init() {
@@ -194,21 +195,30 @@ func (g *basePlanGenerator) calClusterCost(replicaNodeSegments, globalNodeSegmen
 func (g *basePlanGenerator) mergePlans(curr []SegmentAssignPlan, inc []SegmentAssignPlan) []SegmentAssignPlan {
 	// merge plans with the same segment
 	// eg, plan1 is move segment1 from node1 to node2, plan2 is move segment1 from node2 to node3
-	// we should merge plan1 and plan2 to one plan, which is move segment1 from node1 to node3
-	for _, p := range inc {
-		has := false
-		for i := 0; i < len(curr); i++ {
-			if curr[i].Segment.GetID() == p.Segment.GetID() && curr[i].To == p.From {
-				curr[i].To = p.To
-				has = true
-				break
-			}
+	// we should merge plan1 and plan2 to one plan, which is move segment1 from node1 to node2
+	result := make([]SegmentAssignPlan, 0, len(curr)+len(inc))
+	processed := typeutil.NewSet[int]()
+	for _, p := range curr {
+		newPlan, idx, has := lo.FindIndexOf(inc, func(newPlan SegmentAssignPlan) bool {
+			return newPlan.Segment.GetID() == p.Segment.GetID() && newPlan.From == p.To
+		})
+
+		if has {
+			processed.Insert(idx)
+			p.To = newPlan.To
 		}
-		if !has {
-			curr = append(curr, p)
+		// in case of generator 1 move segment from node 1 to node 2 and generator 2 move segment back
+		if p.From != p.To {
+			result = append(result, p)
 		}
 	}
-	return curr
+
+	// add not merged inc plans
+	result = append(result, lo.Filter(inc, func(_ SegmentAssignPlan, idx int) bool {
+		return !processed.Contain(idx)
+	})...)
+
+	return result
 }
 
 type rowCountBasedPlanGenerator struct {
diff --git a/internal/querycoordv2/balance/multi_target_balancer_test.go b/internal/querycoordv2/balance/multi_target_balancer_test.go
index f72bc87e5b..bdf837e3d3 100644
--- a/internal/querycoordv2/balance/multi_target_balancer_test.go
+++ b/internal/querycoordv2/balance/multi_target_balancer_test.go
@@ -329,7 +329,7 @@ func (suite *MultiTargetBalancerTestSuite) TestPlanNoConflict() {
 	totalCount := 0
 	// 10 nodes, at most 100 segments, at most 1000 rows
 	for i := 0; i < 10; i++ {
-		segNum := rand.Intn(100)
+		segNum := rand.Intn(99) + 1
 		for j := 0; j < segNum; j++ {
 			rowCount := rand.Intn(1000)
 			nodeSegments[int64(i)] = append(nodeSegments[int64(i)], &meta.Segment{