mirror of
https://github.com/signalapp/Signal-Server
synced 2026-04-20 20:38:04 +01:00
Add a dedicated size estimation method to MessagesCache
This commit is contained in:
@@ -492,118 +492,64 @@ class MessagesCacheTest {
|
||||
}
|
||||
|
||||
@Test
|
||||
void testMessagesToPersistPagination() throws InterruptedException {
|
||||
void testEstimatePersistedQueueSize() {
|
||||
final UUID destinationUuid = UUID.randomUUID();
|
||||
final ServiceIdentifier serviceId = new AciServiceIdentifier(destinationUuid);
|
||||
final byte deviceId = 1;
|
||||
final byte[] queueKey = MessagesCache.getMessageQueueKey(destinationUuid, deviceId);
|
||||
|
||||
final List<MessageProtos.Envelope> messages = IntStream.range(0, 60)
|
||||
.mapToObj(i -> switch (i % 3) {
|
||||
// Stale MRM
|
||||
case 0 -> generateRandomMessage(UUID.randomUUID(), serviceId, true)
|
||||
// Should count all non-ephemeral, non-stale message bytes
|
||||
long expectedQueueSize = 0L;
|
||||
for (int i = 0; i < 400; i++) {
|
||||
final MessageProtos.Envelope messageToInsert = switch (i % 4) {
|
||||
// An MRM message
|
||||
case 0 -> {
|
||||
|
||||
// First generate a random MRM message
|
||||
final SealedSenderMultiRecipientMessage mrm = generateRandomMrmMessage(serviceId, deviceId);
|
||||
final SealedSenderMultiRecipientMessage.Recipient recepient = mrm.getRecipients()
|
||||
.get(serviceId.toLibsignal());
|
||||
|
||||
// Calculate the size of a message that has the shared content in it
|
||||
final MessageProtos.Envelope message = generateRandomMessage(UUID.randomUUID(), serviceId, true)
|
||||
.toBuilder()
|
||||
.setContent(ByteString.copyFrom(mrm.messageForRecipient(recepient)))
|
||||
.build();
|
||||
expectedQueueSize += message.getSerializedSize();
|
||||
byte[] sharedMrmDataKey = messagesCache.insertSharedMultiRecipientMessagePayload(mrm).join();
|
||||
|
||||
// Insert the MRM message without the content
|
||||
yield message
|
||||
.toBuilder()
|
||||
.clearContent()
|
||||
.setSharedMrmKey(ByteString.copyFrom(sharedMrmDataKey))
|
||||
.build();
|
||||
}
|
||||
|
||||
// A stale MRM message
|
||||
case 1 ->
|
||||
generateRandomMessage(UUID.randomUUID(), serviceId, true)
|
||||
.toBuilder()
|
||||
// clear some things added by the helper
|
||||
.clearContent()
|
||||
.setSharedMrmKey(MessagesCache.STALE_MRM_KEY)
|
||||
.build();
|
||||
// ephemeral message
|
||||
case 1 -> generateRandomMessage(UUID.randomUUID(), serviceId, true)
|
||||
.toBuilder()
|
||||
.setEphemeral(true).build();
|
||||
// standard message
|
||||
case 2 -> generateRandomMessage(UUID.randomUUID(), serviceId, true);
|
||||
default -> throw new IllegalStateException();
|
||||
})
|
||||
.toList();
|
||||
for (MessageProtos.Envelope envelope : messages) {
|
||||
messagesCache.insert(UUID.fromString(envelope.getServerGuid()), destinationUuid, deviceId, envelope).join();
|
||||
|
||||
// An ephemeral message
|
||||
case 2 -> generateRandomMessage(UUID.randomUUID(), serviceId, true).toBuilder().setEphemeral(true).build();
|
||||
|
||||
// A standardard message
|
||||
case 3 -> {
|
||||
final MessageProtos.Envelope message = generateRandomMessage(UUID.randomUUID(), serviceId, true);
|
||||
expectedQueueSize += message.getSerializedSize();
|
||||
yield message;
|
||||
}
|
||||
|
||||
default -> throw new IllegalStateException();
|
||||
};
|
||||
messagesCache.insert(UUID.fromString(messageToInsert.getServerGuid()), destinationUuid, deviceId, messageToInsert).join();
|
||||
}
|
||||
|
||||
final List<UUID> expectedGuidsToPersist = messages.stream()
|
||||
.filter(envelope -> !envelope.getEphemeral() && !envelope.hasSharedMrmKey())
|
||||
.map(envelope -> UUID.fromString(envelope.getServerGuid()))
|
||||
.limit(10)
|
||||
.collect(Collectors.toList());
|
||||
|
||||
// Fetch 10 messages which should discard 20 ephemeral stale messages, and leave the rest
|
||||
final List<UUID> actual = messagesCache.getMessagesToPersist(destinationUuid, deviceId, 10).stream()
|
||||
.map(envelope -> UUID.fromString(envelope.getServerGuid()))
|
||||
.toList();
|
||||
assertIterableEquals(expectedGuidsToPersist, actual);
|
||||
|
||||
// Eventually, the 20 ephemeral/stale messages should be discarded
|
||||
assertTimeoutPreemptively(Duration.ofSeconds(1), () -> {
|
||||
while (REDIS_CLUSTER_EXTENSION.getRedisCluster()
|
||||
.withBinaryCluster(conn -> conn.sync().zcard(queueKey)) != 40) {
|
||||
Thread.sleep(1);
|
||||
}
|
||||
}, "Ephemeral and stale messages should be deleted asynchronously");
|
||||
|
||||
// Let all pending tasks finish and make sure no more stale messages have been deleted
|
||||
sharedExecutorService.shutdown();
|
||||
sharedExecutorService.awaitTermination(1, TimeUnit.SECONDS);
|
||||
assertEquals(REDIS_CLUSTER_EXTENSION.getRedisCluster()
|
||||
.withBinaryCluster(conn -> conn.sync().zcard(queueKey)).longValue(), 40);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testMessagesToPersistReactive() {
|
||||
final UUID destinationUuid = UUID.randomUUID();
|
||||
final ServiceIdentifier serviceId = new AciServiceIdentifier(destinationUuid);
|
||||
final byte deviceId = 1;
|
||||
final byte[] messageQueueKey = MessagesCache.getMessageQueueKey(destinationUuid, deviceId);
|
||||
|
||||
final List<MessageProtos.Envelope> messages = IntStream.range(0, 200)
|
||||
.mapToObj(i -> {
|
||||
if (i % 3 == 0) {
|
||||
final SealedSenderMultiRecipientMessage mrm = generateRandomMrmMessage(serviceId, deviceId);
|
||||
byte[] sharedMrmDataKey = messagesCache.insertSharedMultiRecipientMessagePayload(mrm).join();
|
||||
return generateRandomMessage(UUID.randomUUID(), serviceId, true)
|
||||
.toBuilder()
|
||||
// clear some things added by the helper
|
||||
.clearContent()
|
||||
.setSharedMrmKey(ByteString.copyFrom(sharedMrmDataKey))
|
||||
.build();
|
||||
} else if (i % 13 == 0) {
|
||||
return generateRandomMessage(UUID.randomUUID(), serviceId, true).toBuilder().setEphemeral(true).build();
|
||||
} else if (i % 17 == 0) {
|
||||
return generateRandomMessage(UUID.randomUUID(), serviceId, true)
|
||||
.toBuilder()
|
||||
// clear some things added by the helper
|
||||
.clearContent()
|
||||
.setSharedMrmKey(MessagesCache.STALE_MRM_KEY)
|
||||
.build();
|
||||
} else {
|
||||
return generateRandomMessage(UUID.randomUUID(), serviceId, true);
|
||||
}
|
||||
})
|
||||
.toList();
|
||||
|
||||
for (MessageProtos.Envelope envelope : messages) {
|
||||
messagesCache.insert(UUID.fromString(envelope.getServerGuid()), destinationUuid, deviceId, envelope).join();
|
||||
}
|
||||
|
||||
final List<MessageProtos.Envelope> expected = messages.stream()
|
||||
.filter(envelope -> !envelope.getEphemeral() &&
|
||||
(envelope.getSharedMrmKey() == null || !envelope.getSharedMrmKey().equals(MessagesCache.STALE_MRM_KEY)))
|
||||
.toList();
|
||||
final List<MessageProtos.Envelope> actual = messagesCache
|
||||
.getMessagesToPersistReactive(destinationUuid, deviceId, 7).collectList().block();
|
||||
|
||||
assertEquals(expected.size(), actual.size());
|
||||
for (int i = 0; i < actual.size(); i++) {
|
||||
assertNotNull(actual.get(i).getContent());
|
||||
assertEquals(actual.get(i).getServerGuid(), expected.get(i).getServerGuid());
|
||||
}
|
||||
|
||||
// Ephemeral messages and stale MRM messages are asynchronously deleted, but eventually they should all be removed
|
||||
assertTimeoutPreemptively(Duration.ofSeconds(1), () -> {
|
||||
while (REDIS_CLUSTER_EXTENSION.getRedisCluster()
|
||||
.withBinaryCluster(conn -> conn.sync().zcard(messageQueueKey)) != expected.size()) {
|
||||
Thread.sleep(1);
|
||||
}
|
||||
}, "Ephemeral and stale messages should be deleted asynchronously");
|
||||
long actualQueueSize = messagesCache.estimatePersistedQueueSizeBytes(destinationUuid, deviceId).join();
|
||||
assertEquals(expectedQueueSize, actualQueueSize);
|
||||
}
|
||||
|
||||
@ParameterizedTest
|
||||
|
||||
Reference in New Issue
Block a user