Accquire pessimistic account locks by sets of identifiers instead of lists

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

@@ -8,6 +8,7 @@ import com.amazonaws.services.dynamodbv2.ReleaseLockOptions;
 import com.google.common.annotations.VisibleForTesting;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.Callable;
 import java.util.concurrent.CompletableFuture;
@@ -55,7 +56,7 @@ public class AccountLockManager {
    *
    * @throws Exception if an exception is thrown by the given {@code task}
    */
-  public <V> V withLock(final List<UUID> phoneNumberIdentifiers,
+  public <V> V withLock(final Set<UUID> phoneNumberIdentifiers,
       final Callable<V> task,
       final Executor lockAcquisitionExecutor) throws Exception {
 
@@ -106,7 +107,7 @@ public class AccountLockManager {
    * @return a future that completes normally when the given task has executed successfully and all locks have been
    * released; the returned future may fail with an {@link InterruptedException} if interrupted while acquiring a lock
    */
-  public <T> CompletableFuture<T> withLockAsync(final List<UUID> phoneNumberIdentifiers,
+  public <T> CompletableFuture<T> withLockAsync(final Set<UUID> phoneNumberIdentifiers,
       final Supplier<CompletableFuture<T>> taskSupplier, final Executor executor) {
 
     if (phoneNumberIdentifiers.isEmpty()) {

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

@@ -277,7 +277,7 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
 
     return createTimer.record(() -> {
       try {
-        return accountLockManager.withLock(List.of(pni),
+        return accountLockManager.withLock(Set.of(pni),
             () -> create(number, pni, accountAttributes, accountBadges, aciIdentityKey, pniIdentityKey, primaryDeviceSpec, userAgent), accountLockExecutor);
       } catch (final Exception e) {
         if (e instanceof RuntimeException runtimeException) {
@@ -416,7 +416,7 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
   }
 
   public CompletableFuture<Pair<Account, Device>> addDevice(final Account account, final DeviceSpec deviceSpec, final String linkDeviceToken) {
-    return accountLockManager.withLockAsync(List.of(account.getPhoneNumberIdentifier()),
+    return accountLockManager.withLockAsync(Set.of(account.getPhoneNumberIdentifier()),
         () -> addDevice(account.getIdentifier(IdentityType.ACI), deviceSpec, linkDeviceToken, MAX_UPDATE_ATTEMPTS),
         accountLockExecutor);
   }
@@ -610,7 +610,7 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
       throw new IllegalArgumentException("Cannot remove primary device");
     }
 
-    return accountLockManager.withLockAsync(List.of(account.getPhoneNumberIdentifier()),
+    return accountLockManager.withLockAsync(Set.of(account.getPhoneNumberIdentifier()),
         () -> removeDevice(account.getIdentifier(IdentityType.ACI), deviceId, MAX_UPDATE_ATTEMPTS),
         accountLockExecutor);
   }
@@ -671,7 +671,7 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
     }
 
     try {
-      return accountLockManager.withLock(List.of(account.getPhoneNumberIdentifier(), targetPhoneNumberIdentifier),
+      return accountLockManager.withLock(Set.of(account.getPhoneNumberIdentifier(), targetPhoneNumberIdentifier),
           () -> changeNumber(account, targetNumber, targetPhoneNumberIdentifier, pniIdentityKey, pniSignedPreKeys, pniPqLastResortPreKeys, pniRegistrationIds), accountLockExecutor);
     } catch (final Exception e) {
       if (e instanceof MismatchedDevicesException mismatchedDevicesException) {
@@ -749,7 +749,7 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
       final Map<Byte, Integer> pniRegistrationIds) throws MismatchedDevicesException {
 
     try {
-      return accountLockManager.withLock(List.of(account.getIdentifier(IdentityType.PNI)), () -> {
+      return accountLockManager.withLock(Set.of(account.getIdentifier(IdentityType.PNI)), () -> {
         validateDevices(account, pniSignedPreKeys, pniPqLastResortPreKeys, pniRegistrationIds);
 
         final UUID aci = account.getIdentifier(IdentityType.ACI);
@@ -1263,7 +1263,7 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
   public CompletableFuture<Void> delete(final Account account, final DeletionReason deletionReason) {
     final Timer.Sample sample = Timer.start();
 
-    return accountLockManager.withLockAsync(List.of(account.getPhoneNumberIdentifier()), () -> delete(account),
+    return accountLockManager.withLockAsync(Set.of(account.getPhoneNumberIdentifier()), () -> delete(account),
             accountLockExecutor)
         .whenComplete((ignored, throwable) -> {
           sample.stop(deleteTimer);

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

@@ -2,6 +2,7 @@ package org.whispersystems.textsecuregcm.storage;
 
 import java.util.List;
 import java.util.Optional;
+import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
@@ -41,7 +42,7 @@ public class ClientPublicKeysManager {
    * @return a future that completes when the given key has been stored
    */
   public CompletableFuture<Void> setPublicKey(final Account account, final byte deviceId, final ECPublicKey publicKey) {
-    return accountLockManager.withLockAsync(List.of(account.getPhoneNumberIdentifier()),
+    return accountLockManager.withLockAsync(Set.of(account.getPhoneNumberIdentifier()),
         () -> clientPublicKeys.setPublicKey(account.getIdentifier(IdentityType.ACI), deviceId, publicKey),
         accountLockExecutor);
   }