Shard push scheduling cache

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

@@ -107,6 +107,11 @@ public class WhisperServerConfiguration extends Configuration {
   @JsonProperty
   private RedisConfiguration pushScheduler;
 
+  @NotNull
+  @Valid
+  @JsonProperty
+  private RedisClusterConfiguration pushSchedulerCluster;
+
   @NotNull
   @Valid
   @JsonProperty
@@ -287,6 +292,10 @@ public class WhisperServerConfiguration extends Configuration {
     return pushScheduler;
   }
 
+  public RedisClusterConfiguration getPushSchedulerCluster() {
+    return pushSchedulerCluster;
+  }
+
   public DatabaseConfiguration getMessageStoreConfiguration() {
     return messageStore;
   }

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

@@ -285,6 +285,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     ClientResources messageCacheClientResources       = ClientResources.builder().build();
     ClientResources presenceClientResources           = ClientResources.builder().build();
     ClientResources metricsCacheClientResources       = ClientResources.builder().build();
+    ClientResources pushSchedulerCacheClientResources = ClientResources.builder().ioThreadPoolSize(4).build();
 
     ConnectionEventLogger.logConnectionEvents(generalCacheClientResources);
     ConnectionEventLogger.logConnectionEvents(messageCacheClientResources);
@@ -295,6 +296,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     FaultTolerantRedisCluster messagesCluster          = new FaultTolerantRedisCluster("message_insert_cluster", config.getMessageCacheConfiguration().getRedisClusterConfiguration(), messageCacheClientResources);
     FaultTolerantRedisCluster clientPresenceCluster    = new FaultTolerantRedisCluster("client_presence_cluster", config.getClientPresenceClusterConfiguration(), presenceClientResources);
     FaultTolerantRedisCluster metricsCluster           = new FaultTolerantRedisCluster("metrics_cluster", config.getMetricsClusterConfiguration(), metricsCacheClientResources);
+    FaultTolerantRedisCluster pushSchedulerCluster     = new FaultTolerantRedisCluster("push_scheduler", config.getPushSchedulerCluster(), pushSchedulerCacheClientResources);
 
     BlockingQueue<Runnable> keyspaceNotificationDispatchQueue = new ArrayBlockingQueue<>(10_000);
     Metrics.gaugeCollectionSize(name(getClass(), "keyspaceNotificationDispatchQueueSize"), Collections.emptyList(), keyspaceNotificationDispatchQueue);
@@ -336,7 +338,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     ExternalServiceCredentialGenerator backupCredentialsGenerator    = new ExternalServiceCredentialGenerator(config.getSecureBackupServiceConfiguration().getUserAuthenticationTokenSharedSecret(), new byte[0], false);
     ExternalServiceCredentialGenerator paymentsCredentialsGenerator  = new ExternalServiceCredentialGenerator(config.getPaymentsServiceConfiguration().getUserAuthenticationTokenSharedSecret(), new byte[0], false);
 
-    ApnFallbackManager       apnFallbackManager = new ApnFallbackManager(pushSchedulerClient, apnSender, accountsManager);
+    ApnFallbackManager       apnFallbackManager = new ApnFallbackManager(pushSchedulerClient, pushSchedulerCluster, apnSender, accountsManager);
     TwilioSmsSender          twilioSmsSender    = new TwilioSmsSender(config.getTwilioConfiguration());
     SmsSender                smsSender          = new SmsSender(twilioSmsSender);
     MessageSender            messageSender      = new MessageSender(apnFallbackManager, clientPresenceManager, messagesManager, gcmSender, apnSender, pushLatencyManager);

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

@@ -9,8 +9,14 @@ import com.codahale.metrics.Meter;
 import com.codahale.metrics.MetricRegistry;
 import com.codahale.metrics.RatioGauge;
 import com.codahale.metrics.SharedMetricRegistries;
+import com.google.common.annotations.VisibleForTesting;
+import io.dropwizard.lifecycle.Managed;
+import io.lettuce.core.ScriptOutputType;
+import io.lettuce.core.cluster.SlotHash;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.redis.ClusterLuaScript;
+import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
 import org.whispersystems.textsecuregcm.redis.LuaScript;
 import org.whispersystems.textsecuregcm.redis.RedisException;
 import org.whispersystems.textsecuregcm.redis.ReplicatedJedisPool;
@@ -19,32 +25,37 @@ import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.util.Constants;
 import org.whispersystems.textsecuregcm.util.Pair;
+import org.whispersystems.textsecuregcm.util.RedisClusterUtil;
 import org.whispersystems.textsecuregcm.util.Util;
+import redis.clients.jedis.exceptions.JedisException;
 
 import java.io.IOException;
+import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
+import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.stream.Collectors;
 
 import static com.codahale.metrics.MetricRegistry.name;
-import io.dropwizard.lifecycle.Managed;
-import redis.clients.jedis.Jedis;
-import redis.clients.jedis.exceptions.JedisException;
 
-public class ApnFallbackManager implements Managed, Runnable {
+public class ApnFallbackManager implements Managed {
 
   private static final Logger logger = LoggerFactory.getLogger(ApnFallbackManager.class);
 
-  private static final String PENDING_NOTIFICATIONS_KEY = "PENDING_APN";
+  private static final String SINGLETON_PENDING_NOTIFICATIONS_KEY = "PENDING_APN";
+          static final String NEXT_SLOT_TO_PERSIST_KEY            = "pending_notification_next_slot";
 
   private static final MetricRegistry metricRegistry        = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
   private static final Meter          delivered             = metricRegistry.meter(name(ApnFallbackManager.class, "voip_delivered"));
   private static final Meter          sent                  = metricRegistry.meter(name(ApnFallbackManager.class, "voip_sent"     ));
   private static final Meter          retry                 = metricRegistry.meter(name(ApnFallbackManager.class, "voip_retry"));
   private static final Meter          evicted               = metricRegistry.meter(name(ApnFallbackManager.class, "voip_evicted"));
+  private static final Meter          singletonDestinations = metricRegistry.meter(name(ApnFallbackManager.class, "singleton_destinations"));
+  private static final Meter          clusterDestinations   = metricRegistry.meter(name(ApnFallbackManager.class, "cluster_destinations"));
 
   static {
     metricRegistry.register(name(ApnFallbackManager.class, "voip_ratio"), new VoipRatioGauge(delivered, sent));
@@ -52,118 +63,53 @@ public class ApnFallbackManager implements Managed, Runnable {
 
   private final APNSender                 apnSender;
   private final AccountsManager           accountsManager;
+  private final FaultTolerantRedisCluster cluster;
 
-  private final ReplicatedJedisPool jedisPool;
+  private final LuaScript getSingletonScript;
-  private final InsertOperation     insertOperation;
+  private final LuaScript removeSingletonScript;
-  private final GetOperation        getOperation;
-  private final RemoveOperation     removeOperation;
 
+  private final ClusterLuaScript getClusterScript;
+  private final ClusterLuaScript insertClusterScript;
+  private final ClusterLuaScript removeClusterScript;
 
-  private AtomicBoolean running = new AtomicBoolean(false);
+  private final Thread   singletonWorkerThread;
-  private boolean finished;
+  private final Thread[] clusterWorkerThreads = new Thread[CLUSTER_WORKER_THREAD_COUNT];
 
-  public ApnFallbackManager(ReplicatedJedisPool jedisPool,
+  private static final int CLUSTER_WORKER_THREAD_COUNT = 4;
-                            APNSender apnSender,
-                            AccountsManager accountsManager)
-      throws IOException
-  {
-    this.apnSender       = apnSender;
-    this.accountsManager = accountsManager;
-    this.jedisPool       = jedisPool;
-    this.insertOperation = new InsertOperation(jedisPool);
-    this.getOperation    = new GetOperation(jedisPool);
-    this.removeOperation = new RemoveOperation(jedisPool);
-  }
 
-  public void schedule(Account account, Device device) throws RedisException {
+  private final AtomicBoolean running = new AtomicBoolean(false);
-    try {
-      sent.mark();
-      insertOperation.insert(account, device, System.currentTimeMillis() + (15 * 1000), (15 * 1000));
-    } catch (JedisException e) {
-      throw new RedisException(e);
-    }
-  }
 
-  public boolean isScheduled(Account account, Device device) throws RedisException {
+  private class SingletonCacheWorker implements Runnable {
-    try {
-      String endpoint = "apn_device::" + account.getNumber() + "::" + device.getId();
-
-      try (Jedis jedis = jedisPool.getReadResource()) {
-        return jedis.zscore(PENDING_NOTIFICATIONS_KEY, endpoint) != null;
-      }
-    } catch (JedisException e) {
-      throw new RedisException(e);
-    }
-  }
-
-  public void cancel(Account account, Device device) throws RedisException {
-    try {
-      if (removeOperation.remove(account, device)) {
-        delivered.mark();
-      }
-    } catch (JedisException e) {
-      throw new RedisException(e);
-    }
-  }
-
-  @Override
-  public synchronized void start() {
-    running.set(true);
-    new Thread(this).start();
-  }
-
-  @Override
-  public synchronized void stop() {
-    running.set(false);
-    while (!finished) Util.wait(this);
-  }
 
     @Override
     public void run() {
       while (running.get()) {
         try {
-        List<byte[]> pendingNotifications = getOperation.getPending(100);
+          for (final String numberAndDevice : getPendingDestinationsFromSingletonCache(100)) {
+            singletonDestinations.mark();
 
-        for (byte[] pendingNotification : pendingNotifications) {
-          String                       numberAndDevice = new String(pendingNotification);
             Optional<Pair<String, Long>> separated = getSeparated(numberAndDevice);
 
             if (!separated.isPresent()) {
-            removeOperation.remove(numberAndDevice);
+              removeFromSingleton(numberAndDevice);
               continue;
             }
 
             Optional<Account> account = accountsManager.get(separated.get().first());
 
             if (!account.isPresent()) {
-            removeOperation.remove(numberAndDevice);
+              removeFromSingleton(numberAndDevice);
               continue;
             }
 
             Optional<Device> device = account.get().getDevice(separated.get().second());
 
             if (!device.isPresent()) {
-            removeOperation.remove(numberAndDevice);
+              removeFromSingleton(numberAndDevice);
               continue;
             }
 
-          String apnId = device.get().getVoipApnId();
+            sendNotification(account.get(), device.get());
-
-          if (apnId == null) {
-            removeOperation.remove(account.get(), device.get());
-            continue;
-          }
-
-          long deviceLastSeen = device.get().getLastSeen();
-
-          if (deviceLastSeen < System.currentTimeMillis() - TimeUnit.DAYS.toMillis(90)) {
-            evicted.mark();
-            removeOperation.remove(account.get(), device.get());
-            continue;
-          }
-
-          apnSender.sendMessage(new ApnMessage(apnId, separated.get().first(), separated.get().second(), true, Optional.empty()));
-          retry.mark();
           }
 
         } catch (Exception e) {
@@ -172,14 +118,150 @@ public class ApnFallbackManager implements Managed, Runnable {
 
         Util.sleep(1000);
       }
-
-    synchronized (ApnFallbackManager.this) {
-      finished = true;
-      notifyAll();
     }
   }
 
-  private Optional<Pair<String, Long>> getSeparated(String encoded) {
+  class ClusterCacheWorker implements Runnable {
+
+    @Override
+    public void run() {
+      while (running.get()) {
+        try {
+          final long entriesProcessed = processNextSlot();
+
+          if (entriesProcessed == 0) {
+            Util.sleep(1000);
+          }
+        } catch (Exception e) {
+          logger.warn("Exception while operating", e);
+        }
+      }
+    }
+
+    long processNextSlot() {
+      final int slot = getNextSlot();
+
+      List<String> pendingDestinations;
+      long entriesProcessed = 0;
+
+      do {
+        pendingDestinations = getPendingDestinationsFromClusterCache(slot, 100);
+        entriesProcessed += pendingDestinations.size();
+
+        for (final String uuidAndDevice : pendingDestinations) {
+          clusterDestinations.mark();
+
+          final Optional<Pair<String, Long>> separated = getSeparated(uuidAndDevice);
+
+          final Optional<Account> maybeAccount = separated.map(Pair::first)
+                                                          .map(UUID::fromString)
+                                                          .flatMap(accountsManager::get);
+
+          final Optional<Device> maybeDevice = separated.map(Pair::second)
+                                                        .flatMap(deviceId -> maybeAccount.flatMap(account -> account.getDevice(deviceId)));
+
+          if (maybeAccount.isPresent() && maybeDevice.isPresent()) {
+            sendNotification(maybeAccount.get(), maybeDevice.get());
+          } else {
+            removeFromCluster(uuidAndDevice);
+          }
+        }
+      } while (!pendingDestinations.isEmpty());
+
+      return entriesProcessed;
+    }
+  }
+
+  public ApnFallbackManager(ReplicatedJedisPool jedisPool,
+                            FaultTolerantRedisCluster cluster,
+                            APNSender apnSender,
+                            AccountsManager accountsManager)
+      throws IOException
+  {
+    this.apnSender       = apnSender;
+    this.accountsManager = accountsManager;
+    this.cluster         = cluster;
+
+    this.getSingletonScript    = LuaScript.fromResource(jedisPool, "lua/apn/get.lua");
+    this.removeSingletonScript = LuaScript.fromResource(jedisPool, "lua/apn/remove.lua");
+
+    this.getClusterScript    = ClusterLuaScript.fromResource(cluster, "lua/apn/get.lua", ScriptOutputType.MULTI);
+    this.insertClusterScript = ClusterLuaScript.fromResource(cluster, "lua/apn/insert.lua", ScriptOutputType.VALUE);
+    this.removeClusterScript = ClusterLuaScript.fromResource(cluster, "lua/apn/remove.lua", ScriptOutputType.INTEGER);
+
+    this.singletonWorkerThread = new Thread(new SingletonCacheWorker(), "ApnFallbackManagerSingletonWorker");
+
+    for (int i = 0; i < this.clusterWorkerThreads.length; i++) {
+      this.clusterWorkerThreads[i] = new Thread(new ClusterCacheWorker(), "ApnFallbackManagerClusterWorker-" + i);
+    }
+  }
+
+  public void schedule(Account account, Device device) throws RedisException {
+    schedule(account, device, System.currentTimeMillis());
+  }
+
+  @VisibleForTesting
+  void schedule(Account account, Device device, long timestamp) throws RedisException {
+    try {
+      sent.mark();
+      insert(account, device, timestamp + (15 * 1000), (15 * 1000));
+    } catch (io.lettuce.core.RedisException e) {
+      throw new RedisException(e);
+    }
+  }
+
+  public void cancel(Account account, Device device) throws RedisException {
+    try {
+      if (remove(account, device)) {
+        delivered.mark();
+      }
+    } catch (JedisException | io.lettuce.core.RedisException e) {
+      throw new RedisException(e);
+    }
+  }
+
+  @Override
+  public synchronized void start() {
+    running.set(true);
+    singletonWorkerThread.start();
+
+    for (final Thread clusterWorkerThread : clusterWorkerThreads) {
+      clusterWorkerThread.start();
+    }
+  }
+
+  @Override
+  public synchronized void stop() throws InterruptedException {
+    running.set(false);
+    singletonWorkerThread.join();
+
+    for (final Thread clusterWorkerThread : clusterWorkerThreads) {
+      clusterWorkerThread.join();
+    }
+  }
+
+  private void sendNotification(final Account account, final Device device) {
+    String apnId = device.getVoipApnId();
+
+    if (apnId == null) {
+      remove(account, device);
+      return;
+    }
+
+    long deviceLastSeen = device.getLastSeen();
+
+    if (deviceLastSeen < System.currentTimeMillis() - TimeUnit.DAYS.toMillis(90)) {
+      evicted.mark();
+      remove(account, device);
+      return;
+    }
+
+    apnSender.sendMessage(new ApnMessage(apnId, account.getNumber(), device.getId(), true, Optional.empty()));
+    retry.mark();
+  }
+
+  @VisibleForTesting
+  static Optional<Pair<String, Long>> getSeparated(String encoded) {
     try {
       if (encoded == null) return Optional.empty();
 
@@ -197,66 +279,78 @@ public class ApnFallbackManager implements Managed, Runnable {
     }
   }
 
-  private static class RemoveOperation {
+  private boolean remove(Account account, Device device) {
+    final boolean removedFromSingleton = removeFromSingleton(getSingletonEndpointKey(account, device));
+    final boolean removedFromCluster = removeFromCluster(getClusterEndpointKey(account, device));
 
-    private final LuaScript luaScript;
+    return removedFromSingleton || removedFromCluster;
-
-    RemoveOperation(ReplicatedJedisPool jedisPool) throws IOException {
-      this.luaScript = LuaScript.fromResource(jedisPool, "lua/apn/remove.lua");
   }
 
-    boolean remove(Account account, Device device) {
+  private boolean removeFromSingleton(String endpoint) {
-      String endpoint = "apn_device::" + account.getNumber() + "::" + device.getId();
+    if (!SINGLETON_PENDING_NOTIFICATIONS_KEY.equals(endpoint)) {
-      return remove(endpoint);
+      List<byte[]> keys = Arrays.asList(SINGLETON_PENDING_NOTIFICATIONS_KEY.getBytes(), endpoint.getBytes());
-    }
-
-    boolean remove(String endpoint) {
-      if (!PENDING_NOTIFICATIONS_KEY.equals(endpoint)) {
-        List<byte[]> keys = Arrays.asList(PENDING_NOTIFICATIONS_KEY.getBytes(), endpoint.getBytes());
       List<byte[]> args = Collections.emptyList();
 
-        return ((long)luaScript.execute(keys, args)) > 0;
+      return ((long)removeSingletonScript.execute(keys, args)) > 0;
     }
 
     return false;
   }
 
+  private boolean removeFromCluster(final String endpoint) {
+    final long removed = (long)removeClusterScript.execute(List.of(getClusterPendingNotificationQueueKey(endpoint), endpoint),
+                                                           Collections.emptyList());
+
+    return removed > 0;
   }
 
-  private static class GetOperation {
+  @SuppressWarnings("unchecked")
+  private List<String> getPendingDestinationsFromSingletonCache(final int limit) {
+    List<byte[]> keys = List.of(SINGLETON_PENDING_NOTIFICATIONS_KEY.getBytes());
+    List<byte[]> args = List.of(String.valueOf(System.currentTimeMillis()).getBytes(), String.valueOf(limit).getBytes());
 
-    private final LuaScript luaScript;
+    return ((List<byte[]>) getSingletonScript.execute(keys, args))
-
+            .stream()
-    GetOperation(ReplicatedJedisPool jedisPool) throws IOException {
+            .map(bytes -> new String(bytes, StandardCharsets.UTF_8))
-      this.luaScript = LuaScript.fromResource(jedisPool, "lua/apn/get.lua");
+            .collect(Collectors.toList());
   }
 
-    @SuppressWarnings("SameParameterValue")
+  @SuppressWarnings("unchecked")
-    List<byte[]> getPending(int limit) {
+  @VisibleForTesting
-      List<byte[]> keys = Arrays.asList(PENDING_NOTIFICATIONS_KEY.getBytes());
+  List<String> getPendingDestinationsFromClusterCache(final int slot, final int limit) {
-      List<byte[]> args = Arrays.asList(String.valueOf(System.currentTimeMillis()).getBytes(), String.valueOf(limit).getBytes());
+    return (List<String>)getClusterScript.execute(List.of(getClusterPendingNotificationQueueKey(slot)),
-
+                                                  List.of(String.valueOf(System.currentTimeMillis()), String.valueOf(limit)));
-      return (List<byte[]>) luaScript.execute(keys, args);
-    }
   }
 
-  private static class InsertOperation {
+  private void insert(final Account account, final Device device, final long timestamp, final long interval) {
+    final String endpoint = getClusterEndpointKey(account, device);
 
-    private final LuaScript luaScript;
+    insertClusterScript.execute(List.of(getClusterPendingNotificationQueueKey(endpoint), endpoint),
-
+                                List.of(String.valueOf(timestamp),
-    InsertOperation(ReplicatedJedisPool jedisPool) throws IOException {
+                                        String.valueOf(interval),
-      this.luaScript = LuaScript.fromResource(jedisPool, "lua/apn/insert.lua");
+                                        account.getUuid().toString(),
+                                        String.valueOf(device.getId())));
   }
 
-    public void insert(Account account, Device device, long timestamp, long interval) {
+  private String getSingletonEndpointKey(final Account account, final Device device) {
-      String endpoint = "apn_device::" + account.getNumber() + "::" + device.getId();
+    return "apn_device::" + account.getNumber() + "::" + device.getId();
-
-      List<byte[]> keys = Arrays.asList(PENDING_NOTIFICATIONS_KEY.getBytes(), endpoint.getBytes());
-      List<byte[]> args = Arrays.asList(String.valueOf(timestamp).getBytes(), String.valueOf(interval).getBytes(),
-                                        account.getNumber().getBytes(), String.valueOf(device.getId()).getBytes());
-
-      luaScript.execute(keys, args);
   }
+
+  @VisibleForTesting
+  String getClusterEndpointKey(final Account account, final Device device) {
+    return "apn_device::{" + account.getUuid() + "::" + device.getId() + "}";
+  }
+
+  private String getClusterPendingNotificationQueueKey(final String endpoint) {
+    return getClusterPendingNotificationQueueKey(SlotHash.getSlot(endpoint));
+  }
+
+  private String getClusterPendingNotificationQueueKey(final int slot) {
+    return SINGLETON_PENDING_NOTIFICATIONS_KEY + "::{" + RedisClusterUtil.getMinimalHashTag(slot) + "}";
+  }
+
+  private int getNextSlot() {
+    return (int)(cluster.withCluster(connection -> connection.sync().incr(NEXT_SLOT_TO_PERSIST_KEY)) % SlotHash.SLOT_COUNT);
   }
 
   private static class VoipRatioGauge extends RatioGauge {

service/src/main/resources/lua/apn/get.lua

@@ -1,5 +1,7 @@
--- keys: pending (KEYS[1])
+local pendingNotificationQueue = KEYS[1]
--- argv: max_time (ARGV[1]), limit (ARGV[2])
+
+local maxTime = ARGV[1]
+local limit = ARGV[2]
 
 local hgetall = function (key)
     local bulk = redis.call('HGETALL', key)
@@ -44,7 +46,7 @@ local getNextInterval = function(interval)
 end
 
 
-local results  = redis.call("ZRANGEBYSCORE", KEYS[1], 0, ARGV[1], "LIMIT", 0, ARGV[2])
+local results  = redis.call("ZRANGEBYSCORE", pendingNotificationQueue, 0, maxTime, "LIMIT", 0, limit)
 local collated = {}
 
 if results and next(results) then
@@ -59,12 +61,10 @@ if results and next(results) then
         local nextInterval = getNextInterval(tonumber(lastInterval))
 
         redis.call("HSET", name, "interval", nextInterval)
-        redis.call("ZADD", KEYS[1], tonumber(ARGV[1]) + nextInterval, name)
+        redis.call("ZADD", pendingNotificationQueue, tonumber(maxTime) + nextInterval, name)
 
         collated[i] = pending["account"] .. ":" .. pending["device"]
     end
 end
 
 return collated
-
-

service/src/main/resources/lua/apn/insert.lua

@@ -1,9 +1,14 @@
--- keys: pending (KEYS[1]), user (KEYS[2])
+local pendingNotificationQueue = KEYS[1]
--- args: timestamp (ARGV[1]), interval (ARGV[2]), account (ARGV[3]), device (ARGV[4])
+local endpoint = KEYS[2]
 
-redis.call("HSET", KEYS[2], "created", ARGV[1])
+local timestamp = ARGV[1]
-redis.call("HSET", KEYS[2], "interval", ARGV[2])
+local interval = ARGV[2]
-redis.call("HSET", KEYS[2], "account", ARGV[3])
+local account = ARGV[3]
-redis.call("HSET", KEYS[2], "device", ARGV[4])
+local deviceId = ARGV[4]
 
-redis.call("ZADD", KEYS[1], ARGV[1], KEYS[2])
+redis.call("HSET", endpoint, "created", timestamp)
+redis.call("HSET", endpoint, "interval", interval)
+redis.call("HSET", endpoint, "account", account)
+redis.call("HSET", endpoint, "device", deviceId)
+
+redis.call("ZADD", pendingNotificationQueue, timestamp, endpoint)

service/src/main/resources/lua/apn/remove.lua

@@ -1,4 +1,5 @@
--- keys: queue KEYS[1], endpoint (KEYS[2])
+local pendingNotificationQueue = KEYS[1]
+local endpoint = KEYS[2]
 
-redis.call("DEL", KEYS[2])
+redis.call("DEL", endpoint)
-return redis.call("ZREM", KEYS[1], KEYS[2])
+return redis.call("ZREM", pendingNotificationQueue, endpoint)

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

@@ -0,0 +1,156 @@
+/*
+ * Copyright 2021 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.push;
+
+import io.lettuce.core.cluster.SlotHash;
+import org.apache.commons.lang3.RandomStringUtils;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.mockito.ArgumentCaptor;
+import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
+import org.whispersystems.textsecuregcm.providers.RedisClientFactory;
+import org.whispersystems.textsecuregcm.redis.AbstractRedisClusterTest;
+import org.whispersystems.textsecuregcm.redis.RedisException;
+import org.whispersystems.textsecuregcm.redis.ReplicatedJedisPool;
+import org.whispersystems.textsecuregcm.storage.Account;
+import org.whispersystems.textsecuregcm.storage.AccountsManager;
+import org.whispersystems.textsecuregcm.storage.Device;
+import org.whispersystems.textsecuregcm.util.Pair;
+import redis.clients.jedis.Jedis;
+import redis.embedded.RedisServer;
+
+import java.util.List;
+import java.util.Optional;
+import java.util.UUID;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assume.assumeFalse;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class ApnFallbackManagerTest extends AbstractRedisClusterTest {
+
+    private Account account;
+    private Device device;
+
+    private APNSender apnSender;
+
+    private ApnFallbackManager apnFallbackManager;
+
+    private static RedisServer redisServer;
+
+    private static final UUID   ACCOUNT_UUID   = UUID.randomUUID();
+    private static final String ACCOUNT_NUMBER = "+18005551234";
+    private static final long   DEVICE_ID      = 1L;
+    private static final String VOIP_APN_ID = RandomStringUtils.randomAlphanumeric(32);
+
+    @BeforeClass
+    public static void setUpRedisSingleton() throws Exception {
+        assumeFalse(System.getProperty("os.name").equalsIgnoreCase("windows"));
+
+        redisServer = RedisServer.builder()
+                .setting("appendonly no")
+                .setting("dir " + System.getProperty("java.io.tmpdir"))
+                .port(AbstractRedisClusterTest.getNextRedisClusterPort())
+                .build();
+
+        redisServer.start();
+    }
+
+    @Before
+    public void setUp() throws Exception {
+        super.setUp();
+
+        final String redisUrl = String.format("redis://127.0.0.1:%d", redisServer.ports().get(0));
+
+        final ReplicatedJedisPool replicatedJedisPool = new RedisClientFactory("test-pool",
+                redisUrl,
+                List.of(redisUrl),
+                new CircuitBreakerConfiguration()).getRedisClientPool();
+
+        try (final Jedis jedis = replicatedJedisPool.getWriteResource()) {
+            jedis.flushAll();
+        }
+
+        device = mock(Device.class);
+        when(device.getId()).thenReturn(DEVICE_ID);
+        when(device.getVoipApnId()).thenReturn(VOIP_APN_ID);
+        when(device.getLastSeen()).thenReturn(System.currentTimeMillis());
+
+        account = mock(Account.class);
+        when(account.getUuid()).thenReturn(ACCOUNT_UUID);
+        when(account.getNumber()).thenReturn(ACCOUNT_NUMBER);
+        when(account.getDevice(DEVICE_ID)).thenReturn(Optional.of(device));
+
+        final AccountsManager accountsManager = mock(AccountsManager.class);
+        when(accountsManager.get(ACCOUNT_NUMBER)).thenReturn(Optional.of(account));
+        when(accountsManager.get(ACCOUNT_UUID)).thenReturn(Optional.of(account));
+
+        apnSender = mock(APNSender.class);
+
+        apnFallbackManager = new ApnFallbackManager(replicatedJedisPool, getRedisCluster(), apnSender, accountsManager);
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        super.tearDown();
+    }
+
+    @AfterClass
+    public static void tearDownRedisSingleton() {
+        redisServer.stop();
+    }
+
+    @Test
+    public void testClusterInsert() throws RedisException {
+        final String endpoint = apnFallbackManager.getClusterEndpointKey(account, device);
+
+        assertTrue(apnFallbackManager.getPendingDestinationsFromClusterCache(SlotHash.getSlot(endpoint), 1).isEmpty());
+
+        apnFallbackManager.schedule(account, device, System.currentTimeMillis() - 30_000);
+
+        final List<String> pendingDestinations = apnFallbackManager.getPendingDestinationsFromClusterCache(SlotHash.getSlot(endpoint), 2);
+        assertEquals(1, pendingDestinations.size());
+
+        final Optional<Pair<String, Long>> maybeUuidAndDeviceId = ApnFallbackManager.getSeparated(pendingDestinations.get(0));
+
+        assertTrue(maybeUuidAndDeviceId.isPresent());
+        assertEquals(ACCOUNT_UUID.toString(), maybeUuidAndDeviceId.get().first());
+        assertEquals(DEVICE_ID, (long)maybeUuidAndDeviceId.get().second());
+
+        assertTrue(apnFallbackManager.getPendingDestinationsFromClusterCache(SlotHash.getSlot(endpoint), 1).isEmpty());
+    }
+
+    @Test
+    public void testProcessNextSlot() throws RedisException {
+        final ApnFallbackManager.ClusterCacheWorker worker = apnFallbackManager.new ClusterCacheWorker();
+
+        apnFallbackManager.schedule(account, device, System.currentTimeMillis() - 30_000);
+
+        final int slot = SlotHash.getSlot(apnFallbackManager.getClusterEndpointKey(account, device));
+        final int previousSlot = (slot + SlotHash.SLOT_COUNT - 1) % SlotHash.SLOT_COUNT;
+
+        getRedisCluster().withCluster(connection -> connection.sync().set(ApnFallbackManager.NEXT_SLOT_TO_PERSIST_KEY, String.valueOf(previousSlot)));
+
+        assertEquals(1, worker.processNextSlot());
+
+        final ArgumentCaptor<ApnMessage> messageCaptor = ArgumentCaptor.forClass(ApnMessage.class);
+        verify(apnSender).sendMessage(messageCaptor.capture());
+
+        final ApnMessage message = messageCaptor.getValue();
+
+        assertEquals(VOIP_APN_ID, message.getApnId());
+        assertEquals(ACCOUNT_NUMBER, message.getNumber());
+        assertEquals(DEVICE_ID, message.getDeviceId());
+
+        assertEquals(0, worker.processNextSlot());
+    }
+}