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