From 0fcf28e7e76c8db546263d1a21e9aa178a1bcf4e Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Thu, 6 Aug 2020 17:51:42 -0400
Subject: [PATCH] Use the MessagesManager to actually persist messages.

---
 .../textsecuregcm/WhisperServerService.java   |  2 +-
 .../storage/MessagePersister.java             | 19 +++++++------------
 .../storage/MessagesManager.java              |  6 ++++++
 3 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index 4e130c1f2..d3e88bffb 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -375,7 +375,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     TurnTokenGenerator       turnTokenGenerator  = new TurnTokenGenerator(config.getTurnConfiguration());
     RecaptchaClient          recaptchaClient     = new RecaptchaClient(config.getRecaptchaConfiguration().getSecret());
 
-    MessagePersister             messagePersister        = new MessagePersister(messagesCache, messagesClient, messages, pubSubManager, pushSender, accountsManager,config.getMessageCacheConfiguration().getPersistDelayMinutes(), TimeUnit.MINUTES);
+    MessagePersister             messagePersister        = new MessagePersister(messagesClient, messagesManager, pubSubManager, pushSender, accountsManager,config.getMessageCacheConfiguration().getPersistDelayMinutes(), TimeUnit.MINUTES);
     RedisClusterMessagePersister clusterMessagePersister = new RedisClusterMessagePersister(clusterMessagesCache, messages, pubSubManager, pushSender, accountsManager, Duration.ofMinutes(config.getMessageCacheConfiguration().getPersistDelayMinutes()));
 
     DirectoryReconciliationClient directoryReconciliationClient = new DirectoryReconciliationClient(config.getDirectoryConfiguration().getDirectoryServerConfiguration());
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagePersister.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagePersister.java
index 1d12896a9..9c901f67b 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagePersister.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagePersister.java
@@ -33,7 +33,6 @@ import static com.codahale.metrics.MetricRegistry.name;
 
 public class MessagePersister implements Managed, Runnable {
 
-  private final MessagesCache  messagesCache;
   private final Logger         logger              = LoggerFactory.getLogger(MessagePersister.class);
   private final MetricRegistry metricRegistry      = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
   private final Timer          getQueuesTimer      = metricRegistry.timer(name(MessagesCache.class, "getQueues"   ));
@@ -47,10 +46,10 @@ public class MessagePersister implements Managed, Runnable {
   private final AtomicBoolean running = new AtomicBoolean(true);
 
   private final ReplicatedJedisPool jedisPool;
-  private final Messages            database;
   private final long                delayTime;
   private final TimeUnit            delayTimeUnit;
 
+  private final MessagesManager messagesManager;
   private final PubSubManager   pubSubManager;
   private final PushSender      pushSender;
   private final AccountsManager accountsManager;
@@ -59,9 +58,8 @@ public class MessagePersister implements Managed, Runnable {
 
   private boolean finished = false;
 
-  public MessagePersister(final MessagesCache       messagesCache,
-                          final ReplicatedJedisPool jedisPool,
-                          final Messages            database,
+  public MessagePersister(final ReplicatedJedisPool jedisPool,
+                          final MessagesManager     messagesManager,
                           final PubSubManager       pubSubManager,
                           final PushSender          pushSender,
                           final AccountsManager     accountsManager,
@@ -69,10 +67,9 @@ public class MessagePersister implements Managed, Runnable {
                           final TimeUnit            delayTimeUnit)
       throws IOException
   {
-    this.messagesCache = messagesCache;
-    this.jedisPool = jedisPool;
-    this.database  = database;
+    this.jedisPool     = jedisPool;
 
+    this.messagesManager = messagesManager;
     this.pubSubManager   = pubSubManager;
     this.pushSender      = pushSender;
     this.accountsManager = accountsManager;
@@ -156,16 +153,14 @@ public class MessagePersister implements Managed, Runnable {
   private void persistMessage(Key key, UUID destinationUuid, long score, byte[] message) {
     try {
       MessageProtos.Envelope envelope = MessageProtos.Envelope.parseFrom(message);
-      UUID     guid     = envelope.hasServerGuid() ? UUID.fromString(envelope.getServerGuid()) : null;
+      UUID                   guid     = envelope.hasServerGuid() ? UUID.fromString(envelope.getServerGuid()) : null;
 
       envelope = envelope.toBuilder().clearServerGuid().build();
 
-      database.store(guid, envelope, key.getAddress(), key.getDeviceId());
+      messagesManager.persistMessage(key.getAddress(), destinationUuid, envelope, guid, key.getDeviceId(), score);
     } catch (InvalidProtocolBufferException e) {
       logger.error("Error parsing envelope", e);
     }
-
-    messagesCache.remove(key.getAddress(), destinationUuid, key.getDeviceId(), score);
   }
 
   private List<byte[]> getQueuesToPersist() {
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java
index c27ae5426..cbc14698d 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java
@@ -121,4 +121,10 @@ public class MessagesManager {
     }
   }
 
+  public void persistMessage(String destination, UUID destinationUuid, Envelope envelope, UUID messageGuid, long deviceId, long id) {
+    messages.store(messageGuid, envelope, destination, deviceId);
+
+    final Optional<OutgoingMessageEntity> maybeRemovedMessage = messagesCache.remove(destination, destinationUuid, deviceId, id);
+    removeByIdExperiment.compareSupplierResultAsync(maybeRemovedMessage, () -> clusterMessagesCache.remove(destination, destinationUuid, deviceId, id), experimentExecutor);
+  }
 }