Improve serialization of translog operations (#134243)

Currently translog operations are serialized in a way that is relatively
inefficient. This commit improves performance by re-ordering the
operations to allow most of the header to be serialized without bounds
checks. Additionally, instead of incrementally calculating the checksum
we do a single pass at the end. Finally, we no longer copy the source
twice. Instead serialize it directly into the translog writer.

Author: Tim-Brooks

server/src/main/java/org/elasticsearch/common/bytes/ReleasableBytesReference.java

@@ -51,6 +51,13 @@ public static ReleasableBytesReference wrap(BytesReference reference) {
         return reference.length() == 0 ? empty() : new ReleasableBytesReference(reference, ALWAYS_REFERENCED);
     }
 
+    public static BytesReference unwrap(BytesReference reference) {
+        if (reference instanceof ReleasableBytesReference releasable) {
+            return releasable.delegate;
+        }
+        return reference;
+    }
+
     @Override
     public void incRef() {
         refCounted.incRef();
@@ -278,6 +285,11 @@ public int arrayOffset() {
         return delegate.arrayOffset();
     }
 
+    public BytesReference delegate() {
+        assert hasReferences();
+        return delegate;
+    }
+
     private static final class RefCountedReleasable extends AbstractRefCounted {
 
         private final Releasable releasable;

server/src/main/java/org/elasticsearch/common/io/stream/RecyclerBytesStreamOutput.java

@@ -15,14 +15,12 @@
 import org.elasticsearch.common.bytes.CompositeBytesReference;
 import org.elasticsearch.common.bytes.ReleasableBytesReference;
 import org.elasticsearch.common.recycler.Recycler;
+import org.elasticsearch.common.util.ByteUtils;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
 
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
-import java.lang.invoke.MethodHandles;
-import java.lang.invoke.VarHandle;
-import java.nio.ByteOrder;
 import java.util.ArrayList;
 import java.util.Objects;
 
@@ -33,11 +31,6 @@
  */
 public class RecyclerBytesStreamOutput extends BytesStream implements Releasable {
 
-    static final VarHandle VH_BE_INT = MethodHandles.byteArrayViewVarHandle(int[].class, ByteOrder.BIG_ENDIAN);
-    static final VarHandle VH_LE_INT = MethodHandles.byteArrayViewVarHandle(int[].class, ByteOrder.LITTLE_ENDIAN);
-    static final VarHandle VH_BE_LONG = MethodHandles.byteArrayViewVarHandle(long[].class, ByteOrder.BIG_ENDIAN);
-    static final VarHandle VH_LE_LONG = MethodHandles.byteArrayViewVarHandle(long[].class, ByteOrder.LITTLE_ENDIAN);
-
     private ArrayList<Recycler.V<BytesRef>> pages = new ArrayList<>(8);
     private final Recycler<BytesRef> recycler;
     private final int pageSize;
@@ -148,7 +141,7 @@ public void writeVInt(int i) throws IOException {
         }
     }
 
-    protected static int vIntLength(int value) {
+    public static int vIntLength(int value) {
         int leadingZeros = Integer.numberOfLeadingZeros(value);
         if (leadingZeros >= 25) {
             return 1;
@@ -177,7 +170,7 @@ public void writeInt(int i) throws IOException {
             super.writeInt(i);
         } else {
             BytesRef currentPage = currentBytesRef;
-            VH_BE_INT.set(currentPage.bytes, currentPage.offset + currentPageOffset, i);
+            ByteUtils.writeIntBE(i, currentPage.bytes, currentPage.offset + currentPageOffset);
             this.currentPageOffset = currentPageOffset + 4;
         }
     }
@@ -189,7 +182,7 @@ public void writeIntLE(int i) throws IOException {
             super.writeIntLE(i);
         } else {
             BytesRef currentPage = currentBytesRef;
-            VH_LE_INT.set(currentPage.bytes, currentPage.offset + currentPageOffset, i);
+            ByteUtils.writeIntLE(i, currentPage.bytes, currentPage.offset + currentPageOffset);
             this.currentPageOffset = currentPageOffset + 4;
         }
     }
@@ -201,7 +194,7 @@ public void writeLong(long i) throws IOException {
             super.writeLong(i);
         } else {
             BytesRef currentPage = currentBytesRef;
-            VH_BE_LONG.set(currentPage.bytes, currentPage.offset + currentPageOffset, i);
+            ByteUtils.writeLongBE(i, currentPage.bytes, currentPage.offset + currentPageOffset);
             this.currentPageOffset = currentPageOffset + 8;
         }
     }
@@ -213,7 +206,7 @@ public void writeLongLE(long i) throws IOException {
             super.writeLongLE(i);
         } else {
             BytesRef currentPage = currentBytesRef;
-            VH_LE_LONG.set(currentPage.bytes, currentPage.offset + currentPageOffset, i);
+            ByteUtils.writeLongLE(i, currentPage.bytes, currentPage.offset + currentPageOffset);
             this.currentPageOffset = currentPageOffset + 8;
         }
     }
@@ -241,6 +234,25 @@ public void legacyWriteWithSizePrefix(Writeable writeable) throws IOException {
         }
     }
 
+    /**
+     * Attempt to get one page to perform a write directly into the page. The page will only be returned if the requested bytes can fit.
+     * If requested bytes cannot fit, null will be returned. This will advance the current position in the stream.
+     *
+     * @param bytes the number of bytes for the single write
+     * @return a direct page if there is enough space in current page, otherwise null
+     */
+    public BytesRef tryGetPageForWrite(int bytes) {
+        final int beforePageOffset = this.currentPageOffset;
+        if (bytes <= (pageSize - beforePageOffset)) {
+            BytesRef currentPage = currentBytesRef;
+            BytesRef bytesRef = new BytesRef(currentPage.bytes, currentPage.offset + beforePageOffset, bytes);
+            this.currentPageOffset = beforePageOffset + bytes;
+            return bytesRef;
+        } else {
+            return null;
+        }
+    }
+
     // overridden with some code duplication the same way other write methods in this class are overridden to bypass StreamOutput's
     // intermediary buffers
     @Override

server/src/main/java/org/elasticsearch/index/engine/TranslogDirectoryReader.java

@@ -182,8 +182,9 @@ private static LeafReader createInMemoryReader(
         boolean rootDocOnly,
         Translog.Index operation
     ) {
+        final String id = Uid.decodeId(operation.uid());
         final ParsedDocument parsedDocs = documentParser.parseDocument(
-            new SourceToParse(operation.id(), operation.source(), XContentHelper.xContentType(operation.source()), operation.routing()),
+            new SourceToParse(id, operation.source(), XContentHelper.xContentType(operation.source()), operation.routing()),
             mappingLookup
         );
 
@@ -244,7 +245,7 @@ protected StoredFieldsReader doGetSequentialStoredFieldsReader(StoredFieldsReade
                 }
             };
         } catch (IOException e) {
-            throw new EngineException(shardId, "failed to create an in-memory segment for get [" + operation.id() + "]", e);
+            throw new EngineException(shardId, "failed to create an in-memory segment for get [" + id + "]", e);
         }
     }
 
@@ -339,7 +340,7 @@ private static class TranslogLeafReader extends LeafReader {
             this.engineConfig = engineConfig;
             this.onSegmentCreated = onSegmentCreated;
             this.directory = directory;
-            this.uid = Uid.encodeId(operation.id());
+            this.uid = operation.uid();
         }
 
         private LeafReader getDelegate() {

server/src/main/java/org/elasticsearch/index/mapper/Uid.java

@@ -173,6 +173,12 @@ private static String decodeBase64Id(byte[] idBytes, int offset, int length) {
         return Strings.BASE_64_NO_PADDING_URL_ENCODER.encodeToString(idBytes);
     }
 
+    /** Decode an indexed id back to its original form.
+     *  @see #encodeId */
+    public static String decodeId(BytesRef idBytes) {
+        return decodeId(idBytes.bytes, idBytes.offset, idBytes.length);
+    }
+
     /** Decode an indexed id back to its original form.
      *  @see #encodeId */
     public static String decodeId(byte[] idBytes) {

server/src/main/java/org/elasticsearch/index/shard/IndexShard.java

@@ -2130,7 +2130,12 @@ private Engine.Result applyTranslogOperation(Engine engine, Translog.Operation o
                     index.getAutoGeneratedIdTimestamp(),
                     true,
                     origin,
-                    new SourceToParse(index.id(), index.source(), XContentHelper.xContentType(index.source()), index.routing())
+                    new SourceToParse(
+                        Uid.decodeId(index.uid()),
+                        index.source(),
+                        XContentHelper.xContentType(index.source()),
+                        index.routing()
+                    )
                 );
             }
             case DELETE -> {
@@ -2140,7 +2145,7 @@ private Engine.Result applyTranslogOperation(Engine engine, Translog.Operation o
                     delete.seqNo(),
                     delete.primaryTerm(),
                     delete.version(),
-                    delete.id(),
+                    Uid.decodeId(delete.uid()),
                     versionType,
                     UNASSIGNED_SEQ_NO,
                     0,

server/src/main/java/org/elasticsearch/index/translog/OperationListener.java

@@ -9,18 +9,15 @@
 
 package org.elasticsearch.index.translog;
 
-import org.elasticsearch.common.bytes.BytesReference;
-
 @FunctionalInterface
 public interface OperationListener {
 
     /**
-     * This method is called when a new operation is added to the translog. The BytesReference is a releasable
-     * instance, so it should not be retained beyond the scope of this method.
+     * This method is called when a new operation is added to the translog.
      *
-     * @param data a releasable bytes reference of the data add
+     * @param operation the serialized operation added to the translog
      * @param seqNo the sequence number of the operation
      * @param location the location written
      */
-    void operationAdded(BytesReference data, long seqNo, Translog.Location location);
+    void operationAdded(Translog.Serialized operation, long seqNo, Translog.Location location);
 }