mirror of
https://github.com/signalapp/Signal-Server
synced 2026-04-20 23:58:06 +01:00
Add threadpoool to increase Accounts → Dynamo migration throughput
This commit is contained in:
@@ -13,6 +13,8 @@ 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;
|
||||
@@ -52,7 +54,9 @@ import java.util.Optional;
|
||||
import java.util.concurrent.ArrayBlockingQueue;
|
||||
import java.util.concurrent.BlockingQueue;
|
||||
import java.util.concurrent.ExecutorService;
|
||||
import java.util.concurrent.LinkedBlockingDeque;
|
||||
import java.util.concurrent.ScheduledExecutorService;
|
||||
import java.util.concurrent.ThreadPoolExecutor;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import javax.servlet.DispatcherType;
|
||||
import javax.servlet.FilterRegistration;
|
||||
@@ -286,13 +290,25 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
|
||||
.withRequestTimeout((int) config.getAccountsDynamoDbConfiguration().getClientRequestTimeout().toMillis()))
|
||||
.withCredentials(InstanceProfileCredentialsProvider.getInstance());
|
||||
|
||||
// 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);
|
||||
|
||||
DynamoDB messageDynamoDb = new DynamoDB(messageDynamoDbClientBuilder.build());
|
||||
DynamoDB preKeyDynamoDb = new DynamoDB(keysDynamoDbClientBuilder.build());
|
||||
|
||||
AmazonDynamoDB accountsDynamoDbClient = accountsDynamoDbClientBuilder.build();
|
||||
AmazonDynamoDBAsync accountsDynamodbAsyncClient = accountsDynamoDbAsyncClientBuilder.build();
|
||||
|
||||
Accounts accounts = new Accounts(accountDatabase);
|
||||
AccountsDynamoDb accountsDynamoDb = new AccountsDynamoDb(accountsDynamoDbClient, new DynamoDB(accountsDynamoDbClient), config.getAccountsDynamoDbConfiguration().getTableName(), config.getAccountsDynamoDbConfiguration().getPhoneNumberTableName());
|
||||
AccountsDynamoDb accountsDynamoDb = new AccountsDynamoDb(accountsDynamoDbClient, accountsDynamodbAsyncClient, accountsDynamoDbMigrationThreadPool, new DynamoDB(accountsDynamoDbClient), config.getAccountsDynamoDbConfiguration().getTableName(), config.getAccountsDynamoDbConfiguration().getPhoneNumberTableName());
|
||||
PendingAccounts pendingAccounts = new PendingAccounts(accountDatabase);
|
||||
PendingDevices pendingDevices = new PendingDevices (accountDatabase);
|
||||
Usernames usernames = new Usernames(accountDatabase);
|
||||
|
||||
@@ -8,6 +8,9 @@ public class DynamicAccountsDynamoDbMigrationConfiguration {
|
||||
@JsonProperty
|
||||
boolean backgroundMigrationEnabled;
|
||||
|
||||
@JsonProperty
|
||||
int backgroundMigrationExecutorThreads = 1;
|
||||
|
||||
@JsonProperty
|
||||
boolean deleteEnabled;
|
||||
|
||||
@@ -21,6 +24,10 @@ public class DynamicAccountsDynamoDbMigrationConfiguration {
|
||||
return backgroundMigrationEnabled;
|
||||
}
|
||||
|
||||
public int getBackgroundMigrationExecutorThreads() {
|
||||
return backgroundMigrationExecutorThreads;
|
||||
}
|
||||
|
||||
public void setDeleteEnabled(boolean deleteEnabled) {
|
||||
this.deleteEnabled = deleteEnabled;
|
||||
}
|
||||
|
||||
@@ -2,7 +2,9 @@ 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;
|
||||
@@ -16,17 +18,23 @@ 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;
|
||||
import io.micrometer.core.instrument.Metrics;
|
||||
import io.micrometer.core.instrument.Timer;
|
||||
import java.io.IOException;
|
||||
import java.nio.ByteBuffer;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.Optional;
|
||||
import java.util.UUID;
|
||||
import java.util.concurrent.CompletableFuture;
|
||||
import java.util.concurrent.ThreadPoolExecutor;
|
||||
import java.util.stream.Collectors;
|
||||
import org.whispersystems.textsecuregcm.util.SystemMapper;
|
||||
import org.whispersystems.textsecuregcm.util.UUIDUtil;
|
||||
|
||||
@@ -43,6 +51,9 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
|
||||
|
||||
private final AmazonDynamoDB client;
|
||||
private final Table accountsTable;
|
||||
private final AmazonDynamoDBAsync asyncClient;
|
||||
|
||||
private final ThreadPoolExecutor migrationThreadPool;
|
||||
|
||||
private final String phoneNumbersTableName;
|
||||
|
||||
@@ -52,12 +63,15 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
|
||||
private static final Timer GET_BY_UUID_TIMER = Metrics.timer(name(AccountsDynamoDb.class, "getByUuid"));
|
||||
private static final Timer DELETE_TIMER = Metrics.timer(name(AccountsDynamoDb.class, "delete"));
|
||||
|
||||
public AccountsDynamoDb(AmazonDynamoDB client, DynamoDB dynamoDb, String accountsTableName, String phoneNumbersTableName) {
|
||||
public AccountsDynamoDb(AmazonDynamoDB client, AmazonDynamoDBAsync asyncClient, ThreadPoolExecutor migrationThreadPool, DynamoDB dynamoDb, String accountsTableName, String phoneNumbersTableName) {
|
||||
super(dynamoDb);
|
||||
|
||||
this.client = client;
|
||||
this.accountsTable = dynamoDb.getTable(accountsTableName);
|
||||
this.phoneNumbersTableName = phoneNumbersTableName;
|
||||
|
||||
this.asyncClient = asyncClient;
|
||||
this.migrationThreadPool = migrationThreadPool;
|
||||
}
|
||||
|
||||
@Override
|
||||
@@ -216,7 +230,29 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
|
||||
});
|
||||
}
|
||||
|
||||
public boolean migrate(Account account) {
|
||||
public CompletableFuture<Void> migrate(List<Account> accounts, int threads) {
|
||||
|
||||
migrationThreadPool.setCorePoolSize(threads);
|
||||
migrationThreadPool.setMaximumPoolSize(threads);
|
||||
|
||||
final List<CompletableFuture<?>> futures = accounts.stream()
|
||||
.map(this::migrate)
|
||||
.map(f -> f.whenComplete((migrated, e) -> {
|
||||
if (e == null) {
|
||||
MIGRATED_COUNTER.increment(migrated ? 1 : 0);
|
||||
} else {
|
||||
ERROR_COUNTER.increment();
|
||||
}
|
||||
}))
|
||||
.collect(Collectors.toList());
|
||||
|
||||
return CompletableFuture.allOf(futures.toArray(new CompletableFuture[]{}));
|
||||
}
|
||||
|
||||
private static final Counter MIGRATED_COUNTER = Metrics.counter(name(AccountsDynamoDb.class, "migration", "count"));
|
||||
private static final Counter ERROR_COUNTER = Metrics.counter(name(AccountsDynamoDb.class, "migration", "error"));
|
||||
|
||||
public CompletableFuture<Boolean> migrate(Account account) {
|
||||
try {
|
||||
TransactWriteItem phoneNumberConstraintPut = buildPutWriteItemForPhoneNumberConstraint(account, account.getUuid());
|
||||
|
||||
@@ -233,17 +269,32 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt
|
||||
final TransactWriteItemsRequest request = new TransactWriteItemsRequest()
|
||||
.withTransactItems(phoneNumberConstraintPut, accountPut);
|
||||
|
||||
client.transactWriteItems(request);
|
||||
final CompletableFuture<Boolean> resultFuture = new CompletableFuture<>();
|
||||
|
||||
return true;
|
||||
asyncClient.transactWriteItemsAsync(request,
|
||||
new AsyncHandler<>() {
|
||||
@Override
|
||||
public void onError(Exception exception) {
|
||||
if (exception instanceof TransactionCanceledException) {
|
||||
// account is already migrated
|
||||
resultFuture.complete(false);
|
||||
} else {
|
||||
ERROR_COUNTER.increment();
|
||||
resultFuture.completeExceptionally(exception);
|
||||
}
|
||||
}
|
||||
|
||||
} catch (JsonProcessingException e) {
|
||||
throw new IllegalArgumentException(e);
|
||||
} catch (TransactionCanceledException ignored) {
|
||||
// account is already migrated
|
||||
@Override
|
||||
public void onSuccess(TransactWriteItemsRequest request, TransactWriteItemsResult transactWriteItemsResult) {
|
||||
resultFuture.complete(true);
|
||||
}
|
||||
});
|
||||
|
||||
return resultFuture;
|
||||
|
||||
} catch (Exception e) {
|
||||
return CompletableFuture.failedFuture(e);
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
@VisibleForTesting
|
||||
|
||||
@@ -1,18 +1,12 @@
|
||||
package org.whispersystems.textsecuregcm.storage;
|
||||
|
||||
import static com.codahale.metrics.MetricRegistry.name;
|
||||
|
||||
import io.micrometer.core.instrument.Counter;
|
||||
import io.micrometer.core.instrument.Metrics;
|
||||
import java.util.List;
|
||||
import java.util.Optional;
|
||||
import java.util.UUID;
|
||||
import java.util.concurrent.CompletableFuture;
|
||||
|
||||
public class AccountsDynamoDbMigrator extends AccountDatabaseCrawlerListener {
|
||||
|
||||
private static final Counter MIGRATED_COUNTER = Metrics.counter(name(AccountsDynamoDbMigrator.class, "migrated"));
|
||||
private static final Counter ERROR_COUNTER = Metrics.counter(name(AccountsDynamoDbMigrator.class, "error"));
|
||||
|
||||
private final AccountsDynamoDb accountsDynamoDb;
|
||||
private final DynamicConfigurationManager dynamicConfigurationManager;
|
||||
|
||||
@@ -34,20 +28,14 @@ public class AccountsDynamoDbMigrator extends AccountDatabaseCrawlerListener {
|
||||
@Override
|
||||
protected void onCrawlChunk(Optional<UUID> fromUuid, List<Account> chunkAccounts) {
|
||||
|
||||
if (!dynamicConfigurationManager.getConfiguration().getAccountsDynamoDbMigrationConfiguration().isBackgroundMigrationEnabled()) {
|
||||
if (!dynamicConfigurationManager.getConfiguration().getAccountsDynamoDbMigrationConfiguration()
|
||||
.isBackgroundMigrationEnabled()) {
|
||||
return;
|
||||
}
|
||||
|
||||
for (Account account : chunkAccounts) {
|
||||
try {
|
||||
final boolean migrated = accountsDynamoDb.migrate(account);
|
||||
if (migrated) {
|
||||
MIGRATED_COUNTER.increment();
|
||||
}
|
||||
} catch (final Exception e) {
|
||||
final CompletableFuture<Void> migrationBatch = accountsDynamoDb.migrate(chunkAccounts,
|
||||
dynamicConfigurationManager.getConfiguration().getAccountsDynamoDbMigrationConfiguration().getBackgroundMigrationExecutorThreads());
|
||||
|
||||
ERROR_COUNTER.increment();
|
||||
}
|
||||
}
|
||||
migrationBatch.join();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,6 +10,8 @@ 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;
|
||||
@@ -20,6 +22,9 @@ import io.dropwizard.setup.Environment;
|
||||
import io.lettuce.core.resource.ClientResources;
|
||||
import java.util.Optional;
|
||||
import java.util.concurrent.ExecutorService;
|
||||
import java.util.concurrent.LinkedBlockingDeque;
|
||||
import java.util.concurrent.ThreadPoolExecutor;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import net.sourceforge.argparse4j.inf.Namespace;
|
||||
import net.sourceforge.argparse4j.inf.Subparser;
|
||||
import org.jdbi.v3.core.Jdbi;
|
||||
@@ -111,11 +116,21 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
|
||||
.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);
|
||||
|
||||
DynamoDB messageDynamoDb = new DynamoDB(clientBuilder.build());
|
||||
DynamoDB preKeysDynamoDb = new DynamoDB(keysDynamoDbClientBuilder.build());
|
||||
|
||||
AmazonDynamoDB accountsDynamoDbClient = accountsDynamoDbClientBuilder.build();
|
||||
AmazonDynamoDBAsync accountsDynamoDbAsyncClient = accountsDynamoDbAsyncClientBuilder.build();
|
||||
|
||||
FaultTolerantRedisCluster cacheCluster = new FaultTolerantRedisCluster("main_cache_cluster", configuration.getCacheClusterConfiguration(), redisClusterClientResources);
|
||||
|
||||
@@ -132,7 +147,7 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
|
||||
ExperimentEnrollmentManager experimentEnrollmentManager = new ExperimentEnrollmentManager(dynamicConfigurationManager);
|
||||
|
||||
Accounts accounts = new Accounts(accountDatabase);
|
||||
AccountsDynamoDb accountsDynamoDb = new AccountsDynamoDb(accountsDynamoDbClient, new DynamoDB(accountsDynamoDbClient), configuration.getAccountsDynamoDbConfiguration().getTableName(), configuration.getAccountsDynamoDbConfiguration().getPhoneNumberTableName());
|
||||
AccountsDynamoDb accountsDynamoDb = new AccountsDynamoDb(accountsDynamoDbClient, accountsDynamoDbAsyncClient, accountsDynamoDbMigrationThreadPool, new DynamoDB(accountsDynamoDbClient), configuration.getAccountsDynamoDbConfiguration().getTableName(), configuration.getAccountsDynamoDbConfiguration().getPhoneNumberTableName());
|
||||
Usernames usernames = new Usernames(accountDatabase);
|
||||
Profiles profiles = new Profiles(accountDatabase);
|
||||
ReservedUsernames reservedUsernames = new ReservedUsernames(accountDatabase);
|
||||
|
||||
Reference in New Issue
Block a user