Move to a synchronous, pooled connection model for Redis clusters.

service/pom.xml

@@ -117,6 +117,12 @@
             <version>5.3.3.RELEASE</version>
         </dependency>
 
+        <dependency>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-pool2</artifactId>
+            <version>2.8.1</version>
+        </dependency>
+
         <dependency>
             <groupId>org.postgresql</groupId>
             <artifactId>postgresql</artifactId>

service/src/main/java/org/whispersystems/textsecuregcm/metrics/PushLatencyManager.java

@@ -3,13 +3,15 @@ package org.whispersystems.textsecuregcm.metrics;
 import com.codahale.metrics.MetricRegistry;
 import com.google.common.annotations.VisibleForTesting;
 import io.lettuce.core.SetArgs;
-import io.lettuce.core.cluster.api.async.RedisAdvancedClusterAsyncCommands;
+import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands;
 import io.micrometer.core.instrument.Metrics;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
 
 import java.time.Duration;
+import java.util.Optional;
 import java.util.UUID;
-import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -27,6 +29,8 @@ public class PushLatencyManager {
 
     private final FaultTolerantRedisCluster redisCluster;
 
+    private static final Logger log = LoggerFactory.getLogger(PushLatencyManager.class);
+
     public PushLatencyManager(final FaultTolerantRedisCluster redisCluster) {
         this.redisCluster = redisCluster;
     }
@@ -37,29 +41,33 @@ public class PushLatencyManager {
 
     @VisibleForTesting
     void recordPushSent(final UUID accountUuid, final long deviceId, final long currentTime) {
+        try {
             redisCluster.useCluster(connection ->
-                connection.async().set(getFirstUnacknowledgedPushKey(accountUuid, deviceId), String.valueOf(currentTime), SetArgs.Builder.nx().ex(TTL)));
+                    connection.sync().set(getFirstUnacknowledgedPushKey(accountUuid, deviceId), String.valueOf(currentTime), SetArgs.Builder.nx().ex(TTL)));
+        } catch (final Exception e) {
+            log.warn("Failed to record \"push notification sent\" timestamp", e);
+        }
     }
 
     public void recordQueueRead(final UUID accountUuid, final long deviceId, final String userAgent) {
-        getLatencyAndClearTimestamp(accountUuid, deviceId, System.currentTimeMillis()).thenAccept(latency -> {
+        final Optional<Long> maybeLatency = getLatencyAndClearTimestamp(accountUuid, deviceId, System.currentTimeMillis());
-            if (latency != null) {
+
-                Metrics.timer(TIMER_NAME, UserAgentTagUtil.getUserAgentTags(userAgent)).record(latency, TimeUnit.MILLISECONDS);
+        if (maybeLatency.isPresent()) {
+            Metrics.timer(TIMER_NAME, UserAgentTagUtil.getUserAgentTags(userAgent)).record(maybeLatency.get(), TimeUnit.MILLISECONDS);
         }
-        });
     }
 
     @VisibleForTesting
-    CompletableFuture<Long> getLatencyAndClearTimestamp(final UUID accountUuid, final long deviceId, final long currentTimeMillis) {
+    Optional<Long> getLatencyAndClearTimestamp(final UUID accountUuid, final long deviceId, final long currentTimeMillis) {
         final String key = getFirstUnacknowledgedPushKey(accountUuid, deviceId);
 
         return redisCluster.withCluster(connection -> {
-            final RedisAdvancedClusterAsyncCommands<String, String> commands = connection.async();
+            final RedisAdvancedClusterCommands<String, String> commands = connection.sync();
 
-            final CompletableFuture<String> getFuture = commands.get(key).toCompletableFuture();
+            final String timestampString = commands.get(key);
             commands.del(key);
 
-            return getFuture.thenApply(timestampString -> timestampString != null ? currentTimeMillis - Long.parseLong(timestampString, 10) : null);
+            return timestampString != null ? Optional.of(currentTimeMillis - Long.parseLong(timestampString, 10)) : Optional.empty();
         });
     }
 

service/src/main/java/org/whispersystems/textsecuregcm/providers/RedisClusterHealthCheck.java

@@ -14,10 +14,8 @@ public class RedisClusterHealthCheck extends HealthCheck {
     }
 
     @Override
-    protected Result check() throws Exception {
+    protected Result check() {
-        return CompletableFuture.allOf(redisCluster.withCluster(connection -> connection.async().masters().commands().ping()).futures())
+        redisCluster.withCluster(connection -> connection.sync().masters().commands().ping());
-                .thenApply(v -> Result.healthy())
+        return Result.healthy();
-                .exceptionally(Result::unhealthy)
-                .get();
     }
 }

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

@@ -196,7 +196,7 @@ public class ClientPresenceManager extends RedisClusterPubSubAdapter<String, Str
     }
 
     private void unsubscribeFromRemotePresenceChanges(final String presenceKey) {
-        pubSubConnection.usePubSubConnection(connection -> connection.async().masters().commands().unsubscribe(getKeyspaceNotificationChannel(presenceKey)));
+        pubSubConnection.usePubSubConnection(connection -> connection.sync().masters().commands().unsubscribe(getKeyspaceNotificationChannel(presenceKey)));
     }
 
     void pruneMissingPeers() {

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

@@ -8,6 +8,11 @@ import io.lettuce.core.cluster.RedisClusterClient;
 import io.lettuce.core.cluster.api.StatefulRedisClusterConnection;
 import io.lettuce.core.cluster.pubsub.StatefulRedisClusterPubSubConnection;
 import io.lettuce.core.codec.ByteArrayCodec;
+import io.lettuce.core.support.ConnectionPoolSupport;
+import org.apache.commons.pool2.impl.GenericObjectPool;
+import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
 import org.whispersystems.textsecuregcm.util.CircuitBreakerUtil;
 import org.whispersystems.textsecuregcm.util.Constants;
@@ -20,9 +25,8 @@ import java.util.function.Function;
 import java.util.stream.Collectors;
 
 /**
- * A fault-tolerant access manager for a Redis cluster. A fault-tolerant Redis cluster has separate circuit breakers for
+ * A fault-tolerant access manager for a Redis cluster. A fault-tolerant Redis cluster provides managed,
- * read and write operations because the leader in a Redis cluster shard may fail while its read-only replicas can still
+ * circuit-breaker-protected access to a pool of connections.
- * serve traffic.
  */
 public class FaultTolerantRedisCluster {
 
@@ -30,14 +34,16 @@ public class FaultTolerantRedisCluster {
 
     private final RedisClusterClient clusterClient;
 
-    private final StatefulRedisClusterConnection<String, String> stringClusterConnection;
+    private final GenericObjectPool<StatefulRedisClusterConnection<String, String>> stringConnectionPool;
-    private final StatefulRedisClusterConnection<byte[], byte[]> binaryClusterConnection;
+    private final GenericObjectPool<StatefulRedisClusterConnection<byte[], byte[]>> binaryConnectionPool;
 
     private final List<StatefulRedisClusterPubSubConnection<?, ?>> pubSubConnections = new ArrayList<>();
 
     private final CircuitBreakerConfiguration circuitBreakerConfiguration;
     private final CircuitBreaker circuitBreaker;
 
+    private static final Logger log = LoggerFactory.getLogger(FaultTolerantRedisCluster.class);
+
     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);
     }
@@ -49,8 +55,11 @@ public class FaultTolerantRedisCluster {
         this.clusterClient = clusterClient;
         this.clusterClient.setDefaultTimeout(timeout);
 
-        this.stringClusterConnection = clusterClient.connect();
+        //noinspection unchecked,rawtypes,rawtypes
-        this.binaryClusterConnection = clusterClient.connect(ByteArrayCodec.INSTANCE);
+        this.stringConnectionPool = ConnectionPoolSupport.createGenericObjectPool(clusterClient::connect, new GenericObjectPoolConfig());
+
+        //noinspection unchecked,rawtypes,rawtypes
+        this.binaryConnectionPool = ConnectionPoolSupport.createGenericObjectPool(() -> clusterClient.connect(ByteArrayCodec.INSTANCE), new GenericObjectPoolConfig());
 
         this.circuitBreakerConfiguration = circuitBreakerConfiguration;
         this.circuitBreaker              = CircuitBreaker.of(name + "-read", circuitBreakerConfiguration.toCircuitBreakerConfig());
@@ -61,8 +70,8 @@ public class FaultTolerantRedisCluster {
     }
 
     void shutdown() {
-        stringClusterConnection.close();
+        stringConnectionPool.close();
-        binaryClusterConnection.close();
+        binaryConnectionPool.close();
 
         for (final StatefulRedisClusterPubSubConnection<?, ?> pubSubConnection : pubSubConnections) {
             pubSubConnection.close();
@@ -72,19 +81,55 @@ public class FaultTolerantRedisCluster {
     }
 
     public void useCluster(final Consumer<StatefulRedisClusterConnection<String, String>> consumer) {
-        this.circuitBreaker.executeRunnable(() -> consumer.accept(stringClusterConnection));
+        acceptPooledConnection(stringConnectionPool, consumer);
     }
 
-    public <T> T withCluster(final Function<StatefulRedisClusterConnection<String, String>, T> consumer) {
+    public <T> T withCluster(final Function<StatefulRedisClusterConnection<String, String>, T> function) {
-        return this.circuitBreaker.executeSupplier(() -> consumer.apply(stringClusterConnection));
+        return applyToPooledConnection(stringConnectionPool, function);
     }
 
     public void useBinaryCluster(final Consumer<StatefulRedisClusterConnection<byte[], byte[]>> consumer) {
-        this.circuitBreaker.executeRunnable(() -> consumer.accept(binaryClusterConnection));
+        acceptPooledConnection(binaryConnectionPool, consumer);
     }
 
-    public <T> T withBinaryCluster(final Function<StatefulRedisClusterConnection<byte[], byte[]>, T> consumer) {
+    public <T> T withBinaryCluster(final Function<StatefulRedisClusterConnection<byte[], byte[]>, T> function) {
-        return this.circuitBreaker.executeSupplier(() -> consumer.apply(binaryClusterConnection));
+        return applyToPooledConnection(binaryConnectionPool, function);
+    }
+
+    private <K, V> void acceptPooledConnection(final GenericObjectPool<StatefulRedisClusterConnection<K, V>> pool, final Consumer<StatefulRedisClusterConnection<K, V>> consumer) {
+        try {
+            circuitBreaker.executeCheckedRunnable(() -> {
+                try (final StatefulRedisClusterConnection<K, V> connection = pool.borrowObject()) {
+                    consumer.accept(connection);
+                }
+            });
+        } catch (final Throwable t) {
+            log.warn("Redis operation failure", t);
+
+            if (t instanceof RuntimeException) {
+                throw (RuntimeException) t;
+            } else {
+                throw new RuntimeException(t);
+            }
+        }
+    }
+
+    private <T, K, V> T applyToPooledConnection(final GenericObjectPool<StatefulRedisClusterConnection<K, V>> pool, final Function<StatefulRedisClusterConnection<K, V>, T> function) {
+        try {
+            return circuitBreaker.executeCheckedSupplier(() -> {
+                try (final StatefulRedisClusterConnection<K, V> connection = pool.borrowObject()) {
+                    return function.apply(connection);
+                }
+            });
+        } catch (final Throwable t) {
+            log.warn("Redis operation failure", t);
+
+            if (t instanceof RuntimeException) {
+                throw (RuntimeException) t;
+            } else {
+                throw new RuntimeException(t);
+            }
+        }
     }
 
     public FaultTolerantPubSubConnection<String, String> createPubSubConnection() {

service/src/test/java/org/whispersystems/textsecuregcm/metrics/PushLatencyManagerTest.java

@@ -1,37 +1,17 @@
 package org.whispersystems.textsecuregcm.metrics;
 
-import org.junit.After;
-import org.junit.Before;
 import org.junit.Test;
-import org.mockito.stubbing.Answer;
 import org.whispersystems.textsecuregcm.redis.AbstractRedisClusterTest;
-import org.whispersystems.textsecuregcm.redis.ReplicatedJedisPool;
-import redis.clients.jedis.Jedis;
-import redis.embedded.RedisServer;
 
 import java.util.Optional;
 import java.util.UUID;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.ThreadPoolExecutor;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNull;
-import static org.junit.Assert.fail;
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.Mockito.doAnswer;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
 
 public class PushLatencyManagerTest extends AbstractRedisClusterTest {
 
     @Test
-    public void testGetLatency() throws ExecutionException, InterruptedException {
+    public void testGetLatency() {
         final PushLatencyManager pushLatencyManager  = new PushLatencyManager(getRedisCluster());
         final UUID               accountUuid         = UUID.randomUUID();
         final long               deviceId            = 1;
@@ -39,13 +19,13 @@ public class PushLatencyManagerTest extends AbstractRedisClusterTest {
         final long               pushSentTimestamp   = System.currentTimeMillis();
         final long               clearQueueTimestamp = pushSentTimestamp + expectedLatency;
 
-        assertNull(pushLatencyManager.getLatencyAndClearTimestamp(accountUuid, deviceId, System.currentTimeMillis()).get());
+        assertEquals(Optional.empty(), pushLatencyManager.getLatencyAndClearTimestamp(accountUuid, deviceId, System.currentTimeMillis()));
 
         {
             pushLatencyManager.recordPushSent(accountUuid, deviceId, pushSentTimestamp);
 
-            assertEquals(expectedLatency, (long)pushLatencyManager.getLatencyAndClearTimestamp(accountUuid, deviceId, clearQueueTimestamp).get());
+            assertEquals(Optional.of(expectedLatency), pushLatencyManager.getLatencyAndClearTimestamp(accountUuid, deviceId, clearQueueTimestamp));
-            assertNull(pushLatencyManager.getLatencyAndClearTimestamp(accountUuid, deviceId, System.currentTimeMillis()).get());
+            assertEquals(Optional.empty(), pushLatencyManager.getLatencyAndClearTimestamp(accountUuid, deviceId, System.currentTimeMillis()));
         }
     }
 }