Use reactive streams for WebSocket message queue

Initially, uses `ExperimentEnrollmentManager` to do a safe rollout.

pom.xml

@@ -57,7 +57,7 @@
     <jedis.version>2.9.0</jedis.version>
     <kotlin.version>1.7.10</kotlin.version>
     <kotlinx-serialization.version>1.4.0</kotlinx-serialization.version>
-    <lettuce.version>6.1.9.RELEASE</lettuce.version>
+    <lettuce.version>6.2.0.RELEASE</lettuce.version>
     <libphonenumber.version>8.12.54</libphonenumber.version>
     <logstash.logback.version>7.0.1</logstash.logback.version>
     <micrometer.version>1.9.3</micrometer.version>
@@ -151,6 +151,13 @@
         <type>pom</type>
         <scope>import</scope>
       </dependency>
+      <dependency>
+        <groupId>io.projectreactor</groupId>
+        <artifactId>reactor-bom</artifactId>
+        <version>2020.0.23</version> <!-- 3.4.x, see https://github.com/reactor/reactor#bom-versioning-scheme -->
+        <type>pom</type>
+        <scope>import</scope>
+      </dependency>
       <dependency>
         <groupId>com.eatthepath</groupId>
         <artifactId>pushy</artifactId>

service/pom.xml

@@ -228,6 +228,10 @@
       <groupId>io.github.resilience4j</groupId>
       <artifactId>resilience4j-retry</artifactId>
     </dependency>
+    <dependency>
+      <groupId>io.github.resilience4j</groupId>
+      <artifactId>resilience4j-reactor</artifactId>
+    </dependency>
 
     <dependency>
       <groupId>io.grpc</groupId>
@@ -407,7 +411,6 @@
     <dependency>
       <groupId>io.projectreactor</groupId>
       <artifactId>reactor-core</artifactId>
-      <version>3.3.22.RELEASE</version>
     </dependency>
     <dependency>
       <groupId>io.vavr</groupId>
@@ -420,6 +423,11 @@
       <scope>test</scope>
     </dependency>
 
+    <dependency>
+      <groupId>io.projectreactor</groupId>
+      <artifactId>reactor-test</artifactId>
+    </dependency>
+
     <dependency>
       <groupId>org.signal</groupId>
       <artifactId>embedded-redis</artifactId>

service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java

@@ -116,7 +116,6 @@ import org.whispersystems.textsecuregcm.filters.TimestampResponseFilter;
 import org.whispersystems.textsecuregcm.limits.DynamicRateLimiters;
 import org.whispersystems.textsecuregcm.limits.PushChallengeManager;
 import org.whispersystems.textsecuregcm.limits.RateLimitChallengeManager;
-import org.whispersystems.textsecuregcm.limits.RateLimitChallengeOptionManager;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.mappers.CompletionExceptionMapper;
 import org.whispersystems.textsecuregcm.mappers.DeviceLimitExceededExceptionMapper;
@@ -330,6 +329,13 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
             config.getAppConfig().getConfigurationName(),
             DynamicConfiguration.class);
 
+    BlockingQueue<Runnable> messageDeletionQueue = new ArrayBlockingQueue<>(10_000);
+    Metrics.gaugeCollectionSize(name(getClass(), "messageDeletionQueueSize"), Collections.emptyList(),
+        messageDeletionQueue);
+    ExecutorService messageDeletionAsyncExecutor = environment.lifecycle()
+        .executorService(name(getClass(), "messageDeletionAsyncExecutor-%d")).maxThreads(16)
+        .workQueue(messageDeletionQueue).build();
+
     Accounts accounts = new Accounts(dynamicConfigurationManager,
         dynamoDbClient,
         dynamoDbAsyncClient,
@@ -345,9 +351,10 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     Profiles profiles = new Profiles(dynamoDbClient, dynamoDbAsyncClient,
         config.getDynamoDbTables().getProfiles().getTableName());
     Keys keys = new Keys(dynamoDbClient, config.getDynamoDbTables().getKeys().getTableName());
-    MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbClient,
+    MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbClient, dynamoDbAsyncClient,
         config.getDynamoDbTables().getMessages().getTableName(),
-        config.getDynamoDbTables().getMessages().getExpiration());
+        config.getDynamoDbTables().getMessages().getExpiration(),
+        messageDeletionAsyncExecutor);
     RemoteConfigs remoteConfigs = new RemoteConfigs(dynamoDbClient,
         config.getDynamoDbTables().getRemoteConfig().getTableName());
     PushChallengeDynamoDb pushChallengeDynamoDb = new PushChallengeDynamoDb(dynamoDbClient,
@@ -453,7 +460,8 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     StoredVerificationCodeManager pendingAccountsManager  = new StoredVerificationCodeManager(pendingAccounts);
     StoredVerificationCodeManager pendingDevicesManager   = new StoredVerificationCodeManager(pendingDevices);
     ProfilesManager profilesManager = new ProfilesManager(profiles, cacheCluster);
-    MessagesCache              messagesCache              = new MessagesCache(messagesCluster, messagesCluster, keyspaceNotificationDispatchExecutor);
+    MessagesCache messagesCache = new MessagesCache(messagesCluster, messagesCluster, Clock.systemUTC(),
+        keyspaceNotificationDispatchExecutor, messageDeletionAsyncExecutor);
     PushLatencyManager pushLatencyManager = new PushLatencyManager(metricsCluster, dynamicConfigurationManager);
     ReportMessageManager       reportMessageManager       = new ReportMessageManager(reportMessageDynamoDb, rateLimitersCluster, config.getReportMessageConfiguration().getCounterTtl());
     MessagesManager            messagesManager            = new MessagesManager(messagesDynamoDb, messagesCache, reportMessageManager);
@@ -503,8 +511,6 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     PushChallengeManager pushChallengeManager = new PushChallengeManager(pushNotificationManager, pushChallengeDynamoDb);
     RateLimitChallengeManager rateLimitChallengeManager = new RateLimitChallengeManager(pushChallengeManager,
         recaptchaClient, dynamicRateLimiters);
-    RateLimitChallengeOptionManager rateLimitChallengeOptionManager =
-        new RateLimitChallengeOptionManager(dynamicRateLimiters, dynamicConfigurationManager);
 
     MessagePersister messagePersister = new MessagePersister(messagesCache, messagesManager, accountsManager, dynamicConfigurationManager, Duration.ofMinutes(config.getMessageCacheConfiguration().getPersistDelayMinutes()));
     ChangeNumberManager changeNumberManager = new ChangeNumberManager(messageSender, accountsManager);
@@ -628,8 +634,9 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     webSocketEnvironment.setAuthenticator(new WebSocketAccountAuthenticator(accountAuthenticator));
     webSocketEnvironment.setConnectListener(
         new AuthenticatedConnectListener(receiptSender, messagesManager, pushNotificationManager,
-            clientPresenceManager, websocketScheduledExecutor));
+            clientPresenceManager, websocketScheduledExecutor, experimentEnrollmentManager));
-    webSocketEnvironment.jersey().register(new WebsocketRefreshApplicationEventListener(accountsManager, clientPresenceManager));
+    webSocketEnvironment.jersey()
+        .register(new WebsocketRefreshApplicationEventListener(accountsManager, clientPresenceManager));
     webSocketEnvironment.jersey().register(new ContentLengthFilter(TrafficSource.WEBSOCKET));
     webSocketEnvironment.jersey().register(MultiRecipientMessageProvider.class);
     webSocketEnvironment.jersey().register(new MetricsApplicationEventListener(TrafficSource.WEBSOCKET));

service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2021 Signal Messenger, LLC
+ * Copyright 2013-2022 Signal Messenger, LLC
  * SPDX-License-Identifier: AGPL-3.0-only
  */
 package org.whispersystems.textsecuregcm.controllers;
@@ -30,6 +30,7 @@ import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.Callable;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
@@ -538,14 +539,17 @@ public class MessageController {
   @Timed
   @DELETE
   @Path("/uuid/{uuid}")
-  public void removePendingMessage(@Auth AuthenticatedAccount auth, @PathParam("uuid") UUID uuid) {
+  public CompletableFuture<Void> removePendingMessage(@Auth AuthenticatedAccount auth, @PathParam("uuid") UUID uuid) {
-    messagesManager.delete(
+    return messagesManager.delete(
             auth.getAccount().getUuid(),
             auth.getAuthenticatedDevice().getId(),
             uuid,
-        null).ifPresent(deletedMessage -> {
+            null)
+        .thenAccept(maybeDeletedMessage -> {
+          maybeDeletedMessage.ifPresent(deletedMessage -> {
 
-      WebSocketConnection.recordMessageDeliveryDuration(deletedMessage.getTimestamp(), auth.getAuthenticatedDevice());
+            WebSocketConnection.recordMessageDeliveryDuration(deletedMessage.getTimestamp(),
+                auth.getAuthenticatedDevice());
 
             if (deletedMessage.hasSourceUuid() && deletedMessage.getType() != Type.SERVER_DELIVERY_RECEIPT) {
               try {
@@ -557,6 +561,7 @@ public class MessageController {
               }
             }
           });
+        });
   }
 
   @Timed

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

@@ -6,6 +6,7 @@
 package org.whispersystems.textsecuregcm.redis;
 
 import com.google.common.annotations.VisibleForTesting;
+import io.lettuce.core.RedisException;
 import io.lettuce.core.RedisNoScriptException;
 import io.lettuce.core.ScriptOutputType;
 import io.lettuce.core.cluster.api.StatefulRedisClusterConnection;
@@ -15,9 +16,12 @@ import java.nio.charset.StandardCharsets;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 import org.apache.commons.codec.binary.Hex;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 
 public class ClusterLuaScript {
 
@@ -73,11 +77,31 @@ public class ClusterLuaScript {
         execute(connection, keys.toArray(STRING_ARRAY), args.toArray(STRING_ARRAY)));
   }
 
+  public CompletableFuture<Object> executeAsync(final List<String> keys, final List<String> args) {
+    return redisCluster.withCluster(connection ->
+        executeAsync(connection, keys.toArray(STRING_ARRAY), args.toArray(STRING_ARRAY)));
+  }
+
+  public Flux<Object> executeReactive(final List<String> keys, final List<String> args) {
+    return redisCluster.withCluster(connection ->
+        executeReactive(connection, keys.toArray(STRING_ARRAY), args.toArray(STRING_ARRAY)));
+  }
+
   public Object executeBinary(final List<byte[]> keys, final List<byte[]> args) {
     return redisCluster.withBinaryCluster(connection ->
         execute(connection, keys.toArray(BYTE_ARRAY_ARRAY), args.toArray(BYTE_ARRAY_ARRAY)));
   }
 
+  public CompletableFuture<Object> executeBinaryAsync(final List<byte[]> keys, final List<byte[]> args) {
+    return redisCluster.withBinaryCluster(connection ->
+        executeAsync(connection, keys.toArray(BYTE_ARRAY_ARRAY), args.toArray(BYTE_ARRAY_ARRAY)));
+  }
+
+  public Flux<Object> executeBinaryReactive(final List<byte[]> keys, final List<byte[]> args) {
+    return redisCluster.withBinaryCluster(connection ->
+        executeReactive(connection, keys.toArray(BYTE_ARRAY_ARRAY), args.toArray(BYTE_ARRAY_ARRAY)));
+  }
+
   private <T> Object execute(final StatefulRedisClusterConnection<T, T> connection, final T[] keys, final T[] args) {
     try {
       try {
@@ -90,4 +114,32 @@ public class ClusterLuaScript {
       throw e;
     }
   }
+
+  private <T> CompletableFuture<Object> executeAsync(final StatefulRedisClusterConnection<T, T> connection,
+      final T[] keys, final T[] args) {
+
+    return connection.async().evalsha(sha, scriptOutputType, keys, args)
+        .exceptionallyCompose(throwable -> {
+          if (throwable instanceof RedisNoScriptException) {
+            return connection.async().eval(script, scriptOutputType, keys, args);
+          }
+
+          log.warn("Failed to execute script", throwable);
+          throw new RedisException(throwable);
+        }).toCompletableFuture();
+  }
+
+  private <T> Flux<Object> executeReactive(final StatefulRedisClusterConnection<T, T> connection,
+      final T[] keys, final T[] args) {
+
+    return connection.reactive().evalsha(sha, scriptOutputType, keys, args)
+        .onErrorResume(e -> {
+          if (e instanceof RedisNoScriptException) {
+            return connection.reactive().eval(script, scriptOutputType, keys, args);
+          }
+
+          log.warn("Failed to execute script", e);
+          return Mono.error(e);
+        });
+  }
 }

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

@@ -8,6 +8,8 @@ package org.whispersystems.textsecuregcm.redis;
 import com.codahale.metrics.SharedMetricRegistries;
 import com.google.common.annotations.VisibleForTesting;
 import io.github.resilience4j.circuitbreaker.CircuitBreaker;
+import io.github.resilience4j.reactor.circuitbreaker.operator.CircuitBreakerOperator;
+import io.github.resilience4j.reactor.retry.RetryOperator;
 import io.github.resilience4j.retry.Retry;
 import io.lettuce.core.ClientOptions.DisconnectedBehavior;
 import io.lettuce.core.RedisCommandTimeoutException;
@@ -24,11 +26,13 @@ import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Consumer;
 import java.util.function.Function;
+import org.reactivestreams.Publisher;
 import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
 import org.whispersystems.textsecuregcm.configuration.RedisClusterConfiguration;
 import org.whispersystems.textsecuregcm.configuration.RetryConfiguration;
 import org.whispersystems.textsecuregcm.util.CircuitBreakerUtil;
 import org.whispersystems.textsecuregcm.util.Constants;
+import reactor.core.publisher.Flux;
 
 /**
  * A fault-tolerant access manager for a Redis cluster. A fault-tolerant Redis cluster provides managed,
@@ -111,7 +115,13 @@ public class FaultTolerantRedisCluster {
     return withConnection(binaryConnection, function);
   }
 
-    private <K, V> void useConnection(final StatefulRedisClusterConnection<K, V> connection, final Consumer<StatefulRedisClusterConnection<K, V>> consumer) {
+  public <T> Publisher<T> withBinaryClusterReactive(
+      final Function<StatefulRedisClusterConnection<byte[], byte[]>, Publisher<T>> function) {
+    return withConnectionReactive(binaryConnection, function);
+  }
+
+  private <K, V> void useConnection(final StatefulRedisClusterConnection<K, V> connection,
+      final Consumer<StatefulRedisClusterConnection<K, V>> consumer) {
     try {
       circuitBreaker.executeCheckedRunnable(() -> retry.executeRunnable(() -> consumer.accept(connection)));
     } catch (final Throwable t) {
@@ -123,7 +133,8 @@ public class FaultTolerantRedisCluster {
     }
   }
 
-    private <T, K, V> T withConnection(final StatefulRedisClusterConnection<K, V> connection, final Function<StatefulRedisClusterConnection<K, V>, T> function) {
+  private <T, K, V> T withConnection(final StatefulRedisClusterConnection<K, V> connection,
+      final Function<StatefulRedisClusterConnection<K, V>, T> function) {
     try {
       return circuitBreaker.executeCheckedSupplier(() -> retry.executeCallable(() -> function.apply(connection)));
     } catch (final Throwable t) {
@@ -135,6 +146,14 @@ public class FaultTolerantRedisCluster {
     }
   }
 
+  private <T, K, V> Publisher<T> withConnectionReactive(final StatefulRedisClusterConnection<K, V> connection,
+      final Function<StatefulRedisClusterConnection<K, V>, Publisher<T>> function) {
+
+    return Flux.from(function.apply(connection))
+        .transformDeferred(RetryOperator.of(retry))
+        .transformDeferred(CircuitBreakerOperator.of(circuitBreaker));
+  }
+
   public FaultTolerantPubSubConnection<String, String> createPubSubConnection() {
     final StatefulRedisClusterPubSubConnection<String, String> pubSubConnection = clusterClient.connectPubSub();
     pubSubConnections.add(pubSubConnection);

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

@@ -26,7 +26,7 @@ import software.amazon.awssdk.services.dynamodb.model.BatchWriteItemResponse;
 import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
 import software.amazon.awssdk.services.dynamodb.model.WriteRequest;
 
-public class AbstractDynamoDbStore {
+public abstract class AbstractDynamoDbStore {
 
   private final DynamoDbClient dynamoDbClient;
 

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2021 Signal Messenger, LLC
+ * Copyright 2013-2022 Signal Messenger, LLC
  * SPDX-License-Identifier: AGPL-3.0-only
  */
 
@@ -22,6 +22,7 @@ import io.micrometer.core.instrument.Metrics;
 import io.micrometer.core.instrument.Timer;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
+import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.ArrayList;
@@ -34,23 +35,32 @@ import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
+import org.reactivestreams.Publisher;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.entities.MessageProtos;
 import org.whispersystems.textsecuregcm.redis.ClusterLuaScript;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantPubSubConnection;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
+import org.whispersystems.textsecuregcm.util.Pair;
 import org.whispersystems.textsecuregcm.util.RedisClusterUtil;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Schedulers;
 
 public class MessagesCache extends RedisClusterPubSubAdapter<String, String> implements Managed {
 
   private final FaultTolerantRedisCluster readDeleteCluster;
   private final FaultTolerantPubSubConnection<String, String> pubSubConnection;
+  private final Clock clock;
 
   private final ExecutorService notificationExecutorService;
+  private final ExecutorService messageDeletionExecutorService;
 
   private final ClusterLuaScript insertScript;
   private final ClusterLuaScript removeByGuidScript;
@@ -79,22 +89,23 @@ public class MessagesCache extends RedisClusterPubSubAdapter<String, String> imp
   private static final String QUEUE_KEYSPACE_PREFIX = "__keyspace@0__:user_queue::";
   private static final String PERSISTING_KEYSPACE_PREFIX = "__keyspace@0__:user_queue_persisting::";
 
-  private static final Duration MAX_EPHEMERAL_MESSAGE_DELAY = Duration.ofSeconds(10);
+  @VisibleForTesting
+  static final Duration MAX_EPHEMERAL_MESSAGE_DELAY = Duration.ofSeconds(10);
 
-  private static final String REMOVE_TIMER_NAME = name(MessagesCache.class, "remove");
+  private static final int PAGE_SIZE = 100;
-
-  private static final String REMOVE_METHOD_TAG = "method";
-  private static final String REMOVE_METHOD_UUID = "uuid";
 
   private static final Logger logger = LoggerFactory.getLogger(MessagesCache.class);
 
   public MessagesCache(final FaultTolerantRedisCluster insertCluster, final FaultTolerantRedisCluster readDeleteCluster,
-      final ExecutorService notificationExecutorService) throws IOException {
+      final Clock clock, final ExecutorService notificationExecutorService,
+      final ExecutorService messageDeletionExecutorService) throws IOException {
 
     this.readDeleteCluster = readDeleteCluster;
     this.pubSubConnection = readDeleteCluster.createPubSubConnection();
+    this.clock = clock;
 
     this.notificationExecutorService = notificationExecutorService;
+    this.messageDeletionExecutorService = messageDeletionExecutorService;
 
     this.insertScript = ClusterLuaScript.fromResource(insertCluster, "lua/insert_item.lua", ScriptOutputType.INTEGER);
     this.removeByGuidScript = ClusterLuaScript.fromResource(readDeleteCluster, "lua/remove_item_by_guid.lua",
@@ -147,21 +158,26 @@ public class MessagesCache extends RedisClusterPubSubAdapter<String, String> imp
                 guid.toString().getBytes(StandardCharsets.UTF_8))));
   }
 
-  public Optional<MessageProtos.Envelope> remove(final UUID destinationUuid, final long destinationDevice,
+  public CompletableFuture<Optional<MessageProtos.Envelope>> remove(final UUID destinationUuid,
+      final long destinationDevice,
       final UUID messageGuid) {
-    return remove(destinationUuid, destinationDevice, List.of(messageGuid)).stream().findFirst();
+
+    return remove(destinationUuid, destinationDevice, List.of(messageGuid))
+        .thenApply(removed -> removed.isEmpty() ? Optional.empty() : Optional.of(removed.get(0)));
   }
 
   @SuppressWarnings("unchecked")
-  public List<MessageProtos.Envelope> remove(final UUID destinationUuid, final long destinationDevice,
+  public CompletableFuture<List<MessageProtos.Envelope>> remove(final UUID destinationUuid,
+      final long destinationDevice,
       final List<UUID> messageGuids) {
-    final List<byte[]> serialized = (List<byte[]>) Metrics.timer(REMOVE_TIMER_NAME, REMOVE_METHOD_TAG,
+
-        REMOVE_METHOD_UUID).record(() ->
+    return removeByGuidScript.executeBinaryAsync(List.of(getMessageQueueKey(destinationUuid, destinationDevice),
-        removeByGuidScript.executeBinary(List.of(getMessageQueueKey(destinationUuid, destinationDevice),
                 getMessageQueueMetadataKey(destinationUuid, destinationDevice),
                 getQueueIndexKey(destinationUuid, destinationDevice)),
             messageGuids.stream().map(guid -> guid.toString().getBytes(StandardCharsets.UTF_8))
-                .collect(Collectors.toList())));
+                .collect(Collectors.toList()))
+        .thenApplyAsync(result -> {
+          List<byte[]> serialized = (List<byte[]>) result;
 
           final List<MessageProtos.Envelope> removedMessages = new ArrayList<>(serialized.size());
 
@@ -174,6 +190,7 @@ public class MessagesCache extends RedisClusterPubSubAdapter<String, String> imp
           }
 
           return removedMessages;
+        }, messageDeletionExecutorService);
   }
 
   public boolean hasMessages(final UUID destinationUuid, final long destinationDevice) {
@@ -181,49 +198,109 @@ public class MessagesCache extends RedisClusterPubSubAdapter<String, String> imp
         connection -> connection.sync().zcard(getMessageQueueKey(destinationUuid, destinationDevice)) > 0);
   }
 
-  @SuppressWarnings("unchecked")
+  public Publisher<MessageProtos.Envelope> get(final UUID destinationUuid, final long destinationDevice) {
-  public List<MessageProtos.Envelope> get(final UUID destinationUuid, final long destinationDevice, final int limit) {
-    return getMessagesTimer.record(() -> {
-      final List<byte[]> queueItems = (List<byte[]>) getItemsScript.executeBinary(
-          List.of(getMessageQueueKey(destinationUuid, destinationDevice),
-              getPersistInProgressKey(destinationUuid, destinationDevice)),
-          List.of(String.valueOf(limit).getBytes(StandardCharsets.UTF_8)));
 
     final long earliestAllowableEphemeralTimestamp =
-          System.currentTimeMillis() - MAX_EPHEMERAL_MESSAGE_DELAY.toMillis();
+        clock.millis() - MAX_EPHEMERAL_MESSAGE_DELAY.toMillis();
 
-      final List<MessageProtos.Envelope> messageEntities;
+    final Flux<MessageProtos.Envelope> allMessages = getAllMessages(destinationUuid, destinationDevice)
-      final List<UUID> staleEphemeralMessageGuids = new ArrayList<>();
+        .publish()
+        // We expect exactly two subscribers to this base flux:
+        // 1. the websocket that delivers messages to clients
+        // 2. an internal process to discard stale ephemeral messages
+        // The discard subscriber will subscribe immediately, but we don’t want to do any work if the
+        // websocket never subscribes.
+        .autoConnect(2);
 
-      if (queueItems.size() % 2 == 0) {
+    final Flux<MessageProtos.Envelope> messagesToPublish = allMessages
-        messageEntities = new ArrayList<>(queueItems.size() / 2);
+        .filter(Predicate.not(envelope -> isStaleEphemeralMessage(envelope, earliestAllowableEphemeralTimestamp)));
+
+    final Flux<MessageProtos.Envelope> staleEphemeralMessages = allMessages
+        .filter(envelope -> isStaleEphemeralMessage(envelope, earliestAllowableEphemeralTimestamp));
+
+    discardStaleEphemeralMessages(destinationUuid, destinationDevice, staleEphemeralMessages);
+
+    return messagesToPublish;
+  }
+
+  private static boolean isStaleEphemeralMessage(final MessageProtos.Envelope message,
+      long earliestAllowableTimestamp) {
+    return message.hasEphemeral() && message.getEphemeral() && message.getTimestamp() < earliestAllowableTimestamp;
+  }
+
+  private void discardStaleEphemeralMessages(final UUID destinationUuid, final long destinationDevice,
+      Flux<MessageProtos.Envelope> staleEphemeralMessages) {
+    staleEphemeralMessages
+        .map(e -> UUID.fromString(e.getServerGuid()))
+        .buffer(PAGE_SIZE)
+        .subscribeOn(Schedulers.boundedElastic())
+        .subscribe(staleEphemeralMessageGuids ->
+                remove(destinationUuid, destinationDevice, staleEphemeralMessageGuids)
+                    .thenAccept(removedMessages -> staleEphemeralMessagesCounter.increment(removedMessages.size())),
+            e -> logger.warn("Could not remove stale ephemeral messages from cache", e));
+  }
+
+  @VisibleForTesting
+  Flux<MessageProtos.Envelope> getAllMessages(final UUID destinationUuid, final long destinationDevice) {
+
+    // fetch messages by page
+    return getNextMessagePage(destinationUuid, destinationDevice, -1)
+        .expand(queueItemsAndLastMessageId -> {
+          // expand() is breadth-first, so each page will be published in order
+          if (queueItemsAndLastMessageId.first().isEmpty()) {
+            return Mono.empty();
+          }
+
+          return getNextMessagePage(destinationUuid, destinationDevice, queueItemsAndLastMessageId.second());
+        })
+        .limitRate(1)
+        // we want to ensure we don’t accidentally block the Lettuce/netty i/o executors
+        .publishOn(Schedulers.boundedElastic())
+        .map(Pair::first)
+        .flatMapIterable(queueItems -> {
+          final List<MessageProtos.Envelope> envelopes = new ArrayList<>(queueItems.size() / 2);
 
           for (int i = 0; i < queueItems.size() - 1; i += 2) {
             try {
               final MessageProtos.Envelope message = MessageProtos.Envelope.parseFrom(queueItems.get(i));
-            if (message.getEphemeral() && message.getTimestamp() < earliestAllowableEphemeralTimestamp) {
-              staleEphemeralMessageGuids.add(UUID.fromString(message.getServerGuid()));
-              continue;
-            }
 
-            messageEntities.add(message);
+              envelopes.add(message);
             } catch (InvalidProtocolBufferException e) {
               logger.warn("Failed to parse envelope", e);
             }
           }
-      } else {
+
+          return envelopes;
+        });
+  }
+
+  private Flux<Pair<List<byte[]>, Long>> getNextMessagePage(final UUID destinationUuid, final long destinationDevice,
+      long messageId) {
+
+    return getItemsScript.executeBinaryReactive(
+            List.of(getMessageQueueKey(destinationUuid, destinationDevice),
+                getPersistInProgressKey(destinationUuid, destinationDevice)),
+            List.of(String.valueOf(PAGE_SIZE).getBytes(StandardCharsets.UTF_8),
+                String.valueOf(messageId).getBytes(StandardCharsets.UTF_8)))
+        .map(result -> {
+          logger.trace("Processing page: {}", messageId);
+
+          @SuppressWarnings("unchecked")
+          List<byte[]> queueItems = (List<byte[]>) result;
+
+          if (queueItems.isEmpty()) {
+            return new Pair<>(Collections.emptyList(), null);
+          }
+
+          if (queueItems.size() % 2 != 0) {
             logger.error("\"Get messages\" operation returned a list with a non-even number of elements.");
-        messageEntities = Collections.emptyList();
+            return new Pair<>(Collections.emptyList(), null);
           }
 
-      try {
+          final long lastMessageId = Long.parseLong(
-        remove(destinationUuid, destinationDevice, staleEphemeralMessageGuids);
+              new String(queueItems.get(queueItems.size() - 1), StandardCharsets.UTF_8));
-        staleEphemeralMessagesCounter.increment(staleEphemeralMessageGuids.size());
-      } catch (final Throwable e) {
-        logger.warn("Could not remove stale ephemeral messages from cache", e);
-      }
 
-      return messageEntities;
+          return new Pair<>(queueItems, lastMessageId);
         });
   }
 

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2021 Signal Messenger, LLC
+ * Copyright 2021-2022 Signal Messenger, LLC
  * SPDX-License-Identifier: AGPL-3.0-only
  */
 
@@ -17,19 +17,24 @@ import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.UUID;
-import java.util.stream.Collectors;
+import java.util.concurrent.CompletableFuture;
-import javax.annotation.Nonnull;
+import java.util.concurrent.ExecutorService;
+import java.util.function.Predicate;
+import org.reactivestreams.Publisher;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.entities.MessageProtos;
 import org.whispersystems.textsecuregcm.util.AttributeValues;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 import software.amazon.awssdk.core.SdkBytes;
+import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
-import software.amazon.awssdk.services.dynamodb.model.DeleteItemResponse;
 import software.amazon.awssdk.services.dynamodb.model.DeleteRequest;
 import software.amazon.awssdk.services.dynamodb.model.PutRequest;
 import software.amazon.awssdk.services.dynamodb.model.QueryRequest;
@@ -48,22 +53,25 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
   private static final String KEY_ENVELOPE_BYTES = "EB";
 
   private final Timer storeTimer = timer(name(getClass(), "store"));
-  private final Timer loadTimer = timer(name(getClass(), "load"));
-  private final Timer deleteByGuid = timer(name(getClass(), "delete", "guid"));
-  private final Timer deleteByKey = timer(name(getClass(), "delete", "key"));
   private final Timer deleteByAccount = timer(name(getClass(), "delete", "account"));
   private final Timer deleteByDevice = timer(name(getClass(), "delete", "device"));
 
+  private final DynamoDbAsyncClient dbAsyncClient;
   private final String tableName;
   private final Duration timeToLive;
+  private final ExecutorService messageDeletionExecutor;
 
   private static final Logger logger = LoggerFactory.getLogger(MessagesDynamoDb.class);
 
-  public MessagesDynamoDb(DynamoDbClient dynamoDb, String tableName, Duration timeToLive) {
+  public MessagesDynamoDb(DynamoDbClient dynamoDb, DynamoDbAsyncClient dynamoDbAsyncClient, String tableName,
+      Duration timeToLive, ExecutorService messageDeletionExecutor) {
     super(dynamoDb);
 
+    this.dbAsyncClient = dynamoDbAsyncClient;
     this.tableName = tableName;
     this.timeToLive = timeToLive;
+
+    this.messageDeletionExecutor = messageDeletionExecutor;
   }
 
   public void store(final List<MessageProtos.Envelope> messages, final UUID destinationAccountUuid, final long destinationDeviceId) {
@@ -95,11 +103,11 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
     executeTableWriteItemsUntilComplete(Map.of(tableName, writeItems));
   }
 
-  public List<MessageProtos.Envelope> load(final UUID destinationAccountUuid, final long destinationDeviceId, final int requestedNumberOfMessagesToFetch) {
+  public Publisher<MessageProtos.Envelope> load(final UUID destinationAccountUuid, final long destinationDeviceId,
-    return loadTimer.record(() -> {
+      final Integer limit) {
-      final int numberOfMessagesToFetch = Math.min(requestedNumberOfMessagesToFetch, RESULT_SET_CHUNK_SIZE);
+
     final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid);
-      final QueryRequest queryRequest = QueryRequest.builder()
+    final QueryRequest.Builder queryRequestBuilder = QueryRequest.builder()
         .tableName(tableName)
         .consistentRead(true)
         .keyConditionExpression("#part = :part AND begins_with ( #sort , :sortprefix )")
@@ -108,30 +116,31 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
             "#sort", KEY_SORT))
         .expressionAttributeValues(Map.of(
             ":part", partitionKey,
-              ":sortprefix", convertDestinationDeviceIdToSortKeyPrefix(destinationDeviceId)))
+            ":sortprefix", convertDestinationDeviceIdToSortKeyPrefix(destinationDeviceId)));
-          .limit(numberOfMessagesToFetch)
+
-          .build();
+    if (limit != null) {
-      List<MessageProtos.Envelope> messageEntities = new ArrayList<>(numberOfMessagesToFetch);
+      // some callers don’t take advantage of reactive streams, so we want to support limiting the fetch size. Otherwise,
-      for (Map<String, AttributeValue> message : db().queryPaginator(queryRequest).items()) {
+      // we could fetch up to 1 MB (likely >1,000 messages) and discard 90% of them
+      queryRequestBuilder.limit(Math.min(RESULT_SET_CHUNK_SIZE, limit));
+    }
+
+    final QueryRequest queryRequest = queryRequestBuilder.build();
+
+    return dbAsyncClient.queryPaginator(queryRequest).items()
+        .map(message -> {
           try {
-          messageEntities.add(convertItemToEnvelope(message));
+            return convertItemToEnvelope(message);
           } catch (final InvalidProtocolBufferException e) {
             logger.error("Failed to parse envelope", e);
+            return null;
+          }
+        })
+        .filter(Predicate.not(Objects::isNull));
   }
 
-        if (messageEntities.size() == numberOfMessagesToFetch) {
+  public CompletableFuture<Optional<MessageProtos.Envelope>> deleteMessageByDestinationAndGuid(
-          // queryPaginator() uses limit() as the page size, not as an absolute limit
+      final UUID destinationAccountUuid, final UUID messageUuid) {
-          // …but a page might be smaller than limit, because a page is capped at 1 MB
-          break;
-        }
-      }
-      return messageEntities;
-    });
-  }
 
-  public Optional<MessageProtos.Envelope> deleteMessageByDestinationAndGuid(final UUID destinationAccountUuid,
-      final UUID messageUuid) {
-    return deleteByGuid.record(() -> {
     final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid);
     final QueryRequest queryRequest = QueryRequest.builder()
         .tableName(tableName)
@@ -146,54 +155,53 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
             ":part", partitionKey,
             ":uuid", convertLocalIndexMessageUuidSortKey(messageUuid)))
         .build();
-      return deleteItemsMatchingQueryAndReturnFirstOneActuallyDeleted(partitionKey, queryRequest);
+
-    });
+    // because we are filtering on message UUID, this query should return at most one item,
+    // but it’s simpler to handle the full stream and return the “last” item
+    return Flux.from(dbAsyncClient.queryPaginator(queryRequest).items())
+        .flatMap(item -> Mono.fromCompletionStage(dbAsyncClient.deleteItem(DeleteItemRequest.builder()
+            .tableName(tableName)
+            .key(Map.of(KEY_PARTITION, partitionKey, KEY_SORT,
+                AttributeValues.fromByteArray(item.get(KEY_SORT).b().asByteArray())))
+            .returnValues(ReturnValue.ALL_OLD)
+            .build())))
+        .mapNotNull(deleteItemResponse -> {
+          try {
+            if (deleteItemResponse.attributes() != null && deleteItemResponse.attributes().containsKey(KEY_PARTITION)) {
+              return convertItemToEnvelope(deleteItemResponse.attributes());
+            }
+          } catch (final InvalidProtocolBufferException e) {
+            logger.error("Failed to parse envelope", e);
+          }
+          return null;
+        })
+        .last()
+        .toFuture()
+        .thenApply(Optional::ofNullable);
   }
 
-  public Optional<MessageProtos.Envelope> deleteMessage(final UUID destinationAccountUuid,
+  public CompletableFuture<Optional<MessageProtos.Envelope>> deleteMessage(final UUID destinationAccountUuid,
       final long destinationDeviceId, final UUID messageUuid, final long serverTimestamp) {
-    return deleteByKey.record(() -> {
+
     final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid);
     final AttributeValue sortKey = convertSortKey(destinationDeviceId, serverTimestamp, messageUuid);
     DeleteItemRequest.Builder deleteItemRequest = DeleteItemRequest.builder()
         .tableName(tableName)
         .key(Map.of(KEY_PARTITION, partitionKey, KEY_SORT, sortKey))
         .returnValues(ReturnValue.ALL_OLD);
-      final DeleteItemResponse deleteItemResponse = db().deleteItem(deleteItemRequest.build());
+
+    return dbAsyncClient.deleteItem(deleteItemRequest.build())
+        .thenApplyAsync(deleteItemResponse -> {
           if (deleteItemResponse.attributes() != null && deleteItemResponse.attributes().containsKey(KEY_PARTITION)) {
             try {
               return Optional.of(convertItemToEnvelope(deleteItemResponse.attributes()));
             } catch (final InvalidProtocolBufferException e) {
               logger.error("Failed to parse envelope", e);
-          return Optional.empty();
             }
           }
 
           return Optional.empty();
-    });
+        }, messageDeletionExecutor);
-  }
-
-  @Nonnull
-  private Optional<MessageProtos.Envelope> deleteItemsMatchingQueryAndReturnFirstOneActuallyDeleted(AttributeValue partitionKey, QueryRequest queryRequest) {
-    Optional<MessageProtos.Envelope> result = Optional.empty();
-    for (Map<String, AttributeValue> item : db().queryPaginator(queryRequest).items()) {
-      final byte[] rangeKeyValue = item.get(KEY_SORT).b().asByteArray();
-      DeleteItemRequest.Builder deleteItemRequest = DeleteItemRequest.builder()
-          .tableName(tableName)
-          .key(Map.of(KEY_PARTITION, partitionKey, KEY_SORT, AttributeValues.fromByteArray(rangeKeyValue)));
-      if (result.isEmpty()) {
-        deleteItemRequest.returnValues(ReturnValue.ALL_OLD);
-      }
-      final DeleteItemResponse deleteItemResponse = db().deleteItem(deleteItemRequest.build());
-      if (deleteItemResponse.attributes() != null && deleteItemResponse.attributes().containsKey(KEY_PARTITION)) {
-        try {
-          result = Optional.of(convertItemToEnvelope(deleteItemResponse.attributes()));
-        } catch (final InvalidProtocolBufferException e) {
-          logger.error("Failed to parse envelope", e);
-        }
-      }
-    }
-    return result;
   }
 
   public void deleteAllMessagesForAccount(final UUID destinationAccountUuid) {
@@ -248,7 +256,7 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
                 KEY_PARTITION, partitionKey,
                 KEY_SORT, item.get(KEY_SORT))).build())
             .build())
-        .collect(Collectors.toList());
+        .toList();
     executeTableWriteItemsUntilComplete(Map.of(tableName, deletes));
   }
 

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2021 Signal Messenger, LLC
+ * Copyright 2013-2022 Signal Messenger, LLC
  * SPDX-License-Identifier: AGPL-3.0-only
  */
 package org.whispersystems.textsecuregcm.storage;
@@ -9,19 +9,30 @@ import static com.codahale.metrics.MetricRegistry.name;
 import com.codahale.metrics.Meter;
 import com.codahale.metrics.MetricRegistry;
 import com.codahale.metrics.SharedMetricRegistries;
-import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
 import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.stream.Collectors;
+import javax.annotation.Nullable;
+import org.reactivestreams.Publisher;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.entities.MessageProtos.Envelope;
 import org.whispersystems.textsecuregcm.util.Constants;
 import org.whispersystems.textsecuregcm.util.Pair;
+import reactor.core.publisher.Flux;
 
 public class MessagesManager {
 
   private static final int RESULT_SET_CHUNK_SIZE = 100;
 
+  private static final Logger logger = LoggerFactory.getLogger(MessagesManager.class);
+
   private static final MetricRegistry metricRegistry = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
   private static final Meter cacheHitByGuidMeter = metricRegistry.meter(name(MessagesManager.class, "cacheHitByGuid"));
   private static final Meter cacheMissByGuidMeter = metricRegistry.meter(
@@ -55,18 +66,32 @@ public class MessagesManager {
     return messagesCache.hasMessages(destinationUuid, destinationDevice);
   }
 
-  public Pair<List<Envelope>, Boolean> getMessagesForDevice(UUID destinationUuid, long destinationDevice, final boolean cachedMessagesOnly) {
+  public Pair<List<Envelope>, Boolean> getMessagesForDevice(UUID destinationUuid, long destinationDevice,
-    List<Envelope> messageList = new ArrayList<>();
+      boolean cachedMessagesOnly) {
 
-    if (!cachedMessagesOnly) {
+    final List<Envelope> envelopes = Flux.from(
-      messageList.addAll(messagesDynamoDb.load(destinationUuid, destinationDevice, RESULT_SET_CHUNK_SIZE));
+            getMessagesForDevice(destinationUuid, destinationDevice, RESULT_SET_CHUNK_SIZE, cachedMessagesOnly))
+        .take(RESULT_SET_CHUNK_SIZE, true)
+        .collectList()
+        .blockOptional().orElse(Collections.emptyList());
+
+    return new Pair<>(envelopes, envelopes.size() >= RESULT_SET_CHUNK_SIZE);
   }
 
-    if (messageList.size() < RESULT_SET_CHUNK_SIZE) {
+  public Publisher<Envelope> getMessagesForDeviceReactive(UUID destinationUuid, long destinationDevice,
-      messageList.addAll(messagesCache.get(destinationUuid, destinationDevice, RESULT_SET_CHUNK_SIZE - messageList.size()));
+      final boolean cachedMessagesOnly) {
+
+    return getMessagesForDevice(destinationUuid, destinationDevice, null, cachedMessagesOnly);
   }
 
-    return new Pair<>(messageList, messageList.size() >= RESULT_SET_CHUNK_SIZE);
+  private Publisher<Envelope> getMessagesForDevice(UUID destinationUuid, long destinationDevice,
+      @Nullable Integer limit, final boolean cachedMessagesOnly) {
+
+    final Publisher<Envelope> dynamoPublisher =
+        cachedMessagesOnly ? Flux.empty() : messagesDynamoDb.load(destinationUuid, destinationDevice, limit);
+    final Publisher<Envelope> cachePublisher = messagesCache.get(destinationUuid, destinationDevice);
+
+    return Flux.concat(dynamoPublisher, cachePublisher);
   }
 
   public void clear(UUID destinationUuid) {
@@ -79,21 +104,25 @@ public class MessagesManager {
     messagesDynamoDb.deleteAllMessagesForDevice(destinationUuid, deviceId);
   }
 
-  public Optional<Envelope> delete(UUID destinationUuid, long destinationDeviceId, UUID guid, Long serverTimestamp) {
+  public CompletableFuture<Optional<Envelope>> delete(UUID destinationUuid, long destinationDeviceId, UUID guid,
-    Optional<Envelope> removed = messagesCache.remove(destinationUuid, destinationDeviceId, guid);
+      @Nullable Long serverTimestamp) {
+    return messagesCache.remove(destinationUuid, destinationDeviceId, guid)
+        .thenCompose(removed -> {
 
-    if (removed.isEmpty()) {
+          if (removed.isPresent()) {
-      if (serverTimestamp == null) {
-        removed = messagesDynamoDb.deleteMessageByDestinationAndGuid(destinationUuid, guid);
-      } else {
-        removed = messagesDynamoDb.deleteMessage(destinationUuid, destinationDeviceId, guid, serverTimestamp);
-      }
-      cacheMissByGuidMeter.mark();
-    } else {
             cacheHitByGuidMeter.mark();
+            return CompletableFuture.completedFuture(removed);
           }
 
-    return removed;
+          cacheMissByGuidMeter.mark();
+
+          if (serverTimestamp == null) {
+            return messagesDynamoDb.deleteMessageByDestinationAndGuid(destinationUuid, guid);
+          } else {
+            return messagesDynamoDb.deleteMessage(destinationUuid, destinationDeviceId, guid, serverTimestamp);
+          }
+
+        });
   }
 
   /**
@@ -112,10 +141,15 @@ public class MessagesManager {
 
     final List<UUID> messageGuids = messages.stream().map(message -> UUID.fromString(message.getServerGuid()))
         .collect(Collectors.toList());
-    int messagesRemovedFromCache = messagesCache.remove(destinationUuid, destinationDeviceId, messageGuids).size();
+    int messagesRemovedFromCache = 0;
-
+    try {
+      messagesRemovedFromCache = messagesCache.remove(destinationUuid, destinationDeviceId, messageGuids)
+          .get(30, TimeUnit.SECONDS).size();
       persistMessageMeter.mark(nonEphemeralMessages.size());
 
+    } catch (InterruptedException | ExecutionException | TimeoutException e) {
+      logger.warn("Failed to remove messages from cache", e);
+    }
     return messagesRemovedFromCache;
   }
 
@@ -129,4 +163,5 @@ public class MessagesManager {
   public void removeMessageAvailabilityListener(final MessageAvailabilityListener listener) {
     messagesCache.removeMessageAvailabilityListener(listener);
   }
+
 }

service/src/main/java/org/whispersystems/textsecuregcm/websocket/AuthenticatedConnectListener.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2021 Signal Messenger, LLC
+ * Copyright 2013-2022 Signal Messenger, LLC
  * SPDX-License-Identifier: AGPL-3.0-only
  */
 
@@ -11,14 +11,19 @@ import com.codahale.metrics.Counter;
 import com.codahale.metrics.MetricRegistry;
 import com.codahale.metrics.SharedMetricRegistries;
 import com.codahale.metrics.Timer;
+import io.micrometer.core.instrument.Metrics;
+import io.micrometer.core.instrument.Tags;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import org.eclipse.jetty.websocket.api.UpgradeResponse;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedAccount;
+import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
+import org.whispersystems.textsecuregcm.metrics.MetricsUtil;
 import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
 import org.whispersystems.textsecuregcm.push.NotPushRegisteredException;
 import org.whispersystems.textsecuregcm.push.PushNotificationManager;
@@ -33,12 +38,20 @@ import org.whispersystems.websocket.setup.WebSocketConnectListener;
 public class AuthenticatedConnectListener implements WebSocketConnectListener {
 
   private static final MetricRegistry metricRegistry = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
-  private static final Timer          durationTimer                = metricRegistry.timer(name(WebSocketConnection.class, "connected_duration"                 ));
+  private static final Timer durationTimer = metricRegistry.timer(
-  private static final Timer          unauthenticatedDurationTimer = metricRegistry.timer(name(WebSocketConnection.class, "unauthenticated_connection_duration"));
+      name(WebSocketConnection.class, "connected_duration"));
-  private static final Counter        openWebsocketCounter         = metricRegistry.counter(name(WebSocketConnection.class, "open_websockets"));
+  private static final Timer unauthenticatedDurationTimer = metricRegistry.timer(
+      name(WebSocketConnection.class, "unauthenticated_connection_duration"));
+  private static final Counter openWebsocketCounter = metricRegistry.counter(
+      name(WebSocketConnection.class, "open_websockets"));
+
+  private static final String OPEN_WEBSOCKET_COUNTER_NAME = MetricsUtil.name(WebSocketConnection.class,
+      "openWebsockets");
 
   private static final long RENEW_PRESENCE_INTERVAL_MINUTES = 5;
 
+  private static final String REACTIVE_MESSAGE_QUEUE_EXPERIMENT_NAME = "reactive_message_queue_v1";
+
   private static final Logger log = LoggerFactory.getLogger(AuthenticatedConnectListener.class);
 
   private final ReceiptSender receiptSender;
@@ -46,18 +59,26 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
   private final PushNotificationManager pushNotificationManager;
   private final ClientPresenceManager clientPresenceManager;
   private final ScheduledExecutorService scheduledExecutorService;
+  private final ExperimentEnrollmentManager experimentEnrollmentManager;
+
+  private final AtomicInteger openReactiveWebSockets = new AtomicInteger(0);
+  private final AtomicInteger openStandardWebSockets = new AtomicInteger(0);
 
   public AuthenticatedConnectListener(ReceiptSender receiptSender,
       MessagesManager messagesManager,
       PushNotificationManager pushNotificationManager,
       ClientPresenceManager clientPresenceManager,
-      ScheduledExecutorService scheduledExecutorService)
+      ScheduledExecutorService scheduledExecutorService,
-  {
+      ExperimentEnrollmentManager experimentEnrollmentManager) {
     this.receiptSender = receiptSender;
     this.messagesManager = messagesManager;
     this.pushNotificationManager = pushNotificationManager;
     this.clientPresenceManager = clientPresenceManager;
     this.scheduledExecutorService = scheduledExecutorService;
+    this.experimentEnrollmentManager = experimentEnrollmentManager;
+
+    Metrics.gauge(OPEN_WEBSOCKET_COUNTER_NAME, Tags.of("reactive", String.valueOf(true)), openReactiveWebSockets);
+    Metrics.gauge(OPEN_WEBSOCKET_COUNTER_NAME, Tags.of("reactive", String.valueOf(false)), openStandardWebSockets);
   }
 
   @Override
@@ -66,20 +87,34 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
       final AuthenticatedAccount auth = context.getAuthenticated(AuthenticatedAccount.class);
       final Device device = auth.getAuthenticatedDevice();
       final Timer.Context timer = durationTimer.time();
+      final boolean enrolledInReactiveMessageQueue = experimentEnrollmentManager.isEnrolled(
+          auth.getAccount().getUuid(),
+          REACTIVE_MESSAGE_QUEUE_EXPERIMENT_NAME);
       final WebSocketConnection connection = new WebSocketConnection(receiptSender,
           messagesManager, auth, device,
           context.getClient(),
-          scheduledExecutorService);
+          scheduledExecutorService,
+          enrolledInReactiveMessageQueue);
 
       openWebsocketCounter.inc();
+      if (enrolledInReactiveMessageQueue) {
+        openReactiveWebSockets.incrementAndGet();
+      } else {
+        openStandardWebSockets.incrementAndGet();
+      }
+
       pushNotificationManager.handleMessagesRetrieved(auth.getAccount(), device, context.getClient().getUserAgent());
 
       final AtomicReference<ScheduledFuture<?>> renewPresenceFutureReference = new AtomicReference<>();
 
-      context.addListener(new WebSocketSessionContext.WebSocketEventListener() {
+      context.addListener((closingContext, statusCode, reason) -> {
-        @Override
-        public void onWebSocketClose(WebSocketSessionContext context, int statusCode, String reason) {
         openWebsocketCounter.dec();
+        if (enrolledInReactiveMessageQueue) {
+          openReactiveWebSockets.decrementAndGet();
+        } else {
+          openStandardWebSockets.decrementAndGet();
+        }
+
         timer.stop();
 
         final ScheduledFuture<?> renewPresenceFuture = renewPresenceFutureReference.get();
@@ -102,7 +137,6 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
             }
           }
         });
-        }
       });
 
       try {

service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java

@@ -1,12 +1,11 @@
 /*
- * Copyright 2013-2021 Signal Messenger, LLC
+ * Copyright 2013-2022 Signal Messenger, LLC
  * SPDX-License-Identifier: AGPL-3.0-only
  */
 
 package org.whispersystems.textsecuregcm.websocket;
 
 import static com.codahale.metrics.MetricRegistry.name;
-import static org.whispersystems.textsecuregcm.entities.MessageProtos.Envelope;
 
 import com.codahale.metrics.Histogram;
 import com.codahale.metrics.Meter;
@@ -34,11 +33,13 @@ import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.atomic.LongAdder;
 import javax.ws.rs.WebApplicationException;
 import org.apache.commons.lang3.StringUtils;
+import org.reactivestreams.Publisher;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedAccount;
 import org.whispersystems.textsecuregcm.controllers.MessageController;
 import org.whispersystems.textsecuregcm.controllers.NoSuchUserException;
+import org.whispersystems.textsecuregcm.entities.MessageProtos.Envelope;
 import org.whispersystems.textsecuregcm.metrics.MessageMetrics;
 import org.whispersystems.textsecuregcm.metrics.UserAgentTagUtil;
 import org.whispersystems.textsecuregcm.push.DisplacedPresenceListener;
@@ -49,13 +50,14 @@ import org.whispersystems.textsecuregcm.storage.MessagesManager;
 import org.whispersystems.textsecuregcm.util.Constants;
 import org.whispersystems.textsecuregcm.util.Pair;
 import org.whispersystems.textsecuregcm.util.TimestampHeaderUtil;
-import org.whispersystems.textsecuregcm.util.ua.ClientPlatform;
-import org.whispersystems.textsecuregcm.util.ua.UnrecognizedUserAgentException;
-import org.whispersystems.textsecuregcm.util.ua.UserAgentUtil;
 import org.whispersystems.websocket.WebSocketClient;
 import org.whispersystems.websocket.messages.WebSocketResponseMessage;
+import reactor.core.Disposable;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Scheduler;
+import reactor.core.scheduler.Schedulers;
 
-@SuppressWarnings("OptionalUsedAsFieldOrParameterType")
 public class WebSocketConnection implements MessageAvailabilityListener, DisplacedPresenceListener {
 
   private static final MetricRegistry metricRegistry = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
@@ -70,8 +72,6 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
       name(WebSocketConnection.class, "messagesPersisted"));
   private static final Meter bytesSentMeter = metricRegistry.meter(name(WebSocketConnection.class, "bytes_sent"));
   private static final Meter sendFailuresMeter = metricRegistry.meter(name(WebSocketConnection.class, "send_failures"));
-  private static final Meter discardedMessagesMeter = metricRegistry.meter(
-      name(WebSocketConnection.class, "discardedMessages"));
 
   private static final String INITIAL_QUEUE_LENGTH_DISTRIBUTION_NAME = name(WebSocketConnection.class,
       "initialQueueLength");
@@ -85,11 +85,12 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
       "messageAvailableAfterClientClosed");
   private static final String STATUS_CODE_TAG = "status";
   private static final String STATUS_MESSAGE_TAG = "message";
+  private static final String REACTIVE_TAG = "reactive";
 
   private static final long SLOW_DRAIN_THRESHOLD = 10_000;
 
   @VisibleForTesting
-  static final int MAX_DESKTOP_MESSAGE_SIZE = 1024 * 1024;
+  static final int MESSAGE_PUBLISHER_LIMIT_RATE = 100;
 
   @VisibleForTesting
   static final int MAX_CONSECUTIVE_RETRIES = 5;
@@ -111,8 +112,6 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
 
   private final ScheduledExecutorService scheduledExecutorService;
 
-  private final boolean isDesktopClient;
-
   private final Semaphore processStoredMessagesSemaphore = new Semaphore(1);
   private final AtomicReference<StoredMessageState> storedMessageState = new AtomicReference<>(
       StoredMessageState.PERSISTED_NEW_MESSAGES_AVAILABLE);
@@ -121,8 +120,11 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
   private final AtomicLong queueDrainStartTime = new AtomicLong();
   private final AtomicInteger consecutiveRetries = new AtomicInteger();
   private final AtomicReference<ScheduledFuture<?>> retryFuture = new AtomicReference<>();
+  private final AtomicReference<Disposable> messageSubscription = new AtomicReference<>();
 
   private final Random random = new Random();
+  private final boolean useReactive;
+  private Scheduler reactiveScheduler;
 
   private enum StoredMessageState {
     EMPTY,
@@ -135,7 +137,28 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
       AuthenticatedAccount auth,
       Device device,
       WebSocketClient client,
-      ScheduledExecutorService scheduledExecutorService) {
+      ScheduledExecutorService scheduledExecutorService,
+      boolean useReactive) {
+
+    this(receiptSender,
+        messagesManager,
+        auth,
+        device,
+        client,
+        scheduledExecutorService,
+        useReactive,
+        Schedulers.boundedElastic());
+  }
+
+  @VisibleForTesting
+  WebSocketConnection(ReceiptSender receiptSender,
+      MessagesManager messagesManager,
+      AuthenticatedAccount auth,
+      Device device,
+      WebSocketClient client,
+      ScheduledExecutorService scheduledExecutorService,
+      boolean useReactive,
+      Scheduler reactiveScheduler) {
 
     this(receiptSender,
         messagesManager,
@@ -143,7 +166,9 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
         device,
         client,
         DEFAULT_SEND_FUTURES_TIMEOUT_MILLIS,
-        scheduledExecutorService);
+        scheduledExecutorService,
+        useReactive,
+        reactiveScheduler);
   }
 
   @VisibleForTesting
@@ -153,7 +178,9 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
       Device device,
       WebSocketClient client,
       int sendFuturesTimeoutMillis,
-      ScheduledExecutorService scheduledExecutorService) {
+      ScheduledExecutorService scheduledExecutorService,
+      boolean useReactive,
+      Scheduler reactiveScheduler) {
 
     this.receiptSender = receiptSender;
     this.messagesManager = messagesManager;
@@ -162,16 +189,8 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
     this.client = client;
     this.sendFuturesTimeoutMillis = sendFuturesTimeoutMillis;
     this.scheduledExecutorService = scheduledExecutorService;
-
+    this.useReactive = useReactive;
-    Optional<ClientPlatform> maybePlatform;
+    this.reactiveScheduler = reactiveScheduler;
-
-    try {
-      maybePlatform = Optional.of(UserAgentUtil.parseUserAgentString(client.getUserAgent()).getPlatform());
-    } catch (final UnrecognizedUserAgentException e) {
-      maybePlatform = Optional.empty();
-    }
-
-    this.isDesktopClient = maybePlatform.map(platform -> platform == ClientPlatform.DESKTOP).orElse(false);
   }
 
   public void start() {
@@ -186,10 +205,15 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
       future.cancel(false);
     }
 
+    final Disposable subscription = messageSubscription.get();
+    if (subscription != null) {
+      subscription.dispose();
+    }
+
     client.close(1000, "OK");
   }
 
-  private CompletableFuture<WebSocketResponseMessage> sendMessage(final Envelope message, final Optional<StoredMessageInfo> storedMessageInfo) {
+  private CompletableFuture<?> sendMessage(final Envelope message, StoredMessageInfo storedMessageInfo) {
     // clear ephemeral field from the envelope
     final Optional<byte[]> body = Optional.ofNullable(message.toBuilder().clearEphemeral().build().toByteArray());
 
@@ -199,12 +223,18 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
     MessageMetrics.measureAccountEnvelopeUuidMismatches(auth.getAccount(), message);
 
     // X-Signal-Key: false must be sent until Android stops assuming it missing means true
-    return client.sendRequest("PUT", "/api/v1/message", List.of("X-Signal-Key: false", TimestampHeaderUtil.getTimestampHeader()), body).whenComplete((response, throwable) -> {
+    return client.sendRequest("PUT", "/api/v1/message",
-          if (throwable == null) {
+            List.of("X-Signal-Key: false", TimestampHeaderUtil.getTimestampHeader()), body)
-            if (isSuccessResponse(response)) {
+        .whenComplete((ignored, throwable) -> {
-              if (storedMessageInfo.isPresent()) {
+          if (throwable != null) {
-                messagesManager.delete(auth.getAccount().getUuid(), device.getId(), storedMessageInfo.get().getGuid(), storedMessageInfo.get().getServerTimestamp());
+            sendFailuresMeter.mark();
           }
+        }).thenCompose(response -> {
+          final CompletableFuture<?> result;
+          if (isSuccessResponse(response)) {
+
+            result = messagesManager.delete(auth.getAccount().getUuid(), device.getId(),
+                storedMessageInfo.guid(), storedMessageInfo.serverTimestamp());
 
             if (message.getType() != Envelope.Type.SERVER_DELIVERY_RECEIPT) {
               recordMessageDeliveryDuration(message.getTimestamp(), device);
@@ -212,8 +242,11 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
             }
           } else {
               final List<Tag> tags = new ArrayList<>(
-                  List.of(Tag.of(STATUS_CODE_TAG, String.valueOf(response.getStatus())),
+                  List.of(
-                      UserAgentTagUtil.getPlatformTag(client.getUserAgent())));
+                      Tag.of(STATUS_CODE_TAG, String.valueOf(response.getStatus())),
+                      UserAgentTagUtil.getPlatformTag(client.getUserAgent()),
+                      Tag.of(REACTIVE_TAG, String.valueOf(useReactive))
+                  ));
 
               // TODO Remove this once we've identified the cause of message rejections from desktop clients
               if (StringUtils.isNotBlank(response.getMessage())) {
@@ -221,10 +254,11 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
               }
 
               Metrics.counter(NON_SUCCESS_RESPONSE_COUNTER_NAME, tags).increment();
+
+              result = CompletableFuture.completedFuture(null);
             }
-      } else {
+
-        sendFailuresMeter.mark();
+          return result;
-      }
         });
   }
 
@@ -260,18 +294,37 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
 
   @VisibleForTesting
   void processStoredMessages() {
+    if (useReactive) {
+      processStoredMessages_reactive();
+    } else {
+      processStoredMessage_paged();
+    }
+  }
+
+  private void processStoredMessage_paged() {
+    assert !useReactive;
+
     if (processStoredMessagesSemaphore.tryAcquire()) {
       final StoredMessageState state = storedMessageState.getAndSet(StoredMessageState.EMPTY);
-      final CompletableFuture<Void> queueClearedFuture = new CompletableFuture<>();
+      final CompletableFuture<Void> queueCleared = new CompletableFuture<>();
 
-      sendNextMessagePage(state != StoredMessageState.PERSISTED_NEW_MESSAGES_AVAILABLE, queueClearedFuture);
+      sendNextMessagePage(state != StoredMessageState.PERSISTED_NEW_MESSAGES_AVAILABLE, queueCleared);
 
-      queueClearedFuture.whenComplete((v, cause) -> {
+      setQueueClearedHandler(state, queueCleared);
+    }
+  }
+
+  private void setQueueClearedHandler(final StoredMessageState state, final CompletableFuture<Void> queueCleared) {
+
+    queueCleared.whenComplete((v, cause) -> {
       if (cause == null) {
         consecutiveRetries.set(0);
 
         if (sentInitialQueueEmptyMessage.compareAndSet(false, true)) {
-            final List<Tag> tags = List.of(UserAgentTagUtil.getPlatformTag(client.getUserAgent()));
+          final List<Tag> tags = List.of(
+              UserAgentTagUtil.getPlatformTag(client.getUserAgent()),
+              Tag.of(REACTIVE_TAG, String.valueOf(useReactive))
+          );
           final long drainDuration = System.currentTimeMillis() - queueDrainStartTime.get();
 
           Metrics.summary(INITIAL_QUEUE_LENGTH_DISTRIBUTION_NAME, tags).record(sentMessageCounter.sum());
@@ -316,9 +369,21 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
       }
     });
   }
+
+  private void processStoredMessages_reactive() {
+    assert useReactive;
+
+    if (processStoredMessagesSemaphore.tryAcquire()) {
+      final StoredMessageState state = storedMessageState.getAndSet(StoredMessageState.EMPTY);
+      final CompletableFuture<Void> queueCleared = new CompletableFuture<>();
+
+      sendMessagesReactive(state != StoredMessageState.PERSISTED_NEW_MESSAGES_AVAILABLE, queueCleared);
+
+      setQueueClearedHandler(state, queueCleared);
+    }
   }
 
-  private void sendNextMessagePage(final boolean cachedMessagesOnly, final CompletableFuture<Void> queueClearedFuture) {
+  private void sendNextMessagePage(final boolean cachedMessagesOnly, final CompletableFuture<Void> queueCleared) {
     try {
       final Pair<List<Envelope>, Boolean> messagesAndHasMore = messagesManager.getMessagesForDevice(
           auth.getAccount().getUuid(), device.getId(), cachedMessagesOnly);
@@ -330,25 +395,7 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
 
       for (int i = 0; i < messages.size(); i++) {
         final Envelope envelope = messages.get(i);
-        final UUID messageGuid = UUID.fromString(envelope.getServerGuid());
+        sendFutures[i] = sendMessage(envelope);
-
-        final boolean discard;
-        if (isDesktopClient && envelope.getSerializedSize() > MAX_DESKTOP_MESSAGE_SIZE) {
-          discard = true;
-        } else if (envelope.getStory() && !client.shouldDeliverStories()) {
-          discard = true;
-        } else {
-          discard = false;
-        }
-
-        if (discard) {
-          messagesManager.delete(auth.getAccount().getUuid(), device.getId(), messageGuid, envelope.getServerTimestamp());
-          discardedMessagesMeter.mark();
-
-          sendFutures[i] = CompletableFuture.completedFuture(null);
-        } else {
-          sendFutures[i] = sendMessage(envelope, Optional.of(new StoredMessageInfo(messageGuid, envelope.getServerTimestamp())));
-        }
       }
 
       // Set a large, non-zero timeout, to prevent any failure to acknowledge receipt from blocking indefinitely
@@ -357,16 +404,45 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
           .whenComplete((v, cause) -> {
             if (cause == null) {
               if (hasMore) {
-                sendNextMessagePage(cachedMessagesOnly, queueClearedFuture);
+                sendNextMessagePage(cachedMessagesOnly, queueCleared);
               } else {
-                queueClearedFuture.complete(null);
+                queueCleared.complete(null);
               }
             } else {
-              queueClearedFuture.completeExceptionally(cause);
+              queueCleared.completeExceptionally(cause);
             }
           });
     } catch (final Exception e) {
-      queueClearedFuture.completeExceptionally(e);
+      queueCleared.completeExceptionally(e);
+    }
+  }
+
+  private void sendMessagesReactive(final boolean cachedMessagesOnly, final CompletableFuture<Void> queueCleared) {
+
+    final Publisher<Envelope> messages =
+        messagesManager.getMessagesForDeviceReactive(auth.getAccount().getUuid(), device.getId(), cachedMessagesOnly);
+
+    final Disposable subscription = Flux.from(messages)
+        .limitRate(MESSAGE_PUBLISHER_LIMIT_RATE)
+        .flatMapSequential(envelope ->
+            Mono.fromFuture(sendMessage(envelope).orTimeout(sendFuturesTimeoutMillis, TimeUnit.MILLISECONDS)))
+        .doOnError(queueCleared::completeExceptionally)
+        .doOnComplete(() -> queueCleared.complete(null))
+        .subscribeOn(reactiveScheduler)
+        .subscribe();
+
+    messageSubscription.set(subscription);
+  }
+
+  private CompletableFuture<?> sendMessage(Envelope envelope) {
+    final UUID messageGuid = UUID.fromString(envelope.getServerGuid());
+
+    if (envelope.getStory() && !client.shouldDeliverStories()) {
+      messagesManager.delete(auth.getAccount().getUuid(), device.getId(), messageGuid, envelope.getServerTimestamp());
+
+      return CompletableFuture.completedFuture(null);
+    } else {
+      return sendMessage(envelope, new StoredMessageInfo(messageGuid, envelope.getServerTimestamp()));
     }
   }
 
@@ -381,6 +457,7 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
     messageAvailableMeter.mark();
 
     storedMessageState.compareAndSet(StoredMessageState.EMPTY, StoredMessageState.CACHED_NEW_MESSAGES_AVAILABLE);
+
     processStoredMessages();
 
     return true;
@@ -396,6 +473,7 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
     messagesPersistedMeter.mark();
 
     storedMessageState.set(StoredMessageState.PERSISTED_NEW_MESSAGES_AVAILABLE);
+
     processStoredMessages();
 
     return true;
@@ -405,7 +483,8 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
   public void handleDisplacement(final boolean connectedElsewhere) {
     final Tags tags = Tags.of(
         UserAgentTagUtil.getPlatformTag(client.getUserAgent()),
-        Tag.of("connectedElsewhere", String.valueOf(connectedElsewhere)));
+        Tag.of("connectedElsewhere", String.valueOf(connectedElsewhere)),
+        Tag.of(REACTIVE_TAG, String.valueOf(useReactive)));
 
     Metrics.counter(DISPLACEMENT_COUNTER_NAME, tags).increment();
 
@@ -429,21 +508,7 @@ public class WebSocketConnection implements MessageAvailabilityListener, Displac
     }
   }
 
-  private static class StoredMessageInfo {
+  private record StoredMessageInfo(UUID guid, long serverTimestamp) {
-    private final UUID guid;
-    private final long serverTimestamp;
 
-    public StoredMessageInfo(UUID guid, long serverTimestamp) {
-      this.guid = guid;
-      this.serverTimestamp = serverTimestamp;
-    }
-
-    public UUID getGuid() {
-      return guid;
-    }
-
-    public long getServerTimestamp() {
-      return serverTimestamp;
-    }
   }
 }

service/src/main/java/org/whispersystems/textsecuregcm/workers/AssignUsernameCommand.java

@@ -25,7 +25,6 @@ import net.sourceforge.argparse4j.inf.Subparser;
 import org.whispersystems.textsecuregcm.WhisperServerConfiguration;
 import org.whispersystems.textsecuregcm.auth.ExternalServiceCredentialGenerator;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
-import org.whispersystems.textsecuregcm.push.PushLatencyManager;
 import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
 import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
@@ -45,9 +44,9 @@ import org.whispersystems.textsecuregcm.storage.MessagesManager;
 import org.whispersystems.textsecuregcm.storage.PhoneNumberIdentifiers;
 import org.whispersystems.textsecuregcm.storage.Profiles;
 import org.whispersystems.textsecuregcm.storage.ProfilesManager;
+import org.whispersystems.textsecuregcm.storage.ProhibitedUsernames;
 import org.whispersystems.textsecuregcm.storage.ReportMessageDynamoDb;
 import org.whispersystems.textsecuregcm.storage.ReportMessageManager;
-import org.whispersystems.textsecuregcm.storage.ProhibitedUsernames;
 import org.whispersystems.textsecuregcm.storage.StoredVerificationCodeManager;
 import org.whispersystems.textsecuregcm.storage.UsernameNotAvailableException;
 import org.whispersystems.textsecuregcm.storage.VerificationCodeStore;
@@ -97,6 +96,8 @@ public class AssignUsernameCommand extends EnvironmentCommand<WhisperServerConfi
 
     ExecutorService keyspaceNotificationDispatchExecutor = environment.lifecycle()
         .executorService(name(getClass(), "keyspaceNotification-%d")).maxThreads(4).build();
+    ExecutorService messageDeletionExecutor = environment.lifecycle()
+        .executorService(name(getClass(), "messageDeletion-%d")).maxThreads(4).build();
     ExecutorService backupServiceExecutor = environment.lifecycle()
         .executorService(name(getClass(), "backupService-%d")).maxThreads(8).minThreads(1).build();
     ExecutorService storageServiceExecutor = environment.lifecycle()
@@ -156,15 +157,14 @@ public class AssignUsernameCommand extends EnvironmentCommand<WhisperServerConfi
         configuration.getDynamoDbTables().getReservedUsernames().getTableName());
     Keys keys = new Keys(dynamoDbClient,
         configuration.getDynamoDbTables().getKeys().getTableName());
-    MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbClient,
+    MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbClient, dynamoDbAsyncClient,
         configuration.getDynamoDbTables().getMessages().getTableName(),
-        configuration.getDynamoDbTables().getMessages().getExpiration());
+        configuration.getDynamoDbTables().getMessages().getExpiration(),
+        messageDeletionExecutor);
     FaultTolerantRedisCluster messageInsertCacheCluster = new FaultTolerantRedisCluster("message_insert_cluster",
         configuration.getMessageCacheConfiguration().getRedisClusterConfiguration(), redisClusterClientResources);
     FaultTolerantRedisCluster messageReadDeleteCluster = new FaultTolerantRedisCluster("message_read_delete_cluster",
         configuration.getMessageCacheConfiguration().getRedisClusterConfiguration(), redisClusterClientResources);
-    FaultTolerantRedisCluster metricsCluster = new FaultTolerantRedisCluster("metrics_cluster",
-        configuration.getMetricsClusterConfiguration(), redisClusterClientResources);
     FaultTolerantRedisCluster clientPresenceCluster = new FaultTolerantRedisCluster("client_presence_cluster",
         configuration.getClientPresenceClusterConfiguration(), redisClusterClientResources);
     FaultTolerantRedisCluster rateLimitersCluster = new FaultTolerantRedisCluster("rate_limiters",
@@ -176,8 +176,7 @@ public class AssignUsernameCommand extends EnvironmentCommand<WhisperServerConfi
     ClientPresenceManager clientPresenceManager = new ClientPresenceManager(clientPresenceCluster,
         Executors.newSingleThreadScheduledExecutor(), keyspaceNotificationDispatchExecutor);
     MessagesCache messagesCache = new MessagesCache(messageInsertCacheCluster, messageReadDeleteCluster,
-        keyspaceNotificationDispatchExecutor);
+        Clock.systemUTC(), keyspaceNotificationDispatchExecutor, messageDeletionExecutor);
-    PushLatencyManager pushLatencyManager = new PushLatencyManager(metricsCluster, dynamicConfigurationManager);
     DirectoryQueue directoryQueue = new DirectoryQueue(
         configuration.getDirectoryConfiguration().getSqsConfiguration());
     ProfilesManager profilesManager = new ProfilesManager(profiles, cacheCluster);

service/src/main/java/org/whispersystems/textsecuregcm/workers/DeleteUserCommand.java

@@ -27,7 +27,6 @@ import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.WhisperServerConfiguration;
 import org.whispersystems.textsecuregcm.auth.ExternalServiceCredentialGenerator;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
-import org.whispersystems.textsecuregcm.push.PushLatencyManager;
 import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
 import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
@@ -48,9 +47,9 @@ import org.whispersystems.textsecuregcm.storage.MessagesManager;
 import org.whispersystems.textsecuregcm.storage.PhoneNumberIdentifiers;
 import org.whispersystems.textsecuregcm.storage.Profiles;
 import org.whispersystems.textsecuregcm.storage.ProfilesManager;
+import org.whispersystems.textsecuregcm.storage.ProhibitedUsernames;
 import org.whispersystems.textsecuregcm.storage.ReportMessageDynamoDb;
 import org.whispersystems.textsecuregcm.storage.ReportMessageManager;
-import org.whispersystems.textsecuregcm.storage.ProhibitedUsernames;
 import org.whispersystems.textsecuregcm.storage.StoredVerificationCodeManager;
 import org.whispersystems.textsecuregcm.storage.VerificationCodeStore;
 import org.whispersystems.textsecuregcm.util.DynamoDbFromConfig;
@@ -99,6 +98,8 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
 
       ExecutorService keyspaceNotificationDispatchExecutor = environment.lifecycle()
           .executorService(name(getClass(), "keyspaceNotification-%d")).maxThreads(4).build();
+      ExecutorService messageDeletionExecutor = environment.lifecycle()
+          .executorService(name(getClass(), "messageDeletion-%d")).maxThreads(4).build();
       ExecutorService backupServiceExecutor = environment.lifecycle()
           .executorService(name(getClass(), "backupService-%d")).maxThreads(8).minThreads(1).build();
       ExecutorService storageServiceExecutor = environment.lifecycle()
@@ -158,15 +159,14 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
           configuration.getDynamoDbTables().getReservedUsernames().getTableName());
       Keys keys = new Keys(dynamoDbClient,
           configuration.getDynamoDbTables().getKeys().getTableName());
-      MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbClient,
+      MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbClient, dynamoDbAsyncClient,
           configuration.getDynamoDbTables().getMessages().getTableName(),
-          configuration.getDynamoDbTables().getMessages().getExpiration());
+          configuration.getDynamoDbTables().getMessages().getExpiration(),
+          messageDeletionExecutor);
       FaultTolerantRedisCluster messageInsertCacheCluster = new FaultTolerantRedisCluster("message_insert_cluster",
           configuration.getMessageCacheConfiguration().getRedisClusterConfiguration(), redisClusterClientResources);
       FaultTolerantRedisCluster messageReadDeleteCluster = new FaultTolerantRedisCluster("message_read_delete_cluster",
           configuration.getMessageCacheConfiguration().getRedisClusterConfiguration(), redisClusterClientResources);
-      FaultTolerantRedisCluster metricsCluster = new FaultTolerantRedisCluster("metrics_cluster",
-          configuration.getMetricsClusterConfiguration(), redisClusterClientResources);
       FaultTolerantRedisCluster clientPresenceCluster    = new FaultTolerantRedisCluster("client_presence_cluster",
           configuration.getClientPresenceClusterConfiguration(), redisClusterClientResources);
       FaultTolerantRedisCluster rateLimitersCluster = new FaultTolerantRedisCluster("rate_limiters",
@@ -178,8 +178,7 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
       ClientPresenceManager clientPresenceManager = new ClientPresenceManager(clientPresenceCluster,
           Executors.newSingleThreadScheduledExecutor(), keyspaceNotificationDispatchExecutor);
       MessagesCache messagesCache = new MessagesCache(messageInsertCacheCluster, messageReadDeleteCluster,
-          keyspaceNotificationDispatchExecutor);
+          Clock.systemUTC(), keyspaceNotificationDispatchExecutor, messageDeletionExecutor);
-      PushLatencyManager pushLatencyManager = new PushLatencyManager(metricsCluster, dynamicConfigurationManager);
       DirectoryQueue directoryQueue = new DirectoryQueue(
           configuration.getDirectoryConfiguration().getSqsConfiguration());
       ProfilesManager profilesManager = new ProfilesManager(profiles, cacheCluster);

service/src/main/java/org/whispersystems/textsecuregcm/workers/SetUserDiscoverabilityCommand.java

@@ -26,7 +26,6 @@ import net.sourceforge.argparse4j.inf.Subparser;
 import org.whispersystems.textsecuregcm.WhisperServerConfiguration;
 import org.whispersystems.textsecuregcm.auth.ExternalServiceCredentialGenerator;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
-import org.whispersystems.textsecuregcm.push.PushLatencyManager;
 import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
 import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
@@ -46,9 +45,9 @@ import org.whispersystems.textsecuregcm.storage.MessagesManager;
 import org.whispersystems.textsecuregcm.storage.PhoneNumberIdentifiers;
 import org.whispersystems.textsecuregcm.storage.Profiles;
 import org.whispersystems.textsecuregcm.storage.ProfilesManager;
+import org.whispersystems.textsecuregcm.storage.ProhibitedUsernames;
 import org.whispersystems.textsecuregcm.storage.ReportMessageDynamoDb;
 import org.whispersystems.textsecuregcm.storage.ReportMessageManager;
-import org.whispersystems.textsecuregcm.storage.ProhibitedUsernames;
 import org.whispersystems.textsecuregcm.storage.StoredVerificationCodeManager;
 import org.whispersystems.textsecuregcm.storage.VerificationCodeStore;
 import org.whispersystems.textsecuregcm.util.DynamoDbFromConfig;
@@ -102,6 +101,8 @@ public class SetUserDiscoverabilityCommand extends EnvironmentCommand<WhisperSer
 
       ExecutorService keyspaceNotificationDispatchExecutor = environment.lifecycle()
           .executorService(name(getClass(), "keyspaceNotification-%d")).maxThreads(4).build();
+      ExecutorService messageDeletionExecutor = environment.lifecycle()
+          .executorService(name(getClass(), "messageDeletion-%d")).maxThreads(4).build();
       ExecutorService backupServiceExecutor = environment.lifecycle()
           .executorService(name(getClass(), "backupService-%d")).maxThreads(8).minThreads(1).build();
       ExecutorService storageServiceExecutor = environment.lifecycle()
@@ -161,15 +162,14 @@ public class SetUserDiscoverabilityCommand extends EnvironmentCommand<WhisperSer
           configuration.getDynamoDbTables().getReservedUsernames().getTableName());
       Keys keys = new Keys(dynamoDbClient,
           configuration.getDynamoDbTables().getKeys().getTableName());
-      MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbClient,
+      MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbClient, dynamoDbAsyncClient,
           configuration.getDynamoDbTables().getMessages().getTableName(),
-          configuration.getDynamoDbTables().getMessages().getExpiration());
+          configuration.getDynamoDbTables().getMessages().getExpiration(),
+          messageDeletionExecutor);
       FaultTolerantRedisCluster messageInsertCacheCluster = new FaultTolerantRedisCluster("message_insert_cluster",
           configuration.getMessageCacheConfiguration().getRedisClusterConfiguration(), redisClusterClientResources);
       FaultTolerantRedisCluster messageReadDeleteCluster = new FaultTolerantRedisCluster("message_read_delete_cluster",
           configuration.getMessageCacheConfiguration().getRedisClusterConfiguration(), redisClusterClientResources);
-      FaultTolerantRedisCluster metricsCluster = new FaultTolerantRedisCluster("metrics_cluster",
-          configuration.getMetricsClusterConfiguration(), redisClusterClientResources);
       FaultTolerantRedisCluster clientPresenceCluster = new FaultTolerantRedisCluster("client_presence",
           configuration.getClientPresenceClusterConfiguration(), redisClusterClientResources);
       SecureBackupClient secureBackupClient = new SecureBackupClient(backupCredentialsGenerator, backupServiceExecutor,
@@ -179,8 +179,7 @@ public class SetUserDiscoverabilityCommand extends EnvironmentCommand<WhisperSer
       ClientPresenceManager clientPresenceManager = new ClientPresenceManager(clientPresenceCluster,
           Executors.newSingleThreadScheduledExecutor(), keyspaceNotificationDispatchExecutor);
       MessagesCache messagesCache = new MessagesCache(messageInsertCacheCluster, messageReadDeleteCluster,
-          keyspaceNotificationDispatchExecutor);
+          Clock.systemUTC(), keyspaceNotificationDispatchExecutor, messageDeletionExecutor);
-      PushLatencyManager pushLatencyManager = new PushLatencyManager(metricsCluster, dynamicConfigurationManager);
       DirectoryQueue directoryQueue = new DirectoryQueue(
           configuration.getDirectoryConfiguration().getSqsConfiguration());
       ProfilesManager profilesManager = new ProfilesManager(profiles, cacheCluster);

service/src/main/resources/lua/get_items.lua

@@ -1,6 +1,7 @@
 local queueKey     = KEYS[1]
 local queueLockKey = KEYS[2]
 local limit        = ARGV[1]
+local afterMessageId = ARGV[2]
 
 local locked = redis.call("GET", queueLockKey)
 
@@ -8,7 +9,8 @@ if locked then
     return {}
 end
 
--- The range is inclusive
+if afterMessageId == "null" then
+    -- An index range is inclusive
     local min = 0
     local max = limit - 1
 
@@ -17,3 +19,7 @@ if max < 0 then
     end
 
     return redis.call("ZRANGE", queueKey, min, max, "WITHSCORES")
+else
+    -- note: this is deprecated in Redis 6.2, and should be migrated to zrange after the cluster is updated
+    return redis.call("ZRANGEBYSCORE", queueKey, "("..afterMessageId, "+inf", "WITHSCORES", "LIMIT", 0, limit)
+end

service/src/test/java/org/whispersystems/textsecuregcm/controllers/MessageControllerTest.java

@@ -46,6 +46,7 @@ import java.util.Optional;
 import java.util.Random;
 import java.util.UUID;
 import java.util.concurrent.Callable;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Stream;
@@ -533,17 +534,25 @@ class MessageControllerTest {
     UUID sourceUuid = UUID.randomUUID();
 
     UUID uuid1 = UUID.randomUUID();
-    when(messagesManager.delete(AuthHelper.VALID_UUID, 1, uuid1, null)).thenReturn(Optional.of(generateEnvelope(
+    when(messagesManager.delete(AuthHelper.VALID_UUID, 1, uuid1, null))
-        uuid1, Envelope.Type.CIPHERTEXT_VALUE,
+        .thenReturn(
-        timestamp, sourceUuid, 1, AuthHelper.VALID_UUID, null, "hi".getBytes(), 0)));
+            CompletableFuture.completedFuture(Optional.of(generateEnvelope(uuid1, Envelope.Type.CIPHERTEXT_VALUE,
+                timestamp, sourceUuid, 1, AuthHelper.VALID_UUID, null, "hi".getBytes(), 0))));
 
     UUID uuid2 = UUID.randomUUID();
-    when(messagesManager.delete(AuthHelper.VALID_UUID, 1, uuid2, null)).thenReturn(Optional.of(generateEnvelope(
+    when(messagesManager.delete(AuthHelper.VALID_UUID, 1, uuid2, null))
+        .thenReturn(
+            CompletableFuture.completedFuture(Optional.of(generateEnvelope(
                 uuid2, Envelope.Type.SERVER_DELIVERY_RECEIPT_VALUE,
-        System.currentTimeMillis(), sourceUuid, 1, AuthHelper.VALID_UUID, null, null, 0)));
+                System.currentTimeMillis(), sourceUuid, 1, AuthHelper.VALID_UUID, null, null, 0))));
 
     UUID uuid3 = UUID.randomUUID();
-    when(messagesManager.delete(AuthHelper.VALID_UUID, 1, uuid3, null)).thenReturn(Optional.empty());
+    when(messagesManager.delete(AuthHelper.VALID_UUID, 1, uuid3, null))
+        .thenReturn(CompletableFuture.completedFuture(Optional.empty()));
+
+    UUID uuid4 = UUID.randomUUID();
+    when(messagesManager.delete(AuthHelper.VALID_UUID, 1, uuid4, null))
+        .thenReturn(CompletableFuture.failedFuture(new RuntimeException("Oh No")));
 
     Response response = resources.getJerseyTest()
         .target(String.format("/v1/messages/uuid/%s", uuid1))
@@ -573,6 +582,15 @@ class MessageControllerTest {
     assertThat("Good Response Code", response.getStatus(), is(equalTo(204)));
     verifyNoMoreInteractions(receiptSender);
 
+    response = resources.getJerseyTest()
+        .target(String.format("/v1/messages/uuid/%s", uuid4))
+        .request()
+        .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
+        .delete();
+
+    assertThat("Bad Response Code", response.getStatus(), is(equalTo(500)));
+    verifyNoMoreInteractions(receiptSender);
+
   }
 
   @Test
@@ -700,7 +718,7 @@ class MessageControllerTest {
             .target(String.format("/v1/messages/%s", SINGLE_DEVICE_UUID))
             .request()
             .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID, AuthHelper.VALID_PASSWORD))
-            .header("User-Agent", "FIXME")
+            .header("User-Agent", "Test-UA")
             .put(Entity.entity(SystemMapper.getMapper().readValue(jsonFixture(payloadFilename), IncomingMessageList.class),
                 MediaType.APPLICATION_JSON_TYPE));
 

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

@@ -13,18 +13,29 @@ import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import io.lettuce.core.FlushMode;
+import io.lettuce.core.RedisFuture;
 import io.lettuce.core.RedisNoScriptException;
 import io.lettuce.core.ScriptOutputType;
+import io.lettuce.core.cluster.api.async.RedisAdvancedClusterAsyncCommands;
+import io.lettuce.core.cluster.api.reactive.RedisAdvancedClusterReactiveCommands;
 import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands;
+import io.lettuce.core.protocol.AsyncCommand;
+import io.lettuce.core.protocol.RedisCommand;
 import java.nio.charset.StandardCharsets;
+import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.TimeUnit;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
 import org.whispersystems.textsecuregcm.tests.util.RedisClusterHelper;
+import reactor.core.publisher.Flux;
 
-public class ClusterLuaScriptTest {
+class ClusterLuaScriptTest {
 
   @RegisterExtension
   static final RedisClusterExtension REDIS_CLUSTER_EXTENSION = RedisClusterExtension.builder().build();
@@ -32,7 +43,7 @@ public class ClusterLuaScriptTest {
   @Test
   void testExecute() {
     final RedisAdvancedClusterCommands<String, String> commands = mock(RedisAdvancedClusterCommands.class);
-    final FaultTolerantRedisCluster mockCluster = RedisClusterHelper.buildMockRedisCluster(commands);
+    final FaultTolerantRedisCluster mockCluster = RedisClusterHelper.builder().stringCommands(commands).build();
 
     final String script = "return redis.call(\"SET\", KEYS[1], ARGV[1])";
     final ScriptOutputType scriptOutputType = ScriptOutputType.VALUE;
@@ -51,7 +62,7 @@ public class ClusterLuaScriptTest {
   @Test
   void testExecuteScriptNotLoaded() {
     final RedisAdvancedClusterCommands<String, String> commands = mock(RedisAdvancedClusterCommands.class);
-    final FaultTolerantRedisCluster mockCluster = RedisClusterHelper.buildMockRedisCluster(commands);
+    final FaultTolerantRedisCluster mockCluster = RedisClusterHelper.builder().stringCommands(commands).build();
 
     final String script = "return redis.call(\"SET\", KEYS[1], ARGV[1])";
     final ScriptOutputType scriptOutputType = ScriptOutputType.VALUE;
@@ -71,8 +82,10 @@ public class ClusterLuaScriptTest {
   void testExecuteBinaryScriptNotLoaded() {
     final RedisAdvancedClusterCommands<String, String> stringCommands = mock(RedisAdvancedClusterCommands.class);
     final RedisAdvancedClusterCommands<byte[], byte[]> binaryCommands = mock(RedisAdvancedClusterCommands.class);
-    final FaultTolerantRedisCluster mockCluster =
+    final FaultTolerantRedisCluster mockCluster = RedisClusterHelper.builder()
-        RedisClusterHelper.buildMockRedisCluster(stringCommands, binaryCommands);
+        .stringCommands(stringCommands)
+        .binaryCommands(binaryCommands)
+        .build();
 
     final String script = "return redis.call(\"SET\", KEYS[1], ARGV[1])";
     final ScriptOutputType scriptOutputType = ScriptOutputType.VALUE;
@@ -85,17 +98,85 @@ public class ClusterLuaScriptTest {
     luaScript.executeBinary(keys, values);
 
     verify(binaryCommands).eval(script, scriptOutputType, keys.toArray(new byte[0][]), values.toArray(new byte[0][]));
-    verify(binaryCommands).evalsha(luaScript.getSha(), scriptOutputType, keys.toArray(new byte[0][]), values.toArray(new byte[0][]));
+    verify(binaryCommands).evalsha(luaScript.getSha(), scriptOutputType, keys.toArray(new byte[0][]),
+        values.toArray(new byte[0][]));
   }
 
   @Test
-  public void testExecuteRealCluster() {
+  void testExecuteBinaryAsyncScriptNotLoaded() throws Exception {
+    final RedisAdvancedClusterAsyncCommands<byte[], byte[]> binaryAsyncCommands =
+        mock(RedisAdvancedClusterAsyncCommands.class);
+    final FaultTolerantRedisCluster mockCluster =
+        RedisClusterHelper.builder().binaryAsyncCommands(binaryAsyncCommands).build();
+
+    final String script = "return redis.call(\"SET\", KEYS[1], ARGV[1])";
+    final ScriptOutputType scriptOutputType = ScriptOutputType.VALUE;
+    final List<byte[]> keys = List.of("key".getBytes(StandardCharsets.UTF_8));
+    final List<byte[]> values = List.of("value".getBytes(StandardCharsets.UTF_8));
+
+    final AsyncCommand<?, ?, ?> evalShaFailure = new AsyncCommand<>(mock(RedisCommand.class));
+    evalShaFailure.completeExceptionally(new RedisNoScriptException("OH NO"));
+
+    final AsyncCommand<?, ?, ?> evalSuccess = new AsyncCommand<>(mock(RedisCommand.class));
+    evalSuccess.complete();
+
+    when(binaryAsyncCommands.evalsha(any(), any(), any(), any())).thenReturn((RedisFuture<Object>) evalShaFailure);
+    when(binaryAsyncCommands.eval(anyString(), any(), any(), any())).thenReturn((RedisFuture<Object>) evalSuccess);
+
+    final ClusterLuaScript luaScript = new ClusterLuaScript(mockCluster, script, scriptOutputType);
+    luaScript.executeBinaryAsync(keys, values).get(5, TimeUnit.SECONDS);
+
+    verify(binaryAsyncCommands).eval(script, scriptOutputType, keys.toArray(new byte[0][]),
+        values.toArray(new byte[0][]));
+    verify(binaryAsyncCommands).evalsha(luaScript.getSha(), scriptOutputType, keys.toArray(new byte[0][]),
+        values.toArray(new byte[0][]));
+  }
+
+  @Test
+  void testExecuteBinaryReactiveScriptNotLoaded() {
+    final RedisAdvancedClusterReactiveCommands<byte[], byte[]> binaryReactiveCommands =
+        mock(RedisAdvancedClusterReactiveCommands.class);
+    final FaultTolerantRedisCluster mockCluster = RedisClusterHelper.builder()
+        .binaryReactiveCommands(binaryReactiveCommands).build();
+
+    final String script = "return redis.call(\"SET\", KEYS[1], ARGV[1])";
+    final ScriptOutputType scriptOutputType = ScriptOutputType.VALUE;
+    final List<byte[]> keys = List.of("key".getBytes(StandardCharsets.UTF_8));
+    final List<byte[]> values = List.of("value".getBytes(StandardCharsets.UTF_8));
+
+    when(binaryReactiveCommands.evalsha(any(), any(), any(), any()))
+        .thenReturn(Flux.error(new RedisNoScriptException("OH NO")));
+    when(binaryReactiveCommands.eval(anyString(), any(), any(), any())).thenReturn(Flux.just("ok"));
+
+    final ClusterLuaScript luaScript = new ClusterLuaScript(mockCluster, script, scriptOutputType);
+    luaScript.executeBinaryReactive(keys, values).blockLast(Duration.ofSeconds(5));
+
+    verify(binaryReactiveCommands).eval(script, scriptOutputType, keys.toArray(new byte[0][]),
+        values.toArray(new byte[0][]));
+    verify(binaryReactiveCommands).evalsha(luaScript.getSha(), scriptOutputType, keys.toArray(new byte[0][]),
+        values.toArray(new byte[0][]));
+  }
+
+  @ParameterizedTest
+  @EnumSource(ExecuteMode.class)
+  void testExecuteRealCluster(final ExecuteMode mode) throws Exception {
+    REDIS_CLUSTER_EXTENSION.getRedisCluster().withCluster(c -> c.sync().scriptFlush(FlushMode.SYNC));
+    REDIS_CLUSTER_EXTENSION.getRedisCluster().withCluster(c -> c.sync().configResetstat());
+
     final ClusterLuaScript script = new ClusterLuaScript(REDIS_CLUSTER_EXTENSION.getRedisCluster(),
         "return 2;",
         ScriptOutputType.INTEGER);
 
     for (int i = 0; i < 7; i++) {
-      assertEquals(2L, script.execute(Collections.emptyList(), Collections.emptyList()));
+      final long actual = switch (mode) {
+        case SYNC -> (long) script.execute(Collections.emptyList(), Collections.emptyList());
+        case ASYNC ->
+            (long) script.executeAsync(Collections.emptyList(), Collections.emptyList()).get(5, TimeUnit.SECONDS);
+        case REACTIVE -> (long) script.executeReactive(Collections.emptyList(), Collections.emptyList())
+            .blockLast(Duration.ofSeconds(5));
+      };
+
+      assertEquals(2L, actual);
     }
 
     final int evalCount = REDIS_CLUSTER_EXTENSION.getRedisCluster().withCluster(connection -> {
@@ -120,4 +201,11 @@ public class ClusterLuaScriptTest {
 
     assertEquals(1, evalCount);
   }
+
+  private enum ExecuteMode {
+    SYNC,
+    ASYNC,
+    REACTIVE
+  }
+
 }

service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerConcurrentModificationIntegrationTest.java

@@ -155,7 +155,7 @@ class AccountsManagerConcurrentModificationIntegrationTest {
       accountsManager = new AccountsManager(
           accounts,
           phoneNumberIdentifiers,
-          RedisClusterHelper.buildMockRedisCluster(commands),
+          RedisClusterHelper.builder().stringCommands(commands).build(),
           deletedAccountsManager,
           mock(DirectoryQueue.class),
           mock(Keys.class),

service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerTest.java

@@ -147,7 +147,7 @@ class AccountsManagerTest {
     accountsManager = new AccountsManager(
         accounts,
         phoneNumberIdentifiers,
-        RedisClusterHelper.buildMockRedisCluster(commands),
+        RedisClusterHelper.builder().stringCommands(commands).build(),
         deletedAccountsManager,
         directoryQueue,
         keys,

service/src/test/java/org/whispersystems/textsecuregcm/storage/DynamoDbExtension.java

@@ -78,7 +78,14 @@ public class DynamoDbExtension implements BeforeEachCallback, AfterEachCallback
   }
 
   @Override
-  public void afterEach(ExtensionContext context) throws Exception {
+  public void afterEach(ExtensionContext context) {
+    stopServer();
+  }
+
+  /**
+   * For use in integration tests that want to test resiliency/error handling
+   */
+  public void stopServer() {
     try {
       server.stop();
     } catch (Exception e) {

service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagePersisterIntegrationTest.java

@@ -15,6 +15,7 @@ import com.google.protobuf.ByteString;
 import com.google.protobuf.InvalidProtocolBufferException;
 import io.lettuce.core.cluster.SlotHash;
 import java.nio.ByteBuffer;
+import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.ArrayList;
@@ -32,7 +33,6 @@ import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
 import org.whispersystems.textsecuregcm.entities.MessageProtos;
-import org.whispersystems.textsecuregcm.push.PushLatencyManager;
 import org.whispersystems.textsecuregcm.redis.RedisClusterExtension;
 import org.whispersystems.textsecuregcm.tests.util.MessagesDynamoDbExtension;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
@@ -47,6 +47,7 @@ class MessagePersisterIntegrationTest {
   static final RedisClusterExtension REDIS_CLUSTER_EXTENSION = RedisClusterExtension.builder().build();
 
   private ExecutorService notificationExecutorService;
+  private ExecutorService messageDeletionExecutorService;
   private MessagesCache messagesCache;
   private MessagesManager messagesManager;
   private MessagePersister messagePersister;
@@ -66,13 +67,16 @@ class MessagePersisterIntegrationTest {
 
     when(dynamicConfigurationManager.getConfiguration()).thenReturn(new DynamicConfiguration());
 
+    messageDeletionExecutorService = Executors.newSingleThreadExecutor();
     final MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbExtension.getDynamoDbClient(),
-        MessagesDynamoDbExtension.TABLE_NAME, Duration.ofDays(14));
+        dynamoDbExtension.getDynamoDbAsyncClient(), MessagesDynamoDbExtension.TABLE_NAME, Duration.ofDays(14),
+        messageDeletionExecutorService);
     final AccountsManager accountsManager = mock(AccountsManager.class);
 
     notificationExecutorService = Executors.newSingleThreadExecutor();
     messagesCache = new MessagesCache(REDIS_CLUSTER_EXTENSION.getRedisCluster(),
-        REDIS_CLUSTER_EXTENSION.getRedisCluster(), notificationExecutorService);
+        REDIS_CLUSTER_EXTENSION.getRedisCluster(), Clock.systemUTC(), notificationExecutorService,
+        messageDeletionExecutorService);
     messagesManager = new MessagesManager(messagesDynamoDb, messagesCache, mock(ReportMessageManager.class));
     messagePersister = new MessagePersister(messagesCache, messagesManager, accountsManager,
         dynamicConfigurationManager, PERSIST_DELAY);
@@ -94,6 +98,9 @@ class MessagePersisterIntegrationTest {
   void tearDown() throws Exception {
     notificationExecutorService.shutdown();
     notificationExecutorService.awaitTermination(15, TimeUnit.SECONDS);
+
+    messageDeletionExecutorService.shutdown();
+    messageDeletionExecutorService.awaitTermination(15, TimeUnit.SECONDS);
   }
 
   @Test

service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagePersisterTest.java

@@ -22,6 +22,7 @@ import static org.mockito.Mockito.when;
 import com.google.protobuf.ByteString;
 import io.lettuce.core.cluster.SlotHash;
 import java.nio.charset.StandardCharsets;
+import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.List;
@@ -46,7 +47,7 @@ class MessagePersisterTest {
   @RegisterExtension
   static final RedisClusterExtension REDIS_CLUSTER_EXTENSION = RedisClusterExtension.builder().build();
 
-  private ExecutorService notificationExecutorService;
+  private ExecutorService sharedExecutorService;
   private MessagesCache messagesCache;
   private MessagesDynamoDb messagesDynamoDb;
   private MessagePersister messagePersister;
@@ -74,9 +75,9 @@ class MessagePersisterTest {
     when(account.getNumber()).thenReturn(DESTINATION_ACCOUNT_NUMBER);
     when(dynamicConfigurationManager.getConfiguration()).thenReturn(new DynamicConfiguration());
 
-    notificationExecutorService = Executors.newSingleThreadExecutor();
+    sharedExecutorService = Executors.newSingleThreadExecutor();
     messagesCache = new MessagesCache(REDIS_CLUSTER_EXTENSION.getRedisCluster(),
-        REDIS_CLUSTER_EXTENSION.getRedisCluster(), notificationExecutorService);
+        REDIS_CLUSTER_EXTENSION.getRedisCluster(), Clock.systemUTC(), sharedExecutorService, sharedExecutorService);
     messagePersister = new MessagePersister(messagesCache, messagesManager, accountsManager,
         dynamicConfigurationManager, PERSIST_DELAY);
 
@@ -88,7 +89,7 @@ class MessagePersisterTest {
       messagesDynamoDb.store(messages, destinationUuid, destinationDeviceId);
 
       for (final MessageProtos.Envelope message : messages) {
-        messagesCache.remove(destinationUuid, destinationDeviceId, UUID.fromString(message.getServerGuid()));
+        messagesCache.remove(destinationUuid, destinationDeviceId, UUID.fromString(message.getServerGuid())).get();
       }
 
       return null;
@@ -97,8 +98,8 @@ class MessagePersisterTest {
 
   @AfterEach
   void tearDown() throws Exception {
-    notificationExecutorService.shutdown();
+    sharedExecutorService.shutdown();
-    notificationExecutorService.awaitTermination(1, TimeUnit.SECONDS);
+    sharedExecutorService.awaitTermination(1, TimeUnit.SECONDS);
   }
 
   @Test

service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagesCacheTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2021 Signal Messenger, LLC
+ * Copyright 2013-2022 Signal Messenger, LLC
  * SPDX-License-Identifier: AGPL-3.0-only
  */
 
@@ -7,16 +7,34 @@ package org.whispersystems.textsecuregcm.storage;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.atLeast;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import com.google.protobuf.ByteString;
+import io.lettuce.core.RedisFuture;
 import io.lettuce.core.cluster.SlotHash;
+import io.lettuce.core.cluster.api.async.RedisAdvancedClusterAsyncCommands;
+import io.lettuce.core.cluster.api.reactive.RedisAdvancedClusterReactiveCommands;
+import io.lettuce.core.protocol.AsyncCommand;
+import io.lettuce.core.protocol.RedisCommand;
 import java.nio.charset.StandardCharsets;
+import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
+import java.time.ZoneId;
+import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.Deque;
 import java.util.List;
 import java.util.Optional;
 import java.util.Random;
@@ -26,29 +44,41 @@ import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 import org.apache.commons.lang3.RandomStringUtils;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
+import org.reactivestreams.Publisher;
 import org.whispersystems.textsecuregcm.entities.MessageProtos;
+import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
 import org.whispersystems.textsecuregcm.redis.RedisClusterExtension;
+import org.whispersystems.textsecuregcm.tests.util.RedisClusterHelper;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.FluxSink;
+import reactor.test.StepVerifier;
 
 class MessagesCacheTest {
 
+  private final Random random = new Random();
+  private long serialTimestamp = 0;
+
+  @Nested
+  class WithRealCluster {
+
     @RegisterExtension
     static final RedisClusterExtension REDIS_CLUSTER_EXTENSION = RedisClusterExtension.builder().build();
 
-  private ExecutorService notificationExecutorService;
+    private ExecutorService sharedExecutorService;
     private MessagesCache messagesCache;
 
-  private final Random random = new Random();
-  private long serialTimestamp = 0;
-
     private static final UUID DESTINATION_UUID = UUID.randomUUID();
     private static final int DESTINATION_DEVICE_ID = 7;
 
@@ -60,8 +90,10 @@ class MessagesCacheTest {
         connection.sync().upstream().commands().configSet("notify-keyspace-events", "K$glz");
       });
 
-    notificationExecutorService = Executors.newSingleThreadExecutor();
+      sharedExecutorService = Executors.newSingleThreadExecutor();
-    messagesCache = new MessagesCache(REDIS_CLUSTER_EXTENSION.getRedisCluster(), REDIS_CLUSTER_EXTENSION.getRedisCluster(), notificationExecutorService);
+      messagesCache = new MessagesCache(REDIS_CLUSTER_EXTENSION.getRedisCluster(),
+          REDIS_CLUSTER_EXTENSION.getRedisCluster(), Clock.systemUTC(), sharedExecutorService,
+          sharedExecutorService);
 
       messagesCache.start();
     }
@@ -70,8 +102,8 @@ class MessagesCacheTest {
     void tearDown() throws Exception {
       messagesCache.stop();
 
-    notificationExecutorService.shutdown();
+      sharedExecutorService.shutdown();
-    notificationExecutorService.awaitTermination(1, TimeUnit.SECONDS);
+      sharedExecutorService.awaitTermination(1, TimeUnit.SECONDS);
     }
 
     @ParameterizedTest
@@ -87,7 +119,8 @@ class MessagesCacheTest {
       final UUID duplicateGuid = UUID.randomUUID();
       final MessageProtos.Envelope duplicateMessage = generateRandomMessage(duplicateGuid, false);
 
-    final long firstId = messagesCache.insert(duplicateGuid, DESTINATION_UUID, DESTINATION_DEVICE_ID, duplicateMessage);
+      final long firstId = messagesCache.insert(duplicateGuid, DESTINATION_UUID, DESTINATION_DEVICE_ID,
+          duplicateMessage);
       final long secondId = messagesCache.insert(duplicateGuid, DESTINATION_UUID, DESTINATION_DEVICE_ID,
           duplicateMessage);
 
@@ -96,23 +129,24 @@ class MessagesCacheTest {
 
     @ParameterizedTest
     @ValueSource(booleans = {true, false})
-  void testRemoveByUUID(final boolean sealedSender) {
+    void testRemoveByUUID(final boolean sealedSender) throws Exception {
       final UUID messageGuid = UUID.randomUUID();
 
-    assertEquals(Optional.empty(), messagesCache.remove(DESTINATION_UUID, DESTINATION_DEVICE_ID, messageGuid));
+      assertEquals(Optional.empty(),
+          messagesCache.remove(DESTINATION_UUID, DESTINATION_DEVICE_ID, messageGuid).get(5, TimeUnit.SECONDS));
 
       final MessageProtos.Envelope message = generateRandomMessage(messageGuid, sealedSender);
 
       messagesCache.insert(messageGuid, DESTINATION_UUID, DESTINATION_DEVICE_ID, message);
       final Optional<MessageProtos.Envelope> maybeRemovedMessage = messagesCache.remove(DESTINATION_UUID,
-        DESTINATION_DEVICE_ID, messageGuid);
+          DESTINATION_DEVICE_ID, messageGuid).get(5, TimeUnit.SECONDS);
 
       assertEquals(Optional.of(message), maybeRemovedMessage);
     }
 
     @ParameterizedTest
     @ValueSource(booleans = {true, false})
-  void testRemoveBatchByUUID(final boolean sealedSender) {
+    void testRemoveBatchByUUID(final boolean sealedSender) throws Exception {
       final int messageCount = 10;
 
       final List<MessageProtos.Envelope> messagesToRemove = new ArrayList<>(messageCount);
@@ -125,19 +159,21 @@ class MessagesCacheTest {
 
       assertEquals(Collections.emptyList(), messagesCache.remove(DESTINATION_UUID, DESTINATION_DEVICE_ID,
           messagesToRemove.stream().map(message -> UUID.fromString(message.getServerGuid()))
-            .collect(Collectors.toList())));
+              .collect(Collectors.toList())).get(5, TimeUnit.SECONDS));
 
       for (final MessageProtos.Envelope message : messagesToRemove) {
-      messagesCache.insert(UUID.fromString(message.getServerGuid()), DESTINATION_UUID, DESTINATION_DEVICE_ID, message);
+        messagesCache.insert(UUID.fromString(message.getServerGuid()), DESTINATION_UUID, DESTINATION_DEVICE_ID,
+            message);
       }
 
       for (final MessageProtos.Envelope message : messagesToPreserve) {
-      messagesCache.insert(UUID.fromString(message.getServerGuid()), DESTINATION_UUID, DESTINATION_DEVICE_ID, message);
+        messagesCache.insert(UUID.fromString(message.getServerGuid()), DESTINATION_UUID, DESTINATION_DEVICE_ID,
+            message);
       }
 
       final List<MessageProtos.Envelope> removedMessages = messagesCache.remove(DESTINATION_UUID, DESTINATION_DEVICE_ID,
           messagesToRemove.stream().map(message -> UUID.fromString(message.getServerGuid()))
-            .collect(Collectors.toList()));
+              .collect(Collectors.toList())).get(5, TimeUnit.SECONDS);
 
       assertEquals(messagesToRemove, removedMessages);
       assertEquals(messagesToPreserve,
@@ -157,7 +193,7 @@ class MessagesCacheTest {
 
     @ParameterizedTest
     @ValueSource(booleans = {true, false})
-  void testGetMessages(final boolean sealedSender) {
+    void testGetMessages(final boolean sealedSender) throws Exception {
       final int messageCount = 100;
 
       final List<MessageProtos.Envelope> expectedMessages = new ArrayList<>(messageCount);
@@ -170,7 +206,97 @@ class MessagesCacheTest {
         expectedMessages.add(message);
       }
 
-    assertEquals(expectedMessages, messagesCache.get(DESTINATION_UUID, DESTINATION_DEVICE_ID, messageCount));
+      assertEquals(expectedMessages, get(DESTINATION_UUID, DESTINATION_DEVICE_ID, messageCount));
+
+      messagesCache.remove(DESTINATION_UUID, DESTINATION_DEVICE_ID,
+          expectedMessages.stream()
+              .map(MessageProtos.Envelope::getServerGuid)
+              .map(UUID::fromString)
+              .collect(Collectors.toList()));
+
+      final UUID message1Guid = UUID.randomUUID();
+      final MessageProtos.Envelope message1 = generateRandomMessage(message1Guid, sealedSender);
+      messagesCache.insert(message1Guid, DESTINATION_UUID, DESTINATION_DEVICE_ID, message1);
+      final List<MessageProtos.Envelope> get1 = get(DESTINATION_UUID, DESTINATION_DEVICE_ID,
+          1);
+      assertEquals(List.of(message1), get1);
+
+      messagesCache.remove(DESTINATION_UUID, DESTINATION_DEVICE_ID, message1Guid).get(5, TimeUnit.SECONDS);
+
+      final UUID message2Guid = UUID.randomUUID();
+      final MessageProtos.Envelope message2 = generateRandomMessage(message2Guid, sealedSender);
+
+      messagesCache.insert(message2Guid, DESTINATION_UUID, DESTINATION_DEVICE_ID, message2);
+
+      assertEquals(List.of(message2), get(DESTINATION_UUID, DESTINATION_DEVICE_ID, 1));
+    }
+
+    @ParameterizedTest
+    @ValueSource(booleans = {true, false})
+    void testGetMessagesPublisher(final boolean expectStale) throws Exception {
+      final int messageCount = 214;
+
+      final List<MessageProtos.Envelope> expectedMessages = new ArrayList<>(messageCount);
+
+      for (int i = 0; i < messageCount; i++) {
+        final UUID messageGuid = UUID.randomUUID();
+        final MessageProtos.Envelope message = generateRandomMessage(messageGuid, true);
+        messagesCache.insert(messageGuid, DESTINATION_UUID, DESTINATION_DEVICE_ID, message);
+
+        expectedMessages.add(message);
+      }
+
+      final UUID ephemeralMessageGuid = UUID.randomUUID();
+      final MessageProtos.Envelope ephemeralMessage = generateRandomMessage(ephemeralMessageGuid, true)
+          .toBuilder().setEphemeral(true).build();
+      messagesCache.insert(ephemeralMessageGuid, DESTINATION_UUID, DESTINATION_DEVICE_ID, ephemeralMessage);
+
+      final Clock cacheClock;
+      if (expectStale) {
+        cacheClock = Clock.fixed(Instant.ofEpochMilli(serialTimestamp + 1),
+            ZoneId.of("Etc/UTC"));
+      } else {
+        cacheClock = Clock.fixed(
+            Instant.ofEpochMilli(serialTimestamp + 1).plus(MessagesCache.MAX_EPHEMERAL_MESSAGE_DELAY),
+            ZoneId.of("Etc/UTC"));
+      }
+
+      final MessagesCache messagesCache = new MessagesCache(REDIS_CLUSTER_EXTENSION.getRedisCluster(),
+          REDIS_CLUSTER_EXTENSION.getRedisCluster(),
+          cacheClock,
+          sharedExecutorService,
+          sharedExecutorService);
+
+      final List<MessageProtos.Envelope> actualMessages = Flux.from(
+              messagesCache.get(DESTINATION_UUID, DESTINATION_DEVICE_ID))
+          .collectList()
+          .block(Duration.ofSeconds(5));
+
+      if (expectStale) {
+        final List<MessageProtos.Envelope> expectedAllMessages = new ArrayList<>() {{
+          addAll(expectedMessages);
+          add(ephemeralMessage);
+        }};
+
+        assertEquals(expectedAllMessages, actualMessages);
+
+      } else {
+        assertEquals(expectedMessages, actualMessages);
+
+        // delete all of these messages and call `getAll()`, to confirm that ephemeral messages have been discarded
+        CompletableFuture.allOf(actualMessages.stream()
+                .map(message -> messagesCache.remove(DESTINATION_UUID, DESTINATION_DEVICE_ID,
+                    UUID.fromString(message.getServerGuid())))
+                .toArray(CompletableFuture<?>[]::new))
+            .get(5, TimeUnit.SECONDS);
+
+        final List<MessageProtos.Envelope> messages = messagesCache.getAllMessages(DESTINATION_UUID,
+                DESTINATION_DEVICE_ID)
+            .collectList()
+            .toFuture().get(5, TimeUnit.SECONDS);
+
+        assertTrue(messages.isEmpty());
+      }
     }
 
     @ParameterizedTest
@@ -189,8 +315,8 @@ class MessagesCacheTest {
 
       messagesCache.clear(DESTINATION_UUID, DESTINATION_DEVICE_ID);
 
-    assertEquals(Collections.emptyList(), messagesCache.get(DESTINATION_UUID, DESTINATION_DEVICE_ID, messageCount));
+      assertEquals(Collections.emptyList(), get(DESTINATION_UUID, DESTINATION_DEVICE_ID, messageCount));
-    assertEquals(messageCount, messagesCache.get(DESTINATION_UUID, DESTINATION_DEVICE_ID + 1, messageCount).size());
+      assertEquals(messageCount, get(DESTINATION_UUID, DESTINATION_DEVICE_ID + 1, messageCount).size());
     }
 
     @ParameterizedTest
@@ -209,30 +335,8 @@ class MessagesCacheTest {
 
       messagesCache.clear(DESTINATION_UUID);
 
-    assertEquals(Collections.emptyList(), messagesCache.get(DESTINATION_UUID, DESTINATION_DEVICE_ID, messageCount));
+      assertEquals(Collections.emptyList(), get(DESTINATION_UUID, DESTINATION_DEVICE_ID, messageCount));
-    assertEquals(Collections.emptyList(), messagesCache.get(DESTINATION_UUID, DESTINATION_DEVICE_ID + 1, messageCount));
+      assertEquals(Collections.emptyList(), get(DESTINATION_UUID, DESTINATION_DEVICE_ID + 1, messageCount));
-  }
-
-  private MessageProtos.Envelope generateRandomMessage(final UUID messageGuid, final boolean sealedSender) {
-    return generateRandomMessage(messageGuid, sealedSender, serialTimestamp++);
-  }
-
-  private MessageProtos.Envelope generateRandomMessage(final UUID messageGuid, final boolean sealedSender,
-      final long timestamp) {
-    final MessageProtos.Envelope.Builder envelopeBuilder = MessageProtos.Envelope.newBuilder()
-        .setTimestamp(timestamp)
-        .setServerTimestamp(timestamp)
-        .setContent(ByteString.copyFromUtf8(RandomStringUtils.randomAlphanumeric(256)))
-        .setType(MessageProtos.Envelope.Type.CIPHERTEXT)
-        .setServerGuid(messageGuid.toString())
-        .setDestinationUuid(UUID.randomUUID().toString());
-
-    if (!sealedSender) {
-      envelopeBuilder.setSourceDevice(random.nextInt(256))
-          .setSourceUuid(UUID.randomUUID().toString());
-    }
-
-    return envelopeBuilder.build();
     }
 
     @Test
@@ -359,7 +463,8 @@ class MessagesCacheTest {
 
     /**
      * Helper class that implements {@link MessageAvailabilityListener#handleNewMessagesAvailable()} by always returning
-   * {@code false}. Its {@code counter} field tracks how many times {@code handleNewMessagesAvailable} has been called.
+     * {@code false}. Its {@code counter} field tracks how many times {@code handleNewMessagesAvailable} has been
+     * called.
      * <p>
      * It uses a {@link CompletableFuture} to signal that it has received a “messages available” callback for the first
      * time.
@@ -409,7 +514,7 @@ class MessagesCacheTest {
         // Avoid a race condition by blocking on the message handled future *and* the current notification executor task—
         // the notification executor task includes unsubscribing `listener1`, and, if we don’t wait, sometimes
         // `listener2` will get subscribed before `listener1` is cleaned up
-      notificationExecutorService.submit(() -> listener1.firstMessageHandled.get()).get();
+        sharedExecutorService.submit(() -> listener1.firstMessageHandled.get()).get();
 
         final UUID messageGuid2 = UUID.randomUUID();
         messagesCache.insert(messageGuid2, DESTINATION_UUID, DESTINATION_DEVICE_ID,
@@ -424,4 +529,228 @@ class MessagesCacheTest {
         listener2.firstMessageHandled.get();
       });
     }
+
+    private List<MessageProtos.Envelope> get(final UUID destinationUuid, final long destinationDeviceId,
+        final int messageCount) {
+      return Flux.from(messagesCache.get(destinationUuid, destinationDeviceId))
+          .take(messageCount, true)
+          .collectList()
+          .block();
+    }
+
+  }
+
+  @Nested
+  class WithMockCluster {
+
+    private MessagesCache messagesCache;
+    private RedisAdvancedClusterReactiveCommands<byte[], byte[]> reactiveCommands;
+    private RedisAdvancedClusterAsyncCommands<byte[], byte[]> asyncCommands;
+
+    @SuppressWarnings("unchecked")
+    @BeforeEach
+    void setup() throws Exception {
+      reactiveCommands = mock(RedisAdvancedClusterReactiveCommands.class);
+      asyncCommands = mock(RedisAdvancedClusterAsyncCommands.class);
+
+      final FaultTolerantRedisCluster mockCluster = RedisClusterHelper.builder()
+          .binaryReactiveCommands(reactiveCommands)
+          .binaryAsyncCommands(asyncCommands)
+          .build();
+
+      messagesCache = new MessagesCache(mockCluster, mockCluster, Clock.systemUTC(), mock(ExecutorService.class),
+          mock(ExecutorService.class));
+    }
+
+    @AfterEach
+    void teardown() {
+      StepVerifier.resetDefaultTimeout();
+    }
+
+    @Test
+    void testGetAllMessagesLimitsAndBackpressure() {
+      // this test makes sure that we don’t fetch and buffer all messages from the cache when the publisher
+      // is subscribed. Rather, we should be fetching in pages to satisfy downstream requests, so that memory usage
+      // is limited to few pages of messages
+
+      // we use a combination of Flux.just() and Sinks to control when data is “fetched” from the cache. The initial
+      // Flux.just()s are pages that are readily available, on demand. By design, there are more of these pages than
+      // the initial prefetch. The sinks allow us to create extra demand but defer producing values to satisfy the demand
+      // until later on.
+
+      final AtomicReference<FluxSink<Object>> page4Sink = new AtomicReference<>();
+      final AtomicReference<FluxSink<Object>> page56Sink = new AtomicReference<>();
+      final AtomicReference<FluxSink<Object>> emptyFinalPageSink = new AtomicReference<>();
+
+      final Deque<List<byte[]>> pages = new ArrayDeque<>();
+      pages.add(generatePage());
+      pages.add(generatePage());
+      pages.add(generatePage());
+      pages.add(generatePage());
+      // make sure that stale ephemeral messages are also produced by calls to getAllMessages()
+      pages.add(generateStaleEphemeralPage());
+      pages.add(generatePage());
+
+      when(reactiveCommands.evalsha(any(), any(), any(), any()))
+          .thenReturn(Flux.just(pages.pop()))
+          .thenReturn(Flux.just(pages.pop()))
+          .thenReturn(Flux.just(pages.pop()))
+          .thenReturn(Flux.create(sink -> page4Sink.compareAndSet(null, sink)))
+          .thenReturn(Flux.create(sink -> page56Sink.compareAndSet(null, sink)))
+          .thenReturn(Flux.create(sink -> emptyFinalPageSink.compareAndSet(null, sink)))
+          .thenReturn(Flux.empty());
+
+      final Flux<?> allMessages = messagesCache.getAllMessages(UUID.randomUUID(), 1L);
+
+      // Why initialValue = 3?
+      // 1. messagesCache.getAllMessages() above produces the first call
+      // 2. when we subscribe, the prefetch of 1 results in `expand()`, which produces a second call
+      // 3. there is an implicit “low tide mark” of 1, meaning there will be an extra call to replenish when there is
+      //    1 value remaining
+      final AtomicInteger expectedReactiveCommandInvocations = new AtomicInteger(3);
+
+      StepVerifier.setDefaultTimeout(Duration.ofSeconds(5));
+
+      final int page = 100;
+      final int halfPage = page / 2;
+
+      // in order to fully control demand and separate the prefetch mechanics, initially subscribe with a request of 0
+      StepVerifier.create(allMessages, 0)
+          .expectSubscription()
+          .then(() -> verify(reactiveCommands, times(expectedReactiveCommandInvocations.get())).evalsha(any(), any(),
+              any(), any()))
+          .thenRequest(halfPage) // page 0.5 requested
+          .expectNextCount(halfPage) // page 0.5 produced
+          // page 0.5 produced, 1.5 remain, so no additional interactions with the cache cluster
+          .then(() -> verify(reactiveCommands, times(expectedReactiveCommandInvocations.get())).evalsha(any(),
+              any(), any(), any()))
+          .then(() -> assertNull(page4Sink.get(), "page 4 should not have been fetched yet"))
+          .thenRequest(page) // page 1.5 requested
+          .expectNextCount(page) // page 1.5 produced
+
+          // we now have produced 1.5 pages, have 0.5 buffered, and two more have been prefetched.
+          // after producing more than a full page, we’ll need to replenish from the cache.
+          // future requests will depend on sink emitters.
+          // also NB: times() checks cumulative calls, hence addAndGet
+          .then(() -> verify(reactiveCommands, times(expectedReactiveCommandInvocations.addAndGet(1))).evalsha(any(),
+              any(), any(), any()))
+          .then(() -> assertNotNull(page4Sink.get(), "page 4 should have been fetched"))
+          .thenRequest(page + halfPage) // page 3 requested
+          .expectNextCount(page + halfPage) // page 1.5–3 produced
+
+          .thenRequest(halfPage) // page 3.5 requested
+          .then(() -> assertNull(page56Sink.get(), "page 5 should not have been fetched yet"))
+          .then(() -> page4Sink.get().next(pages.pop()).complete())
+          .expectNextCount(halfPage) // page 3.5 produced
+          .then(() -> verify(reactiveCommands, times(expectedReactiveCommandInvocations.addAndGet(1))).evalsha(any(),
+              any(), any(), any()))
+          .then(() -> assertNotNull(page56Sink.get(), "page 5 should have been fetched"))
+
+          .thenRequest(page) // page 4.5 requested
+          .expectNextCount(halfPage) // page 4 produced
+
+          .thenRequest(page * 4) // request more demand than we will ultimately satisfy
+
+          .then(() -> page56Sink.get().next(pages.pop()).next(pages.pop()).complete())
+          .expectNextCount(page + page) // page 5 and 6 produced
+          .then(() -> emptyFinalPageSink.get().complete())
+          // confirm that cache calls increased by 2: one for page 5-and-6 (we got a two-fer in next(pop()).next(pop()),
+          //   and one for the final, empty page
+          .then(() -> verify(reactiveCommands, times(expectedReactiveCommandInvocations.addAndGet(2))).evalsha(any(),
+              any(), any(),
+              any()))
+          .expectComplete()
+          .log()
+          .verify();
+
+      // make sure that we consumed all the pages, especially in case of refactoring
+      assertTrue(pages.isEmpty());
+    }
+
+    @Test
+    void testGetDiscardsEphemeralMessages() {
+      final Deque<List<byte[]>> pages = new ArrayDeque<>();
+      pages.add(generatePage());
+      pages.add(generatePage());
+      pages.add(generateStaleEphemeralPage());
+
+      when(reactiveCommands.evalsha(any(), any(), any(), any()))
+          .thenReturn(Flux.just(pages.pop()))
+          .thenReturn(Flux.just(pages.pop()))
+          .thenReturn(Flux.just(pages.pop()))
+          .thenReturn(Flux.empty());
+
+      final AsyncCommand<?, ?, ?> removeSuccess = new AsyncCommand<>(mock(RedisCommand.class));
+      removeSuccess.complete();
+
+      when(asyncCommands.evalsha(any(), any(), any(), any()))
+          .thenReturn((RedisFuture) removeSuccess);
+
+      final Publisher<?> allMessages = messagesCache.get(UUID.randomUUID(), 1L);
+
+      StepVerifier.setDefaultTimeout(Duration.ofSeconds(5));
+
+      // async commands are used for remove(), and nothing should happen until we are subscribed
+      verify(asyncCommands, never()).evalsha(any(), any(), any(byte[][].class), any(byte[].class));
+      // the reactive commands will be called once, to prep the first page fetch (but no remote request would actually be sent)
+      verify(reactiveCommands, times(1)).evalsha(any(), any(), any(byte[][].class), any(byte[].class));
+
+      StepVerifier.create(allMessages)
+          .expectSubscription()
+          .expectNextCount(200)
+          .expectComplete()
+          .log()
+          .verify();
+
+      assertTrue(pages.isEmpty());
+      verify(asyncCommands, atLeast(1)).evalsha(any(), any(), any(), any());
+    }
+
+    private List<byte[]> generatePage() {
+      final List<byte[]> messagesAndIds = new ArrayList<>();
+
+      for (int i = 0; i < 100; i++) {
+        final MessageProtos.Envelope envelope = generateRandomMessage(UUID.randomUUID(), true);
+        messagesAndIds.add(envelope.toByteArray());
+        messagesAndIds.add(String.valueOf(serialTimestamp).getBytes());
+      }
+
+      return messagesAndIds;
+    }
+
+    private List<byte[]> generateStaleEphemeralPage() {
+      final List<byte[]> messagesAndIds = new ArrayList<>();
+
+      for (int i = 0; i < 100; i++) {
+        final MessageProtos.Envelope envelope = generateRandomMessage(UUID.randomUUID(), true)
+            .toBuilder().setEphemeral(true).build();
+        messagesAndIds.add(envelope.toByteArray());
+        messagesAndIds.add(String.valueOf(serialTimestamp).getBytes());
+      }
+
+      return messagesAndIds;
+    }
+  }
+
+  private MessageProtos.Envelope generateRandomMessage(final UUID messageGuid, final boolean sealedSender) {
+    return generateRandomMessage(messageGuid, sealedSender, serialTimestamp++);
+  }
+
+  private MessageProtos.Envelope generateRandomMessage(final UUID messageGuid, final boolean sealedSender,
+      final long timestamp) {
+    final MessageProtos.Envelope.Builder envelopeBuilder = MessageProtos.Envelope.newBuilder()
+        .setTimestamp(timestamp)
+        .setServerTimestamp(timestamp)
+        .setContent(ByteString.copyFromUtf8(RandomStringUtils.randomAlphanumeric(256)))
+        .setType(MessageProtos.Envelope.Type.CIPHERTEXT)
+        .setServerGuid(messageGuid.toString())
+        .setDestinationUuid(UUID.randomUUID().toString());
+
+    if (!sealedSender) {
+      envelopeBuilder.setSourceDevice(random.nextInt(256))
+          .setSourceUuid(UUID.randomUUID().toString());
+    }
+
+    return envelopeBuilder.build();
+  }
 }

service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDbTest.java

@@ -9,14 +9,26 @@ import static org.assertj.core.api.Assertions.assertThat;
 
 import com.google.protobuf.ByteString;
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Random;
 import java.util.UUID;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+import org.reactivestreams.Publisher;
 import org.whispersystems.textsecuregcm.entities.MessageProtos;
+import org.whispersystems.textsecuregcm.tests.util.MessageHelper;
 import org.whispersystems.textsecuregcm.tests.util.MessagesDynamoDbExtension;
+import reactor.core.publisher.Flux;
+import reactor.test.StepVerifier;
 
 class MessagesDynamoDbTest {
 
@@ -59,6 +71,7 @@ class MessagesDynamoDbTest {
     MESSAGE3 = builder.build();
   }
 
+  private ExecutorService messageDeletionExecutorService;
   private MessagesDynamoDb messagesDynamoDb;
 
 
@@ -67,8 +80,18 @@ class MessagesDynamoDbTest {
 
   @BeforeEach
   void setup() {
-    messagesDynamoDb = new MessagesDynamoDb(dynamoDbExtension.getDynamoDbClient(), MessagesDynamoDbExtension.TABLE_NAME,
+    messageDeletionExecutorService = Executors.newSingleThreadExecutor();
-        Duration.ofDays(14));
+    messagesDynamoDb = new MessagesDynamoDb(dynamoDbExtension.getDynamoDbClient(),
+        dynamoDbExtension.getDynamoDbAsyncClient(), MessagesDynamoDbExtension.TABLE_NAME, Duration.ofDays(14),
+        messageDeletionExecutorService);
+  }
+
+  @AfterEach
+  void teardown() throws Exception {
+    messageDeletionExecutorService.shutdown();
+    messageDeletionExecutorService.awaitTermination(5, TimeUnit.SECONDS);
+
+    StepVerifier.resetDefaultTimeout();
   }
 
   @Test
@@ -77,7 +100,7 @@ class MessagesDynamoDbTest {
     final int destinationDeviceId = random.nextInt(255) + 1;
     messagesDynamoDb.store(List.of(MESSAGE1, MESSAGE2, MESSAGE3), destinationUuid, destinationDeviceId);
 
-    final List<MessageProtos.Envelope> messagesStored = messagesDynamoDb.load(destinationUuid, destinationDeviceId,
+    final List<MessageProtos.Envelope> messagesStored = load(destinationUuid, destinationDeviceId,
         MessagesDynamoDb.RESULT_SET_CHUNK_SIZE);
     assertThat(messagesStored).isNotNull().hasSize(3);
     final MessageProtos.Envelope firstMessage =
@@ -88,6 +111,73 @@ class MessagesDynamoDbTest {
     assertThat(messagesStored).element(2).isEqualTo(MESSAGE2);
   }
 
+  @ParameterizedTest
+  @ValueSource(ints = {10, 100, 100, 1_000, 3_000})
+  void testLoadManyAfterInsert(final int messageCount) {
+    final UUID destinationUuid = UUID.randomUUID();
+    final int destinationDeviceId = random.nextInt(255) + 1;
+
+    final List<MessageProtos.Envelope> messages = new ArrayList<>(messageCount);
+    for (int i = 0; i < messageCount; i++) {
+      messages.add(MessageHelper.createMessage(UUID.randomUUID(), 1, destinationUuid, (i + 1L) * 1000, "message " + i));
+    }
+
+    messagesDynamoDb.store(messages, destinationUuid, destinationDeviceId);
+
+    final Publisher<?> fetchedMessages = messagesDynamoDb.load(destinationUuid, destinationDeviceId, null);
+
+    final long firstRequest = Math.min(10, messageCount);
+    StepVerifier.setDefaultTimeout(Duration.ofSeconds(15));
+
+    StepVerifier.Step<?> step = StepVerifier.create(fetchedMessages, 0)
+        .expectSubscription()
+        .thenRequest(firstRequest)
+        .expectNextCount(firstRequest);
+
+    if (messageCount > firstRequest) {
+      step = step.thenRequest(messageCount)
+          .expectNextCount(messageCount - firstRequest);
+    }
+
+    step.thenCancel()
+        .verify();
+  }
+
+  @Test
+  void testLimitedLoad() {
+    final int messageCount = 200;
+    final UUID destinationUuid = UUID.randomUUID();
+    final int destinationDeviceId = random.nextInt(255) + 1;
+
+    final List<MessageProtos.Envelope> messages = new ArrayList<>(messageCount);
+    for (int i = 0; i < messageCount; i++) {
+      messages.add(MessageHelper.createMessage(UUID.randomUUID(), 1, destinationUuid, (i + 1L) * 1000, "message " + i));
+    }
+
+    messagesDynamoDb.store(messages, destinationUuid, destinationDeviceId);
+
+    final int messageLoadLimit = 100;
+    final int halfOfMessageLoadLimit = messageLoadLimit / 2;
+    final Publisher<?> fetchedMessages = messagesDynamoDb.load(destinationUuid, destinationDeviceId, messageLoadLimit);
+
+    StepVerifier.setDefaultTimeout(Duration.ofSeconds(10));
+
+    final AtomicInteger messagesRemaining = new AtomicInteger(messageLoadLimit);
+
+    StepVerifier.create(fetchedMessages, 0)
+        .expectSubscription()
+        .thenRequest(halfOfMessageLoadLimit)
+        .expectNextCount(halfOfMessageLoadLimit)
+        // the first 100 should be fetched and buffered, but further requests should fail
+        .then(() -> dynamoDbExtension.stopServer())
+        .thenRequest(halfOfMessageLoadLimit)
+        .expectNextCount(halfOfMessageLoadLimit)
+        // we’ve consumed all the buffered messages, so a single request will fail
+        .thenRequest(1)
+        .expectError()
+        .verify();
+  }
+
   @Test
   void testDeleteForDestination() {
     final UUID destinationUuid = UUID.randomUUID();
@@ -96,18 +186,18 @@ class MessagesDynamoDbTest {
     messagesDynamoDb.store(List.of(MESSAGE2), secondDestinationUuid, 1);
     messagesDynamoDb.store(List.of(MESSAGE3), destinationUuid, 2);
 
-    assertThat(messagesDynamoDb.load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
+    assertThat(load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
         .element(0).isEqualTo(MESSAGE1);
-    assertThat(messagesDynamoDb.load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
+    assertThat(load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
         .element(0).isEqualTo(MESSAGE3);
-    assertThat(messagesDynamoDb.load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull()
+    assertThat(load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull()
         .hasSize(1).element(0).isEqualTo(MESSAGE2);
 
     messagesDynamoDb.deleteAllMessagesForAccount(destinationUuid);
 
-    assertThat(messagesDynamoDb.load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().isEmpty();
+    assertThat(load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().isEmpty();
-    assertThat(messagesDynamoDb.load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().isEmpty();
+    assertThat(load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().isEmpty();
-    assertThat(messagesDynamoDb.load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull()
+    assertThat(load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull()
         .hasSize(1).element(0).isEqualTo(MESSAGE2);
   }
 
@@ -119,71 +209,79 @@ class MessagesDynamoDbTest {
     messagesDynamoDb.store(List.of(MESSAGE2), secondDestinationUuid, 1);
     messagesDynamoDb.store(List.of(MESSAGE3), destinationUuid, 2);
 
-    assertThat(messagesDynamoDb.load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
+    assertThat(load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
         .element(0).isEqualTo(MESSAGE1);
-    assertThat(messagesDynamoDb.load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
+    assertThat(load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
         .element(0).isEqualTo(MESSAGE3);
-    assertThat(messagesDynamoDb.load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull()
+    assertThat(load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull()
         .hasSize(1).element(0).isEqualTo(MESSAGE2);
 
     messagesDynamoDb.deleteAllMessagesForDevice(destinationUuid, 2);
 
-    assertThat(messagesDynamoDb.load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
+    assertThat(load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
         .element(0).isEqualTo(MESSAGE1);
-    assertThat(messagesDynamoDb.load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().isEmpty();
+    assertThat(load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().isEmpty();
-    assertThat(messagesDynamoDb.load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull()
+    assertThat(load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull()
         .hasSize(1).element(0).isEqualTo(MESSAGE2);
   }
 
   @Test
-  void testDeleteMessageByDestinationAndGuid() {
+  void testDeleteMessageByDestinationAndGuid() throws Exception {
     final UUID destinationUuid = UUID.randomUUID();
     final UUID secondDestinationUuid = UUID.randomUUID();
     messagesDynamoDb.store(List.of(MESSAGE1), destinationUuid, 1);
     messagesDynamoDb.store(List.of(MESSAGE2), secondDestinationUuid, 1);
     messagesDynamoDb.store(List.of(MESSAGE3), destinationUuid, 2);
 
-    assertThat(messagesDynamoDb.load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
+    assertThat(load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
         .element(0).isEqualTo(MESSAGE1);
-    assertThat(messagesDynamoDb.load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
+    assertThat(load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
         .element(0).isEqualTo(MESSAGE3);
-    assertThat(messagesDynamoDb.load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull()
+    assertThat(load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull()
         .hasSize(1).element(0).isEqualTo(MESSAGE2);
 
     messagesDynamoDb.deleteMessageByDestinationAndGuid(secondDestinationUuid,
-        UUID.fromString(MESSAGE2.getServerGuid()));
+        UUID.fromString(MESSAGE2.getServerGuid())).get(5, TimeUnit.SECONDS);
 
-    assertThat(messagesDynamoDb.load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
+    assertThat(load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
         .element(0).isEqualTo(MESSAGE1);
-    assertThat(messagesDynamoDb.load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
+    assertThat(load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
         .element(0).isEqualTo(MESSAGE3);
-    assertThat(messagesDynamoDb.load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull()
+    assertThat(load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull()
         .isEmpty();
   }
 
   @Test
-  void testDeleteSingleMessage() {
+  void testDeleteSingleMessage() throws Exception {
     final UUID destinationUuid = UUID.randomUUID();
     final UUID secondDestinationUuid = UUID.randomUUID();
     messagesDynamoDb.store(List.of(MESSAGE1), destinationUuid, 1);
     messagesDynamoDb.store(List.of(MESSAGE2), secondDestinationUuid, 1);
     messagesDynamoDb.store(List.of(MESSAGE3), destinationUuid, 2);
 
-    assertThat(messagesDynamoDb.load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
+    assertThat(load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
         .element(0).isEqualTo(MESSAGE1);
-    assertThat(messagesDynamoDb.load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
+    assertThat(load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
         .element(0).isEqualTo(MESSAGE3);
-    assertThat(messagesDynamoDb.load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull()
+    assertThat(load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull()
         .hasSize(1).element(0).isEqualTo(MESSAGE2);
 
     messagesDynamoDb.deleteMessage(secondDestinationUuid, 1,
-        UUID.fromString(MESSAGE2.getServerGuid()), MESSAGE2.getServerTimestamp());
+        UUID.fromString(MESSAGE2.getServerGuid()), MESSAGE2.getServerTimestamp()).get(1, TimeUnit.SECONDS);
 
-    assertThat(messagesDynamoDb.load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
+    assertThat(load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
         .element(0).isEqualTo(MESSAGE1);
-    assertThat(messagesDynamoDb.load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
+    assertThat(load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().hasSize(1)
         .element(0).isEqualTo(MESSAGE3);
-    assertThat(messagesDynamoDb.load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull()
+    assertThat(load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull()
         .isEmpty();
   }
+
+  private List<MessageProtos.Envelope> load(final UUID destinationUuid, final long destinationDeviceId,
+      final int count) {
+    return Flux.from(messagesDynamoDb.load(destinationUuid, destinationDeviceId, count))
+        .take(count, true)
+        .collectList()
+        .block();
+  }
 }

service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagesManagerTest.java

@@ -14,13 +14,11 @@ import static org.mockito.Mockito.verifyNoMoreInteractions;
 import java.util.UUID;
 import org.junit.jupiter.api.Test;
 import org.whispersystems.textsecuregcm.entities.MessageProtos.Envelope;
-import org.whispersystems.textsecuregcm.push.PushLatencyManager;
 
 class MessagesManagerTest {
 
   private final MessagesDynamoDb messagesDynamoDb = mock(MessagesDynamoDb.class);
   private final MessagesCache messagesCache = mock(MessagesCache.class);
-  private final PushLatencyManager pushLatencyManager = mock(PushLatencyManager.class);
   private final ReportMessageManager reportMessageManager = mock(ReportMessageManager.class);
 
   private final MessagesManager messagesManager = new MessagesManager(messagesDynamoDb, messagesCache,

service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/ProfilesManagerTest.java

@@ -41,7 +41,7 @@ public class ProfilesManagerTest {
   void setUp() {
     //noinspection unchecked
     commands = mock(RedisAdvancedClusterCommands.class);
-    final FaultTolerantRedisCluster cacheCluster = RedisClusterHelper.buildMockRedisCluster(commands);
+    final FaultTolerantRedisCluster cacheCluster = RedisClusterHelper.builder().stringCommands(commands).build();
 
     profiles = mock(Profiles.class);
 

service/src/test/java/org/whispersystems/textsecuregcm/tests/util/MessageHelper.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2022 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.tests.util;
+
+import com.google.protobuf.ByteString;
+import java.nio.charset.StandardCharsets;
+import java.util.UUID;
+import org.whispersystems.textsecuregcm.entities.MessageProtos;
+
+public class MessageHelper {
+
+  public static MessageProtos.Envelope createMessage(UUID senderUuid, final int senderDeviceId, UUID destinationUuid,
+      long timestamp, String content) {
+    return MessageProtos.Envelope.newBuilder()
+        .setServerGuid(UUID.randomUUID().toString())
+        .setType(MessageProtos.Envelope.Type.CIPHERTEXT)
+        .setTimestamp(timestamp)
+        .setServerTimestamp(0)
+        .setSourceUuid(senderUuid.toString())
+        .setSourceDevice(senderDeviceId)
+        .setDestinationUuid(destinationUuid.toString())
+        .setContent(ByteString.copyFrom(content.getBytes(StandardCharsets.UTF_8)))
+        .build();
+  }
+}

service/src/test/java/org/whispersystems/textsecuregcm/tests/util/RedisClusterHelper.java

@@ -5,33 +5,39 @@
 
 package org.whispersystems.textsecuregcm.tests.util;
 
-import io.lettuce.core.cluster.api.StatefulRedisClusterConnection;
-import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands;
-import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
-
-import java.util.function.Consumer;
-import java.util.function.Function;
-
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import io.lettuce.core.cluster.api.StatefulRedisClusterConnection;
+import io.lettuce.core.cluster.api.async.RedisAdvancedClusterAsyncCommands;
+import io.lettuce.core.cluster.api.reactive.RedisAdvancedClusterReactiveCommands;
+import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
+
 public class RedisClusterHelper {
 
-    @SuppressWarnings("unchecked")
+  public static RedisClusterHelper.Builder builder() {
-    public static FaultTolerantRedisCluster buildMockRedisCluster(final RedisAdvancedClusterCommands<String, String> stringCommands) {
+    return new Builder();
-        return buildMockRedisCluster(stringCommands, mock(RedisAdvancedClusterCommands.class));
   }
 
   @SuppressWarnings("unchecked")
-    public static FaultTolerantRedisCluster buildMockRedisCluster(final RedisAdvancedClusterCommands<String, String> stringCommands, final RedisAdvancedClusterCommands<byte[], byte[]> binaryCommands) {
+  private static FaultTolerantRedisCluster buildMockRedisCluster(
+      final RedisAdvancedClusterCommands<String, String> stringCommands,
+      final RedisAdvancedClusterCommands<byte[], byte[]> binaryCommands,
+      final RedisAdvancedClusterAsyncCommands<byte[], byte[]> binaryAsyncCommands,
+      final RedisAdvancedClusterReactiveCommands<byte[], byte[]> binaryReactiveCommands) {
     final FaultTolerantRedisCluster cluster = mock(FaultTolerantRedisCluster.class);
     final StatefulRedisClusterConnection<String, String> stringConnection = mock(StatefulRedisClusterConnection.class);
     final StatefulRedisClusterConnection<byte[], byte[]> binaryConnection = mock(StatefulRedisClusterConnection.class);
 
     when(stringConnection.sync()).thenReturn(stringCommands);
     when(binaryConnection.sync()).thenReturn(binaryCommands);
+    when(binaryConnection.async()).thenReturn(binaryAsyncCommands);
+    when(binaryConnection.reactive()).thenReturn(binaryReactiveCommands);
 
     when(cluster.withCluster(any(Function.class))).thenAnswer(invocation -> {
       return invocation.getArgument(0, Function.class).apply(stringConnection);
@@ -71,4 +77,46 @@ public class RedisClusterHelper {
 
     return cluster;
   }
+
+  @SuppressWarnings("unchecked")
+  public static class Builder {
+
+    private RedisAdvancedClusterCommands<String, String> stringCommands = mock(RedisAdvancedClusterCommands.class);
+    private RedisAdvancedClusterCommands<byte[], byte[]> binaryCommands = mock(RedisAdvancedClusterCommands.class);
+    private RedisAdvancedClusterAsyncCommands<byte[], byte[]> binaryAsyncCommands = mock(
+        RedisAdvancedClusterAsyncCommands.class);
+    private RedisAdvancedClusterReactiveCommands<byte[], byte[]> binaryReactiveCommands = mock(
+        RedisAdvancedClusterReactiveCommands.class);
+
+    private Builder() {
+
+    }
+
+    public Builder stringCommands(final RedisAdvancedClusterCommands<String, String> stringCommands) {
+      this.stringCommands = stringCommands;
+      return this;
+    }
+
+    public Builder binaryCommands(final RedisAdvancedClusterCommands<byte[], byte[]> binaryCommands) {
+      this.binaryCommands = binaryCommands;
+      return this;
+    }
+
+    public Builder binaryAsyncCommands(final RedisAdvancedClusterAsyncCommands<byte[], byte[]> binaryAsyncCommands) {
+      this.binaryAsyncCommands = binaryAsyncCommands;
+      return this;
+    }
+
+    public Builder binaryReactiveCommands(
+        final RedisAdvancedClusterReactiveCommands<byte[], byte[]> binaryReactiveCommands) {
+      this.binaryReactiveCommands = binaryReactiveCommands;
+      return this;
+    }
+
+    public FaultTolerantRedisCluster build() {
+      return RedisClusterHelper.buildMockRedisCluster(stringCommands, binaryCommands, binaryAsyncCommands,
+          binaryReactiveCommands);
+    }
+  }
+
 }

service/src/test/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnectionIntegrationTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2021 Signal Messenger, LLC
+ * Copyright 2013-2022 Signal Messenger, LLC
  * SPDX-License-Identifier: AGPL-3.0-only
  */
 
@@ -22,6 +22,7 @@ import static org.mockito.Mockito.when;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.InvalidProtocolBufferException;
 import java.io.IOException;
+import java.time.Clock;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
@@ -36,8 +37,10 @@ import java.util.concurrent.atomic.AtomicBoolean;
 import org.apache.commons.lang3.RandomStringUtils;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.ArgumentCaptor;
 import org.mockito.stubbing.Answer;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedAccount;
@@ -56,6 +59,7 @@ import org.whispersystems.textsecuregcm.tests.util.MessagesDynamoDbExtension;
 import org.whispersystems.textsecuregcm.util.Pair;
 import org.whispersystems.websocket.WebSocketClient;
 import org.whispersystems.websocket.messages.WebSocketResponseMessage;
+import reactor.core.scheduler.Schedulers;
 
 class WebSocketConnectionIntegrationTest {
 
@@ -65,16 +69,13 @@ class WebSocketConnectionIntegrationTest {
   @RegisterExtension
   static final RedisClusterExtension REDIS_CLUSTER_EXTENSION = RedisClusterExtension.builder().build();
 
-  private static final int SEND_FUTURES_TIMEOUT_MILLIS = 100;
+  private ExecutorService sharedExecutorService;
-
-  private ExecutorService executorService;
   private MessagesDynamoDb messagesDynamoDb;
   private MessagesCache messagesCache;
   private ReportMessageManager reportMessageManager;
   private Account account;
   private Device device;
   private WebSocketClient webSocketClient;
-  private WebSocketConnection webSocketConnection;
   private ScheduledExecutorService retrySchedulingExecutor;
 
   private long serialTimestamp = System.currentTimeMillis();
@@ -82,11 +83,12 @@ class WebSocketConnectionIntegrationTest {
   @BeforeEach
   void setUp() throws Exception {
 
-    executorService = Executors.newSingleThreadExecutor();
+    sharedExecutorService = Executors.newSingleThreadExecutor();
     messagesCache = new MessagesCache(REDIS_CLUSTER_EXTENSION.getRedisCluster(),
-        REDIS_CLUSTER_EXTENSION.getRedisCluster(), executorService);
+        REDIS_CLUSTER_EXTENSION.getRedisCluster(), Clock.systemUTC(), sharedExecutorService, sharedExecutorService);
-    messagesDynamoDb = new MessagesDynamoDb(dynamoDbExtension.getDynamoDbClient(), MessagesDynamoDbExtension.TABLE_NAME,
+    messagesDynamoDb = new MessagesDynamoDb(dynamoDbExtension.getDynamoDbClient(),
-        Duration.ofDays(7));
+        dynamoDbExtension.getDynamoDbAsyncClient(), MessagesDynamoDbExtension.TABLE_NAME, Duration.ofDays(7),
+        sharedExecutorService);
     reportMessageManager = mock(ReportMessageManager.class);
     account = mock(Account.class);
     device = mock(Device.class);
@@ -96,30 +98,36 @@ class WebSocketConnectionIntegrationTest {
     when(account.getNumber()).thenReturn("+18005551234");
     when(account.getUuid()).thenReturn(UUID.randomUUID());
     when(device.getId()).thenReturn(1L);
-
-    webSocketConnection = new WebSocketConnection(
-        mock(ReceiptSender.class),
-        new MessagesManager(messagesDynamoDb, messagesCache, reportMessageManager),
-        new AuthenticatedAccount(() -> new Pair<>(account, device)),
-        device,
-        webSocketClient,
-        SEND_FUTURES_TIMEOUT_MILLIS,
-        retrySchedulingExecutor);
   }
 
   @AfterEach
   void tearDown() throws Exception {
-    executorService.shutdown();
+    sharedExecutorService.shutdown();
-    executorService.awaitTermination(2, TimeUnit.SECONDS);
+    sharedExecutorService.awaitTermination(2, TimeUnit.SECONDS);
 
     retrySchedulingExecutor.shutdown();
     retrySchedulingExecutor.awaitTermination(2, TimeUnit.SECONDS);
   }
 
-  @Test
+  @ParameterizedTest
-  void testProcessStoredMessages() {
+  @CsvSource({
-    final int persistedMessageCount = 207;
+      "207, 173, true",
-    final int cachedMessageCount = 173;
+      "207, 173, false",
+      "323, 0, true",
+      "323, 0, false",
+      "0, 221, true",
+      "0, 221, false",
+  })
+  void testProcessStoredMessages(final int persistedMessageCount, final int cachedMessageCount,
+      final boolean useReactive) {
+    final WebSocketConnection webSocketConnection = new WebSocketConnection(
+        mock(ReceiptSender.class),
+        new MessagesManager(messagesDynamoDb, messagesCache, reportMessageManager),
+        new AuthenticatedAccount(() -> new Pair<>(account, device)),
+        device,
+        webSocketClient,
+        retrySchedulingExecutor,
+        useReactive);
 
     final List<MessageProtos.Envelope> expectedMessages = new ArrayList<>(persistedMessageCount + cachedMessageCount);
 
@@ -150,8 +158,8 @@ class WebSocketConnectionIntegrationTest {
       final AtomicBoolean queueCleared = new AtomicBoolean(false);
 
       when(successResponse.getStatus()).thenReturn(200);
-      when(webSocketClient.sendRequest(eq("PUT"), eq("/api/v1/message"), anyList(), any())).thenReturn(
+      when(webSocketClient.sendRequest(eq("PUT"), eq("/api/v1/message"), anyList(), any()))
-          CompletableFuture.completedFuture(successResponse));
+          .thenReturn(CompletableFuture.completedFuture(successResponse));
 
       when(webSocketClient.sendRequest(eq("PUT"), eq("/api/v1/queue/empty"), anyList(), any())).thenAnswer(
           (Answer<CompletableFuture<WebSocketResponseMessage>>) invocation -> {
@@ -194,8 +202,18 @@ class WebSocketConnectionIntegrationTest {
     });
   }
 
-  @Test
+  @ParameterizedTest
-  void testProcessStoredMessagesClientClosed() {
+  @ValueSource(booleans = {true, false})
+  void testProcessStoredMessagesClientClosed(final boolean useReactive) {
+    final WebSocketConnection webSocketConnection = new WebSocketConnection(
+        mock(ReceiptSender.class),
+        new MessagesManager(messagesDynamoDb, messagesCache, reportMessageManager),
+        new AuthenticatedAccount(() -> new Pair<>(account, device)),
+        device,
+        webSocketClient,
+        retrySchedulingExecutor,
+        useReactive);
+
     final int persistedMessageCount = 207;
     final int cachedMessageCount = 173;
 
@@ -250,8 +268,20 @@ class WebSocketConnectionIntegrationTest {
     });
   }
 
-  @Test
+  @ParameterizedTest
-  void testProcessStoredMessagesSendFutureTimeout() {
+  @ValueSource(booleans = {true, false})
+  void testProcessStoredMessagesSendFutureTimeout(final boolean useReactive) {
+    final WebSocketConnection webSocketConnection = new WebSocketConnection(
+        mock(ReceiptSender.class),
+        new MessagesManager(messagesDynamoDb, messagesCache, reportMessageManager),
+        new AuthenticatedAccount(() -> new Pair<>(account, device)),
+        device,
+        webSocketClient,
+        100, // use a very short timeout, so that this test completes quickly
+        retrySchedulingExecutor,
+        useReactive,
+        Schedulers.boundedElastic());
+
     final int persistedMessageCount = 207;
     final int cachedMessageCount = 173;
 
@@ -346,4 +376,5 @@ class WebSocketConnectionIntegrationTest {
         .setDestinationUuid(UUID.randomUUID().toString())
         .build();
   }
+
 }

service/src/test/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnectionTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2021 Signal Messenger, LLC
+ * Copyright 2013-2022 Signal Messenger, LLC
  * SPDX-License-Identifier: AGPL-3.0-only
  */
 
@@ -12,6 +12,7 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.nullable;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.anyInt;
 import static org.mockito.Mockito.anyLong;
@@ -42,17 +43,20 @@ import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
-import org.apache.commons.lang3.RandomStringUtils;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Stream;
 import org.eclipse.jetty.websocket.api.UpgradeRequest;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
-import org.mockito.ArgumentMatchers;
+import org.junit.jupiter.params.ParameterizedTest;
-import org.mockito.invocation.InvocationOnMock;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.stubbing.Answer;
 import org.whispersystems.textsecuregcm.auth.AccountAuthenticator;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedAccount;
-import org.whispersystems.textsecuregcm.entities.OutgoingMessageEntityList;
+import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
-import org.whispersystems.textsecuregcm.push.ApnPushNotificationScheduler;
 import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
 import org.whispersystems.textsecuregcm.push.PushNotificationManager;
 import org.whispersystems.textsecuregcm.push.ReceiptSender;
@@ -65,6 +69,10 @@ import org.whispersystems.websocket.WebSocketClient;
 import org.whispersystems.websocket.auth.WebSocketAuthenticator.AuthenticationResult;
 import org.whispersystems.websocket.messages.WebSocketResponseMessage;
 import org.whispersystems.websocket.session.WebSocketSessionContext;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.FluxSink;
+import reactor.core.scheduler.Schedulers;
+import reactor.test.StepVerifier;
 
 class WebSocketConnectionTest {
 
@@ -83,7 +91,6 @@ class WebSocketConnectionTest {
   private AuthenticatedAccount auth;
   private UpgradeRequest upgradeRequest;
   private ReceiptSender receiptSender;
-  private PushNotificationManager pushNotificationManager;
   private ScheduledExecutorService retrySchedulingExecutor;
 
   @BeforeEach
@@ -95,17 +102,21 @@ class WebSocketConnectionTest {
     auth = new AuthenticatedAccount(() -> new Pair<>(account, device));
     upgradeRequest = mock(UpgradeRequest.class);
     receiptSender = mock(ReceiptSender.class);
-    pushNotificationManager = mock(PushNotificationManager.class);
     retrySchedulingExecutor = mock(ScheduledExecutorService.class);
   }
 
+  @AfterEach
+  void teardown() {
+    StepVerifier.resetDefaultTimeout();
+  }
+
   @Test
   void testCredentials() {
     MessagesManager storedMessages = mock(MessagesManager.class);
     WebSocketAccountAuthenticator webSocketAuthenticator = new WebSocketAccountAuthenticator(accountAuthenticator);
     AuthenticatedConnectListener connectListener = new AuthenticatedConnectListener(receiptSender, storedMessages,
         mock(PushNotificationManager.class), mock(ClientPresenceManager.class),
-        retrySchedulingExecutor);
+        retrySchedulingExecutor, mock(ExperimentEnrollmentManager.class));
     WebSocketSessionContext sessionContext = mock(WebSocketSessionContext.class);
 
     when(accountAuthenticator.authenticate(eq(new BasicCredentials(VALID_USER, VALID_PASSWORD))))
@@ -114,7 +125,6 @@ class WebSocketConnectionTest {
     when(accountAuthenticator.authenticate(eq(new BasicCredentials(INVALID_USER, INVALID_PASSWORD))))
         .thenReturn(Optional.empty());
 
-
     when(upgradeRequest.getParameterMap()).thenReturn(Map.of(
         "login", List.of(VALID_USER),
         "password", List.of(VALID_PASSWORD)));
@@ -136,8 +146,9 @@ class WebSocketConnectionTest {
     assertTrue(account.isRequired());
   }
 
-  @Test
+  @ParameterizedTest
-  void testOpen() throws Exception {
+  @ValueSource(booleans = {true, false})
+  void testOpen(final boolean useReactive) throws Exception {
     MessagesManager storedMessages = mock(MessagesManager.class);
 
     UUID accountUuid = UUID.randomUUID();
@@ -166,29 +177,31 @@ class WebSocketConnectionTest {
 
     String userAgent = "user-agent";
 
+    if (useReactive) {
+      when(storedMessages.getMessagesForDeviceReactive(account.getUuid(), device.getId(), false))
+          .thenReturn(Flux.fromIterable(outgoingMessages));
+    } else {
       when(storedMessages.getMessagesForDevice(account.getUuid(), device.getId(), false))
           .thenReturn(new Pair<>(outgoingMessages, false));
+    }
 
     final List<CompletableFuture<WebSocketResponseMessage>> futures = new LinkedList<>();
     final WebSocketClient client = mock(WebSocketClient.class);
 
     when(client.getUserAgent()).thenReturn(userAgent);
-    when(client.sendRequest(eq("PUT"), eq("/api/v1/message"), ArgumentMatchers.nullable(List.class), ArgumentMatchers.<Optional<byte[]>>any()))
+    when(client.sendRequest(eq("PUT"), eq("/api/v1/message"), nullable(List.class), any()))
-        .thenAnswer(new Answer<CompletableFuture<WebSocketResponseMessage>>() {
+        .thenAnswer(invocation -> {
-          @Override
-          public CompletableFuture<WebSocketResponseMessage> answer(InvocationOnMock invocationOnMock) {
           CompletableFuture<WebSocketResponseMessage> future = new CompletableFuture<>();
           futures.add(future);
           return future;
-          }
         });
 
     WebSocketConnection connection = new WebSocketConnection(receiptSender, storedMessages,
-        auth, device, client, retrySchedulingExecutor);
+        auth, device, client, retrySchedulingExecutor, useReactive, Schedulers.immediate());
 
     connection.start();
     verify(client, times(3)).sendRequest(eq("PUT"), eq("/api/v1/message"), any(List.class),
-        ArgumentMatchers.<Optional<byte[]>>any());
+        any());
 
     assertEquals(3, futures.size());
 
@@ -208,12 +221,13 @@ class WebSocketConnectionTest {
     verify(client).close(anyInt(), anyString());
   }
 
-  @Test
+  @ParameterizedTest
-  public void testOnlineSend() {
+  @ValueSource(booleans = {true, false})
+  public void testOnlineSend(final boolean useReactive) {
     final MessagesManager messagesManager = mock(MessagesManager.class);
     final WebSocketClient client = mock(WebSocketClient.class);
     final WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor);
+        retrySchedulingExecutor, useReactive, Schedulers.immediate());
 
     final UUID accountUuid = UUID.randomUUID();
 
@@ -222,17 +236,29 @@ class WebSocketConnectionTest {
     when(device.getId()).thenReturn(1L);
     when(client.isOpen()).thenReturn(true);
 
+    if (useReactive) {
+      when(messagesManager.getMessagesForDeviceReactive(eq(accountUuid), eq(1L), anyBoolean()))
+          .thenReturn(Flux.empty())
+          .thenReturn(Flux.just(createMessage(UUID.randomUUID(), UUID.randomUUID(), 1111, "first")))
+          .thenReturn(Flux.just(createMessage(UUID.randomUUID(), UUID.randomUUID(), 2222, "second")))
+          .thenReturn(Flux.empty());
+    } else {
       when(messagesManager.getMessagesForDevice(eq(accountUuid), eq(1L), anyBoolean()))
           .thenReturn(new Pair<>(Collections.emptyList(), false))
-        .thenReturn(new Pair<>(List.of(createMessage(UUID.randomUUID(), UUID.randomUUID(), 1111, "first")), false))
+          .thenReturn(new Pair<>(List.of(createMessage(UUID.randomUUID(), UUID.randomUUID(), 1111, "first")),
-        .thenReturn(new Pair<>(List.of(createMessage(UUID.randomUUID(), UUID.randomUUID(), 2222, "second")), false));
+              false))
+          .thenReturn(new Pair<>(List.of(createMessage(UUID.randomUUID(), UUID.randomUUID(), 2222, "second")),
+              false))
+          .thenReturn(new Pair<>(Collections.emptyList(), false));
+    }
 
     final WebSocketResponseMessage successResponse = mock(WebSocketResponseMessage.class);
     when(successResponse.getStatus()).thenReturn(200);
 
     final AtomicInteger sendCounter = new AtomicInteger(0);
 
-    when(client.sendRequest(eq("PUT"), eq("/api/v1/message"), any(List.class), any(Optional.class))).thenAnswer((Answer<CompletableFuture<WebSocketResponseMessage>>)invocation -> {
+    when(client.sendRequest(eq("PUT"), eq("/api/v1/message"), any(List.class), any(Optional.class)))
+        .thenAnswer(invocation -> {
           synchronized (sendCounter) {
             sendCounter.incrementAndGet();
             sendCounter.notifyAll();
@@ -269,8 +295,9 @@ class WebSocketConnectionTest {
     verify(client, times(2)).sendRequest(eq("PUT"), eq("/api/v1/message"), any(List.class), any(Optional.class));
   }
 
-  @Test
+  @ParameterizedTest
-  void testPendingSend() throws Exception {
+  @ValueSource(booleans = {true, false})
+  void testPendingSend(final boolean useReactive) throws Exception {
     MessagesManager storedMessages = mock(MessagesManager.class);
 
     final UUID accountUuid = UUID.randomUUID();
@@ -311,12 +338,17 @@ class WebSocketConnectionTest {
     when(sender1.getDevices()).thenReturn(sender1devices);
 
     when(accountsManager.getByE164("sender1")).thenReturn(Optional.of(sender1));
-    when(accountsManager.getByE164("sender2")).thenReturn(Optional.<Account>empty());
+    when(accountsManager.getByE164("sender2")).thenReturn(Optional.empty());
 
     String userAgent = "user-agent";
 
+    if (useReactive) {
+      when(storedMessages.getMessagesForDeviceReactive(account.getUuid(), device.getId(), false))
+          .thenReturn(Flux.fromIterable(pendingMessages));
+    } else {
       when(storedMessages.getMessagesForDevice(account.getUuid(), device.getId(), false))
           .thenReturn(new Pair<>(pendingMessages, false));
+    }
 
     final List<CompletableFuture<WebSocketResponseMessage>> futures = new LinkedList<>();
     final WebSocketClient client = mock(WebSocketClient.class);
@@ -330,7 +362,7 @@ class WebSocketConnectionTest {
         });
 
     WebSocketConnection connection = new WebSocketConnection(receiptSender, storedMessages,
-        auth, device, client, retrySchedulingExecutor);
+        auth, device, client, retrySchedulingExecutor, useReactive, Schedulers.immediate());
 
     connection.start();
 
@@ -350,12 +382,13 @@ class WebSocketConnectionTest {
     verify(client).close(anyInt(), anyString());
   }
 
-  @Test
+  @ParameterizedTest
-  void testProcessStoredMessageConcurrency() throws InterruptedException {
+  @ValueSource(booleans = {true, false})
+  void testProcessStoredMessageConcurrency(final boolean useReactive) {
     final MessagesManager messagesManager = mock(MessagesManager.class);
     final WebSocketClient client = mock(WebSocketClient.class);
     final WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor);
+        retrySchedulingExecutor, useReactive, Schedulers.immediate());
 
     when(account.getNumber()).thenReturn("+18005551234");
     when(account.getUuid()).thenReturn(UUID.randomUUID());
@@ -365,8 +398,10 @@ class WebSocketConnectionTest {
     final AtomicBoolean threadWaiting = new AtomicBoolean(false);
     final AtomicBoolean returnMessageList = new AtomicBoolean(false);
 
-    when(messagesManager.getMessagesForDevice(account.getUuid(), 1L, false)).thenAnswer(
+    if (useReactive) {
-        (Answer<OutgoingMessageEntityList>) invocation -> {
+      when(
+          messagesManager.getMessagesForDeviceReactive(account.getUuid(), 1L, false))
+          .thenAnswer(invocation -> {
             synchronized (threadWaiting) {
               threadWaiting.set(true);
               threadWaiting.notifyAll();
@@ -378,13 +413,30 @@ class WebSocketConnectionTest {
               }
             }
 
-      return new OutgoingMessageEntityList(Collections.emptyList(), false);
+            return Flux.empty();
           });
+    } else {
+      when(
+          messagesManager.getMessagesForDevice(account.getUuid(), 1L, false))
+          .thenAnswer(invocation -> {
+            synchronized (threadWaiting) {
+              threadWaiting.set(true);
+              threadWaiting.notifyAll();
+            }
+
+            synchronized (returnMessageList) {
+              while (!returnMessageList.get()) {
+                returnMessageList.wait();
+              }
+            }
+
+            return new Pair<>(Collections.emptyList(), false);
+          });
+    }
 
     final Thread[] threads = new Thread[10];
     final CountDownLatch unblockedThreadsLatch = new CountDownLatch(threads.length - 1);
 
-
     assertTimeoutPreemptively(Duration.ofSeconds(5), () -> {
       for (int i = 0; i < threads.length; i++) {
         threads[i] = new Thread(() -> {
@@ -413,18 +465,24 @@ class WebSocketConnectionTest {
       }
     });
 
+    if (useReactive) {
+      verify(messagesManager).getMessagesForDeviceReactive(any(UUID.class), anyLong(), eq(false));
+    } else {
       verify(messagesManager).getMessagesForDevice(any(UUID.class), anyLong(), eq(false));
     }
+  }
 
-  @Test
+  @ParameterizedTest
-  void testProcessStoredMessagesMultiplePages() {
+  @ValueSource(booleans = {true, false})
+  void testProcessStoredMessagesMultiplePages(final boolean useReactive) {
     final MessagesManager messagesManager = mock(MessagesManager.class);
     final WebSocketClient client = mock(WebSocketClient.class);
     final WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor);
+        retrySchedulingExecutor, useReactive, Schedulers.immediate());
 
     when(account.getNumber()).thenReturn("+18005551234");
-    when(account.getUuid()).thenReturn(UUID.randomUUID());
+    final UUID accountUuid = UUID.randomUUID();
+    when(account.getUuid()).thenReturn(accountUuid);
     when(device.getId()).thenReturn(1L);
     when(client.isOpen()).thenReturn(true);
 
@@ -435,39 +493,56 @@ class WebSocketConnectionTest {
     final List<Envelope> secondPageMessages =
         List.of(createMessage(UUID.randomUUID(), UUID.randomUUID(), 3333, "third"));
 
-    when(messagesManager.getMessagesForDevice(account.getUuid(), 1L, false))
+    if (useReactive) {
+      when(messagesManager.getMessagesForDeviceReactive(eq(accountUuid), eq(1L), eq(false)))
+          .thenReturn(Flux.fromStream(Stream.concat(firstPageMessages.stream(), secondPageMessages.stream())));
+    } else {
+      when(messagesManager.getMessagesForDevice(eq(accountUuid), eq(1L), eq(false)))
           .thenReturn(new Pair<>(firstPageMessages, true))
           .thenReturn(new Pair<>(secondPageMessages, false));
+    }
+
+    when(messagesManager.delete(eq(accountUuid), eq(1L), any(), any()))
+        .thenReturn(CompletableFuture.completedFuture(null));
 
     final WebSocketResponseMessage successResponse = mock(WebSocketResponseMessage.class);
     when(successResponse.getStatus()).thenReturn(200);
 
-    final CountDownLatch sendLatch = new CountDownLatch(firstPageMessages.size() + secondPageMessages.size());
+    final CountDownLatch queueEmptyLatch = new CountDownLatch(1);
 
-    when(client.sendRequest(eq("PUT"), eq("/api/v1/message"), any(List.class), any(Optional.class))).thenAnswer((Answer<CompletableFuture<WebSocketResponseMessage>>)invocation -> {
+    when(client.sendRequest(eq("PUT"), eq("/api/v1/message"), any(List.class), any(Optional.class)))
-      sendLatch.countDown();
+        .thenAnswer(invocation -> {
+          return CompletableFuture.completedFuture(successResponse);
+        });
+
+    when(client.sendRequest(eq("PUT"), eq("/api/v1/queue/empty"), any(List.class), eq(Optional.empty())))
+        .thenAnswer(invocation -> {
+          queueEmptyLatch.countDown();
           return CompletableFuture.completedFuture(successResponse);
         });
 
     assertTimeoutPreemptively(Duration.ofSeconds(5), () -> {
       connection.processStoredMessages();
 
-      sendLatch.await();
+      queueEmptyLatch.await();
     });
 
-    verify(client, times(firstPageMessages.size() + secondPageMessages.size())).sendRequest(eq("PUT"), eq("/api/v1/message"), any(List.class), any(Optional.class));
+    verify(client, times(firstPageMessages.size() + secondPageMessages.size())).sendRequest(eq("PUT"),
+        eq("/api/v1/message"), any(List.class), any(Optional.class));
     verify(client).sendRequest(eq("PUT"), eq("/api/v1/queue/empty"), any(List.class), eq(Optional.empty()));
   }
 
-  @Test
+  @ParameterizedTest
-  void testProcessStoredMessagesContainsSenderUuid() {
+  @ValueSource(booleans = {true, false})
+  void testProcessStoredMessagesContainsSenderUuid(final boolean useReactive) {
     final MessagesManager messagesManager = mock(MessagesManager.class);
     final WebSocketClient client = mock(WebSocketClient.class);
     final WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor);
+        retrySchedulingExecutor, useReactive, Schedulers.immediate());
 
     when(account.getNumber()).thenReturn("+18005551234");
-    when(account.getUuid()).thenReturn(UUID.randomUUID());
+    final UUID accountUuid = UUID.randomUUID();
+    when(account.getUuid()).thenReturn(accountUuid);
     when(device.getId()).thenReturn(1L);
     when(client.isOpen()).thenReturn(true);
 
@@ -475,26 +550,40 @@ class WebSocketConnectionTest {
     final List<Envelope> messages = List.of(
         createMessage(senderUuid, UUID.randomUUID(), 1111L, "message the first"));
 
+    if (useReactive) {
+      when(messagesManager.getMessagesForDeviceReactive(account.getUuid(), 1L, false))
+          .thenReturn(Flux.fromIterable(messages))
+          .thenReturn(Flux.empty());
+    } else {
       when(messagesManager.getMessagesForDevice(account.getUuid(), 1L, false))
-        .thenReturn(new Pair<>(messages, false));
+          .thenReturn(new Pair<>(messages, false))
+          .thenReturn(new Pair<>(Collections.emptyList(), false));
+    }
+
+    when(messagesManager.delete(eq(accountUuid), eq(1L), any(UUID.class), any()))
+        .thenReturn(CompletableFuture.completedFuture(null));
 
     final WebSocketResponseMessage successResponse = mock(WebSocketResponseMessage.class);
     when(successResponse.getStatus()).thenReturn(200);
 
-    final CountDownLatch sendLatch = new CountDownLatch(messages.size());
+    final CountDownLatch queueEmptyLatch = new CountDownLatch(1);
 
-    when(client.sendRequest(eq("PUT"), eq("/api/v1/message"), any(List.class), any(Optional.class))).thenAnswer(invocation -> {
+    when(client.sendRequest(eq("PUT"), eq("/api/v1/message"), any(List.class), any(Optional.class))).thenAnswer(
-      sendLatch.countDown();
+        invocation -> CompletableFuture.completedFuture(successResponse));
+
+    when(client.sendRequest(eq("PUT"), eq("/api/v1/queue/empty"), any(List.class), eq(Optional.empty())))
+        .thenAnswer(invocation -> {
+          queueEmptyLatch.countDown();
           return CompletableFuture.completedFuture(successResponse);
         });
 
     assertTimeoutPreemptively(Duration.ofSeconds(5), () -> {
       connection.processStoredMessages();
-
+      queueEmptyLatch.await();
-      sendLatch.await();
     });
 
-    verify(client, times(messages.size())).sendRequest(eq("PUT"), eq("/api/v1/message"), any(List.class), argThat(argument -> {
+    verify(client, times(messages.size())).sendRequest(eq("PUT"), eq("/api/v1/message"), any(List.class),
+        argThat(argument -> {
           if (argument.isEmpty()) {
             return false;
           }
@@ -513,12 +602,13 @@ class WebSocketConnectionTest {
     verify(client).sendRequest(eq("PUT"), eq("/api/v1/queue/empty"), any(List.class), eq(Optional.empty()));
   }
 
-  @Test
+  @ParameterizedTest
-  void testProcessStoredMessagesSingleEmptyCall() {
+  @ValueSource(booleans = {true, false})
+  void testProcessStoredMessagesSingleEmptyCall(final boolean useReactive) {
     final MessagesManager messagesManager = mock(MessagesManager.class);
     final WebSocketClient client = mock(WebSocketClient.class);
     final WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor);
+        retrySchedulingExecutor, useReactive, Schedulers.immediate());
 
     final UUID accountUuid = UUID.randomUUID();
 
@@ -527,8 +617,13 @@ class WebSocketConnectionTest {
     when(device.getId()).thenReturn(1L);
     when(client.isOpen()).thenReturn(true);
 
+    if (useReactive) {
+      when(messagesManager.getMessagesForDeviceReactive(eq(accountUuid), eq(1L), anyBoolean()))
+          .thenReturn(Flux.empty());
+    } else {
       when(messagesManager.getMessagesForDevice(eq(accountUuid), eq(1L), anyBoolean()))
           .thenReturn(new Pair<>(Collections.emptyList(), false));
+    }
 
     final WebSocketResponseMessage successResponse = mock(WebSocketResponseMessage.class);
     when(successResponse.getStatus()).thenReturn(200);
@@ -543,12 +638,13 @@ class WebSocketConnectionTest {
     verify(client, times(1)).sendRequest(eq("PUT"), eq("/api/v1/queue/empty"), any(List.class), eq(Optional.empty()));
   }
 
-  @Test
+  @ParameterizedTest
-  public void testRequeryOnStateMismatch() {
+  @ValueSource(booleans = {true, false})
+  public void testRequeryOnStateMismatch(final boolean useReactive) {
     final MessagesManager messagesManager = mock(MessagesManager.class);
     final WebSocketClient client = mock(WebSocketClient.class);
     final WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor);
+        retrySchedulingExecutor, useReactive, Schedulers.immediate());
     final UUID accountUuid = UUID.randomUUID();
 
     when(account.getNumber()).thenReturn("+18005551234");
@@ -563,39 +659,57 @@ class WebSocketConnectionTest {
     final List<Envelope> secondPageMessages =
         List.of(createMessage(UUID.randomUUID(), UUID.randomUUID(), 3333, "third"));
 
+    if (useReactive) {
+      when(messagesManager.getMessagesForDeviceReactive(eq(accountUuid), eq(1L), anyBoolean()))
+          .thenReturn(Flux.fromIterable(firstPageMessages))
+          .thenReturn(Flux.fromIterable(secondPageMessages))
+          .thenReturn(Flux.empty());
+    } else {
       when(messagesManager.getMessagesForDevice(eq(accountUuid), eq(1L), anyBoolean()))
           .thenReturn(new Pair<>(firstPageMessages, false))
           .thenReturn(new Pair<>(secondPageMessages, false))
           .thenReturn(new Pair<>(Collections.emptyList(), false));
+    }
+
+    when(messagesManager.delete(eq(accountUuid), eq(1L), any(), any()))
+        .thenReturn(CompletableFuture.completedFuture(null));
 
     final WebSocketResponseMessage successResponse = mock(WebSocketResponseMessage.class);
     when(successResponse.getStatus()).thenReturn(200);
 
-    final CountDownLatch sendLatch = new CountDownLatch(firstPageMessages.size() + secondPageMessages.size());
+    final CountDownLatch queueEmptyLatch = new CountDownLatch(1);
 
-    when(client.sendRequest(eq("PUT"), eq("/api/v1/message"), any(List.class), any(Optional.class))).thenAnswer((Answer<CompletableFuture<WebSocketResponseMessage>>)invocation -> {
+    when(client.sendRequest(eq("PUT"), eq("/api/v1/message"), any(List.class), any(Optional.class)))
+        .thenAnswer(invocation -> {
           connection.handleNewMessagesAvailable();
-      sendLatch.countDown();
 
           return CompletableFuture.completedFuture(successResponse);
         });
 
+    when(client.sendRequest(eq("PUT"), eq("/api/v1/queue/empty"), any(List.class), eq(Optional.empty())))
+        .thenAnswer(invocation -> {
+          queueEmptyLatch.countDown();
+          return CompletableFuture.completedFuture(successResponse);
+        });
+
     assertTimeoutPreemptively(Duration.ofSeconds(5), () -> {
       connection.processStoredMessages();
 
-      sendLatch.await();
+      queueEmptyLatch.await();
     });
 
-    verify(client, times(firstPageMessages.size() + secondPageMessages.size())).sendRequest(eq("PUT"), eq("/api/v1/message"), any(List.class), any(Optional.class));
+    verify(client, times(firstPageMessages.size() + secondPageMessages.size())).sendRequest(eq("PUT"),
+        eq("/api/v1/message"), any(List.class), any(Optional.class));
     verify(client).sendRequest(eq("PUT"), eq("/api/v1/queue/empty"), any(List.class), eq(Optional.empty()));
   }
 
-  @Test
+  @ParameterizedTest
-  void testProcessCachedMessagesOnly() {
+  @ValueSource(booleans = {true, false})
+  void testProcessCachedMessagesOnly(final boolean useReactive) {
     final MessagesManager messagesManager = mock(MessagesManager.class);
     final WebSocketClient client = mock(WebSocketClient.class);
     final WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor);
+        retrySchedulingExecutor, useReactive, Schedulers.immediate());
 
     final UUID accountUuid = UUID.randomUUID();
 
@@ -604,8 +718,13 @@ class WebSocketConnectionTest {
     when(device.getId()).thenReturn(1L);
     when(client.isOpen()).thenReturn(true);
 
+    if (useReactive) {
+      when(messagesManager.getMessagesForDeviceReactive(eq(accountUuid), eq(1L), anyBoolean()))
+          .thenReturn(Flux.empty());
+    } else {
       when(messagesManager.getMessagesForDevice(eq(accountUuid), eq(1L), anyBoolean()))
           .thenReturn(new Pair<>(Collections.emptyList(), false));
+    }
 
     final WebSocketResponseMessage successResponse = mock(WebSocketResponseMessage.class);
     when(successResponse.getStatus()).thenReturn(200);
@@ -616,19 +735,28 @@ class WebSocketConnectionTest {
     // anything.
     connection.processStoredMessages();
 
+    if (useReactive) {
+      verify(messagesManager).getMessagesForDeviceReactive(account.getUuid(), device.getId(), false);
+    } else {
       verify(messagesManager).getMessagesForDevice(account.getUuid(), device.getId(), false);
+    }
 
     connection.handleNewMessagesAvailable();
 
+    if (useReactive) {
+      verify(messagesManager).getMessagesForDeviceReactive(account.getUuid(), device.getId(), true);
+    } else {
       verify(messagesManager).getMessagesForDevice(account.getUuid(), device.getId(), true);
     }
+  }
 
-  @Test
+  @ParameterizedTest
-  void testProcessDatabaseMessagesAfterPersist() {
+  @ValueSource(booleans = {true, false})
+  void testProcessDatabaseMessagesAfterPersist(final boolean useReactive) {
     final MessagesManager messagesManager = mock(MessagesManager.class);
     final WebSocketClient client = mock(WebSocketClient.class);
     final WebSocketConnection connection = new WebSocketConnection(receiptSender, messagesManager, auth, device, client,
-        retrySchedulingExecutor);
+        retrySchedulingExecutor, useReactive, Schedulers.immediate());
 
     final UUID accountUuid = UUID.randomUUID();
 
@@ -637,8 +765,13 @@ class WebSocketConnectionTest {
     when(device.getId()).thenReturn(1L);
     when(client.isOpen()).thenReturn(true);
 
+    if (useReactive) {
+      when(messagesManager.getMessagesForDeviceReactive(eq(accountUuid), eq(1L), anyBoolean()))
+          .thenReturn(Flux.empty());
+    } else {
       when(messagesManager.getMessagesForDevice(eq(accountUuid), eq(1L), anyBoolean()))
           .thenReturn(new Pair<>(Collections.emptyList(), false));
+    }
 
     final WebSocketResponseMessage successResponse = mock(WebSocketResponseMessage.class);
     when(successResponse.getStatus()).thenReturn(200);
@@ -650,151 +783,16 @@ class WebSocketConnectionTest {
     connection.processStoredMessages();
     connection.handleMessagesPersisted();
 
+    if (useReactive) {
+      verify(messagesManager, times(2)).getMessagesForDeviceReactive(account.getUuid(), device.getId(), false);
+    } else {
       verify(messagesManager, times(2)).getMessagesForDevice(account.getUuid(), device.getId(), false);
     }
-
-  @Test
-  void testDiscardOversizedMessagesForDesktop() {
-    MessagesManager storedMessages = mock(MessagesManager.class);
-
-    UUID accountUuid   = UUID.randomUUID();
-    UUID senderOneUuid = UUID.randomUUID();
-    UUID senderTwoUuid = UUID.randomUUID();
-
-    List<Envelope> outgoingMessages = List.of(
-        createMessage(senderOneUuid, UUID.randomUUID(), 1111, "first"),
-        createMessage(senderOneUuid, UUID.randomUUID(), 2222,
-            RandomStringUtils.randomAlphanumeric(WebSocketConnection.MAX_DESKTOP_MESSAGE_SIZE + 1)),
-        createMessage(senderTwoUuid, UUID.randomUUID(), 3333, "third"));
-
-    when(device.getId()).thenReturn(2L);
-
-    when(account.getNumber()).thenReturn("+14152222222");
-    when(account.getUuid()).thenReturn(accountUuid);
-
-    final Device sender1device = mock(Device.class);
-
-    List<Device> sender1devices = List.of(sender1device);
-
-    Account sender1 = mock(Account.class);
-    when(sender1.getDevices()).thenReturn(sender1devices);
-
-    when(accountsManager.getByE164("sender1")).thenReturn(Optional.of(sender1));
-    when(accountsManager.getByE164("sender2")).thenReturn(Optional.empty());
-
-    String userAgent = "Signal-Desktop/1.2.3";
-
-    when(storedMessages.getMessagesForDevice(account.getUuid(), device.getId(), false))
-            .thenReturn(new Pair<>(outgoingMessages, false));
-
-    final List<CompletableFuture<WebSocketResponseMessage>> futures = new LinkedList<>();
-    final WebSocketClient                                   client  = mock(WebSocketClient.class);
-
-    when(client.getUserAgent()).thenReturn(userAgent);
-    when(client.sendRequest(eq("PUT"), eq("/api/v1/message"), ArgumentMatchers.nullable(List.class),
-        ArgumentMatchers.<Optional<byte[]>>any()))
-        .thenAnswer(new Answer<CompletableFuture<WebSocketResponseMessage>>() {
-          @Override
-          public CompletableFuture<WebSocketResponseMessage> answer(InvocationOnMock invocationOnMock) {
-            CompletableFuture<WebSocketResponseMessage> future = new CompletableFuture<>();
-            futures.add(future);
-            return future;
-          }
-        });
-
-    WebSocketConnection connection = new WebSocketConnection(receiptSender, storedMessages, auth, device, client,
-        retrySchedulingExecutor);
-
-    connection.start();
-    verify(client, times(2)).sendRequest(eq("PUT"), eq("/api/v1/message"), ArgumentMatchers.nullable(List.class),
-        ArgumentMatchers.<Optional<byte[]>>any());
-
-    assertEquals(2, futures.size());
-
-    WebSocketResponseMessage response = mock(WebSocketResponseMessage.class);
-    when(response.getStatus()).thenReturn(200);
-    futures.get(0).complete(response);
-    futures.get(1).complete(response);
-
-    // We should delete all three messages even though we only sent two; one got discarded because it was too big for
-    // desktop clients.
-    verify(storedMessages, times(3)).delete(eq(accountUuid), eq(2L), any(UUID.class), any(Long.class));
-
-    connection.stop();
-    verify(client).close(anyInt(), anyString());
   }
 
-  @Test
+  @ParameterizedTest
-  void testSendOversizedMessagesForNonDesktop() {
+  @ValueSource(booleans = {true, false})
-    MessagesManager storedMessages = mock(MessagesManager.class);
+  void testRetrieveMessageException(final boolean useReactive) {
-
-    UUID accountUuid = UUID.randomUUID();
-    UUID senderOneUuid = UUID.randomUUID();
-    UUID senderTwoUuid = UUID.randomUUID();
-
-    List<Envelope> outgoingMessages = List.of(createMessage(senderOneUuid, UUID.randomUUID(), 1111, "first"),
-        createMessage(senderOneUuid, UUID.randomUUID(), 2222,
-            RandomStringUtils.randomAlphanumeric(WebSocketConnection.MAX_DESKTOP_MESSAGE_SIZE + 1)),
-        createMessage(senderTwoUuid, UUID.randomUUID(), 3333, "third"));
-
-    when(device.getId()).thenReturn(2L);
-
-    when(account.getNumber()).thenReturn("+14152222222");
-    when(account.getUuid()).thenReturn(accountUuid);
-
-    final Device sender1device = mock(Device.class);
-
-    List<Device> sender1devices = List.of(sender1device);
-
-    Account sender1 = mock(Account.class);
-    when(sender1.getDevices()).thenReturn(sender1devices);
-
-    when(accountsManager.getByE164("sender1")).thenReturn(Optional.of(sender1));
-    when(accountsManager.getByE164("sender2")).thenReturn(Optional.empty());
-
-    String userAgent = "Signal-Android/4.68.3";
-
-    when(storedMessages.getMessagesForDevice(account.getUuid(), device.getId(), false))
-            .thenReturn(new Pair<>(outgoingMessages, false));
-
-    final List<CompletableFuture<WebSocketResponseMessage>> futures = new LinkedList<>();
-    final WebSocketClient                                   client  = mock(WebSocketClient.class);
-
-    when(client.getUserAgent()).thenReturn(userAgent);
-    when(client.sendRequest(eq("PUT"), eq("/api/v1/message"), ArgumentMatchers.nullable(List.class),
-        ArgumentMatchers.<Optional<byte[]>>any()))
-        .thenAnswer(new Answer<CompletableFuture<WebSocketResponseMessage>>() {
-          @Override
-          public CompletableFuture<WebSocketResponseMessage> answer(InvocationOnMock invocationOnMock) {
-            CompletableFuture<WebSocketResponseMessage> future = new CompletableFuture<>();
-            futures.add(future);
-            return future;
-          }
-        });
-
-    WebSocketConnection connection = new WebSocketConnection(receiptSender, storedMessages, auth, device, client,
-        retrySchedulingExecutor);
-
-    connection.start();
-    verify(client, times(3)).sendRequest(eq("PUT"), eq("/api/v1/message"), ArgumentMatchers.nullable(List.class),
-        ArgumentMatchers.<Optional<byte[]>>any());
-
-    assertEquals(3, futures.size());
-
-    WebSocketResponseMessage response = mock(WebSocketResponseMessage.class);
-    when(response.getStatus()).thenReturn(200);
-    futures.get(0).complete(response);
-    futures.get(1).complete(response);
-    futures.get(2).complete(response);
-
-    verify(storedMessages, times(3)).delete(eq(accountUuid), eq(2L), any(UUID.class), any(Long.class));
-
-    connection.stop();
-    verify(client).close(anyInt(), anyString());
-  }
-
-  @Test
-  void testRetrieveMessageException() {
     MessagesManager storedMessages = mock(MessagesManager.class);
 
     UUID accountUuid = UUID.randomUUID();
@@ -804,10 +802,13 @@ class WebSocketConnectionTest {
     when(account.getNumber()).thenReturn("+14152222222");
     when(account.getUuid()).thenReturn(accountUuid);
 
-    String userAgent = "Signal-Android/4.68.3";
+    if (useReactive) {
-
+      when(storedMessages.getMessagesForDeviceReactive(account.getUuid(), device.getId(), false))
+          .thenReturn(Flux.error(new RedisException("OH NO")));
+    } else {
       when(storedMessages.getMessagesForDevice(account.getUuid(), device.getId(), false))
           .thenThrow(new RedisException("OH NO"));
+    }
 
     when(retrySchedulingExecutor.schedule(any(Runnable.class), anyLong(), any())).thenAnswer(
         (Answer<ScheduledFuture<?>>) invocation -> {
@@ -819,7 +820,7 @@ class WebSocketConnectionTest {
     when(client.isOpen()).thenReturn(true);
 
     WebSocketConnection connection = new WebSocketConnection(receiptSender, storedMessages, auth, device, client,
-        retrySchedulingExecutor);
+        retrySchedulingExecutor, useReactive, Schedulers.immediate());
     connection.start();
 
     verify(retrySchedulingExecutor, times(WebSocketConnection.MAX_CONSECUTIVE_RETRIES)).schedule(any(Runnable.class),
@@ -827,8 +828,9 @@ class WebSocketConnectionTest {
     verify(client).close(eq(1011), anyString());
   }
 
-  @Test
+  @ParameterizedTest
-  void testRetrieveMessageExceptionClientDisconnected() {
+  @ValueSource(booleans = {true, false})
+  void testRetrieveMessageExceptionClientDisconnected(final boolean useReactive) {
     MessagesManager storedMessages = mock(MessagesManager.class);
 
     UUID accountUuid = UUID.randomUUID();
@@ -838,22 +840,143 @@ class WebSocketConnectionTest {
     when(account.getNumber()).thenReturn("+14152222222");
     when(account.getUuid()).thenReturn(accountUuid);
 
-    String userAgent = "Signal-Android/4.68.3";
+    if (useReactive) {
-
+      when(storedMessages.getMessagesForDeviceReactive(account.getUuid(), device.getId(), false))
+          .thenReturn(Flux.error(new RedisException("OH NO")));
+    } else {
       when(storedMessages.getMessagesForDevice(account.getUuid(), device.getId(), false))
           .thenThrow(new RedisException("OH NO"));
+    }
 
     final WebSocketClient client = mock(WebSocketClient.class);
     when(client.isOpen()).thenReturn(false);
 
     WebSocketConnection connection = new WebSocketConnection(receiptSender, storedMessages, auth, device, client,
-        retrySchedulingExecutor);
+        retrySchedulingExecutor, useReactive, Schedulers.immediate());
     connection.start();
 
     verify(retrySchedulingExecutor, never()).schedule(any(Runnable.class), anyLong(), any());
     verify(client, never()).close(anyInt(), anyString());
   }
 
+  @Test
+  @Disabled("This test is flaky")
+  void testReactivePublisherLimitRate() {
+    MessagesManager storedMessages = mock(MessagesManager.class);
+
+    final UUID accountUuid = UUID.randomUUID();
+
+    final long deviceId = 2L;
+    when(device.getId()).thenReturn(deviceId);
+
+    when(account.getNumber()).thenReturn("+14152222222");
+    when(account.getUuid()).thenReturn(accountUuid);
+
+    final int totalMessages = 10;
+    final AtomicReference<FluxSink<Envelope>> sink = new AtomicReference<>();
+
+    final AtomicLong maxRequest = new AtomicLong(-1);
+    final Flux<Envelope> flux = Flux.create(s -> {
+      sink.set(s);
+      s.onRequest(n -> {
+        if (maxRequest.get() < n) {
+          maxRequest.set(n);
+        }
+      });
+    });
+
+    when(storedMessages.getMessagesForDeviceReactive(eq(accountUuid), eq(deviceId), anyBoolean()))
+        .thenReturn(flux);
+
+    final WebSocketClient client = mock(WebSocketClient.class);
+    when(client.isOpen()).thenReturn(true);
+    final WebSocketResponseMessage successResponse = mock(WebSocketResponseMessage.class);
+    when(successResponse.getStatus()).thenReturn(200);
+    when(client.sendRequest(any(), any(), any(), any())).thenReturn(CompletableFuture.completedFuture(successResponse));
+    when(storedMessages.delete(any(), anyLong(), any(), any())).thenReturn(
+        CompletableFuture.completedFuture(Optional.empty()));
+
+    WebSocketConnection connection = new WebSocketConnection(receiptSender, storedMessages, auth, device, client,
+        retrySchedulingExecutor, true);
+
+    connection.start();
+
+    StepVerifier.setDefaultTimeout(Duration.ofSeconds(5));
+
+    StepVerifier.create(flux, 0)
+        .expectSubscription()
+        .thenRequest(totalMessages * 2)
+        .then(() -> {
+          for (long i = 0; i < totalMessages; i++) {
+            sink.get().next(createMessage(UUID.randomUUID(), accountUuid, 1111 * i + 1, "message " + i));
+          }
+          sink.get().complete();
+        })
+        .expectNextCount(totalMessages)
+        .expectComplete()
+        .log()
+        .verify();
+
+    assertEquals(WebSocketConnection.MESSAGE_PUBLISHER_LIMIT_RATE, maxRequest.get());
+  }
+
+  @Test
+  void testReactivePublisherDisposedWhenConnectionStopped() {
+    MessagesManager storedMessages = mock(MessagesManager.class);
+
+    final UUID accountUuid = UUID.randomUUID();
+
+    final long deviceId = 2L;
+    when(device.getId()).thenReturn(deviceId);
+
+    when(account.getNumber()).thenReturn("+14152222222");
+    when(account.getUuid()).thenReturn(accountUuid);
+
+    final AtomicBoolean canceled = new AtomicBoolean();
+
+    final Flux<Envelope> flux = Flux.create(s -> {
+      s.onRequest(n -> {
+        // the subscriber should request more than 1 message, but we will only send one, so that
+        // we are sure the subscriber is waiting for more when we stop the connection
+        assert n > 1;
+        s.next(createMessage(UUID.randomUUID(), UUID.randomUUID(), 1111, "first"));
+      });
+
+      s.onCancel(() -> canceled.set(true));
+    });
+    when(storedMessages.getMessagesForDeviceReactive(eq(accountUuid), eq(deviceId), anyBoolean()))
+        .thenReturn(flux);
+
+    final WebSocketClient client = mock(WebSocketClient.class);
+    when(client.isOpen()).thenReturn(true);
+    final WebSocketResponseMessage successResponse = mock(WebSocketResponseMessage.class);
+    when(successResponse.getStatus()).thenReturn(200);
+    when(client.sendRequest(any(), any(), any(), any())).thenReturn(CompletableFuture.completedFuture(successResponse));
+    when(storedMessages.delete(any(), anyLong(), any(), any())).thenReturn(
+        CompletableFuture.completedFuture(Optional.empty()));
+
+    WebSocketConnection connection = new WebSocketConnection(receiptSender, storedMessages, auth, device, client,
+        retrySchedulingExecutor, true, Schedulers.immediate());
+
+    connection.start();
+
+    verify(client).sendRequest(any(), any(), any(), any());
+
+    // close the connection before the publisher completes
+    connection.stop();
+
+    StepVerifier.setDefaultTimeout(Duration.ofSeconds(2));
+
+    StepVerifier.create(flux)
+        .expectSubscription()
+        .expectNextCount(1)
+        .then(() -> assertTrue(canceled.get()))
+        // this is not entirely intuitive, but expecting a timeout is the recommendation for verifying cancellation
+        .expectTimeout(Duration.ofMillis(100))
+        .log()
+        .verify();
+  }
+
   private Envelope createMessage(UUID senderUuid, UUID destinationUuid, long timestamp, String content) {
     return Envelope.newBuilder()
         .setServerGuid(UUID.randomUUID().toString())

service/src/test/resources/logback-test.xml

@@ -8,4 +8,7 @@
   <root level="warn">
     <appender-ref ref="STDOUT"/>
   </root>
+
+  <!-- uncomment and combine with .log() in StepVerifier for more insight into reactor operations -->
+  <!--  <logger name="reactor" level="debug"/> -->
 </configuration>