Use MessageStream in WebSocketConnection

Author: jon-signal

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

@@ -994,8 +994,9 @@ protected void configureServer(final ServerBuilder<?> serverBuilder) {
         config.idlePrimaryDeviceReminderConfiguration().minIdleDuration(), Clock.systemUTC()));
     webSocketEnvironment.setConnectListener(
         new AuthenticatedConnectListener(accountsManager, receiptSender, messagesManager, messageMetrics, pushNotificationManager,
-            pushNotificationScheduler, redisMessageAvailabilityManager, disconnectionRequestManager,
-            messageDeliveryScheduler, clientReleaseManager, messageDeliveryLoopMonitor, experimentEnrollmentManager));
+            pushNotificationScheduler, disconnectionRequestManager,
+            messageDeliveryScheduler, clientReleaseManager, messageDeliveryLoopMonitor, experimentEnrollmentManager
+        ));
     webSocketEnvironment.jersey().register(new RateLimitByIpFilter(rateLimiters));
     webSocketEnvironment.jersey().register(new RequestStatisticsFilter(TrafficSource.WEBSOCKET));
     webSocketEnvironment.jersey().register(MultiRecipientMessageProvider.class);

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

@@ -7,10 +7,11 @@
 
 import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
 
+import com.google.common.annotations.VisibleForTesting;
 import io.micrometer.core.instrument.Tags;
-
 import java.time.Duration;
 import java.util.Optional;
+import java.util.function.Function;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedDevice;
@@ -23,12 +24,12 @@
 import org.whispersystems.textsecuregcm.push.PushNotificationManager;
 import org.whispersystems.textsecuregcm.push.PushNotificationScheduler;
 import org.whispersystems.textsecuregcm.push.ReceiptSender;
-import org.whispersystems.textsecuregcm.push.RedisMessageAvailabilityManager;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.ClientReleaseManager;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.MessagesManager;
+import org.whispersystems.websocket.WebSocketClient;
 import org.whispersystems.websocket.session.WebSocketSessionContext;
 import org.whispersystems.websocket.setup.WebSocketConnectListener;
 import reactor.core.scheduler.Scheduler;
@@ -45,69 +46,82 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
   private static final Logger log = LoggerFactory.getLogger(AuthenticatedConnectListener.class);
 
   private final AccountsManager accountsManager;
-  private final ReceiptSender receiptSender;
-  private final MessagesManager messagesManager;
-  private final MessageMetrics messageMetrics;
-  private final PushNotificationManager pushNotificationManager;
-  private final PushNotificationScheduler pushNotificationScheduler;
-  private final RedisMessageAvailabilityManager redisMessageAvailabilityManager;
   private final DisconnectionRequestManager disconnectionRequestManager;
-  private final Scheduler messageDeliveryScheduler;
-  private final ClientReleaseManager clientReleaseManager;
-  private final MessageDeliveryLoopMonitor messageDeliveryLoopMonitor;
-  private final ExperimentEnrollmentManager experimentEnrollmentManager;
+  private final WebSocketConnectionBuilder webSocketConnectionBuilder;
 
   private final OpenWebSocketCounter openAuthenticatedWebSocketCounter;
   private final OpenWebSocketCounter openUnauthenticatedWebSocketCounter;
 
+  @VisibleForTesting
+  @FunctionalInterface
+  interface WebSocketConnectionBuilder {
+    WebSocketConnection buildWebSocketConnection(Account account, Device device, WebSocketClient client);
+  }
+
   public AuthenticatedConnectListener(
       final AccountsManager accountsManager,
       final ReceiptSender receiptSender,
       final MessagesManager messagesManager,
       final MessageMetrics messageMetrics,
       final PushNotificationManager pushNotificationManager,
       final PushNotificationScheduler pushNotificationScheduler,
-      final RedisMessageAvailabilityManager redisMessageAvailabilityManager,
       final DisconnectionRequestManager disconnectionRequestManager,
       final Scheduler messageDeliveryScheduler,
       final ClientReleaseManager clientReleaseManager,
       final MessageDeliveryLoopMonitor messageDeliveryLoopMonitor,
       final ExperimentEnrollmentManager experimentEnrollmentManager) {
 
+    this(accountsManager,
+        disconnectionRequestManager,
+        (account, device, client) -> new WebSocketConnection(receiptSender,
+            messagesManager,
+            messageMetrics,
+            pushNotificationManager,
+            pushNotificationScheduler,
+            account,
+            device,
+            client,
+            messageDeliveryScheduler,
+            clientReleaseManager,
+            messageDeliveryLoopMonitor,
+            experimentEnrollmentManager),
+        authenticated -> new OpenWebSocketCounter(OPEN_WEBSOCKET_GAUGE_NAME,
+            NEW_CONNECTION_COUNTER_NAME,
+            CONNECTED_DURATION_TIMER_NAME,
+            Duration.ofHours(3),
+            Tags.of(AUTHENTICATED_TAG_NAME, String.valueOf(authenticated)))
+    );
+  }
+
+  @VisibleForTesting AuthenticatedConnectListener(
+      final AccountsManager accountsManager,
+      final DisconnectionRequestManager disconnectionRequestManager,
+      final WebSocketConnectionBuilder webSocketConnectionBuilder,
+      final Function<Boolean, OpenWebSocketCounter> openWebSocketCounterBuilder) {
+
     this.accountsManager = accountsManager;
-    this.receiptSender = receiptSender;
-    this.messagesManager = messagesManager;
-    this.messageMetrics = messageMetrics;
-    this.pushNotificationManager = pushNotificationManager;
-    this.pushNotificationScheduler = pushNotificationScheduler;
-    this.redisMessageAvailabilityManager = redisMessageAvailabilityManager;
     this.disconnectionRequestManager = disconnectionRequestManager;
-    this.messageDeliveryScheduler = messageDeliveryScheduler;
-    this.clientReleaseManager = clientReleaseManager;
-    this.messageDeliveryLoopMonitor = messageDeliveryLoopMonitor;
-    this.experimentEnrollmentManager = experimentEnrollmentManager;
-
-    openAuthenticatedWebSocketCounter =
-        new OpenWebSocketCounter(OPEN_WEBSOCKET_GAUGE_NAME, NEW_CONNECTION_COUNTER_NAME, CONNECTED_DURATION_TIMER_NAME, Duration.ofHours(3), Tags.of(AUTHENTICATED_TAG_NAME, "true"));
+    this.webSocketConnectionBuilder = webSocketConnectionBuilder;
 
-    openUnauthenticatedWebSocketCounter =
-        new OpenWebSocketCounter(OPEN_WEBSOCKET_GAUGE_NAME, NEW_CONNECTION_COUNTER_NAME, CONNECTED_DURATION_TIMER_NAME, Duration.ofHours(3), Tags.of(AUTHENTICATED_TAG_NAME, "false"));
+    openAuthenticatedWebSocketCounter = openWebSocketCounterBuilder.apply(true);
+    openUnauthenticatedWebSocketCounter = openWebSocketCounterBuilder.apply(false);
   }
 
   @Override
   public void onWebSocketConnect(final WebSocketSessionContext context) {
 
     final boolean authenticated = (context.getAuthenticated() != null);
-    final OpenWebSocketCounter openWebSocketCounter =
-        authenticated ? openAuthenticatedWebSocketCounter : openUnauthenticatedWebSocketCounter;
 
-    openWebSocketCounter.countOpenWebSocket(context);
+    (authenticated ? openAuthenticatedWebSocketCounter : openUnauthenticatedWebSocketCounter).countOpenWebSocket(context);
 
     if (authenticated) {
       final AuthenticatedDevice auth = context.getAuthenticated(AuthenticatedDevice.class);
 
-      final Optional<Account> maybeAuthenticatedAccount = accountsManager.getByAccountIdentifier(auth.accountIdentifier());
-      final Optional<Device> maybeAuthenticatedDevice = maybeAuthenticatedAccount.flatMap(account -> account.getDevice(auth.deviceId()));
+      final Optional<Account> maybeAuthenticatedAccount =
+          accountsManager.getByAccountIdentifier(auth.accountIdentifier());
+
+      final Optional<Device> maybeAuthenticatedDevice =
+          maybeAuthenticatedAccount.flatMap(account -> account.getDevice(auth.deviceId()));
 
       if (maybeAuthenticatedAccount.isEmpty() || maybeAuthenticatedDevice.isEmpty()) {
         log.warn("{}:{} not found when opening authenticated WebSocket", auth.accountIdentifier(), auth.deviceId());
@@ -116,18 +130,10 @@ public void onWebSocketConnect(final WebSocketSessionContext context) {
         return;
       }
 
-      final WebSocketConnection connection = new WebSocketConnection(receiptSender,
-          messagesManager,
-          messageMetrics,
-          pushNotificationManager,
-          pushNotificationScheduler,
-          maybeAuthenticatedAccount.get(),
-          maybeAuthenticatedDevice.get(),
-          context.getClient(),
-          messageDeliveryScheduler,
-          clientReleaseManager,
-          messageDeliveryLoopMonitor,
-          experimentEnrollmentManager);
+      final WebSocketConnection connection =
+          webSocketConnectionBuilder.buildWebSocketConnection(maybeAuthenticatedAccount.get(),
+              maybeAuthenticatedDevice.get(),
+              context.getClient());
 
       disconnectionRequestManager.addListener(maybeAuthenticatedAccount.get().getIdentifier(IdentityType.ACI),
           maybeAuthenticatedDevice.get().getId(),
@@ -138,27 +144,11 @@ public void onWebSocketConnect(final WebSocketSessionContext context) {
             maybeAuthenticatedDevice.get().getId(),
             connection);
 
-        // We begin the shutdown process by removing this client's "presence," which means it will again begin to
-        // receive push notifications for inbound messages. We should do this first because, at this point, the
-        // connection has already closed and attempts to actually deliver a message via the connection will not succeed.
-        // It's preferable to start sending push notifications as soon as possible.
-        redisMessageAvailabilityManager.handleClientDisconnected(auth.accountIdentifier(), auth.deviceId());
-
-        // Finally, stop trying to deliver messages and send a push notification if the connection is aware of any
-        // undelivered messages.
         connection.stop();
       });
 
       try {
-        // Once we "start" the websocket connection, we'll cancel any scheduled "you may have new messages" push
-        // notifications and begin delivering any stored messages for the connected device. We have not yet declared the
-        // client as "present" yet. If a message arrives at this point, we will update the message availability state
-        // correctly, but we may also send a spurious push notification.
         connection.start();
-
-        // Finally, we register this client's presence, which suppresses push notifications. We do this last because
-        // receiving extra push notifications is generally preferable to missing out on a push notification.
-        redisMessageAvailabilityManager.handleClientConnected(auth.accountIdentifier(), auth.deviceId(), connection);
       } catch (final Exception e) {
         log.warn("Failed to initialize websocket", e);
         context.getClient().close(1011, "Unexpected error initializing connection");