Drop pub/sub operations from WebsocketConnection.

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

@@ -404,9 +404,9 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     ///
     WebSocketEnvironment<Account> webSocketEnvironment = new WebSocketEnvironment<>(environment, config.getWebSocketConfiguration(), 90000);
     webSocketEnvironment.setAuthenticator(new WebSocketAccountAuthenticator(accountAuthenticator));
-    webSocketEnvironment.setConnectListener(new AuthenticatedConnectListener(pushSender, receiptSender, messagesManager, pubSubManager, apnFallbackManager, clientPresenceManager));
+    webSocketEnvironment.setConnectListener(new AuthenticatedConnectListener(receiptSender, messagesManager, apnFallbackManager, clientPresenceManager));
     webSocketEnvironment.jersey().register(new MetricsApplicationEventListener(TrafficSource.WEBSOCKET));
-    webSocketEnvironment.jersey().register(new KeepAliveController(pubSubManager));
+    webSocketEnvironment.jersey().register(new KeepAliveController(clientPresenceManager));
     webSocketEnvironment.jersey().register(messageController);
     webSocketEnvironment.jersey().register(profileController);
     webSocketEnvironment.jersey().register(attachmentControllerV1);
@@ -417,7 +417,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     WebSocketEnvironment<Account> provisioningEnvironment = new WebSocketEnvironment<>(environment, webSocketEnvironment.getRequestLog(), 60000);
     provisioningEnvironment.setConnectListener(new ProvisioningConnectListener(pubSubManager));
     provisioningEnvironment.jersey().register(new MetricsApplicationEventListener(TrafficSource.WEBSOCKET));
-    provisioningEnvironment.jersey().register(new KeepAliveController(pubSubManager));
+    provisioningEnvironment.jersey().register(new KeepAliveController(clientPresenceManager));
 
     registerCorsFilter(environment);
     registerExceptionMappers(environment, webSocketEnvironment, provisioningEnvironment);

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

@@ -1,11 +1,11 @@
 package org.whispersystems.textsecuregcm.controllers;
 
 import com.codahale.metrics.annotation.Timed;
+import io.dropwizard.auth.Auth;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
 import org.whispersystems.textsecuregcm.storage.Account;
-import org.whispersystems.textsecuregcm.storage.PubSubManager;
-import org.whispersystems.textsecuregcm.websocket.WebsocketAddress;
 import org.whispersystems.websocket.session.WebSocketSession;
 import org.whispersystems.websocket.session.WebSocketSessionContext;
 
@@ -13,18 +13,16 @@ import javax.ws.rs.GET;
 import javax.ws.rs.Path;
 import javax.ws.rs.core.Response;
 
-import io.dropwizard.auth.Auth;
-
 
 @Path("/v1/keepalive")
 public class KeepAliveController {
 
   private final Logger logger = LoggerFactory.getLogger(KeepAliveController.class);
 
-  private final PubSubManager pubSubManager;
+  private final ClientPresenceManager clientPresenceManager;
 
-  public KeepAliveController(PubSubManager pubSubManager) {
+  public KeepAliveController(final ClientPresenceManager clientPresenceManager) {
-    this.pubSubManager = pubSubManager;
+    this.clientPresenceManager = clientPresenceManager;
   }
 
   @Timed
@@ -33,11 +31,8 @@ public class KeepAliveController {
                                @WebSocketSession WebSocketSessionContext context)
   {
     if (account != null) {
-      WebsocketAddress address = new WebsocketAddress(account.getNumber(),
+      if (!clientPresenceManager.isLocallyPresent(account.getUuid(), account.getAuthenticatedDevice().get().getId())) {
-                                                      account.getAuthenticatedDevice().get().getId());
+        logger.warn("***** No local subscription found for {}::{}", account.getUuid(), account.getAuthenticatedDevice().get().getId());
-
-      if (!pubSubManager.hasLocalSubscription(address)) {
-        logger.warn("***** No local subscription found for: " + address);
         context.getClient().close(1000, "OK");
       }
     }

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

@@ -170,6 +170,10 @@ public class ClientPresenceManager extends RedisClusterPubSubAdapter<String, Str
         }
     }
 
+    public boolean isLocallyPresent(final UUID accountUuid, final long deviceId) {
+        return displacementListenersByPresenceKey.containsKey(getPresenceKey(accountUuid, deviceId));
+    }
+
     public boolean clearPresence(final UUID accountUuid, final long deviceId) {
         return clearPresence(getPresenceKey(accountUuid, deviceId));
     }

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

@@ -4,17 +4,13 @@ import com.codahale.metrics.Counter;
 import com.codahale.metrics.MetricRegistry;
 import com.codahale.metrics.SharedMetricRegistries;
 import com.codahale.metrics.Timer;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.push.ApnFallbackManager;
 import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
-import org.whispersystems.textsecuregcm.push.PushSender;
 import org.whispersystems.textsecuregcm.push.ReceiptSender;
 import org.whispersystems.textsecuregcm.redis.RedisOperation;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.MessagesManager;
-import org.whispersystems.textsecuregcm.storage.PubSubManager;
 import org.whispersystems.textsecuregcm.util.Constants;
 import org.whispersystems.websocket.session.WebSocketSessionContext;
 import org.whispersystems.websocket.setup.WebSocketConnectListener;
@@ -25,30 +21,23 @@ import static com.codahale.metrics.MetricRegistry.name;
 
 public class AuthenticatedConnectListener implements WebSocketConnectListener {
 
-  private static final Logger         logger                       = LoggerFactory.getLogger(WebSocketConnection.class);
   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          unauthenticatedDurationTimer = metricRegistry.timer(name(WebSocketConnection.class, "unauthenticated_connection_duration"));
   private static final Counter        openWebsocketCounter         = metricRegistry.counter(name(WebSocketConnection.class, "open_websockets"));
 
-  private final PushSender            pushSender;
   private final ReceiptSender         receiptSender;
   private final MessagesManager       messagesManager;
-  private final PubSubManager         pubSubManager;
   private final ApnFallbackManager    apnFallbackManager;
   private final ClientPresenceManager clientPresenceManager;
 
-  public AuthenticatedConnectListener(PushSender pushSender,
+  public AuthenticatedConnectListener(ReceiptSender receiptSender,
-                                      ReceiptSender receiptSender,
                                       MessagesManager messagesManager,
-                                      PubSubManager pubSubManager,
                                       ApnFallbackManager apnFallbackManager,
                                       ClientPresenceManager clientPresenceManager)
   {
-    this.pushSender            = pushSender;
     this.receiptSender         = receiptSender;
     this.messagesManager       = messagesManager;
-    this.pubSubManager         = pubSubManager;
     this.apnFallbackManager    = apnFallbackManager;
     this.clientPresenceManager = clientPresenceManager;
   }
@@ -60,8 +49,7 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
       final Device                  device         = account.getAuthenticatedDevice().get();
       final String                  connectionId   = String.valueOf(new SecureRandom().nextLong());
       final Timer.Context           timer          = durationTimer.time();
-      final WebsocketAddress        address        = new WebsocketAddress(account.getNumber(), device.getId());
+      final WebSocketConnection     connection     = new WebSocketConnection(receiptSender,
-      final WebSocketConnection     connection     = new WebSocketConnection(pushSender, receiptSender,
                                                                              messagesManager, account, device,
                                                                              context.getClient(), connectionId);
 
@@ -70,15 +58,16 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
 
       clientPresenceManager.setPresent(account.getUuid(), device.getId(), connection);
       messagesManager.addMessageAvailabilityListener(account.getUuid(), device.getId(), connection);
-      pubSubManager.subscribe(address, connection);
+      connection.start();
 
       context.addListener(new WebSocketSessionContext.WebSocketEventListener() {
         @Override
         public void onWebSocketClose(WebSocketSessionContext context, int statusCode, String reason) {
-          openWebsocketCounter.dec();
-          pubSubManager.unsubscribe(address, connection);
           clientPresenceManager.clearPresence(account.getUuid(), device.getId());
           messagesManager.removeMessageAvailabilityListener(connection);
+          connection.stop();
+
+          openWebsocketCounter.dec();
           timer.stop();
         }
       });

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

@@ -44,7 +44,7 @@ import static org.whispersystems.textsecuregcm.entities.MessageProtos.Envelope;
 import static org.whispersystems.textsecuregcm.storage.PubSubProtos.PubSubMessage;
 
 @SuppressWarnings("OptionalUsedAsFieldOrParameterType")
-public class WebSocketConnection implements DispatchChannel, MessageAvailabilityListener, DisplacedPresenceListener {
+public class WebSocketConnection implements MessageAvailabilityListener, DisplacedPresenceListener {
 
   private static final MetricRegistry metricRegistry                 = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
   public  static final Histogram      messageTime                    = metricRegistry.histogram(name(MessageController.class, "message_delivery_duration"));
@@ -52,14 +52,11 @@ public class WebSocketConnection implements DispatchChannel, MessageAvailability
   private static final Meter          messageAvailableMeter          = metricRegistry.meter(name(WebSocketConnection.class, "messagesAvailable"));
   private static final Meter          ephemeralMessageAvailableMeter = metricRegistry.meter(name(WebSocketConnection.class, "ephemeralMessagesAvailable"));
   private static final Meter          messagesPersistedMeter         = metricRegistry.meter(name(WebSocketConnection.class, "messagesPersisted"));
-  private static final Meter          pubSubNewMessageMeter          = metricRegistry.meter(name(WebSocketConnection.class, "pubSubNewMessage"));
-  private static final Meter          pubSubPersistedMeter           = metricRegistry.meter(name(WebSocketConnection.class, "pubSubPersisted"));
   private static final Meter          displacementMeter              = metricRegistry.meter(name(WebSocketConnection.class, "explicitDisplacement"));
 
   private static final Logger logger = LoggerFactory.getLogger(WebSocketConnection.class);
 
   private final ReceiptSender    receiptSender;
-  private final PushSender       pushSender;
   private final MessagesManager  messagesManager;
 
   private final Account          account;
@@ -77,15 +74,13 @@ public class WebSocketConnection implements DispatchChannel, MessageAvailability
     PERSISTED_NEW_MESSAGES_AVAILABLE
   }
 
-  public WebSocketConnection(PushSender pushSender,
+  public WebSocketConnection(ReceiptSender receiptSender,
-                             ReceiptSender receiptSender,
                              MessagesManager messagesManager,
                              Account account,
                              Device device,
                              WebSocketClient client,
                              String connectionId)
   {
-    this.pushSender      = pushSender;
     this.receiptSender   = receiptSender;
     this.messagesManager = messagesManager;
     this.account         = account;
@@ -94,38 +89,14 @@ public class WebSocketConnection implements DispatchChannel, MessageAvailability
     this.connectionId    = connectionId;
   }
 
-  @Override
+  public void start() {
-  public void onDispatchMessage(String channel, byte[] message) {
-    try {
-      PubSubMessage pubSubMessage = PubSubMessage.parseFrom(message);
-
-      switch (pubSubMessage.getType().getNumber()) {
-        case PubSubMessage.Type.QUERY_DB_VALUE:
-          pubSubPersistedMeter.mark();
-          storedMessageState.set(StoredMessageState.PERSISTED_NEW_MESSAGES_AVAILABLE);
-          processStoredMessages();
-          break;
-        case PubSubMessage.Type.DELIVER_VALUE:
-          pubSubNewMessageMeter.mark();
-          sendMessage(Envelope.parseFrom(pubSubMessage.getContent()), Optional.empty());
-          break;
-        default:
-          logger.warn("Unknown pubsub message: " + pubSubMessage.getType().getNumber());
-      }
-    } catch (InvalidProtocolBufferException e) {
-      logger.warn("Protobuf parse error", e);
-    }
-  }
-
-  @Override
-  public void onDispatchUnsubscribed(String channel) {
-    client.close(1000, "OK");
-  }
-
-  public void onDispatchSubscribed(String channel) {
     processStoredMessages();
   }
 
+  public void stop() {
+    client.close(1000, "OK");
+  }
+
   private CompletableFuture<WebSocketResponseMessage> sendMessage(final Envelope message, final Optional<StoredMessageInfo> storedMessageInfo) {
     try {
       String           header;
@@ -143,20 +114,16 @@ public class WebSocketConnection implements DispatchChannel, MessageAvailability
 
       return client.sendRequest("PUT", "/api/v1/message", List.of(header, TimestampHeaderUtil.getTimestampHeader()), body).whenComplete((response, throwable) -> {
         if (throwable == null) {
-          boolean isReceipt = message.getType() == Envelope.Type.RECEIPT;
-
-          if (isSuccessResponse(response) && !isReceipt) {
-            messageTime.update(System.currentTimeMillis() - message.getTimestamp());
-          }
-
           if (isSuccessResponse(response)) {
-            if (storedMessageInfo.isPresent()) messagesManager.delete(account.getNumber(), account.getUuid(), device.getId(), storedMessageInfo.get().id, storedMessageInfo.get().cached);
+            if (storedMessageInfo.isPresent()) {
-            if (!isReceipt)                    sendDeliveryReceiptFor(message);
+              messagesManager.delete(account.getNumber(), account.getUuid(), device.getId(), storedMessageInfo.get().id, storedMessageInfo.get().cached);
-          } else if (!isSuccessResponse(response) && !storedMessageInfo.isPresent()) {
+            }
-            requeueMessage(message);
+
+            if (message.getType() != Envelope.Type.RECEIPT) {
+              messageTime.update(System.currentTimeMillis() - message.getTimestamp());
+              sendDeliveryReceiptFor(message);
+            }
           }
-        } else {
-          if (!storedMessageInfo.isPresent()) requeueMessage(message);
         }
       });
     } catch (CryptoEncodingException e) {
@@ -165,16 +132,6 @@ public class WebSocketConnection implements DispatchChannel, MessageAvailability
     }
   }
 
-  private void requeueMessage(Envelope message) {
-    pushSender.getWebSocketSender().queueMessage(account, device, message);
-
-    try {
-      pushSender.sendQueuedNotification(account, device);
-    } catch (NotPushRegisteredException e) {
-      logger.warn("requeueMessage", e);
-    }
-  }
-
   private void sendDeliveryReceiptFor(Envelope message) {
     if (!message.hasSource()) return;
 

service/src/test/java/org/whispersystems/textsecuregcm/push/ClientPresenceManagerTest.java

@@ -59,6 +59,19 @@ public class ClientPresenceManagerTest extends AbstractRedisClusterTest {
         assertTrue(clientPresenceManager.isPresent(accountUuid, deviceId));
     }
 
+    @Test
+    public void testIsLocallyPresent() {
+        final UUID accountUuid = UUID.randomUUID();
+        final long deviceId    = 1;
+
+        assertFalse(clientPresenceManager.isLocallyPresent(accountUuid, deviceId));
+
+        clientPresenceManager.setPresent(accountUuid, deviceId, NO_OP);
+        getRedisCluster().useCluster(connection -> connection.sync().flushall());
+
+        assertTrue(clientPresenceManager.isLocallyPresent(accountUuid, deviceId));
+    }
+
     @Test
     public void testLocalDisplacement() {
         final UUID accountUuid = UUID.randomUUID();

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

@@ -14,7 +14,6 @@ import org.mockito.stubbing.Answer;
 import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
 import org.whispersystems.textsecuregcm.entities.MessageProtos;
 import org.whispersystems.textsecuregcm.metrics.PushLatencyManager;
-import org.whispersystems.textsecuregcm.push.PushSender;
 import org.whispersystems.textsecuregcm.push.ReceiptSender;
 import org.whispersystems.textsecuregcm.redis.AbstractRedisClusterTest;
 import org.whispersystems.textsecuregcm.storage.Account;
@@ -78,7 +77,7 @@ public class WebSocketConnectionIntegrationTest extends AbstractRedisClusterTest
         when(account.getUuid()).thenReturn(UUID.randomUUID());
         when(device.getId()).thenReturn(1L);
 
-        webSocketConnection = new WebSocketConnection(mock(PushSender.class),
+        webSocketConnection = new WebSocketConnection(
                 mock(ReceiptSender.class),
                 new MessagesManager(messages, messagesCache, mock(PushLatencyManager.class)),
                 account,

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

@@ -12,15 +12,12 @@ import org.whispersystems.textsecuregcm.entities.OutgoingMessageEntity;
 import org.whispersystems.textsecuregcm.entities.OutgoingMessageEntityList;
 import org.whispersystems.textsecuregcm.push.ApnFallbackManager;
 import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
-import org.whispersystems.textsecuregcm.push.PushSender;
 import org.whispersystems.textsecuregcm.push.ReceiptSender;
 import org.whispersystems.textsecuregcm.push.WebsocketSender;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.MessagesManager;
-import org.whispersystems.textsecuregcm.storage.PubSubManager;
-import org.whispersystems.textsecuregcm.storage.PubSubProtos;
 import org.whispersystems.textsecuregcm.util.Base64;
 import org.whispersystems.websocket.WebSocketClient;
 import org.whispersystems.websocket.auth.WebSocketAuthenticator.AuthenticationResult;
@@ -68,11 +65,9 @@ public class WebSocketConnectionTest {
 
   private static final AccountAuthenticator accountAuthenticator = mock(AccountAuthenticator.class);
   private static final AccountsManager      accountsManager      = mock(AccountsManager.class);
-  private static final PubSubManager        pubSubManager        = mock(PubSubManager.class       );
   private static final Account              account              = mock(Account.class             );
   private static final Device               device               = mock(Device.class              );
   private static final UpgradeRequest       upgradeRequest       = mock(UpgradeRequest.class      );
-  private static final PushSender           pushSender           = mock(PushSender.class);
   private static final ReceiptSender        receiptSender        = mock(ReceiptSender.class);
   private static final ApnFallbackManager   apnFallbackManager   = mock(ApnFallbackManager.class);
 
@@ -80,7 +75,7 @@ public class WebSocketConnectionTest {
   public void testCredentials() throws Exception {
     MessagesManager               storedMessages         = mock(MessagesManager.class);
     WebSocketAccountAuthenticator webSocketAuthenticator = new WebSocketAccountAuthenticator(accountAuthenticator);
-    AuthenticatedConnectListener  connectListener        = new AuthenticatedConnectListener(pushSender, receiptSender, storedMessages, pubSubManager, apnFallbackManager, mock(ClientPresenceManager.class));
+    AuthenticatedConnectListener  connectListener        = new AuthenticatedConnectListener(receiptSender, storedMessages, apnFallbackManager, mock(ClientPresenceManager.class));
     WebSocketSessionContext       sessionContext         = mock(WebSocketSessionContext.class);
 
     when(accountAuthenticator.authenticate(eq(new BasicCredentials(VALID_USER, VALID_PASSWORD))))
@@ -167,11 +162,10 @@ public class WebSocketConnectionTest {
           }
         });
 
-    WebsocketAddress websocketAddress = new WebsocketAddress(account.getNumber(), device.getId());
+    WebSocketConnection connection = new WebSocketConnection(receiptSender, storedMessages,
-    WebSocketConnection connection = new WebSocketConnection(pushSender, receiptSender, storedMessages,
                                                              account, device, client, "someid");
 
-    connection.onDispatchSubscribed(websocketAddress.serialize());
+    connection.start();
     verify(client, times(3)).sendRequest(eq("PUT"), eq("/api/v1/message"), ArgumentMatchers.nullable(List.class), ArgumentMatchers.<Optional<byte[]>>any());
 
     assertTrue(futures.size() == 3);
@@ -186,111 +180,15 @@ public class WebSocketConnectionTest {
     verify(storedMessages, times(1)).delete(eq(account.getNumber()), eq(accountUuid), eq(2L), eq(2L), eq(false));
     verify(receiptSender, times(1)).sendReceipt(eq(account), eq("sender1"), eq(2222L));
 
-    connection.onDispatchUnsubscribed(websocketAddress.serialize());
+    connection.stop();
-    verify(client).close(anyInt(), anyString());
-  }
-
-  @Test
-  public void testOnlineSend() throws Exception {
-    MessagesManager storedMessages = mock(MessagesManager.class);
-    WebsocketSender websocketSender = mock(WebsocketSender.class);
-
-    when(pushSender.getWebSocketSender()).thenReturn(websocketSender);
-
-    Envelope firstMessage = Envelope.newBuilder()
-                                    .setLegacyMessage(ByteString.copyFrom("first".getBytes()))
-                                    .setSource("sender1")
-                                    .setTimestamp(System.currentTimeMillis())
-                                    .setSourceDevice(1)
-                                    .setType(Envelope.Type.CIPHERTEXT)
-                                    .build();
-
-    Envelope secondMessage = Envelope.newBuilder()
-                                     .setLegacyMessage(ByteString.copyFrom("second".getBytes()))
-                                     .setSource("sender2")
-                                     .setTimestamp(System.currentTimeMillis())
-                                     .setSourceDevice(2)
-                                     .setType(Envelope.Type.CIPHERTEXT)
-                                     .build();
-
-    List<OutgoingMessageEntity> pendingMessages     = new LinkedList<>();
-    OutgoingMessageEntityList   pendingMessagesList = new OutgoingMessageEntityList(pendingMessages, false);
-
-    when(device.getId()).thenReturn(2L);
-    when(device.getSignalingKey()).thenReturn(Base64.encodeBytes(new byte[52]));
-
-    when(account.getAuthenticatedDevice()).thenReturn(Optional.of(device));
-    when(account.getNumber()).thenReturn("+14152222222");
-    when(account.getUuid()).thenReturn(UUID.randomUUID());
-
-    final Device sender1device = mock(Device.class);
-
-    Set<Device> sender1devices = new HashSet<Device>() {{
-      add(sender1device);
-    }};
-
-    Account sender1 = mock(Account.class);
-    when(sender1.getDevices()).thenReturn(sender1devices);
-
-    when(accountsManager.get("sender1")).thenReturn(Optional.of(sender1));
-    when(accountsManager.get("sender2")).thenReturn(Optional.<Account>empty());
-
-    String userAgent = "user-agent";
-
-    when(storedMessages.getMessagesForDevice(account.getNumber(), account.getUuid(), device.getId(), userAgent, false))
-        .thenReturn(pendingMessagesList);
-
-    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) throws Throwable {
-            CompletableFuture<WebSocketResponseMessage> future = new CompletableFuture<>();
-            futures.add(future);
-            return future;
-          }
-        });
-
-    WebsocketAddress websocketAddress = new WebsocketAddress(account.getNumber(), device.getId());
-    WebSocketConnection connection = new WebSocketConnection(pushSender, receiptSender, storedMessages,
-                                                             account, device, client, "anotherid");
-
-    connection.onDispatchSubscribed(websocketAddress.serialize());
-    connection.onDispatchMessage(websocketAddress.serialize(), PubSubProtos.PubSubMessage.newBuilder()
-                                                                                         .setType(PubSubProtos.PubSubMessage.Type.DELIVER)
-                                                                                         .setContent(ByteString.copyFrom(firstMessage.toByteArray()))
-                                                                                         .build().toByteArray());
-
-    connection.onDispatchMessage(websocketAddress.serialize(), PubSubProtos.PubSubMessage.newBuilder()
-                                                                                         .setType(PubSubProtos.PubSubMessage.Type.DELIVER)
-                                                                                         .setContent(ByteString.copyFrom(secondMessage.toByteArray()))
-                                                                                         .build().toByteArray());
-
-    verify(client, times(2)).sendRequest(eq("PUT"), eq("/api/v1/message"), ArgumentMatchers.nullable(List.class), ArgumentMatchers.<Optional<byte[]>>any());
-
-    assertEquals(futures.size(), 2);
-
-    WebSocketResponseMessage response = mock(WebSocketResponseMessage.class);
-    when(response.getStatus()).thenReturn(200);
-    futures.get(1).complete(response);
-    futures.get(0).completeExceptionally(new IOException());
-
-    verify(receiptSender, times(1)).sendReceipt(eq(account), eq("sender2"), eq(secondMessage.getTimestamp()));
-    verify(websocketSender, times(1)).queueMessage(eq(account), eq(device), any(Envelope.class));
-    verify(pushSender, times(1)).sendQueuedNotification(eq(account), eq(device));
-
-    connection.onDispatchUnsubscribed(websocketAddress.serialize());
     verify(client).close(anyInt(), anyString());
   }
 
   @Test(timeout = 5_000L)
-  public void testOnlineSendViaKeyspaceNotification() throws Exception {
+  public void testOnlineSend() throws Exception {
     final MessagesManager     messagesManager = mock(MessagesManager.class);
     final WebSocketClient     client          = mock(WebSocketClient.class);
-    final WebSocketConnection connection      = new WebSocketConnection(pushSender, receiptSender, messagesManager, account, device, client, "concurrency");
+    final WebSocketConnection connection      = new WebSocketConnection(receiptSender, messagesManager, account, device, client, "concurrency");
 
     final UUID accountUuid = UUID.randomUUID();
 
@@ -322,7 +220,7 @@ public class WebSocketConnectionTest {
     // messages, the call to CompletableFuture.allOf(...) in processStoredMessages will produce an instantly-succeeded
     // future, and the whenComplete method will get called immediately on THIS thread, so we don't need to synchronize
     // or wait for anything.
-    connection.onDispatchSubscribed("channel");
+    connection.start();
 
     connection.handleNewMessagesAvailable();
 
@@ -349,11 +247,6 @@ public class WebSocketConnectionTest {
     MessagesManager storedMessages  = mock(MessagesManager.class);
     WebsocketSender websocketSender = mock(WebsocketSender.class);
 
-    reset(websocketSender);
-    reset(pushSender);
-
-    when(pushSender.getWebSocketSender()).thenReturn(websocketSender);
-
     final Envelope firstMessage = Envelope.newBuilder()
                                     .setLegacyMessage(ByteString.copyFrom("first".getBytes()))
                                     .setSource("sender1")
@@ -423,11 +316,10 @@ public class WebSocketConnectionTest {
           }
         });
 
-    WebsocketAddress websocketAddress = new WebsocketAddress(account.getNumber(), device.getId());
+    WebSocketConnection connection = new WebSocketConnection(receiptSender, storedMessages,
-    WebSocketConnection connection = new WebSocketConnection(pushSender, receiptSender, storedMessages,
                                                              account, device, client, "onemoreid");
 
-    connection.onDispatchSubscribed(websocketAddress.serialize());
+    connection.start();
 
     verify(client, times(2)).sendRequest(eq("PUT"), eq("/api/v1/message"), ArgumentMatchers.nullable(List.class), ArgumentMatchers.<Optional<byte[]>>any());
 
@@ -440,9 +332,8 @@ public class WebSocketConnectionTest {
 
     verify(receiptSender, times(1)).sendReceipt(eq(account), eq("sender2"), eq(secondMessage.getTimestamp()));
     verifyNoMoreInteractions(websocketSender);
-    verifyNoMoreInteractions(pushSender);
 
-    connection.onDispatchUnsubscribed(websocketAddress.serialize());
+    connection.stop();
     verify(client).close(anyInt(), anyString());
   }
 
@@ -450,7 +341,7 @@ public class WebSocketConnectionTest {
   public void testProcessStoredMessageConcurrency() throws InterruptedException {
     final MessagesManager     messagesManager = mock(MessagesManager.class);
     final WebSocketClient     client          = mock(WebSocketClient.class);
-    final WebSocketConnection connection      = new WebSocketConnection(pushSender, receiptSender, messagesManager, account, device, client, "concurrency");
+    final WebSocketConnection connection      = new WebSocketConnection(receiptSender, messagesManager, account, device, client, "concurrency");
 
     when(account.getNumber()).thenReturn("+18005551234");
     when(account.getUuid()).thenReturn(UUID.randomUUID());
@@ -511,7 +402,7 @@ public class WebSocketConnectionTest {
   public void testProcessStoredMessagesMultiplePages() throws InterruptedException {
     final MessagesManager     messagesManager = mock(MessagesManager.class);
     final WebSocketClient     client          = mock(WebSocketClient.class);
-    final WebSocketConnection connection      = new WebSocketConnection(pushSender, receiptSender, messagesManager, account, device, client, "concurrency");
+    final WebSocketConnection connection      = new WebSocketConnection(receiptSender, messagesManager, account, device, client, "concurrency");
 
     when(account.getNumber()).thenReturn("+18005551234");
     when(account.getUuid()).thenReturn(UUID.randomUUID());
@@ -554,7 +445,7 @@ public class WebSocketConnectionTest {
   public void testProcessStoredMessagesSingleEmptyCall() {
     final MessagesManager     messagesManager = mock(MessagesManager.class);
     final WebSocketClient     client          = mock(WebSocketClient.class);
-    final WebSocketConnection connection      = new WebSocketConnection(pushSender, receiptSender, messagesManager, account, device, client, "concurrency");
+    final WebSocketConnection connection      = new WebSocketConnection(receiptSender, messagesManager, account, device, client, "concurrency");
 
     final UUID accountUuid = UUID.randomUUID();
 
@@ -583,7 +474,7 @@ public class WebSocketConnectionTest {
   public void testRequeryOnStateMismatch() throws InterruptedException {
     final MessagesManager     messagesManager = mock(MessagesManager.class);
     final WebSocketClient     client          = mock(WebSocketClient.class);
-    final WebSocketConnection connection      = new WebSocketConnection(pushSender, receiptSender, messagesManager, account, device, client, "concurrency");
+    final WebSocketConnection connection      = new WebSocketConnection(receiptSender, messagesManager, account, device, client, "concurrency");
     final UUID                accountUuid     = UUID.randomUUID();
 
     when(account.getNumber()).thenReturn("+18005551234");
@@ -609,15 +500,10 @@ public class WebSocketConnectionTest {
     final WebSocketResponseMessage successResponse = mock(WebSocketResponseMessage.class);
     when(successResponse.getStatus()).thenReturn(200);
 
-    final byte[] queryDbMessageBytes = PubSubProtos.PubSubMessage.newBuilder()
-                                                                 .setType(PubSubProtos.PubSubMessage.Type.QUERY_DB)
-                                                                 .build()
-                                                                 .toByteArray();
-
     final CountDownLatch sendLatch = new CountDownLatch(firstPageMessages.size() + secondPageMessages.size());
 
     when(client.sendRequest(eq("PUT"), eq("/api/v1/message"), any(List.class), any(Optional.class))).thenAnswer((Answer<CompletableFuture<WebSocketResponseMessage>>)invocation -> {
-      connection.onDispatchMessage("channel", queryDbMessageBytes);
+      connection.handleNewMessagesAvailable();
       sendLatch.countDown();
 
       return CompletableFuture.completedFuture(successResponse);
@@ -635,7 +521,7 @@ public class WebSocketConnectionTest {
   public void testProcessCachedMessagesOnly() {
     final MessagesManager     messagesManager = mock(MessagesManager.class);
     final WebSocketClient     client          = mock(WebSocketClient.class);
-    final WebSocketConnection connection      = new WebSocketConnection(pushSender, receiptSender, messagesManager, account, device, client, "concurrency");
+    final WebSocketConnection connection      = new WebSocketConnection(receiptSender, messagesManager, account, device, client, "concurrency");
 
     final UUID accountUuid = UUID.randomUUID();
 
@@ -667,7 +553,7 @@ public class WebSocketConnectionTest {
   public void testProcessDatabaseMessagesAfterPersist() {
     final MessagesManager     messagesManager = mock(MessagesManager.class);
     final WebSocketClient     client          = mock(WebSocketClient.class);
-    final WebSocketConnection connection      = new WebSocketConnection(pushSender, receiptSender, messagesManager, account, device, client, "concurrency");
+    final WebSocketConnection connection      = new WebSocketConnection(receiptSender, messagesManager, account, device, client, "concurrency");
 
     final UUID accountUuid = UUID.randomUUID();