Use error-specific retry mechanisms in WebSocketConnection and associated classes

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

@@ -22,6 +22,7 @@ import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import com.google.protobuf.ByteString;
+import io.lettuce.core.RedisCommandTimeoutException;
 import io.lettuce.core.RedisFuture;
 import io.lettuce.core.cluster.SlotHash;
 import io.lettuce.core.cluster.api.async.RedisAdvancedClusterAsyncCommands;
@@ -52,6 +53,7 @@ import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
 import org.apache.commons.lang3.ArrayUtils;
@@ -781,6 +783,47 @@ class MessagesCacheTest {
       verify(asyncCommands, atLeast(1)).evalsha(any(), any(), any(byte[][].class), any(byte[][].class));
     }
 
+    @Test
+    void testGetRetries() {
+      final List<byte[]> page = generatePage();
+
+      final AtomicBoolean emittedError = new AtomicBoolean(false);
+      final AtomicBoolean emittedPage = new AtomicBoolean(false);
+
+      when(reactiveCommands.evalsha(any(), any(), any(byte[][].class), any(byte[][].class)))
+          .thenReturn(Flux.defer(() -> {
+            if (emittedError.compareAndSet(false, true)) {
+              return Flux.error(new RedisCommandTimeoutException("Timeout"));
+            } else if (emittedPage.compareAndSet(false, true)) {
+              return Flux.just(page);
+            }
+
+            return Flux.empty();
+          }));
+
+      final AsyncCommand<?, ?, ?> removeSuccess = new AsyncCommand<>(mock(RedisCommand.class));
+      removeSuccess.complete();
+
+      when(asyncCommands.evalsha(any(), any(), any(byte[][].class), any(byte[][].class)))
+          .thenReturn((RedisFuture) removeSuccess);
+
+      final Publisher<?> allMessages = messagesCache.get(UUID.randomUUID(), Device.PRIMARY_ID);
+
+      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(page.size() / 2)
+          .expectComplete()
+          .log()
+          .verify();
+    }
+
     private List<byte[]> generatePage() {
       final List<byte[]> messagesAndIds = new ArrayList<>();
 

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

@@ -31,7 +31,6 @@ import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
-import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import org.apache.commons.lang3.RandomStringUtils;
@@ -78,7 +77,6 @@ class WebSocketConnectionIntegrationTest {
   static final RedisClusterExtension REDIS_CLUSTER_EXTENSION = RedisClusterExtension.builder().build();
 
   private ExecutorService sharedExecutorService;
-  private ScheduledExecutorService scheduledExecutorService;
   private MessagesDynamoDb messagesDynamoDb;
   private MessagesCache messagesCache;
   private ReportMessageManager reportMessageManager;
@@ -95,7 +93,6 @@ class WebSocketConnectionIntegrationTest {
   @BeforeEach
   void setUp() throws Exception {
     sharedExecutorService = Executors.newSingleThreadExecutor();
-    scheduledExecutorService = Executors.newSingleThreadScheduledExecutor();
     messageDeliveryScheduler = Schedulers.newBoundedElastic(10, 10_000, "messageDelivery");
     dynamicConfigurationManager = mock(DynamicConfigurationManager.class);
     when(dynamicConfigurationManager.getConfiguration()).thenReturn(new DynamicConfiguration());
@@ -118,10 +115,8 @@ class WebSocketConnectionIntegrationTest {
   @AfterEach
   void tearDown() throws Exception {
     sharedExecutorService.shutdown();
+    //noinspection ResultOfMethodCallIgnored
     sharedExecutorService.awaitTermination(2, TimeUnit.SECONDS);
-
-    scheduledExecutorService.shutdown();
-    scheduledExecutorService.awaitTermination(2, TimeUnit.SECONDS);
   }
 
   @ParameterizedTest
@@ -140,7 +135,6 @@ class WebSocketConnectionIntegrationTest {
         account,
         device,
         webSocketClient,
-        scheduledExecutorService,
         messageDeliveryScheduler,
         clientReleaseManager,
         mock(MessageDeliveryLoopMonitor.class),
@@ -230,7 +224,6 @@ class WebSocketConnectionIntegrationTest {
         account,
         device,
         webSocketClient,
-        scheduledExecutorService,
         messageDeliveryScheduler,
         clientReleaseManager,
         mock(MessageDeliveryLoopMonitor.class),
@@ -301,8 +294,7 @@ class WebSocketConnectionIntegrationTest {
         account,
         device,
         webSocketClient,
-        100, // use a very short timeout, so that this test completes quickly
-        scheduledExecutorService,
+        1000, // use a short timeout, so that this test completes quickly
         messageDeliveryScheduler,
         clientReleaseManager,
         mock(MessageDeliveryLoopMonitor.class),
@@ -371,8 +363,8 @@ class WebSocketConnectionIntegrationTest {
       ArgumentCaptor<Optional<byte[]>> messageBodyCaptor = ArgumentCaptor.forClass(Optional.class);
 
       // We expect all of the messages from both pools to be sent, plus one for the future that times out
-      verify(webSocketClient, atMost(persistedMessageCount + cachedMessageCount + 1)).sendRequest(eq("PUT"),
-          eq("/api/v1/message"), anyList(), messageBodyCaptor.capture());
+      verify(webSocketClient, atMost(persistedMessageCount + cachedMessageCount + 1))
+          .sendRequest(eq("PUT"), eq("/api/v1/message"), anyList(), messageBodyCaptor.capture());
 
       verify(webSocketClient).sendRequest(eq("PUT"), eq("/api/v1/queue/empty"), anyList(), eq(Optional.empty()));
 

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

@@ -15,7 +15,6 @@ 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;
 import static org.mockito.Mockito.anyString;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -42,8 +41,6 @@ import java.util.Queue;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.IntStream;
@@ -95,7 +92,6 @@ class WebSocketConnectionTest {
   private UpgradeRequest upgradeRequest;
   private MessagesManager messagesManager;
   private ReceiptSender receiptSender;
-  private ScheduledExecutorService retrySchedulingExecutor;
   private Scheduler messageDeliveryScheduler;
   private ClientReleaseManager clientReleaseManager;
 
@@ -108,7 +104,6 @@ class WebSocketConnectionTest {
     upgradeRequest = mock(UpgradeRequest.class);
     messagesManager = mock(MessagesManager.class);
     receiptSender = mock(ReceiptSender.class);
-    retrySchedulingExecutor = mock(ScheduledExecutorService.class);
     messageDeliveryScheduler = Schedulers.newBoundedElastic(10, 10_000, "messageDelivery");
     clientReleaseManager = mock(ClientReleaseManager.class);
   }
@@ -125,7 +120,7 @@ class WebSocketConnectionTest {
         new WebSocketAccountAuthenticator(accountAuthenticator);
     AuthenticatedConnectListener connectListener = new AuthenticatedConnectListener(accountsManager, receiptSender, messagesManager,
         new MessageMetrics(), mock(PushNotificationManager.class), mock(PushNotificationScheduler.class),
-        mock(RedisMessageAvailabilityManager.class), mock(DisconnectionRequestManager.class), retrySchedulingExecutor,
+        mock(RedisMessageAvailabilityManager.class), mock(DisconnectionRequestManager.class),
         messageDeliveryScheduler, clientReleaseManager, mock(MessageDeliveryLoopMonitor.class),
         mock(ExperimentEnrollmentManager.class));
     WebSocketSessionContext sessionContext = mock(WebSocketSessionContext.class);
@@ -630,7 +625,7 @@ class WebSocketConnectionTest {
   private WebSocketConnection webSocketConnection(final WebSocketClient client) {
     return new WebSocketConnection(receiptSender, messagesManager, new MessageMetrics(),
         mock(PushNotificationManager.class), mock(PushNotificationScheduler.class), account, device, client,
-        retrySchedulingExecutor, Schedulers.immediate(), clientReleaseManager,
+        Schedulers.immediate(), clientReleaseManager,
         mock(MessageDeliveryLoopMonitor.class), mock(ExperimentEnrollmentManager.class));
   }
 
@@ -788,20 +783,12 @@ class WebSocketConnectionTest {
     when(messagesManager.getMessagesForDeviceReactive(account.getIdentifier(IdentityType.ACI), device, false))
         .thenReturn(Flux.error(new RedisException("OH NO")));
 
-    when(retrySchedulingExecutor.schedule(any(Runnable.class), anyLong(), any())).thenAnswer(
-        (Answer<ScheduledFuture<?>>) invocation -> {
-          invocation.getArgument(0, Runnable.class).run();
-          return mock(ScheduledFuture.class);
-        });
-
     final WebSocketClient client = mock(WebSocketClient.class);
     when(client.isOpen()).thenReturn(true);
 
     WebSocketConnection connection = webSocketConnection(client);
     connection.start();
 
-    verify(retrySchedulingExecutor, times(WebSocketConnection.MAX_CONSECUTIVE_RETRIES)).schedule(any(Runnable.class),
-        anyLong(), any());
     verify(client).close(eq(1011), anyString());
   }
 
@@ -823,7 +810,6 @@ class WebSocketConnectionTest {
     WebSocketConnection connection = webSocketConnection(client);
     connection.start();
 
-    verify(retrySchedulingExecutor, never()).schedule(any(Runnable.class), anyLong(), any());
     verify(client, never()).close(anyInt(), anyString());
   }