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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@ -119,7 +119,7 @@ class AccountsManagerChangeNumberIntegrationTest {
disconnectionRequestManager = mock(DisconnectionRequestManager.class); disconnectionRequestManager = mock(DisconnectionRequestManager.class);
final PhoneNumberIdentifiers phoneNumberIdentifiers = 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); final MessagesManager messagesManager = mock(MessagesManager.class);
when(messagesManager.clear(any())).thenReturn(CompletableFuture.completedFuture(null)); when(messagesManager.clear(any())).thenReturn(CompletableFuture.completedFuture(null));

View File

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

View File

@ -195,9 +195,9 @@ class AccountsManagerTest {
final PhoneNumberIdentifiers phoneNumberIdentifiers = mock(PhoneNumberIdentifiers.class); final PhoneNumberIdentifiers phoneNumberIdentifiers = mock(PhoneNumberIdentifiers.class);
phoneNumberIdentifiersByE164 = new HashMap<>(); 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); 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 = @SuppressWarnings("unchecked") final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager =

View File

@ -128,7 +128,7 @@ class AccountsManagerUsernameIntegrationTest {
}); });
final PhoneNumberIdentifiers phoneNumberIdentifiers = 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 MessagesManager messageManager = mock(MessagesManager.class);
final ProfilesManager profileManager = mock(ProfilesManager.class); final ProfilesManager profileManager = mock(ProfilesManager.class);

View File

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

View File

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