Add option to equally likely choose node to read between primary and all replicas (#254)

Author: holyketzer

README.md

@@ -24,7 +24,7 @@ gem 'redis-cluster-client'
 | --- | --- | --- | --- |
 | `:nodes` | String or Hash or Array<String, Hash> | `['redis://127.0.0.1:6379']` | node addresses for startup connection |
 | `:replica` | Boolean | `false` | `true` if client should use scale read feature |
-| `:replica_affinity` | Symbol or String | `:random` | scale reading strategy, `:random` or `:latency` are valid |
+| `:replica_affinity` | Symbol or String | `:random` | scale reading strategy, `:random`, `random_with_primary` or `:latency` are valid |
 | `:fixed_hostname` | String | `nil` | required if client should connect to single endpoint with SSL |
 
 Also, [the other generic options](https://github.com/redis-rb/redis-client#configuration) can be passed.
@@ -44,6 +44,10 @@ RedisClient.cluster.new_client
 RedisClient.cluster(replica: true).new_client
 #=> #<RedisClient::Cluster 172.20.0.2:6379, 172.20.0.3:6379, 172.20.0.4:6379, 172.20.0.5:6379, 172.20.0.6:6379, 172.20.0.7:6379>
 
+# To connect to all nodes to use scale reading feature + make reads equally likely from replicas and primary
+RedisClient.cluster(replica: true, replica_affinity: :random_with_primary).new_client
+#=> #<RedisClient::Cluster 172.20.0.2:6379, 172.20.0.3:6379, 172.20.0.4:6379, 172.20.0.5:6379, 172.20.0.6:6379, 172.20.0.7:6379>
+
 # To connect to all nodes to use scale reading feature prioritizing low-latency replicas
 RedisClient.cluster(replica: true, replica_affinity: :latency).new_client
 #=> #<RedisClient::Cluster 172.20.0.2:6379, 172.20.0.3:6379, 172.20.0.4:6379, 172.20.0.5:6379, 172.20.0.6:6379, 172.20.0.7:6379>

lib/redis_client/cluster/node.rb

@@ -5,6 +5,7 @@
 require 'redis_client/cluster/errors'
 require 'redis_client/cluster/node/primary_only'
 require 'redis_client/cluster/node/random_replica'
+require 'redis_client/cluster/node/random_replica_or_primary'
 require 'redis_client/cluster/node/latency_replica'
 
 class RedisClient
@@ -284,6 +285,8 @@ def update_slot(slot, node_key)
       def make_topology_class(with_replica, replica_affinity)
         if with_replica && replica_affinity == :random
           ::RedisClient::Cluster::Node::RandomReplica
+        elsif with_replica && replica_affinity == :random_with_primary
+          ::RedisClient::Cluster::Node::RandomReplicaOrPrimary
         elsif with_replica && replica_affinity == :latency
           ::RedisClient::Cluster::Node::LatencyReplica
         else

lib/redis_client/cluster/node/random_replica_or_primary.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+require 'redis_client/cluster/node/replica_mixin'
+
+class RedisClient
+  class Cluster
+    class Node
+      class RandomReplicaOrPrimary
+        include ::RedisClient::Cluster::Node::ReplicaMixin
+
+        def replica_clients
+          keys = @replications.values.filter_map(&:sample)
+          @clients.select { |k, _| keys.include?(k) }
+        end
+
+        def clients_for_scanning(seed: nil)
+          random = seed.nil? ? Random : Random.new(seed)
+          keys = @replications.map do |primary_node_key, replica_node_keys|
+            decide_use_primary?(random, replica_node_keys.size) ? primary_node_key : replica_node_keys.sample(random: random)
+          end
+
+          clients.select { |k, _| keys.include?(k) }
+        end
+
+        def find_node_key_of_replica(primary_node_key, seed: nil)
+          random = seed.nil? ? Random : Random.new(seed)
+
+          replica_node_keys = @replications.fetch(primary_node_key, EMPTY_ARRAY)
+          if decide_use_primary?(random, replica_node_keys.size)
+            primary_node_key
+          else
+            replica_node_keys.sample(random: random) || primary_node_key
+          end
+        end
+
+        def any_replica_node_key(seed: nil)
+          random = seed.nil? ? Random : Random.new(seed)
+          @replica_node_keys.sample(random: random) || any_primary_node_key(seed: seed)
+        end
+
+        private
+
+        # Randomly equally likely choose node to read between primary and all replicas
+        # e.g. 1 primary + 1 replica = 50% probability to read from primary
+        # e.g. 1 primary + 2 replica = 33% probability to read from primary
+        # e.g. 1 primary + 0 replica = 100% probability to read from primary
+        def decide_use_primary?(random, replica_nodes)
+          primary_nodes = 1.0
+          total = primary_nodes + replica_nodes
+          random.rand < primary_nodes / total
+        end
+      end
+    end
+  end
+end

test/bench_command.rb

@@ -34,6 +34,22 @@ def new_test_client
     end
   end
 
+  class ScaleReadRandomWithPrimary < BenchmarkWrapper
+    include BenchmarkMixin
+
+    private
+
+    def new_test_client
+      ::RedisClient.cluster(
+        nodes: TEST_NODE_URIS,
+        replica: true,
+        replica_affinity: :random_with_primary,
+        fixed_hostname: TEST_FIXED_HOSTNAME,
+        **TEST_GENERIC_OPTIONS
+      ).new_client
+    end
+  end
+
   class ScaleReadLatency < BenchmarkWrapper
     include BenchmarkMixin
 

test/redis_client/cluster/node/test_random_replica_or_primary.rb

@@ -0,0 +1,84 @@
+# frozen_string_literal: true
+
+require 'testing_helper'
+require 'redis_client/cluster/node/testing_topology_mixin'
+
+class RedisClient
+  class Cluster
+    class Node
+      class TestRandomReplicaWithPrimary < TestingWrapper
+        include TestingTopologyMixin
+
+        def test_clients_with_redis_client
+          got = @test_topology.clients
+          got.each_value { |client| assert_instance_of(::RedisClient, client) }
+          assert_equal(%w[master slave], got.map { |_, v| v.call('ROLE').first }.uniq.sort)
+        end
+
+        def test_clients_with_pooled_redis_client
+          test_topology = ::RedisClient::Cluster::Node::RandomReplicaOrPrimary.new(
+            @replications,
+            @options,
+            { timeout: 3, size: 2 },
+            **TEST_GENERIC_OPTIONS
+          )
+
+          got = test_topology.clients
+          got.each_value { |client| assert_instance_of(::RedisClient::Pooled, client) }
+          assert_equal(%w[master slave], got.map { |_, v| v.call('ROLE').first }.uniq.sort)
+        ensure
+          test_topology&.clients&.each_value(&:close)
+        end
+
+        def test_primary_clients
+          got = @test_topology.primary_clients
+          got.each_value do |client|
+            assert_instance_of(::RedisClient, client)
+            assert_equal('master', client.call('ROLE').first)
+          end
+        end
+
+        def test_replica_clients
+          got = @test_topology.replica_clients
+          got.each_value do |client|
+            assert_instance_of(::RedisClient, client)
+            assert_equal('slave', client.call('ROLE').first)
+          end
+        end
+
+        def test_clients_for_scanning
+          got = @test_topology.clients_for_scanning
+          got.each_value { |client| assert_instance_of(::RedisClient, client) }
+          assert_equal(TEST_SHARD_SIZE, got.size)
+        end
+
+        def test_find_node_key_of_replica
+          want = 'dummy_key'
+          got = @test_topology.find_node_key_of_replica('dummy_key')
+          assert_equal(want, got)
+
+          primary_key = @replications.keys.first
+          replica_keys = @replications.fetch(primary_key)
+          got = @test_topology.find_node_key_of_replica(primary_key)
+          assert_includes(replica_keys + [primary_key], got)
+        end
+
+        def test_any_primary_node_key
+          got = @test_topology.any_primary_node_key
+          assert_includes(@replications.keys, got)
+        end
+
+        def test_any_replica_node_key
+          got = @test_topology.any_replica_node_key
+          assert_includes(@replications.values.flatten, got)
+        end
+
+        private
+
+        def topology_class
+          ::RedisClient::Cluster::Node::RandomReplicaOrPrimary
+        end
+      end
+    end
+  end
+end

test/redis_client/test_cluster.rb

@@ -462,6 +462,21 @@ def new_test_client
       end
     end
 
+    class ScaleReadRandomWithPrimary < TestingWrapper
+      include Mixin
+
+      def new_test_client
+        config = ::RedisClient::ClusterConfig.new(
+          nodes: TEST_NODE_URIS,
+          replica: true,
+          replica_affinity: :random_with_primary,
+          fixed_hostname: TEST_FIXED_HOSTNAME,
+          **TEST_GENERIC_OPTIONS
+        )
+        ::RedisClient::Cluster.new(config)
+      end
+    end
+
     class ScaleReadLatency < TestingWrapper
       include Mixin
 

test/test_against_cluster_state.rb

@@ -129,6 +129,35 @@ def new_test_client
     end
   end
 
+  class ScaleReadRandomWithPrimary < TestingWrapper
+    include Mixin
+
+    def test_the_state_of_cluster_resharding
+      keys = nil
+      do_resharding_test { |ks| keys = ks }
+      keys.each { |key| assert_equal(key, @client.call('GET', key), "Case: GET: #{key}") }
+    end
+
+    def test_the_state_of_cluster_resharding_with_pipelining
+      keys = nil
+      do_resharding_test { |ks| keys = ks }
+      values = @client.pipelined { |pipeline| keys.each { |key| pipeline.call('GET', key) } }
+      keys.each_with_index { |key, i| assert_equal(key, values[i], "Case: GET: #{key}") }
+    end
+
+    private
+
+    def new_test_client
+      ::RedisClient.cluster(
+        nodes: TEST_NODE_URIS,
+        replica: true,
+        replica_affinity: :random_with_primary,
+        fixed_hostname: TEST_FIXED_HOSTNAME,
+        **TEST_GENERIC_OPTIONS
+      ).new_client
+    end
+  end
+
   class ScaleReadLatency < TestingWrapper
     include Mixin