Allow incremental bulk request execution

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -193,6 +193,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion REMOVE_GLOBAL_RETENTION_FROM_TEMPLATES = def(8_723_00_0);
     public static final TransportVersion RANDOM_RERANKER_RETRIEVER = def(8_724_00_0);
     public static final TransportVersion ESQL_PROFILE_SLEEPS = def(8_725_00_0);
+    public static final TransportVersion BULK_INCREMENTAL_STATE = def(8_726_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -91,6 +91,7 @@ final class BulkOperation extends ActionRunnable<BulkResponse> {
     private final OriginSettingClient rolloverClient;
     private final Set<String> failureStoresToBeRolledOver = ConcurrentCollections.newConcurrentSet();
     private final Set<Integer> failedRolloverRequests = ConcurrentCollections.newConcurrentSet();
+    private final Map<ShardId, Exception> shortCircuitShardFailures = ConcurrentCollections.newConcurrentMap();
 
     BulkOperation(
         Task task,
@@ -156,6 +157,7 @@ final class BulkOperation extends ActionRunnable<BulkResponse> {
         this.observer = observer;
         this.failureStoreDocumentConverter = failureStoreDocumentConverter;
         this.rolloverClient = new OriginSettingClient(client, LAZY_ROLLOVER_ORIGIN);
+        this.shortCircuitShardFailures.putAll(bulkRequest.incrementalState().shardLevelFailures());
     }
 
     @Override
@@ -362,7 +364,6 @@ private void executeBulkRequestsByShard(
             for (Map.Entry<ShardId, List<BulkItemRequest>> entry : requestsByShard.entrySet()) {
                 final ShardId shardId = entry.getKey();
                 final List<BulkItemRequest> requests = entry.getValue();
-
                 BulkShardRequest bulkShardRequest = new BulkShardRequest(
                     shardId,
                     bulkRequest.getRefreshPolicy(),
@@ -394,7 +395,12 @@ private void redirectFailuresOrCompleteBulkOperation() {
 
     private void completeBulkOperation() {
         listener.onResponse(
-            new BulkResponse(responses.toArray(new BulkItemResponse[responses.length()]), buildTookInMillis(startTimeNanos))
+            new BulkResponse(
+                responses.toArray(new BulkItemResponse[responses.length()]),
+                buildTookInMillis(startTimeNanos),
+                BulkResponse.NO_INGEST_TOOK,
+                new BulkRequest.IncrementalState(shortCircuitShardFailures)
+            )
         );
         // Allow memory for bulk shard request items to be reclaimed before all items have been completed
         bulkRequest = null;
@@ -420,66 +426,81 @@ private void discardRedirectsAndFinish(Exception exception) {
     }
 
     private void executeBulkShardRequest(BulkShardRequest bulkShardRequest, Releasable releaseOnFinish) {
-        client.executeLocally(TransportShardBulkAction.TYPE, bulkShardRequest, new ActionListener<>() {
+        ShardId shardId = bulkShardRequest.shardId();
+
+        // Short circuit the shark level request with the existing shard failure.
+        if (shortCircuitShardFailures.containsKey(shardId)) {
+            handleShardFailure(bulkShardRequest, clusterService.state(), shortCircuitShardFailures.get(shardId));
+            releaseOnFinish.close();
+        } else {
+            client.executeLocally(TransportShardBulkAction.TYPE, bulkShardRequest, new ActionListener<>() {
 
-            // Lazily get the cluster state to avoid keeping it around longer than it is needed
-            private ClusterState clusterState = null;
+                // Lazily get the cluster state to avoid keeping it around longer than it is needed
+                private ClusterState clusterState = null;
 
-            private ClusterState getClusterState() {
-                if (clusterState == null) {
-                    clusterState = clusterService.state();
+                private ClusterState getClusterState() {
+                    if (clusterState == null) {
+                        clusterState = clusterService.state();
+                    }
+                    return clusterState;
                 }
-                return clusterState;
-            }
 
-            @Override
-            public void onResponse(BulkShardResponse bulkShardResponse) {
-                for (int idx = 0; idx < bulkShardResponse.getResponses().length; idx++) {
-                    // We zip the requests and responses together so that we can identify failed documents and potentially store them
-                    BulkItemResponse bulkItemResponse = bulkShardResponse.getResponses()[idx];
-
-                    if (bulkItemResponse.isFailed()) {
-                        BulkItemRequest bulkItemRequest = bulkShardRequest.items()[idx];
-                        assert bulkItemRequest.id() == bulkItemResponse.getItemId() : "Bulk items were returned out of order";
-
-                        DataStream failureStoreReference = getRedirectTarget(bulkItemRequest.request(), getClusterState().metadata());
-                        if (failureStoreReference != null) {
-                            maybeMarkFailureStoreForRollover(failureStoreReference);
-                            var cause = bulkItemResponse.getFailure().getCause();
-                            addDocumentToRedirectRequests(bulkItemRequest, cause, failureStoreReference.getName());
+                @Override
+                public void onResponse(BulkShardResponse bulkShardResponse) {
+                    for (int idx = 0; idx < bulkShardResponse.getResponses().length; idx++) {
+                        // We zip the requests and responses together so that we can identify failed documents and potentially store them
+                        BulkItemResponse bulkItemResponse = bulkShardResponse.getResponses()[idx];
+
+                        if (bulkItemResponse.isFailed()) {
+                            BulkItemRequest bulkItemRequest = bulkShardRequest.items()[idx];
+                            assert bulkItemRequest.id() == bulkItemResponse.getItemId() : "Bulk items were returned out of order";
+
+                            DataStream failureStoreReference = getRedirectTarget(bulkItemRequest.request(), getClusterState().metadata());
+                            if (failureStoreReference != null) {
+                                maybeMarkFailureStoreForRollover(failureStoreReference);
+                                var cause = bulkItemResponse.getFailure().getCause();
+                                addDocumentToRedirectRequests(bulkItemRequest, cause, failureStoreReference.getName());
+                            }
+                            addFailure(bulkItemResponse);
+                        } else {
+                            bulkItemResponse.getResponse().setShardInfo(bulkShardResponse.getShardInfo());
+                            responses.set(bulkItemResponse.getItemId(), bulkItemResponse);
                         }
-                        addFailure(bulkItemResponse);
-                    } else {
-                        bulkItemResponse.getResponse().setShardInfo(bulkShardResponse.getShardInfo());
-                        responses.set(bulkItemResponse.getItemId(), bulkItemResponse);
                     }
+                    completeShardOperation();
                 }
-                completeShardOperation();
-            }
 
-            @Override
-            public void onFailure(Exception e) {
-                // create failures for all relevant requests
-                for (BulkItemRequest request : bulkShardRequest.items()) {
-                    final String indexName = request.index();
-                    DocWriteRequest<?> docWriteRequest = request.request();
-
-                    DataStream failureStoreReference = getRedirectTarget(docWriteRequest, getClusterState().metadata());
-                    if (failureStoreReference != null) {
-                        maybeMarkFailureStoreForRollover(failureStoreReference);
-                        addDocumentToRedirectRequests(request, e, failureStoreReference.getName());
-                    }
-                    addFailure(docWriteRequest, request.id(), indexName, e);
+                @Override
+                public void onFailure(Exception e) {
+                    assert shortCircuitShardFailures.containsKey(shardId) == false;
+                    shortCircuitShardFailures.put(shardId, e);
+                    handleShardFailure(bulkShardRequest, getClusterState(), e);
+                    completeShardOperation();
                 }
-                completeShardOperation();
-            }
 
-            private void completeShardOperation() {
-                // Clear our handle on the cluster state to allow it to be cleaned up
-                clusterState = null;
-                releaseOnFinish.close();
+                private void completeShardOperation() {
+                    // Clear our handle on the cluster state to allow it to be cleaned up
+                    clusterState = null;
+                    releaseOnFinish.close();
+                }
+            });
+        }
+
+    }
+
+    private void handleShardFailure(BulkShardRequest bulkShardRequest, ClusterState clusterState, Exception e) {
+        // create failures for all relevant requests
+        for (BulkItemRequest request : bulkShardRequest.items()) {
+            final String indexName = request.index();
+            DocWriteRequest<?> docWriteRequest = request.request();
+
+            DataStream failureStoreReference = getRedirectTarget(docWriteRequest, clusterState.metadata());
+            if (failureStoreReference != null) {
+                maybeMarkFailureStoreForRollover(failureStoreReference);
+                addDocumentToRedirectRequests(request, e, failureStoreReference.getName());
             }
-        });
+            addFailure(docWriteRequest, request.id(), indexName, e);
+        }
     }
 
     /**

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

@@ -10,6 +10,7 @@
 
 import org.apache.lucene.util.Accountable;
 import org.apache.lucene.util.RamUsageEstimator;
+import org.elasticsearch.TransportVersions;
 import org.elasticsearch.action.ActionRequest;
 import org.elasticsearch.action.ActionRequestValidationException;
 import org.elasticsearch.action.CompositeIndicesRequest;
@@ -25,9 +26,11 @@
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.RestApiVersion;
 import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
 import org.elasticsearch.transport.RawIndexingDataTransportRequest;
 import org.elasticsearch.xcontent.XContentType;
@@ -37,6 +40,7 @@
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 
@@ -69,6 +73,7 @@ public class BulkRequest extends ActionRequest
     private final Set<String> indices = new HashSet<>();
 
     protected TimeValue timeout = BulkShardRequest.DEFAULT_TIMEOUT;
+    private IncrementalState incrementalState = IncrementalState.EMPTY;
     private ActiveShardCount waitForActiveShards = ActiveShardCount.DEFAULT;
     private RefreshPolicy refreshPolicy = RefreshPolicy.NONE;
     private String globalPipeline;
@@ -90,6 +95,11 @@ public BulkRequest(StreamInput in) throws IOException {
         for (DocWriteRequest<?> request : requests) {
             indices.add(Objects.requireNonNull(request.index(), "request index must not be null"));
         }
+        if (in.getTransportVersion().onOrAfter(TransportVersions.BULK_INCREMENTAL_STATE)) {
+            incrementalState = new BulkRequest.IncrementalState(in);
+        } else {
+            incrementalState = BulkRequest.IncrementalState.EMPTY;
+        }
     }
 
     public BulkRequest(@Nullable String globalIndex) {
@@ -324,6 +334,10 @@ public final BulkRequest timeout(TimeValue timeout) {
         return this;
     }
 
+    public void incrementalState(IncrementalState incrementalState) {
+        this.incrementalState = incrementalState;
+    }
+
     /**
      * Note for internal callers (NOT high level rest client),
      * the global parameter setting is ignored when used with:
@@ -362,6 +376,10 @@ public TimeValue timeout() {
         return timeout;
     }
 
+    public IncrementalState incrementalState() {
+        return incrementalState;
+    }
+
     public String pipeline() {
         return globalPipeline;
     }
@@ -433,6 +451,9 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeCollection(requests, DocWriteRequest::writeDocumentRequest);
         refreshPolicy.writeTo(out);
         out.writeTimeValue(timeout);
+        if (out.getTransportVersion().onOrAfter(TransportVersions.BULK_INCREMENTAL_STATE)) {
+            incrementalState.writeTo(out);
+        }
     }
 
     @Override
@@ -474,4 +495,18 @@ public Set<String> getIndices() {
     public boolean isSimulated() {
         return false; // Always false, but may be overridden by a subclass
     }
+
+    record IncrementalState(Map<ShardId, Exception> shardLevelFailures) implements Writeable {
+
+        static final IncrementalState EMPTY = new IncrementalState(Collections.emptyMap());
+
+        IncrementalState(StreamInput in) throws IOException {
+            this(in.readMap(ShardId::new, input -> input.readException()));
+        }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            out.writeMap(shardLevelFailures, (o, s) -> s.writeTo(o), StreamOutput::writeException);
+        }
+    }
 }

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

@@ -116,7 +116,12 @@ BulkRequest getBulkRequest() {
     ActionListener<BulkResponse> wrapActionListenerIfNeeded(long ingestTookInMillis, ActionListener<BulkResponse> actionListener) {
         if (itemResponses.isEmpty()) {
             return actionListener.map(
-                response -> new BulkResponse(response.getItems(), response.getTook().getMillis(), ingestTookInMillis)
+                response -> new BulkResponse(
+                    response.getItems(),
+                    response.getTook().getMillis(),
+                    ingestTookInMillis,
+                    response.getIncrementalState()
+                )
             );
         } else {
             return actionListener.map(response -> {
@@ -141,7 +146,7 @@ ActionListener<BulkResponse> wrapActionListenerIfNeeded(long ingestTookInMillis,
                     assertResponsesAreCorrect(bulkResponses, allResponses);
                 }
 
-                return new BulkResponse(allResponses, response.getTook().getMillis(), ingestTookInMillis);
+                return new BulkResponse(allResponses, response.getTook().getMillis(), ingestTookInMillis, response.getIncrementalState());
             });
         }
     }

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

@@ -8,6 +8,7 @@
 
 package org.elasticsearch.action.bulk;
 
+import org.elasticsearch.TransportVersions;
 import org.elasticsearch.action.ActionResponse;
 import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -36,22 +37,38 @@ public class BulkResponse extends ActionResponse implements Iterable<BulkItemRes
     private final BulkItemResponse[] responses;
     private final long tookInMillis;
     private final long ingestTookInMillis;
+    private final BulkRequest.IncrementalState incrementalState;
 
     public BulkResponse(StreamInput in) throws IOException {
         super(in);
         responses = in.readArray(BulkItemResponse::new, BulkItemResponse[]::new);
         tookInMillis = in.readVLong();
         ingestTookInMillis = in.readZLong();
+        if (in.getTransportVersion().onOrAfter(TransportVersions.BULK_INCREMENTAL_STATE)) {
+            incrementalState = new BulkRequest.IncrementalState(in);
+        } else {
+            incrementalState = BulkRequest.IncrementalState.EMPTY;
+        }
     }
 
     public BulkResponse(BulkItemResponse[] responses, long tookInMillis) {
         this(responses, tookInMillis, NO_INGEST_TOOK);
     }
 
     public BulkResponse(BulkItemResponse[] responses, long tookInMillis, long ingestTookInMillis) {
+        this(responses, tookInMillis, ingestTookInMillis, BulkRequest.IncrementalState.EMPTY);
+    }
+
+    public BulkResponse(
+        BulkItemResponse[] responses,
+        long tookInMillis,
+        long ingestTookInMillis,
+        BulkRequest.IncrementalState incrementalState
+    ) {
         this.responses = responses;
         this.tookInMillis = tookInMillis;
         this.ingestTookInMillis = ingestTookInMillis;
+        this.incrementalState = incrementalState;
     }
 
     /**
@@ -61,6 +78,10 @@ public TimeValue getTook() {
         return new TimeValue(tookInMillis);
     }
 
+    public long getTookInMillis() {
+        return tookInMillis;
+    }
+
     /**
      * If ingest is enabled returns the bulk ingest preprocessing time, otherwise 0 is returned.
      */
@@ -75,6 +96,10 @@ public long getIngestTookInMillis() {
         return ingestTookInMillis;
     }
 
+    BulkRequest.IncrementalState getIncrementalState() {
+        return incrementalState;
+    }
+
     /**
      * Has anything failed with the execution.
      */
@@ -124,6 +149,9 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeArray(responses);
         out.writeVLong(tookInMillis);
         out.writeZLong(ingestTookInMillis);
+        if (out.getTransportVersion().onOrAfter(TransportVersions.BULK_INCREMENTAL_STATE)) {
+            incrementalState.writeTo(out);
+        }
     }
 
     @Override