Acknowledge messages by ID rather than by passing the whole message around

This commit is contained in:
Jon Chambers
2026-06-11 09:50:48 -04:00
committed by Jon Chambers
parent f524d9fb25
commit 5d624b463a
7 changed files with 34 additions and 37 deletions
@@ -5,6 +5,7 @@
package org.whispersystems.textsecuregcm.storage;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Flow;
import org.whispersystems.textsecuregcm.entities.MessageProtos;
@@ -22,8 +23,9 @@ public interface MessageStream {
/// Acknowledges receipt of the given message. Implementations may delete the message immediately or defer deletion for
/// inclusion in a more efficient batch deletion.
///
/// @param message the message to acknowledge
/// @param messageGuid the identifier of the message to acknowledge
/// @param serverTimestamp the server-issued timestamp of the message to acknowledge
///
/// @return a future that completes when the message stream has processed the acknowledgement
CompletableFuture<Void> acknowledgeMessage(MessageProtos.Envelope message);
CompletableFuture<Void> acknowledgeMessage(UUID messageGuid, long serverTimestamp);
}
@@ -58,13 +58,11 @@ public class RedisDynamoDbMessageStream implements MessageStream {
}
@Override
public CompletableFuture<Void> acknowledgeMessage(final MessageProtos.Envelope message) {
final UUID guid = UUIDUtil.fromByteString(message.getServerGuid());
return messagesCache.remove(accountIdentifier, device.getId(), guid)
public CompletableFuture<Void> acknowledgeMessage(final UUID messageGuid, final long serverTimestamp) {
return messagesCache.remove(accountIdentifier, device.getId(), messageGuid)
.thenCompose(removed -> removed.map(_ -> CompletableFuture.<Void>completedFuture(null))
.orElseGet(() ->
messagesDynamoDb.deleteMessage(accountIdentifier, device, guid, message.getServerTimestamp())
messagesDynamoDb.deleteMessage(accountIdentifier, device, messageGuid, serverTimestamp)
.thenRun(Util.NOOP)))
.whenComplete((_, _) -> messagePublisher.handleMessageAcknowledged());
}
@@ -1,5 +1,7 @@
package org.whispersystems.textsecuregcm.storage.foundationdb;
import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
import com.apple.foundationdb.Database;
import com.apple.foundationdb.KeySelector;
import com.apple.foundationdb.StreamingMode;
@@ -9,28 +11,23 @@ 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 io.micrometer.core.instrument.Counter;
import io.micrometer.core.instrument.Metrics;
import java.util.List;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Flow;
import java.util.function.Consumer;
import io.micrometer.core.instrument.Counter;
import io.micrometer.core.instrument.Metrics;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.whispersystems.textsecuregcm.entities.MessageProtos;
import org.whispersystems.textsecuregcm.storage.MessageStream;
import org.whispersystems.textsecuregcm.storage.MessageStreamEntry;
import org.whispersystems.textsecuregcm.util.Pair;
import org.whispersystems.textsecuregcm.util.UUIDUtil;
import reactor.adapter.JdkFlowAdapter;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
/// A [MessageStream] implementation that fetches messages from FoundationDB
public class FoundationDbMessageStream implements MessageStream {
@@ -158,9 +155,8 @@ public class FoundationDbMessageStream implements MessageStream {
}
@Override
public CompletableFuture<Void> acknowledgeMessage(final MessageProtos.Envelope message) {
acknowledgedMessageBuffer.acknowledgeMessage(
messageGuidCodec.decodeMessageGuid(UUIDUtil.fromByteString(message.getServerGuid())));
public CompletableFuture<Void> acknowledgeMessage(final UUID messageGuid, final long serverTimestamp) {
acknowledgedMessageBuffer.acknowledgeMessage(messageGuidCodec.decodeMessageGuid(messageGuid));
return CompletableFuture.completedFuture(null);
}
@@ -236,7 +236,7 @@ public class WebSocketConnection implements DisconnectionRequestListener {
private CompletableFuture<Void> sendMessage(final Envelope message) {
if (message.getStory() && !client.shouldDeliverStories()) {
return messageStream.acknowledgeMessage(message);
return messageStream.acknowledgeMessage(UUIDUtil.fromByteString(message.getServerGuid()), message.getServerTimestamp());
}
final Optional<byte[]> body = Optional.of(serializeMessage(message));
@@ -266,7 +266,7 @@ public class WebSocketConnection implements DisconnectionRequestListener {
final CompletableFuture<Void> result;
if (isSuccessResponse(response)) {
result = messageStream.acknowledgeMessage(message);
result = messageStream.acknowledgeMessage(UUIDUtil.fromByteString(message.getServerGuid()), message.getServerTimestamp());
if (message.getType() != Envelope.Type.SERVER_DELIVERY_RECEIPT) {
sendDeliveryReceiptFor(message);
@@ -64,7 +64,7 @@ class RedisDynamoDbMessageStreamTest {
when(messagesDynamoDb.deleteMessage(ACCOUNT_IDENTIFIER, device, messageGuid, serverTimestamp))
.thenReturn(CompletableFuture.completedFuture(Optional.of(message)));
redisDynamoDbMessageStream.acknowledgeMessage(message).join();
redisDynamoDbMessageStream.acknowledgeMessage(messageGuid, serverTimestamp).join();
verify(messagesCache).remove(ACCOUNT_IDENTIFIER, DEVICE_ID, messageGuid);
verify(messagesDynamoDb).deleteMessage(ACCOUNT_IDENTIFIER, device, messageGuid, serverTimestamp);
@@ -78,7 +78,7 @@ class RedisDynamoDbMessageStreamTest {
when(messagesCache.remove(ACCOUNT_IDENTIFIER, DEVICE_ID, messageGuid))
.thenReturn(CompletableFuture.completedFuture(Optional.of(RemovedMessage.fromEnvelope(message))));
redisDynamoDbMessageStream.acknowledgeMessage(message).join();
redisDynamoDbMessageStream.acknowledgeMessage(messageGuid, message.getServerTimestamp()).join();
verify(messagesCache).remove(ACCOUNT_IDENTIFIER, DEVICE_ID, messageGuid);
verify(messagesDynamoDb, never()).deleteMessage(any(), any(), any(), anyLong());
@@ -569,7 +569,7 @@ class FoundationDbMessageStoreTest {
.doOnNext(entry -> {
final int messageNum = messageCounter.getAndIncrement();
if (!unacknowledgedMessages.contains(messageNum) && entry instanceof MessageStreamEntry.Envelope(final MessageProtos.Envelope message)) {
acknowledgeFutures.add(messageStream.acknowledgeMessage(message));
acknowledgeFutures.add(messageStream.acknowledgeMessage(UUIDUtil.fromByteString(message.getServerGuid()), message.getServerTimestamp()));
}
}))
.expectNextCount(numMessages)
@@ -8,6 +8,7 @@ package org.whispersystems.textsecuregcm.websocket;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.anyList;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.ArgumentMatchers.argThat;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.any;
@@ -146,7 +147,7 @@ class WebSocketConnectionTest {
new MessageStreamEntry.QueueEmpty(),
new MessageStreamEntry.Envelope(secondSuccessfulMessage))));
when(messageStream.acknowledgeMessage(any())).thenReturn(CompletableFuture.completedFuture(null));
when(messageStream.acknowledgeMessage(any(), anyLong())).thenReturn(CompletableFuture.completedFuture(null));
when(messagesManager.getMessages(account.getIdentifier(IdentityType.ACI), device))
.thenReturn(messageStream);
@@ -172,8 +173,8 @@ class WebSocketConnectionTest {
verify(client).sendRequest(eq("PUT"), eq("/api/v1/message"), anyList(), argThat(body ->
body.isPresent() && Arrays.equals(body.get(), WebSocketConnection.serializeMessage(secondSuccessfulMessage))));
verify(messageStream).acknowledgeMessage(successfulMessage);
verify(messageStream).acknowledgeMessage(secondSuccessfulMessage);
verify(messageStream).acknowledgeMessage(UUIDUtil.fromByteString(successfulMessage.getServerGuid()), successfulMessage.getServerTimestamp());
verify(messageStream).acknowledgeMessage(UUIDUtil.fromByteString(secondSuccessfulMessage.getServerGuid()), secondSuccessfulMessage.getServerTimestamp());
verify(receiptSender)
.sendReceipt(new AciServiceIdentifier(destinationAccountIdentifier),
@@ -215,7 +216,7 @@ class WebSocketConnectionTest {
new MessageStreamEntry.QueueEmpty(),
new MessageStreamEntry.Envelope(secondSuccessfulMessage))));
when(messageStream.acknowledgeMessage(any())).thenReturn(CompletableFuture.completedFuture(null));
when(messageStream.acknowledgeMessage(any(), anyLong())).thenReturn(CompletableFuture.completedFuture(null));
when(messagesManager.getMessages(account.getIdentifier(IdentityType.ACI), device))
.thenReturn(messageStream);
@@ -248,8 +249,8 @@ class WebSocketConnectionTest {
verify(client, never()).sendRequest(eq("PUT"), eq("/api/v1/message"), anyList(), argThat(body ->
body.isPresent() && Arrays.equals(body.get(), WebSocketConnection.serializeMessage(secondSuccessfulMessage))));
verify(messageStream).acknowledgeMessage(successfulMessage);
verify(messageStream, never()).acknowledgeMessage(secondSuccessfulMessage);
verify(messageStream).acknowledgeMessage(UUIDUtil.fromByteString(successfulMessage.getServerGuid()), successfulMessage.getServerTimestamp());
verify(messageStream, never()).acknowledgeMessage(UUIDUtil.fromByteString(secondSuccessfulMessage.getServerGuid()), secondSuccessfulMessage.getServerTimestamp());
verify(receiptSender)
.sendReceipt(new AciServiceIdentifier(destinationAccountIdentifier),
@@ -293,7 +294,7 @@ class WebSocketConnectionTest {
new MessageStreamEntry.QueueEmpty(),
new MessageStreamEntry.Envelope(afterQueueDrainMessage))));
when(messageStream.acknowledgeMessage(any())).thenReturn(CompletableFuture.completedFuture(null));
when(messageStream.acknowledgeMessage(any(), anyLong())).thenReturn(CompletableFuture.completedFuture(null));
when(messagesManager.getMessages(account.getIdentifier(IdentityType.ACI), device))
.thenReturn(messageStream);
@@ -316,7 +317,7 @@ class WebSocketConnectionTest {
// Sending the initial message will succeed after a delay, at which point we'll acknowledge the message. Make sure
// we wait for that process to complete before sending the "queue empty" signal
inOrder.verify(messageStream, timeout(1_000)).acknowledgeMessage(initialMessage);
inOrder.verify(messageStream, timeout(1_000)).acknowledgeMessage(UUIDUtil.fromByteString(initialMessage.getServerGuid()), initialMessage.getServerTimestamp());
inOrder.verify(client, timeout(1_000)).sendRequest(eq("PUT"), eq("/api/v1/queue/empty"), anyList(), eq(Optional.empty()));
webSocketConnection.stop();
@@ -340,7 +341,7 @@ class WebSocketConnectionTest {
when(messageStream.getMessages())
.thenReturn(JdkFlowAdapter.publisherToFlowPublisher(testPublisher));
when(messageStream.acknowledgeMessage(any())).thenReturn(CompletableFuture.completedFuture(null));
when(messageStream.acknowledgeMessage(any(), anyLong())).thenReturn(CompletableFuture.completedFuture(null));
when(messagesManager.getMessages(account.getIdentifier(IdentityType.ACI), device))
.thenReturn(messageStream);
@@ -383,7 +384,7 @@ class WebSocketConnectionTest {
when(messageStream.getMessages())
.thenReturn(JdkFlowAdapter.publisherToFlowPublisher(Flux.just(new MessageStreamEntry.QueueEmpty())));
when(messageStream.acknowledgeMessage(any())).thenReturn(CompletableFuture.completedFuture(null));
when(messageStream.acknowledgeMessage(any(), anyLong())).thenReturn(CompletableFuture.completedFuture(null));
when(messagesManager.getMessages(accountUuid, device)).thenReturn(messageStream);
@@ -412,7 +413,7 @@ class WebSocketConnectionTest {
when(messageStream.getMessages())
.thenReturn(JdkFlowAdapter.publisherToFlowPublisher(Flux.error(new ConflictingMessageConsumerException())));
when(messageStream.acknowledgeMessage(any())).thenReturn(CompletableFuture.completedFuture(null));
when(messageStream.acknowledgeMessage(any(), anyLong())).thenReturn(CompletableFuture.completedFuture(null));
when(messagesManager.getMessages(accountUuid, device)).thenReturn(messageStream);
@@ -438,7 +439,7 @@ class WebSocketConnectionTest {
when(messageStream.getMessages())
.thenReturn(JdkFlowAdapter.publisherToFlowPublisher(Flux.error(new RedisException("OH NO"))));
when(messageStream.acknowledgeMessage(any())).thenReturn(CompletableFuture.completedFuture(null));
when(messageStream.acknowledgeMessage(any(), anyLong())).thenReturn(CompletableFuture.completedFuture(null));
when(messagesManager.getMessages(accountUuid, device)).thenReturn(messageStream);
@@ -474,7 +475,7 @@ class WebSocketConnectionTest {
when(messageStream.getMessages())
.thenReturn(JdkFlowAdapter.publisherToFlowPublisher(flux));
when(messageStream.acknowledgeMessage(any())).thenReturn(CompletableFuture.completedFuture(null));
when(messageStream.acknowledgeMessage(any(), anyLong())).thenReturn(CompletableFuture.completedFuture(null));
when(messagesManager.getMessages(accountUuid, device)).thenReturn(messageStream);
@@ -534,7 +535,7 @@ class WebSocketConnectionTest {
when(messageStream.getMessages())
.thenReturn(JdkFlowAdapter.publisherToFlowPublisher(flux));
when(messageStream.acknowledgeMessage(any())).thenReturn(CompletableFuture.completedFuture(null));
when(messageStream.acknowledgeMessage(any(), anyLong())).thenReturn(CompletableFuture.completedFuture(null));
when(messagesManager.getMessages(accountUuid, device)).thenReturn(messageStream);
when(messagesManager.mayHaveMessages(any(), any())).thenReturn(CompletableFuture.completedFuture(false));