Refactor account locks/deleted account manager

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,
-          DYNAMO_DB_EXTENSION.getLegacyDynamoClient(),
+      final AccountLockManager accountLockManager = new AccountLockManager(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,
-          mock(Keys.class),
+          accountLockManager, deletedAccounts, 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
-        invocation.getArgument(1, Consumer.class).accept(Optional.empty());
+        final Runnable task = invocation.getArgument(1);
+        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,
-          mock(Keys.class),
+          accountLockManager, deletedAccounts, 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 -> {
-      //noinspection unchecked
-      invocation.getArgument(1, Consumer.class).accept(Optional.empty());
-      return null;
-    }).when(deletedAccountsManager).lockAndTake(anyString(), any());
+    when(deletedAccounts.findUuid(anyString())).thenReturn(Optional.empty());
 
     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 -> {
-      //noinspection unchecked
-      invocation.getArgument(1, Consumer.class).accept(Optional.of(recentlyDeletedUuid));
-      return null;
-    }).when(deletedAccountsManager).lockAndTake(anyString(), any());
-
+    when(deletedAccounts.findUuid(anyString())).thenReturn(Optional.of(recentlyDeletedUuid));
     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);
-    doAnswer((final InvocationOnMock invocationOnMock) -> {
-      @SuppressWarnings("unchecked")
-      Consumer<Optional<UUID>> consumer = invocationOnMock.getArgument(1, Consumer.class);
-      consumer.accept(Optional.empty());
+    final AccountLockManager accountLockManager = mock(AccountLockManager.class);
+
+    doAnswer(invocation -> {
+      final Runnable task = invocation.getArgument(1);
+      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));
-  }
-}