From d7447b111da89de2ac5521394a72b2aa4451c49a Mon Sep 17 00:00:00 2001 From: Jonathan Klabunde Tomer <125505367+jkt-signal@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:44:31 -0700 Subject: [PATCH] Add support for pessimistic locking of phone-number-less accounts --- .../storage/AccountLockManager.java | 55 +++++++++++++------ .../storage/AccountLockManagerTest.java | 26 +++++++++ 2 files changed, 65 insertions(+), 16 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 3f0e1d5f4..1bdc72074 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountLockManager.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountLockManager.java @@ -39,22 +39,20 @@ public class AccountLockManager { this.lockClient = lockClient; } - /** - * Acquires a distributed, pessimistic lock for the accounts identified by the given phone number identifiers. 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 identifiers have been acquired, and the locks are released as soon as the task - * completes by any means. - * - * @param phoneNumberIdentifiers the phone number identifiers for which to acquire a distributed, pessimistic lock - * @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. - * - * @return the value returned by the given {@code task} - * - * @throws E if an exception is thrown by the given {@code task} - */ + /// Acquires a distributed, pessimistic lock for the accounts identified by the given phone number identifiers. 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 identifiers have been acquired, and the locks are released as soon as the task + /// completes by any means. + /// + /// @param phoneNumberIdentifiers the phone number identifiers for which to acquire a distributed, pessimistic lock + /// @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. + /// + /// @return the value returned by the given {@code task} + /// + /// @throws E if an exception is thrown by the given {@code task} public V withLock(final Set phoneNumberIdentifiers, final ThrowingSupplier task, final Executor lockAcquisitionExecutor) throws E { @@ -92,4 +90,29 @@ public class AccountLockManager { }, lockAcquisitionExecutor).join(); } } + + /// Acquires a distributed, pessimistic lock for a single account that already exists. The given task + /// runs once a lock for the identifier has been acquired, and the lock is released as soon as + /// the task completes by any means. + /// + /// If the account has a phone number, the lock will be based on its phone-number identifier, and + /// will therefore guard against collisions with other operations that lock the same phone number + /// identifier; if it does not, the lock will be based on its account identifier only. This is + /// safe because accounts without phone numbers can never get phone numbers and accounts with + /// them can never lose them, so we will never see two operations targeting the same account but + /// locking with different identifier types. + /// + /// @param account the account for which to acquire a distributed, pessimistic lock + /// @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. + /// + /// @return the value returned by the given {@code task} + /// + /// @throws E if an exception is thrown by the given {@code task} + public V withSingleAccountLock(Account account, + final ThrowingSupplier task, + final Executor lockAcquisitionExecutor) throws E { + return withLock(Set.of(account.getPhoneNumberIdentifierOptional().orElse(account.getAccountIdentifier())), task, lockAcquisitionExecutor); + } } 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 f2ca703ea..c121e3a2c 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountLockManagerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountLockManagerTest.java @@ -6,10 +6,13 @@ 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 static org.mockito.Mockito.when; +import com.amazonaws.services.dynamodbv2.AcquireLockOptions; import com.amazonaws.services.dynamodbv2.AmazonDynamoDBLockClient; import com.amazonaws.services.dynamodbv2.ReleaseLockOptions; import java.util.Collections; +import java.util.Optional; import java.util.Set; import java.util.UUID; import java.util.concurrent.ExecutorService; @@ -28,6 +31,7 @@ class AccountLockManagerTest { private static final UUID FIRST_PNI = UUID.randomUUID(); private static final UUID SECOND_PNI = UUID.randomUUID(); + private static final UUID ACI = UUID.randomUUID(); @BeforeEach void setUp() { @@ -71,4 +75,26 @@ class AccountLockManagerTest { executor)); verify(task, never()).run(); } + + @Test + void withLockPniAccount() throws Exception { + final Account account = mock(Account.class); + when(account.getAccountIdentifier()).thenReturn(ACI); + when(account.getPhoneNumberIdentifierOptional()).thenReturn(Optional.of(FIRST_PNI)); + + accountLockManager.withSingleAccountLock(account, () -> null, executor); + verify(lockClient, times(1)).acquireLock( + AcquireLockOptions.builder(FIRST_PNI.toString()).withAcquireReleasedLocksConsistently(true).build()); + } + + @Test + void withLockNoPniAccount() throws Exception { + final Account account = mock(Account.class); + when(account.getAccountIdentifier()).thenReturn(ACI); + when(account.getPhoneNumberIdentifierOptional()).thenReturn(Optional.empty()); + + accountLockManager.withSingleAccountLock(account, () -> null, executor); + verify(lockClient, times(1)).acquireLock( + AcquireLockOptions.builder(ACI.toString()).withAcquireReleasedLocksConsistently(true).build()); + } }