Signal/Signal-Server
1
0
Fork 0
mirror of https://github.com/signalapp/Signal-Server synced 2026-04-20 22:48:03 +01:00
Code Issues Packages Projects Releases Wiki Activity

Switch DynamoDB to AWSv2.

Browse Source 
Switch from using com.amazonaws.services.dynamodbv2 to using
software.amazon.awssdk.services.dynamodb for all current DynamoDB uses.
This commit is contained in:
Graeme Connell
2021-05-24 16:43:56 -06:00
committed by gram-signal
parent cbd9681e3e
commit c545cff1b3
31 changed files with 1114 additions and 876 deletions
Download Patch File Download Diff File Expand all files Collapse all files

94
service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
View File

@@ -12,11 +12,6 @@ import com.amazonaws.auth.AWSCredentialsProvider;
import com.amazonaws.auth.AWSStaticCredentialsProvider;
import com.amazonaws.auth.BasicAWSCredentials;
import com.amazonaws.auth.InstanceProfileCredentialsProvider;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBAsync;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBAsyncClientBuilder;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClientBuilder;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.s3.AmazonS3;
import com.amazonaws.services.s3.AmazonS3Client;
import com.codahale.metrics.SharedMetricRegistries;
@@ -194,6 +189,7 @@ import org.whispersystems.textsecuregcm.storage.Usernames;
import org.whispersystems.textsecuregcm.storage.UsernamesManager;
import org.whispersystems.textsecuregcm.util.AsnManager;
import org.whispersystems.textsecuregcm.util.Constants;
import org.whispersystems.textsecuregcm.util.DynamoDbFromConfig;
import org.whispersystems.textsecuregcm.util.TorExitNodeManager;
import org.whispersystems.textsecuregcm.websocket.AuthenticatedConnectListener;
import org.whispersystems.textsecuregcm.websocket.DeadLetterHandler;
@@ -210,6 +206,14 @@ import org.whispersystems.textsecuregcm.workers.VacuumCommand;
import org.whispersystems.textsecuregcm.workers.ZkParamsCommand;
import org.whispersystems.websocket.WebSocketResourceProviderFactory;
import org.whispersystems.websocket.setup.WebSocketEnvironment;
import software.amazon.awssdk.core.client.config.ClientAsyncConfiguration;
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
import software.amazon.awssdk.core.client.config.SdkAdvancedAsyncClientOption;
import software.amazon.awssdk.http.SdkHttpClient;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.DynamoDbClientBuilder;
public class WhisperServerService extends Application<WhisperServerConfiguration> {
@@ -316,80 +320,40 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
FaultTolerantDatabase accountDatabase = new FaultTolerantDatabase("accounts_database", accountJdbi, config.getAccountsDatabaseConfiguration().getCircuitBreakerConfiguration());
FaultTolerantDatabase abuseDatabase = new FaultTolerantDatabase("abuse_database", abuseJdbi, config.getAbuseDatabaseConfiguration().getCircuitBreakerConfiguration());
AmazonDynamoDBClientBuilder messageDynamoDbClientBuilder = AmazonDynamoDBClientBuilder
.standard()
.withRegion(config.getMessageDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) config.getMessageDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) config.getMessageDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
DynamoDbClient messageDynamoDb = DynamoDbFromConfig.client(config.getMessageDynamoDbConfiguration(),
software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
AmazonDynamoDBClientBuilder keysDynamoDbClientBuilder = AmazonDynamoDBClientBuilder
.standard()
.withRegion(config.getKeysDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) config.getKeysDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) config.getKeysDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
DynamoDbClient preKeyDynamoDb = DynamoDbFromConfig.client(config.getKeysDynamoDbConfiguration(),
software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
AmazonDynamoDBClientBuilder accountsDynamoDbClientBuilder = AmazonDynamoDBClientBuilder
.standard()
.withRegion(config.getAccountsDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) config.getAccountsDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) config.getAccountsDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
DynamoDbClient accountsDynamoDbClient = DynamoDbFromConfig.client(config.getAccountsDynamoDbConfiguration(),
software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
// The thread pool core & max sizes are set via dynamic configuration within AccountsDynamoDb
ThreadPoolExecutor accountsDynamoDbMigrationThreadPool = new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS,
new LinkedBlockingDeque<>());
AmazonDynamoDBAsyncClientBuilder accountsDynamoDbAsyncClientBuilder = AmazonDynamoDBAsyncClientBuilder
.standard()
.withRegion(accountsDynamoDbClientBuilder.getRegion())
.withClientConfiguration(accountsDynamoDbClientBuilder.getClientConfiguration())
.withCredentials(accountsDynamoDbClientBuilder.getCredentials())
.withExecutorFactory(() -> accountsDynamoDbMigrationThreadPool);
DynamoDbAsyncClient accountsDynamoDbAsyncClient = DynamoDbFromConfig.asyncClient(config.getAccountsDynamoDbConfiguration(),
software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create(),
accountsDynamoDbMigrationThreadPool);
AmazonDynamoDBClientBuilder migrationDeletedAccountsDynamoDbClientBuilder = AmazonDynamoDBClientBuilder
.standard()
.withRegion(config.getMigrationDeletedAccountsDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) config.getMigrationDeletedAccountsDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) config.getMigrationDeletedAccountsDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
DynamoDbClient recentlyDeletedAccountsDynamoDb = DynamoDbFromConfig.client(config.getMigrationDeletedAccountsDynamoDbConfiguration(),
software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
AmazonDynamoDBClientBuilder migrationRetryAccountsDynamoDbClientBuilder = AmazonDynamoDBClientBuilder
.standard()
.withRegion(config.getMigrationRetryAccountsDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) config.getMigrationRetryAccountsDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) config.getMigrationRetryAccountsDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
DynamoDbClient pushChallengeDynamoDbClient = DynamoDbFromConfig.client(config.getPushChallengeDynamoDbConfiguration(),
software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
AmazonDynamoDBClientBuilder pushChallengeDynamoDbClientBuilder = AmazonDynamoDBClientBuilder
.standard()
.withRegion(config.getPushChallengeDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) config.getPushChallengeDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) config.getPushChallengeDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
DynamoDbClient reportMessageDynamoDbClient = DynamoDbFromConfig.client(config.getReportMessageDynamoDbConfiguration(),
software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
AmazonDynamoDBClientBuilder reportMessageDynamoDbClientBuilder = AmazonDynamoDBClientBuilder
.standard()
.withRegion(config.getReportMessageDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) config.getReportMessageDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) config.getReportMessageDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
DynamoDB messageDynamoDb = new DynamoDB(messageDynamoDbClientBuilder.build());
DynamoDB preKeyDynamoDb = new DynamoDB(keysDynamoDbClientBuilder.build());
AmazonDynamoDB accountsDynamoDbClient = accountsDynamoDbClientBuilder.build();
AmazonDynamoDBAsync accountsDynamodbAsyncClient = accountsDynamoDbAsyncClientBuilder.build();
DynamoDB recentlyDeletedAccountsDynamoDb = new DynamoDB(migrationDeletedAccountsDynamoDbClientBuilder.build());
DynamoDB migrationRetryAccountsDynamoDb = new DynamoDB(migrationRetryAccountsDynamoDbClientBuilder.build());
DynamoDbClient migrationRetryAccountsDynamoDb = DynamoDbFromConfig.client(config.getMigrationRetryAccountsDynamoDbConfiguration(),
software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
MigrationDeletedAccounts migrationDeletedAccounts = new MigrationDeletedAccounts(recentlyDeletedAccountsDynamoDb, config.getMigrationDeletedAccountsDynamoDbConfiguration().getTableName());
MigrationRetryAccounts migrationRetryAccounts = new MigrationRetryAccounts(migrationRetryAccountsDynamoDb, config.getMigrationRetryAccountsDynamoDbConfiguration().getTableName());
Accounts accounts = new Accounts(accountDatabase);
AccountsDynamoDb accountsDynamoDb = new AccountsDynamoDb(accountsDynamoDbClient, accountsDynamodbAsyncClient, accountsDynamoDbMigrationThreadPool, new DynamoDB(accountsDynamoDbClient), config.getAccountsDynamoDbConfiguration().getTableName(), config.getAccountsDynamoDbConfiguration().getPhoneNumberTableName(), migrationDeletedAccounts, migrationRetryAccounts);
AccountsDynamoDb accountsDynamoDb = new AccountsDynamoDb(accountsDynamoDbClient, accountsDynamoDbAsyncClient, accountsDynamoDbMigrationThreadPool, config.getAccountsDynamoDbConfiguration().getTableName(), config.getAccountsDynamoDbConfiguration().getPhoneNumberTableName(), migrationDeletedAccounts, migrationRetryAccounts);
PendingAccounts pendingAccounts = new PendingAccounts(accountDatabase);
PendingDevices pendingDevices = new PendingDevices (accountDatabase);
Usernames usernames = new Usernames(accountDatabase);
@@ -399,8 +363,8 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(messageDynamoDb, config.getMessageDynamoDbConfiguration().getTableName(), config.getMessageDynamoDbConfiguration().getTimeToLive());
AbusiveHostRules abusiveHostRules = new AbusiveHostRules(abuseDatabase);
RemoteConfigs remoteConfigs = new RemoteConfigs(accountDatabase);
PushChallengeDynamoDb pushChallengeDynamoDb = new PushChallengeDynamoDb(new DynamoDB(pushChallengeDynamoDbClientBuilder.build()), config.getPushChallengeDynamoDbConfiguration().getTableName());
ReportMessageDynamoDb reportMessageDynamoDb = new ReportMessageDynamoDb(new DynamoDB(reportMessageDynamoDbClientBuilder.build()), config.getReportMessageDynamoDbConfiguration().getTableName());
PushChallengeDynamoDb pushChallengeDynamoDb = new PushChallengeDynamoDb(pushChallengeDynamoDbClient, config.getPushChallengeDynamoDbConfiguration().getTableName());
ReportMessageDynamoDb reportMessageDynamoDb = new ReportMessageDynamoDb(reportMessageDynamoDbClient, config.getReportMessageDynamoDbConfiguration().getTableName());
RedisClientFactory pubSubClientFactory = new RedisClientFactory("pubsub_cache", config.getPubsubCacheConfiguration().getUrl(), config.getPubsubCacheConfiguration().getReplicaUrls(), config.getPubsubCacheConfiguration().getCircuitBreakerConfiguration());
ReplicatedJedisPool pubsubClient = pubSubClientFactory.getRedisClientPool();

58
service/src/main/java/org/whispersystems/textsecuregcm/storage/AbstractDynamoDbStore.java
View File

@@ -5,21 +5,18 @@
package org.whispersystems.textsecuregcm.storage;
import com.amazonaws.services.dynamodbv2.document.BatchWriteItemOutcome;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.document.Item;
import com.amazonaws.services.dynamodbv2.document.Page;
import com.amazonaws.services.dynamodbv2.document.QueryOutcome;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.TableWriteItems;
import com.amazonaws.services.dynamodbv2.document.spec.QuerySpec;
import io.micrometer.core.instrument.Counter;
import io.micrometer.core.instrument.Timer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.BatchWriteItemRequest;
import software.amazon.awssdk.services.dynamodb.model.BatchWriteItemResponse;
import software.amazon.awssdk.services.dynamodb.model.WriteRequest;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;
@@ -29,7 +26,7 @@ import static io.micrometer.core.instrument.Metrics.timer;
public class AbstractDynamoDbStore {
private final DynamoDB dynamoDb;
private final DynamoDbClient dynamoDbClient;
private final Timer batchWriteItemsFirstPass = timer(name(getClass(), "batchWriteItems"), "firstAttempt", "true");
private final Timer batchWriteItemsRetryPass = timer(name(getClass(), "batchWriteItems"), "firstAttempt", "false");
@@ -41,44 +38,31 @@ public class AbstractDynamoDbStore {
public static final int DYNAMO_DB_MAX_BATCH_SIZE = 25; // This limit comes from Amazon Dynamo DB itself. It will reject batch writes larger than this.
public static final int RESULT_SET_CHUNK_SIZE = 100;
public AbstractDynamoDbStore(final DynamoDB dynamoDb) {
this.dynamoDb = dynamoDb;
public AbstractDynamoDbStore(final DynamoDbClient dynamoDbClient) {
this.dynamoDbClient = dynamoDbClient;
}
protected DynamoDB getDynamoDb() {
return dynamoDb;
protected DynamoDbClient db() {
return dynamoDbClient;
}
protected void executeTableWriteItemsUntilComplete(final TableWriteItems items) {
AtomicReference<BatchWriteItemOutcome> outcome = new AtomicReference<>();
batchWriteItemsFirstPass.record(() -> outcome.set(dynamoDb.batchWriteItem(items)));
protected void executeTableWriteItemsUntilComplete(final Map<String,List<WriteRequest>> items) {
AtomicReference<BatchWriteItemResponse> outcome = new AtomicReference<>();
batchWriteItemsFirstPass.record(() -> outcome.set(dynamoDbClient.batchWriteItem(BatchWriteItemRequest.builder().requestItems(items).build())));
int attemptCount = 0;
while (!outcome.get().getUnprocessedItems().isEmpty() && attemptCount < MAX_ATTEMPTS_TO_SAVE_BATCH_WRITE) {
batchWriteItemsRetryPass.record(() -> outcome.set(dynamoDb.batchWriteItemUnprocessed(outcome.get().getUnprocessedItems())));
while (!outcome.get().unprocessedItems().isEmpty() && attemptCount < MAX_ATTEMPTS_TO_SAVE_BATCH_WRITE) {
batchWriteItemsRetryPass.record(() -> outcome.set(dynamoDbClient.batchWriteItem(BatchWriteItemRequest.builder()
.requestItems(outcome.get().unprocessedItems())
.build())));
++attemptCount;
}
if (!outcome.get().getUnprocessedItems().isEmpty()) {
logger.error("Attempt count ({}) reached max ({}}) before applying all batch writes to dynamo. {} unprocessed items remain.", attemptCount, MAX_ATTEMPTS_TO_SAVE_BATCH_WRITE, outcome.get().getUnprocessedItems().size());
batchWriteItemsUnprocessed.increment(outcome.get().getUnprocessedItems().size());
if (!outcome.get().unprocessedItems().isEmpty()) {
int totalItems = outcome.get().unprocessedItems().values().stream().mapToInt(List::size).sum();
logger.error("Attempt count ({}) reached max ({}}) before applying all batch writes to dynamo. {} unprocessed items remain.", attemptCount, MAX_ATTEMPTS_TO_SAVE_BATCH_WRITE, totalItems);
batchWriteItemsUnprocessed.increment(totalItems);
}
}
protected long countItemsMatchingQuery(final Table table, final QuerySpec querySpec) {
// This is very confusing, but does appear to be the intended behavior. See:
//
// - https://github.com/aws/aws-sdk-java/issues/693
// - https://github.com/aws/aws-sdk-java/issues/915
// - https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Query.html#Query.Count
long matchingItems = 0;
for (final Page<Item, QueryOutcome> page : table.query(querySpec).pages()) {
matchingItems += page.getLowLevelResult().getQueryResult().getCount();
}
return matchingItems;
}
static <T> void writeInBatches(final Iterable<T> items, final Consumer<List<T>> action) {
final List<T> batch = new ArrayList<>(DYNAMO_DB_MAX_BATCH_SIZE);

286
service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsDynamoDb.java
View File

@@ -2,25 +2,6 @@ package org.whispersystems.textsecuregcm.storage;
import static com.codahale.metrics.MetricRegistry.name;
import com.amazonaws.handlers.AsyncHandler;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBAsync;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.document.Item;
import com.amazonaws.services.dynamodbv2.document.PrimaryKey;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.spec.GetItemSpec;
import com.amazonaws.services.dynamodbv2.model.AttributeValue;
import com.amazonaws.services.dynamodbv2.model.CancellationReason;
import com.amazonaws.services.dynamodbv2.model.Delete;
import com.amazonaws.services.dynamodbv2.model.GetItemResult;
import com.amazonaws.services.dynamodbv2.model.Put;
import com.amazonaws.services.dynamodbv2.model.ReturnValuesOnConditionCheckFailure;
import com.amazonaws.services.dynamodbv2.model.TransactWriteItem;
import com.amazonaws.services.dynamodbv2.model.TransactWriteItemsRequest;
import com.amazonaws.services.dynamodbv2.model.TransactWriteItemsResult;
import com.amazonaws.services.dynamodbv2.model.TransactionCanceledException;
import com.amazonaws.services.dynamodbv2.model.UpdateItemRequest;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.google.common.annotations.VisibleForTesting;
import io.micrometer.core.instrument.Counter;
@@ -33,12 +14,27 @@ import java.util.Map;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.whispersystems.textsecuregcm.util.AttributeValues;
import org.whispersystems.textsecuregcm.util.SystemMapper;
import org.whispersystems.textsecuregcm.util.UUIDUtil;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.CancellationReason;
import software.amazon.awssdk.services.dynamodb.model.Delete;
import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
import software.amazon.awssdk.services.dynamodb.model.Put;
import software.amazon.awssdk.services.dynamodb.model.ReturnValuesOnConditionCheckFailure;
import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem;
import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsRequest;
import software.amazon.awssdk.services.dynamodb.model.TransactionCanceledException;
import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountStore {
@@ -51,9 +47,8 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
static final String ATTR_MIGRATION_VERSION = "V";
private final AmazonDynamoDB client;
private final Table accountsTable;
private final AmazonDynamoDBAsync asyncClient;
private final DynamoDbClient client;
private final DynamoDbAsyncClient asyncClient;
private final ThreadPoolExecutor migrationThreadPool;
@@ -61,6 +56,7 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
private final MigrationRetryAccounts migrationRetryAccounts;
private final String phoneNumbersTableName;
private final String accountsTableName;
private static final Timer CREATE_TIMER = Metrics.timer(name(AccountsDynamoDb.class, "create"));
private static final Timer UPDATE_TIMER = Metrics.timer(name(AccountsDynamoDb.class, "update"));
@@ -70,18 +66,17 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
private final Logger logger = LoggerFactory.getLogger(AccountsDynamoDb.class);
public AccountsDynamoDb(AmazonDynamoDB client, AmazonDynamoDBAsync asyncClient,
ThreadPoolExecutor migrationThreadPool, DynamoDB dynamoDb, String accountsTableName, String phoneNumbersTableName,
public AccountsDynamoDb(DynamoDbClient client, DynamoDbAsyncClient asyncClient,
ThreadPoolExecutor migrationThreadPool, String accountsTableName, String phoneNumbersTableName,
MigrationDeletedAccounts migrationDeletedAccounts,
MigrationRetryAccounts accountsMigrationErrors) {
super(dynamoDb);
super(client);
this.client = client;
this.accountsTable = dynamoDb.getTable(accountsTableName);
this.phoneNumbersTableName = phoneNumbersTableName;
this.asyncClient = asyncClient;
this.phoneNumbersTableName = phoneNumbersTableName;
this.accountsTableName = accountsTableName;
this.migrationThreadPool = migrationThreadPool;
this.migrationDeletedAccounts = migrationDeletedAccounts;
@@ -90,32 +85,34 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
@Override
public boolean create(Account account) {
return CREATE_TIMER.record(() -> {
try {
TransactWriteItem phoneNumberConstraintPut = buildPutWriteItemForPhoneNumberConstraint(account, account.getUuid());
TransactWriteItem accountPut = buildPutWriteItemForAccount(account, account.getUuid(), Put.builder()
.conditionExpression("attribute_not_exists(#number) OR #number = :number")
.expressionAttributeNames(Map.of("#number", ATTR_ACCOUNT_E164))
.expressionAttributeValues(Map.of(":number", AttributeValues.fromString(account.getNumber()))));
TransactWriteItem accountPut = buildPutWriteItemForAccount(account, account.getUuid());
final TransactWriteItemsRequest request = new TransactWriteItemsRequest()
.withTransactItems(phoneNumberConstraintPut, accountPut);
final TransactWriteItemsRequest request = TransactWriteItemsRequest.builder()
.transactItems(phoneNumberConstraintPut, accountPut)
.build();
try {
client.transactWriteItems(request);
} catch (TransactionCanceledException e) {
final CancellationReason accountCancellationReason = e.getCancellationReasons().get(1);
final CancellationReason accountCancellationReason = e.cancellationReasons().get(1);
if ("ConditionalCheckFailed".equals(accountCancellationReason.getCode())) {
if ("ConditionalCheckFailed".equals(accountCancellationReason.code())) {
throw new IllegalArgumentException("uuid present with different phone number");
}
final CancellationReason phoneNumberConstraintCancellationReason = e.getCancellationReasons().get(0);
final CancellationReason phoneNumberConstraintCancellationReason = e.cancellationReasons().get(0);
if ("ConditionalCheckFailed".equals(phoneNumberConstraintCancellationReason.getCode())) {
if ("ConditionalCheckFailed".equals(phoneNumberConstraintCancellationReason.code())) {
ByteBuffer actualAccountUuid = phoneNumberConstraintCancellationReason.getItem().get(KEY_ACCOUNT_UUID).getB();
ByteBuffer actualAccountUuid = phoneNumberConstraintCancellationReason.item().get(KEY_ACCOUNT_UUID).b().asByteBuffer();
account.setUuid(UUIDUtil.fromByteBuffer(actualAccountUuid));
update(account);
@@ -134,39 +131,37 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
});
}
private TransactWriteItem buildPutWriteItemForAccount(Account account, UUID uuid) throws JsonProcessingException {
return new TransactWriteItem()
.withPut(
new Put()
.withTableName(accountsTable.getTableName())
.withItem(Map.of(
KEY_ACCOUNT_UUID, new AttributeValue().withB(UUIDUtil.toByteBuffer(uuid)),
ATTR_ACCOUNT_E164, new AttributeValue(account.getNumber()),
ATTR_ACCOUNT_DATA, new AttributeValue()
.withB(ByteBuffer.wrap(SystemMapper.getMapper().writeValueAsBytes(account))),
ATTR_MIGRATION_VERSION, new AttributeValue().withN(
String.valueOf(account.getDynamoDbMigrationVersion()))))
.withConditionExpression("attribute_not_exists(#number) OR #number = :number")
.withExpressionAttributeNames(Map.of("#number", ATTR_ACCOUNT_E164))
.withExpressionAttributeValues(Map.of(":number", new AttributeValue(account.getNumber()))));
private TransactWriteItem buildPutWriteItemForAccount(Account account, UUID uuid, Put.Builder putBuilder) throws JsonProcessingException {
return TransactWriteItem.builder()
.put(putBuilder
.tableName(accountsTableName)
.item(Map.of(
KEY_ACCOUNT_UUID, AttributeValues.fromUUID(uuid),
ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()),
ATTR_ACCOUNT_DATA, AttributeValues.fromByteArray(SystemMapper.getMapper().writeValueAsBytes(account)),
ATTR_MIGRATION_VERSION, AttributeValues.fromInt(account.getDynamoDbMigrationVersion())))
.build())
.build();
}
private TransactWriteItem buildPutWriteItemForPhoneNumberConstraint(Account account, UUID uuid) {
return new TransactWriteItem()
.withPut(
new Put()
.withTableName(phoneNumbersTableName)
.withItem(Map.of(
ATTR_ACCOUNT_E164, new AttributeValue(account.getNumber()),
KEY_ACCOUNT_UUID, new AttributeValue().withB(UUIDUtil.toByteBuffer(uuid))))
.withConditionExpression(
return TransactWriteItem.builder()
.put(
Put.builder()
.tableName(phoneNumbersTableName)
.item(Map.of(
ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()),
KEY_ACCOUNT_UUID, AttributeValues.fromUUID(uuid)))
.conditionExpression(
"attribute_not_exists(#number) OR (attribute_exists(#number) AND #uuid = :uuid)")
.withExpressionAttributeNames(
.expressionAttributeNames(
Map.of("#uuid", KEY_ACCOUNT_UUID,
"#number", ATTR_ACCOUNT_E164))
.withExpressionAttributeValues(
Map.of(":uuid", new AttributeValue().withB(UUIDUtil.toByteBuffer(uuid))))
.withReturnValuesOnConditionCheckFailure(ReturnValuesOnConditionCheckFailure.ALL_OLD));
.expressionAttributeValues(
Map.of(":uuid", AttributeValues.fromUUID(uuid)))
.returnValuesOnConditionCheckFailure(ReturnValuesOnConditionCheckFailure.ALL_OLD)
.build())
.build();
}
@Override
@@ -174,16 +169,18 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
UPDATE_TIMER.record(() -> {
UpdateItemRequest updateItemRequest;
try {
updateItemRequest = new UpdateItemRequest()
.withTableName(accountsTable.getTableName())
.withKey(Map.of(KEY_ACCOUNT_UUID, new AttributeValue().withB(UUIDUtil.toByteBuffer(account.getUuid()))))
.withUpdateExpression("SET #data = :data, #version = :version")
.withConditionExpression("attribute_exists(#number)")
.withExpressionAttributeNames(Map.of("#number", ATTR_ACCOUNT_E164,
updateItemRequest = UpdateItemRequest.builder()
.tableName(accountsTableName)
.key(Map.of(KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid())))
.updateExpression("SET #data = :data, #version = :version")
.conditionExpression("attribute_exists(#number)")
.expressionAttributeNames(Map.of("#number", ATTR_ACCOUNT_E164,
"#data", ATTR_ACCOUNT_DATA,
"#version", ATTR_MIGRATION_VERSION))
.withExpressionAttributeValues(Map.of(":data", new AttributeValue().withB(ByteBuffer.wrap(SystemMapper.getMapper().writeValueAsBytes(account))),
":version", new AttributeValue().withN(String.valueOf(account.getDynamoDbMigrationVersion()))));
.expressionAttributeValues(Map.of(
":data", AttributeValues.fromByteArray(SystemMapper.getMapper().writeValueAsBytes(account)),
":version", AttributeValues.fromInt(account.getDynamoDbMigrationVersion())))
.build();
} catch (JsonProcessingException e) {
throw new IllegalArgumentException(e);
@@ -193,37 +190,42 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
});
}
@Override
public Optional<Account> get(String number) {
return GET_BY_NUMBER_TIMER.record(() -> {
final GetItemResult phoneNumberAndUuid = client.getItem(phoneNumbersTableName,
Map.of(ATTR_ACCOUNT_E164, new AttributeValue(number)), true);
final GetItemResponse response = client.getItem(GetItemRequest.builder()
.tableName(phoneNumbersTableName)
.key(Map.of(ATTR_ACCOUNT_E164, AttributeValues.fromString(number)))
.build());
return Optional.ofNullable(phoneNumberAndUuid.getItem())
.map(item -> item.get(KEY_ACCOUNT_UUID).getB())
.map(uuid -> accountsTable.getItem(new GetItemSpec()
.withPrimaryKey(KEY_ACCOUNT_UUID, uuid.array())
.withConsistentRead(true)))
return Optional.ofNullable(response.item())
.map(item -> item.get(KEY_ACCOUNT_UUID))
.map(uuid -> accountByUuid(uuid))
.map(AccountsDynamoDb::fromItem);
});
}
private Map<String, AttributeValue> accountByUuid(AttributeValue uuid) {
GetItemResponse r = client.getItem(GetItemRequest.builder()
.tableName(accountsTableName)
.key(Map.of(KEY_ACCOUNT_UUID, uuid))
.consistentRead(true)
.build());
return r.item().isEmpty() ? null : r.item();
}
@Override
public Optional<Account> get(UUID uuid) {
Optional<Item> maybeItem = GET_BY_UUID_TIMER.record(() ->
Optional.ofNullable(accountsTable.getItem(new GetItemSpec().
withPrimaryKey(new PrimaryKey(KEY_ACCOUNT_UUID, UUIDUtil.toByteBuffer(uuid)))
.withConsistentRead(true))));
return maybeItem.map(AccountsDynamoDb::fromItem);
return GET_BY_UUID_TIMER.record(() ->
Optional.ofNullable(accountByUuid(AttributeValues.fromUUID(uuid)))
.map(AccountsDynamoDb::fromItem));
}
@Override
public void delete(UUID uuid) {
DELETE_TIMER.record(() -> {
delete(uuid, true);
});
}
@@ -238,18 +240,22 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
maybeAccount.ifPresent(account -> {
TransactWriteItem phoneNumberDelete = new TransactWriteItem()
.withDelete(new Delete()
.withTableName(phoneNumbersTableName)
.withKey(Map.of(ATTR_ACCOUNT_E164, new AttributeValue(account.getNumber()))));
TransactWriteItem phoneNumberDelete = TransactWriteItem.builder()
.delete(Delete.builder()
.tableName(phoneNumbersTableName)
.key(Map.of(ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber())))
.build())
.build();
TransactWriteItem accountDelete = new TransactWriteItem().withDelete(
new Delete()
.withTableName(accountsTable.getTableName())
.withKey(Map.of(KEY_ACCOUNT_UUID, new AttributeValue().withB(UUIDUtil.toByteBuffer(uuid)))));
TransactWriteItem accountDelete = TransactWriteItem.builder()
.delete(Delete.builder()
.tableName(accountsTableName)
.key(Map.of(KEY_ACCOUNT_UUID, AttributeValues.fromUUID(uuid)))
.build())
.build();
TransactWriteItemsRequest request = new TransactWriteItemsRequest()
.withTransactItems(phoneNumberDelete, accountDelete);
TransactWriteItemsRequest request = TransactWriteItemsRequest.builder()
.transactItems(phoneNumberDelete, accountDelete).build();
client.transactWriteItems(request);
});
@@ -299,64 +305,62 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
try {
TransactWriteItem phoneNumberConstraintPut = buildPutWriteItemForPhoneNumberConstraint(account, account.getUuid());
TransactWriteItem accountPut = buildPutWriteItemForAccount(account, account.getUuid());
accountPut.getPut()
.setConditionExpression("attribute_not_exists(#uuid) OR (attribute_exists(#uuid) AND #version < :version)");
accountPut.getPut()
.setExpressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID,
"#version", ATTR_MIGRATION_VERSION));
accountPut.getPut()
.setExpressionAttributeValues(
Map.of(":version", new AttributeValue().withN(String.valueOf(account.getDynamoDbMigrationVersion()))));
TransactWriteItem accountPut = buildPutWriteItemForAccount(account, account.getUuid(), Put.builder()
.conditionExpression("attribute_not_exists(#uuid) OR (attribute_exists(#uuid) AND #version < :version)")
.expressionAttributeNames(Map.of(
"#uuid", KEY_ACCOUNT_UUID,
"#version", ATTR_MIGRATION_VERSION))
.expressionAttributeValues(Map.of(
":version", AttributeValues.fromInt(account.getDynamoDbMigrationVersion()))));
final TransactWriteItemsRequest request = new TransactWriteItemsRequest()
.withTransactItems(phoneNumberConstraintPut, accountPut);
final TransactWriteItemsRequest request = TransactWriteItemsRequest.builder()
.transactItems(phoneNumberConstraintPut, accountPut).build();
final CompletableFuture<Boolean> resultFuture = new CompletableFuture<>();
asyncClient.transactWriteItemsAsync(request,
new AsyncHandler<>() {
@Override
public void onError(Exception exception) {
if (exception instanceof TransactionCanceledException) {
// account is already migrated
resultFuture.complete(false);
} else {
try {
migrationRetryAccounts.put(account.getUuid());
} catch (final Exception e) {
logger.error("Could not store account {}", account.getUuid());
}
resultFuture.completeExceptionally(exception);
}
}
@Override
public void onSuccess(TransactWriteItemsRequest request, TransactWriteItemsResult transactWriteItemsResult) {
resultFuture.complete(true);
}
});
asyncClient.transactWriteItems(request).whenCompleteAsync((result, exception) -> {
if (result != null) {
resultFuture.complete(true);
return;
}
if (exception instanceof CompletionException) {
// whenCompleteAsync can wrap exceptions in a CompletionException; unwrap it to get to the root cause.
exception = exception.getCause();
}
if (exception instanceof TransactionCanceledException) {
// account is already migrated
resultFuture.complete(false);
return;
}
try {
migrationRetryAccounts.put(account.getUuid());
} catch (final Exception e) {
logger.error("Could not store account {}", account.getUuid());
}
resultFuture.completeExceptionally(exception);
});
return resultFuture;
} catch (Exception e) {
return CompletableFuture.failedFuture(e);
}
}
private static String extractCancellationReasonCodes(final TransactionCanceledException exception) {
return exception.getCancellationReasons().stream()
.map(CancellationReason::getCode)
return exception.cancellationReasons().stream()
.map(CancellationReason::code)
.collect(Collectors.joining(", "));
}
@VisibleForTesting
static Account fromItem(Item item) {
static Account fromItem(Map<String, AttributeValue> item) {
if (!item.containsKey(ATTR_ACCOUNT_DATA) ||
!item.containsKey(ATTR_ACCOUNT_E164) ||
!item.containsKey(KEY_ACCOUNT_UUID)) {
throw new RuntimeException("item missing values");
}
try {
Account account = SystemMapper.getMapper().readValue(item.getBinary(ATTR_ACCOUNT_DATA), Account.class);
account.setNumber(item.getString(ATTR_ACCOUNT_E164));
account.setUuid(UUIDUtil.fromByteBuffer(item.getByteBuffer(KEY_ACCOUNT_UUID)));
Account account = SystemMapper.getMapper().readValue(item.get(ATTR_ACCOUNT_DATA).b().asByteArray(), Account.class);
account.setNumber(item.get(ATTR_ACCOUNT_E164).s());
account.setUuid(UUIDUtil.fromByteBuffer(item.get(KEY_ACCOUNT_UUID).b().asByteBuffer()));
return account;

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

@@ -7,7 +7,6 @@ package org.whispersystems.textsecuregcm.storage;
import static com.codahale.metrics.MetricRegistry.name;
import com.amazonaws.services.dynamodbv2.model.ConditionalCheckFailedException;
import com.codahale.metrics.MetricRegistry;
import com.codahale.metrics.SharedMetricRegistries;
import com.codahale.metrics.Timer;
@@ -42,6 +41,7 @@ import org.whispersystems.textsecuregcm.sqs.DirectoryQueue;
import org.whispersystems.textsecuregcm.util.Constants;
import org.whispersystems.textsecuregcm.util.SystemMapper;
import org.whispersystems.textsecuregcm.util.Util;
import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException;
public class AccountsManager {

337
service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDb.java
View File

@@ -7,195 +7,230 @@ package org.whispersystems.textsecuregcm.storage;
import static com.codahale.metrics.MetricRegistry.name;
import com.amazonaws.services.dynamodbv2.document.DeleteItemOutcome;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.document.Item;
import com.amazonaws.services.dynamodbv2.document.PrimaryKey;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.TableWriteItems;
import com.amazonaws.services.dynamodbv2.document.spec.DeleteItemSpec;
import com.amazonaws.services.dynamodbv2.document.spec.QuerySpec;
import com.amazonaws.services.dynamodbv2.model.ReturnValue;
import com.amazonaws.services.dynamodbv2.model.Select;
import com.google.common.annotations.VisibleForTesting;
import io.micrometer.core.instrument.DistributionSummary;
import io.micrometer.core.instrument.Metrics;
import io.micrometer.core.instrument.Timer;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;
import org.whispersystems.textsecuregcm.entities.PreKey;
import org.whispersystems.textsecuregcm.util.UUIDUtil;
import org.whispersystems.textsecuregcm.util.AttributeValues;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
import software.amazon.awssdk.services.dynamodb.model.DeleteItemResponse;
import software.amazon.awssdk.services.dynamodb.model.DeleteRequest;
import software.amazon.awssdk.services.dynamodb.model.PutRequest;
import software.amazon.awssdk.services.dynamodb.model.QueryRequest;
import software.amazon.awssdk.services.dynamodb.model.QueryResponse;
import software.amazon.awssdk.services.dynamodb.model.ReturnValue;
import software.amazon.awssdk.services.dynamodb.model.Select;
import software.amazon.awssdk.services.dynamodb.model.WriteRequest;
public class KeysDynamoDb extends AbstractDynamoDbStore {
private final Table table;
private final String tableName;
static final String KEY_ACCOUNT_UUID = "U";
static final String KEY_DEVICE_ID_KEY_ID = "DK";
static final String KEY_PUBLIC_KEY = "P";
static final String KEY_ACCOUNT_UUID = "U";
static final String KEY_DEVICE_ID_KEY_ID = "DK";
static final String KEY_PUBLIC_KEY = "P";
private static final Timer STORE_KEYS_TIMER = Metrics.timer(name(KeysDynamoDb.class, "storeKeys"));
private static final Timer TAKE_KEY_FOR_DEVICE_TIMER = Metrics.timer(name(KeysDynamoDb.class, "takeKeyForDevice"));
private static final Timer TAKE_KEYS_FOR_ACCOUNT_TIMER = Metrics.timer(name(KeysDynamoDb.class, "takeKeyForAccount"));
private static final Timer GET_KEY_COUNT_TIMER = Metrics.timer(name(KeysDynamoDb.class, "getKeyCount"));
private static final Timer DELETE_KEYS_FOR_DEVICE_TIMER = Metrics.timer(name(KeysDynamoDb.class, "deleteKeysForDevice"));
private static final Timer DELETE_KEYS_FOR_ACCOUNT_TIMER = Metrics.timer(name(KeysDynamoDb.class, "deleteKeysForAccount"));
private static final DistributionSummary CONTESTED_KEY_DISTRIBUTION = Metrics.summary(name(KeysDynamoDb.class, "contestedKeys"));
private static final DistributionSummary KEY_COUNT_DISTRIBUTION = Metrics.summary(name(KeysDynamoDb.class, "keyCount"));
private static final Timer STORE_KEYS_TIMER = Metrics.timer(name(KeysDynamoDb.class, "storeKeys"));
private static final Timer TAKE_KEY_FOR_DEVICE_TIMER = Metrics.timer(name(KeysDynamoDb.class, "takeKeyForDevice"));
private static final Timer TAKE_KEYS_FOR_ACCOUNT_TIMER = Metrics.timer(name(KeysDynamoDb.class, "takeKeyForAccount"));
private static final Timer GET_KEY_COUNT_TIMER = Metrics.timer(name(KeysDynamoDb.class, "getKeyCount"));
private static final Timer DELETE_KEYS_FOR_DEVICE_TIMER = Metrics.timer(name(KeysDynamoDb.class, "deleteKeysForDevice"));
private static final Timer DELETE_KEYS_FOR_ACCOUNT_TIMER = Metrics.timer(name(KeysDynamoDb.class, "deleteKeysForAccount"));
private static final DistributionSummary CONTESTED_KEY_DISTRIBUTION = Metrics.summary(name(KeysDynamoDb.class, "contestedKeys"));
private static final DistributionSummary KEY_COUNT_DISTRIBUTION = Metrics.summary(name(KeysDynamoDb.class, "keyCount"));
public KeysDynamoDb(final DynamoDB dynamoDB, final String tableName) {
super(dynamoDB);
public KeysDynamoDb(final DynamoDbClient dynamoDB, final String tableName) {
super(dynamoDB);
this.tableName = tableName;
}
this.table = dynamoDB.getTable(tableName);
}
public void store(final Account account, final long deviceId, final List<PreKey> keys) {
STORE_KEYS_TIMER.record(() -> {
delete(account, deviceId);
public void store(final Account account, final long deviceId, final List<PreKey> keys) {
STORE_KEYS_TIMER.record(() -> {
delete(account, deviceId);
writeInBatches(keys, batch -> {
List<WriteRequest> items = new ArrayList<>();
for (final PreKey preKey : batch) {
items.add(WriteRequest.builder()
.putRequest(PutRequest.builder()
.item(getItemFromPreKey(account.getUuid(), deviceId, preKey))
.build())
.build());
}
executeTableWriteItemsUntilComplete(Map.of(tableName, items));
});
});
}
writeInBatches(keys, batch -> {
final TableWriteItems items = new TableWriteItems(table.getTableName());
public Optional<PreKey> take(final Account account, final long deviceId) {
return TAKE_KEY_FOR_DEVICE_TIMER.record(() -> {
final AttributeValue partitionKey = getPartitionKey(account.getUuid());
QueryRequest queryRequest = QueryRequest.builder()
.tableName(tableName)
.keyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
.expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID))
.expressionAttributeValues(Map.of(
":uuid", partitionKey,
":sortprefix", getSortKeyPrefix(deviceId)))
.projectionExpression(KEY_DEVICE_ID_KEY_ID)
.consistentRead(false)
.build();
for (final PreKey preKey : batch) {
items.addItemToPut(getItemFromPreKey(account.getUuid(), deviceId, preKey));
}
int contestedKeys = 0;
executeTableWriteItemsUntilComplete(items);
});
});
}
try {
QueryResponse response = db().query(queryRequest);
for (Map<String, AttributeValue> candidate : response.items()) {
DeleteItemRequest deleteItemRequest = DeleteItemRequest.builder()
.tableName(tableName)
.key(Map.of(
KEY_ACCOUNT_UUID, partitionKey,
KEY_DEVICE_ID_KEY_ID, candidate.get(KEY_DEVICE_ID_KEY_ID)))
.returnValues(ReturnValue.ALL_OLD)
.build();
DeleteItemResponse deleteItemResponse = db().deleteItem(deleteItemRequest);
if (deleteItemResponse.attributes() != null) {
return Optional.of(getPreKeyFromItem(deleteItemResponse.attributes()));
}
public Optional<PreKey> take(final Account account, final long deviceId) {
return TAKE_KEY_FOR_DEVICE_TIMER.record(() -> {
final byte[] partitionKey = getPartitionKey(account.getUuid());
contestedKeys++;
}
final QuerySpec querySpec = new QuerySpec().withKeyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
.withNameMap(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID))
.withValueMap(Map.of(":uuid", partitionKey,
":sortprefix", getSortKeyPrefix(deviceId)))
.withProjectionExpression(KEY_DEVICE_ID_KEY_ID)
.withConsistentRead(false);
return Optional.empty();
} finally {
CONTESTED_KEY_DISTRIBUTION.record(contestedKeys);
}
});
}
int contestedKeys = 0;
public Map<Long, PreKey> take(final Account account) {
return TAKE_KEYS_FOR_ACCOUNT_TIMER.record(() -> {
final Map<Long, PreKey> preKeysByDeviceId = new HashMap<>();
try {
for (final Item candidate : table.query(querySpec)) {
final DeleteItemSpec deleteItemSpec = new DeleteItemSpec().withPrimaryKey(KEY_ACCOUNT_UUID, partitionKey, KEY_DEVICE_ID_KEY_ID, candidate.getBinary(KEY_DEVICE_ID_KEY_ID))
.withReturnValues(ReturnValue.ALL_OLD);
for (final Device device : account.getDevices()) {
take(account, device.getId()).ifPresent(preKey -> preKeysByDeviceId.put(device.getId(), preKey));
}
final DeleteItemOutcome outcome = table.deleteItem(deleteItemSpec);
return preKeysByDeviceId;
});
}
if (outcome.getItem() != null) {
return Optional.of(getPreKeyFromItem(outcome.getItem()));
}
public int getCount(final Account account, final long deviceId) {
return GET_KEY_COUNT_TIMER.record(() -> {
QueryRequest queryRequest = QueryRequest.builder()
.tableName(tableName)
.keyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
.expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID))
.expressionAttributeValues(Map.of(
":uuid", getPartitionKey(account.getUuid()),
":sortprefix", getSortKeyPrefix(deviceId)))
.select(Select.COUNT)
.consistentRead(false)
.build();
contestedKeys++;
}
int keyCount = 0;
// This is very confusing, but does appear to be the intended behavior. See:
//
// - https://github.com/aws/aws-sdk-java/issues/693
// - https://github.com/aws/aws-sdk-java/issues/915
// - https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Query.html#Query.Count
for (final QueryResponse page : db().queryPaginator(queryRequest)) {
keyCount += page.count();
}
KEY_COUNT_DISTRIBUTION.record(keyCount);
return keyCount;
});
}
return Optional.empty();
} finally {
CONTESTED_KEY_DISTRIBUTION.record(contestedKeys);
}
});
}
public void delete(final Account account) {
DELETE_KEYS_FOR_ACCOUNT_TIMER.record(() -> {
final QueryRequest queryRequest = QueryRequest.builder()
.tableName(tableName)
.keyConditionExpression("#uuid = :uuid")
.expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID))
.expressionAttributeValues(Map.of(
":uuid", getPartitionKey(account.getUuid())))
.projectionExpression(KEY_DEVICE_ID_KEY_ID)
.consistentRead(true)
.build();
public Map<Long, PreKey> take(final Account account) {
return TAKE_KEYS_FOR_ACCOUNT_TIMER.record(() -> {
final Map<Long, PreKey> preKeysByDeviceId = new HashMap<>();
deleteItemsForAccountMatchingQuery(account, queryRequest);
});
}
for (final Device device : account.getDevices()) {
take(account, device.getId()).ifPresent(preKey -> preKeysByDeviceId.put(device.getId(), preKey));
}
@VisibleForTesting
void delete(final Account account, final long deviceId) {
DELETE_KEYS_FOR_DEVICE_TIMER.record(() -> {
final QueryRequest queryRequest = QueryRequest.builder()
.tableName(tableName)
.keyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
.expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID))
.expressionAttributeValues(Map.of(
":uuid", getPartitionKey(account.getUuid()),
":sortprefix", getSortKeyPrefix(deviceId)))
.projectionExpression(KEY_DEVICE_ID_KEY_ID)
.consistentRead(true)
.build();
return preKeysByDeviceId;
});
}
deleteItemsForAccountMatchingQuery(account, queryRequest);
});
}
public int getCount(final Account account, final long deviceId) {
return GET_KEY_COUNT_TIMER.record(() -> {
final QuerySpec querySpec = new QuerySpec().withKeyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
.withNameMap(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID))
.withValueMap(Map.of(":uuid", getPartitionKey(account.getUuid()),
":sortprefix", getSortKeyPrefix(deviceId)))
.withSelect(Select.COUNT)
.withConsistentRead(false);
private void deleteItemsForAccountMatchingQuery(final Account account, final QueryRequest querySpec) {
final AttributeValue partitionKey = getPartitionKey(account.getUuid());
final int keyCount = (int)countItemsMatchingQuery(table, querySpec);
writeInBatches(db().query(querySpec).items(), batch -> {
List<WriteRequest> deletes = new ArrayList<>();
for (final Map<String, AttributeValue> item : batch) {
deletes.add(WriteRequest.builder()
.deleteRequest(DeleteRequest.builder()
.key(Map.of(
KEY_ACCOUNT_UUID, partitionKey,
KEY_DEVICE_ID_KEY_ID, item.get(KEY_DEVICE_ID_KEY_ID)))
.build())
.build());
}
executeTableWriteItemsUntilComplete(Map.of(tableName, deletes));
});
}
KEY_COUNT_DISTRIBUTION.record(keyCount);
return keyCount;
});
}
private static AttributeValue getPartitionKey(final UUID accountUuid) {
return AttributeValues.fromUUID(accountUuid);
}
public void delete(final Account account) {
DELETE_KEYS_FOR_ACCOUNT_TIMER.record(() -> {
final QuerySpec querySpec = new QuerySpec().withKeyConditionExpression("#uuid = :uuid")
.withNameMap(Map.of("#uuid", KEY_ACCOUNT_UUID))
.withValueMap(Map.of(":uuid", getPartitionKey(account.getUuid())))
.withProjectionExpression(KEY_DEVICE_ID_KEY_ID)
.withConsistentRead(true);
private static AttributeValue getSortKey(final long deviceId, final long keyId) {
final ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[16]);
byteBuffer.putLong(deviceId);
byteBuffer.putLong(keyId);
return AttributeValues.fromByteBuffer(byteBuffer.flip());
}
deleteItemsForAccountMatchingQuery(account, querySpec);
});
}
@VisibleForTesting
static AttributeValue getSortKeyPrefix(final long deviceId) {
final ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[8]);
byteBuffer.putLong(deviceId);
return AttributeValues.fromByteBuffer(byteBuffer.flip());
}
@VisibleForTesting
void delete(final Account account, final long deviceId) {
DELETE_KEYS_FOR_DEVICE_TIMER.record(() -> {
final QuerySpec querySpec = new QuerySpec().withKeyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
.withNameMap(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID))
.withValueMap(Map.of(":uuid", getPartitionKey(account.getUuid()),
":sortprefix", getSortKeyPrefix(deviceId)))
.withProjectionExpression(KEY_DEVICE_ID_KEY_ID)
.withConsistentRead(true);
private Map<String, AttributeValue> getItemFromPreKey(final UUID accountUuid, final long deviceId, final PreKey preKey) {
return Map.of(
KEY_ACCOUNT_UUID, getPartitionKey(accountUuid),
KEY_DEVICE_ID_KEY_ID, getSortKey(deviceId, preKey.getKeyId()),
KEY_PUBLIC_KEY, AttributeValues.fromString(preKey.getPublicKey()));
}
deleteItemsForAccountMatchingQuery(account, querySpec);
});
}
private void deleteItemsForAccountMatchingQuery(final Account account, final QuerySpec querySpec) {
final byte[] partitionKey = getPartitionKey(account.getUuid());
writeInBatches(table.query(querySpec), batch -> {
final TableWriteItems writeItems = new TableWriteItems(table.getTableName());
for (final Item item : batch) {
writeItems.addPrimaryKeyToDelete(new PrimaryKey(KEY_ACCOUNT_UUID, partitionKey, KEY_DEVICE_ID_KEY_ID, item.getBinary(KEY_DEVICE_ID_KEY_ID)));
}
executeTableWriteItemsUntilComplete(writeItems);
});
}
private static byte[] getPartitionKey(final UUID accountUuid) {
return UUIDUtil.toBytes(accountUuid);
}
private static byte[] getSortKey(final long deviceId, final long keyId) {
final ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[16]);
byteBuffer.putLong(deviceId);
byteBuffer.putLong(keyId);
return byteBuffer.array();
}
private static byte[] getSortKeyPrefix(final long deviceId) {
final ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[8]);
byteBuffer.putLong(deviceId);
return byteBuffer.array();
}
private Item getItemFromPreKey(final UUID accountUuid, final long deviceId, final PreKey preKey) {
return new Item().withBinary(KEY_ACCOUNT_UUID, getPartitionKey(accountUuid))
.withBinary(KEY_DEVICE_ID_KEY_ID, getSortKey(deviceId, preKey.getKeyId()))
.withString(KEY_PUBLIC_KEY, preKey.getPublicKey());
}
private PreKey getPreKeyFromItem(final Item item) {
final long keyId = ByteBuffer.wrap(item.getBinary(KEY_DEVICE_ID_KEY_ID)).getLong(8);
return new PreKey(keyId, item.getString(KEY_PUBLIC_KEY));
}
private PreKey getPreKeyFromItem(Map<String, AttributeValue> item) {
final long keyId = item.get(KEY_DEVICE_ID_KEY_ID).b().asByteBuffer().getLong(8);
return new PreKey(keyId, item.get(KEY_PUBLIC_KEY).s());
}
}

256
service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDb.java
View File

@@ -8,17 +8,7 @@ package org.whispersystems.textsecuregcm.storage;
import static com.codahale.metrics.MetricRegistry.name;
import static io.micrometer.core.instrument.Metrics.timer;
import com.amazonaws.services.dynamodbv2.document.DeleteItemOutcome;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.document.Index;
import com.amazonaws.services.dynamodbv2.document.Item;
import com.amazonaws.services.dynamodbv2.document.PrimaryKey;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.TableWriteItems;
import com.amazonaws.services.dynamodbv2.document.api.QueryApi;
import com.amazonaws.services.dynamodbv2.document.spec.DeleteItemSpec;
import com.amazonaws.services.dynamodbv2.document.spec.QuerySpec;
import com.amazonaws.services.dynamodbv2.model.ReturnValue;
import com.google.common.collect.ImmutableMap;
import io.micrometer.core.instrument.Timer;
import java.nio.ByteBuffer;
import java.time.Duration;
@@ -27,11 +17,22 @@ import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;
import java.util.stream.Collectors;
import javax.annotation.Nonnull;
import org.apache.commons.lang3.StringUtils;
import org.whispersystems.textsecuregcm.entities.MessageProtos;
import org.whispersystems.textsecuregcm.entities.OutgoingMessageEntity;
import org.whispersystems.textsecuregcm.util.AttributeValues;
import org.whispersystems.textsecuregcm.util.UUIDUtil;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
import software.amazon.awssdk.services.dynamodb.model.DeleteItemResponse;
import software.amazon.awssdk.services.dynamodb.model.DeleteRequest;
import software.amazon.awssdk.services.dynamodb.model.PutRequest;
import software.amazon.awssdk.services.dynamodb.model.QueryRequest;
import software.amazon.awssdk.services.dynamodb.model.ReturnValue;
import software.amazon.awssdk.services.dynamodb.model.WriteRequest;
public class MessagesDynamoDb extends AbstractDynamoDbStore {
@@ -60,7 +61,7 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
private final String tableName;
private final Duration timeToLive;
public MessagesDynamoDb(DynamoDB dynamoDb, String tableName, Duration timeToLive) {
public MessagesDynamoDb(DynamoDbClient dynamoDb, String tableName, Duration timeToLive) {
super(dynamoDb);
this.tableName = tableName;
@@ -76,54 +77,61 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
throw new IllegalArgumentException("Maximum batch size of " + DYNAMO_DB_MAX_BATCH_SIZE + " execeeded with " + messages.size() + " messages");
}
final byte[] partitionKey = convertPartitionKey(destinationAccountUuid);
TableWriteItems items = new TableWriteItems(tableName);
final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid);
List<WriteRequest> writeItems = new ArrayList<>();
for (MessageProtos.Envelope message : messages) {
final UUID messageUuid = UUID.fromString(message.getServerGuid());
final Item item = new Item().withBinary(KEY_PARTITION, partitionKey)
.withBinary(KEY_SORT, convertSortKey(destinationDeviceId, message.getServerTimestamp(), messageUuid))
.withBinary(LOCAL_INDEX_MESSAGE_UUID_KEY_SORT, convertLocalIndexMessageUuidSortKey(messageUuid))
.withInt(KEY_TYPE, message.getType().getNumber())
.withLong(KEY_TIMESTAMP, message.getTimestamp())
.withLong(KEY_TTL, getTtlForMessage(message));
final ImmutableMap.Builder<String, AttributeValue> item = ImmutableMap.<String, AttributeValue>builder()
.put(KEY_PARTITION, partitionKey)
.put(KEY_SORT, convertSortKey(destinationDeviceId, message.getServerTimestamp(), messageUuid))
.put(LOCAL_INDEX_MESSAGE_UUID_KEY_SORT, convertLocalIndexMessageUuidSortKey(messageUuid))
.put(KEY_TYPE, AttributeValues.fromInt(message.getType().getNumber()))
.put(KEY_TIMESTAMP, AttributeValues.fromLong(message.getTimestamp()))
.put(KEY_TTL, AttributeValues.fromLong(getTtlForMessage(message)));
if (message.hasRelay() && message.getRelay().length() > 0) {
item.withString(KEY_RELAY, message.getRelay());
item.put(KEY_RELAY, AttributeValues.fromString(message.getRelay()));
}
if (message.hasSource()) {
item.withString(KEY_SOURCE, message.getSource());
item.put(KEY_SOURCE, AttributeValues.fromString(message.getSource()));
}
if (message.hasSourceUuid()) {
item.withBinary(KEY_SOURCE_UUID, UUIDUtil.toBytes(UUID.fromString(message.getSourceUuid())));
item.put(KEY_SOURCE_UUID, AttributeValues.fromUUID(UUID.fromString(message.getSourceUuid())));
}
if (message.hasSourceDevice()) {
item.withInt(KEY_SOURCE_DEVICE, message.getSourceDevice());
item.put(KEY_SOURCE_DEVICE, AttributeValues.fromInt(message.getSourceDevice()));
}
if (message.hasLegacyMessage()) {
item.withBinary(KEY_MESSAGE, message.getLegacyMessage().toByteArray());
item.put(KEY_MESSAGE, AttributeValues.fromByteArray(message.getLegacyMessage().toByteArray()));
}
if (message.hasContent()) {
item.withBinary(KEY_CONTENT, message.getContent().toByteArray());
item.put(KEY_CONTENT, AttributeValues.fromByteArray(message.getContent().toByteArray()));
}
items.addItemToPut(item);
writeItems.add(WriteRequest.builder().putRequest(PutRequest.builder()
.item(item.build())
.build()).build());
}
executeTableWriteItemsUntilComplete(items);
executeTableWriteItemsUntilComplete(Map.of(tableName, writeItems));
}
public List<OutgoingMessageEntity> load(final UUID destinationAccountUuid, final long destinationDeviceId, final int requestedNumberOfMessagesToFetch) {
return loadTimer.record(() -> {
final int numberOfMessagesToFetch = Math.min(requestedNumberOfMessagesToFetch, RESULT_SET_CHUNK_SIZE);
final byte[] partitionKey = convertPartitionKey(destinationAccountUuid);
final QuerySpec querySpec = new QuerySpec().withConsistentRead(true)
.withKeyConditionExpression("#part = :part AND begins_with ( #sort , :sortprefix )")
.withNameMap(Map.of("#part", KEY_PARTITION,
"#sort", KEY_SORT))
.withValueMap(Map.of(":part", partitionKey,
":sortprefix", convertDestinationDeviceIdToSortKeyPrefix(destinationDeviceId)))
.withMaxResultSize(numberOfMessagesToFetch);
final Table table = getDynamoDb().getTable(tableName);
final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid);
final QueryRequest queryRequest = QueryRequest.builder()
.tableName(tableName)
.consistentRead(true)
.keyConditionExpression("#part = :part AND begins_with ( #sort , :sortprefix )")
.expressionAttributeNames(Map.of(
"#part", KEY_PARTITION,
"#sort", KEY_SORT))
.expressionAttributeValues(Map.of(
":part", partitionKey,
":sortprefix", convertDestinationDeviceIdToSortKeyPrefix(destinationDeviceId)))
.limit(numberOfMessagesToFetch)
.build();
List<OutgoingMessageEntity> messageEntities = new ArrayList<>(numberOfMessagesToFetch);
for (Item message : table.query(querySpec)) {
for (Map<String, AttributeValue> message : db().query(queryRequest).items()) {
messageEntities.add(convertItemToOutgoingMessageEntity(message));
}
return messageEntities;
@@ -136,53 +144,63 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
throw new IllegalArgumentException("must specify a source");
}
final byte[] partitionKey = convertPartitionKey(destinationAccountUuid);
final QuerySpec querySpec = new QuerySpec().withProjectionExpression(KEY_SORT)
.withConsistentRead(true)
.withKeyConditionExpression("#part = :part AND begins_with ( #sort , :sortprefix )")
.withFilterExpression("#source = :source AND #timestamp = :timestamp")
.withNameMap(Map.of("#part", KEY_PARTITION,
"#sort", KEY_SORT,
"#source", KEY_SOURCE,
"#timestamp", KEY_TIMESTAMP))
.withValueMap(Map.of(":part", partitionKey,
":sortprefix", convertDestinationDeviceIdToSortKeyPrefix(destinationDeviceId),
":source", source,
":timestamp", timestamp));
final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid);
final QueryRequest queryRequest = QueryRequest.builder()
.tableName(tableName)
.projectionExpression(KEY_SORT)
.consistentRead(true)
.keyConditionExpression("#part = :part AND begins_with ( #sort , :sortprefix )")
.filterExpression("#source = :source AND #timestamp = :timestamp")
.expressionAttributeNames(Map.of(
"#part", KEY_PARTITION,
"#sort", KEY_SORT,
"#source", KEY_SOURCE,
"#timestamp", KEY_TIMESTAMP))
.expressionAttributeValues(Map.of(
":part", partitionKey,
":sortprefix", convertDestinationDeviceIdToSortKeyPrefix(destinationDeviceId),
":source", AttributeValues.fromString(source),
":timestamp", AttributeValues.fromLong(timestamp)))
.build();
final Table table = getDynamoDb().getTable(tableName);
return deleteItemsMatchingQueryAndReturnFirstOneActuallyDeleted(table, partitionKey, querySpec, table);
return deleteItemsMatchingQueryAndReturnFirstOneActuallyDeleted(partitionKey, queryRequest);
});
}
public Optional<OutgoingMessageEntity> deleteMessageByDestinationAndGuid(final UUID destinationAccountUuid, final long destinationDeviceId, final UUID messageUuid) {
return deleteByGuid.record(() -> {
final byte[] partitionKey = convertPartitionKey(destinationAccountUuid);
final QuerySpec querySpec = new QuerySpec().withProjectionExpression(KEY_SORT)
.withConsistentRead(true)
.withKeyConditionExpression("#part = :part AND #uuid = :uuid")
.withNameMap(Map.of("#part", KEY_PARTITION,
"#uuid", LOCAL_INDEX_MESSAGE_UUID_KEY_SORT))
.withValueMap(Map.of(":part", partitionKey,
":uuid", convertLocalIndexMessageUuidSortKey(messageUuid)));
final Table table = getDynamoDb().getTable(tableName);
final Index index = table.getIndex(LOCAL_INDEX_MESSAGE_UUID_NAME);
return deleteItemsMatchingQueryAndReturnFirstOneActuallyDeleted(table, partitionKey, querySpec, index);
final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid);
final QueryRequest queryRequest = QueryRequest.builder()
.tableName(tableName)
.indexName(LOCAL_INDEX_MESSAGE_UUID_NAME)
.projectionExpression(KEY_SORT)
.consistentRead(true)
.keyConditionExpression("#part = :part AND #uuid = :uuid")
.expressionAttributeNames(Map.of(
"#part", KEY_PARTITION,
"#uuid", LOCAL_INDEX_MESSAGE_UUID_KEY_SORT))
.expressionAttributeValues(Map.of(
":part", partitionKey,
":uuid", convertLocalIndexMessageUuidSortKey(messageUuid)))
.build();
return deleteItemsMatchingQueryAndReturnFirstOneActuallyDeleted(partitionKey, queryRequest);
});
}
@Nonnull
private Optional<OutgoingMessageEntity> deleteItemsMatchingQueryAndReturnFirstOneActuallyDeleted(Table table, byte[] partitionKey, QuerySpec querySpec, QueryApi queryApi) {
private Optional<OutgoingMessageEntity> deleteItemsMatchingQueryAndReturnFirstOneActuallyDeleted(AttributeValue partitionKey, QueryRequest queryRequest) {
Optional<OutgoingMessageEntity> result = Optional.empty();
for (Item item : queryApi.query(querySpec)) {
final byte[] rangeKeyValue = item.getBinary(KEY_SORT);
DeleteItemSpec deleteItemSpec = new DeleteItemSpec().withPrimaryKey(KEY_PARTITION, partitionKey, KEY_SORT, rangeKeyValue);
for (Map<String, AttributeValue> item : db().query(queryRequest).items()) {
final byte[] rangeKeyValue = item.get(KEY_SORT).b().asByteArray();
DeleteItemRequest.Builder deleteItemRequest = DeleteItemRequest.builder()
.tableName(tableName)
.key(Map.of(KEY_PARTITION, partitionKey, KEY_SORT, AttributeValues.fromByteArray(rangeKeyValue)));
if (result.isEmpty()) {
deleteItemSpec.withReturnValues(ReturnValue.ALL_OLD);
deleteItemRequest.returnValues(ReturnValue.ALL_OLD);
}
final DeleteItemOutcome deleteItemOutcome = table.deleteItem(deleteItemSpec);
if (deleteItemOutcome.getItem() != null && deleteItemOutcome.getItem().hasAttribute(KEY_PARTITION)) {
result = Optional.of(convertItemToOutgoingMessageEntity(deleteItemOutcome.getItem()));
final DeleteItemResponse deleteItemResponse = db().deleteItem(deleteItemRequest.build());
if (deleteItemResponse.attributes() != null && deleteItemResponse.attributes().containsKey(KEY_PARTITION)) {
result = Optional.of(convertItemToOutgoingMessageEntity(deleteItemResponse.attributes()));
}
}
return result;
@@ -190,74 +208,88 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
public void deleteAllMessagesForAccount(final UUID destinationAccountUuid) {
deleteByAccount.record(() -> {
final byte[] partitionKey = convertPartitionKey(destinationAccountUuid);
final QuerySpec querySpec = new QuerySpec().withHashKey(KEY_PARTITION, partitionKey)
.withProjectionExpression(KEY_SORT)
.withConsistentRead(true);
deleteRowsMatchingQuery(partitionKey, querySpec);
final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid);
final QueryRequest queryRequest = QueryRequest.builder()
.tableName(tableName)
.projectionExpression(KEY_SORT)
.consistentRead(true)
.keyConditionExpression("#part = :part")
.expressionAttributeNames(Map.of("#part", KEY_PARTITION))
.expressionAttributeValues(Map.of(":part", partitionKey))
.build();
deleteRowsMatchingQuery(partitionKey, queryRequest);
});
}
public void deleteAllMessagesForDevice(final UUID destinationAccountUuid, final long destinationDeviceId) {
deleteByDevice.record(() -> {
final byte[] partitionKey = convertPartitionKey(destinationAccountUuid);
final QuerySpec querySpec = new QuerySpec().withKeyConditionExpression("#part = :part AND begins_with ( #sort , :sortprefix )")
.withNameMap(Map.of("#part", KEY_PARTITION,
"#sort", KEY_SORT))
.withValueMap(Map.of(":part", partitionKey,
":sortprefix", convertDestinationDeviceIdToSortKeyPrefix(destinationDeviceId)))
.withProjectionExpression(KEY_SORT)
.withConsistentRead(true);
deleteRowsMatchingQuery(partitionKey, querySpec);
final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid);
final QueryRequest queryRequest = QueryRequest.builder()
.tableName(tableName)
.keyConditionExpression("#part = :part AND begins_with ( #sort , :sortprefix )")
.expressionAttributeNames(Map.of(
"#part", KEY_PARTITION,
"#sort", KEY_SORT))
.expressionAttributeValues(Map.of(
":part", partitionKey,
":sortprefix", convertDestinationDeviceIdToSortKeyPrefix(destinationDeviceId)))
.projectionExpression(KEY_SORT)
.consistentRead(true)
.build();
deleteRowsMatchingQuery(partitionKey, queryRequest);
});
}
private OutgoingMessageEntity convertItemToOutgoingMessageEntity(Item message) {
final SortKey sortKey = convertSortKey(message.getBinary(KEY_SORT));
final UUID messageUuid = convertLocalIndexMessageUuidSortKey(message.getBinary(LOCAL_INDEX_MESSAGE_UUID_KEY_SORT));
final int type = message.getInt(KEY_TYPE);
final String relay = message.getString(KEY_RELAY);
final long timestamp = message.getLong(KEY_TIMESTAMP);
final String source = message.getString(KEY_SOURCE);
final UUID sourceUuid = message.hasAttribute(KEY_SOURCE_UUID) ? convertUuidFromBytes(message.getBinary(KEY_SOURCE_UUID), "message source uuid") : null;
final int sourceDevice = message.hasAttribute(KEY_SOURCE_DEVICE) ? message.getInt(KEY_SOURCE_DEVICE) : 0;
final byte[] messageBytes = message.getBinary(KEY_MESSAGE);
final byte[] content = message.getBinary(KEY_CONTENT);
private OutgoingMessageEntity convertItemToOutgoingMessageEntity(Map<String, AttributeValue> message) {
final SortKey sortKey = convertSortKey(message.get(KEY_SORT).b().asByteArray());
final UUID messageUuid = convertLocalIndexMessageUuidSortKey(message.get(LOCAL_INDEX_MESSAGE_UUID_KEY_SORT).b().asByteArray());
final int type = AttributeValues.getInt(message, KEY_TYPE, 0);
final String relay = AttributeValues.getString(message, KEY_RELAY, null);
final long timestamp = AttributeValues.getLong(message, KEY_TIMESTAMP, 0L);
final String source = AttributeValues.getString(message, KEY_SOURCE, null);
final UUID sourceUuid = AttributeValues.getUUID(message, KEY_SOURCE_UUID, null);
final int sourceDevice = AttributeValues.getInt(message, KEY_SOURCE_DEVICE, 0);
final byte[] messageBytes = AttributeValues.getByteArray(message, KEY_MESSAGE, null);
final byte[] content = AttributeValues.getByteArray(message, KEY_CONTENT, null);
return new OutgoingMessageEntity(-1L, false, messageUuid, type, relay, timestamp, source, sourceUuid, sourceDevice, messageBytes, content, sortKey.getServerTimestamp());
}
private void deleteRowsMatchingQuery(byte[] partitionKey, QuerySpec querySpec) {
final Table table = getDynamoDb().getTable(tableName);
writeInBatches(table.query(querySpec), (itemBatch) -> deleteItems(partitionKey, itemBatch));
private void deleteRowsMatchingQuery(AttributeValue partitionKey, QueryRequest querySpec) {
writeInBatches(db().query(querySpec).items(), (itemBatch) -> deleteItems(partitionKey, itemBatch));
}
private void deleteItems(byte[] partitionKey, List<Item> items) {
final TableWriteItems tableWriteItems = new TableWriteItems(tableName);
items.stream().map(item -> new PrimaryKey(KEY_PARTITION, partitionKey, KEY_SORT, item.getBinary(KEY_SORT))).forEach(tableWriteItems::addPrimaryKeyToDelete);
executeTableWriteItemsUntilComplete(tableWriteItems);
private void deleteItems(AttributeValue partitionKey, List<Map<String, AttributeValue>> items) {
List<WriteRequest> deletes = items.stream()
.map(item -> WriteRequest.builder()
.deleteRequest(DeleteRequest.builder().key(Map.of(
KEY_PARTITION, partitionKey,
KEY_SORT, item.get(KEY_SORT))).build())
.build())
.collect(Collectors.toList());
executeTableWriteItemsUntilComplete(Map.of(tableName, deletes));
}
private long getTtlForMessage(MessageProtos.Envelope message) {
return message.getServerTimestamp() / 1000 + timeToLive.getSeconds();
}
private static byte[] convertPartitionKey(final UUID destinationAccountUuid) {
return UUIDUtil.toBytes(destinationAccountUuid);
private static AttributeValue convertPartitionKey(final UUID destinationAccountUuid) {
return AttributeValues.fromUUID(destinationAccountUuid);
}
private static byte[] convertSortKey(final long destinationDeviceId, final long serverTimestamp, final UUID messageUuid) {
private static AttributeValue convertSortKey(final long destinationDeviceId, final long serverTimestamp, final UUID messageUuid) {
ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[32]);
byteBuffer.putLong(destinationDeviceId);
byteBuffer.putLong(serverTimestamp);
byteBuffer.putLong(messageUuid.getMostSignificantBits());
byteBuffer.putLong(messageUuid.getLeastSignificantBits());
return byteBuffer.array();
return AttributeValues.fromByteBuffer(byteBuffer.flip());
}
private static byte[] convertDestinationDeviceIdToSortKeyPrefix(final long destinationDeviceId) {
private static AttributeValue convertDestinationDeviceIdToSortKeyPrefix(final long destinationDeviceId) {
ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[8]);
byteBuffer.putLong(destinationDeviceId);
return byteBuffer.array();
return AttributeValues.fromByteBuffer(byteBuffer.flip());
}
private static SortKey convertSortKey(final byte[] bytes) {
@@ -273,8 +305,8 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
return new SortKey(destinationDeviceId, serverTimestamp, new UUID(mostSigBits, leastSigBits));
}
private static byte[] convertLocalIndexMessageUuidSortKey(final UUID messageUuid) {
return UUIDUtil.toBytes(messageUuid);
private static AttributeValue convertLocalIndexMessageUuidSortKey(final UUID messageUuid) {
return AttributeValues.fromUUID(messageUuid);
}
private static UUID convertLocalIndexMessageUuidSortKey(final byte[] bytes) {

59
service/src/main/java/org/whispersystems/textsecuregcm/storage/MigrationDeletedAccounts.java
View File

@@ -1,42 +1,52 @@
package org.whispersystems.textsecuregcm.storage;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.document.Item;
import com.amazonaws.services.dynamodbv2.document.PrimaryKey;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.TableWriteItems;
import com.amazonaws.services.dynamodbv2.document.spec.ScanSpec;
import com.google.common.annotations.VisibleForTesting;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;
import org.whispersystems.textsecuregcm.util.UUIDUtil;
import java.util.stream.Collectors;
import org.whispersystems.textsecuregcm.util.AttributeValues;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.DeleteRequest;
import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
import software.amazon.awssdk.services.dynamodb.model.ScanResponse;
import software.amazon.awssdk.services.dynamodb.model.WriteRequest;
public class MigrationDeletedAccounts extends AbstractDynamoDbStore {
private final Table table;
private final String tableName;
static final String KEY_UUID = "U";
public MigrationDeletedAccounts(DynamoDB dynamoDb, String tableName) {
public MigrationDeletedAccounts(DynamoDbClient dynamoDb, String tableName) {
super(dynamoDb);
table = dynamoDb.getTable(tableName);
this.tableName = tableName;
}
public void put(UUID uuid) {
table.putItem(new Item()
.withPrimaryKey(primaryKey(uuid)));
db().putItem(PutItemRequest.builder()
.tableName(tableName)
.item(primaryKey(uuid))
.build());
}
public List<UUID> getRecentlyDeletedUuids() {
final List<UUID> uuids = new ArrayList<>();
Optional<ScanResponse> firstPage = db().scanPaginator(ScanRequest.builder()
.tableName(tableName)
.build()).stream().findAny(); // get the first available response
for (Item item : table.scan(new ScanSpec()).firstPage()) {
// only process one page each time. If we have a significant backlog at the end of the migration
// we can handle it separately
uuids.add(UUIDUtil.fromByteBuffer(item.getByteBuffer(KEY_UUID)));
if (firstPage.isPresent()) {
for (Map<String, AttributeValue> item : firstPage.get().items()) {
// only process one page each time. If we have a significant backlog at the end of the migration
// we can handle it separately
uuids.add(AttributeValues.getUUID(item, KEY_UUID, null));
}
}
return uuids;
@@ -45,20 +55,17 @@ public class MigrationDeletedAccounts extends AbstractDynamoDbStore {
public void delete(List<UUID> uuids) {
writeInBatches(uuids, (batch) -> {
List<WriteRequest> deletes = batch.stream().map((uuid) -> WriteRequest.builder().deleteRequest(DeleteRequest.builder()
.key(primaryKey(uuid))
.build()).build()).collect(Collectors.toList());
final TableWriteItems deleteItems = new TableWriteItems(table.getTableName());
for (UUID uuid : batch) {
deleteItems.addPrimaryKeyToDelete(primaryKey(uuid));
}
executeTableWriteItemsUntilComplete(deleteItems);
executeTableWriteItemsUntilComplete(Map.of(tableName, deletes));
});
}
@VisibleForTesting
public static PrimaryKey primaryKey(UUID uuid) {
return new PrimaryKey(KEY_UUID, UUIDUtil.toBytes(uuid));
public static Map<String, AttributeValue> primaryKey(UUID uuid) {
return Map.of(KEY_UUID, AttributeValues.fromUUID(uuid));
}
}

37
service/src/main/java/org/whispersystems/textsecuregcm/storage/MigrationRetryAccounts.java
View File

@@ -1,43 +1,44 @@
package org.whispersystems.textsecuregcm.storage;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.document.Item;
import com.amazonaws.services.dynamodbv2.document.Page;
import com.amazonaws.services.dynamodbv2.document.PrimaryKey;
import com.amazonaws.services.dynamodbv2.document.ScanOutcome;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.spec.ScanSpec;
import com.google.common.annotations.VisibleForTesting;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import org.whispersystems.textsecuregcm.util.UUIDUtil;
import org.whispersystems.textsecuregcm.util.AttributeValues;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
import software.amazon.awssdk.services.dynamodb.model.ScanResponse;
public class MigrationRetryAccounts extends AbstractDynamoDbStore {
private final Table table;
private final String tableName;
static final String KEY_UUID = "U";
public MigrationRetryAccounts(DynamoDB dynamoDb, String tableName) {
public MigrationRetryAccounts(DynamoDbClient dynamoDb, String tableName) {
super(dynamoDb);
table = dynamoDb.getTable(tableName);
this.tableName = tableName;
}
public void put(UUID uuid) {
table.putItem(new Item()
.withPrimaryKey(primaryKey(uuid)));
db().putItem(PutItemRequest.builder()
.tableName(tableName)
.item(primaryKey(uuid))
.build());
}
public List<UUID> getUuids(int max) {
final List<UUID> uuids = new ArrayList<>();
for (Page<Item, ScanOutcome> page : table.scan(new ScanSpec()).pages()) {
for (ScanResponse response : db().scanPaginator(ScanRequest.builder().tableName(tableName).build())) {
for (Item item : page) {
uuids.add(UUIDUtil.fromByteBuffer(item.getByteBuffer(KEY_UUID)));
for (Map<String, AttributeValue> item : response.items()) {
uuids.add(AttributeValues.getUUID(item, KEY_UUID, null));
if (uuids.size() >= max) {
break;
@@ -53,8 +54,8 @@ public class MigrationRetryAccounts extends AbstractDynamoDbStore {
}
@VisibleForTesting
public static PrimaryKey primaryKey(UUID uuid) {
return new PrimaryKey(KEY_UUID, UUIDUtil.toBytes(uuid));
public static Map<String, AttributeValue> primaryKey(UUID uuid) {
return Map.of(KEY_UUID, AttributeValues.fromUUID(uuid));
}
}

48
service/src/main/java/org/whispersystems/textsecuregcm/storage/PushChallengeDynamoDb.java
View File

@@ -5,25 +5,23 @@
package org.whispersystems.textsecuregcm.storage;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.document.Item;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.spec.DeleteItemSpec;
import com.amazonaws.services.dynamodbv2.document.spec.PutItemSpec;
import com.amazonaws.services.dynamodbv2.model.ConditionalCheckFailedException;
import java.time.Clock;
import java.time.Duration;
import java.util.Map;
import java.util.UUID;
import com.google.common.annotations.VisibleForTesting;
import org.whispersystems.textsecuregcm.util.UUIDUtil;
import org.whispersystems.textsecuregcm.util.AttributeValues;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException;
import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
/**
* Stores push challenge tokens. Users may have at most one outstanding push challenge token at a time.
*/
public class PushChallengeDynamoDb extends AbstractDynamoDbStore {
private final Table table;
private final String tableName;
private final Clock clock;
static final String KEY_ACCOUNT_UUID = "U";
@@ -33,15 +31,15 @@ public class PushChallengeDynamoDb extends AbstractDynamoDbStore {
private static final Map<String, String> UUID_NAME_MAP = Map.of("#uuid", KEY_ACCOUNT_UUID);
private static final Map<String, String> CHALLENGE_TOKEN_NAME_MAP = Map.of("#challenge", ATTR_CHALLENGE_TOKEN);
public PushChallengeDynamoDb(final DynamoDB dynamoDB, final String tableName) {
public PushChallengeDynamoDb(final DynamoDbClient dynamoDB, final String tableName) {
this(dynamoDB, tableName, Clock.systemUTC());
}
@VisibleForTesting
PushChallengeDynamoDb(final DynamoDB dynamoDB, final String tableName, final Clock clock) {
PushChallengeDynamoDb(final DynamoDbClient dynamoDB, final String tableName, final Clock clock) {
super(dynamoDB);
this.table = dynamoDB.getTable(tableName);
this.tableName = tableName;
this.clock = clock;
}
@@ -57,13 +55,15 @@ public class PushChallengeDynamoDb extends AbstractDynamoDbStore {
*/
public boolean add(final UUID accountUuid, final byte[] challengeToken, final Duration ttl) {
try {
table.putItem( new PutItemSpec()
.withItem(new Item()
.withBinary(KEY_ACCOUNT_UUID, UUIDUtil.toByteBuffer(accountUuid))
.withBinary(ATTR_CHALLENGE_TOKEN, challengeToken)
.withNumber(ATTR_TTL, getExpirationTimestamp(ttl)))
.withConditionExpression("attribute_not_exists(#uuid)")
.withNameMap(UUID_NAME_MAP));
db().putItem(PutItemRequest.builder()
.tableName(tableName)
.item(Map.of(
KEY_ACCOUNT_UUID, AttributeValues.fromUUID(accountUuid),
ATTR_CHALLENGE_TOKEN, AttributeValues.fromByteArray(challengeToken),
ATTR_TTL, AttributeValues.fromLong(getExpirationTimestamp(ttl))))
.conditionExpression("attribute_not_exists(#uuid)")
.expressionAttributeNames(UUID_NAME_MAP)
.build());
return true;
} catch (final ConditionalCheckFailedException e) {
return false;
@@ -84,11 +84,13 @@ public class PushChallengeDynamoDb extends AbstractDynamoDbStore {
*/
public boolean remove(final UUID accountUuid, final byte[] challengeToken) {
try {
table.deleteItem(new DeleteItemSpec()
.withPrimaryKey(KEY_ACCOUNT_UUID, UUIDUtil.toByteBuffer(accountUuid))
.withConditionExpression("#challenge = :challenge")
.withNameMap(CHALLENGE_TOKEN_NAME_MAP)
.withValueMap(Map.of(":challenge", challengeToken)));
db().deleteItem(DeleteItemRequest.builder()
.tableName(tableName)
.key(Map.of(KEY_ACCOUNT_UUID, AttributeValues.fromUUID(accountUuid)))
.conditionExpression("#challenge = :challenge")
.expressionAttributeNames(CHALLENGE_TOKEN_NAME_MAP)
.expressionAttributeValues(Map.of(":challenge", AttributeValues.fromByteArray(challengeToken)))
.build());
return true;
} catch (final ConditionalCheckFailedException e) {
return false;

52
service/src/main/java/org/whispersystems/textsecuregcm/storage/ReportMessageDynamoDb.java
View File

@@ -1,13 +1,14 @@
package org.whispersystems.textsecuregcm.storage;
import com.amazonaws.services.dynamodbv2.document.DeleteItemOutcome;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.document.Item;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.spec.DeleteItemSpec;
import com.amazonaws.services.dynamodbv2.model.ReturnValue;
import org.whispersystems.textsecuregcm.util.AttributeValues;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
import software.amazon.awssdk.services.dynamodb.model.DeleteItemResponse;
import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
import software.amazon.awssdk.services.dynamodb.model.ReturnValue;
import java.time.Duration;
import java.time.Instant;
import java.util.Map;
public class ReportMessageDynamoDb {
@@ -16,33 +17,30 @@ public class ReportMessageDynamoDb {
static final Duration TIME_TO_LIVE = Duration.ofDays(7);
private final Table table;
private final DynamoDbClient db;
private final String tableName;
public ReportMessageDynamoDb(final DynamoDB dynamoDB, final String tableName) {
this.table = dynamoDB.getTable(tableName);
public ReportMessageDynamoDb(final DynamoDbClient dynamoDB, final String tableName) {
this.db = dynamoDB;
this.tableName = tableName;
}
public void store(byte[] hash) {
table.putItem(buildItemForHash(hash));
}
private Item buildItemForHash(byte[] hash) {
return new Item()
.withBinary(KEY_HASH, hash)
.withLong(ATTR_TTL, Instant.now().plus(TIME_TO_LIVE).getEpochSecond());
db.putItem(PutItemRequest.builder()
.tableName(tableName)
.item(Map.of(
KEY_HASH, AttributeValues.fromByteArray(hash),
ATTR_TTL, AttributeValues.fromLong(Instant.now().plus(TIME_TO_LIVE).getEpochSecond())
))
.build());
}
public boolean remove(byte[] hash) {
final DeleteItemSpec deleteItemSpec = new DeleteItemSpec()
.withPrimaryKey(KEY_HASH, hash)
.withReturnValues(ReturnValue.ALL_OLD);
final DeleteItemOutcome outcome = table.deleteItem(deleteItemSpec);
return outcome.getItem() != null;
final DeleteItemResponse deleteItemResponse = db.deleteItem(DeleteItemRequest.builder()
.tableName(tableName)
.key(Map.of(KEY_HASH, AttributeValues.fromByteArray(hash)))
.returnValues(ReturnValue.ALL_OLD)
.build());
return !deleteItemResponse.attributes().isEmpty();
}
}

89
service/src/main/java/org/whispersystems/textsecuregcm/util/AttributeValues.java Normal file
View File

@@ -0,0 +1,89 @@
/*
* Copyright 2013-2021 Signal Messenger, LLC
* SPDX-License-Identifier: AGPL-3.0-only
*/
package org.whispersystems.textsecuregcm.util;
import software.amazon.awssdk.core.SdkBytes;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import java.nio.ByteBuffer;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;
/** AwsAV provides static helper methods for working with AWS AttributeValues. */
public class AttributeValues {
public static AttributeValue fromString(String value) {
return AttributeValue.builder().s(value).build();
}
public static AttributeValue fromLong(long value) {
return AttributeValue.builder().n(Long.toString(value)).build();
}
public static AttributeValue fromInt(int value) {
return AttributeValue.builder().n(Integer.toString(value)).build();
}
public static AttributeValue fromByteArray(byte[] value) {
return AttributeValues.fromSdkBytes(SdkBytes.fromByteArray(value));
}
public static AttributeValue fromByteBuffer(ByteBuffer value) {
return AttributeValues.fromSdkBytes(SdkBytes.fromByteBuffer(value));
}
public static AttributeValue fromUUID(UUID uuid) {
return AttributeValues.fromSdkBytes(SdkBytes.fromByteArrayUnsafe(UUIDUtil.toBytes(uuid)));
}
public static AttributeValue fromSdkBytes(SdkBytes value) {
return AttributeValue.builder().b(value).build();
}
private static int toInt(AttributeValue av) {
return Integer.parseInt(av.n());
}
private static long toLong(AttributeValue av) {
return Long.parseLong(av.n());
}
private static UUID toUUID(AttributeValue av) {
return UUIDUtil.fromBytes(av.b().asByteArrayUnsafe()); // We're guaranteed not to modify the byte array
}
private static byte[] toByteArray(AttributeValue av) {
return av.b().asByteArray();
}
private static String toString(AttributeValue av) {
return av.s();
}
public static Optional<AttributeValue> get(Map<String, AttributeValue> item, String key) {
return Optional.ofNullable(item.get(key));
}
public static int getInt(Map<String, AttributeValue> item, String key, int defaultValue) {
return AttributeValues.get(item, key).map(AttributeValues::toInt).orElse(defaultValue);
}
public static String getString(Map<String, AttributeValue> item, String key, String defaultValue) {
return AttributeValues.get(item, key).map(AttributeValues::toString).orElse(defaultValue);
}
public static long getLong(Map<String, AttributeValue> item, String key, long defaultValue) {
return AttributeValues.get(item, key).map(AttributeValues::toLong).orElse(defaultValue);
}
public static byte[] getByteArray(Map<String, AttributeValue> item, String key, byte[] defaultValue) {
return AttributeValues.get(item, key).map(AttributeValues::toByteArray).orElse(defaultValue);
}
public static UUID getUUID(Map<String, AttributeValue> item, String key, UUID defaultValue) {
return AttributeValues.get(item, key).map(AttributeValues::toUUID).orElse(defaultValue);
}
}

41
service/src/main/java/org/whispersystems/textsecuregcm/util/DynamoDbFromConfig.java Normal file
View File

@@ -0,0 +1,41 @@
package org.whispersystems.textsecuregcm.util;
import org.whispersystems.textsecuregcm.configuration.DynamoDbConfiguration;
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
import software.amazon.awssdk.core.client.config.ClientAsyncConfiguration;
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
import software.amazon.awssdk.core.client.config.SdkAdvancedAsyncClientOption;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClientBuilder;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import java.util.concurrent.Executor;
public class DynamoDbFromConfig {
private static ClientOverrideConfiguration clientOverrideConfiguration(DynamoDbConfiguration config) {
return ClientOverrideConfiguration.builder()
.apiCallTimeout(config.getClientExecutionTimeout())
.apiCallAttemptTimeout(config.getClientRequestTimeout())
.build();
}
public static DynamoDbClient client(DynamoDbConfiguration config, AwsCredentialsProvider credentialsProvider) {
return DynamoDbClient.builder()
.region(Region.of(config.getRegion()))
.credentialsProvider(credentialsProvider)
.overrideConfiguration(clientOverrideConfiguration(config))
.build();
}
public static DynamoDbAsyncClient asyncClient(DynamoDbConfiguration config, AwsCredentialsProvider credentialsProvider, Executor executor) {
DynamoDbAsyncClientBuilder builder = DynamoDbAsyncClient.builder()
.region(Region.of(config.getRegion()))
.credentialsProvider(credentialsProvider)
.overrideConfiguration(clientOverrideConfiguration(config));
if (executor != null) {
builder.asyncConfiguration(ClientAsyncConfiguration.builder()
.advancedOption(SdkAdvancedAsyncClientOption.FUTURE_COMPLETION_EXECUTOR,
executor)
.build());
}
return builder.build();
}
}

16
service/src/main/java/org/whispersystems/textsecuregcm/util/UUIDUtil.java
View File

@@ -5,6 +5,7 @@
package org.whispersystems.textsecuregcm.util;
import java.nio.BufferUnderflowException;
import java.nio.ByteBuffer;
import java.util.UUID;
@@ -26,12 +27,15 @@ public class UUIDUtil {
}
public static UUID fromByteBuffer(final ByteBuffer byteBuffer) {
if (byteBuffer.array().length != 16) {
throw new IllegalArgumentException("unexpected byte array length; was " + byteBuffer.array().length + " but expected 16");
try {
final long mostSigBits = byteBuffer.getLong();
final long leastSigBits = byteBuffer.getLong();
if (byteBuffer.hasRemaining()) {
throw new IllegalArgumentException("unexpected byte array length; was greater than 16");
}
return new UUID(mostSigBits, leastSigBits);
} catch (BufferUnderflowException e) {
throw new IllegalArgumentException("unexpected byte array length; was less than 16");
}
final long mostSigBits = byteBuffer.getLong();
final long leastSigBits = byteBuffer.getLong();
return new UUID(mostSigBits, leastSigBits);
}
}

82
service/src/main/java/org/whispersystems/textsecuregcm/workers/DeleteUserCommand.java
View File

@@ -9,11 +9,6 @@ import static com.codahale.metrics.MetricRegistry.name;
import com.amazonaws.ClientConfiguration;
import com.amazonaws.auth.InstanceProfileCredentialsProvider;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBAsync;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBAsyncClientBuilder;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClientBuilder;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.fasterxml.jackson.databind.DeserializationFeature;
import io.dropwizard.Application;
import io.dropwizard.cli.EnvironmentCommand;
@@ -59,6 +54,9 @@ import org.whispersystems.textsecuregcm.storage.ReportMessageManager;
import org.whispersystems.textsecuregcm.storage.ReservedUsernames;
import org.whispersystems.textsecuregcm.storage.Usernames;
import org.whispersystems.textsecuregcm.storage.UsernamesManager;
import org.whispersystems.textsecuregcm.util.DynamoDbFromConfig;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfiguration> {
@@ -100,64 +98,20 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
FaultTolerantDatabase accountDatabase = new FaultTolerantDatabase("account_database_delete_user", accountJdbi, configuration.getAccountsDatabaseConfiguration().getCircuitBreakerConfiguration());
ClientResources redisClusterClientResources = ClientResources.builder().build();
AmazonDynamoDBClientBuilder clientBuilder = AmazonDynamoDBClientBuilder
.standard()
.withRegion(configuration.getMessageDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) configuration.getMessageDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) configuration.getMessageDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
AmazonDynamoDBClientBuilder keysDynamoDbClientBuilder = AmazonDynamoDBClientBuilder
.standard()
.withRegion(configuration.getKeysDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) configuration.getKeysDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) configuration.getKeysDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
AmazonDynamoDBClientBuilder accountsDynamoDbClientBuilder = AmazonDynamoDBClientBuilder
.standard()
.withRegion(configuration.getAccountsDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) configuration.getAccountsDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) configuration.getAccountsDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
ThreadPoolExecutor accountsDynamoDbMigrationThreadPool = new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS,
new LinkedBlockingDeque<>());
AmazonDynamoDBAsyncClientBuilder accountsDynamoDbAsyncClientBuilder = AmazonDynamoDBAsyncClientBuilder
.standard()
.withRegion(accountsDynamoDbClientBuilder.getRegion())
.withClientConfiguration(accountsDynamoDbClientBuilder.getClientConfiguration())
.withCredentials(accountsDynamoDbClientBuilder.getCredentials())
.withExecutorFactory(() -> accountsDynamoDbMigrationThreadPool);
AmazonDynamoDBClientBuilder migrationDeletedAccountsDynamoDbClientBuilder = AmazonDynamoDBClientBuilder
.standard()
.withRegion(configuration.getMigrationDeletedAccountsDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) configuration.getMigrationDeletedAccountsDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) configuration.getMigrationDeletedAccountsDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
AmazonDynamoDBClientBuilder migrationRetryAccountsDynamoDbClientBuilder = AmazonDynamoDBClientBuilder
.standard()
.withRegion(configuration.getMigrationRetryAccountsDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) configuration.getMigrationRetryAccountsDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) configuration.getMigrationRetryAccountsDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
AmazonDynamoDBClientBuilder reportMessageDynamoDbClientBuilder = AmazonDynamoDBClientBuilder
.standard()
.withRegion(configuration.getReportMessageDynamoDbConfiguration().getRegion())
.withClientConfiguration(new ClientConfiguration().withClientExecutionTimeout(((int) configuration.getReportMessageDynamoDbConfiguration().getClientExecutionTimeout().toMillis()))
.withRequestTimeout((int) configuration.getReportMessageDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
DynamoDB messageDynamoDb = new DynamoDB(clientBuilder.build());
DynamoDB preKeysDynamoDb = new DynamoDB(keysDynamoDbClientBuilder.build());
DynamoDB reportMessagesDynamoDb = new DynamoDB(reportMessageDynamoDbClientBuilder.build());
AmazonDynamoDB accountsDynamoDbClient = accountsDynamoDbClientBuilder.build();
AmazonDynamoDBAsync accountsDynamoDbAsyncClient = accountsDynamoDbAsyncClientBuilder.build();
DynamoDbClient reportMessagesDynamoDb = DynamoDbFromConfig.client(configuration.getReportMessageDynamoDbConfiguration(),
software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
DynamoDbClient messageDynamoDb = DynamoDbFromConfig.client(configuration.getMessageDynamoDbConfiguration(),
software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
DynamoDbClient preKeysDynamoDb = DynamoDbFromConfig.client(configuration.getKeysDynamoDbConfiguration(),
software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
DynamoDbClient accountsDynamoDbClient = DynamoDbFromConfig.client(configuration.getAccountsDynamoDbConfiguration(),
software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
DynamoDbAsyncClient accountsDynamoDbAsyncClient = DynamoDbFromConfig.asyncClient(configuration.getAccountsDynamoDbConfiguration(),
software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create(),
accountsDynamoDbMigrationThreadPool);
FaultTolerantRedisCluster cacheCluster = new FaultTolerantRedisCluster("main_cache_cluster", configuration.getCacheClusterConfiguration(), redisClusterClientResources);
@@ -173,14 +127,16 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
ExperimentEnrollmentManager experimentEnrollmentManager = new ExperimentEnrollmentManager(dynamicConfigurationManager);
DynamoDB migrationDeletedAccountsDynamoDb = new DynamoDB(migrationDeletedAccountsDynamoDbClientBuilder.build());
DynamoDB migrationRetryAccountsDynamoDb = new DynamoDB(migrationRetryAccountsDynamoDbClientBuilder.build());
DynamoDbClient migrationDeletedAccountsDynamoDb = DynamoDbFromConfig.client(configuration.getMigrationDeletedAccountsDynamoDbConfiguration(),
software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
DynamoDbClient migrationRetryAccountsDynamoDb = DynamoDbFromConfig.client(configuration.getMigrationRetryAccountsDynamoDbConfiguration(),
software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
MigrationDeletedAccounts migrationDeletedAccounts = new MigrationDeletedAccounts(migrationDeletedAccountsDynamoDb, configuration.getMigrationDeletedAccountsDynamoDbConfiguration().getTableName());
MigrationRetryAccounts migrationRetryAccounts = new MigrationRetryAccounts(migrationRetryAccountsDynamoDb, configuration.getMigrationRetryAccountsDynamoDbConfiguration().getTableName());
Accounts accounts = new Accounts(accountDatabase);
AccountsDynamoDb accountsDynamoDb = new AccountsDynamoDb(accountsDynamoDbClient, accountsDynamoDbAsyncClient, accountsDynamoDbMigrationThreadPool, new DynamoDB(accountsDynamoDbClient), configuration.getAccountsDynamoDbConfiguration().getTableName(), configuration.getAccountsDynamoDbConfiguration().getPhoneNumberTableName(), migrationDeletedAccounts, migrationRetryAccounts);
AccountsDynamoDb accountsDynamoDb = new AccountsDynamoDb(accountsDynamoDbClient, accountsDynamoDbAsyncClient, accountsDynamoDbMigrationThreadPool, configuration.getAccountsDynamoDbConfiguration().getTableName(), configuration.getAccountsDynamoDbConfiguration().getPhoneNumberTableName(), migrationDeletedAccounts, migrationRetryAccounts);
Usernames usernames = new Usernames(accountDatabase);
Profiles profiles = new Profiles(accountDatabase);
ReservedUsernames reservedUsernames = new ReservedUsernames(accountDatabase);