Parse FoundationDB envelopes at load time and add GUIDs later

This commit is contained in:
Jon Chambers
2026-06-08 11:59:29 -04:00
committed by Jon Chambers
parent 3e4ad32fce
commit a4381a2617
3 changed files with 21 additions and 14 deletions
@@ -5,13 +5,16 @@ import com.apple.foundationdb.KeySelector;
import com.apple.foundationdb.StreamingMode;
import com.apple.foundationdb.Transaction;
import com.apple.foundationdb.tuple.Versionstamp;
import com.google.protobuf.InvalidProtocolBufferException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.whispersystems.textsecuregcm.entities.MessageProtos;
import org.whispersystems.textsecuregcm.util.Pair;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxSink;
import javax.annotation.Nullable;
import java.io.UncheckedIOException;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
@@ -287,7 +290,12 @@ class FoundationDbMessagePublisher {
final List<FoundationDbMessageStreamEntry.Message> messages = keyValues.stream()
.map(keyValue -> {
final Versionstamp versionstamp = FoundationDbMessageStore.getVersionstamp(keyValue.getKey());
return new FoundationDbMessageStreamEntry.Message(versionstamp, keyValue.getValue());
try {
return new FoundationDbMessageStreamEntry.Message(versionstamp, MessageProtos.Envelope.parseFrom(keyValue.getValue()));
} catch (final InvalidProtocolBufferException e) {
throw new UncheckedIOException(e);
}
})
.toList();
return new Pair<>(lastKeyRead, messages);
@@ -1,8 +1,6 @@
package org.whispersystems.textsecuregcm.storage.foundationdb;
import com.apple.foundationdb.tuple.Versionstamp;
import com.google.protobuf.InvalidProtocolBufferException;
import java.io.UncheckedIOException;
import org.whispersystems.textsecuregcm.entities.MessageProtos;
import org.whispersystems.textsecuregcm.storage.MessageStreamEntry;
import org.whispersystems.textsecuregcm.util.UUIDUtil;
@@ -10,21 +8,16 @@ import org.whispersystems.textsecuregcm.util.UUIDUtil;
sealed interface FoundationDbMessageStreamEntry permits FoundationDbMessageStreamEntry.Message,
FoundationDbMessageStreamEntry.QueueEmpty {
record Message(Versionstamp versionstamp, byte[] payload) implements FoundationDbMessageStreamEntry {}
record Message(Versionstamp versionstamp, MessageProtos.Envelope partialEnvelope) implements FoundationDbMessageStreamEntry {}
record QueueEmpty() implements FoundationDbMessageStreamEntry {}
default MessageStreamEntry toMessageStreamEntry(final MessageGuidCodec messageGuidCodec) {
return switch (this) {
case Message(final Versionstamp versionstamp, final byte[] payload): {
try {
yield new MessageStreamEntry.Envelope(MessageProtos.Envelope.parseFrom(payload)
.toBuilder()
.setServerGuidBinary(UUIDUtil.toByteString(messageGuidCodec.encodeMessageGuid(versionstamp)))
.build());
} catch (final InvalidProtocolBufferException e) {
throw new UncheckedIOException(e);
}
case Message(final Versionstamp versionstamp, final MessageProtos.Envelope partialEnvelope): {
yield new MessageStreamEntry.Envelope(partialEnvelope.toBuilder()
.setServerGuidBinary(UUIDUtil.toByteString(messageGuidCodec.encodeMessageGuid(versionstamp)))
.build());
}
case QueueEmpty():
yield new MessageStreamEntry.QueueEmpty();
@@ -16,6 +16,7 @@ import com.apple.foundationdb.Range;
import com.apple.foundationdb.StreamingMode;
import com.apple.foundationdb.Transaction;
import com.apple.foundationdb.async.AsyncIterable;
import java.io.UncheckedIOException;
import java.nio.ByteBuffer;
import java.security.SecureRandom;
import java.time.Duration;
@@ -325,7 +326,12 @@ class FoundationDbMessagePublisherTest {
}
private FoundationDbMessageStreamEntry.Message getExpectedMessageStreamEntry(final KeyValue keyValue) {
return new FoundationDbMessageStreamEntry.Message(FoundationDbMessageStore.getVersionstamp(keyValue.getKey()), keyValue.getValue());
try {
return new FoundationDbMessageStreamEntry.Message(FoundationDbMessageStore.getVersionstamp(keyValue.getKey()),
MessageProtos.Envelope.parseFrom(keyValue.getValue()));
} catch (final InvalidProtocolBufferException e) {
throw new UncheckedIOException(e);
}
}
private static KeyValue mockKeyValue(final byte key, final MessageProtos.Envelope message) {