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,20 @@ def pubsub
       ::RedisClient::Cluster::PubSub.new(@router, @command_builder)
     end
 
+    def with(key:, write: true, retry_count: 0)
+      raise ArgumentError, 'key must be provided' if key.nil? || key.empty?
+
+      slot = ::RedisClient::Cluster::KeySlotConverter.convert(key)
+      node_key = @router.find_node_key_by_key(key, primary: write)
+      node = @router.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.
+      @router.try_delegate(node, :with, retry_count: retry_count) do |conn|
+        yield ::RedisClient::Cluster::Pinning::ConnectionProxy.new(conn, slot, @router, @command_builder)
+      end
+    end
+
     def close
       @concurrent_worker.close
       @router.close

lib/redis_client/cluster/pinning.rb

@@ -0,0 +1,123 @@
+# frozen_string_literal: true
+
+require 'redis_client'
+require 'delegate'
+
+class RedisClient
+  class Cluster
+    module Pinning
+      EMPTY_ARRAY = [].freeze
+
+      class BaseDelegator < SimpleDelegator
+        def initialize(conn, slot, router, command_builder)
+          @conn = conn
+          @slot = slot
+          @router = router
+          @command_builder = command_builder
+          super(@conn)
+        end
+
+        private
+
+        def ensure_key_not_cross_slot!(key, method_name)
+          slot = ::RedisClient::Cluster::KeySlotConverter.convert(key)
+          return unless slot != @slot
+
+          raise ::RedisClient::Cluster::Transaction::ConsistencyError.new,
+                "Connection pinned to slot #{@slot} but method #{method_name} called with key #{key} in slot #{slot}"
+        end
+
+        def ensure_command_not_cross_slot!(command)
+          command_keys = @router.command_extract_all_keys(command)
+          command_keys.each do |key|
+            slot = ::RedisClient::Cluster::KeySlotConverter.convert(key)
+            if slot != @slot
+              raise ::RedisClient::Cluster::Transaction::ConsistencyError.new,
+                    "Connection pinned to slot #{@slot} but command #{command.inspect} includes key #{key} in slot #{slot}"
+            end
+          end
+        end
+      end
+
+      module CallDelegator
+        def call(*command, **kwargs)
+          command = @command_builder.generate(command, kwargs)
+          ensure_command_not_cross_slot! command
+          super
+        end
+        alias call_once call
+
+        def call_v(command)
+          command = @command_builder.generate(command)
+          ensure_command_not_cross_slot! command
+          super
+        end
+        alias call_once_v call_v
+      end
+
+      module BlockingCallDelegator
+        def blocking_call(timeout, *command, **kwargs)
+          command = @command_builder.generate(command, kwargs)
+          ensure_command_not_cross_slot! command
+          super
+        end
+
+        def blocking_call_v(timeout, command)
+          command = @command_builder.generate(command)
+          ensure_command_not_cross_slot! command
+          super
+        end
+      end
+
+      class PipelineProxy < BaseDelegator
+        include CallDelegator
+        include BlockingCallDelegator
+      end
+
+      class MultiProxy < BaseDelegator
+        include CallDelegator
+      end
+
+      class ConnectionProxy < BaseDelegator
+        include CallDelegator
+        include BlockingCallDelegator
+
+        def sscan(key, *args, **kwargs, &block)
+          ensure_key_not_cross_slot! key, :sscan
+          super
+        end
+
+        def hscan(key, *args, **kwargs, &block)
+          ensure_key_not_cross_slot! key, :hscan
+          super
+        end
+
+        def zscan(key, *args, **kwargs, &block)
+          ensure_key_not_cross_slot! key, :zscan
+          super
+        end
+
+        def pipelined
+          @conn.pipelined do |conn_pipeline|
+            yield PipelineProxy.new(conn_pipeline, @slot, @router, @command_builder)
+          end
+        end
+
+        def multi(watch: nil)
+          call('WATCH', *watch) if watch
+          begin
+            @conn.pipelined do |conn_pipeline|
+              txn = MultiProxy.new(conn_pipeline, @slot, @router, @command_builder)
+              txn.call('MULTI')
+              yield txn
+              txn.call('EXEC')
+            end.last
+          rescue StandardError
+            call('UNWATCH') if connected? && watch
+            raise
+          end
+        end
+      end
+    end
+  end
+end

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