Introduce FaultTolerantRedisClient

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

@@ -0,0 +1,65 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.redis;
+
+import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
+
+import io.github.resilience4j.retry.Retry;
+import io.lettuce.core.RedisException;
+import io.lettuce.core.pubsub.StatefulRedisPubSubConnection;
+import io.micrometer.core.instrument.Metrics;
+import io.micrometer.core.instrument.Timer;
+import java.util.function.Consumer;
+import java.util.function.Function;
+
+abstract class AbstractFaultTolerantPubSubConnection<K, V, C extends StatefulRedisPubSubConnection<K, V>> {
+
+  private final String name;
+  private final C pubSubConnection;
+
+  private final Retry retry;
+
+  private final Timer executeTimer;
+
+  protected AbstractFaultTolerantPubSubConnection(final String name,
+      final C pubSubConnection,
+      final Retry retry) {
+
+    this.name = name;
+    this.pubSubConnection = pubSubConnection;
+    this.retry = retry;
+
+    this.executeTimer = Metrics.timer(name(getClass(), "execute"), "clusterName", name + "-pubsub");
+  }
+
+  protected String getName() {
+    return name;
+  }
+
+  public void usePubSubConnection(final Consumer<C> consumer) {
+    try {
+      retry.executeRunnable(() -> executeTimer.record(() -> consumer.accept(pubSubConnection)));
+    } catch (final Throwable t) {
+      if (t instanceof RedisException) {
+        throw (RedisException) t;
+      } else {
+        throw new RedisException(t);
+      }
+    }
+  }
+
+  public <T> T withPubSubConnection(final Function<C, T> function) {
+    try {
+      return retry.executeCallable(() -> executeTimer.record(() -> function.apply(pubSubConnection)));
+    } catch (final Throwable t) {
+      if (t instanceof RedisException) {
+        throw (RedisException) t;
+      } else {
+        throw new RedisException(t);
+      }
+    }
+  }
+}

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

@@ -25,7 +25,7 @@ import reactor.core.publisher.Mono;
 
 public class ClusterLuaScript {
 
-  private final FaultTolerantRedisCluster redisCluster;
+  private final FaultTolerantRedisClusterClient redisCluster;
   private final ScriptOutputType scriptOutputType;
   private final String script;
   private final String sha;
@@ -35,7 +35,7 @@ public class ClusterLuaScript {
 
   private static final Logger log = LoggerFactory.getLogger(ClusterLuaScript.class);
 
-  public static ClusterLuaScript fromResource(final FaultTolerantRedisCluster redisCluster,
+  public static ClusterLuaScript fromResource(final FaultTolerantRedisClusterClient redisCluster,
       final String resource,
       final ScriptOutputType scriptOutputType) throws IOException {
 
@@ -51,7 +51,7 @@ public class ClusterLuaScript {
   }
 
   @VisibleForTesting
-  ClusterLuaScript(final FaultTolerantRedisCluster redisCluster,
+  ClusterLuaScript(final FaultTolerantRedisClusterClient redisCluster,
       final String script,
       final ScriptOutputType scriptOutputType) {
 

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

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.redis;
+
+import io.github.resilience4j.retry.Retry;
+import io.lettuce.core.cluster.event.ClusterTopologyChangedEvent;
+import io.lettuce.core.cluster.pubsub.StatefulRedisClusterPubSubConnection;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import reactor.core.scheduler.Scheduler;
+
+public class FaultTolerantPubSubClusterConnection<K, V> extends AbstractFaultTolerantPubSubConnection<K, V, StatefulRedisClusterPubSubConnection<K, V>> {
+
+  private final Logger logger = LoggerFactory.getLogger(FaultTolerantPubSubClusterConnection.class);
+
+  private final Retry resubscribeRetry;
+  private final Scheduler topologyChangedEventScheduler;
+
+  protected FaultTolerantPubSubClusterConnection(final String name,
+      final StatefulRedisClusterPubSubConnection<K, V> pubSubConnection,
+      final Retry retry,
+      final Retry resubscribeRetry,
+      final Scheduler topologyChangedEventScheduler) {
+
+    super(name, pubSubConnection, retry);
+
+    pubSubConnection.setNodeMessagePropagation(true);
+
+    this.resubscribeRetry = resubscribeRetry;
+    this.topologyChangedEventScheduler = topologyChangedEventScheduler;
+  }
+
+  public void subscribeToClusterTopologyChangedEvents(final Runnable eventHandler) {
+
+    usePubSubConnection(connection -> connection.getResources().eventBus().get()
+        .filter(event -> event instanceof ClusterTopologyChangedEvent)
+        .subscribeOn(topologyChangedEventScheduler)
+        .subscribe(event -> {
+          logger.info("Got topology change event for {}, resubscribing all keyspace notifications", getName());
+
+          resubscribeRetry.executeRunnable(() -> {
+            try {
+              eventHandler.run();
+            } catch (final RuntimeException e) {
+              logger.warn("Resubscribe for {} failed", getName(), e);
+              throw e;
+            }
+          });
+        }));
+  }
+}

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

@@ -5,90 +5,15 @@
 
 package org.whispersystems.textsecuregcm.redis;
 
-import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
-
 import io.github.resilience4j.retry.Retry;
-import io.lettuce.core.RedisException;
-import io.lettuce.core.cluster.event.ClusterTopologyChangedEvent;
-import io.lettuce.core.cluster.pubsub.StatefulRedisClusterPubSubConnection;
-import io.micrometer.core.instrument.Metrics;
-import io.micrometer.core.instrument.Timer;
-import java.util.function.Consumer;
-import java.util.function.Function;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-import reactor.core.scheduler.Scheduler;
+import io.lettuce.core.pubsub.StatefulRedisPubSubConnection;
 
-public class FaultTolerantPubSubConnection<K, V> {
+public class FaultTolerantPubSubConnection<K, V> extends AbstractFaultTolerantPubSubConnection<K, V, StatefulRedisPubSubConnection<K, V>> {
 
-  private static final Logger logger = LoggerFactory.getLogger(FaultTolerantPubSubConnection.class);
+  protected FaultTolerantPubSubConnection(final String name,
+      final StatefulRedisPubSubConnection<K, V> pubSubConnection,
+      final Retry retry) {
 
-
-  private final String name;
-  private final StatefulRedisClusterPubSubConnection<K, V> pubSubConnection;
-
-  private final Retry retry;
-  private final Retry resubscribeRetry;
-  private final Scheduler topologyChangedEventScheduler;
-
-  private final Timer executeTimer;
-
-  public FaultTolerantPubSubConnection(final String name,
-      final StatefulRedisClusterPubSubConnection<K, V> pubSubConnection,
-      final Retry retry, final Retry resubscribeRetry, final Scheduler topologyChangedEventScheduler) {
-    this.name = name;
-    this.pubSubConnection = pubSubConnection;
-    this.retry = retry;
-    this.resubscribeRetry = resubscribeRetry;
-    this.topologyChangedEventScheduler = topologyChangedEventScheduler;
-
-    this.pubSubConnection.setNodeMessagePropagation(true);
-
-    this.executeTimer = Metrics.timer(name(getClass(), "execute"), "clusterName", name + "-pubsub");
+    super(name, pubSubConnection, retry);
   }
-
-  public void usePubSubConnection(final Consumer<StatefulRedisClusterPubSubConnection<K, V>> consumer) {
-    try {
-      retry.executeRunnable(() -> executeTimer.record(() -> consumer.accept(pubSubConnection)));
-    } catch (final Throwable t) {
-      if (t instanceof RedisException) {
-        throw (RedisException) t;
-      } else {
-        throw new RedisException(t);
-      }
-    }
-  }
-
-  public <T> T withPubSubConnection(final Function<StatefulRedisClusterPubSubConnection<K, V>, T> function) {
-    try {
-      return retry.executeCallable(() -> executeTimer.record(() -> function.apply(pubSubConnection)));
-    } catch (final Throwable t) {
-      if (t instanceof RedisException) {
-        throw (RedisException) t;
-      } else {
-        throw new RedisException(t);
-      }
-    }
-  }
-
-
-  public void subscribeToClusterTopologyChangedEvents(final Runnable eventHandler) {
-
-    usePubSubConnection(connection -> connection.getResources().eventBus().get()
-        .filter(event -> event instanceof ClusterTopologyChangedEvent)
-        .subscribeOn(topologyChangedEventScheduler)
-        .subscribe(event -> {
-          logger.info("Got topology change event for {}, resubscribing all keyspace notifications", name);
-
-          resubscribeRetry.executeRunnable(() -> {
-            try {
-              eventHandler.run();
-            } catch (final RuntimeException e) {
-              logger.warn("Resubscribe for {} failed", name, e);
-              throw e;
-            }
-          });
-        }));
-  }
-
 }

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

@@ -0,0 +1,159 @@
+package org.whispersystems.textsecuregcm.redis;
+
+import io.github.resilience4j.circuitbreaker.CircuitBreaker;
+import io.github.resilience4j.retry.Retry;
+import io.lettuce.core.ClientOptions;
+import io.lettuce.core.RedisClient;
+import io.lettuce.core.RedisCommandTimeoutException;
+import io.lettuce.core.RedisException;
+import io.lettuce.core.RedisURI;
+import io.lettuce.core.TimeoutOptions;
+import io.lettuce.core.api.StatefulRedisConnection;
+import io.lettuce.core.cluster.ClusterClientOptions;
+import io.lettuce.core.cluster.ClusterTopologyRefreshOptions;
+import io.lettuce.core.codec.ByteArrayCodec;
+import io.lettuce.core.pubsub.StatefulRedisPubSubConnection;
+import io.lettuce.core.resource.ClientResources;
+import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
+import org.whispersystems.textsecuregcm.configuration.RedisConfiguration;
+import org.whispersystems.textsecuregcm.configuration.RetryConfiguration;
+import org.whispersystems.textsecuregcm.util.CircuitBreakerUtil;
+import reactor.core.scheduler.Schedulers;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.Consumer;
+import java.util.function.Function;
+
+public class FaultTolerantRedisClient {
+
+  private final String name;
+
+  private final RedisClient redisClient;
+
+  private final StatefulRedisConnection<String, String> stringConnection;
+  private final StatefulRedisConnection<byte[], byte[]> binaryConnection;
+
+  private final List<StatefulRedisPubSubConnection<?, ?>> pubSubConnections = new ArrayList<>();
+
+  private final CircuitBreaker circuitBreaker;
+  private final Retry retry;
+
+  public FaultTolerantRedisClient(final String name,
+                                  final RedisConfiguration redisConfiguration,
+                                  final ClientResources.Builder clientResourcesBuilder) {
+
+    this(name, clientResourcesBuilder,
+        RedisUriUtil.createRedisUriWithTimeout(redisConfiguration.getUri(), redisConfiguration.getTimeout()),
+        redisConfiguration.getTimeout(),
+        redisConfiguration.getCircuitBreakerConfiguration(),
+        redisConfiguration.getRetryConfiguration());
+  }
+
+  FaultTolerantRedisClient(String name,
+                           final ClientResources.Builder clientResourcesBuilder,
+                           final RedisURI redisUri,
+                           final Duration commandTimeout,
+                           final CircuitBreakerConfiguration circuitBreakerConfiguration,
+                           final RetryConfiguration retryConfiguration) {
+
+    this.name = name;
+
+    final LettuceShardCircuitBreaker lettuceShardCircuitBreaker = new LettuceShardCircuitBreaker(name,
+        circuitBreakerConfiguration.toCircuitBreakerConfig(), Schedulers.newSingle("topology-changed-" + name, true));
+    this.redisClient = RedisClient.create(clientResourcesBuilder.build(), redisUri);
+    this.redisClient.setOptions(ClusterClientOptions.builder()
+        .disconnectedBehavior(ClientOptions.DisconnectedBehavior.REJECT_COMMANDS)
+        .validateClusterNodeMembership(false)
+        .topologyRefreshOptions(ClusterTopologyRefreshOptions.builder()
+            .enableAllAdaptiveRefreshTriggers()
+            .build())
+        // for asynchronous commands
+        .timeoutOptions(TimeoutOptions.builder()
+            .fixedTimeout(commandTimeout)
+            .build())
+        .publishOnScheduler(true)
+        .build());
+
+    lettuceShardCircuitBreaker.setEventBus(redisClient.getResources().eventBus());
+
+    this.stringConnection = redisClient.connect();
+    this.binaryConnection = redisClient.connect(ByteArrayCodec.INSTANCE);
+
+    this.circuitBreaker = CircuitBreaker.of(name + "-breaker", circuitBreakerConfiguration.toCircuitBreakerConfig());
+    this.retry = Retry.of(name + "-retry", retryConfiguration.toRetryConfigBuilder()
+        .retryOnException(exception -> exception instanceof RedisCommandTimeoutException).build());
+
+    CircuitBreakerUtil.registerMetrics(retry, FaultTolerantRedisClusterClient.class);
+  }
+
+  public void shutdown() {
+    stringConnection.close();
+
+    for (final StatefulRedisPubSubConnection<?, ?> pubSubConnection : pubSubConnections) {
+      pubSubConnection.close();
+    }
+
+    redisClient.shutdown();
+  }
+
+  public String getName() {
+    return name;
+  }
+
+  public void useConnection(final Consumer<StatefulRedisConnection<String, String>> consumer) {
+    useConnection(stringConnection, consumer);
+  }
+
+  public <T> T withConnection(final Function<StatefulRedisConnection<String, String>, T> function) {
+    return withConnection(stringConnection, function);
+  }
+
+  public void useBinaryConnection(final Consumer<StatefulRedisConnection<byte[], byte[]>> consumer) {
+    useConnection(binaryConnection, consumer);
+  }
+
+  public <T> T withBinaryConnection(final Function<StatefulRedisConnection<byte[], byte[]>, T> function) {
+    return withConnection(binaryConnection, function);
+  }
+
+  public <K, V> void useConnection(final StatefulRedisConnection<K, V> connection,
+      final Consumer<StatefulRedisConnection<K, V>> consumer) {
+    try {
+      circuitBreaker.executeRunnable(() -> retry.executeRunnable(() -> consumer.accept(connection)));
+    } catch (final Throwable t) {
+      if (t instanceof RedisException) {
+        throw (RedisException) t;
+      } else {
+        throw new RedisException(t);
+      }
+    }
+  }
+
+  public <T, K, V> T withConnection(final StatefulRedisConnection<K, V> connection,
+      final Function<StatefulRedisConnection<K, V>, T> function) {
+    try {
+      return circuitBreaker.executeCallable(() -> retry.executeCallable(() -> function.apply(connection)));
+    } catch (final Throwable t) {
+      if (t instanceof RedisException) {
+        throw (RedisException) t;
+      } else {
+        throw new RedisException(t);
+      }
+    }
+  }
+
+  public FaultTolerantPubSubConnection<String, String> createPubSubConnection() {
+    final StatefulRedisPubSubConnection<String, String> pubSubConnection = redisClient.connectPubSub();
+    pubSubConnections.add(pubSubConnection);
+
+    return new FaultTolerantPubSubConnection<>(name, pubSubConnection, retry);
+  }
+
+  public FaultTolerantPubSubConnection<byte[], byte[]> createBinaryPubSubConnection() {
+    final StatefulRedisPubSubConnection<byte[], byte[]> pubSubConnection = redisClient.connectPubSub(ByteArrayCodec.INSTANCE);
+    pubSubConnections.add(pubSubConnection);
+
+    return new FaultTolerantPubSubConnection<>(name, pubSubConnection, retry);
+  }
+}

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

@@ -41,7 +41,7 @@ import reactor.core.scheduler.Schedulers;
  *
  * @see LettuceShardCircuitBreaker
  */
-public class FaultTolerantRedisCluster {
+public class FaultTolerantRedisClusterClient {
 
   private final String name;
 
@@ -56,8 +56,8 @@ public class FaultTolerantRedisCluster {
   private final Retry topologyChangedEventRetry;
 
 
-  public FaultTolerantRedisCluster(final String name, final RedisClusterConfiguration clusterConfiguration,
-      final ClientResources.Builder clientResourcesBuilder) {
+  public FaultTolerantRedisClusterClient(final String name, final RedisClusterConfiguration clusterConfiguration,
+                                         final ClientResources.Builder clientResourcesBuilder) {
 
     this(name, clientResourcesBuilder,
         Collections.singleton(RedisUriUtil.createRedisUriWithTimeout(clusterConfiguration.getConfigurationUri(),
@@ -68,9 +68,9 @@ public class FaultTolerantRedisCluster {
 
   }
 
-  FaultTolerantRedisCluster(String name, final ClientResources.Builder clientResourcesBuilder,
-      Iterable<RedisURI> redisUris, Duration commandTimeout, CircuitBreakerConfiguration circuitBreakerConfig,
-      RetryConfiguration retryConfiguration) {
+  FaultTolerantRedisClusterClient(String name, final ClientResources.Builder clientResourcesBuilder,
+                                  Iterable<RedisURI> redisUris, Duration commandTimeout, CircuitBreakerConfiguration circuitBreakerConfig,
+                                  RetryConfiguration retryConfiguration) {
 
     this.name = name;
 
@@ -112,7 +112,7 @@ public class FaultTolerantRedisCluster {
 
     this.topologyChangedEventRetry = Retry.of(name + "-topologyChangedRetry", topologyChangedEventRetryConfig);
 
-    CircuitBreakerUtil.registerMetrics(retry, FaultTolerantRedisCluster.class);
+    CircuitBreakerUtil.registerMetrics(retry, FaultTolerantRedisClusterClient.class);
   }
 
   public void shutdown() {
@@ -184,11 +184,11 @@ public class FaultTolerantRedisCluster {
         .transformDeferred(RetryOperator.of(retry));
   }
 
-  public FaultTolerantPubSubConnection<String, String> createPubSubConnection() {
+  public FaultTolerantPubSubClusterConnection<String, String> createPubSubConnection() {
     final StatefulRedisClusterPubSubConnection<String, String> pubSubConnection = clusterClient.connectPubSub();
     pubSubConnections.add(pubSubConnection);
 
-    return new FaultTolerantPubSubConnection<>(name, pubSubConnection, retry, topologyChangedEventRetry,
+    return new FaultTolerantPubSubClusterConnection<>(name, pubSubConnection, retry, topologyChangedEventRetry,
         Schedulers.newSingle(name + "-redisPubSubEvents", true));
   }