From d09dcc90fe51a13200ce3ecd6d5ec2b427638591 Mon Sep 17 00:00:00 2001 From: Jon Chambers Date: Fri, 23 Jul 2021 17:09:28 -0400 Subject: [PATCH] Add methods for getting, clearing, locking recently-deleted account records. --- .../storage/AccountsManager.java | 2 +- .../storage/DeletedAccounts.java | 25 ++++++++++- .../storage/DeletedAccountsManager.java | 17 +++++++- .../storage/DeletedAccountsManagerTest.java | 33 +++++++++++++-- .../storage/DeletedAccountsTest.java | 41 ++++++++++++++++--- 5 files changed, 104 insertions(+), 14 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 e7bd77629..d27aa0228 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java @@ -411,7 +411,7 @@ public class AccountsManager { } } - deletedAccountsManager.put(account.getUuid(), account.getNumber()); + deletedAccountsManager.addRecentlyDeletedAccount(account.getUuid(), account.getNumber()); } catch (final RuntimeException | InterruptedException e) { logger.warn("Failed to delete account", e); diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/DeletedAccounts.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/DeletedAccounts.java index a0a9cfc69..116e25f74 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/DeletedAccounts.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/DeletedAccounts.java @@ -13,6 +13,7 @@ import java.util.Collections; import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Queue; import java.util.Set; import java.util.UUID; @@ -20,6 +21,9 @@ import java.util.stream.Collectors; import org.whispersystems.textsecuregcm.util.AttributeValues; import org.whispersystems.textsecuregcm.util.Pair; import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest; +import software.amazon.awssdk.services.dynamodb.model.GetItemRequest; +import software.amazon.awssdk.services.dynamodb.model.GetItemResponse; import software.amazon.awssdk.services.dynamodb.model.AttributeValue; import software.amazon.awssdk.services.dynamodb.model.BatchGetItemRequest; import software.amazon.awssdk.services.dynamodb.model.BatchGetItemResponse; @@ -51,14 +55,31 @@ public class DeletedAccounts extends AbstractDynamoDbStore { this.needsReconciliationIndexName = needsReconciliationIndexName; } - void put(UUID uuid, String e164) { + void put(UUID uuid, String e164, boolean needsReconciliation) { db().putItem(PutItemRequest.builder() .tableName(tableName) .item(Map.of( KEY_ACCOUNT_E164, AttributeValues.fromString(e164), ATTR_ACCOUNT_UUID, AttributeValues.fromUUID(uuid), ATTR_EXPIRES, AttributeValues.fromLong(Instant.now().plus(TIME_TO_LIVE).getEpochSecond()), - ATTR_NEEDS_CDS_RECONCILIATION, AttributeValues.fromInt(1))) + ATTR_NEEDS_CDS_RECONCILIATION, AttributeValues.fromInt(needsReconciliation ? 1 : 0))) + .build()); + } + + Optional findUuid(final String e164) { + final GetItemResponse response = db().getItem(GetItemRequest.builder() + .tableName(tableName) + .consistentRead(true) + .key(Map.of(KEY_ACCOUNT_E164, AttributeValues.fromString(e164))) + .build()); + + return Optional.ofNullable(AttributeValues.getUUID(response.item(), ATTR_ACCOUNT_UUID, null)); + } + + void remove(final String e164) { + db().deleteItem(DeleteItemRequest.builder() + .tableName(tableName) + .key(Map.of(KEY_ACCOUNT_E164, AttributeValues.fromString(e164))) .build()); } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/DeletedAccountsManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/DeletedAccountsManager.java index 0ed64cf6b..6a3d05b67 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/DeletedAccountsManager.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/DeletedAccountsManager.java @@ -15,9 +15,11 @@ import com.amazonaws.services.dynamodbv2.model.LockCurrentlyUnavailableException import java.util.ArrayList; import java.util.Collection; import java.util.List; +import java.util.Optional; import java.util.Set; import java.util.UUID; import java.util.concurrent.TimeUnit; +import java.util.function.Consumer; import java.util.stream.Collectors; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -58,8 +60,19 @@ public class DeletedAccountsManager { .build()); } - public void put(final UUID uuid, final String e164) throws InterruptedException { - withLock(e164, () -> deletedAccounts.put(uuid, e164)); + public void addRecentlyDeletedAccount(final UUID uuid, final String e164) throws InterruptedException { + withLock(e164, () -> deletedAccounts.put(uuid, e164, true)); + } + + public void lockAndTake(final String e164, final Consumer> consumer) throws InterruptedException { + withLock(e164, () -> { + try { + consumer.accept(deletedAccounts.findUuid(e164)); + deletedAccounts.remove(e164); + } catch (final Exception e) { + log.warn("Consumer threw an exception while holding lock on a deleted account record", e); + } + }); } private void withLock(final String e164, final Runnable task) throws InterruptedException { diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/DeletedAccountsManagerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/DeletedAccountsManagerTest.java index 22f7e2d84..2661706c4 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/DeletedAccountsManagerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/DeletedAccountsManagerTest.java @@ -21,6 +21,7 @@ import java.lang.Thread.State; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.UUID; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; @@ -59,11 +60,12 @@ class DeletedAccountsManagerTest { .attributeType(ScalarAttributeType.S).build()) .build(); + private DeletedAccounts deletedAccounts; private DeletedAccountsManager deletedAccountsManager; @BeforeEach void setUp() { - final DeletedAccounts deletedAccounts = new DeletedAccounts(DELETED_ACCOUNTS_DYNAMODB_EXTENSION.getDynamoDbClient(), + deletedAccounts = new DeletedAccounts(DELETED_ACCOUNTS_DYNAMODB_EXTENSION.getDynamoDbClient(), DELETED_ACCOUNTS_DYNAMODB_EXTENSION.getTableName(), NEEDS_RECONCILIATION_INDEX_NAME); @@ -72,6 +74,31 @@ class DeletedAccountsManagerTest { DELETED_ACCOUNTS_LOCK_DYNAMODB_EXTENSION.getTableName()); } + @Test + void testLockAndTake() throws InterruptedException { + final UUID uuid = UUID.randomUUID(); + final String e164 = "+18005551234"; + + deletedAccountsManager.addRecentlyDeletedAccount(uuid, e164); + deletedAccountsManager.lockAndTake(e164, maybeUuid -> assertEquals(Optional.of(uuid), maybeUuid)); + assertEquals(Optional.empty(), deletedAccounts.findUuid(e164)); + } + + @Test + void testLockAndTakeWithException() throws InterruptedException { + final UUID uuid = UUID.randomUUID(); + final String e164 = "+18005551234"; + + deletedAccountsManager.addRecentlyDeletedAccount(uuid, e164); + + deletedAccountsManager.lockAndTake(e164, maybeUuid -> { + assertEquals(Optional.of(uuid), maybeUuid); + throw new RuntimeException("OH NO"); + }); + + assertEquals(Optional.of(uuid), deletedAccounts.findUuid(e164)); + } + @Test void testReconciliationLockContention() throws ChunkProcessingFailedException, InterruptedException { @@ -86,7 +113,7 @@ class DeletedAccountsManagerTest { final Map expectedReconciledAccounts = new HashMap<>(); for (int i = 0; i < uuids.length; i++) { - deletedAccountsManager.put(uuids[i], e164s[i]); + deletedAccountsManager.addRecentlyDeletedAccount(uuids[i], e164s[i]); expectedReconciledAccounts.put(e164s[i], uuids[i]); } @@ -95,7 +122,7 @@ class DeletedAccountsManagerTest { final Thread putThread = new Thread(() -> { try { - deletedAccountsManager.put(replacedUUID, e164s[0]); + deletedAccountsManager.addRecentlyDeletedAccount(replacedUUID, e164s[0]); } catch (InterruptedException e) { throw new RuntimeException(e); } diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/DeletedAccountsTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/DeletedAccountsTest.java index e51629408..3360d004d 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/DeletedAccountsTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/DeletedAccountsTest.java @@ -10,6 +10,7 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import java.util.Collections; import java.util.HashSet; import java.util.List; +import java.util.Optional; import java.util.Set; import java.util.UUID; import org.junit.jupiter.api.BeforeEach; @@ -70,9 +71,9 @@ class DeletedAccountsTest { assertTrue(deletedAccounts.listAccountsToReconcile(1).isEmpty()); - deletedAccounts.put(firstUuid, firstNumber); - deletedAccounts.put(secondUuid, secondNumber); - deletedAccounts.put(thirdUuid, thirdNumber); + deletedAccounts.put(firstUuid, firstNumber, true); + deletedAccounts.put(secondUuid, secondNumber, true); + deletedAccounts.put(thirdUuid, thirdNumber, true); assertEquals(1, deletedAccounts.listAccountsToReconcile(1).size()); @@ -90,6 +91,34 @@ class DeletedAccountsTest { assertTrue(deletedAccounts.listAccountsToReconcile(1).isEmpty()); } + @Test + void testPutFind() { + final UUID uuid = UUID.randomUUID(); + final String e164 = "+18005551234"; + + assertEquals(Optional.empty(), deletedAccounts.findUuid(e164)); + + deletedAccounts.put(uuid, e164, true); + + assertEquals(Optional.of(uuid), deletedAccounts.findUuid(e164)); + } + + @Test + void testRemove() { + final UUID uuid = UUID.randomUUID(); + final String e164 = "+18005551234"; + + assertEquals(Optional.empty(), deletedAccounts.findUuid(e164)); + + deletedAccounts.put(uuid, e164, true); + + assertEquals(Optional.of(uuid), deletedAccounts.findUuid(e164)); + + deletedAccounts.remove(e164); + + assertEquals(Optional.empty(), deletedAccounts.findUuid(e164)); + } + @Test void testGetAccountsNeedingReconciliation() { final UUID firstUuid = UUID.randomUUID(); @@ -102,8 +131,8 @@ class DeletedAccountsTest { assertEquals(Collections.emptySet(), deletedAccounts.getAccountsNeedingReconciliation(List.of(firstNumber, secondNumber, thirdNumber))); - deletedAccounts.put(firstUuid, firstNumber); - deletedAccounts.put(secondUuid, secondNumber); + deletedAccounts.put(firstUuid, firstNumber, true); + deletedAccounts.put(secondUuid, secondNumber, true); assertEquals(Set.of(firstNumber, secondNumber), deletedAccounts.getAccountsNeedingReconciliation(List.of(firstNumber, secondNumber, thirdNumber))); @@ -118,7 +147,7 @@ class DeletedAccountsTest { for (int i = 0; i < itemCount; i++) { final String e164 = String.format("+18000555%04d", i); - deletedAccounts.put(UUID.randomUUID(), e164); + deletedAccounts.put(UUID.randomUUID(), e164, true); expectedAccountsNeedingReconciliation.add(e164); }