From 07c04006df66f2aaf24625f93877fb66fffd6cda Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Mon, 13 Nov 2023 12:14:37 -0500
Subject: [PATCH] Avoid blocking calls in async account updates

---
 .../org/whispersystems/textsecuregcm/storage/Accounts.java | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

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 1c6bf451e..2378e3c16 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java
@@ -807,14 +807,17 @@ public class Accounts extends AbstractDynamoDbStore {
             account.setVersion(AttributeValues.getInt(response.attributes(), "V", account.getVersion() + 1));
             return (Void) null;
           })
-          .exceptionally(throwable -> {
+          .exceptionallyCompose(throwable -> {
             final Throwable unwrapped = ExceptionUtils.unwrap(throwable);
             if (unwrapped instanceof TransactionConflictException) {
               throw new ContestedOptimisticLockException();
             } else if (unwrapped instanceof ConditionalCheckFailedException e) {
               // the exception doesn't give details about which condition failed,
               // but we can infer it was an optimistic locking failure if the UUID is known
-              throw getByAccountIdentifier(account.getUuid()).isPresent() ? new ContestedOptimisticLockException() : e;
+              return getByAccountIdentifierAsync(account.getUuid())
+                  .thenAccept(refreshedAccount -> {
+                    throw refreshedAccount.isPresent() ? new ContestedOptimisticLockException() : e;
+                  });
             } else {
               // rethrow
               throw CompletableFutureUtils.errorAsCompletionException(throwable);