Resharding Add Flush and Refresh functionality (#136880)

Add functionality to forward flush and refresh requests to target shards during a reshard operation.

Author: ankikuma

server/src/main/java/org/elasticsearch/action/admin/indices/flush/ShardFlushRequest.java

@@ -11,6 +11,7 @@
 
 import org.elasticsearch.action.support.ActiveShardCount;
 import org.elasticsearch.action.support.replication.ReplicationRequest;
+import org.elasticsearch.cluster.routing.SplitShardCountSummary;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.index.shard.ShardId;
@@ -21,8 +22,13 @@ public class ShardFlushRequest extends ReplicationRequest<ShardFlushRequest> {
 
     private final FlushRequest request;
 
-    public ShardFlushRequest(FlushRequest request, ShardId shardId) {
-        super(shardId);
+    /**
+     * Creates a request for a resolved shard id and SplitShardCountSummary (used
+     * to determine if the request needs to be executed on a split shard not yet seen by the
+     * coordinator that sent the request)
+     */
+    public ShardFlushRequest(FlushRequest request, ShardId shardId, SplitShardCountSummary reshardSplitShardCountSummary) {
+        super(shardId, reshardSplitShardCountSummary);
         this.request = request;
         this.waitForActiveShards = ActiveShardCount.NONE; // don't wait for any active shards before proceeding, by default
     }

server/src/main/java/org/elasticsearch/action/admin/indices/flush/TransportFlushAction.java

@@ -18,6 +18,7 @@
 import org.elasticsearch.client.internal.node.NodeClient;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
 import org.elasticsearch.cluster.project.ProjectResolver;
+import org.elasticsearch.cluster.routing.SplitShardCountSummary;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.injection.guice.Inject;
@@ -59,8 +60,8 @@ public TransportFlushAction(
     }
 
     @Override
-    protected ShardFlushRequest newShardRequest(FlushRequest request, ShardId shardId) {
-        return new ShardFlushRequest(request, shardId);
+    protected ShardFlushRequest newShardRequest(FlushRequest request, ShardId shardId, SplitShardCountSummary shardCountSummary) {
+        return new ShardFlushRequest(request, shardId, shardCountSummary);
     }
 
     @Override

server/src/main/java/org/elasticsearch/action/admin/indices/flush/TransportShardFlushAction.java

@@ -13,13 +13,16 @@
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionType;
 import org.elasticsearch.action.support.ActionFilters;
+import org.elasticsearch.action.support.replication.ReplicationRequestSplitHelper;
 import org.elasticsearch.action.support.replication.ReplicationResponse;
 import org.elasticsearch.action.support.replication.TransportReplicationAction;
 import org.elasticsearch.cluster.action.shard.ShardStateAction;
+import org.elasticsearch.cluster.project.ProjectResolver;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.core.Tuple;
 import org.elasticsearch.index.shard.IndexShard;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.indices.IndicesService;
@@ -32,12 +35,15 @@
 import org.elasticsearch.transport.TransportService;
 
 import java.io.IOException;
+import java.util.Map;
 
 public class TransportShardFlushAction extends TransportReplicationAction<ShardFlushRequest, ShardFlushRequest, ReplicationResponse> {
 
     public static final String NAME = FlushAction.NAME + "[s]";
     public static final ActionType<ReplicationResponse> TYPE = new ActionType<>(NAME);
 
+    private final ProjectResolver projectResolver;
+
     @Inject
     public TransportShardFlushAction(
         Settings settings,
@@ -46,7 +52,8 @@ public TransportShardFlushAction(
         IndicesService indicesService,
         ThreadPool threadPool,
         ShardStateAction shardStateAction,
-        ActionFilters actionFilters
+        ActionFilters actionFilters,
+        ProjectResolver projectResolver
     ) {
         super(
             settings,
@@ -64,6 +71,7 @@ public TransportShardFlushAction(
             PrimaryActionExecution.RejectOnOverload,
             ReplicaActionExecution.SubjectToCircuitBreaker
         );
+        this.projectResolver = projectResolver;
         transportService.registerRequestHandler(
             PRE_SYNCED_FLUSH_ACTION_NAME,
             threadPool.executor(ThreadPool.Names.FLUSH),
@@ -89,6 +97,27 @@ protected void shardOperationOnPrimary(
         }));
     }
 
+    // We are here because there was a mismatch between the SplitShardCountSummary in the request
+    // and that on the primary shard node. We assume that the request is exactly 1 reshard split behind
+    // the current state.
+    @Override
+    protected Map<ShardId, ShardFlushRequest> splitRequestOnPrimary(ShardFlushRequest request) {
+        return ReplicationRequestSplitHelper.splitRequest(
+            request,
+            projectResolver.getProjectMetadata(clusterService.state()),
+            (targetShard, shardCountSummary) -> new ShardFlushRequest(request.getRequest(), targetShard, shardCountSummary)
+        );
+    }
+
+    @Override
+    protected Tuple<ReplicationResponse, Exception> combineSplitResponses(
+        ShardFlushRequest originalRequest,
+        Map<ShardId, ShardFlushRequest> splitRequests,
+        Map<ShardId, Tuple<ReplicationResponse, Exception>> responses
+    ) {
+        return ReplicationRequestSplitHelper.combineSplitResponses(originalRequest, splitRequests, responses);
+    }
+
     @Override
     protected void shardOperationOnReplica(ShardFlushRequest request, IndexShard replica, ActionListener<ReplicaResult> listener) {
         replica.flush(request.getRequest(), listener.map(flushed -> {

server/src/main/java/org/elasticsearch/action/admin/indices/refresh/TransportRefreshAction.java

@@ -19,6 +19,7 @@
 import org.elasticsearch.client.internal.node.NodeClient;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
 import org.elasticsearch.cluster.project.ProjectResolver;
+import org.elasticsearch.cluster.routing.SplitShardCountSummary;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.injection.guice.Inject;
@@ -60,8 +61,8 @@ public TransportRefreshAction(
     }
 
     @Override
-    protected BasicReplicationRequest newShardRequest(RefreshRequest request, ShardId shardId) {
-        BasicReplicationRequest replicationRequest = new BasicReplicationRequest(shardId);
+    protected BasicReplicationRequest newShardRequest(RefreshRequest request, ShardId shardId, SplitShardCountSummary shardCountSummary) {
+        BasicReplicationRequest replicationRequest = new BasicReplicationRequest(shardId, shardCountSummary);
         replicationRequest.waitForActiveShards(ActiveShardCount.NONE);
         return replicationRequest;
     }

server/src/main/java/org/elasticsearch/action/admin/indices/refresh/TransportShardRefreshAction.java

@@ -16,14 +16,18 @@
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.replication.BasicReplicationRequest;
 import org.elasticsearch.action.support.replication.ReplicationOperation;
+import org.elasticsearch.action.support.replication.ReplicationRequestSplitHelper;
 import org.elasticsearch.action.support.replication.ReplicationResponse;
 import org.elasticsearch.action.support.replication.TransportReplicationAction;
 import org.elasticsearch.cluster.action.shard.ShardStateAction;
+import org.elasticsearch.cluster.project.ProjectResolver;
 import org.elasticsearch.cluster.routing.IndexShardRoutingTable;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.core.Tuple;
 import org.elasticsearch.index.shard.IndexShard;
+import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.indices.IndicesService;
 import org.elasticsearch.injection.guice.Inject;
 import org.elasticsearch.logging.LogManager;
@@ -32,6 +36,7 @@
 import org.elasticsearch.transport.TransportService;
 
 import java.io.IOException;
+import java.util.Map;
 import java.util.concurrent.Executor;
 
 public class TransportShardRefreshAction extends TransportReplicationAction<
@@ -46,6 +51,7 @@ public class TransportShardRefreshAction extends TransportReplicationAction<
     public static final String SOURCE_API = "api";
 
     private final Executor refreshExecutor;
+    private final ProjectResolver projectResolver;
 
     @Inject
     public TransportShardRefreshAction(
@@ -55,7 +61,8 @@ public TransportShardRefreshAction(
         IndicesService indicesService,
         ThreadPool threadPool,
         ShardStateAction shardStateAction,
-        ActionFilters actionFilters
+        ActionFilters actionFilters,
+        ProjectResolver projectResolver
     ) {
         super(
             settings,
@@ -73,6 +80,7 @@ public TransportShardRefreshAction(
             PrimaryActionExecution.RejectOnOverload,
             ReplicaActionExecution.SubjectToCircuitBreaker
         );
+        this.projectResolver = projectResolver;
         // registers the unpromotable version of shard refresh action
         new TransportUnpromotableShardRefreshAction(
             clusterService,
@@ -104,6 +112,27 @@ protected void shardOperationOnPrimary(
         }));
     }
 
+    // We are here because there was mismatch between the SplitShardCountSummary in the request
+    // and that on the primary shard node. We assume that the request is exactly 1 reshard split behind
+    // the current state.
+    @Override
+    protected Map<ShardId, BasicReplicationRequest> splitRequestOnPrimary(BasicReplicationRequest request) {
+        return ReplicationRequestSplitHelper.splitRequest(
+            request,
+            projectResolver.getProjectMetadata(clusterService.state()),
+            (targetShard, shardCountSummary) -> new BasicReplicationRequest(targetShard, shardCountSummary)
+        );
+    }
+
+    @Override
+    protected Tuple<ReplicationResponse, Exception> combineSplitResponses(
+        BasicReplicationRequest originalRequest,
+        Map<ShardId, BasicReplicationRequest> splitRequests,
+        Map<ShardId, Tuple<ReplicationResponse, Exception>> responses
+    ) {
+        return ReplicationRequestSplitHelper.combineSplitResponses(originalRequest, splitRequests, responses);
+    }
+
     @Override
     protected void shardOperationOnReplica(ShardRefreshReplicaRequest request, IndexShard replica, ActionListener<ReplicaResult> listener) {
         replica.externalRefresh(SOURCE_API, listener.safeMap(refreshResult -> {

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

@@ -403,11 +403,9 @@ private void executeBulkRequestsByShard(
                 final List<BulkItemRequest> requests = entry.getValue();
 
                 // Get effective shardCount for shardId and pass it on as parameter to new BulkShardRequest
-                var indexMetadata = project.index(shardId.getIndexName());
-                SplitShardCountSummary reshardSplitShardCountSummary = SplitShardCountSummary.UNSET;
-                if (indexMetadata != null) {
-                    reshardSplitShardCountSummary = SplitShardCountSummary.forIndexing(indexMetadata, shardId.getId());
-                }
+                var indexMetadata = project.getIndexSafe(shardId.getIndex());
+                SplitShardCountSummary reshardSplitShardCountSummary = SplitShardCountSummary.forIndexing(indexMetadata, shardId.getId());
+
                 BulkShardRequest bulkShardRequest = new BulkShardRequest(
                     shardId,
                     reshardSplitShardCountSummary,
@@ -416,7 +414,7 @@ private void executeBulkRequestsByShard(
                     bulkRequest.isSimulated()
                 );
 
-                if (indexMetadata != null && indexMetadata.getInferenceFields().isEmpty() == false) {
+                if (indexMetadata.getInferenceFields().isEmpty() == false) {
                     bulkShardRequest.setInferenceFieldMap(indexMetadata.getInferenceFields());
                 }
                 bulkShardRequest.waitForActiveShards(bulkRequest.waitForActiveShards());

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

@@ -10,8 +10,10 @@
 package org.elasticsearch.action.bulk;
 
 import org.elasticsearch.action.DocWriteRequest;
+import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.metadata.ProjectMetadata;
 import org.elasticsearch.cluster.routing.IndexRouting;
+import org.elasticsearch.cluster.routing.SplitShardCountSummary;
 import org.elasticsearch.core.Tuple;
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.shard.ShardId;
@@ -33,7 +35,9 @@ private ShardBulkSplitHelper() {}
     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));
+        IndexMetadata indexMetadata = project.getIndexSafe(index);
+        IndexRouting indexRouting = IndexRouting.fromIndexMetadata(indexMetadata);
+        SplitShardCountSummary shardCountSummary = SplitShardCountSummary.forIndexing(indexMetadata, request.shardId().getId());
 
         Map<ShardId, List<BulkItemRequest>> requestsByShard = new HashMap<>();
         Map<ShardId, BulkShardRequest> bulkRequestsPerShard = new HashMap<>();
@@ -57,17 +61,26 @@ public static Map<ShardId, BulkShardRequest> splitRequests(BulkShardRequest requ
 
         // 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.
+            // Return the original request if no items were split to target. Note that
+            // this original request still contains the stale SplitShardCountSummary.
+            // This is alright because we hold primary indexing permits while calling this split
+            // method and we execute this request on the primary without letting go of the indexing permits.
+            // This means that a second split cannot occur in the meantime.
             if (requestsByShard.containsKey(sourceShardId)) {
                 return Map.of(sourceShardId, request);
             }
         }
 
+        // Create a new BulkShardRequest(s) with the updated SplitShardCountSummary. This is because
+        // we do not hold primary permits on the target shard, and hence it can proceed with
+        // a second split operation while this request is still pending. We must verify the
+        // SplitShardCountSummary again on the target.
         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,
+                shardCountSummary,
                 request.getRefreshPolicy(),
                 requests.toArray(new BulkItemRequest[0]),
                 request.isSimulated()

server/src/main/java/org/elasticsearch/action/get/TransportShardMultiGetAction.java

@@ -174,6 +174,7 @@ private void handleMultiGetOnUnpromotableShard(
         ShardId shardId = indexShard.shardId();
         if (request.refresh()) {
             logger.trace("send refresh action for shard {}", shardId);
+            // TODO: Do we need to pass in shardCountSummary here ?
             var refreshRequest = new BasicReplicationRequest(shardId);
             refreshRequest.setParentTask(request.getParentTask());
             client.executeLocally(

server/src/main/java/org/elasticsearch/action/support/replication/BasicReplicationRequest.java

@@ -9,6 +9,7 @@
 
 package org.elasticsearch.action.support.replication;
 
+import org.elasticsearch.cluster.routing.SplitShardCountSummary;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.index.shard.ShardId;
 
@@ -24,10 +25,20 @@ public class BasicReplicationRequest extends ReplicationRequest<BasicReplication
     /**
      * Creates a new request with resolved shard id
      */
+    // TODO: Check if callers of this need to be modified to pass in shardCountSummary and eventually remove this constructor
     public BasicReplicationRequest(ShardId shardId) {
         super(shardId);
     }
 
+    /**
+     * Creates a new request with resolved shard id and SplitShardCountSummary (used
+     * to determine if the request needs to be executed on a split shard not yet seen by the
+     * coordinator that sent the request)
+     */
+    public BasicReplicationRequest(ShardId shardId, SplitShardCountSummary reshardSplitShardCountSummary) {
+        super(shardId, reshardSplitShardCountSummary);
+    }
+
     public BasicReplicationRequest(StreamInput in) throws IOException {
         super(in);
     }