Add MRM views experiment to MessagesCache.getMessagesToPersist()

This commit is contained in:
Chris Eager
2024-09-06 17:22:41 -05:00
committed by Chris Eager
parent 5bc6ff0e77
commit 1c617284f3
6 changed files with 90 additions and 50 deletions

View File

@@ -32,11 +32,9 @@ import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
import org.whispersystems.textsecuregcm.entities.MessageProtos;
import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
import org.whispersystems.textsecuregcm.redis.RedisClusterExtension;
import org.whispersystems.textsecuregcm.storage.DynamoDbExtensionSchema.Tables;
import org.whispersystems.textsecuregcm.tests.util.DevicesHelper;
import reactor.core.scheduler.Scheduler;
import reactor.core.scheduler.Schedulers;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
@@ -85,7 +83,7 @@ class MessagePersisterIntegrationTest {
messagesManager = new MessagesManager(messagesDynamoDb, messagesCache, mock(ReportMessageManager.class),
messageDeletionExecutorService);
messagePersister = new MessagePersister(messagesCache, messagesManager, accountsManager,
mock(ClientPresenceManager.class), mock(KeysManager.class), dynamicConfigurationManager, PERSIST_DELAY, 1);
dynamicConfigurationManager, PERSIST_DELAY, 1);
account = mock(Account.class);

View File

@@ -46,7 +46,6 @@ import org.mockito.ArgumentCaptor;
import org.mockito.stubbing.Answer;
import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
import org.whispersystems.textsecuregcm.entities.MessageProtos;
import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
import org.whispersystems.textsecuregcm.redis.RedisClusterExtension;
import org.whispersystems.textsecuregcm.tests.util.DevicesHelper;
import reactor.core.scheduler.Scheduler;
@@ -66,8 +65,6 @@ class MessagePersisterTest {
private MessagesDynamoDb messagesDynamoDb;
private MessagePersister messagePersister;
private AccountsManager accountsManager;
private ClientPresenceManager clientPresenceManager;
private KeysManager keysManager;
private MessagesManager messagesManager;
private Account destinationAccount;
@@ -87,8 +84,6 @@ class MessagePersisterTest {
messagesDynamoDb = mock(MessagesDynamoDb.class);
accountsManager = mock(AccountsManager.class);
clientPresenceManager = mock(ClientPresenceManager.class);
keysManager = mock(KeysManager.class);
destinationAccount = mock(Account.class);;
when(accountsManager.getByAccountIdentifier(DESTINATION_ACCOUNT_UUID)).thenReturn(Optional.of(destinationAccount));
@@ -105,11 +100,10 @@ class MessagePersisterTest {
messageDeliveryScheduler = Schedulers.newBoundedElastic(10, 10_000, "messageDelivery");
messagesCache = new MessagesCache(REDIS_CLUSTER_EXTENSION.getRedisCluster(), sharedExecutorService,
messageDeliveryScheduler, sharedExecutorService, Clock.systemUTC(), dynamicConfigurationManager);
messagePersister = new MessagePersister(messagesCache, messagesManager, accountsManager, clientPresenceManager,
keysManager, dynamicConfigurationManager, PERSIST_DELAY, 1);
messagePersister = new MessagePersister(messagesCache, messagesManager, accountsManager,
dynamicConfigurationManager, PERSIST_DELAY, 1);
when(messagesManager.clear(any(UUID.class), anyByte())).thenReturn(CompletableFuture.completedFuture(null));
when(keysManager.deleteSingleUsePreKeys(any(), anyByte())).thenReturn(CompletableFuture.completedFuture(null));
when(messagesManager.persistMessages(any(UUID.class), any(), any())).thenAnswer(invocation -> {
final UUID destinationUuid = invocation.getArgument(0);

View File

@@ -206,8 +206,7 @@ class MessagesCacheTest {
.collect(Collectors.toList())).get(5, TimeUnit.SECONDS);
assertEquals(messagesToRemove.stream().map(RemovedMessage::fromEnvelope).toList(), removedMessages);
assertEquals(messagesToPreserve,
messagesCache.getMessagesToPersist(DESTINATION_UUID, DESTINATION_DEVICE_ID, messageCount));
assertEquals(messagesToPreserve, get(DESTINATION_UUID, DESTINATION_DEVICE_ID, messageCount));
}
@Test
@@ -620,6 +619,55 @@ class MessagesCacheTest {
}, "Shared MRM data should be deleted asynchronously");
}
@ParameterizedTest
@ValueSource(booleans = {true, false})
void testGetMessagesToPersist(final boolean sharedMrmKeyPresent) throws Exception {
final UUID destinationUuid = UUID.randomUUID();
final byte deviceId = 1;
final UUID messageGuid = UUID.randomUUID();
final MessageProtos.Envelope message = generateRandomMessage(destinationUuid, true);
messagesCache.insert(messageGuid, destinationUuid, deviceId, message);
final SealedSenderMultiRecipientMessage mrm = generateRandomMrmMessage(
new AciServiceIdentifier(destinationUuid), deviceId);
final byte[] sharedMrmDataKey;
if (sharedMrmKeyPresent) {
sharedMrmDataKey = messagesCache.insertSharedMultiRecipientMessagePayload(mrm);
} else {
sharedMrmDataKey = new byte[]{1};
}
final UUID mrmMessageGuid = UUID.randomUUID();
final MessageProtos.Envelope mrmMessage = generateRandomMessage(mrmMessageGuid, true)
.toBuilder()
// clear some things added by the helper
.clearServerGuid()
// mrm views phase 1: messages have content
.setContent(
ByteString.copyFrom(mrm.messageForRecipient(mrm.getRecipients().get(new ServiceId.Aci(destinationUuid)))))
.setSharedMrmKey(ByteString.copyFrom(sharedMrmDataKey))
.build();
messagesCache.insert(mrmMessageGuid, destinationUuid, deviceId, mrmMessage);
final List<MessageProtos.Envelope> messages = get(destinationUuid, deviceId, 100);
assertEquals(2, messages.size());
assertEquals(message.toBuilder()
.setServerGuid(messageGuid.toString())
.build(),
messages.getFirst());
assertEquals(mrmMessage.toBuilder().
clearSharedMrmKey().
setServerGuid(mrmMessageGuid.toString())
.build(),
messages.getLast());
}
private List<MessageProtos.Envelope> get(final UUID destinationUuid, final byte destinationDeviceId,
final int messageCount) {
return Flux.from(messagesCache.get(destinationUuid, destinationDeviceId))