Improve redis-trib replica assignment

This tiny bit of code has gone through so many revisions.  Hopefully
it's more correct now.

Fixes #2204

1 files changed, 42 insertions, 25 deletions

diff --git a/src/redis-trib.rb b/src/redis-trib.rb
index 4002f6309..e5e7f2da7 100755
--- a/src/redis-trib.rb
+++ b/src/redis-trib.rb
@@ -540,7 +540,6 @@ class RedisTrib
         nodes_count = @nodes.length
         masters_count = @nodes.length / (@replicas+1)
         masters = []
-        slaves = []
 
         # The first step is to split instances by IP. This is useful as
         # we'll try to allocate master nodes in different physical machines
@@ -558,16 +557,22 @@ class RedisTrib
 
         # Select master instances
         puts "Using #{masters_count} masters:"
-        while masters.length < masters_count
-            ips.each{|ip,nodes_list|
-                next if nodes_list.length == 0
-                masters << nodes_list.shift
-                puts masters[-1]
-                nodes_count -= 1
-                break if masters.length == masters_count
-            }
+        interleaved = []
+        stop = false
+        while not stop do
+            # Take one node from each IP until we run out of nodes
+            # across every IP.
+            ips.each do |ip,nodes|
+                stop = nodes.empty? and next
+                interleaved.push nodes.shift
+            end
         end
 
+        masters = interleaved.slice!(0, masters_count)
+        nodes_count -= masters.length
+
+        masters.each{|m| puts m}
+
         # Alloc slots on masters
         slots_per_node = ClusterHashSlots.to_f / masters_count
         first = 0
@@ -594,8 +599,8 @@ class RedisTrib
         # all nodes will be used.
         assignment_verbose = false
 
-        [:requested,:unused].each{|assign|
-            masters.each{|m|
+        [:requested,:unused].each do |assign|
+            masters.each do |m|
                 assigned_replicas = 0
                 while assigned_replicas < @replicas
                     break if nodes_count == 0
@@ -609,21 +614,33 @@ class RedisTrib
                                  "role too (#{nodes_count} remaining)."
                         end
                     end
-                    ips.each{|ip,nodes_list|
-                        next if nodes_list.length == 0
-                        # Skip instances with the same IP as the master if we
-                        # have some more IPs available.
-                        next if ip == m.info[:host] && nodes_count > nodes_list.length
-                        slave = nodes_list.shift
-                        slave.set_as_replica(m.info[:name])
-                        nodes_count -= 1
-                        assigned_replicas += 1
-                        puts "Adding replica #{slave} to #{m}"
-                        break
-                    }
+
+                    # Return the first node not matching our current master
+                    node = interleaved.find{|n| n.info[:host] != m.info[:host]}
+
+                    # If we found a node, use it as a best-first match.
+                    # Otherwise, we didn't find a node on a different IP, so we
+                    # go ahead and use a same-IP replica.
+                    if node
+                        slave = node
+                        interleaved.delete node
+                    else
+                        slave = interleaved.shift
+                    end
+                    slave.set_as_replica(m.info[:name])
+                    nodes_count -= 1
+                    assigned_replicas += 1
+                    puts "Adding replica #{slave} to #{m}"
+
+                    # If we are in the "assign extra nodes" loop,
+                    # we want to assign one extra replica to each
+                    # master before repeating masters.
+                    # This break lets us assign extra replicas to masters
+                    # in a round-robin way.
+                    break if assign == :unused
                 end
-            }
-        }
+            end
+        end
     end
 
     def flush_nodes_config