From e12ba6b15b49c1ee601aa79442937e2a9dcf4a4d Mon Sep 17 00:00:00 2001
From: Chris Eager <chris@signal.org>
Date: Tue, 25 Mar 2025 10:31:33 -0500
Subject: [PATCH] Reduce page size in MessagesDynamoDb#mayHaveUrgentMessages

---
 .../storage/MessagesDynamoDb.java             |  7 ++---
 .../storage/MessagesDynamoDbTest.java         | 26 +++++++++++++++----
 2 files changed, 25 insertions(+), 8 deletions(-)

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 5a9bf8937..e884e55c3 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDb.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDb.java
@@ -56,6 +56,9 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
   @VisibleForTesting
   static final String LOCAL_INDEX_MESSAGE_UUID_KEY_SORT = "U";
 
+  @VisibleForTesting
+  static final int MAY_HAVE_URGENT_MESSAGES_QUERY_LIMIT = 20;
+
   private static final String KEY_TTL = "E";
   private static final String KEY_ENVELOPE_BYTES = "EB";
 
@@ -67,8 +70,6 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
   private final ExecutorService messageDeletionExecutor;
   private final Scheduler messageDeletionScheduler;
 
-  private static final CompletableFuture<?>[] EMPTY_FUTURE_ARRAY = new CompletableFuture<?>[0];
-
   private static final Logger logger = LoggerFactory.getLogger(MessagesDynamoDb.class);
 
   public MessagesDynamoDb(DynamoDbClient dynamoDb, DynamoDbAsyncClient dynamoDbAsyncClient, String tableName,
@@ -126,7 +127,7 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
   }
 
   public CompletableFuture<Boolean> mayHaveUrgentMessages(final UUID accountIdentifier, final Device device) {
-    return Flux.from(load(accountIdentifier, device, null))
+    return Flux.from(load(accountIdentifier, device, MAY_HAVE_URGENT_MESSAGES_QUERY_LIMIT))
         .any(MessageProtos.Envelope::getUrgent)
         .toFuture();
   }
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDbTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDbTest.java
index 0ea2206c3..d27fd8c7c 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDbTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDbTest.java
@@ -313,11 +313,14 @@ class MessagesDynamoDbTest {
 
     assertThat(messagesDynamoDb.mayHaveUrgentMessages(destinationUuid, destinationDevice).join()).isFalse();
 
+    // used as the stable sort key, and the urgent message should be sorted last
+    int serverTimestamp = 1;
     {
       final MessageProtos.Envelope nonUrgentMessage = MessageProtos.Envelope.newBuilder()
           .setUrgent(false)
           .setServerGuid(UUID.randomUUID().toString())
-          .setDestinationServiceId(UUID.randomUUID().toString())
+          .setDestinationServiceId(destinationUuid.toString())
+          .setServerTimestamp(serverTimestamp++)
           .build();
 
       messagesDynamoDb.store(List.of(nonUrgentMessage), destinationUuid, destinationDevice);
@@ -326,13 +329,26 @@ class MessagesDynamoDbTest {
     assertThat(messagesDynamoDb.mayHaveUrgentMessages(destinationUuid, destinationDevice).join()).isFalse();
 
     {
-      final MessageProtos.Envelope urgentMessage = MessageProtos.Envelope.newBuilder()
+      final List<MessageProtos.Envelope> messages = new ArrayList<>();
+      // store more non-urgent messages
+      for (int i = 0; i < MessagesDynamoDb.MAY_HAVE_URGENT_MESSAGES_QUERY_LIMIT * 5; i++) {
+        messages.add(MessageProtos.Envelope.newBuilder()
+            .setUrgent(false)
+            .setServerGuid(UUID.randomUUID().toString())
+            .setDestinationServiceId(destinationUuid.toString())
+                .setServerTimestamp(serverTimestamp++)
+            .build());
+      }
+
+      // and one urgent message
+      messages.add(MessageProtos.Envelope.newBuilder()
           .setUrgent(true)
           .setServerGuid(UUID.randomUUID().toString())
-          .setDestinationServiceId(UUID.randomUUID().toString())
-          .build();
+          .setDestinationServiceId(destinationUuid.toString())
+          .setServerTimestamp(serverTimestamp++)
+          .build());
 
-      messagesDynamoDb.store(List.of(urgentMessage), destinationUuid, destinationDevice);
+      messagesDynamoDb.store(messages, destinationUuid, destinationDevice);
     }
 
     assertThat(messagesDynamoDb.mayHaveUrgentMessages(destinationUuid, destinationDevice).join()).isTrue();