Use reactive streams for WebSocket message queue

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

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(
-        uuid1, Envelope.Type.CIPHERTEXT_VALUE,
-        timestamp, sourceUuid, 1, AuthHelper.VALID_UUID, null, "hi".getBytes(), 0)));
+    when(messagesManager.delete(AuthHelper.VALID_UUID, 1, uuid1, null))
+        .thenReturn(
+            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(
-        uuid2, Envelope.Type.SERVER_DELIVERY_RECEIPT_VALUE,
-        System.currentTimeMillis(), sourceUuid, 1, AuthHelper.VALID_UUID, null, null, 0)));
+    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))));
 
     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 =
-        RedisClusterHelper.buildMockRedisCluster(stringCommands, binaryCommands);
+    final FaultTolerantRedisCluster mockCluster = RedisClusterHelper.builder()
+        .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();
-    notificationExecutorService.awaitTermination(1, TimeUnit.SECONDS);
+    sharedExecutorService.shutdown();
+    sharedExecutorService.awaitTermination(1, TimeUnit.SECONDS);
   }
 
   @Test

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,
-        Duration.ofDays(14));
+    messageDeletionExecutorService = Executors.newSingleThreadExecutor();
+    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(messagesDynamoDb.load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().isEmpty();
-    assertThat(messagesDynamoDb.load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull()
+    assertThat(load(destinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().isEmpty();
+    assertThat(load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().isEmpty();
+    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(messagesDynamoDb.load(secondDestinationUuid, 1, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull()
+    assertThat(load(destinationUuid, 2, MessagesDynamoDb.RESULT_SET_CHUNK_SIZE)).isNotNull().isEmpty();
+    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,70 +5,118 @@
 
 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 FaultTolerantRedisCluster buildMockRedisCluster(final RedisAdvancedClusterCommands<String, String> stringCommands) {
-        return buildMockRedisCluster(stringCommands, mock(RedisAdvancedClusterCommands.class));
+  public static RedisClusterHelper.Builder builder() {
+    return new Builder();
+  }
+
+  @SuppressWarnings("unchecked")
+  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);
+    });
+
+    doAnswer(invocation -> {
+      invocation.getArgument(0, Consumer.class).accept(stringConnection);
+      return null;
+    }).when(cluster).useCluster(any(Consumer.class));
+
+    when(cluster.withCluster(any(Function.class))).thenAnswer(invocation -> {
+      return invocation.getArgument(0, Function.class).apply(stringConnection);
+    });
+
+    doAnswer(invocation -> {
+      invocation.getArgument(0, Consumer.class).accept(stringConnection);
+      return null;
+    }).when(cluster).useCluster(any(Consumer.class));
+
+    when(cluster.withBinaryCluster(any(Function.class))).thenAnswer(invocation -> {
+      return invocation.getArgument(0, Function.class).apply(binaryConnection);
+    });
+
+    doAnswer(invocation -> {
+      invocation.getArgument(0, Consumer.class).accept(binaryConnection);
+      return null;
+    }).when(cluster).useBinaryCluster(any(Consumer.class));
+
+    when(cluster.withBinaryCluster(any(Function.class))).thenAnswer(invocation -> {
+      return invocation.getArgument(0, Function.class).apply(binaryConnection);
+    });
+
+    doAnswer(invocation -> {
+      invocation.getArgument(0, Consumer.class).accept(binaryConnection);
+      return null;
+    }).when(cluster).useBinaryCluster(any(Consumer.class));
+
+    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() {
+
     }
 
-    @SuppressWarnings("unchecked")
-    public static FaultTolerantRedisCluster buildMockRedisCluster(final RedisAdvancedClusterCommands<String, String> stringCommands, final RedisAdvancedClusterCommands<byte[], byte[]> binaryCommands) {
-        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(cluster.withCluster(any(Function.class))).thenAnswer(invocation -> {
-            return invocation.getArgument(0, Function.class).apply(stringConnection);
-        });
-        
-        doAnswer(invocation -> {
-            invocation.getArgument(0, Consumer.class).accept(stringConnection);
-            return null;
-        }).when(cluster).useCluster(any(Consumer.class));
-
-        when(cluster.withCluster(any(Function.class))).thenAnswer(invocation -> {
-            return invocation.getArgument(0, Function.class).apply(stringConnection);
-        });
-
-        doAnswer(invocation -> {
-            invocation.getArgument(0, Consumer.class).accept(stringConnection);
-            return null;
-        }).when(cluster).useCluster(any(Consumer.class));
-
-        when(cluster.withBinaryCluster(any(Function.class))).thenAnswer(invocation -> {
-            return invocation.getArgument(0, Function.class).apply(binaryConnection);
-        });
-
-        doAnswer(invocation -> {
-            invocation.getArgument(0, Consumer.class).accept(binaryConnection);
-            return null;
-        }).when(cluster).useBinaryCluster(any(Consumer.class));
-
-        when(cluster.withBinaryCluster(any(Function.class))).thenAnswer(invocation -> {
-            return invocation.getArgument(0, Function.class).apply(binaryConnection);
-        });
-
-        doAnswer(invocation -> {
-            invocation.getArgument(0, Consumer.class).accept(binaryConnection);
-            return null;
-        }).when(cluster).useBinaryCluster(any(Consumer.class));
-        
-        return cluster;
+    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 executorService;
+  private ExecutorService sharedExecutorService;
   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);
-    messagesDynamoDb = new MessagesDynamoDb(dynamoDbExtension.getDynamoDbClient(), MessagesDynamoDbExtension.TABLE_NAME,
-        Duration.ofDays(7));
+        REDIS_CLUSTER_EXTENSION.getRedisCluster(), Clock.systemUTC(), sharedExecutorService, sharedExecutorService);
+    messagesDynamoDb = new MessagesDynamoDb(dynamoDbExtension.getDynamoDbClient(),
+        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();
-    executorService.awaitTermination(2, TimeUnit.SECONDS);
+    sharedExecutorService.shutdown();
+    sharedExecutorService.awaitTermination(2, TimeUnit.SECONDS);
 
     retrySchedulingExecutor.shutdown();
     retrySchedulingExecutor.awaitTermination(2, TimeUnit.SECONDS);
   }
 
-  @Test
-  void testProcessStoredMessages() {
-    final int persistedMessageCount = 207;
-    final int cachedMessageCount = 173;
+  @ParameterizedTest
+  @CsvSource({
+      "207, 173, true",
+      "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(
-          CompletableFuture.completedFuture(successResponse));
+      when(webSocketClient.sendRequest(eq("PUT"), eq("/api/v1/message"), anyList(), any()))
+          .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
-  void testProcessStoredMessagesClientClosed() {
+  @ParameterizedTest
+  @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
-  void testProcessStoredMessagesSendFutureTimeout() {
+  @ParameterizedTest
+  @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();
   }
+
 }