From a4b98f38a67ef39dbdb4d95338f87918bf5e35d3 Mon Sep 17 00:00:00 2001 From: Jon Chambers Date: Mon, 12 May 2025 12:36:51 -0400 Subject: [PATCH] Use a `Callable` for tasks performed within the scope of a pessimistic lock --- .../storage/AccountLockManager.java | 14 +- .../storage/AccountsManager.java | 353 ++++++++++-------- .../storage/AccountLockManagerTest.java | 8 +- ...ConcurrentModificationIntegrationTest.java | 9 +- .../storage/AccountsManagerTest.java | 9 +- ...ccountsManagerUsernameIntegrationTest.java | 11 +- 6 files changed, 221 insertions(+), 183 deletions(-) diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountLockManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountLockManager.java index 7347a62de..b09651c0c 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountLockManager.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountLockManager.java @@ -9,6 +9,7 @@ import com.google.common.annotations.VisibleForTesting; import java.util.ArrayList; import java.util.List; import java.util.UUID; +import java.util.concurrent.Callable; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; import java.util.concurrent.Executor; @@ -49,10 +50,15 @@ public class AccountLockManager { * @param task the task to execute once locks have been acquired * @param lockAcquisitionExecutor the executor on which to run blocking lock acquire/release tasks. this executor * should not use virtual threads. - * @throws InterruptedException if interrupted while acquiring a lock + * + * @return the value returned by the given {@code task} + * + * @throws Exception if an exception is thrown by the given {@code task} */ - public void withLock(final List phoneNumberIdentifiers, final Runnable task, - final Executor lockAcquisitionExecutor) { + public V withLock(final List phoneNumberIdentifiers, + final Callable task, + final Executor lockAcquisitionExecutor) throws Exception { + if (phoneNumberIdentifiers.isEmpty()) { throw new IllegalArgumentException("List of PNIs to lock must not be empty"); } @@ -75,7 +81,7 @@ public class AccountLockManager { } }, lockAcquisitionExecutor).join(); - task.run(); + return task.call(); } finally { CompletableFuture.runAsync(() -> { for (final LockItem lockItem : lockItems) { diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java index 848db12a2..e69e39efb 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java @@ -273,128 +273,148 @@ public class AccountsManager extends RedisPubSubAdapter implemen final DeviceSpec primaryDeviceSpec, @Nullable final String userAgent) throws InterruptedException { - final Account account = new Account(); final UUID pni = phoneNumberIdentifiers.getPhoneNumberIdentifier(number).join(); return createTimer.record(() -> { - accountLockManager.withLock(List.of(pni), () -> { - final Optional maybeRecentlyDeletedAccountIdentifier = - accounts.findRecentlyDeletedAccountIdentifier(pni); - - // Reuse the ACI from any recently-deleted account with this number to cover cases where somebody is - // re-registering. - account.setUuid(maybeRecentlyDeletedAccountIdentifier.orElseGet(UUID::randomUUID)); - account.setNumber(number, pni); - account.setIdentityKey(aciIdentityKey); - account.setPhoneNumberIdentityKey(pniIdentityKey); - 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); - - String accountCreationType = maybeRecentlyDeletedAccountIdentifier.isPresent() ? "recently-deleted" : "new"; - - final String pushTokenType; - - if (primaryDeviceSpec.apnRegistrationId().isPresent()) { - pushTokenType = "apns"; - } else if (primaryDeviceSpec.gcmRegistrationId().isPresent()) { - pushTokenType = "fcm"; - } else { - pushTokenType = "none"; + try { + return accountLockManager.withLock(List.of(pni), + () -> create(number, pni, accountAttributes, accountBadges, aciIdentityKey, pniIdentityKey, primaryDeviceSpec, userAgent), accountLockExecutor); + } catch (final Exception e) { + if (e instanceof RuntimeException runtimeException) { + throw runtimeException; } - String previousPushTokenType = null; - - try { - accounts.create(account, keysManager.buildWriteItemsForNewDevice(account.getIdentifier(IdentityType.ACI), - account.getIdentifier(IdentityType.PNI), - Device.PRIMARY_ID, - primaryDeviceSpec.aciSignedPreKey(), - primaryDeviceSpec.pniSignedPreKey(), - primaryDeviceSpec.aciPqLastResortPreKey(), - primaryDeviceSpec.pniPqLastResortPreKey())); - } catch (final AccountAlreadyExistsException e) { - accountCreationType = "re-registration"; - - if (StringUtils.isNotBlank(e.getExistingAccount().getPrimaryDevice().getApnId())) { - previousPushTokenType = "apns"; - } else if (StringUtils.isNotBlank(e.getExistingAccount().getPrimaryDevice().getGcmId())) { - previousPushTokenType = "fcm"; - } else { - previousPushTokenType = "none"; - } - - final UUID aci = e.getExistingAccount().getIdentifier(IdentityType.ACI); - account.setUuid(aci); - - final List additionalWriteItems = Stream.concat( - keysManager.buildWriteItemsForNewDevice(account.getIdentifier(IdentityType.ACI), - account.getIdentifier(IdentityType.PNI), - Device.PRIMARY_ID, - primaryDeviceSpec.aciSignedPreKey(), - primaryDeviceSpec.pniSignedPreKey(), - primaryDeviceSpec.aciPqLastResortPreKey(), - primaryDeviceSpec.pniPqLastResortPreKey()).stream(), - e.getExistingAccount().getDevices() - .stream() - .map(Device::getId) - // No need to clear the keys for the primary device since we'll just overwrite them in the same - // transaction anyhow - .filter(existingDeviceId -> existingDeviceId != Device.PRIMARY_ID) - .flatMap(existingDeviceId -> - keysManager.buildWriteItemsForRemovedDevice(aci, pni, existingDeviceId).stream())) - .toList(); - - CompletableFuture.allOf( - keysManager.deleteSingleUsePreKeys(aci), - keysManager.deleteSingleUsePreKeys(pni), - messagesManager.clear(aci), - profilesManager.deleteAll(aci)) - .thenCompose(ignored -> disconnectionRequestManager.requestDisconnection(aci)) - .thenCompose(ignored -> accounts.reclaimAccount(e.getExistingAccount(), account, additionalWriteItems)) - .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.deleteSingleUsePreKeys(aci), - keysManager.deleteSingleUsePreKeys(pni), - messagesManager.clear(aci), - profilesManager.deleteAll(aci)); - }) - .join(); - } - - redisSet(account); - - Tags tags = Tags.of(UserAgentTagUtil.getPlatformTag(userAgent), - Tag.of("type", accountCreationType), - Tag.of("hasPushToken", String.valueOf( - primaryDeviceSpec.apnRegistrationId().isPresent() || primaryDeviceSpec.gcmRegistrationId() - .isPresent())), - Tag.of("pushTokenType", pushTokenType)); - - if (StringUtils.isNotBlank(previousPushTokenType)) { - tags = tags.and(Tag.of("previousPushTokenType", previousPushTokenType)); - } - - Metrics.counter(CREATE_COUNTER_NAME, tags).increment(); - - accountAttributes.recoveryPassword().ifPresent(registrationRecoveryPassword -> - registrationRecoveryPasswordsManager.store(account.getIdentifier(IdentityType.PNI), - registrationRecoveryPassword)); - }, accountLockExecutor); - - return account; + logger.error("Unexpected exception while creating account", e); + throw new RuntimeException(e); + } }); } + private Account create(final String number, + final UUID pni, + final AccountAttributes accountAttributes, + final List accountBadges, + final IdentityKey aciIdentityKey, + final IdentityKey pniIdentityKey, + final DeviceSpec primaryDeviceSpec, + @Nullable final String userAgent) { + + final Account account = new Account(); + final Optional maybeRecentlyDeletedAccountIdentifier = + accounts.findRecentlyDeletedAccountIdentifier(pni); + + // Reuse the ACI from any recently-deleted account with this number to cover cases where somebody is + // re-registering. + account.setUuid(maybeRecentlyDeletedAccountIdentifier.orElseGet(UUID::randomUUID)); + account.setNumber(number, pni); + account.setIdentityKey(aciIdentityKey); + account.setPhoneNumberIdentityKey(pniIdentityKey); + 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); + + String accountCreationType = maybeRecentlyDeletedAccountIdentifier.isPresent() ? "recently-deleted" : "new"; + + final String pushTokenType; + + if (primaryDeviceSpec.apnRegistrationId().isPresent()) { + pushTokenType = "apns"; + } else if (primaryDeviceSpec.gcmRegistrationId().isPresent()) { + pushTokenType = "fcm"; + } else { + pushTokenType = "none"; + } + + String previousPushTokenType = null; + + try { + accounts.create(account, keysManager.buildWriteItemsForNewDevice(account.getIdentifier(IdentityType.ACI), + account.getIdentifier(IdentityType.PNI), + Device.PRIMARY_ID, + primaryDeviceSpec.aciSignedPreKey(), + primaryDeviceSpec.pniSignedPreKey(), + primaryDeviceSpec.aciPqLastResortPreKey(), + primaryDeviceSpec.pniPqLastResortPreKey())); + } catch (final AccountAlreadyExistsException e) { + accountCreationType = "re-registration"; + + if (StringUtils.isNotBlank(e.getExistingAccount().getPrimaryDevice().getApnId())) { + previousPushTokenType = "apns"; + } else if (StringUtils.isNotBlank(e.getExistingAccount().getPrimaryDevice().getGcmId())) { + previousPushTokenType = "fcm"; + } else { + previousPushTokenType = "none"; + } + + final UUID aci = e.getExistingAccount().getIdentifier(IdentityType.ACI); + account.setUuid(aci); + + final List additionalWriteItems = Stream.concat( + keysManager.buildWriteItemsForNewDevice(account.getIdentifier(IdentityType.ACI), + account.getIdentifier(IdentityType.PNI), + Device.PRIMARY_ID, + primaryDeviceSpec.aciSignedPreKey(), + primaryDeviceSpec.pniSignedPreKey(), + primaryDeviceSpec.aciPqLastResortPreKey(), + primaryDeviceSpec.pniPqLastResortPreKey()).stream(), + e.getExistingAccount().getDevices() + .stream() + .map(Device::getId) + // No need to clear the keys for the primary device since we'll just overwrite them in the same + // transaction anyhow + .filter(existingDeviceId -> existingDeviceId != Device.PRIMARY_ID) + .flatMap(existingDeviceId -> + keysManager.buildWriteItemsForRemovedDevice(aci, pni, existingDeviceId).stream())) + .toList(); + + CompletableFuture.allOf( + keysManager.deleteSingleUsePreKeys(aci), + keysManager.deleteSingleUsePreKeys(pni), + messagesManager.clear(aci), + profilesManager.deleteAll(aci)) + .thenCompose(ignored -> disconnectionRequestManager.requestDisconnection(aci)) + .thenCompose(ignored -> accounts.reclaimAccount(e.getExistingAccount(), account, additionalWriteItems)) + .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.deleteSingleUsePreKeys(aci), + keysManager.deleteSingleUsePreKeys(pni), + messagesManager.clear(aci), + profilesManager.deleteAll(aci)); + }) + .join(); + } + + redisSet(account); + + Tags tags = Tags.of(UserAgentTagUtil.getPlatformTag(userAgent), + Tag.of("type", accountCreationType), + Tag.of("hasPushToken", String.valueOf( + primaryDeviceSpec.apnRegistrationId().isPresent() || primaryDeviceSpec.gcmRegistrationId() + .isPresent())), + Tag.of("pushTokenType", pushTokenType)); + + if (StringUtils.isNotBlank(previousPushTokenType)) { + tags = tags.and(Tag.of("previousPushTokenType", previousPushTokenType)); + } + + Metrics.counter(CREATE_COUNTER_NAME, tags).increment(); + + accountAttributes.recoveryPassword().ifPresent(registrationRecoveryPassword -> + registrationRecoveryPasswordsManager.store(account.getIdentifier(IdentityType.PNI), + registrationRecoveryPassword)); + + return account; + } + public CompletableFuture> addDevice(final Account account, final DeviceSpec deviceSpec, final String linkDeviceToken) { return accountLockManager.withLockAsync(List.of(account.getPhoneNumberIdentifier()), () -> addDevice(account.getIdentifier(IdentityType.ACI), deviceSpec, linkDeviceToken, MAX_UPDATE_ATTEMPTS), @@ -655,57 +675,74 @@ public class AccountsManager extends RedisPubSubAdapter implemen validateDevices(account, pniSignedPreKeys, pniPqLastResortPreKeys, pniRegistrationIds); - final AtomicReference updatedAccount = new AtomicReference<>(); - - accountLockManager.withLock(List.of(account.getPhoneNumberIdentifier(), targetPhoneNumberIdentifier), () -> { - redisDelete(account); - - // There are three possible states for accounts associated with the target phone number: - // - // 1. An account exists with the target PNI; the caller has proved ownership of the number, so delete the - // account with the target PNI. This will leave a "deleted account" record for the deleted account mapping - // the UUID of the deleted account to the target PNI. We'll then overwrite that so it points to the - // original PNI to facilitate switching back and forth between numbers. - // 2. No account with the target PNI 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 PNI - // 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 PNI exists at all, in which case no additional action is needed. - final Optional recentlyDeletedAci = accounts.findRecentlyDeletedAccountIdentifier(targetPhoneNumberIdentifier); - final Optional maybeExistingAccount = getByE164(targetNumber); - final Optional maybeDisplacedUuid; - - if (maybeExistingAccount.isPresent()) { - delete(maybeExistingAccount.get()).join(); - maybeDisplacedUuid = maybeExistingAccount.map(Account::getUuid); - } else { - maybeDisplacedUuid = recentlyDeletedAci; + try { + return accountLockManager.withLock(List.of(account.getPhoneNumberIdentifier(), targetPhoneNumberIdentifier), + () -> changeNumber(account, targetNumber, targetPhoneNumberIdentifier, pniIdentityKey, pniSignedPreKeys, pniPqLastResortPreKeys, pniRegistrationIds), accountLockExecutor); + } catch (final Exception e) { + if (e instanceof MismatchedDevicesException mismatchedDevicesException) { + throw mismatchedDevicesException; + } if (e instanceof RuntimeException runtimeException) { + throw runtimeException; } - final UUID uuid = account.getUuid(); + logger.error("Unexpected exception when changing phone number", e); + throw new RuntimeException(e); + } + } - CompletableFuture.allOf( - keysManager.deleteSingleUsePreKeys(targetPhoneNumberIdentifier), - keysManager.deleteSingleUsePreKeys(originalPhoneNumberIdentifier)) - .join(); + private Account changeNumber(final Account account, + final String targetNumber, + final UUID targetPhoneNumberIdentifier, + final IdentityKey pniIdentityKey, + final Map pniSignedPreKeys, + final Map pniPqLastResortPreKeys, + final Map pniRegistrationIds) { + + final UUID originalPhoneNumberIdentifier = account.getPhoneNumberIdentifier(); + + redisDelete(account); + + // There are three possible states for accounts associated with the target phone number: + // + // 1. An account exists with the target PNI; the caller has proved ownership of the number, so delete the + // account with the target PNI. This will leave a "deleted account" record for the deleted account mapping + // the UUID of the deleted account to the target PNI. We'll then overwrite that so it points to the + // original PNI to facilitate switching back and forth between numbers. + // 2. No account with the target PNI 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 PNI + // 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 PNI exists at all, in which case no additional action is needed. + final Optional recentlyDeletedAci = accounts.findRecentlyDeletedAccountIdentifier(targetPhoneNumberIdentifier); + final Optional maybeExistingAccount = getByE164(targetNumber); + final Optional maybeDisplacedUuid; + + if (maybeExistingAccount.isPresent()) { + delete(maybeExistingAccount.get()).join(); + maybeDisplacedUuid = maybeExistingAccount.map(Account::getUuid); + } else { + maybeDisplacedUuid = recentlyDeletedAci; + } + + final UUID uuid = account.getUuid(); + + CompletableFuture.allOf( + keysManager.deleteSingleUsePreKeys(targetPhoneNumberIdentifier), + keysManager.deleteSingleUsePreKeys(originalPhoneNumberIdentifier)) + .join(); final Collection keyWriteItems = buildPniKeyWriteItems(targetPhoneNumberIdentifier, pniSignedPreKeys, pniPqLastResortPreKeys); - final Account numberChangedAccount = updateWithRetries( - account, - a -> { - setPniKeys(account, pniIdentityKey, pniRegistrationIds); - return true; - }, - a -> accounts.changeNumber(a, targetNumber, targetPhoneNumberIdentifier, maybeDisplacedUuid, keyWriteItems), - () -> accounts.getByAccountIdentifier(uuid).orElseThrow(), - AccountChangeValidator.NUMBER_CHANGE_VALIDATOR); - - updatedAccount.set(numberChangedAccount); - }, accountLockExecutor); - - return updatedAccount.get(); + return updateWithRetries( + account, + a -> { + setPniKeys(account, pniIdentityKey, pniRegistrationIds); + return true; + }, + a -> accounts.changeNumber(a, targetNumber, targetPhoneNumberIdentifier, maybeDisplacedUuid, keyWriteItems), + () -> accounts.getByAccountIdentifier(uuid).orElseThrow(), + AccountChangeValidator.NUMBER_CHANGE_VALIDATOR); } public Account updatePniKeys(final Account account, diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountLockManagerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountLockManagerTest.java index faf018ff2..3156e24f7 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountLockManagerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountLockManagerTest.java @@ -47,9 +47,8 @@ class AccountLockManagerTest { } @Test - void withLock() throws InterruptedException { - accountLockManager.withLock(List.of(FIRST_PNI, SECOND_PNI), () -> { - }, executor); + void withLock() throws Exception { + accountLockManager.withLock(List.of(FIRST_PNI, SECOND_PNI), () -> null, executor); verify(lockClient, times(2)).acquireLock(any()); verify(lockClient, times(2)).releaseLock(any(ReleaseLockOptions.class)); @@ -69,8 +68,7 @@ class AccountLockManagerTest { void withLockEmptyList() { final Runnable task = mock(Runnable.class); - assertThrows(IllegalArgumentException.class, () -> accountLockManager.withLock(Collections.emptyList(), () -> { - }, + assertThrows(IllegalArgumentException.class, () -> accountLockManager.withLock(Collections.emptyList(), () -> null, executor)); verify(task, never()).run(); } diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerConcurrentModificationIntegrationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerConcurrentModificationIntegrationTest.java index b363cdd38..b04d93397 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerConcurrentModificationIntegrationTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerConcurrentModificationIntegrationTest.java @@ -28,6 +28,7 @@ import java.util.ArrayList; import java.util.Optional; import java.util.Set; import java.util.UUID; +import java.util.concurrent.Callable; import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executor; import java.util.concurrent.LinkedBlockingDeque; @@ -84,7 +85,7 @@ class AccountsManagerConcurrentModificationIntegrationTest { private Executor mutationExecutor = new ThreadPoolExecutor(20, 20, 5, TimeUnit.SECONDS, new LinkedBlockingDeque<>(20)); @BeforeEach - void setup() throws InterruptedException { + void setup() throws Exception { @SuppressWarnings("unchecked") final DynamicConfigurationManager dynamicConfigurationManager = mock(DynamicConfigurationManager.class); @@ -108,10 +109,8 @@ class AccountsManagerConcurrentModificationIntegrationTest { final AccountLockManager accountLockManager = mock(AccountLockManager.class); doAnswer(invocation -> { - final Runnable task = invocation.getArgument(1); - task.run(); - - return null; + final Callable task = invocation.getArgument(1); + return task.call(); }).when(accountLockManager).withLock(anyList(), any(), any()); when(accountLockManager.withLockAsync(anyList(), any(), any())).thenAnswer(invocation -> { diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerTest.java index 097d2b940..1bba9a59a 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerTest.java @@ -52,6 +52,7 @@ import java.util.Objects; import java.util.Optional; import java.util.Set; import java.util.UUID; +import java.util.concurrent.Callable; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; import java.util.concurrent.Executor; @@ -153,7 +154,7 @@ class AccountsManagerTest { }; @BeforeEach - void setup() throws InterruptedException { + void setup() throws Exception { accounts = mock(Accounts.class); keysManager = mock(KeysManager.class); messagesManager = mock(MessagesManager.class); @@ -213,10 +214,8 @@ class AccountsManagerTest { final AccountLockManager accountLockManager = mock(AccountLockManager.class); doAnswer(invocation -> { - final Runnable task = invocation.getArgument(1); - task.run(); - - return null; + final Callable task = invocation.getArgument(1); + return task.call(); }).when(accountLockManager).withLock(anyList(), any(), any()); when(accountLockManager.withLockAsync(anyList(), any(), any())).thenAnswer(invocation -> { diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java index 523015337..4fa4ab6e1 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java @@ -28,6 +28,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; import java.util.UUID; +import java.util.concurrent.Callable; import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executors; import java.util.function.Supplier; @@ -81,12 +82,12 @@ class AccountsManagerUsernameIntegrationTest { private Accounts accounts; @BeforeEach - void setup() throws InterruptedException { + void setup() throws Exception { buildAccountsManager(1, 2, 10); } private void buildAccountsManager(final int initialWidth, int discriminatorMaxWidth, int attemptsPerWidth) - throws InterruptedException { + throws Exception { @SuppressWarnings("unchecked") final DynamicConfigurationManager dynamicConfigurationManager = mock(DynamicConfigurationManager.class); @@ -115,10 +116,8 @@ class AccountsManagerUsernameIntegrationTest { final AccountLockManager accountLockManager = mock(AccountLockManager.class); doAnswer(invocation -> { - final Runnable task = invocation.getArgument(1); - task.run(); - - return null; + final Callable task = invocation.getArgument(1); + return task.call(); }).when(accountLockManager).withLock(anyList(), any(), any()); when(accountLockManager.withLockAsync(anyList(), any(), any())).thenAnswer(invocation -> {