Implement RedisClient::Cluster#with

KJ Tsanaktsidis · KJ Tsanaktsidis · commit 3b5adf4c87cd · 2024-01-16T15:23:53.000+11:00
This PR actually implements the #with method for checking out a
connection directly to a particular Redis node. This allows you to
perform transactional commands directly against the underlying
`RedisClient` instance, so `#watch`, `#multi`, etc all work.

This does NOT yet implement any kind of validation for the commands
issued. If you issue commands involving keys not owned by the node, you
will receive `Redis::CommandError`'s. A future PR will add client-side
validation to hopefully make it possible to catch problems related to
transaction misuse earlier in development. There's a pending test for
this.

This PR also does not yet implement `RedisClient::Cluster#multi` in
terms of `#with`. That will also come in a future PR.
diff --git a/README.md b/README.md
@@ -168,6 +168,84 @@ 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`, specifying the slot by providing a key which hashes to that slot in the `slot_key` argument:
+
+```ruby
+cli.with(slot_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"
+end
+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(slot_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(slot_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(slot_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(slot_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(slot_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.
diff --git a/lib/redis_client/cluster.rb b/lib/redis_client/cluster.rb
@@ -93,6 +93,17 @@ def multi(watch: nil, &block)
       ::RedisClient::Cluster::Transaction.new(@router, @command_builder).execute(watch: watch, &block)
     end
 
+    def with(slot_key:, write: true, retry_count: 0, &block)
+      raise ArgumentError, 'slot_key must be provided' if slot_key.nil? || slot_key.empty?
+
+      node_key = @router.find_node_key_by_key(slot_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)
+    end
+
     def pubsub
       ::RedisClient::Cluster::PubSub.new(@router, @command_builder)
     end
diff --git a/test/redis_client/test_cluster.rb b/test/redis_client/test_cluster.rb
@@ -503,6 +503,207 @@ def test_only_reshards_own_errors
         assert_equal correct_primary_key, router.find_node_key_by_key('testkey', primary: true)
       end
 
+      def test_pinning_single_key
+        got = @client.with(slot_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(slot_key: nil) {}
+        end
+      end
+
+      def test_pinning_two_keys
+        got = @client.with(slot_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
+        skip 'This is not implemented yet!'
+
+        assert_raises(::RedisClient::Cluster::Transaction::ConsistencyError) do
+          @client.with(slot_key: '{slot1}') do |conn|
+            conn.call('GET', '{slot2}')
+          end
+        end
+      end
+
+      def test_pinning_pipeline
+        got = @client.with(slot_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(slot_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(slot_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(slot_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.to_a.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(slot_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.to_a.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(slot_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.to_a.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(slot_key: '{slot}') do |conn|
+          conn.call('WATCH', '{slot}key1')
+          conn.call('UNWATCH')
+        end
+
+        assert_equal('OK', got)
+      end
+
+      def test_pinning_timeouts_update_topology
+        # Create a new test client with a lower timeout for this test so it's fast.
+        captured_commands = CommandCaptureMiddleware::CommandBuffer.new
+        client = new_test_client(capture_buffer: captured_commands, timeout: 0.5)
+
+        client.call('DEL', '{slot}list')
+        captured_commands.clear
+
+        assert_raises(::RedisClient::ReadTimeoutError) do
+          client.with(slot_key: '{slot}') do |conn|
+            conn.call_v(['BLPOP', '{slot}list', 0])
+          end
+        end
+        assert_includes(captured_commands.to_a.map(&:command), %w[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(slot_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(slot_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
