Implement validation of slot numbers in #with

This is achieved by having the #with method "lock" the yielded
connection to a particular slot, and having middleware which implements
that locking.

Author: KJ Tsanaktsidis

lib/redis_client/cluster.rb

@@ -93,15 +93,19 @@ def multi(watch: nil, &block)
       ::RedisClient::Cluster::Transaction.new(@router, @command_builder).execute(watch: watch, &block)
     end
 
-    def with(key: nil, hashtag: nil, write: true, retry_count: 0, &block)
+    def with(key: nil, hashtag: nil, write: true, retry_count: 0)
       key = process_with_arguments(key, hashtag)
 
       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, &block)
+      @router.try_delegate(node, :with, retry_count: retry_count) do |conn|
+        conn.locked_to_key_slot(key) do
+          yield conn
+        end
+      end
     end
 
     def pubsub

lib/redis_client/cluster/node.rb

@@ -8,6 +8,7 @@
 require 'redis_client/cluster/node/random_replica'
 require 'redis_client/cluster/node/random_replica_or_primary'
 require 'redis_client/cluster/node/latency_replica'
+require 'redis_client/cluster/pinning'
 
 class RedisClient
   class Cluster
@@ -79,13 +80,15 @@ def []=(index, element)
       end
 
       class SingleNodeRedisClient < ::RedisClient
+        include Pinning::ClientMixin
       end
 
       class Config < ::RedisClient::Config
         def initialize(cluster_commands:, scale_read: false, middlewares: nil, **kwargs)
           @scale_read = scale_read
           @cluster_commands = cluster_commands
           middlewares ||= []
+          middlewares.unshift Pinning::ClientMiddleware
           middlewares.unshift ErrorIdentification::Middleware
           super(
             middlewares: middlewares,

lib/redis_client/cluster/pinning.rb

@@ -0,0 +1,59 @@
+# frozen_string_literal: true
+
+class RedisClient
+  class Cluster
+    module Pinning
+      module ClientMixin
+        attr_reader :locked_key_slot
+
+        # This gets called when handing out connections in Cluster#with to lock the returned
+        # connections to a given slot.
+        def locked_to_key_slot(key_slot)
+          raise ArgumentError, 'recursive slot locking is not allowed' if @locked_key_slot
+
+          begin
+            @locked_key_slot = key_slot
+            yield
+          ensure
+            @locked_key_slot = nil
+          end
+        end
+      end
+
+      # This middleware is what actually enforces the slot locking above.
+      module ClientMiddleware
+        def initialize(client)
+          @client = client
+          super
+        end
+
+        def assert_slot_valid!(command, config) # rubocop:disable Metrics/AbcSize
+          return unless @client.locked_key_slot
+          return unless config.cluster_commands.loaded?
+
+          keys = config.cluster_commands.extract_all_keys(command)
+          key_slots = keys.map { |k| ::RedisClient::Cluster::KeySlotConverter.convert(k) }
+          locked_slot = ::RedisClient::Cluster::KeySlotConverter.convert(@client.locked_key_slot)
+          return if key_slots.all? { |slot| slot ==  locked_slot }
+
+          key_slot_pairs = keys.zip(key_slots).map { |key, slot| "#{key} => #{slot}" }.join(', ')
+          raise ::RedisClient::Cluster::Transaction::ConsistencyError, <<~MESSAGE
+            Connection is pinned to slot #{locked_slot} (via key #{@client.locked_key_slot}). \
+            However, command #{command.inspect} has keys hashing to slots #{key_slot_pairs}. \
+            Transactions in redis cluster must only refer to keys hashing to the same slot.
+          MESSAGE
+        end
+
+        def call(command, config)
+          assert_slot_valid!(command, config)
+          super
+        end
+
+        def call_pipelined(command, config)
+          assert_slot_valid!(command, config)
+          super
+        end
+      end
+    end
+  end
+end

test/redis_client/test_cluster.rb

@@ -533,8 +533,6 @@ def test_pinning_two_keys
       end
 
       def test_pinning_cross_slot
-        skip 'This is not implemented yet!'
-
         assert_raises(::RedisClient::Cluster::Transaction::ConsistencyError) do
           @client.with(hashtag: 'slot1') do |conn|
             conn.call('GET', '{slot2}')