@@ -250,10 +250,10 @@ func (suite *RowCountBasedBalancerTestSuite) TestBalance() {
250
250
},
251
251
distributionChannels : map [int64 ][]* meta.DmChannel {
252
252
2 : {
253
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 2 },
253
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 2 , View : & meta. LeaderView { ID : 2 , CollectionID : 1 } },
254
254
},
255
255
3 : {
256
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v3" }, Node : 3 },
256
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v3" }, Node : 3 , View : & meta. LeaderView { ID : 3 , CollectionID : 1 } },
257
257
},
258
258
},
259
259
expectPlans : []SegmentAssignPlan {},
@@ -280,10 +280,10 @@ func (suite *RowCountBasedBalancerTestSuite) TestBalance() {
280
280
},
281
281
distributionChannels : map [int64 ][]* meta.DmChannel {
282
282
2 : {
283
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 2 },
283
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 2 , View : & meta. LeaderView { ID : 2 , CollectionID : 1 } },
284
284
},
285
285
1 : {
286
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v3" }, Node : 1 },
286
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v3" }, Node : 1 , View : & meta. LeaderView { ID : 3 , CollectionID : 1 } },
287
287
},
288
288
},
289
289
expectPlans : []SegmentAssignPlan {
@@ -301,8 +301,8 @@ func (suite *RowCountBasedBalancerTestSuite) TestBalance() {
301
301
distributions : map [int64 ][]* meta.Segment {},
302
302
distributionChannels : map [int64 ][]* meta.DmChannel {
303
303
2 : {
304
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 2 },
305
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v3" }, Node : 2 },
304
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 2 , View : & meta. LeaderView { ID : 2 , CollectionID : 1 } },
305
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v3" }, Node : 2 , View : & meta. LeaderView { ID : 3 , CollectionID : 1 } },
306
306
},
307
307
3 : {},
308
308
},
@@ -320,14 +320,14 @@ func (suite *RowCountBasedBalancerTestSuite) TestBalance() {
320
320
distributions : map [int64 ][]* meta.Segment {},
321
321
distributionChannels : map [int64 ][]* meta.DmChannel {
322
322
1 : {
323
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v1" }, Node : 1 },
324
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 1 },
323
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v1" }, Node : 1 , View : & meta. LeaderView { ID : 1 , CollectionID : 1 , Channel : "v1" , Status : & querypb. LeaderViewStatus { Serviceable : true }} },
324
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 1 , View : & meta. LeaderView { ID : 2 , CollectionID : 1 , Channel : "v2" , Status : & querypb. LeaderViewStatus { Serviceable : true }} },
325
325
},
326
326
2 : {
327
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v3" }, Node : 2 },
327
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v3" }, Node : 2 , View : & meta. LeaderView { ID : 3 , CollectionID : 1 } },
328
328
},
329
329
3 : {
330
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v4" }, Node : 3 },
330
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v4" }, Node : 3 , View : & meta. LeaderView { ID : 4 , CollectionID : 1 } },
331
331
},
332
332
},
333
333
expectPlans : []SegmentAssignPlan {},
@@ -342,8 +342,8 @@ func (suite *RowCountBasedBalancerTestSuite) TestBalance() {
342
342
distributions : map [int64 ][]* meta.Segment {},
343
343
distributionChannels : map [int64 ][]* meta.DmChannel {
344
344
1 : {
345
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v1" }, Node : 1 },
346
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 1 },
345
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v1" }, Node : 1 , View : & meta. LeaderView { ID : 1 , CollectionID : 1 , Channel : "v1" , Status : & querypb. LeaderViewStatus { Serviceable : true }} },
346
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 1 , View : & meta. LeaderView { ID : 2 , CollectionID : 1 , Channel : "v2" , Status : & querypb. LeaderViewStatus { Serviceable : true }} },
347
347
},
348
348
2 : {},
349
349
3 : {},
@@ -443,8 +443,6 @@ func (suite *RowCountBasedBalancerTestSuite) TestBalance() {
443
443
assertChannelAssignPlanElementMatch (& suite .Suite , c .expectChannelPlans , channelPlans , true )
444
444
}
445
445
446
- // clear distribution
447
-
448
446
for _ , node := range c .nodes {
449
447
balancer .meta .ResourceManager .HandleNodeDown (ctx , node )
450
448
balancer .nodeManager .Remove (node )
@@ -535,7 +533,7 @@ func (suite *RowCountBasedBalancerTestSuite) TestBalanceOnPartStopping() {
535
533
},
536
534
distributionChannels : map [int64 ][]* meta.DmChannel {
537
535
2 : {
538
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 2 },
536
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 2 , View : & meta. LeaderView { ID : 2 , CollectionID : 1 } },
539
537
},
540
538
},
541
539
expectPlans : []SegmentAssignPlan {
@@ -595,10 +593,10 @@ func (suite *RowCountBasedBalancerTestSuite) TestBalanceOnPartStopping() {
595
593
},
596
594
distributionChannels : map [int64 ][]* meta.DmChannel {
597
595
2 : {
598
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 2 },
596
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 2 , View : & meta. LeaderView { ID : 2 , CollectionID : 1 } },
599
597
},
600
598
3 : {
601
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v3" }, Node : 3 },
599
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v3" }, Node : 3 , View : & meta. LeaderView { ID : 3 , CollectionID : 1 } },
602
600
},
603
601
},
604
602
expectPlans : []SegmentAssignPlan {},
@@ -688,10 +686,10 @@ func (suite *RowCountBasedBalancerTestSuite) TestBalanceOutboundNodes() {
688
686
},
689
687
distributionChannels : map [int64 ][]* meta.DmChannel {
690
688
2 : {
691
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 2 },
689
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 2 , View : & meta. LeaderView { ID : 2 , CollectionID : 1 } },
692
690
},
693
691
3 : {
694
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v3" }, Node : 3 },
692
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v3" }, Node : 3 , View : & meta. LeaderView { ID : 3 , CollectionID : 1 } },
695
693
},
696
694
},
697
695
expectPlans : []SegmentAssignPlan {},
@@ -718,10 +716,10 @@ func (suite *RowCountBasedBalancerTestSuite) TestBalanceOutboundNodes() {
718
716
},
719
717
distributionChannels : map [int64 ][]* meta.DmChannel {
720
718
2 : {
721
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 2 },
719
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 2 , View : & meta. LeaderView { ID : 2 , CollectionID : 1 } },
722
720
},
723
721
1 : {
724
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v3" }, Node : 1 },
722
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v3" }, Node : 1 , View : & meta. LeaderView { ID : 3 , CollectionID : 1 } },
725
723
},
726
724
},
727
725
expectPlans : []SegmentAssignPlan {
@@ -897,12 +895,18 @@ func (suite *RowCountBasedBalancerTestSuite) TestAssignSegmentWithGrowing() {
897
895
}
898
896
899
897
// mock 50 growing row count in node 1, which is delegator, expect all segment assign to node 2
900
- leaderView := & meta.LeaderView {
901
- ID : 1 ,
902
- CollectionID : 1 ,
903
- NumOfGrowingRows : 50 ,
904
- }
905
- suite .balancer .dist .LeaderViewManager .Update (1 , leaderView )
898
+ suite .balancer .dist .ChannelDistManager .Update (1 , & meta.DmChannel {
899
+ VchannelInfo : & datapb.VchannelInfo {
900
+ CollectionID : 1 ,
901
+ ChannelName : "v1" ,
902
+ },
903
+ Node : 1 ,
904
+ View : & meta.LeaderView {
905
+ ID : 1 ,
906
+ CollectionID : 1 ,
907
+ NumOfGrowingRows : 50 ,
908
+ },
909
+ })
906
910
plans := balancer .AssignSegment (ctx , 1 , toAssign , lo .Keys (distributions ), false )
907
911
for _ , p := range plans {
908
912
suite .Equal (int64 (2 ), p .To )
@@ -934,14 +938,18 @@ func (suite *RowCountBasedBalancerTestSuite) TestDisableBalanceChannel() {
934
938
distributions : map [int64 ][]* meta.Segment {},
935
939
distributionChannels : map [int64 ][]* meta.DmChannel {
936
940
2 : {
937
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 2 },
938
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v3" }, Node : 2 },
941
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 2 , View : & meta. LeaderView { ID : 2 , CollectionID : 1 } },
942
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v3" }, Node : 2 , View : & meta. LeaderView { ID : 2 , CollectionID : 1 } },
939
943
},
940
944
3 : {},
941
945
},
942
946
expectPlans : []SegmentAssignPlan {},
943
947
expectChannelPlans : []ChannelAssignPlan {
944
- {Channel : & meta.DmChannel {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v3" }, Node : 2 }, From : 2 , To : 3 , Replica : newReplicaDefaultRG (1 )},
948
+ {Channel : & meta.DmChannel {
949
+ VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v3" },
950
+ Node : 2 ,
951
+ View : & meta.LeaderView {ID : 3 , CollectionID : 1 },
952
+ }, From : 2 , To : 3 , Replica : newReplicaDefaultRG (1 )},
945
953
},
946
954
enableBalanceChannel : true ,
947
955
},
@@ -955,8 +963,8 @@ func (suite *RowCountBasedBalancerTestSuite) TestDisableBalanceChannel() {
955
963
distributions : map [int64 ][]* meta.Segment {},
956
964
distributionChannels : map [int64 ][]* meta.DmChannel {
957
965
2 : {
958
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 2 },
959
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v3" }, Node : 2 },
966
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v2" }, Node : 2 , View : & meta. LeaderView { ID : 2 , CollectionID : 1 } },
967
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "v3" }, Node : 2 , View : & meta. LeaderView { ID : 2 , CollectionID : 1 } },
960
968
},
961
969
3 : {},
962
970
},
@@ -1099,12 +1107,12 @@ func (suite *RowCountBasedBalancerTestSuite) TestMultiReplicaBalance() {
1099
1107
},
1100
1108
channelDist : map [int64 ][]* meta.DmChannel {
1101
1109
1 : {
1102
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "channel1" }, Node : 1 },
1103
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "channel2" }, Node : 1 },
1110
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "channel1" }, Node : 1 , View : & meta. LeaderView { ID : 1 , CollectionID : 1 } },
1111
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "channel2" }, Node : 1 , View : & meta. LeaderView { ID : 1 , CollectionID : 1 } },
1104
1112
},
1105
1113
3 : {
1106
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "channel3" }, Node : 3 },
1107
- {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "channel4" }, Node : 3 },
1114
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "channel3" }, Node : 3 , View : & meta. LeaderView { ID : 3 , CollectionID : 1 } },
1115
+ {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "channel4" }, Node : 3 , View : & meta. LeaderView { ID : 3 , CollectionID : 1 } },
1108
1116
},
1109
1117
},
1110
1118
expectPlans : []SegmentAssignPlan {},
@@ -1162,10 +1170,26 @@ func (suite *RowCountBasedBalancerTestSuite) TestMultiReplicaBalance() {
1162
1170
suite .Len (channelPlans , 2 )
1163
1171
1164
1172
// mock new distribution after channel balance
1165
- balancer .dist .ChannelDistManager .Update (1 , & meta.DmChannel {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "channel1" }, Node : 1 })
1166
- balancer .dist .ChannelDistManager .Update (2 , & meta.DmChannel {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "channel2" }, Node : 2 })
1167
- balancer .dist .ChannelDistManager .Update (3 , & meta.DmChannel {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "channel3" }, Node : 3 })
1168
- balancer .dist .ChannelDistManager .Update (4 , & meta.DmChannel {VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "channel4" }, Node : 4 })
1173
+ balancer .dist .ChannelDistManager .Update (1 , & meta.DmChannel {
1174
+ VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "channel1" },
1175
+ Node : 1 ,
1176
+ View : & meta.LeaderView {ID : 1 , CollectionID : 1 },
1177
+ })
1178
+ balancer .dist .ChannelDistManager .Update (2 , & meta.DmChannel {
1179
+ VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "channel2" },
1180
+ Node : 2 ,
1181
+ View : & meta.LeaderView {ID : 2 , CollectionID : 1 },
1182
+ })
1183
+ balancer .dist .ChannelDistManager .Update (3 , & meta.DmChannel {
1184
+ VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "channel3" },
1185
+ Node : 3 ,
1186
+ View : & meta.LeaderView {ID : 3 , CollectionID : 1 },
1187
+ })
1188
+ balancer .dist .ChannelDistManager .Update (4 , & meta.DmChannel {
1189
+ VchannelInfo : & datapb.VchannelInfo {CollectionID : 1 , ChannelName : "channel4" },
1190
+ Node : 4 ,
1191
+ View : & meta.LeaderView {ID : 4 , CollectionID : 1 },
1192
+ })
1169
1193
1170
1194
// expected to balance segment
1171
1195
segmentPlans , channelPlans = suite .getCollectionBalancePlans (balancer , c .collectionID )
0 commit comments