perf: optimize memory allocation for slot-node-mapping array (#143)

Author: supercaracal

lib/redis_client/cluster/node.rb

@@ -35,6 +35,26 @@ def replica?
         end
       end
 
+      SLOT_OPTIMIZATION_MAX_SHARD_SIZE = 256
+      SLOT_OPTIMIZATION_STRING = '0' * SLOT_SIZE
+      Slot = Struct.new('RedisSlot', :slots, :primary_node_keys, keyword_init: true) do
+        def [](slot)
+          primary_node_keys[slots.getbyte(slot)]
+        end
+
+        def []=(slot, primary_node_key)
+          index = primary_node_keys.find_index(primary_node_key)
+          if index.nil?
+            raise(::RedisClient::Cluster::Node::ReloadNeeded, primary_node_key) if primary_node_keys.size >= SLOT_OPTIMIZATION_MAX_SHARD_SIZE
+
+            index = primary_node_keys.size
+            primary_node_keys << primary_node_key
+          end
+
+          slots.setbyte(slot, index)
+        end
+      end
+
       class Config < ::RedisClient::Config
         def initialize(scale_read: false, **kwargs)
           @scale_read = scale_read
@@ -225,7 +245,7 @@ def make_topology_class(with_replica, replica_affinity)
       end
 
       def build_slot_node_mappings(node_info_list)
-        slots = Array.new(SLOT_SIZE)
+        slots = make_array_for_slot_node_mappings(node_info_list)
         node_info_list.each do |info|
           next if info.slots.nil? || info.slots.empty?
 
@@ -235,6 +255,15 @@ def build_slot_node_mappings(node_info_list)
         slots
       end
 
+      def make_array_for_slot_node_mappings(node_info_list)
+        return Array.new(SLOT_SIZE) if node_info_list.count(&:primary?) > SLOT_OPTIMIZATION_MAX_SHARD_SIZE
+
+        ::RedisClient::Cluster::Node::Slot.new(
+          slots: String.new(SLOT_OPTIMIZATION_STRING, encoding: Encoding::BINARY, capacity: SLOT_SIZE),
+          primary_node_keys: node_info_list.select(&:primary?).map(&:node_key)
+        )
+      end
+
       def build_replication_mappings(node_info_list) # rubocop:disable Metrics/AbcSize
         dict = node_info_list.to_h { |info| [info.id, info] }
         node_info_list.each_with_object(Hash.new { |h, k| h[k] = [] }) do |info, acc|

test/redis_client/cluster/test_node.rb

@@ -334,7 +334,7 @@ def test_any_replica_node_key
       def test_update_slot
         sample_slot = 0
         base_node_key = @test_node.find_node_key_of_primary(sample_slot)
-        another_node_key = @test_node_info_list.find { |info| info.node_key != base_node_key && info.primary? }
+        another_node_key = @test_node_info_list.find { |info| info.node_key != base_node_key && info.primary? }&.node_key
         @test_node.update_slot(sample_slot, another_node_key)
         assert_equal(another_node_key, @test_node.find_node_key_of_primary(sample_slot))
       end
@@ -353,12 +353,12 @@ def test_make_topology_class
 
       def test_build_slot_node_mappings
         node_info_list = [
-          { node_key: '127.0.0.1:7001', slots: [[0, 3000], [3002, 5460], [15_001, 15_001]] },
-          { node_key: '127.0.0.1:7002', slots: [[3001, 3001], [5461, 7000], [7002, 10_922]] },
-          { node_key: '127.0.0.1:7003', slots: [[7001, 7001], [10_923, 15_000], [15_002, 16_383]] },
-          { node_key: '127.0.0.1:7004', slots: [] },
-          { node_key: '127.0.0.1:7005', slots: [] },
-          { node_key: '127.0.0.1:7006', slots: [] }
+          { node_key: '127.0.0.1:7001', role: 'master', slots: [[0, 3000], [3002, 5460], [15_001, 15_001]] },
+          { node_key: '127.0.0.1:7002', role: 'master', slots: [[3001, 3001], [5461, 7000], [7002, 10_922]] },
+          { node_key: '127.0.0.1:7003', role: 'master', slots: [[7001, 7001], [10_923, 15_000], [15_002, 16_383]] },
+          { node_key: '127.0.0.1:7004', role: 'slave', slots: [] },
+          { node_key: '127.0.0.1:7005', role: 'slave', slots: [] },
+          { node_key: '127.0.0.1:7006', role: 'slave', slots: [] }
         ].map { |info| ::RedisClient::Cluster::Node::Info.new(**info) }
 
         got = @test_node.send(:build_slot_node_mappings, node_info_list)
@@ -372,6 +372,58 @@ def test_build_slot_node_mappings
         end
       end
 
+      def test_make_array_for_slot_node_mappings_optimized
+        node_info_list = Array.new(256) do |i|
+          ::RedisClient::Cluster::Node::Info.new(
+            node_key: "127.0.0.1:#{1024 + i + 1}",
+            role: 'master'
+          )
+        end
+
+        want = node_info_list.first.node_key
+        got = @test_node.send(:make_array_for_slot_node_mappings, node_info_list)
+        assert_instance_of(Struct::RedisSlot, got)
+        ::RedisClient::Cluster::Node::SLOT_SIZE.times do |i|
+          got[i] = want
+          assert_equal(want, got[i], "Case: #{i}")
+        end
+      end
+
+      def test_make_array_for_slot_node_mappings_unoptimized
+        node_info_list = Array.new(257) do |i|
+          ::RedisClient::Cluster::Node::Info.new(
+            node_key: "127.0.0.1:#{1024 + i + 1}",
+            role: 'master'
+          )
+        end
+
+        want = node_info_list.first.node_key
+        got = @test_node.send(:make_array_for_slot_node_mappings, node_info_list)
+        assert_instance_of(Array, got)
+        ::RedisClient::Cluster::Node::SLOT_SIZE.times do |i|
+          got[i] = want
+          assert_equal(want, got[i], "Case: #{i}")
+        end
+      end
+
+      def test_make_array_for_slot_node_mappings_max_shard_size
+        node_info_list = Array.new(255) do |i|
+          ::RedisClient::Cluster::Node::Info.new(
+            node_key: "127.0.0.1:#{1024 + i + 1}",
+            role: 'master'
+          )
+        end
+
+        got = @test_node.send(:make_array_for_slot_node_mappings, node_info_list)
+        assert_instance_of(Struct::RedisSlot, got)
+
+        ::RedisClient::Cluster::Node::SLOT_SIZE.times { |i| got[i] = node_info_list.first.node_key }
+
+        got[0] = 'newbie:6379'
+        assert_equal('newbie:6379', got[0])
+        assert_raises(::RedisClient::Cluster::Node::ReloadNeeded) { got[0] = 'zombie:6379' }
+      end
+
       def test_build_replication_mappings_regular
         node_key1 = '127.0.0.1:7001'
         node_key2 = '127.0.0.1:7002'