Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit db2cd20

Browse files
committedMar 21, 2025·
Skip shared multi-recipient message payloads for small messages
1 parent 9ef6024 commit db2cd20

File tree

2 files changed

+85
-21
lines changed

2 files changed

+85
-21
lines changed
 

‎service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java

Lines changed: 48 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import java.util.concurrent.ExecutorService;
2222
import java.util.concurrent.TimeUnit;
2323
import java.util.concurrent.TimeoutException;
24+
import java.util.function.BiFunction;
2425
import java.util.stream.Collectors;
2526
import java.util.stream.IntStream;
2627
import javax.annotation.Nullable;
@@ -41,7 +42,9 @@
4142
public class MessagesManager {
4243

4344
private static final int RESULT_SET_CHUNK_SIZE = 100;
44-
final String GET_MESSAGES_FOR_DEVICE_FLUX_NAME = name(MessagesManager.class, "getMessagesForDevice");
45+
private final static String GET_MESSAGES_FOR_DEVICE_FLUX_NAME = name(MessagesManager.class, "getMessagesForDevice");
46+
// shared payloads have some overhead, which sometimes exceeds the size if we just wrote the content directly
47+
private static final int MULTI_RECIPIENT_MESSAGE_MINIMUM_SIZE_FOR_SHARED_PAYLOAD = 150;
4548

4649
private static final Logger logger = LoggerFactory.getLogger(MessagesManager.class);
4750

@@ -139,17 +142,50 @@ public CompletableFuture<Map<Account, Map<Byte, Boolean>>> insertMultiRecipientM
139142

140143
final long serverTimestamp = clock.millis();
141144

142-
return insertSharedMultiRecipientMessagePayload(multiRecipientMessage)
143-
.thenCompose(sharedMrmKey -> {
144-
final Envelope prototypeMessage = Envelope.newBuilder()
145-
.setType(Envelope.Type.UNIDENTIFIED_SENDER)
146-
.setClientTimestamp(clientTimestamp == 0 ? serverTimestamp : clientTimestamp)
147-
.setServerTimestamp(serverTimestamp)
148-
.setStory(isStory)
149-
.setEphemeral(isEphemeral)
150-
.setUrgent(isUrgent)
151-
.setSharedMrmKey(ByteString.copyFrom(sharedMrmKey))
145+
final Envelope.Builder prototypeMessageBuilder = Envelope.newBuilder()
146+
.setType(Envelope.Type.UNIDENTIFIED_SENDER)
147+
.setClientTimestamp(clientTimestamp == 0 ? serverTimestamp : clientTimestamp)
148+
.setServerTimestamp(serverTimestamp)
149+
.setStory(isStory)
150+
.setEphemeral(isEphemeral)
151+
.setUrgent(isUrgent);
152+
153+
final CompletableFuture<Envelope> prototypeMessageFuture;
154+
final BiFunction<ServiceIdentifier, Envelope, Envelope> recipientEnvelopeBuilder;
155+
156+
// A shortcut -- message sizes do not vary by recipient in the current SealedSenderMultiRecipientMessage version
157+
final int perRecipientMessageSize = multiRecipientMessage.getRecipients().values().stream().findAny()
158+
.map(multiRecipientMessage::messageSizeForRecipient)
159+
.orElse(0);
160+
161+
multiRecipientMessage.messageSizeForRecipient(
162+
multiRecipientMessage.getRecipients().values().iterator().next());
163+
if (perRecipientMessageSize >= MULTI_RECIPIENT_MESSAGE_MINIMUM_SIZE_FOR_SHARED_PAYLOAD) {
164+
165+
// the message is large enough that the shared payload overhead is worth it, so insert into the cache
166+
prototypeMessageFuture = insertSharedMultiRecipientMessagePayload((multiRecipientMessage))
167+
.thenApply(sharedMrmKey -> prototypeMessageBuilder
168+
.setSharedMrmKey(ByteString.copyFrom(sharedMrmKey))
169+
.build());
170+
171+
recipientEnvelopeBuilder = (serviceIdentifier, prototype) -> prototype.toBuilder()
172+
.setDestinationServiceId(serviceIdentifier.toServiceIdentifierString())
173+
.build();
174+
175+
} else {
176+
177+
prototypeMessageFuture = CompletableFuture.completedFuture(prototypeMessageBuilder.build());
178+
179+
recipientEnvelopeBuilder = (serviceIdentifier, prototype) ->
180+
prototype.toBuilder()
181+
.setDestinationServiceId(serviceIdentifier.toServiceIdentifierString())
182+
.setContent(ByteString.copyFrom(multiRecipientMessage.messageForRecipient(
183+
multiRecipientMessage.getRecipients().get(serviceIdentifier.toLibsignal()))))
152184
.build();
185+
}
186+
187+
return prototypeMessageFuture
188+
.thenCompose(prototypeMessage -> {
153189

154190
final Map<Account, Map<Byte, Boolean>> clientPresenceByAccountAndDevice = new ConcurrentHashMap<>();
155191

@@ -162,9 +198,7 @@ public CompletableFuture<Map<Account, Map<Byte, Boolean>>> insertMultiRecipientM
162198

163199
return insertAsync(resolvedRecipients.get(recipient).getIdentifier(IdentityType.ACI),
164200
IntStream.range(0, devices.length).mapToObj(i -> devices[i])
165-
.collect(Collectors.toMap(deviceId -> deviceId, deviceId -> prototypeMessage.toBuilder()
166-
.setDestinationServiceId(serviceIdentifier.toServiceIdentifierString())
167-
.build())))
201+
.collect(Collectors.toMap(deviceId -> deviceId, deviceId -> recipientEnvelopeBuilder.apply(serviceIdentifier, prototypeMessage))))
168202
.thenAccept(clientPresenceByDeviceId ->
169203
clientPresenceByAccountAndDevice.put(resolvedRecipients.get(recipient),
170204
clientPresenceByDeviceId));

‎service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagesManagerTest.java

Lines changed: 37 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,9 @@
2828
import java.util.concurrent.CompletableFuture;
2929
import java.util.concurrent.Executors;
3030
import java.util.concurrent.ThreadLocalRandom;
31+
import java.util.function.Function;
32+
import java.util.stream.Collectors;
33+
import java.util.stream.Stream;
3134
import org.junit.jupiter.api.BeforeEach;
3235
import org.junit.jupiter.api.Test;
3336
import org.junit.jupiter.params.ParameterizedTest;
@@ -83,8 +86,15 @@ void insert() {
8386
verifyNoMoreInteractions(reportMessageManager);
8487
}
8588

86-
@Test
87-
void insertMultiRecipientMessage() throws InvalidMessageException, InvalidVersionException {
89+
@ParameterizedTest
90+
@CsvSource({
91+
"32, false",
92+
"99, false",
93+
"100, true",
94+
"200, true",
95+
"1024, true",
96+
})
97+
void insertMultiRecipientMessage(final int sharedPayloadSize, final boolean expectSharedMrm) throws InvalidMessageException, InvalidVersionException {
8898
final ServiceIdentifier singleDeviceAccountAciServiceIdentifier = new AciServiceIdentifier(UUID.randomUUID());
8999
final ServiceIdentifier singleDeviceAccountPniServiceIdentifier = new PniServiceIdentifier(UUID.randomUUID());
90100
final ServiceIdentifier multiDeviceAccountAciServiceIdentifier = new AciServiceIdentifier(UUID.randomUUID());
@@ -105,7 +115,7 @@ void insertMultiRecipientMessage() throws InvalidMessageException, InvalidVersio
105115
new TestRecipient(multiDeviceAccountAciServiceIdentifier, (byte) (Device.PRIMARY_ID + 1), 3, new byte[48]),
106116
new TestRecipient(unresolvedAccountAciServiceIdentifier, Device.PRIMARY_ID, 4, new byte[48]),
107117
new TestRecipient(singleDeviceAccountPniServiceIdentifier, Device.PRIMARY_ID, 5, new byte[48])
108-
));
118+
), sharedPayloadSize);
109119

110120
final SealedSenderMultiRecipientMessage multiRecipientMessage =
111121
SealedSenderMultiRecipientMessage.parse(multiRecipientMessageBytes);
@@ -158,26 +168,46 @@ void insertMultiRecipientMessage() throws InvalidMessageException, InvalidVersio
158168
.setStory(isStory)
159169
.setEphemeral(isEphemeral)
160170
.setUrgent(isUrgent)
161-
.setSharedMrmKey(ByteString.copyFrom(sharedMrmKey))
162171
.build();
163172

173+
final Map<ServiceIdentifier, Envelope> expectedEnvelopesByServiceIdentifier = Stream.of(singleDeviceAccountAciServiceIdentifier, singleDeviceAccountPniServiceIdentifier, multiDeviceAccountAciServiceIdentifier)
174+
.collect(Collectors.toMap(
175+
Function.identity(),
176+
serviceIdentifier -> {
177+
178+
final Envelope.Builder envelopeBuilder = prototypeExpectedMessage.toBuilder()
179+
.setDestinationServiceId(serviceIdentifier.toServiceIdentifierString());
180+
181+
if (expectSharedMrm) {
182+
return envelopeBuilder
183+
.setSharedMrmKey(ByteString.copyFrom(sharedMrmKey))
184+
.build();
185+
}
186+
187+
return envelopeBuilder.setContent(ByteString.copyFrom(multiRecipientMessage.messageForRecipient(
188+
multiRecipientMessage.getRecipients().get(serviceIdentifier.toLibsignal()))))
189+
.build();
190+
191+
}
192+
));
193+
164194
assertEquals(expectedPresenceByAccountAndDeviceId,
165195
messagesManager.insertMultiRecipientMessage(multiRecipientMessage, resolvedRecipients, clientTimestamp, isStory, isEphemeral, isUrgent).join());
166196

167197
verify(messagesCache).insert(any(),
168198
eq(singleDeviceAccountAciServiceIdentifier.uuid()),
169199
eq(Device.PRIMARY_ID),
170-
eq(prototypeExpectedMessage.toBuilder().setDestinationServiceId(singleDeviceAccountAciServiceIdentifier.toServiceIdentifierString()).build()));
200+
eq(expectedEnvelopesByServiceIdentifier.get(singleDeviceAccountAciServiceIdentifier)));
171201

172202
verify(messagesCache).insert(any(),
173203
eq(singleDeviceAccountAciServiceIdentifier.uuid()),
174204
eq(Device.PRIMARY_ID),
175-
eq(prototypeExpectedMessage.toBuilder().setDestinationServiceId(singleDeviceAccountPniServiceIdentifier.toServiceIdentifierString()).build()));
205+
eq(expectedEnvelopesByServiceIdentifier.get(singleDeviceAccountPniServiceIdentifier)));
176206

177207
verify(messagesCache).insert(any(),
178208
eq(multiDeviceAccountAciServiceIdentifier.uuid()),
179209
eq((byte) (Device.PRIMARY_ID + 1)),
180-
eq(prototypeExpectedMessage.toBuilder().setDestinationServiceId(multiDeviceAccountAciServiceIdentifier.toServiceIdentifierString()).build()));
210+
eq(expectedEnvelopesByServiceIdentifier.get(multiDeviceAccountAciServiceIdentifier)));
181211

182212
verify(messagesCache, never()).insert(any(),
183213
eq(unresolvedAccountAciServiceIdentifier.uuid()),

0 commit comments

Comments
 (0)
Please sign in to comment.