Fix distribution of masters.

Author: doyoubi

ruskit/distribute.py

@@ -14,6 +14,13 @@ def __init__(self, node, tag, host_index, master=None):
         self.master = master
         self.slaves = []
 
+    @classmethod
+    def divide_by_host(cls, nodes, host_num):
+        nodes_per_host = [[] for _ in xrange(host_num)]
+        for n in nodes:
+            nodes_per_host[n.host_index].append(n)
+        return nodes_per_host
+
     def __getattr__(self, attr):
         return getattr(self.node, attr)
 

ruskit/failover.py

@@ -1,10 +1,11 @@
+from copy import copy
 import logging
 from itertools import izip_longest
 
 import redis
 
 from .cluster import Cluster, ClusterNode
-from .distribute import gen_distribution
+from .distribute import gen_distribution, NodeWrapper
 
 
 logger = logging.getLogger(__name__)
@@ -16,13 +17,43 @@ def __init__(self, cluster, new_nodes):
         self.cluster = cluster
         self.new_nodes = new_nodes
         self.result = None
+        self.dis = gen_distribution(self.cluster.nodes, new_nodes)
 
     def get_distribution(self):
-        return gen_distribution(self.cluster.nodes, self.new_nodes)
+        hosts = map(copy, self.dis['hosts'])
+        masters = self.dis['masters']
+        slaves = self.dis['slaves']
+        result = self.peek_result()
+        plan = result['plan']
+        deleted_masters = [p['master'] for p in plan]
+        new_slaves = []
+        new_masters = []
+        for p in plan:
+            m = copy(p['master'])
+            s = copy(p['slave'])
+            # After failover, their roles will change
+            m.slaves = []
+            m.master = s
+            s.slaves.append(m)
+            s.master = None
+            new_masters.append(s)
+            new_slaves.append(m)
+        new_slaves = NodeWrapper.divide_by_host(new_slaves, len(hosts))
+        new_masters = NodeWrapper.divide_by_host(new_masters, len(hosts))
+
+        masters = [list(set(m) - set(deleted_masters)) for m in masters]
+        masters = [o + n for o, n in zip(masters, new_masters)]
+        slaves = [o + n for o, n in zip(slaves, new_slaves)]
+        return {
+            'hosts': hosts,
+            'masters': masters,
+            'slaves': slaves,
+            'frees': NodeWrapper.divide_by_host([], len(hosts)),
+        }
 
     def peek_result(self):
         if self.result:
-            return result
+            return self.result
         return self.gen_plan(self.new_nodes)
 
     def move_masters_to_new_hosts(self, fast_mode=False):
@@ -64,24 +95,35 @@ def promote_new_masters(self, plan):
 
     def gen_plan(self, new_nodes):
         plan = []
-        dis = gen_distribution(self.cluster.nodes, new_nodes)
-        masters = dis['masters']
-        frees = filter(None, dis['frees'])
-        hosts_len = len(dis['hosts'])
+        masters = map(copy, self.dis['masters'])
+        frees = filter(None, map(copy, self.dis['frees']))
+        hosts_len = len(self.dis['hosts'])
         masters_sum = len(sum(masters, []))
         aver = (masters_sum + hosts_len - 1) / hosts_len
-        failover_masters = sum([h[aver:] for h in masters], [])
-        # merge [[a,b], [c], [d,e,f]] to [a,c,e,b,e,f] for example
+        new_hosts_num = len(set(n.host_index for n in sum(frees, [])))
+        num_in_old_hosts = masters_sum / \
+            hosts_len * (hosts_len - new_hosts_num) + masters_sum % hosts_len
+        moved_masters = masters_sum - num_in_old_hosts
+        # merge [[a,b], [c], [d,e,f]] to [a,c,d,b,e,f] for example
         frees = filter(None, sum(map(list, izip_longest(*frees)), []))
-        while len(failover_masters) > 0 and len(frees) > 0:
-            m = failover_masters.pop()
-            f = frees.pop()
+        while moved_masters > 0 and len(frees) > 0:
+            m = self.select_master(masters)
+            f = frees.pop(0)
+            f.master = m
+            m.slaves.append(f)
             plan.append({
                 'slave': f,
                 'master': m,
             })
+            moved_masters -= 1
         self.result = {
             'plan': plan,
             'frees': frees,
         }
         return self.result
+
+    def select_master(self, masters):
+        nums = map(len, masters)
+        m = max(nums)
+        i = nums.index(m)
+        return masters[i].pop()

tests/test_move_nodes.py

@@ -15,17 +15,18 @@ def __init__(self):
 def dummy_gen_distribution_for_move_masters(nodes, new_nodes):
     hosts = ['host1', 'host2', 'host3', 'host4']
     masters = [
-        ['m1', 'm2', 'm3'],
-        ['m4', 'm5', 'm6', 'm10', 'm11'],
-        ['m7', 'm8', 'm9'],
+        [NodeWrapper(i, 'm{}'.format(i), 0) for i in range(1, 4)],
+        [NodeWrapper(i, 'm{}'.format(i), 1) for i in range(4, 9)],
+        [NodeWrapper(i, 'm{}'.format(i), 2) for i in range(9, 11)],
         [],
     ]
     slaves = [
-        ['s1'],
-        ['s2'],
-        ['s3'],
+        [NodeWrapper(None, 's1', 1)],
+        [NodeWrapper(None, 's2', 2)],
+        [NodeWrapper(None, 's3', 3)],
+        [],
     ]
-    frees = [[], [], [], ['f1', 'f2']]
+    frees = [[], [], [], [NodeWrapper(i, 'f{}'.format(i), 3) for i in range(1, 3)]]
     return {
         'hosts': hosts,
         'masters': masters,
@@ -92,12 +93,20 @@ def test_move_master(self):
         plan = result['plan']
         p1 = plan[0]
         p2 = plan[1]
-        if p1['master'] != 'm11':
+        if p1['master'].tag != 'm8':
             p1, p2 = p2, p1
-        self.assertEqual(p1['master'], 'm11')
-        self.assertEqual(p1['slave'], 'f2')
-        self.assertEqual(p2['master'], 'm10')
-        self.assertEqual(p2['slave'], 'f1')
+        slaves = ('f1', 'f2')
+        self.assertEqual(p1['master'].tag, 'm8')
+        self.assertTrue(p1['slave'].tag in slaves)
+        self.assertEqual(p2['master'].tag, 'm7')
+        self.assertTrue(p2['slave'].tag in slaves)
+
+        dis = manager.get_distribution()
+        masters = dis['masters']
+        masters_num = list(set(map(len, masters)))
+        self.assertTrue(len(masters_num) <= 2)
+        if len(masters_num) == 2:
+            self.assertTrue(abs(masters_num[0] - masters_num[1]) <= 1)
 
     @patch('ruskit.distribute.gen_distribution',
         MoveSlavesMockData.dummy_gen_distribution)