Repository: redis/lettuce (java)

Issue:
`resubscribe` logic looks ignoring sharded Pub/Sub (`shardChannels` Set of `PubSubEndpoint`)
## Bug Report

Thank you for adding support for sharded PubSub in the 6.4.0.
Related to that implementation, I would like to confirm about `resubscribe` (subscription recovery after reconnection) behavior.
It looks like `resubscribe` handling only non-sharded PubSub subscriptions.

#### Current Behavior

Sharded PubSub subscriptions looks not recovered after disconnection + re-connection, despite Lettuce automatically recovers non-sharded PubSub subscriptions (the recovery behavior is very helpful for applications, and also difficult / error-prone to implement by ourselves).

In the https://github.com/redis/lettuce/commit/1457486165edb8e7184c0214dae9973c7f6ff567#diff-73f0c5ac0f7d6d94f8d47e1a75564aa344cbbbd4c91ce786711b5c3673044956, `PubSubEndpoint.shardChannels` (`private`) have been added but it looks that variable is private inside `PubSubEndpoint` and no any code of `PubSubEndpoint` *reading* that variable's data. Also implementations of `StatefulRedisPubSubConnectionImpl.resubscribe` and `StatefulRedisClusterPubSubConnectionImpl.resubscribe` does not accessing `shardChannels`.

#### Input Code

Just make sharded Pub/Sub connection, and wait until some disconnection + re-connection occurs on that environment...

<!-- Java/Kotlin/Scala/Groovy/… or Repo link if applicable: -->

<details>
<summary>Input Code</summary>

```java
StatefulRedisPubSubConnection<String, String> connection = client.connectPubSub()
connection.addListener(new RedisPubSubAdapter<String, String>() { ... })

RedisPubSubCommands<String, String> sync = connection.sync();
sync.ssubscribe("channel"); // Sharded subscription
```

</details>

#### Expected behavior/code

As same as non-sharded Pub/Sub, automatically recover sharded Pub/Sub subscriptions when possible to connect to that shard/node.
( May be a bit different logic from non-sharded Pub/Sub needed. Because the SSUBSCRIBE requires living connection to the particular shard's node(s) )

#### Environment

- Lettuce version(s): 6.4.0
- Redis version: 7.1 (I belive this issue happens even with Redis 7.0)

#### Possible Solution

(Not tested) Watch the connection's state somehow and manually call `ssubscribe` again by application's logic.

#### Additional context

Because currently I don't have full access to a target environment / Redis cluster, this issue is based on best guess by reading code. I hope this code reading makes some sense...


Pull request:
Resubscribe logic ignoring Pub/Sub shard channels after reconnect. #2940
Closes #2940 

<!--
Thank you for proposing a pull request. This template will guide you through the essential steps necessary for a pull request.
-->
Make sure that:

- [ ] You have read the [contribution guidelines](https://github.com/lettuce-io/lettuce-core/blob/main/.github/CONTRIBUTING.md).
- [ ] You have created a feature request first to discuss your contribution intent. Please reference the feature request ticket number in the pull request.
- [ ] You applied code formatting rules using the `mvn formatter:format` target. Don’t submit any formatting related changes.
- [ ] You submit test cases (unit or integration tests) that back your changes.

 <!--
Great! Live long and prosper.
-->


Patch:
diff --git a/src/main/java/io/lettuce/core/pubsub/PubSubEndpoint.java b/src/main/java/io/lettuce/core/pubsub/PubSubEndpoint.java
index 671b24aed..36c511d52 100644
--- a/src/main/java/io/lettuce/core/pubsub/PubSubEndpoint.java
+++ b/src/main/java/io/lettuce/core/pubsub/PubSubEndpoint.java
@@ -128,6 +128,14 @@ public class PubSubEndpoint<K, V> extends DefaultEndpoint {
         return unwrap(this.channels);
     }
 
+    public boolean hasShardChannelSubscriptions() {
+        return !shardChannels.isEmpty();
+    }
+
+    public Set<K> getShardChannels() {
+        return unwrap(this.shardChannels);
+    }
+
     public boolean hasPatternSubscriptions() {
         return !patterns.isEmpty();
     }
diff --git a/src/main/java/io/lettuce/core/pubsub/StatefulRedisPubSubConnectionImpl.java b/src/main/java/io/lettuce/core/pubsub/StatefulRedisPubSubConnectionImpl.java
index 09c802e80..6e012f432 100644
--- a/src/main/java/io/lettuce/core/pubsub/StatefulRedisPubSubConnectionImpl.java
+++ b/src/main/java/io/lettuce/core/pubsub/StatefulRedisPubSubConnectionImpl.java
@@ -130,6 +130,10 @@ public class StatefulRedisPubSubConnectionImpl<K, V> extends StatefulRedisConnec
             result.add(async().subscribe(toArray(endpoint.getChannels())));
         }
 
+        if (endpoint.hasShardChannelSubscriptions()) {
+            result.add(async().ssubscribe(toArray(endpoint.getShardChannels())));
+        }
+
         if (endpoint.hasPatternSubscriptions()) {
             result.add(async().psubscribe(toArray(endpoint.getPatterns())));
         }


Test patch:
diff --git a/src/test/java/io/lettuce/core/pubsub/PubSubCommandIntegrationTests.java b/src/test/java/io/lettuce/core/pubsub/PubSubCommandIntegrationTests.java
index fdd778055..9936ba428 100644
--- a/src/test/java/io/lettuce/core/pubsub/PubSubCommandIntegrationTests.java
+++ b/src/test/java/io/lettuce/core/pubsub/PubSubCommandIntegrationTests.java
@@ -86,6 +86,8 @@ class PubSubCommandIntegrationTests extends AbstractRedisClientTest {
 
     BlockingQueue<Long> counts = listener.getCounts();
 
+    BlockingQueue<Long> shardCounts = listener.getShardCounts();
+
     String channel = "channel0";
 
     String shardChannel = "shard-channel";
@@ -523,6 +525,24 @@ class PubSubCommandIntegrationTests extends AbstractRedisClientTest {
         assertThat(messages.take()).isEqualTo(message);
     }
 
+    @Test
+    void resubscribeShardChannelsOnReconnect() throws Exception {
+        pubsub.ssubscribe(shardChannel);
+        assertThat(shardChannels.take()).isEqualTo(shardChannel);
+        assertThat((long) shardCounts.take()).isEqualTo(1);
+
+        pubsub.quit();
+
+        assertThat(shardChannels.take()).isEqualTo(shardChannel);
+        assertThat((long) shardCounts.take()).isEqualTo(1);
+
+        Wait.untilTrue(pubsub::isOpen).waitOrTimeout();
+
+        redis.spublish(shardChannel, shardMessage);
+        assertThat(shardChannels.take()).isEqualTo(shardChannel);
+        assertThat(messages.take()).isEqualTo(shardMessage);
+    }
+
     @Test
     void adapter() throws Exception {
         final BlockingQueue<Long> localCounts = LettuceFactories.newBlockingQueue();
diff --git a/src/test/java/io/lettuce/core/pubsub/StatefulRedisPubSubConnectionImplUnitTests.java b/src/test/java/io/lettuce/core/pubsub/StatefulRedisPubSubConnectionImplUnitTests.java
index 8cd91df26..9eb552824 100644
--- a/src/test/java/io/lettuce/core/pubsub/StatefulRedisPubSubConnectionImplUnitTests.java
+++ b/src/test/java/io/lettuce/core/pubsub/StatefulRedisPubSubConnectionImplUnitTests.java
@@ -12,7 +12,7 @@ import org.junit.jupiter.api.Tag;
 import org.junit.jupiter.api.Test;
 
 import static io.lettuce.TestTags.UNIT_TEST;
-import static org.junit.Assert.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.mockito.Mockito.*;
 
@@ -81,6 +81,7 @@ class StatefulRedisPubSubConnectionImplUnitTests {
         when(mockedEndpoint.hasChannelSubscriptions()).thenReturn(true);
         when(mockedEndpoint.getChannels()).thenReturn(new HashSet<>(Arrays.asList(new String[] { "channel1", "channel2" })));
         when(mockedEndpoint.hasPatternSubscriptions()).thenReturn(false);
+        when(mockedEndpoint.hasShardChannelSubscriptions()).thenReturn(false);
 
         List<RedisFuture<Void>> subscriptions = connection.resubscribe();
         RedisFuture<Void> commandFuture = subscriptions.get(0);
@@ -90,17 +91,35 @@ class StatefulRedisPubSubConnectionImplUnitTests {
     }
 
     @Test
-    void resubscribeChannelAndPatternSubscription() {
+    void resubscribeShardChannelSubscription() {
+        when(mockedEndpoint.hasShardChannelSubscriptions()).thenReturn(true);
+        when(mockedEndpoint.getShardChannels())
+                .thenReturn(new HashSet<>(Arrays.asList(new String[] { "shard_channel1", "shard_channel2" })));
+        when(mockedEndpoint.hasChannelSubscriptions()).thenReturn(false);
+        when(mockedEndpoint.hasPatternSubscriptions()).thenReturn(false);
+
+        List<RedisFuture<Void>> subscriptions = connection.resubscribe();
+        RedisFuture<Void> commandFuture = subscriptions.get(0);
+
+        assertEquals(1, subscriptions.size());
+        assertInstanceOf(AsyncCommand.class, commandFuture);
+    }
+
+    @Test
+    void resubscribeChannelAndPatternAndShardChanelSubscription() {
         when(mockedEndpoint.hasChannelSubscriptions()).thenReturn(true);
-        when(mockedEndpoint.getChannels()).thenReturn(new HashSet<>(Arrays.asList(new String[] { "channel1", "channel2" })));
         when(mockedEndpoint.hasPatternSubscriptions()).thenReturn(true);
+        when(mockedEndpoint.hasShardChannelSubscriptions()).thenReturn(true);
+        when(mockedEndpoint.getChannels()).thenReturn(new HashSet<>(Arrays.asList(new String[] { "channel1", "channel2" })));
         when(mockedEndpoint.getPatterns()).thenReturn(new HashSet<>(Arrays.asList(new String[] { "bcast*", "echo" })));
-
+        when(mockedEndpoint.getShardChannels())
+                .thenReturn(new HashSet<>(Arrays.asList(new String[] { "shard_channel1", "shard_channel2" })));
         List<RedisFuture<Void>> subscriptions = connection.resubscribe();
 
-        assertEquals(2, subscriptions.size());
+        assertEquals(3, subscriptions.size());
         assertInstanceOf(AsyncCommand.class, subscriptions.get(0));
         assertInstanceOf(AsyncCommand.class, subscriptions.get(1));
+        assertInstanceOf(AsyncCommand.class, subscriptions.get(1));
     }
 
 }