Use the LuceneSerializer to encode store fields, too.

Since existing fields will be straight protobuf without the serializer's prefix byte, recognize that, too.

Author: MMcM

fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/directory/FDBDirectory.java

@@ -444,10 +444,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.encode(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 +543,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,7 +554,7 @@ 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.decodePossiblyWithoutPrefix(rawBytes));
     }
 
     @Nonnull

fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/directory/LuceneSerializer.java

@@ -131,7 +131,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 +270,36 @@ private void decryptIfNeeded(@Nonnull CompressedAndEncryptedSerializerState stat
         }
         encodedDataInput.reset(decrypted);
     }
+
+    @Nullable
+    public byte[] decodePossiblyWithoutPrefix(@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.
+    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;
+    }
 }

fdb-record-layer-lucene/src/test/java/com/apple/foundationdb/record/lucene/LuceneSerializerTest.java

@@ -32,10 +32,13 @@
 import org.junit.jupiter.api.Tag;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
 
 import javax.crypto.KeyGenerator;
 import javax.crypto.SecretKey;
 import java.util.concurrent.ThreadLocalRandom;
+import java.util.stream.Stream;
 
 /**
  * Test for Lucene data compression/decompression and encryption/decryption validation.
@@ -122,4 +125,33 @@ void testEncodingWithEncryption(boolean compressToo) throws Exception {
         final LuceneFileSystemProto.LuceneFileReference decryptedReference = LuceneFileSystemProto.LuceneFileReference.parseFrom(decodedValue);
         Assertions.assertEquals(content, decryptedReference.getContent());
     }
+
+    @ParameterizedTest
+    @MethodSource
+    void testProtobufMessageWithoutPrefix(boolean encode, boolean encrypt, boolean compress) throws Exception {
+        KeyGenerator keyGen = KeyGenerator.getInstance("AES");
+        keyGen.init(128);
+        SecretKey key = keyGen.generateKey();
+        final SerializationKeyManager keyManager = new FixedZeroKeyManager(key, null, null);
+        final LuceneSerializer serializer = new LuceneSerializer(compress, encrypt, keyManager);
+        final String storedField = RandomUtil.randomAlphanumericString(ThreadLocalRandom.current(), 20);
+        final LuceneStoredFieldsProto.LuceneStoredFields.Builder builder = LuceneStoredFieldsProto.LuceneStoredFields.newBuilder();
+        builder.addStoredFieldsBuilder().setFieldNumber(5).setStringValue(storedField);
+        final byte[] value = builder.build().toByteArray();
+        final byte[] encodedValue = encode ? serializer.encode(value) : value;
+        final byte[] decodedValue = serializer.decodePossiblyWithoutPrefix(encodedValue);
+        Assertions.assertArrayEquals(value, decodedValue);
+        final LuceneStoredFieldsProto.LuceneStoredFields storedFields = LuceneStoredFieldsProto.LuceneStoredFields.parseFrom(decodedValue);
+        Assertions.assertEquals(storedField, storedFields.getStoredFields(0).getStringValue());
+    }
+
+    static Stream<Arguments> testProtobufMessageWithoutPrefix() {
+        return Stream.of(
+               Arguments.of(false, false, false),
+               Arguments.of(true, false, false),
+               Arguments.of(true, false, true),
+               Arguments.of(true, true, false),
+               Arguments.of(true, true, true)
+        );
+    }
 }