From 3f4e1522ebe1606d54285df997e2c744e3c88a94 Mon Sep 17 00:00:00 2001
From: Chris Eager <chris@signal.org>
Date: Mon, 13 Sep 2021 14:08:47 -0700
Subject: [PATCH] Only put accounts that exhaust optimistic lock retries in
 migration retry table

---
 .../textsecuregcm/storage/AccountsDynamoDb.java  | 16 +++++++++++++---
 .../textsecuregcm/storage/AccountsManager.java   | 11 +++++++++--
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsDynamoDb.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsDynamoDb.java
index 6f5e00fc7..3737561f9 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsDynamoDb.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsDynamoDb.java
@@ -230,9 +230,11 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
           throw get(account.getUuid()).isPresent() ? new ContestedOptimisticLockException() : e;
         }
       } catch (final Exception e) {
-        // the Dynamo account now lags the Postgres account version. Put it in the migration retry table so that it will
-        // get updated faster—otherwise it will be stale until the accounts crawler runs again
-        migrationRetryAccounts.put(account.getUuid());
+        if (!(e instanceof ContestedOptimisticLockException)) {
+          // the Dynamo account now lags the Postgres account version. Put it in the migration retry table so that it will
+          // get updated faster—otherwise it will be stale until the accounts crawler runs again
+          migrationRetryAccounts.put(account.getUuid());
+        }
         throw e;
       }
 
@@ -430,6 +432,14 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
     }
   }
 
+  void putUuidForMigrationRetry(final UUID uuid) {
+    try {
+      migrationRetryAccounts.put(uuid);
+    } catch (final Exception e) {
+      logger.error("Failed to store for retry: {}", uuid, e);
+    }
+  }
+
   private static String extractCancellationReasonCodes(final TransactionCanceledException exception) {
     return exception.cancellationReasons().stream()
         .map(CancellationReason::code)
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 cf149310b..b17c12e0c 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java
@@ -307,8 +307,15 @@ public class AccountsManager {
       updatedAccount = updateWithRetries(account, updater, this::databaseUpdate, () -> databaseGet(uuid).get());
 
       if (dynamoWriteEnabled()) {
-        runSafelyAndRecordMetrics(() -> dynamoGet(uuid).map(dynamoAccount ->
-                updateWithRetries(dynamoAccount, updater, this::dynamoUpdate, () -> dynamoGet(uuid).get())),
+        runSafelyAndRecordMetrics(() -> dynamoGet(uuid).map(dynamoAccount -> {
+              try {
+                return updateWithRetries(dynamoAccount, updater, this::dynamoUpdate, () -> dynamoGet(uuid).get());
+              } catch (final OptimisticLockRetryLimitExceededException e) {
+                accountsDynamoDb.putUuidForMigrationRetry(uuid);
+
+                throw e;
+              }
+            }),
             Optional.of(uuid),
             Optional.of(updatedAccount),
             this::compareAccounts,