Merge pull request #894 from rahul342/master

Extend SlotLoader to work with multiple slot ranges

Author: byroot

lib/redis/cluster/slot.rb

@@ -50,9 +50,12 @@ def slave?(node_key)
         @node_flags[node_key] == ROLE_SLAVE
       end
 
+      # available_slots is mapping of node_key to list of slot ranges
       def build_slot_node_key_map(available_slots)
-        available_slots.each_with_object({}) do |(node_key, slots), acc|
-          slots.each { |slot| assign_node_key(acc, slot, node_key) }
+        available_slots.each_with_object({}) do |(node_key, slots_arr), acc|
+          slots_arr.each do |slots|
+            slots.each { |slot| assign_node_key(acc, slot, node_key) }
+          end
         end
       end
 

lib/redis/cluster/slot_loader.rb

@@ -13,7 +13,7 @@ def load(nodes)
         info = {}
 
         nodes.each do |node|
-          info = Hash[*fetch_slot_info(node)]
+          info = fetch_slot_info(node)
           info.empty? ? next : break
         end
 
@@ -23,9 +23,11 @@ def load(nodes)
       end
 
       def fetch_slot_info(node)
+        hash_with_default_arr = Hash.new { |h, k| h[k] = [] }
         node.call(%i[cluster slots])
-            .map { |arr| parse_slot_info(arr, default_ip: node.host) }
-            .flatten
+          .flat_map { |arr| parse_slot_info(arr, default_ip: node.host) }
+          .each_with_object(hash_with_default_arr) { |arr, h| h[arr[0]] << arr[1] }
+
       rescue CannotConnectError, ConnectionError, CommandError
         {} # can retry on another node
       end
@@ -34,7 +36,6 @@ def parse_slot_info(arr, default_ip:)
         first_slot, last_slot = arr[0..1]
         slot_range = (first_slot..last_slot).freeze
         arr[2..-1].map { |addr| [stringify_node_key(addr, default_ip), slot_range] }
-                  .flatten
       end
 
       def stringify_node_key(arr, default_ip)

test/cluster_client_slots_test.rb

@@ -7,7 +7,7 @@ class TestClusterClientSlots < Minitest::Test
   include Helper::Cluster
 
   def test_slot_class
-    slot = Redis::Cluster::Slot.new('127.0.0.1:7000' => 1..10)
+    slot = Redis::Cluster::Slot.new('127.0.0.1:7000' => [1..10])
 
     assert_equal false, slot.exists?(0)
     assert_equal true, slot.exists?(1)
@@ -26,15 +26,67 @@ def test_slot_class
     assert_nil slot.put(1, '127.0.0.1:7001')
   end
 
+  def test_slot_class_with_multiple_slot_ranges
+    slot = Redis::Cluster::Slot.new('127.0.0.1:7000' => [1..10, 30..40])
+
+    assert_equal false, slot.exists?(0)
+    assert_equal true, slot.exists?(1)
+    assert_equal true, slot.exists?(10)
+    assert_equal false, slot.exists?(11)
+    assert_equal true, slot.exists?(30)
+    assert_equal true, slot.exists?(40)
+    assert_equal false, slot.exists?(41)
+
+    assert_nil slot.find_node_key_of_master(0)
+    assert_nil slot.find_node_key_of_slave(0)
+    assert_equal '127.0.0.1:7000', slot.find_node_key_of_master(1)
+    assert_equal '127.0.0.1:7000', slot.find_node_key_of_slave(1)
+    assert_equal '127.0.0.1:7000', slot.find_node_key_of_master(10)
+    assert_equal '127.0.0.1:7000', slot.find_node_key_of_slave(10)
+    assert_equal '127.0.0.1:7000', slot.find_node_key_of_slave(30)
+    assert_equal '127.0.0.1:7000', slot.find_node_key_of_slave(40)
+    assert_nil slot.find_node_key_of_master(11)
+    assert_nil slot.find_node_key_of_slave(11)
+    assert_nil slot.find_node_key_of_master(41)
+    assert_nil slot.find_node_key_of_slave(41)
+
+    assert_nil slot.put(1, '127.0.0.1:7001')
+    assert_nil slot.put(30, '127.0.0.1:7001')
+  end
+
   def test_slot_class_with_node_flags_and_replicas
-    slot = Redis::Cluster::Slot.new({ '127.0.0.1:7000' => 1..10, '127.0.0.1:7001' => 1..10 },
+    slot = Redis::Cluster::Slot.new({ '127.0.0.1:7000' => [1..10], '127.0.0.1:7001' => [1..10] },
+                                    { '127.0.0.1:7000' => 'master', '127.0.0.1:7001' => 'slave' },
+                                    true)
+
+    assert_equal false, slot.exists?(0)
+    assert_equal true, slot.exists?(1)
+    assert_equal true, slot.exists?(10)
+    assert_equal false, slot.exists?(11)
+
+    assert_nil slot.find_node_key_of_master(0)
+    assert_nil slot.find_node_key_of_slave(0)
+    assert_equal '127.0.0.1:7000', slot.find_node_key_of_master(1)
+    assert_equal '127.0.0.1:7001', slot.find_node_key_of_slave(1)
+    assert_equal '127.0.0.1:7000', slot.find_node_key_of_master(10)
+    assert_equal '127.0.0.1:7001', slot.find_node_key_of_slave(10)
+    assert_nil slot.find_node_key_of_master(11)
+    assert_nil slot.find_node_key_of_slave(11)
+
+    assert_nil slot.put(1, '127.0.0.1:7002')
+  end
+
+  def test_slot_class_with_node_flags_replicas_and_slot_range
+    slot = Redis::Cluster::Slot.new({ '127.0.0.1:7000' => [1..10, 30..40], '127.0.0.1:7001' => [1..10, 30..40] },
                                     { '127.0.0.1:7000' => 'master', '127.0.0.1:7001' => 'slave' },
                                     true)
 
     assert_equal false, slot.exists?(0)
     assert_equal true, slot.exists?(1)
     assert_equal true, slot.exists?(10)
     assert_equal false, slot.exists?(11)
+    assert_equal true, slot.exists?(30)
+    assert_equal false, slot.exists?(41)
 
     assert_nil slot.find_node_key_of_master(0)
     assert_nil slot.find_node_key_of_slave(0)
@@ -44,12 +96,16 @@ def test_slot_class_with_node_flags_and_replicas
     assert_equal '127.0.0.1:7001', slot.find_node_key_of_slave(10)
     assert_nil slot.find_node_key_of_master(11)
     assert_nil slot.find_node_key_of_slave(11)
+    assert_equal '127.0.0.1:7000', slot.find_node_key_of_master(30)
+    assert_equal '127.0.0.1:7001', slot.find_node_key_of_slave(30)
+    assert_nil slot.find_node_key_of_master(41)
+    assert_nil slot.find_node_key_of_slave(41)
 
     assert_nil slot.put(1, '127.0.0.1:7002')
   end
 
   def test_slot_class_with_node_flags_and_without_replicas
-    slot = Redis::Cluster::Slot.new({ '127.0.0.1:7000' => 1..10, '127.0.0.1:7001' => 1..10 },
+    slot = Redis::Cluster::Slot.new({ '127.0.0.1:7000' => [1..10], '127.0.0.1:7001' => [1..10] },
                                     { '127.0.0.1:7000' => 'master', '127.0.0.1:7001' => 'slave' },
                                     false)