Refactor/clarify account creation/reclamation process

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

@@ -0,0 +1,13 @@
+package org.whispersystems.textsecuregcm.storage;
+
+class AccountAlreadyExistsException extends Exception {
+  private final Account existingAccount;
+
+  public AccountAlreadyExistsException(final Account existingAccount) {
+    this.existingAccount = existingAccount;
+  }
+
+  public Account getExistingAccount() {
+    return existingAccount;
+  }
+}

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

@@ -14,7 +14,6 @@ import com.google.common.base.Throwables;
 import io.micrometer.core.instrument.Metrics;
 import io.micrometer.core.instrument.Timer;
 import java.io.IOException;
-import java.nio.ByteBuffer;
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
@@ -29,15 +28,12 @@ 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;
 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;
@@ -186,84 +182,86 @@ public class Accounts extends AbstractDynamoDbStore {
         deletedAccountsTableName);
   }
 
-  public boolean create(final Account account,
+  boolean create(final Account account, final List<TransactWriteItem> additionalWriteItems)
-      final Function<Account, Collection<TransactWriteItem>> additionalWriteItemsFunction,
+      throws AccountAlreadyExistsException {
-      final BiFunction<UUID, UUID, CompletableFuture<Void>> existingAccountCleanupOperation) {
+
+    final Timer.Sample sample = Timer.start();
+
+    try {
+      final AttributeValue uuidAttr = AttributeValues.fromUUID(account.getUuid());
+      final AttributeValue numberAttr = AttributeValues.fromString(account.getNumber());
+      final AttributeValue pniUuidAttr = AttributeValues.fromUUID(account.getPhoneNumberIdentifier());
+
+      final TransactWriteItem phoneNumberConstraintPut = buildConstraintTablePutIfAbsent(
+          phoneNumberConstraintTableName, uuidAttr, ATTR_ACCOUNT_E164, numberAttr);
+
+      final TransactWriteItem phoneNumberIdentifierConstraintPut = buildConstraintTablePutIfAbsent(
+          phoneNumberIdentifierConstraintTableName, uuidAttr, ATTR_PNI_UUID, pniUuidAttr);
+
+      final TransactWriteItem accountPut = buildAccountPut(account, uuidAttr, numberAttr, pniUuidAttr);
+
+      // Clear any "recently deleted account" record for this number since, if it existed, we've used its old ACI for
+      // the newly-created account.
+      final TransactWriteItem deletedAccountDelete = buildRemoveDeletedAccount(account.getNumber());
+
+      final Collection<TransactWriteItem> writeItems = new ArrayList<>(
+          List.of(phoneNumberConstraintPut, phoneNumberIdentifierConstraintPut, accountPut, deletedAccountDelete));
+
+      writeItems.addAll(additionalWriteItems);
+
+      final TransactWriteItemsRequest request = TransactWriteItemsRequest.builder()
+          .transactItems(writeItems)
+          .build();
 
-    return CREATE_TIMER.record(() -> {
       try {
-        final AttributeValue uuidAttr = AttributeValues.fromUUID(account.getUuid());
+        db().transactWriteItems(request);
-        final AttributeValue numberAttr = AttributeValues.fromString(account.getNumber());
+      } catch (final TransactionCanceledException e) {
-        final AttributeValue pniUuidAttr = AttributeValues.fromUUID(account.getPhoneNumberIdentifier());
 
-        final TransactWriteItem phoneNumberConstraintPut = buildConstraintTablePutIfAbsent(
+        final CancellationReason accountCancellationReason = e.cancellationReasons().get(2);
-            phoneNumberConstraintTableName, uuidAttr, ATTR_ACCOUNT_E164, numberAttr);
 
-        final TransactWriteItem phoneNumberIdentifierConstraintPut = buildConstraintTablePutIfAbsent(
+        if (conditionalCheckFailed(accountCancellationReason)) {
-            phoneNumberIdentifierConstraintTableName, uuidAttr, ATTR_PNI_UUID, pniUuidAttr);
+          throw new IllegalArgumentException("account identifier present with different phone number");
-
-        final TransactWriteItem accountPut = buildAccountPut(account, uuidAttr, numberAttr, pniUuidAttr);
-
-        // Clear any "recently deleted account" record for this number since, if it existed, we've used its old ACI for
-        // the newly-created account.
-        final TransactWriteItem deletedAccountDelete = buildRemoveDeletedAccount(account.getNumber());
-
-        final Collection<TransactWriteItem> writeItems = new ArrayList<>(
-            List.of(phoneNumberConstraintPut, phoneNumberIdentifierConstraintPut, accountPut, deletedAccountDelete));
-
-        writeItems.addAll(additionalWriteItemsFunction.apply(account));
-
-        final TransactWriteItemsRequest request = TransactWriteItemsRequest.builder()
-            .transactItems(writeItems)
-            .build();
-
-        try {
-          db().transactWriteItems(request);
-        } catch (final TransactionCanceledException e) {
-
-          final CancellationReason accountCancellationReason = e.cancellationReasons().get(2);
-
-          if (conditionalCheckFailed(accountCancellationReason)) {
-            throw new IllegalArgumentException("account identifier present with different phone number");
-          }
-
-          final CancellationReason phoneNumberConstraintCancellationReason = e.cancellationReasons().get(0);
-          final CancellationReason phoneNumberIdentifierConstraintCancellationReason = e.cancellationReasons().get(1);
-
-          if (conditionalCheckFailed(phoneNumberConstraintCancellationReason)
-              || conditionalCheckFailed(phoneNumberIdentifierConstraintCancellationReason)) {
-
-            // In theory, both reasons should trip in tandem and either should give us the information we need. Even so,
-            // we'll be cautious here and make sure we're choosing a condition check that really failed.
-            final CancellationReason reason = conditionalCheckFailed(phoneNumberConstraintCancellationReason)
-                ? phoneNumberConstraintCancellationReason
-                : phoneNumberIdentifierConstraintCancellationReason;
-
-            final ByteBuffer actualAccountUuid = reason.item().get(KEY_ACCOUNT_UUID).b().asByteBuffer();
-            account.setUuid(UUIDUtil.fromByteBuffer(actualAccountUuid));
-            final Account existingAccount = getByAccountIdentifier(account.getUuid()).orElseThrow();
-            account.setNumber(existingAccount.getNumber(), existingAccount.getPhoneNumberIdentifier());
-
-            existingAccountCleanupOperation.apply(existingAccount.getIdentifier(IdentityType.ACI), existingAccount.getIdentifier(IdentityType.PNI))
-                .thenCompose(ignored -> reclaimAccount(existingAccount, account, additionalWriteItemsFunction.apply(account)))
-                .join();
-
-            return false;
-          }
-
-          if (TRANSACTION_CONFLICT.equals(accountCancellationReason.code())) {
-            // this should only happen if two clients manage to make concurrent create() calls
-            throw new ContestedOptimisticLockException();
-          }
-
-          // this shouldn't happen
-          throw new RuntimeException("could not create account: " + extractCancellationReasonCodes(e));
         }
-      } catch (final JsonProcessingException e) {
+
-        throw new IllegalArgumentException(e);
+        final CancellationReason phoneNumberConstraintCancellationReason = e.cancellationReasons().get(0);
+        final CancellationReason phoneNumberIdentifierConstraintCancellationReason = e.cancellationReasons().get(1);
+
+        if (conditionalCheckFailed(phoneNumberConstraintCancellationReason)
+            || conditionalCheckFailed(phoneNumberIdentifierConstraintCancellationReason)) {
+
+          // In theory, both reasons should trip in tandem and either should give us the information we need. Even so,
+          // we'll be cautious here and make sure we're choosing a condition check that really failed.
+          final CancellationReason reason = conditionalCheckFailed(phoneNumberConstraintCancellationReason)
+              ? phoneNumberConstraintCancellationReason
+              : phoneNumberIdentifierConstraintCancellationReason;
+
+          final UUID existingAccountUuid =
+              UUIDUtil.fromByteBuffer(reason.item().get(KEY_ACCOUNT_UUID).b().asByteBuffer());
+
+          // This is unlikely, but it could be that the existing account was deleted in between the time the transaction
+          // happened and when we tried to read the full existing account. If that happens, we can just consider this a
+          // contested lock, and retrying is likely to succeed.
+          final Account existingAccount = getByAccountIdentifier(existingAccountUuid)
+              .orElseThrow(ContestedOptimisticLockException::new);
+
+          throw new AccountAlreadyExistsException(existingAccount);
+        }
+
+        if (TRANSACTION_CONFLICT.equals(accountCancellationReason.code())) {
+          // this should only happen if two clients manage to make concurrent create() calls
+          throw new ContestedOptimisticLockException();
+        }
+
+        // this shouldn't happen
+        throw new RuntimeException("could not create account: " + extractCancellationReasonCodes(e));
       }
-      return true;
+    } catch (final JsonProcessingException e) {
-    });
+      throw new IllegalArgumentException(e);
+    } finally {
+      sample.stop(CREATE_TIMER);
+    }
+
+    return true;
   }
 
   /**
@@ -272,9 +270,13 @@ public class Accounts extends AbstractDynamoDbStore {
    * @param existingAccount the existing account in the accounts table
    * @param accountToCreate a new account, with the same number and identifier as existingAccount
    */
-  private CompletionStage<Void> reclaimAccount(final Account existingAccount, final Account accountToCreate, final Collection<TransactWriteItem> additionalWriteItems) {
+  CompletionStage<Void> reclaimAccount(final Account existingAccount,
+      final Account accountToCreate,
+      final Collection<TransactWriteItem> additionalWriteItems) {
+
     if (!existingAccount.getUuid().equals(accountToCreate.getUuid()) ||
         !existingAccount.getNumber().equals(accountToCreate.getNumber())) {
+
       throw new IllegalArgumentException("reclaimed accounts must match");
     }
 

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

@@ -18,7 +18,6 @@ import com.google.common.base.Preconditions;
 import io.lettuce.core.RedisException;
 import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands;
 import io.micrometer.core.instrument.Metrics;
-import io.micrometer.core.instrument.Tags;
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.time.Clock;
@@ -182,80 +181,79 @@ public class AccountsManager {
       final IdentityKey pniIdentityKey,
       final DeviceSpec primaryDeviceSpec) throws InterruptedException {
 
-    try (Timer.Context ignored = createTimer.time()) {
+    final Account account = new Account();
-      final Account account = new Account();
 
+    try (Timer.Context ignoredTimerContext = createTimer.time()) {
       accountLockManager.withLock(List.of(number), () -> {
-        final Device device = primaryDeviceSpec.toDevice(Device.PRIMARY_ID, clock);
-
-        account.setNumber(number, phoneNumberIdentifiers.getPhoneNumberIdentifier(number));
-
         final Optional<UUID> maybeRecentlyDeletedAccountIdentifier =
             accounts.findRecentlyDeletedAccountIdentifier(number);
 
         // Reuse the ACI from any recently-deleted account with this number to cover cases where somebody is
         // re-registering.
+        account.setNumber(number, phoneNumberIdentifiers.getPhoneNumberIdentifier(number));
         account.setUuid(maybeRecentlyDeletedAccountIdentifier.orElseGet(UUID::randomUUID));
         account.setIdentityKey(aciIdentityKey);
         account.setPhoneNumberIdentityKey(pniIdentityKey);
-        account.addDevice(device);
+        account.addDevice(primaryDeviceSpec.toDevice(Device.PRIMARY_ID, clock));
         account.setRegistrationLockFromAttributes(accountAttributes);
         account.setUnidentifiedAccessKey(accountAttributes.getUnidentifiedAccessKey());
         account.setUnrestrictedUnidentifiedAccess(accountAttributes.isUnrestrictedUnidentifiedAccess());
         account.setDiscoverableByPhoneNumber(accountAttributes.isDiscoverableByPhoneNumber());
         account.setBadges(clock, accountBadges);
 
-        final UUID originalUuid = account.getUuid();
+        String accountCreationType = maybeRecentlyDeletedAccountIdentifier.isPresent() ? "recently-deleted" : "new";
 
-        final boolean freshUser = accounts.create(account,
+        try {
-            a -> keysManager.buildWriteItemsForRepeatedUseKeys(a.getIdentifier(IdentityType.ACI),
+          accounts.create(account, keysManager.buildWriteItemsForRepeatedUseKeys(account.getIdentifier(IdentityType.ACI),
-                a.getIdentifier(IdentityType.PNI),
+              account.getIdentifier(IdentityType.PNI),
-                Device.PRIMARY_ID,
+              Device.PRIMARY_ID,
-                primaryDeviceSpec.aciSignedPreKey(),
+              primaryDeviceSpec.aciSignedPreKey(),
-                primaryDeviceSpec.pniSignedPreKey(),
+              primaryDeviceSpec.pniSignedPreKey(),
-                primaryDeviceSpec.aciPqLastResortPreKey(),
+              primaryDeviceSpec.aciPqLastResortPreKey(),
-                primaryDeviceSpec.pniPqLastResortPreKey()),
+              primaryDeviceSpec.pniPqLastResortPreKey()));
-            (aci, pni) -> CompletableFuture.allOf(
+        } catch (final AccountAlreadyExistsException e) {
-                keysManager.delete(aci),
+          accountCreationType = "re-registration";
-                keysManager.delete(pni),
-                messagesManager.clear(aci),
-                profilesManager.deleteAll(aci)
-            ).thenRunAsync(() -> clientPresenceManager.disconnectAllPresencesForUuid(aci), clientPresenceExecutor));
 
-        if (!account.getUuid().equals(originalUuid)) {
+          final UUID aci = e.getExistingAccount().getIdentifier(IdentityType.ACI);
-          // If the UUID changed, then we overwrote an existing account. We should have cleared all messages before
+          final UUID pni = e.getExistingAccount().getIdentifier(IdentityType.PNI);
-          // 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.setUuid(aci);
-          // account to make sure we really REALLY got everything.
+          account.setNumber(e.getExistingAccount().getNumber(), pni);
-          //
+
-          // We exclude the primary device's repeated-use keys from deletion because new keys were provided as
+          CompletableFuture.allOf(
-          // part of the account creation process, and we don't want to delete the keys that just got added.
+                  keysManager.delete(aci),
-          CompletableFuture.allOf(keysManager.delete(account.getIdentifier(IdentityType.ACI), true),
+                  keysManager.delete(pni),
-              keysManager.delete(account.getIdentifier(IdentityType.PNI), true),
+                  messagesManager.clear(aci),
-              messagesManager.clear(account.getIdentifier(IdentityType.ACI)),
+                  profilesManager.deleteAll(aci))
-              profilesManager.deleteAll(account.getIdentifier(IdentityType.ACI)))
+              .thenRunAsync(() -> clientPresenceManager.disconnectAllPresencesForUuid(aci), clientPresenceExecutor)
+              .thenCompose(ignored -> accounts.reclaimAccount(e.getExistingAccount(),
+                  account,
+                  keysManager.buildWriteItemsForRepeatedUseKeys(account.getIdentifier(IdentityType.ACI),
+                      account.getIdentifier(IdentityType.PNI),
+                      Device.PRIMARY_ID,
+                      primaryDeviceSpec.aciSignedPreKey(),
+                      primaryDeviceSpec.pniSignedPreKey(),
+                      primaryDeviceSpec.aciPqLastResortPreKey(),
+                      primaryDeviceSpec.pniPqLastResortPreKey())))
+              .thenCompose(ignored -> {
+                // 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.
+                return CompletableFuture.allOf(keysManager.delete(aci, true),
+                        keysManager.delete(pni, true),
+                        messagesManager.clear(aci),
+                        profilesManager.deleteAll(aci));
+              })
               .join();
         }
 
         redisSet(account);
 
-        final Tags tags;
+        Metrics.counter(CREATE_COUNTER_NAME, "type", accountCreationType).increment();
-
-        // 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
-        // 2. This is a re-registration of an existing account. The storage layer will update the existing account in
-        //    place to match the account record created above, and will update the UUID of the newly-created account
-        //    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.
-        if (freshUser) {
-          tags = Tags.of("type", maybeRecentlyDeletedAccountIdentifier.isPresent() ? "recently-deleted" : "new");
-        } else {
-          tags = Tags.of("type", "re-registration");
-        }
-
-        Metrics.counter(CREATE_COUNTER_NAME, tags).increment();
 
         accountAttributes.recoveryPassword().ifPresent(registrationRecoveryPassword ->
             registrationRecoveryPasswordsManager.storeForCurrentNumber(account.getNumber(), registrationRecoveryPassword));

service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerTest.java

@@ -50,7 +50,6 @@ import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ThreadLocalRandom;
-import java.util.function.BiFunction;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
@@ -868,14 +867,14 @@ class AccountsManagerTest {
   }
 
   @Test
-  void testUpdate_dynamoOptimisticLockingFailureDuringCreate() {
+  void testUpdate_dynamoOptimisticLockingFailureDuringCreate() throws AccountAlreadyExistsException {
     UUID uuid = UUID.randomUUID();
     Account account = AccountsHelper.generateTestAccount("+14152222222", uuid, UUID.randomUUID(), new ArrayList<>(), new byte[UnidentifiedAccessUtil.UNIDENTIFIED_ACCESS_KEY_LENGTH]);
 
     when(commands.get(eq("Account3::" + uuid))).thenReturn(null);
     when(accounts.getByAccountIdentifier(uuid)).thenReturn(Optional.empty())
         .thenReturn(Optional.of(account));
-    when(accounts.create(any(), any(), any())).thenThrow(ContestedOptimisticLockException.class);
+    when(accounts.create(any(), any())).thenThrow(ContestedOptimisticLockException.class);
 
     accountsManager.update(account, a -> {
     });
@@ -992,42 +991,49 @@ class AccountsManagerTest {
   }
 
   @Test
-  void testCreateFreshAccount() throws InterruptedException {
+  void testCreateFreshAccount() throws InterruptedException, AccountAlreadyExistsException {
-    when(accounts.create(any(), any(), any())).thenReturn(true);
+    when(accounts.create(any(), any())).thenReturn(true);
 
     final String e164 = "+18005550123";
     final AccountAttributes attributes = new AccountAttributes(false, 1, 2, null, null, true, null);
 
     createAccount(e164, attributes);
 
-    verify(accounts).create(argThat(account -> e164.equals(account.getNumber())), any(), any());
+    verify(accounts).create(argThat(account -> e164.equals(account.getNumber())), any());
 
     verifyNoInteractions(messagesManager);
     verifyNoInteractions(profilesManager);
   }
 
   @Test
-  void testReregisterAccount() throws InterruptedException {
+  void testReregisterAccount() throws InterruptedException, AccountAlreadyExistsException {
     final UUID existingUuid = UUID.randomUUID();
 
     final String e164 = "+18005550123";
     final AccountAttributes attributes = new AccountAttributes(false, 1, 2, null, null, true, null);
 
-    when(accounts.create(any(), any(), any())).thenAnswer(invocation -> {
+    when(accounts.create(any(), any()))
-      invocation.getArgument(0, Account.class).setUuid(existingUuid);
+        .thenAnswer(invocation -> {
+          final Account requestedAccount = invocation.getArgument(0);
 
-      final BiFunction<UUID, UUID, CompletableFuture<Void>> cleanupOperation = invocation.getArgument(2);
+          final Account existingAccount = mock(Account.class);
-      cleanupOperation.apply(existingUuid, phoneNumberIdentifiersByE164.get(e164));
+          when(existingAccount.getUuid()).thenReturn(existingUuid);
+          when(existingAccount.getIdentifier(IdentityType.ACI)).thenReturn(existingUuid);
+          when(existingAccount.getNumber()).thenReturn(e164);
+          when(existingAccount.getPhoneNumberIdentifier()).thenReturn(requestedAccount.getIdentifier(IdentityType.PNI));
+          when(existingAccount.getIdentifier(IdentityType.PNI)).thenReturn(requestedAccount.getIdentifier(IdentityType.PNI));
 
-      return false;
+          throw new AccountAlreadyExistsException(existingAccount);
-    });
+        });
+
+    when(accounts.reclaimAccount(any(), any(), any())).thenReturn(CompletableFuture.completedFuture(null));
 
     createAccount(e164, attributes);
 
     assertTrue(phoneNumberIdentifiersByE164.containsKey(e164));
 
     verify(accounts)
-        .create(argThat(account -> e164.equals(account.getNumber()) && existingUuid.equals(account.getUuid())), any(), any());
+        .create(argThat(account -> e164.equals(account.getNumber()) && existingUuid.equals(account.getUuid())), any());
 
     verify(keysManager).delete(existingUuid);
     verify(keysManager).delete(phoneNumberIdentifiersByE164.get(e164));
@@ -1039,23 +1045,30 @@ class AccountsManagerTest {
   }
 
   @Test
-  void testCreateAccountRecentlyDeleted() throws InterruptedException {
+  void testCreateAccountRecentlyDeleted() throws InterruptedException, AccountAlreadyExistsException {
     final UUID recentlyDeletedUuid = UUID.randomUUID();
 
     when(accounts.findRecentlyDeletedAccountIdentifier(anyString())).thenReturn(Optional.of(recentlyDeletedUuid));
-    when(accounts.create(any(), any(), any())).thenReturn(true);
+    when(accounts.create(any(), any())).thenReturn(true);
 
     final String e164 = "+18005550123";
     final AccountAttributes attributes = new AccountAttributes(false, 1, 2, null, null, true, null);
 
-    createAccount(e164, attributes);
+    final Account account = createAccount(e164, attributes);
 
     verify(accounts).create(
-        argThat(account -> e164.equals(account.getNumber()) && recentlyDeletedUuid.equals(account.getUuid())),
+        argThat(a -> e164.equals(a.getNumber()) && recentlyDeletedUuid.equals(a.getUuid())),
+        any());
+
+    verify(keysManager).buildWriteItemsForRepeatedUseKeys(eq(account.getIdentifier(IdentityType.ACI)),
+        eq(account.getIdentifier(IdentityType.PNI)),
+        eq(Device.PRIMARY_ID),
+        any(),
+        any(),
         any(),
         any());
 
-    verifyNoInteractions(keysManager);
+    verifyNoMoreInteractions(keysManager);
     verifyNoInteractions(messagesManager);
     verifyNoInteractions(profilesManager);
   }

service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java

@@ -309,12 +309,11 @@ class AccountsManagerUsernameIntegrationTest {
   }
 
   @Test
-  public void testUsernameLinks() throws InterruptedException {
+  public void testUsernameLinks() throws InterruptedException, AccountAlreadyExistsException {
     final Account account = AccountsHelper.createAccount(accountsManager, "+18005551111");
 
     account.setUsernameHash(TestRandomUtil.nextBytes(16));
-    accounts.create(account, ignored -> Collections.emptyList(),
+    accounts.create(account, Collections.emptyList());
-        (ignoredAci, ignoredPni) -> CompletableFuture.completedFuture(null));
 
     final UUID linkHandle = UUID.randomUUID();
     final byte[] encryptedUsername = TestRandomUtil.nextBytes(32);

service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java

@@ -9,6 +9,7 @@ import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatNoException;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertInstanceOf;
@@ -214,34 +215,6 @@ class AccountsTest {
     assertThat(accounts.findRecentlyDeletedAccountIdentifier(account.getNumber())).isEmpty();
   }
 
-  @Test
-  void testStoreCleanupFailure() {
-    final Account existingAccount = nextRandomAccount();
-    createAccount(existingAccount);
-
-    verifyStoredState(existingAccount.getNumber(),
-        existingAccount.getUuid(),
-        existingAccount.getPhoneNumberIdentifier(),
-        existingAccount.getUsernameHash().orElse(null),
-        existingAccount,
-        true);
-
-    final CompletionException completionException = assertThrows(CompletionException.class,
-        () -> accounts.create(generateAccount(existingAccount.getNumber(), UUID.randomUUID(), UUID.randomUUID()),
-            ignored -> Collections.emptyList(),
-            (aci, pni) -> CompletableFuture.failedFuture(new RuntimeException("OH NO"))));
-
-    assertTrue(completionException.getCause() instanceof RuntimeException);
-
-    // If the existing account cleanup task failed, we should not overwrite the existing account record
-    verifyStoredState(existingAccount.getNumber(),
-        existingAccount.getUuid(),
-        existingAccount.getPhoneNumberIdentifier(),
-        existingAccount.getUsernameHash().orElse(null),
-        existingAccount,
-        true);
-  }
-
   @Test
   void testStoreMulti() {
     final List<Device> devices = List.of(generateDevice(DEVICE_ID_1), generateDevice(DEVICE_ID_2));
@@ -388,10 +361,10 @@ class AccountsTest {
         accounts.reserveUsernameHash(account, usernameHash, Duration.ofMinutes(1)).join();
         accounts.confirmUsernameHash(account, usernameHash, encryptedUsername).join();
 
-        // simulate a failed re-reg: we give the account a reclaimable username, but we'll try
+        // simulate a partially-completed re-reg: we give the account a reclaimable username, but we'll try
         // re-registering again later in the test case
         account = generateAccount("+14151112222", UUID.randomUUID(), UUID.randomUUID(), List.of(generateDevice(DEVICE_ID_1)));
-        createAccount(account);
+        reclaimAccount(account);
         break;
       case CONFIRMED:
         accounts.reserveUsernameHash(account, usernameHash, Duration.ofMinutes(1)).join();
@@ -403,7 +376,7 @@ class AccountsTest {
 
     // re-register the account
     account = generateAccount("+14151112222", UUID.randomUUID(), UUID.randomUUID(), List.of(generateDevice(DEVICE_ID_1)));
-    createAccount(account);
+    reclaimAccount(account);
 
     // If we had a username link, or we had previously saved a username link from another re-registration, make sure
     // we preserve it
@@ -418,50 +391,63 @@ class AccountsTest {
     assertThat(account.getUsernameLinkHandle().equals(preservedLink.orElse(null)))
         .isEqualTo(shouldReuseLink);
 
-
     // in all cases, we should now have usernameHash, usernameLink, and encryptedUsername set
     assertThat(account.getUsernameHash()).isNotEmpty();
     assertThat(account.getEncryptedUsername()).isNotEmpty();
     assertThat(account.getUsernameLinkHandle()).isNotNull();
     assertThat(account.getReservedUsernameHash()).isEmpty();
+  }
 
+  private void reclaimAccount(final Account reregisteredAccount) {
+    final AccountAlreadyExistsException accountAlreadyExistsException =
+        assertThrows(AccountAlreadyExistsException.class,
+            () -> accounts.create(reregisteredAccount, Collections.emptyList()));
+
+    reregisteredAccount.setUuid(accountAlreadyExistsException.getExistingAccount().getUuid());
+    reregisteredAccount.setNumber(accountAlreadyExistsException.getExistingAccount().getNumber(),
+        accountAlreadyExistsException.getExistingAccount().getPhoneNumberIdentifier());
+
+    assertDoesNotThrow(() -> accounts.reclaimAccount(accountAlreadyExistsException.getExistingAccount(),
+        reregisteredAccount,
+        Collections.emptyList()).toCompletableFuture().join());
   }
 
   @Test
-  void testOverwrite() {
+  void testReclaimAccount() {
-    Device device = generateDevice(DEVICE_ID_1);
+    final String e164 = "+14151112222";
-    UUID firstUuid = UUID.randomUUID();
+    final Device device = generateDevice(DEVICE_ID_1);
-    UUID firstPni = UUID.randomUUID();
+    final UUID existingUuid = UUID.randomUUID();
-    Account account = generateAccount("+14151112222", firstUuid, firstPni, List.of(device));
+    final UUID existingPni = UUID.randomUUID();
+    final Account existingAccount = generateAccount(e164, existingUuid, existingPni, List.of(device));
 
-    createAccount(account);
+    createAccount(existingAccount);
 
     final byte[] usernameHash = randomBytes(32);
     final byte[] encryptedUsername = randomBytes(16);
 
     // Set up the existing account to have a username hash
-    accounts.confirmUsernameHash(account, usernameHash, encryptedUsername).join();
+    accounts.confirmUsernameHash(existingAccount, usernameHash, encryptedUsername).join();
-    final UUID usernameLinkHandle = account.getUsernameLinkHandle();
+    final UUID usernameLinkHandle = existingAccount.getUsernameLinkHandle();
 
-    verifyStoredState("+14151112222", account.getUuid(), account.getPhoneNumberIdentifier(), usernameHash, account, true);
+    verifyStoredState(e164, existingAccount.getUuid(), existingAccount.getPhoneNumberIdentifier(), usernameHash, existingAccount, true);
 
-    assertPhoneNumberConstraintExists("+14151112222", firstUuid);
+    assertPhoneNumberConstraintExists(e164, existingUuid);
-    assertPhoneNumberIdentifierConstraintExists(firstPni, firstUuid);
+    assertPhoneNumberIdentifierConstraintExists(existingPni, existingUuid);
 
-    accounts.update(account);
+    assertDoesNotThrow(() -> accounts.update(existingAccount));
 
-    UUID secondUuid = UUID.randomUUID();
+    final UUID secondUuid = UUID.randomUUID();
 
-    device = generateDevice(DEVICE_ID_1);
+    final Device secondDevice = generateDevice(DEVICE_ID_1);
-    account = generateAccount("+14151112222", secondUuid, UUID.randomUUID(), List.of(device));
+    final Account secondAccount = generateAccount(e164, secondUuid, UUID.randomUUID(), List.of(secondDevice));
+
+    reclaimAccount(secondAccount);
 
-    final boolean freshUser = createAccount(account);
-    assertThat(freshUser).isFalse();
     // usernameHash should be unset
-    verifyStoredState("+14151112222", firstUuid, firstPni, null, account, true);
+    verifyStoredState("+14151112222", existingUuid, existingPni, null, secondAccount, true);
 
     // username should become 'reclaimable'
-    Map<String, AttributeValue> item = readAccount(firstUuid);
+    Map<String, AttributeValue> item = readAccount(existingUuid);
     Account result = Accounts.fromItem(item);
     assertThat(AttributeValues.getUUID(item, Accounts.ATTR_USERNAME_LINK_UUID, null))
         .isEqualTo(usernameLinkHandle)
@@ -472,7 +458,7 @@ class AccountsTest {
 
     // should keep the same usernameLink, now encryptedUsername should be set
     accounts.confirmUsernameHash(result, usernameHash, encryptedUsername).join();
-    item = readAccount(firstUuid);
+    item = readAccount(existingUuid);
     result = Accounts.fromItem(item);
     assertThat(AttributeValues.getUUID(item, Accounts.ATTR_USERNAME_LINK_UUID, null))
         .isEqualTo(usernameLinkHandle)
@@ -481,11 +467,10 @@ class AccountsTest {
     assertArrayEquals(result.getUsernameHash().get(), usernameHash);
     assertThat(result.getReservedUsernameHash()).isEmpty();
 
-    assertPhoneNumberConstraintExists("+14151112222", firstUuid);
+    assertPhoneNumberConstraintExists("+14151112222", existingUuid);
-    assertPhoneNumberIdentifierConstraintExists(firstPni, firstUuid);
+    assertPhoneNumberIdentifierConstraintExists(existingPni, existingUuid);
 
-    device = generateDevice(DEVICE_ID_1);
+    Account invalidAccount = generateAccount("+14151113333", existingUuid, UUID.randomUUID(), List.of(generateDevice(DEVICE_ID_1)));
-    Account invalidAccount = generateAccount("+14151113333", firstUuid, UUID.randomUUID(), List.of(device));
 
     assertThatThrownBy(() -> createAccount(invalidAccount));
   }
@@ -1237,8 +1222,11 @@ class AccountsTest {
   }
 
   private boolean createAccount(final Account account) {
-    return accounts.create(account, ignored -> Collections.emptyList(),
+    try {
-        (ignoredAci, ignoredPni) -> CompletableFuture.completedFuture(null));
+      return accounts.create(account, Collections.emptyList());
+    } catch (AccountAlreadyExistsException e) {
+      throw new IllegalStateException(e);
+    }
   }
 
   private static Account nextRandomAccount() {