Perform cleanup operations before overwriting an existing account record

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

@@ -468,6 +468,11 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         .minThreads(8)
         .maxThreads(8)
         .build();
+    ExecutorService clientPresenceExecutor = environment.lifecycle()
+        .executorService(name(getClass(), "clientPresence-%d"))
+        .minThreads(8)
+        .maxThreads(8)
+        .build();
     ScheduledExecutorService subscriptionProcessorRetryExecutor = environment.lifecycle()
         .scheduledExecutorService(name(getClass(), "subscriptionProcessorRetry-%d")).threads(1).build();
 
@@ -540,7 +545,8 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         accountLockManager, keysManager, messagesManager, profilesManager,
         secureStorageClient, secureValueRecovery2Client,
         clientPresenceManager,
-        experimentEnrollmentManager, registrationRecoveryPasswordsManager, accountLockExecutor, clock);
+        experimentEnrollmentManager, registrationRecoveryPasswordsManager, accountLockExecutor, clientPresenceExecutor,
+        clock);
     RemoteConfigsManager remoteConfigsManager = new RemoteConfigsManager(remoteConfigs);
     APNSender apnSender = new APNSender(apnSenderExecutor, config.getApnConfiguration());
     FcmSender fcmSender = new FcmSender(fcmSenderExecutor, config.getFcmConfiguration().credentials().value());

service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java

@@ -29,6 +29,7 @@ import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.CompletionStage;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
@@ -36,6 +37,7 @@ import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.identity.IdentityType;
 import org.whispersystems.textsecuregcm.util.AsyncTimerUtil;
 import org.whispersystems.textsecuregcm.util.AttributeValues;
 import org.whispersystems.textsecuregcm.util.ExceptionUtils;
@@ -184,7 +186,9 @@ public class Accounts extends AbstractDynamoDbStore {
         deletedAccountsTableName);
   }
 
-  public boolean create(final Account account, final Function<Account, Collection<TransactWriteItem>> additionalWriteItemsFunction) {
+  public boolean create(final Account account,
+      final Function<Account, Collection<TransactWriteItem>> additionalWriteItemsFunction,
+      final BiFunction<UUID, UUID, CompletableFuture<Void>> existingAccountCleanupOperation) {
 
     return CREATE_TIMER.record(() -> {
       try {
@@ -239,7 +243,10 @@ public class Accounts extends AbstractDynamoDbStore {
             account.setUuid(UUIDUtil.fromByteBuffer(actualAccountUuid));
             final Account existingAccount = getByAccountIdentifier(account.getUuid()).orElseThrow();
             account.setNumber(existingAccount.getNumber(), existingAccount.getPhoneNumberIdentifier());
-            joinAndUnwrapUpdateFuture(reclaimAccount(existingAccount, account, additionalWriteItemsFunction.apply(account)));
+
+            existingAccountCleanupOperation.apply(existingAccount.getIdentifier(IdentityType.ACI), existingAccount.getIdentifier(IdentityType.PNI))
+                .thenCompose(ignored -> reclaimAccount(existingAccount, account, additionalWriteItemsFunction.apply(account)))
+                .join();
 
             return false;
           }

service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java

@@ -114,6 +114,7 @@ public class AccountsManager {
   private final ExperimentEnrollmentManager experimentEnrollmentManager;
   private final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager;
   private final Executor accountLockExecutor;
+  private final Executor clientPresenceExecutor;
   private final Clock clock;
 
   private static final ObjectWriter ACCOUNT_REDIS_JSON_WRITER = SystemMapper.jsonMapper()
@@ -159,6 +160,7 @@ public class AccountsManager {
       final ExperimentEnrollmentManager experimentEnrollmentManager,
       final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager,
       final Executor accountLockExecutor,
+      final Executor clientPresenceExecutor,
       final Clock clock) {
     this.accounts = accounts;
     this.phoneNumberIdentifiers = phoneNumberIdentifiers;
@@ -173,6 +175,7 @@ public class AccountsManager {
     this.experimentEnrollmentManager = experimentEnrollmentManager;
     this.registrationRecoveryPasswordsManager = requireNonNull(registrationRecoveryPasswordsManager);
     this.accountLockExecutor = accountLockExecutor;
+    this.clientPresenceExecutor = clientPresenceExecutor;
     this.clock = requireNonNull(clock);
   }
 
@@ -243,14 +246,33 @@ public class AccountsManager {
                 aciSignedPreKey,
                 pniSignedPreKey,
                 aciPqLastResortPreKey,
-                pniPqLastResortPreKey));
+                pniPqLastResortPreKey),
+            (aci, pni) -> CompletableFuture.allOf(
+                keysManager.delete(aci),
+                keysManager.delete(pni),
+                messagesManager.clear(aci),
+                profilesManager.deleteAll(aci)
+            ).thenRunAsync(() -> clientPresenceManager.disconnectAllPresencesForUuid(aci), clientPresenceExecutor));
 
-        // create() sometimes updates the UUID, if there was a number conflict.
-        // for metrics, we want secondary to run with the same original UUID
-        final UUID actualUuid = account.getUuid();
+        if (!account.getUuid().equals(originalUuid)) {
+          // If the UUID changed, then we overwrote an existing account. We should have cleared all messages before
+          // overwriting the old account, but more may have arrived while we were working. Similarly, the old account
+          // holder could have added keys or profiles. We'll largely repeat the cleanup process after creating the
+          // account to make sure we really REALLY got everything.
+          //
+          // We exclude the primary device's repeated-use keys from deletion because new keys were provided as
+          // part of the account creation process, and we don't want to delete the keys that just got added.
+          CompletableFuture.allOf(keysManager.delete(account.getIdentifier(IdentityType.ACI), true),
+              keysManager.delete(account.getIdentifier(IdentityType.PNI), true),
+              messagesManager.clear(account.getIdentifier(IdentityType.ACI)),
+              profilesManager.deleteAll(account.getIdentifier(IdentityType.ACI)))
+              .join();
+        }
 
         redisSet(account);
 
+        final Tags tags;
+
         // In terms of previously-existing accounts, there are three possible cases:
         //
         // 1. This is a completely new account; there was no pre-existing account and no recently-deleted account
@@ -259,27 +281,6 @@ public class AccountsManager {
         //    instance to match the stored account record (i.e. originalUuid != actualUuid).
         // 3. This is a re-registration of a recently-deleted account, in which case maybeRecentlyDeletedUuid is
         //    present.
-        //
-        // All cases are mutually-exclusive. In the first case, we don't need to do anything. In the third, we can be
-        // confident that everything has already been deleted. In the second case, though, we're taking over an existing
-        // account and need to clear out messages and keys that may have been stored for the old account.
-        if (!originalUuid.equals(actualUuid)) {
-          // We exclude the primary device's repeated-use keys from deletion because new keys were provided as part of
-          // the account creation process, and we don't want to delete the keys that just got added.
-          final CompletableFuture<Void> deleteKeysFuture = CompletableFuture.allOf(
-              keysManager.delete(actualUuid, true),
-              keysManager.delete(account.getPhoneNumberIdentifier(), true));
-
-          messagesManager.clear(actualUuid).join();
-          profilesManager.deleteAll(actualUuid).join();
-
-          deleteKeysFuture.join();
-
-          clientPresenceManager.disconnectAllPresencesForUuid(actualUuid);
-        }
-
-        final Tags tags;
-
         if (freshUser) {
           tags = Tags.of("type", maybeRecentlyDeletedAccountIdentifier.isPresent() ? "recently-deleted" : "new");
         } else {

service/src/main/java/org/whispersystems/textsecuregcm/workers/AssignUsernameCommand.java

@@ -116,6 +116,8 @@ public class AssignUsernameCommand extends EnvironmentCommand<WhisperServerConfi
         .executorService(name(getClass(), "storageService-%d")).maxThreads(1).minThreads(1).build();
     ExecutorService accountLockExecutor = environment.lifecycle()
         .executorService(name(getClass(), "accountLock-%d")).minThreads(1).maxThreads(1).build();
+    ExecutorService clientPresenceExecutor = environment.lifecycle()
+        .executorService(name(getClass(), "clientPresence-%d")).minThreads(1).maxThreads(1).build();
     ScheduledExecutorService secureValueRecoveryServiceRetryExecutor = environment.lifecycle()
         .scheduledExecutorService(name(getClass(), "secureValueRecoveryServiceRetry-%d")).threads(1).build();
     ScheduledExecutorService storageServiceRetryExecutor = environment.lifecycle()
@@ -202,7 +204,8 @@ public class AssignUsernameCommand extends EnvironmentCommand<WhisperServerConfi
     AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
         accountLockManager, keys, messagesManager, profilesManager,
             secureStorageClient, secureValueRecovery2Client, clientPresenceManager,
-        experimentEnrollmentManager, registrationRecoveryPasswordsManager, accountLockExecutor, Clock.systemUTC());
+        experimentEnrollmentManager, registrationRecoveryPasswordsManager, accountLockExecutor, clientPresenceExecutor,
+        Clock.systemUTC());
 
     final String usernameHash = namespace.getString("usernameHash");
     final String encryptedUsername = namespace.getString("encryptedUsername");

service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java

@@ -88,6 +88,8 @@ record CommandDependencies(
         .executorService(name(name, "storageService-%d")).maxThreads(8).minThreads(8).build();
     ExecutorService accountLockExecutor = environment.lifecycle()
         .executorService(name(name, "accountLock-%d")).minThreads(8).maxThreads(8).build();
+    ExecutorService clientPresenceExecutor = environment.lifecycle()
+        .executorService(name(name, "clientPresence-%d")).minThreads(8).maxThreads(8).build();
 
     ScheduledExecutorService secureValueRecoveryServiceRetryExecutor = environment.lifecycle()
         .scheduledExecutorService(name(name, "secureValueRecoveryServiceRetry-%d")).threads(1).build();
@@ -177,7 +179,8 @@ record CommandDependencies(
     AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
         accountLockManager, keys, messagesManager, profilesManager,
         secureStorageClient, secureValueRecovery2Client, clientPresenceManager,
-        experimentEnrollmentManager, registrationRecoveryPasswordsManager, accountLockExecutor, clock);
+        experimentEnrollmentManager, registrationRecoveryPasswordsManager, accountLockExecutor, clientPresenceExecutor,
+        clock);
 
     environment.lifecycle().manage(messagesCache);
     environment.lifecycle().manage(clientPresenceManager);