Break out FaultTolerantPubSubConnection as its own thing so different use cases can have their own subscription space.

service/src/main/java/org/whispersystems/textsecuregcm/push/ClientPresenceManager.java

@@ -15,6 +15,7 @@ import io.lettuce.core.cluster.pubsub.RedisClusterPubSubAdapter;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.redis.ClusterLuaScript;
+import org.whispersystems.textsecuregcm.redis.FaultTolerantPubSubConnection;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
 import org.whispersystems.textsecuregcm.util.Constants;
 
@@ -42,8 +43,10 @@ public class ClientPresenceManager extends RedisClusterPubSubAdapter<String, Str
     private final String managerId             = UUID.randomUUID().toString();
     private final String connectedClientSetKey = getConnectedClientSetKey(managerId);
 
-    private final FaultTolerantRedisCluster presenceCluster;
+    private final FaultTolerantRedisCluster                     presenceCluster;
-    private final ClusterLuaScript          clearPresenceScript;
+    private final FaultTolerantPubSubConnection<String, String> pubSubConnection;
+
+    private final ClusterLuaScript clearPresenceScript;
 
     private final ScheduledExecutorService scheduledExecutorService;
     private       ScheduledFuture<?>       pruneMissingPeersFuture;
@@ -65,8 +68,9 @@ public class ClientPresenceManager extends RedisClusterPubSubAdapter<String, Str
 
     public ClientPresenceManager(final FaultTolerantRedisCluster presenceCluster, final ScheduledExecutorService scheduledExecutorService) throws IOException {
         this.presenceCluster          = presenceCluster;
-        this.scheduledExecutorService = scheduledExecutorService;
+        this.pubSubConnection         = this.presenceCluster.createPubSubConnection();
         this.clearPresenceScript      = ClusterLuaScript.fromResource(presenceCluster, "lua/clear_presence.lua", ScriptOutputType.INTEGER);
+        this.scheduledExecutorService = scheduledExecutorService;
 
         final MetricRegistry metricRegistry = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
         metricRegistry.gauge(name(getClass(), "localClientCount"), () -> displacementListenersByPresenceKey::size);
@@ -79,9 +83,14 @@ public class ClientPresenceManager extends RedisClusterPubSubAdapter<String, Str
         this.remoteDisplacementMeter = metricRegistry.meter(name(getClass(), "remoteDisplacement"));
     }
 
+    @VisibleForTesting
+    FaultTolerantPubSubConnection<String, String> getPubSubConnection() {
+        return pubSubConnection;
+    }
+
     @Override
     public void start() {
-        presenceCluster.usePubSubConnection(connection -> {
+        pubSubConnection.usePubSubConnection(connection -> {
             connection.addListener(this);
             connection.getResources().eventBus().get()
                                                 .filter(event -> event instanceof ClusterTopologyChangedEvent)
@@ -103,7 +112,7 @@ public class ClientPresenceManager extends RedisClusterPubSubAdapter<String, Str
 
     @Override
     public void stop() {
-        presenceCluster.usePubSubConnection(connection -> connection.removeListener(this));
+        pubSubConnection.usePubSubConnection(connection -> connection.removeListener(this));
 
         if (pruneMissingPeersFuture != null) {
             pruneMissingPeersFuture.cancel(false);
@@ -118,7 +127,7 @@ public class ClientPresenceManager extends RedisClusterPubSubAdapter<String, Str
             connection.sync().del(getConnectedClientSetKey(managerId));
         });
 
-        presenceCluster.usePubSubConnection(connection -> connection.sync().masters().commands().unsubscribe(getManagerPresenceChannel(managerId)));
+        pubSubConnection.usePubSubConnection(connection -> connection.sync().masters().commands().unsubscribe(getManagerPresenceChannel(managerId)));
     }
 
     public void setPresent(final UUID accountUuid, final long deviceId, final DisplacedPresenceListener displacementListener) {
@@ -175,7 +184,7 @@ public class ClientPresenceManager extends RedisClusterPubSubAdapter<String, Str
     private void subscribeForRemotePresenceChanges(final String presenceKey) {
         final int slot = SlotHash.getSlot(presenceKey);
 
-        presenceCluster.usePubSubConnection(connection -> connection.sync().nodes(node -> node.is(RedisClusterNode.NodeFlag.MASTER) && node.hasSlot(slot))
+        pubSubConnection.usePubSubConnection(connection -> connection.sync().nodes(node -> node.is(RedisClusterNode.NodeFlag.MASTER) && node.hasSlot(slot))
                                                                            .commands()
                                                                            .subscribe(getKeyspaceNotificationChannel(presenceKey)));
     }
@@ -187,7 +196,7 @@ public class ClientPresenceManager extends RedisClusterPubSubAdapter<String, Str
     }
 
     private void unsubscribeFromRemotePresenceChanges(final String presenceKey) {
-        presenceCluster.usePubSubConnection(connection -> connection.async().masters().commands().unsubscribe(getKeyspaceNotificationChannel(presenceKey)));
+        pubSubConnection.usePubSubConnection(connection -> connection.async().masters().commands().unsubscribe(getKeyspaceNotificationChannel(presenceKey)));
     }
 
     void pruneMissingPeers() {

service/src/main/java/org/whispersystems/textsecuregcm/redis/FaultTolerantPubSubConnection.java

@@ -0,0 +1,36 @@
+package org.whispersystems.textsecuregcm.redis;
+
+import com.codahale.metrics.SharedMetricRegistries;
+import io.github.resilience4j.circuitbreaker.CircuitBreaker;
+import io.lettuce.core.cluster.pubsub.StatefulRedisClusterPubSubConnection;
+import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
+import org.whispersystems.textsecuregcm.util.CircuitBreakerUtil;
+import org.whispersystems.textsecuregcm.util.Constants;
+
+import java.util.function.Consumer;
+import java.util.function.Function;
+
+public class FaultTolerantPubSubConnection<K, V> {
+
+    private final StatefulRedisClusterPubSubConnection<K, V> pubSubConnection;
+    private final CircuitBreaker circuitBreaker;
+
+    public FaultTolerantPubSubConnection(final String name, final StatefulRedisClusterPubSubConnection<K, V> pubSubConnection, final CircuitBreakerConfiguration circuitBreakerConfiguration) {
+        this.pubSubConnection = pubSubConnection;
+        this.circuitBreaker   = CircuitBreaker.of(name + "-pubsub", circuitBreakerConfiguration.toCircuitBreakerConfig());
+
+        CircuitBreakerUtil.registerMetrics(SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME),
+                this.circuitBreaker,
+                FaultTolerantRedisCluster.class);
+
+        this.pubSubConnection.setNodeMessagePropagation(true);
+    }
+
+    public void usePubSubConnection(final Consumer<StatefulRedisClusterPubSubConnection<K, V>> consumer) {
+        this.circuitBreaker.executeRunnable(() -> consumer.accept(pubSubConnection));
+    }
+
+    public <T> T withPubSubConnection(final Function<StatefulRedisClusterPubSubConnection<K, V>, T> consumer) {
+        return this.circuitBreaker.executeSupplier(() -> consumer.apply(pubSubConnection));
+    }
+}

service/src/main/java/org/whispersystems/textsecuregcm/redis/FaultTolerantRedisCluster.java

@@ -13,6 +13,7 @@ import org.whispersystems.textsecuregcm.util.CircuitBreakerUtil;
 import org.whispersystems.textsecuregcm.util.Constants;
 
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -25,15 +26,18 @@ import java.util.stream.Collectors;
  */
 public class FaultTolerantRedisCluster {
 
+    private final String name;
+
     private final RedisClusterClient clusterClient;
 
-    private final StatefulRedisClusterConnection<String, String>       stringClusterConnection;
+    private final StatefulRedisClusterConnection<String, String> stringClusterConnection;
-    private final StatefulRedisClusterConnection<byte[], byte[]>       binaryClusterConnection;
+    private final StatefulRedisClusterConnection<byte[], byte[]> binaryClusterConnection;
-    private final StatefulRedisClusterPubSubConnection<String, String> pubSubClusterConnection;
 
-    private final CircuitBreaker readCircuitBreaker;
+    private final List<StatefulRedisClusterPubSubConnection<?, ?>> pubSubConnections = new ArrayList<>();
-    private final CircuitBreaker writeCircuitBreaker;
+
-    private final CircuitBreaker pubSubCircuitBreaker;
+    private final CircuitBreakerConfiguration circuitBreakerConfiguration;
+    private final CircuitBreaker              readCircuitBreaker;
+    private final CircuitBreaker              writeCircuitBreaker;
 
     public FaultTolerantRedisCluster(final String name, final List<String> urls, final Duration timeout, final CircuitBreakerConfiguration circuitBreakerConfiguration) {
         this(name, RedisClusterClient.create(urls.stream().map(RedisURI::create).collect(Collectors.toList())), timeout, circuitBreakerConfiguration);
@@ -41,17 +45,17 @@ public class FaultTolerantRedisCluster {
 
     @VisibleForTesting
     FaultTolerantRedisCluster(final String name, final RedisClusterClient clusterClient, final Duration timeout, final CircuitBreakerConfiguration circuitBreakerConfiguration) {
+        this.name = name;
+
         this.clusterClient = clusterClient;
         this.clusterClient.setDefaultTimeout(timeout);
 
         this.stringClusterConnection = clusterClient.connect();
         this.binaryClusterConnection = clusterClient.connect(ByteArrayCodec.INSTANCE);
-        this.pubSubClusterConnection = clusterClient.connectPubSub();
-        this.readCircuitBreaker      = CircuitBreaker.of(name + "-read", circuitBreakerConfiguration.toCircuitBreakerConfig());
-        this.writeCircuitBreaker     = CircuitBreaker.of(name + "-write", circuitBreakerConfiguration.toCircuitBreakerConfig());
-        this.pubSubCircuitBreaker    = CircuitBreaker.of(name + "-pubsub", circuitBreakerConfiguration.toCircuitBreakerConfig());
 
-        this.pubSubClusterConnection.setNodeMessagePropagation(true);
+        this.circuitBreakerConfiguration = circuitBreakerConfiguration;
+        this.readCircuitBreaker          = CircuitBreaker.of(name + "-read", circuitBreakerConfiguration.toCircuitBreakerConfig());
+        this.writeCircuitBreaker         = CircuitBreaker.of(name + "-write", circuitBreakerConfiguration.toCircuitBreakerConfig());
 
         CircuitBreakerUtil.registerMetrics(SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME),
                 readCircuitBreaker,
@@ -60,16 +64,15 @@ public class FaultTolerantRedisCluster {
         CircuitBreakerUtil.registerMetrics(SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME),
                 writeCircuitBreaker,
                 FaultTolerantRedisCluster.class);
-
-        CircuitBreakerUtil.registerMetrics(SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME),
-                pubSubCircuitBreaker,
-                FaultTolerantRedisCluster.class);
     }
 
     void shutdown() {
         stringClusterConnection.close();
         binaryClusterConnection.close();
-        pubSubClusterConnection.close();
+
+        for (final StatefulRedisClusterPubSubConnection<?, ?> pubSubConnection : pubSubConnections) {
+            pubSubConnection.close();
+        }
 
         clusterClient.shutdown();
     }
@@ -106,11 +109,10 @@ public class FaultTolerantRedisCluster {
         return this.writeCircuitBreaker.executeSupplier(() -> consumer.apply(binaryClusterConnection));
     }
 
-    public void usePubSubConnection(final Consumer<StatefulRedisClusterPubSubConnection<String, String>> consumer) {
+    public FaultTolerantPubSubConnection<String, String> createPubSubConnection() {
-        this.pubSubCircuitBreaker.executeRunnable(() -> consumer.accept(pubSubClusterConnection));
+        final StatefulRedisClusterPubSubConnection<String, String> pubSubConnection = clusterClient.connectPubSub();
-    }
+        pubSubConnections.add(pubSubConnection);
 
-    public <T> T withPubSubConnection(final Function<StatefulRedisClusterPubSubConnection<String, String>, T> consumer) {
+        return new FaultTolerantPubSubConnection<>(name, pubSubConnection, circuitBreakerConfiguration);
-        return this.pubSubCircuitBreaker.executeSupplier(() -> consumer.apply(pubSubClusterConnection));
     }
 }

service/src/main/java/org/whispersystems/textsecuregcm/storage/RedisClusterMessagesCache.java

@@ -13,6 +13,7 @@ import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.entities.MessageProtos;
 import org.whispersystems.textsecuregcm.entities.OutgoingMessageEntity;
 import org.whispersystems.textsecuregcm.redis.ClusterLuaScript;
+import org.whispersystems.textsecuregcm.redis.FaultTolerantPubSubConnection;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
 import org.whispersystems.textsecuregcm.util.RedisClusterUtil;
 
@@ -33,8 +34,10 @@ import static com.codahale.metrics.MetricRegistry.name;
 
 public class RedisClusterMessagesCache extends RedisClusterPubSubAdapter<String, String> implements UserMessagesCache {
 
-    private final FaultTolerantRedisCluster redisCluster;
+    private final FaultTolerantRedisCluster                     redisCluster;
-    private final ExecutorService           notificationExecutorService;
+    private final FaultTolerantPubSubConnection<String, String> pubSubConnection;
+
+    private final ExecutorService notificationExecutorService;
 
     private final ClusterLuaScript insertScript;
     private final ClusterLuaScript removeByIdScript;
@@ -67,7 +70,9 @@ public class RedisClusterMessagesCache extends RedisClusterPubSubAdapter<String,
 
     public RedisClusterMessagesCache(final FaultTolerantRedisCluster redisCluster, final ExecutorService notificationExecutorService) throws IOException {
 
-        this.redisCluster                = redisCluster;
+        this.redisCluster     = redisCluster;
+        this.pubSubConnection = redisCluster.createPubSubConnection();
+
         this.notificationExecutorService = notificationExecutorService;
 
         this.insertScript             = ClusterLuaScript.fromResource(redisCluster, "lua/insert_item.lua",           ScriptOutputType.INTEGER);
@@ -78,21 +83,21 @@ public class RedisClusterMessagesCache extends RedisClusterPubSubAdapter<String,
         this.removeQueueScript        = ClusterLuaScript.fromResource(redisCluster, "lua/remove_queue.lua",          ScriptOutputType.STATUS);
         this.getQueuesToPersistScript = ClusterLuaScript.fromResource(redisCluster, "lua/get_queues_to_persist.lua", ScriptOutputType.MULTI);
 
-        redisCluster.usePubSubConnection(connection -> {
+        pubSubConnection.usePubSubConnection(connection -> {
             connection.addListener(this);
             connection.getResources().eventBus().get()
                     .filter(event -> event instanceof ClusterTopologyChangedEvent)
                     .handle((event, sink) -> {
-                        resubscribeAll();
+                        subscribeForKeyspaceNotifications();
                         sink.next(event);
                     });
-
-            connection.sync().masters().commands().psubscribe(QUEUE_KEYSPACE_PATTERN, PERSISTING_KEYSPACE_PATTERN);
         });
+
+        subscribeForKeyspaceNotifications();
     }
 
-    private void resubscribeAll() {
+    private void subscribeForKeyspaceNotifications() {
-        redisCluster.usePubSubConnection(connection -> connection.sync().masters().commands().psubscribe(QUEUE_KEYSPACE_PATTERN, PERSISTING_KEYSPACE_PATTERN));
+        pubSubConnection.usePubSubConnection(connection -> connection.sync().masters().commands().psubscribe(QUEUE_KEYSPACE_PATTERN, PERSISTING_KEYSPACE_PATTERN));
     }
 
     @Override

service/src/test/java/org/whispersystems/textsecuregcm/push/ClientPresenceManagerTest.java

@@ -123,7 +123,7 @@ public class ClientPresenceManagerTest extends AbstractRedisClusterTest {
                 }
             });
 
-            getRedisCluster().usePubSubConnection(connection -> connection.getResources().eventBus().publish(new ClusterTopologyChangedEvent(List.of(), List.of())));
+            clientPresenceManager.getPubSubConnection().usePubSubConnection(connection -> connection.getResources().eventBus().publish(new ClusterTopologyChangedEvent(List.of(), List.of())));
 
             getRedisCluster().useWriteCluster(connection -> connection.sync().set(ClientPresenceManager.getPresenceKey(accountUuid, deviceId),
                     UUID.randomUUID().toString()));
@@ -170,8 +170,7 @@ public class ClientPresenceManagerTest extends AbstractRedisClusterTest {
             addClientPresence(missingPeerId);
         }
 
-        getRedisCluster().usePubSubConnection(connection -> connection.sync().masters().commands().subscribe(ClientPresenceManager.getManagerPresenceChannel(presentPeerId)));
+        clientPresenceManager.getPubSubConnection().usePubSubConnection(connection -> connection.sync().masters().commands().subscribe(ClientPresenceManager.getManagerPresenceChannel(presentPeerId)));
-
         clientPresenceManager.pruneMissingPeers();
 
         assertEquals(1, (long)getRedisCluster().withWriteCluster(connection -> connection.sync().exists(ClientPresenceManager.getConnectedClientSetKey(presentPeerId))));

service/src/test/java/org/whispersystems/textsecuregcm/redis/FaultTolerantPubSubConnectionTest.java

@@ -0,0 +1,53 @@
+package org.whispersystems.textsecuregcm.redis;
+
+import io.github.resilience4j.circuitbreaker.CircuitBreakerOpenException;
+import io.lettuce.core.RedisException;
+import io.lettuce.core.cluster.pubsub.StatefulRedisClusterPubSubConnection;
+import io.lettuce.core.cluster.pubsub.api.sync.RedisClusterPubSubCommands;
+import org.junit.Before;
+import org.junit.Test;
+import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class FaultTolerantPubSubConnectionTest {
+
+    private RedisClusterPubSubCommands<String, String>    pubSubCommands;
+    private FaultTolerantPubSubConnection<String, String> faultTolerantPubSubConnection;
+
+    @SuppressWarnings("unchecked")
+    @Before
+    public void setUp() {
+        final StatefulRedisClusterPubSubConnection<String, String> pubSubConnection = mock(StatefulRedisClusterPubSubConnection.class);
+
+        pubSubCommands = mock(RedisClusterPubSubCommands.class);
+
+        when(pubSubConnection.sync()).thenReturn(pubSubCommands);
+
+        final CircuitBreakerConfiguration breakerConfiguration = new CircuitBreakerConfiguration();
+        breakerConfiguration.setFailureRateThreshold(100);
+        breakerConfiguration.setRingBufferSizeInClosedState(1);
+        breakerConfiguration.setWaitDurationInOpenStateInSeconds(Integer.MAX_VALUE);
+
+        faultTolerantPubSubConnection = new FaultTolerantPubSubConnection<>("test", pubSubConnection, breakerConfiguration);
+    }
+
+    @Test
+    public void testBreaker() {
+        when(pubSubCommands.get(anyString()))
+                .thenReturn("value")
+                .thenThrow(new io.lettuce.core.RedisException("Badness has ensued."));
+
+        assertEquals("value", faultTolerantPubSubConnection.withPubSubConnection(connection -> connection.sync().get("key")));
+
+        assertThrows(RedisException.class,
+                () -> faultTolerantPubSubConnection.withPubSubConnection(connection -> connection.sync().get("OH NO")));
+
+        assertThrows(CircuitBreakerOpenException.class,
+                () -> faultTolerantPubSubConnection.withPubSubConnection(connection -> connection.sync().get("OH NO")));
+    }
+}

service/src/test/java/org/whispersystems/textsecuregcm/redis/FaultTolerantRedisClusterTest.java

@@ -6,7 +6,6 @@ import io.lettuce.core.cluster.RedisClusterClient;
 import io.lettuce.core.cluster.api.StatefulRedisClusterConnection;
 import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands;
 import io.lettuce.core.cluster.pubsub.StatefulRedisClusterPubSubConnection;
-import io.lettuce.core.cluster.pubsub.api.sync.RedisClusterPubSubCommands;
 import org.junit.Before;
 import org.junit.Test;
 import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
@@ -21,9 +20,7 @@ import static org.mockito.Mockito.when;
 public class FaultTolerantRedisClusterTest {
 
     private RedisAdvancedClusterCommands<String, String> clusterCommands;
-    private RedisClusterPubSubCommands<String, String>   pubSubCommands;
+    private FaultTolerantRedisCluster                    faultTolerantCluster;
-
-    private FaultTolerantRedisCluster faultTolerantCluster;
 
     @SuppressWarnings("unchecked")
     @Before
@@ -33,12 +30,10 @@ public class FaultTolerantRedisClusterTest {
         final StatefulRedisClusterPubSubConnection<String, String> pubSubConnection  = mock(StatefulRedisClusterPubSubConnection.class);
 
         clusterCommands = mock(RedisAdvancedClusterCommands.class);
-        pubSubCommands  = mock(RedisClusterPubSubCommands.class);
 
         when(clusterClient.connect()).thenReturn(clusterConnection);
         when(clusterClient.connectPubSub()).thenReturn(pubSubConnection);
         when(clusterConnection.sync()).thenReturn(clusterCommands);
-        when(pubSubConnection.sync()).thenReturn(pubSubCommands);
 
         final CircuitBreakerConfiguration breakerConfiguration = new CircuitBreakerConfiguration();
         breakerConfiguration.setFailureRateThreshold(100);
@@ -92,19 +87,4 @@ public class FaultTolerantRedisClusterTest {
         assertThrows(CircuitBreakerOpenException.class,
                 () -> faultTolerantCluster.withWriteCluster(connection -> connection.sync().get("OH NO")));
     }
-
-    @Test
-    public void testPubSubBreaker() {
-        when(pubSubCommands.publish(anyString(), anyString()))
-                .thenReturn(1L)
-                .thenThrow(new RedisException("Badness has ensued."));
-
-        assertEquals(1L, (long)faultTolerantCluster.withPubSubConnection(connection -> connection.sync().publish("channel", "message")));
-
-        assertThrows(RedisException.class,
-                () -> faultTolerantCluster.withPubSubConnection(connection -> connection.sync().publish("channel", "OH NO")));
-
-        assertThrows(CircuitBreakerOpenException.class,
-                () -> faultTolerantCluster.withPubSubConnection(connection -> connection.sync().publish("channel", "OH NO")));
-    }
 }