diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java index c414810ac..bf5a3eb2d 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java @@ -404,7 +404,7 @@ public class WhisperServerService extends Application 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 implemen validateDevices(account, pniSignedPreKeys, pniPqLastResortPreKeys, pniRegistrationIds); final AtomicReference 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 implemen } public UUID getPhoneNumberIdentifier(String e164) { - return phoneNumberIdentifiers.getPhoneNumberIdentifier(e164); + return phoneNumberIdentifiers.getPhoneNumberIdentifier(e164).join(); } public Optional findRecentlyDeletedAccountIdentifier(final String e164) { diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/PhoneNumberIdentifiers.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/PhoneNumberIdentifiers.java index 384b7c71f..2a4b0f087 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/PhoneNumberIdentifiers.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/PhoneNumberIdentifiers.java @@ -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 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 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> 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 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)); } } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java index df2c3e155..745451b07 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java @@ -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()); diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountCreationDeletionIntegrationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountCreationDeletionIntegrationTest.java index 146d849cc..8dc5ae284 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountCreationDeletionIntegrationTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountCreationDeletionIntegrationTest.java @@ -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); diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerChangeNumberIntegrationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerChangeNumberIntegrationTest.java index 74c8227f7..180d816c4 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerChangeNumberIntegrationTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerChangeNumberIntegrationTest.java @@ -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)); diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerConcurrentModificationIntegrationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerConcurrentModificationIntegrationTest.java index cc981c16c..6d9706a04 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerConcurrentModificationIntegrationTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerConcurrentModificationIntegrationTest.java @@ -122,7 +122,7 @@ class AccountsManagerConcurrentModificationIntegrationTest { final PhoneNumberIdentifiers phoneNumberIdentifiers = mock(PhoneNumberIdentifiers.class); when(phoneNumberIdentifiers.getPhoneNumberIdentifier(anyString())) - .thenAnswer((Answer) invocation -> UUID.randomUUID()); + .thenAnswer((Answer>) invocation -> CompletableFuture.completedFuture(UUID.randomUUID())); accountsManager = new AccountsManager( accounts, diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerTest.java index 7a2a3be1c..7e43a6860 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerTest.java @@ -195,9 +195,9 @@ class AccountsManagerTest { final PhoneNumberIdentifiers phoneNumberIdentifiers = mock(PhoneNumberIdentifiers.class); phoneNumberIdentifiersByE164 = new HashMap<>(); - when(phoneNumberIdentifiers.getPhoneNumberIdentifier(anyString())).thenAnswer((Answer) invocation -> { + when(phoneNumberIdentifiers.getPhoneNumberIdentifier(anyString())).thenAnswer((Answer>) 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 dynamicConfigurationManager = diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java index 7c6e1868c..0b25deb32 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java @@ -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); diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AddRemoveDeviceIntegrationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AddRemoveDeviceIntegrationTest.java index 96f6aa0b8..7d08be0c8 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AddRemoveDeviceIntegrationTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AddRemoveDeviceIntegrationTest.java @@ -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); diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/PhoneNumberIdentifiersTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/PhoneNumberIdentifiersTest.java index 174465a3b..e0ba6d0df 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/PhoneNumberIdentifiersTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/PhoneNumberIdentifiersTest.java @@ -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()); } }