From bcb89924b43422d24383a9934a1521f86cd0f1f5 Mon Sep 17 00:00:00 2001
From: Jon Chambers <63609320+jon-signal@users.noreply.github.com>
Date: Tue, 3 Aug 2021 11:54:26 -0400
Subject: [PATCH] Simplify optimistic write logic

---
 .../storage/AccountsManager.java              | 40 +++++--------------
 .../tests/storage/AccountsManagerTest.java    | 10 ++---
 2 files changed, 13 insertions(+), 37 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java
index 725fd7f0f..ca084b4bc 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java
@@ -293,47 +293,21 @@ public class AccountsManager {
 
     try (Timer.Context ignored = updateTimer.time()) {
 
-      if (!updater.apply(account)) {
-        return account;
-      }
-
-      {
-        // optimistically increment version
-        final int originalVersion = account.getVersion();
-        account.setVersion(originalVersion + 1);
-        redisSet(account);
-        account.setVersion(originalVersion);
-      }
+      redisDelete(account);
 
       final UUID uuid = account.getUuid();
 
       updatedAccount = updateWithRetries(account, updater, this::databaseUpdate, () -> databaseGet(uuid).get());
 
       if (dynamoWriteEnabled()) {
-        runSafelyAndRecordMetrics(() -> {
-
-              final Optional<Account> dynamoAccount = dynamoGet(uuid);
-              if (dynamoAccount.isPresent()) {
-
-                if (!updater.apply(dynamoAccount.get())) {
-                  return dynamoAccount;
-                }
-
-                Account dynamoUpdatedAccount = updateWithRetries(dynamoAccount.get(),
-                    updater,
-                    this::dynamoUpdate,
-                    () -> dynamoGet(uuid).get());
-
-                return Optional.of(dynamoUpdatedAccount);
-              }
-
-              return Optional.empty();
-            }, Optional.of(uuid), Optional.of(updatedAccount),
+        runSafelyAndRecordMetrics(() -> dynamoGet(uuid).map(dynamoAccount ->
+                updateWithRetries(dynamoAccount, updater, this::dynamoUpdate, () -> dynamoGet(uuid).get())),
+            Optional.of(uuid),
+            Optional.of(updatedAccount),
             this::compareAccounts,
             "update");
       }
 
-      // set the cache again, so that all updates are coalesced
       redisSet(updatedAccount);
     }
 
@@ -349,6 +323,10 @@ public class AccountsManager {
   private Account updateWithRetries(Account account, Function<Account, Boolean> updater, Consumer<Account> persister,
       Supplier<Account> retriever) {
 
+    if (!updater.apply(account)) {
+      return account;
+    }
+
     final int maxTries = 10;
     int tries = 0;
 
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountsManagerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountsManagerTest.java
index 93ace5d4a..40c75bff1 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountsManagerTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountsManagerTest.java
@@ -338,16 +338,14 @@ class AccountsManagerTest {
 
     ArgumentCaptor<String> redisSetArgumentCapture = ArgumentCaptor.forClass(String.class);
 
-    verify(commands, times(4)).set(anyString(), redisSetArgumentCapture.capture());
+    verify(commands, times(2)).set(anyString(), redisSetArgumentCapture.capture());
 
-    Account firstAccountCached = JsonHelpers.fromJson(redisSetArgumentCapture.getAllValues().get(1), Account.class);
-    Account secondAccountCached = JsonHelpers.fromJson(redisSetArgumentCapture.getAllValues().get(3), Account.class);
+    Account accountCached = JsonHelpers.fromJson(redisSetArgumentCapture.getAllValues().get(1), Account.class);
 
     // uuid is @JsonIgnore, so we need to set it for compareAccounts to work
-    firstAccountCached.setUuid(uuid);
-    secondAccountCached.setUuid(uuid);
+    accountCached.setUuid(uuid);
 
-    assertEquals(Optional.empty(), accountsManager.compareAccounts(Optional.of(firstAccountCached), Optional.of(secondAccountCached)));
+    assertEquals(Optional.empty(), accountsManager.compareAccounts(Optional.of(updatedAccount), Optional.of(accountCached)));
   }
 
   @Test