Refactor account locks/deleted account manager

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

@@ -169,11 +169,11 @@ import org.whispersystems.textsecuregcm.storage.AccountCleaner;
 import org.whispersystems.textsecuregcm.storage.AccountDatabaseCrawler;
 import org.whispersystems.textsecuregcm.storage.AccountDatabaseCrawlerCache;
 import org.whispersystems.textsecuregcm.storage.AccountDatabaseCrawlerListener;
+import org.whispersystems.textsecuregcm.storage.AccountLockManager;
 import org.whispersystems.textsecuregcm.storage.Accounts;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.ChangeNumberManager;
 import org.whispersystems.textsecuregcm.storage.DeletedAccounts;
-import org.whispersystems.textsecuregcm.storage.DeletedAccountsManager;
 import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
 import org.whispersystems.textsecuregcm.storage.IssuedReceiptsManager;
 import org.whispersystems.textsecuregcm.storage.Keys;
@@ -506,10 +506,9 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         config.getReportMessageConfiguration().getCounterTtl());
     MessagesManager messagesManager = new MessagesManager(messagesDynamoDb, messagesCache, reportMessageManager,
         messageDeletionAsyncExecutor);
-    DeletedAccountsManager deletedAccountsManager = new DeletedAccountsManager(deletedAccounts,
+    AccountLockManager accountLockManager = new AccountLockManager(deletedAccountsLockDynamoDbClient, config.getDynamoDbTables().getDeletedAccountsLock().getTableName());
-        deletedAccountsLockDynamoDbClient, config.getDynamoDbTables().getDeletedAccountsLock().getTableName());
     AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
-        deletedAccountsManager, keys, messagesManager, profilesManager,
+        accountLockManager, deletedAccounts, keys, messagesManager, profilesManager,
         pendingAccountsManager, secureStorageClient, secureBackupClient, secureValueRecovery2Client,
         clientPresenceManager,
         experimentEnrollmentManager, registrationRecoveryPasswordsManager, clock);
@@ -744,7 +743,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         new DirectoryV2Controller(directoryV2CredentialsGenerator),
         new DonationController(clock, zkReceiptOperations, redeemedReceiptsManager, accountsManager, config.getBadges(),
             ReceiptCredentialPresentation::new),
-        new MessageController(rateLimiters, messageSender, receiptSender, accountsManager, deletedAccountsManager,
+        new MessageController(rateLimiters, messageSender, receiptSender, accountsManager, deletedAccounts,
             messagesManager, pushNotificationManager, reportMessageManager, multiRecipientMessageExecutor,
             messageDeliveryScheduler, reportSpamTokenProvider),
         new PaymentsController(currencyManager, paymentsCredentialsGenerator),

service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java

@@ -93,7 +93,7 @@ import org.whispersystems.textsecuregcm.spam.FilterSpam;
 import org.whispersystems.textsecuregcm.spam.ReportSpamTokenProvider;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
-import org.whispersystems.textsecuregcm.storage.DeletedAccountsManager;
+import org.whispersystems.textsecuregcm.storage.DeletedAccounts;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.MessagesManager;
 import org.whispersystems.textsecuregcm.storage.ReportMessageManager;
@@ -117,7 +117,7 @@ public class MessageController {
   private final MessageSender messageSender;
   private final ReceiptSender receiptSender;
   private final AccountsManager accountsManager;
-  private final DeletedAccountsManager deletedAccountsManager;
+  private final DeletedAccounts deletedAccounts;
   private final MessagesManager messagesManager;
   private final PushNotificationManager pushNotificationManager;
   private final ReportMessageManager reportMessageManager;
@@ -150,7 +150,7 @@ public class MessageController {
       MessageSender messageSender,
       ReceiptSender receiptSender,
       AccountsManager accountsManager,
-      DeletedAccountsManager deletedAccountsManager,
+      DeletedAccounts deletedAccounts,
       MessagesManager messagesManager,
       PushNotificationManager pushNotificationManager,
       ReportMessageManager reportMessageManager,
@@ -161,7 +161,7 @@ public class MessageController {
     this.messageSender = messageSender;
     this.receiptSender = receiptSender;
     this.accountsManager = accountsManager;
-    this.deletedAccountsManager = deletedAccountsManager;
+    this.deletedAccounts = deletedAccounts;
     this.messagesManager = messagesManager;
     this.pushNotificationManager = pushNotificationManager;
     this.reportMessageManager = reportMessageManager;
@@ -621,7 +621,7 @@ public class MessageController {
         sourceAci = maybeAccount.map(Account::getUuid);
         sourcePni = maybeAccount.map(Account::getPhoneNumberIdentifier);
       } else {
-        sourceAci = deletedAccountsManager.findDeletedAccountAci(source);
+        sourceAci = deletedAccounts.findUuid(source);
         sourcePni = Optional.ofNullable(accountsManager.getPhoneNumberIdentifier(source));
       }
     } else {
@@ -631,7 +631,7 @@ public class MessageController {
 
       if (sourceAccount.isEmpty()) {
         logger.warn("Could not find source: {}", sourceAci.get());
-        sourceNumber = deletedAccountsManager.findDeletedAccountE164(sourceAci.get());
+        sourceNumber = deletedAccounts.findE164(sourceAci.get());
         sourcePni = sourceNumber.map(accountsManager::getPhoneNumberIdentifier);
       } else {
         sourceNumber = sourceAccount.map(Account::getNumber);

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

@@ -0,0 +1,66 @@
+package org.whispersystems.textsecuregcm.storage;
+
+import com.amazonaws.services.dynamodbv2.AcquireLockOptions;
+import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
+import com.amazonaws.services.dynamodbv2.AmazonDynamoDBLockClient;
+import com.amazonaws.services.dynamodbv2.AmazonDynamoDBLockClientOptions;
+import com.amazonaws.services.dynamodbv2.LockItem;
+import com.amazonaws.services.dynamodbv2.ReleaseLockOptions;
+import com.google.common.annotations.VisibleForTesting;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+
+public class AccountLockManager {
+
+  private final AmazonDynamoDBLockClient lockClient;
+
+  public AccountLockManager(final AmazonDynamoDB lockDynamoDb, final String lockTableName) {
+    this(new AmazonDynamoDBLockClient(
+        AmazonDynamoDBLockClientOptions.builder(lockDynamoDb, lockTableName)
+            .withPartitionKeyName(DeletedAccounts.KEY_ACCOUNT_E164)
+            .withLeaseDuration(15L)
+            .withHeartbeatPeriod(2L)
+            .withTimeUnit(TimeUnit.SECONDS)
+            .withCreateHeartbeatBackgroundThread(true)
+            .build()));
+  }
+
+  @VisibleForTesting
+  AccountLockManager(final AmazonDynamoDBLockClient lockClient) {
+    this.lockClient = lockClient;
+  }
+
+  /**
+   * Acquires a distributed, pessimistic lock for the accounts identified by the given phone numbers. By design, the
+   * accounts need not actually exist in order to acquire a lock; this allows lock acquisition for operations that span
+   * account lifecycle changes (like deleting an account or changing a phone number). The given task runs once locks for
+   * all given phone numbers have been acquired, and the locks are released as soon as the task completes by any means.
+   *
+   * @param e164s the phone numbers for which to acquire a distributed, pessimistic lock
+   * @param task the task to execute once locks have been acquired
+   *
+   * @throws InterruptedException if interrupted while acquiring a lock
+   */
+  public void withLock(final List<String> e164s, final Runnable task) throws InterruptedException {
+    if (e164s.isEmpty()) {
+      throw new IllegalArgumentException("List of e164s to lock must not be empty");
+    }
+
+    final List<LockItem> lockItems = new ArrayList<>(e164s.size());
+
+    try {
+      for (final String e164 : e164s) {
+        lockItems.add(lockClient.acquireLock(AcquireLockOptions.builder(e164)
+            .withAcquireReleasedLocksConsistently(true)
+            .build()));
+      }
+
+      task.run();
+    } finally {
+      lockItems.forEach(lockItem -> lockClient.releaseLock(ReleaseLockOptions.builder(lockItem)
+          .withBestEffort(true)
+          .build()));
+    }
+  }
+}

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

@@ -87,7 +87,8 @@ public class AccountsManager {
   private final Accounts accounts;
   private final PhoneNumberIdentifiers phoneNumberIdentifiers;
   private final FaultTolerantRedisCluster cacheCluster;
-  private final DeletedAccountsManager deletedAccountsManager;
+  private final AccountLockManager accountLockManager;
+  private final DeletedAccounts deletedAccounts;
   private final Keys keys;
   private final MessagesManager messagesManager;
   private final ProfilesManager profilesManager;
@@ -130,7 +131,8 @@ public class AccountsManager {
   public AccountsManager(final Accounts accounts,
       final PhoneNumberIdentifiers phoneNumberIdentifiers,
       final FaultTolerantRedisCluster cacheCluster,
-      final DeletedAccountsManager deletedAccountsManager,
+      final AccountLockManager accountLockManager,
+      final DeletedAccounts deletedAccounts,
       final Keys keys,
       final MessagesManager messagesManager,
       final ProfilesManager profilesManager,
@@ -145,7 +147,8 @@ public class AccountsManager {
     this.accounts = accounts;
     this.phoneNumberIdentifiers = phoneNumberIdentifiers;
     this.cacheCluster = cacheCluster;
-    this.deletedAccountsManager = deletedAccountsManager;
+    this.accountLockManager = accountLockManager;
+    this.deletedAccounts = deletedAccounts;
     this.keys = keys;
     this.messagesManager = messagesManager;
     this.profilesManager = profilesManager;
@@ -168,7 +171,7 @@ public class AccountsManager {
     try (Timer.Context ignored = createTimer.time()) {
       final Account account = new Account();
 
-      deletedAccountsManager.lockAndTake(number, maybeRecentlyDeletedUuid -> {
+      accountLockManager.withLock(List.of(number), () -> {
         Device device = new Device();
         device.setId(Device.MASTER_ID);
         device.setAuthTokenHash(SaltedTokenHash.generateFor(password));
@@ -182,7 +185,10 @@ public class AccountsManager {
         device.setUserAgent(signalAgent);
 
         account.setNumber(number, phoneNumberIdentifiers.getPhoneNumberIdentifier(number));
-        account.setUuid(maybeRecentlyDeletedUuid.orElseGet(UUID::randomUUID));
+
+        // Reuse the ACI from any recently-deleted account with this number to cover cases where somebody is
+        // re-registering.
+        account.setUuid(deletedAccounts.findUuid(number).orElseGet(UUID::randomUUID));
         account.addDevice(device);
         account.setRegistrationLockFromAttributes(accountAttributes);
         account.setUnidentifiedAccessKey(accountAttributes.getUnidentifiedAccessKey());
@@ -236,6 +242,10 @@ public class AccountsManager {
 
         accountAttributes.recoveryPassword().ifPresent(registrationRecoveryPassword ->
             registrationRecoveryPasswordsManager.storeForCurrentNumber(account.getNumber(), registrationRecoveryPassword));
+
+        // Clear any "recently deleted account" record for this number since, if it existed, we've used its old ACI for
+        // the newly-created account.
+        deletedAccounts.remove(number);
       });
 
       return account;
@@ -243,7 +253,7 @@ public class AccountsManager {
   }
 
   public Account changeNumber(final Account account,
-                              final String number,
+                              final String targetNumber,
                               @Nullable final byte[] pniIdentityKey,
                               @Nullable final Map<Long, SignedPreKey> pniSignedPreKeys,
                               @Nullable final Map<Long, SignedPreKey> pniPqLastResortPreKeys,
@@ -252,7 +262,7 @@ public class AccountsManager {
     final String originalNumber = account.getNumber();
     final UUID originalPhoneNumberIdentifier = account.getPhoneNumberIdentifier();
 
-    if (originalNumber.equals(number)) {
+    if (originalNumber.equals(targetNumber)) {
       if (pniIdentityKey != null) {
         throw new IllegalArgumentException("change number must supply a changed phone number; otherwise use updatePniKeys");
       }
@@ -263,28 +273,42 @@ public class AccountsManager {
 
     final AtomicReference<Account> updatedAccount = new AtomicReference<>();
 
-    deletedAccountsManager.lockAndPut(account.getNumber(), number, (originalAci, deletedAci) -> {
+    accountLockManager.withLock(List.of(account.getNumber(), targetNumber), () -> {
       redisDelete(account);
 
-      final Optional<Account> maybeExistingAccount = getByE164(number);
+      // There are three possible states for accounts associated with the target phone number:
-      final Optional<UUID> displacedUuid;
+      //
+      // 1. An account exists with the target number; the caller has proved ownership of the number, so delete the
+      //    account with the target number. This will leave a "deleted account" record for the deleted account mapping
+      //    the UUID of the deleted account to the target phone number. We'll then overwrite that so it points to the
+      //    original number to facilitate switching back and forth between numbers.
+      // 2. No account with the target number exists, but one has recently been deleted. In that case, add a "deleted
+      //    account" record that maps the ACI of the recently-deleted account to the now-abandoned original phone number
+      //    of the account changing its number (which facilitates ACI consistency in cases that a party is switching
+      //    back and forth between numbers).
+      // 3. No account with the target number exists at all, in which case no additional action is needed.
+      final Optional<UUID> recentlyDeletedAci = deletedAccounts.findUuid(targetNumber);
+      final Optional<Account> maybeExistingAccount = getByE164(targetNumber);
+      final Optional<UUID> maybeDisplacedUuid;
 
       if (maybeExistingAccount.isPresent()) {
         delete(maybeExistingAccount.get());
-        displacedUuid = maybeExistingAccount.map(Account::getUuid);
+        maybeDisplacedUuid = maybeExistingAccount.map(Account::getUuid);
       } else {
-        displacedUuid = deletedAci;
+        maybeDisplacedUuid = recentlyDeletedAci;
       }
 
+      maybeDisplacedUuid.ifPresent(displacedUuid -> deletedAccounts.put(displacedUuid, originalNumber));
+
       final UUID uuid = account.getUuid();
-      final UUID phoneNumberIdentifier = phoneNumberIdentifiers.getPhoneNumberIdentifier(number);
+      final UUID phoneNumberIdentifier = phoneNumberIdentifiers.getPhoneNumberIdentifier(targetNumber);
 
       final Account numberChangedAccount;
 
       numberChangedAccount = updateWithRetries(
           account,
           a -> { setPniKeys(account, pniIdentityKey, pniSignedPreKeys, pniRegistrationIds); return true; },
-          a -> accounts.changeNumber(a, number, phoneNumberIdentifier),
+          a -> accounts.changeNumber(a, targetNumber, phoneNumberIdentifier),
           () -> accounts.getByAccountIdentifier(uuid).orElseThrow(),
           AccountChangeValidator.NUMBER_CHANGE_VALIDATOR);
 
@@ -301,8 +325,6 @@ public class AccountsManager {
                     Function.identity(),
                     pniPqLastResortPreKeys::get)));
       }
-      
-      return displacedUuid;
     });
 
     return updatedAccount.get();
@@ -675,10 +697,13 @@ public class AccountsManager {
 
   public void delete(final Account account, final DeletionReason deletionReason) throws InterruptedException {
     try (final Timer.Context ignored = deleteTimer.time()) {
-      deletedAccountsManager.lockAndPut(account.getNumber(), () -> {
+      accountLockManager.withLock(List.of(account.getNumber()), () -> {
+        final UUID accountIdentifier = account.getUuid();
+        final String e164 = account.getNumber();
+
         delete(account);
 
-        return account.getUuid();
+        deletedAccounts.put(accountIdentifier, e164);
       });
     } catch (final RuntimeException | InterruptedException e) {
       logger.warn("Failed to delete account", e);

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

@@ -18,7 +18,7 @@ import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.QueryRequest;
 import software.amazon.awssdk.services.dynamodb.model.QueryResponse;
 
-public class DeletedAccounts extends AbstractDynamoDbStore {
+public class DeletedAccounts {
 
   // e164, primary key
   static final String KEY_ACCOUNT_E164 = "P";
@@ -29,19 +29,16 @@ public class DeletedAccounts extends AbstractDynamoDbStore {
 
   static final Duration TIME_TO_LIVE = Duration.ofDays(30);
 
-  // Note that this limit is imposed by DynamoDB itself; going above 100 will result in errors
+  private final DynamoDbClient dynamoDbClient;
-  static final int GET_BATCH_SIZE = 100;
-
   private final String tableName;
 
-  public DeletedAccounts(final DynamoDbClient dynamoDb, final String tableName) {
+  public DeletedAccounts(final DynamoDbClient dynamoDbClient, final String tableName) {
-
+    this.dynamoDbClient = dynamoDbClient;
-    super(dynamoDb);
     this.tableName = tableName;
   }
 
-  void put(UUID uuid, String e164) {
+  public void put(UUID uuid, String e164) {
-    db().putItem(PutItemRequest.builder()
+    dynamoDbClient.putItem(PutItemRequest.builder()
         .tableName(tableName)
         .item(Map.of(
             KEY_ACCOUNT_E164, AttributeValues.fromString(e164),
@@ -50,8 +47,8 @@ public class DeletedAccounts extends AbstractDynamoDbStore {
         .build());
   }
 
-  Optional<UUID> findUuid(final String e164) {
+  public Optional<UUID> findUuid(final String e164) {
-    final GetItemResponse response = db().getItem(GetItemRequest.builder()
+    final GetItemResponse response = dynamoDbClient.getItem(GetItemRequest.builder()
         .tableName(tableName)
         .consistentRead(true)
         .key(Map.of(KEY_ACCOUNT_E164, AttributeValues.fromString(e164)))
@@ -60,8 +57,8 @@ public class DeletedAccounts extends AbstractDynamoDbStore {
     return Optional.ofNullable(AttributeValues.getUUID(response.item(), ATTR_ACCOUNT_UUID, null));
   }
 
-  Optional<String> findE164(final UUID uuid) {
+  public Optional<String> findE164(final UUID uuid) {
-    final QueryResponse response = db().query(QueryRequest.builder()
+    final QueryResponse response = dynamoDbClient.query(QueryRequest.builder()
         .tableName(tableName)
         .indexName(UUID_TO_E164_INDEX_NAME)
         .keyConditionExpression("#uuid = :uuid")
@@ -75,15 +72,14 @@ public class DeletedAccounts extends AbstractDynamoDbStore {
     }
 
     if (response.count() > 1) {
-      throw new RuntimeException(
+      throw new RuntimeException("Impossible result: more than one phone number returned for UUID: " + uuid);
-          "Impossible result: more than one phone number returned for UUID: " + uuid);
     }
 
     return Optional.ofNullable(response.items().get(0).get(KEY_ACCOUNT_E164).s());
   }
 
-  void remove(final String e164) {
+  public void remove(final String e164) {
-    db().deleteItem(DeleteItemRequest.builder()
+    dynamoDbClient.deleteItem(DeleteItemRequest.builder()
         .tableName(tableName)
         .key(Map.of(KEY_ACCOUNT_E164, AttributeValues.fromString(e164)))
         .build());

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

@@ -1,145 +0,0 @@
-/*
- * Copyright 2013-2021 Signal Messenger, LLC
- * SPDX-License-Identifier: AGPL-3.0-only
- */
-
-package org.whispersystems.textsecuregcm.storage;
-
-import com.amazonaws.services.dynamodbv2.AcquireLockOptions;
-import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
-import com.amazonaws.services.dynamodbv2.AmazonDynamoDBLockClient;
-import com.amazonaws.services.dynamodbv2.AmazonDynamoDBLockClientOptions;
-import com.amazonaws.services.dynamodbv2.LockItem;
-import com.amazonaws.services.dynamodbv2.ReleaseLockOptions;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Optional;
-import java.util.UUID;
-import java.util.concurrent.TimeUnit;
-import java.util.function.BiFunction;
-import java.util.function.Consumer;
-import java.util.function.Supplier;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-public class DeletedAccountsManager {
-
-  private final DeletedAccounts deletedAccounts;
-
-  private final AmazonDynamoDBLockClient lockClient;
-
-  private static final Logger log = LoggerFactory.getLogger(DeletedAccountsManager.class);
-
-  public DeletedAccountsManager(final DeletedAccounts deletedAccounts, final AmazonDynamoDB lockDynamoDb, final String lockTableName) {
-    this.deletedAccounts = deletedAccounts;
-
-    lockClient = new AmazonDynamoDBLockClient(
-        AmazonDynamoDBLockClientOptions.builder(lockDynamoDb, lockTableName)
-            .withPartitionKeyName(DeletedAccounts.KEY_ACCOUNT_E164)
-            .withLeaseDuration(15L)
-            .withHeartbeatPeriod(2L)
-            .withTimeUnit(TimeUnit.SECONDS)
-            .withCreateHeartbeatBackgroundThread(true)
-            .build());
-  }
-
-  /**
-   * Acquires a pessimistic lock for the given phone number and performs the given action, passing the UUID of the
-   * recently-deleted account (if any) that previously held the given number.
-   *
-   * @param e164 the phone number to lock and with which to perform an action
-   * @param consumer the action to take; accepts the UUID of the account that previously held the given e164, if any,
-   * as an argument
-   *
-   * @throws InterruptedException if interrupted while waiting to acquire a lock on the given phone number
-   */
-  public void lockAndTake(final String e164, final Consumer<Optional<UUID>> consumer) throws InterruptedException {
-    withLock(List.of(e164), acis -> {
-      try {
-        consumer.accept(acis.get(0));
-        deletedAccounts.remove(e164);
-      } catch (final Exception e) {
-        log.warn("Consumer threw an exception while holding lock on a deleted account record", e);
-        throw new RuntimeException(e);
-      }
-    });
-  }
-
-  /**
-   * Acquires a pessimistic lock for the given phone number and performs an action that deletes an account, returning
-   * the UUID of the deleted account. The UUID of the deleted account will be stored in the list of recently-deleted
-   * e164-to-UUID mappings.
-   *
-   * @param e164 the phone number to lock and with which to perform an action
-   * @param supplier the deletion action to take on the account associated with the given number; must return the UUID
-   * of the deleted account
-   *
-   * @throws InterruptedException if interrupted while waiting to acquire a lock on the given phone number
-   */
-  public void lockAndPut(final String e164, final Supplier<UUID> supplier) throws InterruptedException {
-    withLock(List.of(e164), ignored -> {
-      try {
-        deletedAccounts.put(supplier.get(), e164);
-      } catch (final Exception e) {
-        log.warn("Supplier threw an exception while holding lock on a deleted account record", e);
-        throw new RuntimeException(e);
-      }
-    });
-  }
-
-  /**
-   * Acquires a pessimistic lock for the given phone numbers and performs an action that may or may not delete an
-   * account associated with the target number. The UUID of the deleted account (if any) will be stored in the list of
-   * recently-deleted e164-to-UUID mappings.
-   *
-   * @param original the phone number of an existing account to lock and with which to perform an action
-   * @param target   the phone number of an account that may or may not exist with which to perform an action
-   * @param function the action to take on the given phone numbers and ACIs, if known; the action may delete the account
-   *                 identified by the target number, in which case it must return the ACI of that account
-   * @throws InterruptedException if interrupted while waiting to acquire a lock on the given phone numbers
-   */
-  public void lockAndPut(final String original, final String target,
-      final BiFunction<Optional<UUID>, Optional<UUID>, Optional<UUID>> function)
-      throws InterruptedException {
-
-    withLock(List.of(original, target), acis -> {
-      try {
-        function.apply(acis.get(0), acis.get(1)).ifPresent(aci -> deletedAccounts.put(aci, original));
-      } catch (final Exception e) {
-        log.warn("Supplier threw an exception while holding lock on a deleted account record", e);
-        throw new RuntimeException(e);
-      }
-    });
-  }
-
-  private void withLock(final List<String> e164s, final Consumer<List<Optional<UUID>>> task)
-      throws InterruptedException {
-    final List<LockItem> lockItems = new ArrayList<>(e164s.size());
-
-    try {
-      final List<Optional<UUID>> previouslyDeletedUuids = new ArrayList<>(e164s.size());
-      for (final String e164 : e164s) {
-        lockItems.add(lockClient.acquireLock(AcquireLockOptions.builder(e164)
-            .withAcquireReleasedLocksConsistently(true)
-            .build()));
-        previouslyDeletedUuids.add(deletedAccounts.findUuid(e164));
-      }
-
-      task.accept(previouslyDeletedUuids);
-    } finally {
-      for (final LockItem lockItem : lockItems) {
-        lockClient.releaseLock(ReleaseLockOptions.builder(lockItem)
-            .withBestEffort(true)
-            .build());
-      }
-    }
-  }
-
-  public Optional<UUID> findDeletedAccountAci(final String e164) {
-    return deletedAccounts.findUuid(e164);
-  }
-
-  public Optional<String> findDeletedAccountE164(final UUID uuid) {
-    return deletedAccounts.findE164(uuid);
-  }
-}

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

@@ -37,10 +37,10 @@ import org.whispersystems.textsecuregcm.securebackup.SecureBackupClient;
 import org.whispersystems.textsecuregcm.securestorage.SecureStorageClient;
 import org.whispersystems.textsecuregcm.securevaluerecovery.SecureValueRecovery2Client;
 import org.whispersystems.textsecuregcm.storage.Account;
+import org.whispersystems.textsecuregcm.storage.AccountLockManager;
 import org.whispersystems.textsecuregcm.storage.Accounts;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.DeletedAccounts;
-import org.whispersystems.textsecuregcm.storage.DeletedAccountsManager;
 import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
 import org.whispersystems.textsecuregcm.storage.Keys;
 import org.whispersystems.textsecuregcm.storage.MessagesCache;
@@ -205,12 +205,10 @@ public class AssignUsernameCommand extends EnvironmentCommand<WhisperServerConfi
         configuration.getReportMessageConfiguration().getCounterTtl());
     MessagesManager messagesManager = new MessagesManager(messagesDynamoDb, messagesCache,
         reportMessageManager, messageDeletionExecutor);
-    DeletedAccountsManager deletedAccountsManager = new DeletedAccountsManager(deletedAccounts,
+    AccountLockManager accountLockManager = new AccountLockManager(deletedAccountsLockDynamoDbClient, configuration.getDynamoDbTables().getDeletedAccountsLock().getTableName());
-        deletedAccountsLockDynamoDbClient,
-        configuration.getDynamoDbTables().getDeletedAccountsLock().getTableName());
     StoredVerificationCodeManager pendingAccountsManager = new StoredVerificationCodeManager(pendingAccounts);
     AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
-        deletedAccountsManager, keys, messagesManager, profilesManager,
+        accountLockManager, deletedAccounts, keys, messagesManager, profilesManager,
         pendingAccountsManager, secureStorageClient, secureBackupClient, secureValueRecovery2Client, clientPresenceManager,
         experimentEnrollmentManager, registrationRecoveryPasswordsManager, Clock.systemUTC());
 

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

@@ -31,10 +31,10 @@ import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
 import org.whispersystems.textsecuregcm.securebackup.SecureBackupClient;
 import org.whispersystems.textsecuregcm.securestorage.SecureStorageClient;
 import org.whispersystems.textsecuregcm.securevaluerecovery.SecureValueRecovery2Client;
+import org.whispersystems.textsecuregcm.storage.AccountLockManager;
 import org.whispersystems.textsecuregcm.storage.Accounts;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.DeletedAccounts;
-import org.whispersystems.textsecuregcm.storage.DeletedAccountsManager;
 import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
 import org.whispersystems.textsecuregcm.storage.Keys;
 import org.whispersystems.textsecuregcm.storage.MessagesCache;
@@ -64,7 +64,6 @@ record CommandDependencies(
     ReportMessageManager reportMessageManager,
     MessagesCache messagesCache,
     MessagesManager messagesManager,
-    DeletedAccountsManager deletedAccountsManager,
     StoredVerificationCodeManager pendingAccountsManager,
     ClientPresenceManager clientPresenceManager,
     Keys keys,
@@ -190,12 +189,10 @@ record CommandDependencies(
         configuration.getReportMessageConfiguration().getCounterTtl());
     MessagesManager messagesManager = new MessagesManager(messagesDynamoDb, messagesCache,
         reportMessageManager, messageDeletionExecutor);
-    DeletedAccountsManager deletedAccountsManager = new DeletedAccountsManager(deletedAccounts,
+    AccountLockManager accountLockManager = new AccountLockManager(deletedAccountsLockDynamoDbClient, configuration.getDynamoDbTables().getDeletedAccountsLock().getTableName());
-        deletedAccountsLockDynamoDbClient,
-        configuration.getDynamoDbTables().getDeletedAccountsLock().getTableName());
     StoredVerificationCodeManager pendingAccountsManager = new StoredVerificationCodeManager(pendingAccounts);
     AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
-        deletedAccountsManager, keys, messagesManager, profilesManager,
+        accountLockManager, deletedAccounts, keys, messagesManager, profilesManager,
         pendingAccountsManager, secureStorageClient, secureBackupClient, secureValueRecovery2Client, clientPresenceManager,
         experimentEnrollmentManager, registrationRecoveryPasswordsManager, clock);
 
@@ -205,7 +202,6 @@ record CommandDependencies(
         reportMessageManager,
         messagesCache,
         messagesManager,
-        deletedAccountsManager,
         pendingAccountsManager,
         clientPresenceManager,
         keys,

service/src/test/java/org/whispersystems/textsecuregcm/controllers/MessageControllerTest.java

@@ -97,7 +97,7 @@ import org.whispersystems.textsecuregcm.push.ReceiptSender;
 import org.whispersystems.textsecuregcm.spam.ReportSpamTokenProvider;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
-import org.whispersystems.textsecuregcm.storage.DeletedAccountsManager;
+import org.whispersystems.textsecuregcm.storage.DeletedAccounts;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.MessagesManager;
 import org.whispersystems.textsecuregcm.storage.ReportMessageManager;
@@ -141,7 +141,7 @@ class MessageControllerTest {
   private static final MessageSender messageSender = mock(MessageSender.class);
   private static final ReceiptSender receiptSender = mock(ReceiptSender.class);
   private static final AccountsManager accountsManager = mock(AccountsManager.class);
-  private static final DeletedAccountsManager deletedAccountsManager = mock(DeletedAccountsManager.class);
+  private static final DeletedAccounts deletedAccounts = mock(DeletedAccounts.class);
   private static final MessagesManager messagesManager = mock(MessagesManager.class);
   private static final RateLimiters rateLimiters = mock(RateLimiters.class);
   private static final RateLimiter rateLimiter = mock(RateLimiter.class);
@@ -159,7 +159,7 @@ class MessageControllerTest {
       .addProvider(MultiRecipientMessageProvider.class)
       .setTestContainerFactory(new GrizzlyWebTestContainerFactory())
       .addResource(
-          new MessageController(rateLimiters, messageSender, receiptSender, accountsManager, deletedAccountsManager,
+          new MessageController(rateLimiters, messageSender, receiptSender, accountsManager, deletedAccounts,
               messagesManager, pushNotificationManager, reportMessageManager, multiRecipientMessageExecutor,
               messageDeliveryScheduler, ReportSpamTokenProvider.noop()))
       .build();
@@ -214,7 +214,7 @@ class MessageControllerTest {
         messageSender,
         receiptSender,
         accountsManager,
-        deletedAccountsManager,
+        deletedAccounts,
         messagesManager,
         rateLimiters,
         rateLimiter,
@@ -634,7 +634,7 @@ class MessageControllerTest {
     when(account.getPhoneNumberIdentifier()).thenReturn(senderPni);
 
     when(accountsManager.getByE164(senderNumber)).thenReturn(Optional.of(account));
-    when(deletedAccountsManager.findDeletedAccountAci(senderNumber)).thenReturn(Optional.of(senderAci));
+    when(deletedAccounts.findUuid(senderNumber)).thenReturn(Optional.of(senderAci));
     when(accountsManager.getPhoneNumberIdentifier(senderNumber)).thenReturn(senderPni);
 
     Response response =
@@ -649,7 +649,7 @@ class MessageControllerTest {
 
     verify(reportMessageManager).report(Optional.of(senderNumber), Optional.of(senderAci), Optional.of(senderPni),
         messageGuid, AuthHelper.VALID_UUID, Optional.empty(), userAgent);
-    verify(deletedAccountsManager, never()).findDeletedAccountE164(any(UUID.class));
+    verify(deletedAccounts, never()).findE164(any(UUID.class));
     verify(accountsManager, never()).getPhoneNumberIdentifier(anyString());
 
     when(accountsManager.getByE164(senderNumber)).thenReturn(Optional.empty());
@@ -684,7 +684,7 @@ class MessageControllerTest {
     when(account.getPhoneNumberIdentifier()).thenReturn(senderPni);
 
     when(accountsManager.getByAccountIdentifier(senderAci)).thenReturn(Optional.of(account));
-    when(deletedAccountsManager.findDeletedAccountE164(senderAci)).thenReturn(Optional.of(senderNumber));
+    when(deletedAccounts.findE164(senderAci)).thenReturn(Optional.of(senderNumber));
     when(accountsManager.getPhoneNumberIdentifier(senderNumber)).thenReturn(senderPni);
 
     Response response =
@@ -699,7 +699,7 @@ class MessageControllerTest {
 
     verify(reportMessageManager).report(Optional.of(senderNumber), Optional.of(senderAci), Optional.of(senderPni),
         messageGuid, AuthHelper.VALID_UUID, Optional.empty(), userAgent);
-    verify(deletedAccountsManager, never()).findDeletedAccountE164(any(UUID.class));
+    verify(deletedAccounts, never()).findE164(any(UUID.class));
     verify(accountsManager, never()).getPhoneNumberIdentifier(anyString());
 
     when(accountsManager.getByAccountIdentifier(senderAci)).thenReturn(Optional.empty());
@@ -734,7 +734,7 @@ class MessageControllerTest {
     when(account.getPhoneNumberIdentifier()).thenReturn(senderPni);
 
     when(accountsManager.getByAccountIdentifier(senderAci)).thenReturn(Optional.of(account));
-    when(deletedAccountsManager.findDeletedAccountE164(senderAci)).thenReturn(Optional.of(senderNumber));
+    when(deletedAccounts.findE164(senderAci)).thenReturn(Optional.of(senderNumber));
     when(accountsManager.getPhoneNumberIdentifier(senderNumber)).thenReturn(senderPni);
 
     Entity<SpamReport> entity = Entity.entity(new SpamReport(new byte[3]), "application/json");
@@ -753,7 +753,7 @@ class MessageControllerTest {
         eq(AuthHelper.VALID_UUID),
         argThat(maybeBytes -> maybeBytes.map(bytes -> Arrays.equals(bytes, new byte[3])).orElse(false)),
         any());
-    verify(deletedAccountsManager, never()).findDeletedAccountE164(any(UUID.class));
+    verify(deletedAccounts, never()).findE164(any(UUID.class));
     verify(accountsManager, never()).getPhoneNumberIdentifier(anyString());
     when(accountsManager.getByAccountIdentifier(senderAci)).thenReturn(Optional.empty());
 
@@ -792,7 +792,7 @@ class MessageControllerTest {
     when(account.getPhoneNumberIdentifier()).thenReturn(senderPni);
 
     when(accountsManager.getByAccountIdentifier(senderAci)).thenReturn(Optional.of(account));
-    when(deletedAccountsManager.findDeletedAccountE164(senderAci)).thenReturn(Optional.of(senderNumber));
+    when(deletedAccounts.findE164(senderAci)).thenReturn(Optional.of(senderNumber));
     when(accountsManager.getPhoneNumberIdentifier(senderNumber)).thenReturn(senderPni);
 
     Response response =

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

@@ -0,0 +1,62 @@
+package org.whispersystems.textsecuregcm.storage;
+
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import com.amazonaws.services.dynamodbv2.AmazonDynamoDBLockClient;
+import com.amazonaws.services.dynamodbv2.ReleaseLockOptions;
+import com.google.i18n.phonenumbers.PhoneNumberUtil;
+import java.util.Collections;
+import java.util.List;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+class AccountLockManagerTest {
+
+  private AmazonDynamoDBLockClient lockClient;
+
+  private AccountLockManager accountLockManager;
+
+  private static final String FIRST_NUMBER = PhoneNumberUtil.getInstance().format(
+      PhoneNumberUtil.getInstance().getExampleNumber("US"), PhoneNumberUtil.PhoneNumberFormat.E164);
+
+  private static final String SECOND_NUMBER = PhoneNumberUtil.getInstance().format(
+      PhoneNumberUtil.getInstance().getExampleNumber("JP"), PhoneNumberUtil.PhoneNumberFormat.E164);
+
+  @BeforeEach
+  void setUp() {
+    lockClient = mock(AmazonDynamoDBLockClient.class);
+
+    accountLockManager = new AccountLockManager(lockClient);
+  }
+
+  @Test
+  void withLock() throws InterruptedException {
+    accountLockManager.withLock(List.of(FIRST_NUMBER, SECOND_NUMBER), () -> {});
+
+    verify(lockClient, times(2)).acquireLock(any());
+    verify(lockClient, times(2)).releaseLock(any(ReleaseLockOptions.class));
+  }
+
+  @Test
+  void withLockTaskThrowsException() throws InterruptedException {
+    assertThrows(RuntimeException.class, () -> accountLockManager.withLock(List.of(FIRST_NUMBER, SECOND_NUMBER), () -> {
+      throw new RuntimeException();
+    }));
+
+    verify(lockClient, times(2)).acquireLock(any());
+    verify(lockClient, times(2)).releaseLock(any(ReleaseLockOptions.class));
+  }
+
+  @Test
+  void withLockEmptyList() {
+    final Runnable task = mock(Runnable.class);
+
+    assertThrows(IllegalArgumentException.class, () -> accountLockManager.withLock(Collections.emptyList(), () -> {}));
+    verify(task, never()).run();
+  }
+}

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

@@ -80,8 +80,7 @@ class AccountsManagerChangeNumberIntegrationTest {
       deletedAccounts = new DeletedAccounts(DYNAMO_DB_EXTENSION.getDynamoDbClient(),
           Tables.DELETED_ACCOUNTS.tableName());
 
-      final DeletedAccountsManager deletedAccountsManager = new DeletedAccountsManager(deletedAccounts,
+      final AccountLockManager accountLockManager = new AccountLockManager(DYNAMO_DB_EXTENSION.getLegacyDynamoClient(),
-          DYNAMO_DB_EXTENSION.getLegacyDynamoClient(),
           Tables.DELETED_ACCOUNTS_LOCK.tableName());
 
       final SecureStorageClient secureStorageClient = mock(SecureStorageClient.class);
@@ -102,8 +101,7 @@ class AccountsManagerChangeNumberIntegrationTest {
           accounts,
           phoneNumberIdentifiers,
           CACHE_CLUSTER_EXTENSION.getRedisCluster(),
-          deletedAccountsManager,
+          accountLockManager, deletedAccounts, mock(Keys.class),
-          mock(Keys.class),
           mock(MessagesManager.class),
           mock(ProfilesManager.class),
           mock(StoredVerificationCodeManager.class),

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

@@ -92,13 +92,17 @@ class AccountsManagerConcurrentModificationIntegrationTest {
       //noinspection unchecked
       commands = mock(RedisAdvancedClusterCommands.class);
 
-      final DeletedAccountsManager deletedAccountsManager = mock(DeletedAccountsManager.class);
+      final AccountLockManager accountLockManager = mock(AccountLockManager.class);
 
       doAnswer(invocation -> {
-        //noinspection unchecked
+        final Runnable task = invocation.getArgument(1);
-        invocation.getArgument(1, Consumer.class).accept(Optional.empty());
+        task.run();
+
         return null;
-      }).when(deletedAccountsManager).lockAndTake(anyString(), any());
+      }).when(accountLockManager).withLock(any(), any());
+
+      final DeletedAccounts deletedAccounts = mock(DeletedAccounts.class);
+      when(deletedAccounts.findUuid(any())).thenReturn(Optional.empty());
 
       final PhoneNumberIdentifiers phoneNumberIdentifiers = mock(PhoneNumberIdentifiers.class);
       when(phoneNumberIdentifiers.getPhoneNumberIdentifier(anyString()))
@@ -108,8 +112,7 @@ class AccountsManagerConcurrentModificationIntegrationTest {
           accounts,
           phoneNumberIdentifiers,
           RedisClusterHelper.builder().stringCommands(commands).build(),
-          deletedAccountsManager,
+          accountLockManager, deletedAccounts, mock(Keys.class),
-          mock(Keys.class),
           mock(MessagesManager.class),
           mock(ProfilesManager.class),
           mock(StoredVerificationCodeManager.class),

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

@@ -36,7 +36,6 @@ import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
-import java.util.function.BiFunction;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -71,7 +70,7 @@ class AccountsManagerTest {
   private static final byte[] USERNAME_HASH_2 = Base64.getUrlDecoder().decode(BASE_64_URL_USERNAME_HASH_2);
 
   private Accounts accounts;
-  private DeletedAccountsManager deletedAccountsManager;
+  private DeletedAccounts deletedAccounts;
   private Keys keys;
   private MessagesManager messagesManager;
   private ProfilesManager profilesManager;
@@ -94,7 +93,7 @@ class AccountsManagerTest {
   @BeforeEach
   void setup() throws InterruptedException {
     accounts = mock(Accounts.class);
-    deletedAccountsManager = mock(DeletedAccountsManager.class);
+    deletedAccounts = mock(DeletedAccounts.class);
     keys = mock(Keys.class);
     messagesManager = mock(MessagesManager.class);
     profilesManager = mock(ProfilesManager.class);
@@ -113,11 +112,7 @@ class AccountsManagerTest {
       return null;
     }).when(accounts).changeNumber(any(), anyString(), any());
 
-    doAnswer(invocation -> {
+    when(deletedAccounts.findUuid(anyString())).thenReturn(Optional.empty());
-      //noinspection unchecked
-      invocation.getArgument(1, Consumer.class).accept(Optional.empty());
-      return null;
-    }).when(deletedAccountsManager).lockAndTake(anyString(), any());
 
     final SecureStorageClient storageClient = mock(SecureStorageClient.class);
     when(storageClient.deleteStoredData(any())).thenReturn(CompletableFuture.completedFuture(null));
@@ -147,11 +142,21 @@ class AccountsManagerTest {
     when(enrollmentManager.isEnrolled(any(UUID.class), eq(AccountsManager.USERNAME_EXPERIMENT_NAME))).thenReturn(true);
     when(accounts.usernameHashAvailable(any())).thenReturn(true);
 
+    final AccountLockManager accountLockManager = mock(AccountLockManager.class);
+
+    doAnswer(invocation -> {
+      final Runnable task = invocation.getArgument(1);
+      task.run();
+
+      return null;
+    }).when(accountLockManager).withLock(any(), any());
+
     accountsManager = new AccountsManager(
         accounts,
         phoneNumberIdentifiers,
         RedisClusterHelper.builder().stringCommands(commands).build(),
-        deletedAccountsManager,
+        accountLockManager,
+        deletedAccounts,
         keys,
         messagesManager,
         profilesManager,
@@ -571,12 +576,7 @@ class AccountsManagerTest {
   void testCreateAccountRecentlyDeleted() throws InterruptedException {
     final UUID recentlyDeletedUuid = UUID.randomUUID();
 
-    doAnswer(invocation -> {
+    when(deletedAccounts.findUuid(anyString())).thenReturn(Optional.of(recentlyDeletedUuid));
-      //noinspection unchecked
-      invocation.getArgument(1, Consumer.class).accept(Optional.of(recentlyDeletedUuid));
-      return null;
-    }).when(deletedAccountsManager).lockAndTake(anyString(), any());
-
     when(accounts.create(any())).thenReturn(true);
 
     final String e164 = "+18005550123";
@@ -634,9 +634,6 @@ class AccountsManagerTest {
 
   @Test
   void testChangePhoneNumber() throws InterruptedException, MismatchedDevicesException {
-    doAnswer(invocation -> invocation.getArgument(2, BiFunction.class).apply(Optional.empty(), Optional.empty()))
-        .when(deletedAccountsManager).lockAndPut(anyString(), anyString(), any());
-
     final String originalNumber = "+14152222222";
     final String targetNumber = "+14153333333";
     final UUID uuid = UUID.randomUUID();
@@ -661,7 +658,7 @@ class AccountsManagerTest {
     account = accountsManager.changeNumber(account, number, null, null, null, null);
 
     assertEquals(number, account.getNumber());
-    verify(deletedAccountsManager, never()).lockAndPut(anyString(), anyString(), any());
+    verify(deletedAccounts, never()).put(any(), any());
     verify(keys, never()).delete(any());
   }
 
@@ -676,15 +673,12 @@ class AccountsManagerTest {
         "AccountsManager should not allow use of changeNumber with new PNI keys but without changing number");
 
     verify(accounts, never()).update(any());
-    verifyNoInteractions(deletedAccountsManager);
+    verifyNoInteractions(deletedAccounts);
     verifyNoInteractions(keys);
   }
 
   @Test
   void testChangePhoneNumberExistingAccount() throws InterruptedException, MismatchedDevicesException {
-    doAnswer(invocation -> invocation.getArgument(2, BiFunction.class).apply(Optional.empty(), Optional.empty()))
-        .when(deletedAccountsManager).lockAndPut(anyString(), anyString(), any());
-
     final String originalNumber = "+14152222222";
     final String targetNumber = "+14153333333";
     final UUID existingAccountUuid = UUID.randomUUID();
@@ -712,9 +706,6 @@ class AccountsManagerTest {
 
   @Test
   void testChangePhoneNumberWithPqKeysExistingAccount() throws InterruptedException, MismatchedDevicesException {
-    doAnswer(invocation -> invocation.getArgument(2, BiFunction.class).apply(Optional.empty(), Optional.empty()))
-        .when(deletedAccountsManager).lockAndPut(anyString(), anyString(), any());
-
     final String originalNumber = "+14152222222";
     final String targetNumber = "+14153333333";
     final UUID existingAccountUuid = UUID.randomUUID();
@@ -799,7 +790,7 @@ class AccountsManagerTest {
         updatedAccount.getDevices().stream().collect(Collectors.toMap(Device::getId, d -> d.getPhoneNumberIdentityRegistrationId().getAsInt())));
 
     verify(accounts).update(any());
-    verifyNoInteractions(deletedAccountsManager);
+    verifyNoInteractions(deletedAccounts);
 
     verify(keys).delete(oldPni);
   }
@@ -846,7 +837,7 @@ class AccountsManagerTest {
         updatedAccount.getDevices().stream().collect(Collectors.toMap(Device::getId, d -> d.getPhoneNumberIdentityRegistrationId().getAsInt())));
 
     verify(accounts).update(any());
-    verifyNoInteractions(deletedAccountsManager);
+    verifyNoInteractions(deletedAccounts);
 
     verify(keys).delete(oldPni);
 

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

@@ -30,12 +30,10 @@ import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
-import java.util.function.Consumer;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.mockito.Mockito;
-import org.mockito.invocation.InvocationOnMock;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
 import org.whispersystems.textsecuregcm.entities.AccountAttributes;
 import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
@@ -94,13 +92,17 @@ class AccountsManagerUsernameIntegrationTest {
         Tables.USERNAMES.tableName(),
         SCAN_PAGE_SIZE));
 
-    final DeletedAccountsManager deletedAccountsManager = mock(DeletedAccountsManager.class);
+    final AccountLockManager accountLockManager = mock(AccountLockManager.class);
-    doAnswer((final InvocationOnMock invocationOnMock) -> {
+
-      @SuppressWarnings("unchecked")
+    doAnswer(invocation -> {
-      Consumer<Optional<UUID>> consumer = invocationOnMock.getArgument(1, Consumer.class);
+      final Runnable task = invocation.getArgument(1);
-      consumer.accept(Optional.empty());
+      task.run();
+
       return null;
-    }).when(deletedAccountsManager).lockAndTake(any(), any());
+    }).when(accountLockManager).withLock(any(), any());
+
+    final DeletedAccounts deletedAccounts = mock(DeletedAccounts.class);
+    when(deletedAccounts.findUuid(any())).thenReturn(Optional.empty());
 
     final PhoneNumberIdentifiers phoneNumberIdentifiers =
         new PhoneNumberIdentifiers(DYNAMO_DB_EXTENSION.getDynamoDbClient(), Tables.PNI.tableName());
@@ -112,7 +114,8 @@ class AccountsManagerUsernameIntegrationTest {
         accounts,
         phoneNumberIdentifiers,
         CACHE_CLUSTER_EXTENSION.getRedisCluster(),
-        deletedAccountsManager,
+        accountLockManager,
+        deletedAccounts,
         mock(Keys.class),
         mock(MessagesManager.class),
         mock(ProfilesManager.class),

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

@@ -1,61 +0,0 @@
-/*
- * Copyright 2013-2021 Signal Messenger, LLC
- * SPDX-License-Identifier: AGPL-3.0-only
- */
-
-package org.whispersystems.textsecuregcm.storage;
-
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertThrows;
-
-import java.util.Optional;
-import java.util.UUID;
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
-import org.junit.jupiter.api.extension.RegisterExtension;
-import org.whispersystems.textsecuregcm.storage.DynamoDbExtensionSchema.Tables;
-
-class DeletedAccountsManagerTest {
-
-  @RegisterExtension
-  static final DynamoDbExtension DYNAMO_DB_EXTENSION =
-      new DynamoDbExtension(Tables.DELETED_ACCOUNTS, Tables.DELETED_ACCOUNTS_LOCK);
-
-  private DeletedAccounts deletedAccounts;
-  private DeletedAccountsManager deletedAccountsManager;
-
-  @BeforeEach
-  void setUp() {
-    deletedAccounts = new DeletedAccounts(DYNAMO_DB_EXTENSION.getDynamoDbClient(),
-        Tables.DELETED_ACCOUNTS.tableName());
-
-    deletedAccountsManager = new DeletedAccountsManager(deletedAccounts,
-        DYNAMO_DB_EXTENSION.getLegacyDynamoClient(),
-        Tables.DELETED_ACCOUNTS_LOCK.tableName());
-  }
-
-  @Test
-  void testLockAndTake() throws InterruptedException {
-    final UUID uuid = UUID.randomUUID();
-    final String e164 = "+18005551234";
-
-    deletedAccounts.put(uuid, e164);
-    deletedAccountsManager.lockAndTake(e164, maybeUuid -> assertEquals(Optional.of(uuid), maybeUuid));
-    assertEquals(Optional.empty(), deletedAccounts.findUuid(e164));
-  }
-
-  @Test
-  void testLockAndTakeWithException() {
-    final UUID uuid = UUID.randomUUID();
-    final String e164 = "+18005551234";
-
-    deletedAccounts.put(uuid, e164);
-
-    assertThrows(RuntimeException.class, () -> deletedAccountsManager.lockAndTake(e164, maybeUuid -> {
-      assertEquals(Optional.of(uuid), maybeUuid);
-      throw new RuntimeException("OH NO");
-    }));
-
-    assertEquals(Optional.of(uuid), deletedAccounts.findUuid(e164));
-  }
-}