Use a custom redis pubsub implementation rather than Jedis.

// FREEBIE

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

@@ -1,5 +1,8 @@
 package org.whispersystems.textsecuregcm.websocket;
 
+import com.codahale.metrics.Histogram;
+import com.codahale.metrics.MetricRegistry;
+import com.codahale.metrics.SharedMetricRegistries;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.push.PushSender;
@@ -8,14 +11,19 @@ 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.Constants;
 import org.whispersystems.textsecuregcm.util.Util;
 import org.whispersystems.websocket.session.WebSocketSessionContext;
 import org.whispersystems.websocket.setup.WebSocketConnectListener;
 
+import static com.codahale.metrics.MetricRegistry.name;
+
 public class AuthenticatedConnectListener implements WebSocketConnectListener {
 
-  private static final Logger logger = LoggerFactory.getLogger(WebSocketConnection.class);
+  private static final Logger         logger            = LoggerFactory.getLogger(WebSocketConnection.class);
+  private static final MetricRegistry metricRegistry    = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
+  private static final Histogram      durationHistogram = metricRegistry.histogram(name(WebSocketConnection.class, "connected_duration"));
+
 
   private final AccountsManager accountsManager;
   private final PushSender      pushSender;
@@ -33,23 +41,22 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
 
   @Override
   public void onWebSocketConnect(WebSocketSessionContext context) {
-    Account account = context.getAuthenticated(Account.class).get();
-    Device  device  = account.getAuthenticatedDevice().get();
+    final Account             account     = context.getAuthenticated(Account.class).get();
+    final Device              device      = account.getAuthenticatedDevice().get();
+    final long                connectTime = System.currentTimeMillis();
+    final WebsocketAddress    address     = new WebsocketAddress(account.getNumber(), device.getId());
+    final WebSocketConnection connection  = new WebSocketConnection(accountsManager, pushSender,
+                                                                    messagesManager, account, device,
+                                                                    context.getClient());
 
     updateLastSeen(account, device);
-    closeExistingDeviceConnection(account, device);
-
-    final WebSocketConnection connection = new WebSocketConnection(accountsManager, pushSender,
-                                                                   messagesManager, pubSubManager,
-                                                                   account, device,
-                                                                   context.getClient());
-
-    connection.onConnected();
+    pubSubManager.subscribe(address, connection);
 
     context.addListener(new WebSocketSessionContext.WebSocketEventListener() {
       @Override
       public void onWebSocketClose(WebSocketSessionContext context, int statusCode, String reason) {
-        connection.onConnectionLost();
+        pubSubManager.unsubscribe(address, connection);
+        durationHistogram.update(System.currentTimeMillis() - connectTime);
       }
     });
   }
@@ -60,12 +67,5 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
       accountsManager.update(account);
     }
   }
-
-  private void closeExistingDeviceConnection(Account account, Device device) {
-    pubSubManager.publish(new WebsocketAddress(account.getNumber(), device.getId()),
-                          PubSubProtos.PubSubMessage.newBuilder()
-                                                    .setType(PubSubProtos.PubSubMessage.Type.CLOSE)
-                                                    .build());
-  }
 }
 

src/main/java/org/whispersystems/textsecuregcm/websocket/DeadLetterHandler.java

@@ -3,11 +3,13 @@ package org.whispersystems.textsecuregcm.websocket;
 import com.google.protobuf.InvalidProtocolBufferException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.whispersystems.dispatch.DispatchChannel;
 import org.whispersystems.textsecuregcm.entities.MessageProtos.OutgoingMessageSignal;
 import org.whispersystems.textsecuregcm.storage.MessagesManager;
 import org.whispersystems.textsecuregcm.storage.PubSubProtos;
+import org.whispersystems.textsecuregcm.storage.PubSubProtos.PubSubMessage;
 
-public class DeadLetterHandler {
+public class DeadLetterHandler implements DispatchChannel {
 
   private final Logger logger = LoggerFactory.getLogger(DeadLetterHandler.class);
 
@@ -17,14 +19,16 @@ public class DeadLetterHandler {
     this.messagesManager = messagesManager;
   }
 
-  public void handle(byte[] channel, PubSubProtos.PubSubMessage pubSubMessage) {
+  @Override
+  public void onDispatchMessage(String channel, byte[] data) {
     try {
-      WebsocketAddress address = new WebsocketAddress(new String(channel));
+      logger.warn("Handling dead letter to: " + channel);
 
-      logger.warn("Handling dead letter to: " + address);
+      WebsocketAddress address       = new WebsocketAddress(channel);
+      PubSubMessage    pubSubMessage = PubSubMessage.parseFrom(data);
 
       switch (pubSubMessage.getType().getNumber()) {
-        case PubSubProtos.PubSubMessage.Type.DELIVER_VALUE:
+        case PubSubMessage.Type.DELIVER_VALUE:
           OutgoingMessageSignal message = OutgoingMessageSignal.parseFrom(pubSubMessage.getContent());
           messagesManager.insert(address.getNumber(), address.getDeviceId(), message);
           break;
@@ -36,4 +40,13 @@ public class DeadLetterHandler {
     }
   }
 
+  @Override
+  public void onDispatchSubscribed(String channel) {
+    logger.warn("DeadLetterHandler subscription notice! " + channel);
+  }
+
+  @Override
+  public void onDispatchUnsubscribed(String channel) {
+    logger.warn("DeadLetterHandler unsubscribe notice! " + channel);
+  }
 }

src/main/java/org/whispersystems/textsecuregcm/websocket/ProvisioningConnectListener.java

@@ -14,13 +14,15 @@ public class ProvisioningConnectListener implements WebSocketConnectListener {
 
   @Override
   public void onWebSocketConnect(WebSocketSessionContext context) {
-    final ProvisioningConnection connection = new ProvisioningConnection(pubSubManager, context.getClient());
-    connection.onConnected();
+    final ProvisioningConnection connection          = new ProvisioningConnection(context.getClient());
+    final ProvisioningAddress    provisioningAddress = ProvisioningAddress.generate();
+
+    pubSubManager.subscribe(provisioningAddress, connection);
 
     context.addListener(new WebSocketSessionContext.WebSocketEventListener() {
       @Override
       public void onWebSocketClose(WebSocketSessionContext context, int statusCode, String reason) {
-        connection.onConnectionLost();
+        pubSubManager.unsubscribe(provisioningAddress, connection);
       }
     });
   }

src/main/java/org/whispersystems/textsecuregcm/websocket/ProvisioningConnection.java

@@ -4,58 +4,62 @@ import com.google.common.base.Optional;
 import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
+import com.google.protobuf.InvalidProtocolBufferException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.whispersystems.dispatch.DispatchChannel;
 import org.whispersystems.textsecuregcm.entities.MessageProtos.ProvisioningUuid;
-import org.whispersystems.textsecuregcm.storage.PubSubListener;
-import org.whispersystems.textsecuregcm.storage.PubSubManager;
 import org.whispersystems.textsecuregcm.storage.PubSubProtos.PubSubMessage;
 import org.whispersystems.websocket.WebSocketClient;
 import org.whispersystems.websocket.messages.WebSocketResponseMessage;
 
-public class ProvisioningConnection implements PubSubListener {
+public class ProvisioningConnection implements DispatchChannel {
 
-  private final PubSubManager       pubSubManager;
-  private final ProvisioningAddress provisioningAddress;
-  private final WebSocketClient     client;
+  private final Logger logger = LoggerFactory.getLogger(ProvisioningConnection.class);
 
-  public ProvisioningConnection(PubSubManager pubSubManager, WebSocketClient client) {
-    this.pubSubManager       = pubSubManager;
-    this.client              = client;
-    this.provisioningAddress = ProvisioningAddress.generate();
+  private final WebSocketClient client;
+
+  public ProvisioningConnection(WebSocketClient client) {
+    this.client = client;
   }
 
   @Override
-  public void onPubSubMessage(PubSubMessage outgoingMessage) {
-    if (outgoingMessage.getType() == PubSubMessage.Type.DELIVER) {
-      Optional<byte[]> body = Optional.of(outgoingMessage.getContent().toByteArray());
+  public void onDispatchMessage(String channel, byte[] message) {
+    try {
+      PubSubMessage outgoingMessage = PubSubMessage.parseFrom(message);
 
-      ListenableFuture<WebSocketResponseMessage> response = client.sendRequest("PUT", "/v1/message", body);
+      if (outgoingMessage.getType() == PubSubMessage.Type.DELIVER) {
+        Optional<byte[]> body = Optional.of(outgoingMessage.getContent().toByteArray());
 
-      Futures.addCallback(response, new FutureCallback<WebSocketResponseMessage>() {
-        @Override
-        public void onSuccess(WebSocketResponseMessage webSocketResponseMessage) {
-          pubSubManager.unsubscribe(provisioningAddress, ProvisioningConnection.this);
-          client.close(1001, "All you get.");
-        }
+        ListenableFuture<WebSocketResponseMessage> response = client.sendRequest("PUT", "/v1/message", body);
 
-        @Override
-        public void onFailure(Throwable throwable) {
-          pubSubManager.unsubscribe(provisioningAddress, ProvisioningConnection.this);
-          client.close(1001, "That's all!");
-        }
-      });
+        Futures.addCallback(response, new FutureCallback<WebSocketResponseMessage>() {
+          @Override
+          public void onSuccess(WebSocketResponseMessage webSocketResponseMessage) {
+            client.close(1001, "All you get.");
+          }
+
+          @Override
+          public void onFailure(Throwable throwable) {
+            client.close(1001, "That's all!");
+          }
+        });
+      }
+    } catch (InvalidProtocolBufferException e) {
+      logger.warn("Protobuf Error: ", e);
     }
   }
 
-  public void onConnected() {
-    this.pubSubManager.subscribe(provisioningAddress, this);
+  @Override
+  public void onDispatchSubscribed(String channel) {
     this.client.sendRequest("PUT", "/v1/address", Optional.of(ProvisioningUuid.newBuilder()
-                                                                              .setUuid(provisioningAddress.getAddress())
+                                                                              .setUuid(channel)
                                                                               .build()
                                                                               .toByteArray()));
   }
 
-  public void onConnectionLost() {
-    this.pubSubManager.unsubscribe(provisioningAddress, this);
-    this.client.close(1001, "Done");
+  @Override
+  public void onDispatchUnsubscribed(String channel) {
+    this.client.close(1001, "Closed");
   }
 }

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

@@ -10,6 +10,7 @@ import com.google.common.util.concurrent.ListenableFuture;
 import com.google.protobuf.InvalidProtocolBufferException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.whispersystems.dispatch.DispatchChannel;
 import org.whispersystems.textsecuregcm.entities.CryptoEncodingException;
 import org.whispersystems.textsecuregcm.entities.EncryptedOutgoingMessage;
 import org.whispersystems.textsecuregcm.push.NotPushRegisteredException;
@@ -19,8 +20,6 @@ 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.PubSubListener;
-import org.whispersystems.textsecuregcm.storage.PubSubManager;
 import org.whispersystems.textsecuregcm.util.Constants;
 import org.whispersystems.textsecuregcm.util.Pair;
 import org.whispersystems.websocket.WebSocketClient;
@@ -34,21 +33,16 @@ import static com.codahale.metrics.MetricRegistry.name;
 import static org.whispersystems.textsecuregcm.entities.MessageProtos.OutgoingMessageSignal;
 import static org.whispersystems.textsecuregcm.storage.PubSubProtos.PubSubMessage;
 
-public class WebSocketConnection implements PubSubListener {
+public class WebSocketConnection implements DispatchChannel {
 
   private static final Logger logger = LoggerFactory.getLogger(WebSocketConnection.class);
 
-  private static final MetricRegistry metricRegistry    = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
-  private static final Histogram      durationHistogram = metricRegistry.histogram(name(WebSocketConnection.class, "connected_duration"));
-
   private final AccountsManager  accountsManager;
   private final PushSender       pushSender;
   private final MessagesManager  messagesManager;
-  private final PubSubManager    pubSubManager;
 
   private final Account          account;
   private final Device           device;
-  private final WebsocketAddress address;
   private final WebSocketClient  client;
 
   private long connectionStartTime;
@@ -56,7 +50,6 @@ public class WebSocketConnection implements PubSubListener {
   public WebSocketConnection(AccountsManager accountsManager,
                              PushSender pushSender,
                              MessagesManager messagesManager,
-                             PubSubManager pubSubManager,
                              Account account,
                              Device device,
                              WebSocketClient client)
@@ -64,27 +57,16 @@ public class WebSocketConnection implements PubSubListener {
     this.accountsManager = accountsManager;
     this.pushSender      = pushSender;
     this.messagesManager = messagesManager;
-    this.pubSubManager   = pubSubManager;
     this.account         = account;
     this.device          = device;
     this.client          = client;
-    this.address         = new WebsocketAddress(account.getNumber(), device.getId());
-  }
-
-  public void onConnected() {
-    connectionStartTime = System.currentTimeMillis();
-    pubSubManager.subscribe(address, this);
-    processStoredMessages();
-  }
-
-  public void onConnectionLost() {
-    durationHistogram.update(System.currentTimeMillis() - connectionStartTime);
-    pubSubManager.unsubscribe(address, this);
   }
 
   @Override
-  public void onPubSubMessage(PubSubMessage pubSubMessage) {
+  public void onDispatchMessage(String channel, byte[] message) {
     try {
+      PubSubMessage pubSubMessage = PubSubMessage.parseFrom(message);
+
       switch (pubSubMessage.getType().getNumber()) {
         case PubSubMessage.Type.QUERY_DB_VALUE:
           processStoredMessages();
@@ -92,10 +74,6 @@ public class WebSocketConnection implements PubSubListener {
         case PubSubMessage.Type.DELIVER_VALUE:
           sendMessage(OutgoingMessageSignal.parseFrom(pubSubMessage.getContent()), Optional.<Long>absent());
           break;
-        case PubSubMessage.Type.CLOSE_VALUE:
-          client.close(1000, "OK");
-          pubSubManager.unsubscribe(address, this);
-          break;
         default:
           logger.warn("Unknown pubsub message: " + pubSubMessage.getType().getNumber());
       }
@@ -104,6 +82,15 @@ public class WebSocketConnection implements PubSubListener {
     }
   }
 
+  @Override
+  public void onDispatchUnsubscribed(String channel) {
+    client.close(1000, "OK");
+  }
+
+  public void onDispatchSubscribed(String channel) {
+    processStoredMessages();
+  }
+
   private void sendMessage(final OutgoingMessageSignal message,
                            final Optional<Long> storedMessageId)
   {
@@ -180,4 +167,6 @@ public class WebSocketConnection implements PubSubListener {
       sendMessage(message.second(), Optional.of(message.first()));
     }
   }
+
+
 }