Reshard add shard count (#133985)

Add a new field to ReplicationRequest to be used once the resharding feature is enabled. This is the reshardShardCountSummary - the effective shard count as seen by the coordination node.

Author: ankikuma

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

@@ -401,13 +401,20 @@ private void executeBulkRequestsByShard(
                 final ShardId shardId = entry.getKey();
                 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());
+                int reshardSplitShardCountSummary = 0;
+                if (indexMetadata != null) {
+                    reshardSplitShardCountSummary = indexMetadata.getReshardSplitShardCountSummaryForIndexing(shardId.getId());
+                }
                 BulkShardRequest bulkShardRequest = new BulkShardRequest(
                     shardId,
+                    reshardSplitShardCountSummary,
                     bulkRequest.getRefreshPolicy(),
                     requests.toArray(new BulkItemRequest[0]),
                     bulkRequest.isSimulated()
                 );
-                var indexMetadata = project.index(shardId.getIndexName());
+
                 if (indexMetadata != null && indexMetadata.getInferenceFields().isEmpty() == false) {
                     bulkShardRequest.setInferenceFieldMap(indexMetadata.getInferenceFields());
                 }

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

@@ -50,12 +50,26 @@ public BulkShardRequest(StreamInput in) throws IOException {
         }
     }
 
+    public BulkShardRequest(ShardId shardId, int reshardSplitShardCountSummary, RefreshPolicy refreshPolicy, BulkItemRequest[] items) {
+        this(shardId, reshardSplitShardCountSummary, refreshPolicy, items, false);
+    }
+
     public BulkShardRequest(ShardId shardId, RefreshPolicy refreshPolicy, BulkItemRequest[] items) {
-        this(shardId, refreshPolicy, items, false);
+        this(shardId, 0, refreshPolicy, items, false);
     }
 
     public BulkShardRequest(ShardId shardId, RefreshPolicy refreshPolicy, BulkItemRequest[] items, boolean isSimulated) {
-        super(shardId);
+        this(shardId, 0, refreshPolicy, items, isSimulated);
+    }
+
+    public BulkShardRequest(
+        ShardId shardId,
+        int reshardSplitShardCountSummary,
+        RefreshPolicy refreshPolicy,
+        BulkItemRequest[] items,
+        boolean isSimulated
+    ) {
+        super(shardId, reshardSplitShardCountSummary);
         this.items = items;
         setRefreshPolicy(refreshPolicy);
         this.isSimulated = isSimulated;

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

@@ -47,6 +47,10 @@ public ReplicatedWriteRequest(@Nullable ShardId shardId) {
         super(shardId);
     }
 
+    public ReplicatedWriteRequest(@Nullable ShardId shardId, int reshardSplitShardCountSummary) {
+        super(shardId, reshardSplitShardCountSummary);
+    }
+
     @Override
     @SuppressWarnings("unchecked")
     public R setRefreshPolicy(RefreshPolicy refreshPolicy) {

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

@@ -9,6 +9,7 @@
 
 package org.elasticsearch.action.support.replication;
 
+import org.elasticsearch.TransportVersion;
 import org.elasticsearch.action.ActionRequestValidationException;
 import org.elasticsearch.action.IndicesRequest;
 import org.elasticsearch.action.LegacyActionRequest;
@@ -37,6 +38,8 @@ public abstract class ReplicationRequest<Request extends ReplicationRequest<Requ
     implements
         IndicesRequest {
 
+    private static final TransportVersion INDEX_RESHARD_SHARDCOUNT_SUMMARY = TransportVersion.fromName("index_reshard_shardcount_summary");
+
     public static final TimeValue DEFAULT_TIMEOUT = TimeValue.timeValueMinutes(1);
 
     /**
@@ -49,6 +52,45 @@ public abstract class ReplicationRequest<Request extends ReplicationRequest<Requ
     protected TimeValue timeout;
     protected String index;
 
+    /**
+     * The reshardSplitShardCountSummary has been added to accommodate the Resharding feature.
+     * This is populated when the coordinator is deciding which shards a request applies to.
+     * For example, {@link org.elasticsearch.action.bulk.BulkOperation} splits
+     * an incoming bulk request into shard level {@link org.elasticsearch.action.bulk.BulkShardRequest}
+     * based on its cluster state view of the number of shards that are ready for indexing.
+     * The purpose of this metadata is to reconcile the cluster state visible at the coordinating
+     * node with that visible at the source shard node. (w.r.t resharding).
+     * When an index is being split, there is a point in time when the newly created shard (target shard)
+     * takes over its portion of the document space from the original shard (source shard).
+     * Although the handoff is atomic at the original (source shard) and new shards (target shard),
+     * there is a window of time between the coordinating node creating a shard request and the shard receiving and processing it.
+     * This field is used by the original shard (source shard) when it processes the request to detect whether
+     * the coordinator's view of the new shard's state when it created the request matches the shard's current state,
+     * or whether the request must be reprocessed taking into account the current shard states.
+     *
+     * Note that we are able to get away with a single number, instead of an array of target shard states,
+     * because we only allow splits in increments of 2x.
+     *
+     * Example 1:
+     * Suppose we are resharding an index from 2 -> 4 shards. While splitting a bulk request, the coordinator observes
+     * that target shards are not ready for indexing. So requests that are meant for shard 0 and 2 are bundled together,
+     * sent to shard 0 with “reshardSplitShardCountSummary” 2 in the request.
+     * Requests that are meant for shard 1 and 3 are bundled together,
+     * sent to shard 1 with “reshardSplitShardCountSummary” 2 in the request.
+     *
+     * Example 2:
+     * Suppose we are resharding an index from 4 -> 8 shards. While splitting a bulk request, the coordinator observes
+     * that source shard 0 has completed HANDOFF but source shards 1, 2, 3 have not completed handoff.
+     * So, the shard-bulk-request it sends to shard 0 and 4 has the "reshardSplitShardCountSummary" 8,
+     * while the shard-bulk-request it sends to shard 1,2,3 has the "reshardSplitShardCountSummary" 4.
+     * Note that in this case no shard-bulk-request is sent to shards 5, 6, 7 and the requests that were meant for these target shards
+     * are bundled together with and sent to their source shards.
+     *
+     * A value of 0 indicates an INVALID reshardSplitShardCountSummary. Hence, a request with INVALID reshardSplitShardCountSummary
+     * will be treated as a Summary mismatch on the source shard node.
+     */
+    protected final int reshardSplitShardCountSummary;
+
     /**
      * The number of shard copies that must be active before proceeding with the replication action.
      */
@@ -61,6 +103,10 @@ public ReplicationRequest(StreamInput in) throws IOException {
     }
 
     public ReplicationRequest(@Nullable ShardId shardId, StreamInput in) throws IOException {
+        this(shardId, 0, in);
+    }
+
+    public ReplicationRequest(@Nullable ShardId shardId, int reshardSplitShardCountSummary, StreamInput in) throws IOException {
         super(in);
         final boolean thinRead = shardId != null;
         if (thinRead) {
@@ -80,15 +126,32 @@ public ReplicationRequest(@Nullable ShardId shardId, StreamInput in) throws IOEx
             index = in.readString();
         }
         routedBasedOnClusterVersion = in.readVLong();
+        if (in.getTransportVersion().supports(INDEX_RESHARD_SHARDCOUNT_SUMMARY)) {
+            if (thinRead) {
+                this.reshardSplitShardCountSummary = reshardSplitShardCountSummary;
+            } else {
+                this.reshardSplitShardCountSummary = in.readInt();
+            }
+        } else {
+            this.reshardSplitShardCountSummary = 0;
+        }
     }
 
     /**
      * Creates a new request with resolved shard id
      */
     public ReplicationRequest(@Nullable ShardId shardId) {
+        this(shardId, 0);
+    }
+
+    /**
+     * Creates a new request with resolved shard id and reshardSplitShardCountSummary
+     */
+    public ReplicationRequest(@Nullable ShardId shardId, int reshardSplitShardCountSummary) {
         this.index = shardId == null ? null : shardId.getIndexName();
         this.shardId = shardId;
         this.timeout = DEFAULT_TIMEOUT;
+        this.reshardSplitShardCountSummary = reshardSplitShardCountSummary;
     }
 
     /**
@@ -137,6 +200,14 @@ public ShardId shardId() {
         return shardId;
     }
 
+    /**
+     * @return The effective shard count as seen by the coordinator when creating this request.
+     * can be 0 if this has not yet been resolved.
+     */
+    public int reshardSplitShardCountSummary() {
+        return reshardSplitShardCountSummary;
+    }
+
     /**
      * Sets the number of shard copies that must be active before proceeding with the replication
      * operation. Defaults to {@link ActiveShardCount#DEFAULT}, which requires one shard copy
@@ -191,11 +262,14 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeTimeValue(timeout);
         out.writeString(index);
         out.writeVLong(routedBasedOnClusterVersion);
+        if (out.getTransportVersion().supports(INDEX_RESHARD_SHARDCOUNT_SUMMARY)) {
+            out.writeInt(reshardSplitShardCountSummary);
+        }
     }
 
     /**
      * Thin serialization that does not write {@link #shardId} and will only write {@link #index} if it is different from the index name in
-     * {@link #shardId}.
+     * {@link #shardId}. Since we do not write {@link #shardId}, we also do not write {@link #reshardSplitShardCountSummary}.
      */
     public void writeThin(StreamOutput out) throws IOException {
         super.writeTo(out);

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

@@ -459,14 +459,21 @@ void runWithPrimaryShardReference(final PrimaryShardReference primaryShardRefere
             try {
                 final ClusterState clusterState = clusterService.state();
                 final Index index = primaryShardReference.routingEntry().index();
-                final ProjectId projectId = clusterState.metadata().projectFor(index).id();
+                final ProjectMetadata project = clusterState.metadata().projectFor(index);
+                final ProjectId projectId = project.id();
+                final IndexMetadata indexMetadata = project.index(index);
 
                 final ClusterBlockException blockException = blockExceptions(clusterState, projectId, index.getName());
                 if (blockException != null) {
                     logger.trace("cluster is blocked, action failed on primary", blockException);
                     throw blockException;
                 }
 
+                int reshardSplitShardCountSummary = primaryRequest.getRequest().reshardSplitShardCountSummary();
+                assert (reshardSplitShardCountSummary == 0
+                    || reshardSplitShardCountSummary == indexMetadata.getReshardSplitShardCountSummaryForIndexing(
+                        primaryRequest.getRequest().shardId().getId()
+                    ));
                 if (primaryShardReference.isRelocated()) {
                     primaryShardReference.close(); // release shard operation lock as soon as possible
                     setPhase(replicationTask, "primary_delegation");

server/src/main/java/org/elasticsearch/cluster/metadata/IndexMetadata.java

@@ -1181,6 +1181,71 @@ public int getNumberOfShards() {
         return numberOfShards;
     }
 
+    /**
+     * This method is used in the context of the resharding feature.
+     * Given a {@code shardId} and {@code minShardState} i.e. the minimum target shard state required for
+     * an operation to be routed to target shards,
+     * this method returns the "effective" shard count as seen by this IndexMetadata.
+     *
+     * The reshardSplitShardCountSummary tells us whether the coordinator routed requests to the source shard or
+     * to both source and target shards. Requests are routed to both source and target shards
+     * once the target shards are ready for an operation.
+     *
+     * The coordinator routes requests to source and target shards, based on its cluster state view of the state of shards
+     * undergoing a resharding operation. This method is used to populate a field in the shard level requests sent to
+     * source and target shards, as a proxy for the cluster state version. The same calculation is then done at the source shard
+     * to verify if the coordinator and source node's view of the resharding state have a mismatch.
+     * See {@link org.elasticsearch.action.support.replication.ReplicationRequest#reshardSplitShardCountSummary}
+     * for a detailed description of how this value is used.
+     *
+     * @param shardId  Input shardId for which we want to calculate the effective shard count
+     * @param minShardState Minimum target shard state required for the target to be considered ready
+     * @return Effective shard count as seen by an operation using this IndexMetadata
+     */
+    private int getReshardSplitShardCountSummary(int shardId, IndexReshardingState.Split.TargetShardState minShardState) {
+        assert shardId >= 0 && shardId < getNumberOfShards() : "shardId is out of bounds";
+        int shardCount = getNumberOfShards();
+        if (reshardingMetadata != null) {
+            if (reshardingMetadata.getSplit().isTargetShard(shardId)) {
+                int sourceShardId = reshardingMetadata.getSplit().sourceShard(shardId);
+                // Requests cannot be routed to target shards until they are ready
+                assert reshardingMetadata.getSplit().allTargetStatesAtLeast(sourceShardId, minShardState) : "unexpected target state";
+                shardCount = reshardingMetadata.getSplit().shardCountAfter();
+            } else if (reshardingMetadata.getSplit().isSourceShard(shardId)) {
+                if (reshardingMetadata.getSplit().allTargetStatesAtLeast(shardId, minShardState)) {
+                    shardCount = reshardingMetadata.getSplit().shardCountAfter();
+                } else {
+                    shardCount = reshardingMetadata.getSplit().shardCountBefore();
+                }
+            }
+        }
+        return shardCount;
+    }
+
+    /**
+     * This method is used in the context of the resharding feature.
+     * Given a {@code shardId}, this method returns the "effective" shard count
+     * as seen by this IndexMetadata, for indexing operations.
+     *
+     * See {@code getReshardSplitShardCountSummary} for more details.
+     * @param shardId  Input shardId for which we want to calculate the effective shard count
+     */
+    public int getReshardSplitShardCountSummaryForIndexing(int shardId) {
+        return (getReshardSplitShardCountSummary(shardId, IndexReshardingState.Split.TargetShardState.HANDOFF));
+    }
+
+    /**
+     * This method is used in the context of the resharding feature.
+     * Given a {@code shardId}, this method returns the "effective" shard count
+     * as seen by this IndexMetadata, for search operations.
+     *
+     * See {@code getReshardSplitShardCount} for more details.
+     * @param shardId  Input shardId for which we want to calculate the effective shard count
+     */
+    public int getReshardSplitShardCountSummaryForSearch(int shardId) {
+        return (getReshardSplitShardCountSummary(shardId, IndexReshardingState.Split.TargetShardState.SPLIT));
+    }
+
     public int getNumberOfReplicas() {
         return numberOfReplicas;
     }
@@ -3022,7 +3087,7 @@ public IndexMetadata fromXContent(XContentParser parser) throws IOException {
      * Returns the number of shards that should be used for routing. This basically defines the hash space we use in
      * {@link IndexRouting#indexShard} to route documents
      * to shards based on their ID or their specific routing value. The default value is {@link #getNumberOfShards()}. This value only
-     * changes if and index is shrunk.
+     * changes if an index is shrunk.
      */
     public int getRoutingNumShards() {
         return routingNumShards;
@@ -3042,7 +3107,7 @@ public int getRoutingFactor() {
      * @param shardId the id of the target shard to split into
      * @param sourceIndexMetadata the source index metadata
      * @param numTargetShards the total number of shards in the target index
-     * @return a the source shard ID to split off from
+     * @return the source shard ID to split off from
      */
     public static ShardId selectSplitShard(int shardId, IndexMetadata sourceIndexMetadata, int numTargetShards) {
         int numSourceShards = sourceIndexMetadata.getNumberOfShards();

server/src/main/java/org/elasticsearch/cluster/metadata/IndexReshardingState.java

@@ -378,6 +378,16 @@ public boolean targetStateAtLeast(int shardNum, TargetShardState targetShardStat
             return getTargetShardState(shardNum).ordinal() >= targetShardState.ordinal();
         }
 
+        public boolean allTargetStatesAtLeast(int sourceShardId, TargetShardState targetShardState) {
+            var targets = getTargetStatesFor(sourceShardId);
+            for (TargetShardState state : targets) {
+                if (state.ordinal() < targetShardState.ordinal()) {
+                    return false;
+                }
+            }
+            return true;
+        }
+
         public Stream<TargetShardState> targetStates() {
             return Arrays.stream(targetShards);
         }

server/src/main/java/org/elasticsearch/cluster/routing/IndexRouting.java

@@ -72,6 +72,7 @@ public static IndexRouting fromIndexMetadata(IndexMetadata metadata) {
     protected final String indexName;
     private final int routingNumShards;
     private final int routingFactor;
+    @Nullable
     private final IndexReshardingMetadata indexReshardingMetadata;
 
     private IndexRouting(IndexMetadata metadata) {

server/src/main/resources/transport/definitions/referable/index_reshard_shardcount_summary.csv

@@ -0,0 +1 @@
+9172000

server/src/main/resources/transport/upper_bounds/9.2.csv

@@ -1 +1 @@
-esql_fuse_linear_operator_status,9171000
+index_reshard_shardcount_summary,9172000