Add support for delegating write to split target (#136241)

This commit adds the logic to delegate bulk shard requests to the split
target when a primary receives a request from a stale coordinator.

Author: Tim-Brooks

server/src/main/java/org/elasticsearch/action/DocWriteRequest.java

@@ -176,6 +176,11 @@ default void postRoutingProcess(IndexRouting indexRouting) {}
      */
     int route(IndexRouting indexRouting);
 
+    /**
+     * Pick the appropriate target shard id this request should be routed to during resharding.
+     */
+    int rerouteAtSourceDuringResharding(IndexRouting indexRouting);
+
     /**
      * Resolves the write index that should receive this request
      * based on the provided index abstraction.

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

@@ -0,0 +1,108 @@
+/*
+ * 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.action.bulk;
+
+import org.elasticsearch.action.DocWriteRequest;
+import org.elasticsearch.cluster.metadata.ProjectMetadata;
+import org.elasticsearch.cluster.routing.IndexRouting;
+import org.elasticsearch.core.Tuple;
+import org.elasticsearch.index.Index;
+import org.elasticsearch.index.shard.ShardId;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public final class ShardBulkSplitHelper {
+
+    private ShardBulkSplitHelper() {}
+
+    /**
+     * Splits a bulk request into multiple requests for each shard. If the items in the request only route to the source shard it will
+     * return the original request. If the items only route to the target shard it will return a map with one request. If the requests
+     * route to both the map will have a request for each shard.
+     */
+    public static Map<ShardId, BulkShardRequest> splitRequests(BulkShardRequest request, ProjectMetadata project) {
+        final ShardId sourceShardId = request.shardId();
+        final Index index = sourceShardId.getIndex();
+        IndexRouting indexRouting = IndexRouting.fromIndexMetadata(project.getIndexSafe(index));
+
+        Map<ShardId, List<BulkItemRequest>> requestsByShard = new HashMap<>();
+        Map<ShardId, BulkShardRequest> bulkRequestsPerShard = new HashMap<>();
+
+        // Iterate through the items in the input request and split them based on the
+        // current resharding-split state.
+        BulkItemRequest[] items = request.items();
+        if (items.length == 0) {  // Nothing to split
+            return Map.of(sourceShardId, request);
+        }
+
+        for (BulkItemRequest bulkItemRequest : items) {
+            DocWriteRequest<?> docWriteRequest = bulkItemRequest.request();
+            int newShardId = docWriteRequest.rerouteAtSourceDuringResharding(indexRouting);
+            List<BulkItemRequest> shardRequests = requestsByShard.computeIfAbsent(
+                new ShardId(index, newShardId),
+                shardNum -> new ArrayList<>()
+            );
+            shardRequests.add(new BulkItemRequest(bulkItemRequest.id(), bulkItemRequest.request()));
+        }
+
+        // All items belong to either the source shard or target shard.
+        if (requestsByShard.size() == 1) {
+            // Return the original request if no items were split to target.
+            if (requestsByShard.containsKey(sourceShardId)) {
+                return Map.of(sourceShardId, request);
+            }
+        }
+
+        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,
+                request.getRefreshPolicy(),
+                requests.toArray(new BulkItemRequest[0]),
+                request.isSimulated()
+            );
+            bulkRequestsPerShard.put(shardId, bulkShardRequest);
+        }
+        return bulkRequestsPerShard;
+    }
+
+    public static Tuple<BulkShardResponse, Exception> combineResponses(
+        BulkShardRequest originalRequest,
+        Map<ShardId, BulkShardRequest> splitRequests,
+        Map<ShardId, Tuple<BulkShardResponse, Exception>> responses
+    ) {
+        BulkItemResponse[] bulkItemResponses = new BulkItemResponse[originalRequest.items().length];
+        for (Map.Entry<ShardId, Tuple<BulkShardResponse, Exception>> entry : responses.entrySet()) {
+            ShardId shardId = entry.getKey();
+            Tuple<BulkShardResponse, Exception> value = entry.getValue();
+            Exception exception = value.v2();
+            if (exception != null) {
+                BulkShardRequest bulkShardRequest = splitRequests.get(shardId);
+                for (BulkItemRequest item : bulkShardRequest.items()) {
+                    DocWriteRequest<?> request = item.request();
+                    BulkItemResponse.Failure failure = new BulkItemResponse.Failure(item.index(), request.id(), exception);
+                    bulkItemResponses[item.id()] = BulkItemResponse.failure(item.id(), request.opType(), failure);
+                }
+            } else {
+                for (BulkItemResponse bulkItemResponse : value.v1().getResponses()) {
+                    bulkItemResponses[bulkItemResponse.getItemId()] = bulkItemResponse;
+                }
+            }
+        }
+        BulkShardResponse bulkShardResponse = new BulkShardResponse(originalRequest.shardId(), bulkItemResponses);
+        // TODO: Decide how to handle
+        bulkShardResponse.setShardInfo(responses.get(originalRequest.shardId()).v1().getShardInfo());
+        return new Tuple<>(bulkShardResponse, null);
+    }
+}

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

@@ -55,6 +55,7 @@
 import org.elasticsearch.index.mapper.SourceToParse;
 import org.elasticsearch.index.seqno.SequenceNumbers;
 import org.elasticsearch.index.shard.IndexShard;
+import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.index.translog.Translog;
 import org.elasticsearch.indices.ExecutorSelector;
 import org.elasticsearch.indices.IndicesService;
@@ -163,6 +164,20 @@ protected void shardOperationOnPrimary(
         primary.ensureMutable(listener.delegateFailure((l, ignored) -> super.shardOperationOnPrimary(request, primary, l)), true);
     }
 
+    @Override
+    protected Map<ShardId, BulkShardRequest> splitRequestOnPrimary(BulkShardRequest request) {
+        return ShardBulkSplitHelper.splitRequests(request, projectResolver.getProjectMetadata(clusterService.state()));
+    }
+
+    @Override
+    protected Tuple<BulkShardResponse, Exception> combineSplitResponses(
+        BulkShardRequest originalRequest,
+        Map<ShardId, BulkShardRequest> splitRequests,
+        Map<ShardId, Tuple<BulkShardResponse, Exception>> responses
+    ) {
+        return ShardBulkSplitHelper.combineResponses(originalRequest, splitRequests, responses);
+    }
+
     @Override
     protected void dispatchedShardOperationOnPrimary(
         BulkShardRequest request,

server/src/main/java/org/elasticsearch/action/delete/DeleteRequest.java

@@ -236,6 +236,11 @@ public int route(IndexRouting indexRouting) {
         return indexRouting.deleteShard(id, routing);
     }
 
+    @Override
+    public int rerouteAtSourceDuringResharding(IndexRouting indexRouting) {
+        return indexRouting.deleteShard(id, routing);
+    }
+
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);

server/src/main/java/org/elasticsearch/action/index/IndexRequest.java

@@ -928,6 +928,11 @@ public int route(IndexRouting indexRouting) {
         return indexRouting.indexShard(this);
     }
 
+    @Override
+    public int rerouteAtSourceDuringResharding(IndexRouting indexRouting) {
+        return indexRouting.rerouteToTarget(this);
+    }
+
     public IndexRequest setRequireAlias(boolean requireAlias) {
         this.requireAlias = requireAlias;
         return this;

server/src/main/java/org/elasticsearch/action/support/master/TermOverridingMasterNodeRequest.java

@@ -33,7 +33,7 @@
  * which it may be necessary to use the test utility {@code MasterNodeRequestHelper#unwrapTermOverride} to remove the wrapper and access the
  * inner request.
  */
-class TermOverridingMasterNodeRequest extends AbstractTransportRequest {
+public class TermOverridingMasterNodeRequest extends AbstractTransportRequest {
 
     private static final Logger logger = LogManager.getLogger(TermOverridingMasterNodeRequest.class);