Introduce AcknowledgedMessageBuffer

service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/AcknowledgedMessageBuffer.java

@@ -0,0 +1,102 @@
+/*
+ * Copyright 2026 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.storage.foundationdb;
+
+import com.apple.foundationdb.tuple.Versionstamp;
+import com.google.common.annotations.VisibleForTesting;
+import org.whispersystems.textsecuregcm.util.Pair;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.NavigableMap;
+import java.util.TreeMap;
+
+/// An acknowledged message buffer keeps track of messages that have been sent, whether they have been acknowledged by
+/// the recipient, and whether those acknowledgements have been flushed to the underlying data store. It also batches
+/// acknowledgements into contiguous ranges in the interest of efficiency.
+class AcknowledgedMessageBuffer {
+
+  /// The maximum number of unacknowledged messages we will send before closing the stream with an error
+  private final int maxUnacknowledgedMessages;
+
+  /// Map of versionstamp -> acknowledged? to keep track of acknowledged versionstamp ranges to clear
+  private final NavigableMap<Versionstamp, Boolean> sentVersionstamps = new TreeMap<>();
+  private int unacknowledgedMessages = 0;
+
+  AcknowledgedMessageBuffer(final int maxUnacknowledgedMessages) {
+    this.maxUnacknowledgedMessages = maxUnacknowledgedMessages;
+  }
+
+  public synchronized void addUnacknowledgedMessage(final Versionstamp versionstamp) throws TooManyUnacknowledgedMessagesException {
+    if (++unacknowledgedMessages > maxUnacknowledgedMessages) {
+      throw new TooManyUnacknowledgedMessagesException();
+    }
+
+    sentVersionstamps.put(versionstamp, false);
+  }
+
+  public synchronized void acknowledgeMessage(final Versionstamp versionstamp) {
+    // If we don't know about this versionstamp, or it's already acknowledged, there's nothing to do
+    if (!sentVersionstamps.containsKey(versionstamp) || sentVersionstamps.get(versionstamp)) {
+      return;
+    }
+
+    unacknowledgedMessages--;
+    sentVersionstamps.put(versionstamp, true);
+  }
+
+  @VisibleForTesting
+  int getUnacknowledgedMessageCount() {
+    return unacknowledgedMessages;
+  }
+
+  /// Computes a list of acknowledged contiguous versionstamp ranges that can be cleared from the database and removes
+  /// them from ths unacknowledged mesage buffer.
+  ///
+  /// @return a list of acknowledged contiguous versionstamp ranges that can be cleared from the database.
+  public synchronized List<Pair<Versionstamp, Versionstamp>> takeFlushableRanges() {
+    final List<Pair<Versionstamp, Versionstamp>> flushableRanges = computeFlushableRanges();
+    flushableRanges.forEach(range -> sentVersionstamps.subMap(range.first(), true, range.second(), true).clear());
+
+    return flushableRanges;
+  }
+
+  /// Computes a list of acknowledged contiguous versionstamp ranges that can be cleared from the database.
+  ///
+  /// @return a list of acknowledged contiguous versionstamp ranges that can be cleared from the database.
+  @VisibleForTesting
+  synchronized List<Pair<Versionstamp, Versionstamp>> computeFlushableRanges() {
+    final List<Pair<Versionstamp, Versionstamp>> flushableRanges = new ArrayList<>();
+    Versionstamp startInclusive = null;
+    Versionstamp endInclusive = null;
+
+    for (final Map.Entry<Versionstamp, Boolean> entry : sentVersionstamps.entrySet()) {
+      if (entry.getValue()) {
+        // Message is acknowledged, so we can either start tracking a new range if we aren't already, or extend our
+        // current range
+        if (startInclusive == null) {
+          startInclusive = entry.getKey();
+        }
+        endInclusive = entry.getKey();
+      } else {
+        // Message is un-acknowledged, which means either it is a "range-breaker" or we never started tracking an
+        // acknowledged range. Mark the currently tracked range flushable, if it exists.
+        if (startInclusive != null) {
+          assert endInclusive != null;
+          flushableRanges.add(new Pair<>(startInclusive, endInclusive));
+          startInclusive = null;
+        }
+      }
+    }
+
+    if (startInclusive != null) {
+      assert endInclusive != null;
+      flushableRanges.add(new Pair<>(startInclusive, endInclusive));
+    }
+
+    return flushableRanges;
+  }
+}

service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStream.java

@@ -5,22 +5,17 @@ import com.apple.foundationdb.KeySelector;
 import com.apple.foundationdb.StreamingMode;
 import com.apple.foundationdb.Transaction;
 import com.apple.foundationdb.subspace.Subspace;
-import java.time.Clock;
-import java.time.Duration;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
-import java.util.NavigableMap;
-import java.util.Optional;
-import java.util.TreeMap;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.Flow;
-import java.util.function.Consumer;
-
 import com.apple.foundationdb.tuple.ByteArrayUtil;
 import com.apple.foundationdb.tuple.Tuple;
 import com.apple.foundationdb.tuple.Versionstamp;
 import com.google.common.annotations.VisibleForTesting;
+import java.time.Clock;
+import java.time.Duration;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Flow;
+import java.util.function.Consumer;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.entities.MessageProtos;
@@ -41,15 +36,11 @@ public class FoundationDbMessageStream implements MessageStream {
   private final MessageGuidCodec messageGuidCodec;
   /// The maximum number of messages we will fetch per range query operation to avoid excessive memory consumption
   private final int maxMessagesPerScan;
-  /// The maximum number of unacknowledged messages we will send before closing the stream with an error
-  private final int maxUnacknowledgedMessages;
   private final Flow.Publisher<MessageStreamEntry> messageStreamPublisher;
   private final Runnable doAfterCleanup;
   private final Clock clock;
 
-  /// Map of versionstamp -> acknowledged? to keep track of acknowledged versionstamp ranges to clear
-  private final NavigableMap<Versionstamp, Boolean> sentVersionstamps = new TreeMap<>();
-  private int unacknowledgedMessages = 0;
+  private final AcknowledgedMessageBuffer acknowledgedMessageBuffer;
 
   @VisibleForTesting
   static final Duration MAX_EPHEMERAL_MESSAGE_DELAY = Duration.ofSeconds(10);
@@ -72,10 +63,11 @@ public class FoundationDbMessageStream implements MessageStream {
     this.database = database;
     this.messageGuidCodec = messageGuidCodec;
     this.maxMessagesPerScan = maxMessagesPerScan;
-    this.maxUnacknowledgedMessages = maxUnacknowledgedMessages;
     this.clock = clock;
     this.messageStreamPublisher = JdkFlowAdapter.publisherToFlowPublisher(createMessagePublisher());
     this.doAfterCleanup = doAfterCleanup;
+
+    this.acknowledgedMessageBuffer = new AcknowledgedMessageBuffer(maxUnacknowledgedMessages);
   }
 
   @Override
@@ -93,12 +85,14 @@ public class FoundationDbMessageStream implements MessageStream {
     return createFoundationDbMessagePublisher()
         .<FoundationDbMessageStreamEntry>handle((messageStreamEntry, sink) -> {
           if (messageStreamEntry instanceof final FoundationDbMessageStreamEntry.Message message) {
-            final int numUnacknowledgedMessages = onVersionstampSent(message.versionstamp());
-            if (numUnacknowledgedMessages > maxUnacknowledgedMessages) {
-              sink.error(new TooManyUnacknowledgedMessagesException());
+            try {
+              acknowledgedMessageBuffer.addUnacknowledgedMessage(message.versionstamp());
+            } catch (final TooManyUnacknowledgedMessagesException e) {
+              sink.error(e);
               return;
             }
           }
+
           sink.next(messageStreamEntry);
         })
         .map(fdbMessageStreamEntry -> fdbMessageStreamEntry.toMessageStreamEntry(messageGuidCodec))
@@ -166,30 +160,12 @@ public class FoundationDbMessageStream implements MessageStream {
 
   @Override
   public CompletableFuture<Void> acknowledgeMessage(final MessageProtos.Envelope message) {
-    handleAcknowledged(messageGuidCodec.decodeMessageGuid(UUIDUtil.fromByteString(message.getServerGuidBinary())));
+    acknowledgedMessageBuffer.acknowledgeMessage(
+        messageGuidCodec.decodeMessageGuid(UUIDUtil.fromByteString(message.getServerGuidBinary())));
+
     return CompletableFuture.completedFuture(null);
   }
 
-  @VisibleForTesting
-  synchronized void handleAcknowledged(final Versionstamp versionstamp) {
-    // If we don't know about this versionstamp, or it's already acknowledged, there's nothing to do
-    if (!sentVersionstamps.containsKey(versionstamp) || sentVersionstamps.get(versionstamp)) {
-      return;
-    }
-
-    unacknowledgedMessages--;
-    sentVersionstamps.put(versionstamp, true);
-  }
-
-  /// Returns the number of unacknowledged messages after sending the provided versionstamp
-  ///
-  /// @param versionstamp The versionstamp being sent
-  @VisibleForTesting
-  synchronized int onVersionstampSent(final Versionstamp versionstamp) {
-    sentVersionstamps.put(versionstamp, false);
-    return ++unacknowledgedMessages;
-  }
-
   /// Clear the versionstamp range (startInclusive, endInclusive) in a single FoundationDB operation.
   ///
   /// @param transaction    The FoundationDB transaction in which to perform the range clear
@@ -216,50 +192,13 @@ public class FoundationDbMessageStream implements MessageStream {
   }
 
   private synchronized Consumer<Transaction> clearAcknowledgedMessages() {
-    final List<Pair<Versionstamp, Versionstamp>> flushableRanges = computeFlushableRanges();
-    flushableRanges.forEach(range -> sentVersionstamps.subMap(range.first(), true, range.second(), true).clear());
+    final List<Pair<Versionstamp, Versionstamp>> flushableRanges = acknowledgedMessageBuffer.takeFlushableRanges();
+
     return transaction -> flushableRanges.forEach(range -> clearRange(transaction, range.first(), range.second()));
   }
 
-  /// Computes a list of acknowledged contiguous versionstamp ranges that can be cleared from the database.
-  ///
-  /// @return a list of acknowledged contiguous versionstamp ranges that can be cleared from the database.
-  @VisibleForTesting
-  synchronized List<Pair<Versionstamp, Versionstamp>> computeFlushableRanges() {
-    final List<Pair<Versionstamp, Versionstamp>> flushableRanges = new ArrayList<>();
-      Versionstamp startInclusive = null;
-      Versionstamp endInclusive = null;
-
-      for (final Map.Entry<Versionstamp, Boolean> entry : sentVersionstamps.entrySet()) {
-        if (entry.getValue()) {
-          // Message is acknowledged, so we can either start tracking a new range if we aren't already, or extend our
-          // current range
-          if (startInclusive == null) {
-            startInclusive = entry.getKey();
-          }
-          endInclusive = entry.getKey();
-        } else {
-          // Message is un-acknowledged, which means either it is a "range-breaker" or we never started tracking an
-          // acknowledged range. Mark the currently tracked range flushable, if it exists.
-          if (startInclusive != null) {
-            assert endInclusive != null;
-            flushableRanges.add(new Pair<>(startInclusive, endInclusive));
-            startInclusive = null;
-          }
-        }
-
-      }
-
-      if (startInclusive != null) {
-        assert endInclusive != null;
-        flushableRanges.add(new Pair<>(startInclusive, endInclusive));
-      }
-
-    return flushableRanges;
-  }
-
   private static boolean isStaleEphemeralMessage(final MessageProtos.Envelope message,
-      long earliestAllowableTimestamp) {
+      final long earliestAllowableTimestamp) {
     return message.getEphemeral() && message.getClientTimestamp() < earliestAllowableTimestamp;
   }
 }

service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/TooManyUnacknowledgedMessagesException.java

@@ -5,5 +5,7 @@
 
 package org.whispersystems.textsecuregcm.storage.foundationdb;
 
-public class TooManyUnacknowledgedMessagesException extends RuntimeException {
+import org.whispersystems.textsecuregcm.util.NoStackTraceException;
+
+class TooManyUnacknowledgedMessagesException extends NoStackTraceException {
 }

service/src/test/java/org/whispersystems/textsecuregcm/storage/foundationdb/AcknowledgedMessageBufferTest.java

@@ -1,46 +1,45 @@
+/*
+ * Copyright 2026 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
 package org.whispersystems.textsecuregcm.storage.foundationdb;
 
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.mockito.Mockito.mock;
 
-import com.apple.foundationdb.Database;
-import com.apple.foundationdb.subspace.Subspace;
 import com.apple.foundationdb.tuple.Versionstamp;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
-import java.time.Clock;
 import java.util.Collections;
 import java.util.List;
+import java.util.stream.IntStream;
 import java.util.stream.Stream;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
 import org.whispersystems.textsecuregcm.util.Pair;
-import org.whispersystems.textsecuregcm.util.Util;
 
-class FoundationDbMessageStreamTest {
+class AcknowledgedMessageBufferTest {
 
   @ParameterizedTest
   @MethodSource
-  void computeFlushableRanges(final List<Integer> versionstampsSent, final List<Integer> versionstampsAcknowledged,
+  void computeFlushableRanges(final List<Integer> versionstampsSent,
+      final List<Integer> versionstampsAcknowledged,
       final List<Pair<Integer, Integer>> expectedRanges) {
-    final FoundationDbMessageStream foundationDbMessageStream = new FoundationDbMessageStream(
-        mock(Subspace.class),
-        new byte[]{0},
-        mock(Database.class),
-        mock(MessageGuidCodec.class),
-        100,
-        FoundationDbMessageStream.DEFAULT_MAX_UNACKNOWLEDGED_MESSAGES,
-        Util.NOOP,
-        Clock.systemUTC()
-    );
 
-    versionstampsSent.forEach(
-        versionstamp -> foundationDbMessageStream.onVersionstampSent(versionstampFromInt(versionstamp)));
-    versionstampsAcknowledged.forEach(
-        versionstamp -> foundationDbMessageStream.handleAcknowledged(versionstampFromInt(versionstamp)));
+    final AcknowledgedMessageBuffer acknowledgedMessageBuffer =
+        new AcknowledgedMessageBuffer(FoundationDbMessageStream.DEFAULT_MAX_UNACKNOWLEDGED_MESSAGES);
+
+    versionstampsSent.forEach(versionstamp -> assertDoesNotThrow(() ->
+        acknowledgedMessageBuffer.addUnacknowledgedMessage(versionstampFromInt(versionstamp))));
+
+    versionstampsAcknowledged.forEach(versionstamp ->
+        acknowledgedMessageBuffer.acknowledgeMessage(versionstampFromInt(versionstamp)));
+
+    final List<Pair<Versionstamp, Versionstamp>> flushableRanges = acknowledgedMessageBuffer.computeFlushableRanges();
 
-    final List<Pair<Versionstamp, Versionstamp>> flushableRanges = foundationDbMessageStream.computeFlushableRanges();
     assertEquals(expectedRanges
             .stream()
             .map(range -> new Pair<>(versionstampFromInt(range.first()), versionstampFromInt(range.second())))
@@ -48,6 +47,28 @@ class FoundationDbMessageStreamTest {
         flushableRanges);
   }
 
+  @Test
+  void takeFlushableRanges() {
+    final AcknowledgedMessageBuffer acknowledgedMessageBuffer =
+        new AcknowledgedMessageBuffer(FoundationDbMessageStream.DEFAULT_MAX_UNACKNOWLEDGED_MESSAGES);
+
+    IntStream.range(0, 10)
+        .mapToObj(AcknowledgedMessageBufferTest::versionstampFromInt)
+        .forEach(versionstamp -> assertDoesNotThrow(
+            () -> acknowledgedMessageBuffer.addUnacknowledgedMessage(versionstamp)));
+
+    assertEquals(10, acknowledgedMessageBuffer.getUnacknowledgedMessageCount());
+
+    IntStream.range(2, 8)
+        .mapToObj(AcknowledgedMessageBufferTest::versionstampFromInt)
+        .forEach(acknowledgedMessageBuffer::acknowledgeMessage);
+
+    assertEquals(List.of(new Pair<>(versionstampFromInt(2), versionstampFromInt(7))),
+        acknowledgedMessageBuffer.takeFlushableRanges());
+
+    assertEquals(4, acknowledgedMessageBuffer.getUnacknowledgedMessageCount());
+  }
+
   static Stream<Arguments> computeFlushableRanges() {
     return Stream.of(
         Arguments.argumentSet("Single acknowledged message range", List.of(1, 2, 3), List.of(1),
@@ -65,12 +86,10 @@ class FoundationDbMessageStreamTest {
     );
   }
 
-
   private static Versionstamp versionstampFromInt(final int version) {
     final ByteBuffer buf = ByteBuffer.allocate(10).order(ByteOrder.BIG_ENDIAN);
-    buf.putLong(version);     // 8 bytes: transaction version
+    buf.putLong(version);      // 8 bytes: transaction version
     buf.putShort((short) 2);   // 2 bytes: batch order within transaction
     return Versionstamp.complete(buf.array());
   }
-
 }