feat(cluster): split a cross-slot SSUBSCRIBE per slot

CHANGELOG.md

@@ -8,6 +8,15 @@ Versions up to and including `0.19.3` are documented in the
 
 ## [Unreleased]
 
+### Changed
+
+- **`ssubscribe` accepts shard channels of several slots in a cluster.** The
+  server requires one slot per `SSUBSCRIBE`, and the client refused a
+  cross-slot one with `MismatchedKeySlots`. A shard subscription is per
+  channel, so nothing is lost by splitting it: one `SSUBSCRIBE` is now sent per
+  slot, to the master owning it, and the caller is answered once every master
+  has confirmed.
+
 ### Fixed
 
 - **A sharded `PubSubStream` spanning several slots is released in a

src/commands/pub_sub_commands.rs

@@ -186,10 +186,15 @@ pub trait PubSubCommands<'a>: Sized {
         prepare_command(self, cmd("SPUBLISH").key(shardchannel).arg(message))
     }
 
-    /// Subscribes the client to the specified channels.
+    /// Subscribes the client to the specified shard channels.
+    ///
+    /// On a cluster, the shard channels may hash to different slots: one
+    /// `SSUBSCRIBE` is sent per slot, to the master owning it, and the call
+    /// returns once every one of them is confirmed. The same goes for the
+    /// `SUNSUBSCRIBE` the stream sends when it is closed or dropped.
     ///
     /// # See Also
-    /// [<https://redis.io/commands/subscribe/>](https://redis.io/commands/subscribe/)
+    /// [<https://redis.io/commands/ssubscribe/>](https://redis.io/commands/ssubscribe/)
     #[allow(async_fn_in_trait)]
     async fn ssubscribe(self, shardchannels: impl Serialize) -> Result<PubSubStream>;
 

src/network/cluster_connection.rs

@@ -151,18 +151,21 @@ fn is_broadcast_pub_sub_command(command: &Command) -> bool {
     )
 }
 
-/// Whether the command is a `SUNSUBSCRIBE` naming shard channels of more than
-/// one slot.
+/// Whether the command is a `SSUBSCRIBE` or a `SUNSUBSCRIBE` naming shard
+/// channels of more than one slot.
 ///
-/// A single `SSUBSCRIBE` must name one slot, but a `PubSubStream` may add shard
-/// channels of other slots afterwards, and its `close` and `drop` release them
-/// all in one `SUNSUBSCRIBE`. Routed as a single-slot command, that cleanup was
-/// refused before reaching any node: every subscription stayed on its master
-/// and in the client's registry. Cut per slot like a multi-shard command, each
-/// master hears about its own shard channels — the acknowledgements are matched
-/// per channel, so the caller is answered once all of them are in.
-fn is_shard_unsubscribe_spanning_slots(command: &Command) -> bool {
-    if command.name() != b"SUNSUBSCRIBE" {
+/// The server requires a single slot per command, but a shard subscription is
+/// per channel: nothing is atomic about subscribing to two of them at once.
+/// Cut per slot like a multi-shard command, each master hears about its own
+/// shard channels, and the acknowledgements are matched per channel by the
+/// router, so the caller is answered once all of them are in.
+///
+/// For `SUNSUBSCRIBE` this is not a convenience: a `PubSubStream` collects
+/// shard channels one `ssubscribe` at a time, and its `close` and `drop`
+/// release them all in one command. Routed as a single-slot command, that
+/// cleanup was refused before reaching any node.
+fn is_shard_pub_sub_spanning_slots(command: &Command) -> bool {
+    if !matches!(command.name(), b"SSUBSCRIBE" | b"SUNSUBSCRIBE") {
         return false;
     }
 
@@ -205,7 +208,7 @@ fn route(command: &Command) -> Result<Route> {
         });
     }
 
-    if is_shard_unsubscribe_spanning_slots(command) {
+    if is_shard_pub_sub_spanning_slots(command) {
         return Ok(Route::MultiShard);
     }
 

src/tests/pub_sub_commands.rs

@@ -1597,39 +1597,69 @@ async fn closing_a_shard_stream_spanning_slots_succeeds_in_cluster() -> Result<(
     Ok(())
 }
 
-/// A SSUBSCRIBE refused before being sent -- its shard channels span two
-/// slots -- must leave nothing behind: the client's next subscriptions are
-/// confirmed and receive their messages.
+/// A single `ssubscribe` may name shard channels of several slots: the cluster
+/// connection sends one `SSUBSCRIBE` per slot, to the master owning it. The
+/// server refuses a cross-slot one with `CROSSSLOT`, and the client used to
+/// refuse it itself with `MismatchedKeySlots`.
+///
+/// This is the reproduction of the original report, run to the end: subscribe,
+/// receive from both masters, drop, subscribe again.
 #[tokio::test]
 #[serial]
-async fn a_refused_cross_slot_ssubscribe_leaves_the_client_usable() -> Result<()> {
+async fn a_cross_slot_ssubscribe_is_split_per_slot_in_cluster() -> Result<()> {
     log_try_init();
 
+    // `{a}` and `{b}` hash to slots served by two different masters of the
+    // three-node test cluster.
+    const CHANNELS: [&str; 2] = ["cross_slot_ssubscribe{a}", "cross_slot_ssubscribe{b}"];
+
     let client = get_cluster_test_client().await?;
     let publisher = get_cluster_test_client().await?;
 
-    let refused = client
-        .ssubscribe(["refused_ssubscribe{a}", "refused_ssubscribe{b}"])
-        .await;
-    assert!(matches!(
-        refused.map(|_| ()).unwrap_err().kind(),
-        ErrorKind::Client(ClientError::MismatchedKeySlots)
-    ));
+    // A shard subscription lives on the master owning its slot only, so each
+    // master is asked about its own.
+    let shard_results: Vec<ClusterShardResult> = client.cluster_shards().await?;
+    let mut master_clients = Vec::new();
+    for shard in &shard_results {
+        let master_node = shard.nodes.iter().find(|n| n.role == "master").unwrap();
+        master_clients.push(
+            Client::connect((master_node.ip.clone(), master_node.port.unwrap()).into_config()?)
+                .await?,
+        );
+    }
 
-    let mut stream = timeout(
-        Duration::from_secs(5),
-        TimeoutKind::Command,
-        client.ssubscribe("refused_ssubscribe{a}"),
-    )
-    .await??;
-    publisher
-        .spublish("refused_ssubscribe{a}", "payload")
+    for round in 0..2 {
+        let mut stream = timeout(
+            Duration::from_secs(5),
+            TimeoutKind::Command,
+            client.ssubscribe(CHANNELS),
+        )
+        .await??;
+
+        for channel in CHANNELS {
+            publisher.spublish(channel, round.to_string()).await?;
+            let message = timeout(Duration::from_secs(5), TimeoutKind::Command, stream.next())
+                .await?
+                .unwrap()?;
+            assert_eq!(channel.as_bytes(), message.channel());
+        }
+
+        drop(stream);
+
+        // The client clears its registry only when each master confirms the
+        // unsubscription, which `drop` does not wait for.
+        wait_for(Duration::from_secs(5), || async {
+            let mut subscribed = HashSet::<String>::new();
+            for master_client in &master_clients {
+                let channels: HashSet<String> = master_client
+                    .pub_sub_shardchannels("cross_slot_ssubscribe*")
+                    .await?;
+                subscribed.extend(channels);
+            }
+            Ok(subscribed.is_empty())
+        })
         .await?;
-    let message = timeout(Duration::from_secs(5), TimeoutKind::Command, stream.next())
-        .await?
-        .unwrap()?;
-    assert_eq!(b"payload", message.payload());
+    }
 
-    stream.close().await?;
     Ok(())
 }