Treat transaction conflicts during transactional account updates as c…

…ontested optimistic locks
signalapp · Dec 8, 2023 · cca747a · cca747a
1 parent 417d99a
commit cca747a
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 1 deletion.
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java
@@ -887,7 +887,14 @@ public CompletionStage<Void> updateTransactionallyAsync(final Account account,
             final Throwable unwrapped = ExceptionUtils.unwrap(throwable);
 
             if (unwrapped instanceof TransactionCanceledException transactionCanceledException) {
-              if ("ConditionalCheckFailed".equals(transactionCanceledException.cancellationReasons().get(0).code())) {
+              if (CONDITIONAL_CHECK_FAILED.equals(transactionCanceledException.cancellationReasons().get(0).code())) {
+                throw new ContestedOptimisticLockException();
+              }
+
+              if (transactionCanceledException.cancellationReasons()
+                  .stream()
+                  .anyMatch(reason -> TRANSACTION_CONFLICT.equals(reason.code()))) {
+
                 throw new ContestedOptimisticLockException();
               }
             }

diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java
@@ -64,6 +64,7 @@
 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
+import software.amazon.awssdk.services.dynamodb.model.CancellationReason;
 import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException;
 import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
@@ -627,6 +628,32 @@ void testUpdateTransactionallyContestedLock() {
     assertTrue(completionException.getCause() instanceof ContestedOptimisticLockException);
   }
 
+  @Test
+  void testUpdateTransactionallyWithMockTransactionConflictException() {
+    final DynamoDbAsyncClient dynamoDbAsyncClient = mock(DynamoDbAsyncClient.class);
+
+    accounts = new Accounts(mock(DynamoDbClient.class),
+        dynamoDbAsyncClient,
+        Tables.ACCOUNTS.tableName(),
+        Tables.NUMBERS.tableName(),
+        Tables.PNI_ASSIGNMENTS.tableName(),
+        Tables.USERNAMES.tableName(),
+        Tables.DELETED_ACCOUNTS.tableName());
+
+    when(dynamoDbAsyncClient.transactWriteItems(any(TransactWriteItemsRequest.class)))
+        .thenReturn(CompletableFuture.failedFuture(TransactionCanceledException.builder()
+            .cancellationReasons(CancellationReason.builder()
+                .code("TransactionConflict")
+                .build())
+            .build()));
+
+    Account account = generateAccount("+14151112222", UUID.randomUUID(), UUID.randomUUID());
+
+    assertThatThrownBy(() -> accounts.updateTransactionallyAsync(account, Collections.emptyList()).toCompletableFuture().join())
+        .isInstanceOfAny(CompletionException.class)
+        .hasCauseInstanceOf(ContestedOptimisticLockException.class);
+  }
+
   @Test
   void testGetAll() {
     final List<Account> expectedAccounts = new ArrayList<>();