fix: Pub/Sub command bugs (#284)

Author: supercaracal

lib/redis_client/cluster/pub_sub.rb

@@ -52,10 +52,12 @@ def initialize(router, command_builder)
 
       def call(*args, **kwargs)
         _call(@command_builder.generate(args, kwargs))
+        nil
       end
 
       def call_v(command)
         _call(@command_builder.generate(command))
+        nil
       end
 
       def close
@@ -86,10 +88,31 @@ def next_event(timeout = nil)
       private
 
       def _call(command)
+        case ::RedisClient::Cluster::NormalizedCmdName.instance.get_by_command(command)
+        when 'subscribe', 'psubscribe', 'ssubscribe' then call_to_single_state(command)
+        when 'unsubscribe', 'punsubscribe' then call_to_all_states(command)
+        when 'sunsubscribe' then call_for_sharded_states(command)
+        else call_to_single_state(command)
+        end
+      end
+
+      def call_to_single_state(command)
         node_key = @router.find_node_key(command)
         try_call(node_key, command)
       end
 
+      def call_to_all_states(command)
+        @state_dict.each_value { |s| s.call_v(command) }
+      end
+
+      def call_for_sharded_states(command)
+        if command.size == 1
+          call_to_all_states(command)
+        else
+          call_to_single_state(command)
+        end
+      end
+
       def try_call(node_key, command, retry_count: 1)
         add_state(node_key).call(command)
       rescue ::RedisClient::CommandError => e

lib/redis_client/cluster/router.rb

@@ -288,11 +288,18 @@ def send_script_command(method, command, args, &block) # rubocop:disable Metrics
 
       def send_pubsub_command(method, command, args, &block) # rubocop:disable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
         case ::RedisClient::Cluster::NormalizedCmdName.instance.get_by_subcommand(command)
-        when 'channels' then @node.call_all(method, command, args).flatten.uniq.sort_by(&:to_s).then(&TSF.call(block))
+        when 'channels'
+          @node.call_all(method, command, args).flatten.uniq.sort_by(&:to_s).then(&TSF.call(block))
+        when 'shardchannels'
+          @node.call_replicas(method, command, args).flatten.uniq.sort_by(&:to_s).then(&TSF.call(block))
+        when 'numpat'
+          @node.call_all(method, command, args).select { |e| e.is_a?(Integer) }.sum.then(&TSF.call(block))
         when 'numsub'
           @node.call_all(method, command, args).reject(&:empty?).map { |e| Hash[*e] }
                .reduce({}) { |a, e| a.merge(e) { |_, v1, v2| v1 + v2 } }.then(&TSF.call(block))
-        when 'numpat' then @node.call_all(method, command, args).select { |e| e.is_a?(Integer) }.sum.then(&TSF.call(block))
+        when 'shardnumsub'
+          @node.call_replicas(method, command, args).reject(&:empty?).map { |e| Hash[*e] }
+               .reduce({}) { |a, e| a.merge(e) { |_, v1, v2| v1 + v2 } }.then(&TSF.call(block))
         else assign_node(command).public_send(method, *args, command, &block)
         end
       end

test/redis_client/test_cluster.rb

@@ -273,6 +273,7 @@ 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
 
@@ -289,6 +290,7 @@ def test_global_pubsub_without_timeout
           assert_equal(want, got)
           Fiber.yield('my-global-published-channel')
           Fiber.yield(collect_messages(pubsub, size: 1, timeout: nil).first)
+          pubsub.call('UNSUBSCRIBE')
           pubsub.close
         end
 
@@ -304,6 +306,7 @@ def test_global_pubsub_with_multiple_channels
           10.times { |i| assert_equal(['subscribe', "g-chan#{i}", i + 1], got[i]) }
           Fiber.yield
           Fiber.yield(collect_messages(pubsub, size: 10))
+          pubsub.call('UNSUBSCRIBE')
           pubsub.close
         end
 
@@ -325,6 +328,7 @@ 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
 
@@ -347,6 +351,7 @@ def test_sharded_pubsub_without_timeout
           assert_equal(want, got)
           Fiber.yield('my-sharded-published-channel')
           Fiber.yield(collect_messages(pubsub, size: 1, timeout: nil).first)
+          pubsub.call('SUNSUBSCRIBE')
           pubsub.close
         end
 
@@ -367,6 +372,7 @@ def test_sharded_pubsub_with_multiple_channels
           10.times { |i| assert_equal(['ssubscribe', "s-chan#{i}"], got[i].take(2)) }
           Fiber.yield
           Fiber.yield(collect_messages(pubsub, size: 10))
+          pubsub.call('SUNSUBSCRIBE')
           pubsub.close
         end
 
@@ -376,6 +382,19 @@ def test_sharded_pubsub_with_multiple_channels
         10.times { |i| assert_equal(['smessage', "s-chan#{i}", i.to_s], got[i]) }
       end
 
+      def test_other_pubsub_commands
+        assert_instance_of(Array, @client.call('pubsub', 'channels'))
+        assert_instance_of(Integer, @client.call('pubsub', 'numpat'))
+        assert_instance_of(Hash, @client.call('pubsub', 'numsub'))
+        assert_instance_of(Array, @client.call('pubsub', 'shardchannels')) if TEST_REDIS_MAJOR_VERSION >= 7
+        assert_instance_of(Hash, @client.call('pubsub', 'shardnumsub')) if TEST_REDIS_MAJOR_VERSION >= 7
+        ps = @client.pubsub
+        assert_nil(ps.call('unsubscribe'))
+        assert_nil(ps.call('punsubscribe'))
+        assert_nil(ps.call('sunsubscribe')) if TEST_REDIS_MAJOR_VERSION >= 7
+        ps.close
+      end
+
       def test_dedicated_commands # rubocop:disable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
         10.times { |i| @client.call('SET', "key#{i}", i) }
         wait_for_replication