Complete sharded Pub/Sub support

spec/cluster_spec.cr

@@ -75,13 +75,28 @@ describe Redis::Cluster do
   describe "pub/sub" do
     it "publishes only to a single shard" do
       messages = Channel(String).new
+      subscribed = Channel(Nil).new
+      done = Channel(Exception?).new
 
       spawn do
         cluster.ssubscribe "a" do |subscription|
+          subscription.on_subscribe do
+            subscribed.send nil
+          end
           subscription.on_message do |_channel, msg|
             messages.send msg
+            subscription.close
           end
         end
+        done.send nil
+      rescue ex
+        done.send ex
+      end
+
+      select
+      when subscribed.receive
+      when timeout(1.second)
+        raise "Timed out while waiting for sharded subscription"
       end
 
       # "a" goes to hash slot 15495
@@ -92,10 +107,8 @@ describe Redis::Cluster do
         cluster.ssubscribe "a", "b" { }
       end
 
-      sleep 10.milliseconds
-
-      cluster.spublish "a", "yep"
-      cluster.spublish "b", "nope"
+      cluster.spublish("a", "yep").should eq 1
+      cluster.spublish("b", "nope").should eq 0
 
       select
       when msg = messages.receive
@@ -105,34 +118,57 @@ describe Redis::Cluster do
       end
 
       select
-      when msg = messages.receive
-        raise "Expected not to receive a message, received: #{msg.inspect}"
-      when timeout(10.milliseconds)
+      when error = done.receive
+        raise error if error
+      when timeout(1.second)
+        raise "Timed out while waiting for sharded subscription to finish"
       end
     end
 
     it "publishes to all shards" do
       messages = Channel({String, String}).new
+      subscribed = Channel(Nil).new
+      done = Channel(Exception?).new
 
       spawn do
         cluster.subscribe "a" do |subscription|
+          subscription.on_subscribe do
+            subscribed.send nil
+          end
           subscription.on_message do |channel, msg|
             messages.send({channel, msg})
+            subscription.close
           end
         end
+        done.send nil
+      rescue ex
+        done.send ex
       end
 
       spawn do
         # This has to be in a separate fiber because it's going to a different
         # node in the cluster.
         cluster.subscribe "b" do |subscription|
+          subscription.on_subscribe do
+            subscribed.send nil
+          end
           subscription.on_message do |channel, msg|
             messages.send({channel, msg})
+            subscription.close
           end
         end
+        done.send nil
+      rescue ex
+        done.send ex
       end
 
-      sleep 10.milliseconds
+      2.times do
+        select
+        when subscribed.receive
+        when timeout(1.second)
+          raise "Timed out while waiting for subscription"
+        end
+      end
 
       cluster.publish "a", "1"
       cluster.publish "b", "2"
@@ -149,6 +185,15 @@ describe Redis::Cluster do
 
       received.should contain({"a", "1"})
       received.should contain({"b", "2"})
+
+      2.times do
+        select
+        when error = done.receive
+          raise error if error
+        when timeout(1.second)
+          raise "Timed out while waiting for subscription to finish"
+        end
+      end
     end
   end
 end

spec/redis_spec.cr

@@ -666,6 +666,120 @@ describe Redis::Client do
         end
       end
     end
+
+    if test.server_version >= Version["7.0.0"]
+      it "can publish and subscribe to sharded channels" do
+        channel = random_key
+        subscribed = Channel(Nil).new
+        messages = Channel({String, String}).new
+        unsubscribed = Channel({String, Int64}).new
+        done = Channel(Exception?).new
+
+        spawn do
+          redis.ssubscribe channel do |subscription, connection|
+            subscription.on_subscribe do
+              subscribed.send nil
+            end
+            subscription.on_message do |received_channel, message|
+              messages.send({received_channel, message})
+              connection.sunsubscribe received_channel
+            end
+            subscription.on_unsubscribe do |unsubscribed_channel, count|
+              unsubscribed.send({unsubscribed_channel, count})
+            end
+          end
+          done.send nil
+        rescue ex
+          done.send ex
+        end
+
+        select
+        when subscribed.receive
+        when timeout(1.second)
+          raise "Timed out while waiting for sharded subscription"
+        end
+
+        redis.spublish(channel, "hello").should eq 1
+
+        select
+        when message = messages.receive
+          message.should eq({channel, "hello"})
+        when timeout(1.second)
+          raise "Timed out while waiting for sharded message"
+        end
+
+        select
+        when event = unsubscribed.receive
+          event.should eq({channel, 0_i64})
+        when timeout(1.second)
+          raise "Timed out while waiting for sharded unsubscription"
+        end
+
+        select
+        when error = done.receive
+          raise error if error
+        when timeout(1.second)
+          raise "Timed out while waiting for sharded subscription to finish"
+        end
+      end
+
+      it "closes each subscription type with the matching command" do
+        channel = random_key
+        pattern = "#{random_key}*"
+        shard_channel = random_key
+        subscribed = Channel(String).new
+        unsubscribed = Channel(String).new
+        done = Channel(Exception?).new
+
+        spawn do
+          redis.subscribe channel do |subscription, connection|
+            subscriptions = Set(String).new
+            subscription.on_subscribe do |subscribed_channel|
+              subscriptions << subscribed_channel
+              subscribed.send subscribed_channel
+              subscription.close if subscriptions.size == 3
+            end
+            subscription.on_unsubscribe do |unsubscribed_channel|
+              unsubscribed.send unsubscribed_channel
+            end
+            connection.psubscribe pattern
+            connection.ssubscribe shard_channel
+          end
+          done.send nil
+        rescue ex
+          done.send ex
+        end
+
+        subscriptions = Set(String).new
+        3.times do
+          select
+          when subscribed_channel = subscribed.receive
+            subscriptions << subscribed_channel
+          when timeout(1.second)
+            raise "Timed out while waiting for subscriptions"
+          end
+        end
+        subscriptions.should eq Set{channel, pattern, shard_channel}
+
+        unsubscriptions = Set(String).new
+        3.times do
+          select
+          when unsubscribed_channel = unsubscribed.receive
+            unsubscriptions << unsubscribed_channel
+          when timeout(1.second)
+            raise "Timed out while waiting for unsubscriptions"
+          end
+        end
+        unsubscriptions.should eq Set{channel, pattern, shard_channel}
+
+        select
+        when error = done.receive
+          raise error if error
+        when timeout(1.second)
+          raise "Timed out while waiting for subscriptions to finish"
+        end
+      end
+    end
   end
 
   it "gets info about the Redis server" do

src/client.cr

@@ -177,6 +177,10 @@ module Redis
       checkout(&.psubscribe(*channels) { |subscription, conn| yield subscription, conn })
     end
 
+    def ssubscribe(*channels, &)
+      checkout(&.ssubscribe(*channels) { |subscription, conn| yield subscription, conn })
+    end
+
     def close
       @pool.close
     end

src/cluster.cr

@@ -167,13 +167,6 @@ module Redis
       each_master(&.run({"flushdb"}))
     end
 
-    # Publish a message to a sharded pubsub channel. The message is delivered
-    # only to subscribers on the shard that owns the channel's hash slot —
-    # unlike `publish`, it is not propagated across the cluster.
-    def spublish(channel : String, message : String)
-      write_pool_for(channel).checkout(&.run({"spublish", channel, message}))
-    end
-
     # Subscribe to one or more pubsub channels. Regular (non-sharded) pub/sub
     # messages are propagated across the cluster, so the subscriber will
     # receive messages regardless of which node a publisher targets.
@@ -190,7 +183,7 @@ module Redis
     # Subscribe to one or more sharded pubsub channels. All channels must hash
     # to the same slot (use `{}` to force co-location). The block yields a
     # `Subscription` and the underlying `Connection`, which holds the
-    # subscription for its duration.
+    # subscription for its duration. Requires Redis 7.0 or later.
     #
     # ```
     # cluster.ssubscribe "orders" do |subscription, connection|
@@ -203,15 +196,6 @@ module Redis
       write_pool_for(channels.first).checkout(&.ssubscribe(*channels) { |subscription, conn| yield subscription, conn })
     end
 
-    # Unsubscribe from sharded pubsub channels on the shard that owns their
-    # hash slot. Typically called from inside an `ssubscribe` block via the
-    # yielded `Connection`; calling this on the cluster directly routes a
-    # standalone `SUNSUBSCRIBE` to the shard and has no effect on a
-    # subscription held by a different connection.
-    def sunsubscribe(*channels : String)
-      write_pool_for(channels.first).checkout(&.sunsubscribe(*channels))
-    end
-
     def run(command full_command)
       if full_command.empty?
         raise ArgumentError.new("Redis commands must have at least one component")

src/commands.cr

@@ -330,6 +330,15 @@ module Redis
       run({"publish", channel, message})
     end
 
+    # Publish a message to a sharded pubsub channel. In a Redis cluster, the
+    # message is delivered only to subscribers on the shard that owns the
+    # channel's hash slot.
+    #
+    # Requires Redis 7.0 or later.
+    def spublish(channel : String, message : String)
+      run({"spublish", channel, message})
+    end
+
     # Preload a Lua script, returning the SHA of the script to pass to `evalsha`.
     # ```
     # sha = redis.script_load(<<-LUA)

src/commands/immediate.cr

@@ -57,6 +57,7 @@ module Redis::Commands::Immediate
       script_load:   String,
       script_exists: Array,
       info:          String,
+      spublish:      Int64,
 
       # String commands
       append:      Int64,