From cb50b44d8f4f5e066b9602af3f9ec387a4224c16 Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Thu, 10 Nov 2022 12:53:52 -0500
Subject: [PATCH] Allow the account cleaner to operate on multiple accounts in
 parallel

---
 .../textsecuregcm/WhisperServerService.java   |  3 +-
 .../textsecuregcm/storage/AccountCleaner.java | 36 +++++++++++++++----
 .../storage/AccountCleanerTest.java           | 15 +++++++-
 3 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index 27f5b8d1..10131462 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -403,6 +403,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     ExecutorService          fcmSenderExecutor                    = environment.lifecycle().executorService(name(getClass(), "fcmSender-%d")).maxThreads(32).minThreads(32).workQueue(fcmSenderQueue).build();
     ExecutorService          backupServiceExecutor                = environment.lifecycle().executorService(name(getClass(), "backupService-%d")).maxThreads(1).minThreads(1).build();
     ExecutorService          storageServiceExecutor               = environment.lifecycle().executorService(name(getClass(), "storageService-%d")).maxThreads(1).minThreads(1).build();
+    ExecutorService          accountDeletionExecutor              = environment.lifecycle().executorService(name(getClass(), "accountCleaner-%d")).maxThreads(16).minThreads(16).build();
 
     // TODO: generally speaking this is a DynamoDB I/O executor for the accounts table; we should eventually have a general executor for speaking to the accounts table, but most of the server is still synchronous so this isn't widely useful yet
     ExecutorService batchIdentityCheckExecutor = environment.lifecycle().executorService(name(getClass(), "batchIdentityCheck-%d")).minThreads(32).maxThreads(32).build();
@@ -552,7 +553,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         new AccountDatabaseCrawlerCache(cacheCluster, AccountDatabaseCrawlerCache.ACCOUNT_CLEANER_PREFIX);
     AccountDatabaseCrawler accountCleanerAccountDatabaseCrawler = new AccountDatabaseCrawler("Account cleaner crawler",
         accountsManager,
-        accountCleanerAccountDatabaseCrawlerCache, List.of(new AccountCleaner(accountsManager)),
+        accountCleanerAccountDatabaseCrawlerCache, List.of(new AccountCleaner(accountsManager, accountDeletionExecutor)),
         config.getAccountDatabaseCrawlerConfiguration().getChunkSize(),
         config.getAccountDatabaseCrawlerConfiguration().getChunkIntervalMs()
     );
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountCleaner.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountCleaner.java
index 2574aaaa..051f73c2 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountCleaner.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountCleaner.java
@@ -7,9 +7,13 @@ package org.whispersystems.textsecuregcm.storage;
 import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
 
 import io.micrometer.core.instrument.Metrics;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -23,9 +27,11 @@ public class AccountCleaner extends AccountDatabaseCrawlerListener {
   private static final String DELETION_REASON_TAG_NAME = "reason";
 
   private final AccountsManager accountsManager;
+  private final Executor deletionExecutor;
 
-  public AccountCleaner(AccountsManager accountsManager) {
+  public AccountCleaner(final AccountsManager accountsManager, final Executor deletionExecutor) {
     this.accountsManager = accountsManager;
+    this.deletionExecutor = deletionExecutor;
   }
 
   @Override
@@ -38,18 +44,34 @@ public class AccountCleaner extends AccountDatabaseCrawlerListener {
 
   @Override
   protected void onCrawlChunk(Optional<UUID> fromUuid, List<Account> chunkAccounts) {
+    final List<CompletableFuture<Void>> deletionFutures = new ArrayList<>();
+
     for (Account account : chunkAccounts) {
       if (isExpired(account) || needsExplicitRemoval(account)) {
         final String deletionReason = needsExplicitRemoval(account) ? "newlyExpired" : "previouslyExpired";
 
-        try {
-          accountsManager.delete(account, AccountsManager.DeletionReason.EXPIRED);
-          Metrics.counter(DELETED_ACCOUNT_COUNTER_NAME, DELETION_REASON_TAG_NAME, deletionReason).increment();
-        } catch (final Exception e) {
-          log.warn("Failed to delete account {}", account.getUuid(), e);
-        }
+        deletionFutures.add(CompletableFuture.runAsync(() -> {
+              try {
+                accountsManager.delete(account, AccountsManager.DeletionReason.EXPIRED);
+              } catch (final InterruptedException e) {
+                throw new CompletionException(e);
+              }
+            }, deletionExecutor)
+            .whenComplete((ignored, throwable) -> {
+              if (throwable != null) {
+                log.warn("Failed to delete account {}", account.getUuid(), throwable);
+              } else {
+                Metrics.counter(DELETED_ACCOUNT_COUNTER_NAME, DELETION_REASON_TAG_NAME, deletionReason).increment();
+              }
+            }));
       }
     }
+
+    try {
+      CompletableFuture.allOf(deletionFutures.toArray(new CompletableFuture[0])).join();
+    } catch (final Exception e) {
+      log.debug("Failed to delete one or more accounts in chunk", e);
+    }
   }
 
   private boolean needsExplicitRemoval(Account account) {
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountCleanerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountCleanerTest.java
index ef64e4a8..0100983f 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountCleanerTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountCleanerTest.java
@@ -15,7 +15,10 @@ import static org.mockito.Mockito.when;
 import java.util.Arrays;
 import java.util.Optional;
 import java.util.UUID;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.whispersystems.textsecuregcm.storage.AccountsManager.DeletionReason;
@@ -32,6 +35,8 @@ class AccountCleanerTest {
   private final Device  undeletedDisabledDevice  = mock(Device.class );
   private final Device  undeletedEnabledDevice   = mock(Device.class );
 
+  private ExecutorService deletionExecutor;
+
 
   @BeforeEach
   void setup() {
@@ -61,11 +66,19 @@ class AccountCleanerTest {
     when(undeletedEnabledAccount.getNumber()).thenReturn("+14153333333");
     when(undeletedEnabledAccount.getLastSeen()).thenReturn(System.currentTimeMillis() - TimeUnit.DAYS.toMillis(364));
     when(undeletedEnabledAccount.getUuid()).thenReturn(UUID.randomUUID());
+
+    deletionExecutor = Executors.newFixedThreadPool(2);
+  }
+
+  @AfterEach
+  void tearDown() throws InterruptedException {
+    deletionExecutor.shutdown();
+    deletionExecutor.awaitTermination(2, TimeUnit.SECONDS);
   }
 
   @Test
   void testAccounts() throws AccountDatabaseCrawlerRestartException, InterruptedException {
-    AccountCleaner accountCleaner = new AccountCleaner(accountsManager);
+    AccountCleaner accountCleaner = new AccountCleaner(accountsManager, deletionExecutor);
     accountCleaner.onCrawlStart();
     accountCleaner.timeAndProcessCrawlChunk(Optional.empty(), Arrays.asList(deletedDisabledAccount, undeletedDisabledAccount, undeletedEnabledAccount));
     accountCleaner.onCrawlEnd(Optional.empty());