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

Make `PhoneNumberIdentifiers` operations asynchronous

This commit is contained in:
Jon Chambers 2024-11-22 12:33:09 -05:00 committed by Jon Chambers
parent 0023cb2521
commit 8c9cc4cce5
11 changed files with 70 additions and 77 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -404,7 +404,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
config.getDynamoDbTables().getAccounts().getUsedLinkDeviceTokensTableName());
ClientReleases clientReleases = new ClientReleases(dynamoDbAsyncClient,
config.getDynamoDbTables().getClientReleases().getTableName());
PhoneNumberIdentifiers phoneNumberIdentifiers = new PhoneNumberIdentifiers(dynamoDbClient,
PhoneNumberIdentifiers phoneNumberIdentifiers = new PhoneNumberIdentifiers(dynamoDbAsyncClient,
config.getDynamoDbTables().getPhoneNumberIdentifiers().getTableName());
Profiles profiles = new Profiles(dynamoDbClient, dynamoDbAsyncClient,
config.getDynamoDbTables().getProfiles().getTableName());

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

@ -271,7 +271,7 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
@Nullable final String userAgent) throws InterruptedException {
final Account account = new Account();
final UUID phoneNumberIdentifier = phoneNumberIdentifiers.getPhoneNumberIdentifier(number);
final UUID phoneNumberIdentifier = phoneNumberIdentifiers.getPhoneNumberIdentifier(number).join();
return createTimer.record(() -> {
accountLockManager.withLock(List.of(number), List.of(phoneNumberIdentifier), () -> {
@ -651,7 +651,7 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
validateDevices(account, pniSignedPreKeys, pniPqLastResortPreKeys, pniRegistrationIds);
final AtomicReference<Account> updatedAccount = new AtomicReference<>();
final UUID targetPhoneNumberIdentifier = phoneNumberIdentifiers.getPhoneNumberIdentifier(targetNumber);
final UUID targetPhoneNumberIdentifier = phoneNumberIdentifiers.getPhoneNumberIdentifier(targetNumber).join();
accountLockManager.withLock(List.of(account.getNumber(), targetNumber),
List.of(account.getPhoneNumberIdentifier(), targetPhoneNumberIdentifier), () -> {
@ -1207,7 +1207,7 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
}
public UUID getPhoneNumberIdentifier(String e164) {
return phoneNumberIdentifiers.getPhoneNumberIdentifier(e164);
return phoneNumberIdentifiers.getPhoneNumberIdentifier(e164).join();
}
public Optional<UUID> findRecentlyDeletedAccountIdentifier(final String e164) {

105
service/src/main/java/org/whispersystems/textsecuregcm/storage/PhoneNumberIdentifiers.java
View File

@ -13,15 +13,13 @@ import io.micrometer.core.instrument.Timer;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import org.whispersystems.textsecuregcm.util.AttributeValues;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
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.UpdateItemRequest;
import software.amazon.awssdk.services.dynamodb.model.UpdateItemResponse;
/**
* Manages a global, persistent mapping of phone numbers to phone number identifiers regardless of whether those
@ -29,7 +27,7 @@ import software.amazon.awssdk.services.dynamodb.model.UpdateItemResponse;
*/
public class PhoneNumberIdentifiers {
private final DynamoDbClient dynamoDbClient;
private final DynamoDbAsyncClient dynamoDbClient;
private final String tableName;
@VisibleForTesting
@ -42,7 +40,7 @@ public class PhoneNumberIdentifiers {
private static final Timer GET_PNI_TIMER = Metrics.timer(name(PhoneNumberIdentifiers.class, "get"));
private static final Timer SET_PNI_TIMER = Metrics.timer(name(PhoneNumberIdentifiers.class, "set"));
public PhoneNumberIdentifiers(final DynamoDbClient dynamoDbClient, final String tableName) {
public PhoneNumberIdentifiers(final DynamoDbAsyncClient dynamoDbClient, final String tableName) {
this.dynamoDbClient = dynamoDbClient;
this.tableName = tableName;
}
@ -53,67 +51,62 @@ public class PhoneNumberIdentifiers {
* @param phoneNumber the phone number for which to retrieve a phone number identifier
* @return the phone number identifier associated with the given phone number
*/
public UUID getPhoneNumberIdentifier(final String phoneNumber) {
final GetItemResponse response = GET_PNI_TIMER.record(() -> dynamoDbClient.getItem(GetItemRequest.builder()
public CompletableFuture<UUID> getPhoneNumberIdentifier(final String phoneNumber) {
final Timer.Sample sample = Timer.start();
return dynamoDbClient.getItem(GetItemRequest.builder()
.tableName(tableName)
.key(Map.of(KEY_E164, AttributeValues.fromString(phoneNumber)))
.projectionExpression(ATTR_PHONE_NUMBER_IDENTIFIER)
.build()));
final UUID phoneNumberIdentifier;
if (response.hasItem()) {
phoneNumberIdentifier = AttributeValues.getUUID(response.item(), ATTR_PHONE_NUMBER_IDENTIFIER, null);
} else {
phoneNumberIdentifier = generatePhoneNumberIdentifierIfNotExists(phoneNumber);
}
if (phoneNumberIdentifier == null) {
throw new RuntimeException("Could not retrieve phone number identifier from stored item");
}
return phoneNumberIdentifier;
.build())
.thenCompose(response -> response.hasItem()
? CompletableFuture.completedFuture(AttributeValues.getUUID(response.item(), ATTR_PHONE_NUMBER_IDENTIFIER, null))
: generatePhoneNumberIdentifierIfNotExists(phoneNumber))
.whenComplete((ignored, throwable) -> sample.stop(GET_PNI_TIMER));
}
public Optional<String> getPhoneNumber(final UUID phoneNumberIdentifier) {
final QueryResponse response = dynamoDbClient.query(QueryRequest.builder()
.tableName(tableName)
.indexName(INDEX_NAME)
.keyConditionExpression("#pni = :pni")
.projectionExpression("#phone_number")
.expressionAttributeNames(Map.of(
"#phone_number", KEY_E164,
"#pni", ATTR_PHONE_NUMBER_IDENTIFIER
))
.expressionAttributeValues(Map.of(
":pni", AttributeValues.fromUUID(phoneNumberIdentifier)
))
.build());
public CompletableFuture<Optional<String>> getPhoneNumber(final UUID phoneNumberIdentifier) {
return dynamoDbClient.query(QueryRequest.builder()
.tableName(tableName)
.indexName(INDEX_NAME)
.keyConditionExpression("#pni = :pni")
.projectionExpression("#phone_number")
.expressionAttributeNames(Map.of(
"#phone_number", KEY_E164,
"#pni", ATTR_PHONE_NUMBER_IDENTIFIER
))
.expressionAttributeValues(Map.of(
":pni", AttributeValues.fromUUID(phoneNumberIdentifier)
))
.build())
.thenApply(response -> {
if (response.count() == 0) {
return Optional.empty();
}
if (response.count() == 0) {
return Optional.empty();
}
if (response.count() > 1) {
throw new RuntimeException(
"Impossible result: more than one phone number returned for PNI: " + phoneNumberIdentifier);
}
if (response.count() > 1) {
throw new RuntimeException(
"Impossible result: more than one phone number returned for PNI: " + phoneNumberIdentifier);
}
return Optional.ofNullable(response.items().get(0).get(KEY_E164).s());
return Optional.ofNullable(response.items().getFirst().get(KEY_E164).s());
});
}
@VisibleForTesting
UUID generatePhoneNumberIdentifierIfNotExists(final String phoneNumber) {
final UpdateItemResponse response = SET_PNI_TIMER.record(() -> dynamoDbClient.updateItem(UpdateItemRequest.builder()
.tableName(tableName)
.key(Map.of(KEY_E164, AttributeValues.fromString(phoneNumber)))
.updateExpression("SET #pni = if_not_exists(#pni, :pni)")
.expressionAttributeNames(Map.of("#pni", ATTR_PHONE_NUMBER_IDENTIFIER))
.expressionAttributeValues(Map.of(":pni", AttributeValues.fromUUID(UUID.randomUUID())))
.returnValues(ReturnValue.ALL_NEW)
.build()));
CompletableFuture<UUID> generatePhoneNumberIdentifierIfNotExists(final String phoneNumber) {
final Timer.Sample sample = Timer.start();
return AttributeValues.getUUID(response.attributes(), ATTR_PHONE_NUMBER_IDENTIFIER, null);
return dynamoDbClient.updateItem(UpdateItemRequest.builder()
.tableName(tableName)
.key(Map.of(KEY_E164, AttributeValues.fromString(phoneNumber)))
.updateExpression("SET #pni = if_not_exists(#pni, :pni)")
.expressionAttributeNames(Map.of("#pni", ATTR_PHONE_NUMBER_IDENTIFIER))
.expressionAttributeValues(Map.of(":pni", AttributeValues.fromUUID(UUID.randomUUID())))
.returnValues(ReturnValue.ALL_NEW)
.build())
.thenApply(response -> AttributeValues.getUUID(response.attributes(), ATTR_PHONE_NUMBER_IDENTIFIER, null))
.whenComplete((ignored, throwable) -> sample.stop(SET_PNI_TIMER));
}
}

2
service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java
View File

@ -185,7 +185,7 @@ record CommandDependencies(
configuration.getDynamoDbTables().getAccounts().getUsernamesTableName(),
configuration.getDynamoDbTables().getDeletedAccounts().getTableName(),
configuration.getDynamoDbTables().getAccounts().getUsedLinkDeviceTokensTableName());
PhoneNumberIdentifiers phoneNumberIdentifiers = new PhoneNumberIdentifiers(dynamoDbClient,
PhoneNumberIdentifiers phoneNumberIdentifiers = new PhoneNumberIdentifiers(dynamoDbAsyncClient,
configuration.getDynamoDbTables().getPhoneNumberIdentifiers().getTableName());
Profiles profiles = new Profiles(dynamoDbClient, dynamoDbAsyncClient,
configuration.getDynamoDbTables().getProfiles().getTableName());

2
service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountCreationDeletionIntegrationTest.java
View File

@ -125,7 +125,7 @@ public class AccountCreationDeletionIntegrationTest {
when(svr2Client.deleteBackups(any())).thenReturn(CompletableFuture.completedFuture(null));
final PhoneNumberIdentifiers phoneNumberIdentifiers =
new PhoneNumberIdentifiers(DYNAMO_DB_EXTENSION.getDynamoDbClient(),
new PhoneNumberIdentifiers(DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
DynamoDbExtensionSchema.Tables.PNI.tableName());
final MessagesManager messagesManager = mock(MessagesManager.class);

2
service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerChangeNumberIntegrationTest.java
View File

@ -119,7 +119,7 @@ class AccountsManagerChangeNumberIntegrationTest {
disconnectionRequestManager = mock(DisconnectionRequestManager.class);
final PhoneNumberIdentifiers phoneNumberIdentifiers =
new PhoneNumberIdentifiers(DYNAMO_DB_EXTENSION.getDynamoDbClient(), Tables.PNI.tableName());
new PhoneNumberIdentifiers(DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(), Tables.PNI.tableName());
final MessagesManager messagesManager = mock(MessagesManager.class);
when(messagesManager.clear(any())).thenReturn(CompletableFuture.completedFuture(null));

2
service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerConcurrentModificationIntegrationTest.java
View File

@ -122,7 +122,7 @@ class AccountsManagerConcurrentModificationIntegrationTest {
final PhoneNumberIdentifiers phoneNumberIdentifiers = mock(PhoneNumberIdentifiers.class);
when(phoneNumberIdentifiers.getPhoneNumberIdentifier(anyString()))
.thenAnswer((Answer<UUID>) invocation -> UUID.randomUUID());
.thenAnswer((Answer<CompletableFuture<UUID>>) invocation -> CompletableFuture.completedFuture(UUID.randomUUID()));
accountsManager = new AccountsManager(
accounts,

4
service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerTest.java
View File

@ -195,9 +195,9 @@ class AccountsManagerTest {
final PhoneNumberIdentifiers phoneNumberIdentifiers = mock(PhoneNumberIdentifiers.class);
phoneNumberIdentifiersByE164 = new HashMap<>();
when(phoneNumberIdentifiers.getPhoneNumberIdentifier(anyString())).thenAnswer((Answer<UUID>) invocation -> {
when(phoneNumberIdentifiers.getPhoneNumberIdentifier(anyString())).thenAnswer((Answer<CompletableFuture<UUID>>) invocation -> {
final String number = invocation.getArgument(0, String.class);
return phoneNumberIdentifiersByE164.computeIfAbsent(number, n -> UUID.randomUUID());
return CompletableFuture.completedFuture(phoneNumberIdentifiersByE164.computeIfAbsent(number, n -> UUID.randomUUID()));
});
@SuppressWarnings("unchecked") final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager =

2
service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java
View File

@ -128,7 +128,7 @@ class AccountsManagerUsernameIntegrationTest {
});
final PhoneNumberIdentifiers phoneNumberIdentifiers =
new PhoneNumberIdentifiers(DYNAMO_DB_EXTENSION.getDynamoDbClient(), Tables.PNI.tableName());
new PhoneNumberIdentifiers(DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(), Tables.PNI.tableName());
final MessagesManager messageManager = mock(MessagesManager.class);
final ProfilesManager profileManager = mock(ProfilesManager.class);

2
service/src/test/java/org/whispersystems/textsecuregcm/storage/AddRemoveDeviceIntegrationTest.java
View File

@ -125,7 +125,7 @@ public class AddRemoveDeviceIntegrationTest {
when(svr2Client.deleteBackups(any())).thenReturn(CompletableFuture.completedFuture(null));
final PhoneNumberIdentifiers phoneNumberIdentifiers =
new PhoneNumberIdentifiers(DYNAMO_DB_EXTENSION.getDynamoDbClient(),
new PhoneNumberIdentifiers(DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
DynamoDbExtensionSchema.Tables.PNI.tableName());
messagesManager = mock(MessagesManager.class);

18
service/src/test/java/org/whispersystems/textsecuregcm/storage/PhoneNumberIdentifiersTest.java
View File

@ -25,7 +25,7 @@ class PhoneNumberIdentifiersTest {
@BeforeEach
void setUp() {
phoneNumberIdentifiers = new PhoneNumberIdentifiers(DYNAMO_DB_EXTENSION.getDynamoDbClient(),
phoneNumberIdentifiers = new PhoneNumberIdentifiers(DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
Tables.PNI.tableName());
}
@ -34,28 +34,28 @@ class PhoneNumberIdentifiersTest {
final String number = "+18005551234";
final String differentNumber = "+18005556789";
final UUID firstPni = phoneNumberIdentifiers.getPhoneNumberIdentifier(number);
final UUID secondPni = phoneNumberIdentifiers.getPhoneNumberIdentifier(number);
final UUID firstPni = phoneNumberIdentifiers.getPhoneNumberIdentifier(number).join();
final UUID secondPni = phoneNumberIdentifiers.getPhoneNumberIdentifier(number).join();
assertEquals(firstPni, secondPni);
assertNotEquals(firstPni, phoneNumberIdentifiers.getPhoneNumberIdentifier(differentNumber));
assertNotEquals(firstPni, phoneNumberIdentifiers.getPhoneNumberIdentifier(differentNumber).join());
}
@Test
void generatePhoneNumberIdentifierIfNotExists() {
final String number = "+18005551234";
assertEquals(phoneNumberIdentifiers.generatePhoneNumberIdentifierIfNotExists(number),
phoneNumberIdentifiers.generatePhoneNumberIdentifierIfNotExists(number));
assertEquals(phoneNumberIdentifiers.generatePhoneNumberIdentifierIfNotExists(number).join(),
phoneNumberIdentifiers.generatePhoneNumberIdentifierIfNotExists(number).join());
}
@Test
void getPhoneNumber() {
final String number = "+18005551234";
assertFalse(phoneNumberIdentifiers.getPhoneNumber(UUID.randomUUID()).isPresent());
assertFalse(phoneNumberIdentifiers.getPhoneNumber(UUID.randomUUID()).join().isPresent());
final UUID pni = phoneNumberIdentifiers.getPhoneNumberIdentifier(number);
assertEquals(Optional.of(number), phoneNumberIdentifiers.getPhoneNumber(pni));
final UUID pni = phoneNumberIdentifiers.getPhoneNumberIdentifier(number).join();
assertEquals(Optional.of(number), phoneNumberIdentifiers.getPhoneNumber(pni).join());
}
}