MirrorMirror
/
Signal-Server
mirror of https://github.com/signalapp/Signal-Server.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add methods for getting, clearing, locking recently-deleted account records.

This commit is contained in:
Jon Chambers 2021-07-23 17:09:28 -04:00 committed by Jon Chambers
parent 1fd1207bf6
commit d09dcc90fe
5 changed files with 104 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java
View File

@ -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);

25
service/src/main/java/org/whispersystems/textsecuregcm/storage/DeletedAccounts.java
View File

@ -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<UUID> 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());
}

17
service/src/main/java/org/whispersystems/textsecuregcm/storage/DeletedAccountsManager.java
View File

@ -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<Optional<UUID>> 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 {

33
service/src/test/java/org/whispersystems/textsecuregcm/storage/DeletedAccountsManagerTest.java
View File

@ -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<String, UUID> 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);
}

41
service/src/test/java/org/whispersystems/textsecuregcm/storage/DeletedAccountsTest.java
View File

@ -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);
}