fix: enbuged for Pub/Sub (#222)

Author: supercaracal

lib/redis_client/cluster/pub_sub.rb

@@ -3,13 +3,45 @@
 class RedisClient
   class Cluster
     class PubSub
-      MAX_THREADS = Integer(ENV.fetch('REDIS_CLIENT_MAX_THREADS', 5))
+      class State
+        def initialize(client)
+          @client = client
+          @worker = nil
+        end
+
+        def call(command)
+          @client.call_v(command)
+        end
+
+        def close
+          @worker.exit if @worker&.alive?
+          @client.close
+        end
+
+        def take_message(timeout)
+          @worker = subscribe(@client, timeout) if @worker.nil?
+          return if @worker.join(0.01).nil?
+
+          message = @worker[:reply]
+          @worker = nil
+          message
+        end
+
+        private
+
+        def subscribe(client, timeout)
+          Thread.new(client, timeout) do |pubsub, to|
+            Thread.current[:reply] = pubsub.next_event(to)
+          rescue StandardError => e
+            Thread.current[:reply] = e
+          end
+        end
+      end
 
       def initialize(router, command_builder)
         @router = router
         @command_builder = command_builder
-        @pubsub_states = {}
-        @messages = []
+        @states = {}
       end
 
       def call(*args, **kwargs)
@@ -21,46 +53,39 @@ def call_v(command)
       end
 
       def close
-        @pubsub_states.each_value(&:close)
-        @pubsub_states.clear
-        @messages.clear
+        @states.each_value(&:close)
+        @states.clear
       end
 
       def next_event(timeout = nil)
-        return if @pubsub_states.empty?
-        return @messages.shift unless @messages.empty?
+        return if @states.empty?
 
-        collect_messages(timeout)
-        @messages.shift
+        max_duration = calc_max_duration(timeout)
+        starting = obtain_current_time
+        loop do
+          break if max_duration > 0 && obtain_current_time - starting > max_duration
+
+          @states.each_value do |pubsub|
+            message = pubsub.take_message(timeout)
+            return message if message
+          end
+        end
       end
 
       private
 
       def _call(command)
         node_key = @router.find_node_key(command)
-        pubsub = if @pubsub_states.key?(node_key)
-                   @pubsub_states[node_key]
-                 else
-                   @pubsub_states[node_key] = @router.find_node(node_key).pubsub
-                 end
-        pubsub.call_v(command)
+        @states[node_key] = State.new(@router.find_node(node_key).pubsub) unless @states.key?(node_key)
+        @states[node_key].call(command)
       end
 
-      def collect_messages(timeout)
-        @pubsub_states.each_slice(MAX_THREADS) do |chuncked_pubsub_states|
-          threads = chuncked_pubsub_states.map do |_, v|
-            Thread.new(v) do |pubsub|
-              Thread.current[:reply] = pubsub.next_event(timeout)
-            rescue StandardError => e
-              Thread.current[:reply] = e
-            end
-          end
+      def obtain_current_time
+        Process.clock_gettime(Process::CLOCK_MONOTONIC, :microsecond)
+      end
 
-          threads.each do |t|
-            t.join
-            @messages << t[:reply] unless t[:reply].nil?
-          end
-        end
+      def calc_max_duration(timeout)
+        timeout.nil? || timeout < 0 ? 0 : timeout * 1_000_000
       end
     end
   end

test/redis_client/test_cluster.rb

@@ -194,16 +194,28 @@ def test_global_pubsub
           assert_equal(['subscribe', channel, 1], pubsub.next_event(TEST_TIMEOUT_SEC))
           Fiber.yield(channel)
           Fiber.yield(pubsub.next_event(TEST_TIMEOUT_SEC))
-          pubsub.call('UNSUBSCRIBE')
           pubsub.close
         end
 
         channel = sub.resume(@client.pubsub)
-        publish_messages do |cli|
-          cli.call('PUBLISH', channel, 'hello global world')
+        publish_messages { |cli| cli.call('PUBLISH', channel, 'hello global world') }
+        assert_equal(['message', channel, 'hello global world'], sub.resume)
+      end
+
+      def test_global_pubsub_without_timeout
+        sub = Fiber.new do |pubsub|
+          pubsub.call('SUBSCRIBE', 'my-global-not-published-channel', 'my-global-published-channel')
+          want = [%w[subscribe my-global-not-published-channel], %w[subscribe my-global-published-channel]]
+          got = collect_messages(pubsub, size: 2, timeout: nil).map { |e| e.take(2) }.sort_by { |e| e[1].to_s }
+          assert_equal(want, got)
+          Fiber.yield('my-global-published-channel')
+          Fiber.yield(collect_messages(pubsub, size: 1, timeout: nil).first)
+          pubsub.close
         end
 
-        assert_equal(['message', channel, 'hello global world'], sub.resume)
+        channel = sub.resume(@client.pubsub)
+        publish_messages { |cli| cli.call('PUBLISH', channel, 'hello global published world') }
+        assert_equal(['message', channel, 'hello global published world'], sub.resume)
       end
 
       def test_global_pubsub_with_multiple_channels
@@ -214,19 +226,15 @@ def test_global_pubsub_with_multiple_channels
 
         sub = Fiber.new do |pubsub|
           pubsub.call('SUBSCRIBE', *Array.new(10) { |i| "g-chan#{i}" })
-          got = collect_messages(pubsub).sort_by { |e| e[1].to_s }
+          got = collect_messages(pubsub, size: 10).sort_by { |e| e[1].to_s }
           10.times { |i| assert_equal(['subscribe', "g-chan#{i}", i + 1], got[i]) }
           Fiber.yield
-          Fiber.yield(collect_messages(pubsub))
-          pubsub.call('UNSUBSCRIBE')
+          Fiber.yield(collect_messages(pubsub, size: 10))
           pubsub.close
         end
 
         sub.resume(@client.pubsub)
-        publish_messages do |cli|
-          cli.pipelined { |pi| 10.times { |i| pi.call('PUBLISH', "g-chan#{i}", i) } }
-        end
-
+        publish_messages { |cli| cli.pipelined { |pi| 10.times { |i| pi.call('PUBLISH', "g-chan#{i}", i) } } }
         got = sub.resume.sort_by { |e| e[1].to_s }
         10.times { |i| assert_equal(['message', "g-chan#{i}", i.to_s], got[i]) }
       end
@@ -243,16 +251,34 @@ def test_sharded_pubsub
           assert_equal(['ssubscribe', channel, 1], pubsub.next_event(TEST_TIMEOUT_SEC))
           Fiber.yield(channel)
           Fiber.yield(pubsub.next_event(TEST_TIMEOUT_SEC))
-          pubsub.call('SUNSUBSCRIBE')
           pubsub.close
         end
 
         channel = sub.resume(@client.pubsub)
-        publish_messages do |cli|
-          cli.call('SPUBLISH', channel, 'hello sharded world')
+        publish_messages { |cli| cli.call('SPUBLISH', channel, 'hello sharded world') }
+        assert_equal(['smessage', channel, 'hello sharded world'], sub.resume)
+      end
+
+      def test_sharded_pubsub_without_timeout
+        if TEST_REDIS_MAJOR_VERSION < 7
+          skip('Sharded Pub/Sub is supported by Redis 7+.')
+          return
         end
 
-        assert_equal(['smessage', channel, 'hello sharded world'], sub.resume)
+        sub = Fiber.new do |pubsub|
+          pubsub.call('SSUBSCRIBE', 'my-sharded-not-published-channel')
+          pubsub.call('SSUBSCRIBE', 'my-sharded-published-channel')
+          want = [%w[ssubscribe my-sharded-not-published-channel], %w[ssubscribe my-sharded-published-channel]]
+          got = collect_messages(pubsub, size: 2, timeout: nil).map { |e| e.take(2) }.sort_by { |e| e[1].to_s }
+          assert_equal(want, got)
+          Fiber.yield('my-sharded-published-channel')
+          Fiber.yield(collect_messages(pubsub, size: 1, timeout: nil).first)
+          pubsub.close
+        end
+
+        channel = sub.resume(@client.pubsub)
+        publish_messages { |cli| cli.call('SPUBLISH', channel, 'hello sharded published world') }
+        assert_equal(['smessage', channel, 'hello sharded published world'], sub.resume)
       end
 
       def test_sharded_pubsub_with_multiple_channels
@@ -268,19 +294,15 @@ def test_sharded_pubsub_with_multiple_channels
 
         sub = Fiber.new do |pubsub|
           10.times { |i| pubsub.call('SSUBSCRIBE', "s-chan#{i}") }
-          got = collect_messages(pubsub).sort_by { |e| e[1].to_s }
+          got = collect_messages(pubsub, size: 10).sort_by { |e| e[1].to_s }
           10.times { |i| assert_equal(['ssubscribe', "s-chan#{i}"], got[i].take(2)) }
           Fiber.yield
-          Fiber.yield(collect_messages(pubsub))
-          pubsub.call('SUNSUBSCRIBE')
+          Fiber.yield(collect_messages(pubsub, size: 10))
           pubsub.close
         end
 
         sub.resume(@client.pubsub)
-        publish_messages do |cli|
-          cli.pipelined { |pi| 10.times { |i| pi.call('SPUBLISH', "s-chan#{i}", i) } }
-        end
-
+        publish_messages { |cli| cli.pipelined { |pi| 10.times { |i| pi.call('SPUBLISH', "s-chan#{i}", i) } } }
         got = sub.resume.sort_by { |e| e[1].to_s }
         10.times { |i| assert_equal(['smessage', "s-chan#{i}", i.to_s], got[i]) }
       end
@@ -386,7 +408,7 @@ def wait_for_replication
         @client&.blocking_call(client_side_timeout, 'WAIT', TEST_REPLICA_SIZE, server_side_timeout)
       end
 
-      def collect_messages(pubsub, max_attempts: 30, timeout: 1.0)
+      def collect_messages(pubsub, size:, max_attempts: 30, timeout: 1.0)
         messages = []
         attempts = 0
         loop do
@@ -396,14 +418,9 @@ def collect_messages(pubsub, max_attempts: 30, timeout: 1.0)
           reply = pubsub.next_event(timeout)
           break if reply.nil?
 
-          if reply.first.is_a?(Array)
-            messages += reply
-          else
-            messages << reply
-          end
+          messages << reply
+          break messages if messages.size == size
         end
-
-        messages
       end
 
       def publish_messages