diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/configuration/FoundationDbClusterConfiguration.java b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/FoundationDbClusterConfiguration.java new file mode 100644 index 000000000..88febdf71 --- /dev/null +++ b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/FoundationDbClusterConfiguration.java @@ -0,0 +1,19 @@ +/* + * Copyright 2026 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.whispersystems.textsecuregcm.configuration; + +import jakarta.validation.constraints.AssertTrue; +import org.apache.commons.lang3.StringUtils; +import javax.annotation.Nullable; + +public record FoundationDbClusterConfiguration(@Nullable String clusterFileUrl, + @Nullable String clusterFileContents) { + + @AssertTrue + public boolean isSingleClusterFileSourceSpecified() { + return StringUtils.isBlank(clusterFileUrl) ^ StringUtils.isBlank(clusterFileContents); + } +} diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/configuration/FoundationDbMessagesConfiguration.java b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/FoundationDbMessagesConfiguration.java new file mode 100644 index 000000000..1cf82f371 --- /dev/null +++ b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/FoundationDbMessagesConfiguration.java @@ -0,0 +1,45 @@ +/* + * Copyright 2026 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.whispersystems.textsecuregcm.configuration; + +import jakarta.validation.Valid; +import jakarta.validation.constraints.AssertTrue; +import jakarta.validation.constraints.Max; +import jakarta.validation.constraints.NotEmpty; +import jakarta.validation.constraints.PositiveOrZero; +import jakarta.validation.constraints.Size; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import org.whispersystems.textsecuregcm.storage.foundationdb.FoundationDbMessageStore; + +public record FoundationDbMessagesConfiguration(@NotEmpty Map clusters, + @NotEmpty Map<@PositiveOrZero @Max(FoundationDbMessageStore.MAX_EPOCHS - 1) Integer, @Size(min = 1, max = FoundationDbMessageStore.MAX_SHARDS - 1) List> epochs) { + + @AssertTrue + boolean isEveryEpochClusterConfigured() { + for (final List clustersInEpoch : epochs().values()) { + for (final String cluster : clustersInEpoch) { + if (!clusters.containsKey(cluster)) { + return false; + } + } + } + + return true; + } + + @AssertTrue + boolean isEveryEpochFreeOfDuplicates() { + for (final List clustersInEpoch : epochs().values()) { + if (new HashSet<>(clustersInEpoch).size() != clustersInEpoch.size()) { + return false; + } + } + + return true; + } +} diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicConfiguration.java b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicConfiguration.java index 2a9b90bfe..aa038a9d6 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicConfiguration.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicConfiguration.java @@ -84,6 +84,11 @@ public class DynamicConfiguration { @Valid private DynamicTurnConfiguration turn = new DynamicTurnConfiguration(); + @JsonProperty + @Valid + private DynamicFoundationDbMessagesConfiguration foundationDbMessages = + new DynamicFoundationDbMessagesConfiguration(0); + public Optional getExperimentEnrollmentConfiguration( final String experimentName) { return Optional.ofNullable(experiments.get(experimentName)); @@ -153,4 +158,8 @@ public class DynamicConfiguration { public DynamicTurnConfiguration getTurnConfiguration() { return turn; } + + public DynamicFoundationDbMessagesConfiguration getFoundationDbMessagesConfiguration() { + return foundationDbMessages; + } } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicFoundationDbMessagesConfiguration.java b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicFoundationDbMessagesConfiguration.java new file mode 100644 index 000000000..9d1f629ae --- /dev/null +++ b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicFoundationDbMessagesConfiguration.java @@ -0,0 +1,13 @@ +/* + * Copyright 2026 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.whispersystems.textsecuregcm.configuration.dynamic; + +import jakarta.validation.constraints.Max; +import jakarta.validation.constraints.PositiveOrZero; +import org.whispersystems.textsecuregcm.storage.foundationdb.FoundationDbMessageStore; + +public record DynamicFoundationDbMessagesConfiguration(@PositiveOrZero @Max(FoundationDbMessageStore.MAX_EPOCHS - 1) int activeEpoch) { +} diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStore.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStore.java index 0bb1c3c02..f2ceff7c9 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStore.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/foundationdb/FoundationDbMessageStore.java @@ -17,12 +17,13 @@ import java.util.List; import java.util.Map; import java.util.Optional; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.Executor; import java.util.function.Function; import java.util.stream.Collectors; +import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration; import org.whispersystems.textsecuregcm.entities.MessageProtos; import org.whispersystems.textsecuregcm.identity.AciServiceIdentifier; import org.whispersystems.textsecuregcm.storage.Device; +import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager; import org.whispersystems.textsecuregcm.storage.MessageStream; import org.whispersystems.textsecuregcm.util.Conversions; import org.whispersystems.textsecuregcm.util.Util; @@ -40,8 +41,9 @@ import org.whispersystems.textsecuregcm.util.Util; /// * {versionstamp_2} => envelope_2 public class FoundationDbMessageStore { - private final Database[] databases; + private final Database[][] databasesByEpoch; private final VersionstampUUIDCipher versionstampUUIDCipher; + private final DynamicConfigurationManager dynamicConfigurationManager; private final Clock clock; private static final Subspace MESSAGES_SUBSPACE = new Subspace(Tuple.from("M")); @@ -53,6 +55,11 @@ public class FoundationDbMessageStore { /// suggest a limit of 1MB to avoid performance issues, although the hard limit is 10MB private static final long MAX_MESSAGE_CHUNK_SIZE = DataSize.megabytes(1).toBytes(); + // We pack the current configuration epoch and shard ID into a single byte of "user data" in each message + // versionstamp. We use two bits for the epoch and six for the shard ID. + public static final int MAX_EPOCHS = 4; + public static final int MAX_SHARDS = 64; + /// Result of inserting a message for a particular device /// /// @param versionstamp the versionstamp of the transaction in which this device's message was inserted, empty @@ -61,12 +68,19 @@ public class FoundationDbMessageStore { public record InsertResult(Optional versionstamp, boolean present) { } - public FoundationDbMessageStore(final Database[] databases, + public FoundationDbMessageStore(final Map> databasesByEpoch, final VersionstampUUIDCipher versionstampUUIDCipher, + final DynamicConfigurationManager dynamicConfigurationManager, final Clock clock) { - this.databases = databases; + final Database[][] databasesByEpochArray = new Database[MAX_EPOCHS][]; + + databasesByEpoch.forEach((epoch, databases) -> + databasesByEpochArray[epoch] = databases.toArray(Database[]::new)); + + this.databasesByEpoch = databasesByEpochArray; this.versionstampUUIDCipher = versionstampUUIDCipher; + this.dynamicConfigurationManager = dynamicConfigurationManager; this.clock = clock; } @@ -115,9 +129,12 @@ public class FoundationDbMessageStore { throw new IllegalArgumentException("Messages must not have pre-set server GUIDs"); } + final int activeEpoch = + dynamicConfigurationManager.getConfiguration().getFoundationDbMessagesConfiguration().activeEpoch(); + final Map>>> messagesByShardId = messagesByServiceIdentifier.entrySet().stream() - .collect(Collectors.groupingBy(entry -> hashAciToShardNumber(entry.getKey()))); + .collect(Collectors.groupingBy(entry -> hashAciToShardNumber(entry.getKey(), activeEpoch))); final List>>> chunkFutures = new ArrayList<>(); @@ -133,7 +150,7 @@ public class FoundationDbMessageStore { .sum(); if (estimatedTransactionSize > MAX_MESSAGE_CHUNK_SIZE) { - chunkFutures.add(insertChunk(shardId, messagesForShard.subList(start, current))); + chunkFutures.add(insertChunk(shardId, activeEpoch, messagesForShard.subList(start, current))); start = current; estimatedTransactionSize = 0; @@ -143,7 +160,7 @@ public class FoundationDbMessageStore { } assert start < messagesForShard.size(); - chunkFutures.add(insertChunk(shardId, messagesForShard.subList(start, messagesForShard.size()))); + chunkFutures.add(insertChunk(shardId, activeEpoch, messagesForShard.subList(start, messagesForShard.size()))); }); return CompletableFuture.allOf(chunkFutures.toArray(CompletableFuture[]::new)) @@ -157,6 +174,7 @@ public class FoundationDbMessageStore { private CompletableFuture>> insertChunk( final int shardId, + final int epoch, final List>> messagesByAccountIdentifier) { final Map>> insertFuturesByAci = new HashMap<>(); @@ -168,9 +186,9 @@ public class FoundationDbMessageStore { .map(MessageProtos.Envelope::getEphemeral) .orElseThrow(() -> new IllegalStateException("One or more bundles is empty")); - return databases[shardId].runAsync(transaction -> { + return getDatabases(epoch)[shardId].runAsync(transaction -> { messagesByAccountIdentifier.forEach(entry -> - insertFuturesByAci.put(entry.getKey(), insert(entry.getKey(), entry.getValue(), transaction))); + insertFuturesByAci.put(entry.getKey(), insert(entry.getKey(), entry.getValue(), epoch, shardId, transaction))); return CompletableFuture.allOf(insertFuturesByAci.values().toArray(CompletableFuture[]::new)) .thenApply(_ -> { @@ -181,7 +199,8 @@ public class FoundationDbMessageStore { .anyMatch(isPresent -> isPresent); if (anyClientPresent || !ephemeral) { return transaction.getVersionstamp() - .thenApply(versionstampBytes -> Optional.of(Versionstamp.complete(versionstampBytes, shardId))); + .thenApply(versionstampBytes -> Optional.of(Versionstamp.complete(versionstampBytes, + packUserData(epoch, shardId)))); } return CompletableFuture.completedFuture(Optional.empty()); }); @@ -217,6 +236,8 @@ public class FoundationDbMessageStore { /// @return a future that yields the presence state of each destination device private CompletableFuture> insert(final AciServiceIdentifier aci, final Map messagesByDeviceId, + final int epoch, + final int shardId, final Transaction transaction) { final Map> messageInsertFuturesByDeviceId = messagesByDeviceId.entrySet() @@ -232,7 +253,7 @@ public class FoundationDbMessageStore { if (isPresent || !message.getEphemeral()) { transaction.mutate(MutationType.SET_VERSIONSTAMPED_KEY, getDeviceQueueSubspace(aci, deviceId) - .packWithVersionstamp(Tuple.from(Versionstamp.incomplete(hashAciToShardNumber(aci)))), message.toByteArray()); + .packWithVersionstamp(Tuple.from(Versionstamp.incomplete(packUserData(epoch, shardId)))), message.toByteArray()); } return isPresent; @@ -251,7 +272,7 @@ public class FoundationDbMessageStore { if (anyClientPresent) { transaction.mutate(MutationType.SET_VERSIONSTAMPED_VALUE, getMessagesAvailableWatchKey(aci), - Tuple.from(Versionstamp.incomplete(hashAciToShardNumber(aci))).packWithVersionstamp()); + Tuple.from(Versionstamp.incomplete(packUserData(epoch, shardId))).packWithVersionstamp()); } return presenceByDeviceId; @@ -264,11 +285,24 @@ public class FoundationDbMessageStore { } @VisibleForTesting - MessageStream getMessages(final AciServiceIdentifier aci, final Device destinationDevice, - final int maxMessagesPerScan, final int maxUnacknowledgedMessages, final Runnable doAfterCleanup) { + MessageStream getMessages(final AciServiceIdentifier aci, + final Device destinationDevice, + final int maxMessagesPerScan, + final int maxUnacknowledgedMessages, + final Runnable doAfterCleanup) { + + // For each configured database epoch, which database held (or holds) the messages for this ACI/device pair? + final Database[] databasesForQueueByEpoch = new Database[databasesByEpoch.length]; + + for (int epoch = 0; epoch < databasesByEpoch.length; epoch++) { + databasesForQueueByEpoch[epoch] = databasesByEpoch[epoch] != null + ? getDatabases(epoch)[hashAciToShardNumber(aci, epoch)] + : null; + } + return new FoundationDbMessageStream(getDeviceQueueSubspace(aci, destinationDevice.getId()), getMessagesAvailableWatchKey(aci), - getShardForAci(aci), + databasesForQueueByEpoch, new MessageGuidCodec(aci.uuid(), destinationDevice.getId(), versionstampUUIDCipher), maxMessagesPerScan, maxUnacknowledgedMessages, @@ -280,14 +314,46 @@ public class FoundationDbMessageStore { } @VisibleForTesting - Database getShardForAci(final AciServiceIdentifier aci) { - return databases[hashAciToShardNumber(aci)]; + Database getShardForAci(final AciServiceIdentifier aci, final int epoch) { + return getDatabases(epoch)[hashAciToShardNumber(aci, epoch)]; + } + + private Database[] getDatabases(final int epoch) { + if (databasesByEpoch[epoch] == null) { + throw new IllegalStateException("Epoch (%d) not in static configuration".formatted(epoch)); + } + + return databasesByEpoch[epoch]; } @VisibleForTesting - int hashAciToShardNumber(final AciServiceIdentifier aci) { + int hashAciToShardNumber(final AciServiceIdentifier aci, final int epoch) { // We use a consistent hash here to reduce the number of key remappings if we increase the number of shards - return Hashing.consistentHash(aci.uuid().getLeastSignificantBits(), databases.length); + return Hashing.consistentHash(aci.uuid().getLeastSignificantBits(), getDatabases(epoch).length); + } + + @VisibleForTesting + static int packUserData(final int epoch, final int shardId) { + if (epoch < 0 || epoch >= MAX_EPOCHS) { + throw new IllegalArgumentException("Epoch (%d) outside of allowable range (0 to %d, exclusive)".formatted( + epoch, MAX_EPOCHS)); + } + + if (shardId < 0 || shardId >= MAX_SHARDS) { + throw new IllegalArgumentException("Shard ID (%d) outside of allowable range (0 to %d, exclusive)".formatted( + epoch, MAX_SHARDS)); + } + + return epoch << 6 | shardId; + } + + static int getConfigurationEpoch(final Versionstamp versionstamp) { + return versionstamp.getUserVersion() >> 6 & 0x03; + } + + @VisibleForTesting + static int getShardId(final Versionstamp versionstamp) { + return versionstamp.getUserVersion() & 0x3f; } @VisibleForTesting 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 b03a7d632..2a713a001 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 @@ -13,12 +13,18 @@ import com.apple.foundationdb.tuple.Versionstamp; import com.google.common.annotations.VisibleForTesting; import io.micrometer.core.instrument.Counter; import io.micrometer.core.instrument.Metrics; +import java.util.Arrays; +import java.util.Comparator; +import java.util.IdentityHashMap; import java.util.List; +import java.util.Map; +import java.util.Objects; import java.util.Optional; import java.util.UUID; import java.util.concurrent.CompletableFuture; import java.util.concurrent.Flow; import java.util.function.Consumer; +import java.util.stream.Collectors; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.whispersystems.textsecuregcm.storage.MessageStream; @@ -27,43 +33,62 @@ import org.whispersystems.textsecuregcm.util.Pair; import reactor.adapter.JdkFlowAdapter; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.util.function.Tuple2; +import reactor.util.function.Tuples; + /// A [MessageStream] implementation that fetches messages from FoundationDB public class FoundationDbMessageStream implements MessageStream { private final Subspace deviceQueueSubspace; private final byte[] messagesAvailableWatchKey; - private final Database database; + private final Database[] databasesByEpoch; private final MessageGuidCodec messageGuidCodec; /// The maximum number of messages we will fetch per range query operation to avoid excessive memory consumption private final int maxMessagesPerScan; private final Flow.Publisher messageStreamPublisher; private final Runnable doAfterCleanup; - private final AcknowledgedMessageBuffer acknowledgedMessageBuffer; + + private final Map acknowledgedMessageBuffersByDatabase; + private final Counter staleEphemeralMessagesCounter = Metrics.counter( name(FoundationDbMessageStream.class, "staleEphemeralMessages")); static final int DEFAULT_MAX_MESSAGES_PER_SCAN = 1024; @VisibleForTesting static final int DEFAULT_MAX_UNACKNOWLEDGED_MESSAGES = 16_384; + private static final Comparator STREAM_ENTRY_TIMESTAMP_COMPARATOR = + Comparator.comparingLong(streamEntry -> streamEntry.partialEnvelope().getServerTimestamp()); + private static final Logger LOGGER = LoggerFactory.getLogger(FoundationDbMessageStream.class); FoundationDbMessageStream(final Subspace deviceQueueSubspace, final byte[] messagesAvailableWatchKey, - final Database database, + final Database[] databasesByEpoch, final MessageGuidCodec messageGuidCodec, final int maxMessagesPerScan, final int maxUnacknowledgedMessages, final Runnable doAfterCleanup) { this.deviceQueueSubspace = deviceQueueSubspace; this.messagesAvailableWatchKey = messagesAvailableWatchKey; - this.database = database; + this.databasesByEpoch = databasesByEpoch; this.messageGuidCodec = messageGuidCodec; this.maxMessagesPerScan = maxMessagesPerScan; this.messageStreamPublisher = JdkFlowAdapter.publisherToFlowPublisher(createMessagePublisher()); this.doAfterCleanup = doAfterCleanup; - this.acknowledgedMessageBuffer = new AcknowledgedMessageBuffer(maxUnacknowledgedMessages); + // Not all epochs may be in use (this is true most of the time) and if we DO have multiple epochs in play, it's + // possible/likely that a given queue will be on the same shard in multiple epochs. We only want acknowledgement + // buffer per distinct shard, so find the distinct shards for this queue and create a buffer for each. + this.acknowledgedMessageBuffersByDatabase = Arrays.stream(databasesByEpoch) + .filter(Objects::nonNull) + .distinct() + .collect(Collectors.toMap(database -> database, + _ -> new AcknowledgedMessageBuffer(maxUnacknowledgedMessages), + (_, _) -> { + throw new AssertionError("Duplicate database in distinct stream"); + }, + IdentityHashMap::new)); } @Override @@ -81,7 +106,7 @@ public class FoundationDbMessageStream implements MessageStream { .handle((messageStreamEntry, sink) -> { if (messageStreamEntry instanceof final FoundationDbMessageStreamEntry.Message message) { try { - acknowledgedMessageBuffer.addUnacknowledgedMessage(message.versionstamp()); + getAcknowledgedMessageBuffer(message.versionstamp()).addUnacknowledgedMessage(message.versionstamp()); } catch (final TooManyUnacknowledgedMessagesException e) { sink.error(e); return; @@ -107,34 +132,75 @@ public class FoundationDbMessageStream implements MessageStream { /// on [#messagesAvailableWatchKey] which is updated when a new message is available. /// See [FoundationDbMessageStore] for more details on the message insert process. private Flux createFoundationDbMessagePublisher() { - return Mono.fromFuture(this::getEndOfQueueKeyExclusive) - .flatMapMany(maybeEndOfQueueKeyExclusive -> { - final Flux finitePublisher = maybeEndOfQueueKeyExclusive - .map(endOfQueueKeyExclusive -> FoundationDbMessagePublisher.createFinitePublisher( - KeySelector.firstGreaterOrEqual(deviceQueueSubspace.range().begin), - endOfQueueKeyExclusive, database, maxMessagesPerScan, this::clearAcknowledgedMessages).getMessages()) - .orElseGet(Flux::empty) - .handle((fdbMessageStreamEntry, sink) -> { - // Ephemeral messages from the finite stream are considered stale and automatically discarded - if (fdbMessageStreamEntry.partialEnvelope().getEphemeral()) { - acknowledgedMessageBuffer.acknowledgeStaleEphemeralMessage(fdbMessageStreamEntry.versionstamp()); - staleEphemeralMessagesCounter.increment(); - return; - } - sink.next(fdbMessageStreamEntry); - }); - final KeySelector infinitePublisherBeginKey = maybeEndOfQueueKeyExclusive.orElseGet( - () -> KeySelector.firstGreaterOrEqual(deviceQueueSubspace.range().begin)); - final Flux infinitePublisher = FoundationDbMessagePublisher.createInfinitePublisher( - infinitePublisherBeginKey, KeySelector.firstGreaterThan(deviceQueueSubspace.range().end), - database, maxMessagesPerScan, messagesAvailableWatchKey, this::clearAcknowledgedMessages).getMessages(); - return Flux.concat( - finitePublisher, - Mono.just(new FoundationDbMessageStreamEntry.QueueEmpty()), - infinitePublisher - ); - }); + // This may seem like an odd construction since it looks like we could also just do `Flux#fromArray`, but + // `Flux#fromArray` cannot handle `null` elements + final List databases = Arrays.stream(databasesByEpoch).filter(Objects::nonNull).distinct().toList(); + return Flux.fromIterable(databases) + .flatMap(database -> Mono.fromFuture(getEndOfQueueKeyExclusive(database)) + .map(maybeEndOfQueueKeyExclusive -> Tuples.of(database, maybeEndOfQueueKeyExclusive))) + .collectMap(Tuple2::getT1, Tuple2::getT2) + .flatMapMany(endOfQueueKeysByDatabase -> { + @SuppressWarnings("unchecked") final Flux[] finitePublishers = + endOfQueueKeysByDatabase.entrySet().stream() + .map(entry -> { + final Database database = entry.getKey(); + final Optional maybeEndOfQueueKeyExclusive = entry.getValue(); + + return maybeEndOfQueueKeyExclusive + .map(endOfQueueKeyExclusive -> FoundationDbMessagePublisher.createFinitePublisher( + KeySelector.firstGreaterOrEqual(deviceQueueSubspace.range().begin), + endOfQueueKeyExclusive, + database, + maxMessagesPerScan, + () -> this.clearAcknowledgedMessages(database)) + .getMessages()) + .orElseGet(Flux::empty) + .handle((fdbMessageStreamEntry, sink) -> { + // Ephemeral messages from the finite stream are considered stale and automatically discarded + if (fdbMessageStreamEntry.partialEnvelope().getEphemeral()) { + acknowledgedMessageBuffersByDatabase.get(database).acknowledgeStaleEphemeralMessage(fdbMessageStreamEntry.versionstamp()); + staleEphemeralMessagesCounter.increment(); + return; + } + sink.next(fdbMessageStreamEntry); + }); + }) + .toArray(Flux[]::new); + + @SuppressWarnings("unchecked") final Flux[] infinitePublishers = + endOfQueueKeysByDatabase.entrySet().stream() + .map(entry -> { + final Database database = entry.getKey(); + final Optional maybeEndOfQueueKeyExclusive = entry.getValue(); + + final KeySelector infinitePublisherBeginKey = maybeEndOfQueueKeyExclusive + .orElseGet(() -> KeySelector.firstGreaterOrEqual(deviceQueueSubspace.range().begin)); + + return FoundationDbMessagePublisher.createInfinitePublisher( + infinitePublisherBeginKey, + KeySelector.firstGreaterThan(deviceQueueSubspace.range().end), + database, + maxMessagesPerScan, + messagesAvailableWatchKey, + () -> clearAcknowledgedMessages(database)).getMessages(); + }) + .toArray(Flux[]::new); + + return Flux.concat( + Flux.mergeComparing(maxMessagesPerScan, STREAM_ENTRY_TIMESTAMP_COMPARATOR, finitePublishers), + Mono.just(new FoundationDbMessageStreamEntry.QueueEmpty()), + + // Note that we use `mergePriority` instead of `mergeComparing` for the "live"/non-terminating publishers + // because `mergePriority` sorts messages _as they arrive._ If we used `mergeComparing` for the live + // streams and one of the streams had no new messages (which will be true most of the time), then the + // merged publisher would never emit any signals because it'd be waiting to have something to compare + // against. This does mean that we risk some slightly out-of-order messages in the exceedingly rare cases + // where messages arrive at different servers while somebody is connected and a migration is in progress, + // but that should be (again) exceedingly rare and also minimally-disruptive (i.e. it would self-correct + // so quickly that end users would likely never even notice). + Flux.mergePriority(maxMessagesPerScan, STREAM_ENTRY_TIMESTAMP_COMPARATOR, infinitePublishers)); + }); } /// Gets a [KeySelector] for the first key greater than the current greatest key in the device queue. This allows us @@ -142,7 +208,7 @@ public class FoundationDbMessageStream implements MessageStream { /// subsequent scan. /// /// @return a [KeySelector] for the first key greater than the current greatest key in the device queue. - private CompletableFuture> getEndOfQueueKeyExclusive() { + private CompletableFuture> getEndOfQueueKeyExclusive(final Database database) { return database.runAsync( transaction -> transaction.getRange(deviceQueueSubspace.range(), 1, true, StreamingMode.EXACT).asList() .thenApply(items -> { @@ -156,7 +222,8 @@ public class FoundationDbMessageStream implements MessageStream { @Override public CompletableFuture acknowledgeMessage(final UUID messageGuid, final long serverTimestamp) { - acknowledgedMessageBuffer.acknowledgeMessage(messageGuidCodec.decodeMessageGuid(messageGuid)); + final Versionstamp versionstamp = messageGuidCodec.decodeMessageGuid(messageGuid); + getAcknowledgedMessageBuffer(versionstamp).acknowledgeMessage(versionstamp); return CompletableFuture.completedFuture(null); } @@ -174,20 +241,39 @@ public class FoundationDbMessageStream implements MessageStream { /// Clear all outstanding acknowledged messages. Called when the stream ends private CompletableFuture flushAllAcknowledgedMessages() { - final Consumer clearAllAcknowlegedMessagedConsumer = clearAcknowledgedMessages(); - return database.runAsync(transaction -> { - clearAllAcknowlegedMessagedConsumer.accept(transaction); - return CompletableFuture.completedFuture((Void) null); + return CompletableFuture.allOf(Arrays.stream(databasesByEpoch) + .filter(Objects::nonNull) + .distinct() + .map(database -> { + final Consumer clearAllAcknowlegedMessagedConsumer = clearAcknowledgedMessages(database); + + return database.runAsync(transaction -> { + clearAllAcknowlegedMessagedConsumer.accept(transaction); + return CompletableFuture.completedFuture((Void) null); + }) + .whenComplete((_, throwable) -> { + if (throwable != null) { + LOGGER.warn("Failed to clear acknowledged messages", throwable); + } + }); }) - .whenComplete((_, throwable) -> { - if (throwable != null) { - LOGGER.warn("Failed to clear acknowledged messages", throwable); - } - }); + .toArray(CompletableFuture[]::new)); } - private synchronized Consumer clearAcknowledgedMessages() { - final List> flushableRanges = acknowledgedMessageBuffer.takeFlushableRanges(); + private AcknowledgedMessageBuffer getAcknowledgedMessageBuffer(final Versionstamp versionstamp) { + final int epoch = FoundationDbMessageStore.getConfigurationEpoch(versionstamp); + final Database database = databasesByEpoch[epoch]; + + if (database == null) { + throw new IllegalStateException("Read message for unrecognized epoch"); + } + + return acknowledgedMessageBuffersByDatabase.get(database); + } + + private synchronized Consumer clearAcknowledgedMessages(final Database database) { + final List> flushableRanges = + acknowledgedMessageBuffersByDatabase.get(database).takeFlushableRanges(); return transaction -> flushableRanges.forEach(range -> clearRange(transaction, range.first(), range.second())); } diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/configuration/FoundationDbClusterConfigurationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/configuration/FoundationDbClusterConfigurationTest.java new file mode 100644 index 000000000..999ac0758 --- /dev/null +++ b/service/src/test/java/org/whispersystems/textsecuregcm/configuration/FoundationDbClusterConfigurationTest.java @@ -0,0 +1,40 @@ +/* + * Copyright 2026 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.whispersystems.textsecuregcm.configuration; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import java.util.List; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +class FoundationDbClusterConfigurationTest { + + @ParameterizedTest + @MethodSource + void isSingleClusterFileSourceSpecified(final FoundationDbClusterConfiguration clusterConfiguration, + final boolean expectSingleClusterFileSourceSpecified) { + + assertEquals(expectSingleClusterFileSourceSpecified, clusterConfiguration.isSingleClusterFileSourceSpecified()); + } + + private static List isSingleClusterFileSourceSpecified() { + return List.of( + Arguments.argumentSet("Cluster file URL only", + new FoundationDbClusterConfiguration("test-url", null), true), + + Arguments.argumentSet("Cluster file contents only", + new FoundationDbClusterConfiguration(null, "test-contents"), true), + + Arguments.argumentSet("Both", + new FoundationDbClusterConfiguration("test-url", "test-contents"), false), + + Arguments.argumentSet("Neither", + new FoundationDbClusterConfiguration(null, null), false) + ); + } +} diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/configuration/FoundationDbMessagesConfigurationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/configuration/FoundationDbMessagesConfigurationTest.java new file mode 100644 index 000000000..dd7412591 --- /dev/null +++ b/service/src/test/java/org/whispersystems/textsecuregcm/configuration/FoundationDbMessagesConfigurationTest.java @@ -0,0 +1,42 @@ +/* + * Copyright 2026 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.whispersystems.textsecuregcm.configuration; + +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.*; + +class FoundationDbMessagesConfigurationTest { + + @Test + void isEveryEpochClusterConfigured() { + assertTrue(new FoundationDbMessagesConfiguration( + Map.of("messages-0", new FoundationDbClusterConfiguration("test-url", null)), + Map.of(0, List.of("messages-0")) + ).isEveryEpochClusterConfigured()); + + assertFalse(new FoundationDbMessagesConfiguration( + Map.of("messages-0", new FoundationDbClusterConfiguration("test-url", null)), + Map.of(0, List.of("messages-0", "unconfigured-cluster")) + ).isEveryEpochClusterConfigured()); + } + + @Test + void isEveryEpochFreeOfDuplicates() { + assertTrue(new FoundationDbMessagesConfiguration( + Map.of("messages-0", new FoundationDbClusterConfiguration("test-url", null)), + Map.of(0, List.of("messages-0")) + ).isEveryEpochFreeOfDuplicates()); + + assertFalse(new FoundationDbMessagesConfiguration( + Map.of("messages-0", new FoundationDbClusterConfiguration("test-url", null)), + Map.of(0, List.of("messages-0", "messages-0")) + ).isEveryEpochFreeOfDuplicates()); + } +} 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 75ffc31e6..9c7002e1d 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 @@ -8,7 +8,10 @@ import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import com.apple.foundationdb.Database; import com.apple.foundationdb.KeyValue; import com.apple.foundationdb.async.AsyncUtil; import com.apple.foundationdb.tuple.Tuple; @@ -23,6 +26,7 @@ import java.time.Duration; import java.time.Instant; import java.time.ZoneId; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -36,7 +40,9 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; import java.util.function.Function; +import java.util.function.Supplier; import java.util.stream.Collectors; import java.util.stream.IntStream; import java.util.stream.Stream; @@ -49,9 +55,14 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; import org.junit.jupiter.params.provider.ValueSource; +import org.junitpioneer.jupiter.cartesian.CartesianTest; +import org.junitpioneer.jupiter.params.IntRangeSource; +import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration; +import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicFoundationDbMessagesConfiguration; import org.whispersystems.textsecuregcm.entities.MessageProtos; import org.whispersystems.textsecuregcm.identity.AciServiceIdentifier; import org.whispersystems.textsecuregcm.storage.Device; +import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager; import org.whispersystems.textsecuregcm.storage.FoundationDbClusterExtension; import org.whispersystems.textsecuregcm.storage.MessageStream; import org.whispersystems.textsecuregcm.storage.MessageStreamEntry; @@ -69,10 +80,14 @@ class FoundationDbMessageStoreTest { static FoundationDbClusterExtension FOUNDATION_DB_EXTENSION = new FoundationDbClusterExtension(2); private VersionstampUUIDCipher versionstampUUIDCipher; + private DynamicFoundationDbMessagesConfiguration foundationDbMessagesConfiguration; private FoundationDbMessageStore foundationDbMessageStore; private static final Clock CLOCK = Clock.fixed(Instant.ofEpochSecond(500), ZoneId.of("UTC")); + private static final int DEFAULT_EPOCH = 0; + private static final int FUTURE_EPOCH = 2; + @BeforeEach void setup() { final byte[] versionstampCipherKey = new byte[16]; @@ -80,9 +95,29 @@ class FoundationDbMessageStoreTest { versionstampUUIDCipher = new VersionstampUUIDCipher(0, versionstampCipherKey); + foundationDbMessagesConfiguration = mock(DynamicFoundationDbMessagesConfiguration.class); + when(foundationDbMessagesConfiguration.activeEpoch()).thenReturn(DEFAULT_EPOCH); + + final DynamicConfiguration dynamicConfiguration = mock(DynamicConfiguration.class); + when(dynamicConfiguration.getFoundationDbMessagesConfiguration()).thenReturn(foundationDbMessagesConfiguration); + + @SuppressWarnings("unchecked") final DynamicConfigurationManager dynamicConfigurationManager = + mock(DynamicConfigurationManager.class); + + when(dynamicConfigurationManager.getConfiguration()).thenReturn(dynamicConfiguration); + + final List databases = Arrays.asList(FOUNDATION_DB_EXTENSION.getDatabases()); + foundationDbMessageStore = new FoundationDbMessageStore( - FOUNDATION_DB_EXTENSION.getDatabases(), + // Simulate a topology change by reversing the order of the (two) databases in the group in a second epoch. This + // construction ensures that queues will land in different databases in different epochs, and defining + // non-contiguous epochs forces us to deal with "holes" and null values in the epoch list. + Map.of( + DEFAULT_EPOCH, databases, + FUTURE_EPOCH, databases.reversed() + ), versionstampUUIDCipher, + dynamicConfigurationManager, CLOCK); } @@ -157,6 +192,53 @@ class FoundationDbMessageStoreTest { ); } + @Test + void insertEpochChange() throws InvalidProtocolBufferException { + final AciServiceIdentifier aci = new AciServiceIdentifier(UUID.randomUUID()); + final byte deviceId = Device.PRIMARY_ID; + + final MessageProtos.Envelope defaultEpochMessage = generateRandomMessage(false); + final MessageProtos.Envelope futureEpochMessage = generateRandomMessage(false); + + when(foundationDbMessagesConfiguration.activeEpoch()).thenReturn(DEFAULT_EPOCH); + + final Map defaultEpochInsertResult = + foundationDbMessageStore.insert(aci, Map.of(deviceId, defaultEpochMessage)).join(); + + when(foundationDbMessagesConfiguration.activeEpoch()).thenReturn(FUTURE_EPOCH); + + final Map futureEpochInsertResult = + foundationDbMessageStore.insert(aci, Map.of(deviceId, futureEpochMessage)).join(); + + for (int epoch : new int[] { DEFAULT_EPOCH, FUTURE_EPOCH }) { + final List itemsInDeviceQueue = getItemsInDeviceQueue(aci, deviceId, epoch); + assertEquals(1, itemsInDeviceQueue.size()); + + final Versionstamp expectedVersionstamp = switch (epoch) { + case DEFAULT_EPOCH -> defaultEpochInsertResult.get(deviceId).versionstamp().orElseThrow(); + case FUTURE_EPOCH -> futureEpochInsertResult.get(deviceId).versionstamp().orElseThrow(); + default -> throw new AssertionError("Unexpected epoch"); + }; + + final Versionstamp retrievedVersionstamp = + FoundationDbMessageStore.getVersionstamp(itemsInDeviceQueue.getFirst().getKey()); + + assertEquals(expectedVersionstamp, retrievedVersionstamp); + assertEquals(epoch, FoundationDbMessageStore.getConfigurationEpoch(retrievedVersionstamp)); + + final MessageProtos.Envelope expectedEnvelope = switch (epoch) { + case DEFAULT_EPOCH -> defaultEpochMessage; + case FUTURE_EPOCH -> futureEpochMessage; + default -> throw new AssertionError("Unexpected epoch"); + }; + + final MessageProtos.Envelope retrievedEnvelope = + MessageProtos.Envelope.parseFrom(itemsInDeviceQueue.getFirst().getValue()); + + assertEquals(expectedEnvelope, retrievedEnvelope); + } + } + @Test void versionstampCorrectlyUpdatedOnMultipleInserts() { final AciServiceIdentifier aci = new AciServiceIdentifier(UUID.randomUUID()); @@ -286,7 +368,7 @@ class FoundationDbMessageStoreTest { // assert that each shard has the expected number of committed transactions. final Map> returnedVersionstampsByShard = new HashMap<>(); result.forEach((aci, deviceResults) -> { - final int shardNum = foundationDbMessageStore.hashAciToShardNumber(aci); + final int shardNum = foundationDbMessageStore.hashAciToShardNumber(aci, DEFAULT_EPOCH); final Set versionstampSet = returnedVersionstampsByShard.computeIfAbsent(shardNum, _ -> new HashSet<>()); deviceResults.forEach((_, deviceResult) -> deviceResult.versionstamp().ifPresent(versionstampSet::add)); }); @@ -438,6 +520,69 @@ class FoundationDbMessageStoreTest { ); } + @Test + void getMessagesEpochChange() { + final AciServiceIdentifier aci = new AciServiceIdentifier(UUID.randomUUID()); + final byte deviceId = Device.PRIMARY_ID; + + final Device device = new Device(); + device.setId(deviceId); + + final int messagesPerBatch = 8; + + final AtomicLong serialTimestamp = new AtomicLong(); + + final List existingDefaultEpochMessages = + generateAndInsertMessages(aci, deviceId, DEFAULT_EPOCH, messagesPerBatch, serialTimestamp::getAndIncrement); + + final List existingFutureEpochMessages = + generateAndInsertMessages(aci, deviceId, FUTURE_EPOCH, messagesPerBatch, serialTimestamp::getAndIncrement); + + final List liveDefaultEpochMessages = new ArrayList<>(); + final List liveFutureEpochMessages = new ArrayList<>(); + + final MessageStream messageStream = foundationDbMessageStore.getMessages(aci, device); + final List retrievedEntries = new ArrayList<>(); + final CountDownLatch queueEmptyLatch = new CountDownLatch(1); + + Thread.ofVirtual().start(() -> { + try { + // Wait until queue is empty + assertTrue(queueEmptyLatch.await(1000, TimeUnit.MILLISECONDS)); + + // Then publish more messages + liveDefaultEpochMessages.addAll(generateAndInsertMessages(aci, deviceId, DEFAULT_EPOCH, messagesPerBatch, serialTimestamp::getAndIncrement)); + liveFutureEpochMessages.addAll(generateAndInsertMessages(aci, deviceId, FUTURE_EPOCH, messagesPerBatch, serialTimestamp::getAndIncrement)); + } catch (final InterruptedException e) { + fail(e); + } + }); + + writePresenceKey(aci, deviceId, 1, 5L, DEFAULT_EPOCH); + writePresenceKey(aci, deviceId, 1, 5L, FUTURE_EPOCH); + + StepVerifier.create(JdkFlowAdapter.flowPublisherToFlux(messageStream.getMessages())) + .recordWith(() -> retrievedEntries) + .expectNextCount(existingDefaultEpochMessages.size() + existingFutureEpochMessages.size()) + .expectNext(new MessageStreamEntry.QueueEmpty()) + .then(queueEmptyLatch::countDown) + .expectNextCount(2 * messagesPerBatch) + .verifyTimeout(Duration.ofSeconds(1)); + + final List retrievedMessages = retrievedEntries.stream() + .filter(messageStreamEntry -> messageStreamEntry instanceof MessageStreamEntry.Envelope) + .map(messageStreamEntry -> ((MessageStreamEntry.Envelope) messageStreamEntry).message()) + .toList(); + + assertEquals(4 * messagesPerBatch, retrievedMessages.size()); + assertEquals(existingDefaultEpochMessages, retrievedMessages.subList(0, existingDefaultEpochMessages.size())); + assertEquals(existingFutureEpochMessages, retrievedMessages.subList(existingDefaultEpochMessages.size(), existingDefaultEpochMessages.size() + existingFutureEpochMessages.size())); + + // Order is not strictly defined for "competing" live messages + assertTrue(retrievedMessages.containsAll(liveDefaultEpochMessages)); + assertTrue(retrievedMessages.containsAll(liveFutureEpochMessages)); + } + @Test void getMessagesPublishMoreAfterQueueEmpty() { final AciServiceIdentifier aci = new AciServiceIdentifier(UUID.randomUUID()); @@ -616,6 +761,68 @@ class FoundationDbMessageStoreTest { ); } + @Test + void acknowledgeMessagesEpochChange() throws InterruptedException { + final AciServiceIdentifier aci = new AciServiceIdentifier(UUID.randomUUID()); + final byte deviceId = Device.PRIMARY_ID; + + final Device device = new Device(); + device.setId(deviceId); + + final int messagesPerBatch = 8; + + final AtomicLong serialTimestamp = new AtomicLong(); + + generateAndInsertMessages(aci, deviceId, DEFAULT_EPOCH, messagesPerBatch, serialTimestamp::getAndIncrement); + generateAndInsertMessages(aci, deviceId, FUTURE_EPOCH, messagesPerBatch, serialTimestamp::getAndIncrement); + + { + final CountDownLatch cleanupLatch = new CountDownLatch(1); + + final MessageStream messageStream = foundationDbMessageStore.getMessages(aci, + device, + FoundationDbMessageStream.DEFAULT_MAX_MESSAGES_PER_SCAN, + FoundationDbMessageStream.DEFAULT_MAX_UNACKNOWLEDGED_MESSAGES, + cleanupLatch::countDown); + + final List retrievedEntries = new ArrayList<>(); + + StepVerifier.create(JdkFlowAdapter.flowPublisherToFlux(messageStream.getMessages())) + .recordWith(() -> retrievedEntries) + .expectNextCount(2 * messagesPerBatch) + .expectNext(new MessageStreamEntry.QueueEmpty()) + .then(() -> retrievedEntries.stream() + .filter(messageStreamEntry -> messageStreamEntry instanceof MessageStreamEntry.Envelope) + .map(messageStreamEntry -> ((MessageStreamEntry.Envelope) messageStreamEntry).message()) + // Acknowledge messages with even-numbered timestamps; this will spread acknowledgements across both epochs + .filter(message -> message.getServerTimestamp() % 2 == 0) + .forEach(message -> messageStream.acknowledgeMessage(UUIDUtil.fromByteString(message.getServerGuid()), message.getServerTimestamp()).join())) + .verifyTimeout(Duration.ofSeconds(1)); + + cleanupLatch.await(); + } + + { + final MessageStream messageStream = foundationDbMessageStore.getMessages(aci, device); + final List retrievedEntries = new ArrayList<>(); + + StepVerifier.create(JdkFlowAdapter.flowPublisherToFlux(messageStream.getMessages())) + .recordWith(() -> retrievedEntries) + .expectNextCount(messagesPerBatch) + .expectNext(new MessageStreamEntry.QueueEmpty()) + .verifyTimeout(Duration.ofSeconds(1)); + + final List retrievedMessages = retrievedEntries.stream() + .filter(messageStreamEntry -> messageStreamEntry instanceof MessageStreamEntry.Envelope) + .map(messageStreamEntry -> ((MessageStreamEntry.Envelope) messageStreamEntry).message()) + .toList(); + + assertEquals(messagesPerBatch, retrievedMessages.size()); + assertTrue(retrievedMessages.stream().noneMatch(message -> message.getServerTimestamp() % 2 == 0), + "All messages with even-numbered timestamps should be acknowledged and removed"); + } + } + @Test void outstandingUnacknowledgedMessages() { final int numMessages = 5; @@ -624,13 +831,9 @@ class FoundationDbMessageStoreTest { final AciServiceIdentifier aci = new AciServiceIdentifier(UUID.randomUUID()); writePresenceKey(aci, Device.PRIMARY_ID, 1, 5L); - final List versionstamps = IntStream.range(0, numMessages) - .mapToObj( - _ -> foundationDbMessageStore.insert(aci, Map.of(Device.PRIMARY_ID, generateRandomMessage(false))).join() - .get(Device.PRIMARY_ID) - .versionstamp() - .orElseThrow()) - .toList(); + for (int i = 0; i < numMessages; i++) { + foundationDbMessageStore.insert(aci, Map.of(Device.PRIMARY_ID, generateRandomMessage(false))).join(); + } final Device device = new Device(); device.setId(Device.PRIMARY_ID); @@ -701,6 +904,17 @@ class FoundationDbMessageStoreTest { assertNotNull(getMessageByVersionstamp(aci, Device.PRIMARY_ID, deliveredUnacknowledgedVersionstamp.join())); } + @CartesianTest + void packUserVersion(@IntRangeSource(from = 0, to = FoundationDbMessageStore.MAX_EPOCHS) final int epoch, + @IntRangeSource(from = 0, to = FoundationDbMessageStore.MAX_SHARDS) final int shardId) { + + final Versionstamp versionstamp = + Versionstamp.complete(new byte[10], FoundationDbMessageStore.packUserData(epoch, shardId)); + + assertEquals(epoch, FoundationDbMessageStore.getConfigurationEpoch(versionstamp)); + assertEquals(shardId, FoundationDbMessageStore.getShardId(versionstamp)); + } + static MessageProtos.Envelope generateRandomMessage(final boolean ephemeral, final int contentSize) { return generateRandomMessage(ephemeral, contentSize, CLOCK.millis()); } @@ -718,10 +932,48 @@ class FoundationDbMessageStoreTest { .build(); } + private List generateAndInsertMessages(final AciServiceIdentifier aci, + final byte deviceId, + final int epoch, + final int messageCount, + final Supplier timestampSupplier) { + + final MessageGuidCodec messageGuidCodec = + new MessageGuidCodec(aci.uuid(), deviceId, versionstampUUIDCipher); + + when(foundationDbMessagesConfiguration.activeEpoch()).thenReturn(epoch); + + return IntStream.range(0, messageCount) + .mapToObj(_ -> { + final MessageProtos.Envelope message = + generateRandomMessage(false, 16, timestampSupplier.get()); + + final FoundationDbMessageStore.InsertResult insertResult = + foundationDbMessageStore.insert(aci, Map.of(deviceId, message)).join().get(deviceId); + + final Versionstamp versionstamp = insertResult.versionstamp().orElseThrow(); + final UUID messageGuid = messageGuidCodec.encodeMessageGuid(versionstamp); + + return message.toBuilder().setServerGuid(UUIDUtil.toByteString(messageGuid)).build(); + }) + .toList(); + } + @Nullable - private byte[] getMessageByVersionstamp(final AciServiceIdentifier aci, final byte deviceId, + private byte[] getMessageByVersionstamp(final AciServiceIdentifier aci, + final byte deviceId, final Versionstamp versionstamp) { - return foundationDbMessageStore.getShardForAci(aci).read(transaction -> { + + return getMessageByVersionstamp(aci, deviceId, versionstamp, DEFAULT_EPOCH); + } + + @Nullable + private byte[] getMessageByVersionstamp(final AciServiceIdentifier aci, + final byte deviceId, + final Versionstamp versionstamp, + final int epoch) { + + return foundationDbMessageStore.getShardForAci(aci, epoch).read(transaction -> { final byte[] key = FoundationDbMessageStore.getDeviceQueueSubspace(aci, deviceId) .pack(Tuple.from(versionstamp)); return transaction.get(key); @@ -729,16 +981,32 @@ class FoundationDbMessageStoreTest { } private Optional getMessagesAvailableWatch(final AciServiceIdentifier aci) { - return foundationDbMessageStore.getShardForAci(aci) + return getMessagesAvailableWatch(aci, DEFAULT_EPOCH); + } + + private Optional getMessagesAvailableWatch(final AciServiceIdentifier aci, final int epoch) { + return foundationDbMessageStore.getShardForAci(aci, epoch) .read(transaction -> transaction.get(FoundationDbMessageStore.getMessagesAvailableWatchKey(aci)) .thenApply(value -> value == null ? null : Tuple.fromBytes(value).getVersionstamp(0)) .thenApply(Optional::ofNullable)) .join(); } - private void writePresenceKey(final AciServiceIdentifier aci, final byte deviceId, final int serverId, + private void writePresenceKey(final AciServiceIdentifier aci, + final byte deviceId, + final int serverId, final long secondsBeforeCurrentTime) { - foundationDbMessageStore.getShardForAci(aci).run(transaction -> { + + writePresenceKey(aci, deviceId, serverId, secondsBeforeCurrentTime, DEFAULT_EPOCH); + } + + private void writePresenceKey(final AciServiceIdentifier aci, + final byte deviceId, + final int serverId, + final long secondsBeforeCurrentTime, + final int epoch) { + + foundationDbMessageStore.getShardForAci(aci, epoch).run(transaction -> { final byte[] presenceKey = foundationDbMessageStore.getPresenceKey(aci, deviceId); final long presenceUpdateEpochSeconds = getEpochSecondsBeforeClock(secondsBeforeCurrentTime); final long presenceValue = constructPresenceValue(serverId, presenceUpdateEpochSeconds); @@ -756,18 +1024,26 @@ class FoundationDbMessageStoreTest { } private AciServiceIdentifier generateRandomAciForShard(final int shardNumber) { + return generateRandomAciForShard(shardNumber, DEFAULT_EPOCH); + } + + private AciServiceIdentifier generateRandomAciForShard(final int shardNumber, final int epoch) { assert shardNumber < FOUNDATION_DB_EXTENSION.getDatabases().length; while (true) { final AciServiceIdentifier aci = new AciServiceIdentifier(UUID.randomUUID()); - if (foundationDbMessageStore.hashAciToShardNumber(aci) == shardNumber) { + if (foundationDbMessageStore.hashAciToShardNumber(aci, epoch) == shardNumber) { return aci; } } } private List getItemsInDeviceQueue(final AciServiceIdentifier aci, final byte deviceId) { - return foundationDbMessageStore.getShardForAci(aci).readAsync(transaction -> AsyncUtil.collect(transaction.getRange( - FoundationDbMessageStore.getDeviceQueueSubspace(aci, deviceId).range()))).join(); + return getItemsInDeviceQueue(aci, deviceId, DEFAULT_EPOCH); } + private List getItemsInDeviceQueue(final AciServiceIdentifier aci, final byte deviceId, final int epoch) { + return foundationDbMessageStore.getShardForAci(aci, epoch).readAsync(transaction -> + AsyncUtil.collect(transaction.getRange(FoundationDbMessageStore.getDeviceQueueSubspace(aci, deviceId).range()))) + .join(); + } }