Implement RedisClient::Cluster#with

This implements support for "pinning" a cluster client to a particular
keyslot. This allows the use of WATCH/MULTI/EXEC transactions on that
node without any ambiguity.

Because RedisClient also implements watch as "yield self" and ignores
all arguments, this means that you can easily write transactions that
are compatible with both cluster and non-clustered redis (provided you
use hashtags for your keys).

Author: KJ Tsanaktsidis

.rubocop.yml

@@ -19,6 +19,8 @@ Metrics/ClassLength:
 
 Metrics/ModuleLength:
   Max: 500
+  Exclude:
+    - 'test/**/*'
 
 Metrics/MethodLength:
   Max: 50

README.md

@@ -168,6 +168,94 @@ cli.call('MGET', '{key}1', '{key}2', '{key}3')
 #=> [nil, nil, nil]
 ```
 
+## Transactions
+This gem supports [Redis transactions](https://redis.io/topics/transactions), including atomicity with `MULTI`/`EXEC`,
+and conditional execution with `WATCH`. Redis does not support cross-slot transactions, so all keys used within a
+transaction must live in the same key slot. To use transactions, you must thus "pin" your client to a single slot using
+`#with`. Pass a key to `#with`, and the yielded connection object will enforce that all keys used in the block must
+belong to the same slot.
+
+```ruby
+cli.with(key: '{key}1') do |conn|
+  conn.call('SET', '{key}1', 'This is OK')
+  #=> "OK"
+  conn.call('SET', '{key}2', 'This is also OK; it has the same hashtag, and so the same slot')
+  #=> "OK"
+  conn.call('SET', '{otherslot}3', 'This will raise an exception; this key is not in the same slot as {key}1')
+  #=> Connection pinned to slot 12539 but command ["SET", "{otherslot}3", "..."] includes key {otherslot}3 in slot 10271 (RedisClient::Cluster::Transaction::ConsistencyError)
+end
+```
+
+When using hash tags to enforce same-slot key hashing, it's often neat to simply pass the hashtag _only_ to `#with`:
+
+```ruby
+cli.with(key: '{myslot}') do |conn|
+  # You can use any key starting with {myslot} in this block
+end
+```
+
+Once you have pinned a client to a particular slot, you can use the same transaction APIs as the
+[redis-client](https://github.com/redis-rb/redis-client#usage) gem allows.
+
+```ruby
+# No concurrent client will ever see the value 1 in 'mykey'; it will see either zero or two.
+cli.call('SET', 'key', 0)
+cli.with(key: 'key') do |conn|
+  conn.multi do |txn|
+    txn.call('INCR', 'key')
+    txn.call('INCR', 'key')
+  end
+  #=> ['OK', 'OK']
+end
+
+# Conditional execution with WATCH can be used to e.g. atomically swap two keys
+cli.call('MSET', '{key}1', 'v1', '{key}2', 'v2')
+cli.with(key: '{key}') do |conn|
+  conn.call('WATCH', '{key}1', '{key}2')
+  conn.multi do |txn|
+    old_key1 = conn.call('GET', '{key}1')
+    old_key2 = conn.call('GET', '{key}2')
+    txn.call('SET', '{key}1', old_key2)
+    txn.call('SET', '{key}2', old_key1)
+  end
+  # This transaction will swap the values of {key}1 and {key}2 only if no concurrent connection modified
+  # either of the values
+end
+
+# You can also pass watch: to #multi as a shortcut
+cli.call('MSET', '{key}1', 'v1', '{key}2', 'v2')
+cli.with(key: '{key}') do |conn|
+  conn.multi(watch: ['{key}1', '{key}2']) do |txn|
+    old_key1, old_key2 = conn.call('MGET', '{key}1', '{key}2')
+    txn.call('MSET', '{key}1', old_key2, '{key}2', old_key1)
+  end
+end
+```
+
+Pinned connections are aware of redirections and node failures like ordinary calls to `RedisClient::Cluster`, but because
+you may have written non-idempotent code inside your block, the block is not automatically retried if e.g. the slot
+it is operating on moves to a different node. If you want this, you can opt-in to retries by passing nonzero
+`retry_count` to `#with`.
+
+```ruby
+cli.with(key: '{key}', retry_count: 1) do |conn|
+  conn.call('GET', '{key}1')
+  #=> "value1"
+
+  # Now, some changes in cluster topology mean that {key} is moved to a different node!
+
+  conn.call('GET', '{key}2')
+  #=> MOVED 9039 127.0.0.1:16381 (RedisClient::CommandError)
+
+  # Luckily, the block will get retried (once) and so both GETs will be re-executed on the newly-discovered
+  # correct node.
+end
+```
+
+Because `RedisClient` from the redis-client gem implements `#with` as simply `yield self` and ignores all of its
+arguments, it's possible to write code which is compatible with both redis-client and redis-cluster-client; the `#with`
+call will pin the connection to a slot when using clustering, or be a no-op when not.
+
 ## ACL
 The cluster client internally calls [COMMAND](https://redis.io/commands/command/) and [CLUSTER NODES](https://redis.io/commands/cluster-nodes/) commands to operate correctly.
 So please permit it like the followings.

lib/redis_client/cluster.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require 'redis_client/cluster/concurrent_worker'
+require 'redis_client/cluster/pinning'
 require 'redis_client/cluster/pipeline'
 require 'redis_client/cluster/pub_sub'
 require 'redis_client/cluster/router'
@@ -97,6 +98,12 @@ def pubsub
       ::RedisClient::Cluster::PubSub.new(@router, @command_builder)
     end
 
+    def with(key:, write: true, retry_count: 0)
+      @router.with(key: key, write: write, retry_count: retry_count) do |conn, slot|
+        yield ::RedisClient::Cluster::Pinning::ConnectionProxy.new(conn, slot, @router, @command_builder)
+      end
+    end
+
     def close
       @concurrent_worker.close
       @router.close

lib/redis_client/cluster/router.rb

@@ -199,6 +199,10 @@ def command_exists?(name)
         @command.exists?(name)
       end
 
+      def command_extract_all_keys(command)
+        @command.extract_all_keys(command)
+      end
+
       def assign_redirection_node(err_msg)
         _, slot, node_key = err_msg.split
         slot = slot.to_i
@@ -219,6 +223,26 @@ def close
         @node.each(&:close)
       end
 
+      def with(key:, write: true, retry_count: 0)
+        raise ArgumentError, 'key must be provided' if key.nil?
+
+        slot = ::RedisClient::Cluster::KeySlotConverter.convert_with_hashtag(key)
+        node_key = if write
+                     @node.find_node_key_of_primary(slot)
+                   else
+                     @node.find_node_key_of_replica(slot)
+                   end
+        node = find_node(node_key)
+        # Calling #with checks out the underlying connection if this is a pooled connection
+        # Calling it through #try_delegate ensures we handle any redirections and retry the entire
+        # transaction if so.
+        try_delegate(node, :with, retry_count: retry_count) do |conn|
+          conn.disable_reconnection do
+            yield conn, slot
+          end
+        end
+      end
+
       private
 
       def send_wait_command(method, command, args, retry_count: 3, &block) # rubocop:disable Metrics/AbcSize

test/redis_client/test_cluster.rb

@@ -480,6 +480,213 @@ def test_circuit_breakers
         cli&.close
       end
 
+      def test_pinning_single_key
+        got = @client.with(key: 'key1') do |conn|
+          conn.call('SET', 'key1', 'hello')
+          conn.call('GET', 'key1')
+        end
+        assert_equal('hello', got)
+      end
+
+      def test_pinning_no_key
+        assert_raises(ArgumentError) do
+          @client.with(key: nil) {}
+        end
+      end
+
+      def test_pinning_two_keys
+        got = @client.with(key: '{slot}') do |conn|
+          conn.call('SET', '{slot}key1', 'v1')
+          conn.call('SET', '{slot}key2', 'v2')
+          conn.call('MGET', '{slot}key1', '{slot}key2')
+        end
+        assert_equal(%w[v1 v2], got)
+      end
+
+      def test_pinning_cross_slot
+        assert_raises(::RedisClient::Cluster::Transaction::ConsistencyError) do
+          @client.with(key: '{slot1}') do |conn|
+            conn.call('GET', '{slot2}')
+          end
+        end
+      end
+
+      def test_pinning_pipeline
+        got = @client.with(key: '{slot}') do |conn|
+          conn.call_v(['SET', '{slot}counter', 0])
+          conn.pipelined do |pipe|
+            pipe.call_v(['INCR', '{slot}counter'])
+            pipe.call_v(['INCR', '{slot}counter'])
+            pipe.call_v(['INCR', '{slot}counter'])
+          end
+          conn.call_v(['GET', '{slot}counter']).to_i
+        end
+
+        assert_equal(3, got)
+      end
+
+      def test_pinning_pipeline_with_error
+        assert_raises(RedisClient::CommandError) do
+          @client.with(key: '{slot}') do |conn|
+            conn.pipelined do |pipeline|
+              pipeline.call('SET', '{slot}key', 'first')
+              pipeline.call('SET', '{slot}key', 'second', 'too many args')
+              pipeline.call('SET', '{slot}key', 'third')
+            end
+          end
+        end
+
+        wait_for_replication
+        assert_equal('third', @client.call('GET', '{slot}key'))
+      end
+
+      def test_pinning_transaction
+        got = @client.with(key: '{slot}') do |conn|
+          conn.multi do |txn|
+            txn.call('SET', '{slot}key1', 'value1')
+            txn.call('SET', '{slot}key2', 'value2')
+          end
+        end
+
+        assert_equal(%w[OK OK], got)
+      end
+
+      def test_pinning_transaction_watch_arg
+        @client.call('MSET', '{slot}key1', 'val1', '{slot}key2', 'val2')
+        @captured_commands.clear
+
+        got = @client.with(key: '{slot}') do |conn|
+          conn.multi(watch: ['{slot}key1', '{slot}key2']) do |txn|
+            old_key1, old_key2 = conn.call('MGET', '{slot}key1', '{slot}key2')
+            txn.call('MSET', '{slot}key1', old_key2, '{slot}key2', old_key1)
+          end
+        end
+
+        assert_equal([
+                       %w[WATCH {slot}key1 {slot}key2],
+                       %w[MGET {slot}key1 {slot}key2],
+                       %w[MULTI],
+                       %w[MSET {slot}key1 val2 {slot}key2 val1],
+                       %w[EXEC]
+                     ], @captured_commands.map(&:command))
+
+        wait_for_replication
+        assert_equal(['OK'], got)
+        assert_equal(%w[val2 val1], @client.call('MGET', '{slot}key1', '{slot}key2'))
+      end
+
+      def test_pinning_transaction_watch_arg_unwatches_on_raise
+        ex = Class.new(StandardError)
+        @captured_commands.clear
+
+        assert_raises(ex) do
+          @client.with(key: '{slot}') do |conn|
+            conn.multi(watch: ['{slot}key1']) do |_txn|
+              conn.call('GET', '{slot}key1')
+              raise ex, 'boom'
+            end
+          end
+        end
+
+        assert_equal([
+                       %w[WATCH {slot}key1],
+                       %w[GET {slot}key1],
+                       %w[UNWATCH]
+                     ], @captured_commands.map(&:command))
+      end
+
+      def test_pinning_transaction_can_watch_manually
+        @client.call('MSET', '{slot}key1', 'val1', '{slot}key2', 'val2')
+        @captured_commands.clear
+
+        got = @client.with(key: '{slot}') do |conn|
+          conn.call('WATCH', '{slot}key1', '{slot}key2')
+          old_key1, old_key2 = conn.call('MGET', '{slot}key1', '{slot}key2')
+          conn.multi do |txn|
+            txn.call('MSET', '{slot}key1', old_key2, '{slot}key2', old_key1)
+          end
+        end
+
+        assert_equal([
+                       %w[WATCH {slot}key1 {slot}key2],
+                       %w[MGET {slot}key1 {slot}key2],
+                       %w[MULTI],
+                       %w[MSET {slot}key1 val2 {slot}key2 val1],
+                       %w[EXEC]
+                     ], @captured_commands.map(&:command))
+
+        wait_for_replication
+        assert_equal(['OK'], got)
+        assert_equal(%w[val2 val1], @client.call('MGET', '{slot}key1', '{slot}key2'))
+      end
+
+      def test_pinning_transaction_can_unwatch_manually
+        got = @client.with(key: '{slot}') do |conn|
+          conn.call('WATCH', '{slot}key1')
+          conn.call('UNWATCH')
+        end
+
+        assert_equal('OK', got)
+      end
+
+      def test_pinning_nonblocking_timeouts_update_topology
+        @client.call('DEL', '{slot}list')
+        @captured_commands.clear
+
+        assert_raises(::RedisClient::ReadTimeoutError) do
+          @client.with(key: '{slot}') do |conn|
+            conn.call_v(['BLPOP', '{slot}list', 0])
+          end
+        end
+        assert_includes(@captured_commands.map(&:command), ['CLUSTER', 'NODES'])
+      end
+
+      def test_pinning_blocking_timeouts_do_not_update_topology
+        @client.call('DEL', '{slot}list')
+        @captured_commands.clear
+
+        assert_raises(::RedisClient::ReadTimeoutError) do
+          @client.with(key: '{slot}') do |conn|
+            conn.blocking_call_v(1, ['BLPOP', '{slot}list', 0])
+          end
+        end
+        refute_includes(@captured_commands.map(&:command), ['CLUSTER', 'NODES'])
+      end
+
+      def test_pinning_sscan
+        @client.call('DEL', '{slot}set')
+        expected_set = Set.new
+        scanned_set = Set.new
+        1000.times do |i|
+          expected_set << i
+          @client.call('SADD', '{slot}set', i)
+        end
+        @client.with(key: '{slot}') do |conn|
+          conn.sscan('{slot}set') do |i|
+            scanned_set << i.to_i
+          end
+        end
+
+        assert_equal(expected_set, scanned_set)
+      end
+
+      def test_pinning_zscan
+        @client.call('DEL', '{slot}set')
+        expected_set = Set.new
+        scanned_set = Set.new
+        1000.times do |i|
+          expected_set << "member#{i}"
+          @client.call('ZADD', '{slot}set', i, "member#{i}")
+        end
+        @client.with(key: '{slot}') do |conn|
+          conn.zscan('{slot}set') do |i|
+            scanned_set << i
+          end
+        end
+
+        assert_equal(expected_set, scanned_set)
+      end
+
       private
 
       def wait_for_replication