Use the LuceneSerializer to encode stored fields and field infos, too. (#3669)

MMcM · web-flow · commit 1584f34fdc57 · 2025-10-16T16:09:45.000-07:00
Since existing fields will be straight protobuf without the serializer's
prefix byte, recognize that, too.
diff --git a/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneLogMessageKeys.java b/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneLogMessageKeys.java
@@ -47,6 +47,7 @@ public enum LuceneLogMessageKeys {
     ENCODED_DATA_SIZE,
     ENCRYPTED_EVENTUALLY,
     ENCRYPTION_SUPPOSED,
+    FIELD_PROTOBUF_ENCODED,
     FILE_ACTUAL_TOTAL_SIZE,
     FILE_COUNT,
     FILE_ID,
diff --git a/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneRecordContextProperties.java b/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneRecordContextProperties.java
@@ -49,6 +49,11 @@ public final class LuceneRecordContextProperties {
 
     public static final RecordLayerPropertyKey<SerializationKeyManager> LUCENE_INDEX_KEY_MANAGER = new RecordLayerPropertyKey<>("com.apple.foundationdb.record.lucene.keyManager", null, SerializationKeyManager.class);
 
+    /**
+     * Whether {@code StoredField} and {@code FieldInfo} are also encoded, allowing compression and encryption.
+     */
+    public static final RecordLayerPropertyKey<Boolean> LUCENE_FIELD_PROTOBUF_PREFIX_ENABLED = RecordLayerPropertyKey.booleanPropertyKey("com.apple.foundationdb.record.lucene.fieldProtobufPrefixEnabled", false);
+
     /**
      * An {@link ExecutorService} to use for parallel execution in {@link LuceneRecordCursor}.
      */
diff --git a/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/codec/LuceneOptimizedStoredFieldsReader.java b/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/codec/LuceneOptimizedStoredFieldsReader.java
@@ -20,7 +20,6 @@
 
 package com.apple.foundationdb.record.lucene.codec;
 
-import com.apple.foundationdb.KeyValue;
 import com.apple.foundationdb.annotation.SpotBugsSuppressWarnings;
 import com.apple.foundationdb.record.RecordCoreException;
 import com.apple.foundationdb.record.lucene.LuceneExceptions;
@@ -65,10 +64,10 @@ public LuceneOptimizedStoredFieldsReader(final FDBDirectory directory, final Seg
 
     public static List<byte[]> getPrimaryKeys(final String segmentName, final FDBDirectory directory) throws IOException {
         try {
-            final List<KeyValue> rawStoredFields = directory.readAllStoredFields(segmentName);
+            final List<byte[]> rawStoredFields = directory.readAllStoredFields(segmentName);
             List<byte[]> primaryKeys = new ArrayList<>();
-            for (final KeyValue rawStoredField : rawStoredFields) {
-                final var storedFields = LuceneStoredFieldsProto.LuceneStoredFields.parseFrom(rawStoredField.getValue());
+            for (final byte[] rawStoredField : rawStoredFields) {
+                final var storedFields = LuceneStoredFieldsProto.LuceneStoredFields.parseFrom(rawStoredField);
                 primaryKeys.add(storedFields.getPrimaryKey().toByteArray());
             }
             return primaryKeys;
diff --git a/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/directory/FDBDirectory.java b/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/directory/FDBDirectory.java
@@ -92,6 +92,7 @@
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import java.util.zip.CRC32;
 
@@ -229,7 +230,8 @@ private FDBDirectory(@Nonnull Subspace subspace, @Nullable Map<String, String> i
         this.fileSequenceCounter = new AtomicLong(-1);
         this.serializer = new LuceneSerializer(Objects.requireNonNullElse(agilityContext.getPropertyValue(LuceneRecordContextProperties.LUCENE_INDEX_COMPRESSION_ENABLED), false),
                 Objects.requireNonNullElse(agilityContext.getPropertyValue(LuceneRecordContextProperties.LUCENE_INDEX_ENCRYPTION_ENABLED), false),
-                agilityContext.getPropertyValue(LuceneRecordContextProperties.LUCENE_INDEX_KEY_MANAGER));
+                agilityContext.getPropertyValue(LuceneRecordContextProperties.LUCENE_INDEX_KEY_MANAGER),
+                Objects.requireNonNullElse(agilityContext.getPropertyValue(LuceneRecordContextProperties.LUCENE_FIELD_PROTOBUF_PREFIX_ENABLED), false));
         this.fileReferenceMapSupplier = Suppliers.memoize(this::loadFileReferenceCacheForMemoization);
         this.sharedCacheManager = sharedCacheManager;
         this.sharedCacheKey = sharedCacheKey;
@@ -338,11 +340,12 @@ public void setFieldInfoId(final String filename, final long id, final ByteStrin
         writeFDBLuceneFileReference(filename, reference);
     }
 
-    void writeFieldInfos(long id, byte[] value) {
+    void writeFieldInfos(long id, byte[] rawBytes) {
         if (id == 0) {
             throw new RecordCoreArgumentException("FieldInfo id should never be 0");
         }
         byte[] key = fieldInfosSubspace.pack(id);
+        byte[] value = serializer.encodeFieldProtobuf(rawBytes);
         agilityContext.recordSize(LuceneEvents.SizeEvents.LUCENE_WRITE, key.length + value.length);
         if (LOGGER.isTraceEnabled()) {
             LOGGER.trace(getLogMessage("Write lucene stored field infos data",
@@ -357,7 +360,9 @@ Stream<NonnullPair<Long, byte[]>> getAllFieldInfosStream() {
                 LuceneEvents.Waits.WAIT_LUCENE_READ_FIELD_INFOS,
                 agilityContext.apply(aContext -> aContext.ensureActive().getRange(fieldInfosSubspace.range()).asList()))
                 .stream()
-                .map(keyValue -> NonnullPair.of(fieldInfosSubspace.unpack(keyValue.getKey()).getLong(0), keyValue.getValue()));
+                .map(keyValue -> NonnullPair.of(
+                        fieldInfosSubspace.unpack(keyValue.getKey()).getLong(0),
+                        serializer.decodeFieldProtobuf(keyValue.getValue())));
     }
 
     public CompletableFuture<Integer> getFieldInfosCount() {
@@ -444,10 +449,11 @@ public int writeData(final long id, final int block, @Nonnull final byte[] value
      * Write stored fields document to the DB.
      * @param segmentName the segment name writing to
      * @param docID the document ID to write
-     * @param value the bytes value of the stored fields
+     * @param rawBytes the bytes value of the stored fields
      */
-    public void writeStoredFields(@Nonnull String segmentName, int docID, @Nonnull final byte[] value) {
+    public void writeStoredFields(@Nonnull String segmentName, int docID, @Nonnull final byte[] rawBytes) {
         byte[] key = storedFieldsSubspace.pack(Tuple.from(segmentName, docID));
+        byte[] value = serializer.encodeFieldProtobuf(rawBytes);
         agilityContext.recordSize(LuceneEvents.SizeEvents.LUCENE_WRITE_STORED_FIELDS, key.length + value.length);
         if (LOGGER.isTraceEnabled()) {
             LOGGER.trace(getLogMessage("Write lucene stored fields data",
@@ -542,7 +548,7 @@ private CompletableFuture<byte[]> readData(long id, int block) {
     }
 
     @Nonnull
-    public byte[] readStoredFields(String segmentName, int docId) throws IOException {
+    public byte[] readStoredFields(String segmentName, int docId) {
         final byte[] key = storedFieldsSubspace.pack(Tuple.from(segmentName, docId));
         final byte[] rawBytes = asyncToSync(LuceneEvents.Waits.WAIT_LUCENE_GET_STORED_FIELDS,
                 agilityContext.instrument(LuceneEvents.Events.LUCENE_READ_STORED_FIELDS,
@@ -553,11 +559,11 @@ public byte[] readStoredFields(String segmentName, int docId) throws IOException
                     .addLogInfo(LuceneLogMessageKeys.DOC_ID, docId)
                     .addLogInfo(LogMessageKeys.KEY, ByteArrayUtil2.loggable(key));
         }
-        return rawBytes;
+        return Objects.requireNonNull(serializer.decodeFieldProtobuf(rawBytes));
     }
 
     @Nonnull
-    public List<KeyValue> readAllStoredFields(String segmentName) {
+    public List<byte[]> readAllStoredFields(String segmentName) {
         final Range range = storedFieldsSubspace.range(Tuple.from(segmentName));
         final List<KeyValue> list = asyncToSync(LuceneEvents.Waits.WAIT_LUCENE_GET_ALL_STORED_FIELDS,
                 agilityContext.getRange(range.begin, range.end));
@@ -567,7 +573,7 @@ public List<KeyValue> readAllStoredFields(String segmentName) {
                     .addLogInfo(LogMessageKeys.RANGE_START, ByteArrayUtil2.loggable(range.begin))
                     .addLogInfo(LogMessageKeys.RANGE_END, ByteArrayUtil2.loggable(range.end));
         }
-        return list;
+        return list.stream().map(KeyValue::getValue).map(serializer::decodeFieldProtobuf).collect(Collectors.toList());
     }
 
     /**
@@ -1055,7 +1061,8 @@ private KeyValueLogMessage getKeyValueLogMessage(final @Nonnull String staticMsg
         return KeyValueLogMessage.build(staticMsg, keysAndValues)
                 .addKeyAndValue(LogMessageKeys.SUBSPACE, subspace)
                 .addKeyAndValue(LuceneLogMessageKeys.COMPRESSION_SUPPOSED, serializer.isCompressionEnabled())
-                .addKeyAndValue(LuceneLogMessageKeys.ENCRYPTION_SUPPOSED, serializer.isEncryptionEnabled());
+                .addKeyAndValue(LuceneLogMessageKeys.ENCRYPTION_SUPPOSED, serializer.isEncryptionEnabled())
+                .addKeyAndValue(LuceneLogMessageKeys.FIELD_PROTOBUF_ENCODED, serializer.isFieldProtobufPrefixEnabled());
     }
 
     /**
diff --git a/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/directory/LuceneSerializer.java b/fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/directory/LuceneSerializer.java
@@ -57,12 +57,15 @@ public class LuceneSerializer {
     private final boolean encryptionEnabled;
     @Nullable
     private final SerializationKeyManager keyManager;
+    private final boolean fieldProtobufPrefixEnabled;
 
     public LuceneSerializer(boolean compressionEnabled,
-                            boolean encryptionEnabled, @Nullable SerializationKeyManager keyManager) {
+                            boolean encryptionEnabled, @Nullable SerializationKeyManager keyManager,
+                            boolean fieldProtobufPrefixEnabled) {
         this.compressionEnabled = compressionEnabled;
         this.encryptionEnabled = encryptionEnabled;
         this.keyManager = keyManager;
+        this.fieldProtobufPrefixEnabled = fieldProtobufPrefixEnabled;
     }
 
     public boolean isCompressionEnabled() {
@@ -73,6 +76,10 @@ public boolean isEncryptionEnabled() {
         return encryptionEnabled;
     }
 
+    public boolean isFieldProtobufPrefixEnabled() {
+        return fieldProtobufPrefixEnabled;
+    }
+
     @Nullable
     public SerializationKeyManager getKeyManager() {
         return keyManager;
@@ -131,7 +138,7 @@ public byte[] decode(@Nullable byte[] data) {
             return null;
         }
 
-        if (data.length < 2) {
+        if (data.length < 1) {
             throw new RecordCoreException("Invalid data")
                     .addLogInfo(LuceneLogMessageKeys.DATA_VALUE, data);
         }
@@ -270,4 +277,46 @@ private void decryptIfNeeded(@Nonnull CompressedAndEncryptedSerializerState stat
         }
         encodedDataInput.reset(decrypted);
     }
+
+    @Nullable
+    public byte[] encodeFieldProtobuf(@Nullable byte[] bytes) {
+        if (fieldProtobufPrefixEnabled) {
+            return encode(bytes);
+        } else {
+            return bytes;
+        }
+    }
+
+    @Nullable
+    public byte[] decodeFieldProtobuf(@Nullable byte[] bytes) {
+        if (bytes == null) {
+            return null;
+        }
+
+        if (isProtobufMessageWithoutPrefix(bytes)) {
+            return bytes;
+        }
+        return decode(bytes);
+    }
+
+    // This can be removed once it is guaranteed that all indexes are using the encoded format.
+    // Only works for Protobuf messages all of whose fields are themselves length-delimited,
+    // such as LuceneStoredFields (StoredField or bytes) or FieldInfos (FieldInfo).
+    private boolean isProtobufMessageWithoutPrefix(@Nonnull byte[] bytes) {
+        if (bytes.length < 1) {
+            return true;    // No room for prefix; empty message.
+        }
+        final int byte0 = bytes[0];
+        final int fieldTypeOrPrefixFlags = byte0 & 7;
+        // Either Protobuf LEN or ENCODING_COMPRESSED.
+        if (fieldTypeOrPrefixFlags != 2) {
+            return false;
+        }
+        final int fieldNumberOrKeyNumber = byte0 >> 3;
+        // ENCODING_COMPRESSED will never have a key; 0 is not a valid field number.
+        if (fieldNumberOrKeyNumber == 0) {
+            return false;
+        }
+        return true;
+    }
 }
diff --git a/fdb-record-layer-lucene/src/test/java/com/apple/foundationdb/record/lucene/FDBLuceneQueryTest.java b/fdb-record-layer-lucene/src/test/java/com/apple/foundationdb/record/lucene/FDBLuceneQueryTest.java
@@ -235,7 +235,8 @@ protected RecordLayerPropertyStorage.Builder addDefaultProps(final RecordLayerPr
         return super.addDefaultProps(props)
                 .addProp(LuceneRecordContextProperties.LUCENE_EXECUTOR_SERVICE, (Supplier<ExecutorService>)() -> executorService)
                 .addProp(LuceneRecordContextProperties.LUCENE_INDEX_ENCRYPTION_ENABLED, (Supplier<Boolean>)() -> keyManager != null)
-                .addProp(LuceneRecordContextProperties.LUCENE_INDEX_KEY_MANAGER, (Supplier<SerializationKeyManager>)() -> keyManager);
+                .addProp(LuceneRecordContextProperties.LUCENE_INDEX_KEY_MANAGER, (Supplier<SerializationKeyManager>)() -> keyManager)
+                .addProp(LuceneRecordContextProperties.LUCENE_FIELD_PROTOBUF_PREFIX_ENABLED, (Supplier<Boolean>)() -> keyManager != null);
     }
 
     @SuppressWarnings("SameParameterValue")
diff --git a/fdb-record-layer-lucene/src/test/java/com/apple/foundationdb/record/lucene/LuceneIndexMaintenanceTest.java b/fdb-record-layer-lucene/src/test/java/com/apple/foundationdb/record/lucene/LuceneIndexMaintenanceTest.java
@@ -347,7 +347,8 @@ void manyDocuments(boolean isGrouped,
                 .addProp(LuceneRecordContextProperties.LUCENE_MAX_DOCUMENTS_TO_MOVE_DURING_REPARTITIONING, dataModel.nextInt(1000) + repartitionCount)
                 .addProp(LuceneRecordContextProperties.LUCENE_MERGE_SEGMENTS_PER_TIER, (double)dataModel.nextInt(10) + 2) // it must be at least 2.0
                 .addProp(LuceneRecordContextProperties.LUCENE_INDEX_COMPRESSION_ENABLED, compressed)
-                .addProp(LuceneRecordContextProperties.LUCENE_INDEX_ENCRYPTION_ENABLED, encrypted);
+                .addProp(LuceneRecordContextProperties.LUCENE_INDEX_ENCRYPTION_ENABLED, encrypted)
+                .addProp(LuceneRecordContextProperties.LUCENE_FIELD_PROTOBUF_PREFIX_ENABLED, encrypted);
         if (encrypted) {
             contextPropsBuilder.addProp(LuceneRecordContextProperties.LUCENE_INDEX_KEY_MANAGER, new RollingTestKeyManager(seed));
         }
@@ -917,7 +918,8 @@ void changingEncryptionKey(boolean isSynthetic, boolean isGrouped, long seed) th
 
         final RecordLayerPropertyStorage.Builder contextPropsBuilder = RecordLayerPropertyStorage.newBuilder()
                 .addProp(LuceneRecordContextProperties.LUCENE_INDEX_COMPRESSION_ENABLED, true)
-                .addProp(LuceneRecordContextProperties.LUCENE_INDEX_ENCRYPTION_ENABLED, true);
+                .addProp(LuceneRecordContextProperties.LUCENE_INDEX_ENCRYPTION_ENABLED, true)
+                .addProp(LuceneRecordContextProperties.LUCENE_FIELD_PROTOBUF_PREFIX_ENABLED, true);
         final KeyGenerator keyGen = KeyGenerator.getInstance("AES");
         keyGen.init(128);
         final SecretKey key1 = keyGen.generateKey();
diff --git a/fdb-record-layer-lucene/src/test/java/com/apple/foundationdb/record/lucene/LuceneSerializerTest.java b/fdb-record-layer-lucene/src/test/java/com/apple/foundationdb/record/lucene/LuceneSerializerTest.java
diff --git a/fdb-record-layer-lucene/src/test/java/com/apple/foundationdb/record/lucene/codec/TestFDBDirectory.java b/fdb-record-layer-lucene/src/test/java/com/apple/foundationdb/record/lucene/codec/TestFDBDirectory.java

Original file line number	Diff line number	Diff line change
`@@ -235,7 +235,8 @@ protected RecordLayerPropertyStorage.Builder addDefaultProps(final RecordLayerPr`
`235`	`235`	`return super.addDefaultProps(props)`
`236`	`236`	`.addProp(LuceneRecordContextProperties.LUCENE_EXECUTOR_SERVICE, (Supplier<ExecutorService>)() -> executorService)`
`237`	`237`	`.addProp(LuceneRecordContextProperties.LUCENE_INDEX_ENCRYPTION_ENABLED, (Supplier<Boolean>)() -> keyManager != null)`
`238`		`- .addProp(LuceneRecordContextProperties.LUCENE_INDEX_KEY_MANAGER, (Supplier<SerializationKeyManager>)() -> keyManager);`
	`238`	`+ .addProp(LuceneRecordContextProperties.LUCENE_INDEX_KEY_MANAGER, (Supplier<SerializationKeyManager>)() -> keyManager)`
	`239`	`+ .addProp(LuceneRecordContextProperties.LUCENE_FIELD_PROTOBUF_PREFIX_ENABLED, (Supplier<Boolean>)() -> keyManager != null);`
`239`	`240`	`}`
`240`	`241`
`241`	`242`	`@SuppressWarnings("SameParameterValue")`