Signal/Signal-Server
1
0
Fork 0
mirror of https://github.com/signalapp/Signal-Server synced 2026-04-21 02:48:03 +01:00
Code Issues Packages Projects Releases Wiki Activity

Implement MRM insert in FoundationDB message store

Browse Source
This commit is contained in:
Ameya Lokare
2025-08-18 13:47:47 -07:00
parent a5423b6e21
commit b8e8fd3313
6 changed files with 499 additions and 177 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
service/src/test/java/org/whispersystems/textsecuregcm/storage/FoundationDbClusterExtension.java Normal file
View File

@@ -0,0 +1,54 @@
/*
* Copyright 2025 Signal Messenger, LLC
* SPDX-License-Identifier: AGPL-3.0-only
*/
package org.whispersystems.textsecuregcm.storage;
import com.apple.foundationdb.Database;
import com.apple.foundationdb.FDB;
import org.junit.jupiter.api.extension.BeforeAllCallback;
import org.junit.jupiter.api.extension.ExtensionContext;
import java.io.IOException;
public class FoundationDbClusterExtension implements BeforeAllCallback, ExtensionContext.Store.CloseableResource {
private FoundationDbDatabaseLifecycleManager[] databaseLifecycleManagers;
private Database[] databases;
public FoundationDbClusterExtension(final int numInstances) {
this.databaseLifecycleManagers = new FoundationDbDatabaseLifecycleManager[numInstances];
this.databases = new Database[numInstances];
}
@Override
public void beforeAll(final ExtensionContext context) throws IOException {
if (databaseLifecycleManagers[0] == null) {
final String serviceContainerNamePrefix = System.getProperty("foundationDb.serviceContainerNamePrefix");
for (int i = 0; i < databaseLifecycleManagers.length; i++) {
final FoundationDbDatabaseLifecycleManager databaseLifecycleManager = serviceContainerNamePrefix != null
? new ServiceContainerFoundationDbDatabaseLifecycleManager(serviceContainerNamePrefix + i)
: new TestcontainersFoundationDbDatabaseLifecycleManager();
databaseLifecycleManager.initializeDatabase(FDB.selectAPIVersion(FoundationDbVersion.getFoundationDbApiVersion()));
databaseLifecycleManagers[i] = databaseLifecycleManager;
databases[i] = databaseLifecycleManager.getDatabase();
}
}
}
public Database[] getDatabases() {
return databases;
}
@Override
public void close() throws Throwable {
if (databaseLifecycleManagers[0] != null) {
for (final FoundationDbDatabaseLifecycleManager databaseLifecycleManager : databaseLifecycleManagers) {
databaseLifecycleManager.closeDatabase();
}
}
}
}

43
service/src/test/java/org/whispersystems/textsecuregcm/storage/FoundationDbExtension.java
View File

@@ -1,43 +0,0 @@
/*
* Copyright 2025 Signal Messenger, LLC
* SPDX-License-Identifier: AGPL-3.0-only
*/
package org.whispersystems.textsecuregcm.storage;
import com.apple.foundationdb.Database;
import com.apple.foundationdb.FDB;
import java.io.IOException;
import org.junit.jupiter.api.extension.BeforeAllCallback;
import org.junit.jupiter.api.extension.ExtensionContext;
public class FoundationDbExtension implements BeforeAllCallback, ExtensionContext.Store.CloseableResource {
private static FoundationDbDatabaseLifecycleManager databaseLifecycleManager;
@Override
public void beforeAll(final ExtensionContext context) throws IOException {
if (databaseLifecycleManager == null) {
final String serviceContainerName = System.getProperty("foundationDb.serviceContainerName");
databaseLifecycleManager = serviceContainerName != null
? new ServiceContainerFoundationDbDatabaseLifecycleManager(serviceContainerName)
: new TestcontainersFoundationDbDatabaseLifecycleManager();
databaseLifecycleManager.initializeDatabase(FDB.selectAPIVersion(FoundationDbVersion.getFoundationDbApiVersion()));
context.getRoot().getStore(ExtensionContext.Namespace.GLOBAL).put(getClass().getName(), this);
}
}
public Database getDatabase() {
return databaseLifecycleManager.getDatabase();
}
@Override
public void close() throws Throwable {
if (databaseLifecycleManager != null) {
databaseLifecycleManager.closeDatabase();
}
}
}

34
service/src/test/java/org/whispersystems/textsecuregcm/storage/FoundationDbTest.java
View File

@@ -1,34 +0,0 @@
/*
* Copyright 2025 Signal Messenger, LLC
* SPDX-License-Identifier: AGPL-3.0-only
*/
package org.whispersystems.textsecuregcm.storage;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
import org.whispersystems.textsecuregcm.util.TestRandomUtil;
import java.nio.charset.StandardCharsets;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
public class FoundationDbTest {
@RegisterExtension
static FoundationDbExtension FOUNDATION_DB_EXTENSION = new FoundationDbExtension();
@Test
void setGetValue() {
final byte[] key = "test".getBytes(StandardCharsets.UTF_8);
final byte[] value = TestRandomUtil.nextBytes(16);
FOUNDATION_DB_EXTENSION.getDatabase().run(transaction -> {
transaction.set(key, value);
return null;
});
final byte[] retrievedValue = FOUNDATION_DB_EXTENSION.getDatabase().run(transaction -> transaction.get(key).join());
assertArrayEquals(value, retrievedValue);
}
}

261
service/src/test/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStoreTest.java
View File

@@ -3,9 +3,11 @@ package org.whispersystems.textsecuregcm.storage.foundationdb;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import com.apple.foundationdb.Database;
import com.apple.foundationdb.KeyValue;
import com.apple.foundationdb.async.AsyncUtil;
import com.apple.foundationdb.tuple.Tuple;
import com.apple.foundationdb.tuple.Versionstamp;
import com.google.protobuf.ByteString;
@@ -14,16 +16,21 @@ import java.io.UncheckedIOException;
import java.time.Clock;
import java.time.Instant;
import java.time.ZoneId;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.Executors;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;
import io.dropwizard.util.DataSize;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.Timeout;
@@ -35,7 +42,7 @@ import org.junit.jupiter.params.provider.ValueSource;
import org.whispersystems.textsecuregcm.entities.MessageProtos;
import org.whispersystems.textsecuregcm.identity.AciServiceIdentifier;
import org.whispersystems.textsecuregcm.storage.Device;
import org.whispersystems.textsecuregcm.storage.FoundationDbExtension;
import org.whispersystems.textsecuregcm.storage.FoundationDbClusterExtension;
import org.whispersystems.textsecuregcm.util.Conversions;
import org.whispersystems.textsecuregcm.util.TestRandomUtil;
@@ -43,7 +50,7 @@ import org.whispersystems.textsecuregcm.util.TestRandomUtil;
class FoundationDbMessageStoreTest {
@RegisterExtension
static FoundationDbExtension FOUNDATION_DB_EXTENSION = new FoundationDbExtension();
static FoundationDbClusterExtension FOUNDATION_DB_EXTENSION = new FoundationDbClusterExtension(2);
private FoundationDbMessageStore foundationDbMessageStore;
@@ -52,7 +59,7 @@ class FoundationDbMessageStoreTest {
@BeforeEach
void setup() {
foundationDbMessageStore = new FoundationDbMessageStore(
new Database[]{FOUNDATION_DB_EXTENSION.getDatabase()},
FOUNDATION_DB_EXTENSION.getDatabases(),
Executors.newVirtualThreadPerTaskExecutor(),
CLOCK);
}
@@ -60,7 +67,7 @@ class FoundationDbMessageStoreTest {
@ParameterizedTest
@MethodSource
void insert(final long presenceUpdatedBeforeSeconds, final boolean ephemeral, final boolean expectMessagesInserted,
final boolean expectVersionstampUpdated) {
final boolean expectVersionstampUpdated, final boolean expectPresenceState) {
final AciServiceIdentifier aci = new AciServiceIdentifier(UUID.randomUUID());
final List<Byte> deviceIds = IntStream.range(Device.PRIMARY_ID, Device.PRIMARY_ID + 6)
.mapToObj(i -> (byte) i)
@@ -68,15 +75,19 @@ class FoundationDbMessageStoreTest {
deviceIds.forEach(deviceId -> writePresenceKey(aci, deviceId, 1, presenceUpdatedBeforeSeconds));
final Map<Byte, MessageProtos.Envelope> messagesByDeviceId = deviceIds.stream()
.collect(Collectors.toMap(Function.identity(), _ -> generateRandomMessage(ephemeral)));
final Optional<Versionstamp> versionstamp = foundationDbMessageStore.insert(aci, messagesByDeviceId).join();
assertNotNull(versionstamp);
final Map<Byte, FoundationDbMessageStore.InsertResult> result = foundationDbMessageStore.insert(aci, messagesByDeviceId).join();
assertNotNull(result);
final Optional<Versionstamp> returnedVersionstamp = result.values().stream().findFirst()
.flatMap(FoundationDbMessageStore.InsertResult::versionstamp);
if (expectMessagesInserted) {
assertTrue(versionstamp.isPresent());
assertTrue(returnedVersionstamp.isPresent());
assertTrue(result.values().stream().allMatch(insertResult -> returnedVersionstamp.equals(insertResult.versionstamp())));
final Map<Byte, MessageProtos.Envelope> storedMessagesByDeviceId = deviceIds.stream()
.collect(Collectors.toMap(Function.identity(), deviceId -> {
try {
return MessageProtos.Envelope.parseFrom(getMessageByVersionstamp(aci, deviceId, versionstamp.get()));
return MessageProtos.Envelope.parseFrom(
getMessageByVersionstamp(aci, deviceId, returnedVersionstamp.get()));
} catch (final InvalidProtocolBufferException e) {
throw new UncheckedIOException(e);
}
@@ -84,28 +95,32 @@ class FoundationDbMessageStoreTest {
assertEquals(messagesByDeviceId, storedMessagesByDeviceId);
} else {
assertTrue(versionstamp.isEmpty());
assertTrue(result.values().stream().allMatch(insertResult -> insertResult.versionstamp().isEmpty()));
}
if (expectVersionstampUpdated) {
assertEquals(versionstamp, getMessagesAvailableWatch(aci),
final Optional<Versionstamp> messagesAvailableWatchVersionstamp = getMessagesAvailableWatch(aci);
assertTrue(messagesAvailableWatchVersionstamp.isPresent());
assertEquals(returnedVersionstamp, messagesAvailableWatchVersionstamp,
"messages available versionstamp should be the versionstamp of the last insert transaction");
} else {
assertTrue(getMessagesAvailableWatch(aci).isEmpty());
}
assertTrue(result.values().stream().allMatch(insertResult -> insertResult.present() == expectPresenceState));
}
private static Stream<Arguments> insert() {
return Stream.of(
Arguments.argumentSet("Non-ephemeral messages with all devices online",
10L, false, true, true),
10L, false, true, true, true),
Arguments.argumentSet(
"Ephemeral messages with presence updated exactly at the second before which the device would be considered offline",
300L, true, true, true),
300L, true, true, true, true),
Arguments.argumentSet("Non-ephemeral messages for with all devices offline",
310L, false, true, false),
310L, false, true, false, false),
Arguments.argumentSet("Ephemeral messages with all devices offline",
310L, true, false, false)
310L, true, false, false, false)
);
}
@@ -113,10 +128,15 @@ class FoundationDbMessageStoreTest {
void versionstampCorrectlyUpdatedOnMultipleInserts() {
final AciServiceIdentifier aci = new AciServiceIdentifier(UUID.randomUUID());
writePresenceKey(aci, Device.PRIMARY_ID, 1, 10L);
foundationDbMessageStore.insert(aci, Map.of(Device.PRIMARY_ID, generateRandomMessage(false))).join();
final Optional<Versionstamp> secondMessageVersionstamp = foundationDbMessageStore.insert(aci,
foundationDbMessageStore.insert(Map.of(aci, Map.of(Device.PRIMARY_ID, generateRandomMessage(false)))).join();
final Map<Byte, FoundationDbMessageStore.InsertResult> secondMessageInsertResult = foundationDbMessageStore.insert(aci,
Map.of(Device.PRIMARY_ID, generateRandomMessage(false))).join();
assertEquals(secondMessageVersionstamp, getMessagesAvailableWatch(aci));
final Optional<Versionstamp> messagesAvailableWatchVersionstamp = getMessagesAvailableWatch(aci);
assertTrue(messagesAvailableWatchVersionstamp.isPresent());
assertEquals(
secondMessageInsertResult.get(Device.PRIMARY_ID).versionstamp(),
messagesAvailableWatchVersionstamp);
}
@ParameterizedTest
@@ -130,24 +150,29 @@ class FoundationDbMessageStoreTest {
writePresenceKey(aci, Device.PRIMARY_ID, 1, 10L);
final Map<Byte, MessageProtos.Envelope> messagesByDeviceId = deviceIds.stream()
.collect(Collectors.toMap(Function.identity(), _ -> generateRandomMessage(ephemeral)));
final Optional<Versionstamp> versionstamp = foundationDbMessageStore.insert(aci, messagesByDeviceId).join();
assertNotNull(versionstamp);
assertTrue(versionstamp.isPresent(),
"versionstamp should be present since at least one message should be inserted");
final Map<Byte, FoundationDbMessageStore.InsertResult> result = foundationDbMessageStore.insert(aci, messagesByDeviceId).join();
assertNotNull(result);
final Optional<Versionstamp> returnedVersionstamp = result.get(Device.PRIMARY_ID).versionstamp();
assertTrue(returnedVersionstamp.isPresent(),
"versionstamp should be present for online device");
assertArrayEquals(
messagesByDeviceId.get(Device.PRIMARY_ID).toByteArray(),
getMessageByVersionstamp(aci, Device.PRIMARY_ID, versionstamp.get()),
getMessageByVersionstamp(aci, Device.PRIMARY_ID, returnedVersionstamp.get()),
"Message for primary should always be stored since it has a recently updated presence");
if (ephemeral) {
assertTrue(IntStream.range(Device.PRIMARY_ID + 1, Device.PRIMARY_ID + 6)
.mapToObj(deviceId -> getMessageByVersionstamp(aci, (byte) deviceId, versionstamp.get()))
.mapToObj(deviceId -> getMessageByVersionstamp(aci, (byte) deviceId, returnedVersionstamp.get()))
.allMatch(Objects::isNull), "Ephemeral messages for non-present devices must not be stored");
assertTrue(IntStream.range(Device.PRIMARY_ID + 1, Device.PRIMARY_ID + 6)
.mapToObj(deviceId -> result.get((byte) deviceId).versionstamp())
.allMatch(Optional::isEmpty),
"Unexpected versionstamp found for one or more devices that didn't have any messages inserted");
} else {
IntStream.range(Device.PRIMARY_ID + 1, Device.PRIMARY_ID)
.forEach(deviceId -> {
final byte[] messageBytes = getMessageByVersionstamp(aci, (byte) deviceId, versionstamp.get());
final byte[] messageBytes = getMessageByVersionstamp(aci, (byte) deviceId, returnedVersionstamp.get());
assertEquals(messagesByDeviceId.get((byte) deviceId).toByteArray(), messageBytes,
"Non-ephemeral messages must always be stored");
});
@@ -169,23 +194,184 @@ class FoundationDbMessageStoreTest {
Conversions.longToByteArray(constructPresenceValue(42, getEpochSecondsBeforeClock(5))), true),
Arguments.argumentSet("Presence updated same second as current time",
Conversions.longToByteArray(constructPresenceValue(42, getEpochSecondsBeforeClock(0))), true),
Arguments.argumentSet("Presence updated exactly at the second before which it would have been considered offline",
Arguments.argumentSet(
"Presence updated exactly at the second before which it would have been considered offline",
Conversions.longToByteArray(constructPresenceValue(42, getEpochSecondsBeforeClock(300))), true),
Arguments.argumentSet("Presence expired",
Conversions.longToByteArray(constructPresenceValue(42, getEpochSecondsBeforeClock(400))), false)
);
}
/// Represents a cohort of recipients with the same config
record MultiRecipientTestConfig(int shardNum, int numRecipients, boolean devicePresent,
boolean generateEphemeralMessages, boolean expectMessagesInserted) {}
@ParameterizedTest
@MethodSource
void insertMultiRecipient(final List<MultiRecipientTestConfig> testConfigs, final DataSize contentSize,
final int[] expectedNumTransactionsByShard) {
// Generate a list of ACIs for each test config
final List<List<AciServiceIdentifier>> acisByConfig = testConfigs.stream()
.map(testConfig -> IntStream.range(0, testConfig.numRecipients())
.mapToObj(_ -> generateRandomAciForShard(testConfig.shardNum()))
.toList())
.toList();
// Generate MRM bundles for each ACI, for each test config. Later, we'll assert if the stored messages (if expected)
// are the same as those we generated.
final List<Map<AciServiceIdentifier, Map<Byte, MessageProtos.Envelope>>> mrmByConfig = IntStream.range(0,
testConfigs.size())
.mapToObj(i -> {
final List<AciServiceIdentifier> acis = acisByConfig.get(i);
final MultiRecipientTestConfig testConfig = testConfigs.get(i);
return acis.stream()
.collect(Collectors.toMap(
Function.identity(),
_ -> Map.of(Device.PRIMARY_ID,
generateRandomMessage(testConfig.generateEphemeralMessages(), (int) contentSize.toBytes()))));
})
.toList();
// Create the consolidated MRM bundle by ACI.
final Map<AciServiceIdentifier, Map<Byte, MessageProtos.Envelope>> mrmBundle = new HashMap<>();
mrmByConfig.forEach(mrmBundle::putAll);
// Write a presence key for the cohort of recipients if the config indicates that the device must be present.
for (int i = 0; i < testConfigs.size(); i++) {
final List<AciServiceIdentifier> acis = acisByConfig.get(i);
final MultiRecipientTestConfig testConfig = testConfigs.get(i);
if (testConfig.devicePresent()) {
acis.forEach(aci -> writePresenceKey(aci, Device.PRIMARY_ID, 1, 10L));
}
}
final Map<AciServiceIdentifier, Map<Byte, FoundationDbMessageStore.InsertResult>> result = foundationDbMessageStore.insert(mrmBundle).join();
assertNotNull(result);
// Compute the set of versionstamps by shard number from the individual device insert results, so that we can
// assert that each shard has the expected number of committed transactions.
final Map<Integer, Set<Versionstamp>> returnedVersionstampsByShard = new HashMap<>();
result.forEach((aci, deviceResults) -> {
final int shardNum = foundationDbMessageStore.hashAciToShardNumber(aci);
final Set<Versionstamp> versionstampSet = returnedVersionstampsByShard.computeIfAbsent(shardNum, _ -> new HashSet<>());
deviceResults.forEach((_, deviceResult) -> deviceResult.versionstamp().ifPresent(versionstampSet::add));
});
final int[] returnedNumVersionstampsByShard = new int[FOUNDATION_DB_EXTENSION.getDatabases().length];
for (int i = 0; i < returnedNumVersionstampsByShard.length; i++) {
returnedNumVersionstampsByShard[i] = returnedVersionstampsByShard.getOrDefault(i, Collections.emptySet()).size();
}
assertArrayEquals(expectedNumTransactionsByShard, returnedNumVersionstampsByShard);
// For each cohort of recipients, check whether the stored messages (if expected) are the same as those we inserted
// and whether the returned device presence states are the same as the configured states.
IntStream.range(0, testConfigs.size()).forEach(i -> {
final List<AciServiceIdentifier> acis = acisByConfig.get(i);
final MultiRecipientTestConfig shardConfig = testConfigs.get(i);
if (shardConfig.expectMessagesInserted()) {
final Map<AciServiceIdentifier, Map<Byte, MessageProtos.Envelope>> storedMrmBundle = acis.stream()
.collect(Collectors.toMap(Function.identity(), aci -> {
final List<KeyValue> items = getItemsInDeviceQueue(aci, Device.PRIMARY_ID);
assertEquals(1, items.size());
try {
final MessageProtos.Envelope envelope = MessageProtos.Envelope.parseFrom(items.getFirst().getValue());
return Map.of(Device.PRIMARY_ID, envelope);
} catch (final InvalidProtocolBufferException e) {
throw new UncheckedIOException(e);
}
}));
assertEquals(mrmByConfig.get(i), storedMrmBundle,
"Stored message bundle does not match inserted message bundle");
} else {
assertEquals(0, acis
.stream()
.mapToInt(aci -> getItemsInDeviceQueue(aci, Device.PRIMARY_ID).size())
.sum(), "Unexpected messages found in device queue");
}
assertTrue(acis
.stream()
.allMatch(
aci -> result.get(aci).get(Device.PRIMARY_ID).present() == shardConfig.devicePresent()),
"Device presence state from insert result does not match expected state");
});
}
static Stream<Arguments> insertMultiRecipient() {
return Stream.of(
Arguments.argumentSet("Multiple recipients on a single shard should result in a single transaction",
List.of(
new MultiRecipientTestConfig(0, 5, true, false, true)),
DataSize.bytes(128), new int[] {1, 0}),
Arguments.argumentSet(
"Multiple recipients on a single shard exceeding the transaction limit should be broken up into multiple transactions",
List.of(
new MultiRecipientTestConfig(0, 15, true, false, true)),
DataSize.kilobytes(90), new int[] {2, 0}),
Arguments.argumentSet("Multiple recipients on different shards should result in multiple transactions",
List.of(
new MultiRecipientTestConfig(0, 5, true, false, true),
new MultiRecipientTestConfig(1, 5, true, false, true)),
DataSize.bytes(128), new int[] {1, 1}),
Arguments.argumentSet(
"Multiple recipients on different shards each exceeding the transaction limit should be broken up into multiple transactions on each shard",
List.of(
new MultiRecipientTestConfig(0, 15, true, false, true),
new MultiRecipientTestConfig(1, 15, true, false, true)),
DataSize.kilobytes(90), new int[] {2, 2}),
Arguments.argumentSet(
"Multiple recipients on a single shard with ephemeral messages and no devices present should result in no transactions committed",
List.of(
new MultiRecipientTestConfig(0, 5, false, true, false)),
DataSize.bytes(128), new int[] {0, 0}),
Arguments.argumentSet(
"Multiple recipients on different shards with ephemeral messages and no devices present should result in no transactions committed",
List.of(
new MultiRecipientTestConfig(0, 5, false, true, false),
new MultiRecipientTestConfig(1, 5, false, true, false)),
DataSize.bytes(128), new int[] {0, 0}),
Arguments.argumentSet(
"Multiple recipients on two shards with one shard having no devices present should result in only one transaction",
List.of(
new MultiRecipientTestConfig(0, 5, false, true, false),
new MultiRecipientTestConfig(1, 5, true, true, true)),
DataSize.bytes(128), new int[] {0, 1}),
Arguments.argumentSet(
"Multiple recipients on a single shard with some recipients having no devices present should result in only one transaction",
List.of(
new MultiRecipientTestConfig(0, 3, false, true, false),
new MultiRecipientTestConfig(0, 3, true, true, true)),
DataSize.bytes(128), new int[] {1, 0}),
Arguments.argumentSet(
"Multiple recipients on a single shard with total size just exceeding 2 chunks should result in 3 transactions",
List.of(
new MultiRecipientTestConfig(0, 23, true, false, true)),
DataSize.kilobytes(90), new int[] {3, 0})
);
}
@Test
void insertEmptyBundle() {
assertThrows(IllegalArgumentException.class, () -> foundationDbMessageStore.insert(
Map.of(generateRandomAciForShard(0), Collections.emptyMap())));
}
private static MessageProtos.Envelope generateRandomMessage(final boolean ephemeral) {
return generateRandomMessage(ephemeral, 16);
}
private static MessageProtos.Envelope generateRandomMessage(final boolean ephemeral, final int contentSize) {
return MessageProtos.Envelope.newBuilder()
.setContent(ByteString.copyFrom(TestRandomUtil.nextBytes(16)))
.setContent(ByteString.copyFrom(TestRandomUtil.nextBytes(contentSize)))
.setEphemeral(ephemeral)
.build();
}
private byte[] getMessageByVersionstamp(final AciServiceIdentifier aci, final byte deviceId,
final Versionstamp versionstamp) {
return FOUNDATION_DB_EXTENSION.getDatabase().read(transaction -> {
return foundationDbMessageStore.getShardForAci(aci).read(transaction -> {
final byte[] key = foundationDbMessageStore.getDeviceQueueSubspace(aci, deviceId)
.pack(Tuple.from(versionstamp));
return transaction.get(key);
@@ -193,7 +379,7 @@ class FoundationDbMessageStoreTest {
}
private Optional<Versionstamp> getMessagesAvailableWatch(final AciServiceIdentifier aci) {
return FOUNDATION_DB_EXTENSION.getDatabase()
return foundationDbMessageStore.getShardForAci(aci)
.read(transaction -> transaction.get(foundationDbMessageStore.getMessagesAvailableWatchKey(aci))
.thenApply(value -> value == null ? null : Tuple.fromBytes(value).getVersionstamp(0))
.thenApply(Optional::ofNullable))
@@ -202,7 +388,7 @@ class FoundationDbMessageStoreTest {
private void writePresenceKey(final AciServiceIdentifier aci, final byte deviceId, final int serverId,
final long secondsBeforeCurrentTime) {
FOUNDATION_DB_EXTENSION.getDatabase().run(transaction -> {
foundationDbMessageStore.getShardForAci(aci).run(transaction -> {
final byte[] presenceKey = foundationDbMessageStore.getPresenceKey(aci, deviceId);
final long presenceUpdateEpochSeconds = getEpochSecondsBeforeClock(secondsBeforeCurrentTime);
final long presenceValue = constructPresenceValue(serverId, presenceUpdateEpochSeconds);
@@ -219,4 +405,19 @@ class FoundationDbMessageStoreTest {
return (long) (serverId & 0x0ffff) << 48 | (presenceUpdateEpochSeconds & 0x0000ffffffffffffL);
}
private AciServiceIdentifier generateRandomAciForShard(final int shardNumber) {
assert shardNumber < FOUNDATION_DB_EXTENSION.getDatabases().length;
while (true) {
final AciServiceIdentifier aci = new AciServiceIdentifier(UUID.randomUUID());
if (foundationDbMessageStore.hashAciToShardNumber(aci) == shardNumber) {
return aci;
}
}
}
private List<KeyValue> getItemsInDeviceQueue(final AciServiceIdentifier aci, final byte deviceId) {
return foundationDbMessageStore.getShardForAci(aci).readAsync(transaction -> AsyncUtil.collect(transaction.getRange(
foundationDbMessageStore.getDeviceQueueSubspace(aci, deviceId).range()))).join();
}
}