@@ -564,9 +564,11 @@ func (s *Topom) SlotsRebalance(confirm bool) (map[int]int, error) {
564564 var (
565565 assigned = make (map [int ]int )
566566 pendings = make (map [int ][]int )
567+ moveout = make (map [int ]int )
568+ docking []int
567569 )
568570 var groupSize = func (gid int ) int {
569- return assigned [gid ] + len (pendings [gid ])
571+ return assigned [gid ] + len (pendings [gid ]) - moveout [ gid ]
570572 }
571573
572574 // don't migrate slot if it's being migrated
@@ -592,11 +594,6 @@ func (s *Topom) SlotsRebalance(confirm bool) (map[int]int, error) {
592594 }
593595 }
594596
595- // reverse pending list for each group
596- for _ , list := range pendings {
597- sort .Sort (sort .Reverse (sort .IntSlice (list )))
598- }
599-
600597 var tree = rbtree .NewWith (func (x , y interface {}) int {
601598 var gid1 = x .(int )
602599 var gid2 = y .(int )
@@ -612,8 +609,6 @@ func (s *Topom) SlotsRebalance(confirm bool) (map[int]int, error) {
612609 tree .Put (gid , nil )
613610 }
614611
615- var offline []int
616-
617612 // assign offline slots to the smallest group
618613 for _ , m := range ctx .slots {
619614 if m .Action .State != models .ActionNothing {
@@ -622,79 +617,73 @@ func (s *Topom) SlotsRebalance(confirm bool) (map[int]int, error) {
622617 if m .GroupId != 0 {
623618 continue
624619 }
625- gid := tree .Left ().Key .(int )
626- tree .Remove (gid )
627-
628- assigned [gid ]++
629- tree .Put (gid , nil )
630-
631- offline = append (offline , gid )
632- }
633- sort .Ints (offline )
620+ dest := tree .Left ().Key .(int )
621+ tree .Remove (dest )
634622
635- var plans = make (map [int ]int )
623+ docking = append (docking , m .Id )
624+ moveout [dest ]--
636625
637- // create migration plans for offline slots
638- for _ , m := range ctx .slots {
639- if m .Action .State != models .ActionNothing {
640- continue
641- }
642- if m .GroupId != 0 {
643- continue
644- }
645- if len (offline ) != 0 {
646- plans [m .Id ], offline = offline [0 ], offline [1 :]
647- }
626+ tree .Put (dest , nil )
648627 }
649628
650629 var upperBound = (MaxSlotNum + len (groupIds ) - 1 ) / len (groupIds )
651630
652- var newPlan = func (from , dest int ) bool {
653- var fromSize = groupSize (from )
654- var destSize = groupSize (dest )
655- if fromSize <= lowerBound {
656- return false
657- }
658- if destSize >= upperBound {
659- return false
660- }
661- if d := fromSize - destSize ; d <= 1 {
662- return false
663- }
664- var list = pendings [from ]
665- if len (list ) == 0 {
666- return false
667- }
668- plans [list [0 ]] = dest
669- pendings [from ] = list [1 :]
670- assigned [dest ]++
671- return true
672- }
673-
674631 // rebalance between different server groups
675-
676632 for tree .Size () >= 2 {
677633 from := tree .Right ().Key .(int )
678634 tree .Remove (from )
679635
680- if len (pendings [from ]) == 0 {
636+ if len (pendings [from ]) == moveout [ from ] {
681637 continue
682638 }
683-
684639 dest := tree .Left ().Key .(int )
685640 tree .Remove (dest )
686641
687- var updated bool
688- for newPlan (from , dest ) {
689- updated = true
642+ var (
643+ fromSize = groupSize (from )
644+ destSize = groupSize (dest )
645+ )
646+ if fromSize <= lowerBound {
647+ break
648+ }
649+ if destSize >= upperBound {
650+ break
690651 }
691- if ! updated {
652+ if d := fromSize - destSize ; d <= 1 {
692653 break
693654 }
655+ moveout [from ]++
656+ moveout [dest ]--
657+
694658 tree .Put (from , nil )
695659 tree .Put (dest , nil )
696660 }
697661
662+ for gid , n := range moveout {
663+ if n < 0 {
664+ continue
665+ }
666+ if n > 0 {
667+ sids := pendings [gid ]
668+ sort .Sort (sort .Reverse (sort .IntSlice (sids )))
669+
670+ docking = append (docking , sids [0 :n ]... )
671+ pendings [gid ] = sids [n :]
672+ }
673+ delete (moveout , gid )
674+ }
675+ sort .Ints (docking )
676+
677+ var plans = make (map [int ]int )
678+
679+ for _ , gid := range groupIds {
680+ var in = - moveout [gid ]
681+ for i := 0 ; i < in && len (docking ) != 0 ; i ++ {
682+ plans [docking [0 ]] = gid
683+ docking = docking [1 :]
684+ }
685+ }
686+
698687 if ! confirm {
699688 return plans , nil
700689 }
0 commit comments