feat(pubsub): sharded Pub/Sub for cluster mode

CHANGELOG.md

@@ -30,6 +30,7 @@
   * `ft_search` is unsupported inside `pipelined`/`multi` (a queued command yields a `Valkey::Future`, not a reply) and with the `flatten_map: true` compatibility option; both raise `ArgumentError`. Use `call` in those cases.
 * Ruby: PubSub: added support for `subscribe(*channels, timeout_ms:)`, `unsubscribe(*channels, timeout_ms:)`, `publish(message, channel)`, `get_pubsub_message`, `try_get_pubsub_message`, and the connection-time `pubsub: { subscriptions: { exact: [...] } }` option. Requires `protocol: :resp3` ([#308](https://github.com/valkey-io/valkey-glide-ruby/pull/308))
 * Ruby: PubSub: added support for `psubscribe(*patterns, timeout_ms:)`, `punsubscribe(*patterns, timeout_ms:)`, `subscribe_lazy(*channels)`, `unsubscribe_lazy(*channels)`, `psubscribe_lazy(*patterns)`, `punsubscribe_lazy(*patterns)`, and the `pubsub: { callback:, context: }` connection option for callback mode delivery ([#316](https://github.com/valkey-io/valkey-glide-ruby/pull/316))
+* Ruby: PubSub: added support for Sharded pubsub commands([#317](https://github.com/valkey-io/valkey-glide-ruby/pull/317))
 * Ruby: fixed cd workflow to correctly build the ffi with **glibc 2.17** ([#223](https://github.com/valkey-io/valkey-glide-ruby/issues/223))
 * Ruby: scripting commands now dispatch real `EVAL` / `EVALSHA` / `SCRIPT LOAD` to the server instead of a client-side script container ([#213](https://github.com/valkey-io/valkey-glide-ruby/issues/213)). Three behavior changes:
   * `eval` / `evalsha` (and the `_ro` variants) now accept the standard integer key-count form used by `valkey-cli` and the Valkey docs — `eval(script, 1, "mykey", "myarg")`. It previously made the count `KEYS[1]`, shifted the real key into `ARGV[1]`, and dropped the remaining arguments without raising.

lib/valkey.rb

@@ -325,7 +325,8 @@ def initialize(options = {})
       }
     end
 
-    pubsub_config = parse_pubsub_configs(options[:pubsub], protocol: options[:protocol])
+    pubsub_config = parse_pubsub_configs(options[:pubsub], protocol: options[:protocol],
+                                                           cluster_mode: options[:cluster_mode] ? true : false)
     json_options.merge!(pubsub_config)
 
     @pubsub_receiver = Valkey::Glide::PubSubReceiver.make(pubsub_configs: options[:pubsub])
@@ -373,6 +374,13 @@ def initialize(options = {})
     @queued_commands = []
   end
 
+  # True if client is in cluster mode.
+  #
+  # @return [Boolean]
+  def cluster_mode?
+    @cluster_mode
+  end
+
   # Closes the client and frees the native connection.
   def close
     return unless @close_lock&.try_lock
@@ -876,19 +884,23 @@ def convert_response(res, &block)
   #     callback: ->(message, context) { ... },  # callback handler
   #     context: my_app_state                   # callback context
   #   }
-  def parse_pubsub_configs(pubsub_configs, protocol: nil)
+  def parse_pubsub_configs(pubsub_configs, protocol: nil, cluster_mode: false)
     subscriptions = (pubsub_configs || {})[:subscriptions] || {}
     return {} if subscriptions.empty?
 
-    validate_pubsub_subscriptions!(subscriptions, protocol: protocol)
+    validate_pubsub_subscriptions!(subscriptions, protocol: protocol, cluster_mode: cluster_mode)
 
     { "pubsub_subscriptions" => pubsub_subscriptions_to_ffi(subscriptions) }
   end
 
-  def validate_pubsub_subscriptions!(subscriptions, protocol:)
+  def validate_pubsub_subscriptions!(subscriptions, protocol:, cluster_mode: false)
     unknown_modes = subscriptions.keys - SUBSCRIPTION_MODES.keys
     raise ArgumentError, unknown_pubsub_mode_message(unknown_modes) if unknown_modes.any?
     raise Resp3RequiredError, protocol unless RESP3_VALUES.include?(protocol)
+
+    return unless Array(subscriptions[:sharded]).any? && !cluster_mode
+
+    raise ArgumentError, "Sharded Pub/Sub subscriptions are only available in cluster mode."
   end
 
   def pubsub_subscriptions_to_ffi(subscriptions)

lib/valkey/commands/pubsub_commands.rb

@@ -160,12 +160,18 @@ def punsubscribe(*patterns, timeout_ms: 0)
       # @param [Integer, Float] timeout_ms maximum time in milliseconds to wait for the server to
       #   confirm; `0` blocks indefinitely
       # @return [void] returns once the server has confirmed the subscription
-      # @raise [ArgumentError] if timeout_ms is negative
+      # @raise [ArgumentError] if the client is not in cluster mode, or if timeout_ms is negative
+      # @raise [Valkey::Resp3RequiredError] GLIDE Pub/Sub requires RESP3
       # @raise [Valkey::TimeoutError] if the timeout expires before the server confirms
-      # @raise [NotImplementedError] this method is not implemented yet
       #
       # @see https://valkey.io/commands/ssubscribe/
-      def ssubscribe(*channels, timeout_ms: 0) = raise(NotImplementedError, "#{__method__} is not implemented yet")
+      def ssubscribe(*channels, timeout_ms: 0)
+        validate_resp3!
+        validate_cluster_mode!(__method__)
+        raise ArgumentError, "No channels provided for subscription" if channels.empty?
+
+        send_command(RequestType::SSUBSCRIBE_BLOCKING, channels.map(&:to_s) + [parse_timeout(timeout_ms)])
+      end
 
       # Unsubscribe from sharded channels, waiting for the server to confirm the change.
       #
@@ -181,12 +187,17 @@ def ssubscribe(*channels, timeout_ms: 0) = raise(NotImplementedError, "#{__metho
       # @param [Integer, Float] timeout_ms maximum time in milliseconds to wait for the server to
       #   confirm; `0` blocks indefinitely
       # @return [void] returns once the server has confirmed the change
-      # @raise [ArgumentError] if timeout_ms is negative
+      # @raise [ArgumentError] if the client is not in cluster mode, or if timeout_ms is negative
+      # @raise [Valkey::Resp3RequiredError] GLIDE Pub/Sub requires RESP3
       # @raise [Valkey::TimeoutError] if the timeout expires before the server confirms
-      # @raise [NotImplementedError] this method is not implemented yet
       #
       # @see https://valkey.io/commands/sunsubscribe/
-      def sunsubscribe(*channels, timeout_ms: 0) = raise(NotImplementedError, "#{__method__} is not implemented yet")
+      def sunsubscribe(*channels, timeout_ms: 0)
+        validate_resp3!
+        validate_cluster_mode!(__method__)
+
+        send_command(RequestType::SUNSUBSCRIBE_BLOCKING, channels.map(&:to_s) + [parse_timeout(timeout_ms)])
+      end
 
       # Subscribe to exact channels without waiting for the server to confirm.
       #
@@ -272,10 +283,14 @@ def punsubscribe_lazy(*patterns)
       #
       # @param [Array<String>] channels the sharded channels to subscribe to; an empty list is rejected
       # @return [void] returns as soon as the desired subscription state is updated
-      # @raise [NotImplementedError] this method is not implemented yet
+      # @raise [ArgumentError] if the client is not in cluster mode, or if the channel list is empty
+      # @raise [Valkey::Resp3RequiredError] GLIDE Pub/Sub requires RESP3
       #
       # @see https://valkey.io/commands/ssubscribe/
-      def ssubscribe_lazy(*channels) = raise(NotImplementedError, "#{__method__} is not implemented yet")
+      def ssubscribe_lazy(*channels)
+        validate_cluster_mode!(__method__)
+        send_lazy_subscription(RequestType::SSUBSCRIBE, channels, reject_empty: true)
+      end
 
       # Unsubscribe from sharded channels without waiting for the server to confirm.
       #
@@ -289,10 +304,14 @@ def ssubscribe_lazy(*channels) = raise(NotImplementedError, "#{__method__} is no
       # @param [Array<String>] channels the sharded channels to unsubscribe from; an empty list unsubscribes
       #   from all sharded channels
       # @return [void] returns as soon as the desired subscription state is updated
-      # @raise [NotImplementedError] this method is not implemented yet
+      # @raise [ArgumentError] if the client is not in cluster mode
+      # @raise [Valkey::Resp3RequiredError] GLIDE Pub/Sub requires RESP3
       #
       # @see https://valkey.io/commands/sunsubscribe/
-      def sunsubscribe_lazy(*channels) = raise(NotImplementedError, "#{__method__} is not implemented yet")
+      def sunsubscribe_lazy(*channels)
+        validate_cluster_mode!(__method__)
+        send_lazy_subscription(RequestType::SUNSUBSCRIBE, channels)
+      end
 
       # Publish a message on a Pub/Sub channel.
       #
@@ -314,14 +333,12 @@ def sunsubscribe_lazy(*channels) = raise(NotImplementedError, "#{__method__} is
       # @return [Integer] the number of subscriptions that received the message: in cluster mode the
       #   subscriptions on the node the request was routed to, in standalone the subscriptions on the primary
       #   node, which excludes subscriptions configured on replicas
-      # @raise [NotImplementedError] sharded publish is not implemented yet
       #
       # @see https://valkey.io/commands/publish/
       # @see https://valkey.io/commands/spublish/
       def publish(message, channel, sharded: false)
-        raise NotImplementedError, "Sharded publish is not implemented yet" if sharded
-
-        send_command(RequestType::PUBLISH, [channel.to_s, message.to_s])
+        request_type = sharded ? RequestType::SPUBLISH : RequestType::PUBLISH
+        send_command(request_type, [channel.to_s, message.to_s])
       end
 
       # Get this connection's subscription state: what the client asked for and what the server confirmed.
@@ -474,6 +491,12 @@ def validate_resp3!
         raise Resp3RequiredError, protocol unless RESP3_VALUES.include?(protocol)
       end
 
+      def validate_cluster_mode!(command)
+        return if cluster_mode?
+
+        raise ArgumentError, "#{command} is only available in cluster mode."
+      end
+
       def send_lazy_subscription(request_type, channels, reject_empty: false, noun: "channels")
         validate_resp3!
         raise ArgumentError, "No #{noun} provided for subscription" if reject_empty && channels.empty?

test/integration/valkey/pubsub_test.rb

@@ -461,8 +461,189 @@ def test_callback_exception_is_contained
       subscriber&.close
     end
 
+    # --- Sharded Pub/Sub (cluster mode) --------------------------------------
+
+    def test_sharded_message_round_trip
+      skip_unless_sharded_pubsub
+
+      channel = unique_channel
+
+      with_client do |subscriber|
+        subscriber.ssubscribe(channel)
+        r.publish("sharded-msg", channel, sharded: true)
+        received = wait_for_message(subscriber)
+
+        assert_equal "sharded-msg", received.message
+        assert_equal channel,       received.channel
+        assert_nil                  received.pattern
+      end
+    end
+
+    def test_spublish_returns_the_receiver_count
+      skip_unless_sharded_pubsub
+
+      channel = unique_channel
+
+      with_client do |subscriber|
+        assert_equal 0, r.publish("nobody", channel, sharded: true)
+
+        subscriber.ssubscribe(channel)
+        assert_equal 1, r.publish("counted", channel, sharded: true)
+      end
+    end
+
+    def test_sunsubscribe_stops_sharded_delivery
+      skip_unless_sharded_pubsub
+
+      channel = unique_channel
+
+      with_client do |subscriber|
+        subscriber.ssubscribe(channel)
+        r.publish("before", channel, sharded: true)
+        assert_equal "before", wait_for_message(subscriber).message
+
+        subscriber.sunsubscribe(channel)
+
+        assert_equal 0, r.publish("after", channel, sharded: true)
+        assert_nil wait_for_message(subscriber, timeout: UNSUB_WAIT_TIME)
+      end
+    end
+
+    def test_sunsubscribe_all_sharded_channels
+      skip_unless_sharded_pubsub
+
+      # Same hash tag keeps both channels in one slot, so a single ssubscribe
+      # call is routed to one node and both land in the actual subscriptions.
+      channels = Array.new(2) { |i| unique_channel("{shardtag}-#{i}") }
+
+      with_client do |subscriber|
+        subscriber.ssubscribe(*channels)
+
+        subscriber.sunsubscribe
+
+        channels.each { |channel| assert_equal 0, r.publish("orphan", channel, sharded: true) }
+        assert_nil wait_for_message(subscriber, timeout: UNSUB_WAIT_TIME)
+      end
+    end
+
+    def test_ssubscribe_lazy_eventually_delivers
+      skip_unless_sharded_pubsub
+
+      channel = unique_channel
+
+      with_client do |subscriber|
+        subscriber.ssubscribe_lazy(channel)
+
+        received = publish_until_received_sharded("lazy-sharded-msg", channel, subscriber)
+
+        assert_equal "lazy-sharded-msg", received.message
+        assert_equal channel,            received.channel
+        assert_nil                       received.pattern
+      end
+    end
+
+    def test_sunsubscribe_lazy_eventually_stops_delivery
+      skip_unless_sharded_pubsub
+
+      channel = unique_channel
+
+      with_client do |subscriber|
+        subscriber.ssubscribe_lazy(channel)
+        publish_until_received_sharded("before", channel, subscriber)
+
+        subscriber.sunsubscribe_lazy(channel)
+
+        assert_delivery_ceased(channel, subscriber, sharded: true)
+      end
+    end
+
+    def test_connection_time_sharded_subscription
+      skip_unless_sharded_pubsub
+
+      channel = unique_channel
+
+      with_client(pubsub: { subscriptions: { sharded: [channel] } }) do |client|
+        received = publish_until_received_sharded("connect-time-sharded", channel, client)
+
+        assert_equal "connect-time-sharded", received.message
+        assert_equal channel, received.channel
+        assert_nil received.pattern
+      end
+    end
+
+    def test_sharded_publish_reaches_a_subscriber_in_a_different_slot
+      skip_unless_sharded_pubsub
+
+      # Distinct hash tags force the two channels onto different slots (and thus
+      # likely different owning nodes). A same-slot pair would pass even if the
+      # SPUBLISH were misrouted, so the different-slot case is what proves the
+      # publish is routed by the channel it names, not the subscriber's node.
+      subscribed_channel = unique_channel("{slot-a}")
+      other_slot_channel = unique_channel("{slot-b}")
+
+      with_client do |subscriber|
+        subscriber.ssubscribe(subscribed_channel)
+
+        # A message on a channel in a different slot must not arrive here.
+        r.publish("wrong-slot", other_slot_channel, sharded: true)
+        assert_nil wait_for_message(subscriber, timeout: UNSUB_WAIT_TIME)
+
+        # A message on the subscribed channel must arrive, proving the SPUBLISH
+        # was routed to the node owning that channel's slot.
+        r.publish("right-slot", subscribed_channel, sharded: true)
+        assert_equal "right-slot", wait_for_message(subscriber).message
+      end
+    end
+
+    def test_sharded_publish_is_batchable_in_a_pipeline
+      skip_unless_sharded_pubsub
+
+      channel = unique_channel
+
+      with_client do |subscriber|
+        subscriber.ssubscribe(channel)
+
+        counts = r.pipelined { |p| p.publish("piped", channel, sharded: true) }
+
+        assert_equal [1], counts
+        assert_equal "piped", wait_for_message(subscriber).message
+      end
+    end
+
+    def test_sharded_verbs_require_cluster_mode_in_standalone
+      skip "covers the standalone rejection path" if cluster_mode?
+
+      with_client do |client|
+        {
+          ssubscribe: -> { client.ssubscribe("shard1") },
+          sunsubscribe: -> { client.sunsubscribe },
+          ssubscribe_lazy: -> { client.ssubscribe_lazy("shard1") },
+          sunsubscribe_lazy: -> { client.sunsubscribe_lazy }
+        }.each do |name, call|
+          error = assert_raises(ArgumentError, "#{name} must require cluster mode") { call.call }
+          assert_match(/cluster mode/, error.message)
+        end
+      end
+    end
+
+    def test_sharded_connection_config_requires_cluster_mode_in_standalone
+      skip "covers the standalone rejection path" if cluster_mode?
+
+      error = assert_raises(ArgumentError) do
+        _new_client(protocol: :resp3, pubsub: { subscriptions: { sharded: [unique_channel] } })
+      end
+
+      assert_match(/cluster mode/, error.message)
+    end
+
     private
 
+    def skip_unless_sharded_pubsub
+      skip "sharded Pub/Sub is cluster-only" unless cluster_mode?
+
+      omit_version("7.0")
+    end
+
     def with_client(options = {})
       subscriber = _new_client(options.merge(protocol: :resp3))
       yield subscriber
@@ -490,6 +671,18 @@ def publish_until_received(message, channel, subscriber, timeout: MESSAGE_WAIT_S
       end
     end
 
+    def publish_until_received_sharded(message, channel, subscriber, timeout: MESSAGE_WAIT_SECONDS)
+      deadline = monotonic_now + timeout
+
+      loop do
+        r.publish(message, channel, sharded: true)
+        received = wait_for_message(subscriber, timeout: POLL_INTERVAL_SECONDS)
+        return received if received
+
+        flunk("no sharded message on #{channel} within #{timeout}s") if monotonic_now >= deadline
+      end
+    end
+
     def wait_until(timeout: MESSAGE_WAIT_SECONDS)
       deadline = monotonic_now + timeout
 
@@ -523,15 +716,15 @@ def collect_callback_message(queue, payload, channel, timeout: MESSAGE_WAIT_SECO
       end
     end
 
-    def assert_delivery_ceased(channel, subscriber)
+    def assert_delivery_ceased(channel, subscriber, sharded: false)
       quiet_windows = 2
       deadline = monotonic_now + MESSAGE_WAIT_SECONDS
       quiet_count = 0
 
       until quiet_count >= quiet_windows
         flunk("delivery did not cease on #{channel} within #{MESSAGE_WAIT_SECONDS}s") if monotonic_now >= deadline
 
-        r.publish("probe-ceased-#{SecureRandom.hex(4)}", channel)
+        r.publish("probe-ceased-#{SecureRandom.hex(4)}", channel, sharded: sharded)
         message = wait_for_message(subscriber, timeout: UNSUB_WAIT_TIME)
 
         if message