Support for APN fallback retries.

// FREEBIE

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

@@ -55,6 +55,7 @@ import org.whispersystems.textsecuregcm.metrics.NetworkSentGauge;
 import org.whispersystems.textsecuregcm.providers.RedisClientFactory;
 import org.whispersystems.textsecuregcm.providers.RedisHealthCheck;
 import org.whispersystems.textsecuregcm.providers.TimeProvider;
+import org.whispersystems.textsecuregcm.push.ApnFallbackManager;
 import org.whispersystems.textsecuregcm.push.FeedbackHandler;
 import org.whispersystems.textsecuregcm.push.PushSender;
 import org.whispersystems.textsecuregcm.push.PushServiceClient;
@@ -170,15 +171,17 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     AccountAuthenticator   deviceAuthenticator    = new AccountAuthenticator(accountsManager);
     RateLimiters           rateLimiters           = new RateLimiters(config.getLimitsConfiguration(), cacheClient);
 
+    ApnFallbackManager       apnFallbackManager  = new ApnFallbackManager(pushServiceClient);
     TwilioSmsSender          twilioSmsSender     = new TwilioSmsSender(config.getTwilioConfiguration());
     Optional<NexmoSmsSender> nexmoSmsSender      = initializeNexmoSmsSender(config.getNexmoConfiguration());
     SmsSender                smsSender           = new SmsSender(twilioSmsSender, nexmoSmsSender, config.getTwilioConfiguration().isInternational());
     UrlSigner                urlSigner           = new UrlSigner(config.getS3Configuration());
-    PushSender               pushSender          = new PushSender(pushServiceClient, websocketSender);
+    PushSender               pushSender          = new PushSender(apnFallbackManager, pushServiceClient, websocketSender);
     ReceiptSender            receiptSender       = new ReceiptSender(accountsManager, pushSender, federatedClientManager);
     FeedbackHandler          feedbackHandler     = new FeedbackHandler(pushServiceClient, accountsManager);
     Optional<byte[]>         authorizationKey    = config.getRedphoneConfiguration().getAuthorizationKey();
 
+    environment.lifecycle().manage(apnFallbackManager);
     environment.lifecycle().manage(pubSubManager);
     environment.lifecycle().manage(feedbackHandler);
 
@@ -207,7 +210,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     if (config.getWebsocketConfiguration().isEnabled()) {
       WebSocketEnvironment webSocketEnvironment = new WebSocketEnvironment(environment, config, 90000);
       webSocketEnvironment.setAuthenticator(new WebSocketAccountAuthenticator(deviceAuthenticator));
-      webSocketEnvironment.setConnectListener(new AuthenticatedConnectListener(accountsManager, pushSender, receiptSender, messagesManager, pubSubManager));
+      webSocketEnvironment.setConnectListener(new AuthenticatedConnectListener(accountsManager, pushSender, receiptSender, messagesManager, pubSubManager, apnFallbackManager));
       webSocketEnvironment.jersey().register(new KeepAliveController(pubSubManager));
 
       WebSocketEnvironment provisioningEnvironment = new WebSocketEnvironment(environment, config);

src/main/java/org/whispersystems/textsecuregcm/entities/ApnMessage.java

@@ -37,4 +37,22 @@ public class ApnMessage {
     this.message  = message;
     this.voip     = voip;
   }
+
+  public ApnMessage(ApnMessage copy, String apnId, boolean voip) {
+    this.apnId    = apnId;
+    this.number   = copy.number;
+    this.deviceId = copy.deviceId;
+    this.message  = copy.message;
+    this.voip     = voip;
+  }
+
+  @VisibleForTesting
+  public String getApnId() {
+    return apnId;
+  }
+
+  @VisibleForTesting
+  public boolean isVoip() {
+    return voip;
+  }
 }

src/main/java/org/whispersystems/textsecuregcm/push/ApnFallbackManager.java

@@ -0,0 +1,150 @@
+package org.whispersystems.textsecuregcm.push;
+
+import com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.entities.ApnMessage;
+import org.whispersystems.textsecuregcm.util.Util;
+import org.whispersystems.textsecuregcm.websocket.WebsocketAddress;
+
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.Map.Entry;
+import java.util.concurrent.TimeUnit;
+
+import io.dropwizard.lifecycle.Managed;
+
+public class ApnFallbackManager implements Managed, Runnable {
+
+  private static final Logger logger = LoggerFactory.getLogger(ApnFallbackManager.class);
+
+  private final ApnFallbackTaskQueue taskQueue = new ApnFallbackTaskQueue();
+  private final PushServiceClient pushServiceClient;
+
+  public ApnFallbackManager(PushServiceClient pushServiceClient) {
+    this.pushServiceClient = pushServiceClient;
+  }
+
+  public void schedule(final WebsocketAddress address, ApnFallbackTask task) {
+    taskQueue.put(address, task);
+  }
+
+  public void cancel(WebsocketAddress address) {
+    taskQueue.remove(address);
+  }
+
+  @Override
+  public void start() throws Exception {
+    new Thread(this).start();
+  }
+
+  @Override
+  public void stop() throws Exception {
+
+  }
+
+  @Override
+  public void run() {
+    while (true) {
+      try {
+        Entry<WebsocketAddress, ApnFallbackTask> taskEntry  = taskQueue.get();
+        ApnFallbackTask                          task       = taskEntry.getValue();
+        int                                      retryCount = task.getRetryCount();
+
+        if (retryCount == 0) {
+          pushServiceClient.send(task.getMessage());
+          schedule(taskEntry.getKey(), new ApnFallbackTask(task.getApnId(), task.getMessage(),
+                                                           retryCount + 1, task.getDelay()));
+        } else if (retryCount == 1) {
+          pushServiceClient.send(new ApnMessage(task.getMessage(), task.getApnId(), false));
+        }
+      } catch (Throwable e) {
+        logger.warn("ApnFallbackThread", e);
+      }
+    }
+  }
+
+  public static class ApnFallbackTask {
+
+    private final long       delay;
+    private final long       executionTime;
+    private final String     apnId;
+    private final ApnMessage message;
+    private final int        retryCount;
+
+    public ApnFallbackTask(String apnId, ApnMessage message, int retryCount) {
+      this(apnId, message, retryCount, TimeUnit.SECONDS.toMillis(15));
+    }
+
+    @VisibleForTesting
+    public ApnFallbackTask(String apnId, ApnMessage message, int retryCount, long delay) {
+      this.executionTime = System.currentTimeMillis() + delay;
+      this.delay         = delay;
+      this.apnId         = apnId;
+      this.message       = message;
+      this.retryCount    = retryCount;
+    }
+
+    public String getApnId() {
+      return apnId;
+    }
+
+    public ApnMessage getMessage() {
+      return message;
+    }
+
+    public int getRetryCount() {
+      return retryCount;
+    }
+
+    public long getExecutionTime() {
+      return executionTime;
+    }
+
+    public long getDelay() {
+      return delay;
+    }
+  }
+
+  @VisibleForTesting
+  public static class ApnFallbackTaskQueue {
+
+    private final LinkedHashMap<WebsocketAddress, ApnFallbackTask> tasks = new LinkedHashMap<>();
+
+    public Entry<WebsocketAddress, ApnFallbackTask> get() {
+      while (true) {
+        long timeDelta;
+
+        synchronized (tasks) {
+          while (tasks.isEmpty()) Util.wait(tasks);
+
+          Iterator<Entry<WebsocketAddress, ApnFallbackTask>> iterator  = tasks.entrySet().iterator();
+          Entry<WebsocketAddress, ApnFallbackTask>           nextTask  = iterator.next();
+
+          timeDelta = nextTask.getValue().getExecutionTime() - System.currentTimeMillis();
+
+          if (timeDelta <= 0) {
+            iterator.remove();
+            return nextTask;
+          }
+        }
+
+        Util.sleep(timeDelta);
+      }
+    }
+
+    public void put(WebsocketAddress address, ApnFallbackTask task) {
+      synchronized (tasks) {
+        tasks.put(address, task);
+        tasks.notifyAll();
+      }
+    }
+
+    public void remove(WebsocketAddress address) {
+      synchronized (tasks) {
+        tasks.remove(address);
+      }
+    }
+  }
+
+}

src/main/java/org/whispersystems/textsecuregcm/push/PushSender.java

@@ -22,11 +22,16 @@ import org.whispersystems.textsecuregcm.entities.ApnMessage;
 import org.whispersystems.textsecuregcm.entities.CryptoEncodingException;
 import org.whispersystems.textsecuregcm.entities.EncryptedOutgoingMessage;
 import org.whispersystems.textsecuregcm.entities.GcmMessage;
+import org.whispersystems.textsecuregcm.push.ApnFallbackManager.ApnFallbackTask;
 import org.whispersystems.textsecuregcm.push.WebsocketSender.DeliveryStatus;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.util.Util;
+import org.whispersystems.textsecuregcm.websocket.WebsocketAddress;
 
+import java.util.LinkedHashMap;
+
+import io.dropwizard.lifecycle.Managed;
 import static org.whispersystems.textsecuregcm.entities.MessageProtos.OutgoingMessageSignal;
 
 public class PushSender {
@@ -35,12 +40,14 @@ public class PushSender {
 
   private static final String APN_PAYLOAD = "{\"aps\":{\"sound\":\"default\",\"badge\":%d,\"alert\":{\"loc-key\":\"APN_Message\"}}}";
 
-  private final PushServiceClient pushServiceClient;
+  private final ApnFallbackManager apnFallbackManager;
-  private final WebsocketSender   webSocketSender;
+  private final PushServiceClient  pushServiceClient;
+  private final WebsocketSender    webSocketSender;
 
-  public PushSender(PushServiceClient pushServiceClient, WebsocketSender websocketSender) {
+  public PushSender(ApnFallbackManager apnFallbackManager, PushServiceClient pushServiceClient, WebsocketSender websocketSender) {
-    this.pushServiceClient = pushServiceClient;
+    this.apnFallbackManager = apnFallbackManager;
-    this.webSocketSender   = websocketSender;
+    this.pushServiceClient  = pushServiceClient;
+    this.webSocketSender    = websocketSender;
   }
 
   public void sendMessage(Account account, Device device, OutgoingMessageSignal message)
@@ -106,6 +113,12 @@ public class PushSender {
       ApnMessage apnMessage = new ApnMessage(apnId, account.getNumber(), (int)device.getId(),
                                              String.format(APN_PAYLOAD, deliveryStatus.getMessageQueueDepth()),
                                              isVoip);
+
+      if (isVoip) {
+        apnFallbackManager.schedule(new WebsocketAddress(account.getNumber(), device.getId()),
+                                    new ApnFallbackTask(device.getApnId(), apnMessage, 0));
+      }
+
       pushServiceClient.send(apnMessage);
     }
   }

src/main/java/org/whispersystems/textsecuregcm/util/Util.java

@@ -115,12 +115,20 @@ public class Util {
     return parts;
   }
 
-  public static void sleep(int i) {
+  public static void sleep(long i) {
     try {
       Thread.sleep(i);
     } catch (InterruptedException ie) {}
   }
 
+  public static void wait(Object object) {
+    try {
+      object.wait();
+    } catch (InterruptedException e) {
+      throw new AssertionError(e);
+    }
+  }
+
   public static long todayInMillis() {
     return TimeUnit.DAYS.toMillis(TimeUnit.MILLISECONDS.toDays(System.currentTimeMillis()));
   }

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

@@ -5,6 +5,7 @@ import com.codahale.metrics.MetricRegistry;
 import com.codahale.metrics.SharedMetricRegistries;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.push.ApnFallbackManager;
 import org.whispersystems.textsecuregcm.push.PushSender;
 import org.whispersystems.textsecuregcm.push.ReceiptSender;
 import org.whispersystems.textsecuregcm.storage.Account;
@@ -25,22 +26,23 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
   private static final MetricRegistry metricRegistry    = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
   private static final Histogram      durationHistogram = metricRegistry.histogram(name(WebSocketConnection.class, "connected_duration"));
 
-
+  private final ApnFallbackManager apnFallbackManager;
-  private final AccountsManager accountsManager;
+  private final AccountsManager    accountsManager;
-  private final PushSender      pushSender;
+  private final PushSender         pushSender;
-  private final ReceiptSender   receiptSender;
+  private final ReceiptSender      receiptSender;
-  private final MessagesManager messagesManager;
+  private final MessagesManager    messagesManager;
-  private final PubSubManager   pubSubManager;
+  private final PubSubManager      pubSubManager;
 
   public AuthenticatedConnectListener(AccountsManager accountsManager, PushSender pushSender,
                                       ReceiptSender receiptSender,  MessagesManager messagesManager,
-                                      PubSubManager pubSubManager)
+                                      PubSubManager pubSubManager, ApnFallbackManager apnFallbackManager)
   {
-    this.accountsManager = accountsManager;
+    this.accountsManager    = accountsManager;
-    this.pushSender      = pushSender;
+    this.pushSender         = pushSender;
-    this.receiptSender   = receiptSender;
+    this.receiptSender      = receiptSender;
-    this.messagesManager = messagesManager;
+    this.messagesManager    = messagesManager;
-    this.pubSubManager   = pubSubManager;
+    this.pubSubManager      = pubSubManager;
+    this.apnFallbackManager = apnFallbackManager;
   }
 
   @Override
@@ -53,6 +55,7 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
                                                                     messagesManager, account, device,
                                                                     context.getClient());
 
+    apnFallbackManager.cancel(address);
     updateLastSeen(account, device);
     pubSubManager.subscribe(address, connection);
 

src/test/java/org/whispersystems/textsecuregcm/tests/push/ApnFallbackManagerTest.java

@@ -0,0 +1,83 @@
+package org.whispersystems.textsecuregcm.tests.push;
+
+import org.junit.Test;
+import org.mockito.ArgumentCaptor;
+import org.whispersystems.textsecuregcm.entities.ApnMessage;
+import org.whispersystems.textsecuregcm.push.ApnFallbackManager;
+import org.whispersystems.textsecuregcm.push.ApnFallbackManager.ApnFallbackTask;
+import org.whispersystems.textsecuregcm.push.PushServiceClient;
+import org.whispersystems.textsecuregcm.util.Util;
+import org.whispersystems.textsecuregcm.websocket.WebsocketAddress;
+
+import java.util.List;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.*;
+
+public class ApnFallbackManagerTest {
+
+  @Test
+  public void testFullFallback() throws Exception {
+    PushServiceClient  pushServiceClient  = mock(PushServiceClient.class);
+    WebsocketAddress   address            = mock(WebsocketAddress.class );
+    ApnMessage         message            = new ApnMessage("bar", "123", 1, "hmm", true);
+    ApnFallbackTask    task               = new ApnFallbackTask("foo", message, 0, 500);
+
+    ApnFallbackManager apnFallbackManager = new ApnFallbackManager(pushServiceClient);
+    apnFallbackManager.start();
+
+    apnFallbackManager.schedule(address, task);
+
+    Util.sleep(1100);
+
+    ArgumentCaptor<ApnMessage> captor = ArgumentCaptor.forClass(ApnMessage.class);
+    verify(pushServiceClient, times(2)).send(captor.capture());
+
+    List<ApnMessage> messages = captor.getAllValues();
+    assertEquals(messages.get(0), message);
+    assertEquals(messages.get(1).getApnId(), task.getApnId());
+    assertFalse(messages.get(1).isVoip());
+  }
+
+  @Test
+  public void testPartialFallback() throws Exception {
+    PushServiceClient  pushServiceClient  = mock(PushServiceClient.class);
+    WebsocketAddress   address            = mock(WebsocketAddress.class );
+    ApnMessage         message            = new ApnMessage("bar", "123", 1, "hmm", true);
+    ApnFallbackTask    task               = new ApnFallbackTask   ("foo", message, 0, 500);
+
+    ApnFallbackManager apnFallbackManager = new ApnFallbackManager(pushServiceClient);
+    apnFallbackManager.start();
+
+    apnFallbackManager.schedule(address, task);
+
+    Util.sleep(600);
+
+    apnFallbackManager.cancel(address);
+
+    Util.sleep(600);
+
+    verify(pushServiceClient, times(1)).send(eq(message));
+  }
+
+  @Test
+  public void testNoFallback() throws Exception {
+    PushServiceClient  pushServiceClient  = mock(PushServiceClient.class);
+    WebsocketAddress   address            = mock(WebsocketAddress.class );
+    ApnMessage         message            = new ApnMessage("bar", "123", 1, "hmm", true);
+    ApnFallbackTask    task               = new ApnFallbackTask   ("foo", message, 0, 500);
+
+    ApnFallbackManager apnFallbackManager = new ApnFallbackManager(pushServiceClient);
+    apnFallbackManager.start();
+
+    apnFallbackManager.schedule(address, task);
+    apnFallbackManager.cancel(address);
+
+    Util.sleep(1100);
+
+    verifyNoMoreInteractions(pushServiceClient);
+  }
+
+}

src/test/java/org/whispersystems/textsecuregcm/tests/push/ApnFallbackTaskQueueTest.java

@@ -0,0 +1,93 @@
+package org.whispersystems.textsecuregcm.tests.push;
+
+
+import org.junit.Test;
+import org.whispersystems.textsecuregcm.push.ApnFallbackManager;
+import org.whispersystems.textsecuregcm.push.ApnFallbackManager.ApnFallbackTask;
+import org.whispersystems.textsecuregcm.push.ApnFallbackManager.ApnFallbackTaskQueue;
+import org.whispersystems.textsecuregcm.util.Util;
+import org.whispersystems.textsecuregcm.websocket.WebsocketAddress;
+
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class ApnFallbackTaskQueueTest {
+
+  @Test
+  public void testBlocking() {
+    final ApnFallbackTaskQueue taskQueue = new ApnFallbackTaskQueue();
+
+    final WebsocketAddress address = mock(WebsocketAddress.class);
+    final ApnFallbackTask  task    = mock(ApnFallbackTask.class );
+
+    when(task.getExecutionTime()).thenReturn(System.currentTimeMillis() - 1000);
+
+    new Thread() {
+      @Override
+      public void run() {
+        Util.sleep(500);
+        taskQueue.put(address, task);
+      }
+    }.start();
+
+    Map.Entry<WebsocketAddress, ApnFallbackTask> result = taskQueue.get();
+
+    assertEquals(result.getKey(), address);
+    assertEquals(result.getValue(), task);
+  }
+
+  @Test
+  public void testElapsedTime() {
+    final ApnFallbackTaskQueue taskQueue = new ApnFallbackTaskQueue();
+    final WebsocketAddress     address   = mock(WebsocketAddress.class);
+    final ApnFallbackTask      task      = mock(ApnFallbackTask.class );
+
+    long currentTime = System.currentTimeMillis();
+
+    when(task.getExecutionTime()).thenReturn(currentTime + 1000);
+
+    taskQueue.put(address, task);
+    Map.Entry<WebsocketAddress, ApnFallbackTask> result = taskQueue.get();
+
+    assertTrue(System.currentTimeMillis() >= currentTime + 1000);
+    assertEquals(result.getKey(), address);
+    assertEquals(result.getValue(), task);
+  }
+
+  @Test
+  public void testCanceled() {
+    final ApnFallbackTaskQueue taskQueue = new ApnFallbackTaskQueue();
+    final WebsocketAddress     addressOne   = mock(WebsocketAddress.class);
+    final ApnFallbackTask      taskOne      = mock(ApnFallbackTask.class );
+    final WebsocketAddress     addressTwo   = mock(WebsocketAddress.class);
+    final ApnFallbackTask      taskTwo      = mock(ApnFallbackTask.class );
+
+    long currentTime = System.currentTimeMillis();
+
+    when(taskOne.getExecutionTime()).thenReturn(currentTime + 1000);
+    when(taskTwo.getExecutionTime()).thenReturn(currentTime + 2000);
+
+    taskQueue.put(addressOne, taskOne);
+    taskQueue.put(addressTwo, taskTwo);
+
+    new Thread() {
+      @Override
+      public void run() {
+        Util.sleep(300);
+        taskQueue.remove(addressOne);
+      }
+    }.start();
+
+    Map.Entry<WebsocketAddress, ApnFallbackTask> result = taskQueue.get();
+
+    assertTrue(System.currentTimeMillis() >= currentTime + 2000);
+    assertEquals(result.getKey(), addressTwo);
+    assertEquals(result.getValue(), taskTwo);
+  }
+
+
+}

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

@@ -10,6 +10,7 @@ import org.mockito.invocation.InvocationOnMock;
 import org.mockito.stubbing.Answer;
 import org.whispersystems.textsecuregcm.auth.AccountAuthenticator;
 import org.whispersystems.textsecuregcm.entities.OutgoingMessageEntity;
+import org.whispersystems.textsecuregcm.push.ApnFallbackManager;
 import org.whispersystems.textsecuregcm.push.PushSender;
 import org.whispersystems.textsecuregcm.push.ReceiptSender;
 import org.whispersystems.textsecuregcm.storage.Account;
@@ -60,12 +61,13 @@ public class WebSocketConnectionTest {
   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);
 
   @Test
   public void testCredentials() throws Exception {
     MessagesManager               storedMessages         = mock(MessagesManager.class);
     WebSocketAccountAuthenticator webSocketAuthenticator = new WebSocketAccountAuthenticator(accountAuthenticator);
-    AuthenticatedConnectListener  connectListener        = new AuthenticatedConnectListener(accountsManager, pushSender, receiptSender, storedMessages, pubSubManager);
+    AuthenticatedConnectListener  connectListener        = new AuthenticatedConnectListener(accountsManager, pushSender, receiptSender, storedMessages, pubSubManager, apnFallbackManager);
     WebSocketSessionContext       sessionContext         = mock(WebSocketSessionContext.class);
 
     when(accountAuthenticator.authenticate(eq(new BasicCredentials(VALID_USER, VALID_PASSWORD))))