diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/AcknowledgedMessageBuffer.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/AcknowledgedMessageBuffer.java new file mode 100644 index 000000000..91007ad06 --- /dev/null +++ b/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 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> takeFlushableRanges() { + final List> 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> computeFlushableRanges() { + final List> flushableRanges = new ArrayList<>(); + Versionstamp startInclusive = null; + Versionstamp endInclusive = null; + + for (final Map.Entry 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; + } +} diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStream.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStream.java index e9ce346b5..4061b1a6e 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStream.java +++ b/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 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 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() .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 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 clearAcknowledgedMessages() { - final List> flushableRanges = computeFlushableRanges(); - flushableRanges.forEach(range -> sentVersionstamps.subMap(range.first(), true, range.second(), true).clear()); + final List> 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> computeFlushableRanges() { - final List> flushableRanges = new ArrayList<>(); - Versionstamp startInclusive = null; - Versionstamp endInclusive = null; - - for (final Map.Entry 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; } } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/TooManyUnacknowledgedMessagesException.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/TooManyUnacknowledgedMessagesException.java index 3ab206477..638c8570e 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/TooManyUnacknowledgedMessagesException.java +++ b/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 { } diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStreamTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/foundationdb/AcknowledgedMessageBufferTest.java similarity index 55% rename from service/src/test/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStreamTest.java rename to service/src/test/java/org/whispersystems/textsecuregcm/storage/foundationdb/AcknowledgedMessageBufferTest.java index 68649516f..34693352a 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStreamTest.java +++ b/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 versionstampsSent, final List versionstampsAcknowledged, + void computeFlushableRanges(final List versionstampsSent, + final List versionstampsAcknowledged, final List> 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> flushableRanges = acknowledgedMessageBuffer.computeFlushableRanges(); - final List> 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 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()); } - }