Merge pull request #789 from supercaracal/fix-resharding-helper

Fix resharding test helper method for cluster mode and enhance migrate command method

Author: byroot

lib/redis.rb

@@ -499,22 +499,27 @@ def restore(key, ttl, serialized_value, options = {})
 
   # Transfer a key from the connected instance to another instance.
   #
-  # @param [String] key
+  # @param [String, Array<String>] key
   # @param [Hash] options
   #   - `:host => String`: host of instance to migrate to
   #   - `:port => Integer`: port of instance to migrate to
   #   - `:db => Integer`: database to migrate to (default: same as source)
   #   - `:timeout => Integer`: timeout (default: same as connection timeout)
+  #   - `:copy => Boolean`: Do not remove the key from the local instance.
+  #   - `:replace => Boolean`: Replace existing key on the remote instance.
   # @return [String] `"OK"`
   def migrate(key, options)
-    host = options[:host] || raise(RuntimeError, ":host not specified")
-    port = options[:port] || raise(RuntimeError, ":port not specified")
-    db = (options[:db] || @client.db).to_i
-    timeout = (options[:timeout] || @client.timeout).to_i
+    args = [:migrate]
+    args << (options[:host] || raise(':host not specified'))
+    args << (options[:port] || raise(':port not specified'))
+    args << (key.is_a?(String) ? key : '')
+    args << (options[:db] || @client.db).to_i
+    args << (options[:timeout] || @client.timeout).to_i
+    args << 'COPY' if options[:copy]
+    args << 'REPLACE' if options[:replace]
+    args += ['KEYS', *key] if key.is_a?(Array)
 
-    synchronize do |client|
-      client.call([:migrate, host, port, key, db, timeout])
-    end
+    synchronize { |client| client.call(args) }
   end
 
   # Delete one or more keys.

test/commands_on_value_types_test.rb

@@ -202,6 +202,21 @@ def test_migrate
       actual = redis.migrate("foo", options.merge(:timeout => default_timeout + 1))
       expected = ["127.0.0.1", "1234", "foo", default_db.to_s, (default_timeout + 1).to_s]
       assert_equal expected, actual
+
+      # Test copy override
+      actual = redis.migrate('foo', options.merge(copy: true))
+      expected = ['127.0.0.1', '1234', 'foo', default_db.to_s, default_timeout.to_s, 'COPY']
+      assert_equal expected, actual
+
+      # Test replace override
+      actual = redis.migrate('foo', options.merge(replace: true))
+      expected = ['127.0.0.1', '1234', 'foo', default_db.to_s, default_timeout.to_s, 'REPLACE']
+      assert_equal expected, actual
+
+      # Test multiple keys
+      actual = redis.migrate(%w[foo bar baz], options)
+      expected = ['127.0.0.1', '1234', '', default_db.to_s, default_timeout.to_s, 'KEYS', 'foo', 'bar', 'baz']
+      assert_equal expected, actual
     end
   end
 end

test/support/cluster/orchestrator.rb

@@ -39,7 +39,7 @@ def failover
     wait_cluster_recovering(@clients)
   end
 
-  def start_resharding(slot, src_node_key, dest_node_key)
+  def start_resharding(slot, src_node_key, dest_node_key, slice_size: 10)
     node_map = hashify_node_map(@clients.first)
     src_node_id = node_map.fetch(src_node_key)
     src_client = find_client(@clients, src_node_key)
@@ -50,11 +50,21 @@ def start_resharding(slot, src_node_key, dest_node_key)
     dest_client.cluster(:setslot, slot, 'IMPORTING', src_node_id)
     src_client.cluster(:setslot, slot, 'MIGRATING', dest_node_id)
 
+    keys_count = src_client.cluster(:countkeysinslot, slot)
     loop do
-      keys = src_client.cluster(:getkeysinslot, slot, 100)
+      break if keys_count <= 0
+      keys = src_client.cluster(:getkeysinslot, slot, slice_size)
       break if keys.empty?
-      keys.each { |k| src_client.migrate(k, host: dest_host, port: dest_port) }
-      sleep 0.1
+      keys.each do |k|
+        begin
+          src_client.migrate(k, host: dest_host, port: dest_port)
+        rescue Redis::CommandError => err
+          raise unless err.message.start_with?('IOERR')
+          src_client.migrate(k, host: dest_host, port: dest_port, replace: true) # retry once
+        ensure
+          keys_count -= 1
+        end
+      end
     end
   end