Introduce EngineBatch abstraction

server/src/internalClusterTest/java/org/elasticsearch/plugins/internal/XContentMeteringParserDecoratorIT.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.action.index.IndexRequest;
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.IndexSettings;
+import org.elasticsearch.index.engine.EngineBatch;
 import org.elasticsearch.index.engine.EngineFactory;
 import org.elasticsearch.index.engine.InternalEngine;
 import org.elasticsearch.index.mapper.MapperService;
@@ -20,7 +21,6 @@
 import org.elasticsearch.plugins.EnginePlugin;
 import org.elasticsearch.plugins.IngestPlugin;
 import org.elasticsearch.plugins.Plugin;
-import org.elasticsearch.sourcebatch.SourceBatch;
 import org.elasticsearch.test.ESIntegTestCase;
 import org.elasticsearch.xcontent.FilterXContentParserWrapper;
 import org.elasticsearch.xcontent.XContentParser;
@@ -110,9 +110,9 @@ public IndexResult index(Index index) throws IOException {
                 }
 
                 @Override
-                public List<IndexResult> indexBatch(List<Index> operations, SourceBatch batch) throws IOException {
-                    List<IndexResult> results = super.indexBatch(operations, batch);
-                    for (Index op : operations) {
+                public List<IndexResult> indexBatch(EngineBatch batch) throws IOException {
+                    List<IndexResult> results = super.indexBatch(batch);
+                    for (Index op : batch.operations()) {
                         reportDocumentSize(op.parsedDoc());
                     }
                     return results;

server/src/main/java/org/elasticsearch/action/bulk/ShardBatchIndexer.java

@@ -19,6 +19,7 @@
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.eirf.EirfRowXContentParser;
 import org.elasticsearch.index.engine.Engine;
+import org.elasticsearch.index.engine.EngineBatch;
 import org.elasticsearch.index.mapper.ShardBatchMapper;
 import org.elasticsearch.index.mapper.SourceToParse;
 import org.elasticsearch.index.seqno.SequenceNumbers;
@@ -135,7 +136,7 @@ private static void doBatchIndexOnPrimary(
             // The chunk's operations map 1:1 to the rows [chunkStart, chunkEnd); pass the matching slice so the
             // engine can write them as a single Translog.IndexBatch record.
             final SourceBatch chunkBatch = batch.slice(chunkStart, chunkEnd);
-            final List<Engine.IndexResult> results = primary.applyIndexOperationBatchOnPrimary(operations, chunkBatch);
+            final List<Engine.IndexResult> results = primary.applyIndexOperationBatchOnPrimary(new EngineBatch(operations, chunkBatch));
 
             for (Engine.IndexResult result : results) {
                 assert context.hasMoreOperationsToExecute();
@@ -225,7 +226,7 @@ static ReplicaBatchResult performBatchIndexOnReplica(BulkItemRequest[] items, So
                 // operations are the contiguous run [chunkStart, chunkStart + operations.size()); pass the matching slice
                 // so the engine writes them as a single Translog.IndexBatch record.
                 final SourceBatch chunkBatch = batch.slice(chunkStart, chunkStart + operations.size());
-                final List<Engine.IndexResult> results = replica.applyIndexOperationBatchOnReplica(operations, chunkBatch);
+                final List<Engine.IndexResult> results = replica.applyIndexOperationBatchOnReplica(new EngineBatch(operations, chunkBatch));
                 for (Engine.IndexResult result : results) {
                     if (result.getFailure() != null) {
                         throw result.getFailure();

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

@@ -102,7 +102,6 @@
 import org.elasticsearch.indices.IndexingMemoryController;
 import org.elasticsearch.indices.recovery.RecoverySettings;
 import org.elasticsearch.search.suggest.completion.CompletionStats;
-import org.elasticsearch.sourcebatch.SourceBatch;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.Transports;
 
@@ -744,9 +743,9 @@ public Condition newCondition() {
      */
     public abstract IndexResult index(Index index) throws IOException;
 
-    public List<IndexResult> indexBatch(List<Index> operations, SourceBatch batch) throws IOException {
-        ArrayList<IndexResult> results = new ArrayList<>(operations.size());
-        for (Index index : operations) {
+    public List<IndexResult> indexBatch(EngineBatch batch) throws IOException {
+        ArrayList<IndexResult> results = new ArrayList<>(batch.operations().size());
+        for (Index index : batch.operations()) {
             results.add(index(index));
         }
         return results;

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

@@ -0,0 +1,23 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.engine;
+
+import org.elasticsearch.sourcebatch.SourceBatch;
+
+import java.util.List;
+
+/**
+ * A batch of index operations ready for engine-level processing, produced by the bulk batch
+ * indexing path and consumed by {@link Engine#indexBatch}.
+ *
+ * @param operations  per-document {@link Engine.Index} operations for this batch
+ * @param sourceBatch the raw encoded source data backing the operations
+ */
+public record EngineBatch(List<Engine.Index> operations, SourceBatch sourceBatch) {}

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

@@ -1385,7 +1385,9 @@ public IndexResult index(Index index) throws IOException {
     }
 
     @Override
-    public List<IndexResult> indexBatch(List<Index> operations, SourceBatch batch) throws IOException {
+    public List<IndexResult> indexBatch(EngineBatch engineBatch) throws IOException {
+        final List<Index> operations = engineBatch.operations();
+        final SourceBatch batch = engineBatch.sourceBatch();
         assert operations.size() == batch.docCount()
             : "operations [" + operations.size() + "] must map 1:1 to batch rows [" + batch.docCount() + "]";
         try (var ignored = acquireEnsureOpenRef()) {

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

@@ -96,6 +96,7 @@
 import org.elasticsearch.index.engine.CommitStats;
 import org.elasticsearch.index.engine.Engine;
 import org.elasticsearch.index.engine.Engine.GetResult;
+import org.elasticsearch.index.engine.EngineBatch;
 import org.elasticsearch.index.engine.EngineConfig;
 import org.elasticsearch.index.engine.EngineException;
 import org.elasticsearch.index.engine.EngineFactory;
@@ -165,7 +166,6 @@
 import org.elasticsearch.search.internal.FieldUsageTrackingDirectoryReader;
 import org.elasticsearch.search.suggest.completion.CompletionStats;
 import org.elasticsearch.snapshots.Snapshot;
-import org.elasticsearch.sourcebatch.SourceBatch;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.Transports;
 
@@ -1144,29 +1144,30 @@ public static Engine.Index prepareIndex(
      * Applies a batch of index operations on the primary. Returns null if any operation requires a mapping update,
      * signaling the caller to fall back to the item-by-item path.
      */
-    public List<Engine.IndexResult> applyIndexOperationBatchOnPrimary(List<Engine.Index> operations, SourceBatch batch) throws IOException {
+    public List<Engine.IndexResult> applyIndexOperationBatchOnPrimary(EngineBatch batch) throws IOException {
         ensureWriteAllowed(Engine.Operation.Origin.PRIMARY);
         final Engine engine = getEngine();
-        return indexBatch(engine, operations, batch);
+        return indexBatch(engine, batch);
     }
 
     /**
      * Applies a batch of index operations on a replica.
      */
-    public List<Engine.IndexResult> applyIndexOperationBatchOnReplica(List<Engine.Index> operations, SourceBatch batch) throws IOException {
+    public List<Engine.IndexResult> applyIndexOperationBatchOnReplica(EngineBatch batch) throws IOException {
         ensureWriteAllowed(Engine.Operation.Origin.REPLICA);
         final Engine engine = getEngine();
-        return indexBatch(engine, operations, batch);
+        return indexBatch(engine, batch);
     }
 
-    private List<Engine.IndexResult> indexBatch(Engine engine, List<Engine.Index> operations, SourceBatch batch) throws IOException {
+    private List<Engine.IndexResult> indexBatch(Engine engine, EngineBatch batch) throws IOException {
+        final List<Engine.Index> operations = batch.operations();
         List<Engine.Index> preIndexOps = new ArrayList<>(operations.size());
         // TODO: Right now the only production users are stats. Should add batch listener.
         for (Engine.Index op : operations) {
             preIndexOps.add(indexingOperationListeners.preIndex(shardId, op));
         }
         try {
-            List<Engine.IndexResult> results = engine.indexBatch(preIndexOps, batch);
+            List<Engine.IndexResult> results = engine.indexBatch(new EngineBatch(preIndexOps, batch.sourceBatch()));
             // TODO: Look at if these can be batch optimized
             for (int i = 0; i < results.size(); i++) {
                 indexingOperationListeners.postIndex(shardId, preIndexOps.get(i), results.get(i));

server/src/test/java/org/elasticsearch/index/engine/InternalEngineTests.java

@@ -237,15 +237,15 @@ private static Engine.IndexResult indexDoc(Engine engine, Engine.Index operation
         if (randomBoolean()) {
             SourceBatch batch = tryEncodeAsBatch(List.of(operation));
             if (batch != null) {
-                return engine.indexBatch(List.of(operation), batch).getFirst();
+                return engine.indexBatch(new EngineBatch(List.of(operation), batch)).getFirst();
             }
         }
         return engine.index(operation);
     }
 
     /**
      * Encodes the given ops' sources into a {@link SourceBatch} that can be passed to
-     * {@link Engine#indexBatch(List, SourceBatch)}.
+     * {@link Engine#indexBatch(EngineBatch)}.
      */
     private static SourceBatch encodeAsBatch(List<Engine.Index> operations) throws IOException {
         List<BytesReference> sources = new ArrayList<>(operations.size());
@@ -8077,7 +8077,7 @@ public void testIndexBatchMultipleDocuments() throws IOException {
             ParsedDocument doc = createParsedDoc(Integer.toString(i), null);
             ops.add(indexForDoc(doc));
         }
-        List<Engine.IndexResult> results = engine.indexBatch(ops, encodeAsBatch(ops));
+        List<Engine.IndexResult> results = engine.indexBatch(new EngineBatch(ops, encodeAsBatch(ops)));
         assertThat(results, hasSize(ops.size()));
         for (int i = 0; i < results.size(); i++) {
             Engine.IndexResult result = results.get(i);
@@ -8096,7 +8096,7 @@ public void testIndexBatchSequenceNumbersAreMonotonic() throws IOException {
         for (int i = 0; i < batchSize; i++) {
             ops.add(indexForDoc(createParsedDoc(Integer.toString(i), null)));
         }
-        List<Engine.IndexResult> results = engine.indexBatch(ops, encodeAsBatch(ops));
+        List<Engine.IndexResult> results = engine.indexBatch(new EngineBatch(ops, encodeAsBatch(ops)));
         long prevSeqNo = -1;
         for (Engine.IndexResult result : results) {
             assertThat(result.getResultType(), equalTo(Engine.Result.Type.SUCCESS));
@@ -8109,7 +8109,7 @@ public void testIndexBatchSingletonEquivalentToIndex() throws IOException {
         ParsedDocument doc = createParsedDoc("1", null);
         Engine.Index op = indexForDoc(doc);
         List<Engine.Index> ops = List.of(op);
-        List<Engine.IndexResult> batchResults = engine.indexBatch(ops, encodeAsBatch(ops));
+        List<Engine.IndexResult> batchResults = engine.indexBatch(new EngineBatch(ops, encodeAsBatch(ops)));
         assertThat(batchResults, hasSize(1));
         Engine.IndexResult result = batchResults.getFirst();
         assertThat(result.getResultType(), equalTo(Engine.Result.Type.SUCCESS));
@@ -8143,7 +8143,7 @@ public void testIndexBatchVersionConflict() throws IOException {
             firstResult.getTerm()
         );
         List<Engine.Index> ops = List.of(conflictingOp);
-        List<Engine.IndexResult> results = engine.indexBatch(ops, encodeAsBatch(ops));
+        List<Engine.IndexResult> results = engine.indexBatch(new EngineBatch(ops, encodeAsBatch(ops)));
         assertThat(results, hasSize(1));
         assertThat(results.getFirst().getResultType(), equalTo(Engine.Result.Type.FAILURE));
         assertThat(results.getFirst().getFailure(), instanceOf(VersionConflictEngineException.class));
@@ -8159,7 +8159,7 @@ public void testIndexBatchMixedNewAndExisting() throws IOException {
             indexForDoc(createParsedDoc("1", null)),
             indexForDoc(createParsedDoc("2", null))
         );
-        List<Engine.IndexResult> results = engine.indexBatch(ops, encodeAsBatch(ops));
+        List<Engine.IndexResult> results = engine.indexBatch(new EngineBatch(ops, encodeAsBatch(ops)));
         assertThat(results, hasSize(3));
         for (Engine.IndexResult result : results) {
             assertThat(result.getResultType(), equalTo(Engine.Result.Type.SUCCESS));
@@ -8187,7 +8187,7 @@ public void testIndexBatchVersionLookupFromLucene() throws IOException {
         for (ParsedDocument doc : docs) {
             updates.add(indexForDoc(doc));
         }
-        List<Engine.IndexResult> results = engine.indexBatch(updates, encodeAsBatch(updates));
+        List<Engine.IndexResult> results = engine.indexBatch(new EngineBatch(updates, encodeAsBatch(updates)));
         assertThat(results, hasSize(count));
         for (Engine.IndexResult result : results) {
             assertThat(result.getResultType(), equalTo(Engine.Result.Type.SUCCESS));
@@ -8217,7 +8217,7 @@ public void testIndexBatchVersionConflictFromLucene() throws IOException {
             firstResult.getTerm()
         );
         List<Engine.Index> ops = List.of(conflictingOp);
-        List<Engine.IndexResult> results = engine.indexBatch(ops, encodeAsBatch(ops));
+        List<Engine.IndexResult> results = engine.indexBatch(new EngineBatch(ops, encodeAsBatch(ops)));
         assertThat(results, hasSize(1));
         assertThat(results.getFirst().getResultType(), equalTo(Engine.Result.Type.FAILURE));
         assertThat(results.getFirst().getFailure(), instanceOf(VersionConflictEngineException.class));
@@ -8234,7 +8234,7 @@ public void testIndexBatchMixedLuceneAndVersionMap() throws IOException {
         indexDoc(engine, indexForDoc(doc2));
 
         List<Engine.Index> ops = List.of(indexForDoc(doc1), indexForDoc(doc2));
-        List<Engine.IndexResult> results = engine.indexBatch(ops, encodeAsBatch(ops));
+        List<Engine.IndexResult> results = engine.indexBatch(new EngineBatch(ops, encodeAsBatch(ops)));
         assertThat(results, hasSize(2));
         for (Engine.IndexResult result : results) {
             assertThat(result.getResultType(), equalTo(Engine.Result.Type.SUCCESS));
@@ -8273,7 +8273,7 @@ public void testIndexBatchGcExpiredTombstoneTreatedAsNotFound() throws IOExcepti
         // where the stale live document is found, and the operation is incorrectly treated as an
         // update rather than a create.
         List<Engine.Index> ops = List.of(indexForDoc(doc));
-        List<Engine.IndexResult> results = engine.indexBatch(ops, encodeAsBatch(ops));
+        List<Engine.IndexResult> results = engine.indexBatch(new EngineBatch(ops, encodeAsBatch(ops)));
         assertThat(results, hasSize(1));
         assertThat(results.getFirst().getResultType(), equalTo(Engine.Result.Type.SUCCESS));
         assertThat(results.getFirst().isCreated(), equalTo(true));
@@ -8288,7 +8288,7 @@ public void testIndexBatchFastPathOnly() throws IOException {
         for (int i = 0; i < count; i++) {
             ops.add(appendOnlyPrimary(createParsedDoc(Integer.toString(i), null), false, timestamp + i));
         }
-        List<Engine.IndexResult> results = engine.indexBatch(ops, encodeAsBatch(ops));
+        List<Engine.IndexResult> results = engine.indexBatch(new EngineBatch(ops, encodeAsBatch(ops)));
         assertThat(results, hasSize(count));
         for (Engine.IndexResult result : results) {
             assertThat(result.getResultType(), equalTo(Engine.Result.Type.SUCCESS));
@@ -8338,7 +8338,7 @@ public void testIndexBatchTimeSeriesPhase2() throws IOException {
                 for (ParsedDocument doc : docs) {
                     updates.add(indexForDoc(doc));
                 }
-                List<Engine.IndexResult> results = engine.indexBatch(updates, encodeAsBatch(updates));
+                List<Engine.IndexResult> results = engine.indexBatch(new EngineBatch(updates, encodeAsBatch(updates)));
                 assertThat(results, hasSize(count));
                 for (Engine.IndexResult result : results) {
                     assertThat(result.getResultType(), equalTo(Engine.Result.Type.SUCCESS));
@@ -8354,7 +8354,7 @@ public void testMixedPrimaryTermThrows() throws IOException {
         Engine.Index op1 = new Engine.Index(newUid(doc1), primaryTerm.get(), doc1);
         Engine.Index op2 = new Engine.Index(newUid(doc2), primaryTerm.get() + 1, doc2);
         var updates = List.of(op1, op2);
-        expectThrows(AssertionError.class, () -> engine.indexBatch(updates, encodeAsBatch(updates)));
+        expectThrows(AssertionError.class, () -> engine.indexBatch(new EngineBatch(updates, encodeAsBatch(updates))));
     }
 
     public void testIndexBatchSeqNosAreContiguous() throws IOException {
@@ -8365,7 +8365,7 @@ public void testIndexBatchSeqNosAreContiguous() throws IOException {
             ops.add(indexForDoc(createParsedDoc(Integer.toString(i), null)));
         }
         long seqNoBefore = engine.getLocalCheckpointTracker().getMaxSeqNo();
-        List<Engine.IndexResult> results = engine.indexBatch(ops, encodeAsBatch(ops));
+        List<Engine.IndexResult> results = engine.indexBatch(new EngineBatch(ops, encodeAsBatch(ops)));
 
         assertThat(results, hasSize(batchSize));
         long firstSeqNo = seqNoBefore + 1;
@@ -8382,7 +8382,7 @@ public void testIndexBatchProcessedCheckpointAdvancesAfterBatch() throws IOExcep
             ops.add(indexForDoc(createParsedDoc(Integer.toString(i), null)));
         }
         long checkpointBefore = engine.getProcessedLocalCheckpoint();
-        List<Engine.IndexResult> results = engine.indexBatch(ops, encodeAsBatch(ops));
+        List<Engine.IndexResult> results = engine.indexBatch(new EngineBatch(ops, encodeAsBatch(ops)));
 
         assertThat(results, hasSize(batchSize));
         long expectedCheckpoint = checkpointBefore + batchSize;
@@ -8421,7 +8421,7 @@ public void testIndexBatchCheckpointWithVersionConflicts() throws IOException {
 
         long checkpointBefore = engine.getProcessedLocalCheckpoint();
         List<Engine.Index> ops = List.of(conflicting, goodOp);
-        List<Engine.IndexResult> results = engine.indexBatch(ops, encodeAsBatch(ops));
+        List<Engine.IndexResult> results = engine.indexBatch(new EngineBatch(ops, encodeAsBatch(ops)));
 
         assertThat(results, hasSize(2));
         // conflicting op: failure, no seq no assigned

test/framework/src/main/java/org/elasticsearch/index/engine/InternalTestEngine.java

@@ -12,7 +12,6 @@
 import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
 import org.elasticsearch.index.seqno.LocalCheckpointTracker;
 import org.elasticsearch.index.seqno.SequenceNumbers;
-import org.elasticsearch.sourcebatch.SourceBatch;
 
 import java.io.IOException;
 import java.util.List;
@@ -54,8 +53,8 @@ public IndexResult index(Index index) throws IOException {
     }
 
     @Override
-    public List<IndexResult> indexBatch(List<Index> operations, SourceBatch batch) throws IOException {
-        for (Index index : operations) {
+    public List<IndexResult> indexBatch(EngineBatch batch) throws IOException {
+        for (Index index : batch.operations()) {
             if (index.seqNo() != SequenceNumbers.UNASSIGNED_SEQ_NO) {
                 idToMaxSeqNo.compute(index.id(), (id, existing) -> {
                     if (existing == null) {
@@ -68,7 +67,7 @@ public List<IndexResult> indexBatch(List<Index> operations, SourceBatch batch) t
                 });
             }
         }
-        return super.indexBatch(operations, batch);
+        return super.indexBatch(batch);
     }
 
     @Override