Map TransactionConflict to ContestedOptimisticLockException in username flows

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

@@ -455,7 +455,8 @@ public class Accounts extends AbstractDynamoDbStore {
   /**
    * Reserve a username hash under the account UUID
    * @return a future that completes once the username hash has been reserved; may fail with an
-   * {@link ContestedOptimisticLockException} if the account has been updated or
+   * {@link ContestedOptimisticLockException} if the account has been updated or there are concurrent updates to the
+   * account or constraint records, and with an
    * {@link UsernameHashNotAvailableException} if the username was taken by someone else
    */
   public CompletableFuture<Void> reserveUsernameHash(
@@ -523,11 +524,12 @@ public class Accounts extends AbstractDynamoDbStore {
         .exceptionally(throwable -> {
           if (ExceptionUtils.unwrap(throwable) instanceof TransactionCanceledException e) {
             // If the constraint table update failed the condition check, the username's taken and we should stop
-            // trying. However if it was only in the accounts table that the condition check update failed, it's an
+            // trying. However, if the accounts table fails the conditional check or
-            // optimistic locking failure (the account was concurrently updated) and we should try again.
+            // either table was concurrently updated, it's an optimistic locking failure and we should try again.
             if (conditionalCheckFailed(e.cancellationReasons().get(0))) {
               throw ExceptionUtils.wrap(new UsernameHashNotAvailableException());
-            } else if (conditionalCheckFailed(e.cancellationReasons().get(1))) {
+            } else if (conditionalCheckFailed(e.cancellationReasons().get(1)) ||
+                e.cancellationReasons().stream().anyMatch(Accounts::isTransactionConflict)) {
               throw new ContestedOptimisticLockException();
             }
           }
@@ -551,8 +553,10 @@ public class Accounts extends AbstractDynamoDbStore {
    *
    * @param account to update
    * @param usernameHash believed to be available
+   * @param encryptedUsername the encrypted form of the previously reserved username; used for the username link
    * @return a future that completes once the username hash has been confirmed; may fail with an
-   * {@link ContestedOptimisticLockException} if the account has been updated or
+   * {@link ContestedOptimisticLockException} if the account has been updated or there are concurrent updates to the
+   * account or constraint records, and with an
    * {@link UsernameHashNotAvailableException} if the username was taken by someone else
    */
   public CompletableFuture<Void> confirmUsernameHash(final Account account, final byte[] usernameHash, @Nullable final byte[] encryptedUsername) {
@@ -585,13 +589,14 @@ public class Accounts extends AbstractDynamoDbStore {
         .exceptionally(throwable -> {
           if (ExceptionUtils.unwrap(throwable) instanceof TransactionCanceledException e) {
             // If the constraint table update failed the condition check, the username's taken and we should stop
-            // trying. However if it was only in the accounts table that the condition check update failed, it's an
+            // trying. However, if the accounts table fails the conditional check or
-            // optimistic locking failure (the account was concurrently updated) and we should try again.
+            // either table was concurrently updated, it's an optimistic locking failure and we should try again.
             // NOTE: the fixed indices here must be kept in sync with the creation of the TransactWriteItems in
             // buildConfirmUsernameHashRequest!
             if (conditionalCheckFailed(e.cancellationReasons().get(0))) {
               throw ExceptionUtils.wrap(new UsernameHashNotAvailableException());
-            } else if (conditionalCheckFailed(e.cancellationReasons().get(1))) {
+            } else if (conditionalCheckFailed(e.cancellationReasons().get(1)) ||
+                e.cancellationReasons().stream().anyMatch(Accounts::isTransactionConflict)) {
               throw new ContestedOptimisticLockException();
             }
           }
@@ -690,6 +695,14 @@ public class Accounts extends AbstractDynamoDbStore {
         .build();
   }
 
+  /**
+   * Clear the username hash and link from the given account
+   *
+   * @param account to update
+   * @return a future that completes once the username data has been cleared;
+   * it can fail with a {@link ContestedOptimisticLockException} if there are concurrent updates
+   * to the account or username constraint records.
+   */
   public CompletableFuture<Void> clearUsernameHash(final Account account) {
     return account.getUsernameHash().map(usernameHash -> {
       final Timer.Sample sample = Timer.start();
@@ -714,8 +727,9 @@ public class Accounts extends AbstractDynamoDbStore {
             account.setVersion(account.getVersion() + 1);
           })
           .exceptionally(throwable -> {
-            if (ExceptionUtils.unwrap(throwable) instanceof TransactionCanceledException transactionCanceledException) {
+            if (ExceptionUtils.unwrap(throwable) instanceof TransactionCanceledException e) {
-              if (conditionalCheckFailed(transactionCanceledException.cancellationReasons().get(0))) {
+              if (conditionalCheckFailed(e.cancellationReasons().get(0)) ||
+                  e.cancellationReasons().stream().anyMatch(Accounts::isTransactionConflict)) {
                 throw new ContestedOptimisticLockException();
               }
             }
@@ -1371,4 +1385,8 @@ public class Accounts extends AbstractDynamoDbStore {
   private static boolean conditionalCheckFailed(final CancellationReason reason) {
     return CONDITIONAL_CHECK_FAILED.equals(reason.code());
   }
+
+  private static boolean isTransactionConflict(final CancellationReason reason) {
+    return TRANSACTION_CONFLICT.equals(reason.code());
+  }
 }

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

@@ -73,6 +73,7 @@ import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.ReturnValuesOnConditionCheckFailure;
 import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem;
 import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsRequest;
+import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsResponse;
 import software.amazon.awssdk.services.dynamodb.model.TransactionCanceledException;
 import software.amazon.awssdk.services.dynamodb.model.TransactionConflictException;
 import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
@@ -960,7 +961,7 @@ class AccountsTest {
   }
 
   @Test
-  void testUsernameHashConflict() {
+  void testUsernameHashNotAvailable() {
     final Account firstAccount = generateAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID());
     final Account secondAccount = generateAccount("+18005559876", UUID.randomUUID(), UUID.randomUUID());
 
@@ -994,6 +995,92 @@ class AccountsTest {
     assertThat(secondAccount.getUsernameHash()).isEmpty();
   }
 
+  @ParameterizedTest
+  @MethodSource
+  void testReserveUsernameHashTransactionConflict(final Optional<String> constraintCancellationString,
+      final Optional<String> accountsCancellationString,
+      final Class<Exception> expectedException) {
+    final DynamoDbAsyncClient dbAsyncClient = mock(DynamoDbAsyncClient.class);
+
+    accounts = new Accounts(mock(DynamoDbClient.class),
+        dbAsyncClient,
+        Tables.ACCOUNTS.tableName(),
+        Tables.NUMBERS.tableName(),
+        Tables.PNI_ASSIGNMENTS.tableName(),
+        Tables.USERNAMES.tableName(),
+        Tables.DELETED_ACCOUNTS.tableName());
+    final Account account = generateAccount("+14155551111", UUID.randomUUID(), UUID.randomUUID());
+    createAccount(account);
+
+    final CancellationReason constraintCancellationReason = constraintCancellationString.map(
+        reason -> CancellationReason.builder().code(reason).build()
+    ).orElse(CancellationReason.builder().build());
+
+    final CancellationReason accountsCancellationReason = accountsCancellationString.map(
+        reason -> CancellationReason.builder().code(reason).build()
+    ).orElse(CancellationReason.builder().build());
+
+    when(dbAsyncClient.transactWriteItems(any(TransactWriteItemsRequest.class)))
+        .thenReturn(CompletableFuture.failedFuture(TransactionCanceledException.builder()
+            .cancellationReasons(constraintCancellationReason, accountsCancellationReason)
+            .build()));
+
+    CompletableFutureTestUtil.assertFailsWithCause(expectedException,
+        accounts.reserveUsernameHash(account, USERNAME_HASH_1, Duration.ofDays(1)));
+  }
+
+  private static Stream<Arguments> testReserveUsernameHashTransactionConflict() {
+    return Stream.of(
+        Arguments.of(Optional.of("TransactionConflict"), Optional.empty(), ContestedOptimisticLockException.class),
+        Arguments.of(Optional.empty(), Optional.of("TransactionConflict"), ContestedOptimisticLockException.class),
+        Arguments.of(Optional.of("ConditionalCheckFailed"), Optional.of("TransactionConflict"), UsernameHashNotAvailableException.class)
+    );
+  }
+
+  @ParameterizedTest
+  @MethodSource
+  void testConfirmUsernameHashTransactionConflict(final Optional<String> constraintCancellationString,
+      final Optional<String> accountsCancellationString,
+      final Class<Exception> expectedException) {
+    final DynamoDbAsyncClient dbAsyncClient = mock(DynamoDbAsyncClient.class);
+
+    accounts = new Accounts(mock(DynamoDbClient.class),
+        dbAsyncClient,
+        Tables.ACCOUNTS.tableName(),
+        Tables.NUMBERS.tableName(),
+        Tables.PNI_ASSIGNMENTS.tableName(),
+        Tables.USERNAMES.tableName(),
+        Tables.DELETED_ACCOUNTS.tableName());
+    final Account account = generateAccount("+14155551111", UUID.randomUUID(), UUID.randomUUID());
+    createAccount(account);
+
+    final CancellationReason constraintCancellationReason = constraintCancellationString.map(
+        reason -> CancellationReason.builder().code(reason).build()
+    ).orElse(CancellationReason.builder().build());
+
+    final CancellationReason accountsCancellationReason = accountsCancellationString.map(
+        reason -> CancellationReason.builder().code(reason).build()
+    ).orElse(CancellationReason.builder().build());
+
+    when(dbAsyncClient.transactWriteItems(any(TransactWriteItemsRequest.class)))
+        .thenReturn(CompletableFuture.failedFuture(TransactionCanceledException.builder()
+            .cancellationReasons(constraintCancellationReason,
+                accountsCancellationReason,
+                CancellationReason.builder().build())
+            .build()));
+
+    CompletableFutureTestUtil.assertFailsWithCause(expectedException,
+        accounts.confirmUsernameHash(account, USERNAME_HASH_1, ENCRYPTED_USERNAME_1));
+  }
+
+  private static Stream<Arguments> testConfirmUsernameHashTransactionConflict() {
+    return Stream.of(
+        Arguments.of(Optional.of("TransactionConflict"), Optional.empty(), ContestedOptimisticLockException.class),
+        Arguments.of(Optional.empty(), Optional.of("TransactionConflict"), ContestedOptimisticLockException.class),
+        Arguments.of(Optional.of("ConditionalCheckFailed"), Optional.of("TransactionConflict"), UsernameHashNotAvailableException.class)
+    );
+  }
+
   @Test
   void testConfirmUsernameHashVersionMismatch() {
     final Account account = generateAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID());
@@ -1051,6 +1138,55 @@ class AccountsTest {
     assertArrayEquals(account.getUsernameHash().orElseThrow(), USERNAME_HASH_1);
   }
 
+  @ParameterizedTest
+  @MethodSource
+  void testClearUsernameTransactionConflict(final Optional<String> constraintCancellationString,
+      final Optional<String> accountsCancellationString) {
+    final DynamoDbAsyncClient dbAsyncClient = mock(DynamoDbAsyncClient.class);
+
+    accounts = new Accounts(mock(DynamoDbClient.class),
+        dbAsyncClient,
+        Tables.ACCOUNTS.tableName(),
+        Tables.NUMBERS.tableName(),
+        Tables.PNI_ASSIGNMENTS.tableName(),
+        Tables.USERNAMES.tableName(),
+        Tables.DELETED_ACCOUNTS.tableName());
+
+    final Account account = generateAccount("+14155551111", UUID.randomUUID(), UUID.randomUUID());
+    createAccount(account);
+
+    when(dbAsyncClient.transactWriteItems(any(TransactWriteItemsRequest.class)))
+        .thenReturn(CompletableFuture.completedFuture(mock(TransactWriteItemsResponse.class)));
+
+    accounts.reserveUsernameHash(account, USERNAME_HASH_1, Duration.ofDays(1)).join();
+    accounts.confirmUsernameHash(account, USERNAME_HASH_1, ENCRYPTED_USERNAME_1).join();
+
+    final CancellationReason constraintCancellationReason = constraintCancellationString.map(
+        reason -> CancellationReason.builder().code(reason).build()
+    ).orElse(CancellationReason.builder().build());
+
+    final CancellationReason accountsCancellationReason = accountsCancellationString.map(
+        reason -> CancellationReason.builder().code(reason).build()
+    ).orElse(CancellationReason.builder().build());
+
+    when(dbAsyncClient.transactWriteItems(any(TransactWriteItemsRequest.class)))
+        .thenReturn(CompletableFuture.failedFuture(TransactionCanceledException.builder()
+            .cancellationReasons(accountsCancellationReason, constraintCancellationReason)
+            .build()));
+
+    CompletableFutureTestUtil.assertFailsWithCause(ContestedOptimisticLockException.class,
+        accounts.clearUsernameHash(account));
+
+    assertArrayEquals(account.getUsernameHash().orElseThrow(), USERNAME_HASH_1);
+  }
+
+  private static Stream<Arguments> testClearUsernameTransactionConflict() {
+    return Stream.of(
+        Arguments.of(Optional.empty(), Optional.of("TransactionConflict"), ContestedOptimisticLockException.class),
+        Arguments.of(Optional.of("TransactionConflict"), Optional.empty(), ContestedOptimisticLockException.class)
+    );
+  }
+
   @Test
   void testReservedUsernameHash() {
     final Account account1 = generateAccount("+18005551111", UUID.randomUUID(), UUID.randomUUID());