copy slaves to new slots as well

Author: kefu-reddit

rediscluster/client.py

@@ -680,8 +680,8 @@ def _execute_command(self, *args, **kwargs):
                 self.refresh_table_asap = True
                 self.connection_pool.nodes.increment_reinitialize_counter()
 
-                node = self.connection_pool.nodes.set_node(e.host, e.port, server_type='master')
-                self.connection_pool.nodes.slots[e.slot_id][0] = node
+                node = self.connection_pool.nodes.get_node(e.host, e.port, server_type='master')
+                self.connection_pool.nodes.move_slot_to_node(e.slot_id, node)
             except TryAgainError as e:
                 log.exception("TryAgainError")
 

rediscluster/nodemanager.py

@@ -39,6 +39,7 @@ def __init__(self, startup_nodes=None, reinitialize_steps=None, skip_full_covera
         self.connection_kwargs = connection_kwargs
         self.nodes = {}
         self.slots = {}
+        self.slave_nodes_by_master = {}
         self.startup_nodes = [] if startup_nodes is None else startup_nodes
         self.orig_startup_nodes = [node for node in self.startup_nodes]
         self.reinitialize_counter = 0
@@ -257,6 +258,8 @@ def initialize(self):
                 node, node_name = self.make_node_obj(master_node[0], master_node[1], 'master')
                 nodes_cache[node_name] = node
 
+                self.slave_nodes_by_master[node_name] = set()
+
                 for i in range(int(slot[0]), int(slot[1]) + 1):
                     if i not in tmp_slots:
                         tmp_slots[i] = [node]
@@ -267,6 +270,7 @@ def initialize(self):
                             target_slave_node, slave_node_name = self.make_node_obj(slave_node[0], slave_node[1], 'slave')
                             nodes_cache[slave_node_name] = target_slave_node
                             tmp_slots[i].append(target_slave_node)
+                            self.slave_nodes_by_master[node_name].add(slave_node_name)
                     else:
                         # Validate that 2 nodes want to use the same slot cache setup
                         if tmp_slots[i][0]['name'] != node['name']:
@@ -409,6 +413,22 @@ def set_node(self, host, port, server_type=None):
 
         return node
 
+    def get_node(self, host, port, server_type=None):
+        node, node_name = self.make_node_obj(host, port, server_type)
+        if node_name not in self.nodes:
+            self.nodes[node_name] = node
+        return self.nodes[node_name]
+
+    def move_slot_to_node(self, slot, node):
+        node_name = node['name']
+        self.slots[slot] = [node]
+        slave_nodes = self.slave_nodes_by_master.get(node_name)
+        if slave_nodes:
+            for slave_name in slave_nodes:
+                slave_node = self.nodes.get(slave_name)
+                if slave_node:
+                    self.slots[slot].append(slave_node)
+
     def populate_startup_nodes(self):
         """
         Do something with all startup nodes and filters out any duplicates

setup.py

@@ -20,7 +20,7 @@
 
 setup(
     name="redis-py-cluster-patched",
-    version="2.1.0.999.6",
+    version="2.1.0.999.7",
     description="Library for communicating with Redis Clusters. Built on top of redis-py lib",
     long_description=readme + '\n\n' + history,
     long_description_content_type="text/markdown",
@@ -30,7 +30,7 @@
     maintainer_email='Grokzen@gmail.com',
     packages=["rediscluster"],
     url='http://github.com/grokzen/redis-py-cluster',
-    download_url="https://github.com/duke-cliff/redis-py-cluster/archive/2.1.0.999.6.tar.gz",
+    download_url="https://github.com/duke-cliff/redis-py-cluster/archive/2.1.0.999.7.tar.gz",
     license='MIT',
     install_requires=[
         'redis>=3.0.0,<4.0.0'