diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java index 6d753b852..4bedeb482 100644 --- 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 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 index 7128c9360..c35dcf717 100644 --- 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 expectedAccounts = new ArrayList<>();