@@ -77,6 +77,7 @@ func (suite *BalanceCheckerTestSuite) SetupTest() {
77
77
suite .meta = meta .NewMeta (idAllocator , store , suite .nodeMgr )
78
78
suite .broker = meta .NewMockBroker (suite .T ())
79
79
suite .scheduler = task .NewMockScheduler (suite .T ())
80
+ suite .scheduler .EXPECT ().Add (mock .Anything ).Return (nil ).Maybe ()
80
81
suite .targetMgr = meta .NewTargetManager (suite .broker , suite .meta )
81
82
82
83
suite .balancer = balance .NewMockBalancer (suite .T ())
@@ -326,8 +327,15 @@ func (suite *BalanceCheckerTestSuite) TestStoppingBalance() {
326
327
}
327
328
segPlans = append (segPlans , mockPlan )
328
329
suite .balancer .EXPECT ().BalanceReplica (mock .Anything , mock .Anything ).Return (segPlans , chanPlans )
329
- tasks := suite .checker .Check (context .TODO ())
330
- suite .Len (tasks , 1 )
330
+
331
+ tasks := make ([]task.Task , 0 )
332
+ suite .scheduler .ExpectedCalls = nil
333
+ suite .scheduler .EXPECT ().Add (mock .Anything ).RunAndReturn (func (task task.Task ) error {
334
+ tasks = append (tasks , task )
335
+ return nil
336
+ })
337
+ suite .checker .Check (context .TODO ())
338
+ suite .Len (tasks , 2 )
331
339
}
332
340
333
341
func (suite * BalanceCheckerTestSuite ) TestTargetNotReady () {
@@ -850,6 +858,156 @@ func (suite *BalanceCheckerTestSuite) TestHasUnbalancedCollectionFlag() {
850
858
"stoppingBalanceCollectionsCurrentRound should contain the collection when it has RO nodes" )
851
859
}
852
860
861
+ func (suite * BalanceCheckerTestSuite ) TestCheckBatchSizesAndMultiCollection () {
862
+ ctx := context .Background ()
863
+
864
+ // Set up nodes
865
+ nodeID1 , nodeID2 := int64 (1 ), int64 (2 )
866
+ suite .nodeMgr .Add (session .NewNodeInfo (session.ImmutableNodeInfo {
867
+ NodeID : nodeID1 ,
868
+ Address : "localhost" ,
869
+ Hostname : "localhost" ,
870
+ }))
871
+ suite .nodeMgr .Add (session .NewNodeInfo (session.ImmutableNodeInfo {
872
+ NodeID : nodeID2 ,
873
+ Address : "localhost" ,
874
+ Hostname : "localhost" ,
875
+ }))
876
+ suite .checker .meta .ResourceManager .HandleNodeUp (ctx , nodeID1 )
877
+ suite .checker .meta .ResourceManager .HandleNodeUp (ctx , nodeID2 )
878
+
879
+ // Create 3 collections
880
+ for i := 1 ; i <= 3 ; i ++ {
881
+ cid := int64 (i )
882
+ replicaID := int64 (100 + i )
883
+
884
+ collection := utils .CreateTestCollection (cid , int32 (replicaID ))
885
+ collection .Status = querypb .LoadStatus_Loaded
886
+ replica := utils .CreateTestReplica (replicaID , cid , []int64 {})
887
+ mutableReplica := replica .CopyForWrite ()
888
+ mutableReplica .AddRWNode (nodeID1 )
889
+ mutableReplica .AddRONode (nodeID2 )
890
+
891
+ suite .checker .meta .CollectionManager .PutCollection (ctx , collection )
892
+ suite .checker .meta .ReplicaManager .Put (ctx , mutableReplica .IntoReplica ())
893
+ }
894
+
895
+ // Mock target manager
896
+ mockTargetManager := meta .NewMockTargetManager (suite .T ())
897
+ suite .checker .targetMgr = mockTargetManager
898
+
899
+ // All collections have same row count for simplicity
900
+ mockTargetManager .EXPECT ().GetCollectionRowCount (mock .Anything , mock .Anything , mock .Anything ).Return (int64 (100 )).Maybe ()
901
+ mockTargetManager .EXPECT ().IsCurrentTargetReady (mock .Anything , mock .Anything ).Return (true ).Maybe ()
902
+ mockTargetManager .EXPECT ().IsNextTargetExist (mock .Anything , mock .Anything ).Return (true ).Maybe ()
903
+ mockTargetManager .EXPECT ().IsCurrentTargetExist (mock .Anything , mock .Anything , mock .Anything ).Return (true ).Maybe ()
904
+
905
+ // For each collection, return different segment plans
906
+ suite .balancer .EXPECT ().BalanceReplica (mock .Anything , mock .AnythingOfType ("*meta.Replica" )).RunAndReturn (
907
+ func (ctx context.Context , replica * meta.Replica ) ([]balance.SegmentAssignPlan , []balance.ChannelAssignPlan ) {
908
+ // Create 2 segment plans and 1 channel plan per replica
909
+ collID := replica .GetCollectionID ()
910
+ segPlans := make ([]balance.SegmentAssignPlan , 0 )
911
+ chanPlans := make ([]balance.ChannelAssignPlan , 0 )
912
+
913
+ // Create 2 segment plans
914
+ for j := 1 ; j <= 2 ; j ++ {
915
+ segID := collID * 100 + int64 (j )
916
+ segPlan := balance.SegmentAssignPlan {
917
+ Segment : utils .CreateTestSegment (segID , collID , 1 , 1 , 1 , "test-channel" ),
918
+ Replica : replica ,
919
+ From : nodeID1 ,
920
+ To : nodeID2 ,
921
+ }
922
+ segPlans = append (segPlans , segPlan )
923
+ }
924
+
925
+ // Create 1 channel plan
926
+ chanPlan := balance.ChannelAssignPlan {
927
+ Channel : & meta.DmChannel {
928
+ VchannelInfo : & datapb.VchannelInfo {
929
+ CollectionID : collID ,
930
+ ChannelName : "test-channel" ,
931
+ },
932
+ },
933
+ Replica : replica ,
934
+ From : nodeID1 ,
935
+ To : nodeID2 ,
936
+ }
937
+ chanPlans = append (chanPlans , chanPlan )
938
+
939
+ return segPlans , chanPlans
940
+ }).Maybe ()
941
+
942
+ // Add tasks to check batch size limits
943
+ var addedTasks []task.Task
944
+ suite .scheduler .ExpectedCalls = nil
945
+ suite .scheduler .EXPECT ().Add (mock .Anything ).RunAndReturn (func (t task.Task ) error {
946
+ addedTasks = append (addedTasks , t )
947
+ return nil
948
+ }).Maybe ()
949
+
950
+ // Test 1: Balance with multiple collections disabled
951
+ paramtable .Get ().Save (Params .QueryCoordCfg .AutoBalance .Key , "true" )
952
+ paramtable .Get ().Save (Params .QueryCoordCfg .EnableBalanceOnMultipleCollections .Key , "false" )
953
+ // Set batch sizes to large values to test single-collection case
954
+ paramtable .Get ().Save (Params .QueryCoordCfg .BalanceSegmentBatchSize .Key , "10" )
955
+ paramtable .Get ().Save (Params .QueryCoordCfg .BalanceChannelBatchSize .Key , "10" )
956
+
957
+ // Reset test state
958
+ suite .checker .stoppingBalanceCollectionsCurrentRound .Clear ()
959
+ suite .checker .autoBalanceTs = time.Time {} // Reset to trigger auto balance
960
+ addedTasks = nil
961
+
962
+ // Run the Check method
963
+ suite .checker .Check (ctx )
964
+
965
+ // Should have tasks for a single collection (2 segment tasks + 1 channel task)
966
+ suite .Equal (3 , len (addedTasks ), "Should have tasks for a single collection when multiple collections balance is disabled" )
967
+
968
+ // Test 2: Balance with multiple collections enabled
969
+ paramtable .Get ().Save (Params .QueryCoordCfg .EnableBalanceOnMultipleCollections .Key , "true" )
970
+
971
+ // Reset test state
972
+ suite .checker .autoBalanceTs = time.Time {}
973
+ suite .checker .stoppingBalanceCollectionsCurrentRound .Clear ()
974
+ addedTasks = nil
975
+
976
+ // Run the Check method
977
+ suite .checker .Check (ctx )
978
+
979
+ // Should have tasks for all collections (3 collections * (2 segment tasks + 1 channel task) = 9 tasks)
980
+ suite .Equal (9 , len (addedTasks ), "Should have tasks for all collections when multiple collections balance is enabled" )
981
+
982
+ // Test 3: Batch size limits
983
+ paramtable .Get ().Save (Params .QueryCoordCfg .BalanceSegmentBatchSize .Key , "2" )
984
+ paramtable .Get ().Save (Params .QueryCoordCfg .BalanceChannelBatchSize .Key , "1" )
985
+
986
+ // Reset test state
987
+ suite .checker .stoppingBalanceCollectionsCurrentRound .Clear ()
988
+ addedTasks = nil
989
+
990
+ // Run the Check method
991
+ suite .checker .Check (ctx )
992
+
993
+ // Should respect batch size limits: 2 segment tasks + 1 channel task = 3 tasks
994
+ suite .Equal (3 , len (addedTasks ), "Should respect batch size limits" )
995
+
996
+ // Count segment tasks and channel tasks
997
+ segmentTaskCount := 0
998
+ channelTaskCount := 0
999
+ for _ , t := range addedTasks {
1000
+ if _ , ok := t .(* task.SegmentTask ); ok {
1001
+ segmentTaskCount ++
1002
+ } else {
1003
+ channelTaskCount ++
1004
+ }
1005
+ }
1006
+
1007
+ suite .LessOrEqual (segmentTaskCount , 2 , "Should have at most 2 segment tasks due to batch size limit" )
1008
+ suite .LessOrEqual (channelTaskCount , 1 , "Should have at most 1 channel task due to batch size limit" )
1009
+ }
1010
+
853
1011
func TestBalanceCheckerSuite (t * testing.T ) {
854
1012
suite .Run (t , new (BalanceCheckerTestSuite ))
855
1013
}
0 commit comments