From 5caa951c612d049d3b3aea3764a26ec247c12eba Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Tue, 8 Aug 2023 10:34:58 -0400
Subject: [PATCH] Make `MessagesCache#clear` methods asynchronous

---
 .../textsecuregcm/storage/MessagesCache.java  | 23 +++++++++++--------
 .../storage/MessagesManager.java              |  4 ++--
 .../storage/MessagesCacheTest.java            | 10 ++------
 3 files changed, 17 insertions(+), 20 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesCache.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesCache.java
index 52e04b858..626e6fd8a 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesCache.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesCache.java
@@ -329,21 +329,24 @@ public class MessagesCache extends RedisClusterPubSubAdapter<String, String> imp
     });
   }
 
-  public void clear(final UUID destinationUuid) {
-    // TODO Remove null check in a fully UUID-based world
-    if (destinationUuid != null) {
-      for (int i = 1; i < Device.MAXIMUM_DEVICE_ID; i++) {
-        clear(destinationUuid, i);
-      }
+  public CompletableFuture<Void> clear(final UUID destinationUuid) {
+    final CompletableFuture<?>[] clearFutures = new CompletableFuture[Device.MAXIMUM_DEVICE_ID];
+
+    for (int deviceId = 0; deviceId < Device.MAXIMUM_DEVICE_ID; deviceId++) {
+      clearFutures[deviceId] = clear(destinationUuid, deviceId);
     }
+
+    return CompletableFuture.allOf(clearFutures);
   }
 
-  public void clear(final UUID destinationUuid, final long deviceId) {
-    clearQueueTimer.record(() ->
-        removeQueueScript.executeBinary(List.of(getMessageQueueKey(destinationUuid, deviceId),
+  public CompletableFuture<Void> clear(final UUID destinationUuid, final long deviceId) {
+    final Timer.Sample sample = Timer.start();
+
+    return removeQueueScript.executeBinaryAsync(List.of(getMessageQueueKey(destinationUuid, deviceId),
                 getMessageQueueMetadataKey(destinationUuid, deviceId),
                 getQueueIndexKey(destinationUuid, deviceId)),
-            Collections.emptyList()));
+            Collections.emptyList())
+        .thenRun(() -> sample.stop(clearQueueTimer));
   }
 
   int getNextSlotToPersist() {
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 3810f26ca..1af48d345 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java
@@ -103,12 +103,12 @@ public class MessagesManager {
   }
 
   public void clear(UUID destinationUuid) {
-    messagesCache.clear(destinationUuid);
+    messagesCache.clear(destinationUuid).join();
     messagesDynamoDb.deleteAllMessagesForAccount(destinationUuid);
   }
 
   public void clear(UUID destinationUuid, long deviceId) {
-    messagesCache.clear(destinationUuid, deviceId);
+    messagesCache.clear(destinationUuid, deviceId).join();
     messagesDynamoDb.deleteAllMessagesForDevice(destinationUuid, deviceId);
   }
 
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagesCacheTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagesCacheTest.java
index 2bae8bec7..c9aff667d 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagesCacheTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagesCacheTest.java
@@ -320,7 +320,7 @@ class MessagesCacheTest {
         }
       }
 
-      messagesCache.clear(DESTINATION_UUID, DESTINATION_DEVICE_ID);
+      messagesCache.clear(DESTINATION_UUID, DESTINATION_DEVICE_ID).join();
 
       assertEquals(Collections.emptyList(), get(DESTINATION_UUID, DESTINATION_DEVICE_ID, messageCount));
       assertEquals(messageCount, get(DESTINATION_UUID, DESTINATION_DEVICE_ID + 1, messageCount).size());
@@ -340,18 +340,12 @@ class MessagesCacheTest {
         }
       }
 
-      messagesCache.clear(DESTINATION_UUID);
+      messagesCache.clear(DESTINATION_UUID).join();
 
       assertEquals(Collections.emptyList(), get(DESTINATION_UUID, DESTINATION_DEVICE_ID, messageCount));
       assertEquals(Collections.emptyList(), get(DESTINATION_UUID, DESTINATION_DEVICE_ID + 1, messageCount));
     }
 
-    @Test
-    void testClearNullUuid() {
-      // We're happy as long as this doesn't throw an exception
-      messagesCache.clear(null);
-    }
-
     @Test
     void testGetAccountFromQueueName() {
       assertEquals(DESTINATION_UUID,