From 5d624b463a0a4548014a20abd2ca9502abfc98c2 Mon Sep 17 00:00:00 2001 From: Jon Chambers Date: Thu, 11 Jun 2026 09:50:48 -0400 Subject: [PATCH] Acknowledge messages by ID rather than by passing the whole message around --- .../textsecuregcm/storage/MessageStream.java | 6 ++-- .../storage/RedisDynamoDbMessageStream.java | 8 ++--- .../FoundationDbMessageStream.java | 18 +++++------- .../websocket/WebSocketConnection.java | 4 +-- .../RedisDynamoDbMessageStreamTest.java | 4 +-- .../FoundationDbMessageStoreTest.java | 2 +- .../websocket/WebSocketConnectionTest.java | 29 ++++++++++--------- 7 files changed, 34 insertions(+), 37 deletions(-) diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessageStream.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessageStream.java index 03267fd83..38f7f8959 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessageStream.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessageStream.java @@ -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 acknowledgeMessage(MessageProtos.Envelope message); + CompletableFuture acknowledgeMessage(UUID messageGuid, long serverTimestamp); } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/RedisDynamoDbMessageStream.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/RedisDynamoDbMessageStream.java index f567b4141..611ad518d 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/RedisDynamoDbMessageStream.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/RedisDynamoDbMessageStream.java @@ -58,13 +58,11 @@ public class RedisDynamoDbMessageStream implements MessageStream { } @Override - public CompletableFuture acknowledgeMessage(final MessageProtos.Envelope message) { - final UUID guid = UUIDUtil.fromByteString(message.getServerGuid()); - - return messagesCache.remove(accountIdentifier, device.getId(), guid) + public CompletableFuture acknowledgeMessage(final UUID messageGuid, final long serverTimestamp) { + return messagesCache.remove(accountIdentifier, device.getId(), messageGuid) .thenCompose(removed -> removed.map(_ -> CompletableFuture.completedFuture(null)) .orElseGet(() -> - messagesDynamoDb.deleteMessage(accountIdentifier, device, guid, message.getServerTimestamp()) + messagesDynamoDb.deleteMessage(accountIdentifier, device, messageGuid, serverTimestamp) .thenRun(Util.NOOP))) .whenComplete((_, _) -> messagePublisher.handleMessageAcknowledged()); } 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 f97958585..b03a7d632 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 @@ -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 acknowledgeMessage(final MessageProtos.Envelope message) { - acknowledgedMessageBuffer.acknowledgeMessage( - messageGuidCodec.decodeMessageGuid(UUIDUtil.fromByteString(message.getServerGuid()))); + public CompletableFuture acknowledgeMessage(final UUID messageGuid, final long serverTimestamp) { + acknowledgedMessageBuffer.acknowledgeMessage(messageGuidCodec.decodeMessageGuid(messageGuid)); return CompletableFuture.completedFuture(null); } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java b/service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java index 83c21667a..43ae369c6 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java @@ -236,7 +236,7 @@ public class WebSocketConnection implements DisconnectionRequestListener { private CompletableFuture sendMessage(final Envelope message) { if (message.getStory() && !client.shouldDeliverStories()) { - return messageStream.acknowledgeMessage(message); + return messageStream.acknowledgeMessage(UUIDUtil.fromByteString(message.getServerGuid()), message.getServerTimestamp()); } final Optional body = Optional.of(serializeMessage(message)); @@ -266,7 +266,7 @@ public class WebSocketConnection implements DisconnectionRequestListener { final CompletableFuture 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); diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/RedisDynamoDbMessageStreamTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/RedisDynamoDbMessageStreamTest.java index 59dd10412..44e710d50 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/RedisDynamoDbMessageStreamTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/RedisDynamoDbMessageStreamTest.java @@ -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()); diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStoreTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStoreTest.java index 5a4965a2c..75ffc31e6 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStoreTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStoreTest.java @@ -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) diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnectionTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnectionTest.java index 2c921a68a..50f5d8caf 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnectionTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnectionTest.java @@ -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));