From cd49ea43c003b0789c1185796b0ef32f3dbc9437 Mon Sep 17 00:00:00 2001
From: Chris Eager <chris@signal.org>
Date: Tue, 14 Sep 2021 15:51:42 -0700
Subject: [PATCH] Use `queryPaginator` when loading messages

---
 .../textsecuregcm/storage/MessagesDynamoDb.java            | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDb.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDb.java
index 771cbc94a..e97647e72 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDb.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDb.java
@@ -129,8 +129,13 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
           .limit(numberOfMessagesToFetch)
           .build();
       List<OutgoingMessageEntity> messageEntities = new ArrayList<>(numberOfMessagesToFetch);
-      for (Map<String, AttributeValue> message : db().query(queryRequest).items()) {
+      for (Map<String, AttributeValue> message : db().queryPaginator(queryRequest).items()) {
         messageEntities.add(convertItemToOutgoingMessageEntity(message));
+        if (messageEntities.size() == numberOfMessagesToFetch) {
+          // queryPaginator() uses limit() as the page size, not as an absolute limit
+          // …but a page might be smaller than limit, because a page is capped at 1 MB
+          break;
+        }
       }
       return messageEntities;
     });