From a308331b8169af892717dc67661a14c4ec327c37 Mon Sep 17 00:00:00 2001 From: wei liu Date: Thu, 4 Dec 2025 10:19:11 +0800 Subject: [PATCH] fix: Set replica field in balance plans to prevent panic (#45722) issue: #45598 The MultiTargetBalancer was missing replica field assignment in the generated segment and channel plans, which caused panic during balance operations. This change ensures that all balance plans have the replica field properly set to fix the panic issue. Also refactored the balance test to extract common test logic into a reusable helper function and added a new integration test specifically for MultipleTargetBalancer policy. Signed-off-by: Wei Liu --- .../querycoordv2/balance/multi_target_balance.go | 6 +++++- tests/integration/balance/balance_test.go | 14 ++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/internal/querycoordv2/balance/multi_target_balance.go b/internal/querycoordv2/balance/multi_target_balance.go index e3ae7725d8..356fca53df 100644 --- a/internal/querycoordv2/balance/multi_target_balance.go +++ b/internal/querycoordv2/balance/multi_target_balance.go @@ -559,7 +559,11 @@ func (b *MultiTargetBalancer) genSegmentPlan(ctx context.Context, replica *meta. globalNodeSegments[node] = b.dist.SegmentDistManager.GetByFilter(meta.WithNodeID(node)) } - return b.genPlanByDistributions(nodeSegments, globalNodeSegments) + plans := b.genPlanByDistributions(nodeSegments, globalNodeSegments) + for i := range plans { + plans[i].Replica = replica + } + return plans } func (b *MultiTargetBalancer) genPlanByDistributions(nodeSegments, globalNodeSegments map[int64][]*meta.Segment) []SegmentAssignPlan { diff --git a/tests/integration/balance/balance_test.go b/tests/integration/balance/balance_test.go index 53f41c619f..a2a1d17a5b 100644 --- a/tests/integration/balance/balance_test.go +++ b/tests/integration/balance/balance_test.go @@ -170,6 +170,10 @@ func (s *BalanceTestSuit) initCollection(collectionName string, replica int, cha } func (s *BalanceTestSuit) TestBalanceOnSingleReplica() { + testBalanceOnSingleReplica(s) +} + +func testBalanceOnSingleReplica(s *BalanceTestSuit) { name := "test_balance_" + funcutil.GenRandomStr() s.initCollection(name, 1, 2, 2, 2000, 500) @@ -414,6 +418,16 @@ func (s *BalanceTestSuit) TestConcurrentBalanceChannelAndSegment() { s.Equal(int64(0), failCounter.Load()) } +func (s *BalanceTestSuit) TestMultiTargetBalancePolicy() { + // Set balance policy to MultipleTargetBalancer + revertGuard := s.Cluster.MustModifyMilvusConfig(map[string]string{ + paramtable.Get().QueryCoordCfg.Balancer.Key: "MultipleTargetBalancer", + }) + defer revertGuard() + + testBalanceOnSingleReplica(s) +} + func TestBalance(t *testing.T) { suite.Run(t, new(BalanceTestSuit)) }