diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AbstractDynamoDbStore.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AbstractDynamoDbStore.java index 401b3979b..bfd3baf62 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AbstractDynamoDbStore.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AbstractDynamoDbStore.java @@ -7,7 +7,13 @@ package org.whispersystems.textsecuregcm.storage; import com.amazonaws.services.dynamodbv2.document.BatchWriteItemOutcome; import com.amazonaws.services.dynamodbv2.document.DynamoDB; +import com.amazonaws.services.dynamodbv2.document.Item; +import com.amazonaws.services.dynamodbv2.document.ItemCollection; +import com.amazonaws.services.dynamodbv2.document.Page; +import com.amazonaws.services.dynamodbv2.document.QueryOutcome; +import com.amazonaws.services.dynamodbv2.document.Table; import com.amazonaws.services.dynamodbv2.document.TableWriteItems; +import com.amazonaws.services.dynamodbv2.document.spec.QuerySpec; import io.micrometer.core.instrument.Counter; import io.micrometer.core.instrument.Timer; import org.slf4j.Logger; @@ -58,6 +64,22 @@ public class AbstractDynamoDbStore { } } + protected long countItemsMatchingQuery(final Table table, final QuerySpec querySpec) { + // This is very confusing, but does appear to be the intended behavior. See: + // + // - https://github.com/aws/aws-sdk-java/issues/693 + // - https://github.com/aws/aws-sdk-java/issues/915 + // - https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Query.html#Query.Count + + long matchingItems = 0; + + for (final Page page : table.query(querySpec).pages()) { + matchingItems += page.getLowLevelResult().getQueryResult().getCount(); + } + + return matchingItems; + } + static void writeInBatches(final Iterable items, final Consumer> action) { final List batch = new ArrayList<>(DYNAMO_DB_MAX_BATCH_SIZE); diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDb.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDb.java index 1063a325d..65d242f39 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDb.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDb.java @@ -128,11 +128,7 @@ public class KeysDynamoDb extends AbstractDynamoDbStore implements PreKeyStore { .withSelect(Select.COUNT) .withConsistentRead(false); - // This is very confusing, but does appear to be the intended behavior. See: - // - // - https://github.com/aws/aws-sdk-java/issues/693 - // - https://github.com/aws/aws-sdk-java/issues/915 - return table.query(querySpec).firstPage().getLowLevelResult().getQueryResult().getCount(); + return (int)countItemsMatchingQuery(table, querySpec); }); }