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

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

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

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 {