Refactor key-slot conversion logic (#301)

Author: supercaracal

lib/redis_client/cluster/command.rb

@@ -2,14 +2,13 @@
 
 require 'redis_client'
 require 'redis_client/cluster/errors'
+require 'redis_client/cluster/key_slot_converter'
 require 'redis_client/cluster/normalized_cmd_name'
 
 class RedisClient
   class Cluster
     class Command
       EMPTY_STRING = ''
-      LEFT_BRACKET = '{'
-      RIGHT_BRACKET = '}'
       EMPTY_HASH = {}.freeze
 
       Detail = Struct.new(
@@ -65,9 +64,7 @@ def extract_first_key(command)
         i = determine_first_key_position(command)
         return EMPTY_STRING if i == 0
 
-        key = (command[i].is_a?(Array) ? command[i].flatten.first : command[i]).to_s
-        hash_tag = extract_hash_tag(key)
-        hash_tag.empty? ? key : hash_tag
+        (command[i].is_a?(Array) ? command[i].flatten.first : command[i]).to_s
       end
 
       def should_send_to_primary?(command)
@@ -105,18 +102,6 @@ def determine_optional_key_position(command, option_name) # rubocop:disable Metr
         idx = command&.flatten&.map(&:to_s)&.map(&:downcase)&.index(option_name&.downcase)
         idx.nil? ? 0 : idx + 1
       end
-
-      # @see https://redis.io/topics/cluster-spec#keys-hash-tags Keys hash tags
-      def extract_hash_tag(key)
-        key = key.to_s
-        s = key.index(LEFT_BRACKET)
-        return EMPTY_STRING if s.nil?
-
-        e = key.index(RIGHT_BRACKET, s + 1)
-        return EMPTY_STRING if e.nil?
-
-        key[s + 1..e - 1]
-      end
     end
   end
 end

lib/redis_client/cluster/key_slot_converter.rb

@@ -3,6 +3,9 @@
 class RedisClient
   class Cluster
     module KeySlotConverter
+      EMPTY_STRING = ''
+      LEFT_BRACKET = '{'
+      RIGHT_BRACKET = '}'
       XMODEM_CRC16_LOOKUP = [
         0x0000, 0x1021, 0x2042, 0x3063, 0x4084, 0x50a5, 0x60c6, 0x70e7,
         0x8108, 0x9129, 0xa14a, 0xb16b, 0xc18c, 0xd1ad, 0xe1ce, 0xf1ef,
@@ -45,13 +48,28 @@ module KeySlotConverter
       def convert(key)
         return nil if key.nil?
 
+        hash_tag = extract_hash_tag(key)
+        key = hash_tag unless hash_tag.empty?
+
         crc = 0
         key.each_byte do |b|
           crc = ((crc << 8) & 0xffff) ^ XMODEM_CRC16_LOOKUP[((crc >> 8) ^ b) & 0xff]
         end
 
         crc % HASH_SLOTS
       end
+
+      # @see https://redis.io/topics/cluster-spec#keys-hash-tags Keys hash tags
+      def extract_hash_tag(key)
+        key = key.to_s
+        s = key.index(LEFT_BRACKET)
+        return EMPTY_STRING if s.nil?
+
+        e = key.index(RIGHT_BRACKET, s + 1)
+        return EMPTY_STRING if e.nil?
+
+        key[s + 1..e - 1]
+      end
     end
   end
 end

lib/redis_client/cluster/router.rb

@@ -172,21 +172,25 @@ def assign_node(command)
         find_node(node_key)
       end
 
-      def find_node_key(command, seed: nil)
-        key = @command.extract_first_key(command)
-        slot = key.empty? ? nil : ::RedisClient::Cluster::KeySlotConverter.convert(key)
-
-        if @command.should_send_to_primary?(command)
-          @node.find_node_key_of_primary(slot) || @node.any_primary_node_key(seed: seed)
+      def find_node_key_by_key(key, seed: nil, primary: false)
+        if key && !key.empty?
+          slot = ::RedisClient::Cluster::KeySlotConverter.convert(key)
+          primary ? @node.find_node_key_of_primary(slot) : @node.find_node_key_of_replica(slot)
         else
-          @node.find_node_key_of_replica(slot, seed: seed) || @node.any_replica_node_key(seed: seed)
+          primary ? @node.any_primary_node_key(seed: seed) : @node.any_replica_node_key(seed: seed)
         end
       end
 
+      def find_node_key(command, seed: nil)
+        key = @command.extract_first_key(command)
+        find_node_key_by_key(key, seed: seed, primary: @command.should_send_to_primary?(command))
+      end
+
       def find_primary_node_key(command)
         key = @command.extract_first_key(command)
-        slot = key.empty? ? nil : ::RedisClient::Cluster::KeySlotConverter.convert(key)
-        @node.find_node_key_of_primary(slot)
+        return nil unless key&.size&.> 0
+
+        find_node_key_by_key(key, primary: true)
       end
 
       def find_node(node_key, retry_count: 3)

test/redis_client/cluster/test_command.rb

@@ -84,7 +84,7 @@ def test_extract_first_key
         [
           { command: %w[SET foo 1], want: 'foo' },
           { command: %w[GET foo], want: 'foo' },
-          { command: %w[GET foo{bar}baz], want: 'bar' },
+          { command: %w[GET foo{bar}baz], want: 'foo{bar}baz' },
           { command: %w[MGET foo bar baz], want: 'foo' },
           { command: %w[UNKNOWN foo bar], want: '' },
           { command: [['GET'], 'foo'], want: 'foo' },
@@ -190,28 +190,6 @@ def test_determine_optional_key_position
           assert_equal(c[:want], got, msg)
         end
       end
-
-      def test_extract_hash_tag
-        cmd = ::RedisClient::Cluster::Command.load(@raw_clients)
-        [
-          { key: 'foo', want: '' },
-          { key: 'foo{bar}baz', want: 'bar' },
-          { key: 'foo{bar}baz{qux}quuc', want: 'bar' },
-          { key: 'foo}bar{baz', want: '' },
-          { key: 'foo{bar', want: '' },
-          { key: 'foo}bar', want: '' },
-          { key: 'foo{}bar', want: '' },
-          { key: '{}foo', want: '' },
-          { key: 'foo{}', want: '' },
-          { key: '{}', want: '' },
-          { key: '', want: '' },
-          { key: nil, want: '' }
-        ].each_with_index do |c, idx|
-          msg = "Case: #{idx}"
-          got = cmd.send(:extract_hash_tag, c[:key])
-          assert_equal(c[:want], got, msg)
-        end
-      end
     end
   end
 end

test/redis_client/cluster/test_key_slot_converter.rb

@@ -27,6 +27,27 @@ def test_convert
         got = ::RedisClient::Cluster::KeySlotConverter.convert(multi_byte_key)
         assert_equal(want, got, "Case: #{multi_byte_key}")
       end
+
+      def test_extract_hash_tag
+        [
+          { key: 'foo', want: '' },
+          { key: 'foo{bar}baz', want: 'bar' },
+          { key: 'foo{bar}baz{qux}quuc', want: 'bar' },
+          { key: 'foo}bar{baz', want: '' },
+          { key: 'foo{bar', want: '' },
+          { key: 'foo}bar', want: '' },
+          { key: 'foo{}bar', want: '' },
+          { key: '{}foo', want: '' },
+          { key: 'foo{}', want: '' },
+          { key: '{}', want: '' },
+          { key: '', want: '' },
+          { key: nil, want: '' }
+        ].each_with_index do |c, idx|
+          msg = "Case: #{idx}"
+          got = ::RedisClient::Cluster::KeySlotConverter.extract_hash_tag(c[:key])
+          assert_equal(c[:want], got, msg)
+        end
+      end
     end
   end
 end