Add parallel pathways for getting accounts asyncronously to `Accounts`

service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java

@@ -23,6 +23,7 @@ import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.TimeUnit;
@@ -631,12 +632,23 @@ public class Accounts extends AbstractDynamoDbStore {
         GET_BY_NUMBER_TIMER, phoneNumberConstraintTableName, ATTR_ACCOUNT_E164, AttributeValues.fromString(number));
   }
 
+  @Nonnull
+  public CompletableFuture<Optional<Account>> getByE164Async(final String number) {
+    return getByIndirectLookupAsync(
+        GET_BY_NUMBER_TIMER, phoneNumberConstraintTableName, ATTR_ACCOUNT_E164, AttributeValues.fromString(number));
+  }
+
   @Nonnull
   public Optional<Account> getByPhoneNumberIdentifier(final UUID phoneNumberIdentifier) {
     return getByIndirectLookup(
         GET_BY_PNI_TIMER, phoneNumberIdentifierConstraintTableName, ATTR_PNI_UUID, AttributeValues.fromUUID(phoneNumberIdentifier));
   }
 
+  @Nonnull
+  public CompletableFuture<Optional<Account>> getByPhoneNumberIdentifierAsync(final UUID phoneNumberIdentifier) {
+    return getByIndirectLookupAsync(GET_BY_PNI_TIMER, phoneNumberIdentifierConstraintTableName, ATTR_PNI_UUID, AttributeValues.fromUUID(phoneNumberIdentifier));
+  }
+
   @Nonnull
   public Optional<Account> getByUsernameHash(final byte[] usernameHash) {
     return getByIndirectLookup(
@@ -662,6 +674,13 @@ public class Accounts extends AbstractDynamoDbStore {
             .map(Accounts::fromItem)));
   }
 
+  @Nonnull
+  public CompletableFuture<Optional<Account>> getByAccountIdentifierAsync(final UUID uuid) {
+    return record(GET_BY_UUID_TIMER, () -> itemByKeyAsync(accountsTableName, KEY_ACCOUNT_UUID, AttributeValues.fromUUID(uuid))
+        .thenApply(maybeItem -> maybeItem.map(Accounts::fromItem)))
+        .toCompletableFuture();
+  }
+
   public void delete(final UUID uuid) {
     DELETE_TIMER.record(() -> getByAccountIdentifier(uuid).ifPresent(account -> {
 
@@ -724,6 +743,16 @@ public class Accounts extends AbstractDynamoDbStore {
     return getByIndirectLookup(timer, tableName, keyName, keyValue, i -> true);
   }
 
+  @Nonnull
+  private CompletableFuture<Optional<Account>> getByIndirectLookupAsync(
+      final Timer timer,
+      final String tableName,
+      final String keyName,
+      final AttributeValue keyValue) {
+
+    return getByIndirectLookupAsync(timer, tableName, keyName, keyValue, i -> true);
+  }
+
   @Nonnull
   private Optional<Account> getByIndirectLookup(
       final Timer timer,
@@ -739,6 +768,24 @@ public class Accounts extends AbstractDynamoDbStore {
         .map(Accounts::fromItem)));
   }
 
+  @Nonnull
+  private CompletableFuture<Optional<Account>> getByIndirectLookupAsync(
+      final Timer timer,
+      final String tableName,
+      final String keyName,
+      final AttributeValue keyValue,
+      final Predicate<? super Map<String, AttributeValue>> predicate) {
+
+    return record(timer, () -> itemByKeyAsync(tableName, keyName, keyValue)
+        .thenCompose(maybeItem -> maybeItem
+            .filter(predicate)
+            .map(item -> item.get(KEY_ACCOUNT_UUID))
+            .map(uuid -> itemByKeyAsync(accountsTableName, KEY_ACCOUNT_UUID, uuid)
+                .thenApply(maybeAccountItem -> maybeAccountItem.map(Accounts::fromItem)))
+            .orElse(CompletableFuture.completedFuture(Optional.empty()))))
+        .toCompletableFuture();
+  }
+
   @Nonnull
   private Optional<Map<String, AttributeValue>> itemByKey(final String table, final String keyName, final AttributeValue keyValue) {
     final GetItemResponse response = db().getItem(GetItemRequest.builder()
@@ -749,6 +796,16 @@ public class Accounts extends AbstractDynamoDbStore {
     return Optional.ofNullable(response.item()).filter(m -> !m.isEmpty());
   }
 
+  @Nonnull
+  private CompletableFuture<Optional<Map<String, AttributeValue>>> itemByKeyAsync(final String table, final String keyName, final AttributeValue keyValue) {
+    return asyncClient.getItem(GetItemRequest.builder()
+            .tableName(table)
+            .key(Map.of(keyName, keyValue))
+            .consistentRead(true)
+            .build())
+        .thenApply(response -> Optional.ofNullable(response.item()).filter(item -> !item.isEmpty()));
+  }
+
   @Nonnull
   private Optional<Map<String, AttributeValue>> itemByGsiKey(final String table, final String indexName, final String keyName, final AttributeValue keyValue) {
     final QueryResponse response = db().query(QueryRequest.builder()

service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java

@@ -41,6 +41,7 @@ import java.util.stream.Collectors;
 import org.apache.commons.lang3.RandomUtils;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
@@ -73,6 +74,7 @@ import software.amazon.awssdk.services.dynamodb.model.TransactionCanceledExcepti
 import software.amazon.awssdk.services.dynamodb.model.TransactionConflictException;
 import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
 
+@Timeout(value = 10, threadMode = Timeout.ThreadMode.SEPARATE_THREAD)
 class AccountsTest {
 
   private static final String BASE_64_URL_USERNAME_HASH_1 = "9p6Tip7BFefFOJzv4kv4GyXEYsBVfk_WbjNejdlOvQE";
@@ -573,6 +575,44 @@ class AccountsTest {
     assertThat(retrieved.isPresent()).isFalse();
   }
 
+  @Test
+  void getByAccountIdentifierAsync() {
+    assertThat(accounts.getByAccountIdentifierAsync(UUID.randomUUID()).join()).isEmpty();
+
+    final Account account =
+        generateAccount("+14151112222", UUID.randomUUID(), UUID.randomUUID(), List.of(generateDevice(1)));
+
+    accounts.create(account);
+
+    assertThat(accounts.getByAccountIdentifierAsync(account.getUuid()).join()).isPresent();
+  }
+
+  @Test
+  void getByPhoneNumberIdentifierAsync() {
+    assertThat(accounts.getByPhoneNumberIdentifierAsync(UUID.randomUUID()).join()).isEmpty();
+
+    final Account account =
+        generateAccount("+14151112222", UUID.randomUUID(), UUID.randomUUID(), List.of(generateDevice(1)));
+
+    accounts.create(account);
+
+    assertThat(accounts.getByPhoneNumberIdentifierAsync(account.getPhoneNumberIdentifier()).join()).isPresent();
+  }
+
+  @Test
+  void getByE164Async() {
+    final String e164 = "+14151112222";
+
+    assertThat(accounts.getByE164Async(e164).join()).isEmpty();
+
+    final Account account =
+        generateAccount(e164, UUID.randomUUID(), UUID.randomUUID(), List.of(generateDevice(1)));
+
+    accounts.create(account);
+
+    assertThat(accounts.getByE164Async(e164).join()).isPresent();
+  }
+
   @Test
   void testCanonicallyDiscoverableSet() {
     Device device = generateDevice(1);