Adjust method to transition split state

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

@@ -12,6 +12,7 @@
 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.index.shard.ShardId;
 import org.elasticsearch.xcontent.ConstructingObjectParser;
 import org.elasticsearch.xcontent.ParseField;
@@ -20,7 +21,6 @@
 import org.elasticsearch.xcontent.XContentParser;
 
 import java.io.IOException;
-import java.util.Arrays;
 import java.util.Objects;
 
 /**
@@ -207,16 +207,18 @@ public static IndexReshardingMetadata newSplitByMultiple(int shardCount, int mul
         return new IndexReshardingMetadata(IndexReshardingState.Split.newSplitByMultiple(shardCount, multiple));
     }
 
-    public IndexReshardingMetadata transitionSplitTargetToHandoff(ShardId shardId) {
+    public static boolean isSplitTarget(ShardId shardId, @Nullable IndexReshardingMetadata reshardingMetadata) {
+        return reshardingMetadata != null && reshardingMetadata.isSplit() && reshardingMetadata.getSplit().isTargetShard(shardId.id());
+    }
+
+    public IndexReshardingMetadata transitionSplitTargetToNewState(
+        ShardId shardId,
+        IndexReshardingState.Split.TargetShardState newTargetState
+    ) {
         assert state instanceof IndexReshardingState.Split;
-        IndexReshardingState.Split splitState = (IndexReshardingState.Split) state;
-        IndexReshardingState.Split.TargetShardState[] newTargets = Arrays.copyOf(
-            splitState.targetShards(),
-            splitState.targetShards().length
-        );
-        int i = shardId.getId() - state.shardCountBefore();
-        newTargets[i] = IndexReshardingState.Split.TargetShardState.HANDOFF;
-        return new IndexReshardingMetadata(new IndexReshardingState.Split(splitState.sourceShards(), newTargets));
+        IndexReshardingState.Split.Builder builder = new IndexReshardingState.Split.Builder((IndexReshardingState.Split) state);
+        builder.setTargetShardState(shardId.getId(), newTargetState);
+        return new IndexReshardingMetadata(builder.build());
     }
 
     /**

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

@@ -23,6 +23,7 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
+import java.util.stream.Stream;
 
 /**
  * IndexReshardingState is an abstract class holding the persistent state of a generic resharding operation. It contains
@@ -366,6 +367,10 @@ public TargetShardState getTargetShardState(int shardNum) {
             return targetShards[targetShardNum];
         }
 
+        public boolean targetStateAtLeast(int shardNum, TargetShardState targetShardState) {
+            return getTargetShardState(shardNum).ordinal() >= targetShardState.ordinal();
+        }
+
         /**
          * Check whether this metadata represents an incomplete split
          * @return true if the split is incomplete (not all source shards are DONE)
@@ -380,6 +385,10 @@ public boolean inProgress() {
             return false;
         }
 
+        public Stream<TargetShardState> targetStates() {
+            return Arrays.stream(targetShards);
+        }
+
         /**
          * Check whether all target shards for the given source shard are done.
          * @param shardNum a source shard index greater than or equal to 0 and less than the original shard count

server/src/test/java/org/elasticsearch/cluster/metadata/IndexReshardingMetadataTests.java

@@ -21,13 +21,15 @@ public void testSplit() {
         var split = metadata.getSplit();
 
         // starting state is as expected
-        assert metadata.shardCountBefore() == numShards;
-        assert metadata.shardCountAfter() == numShards * multiple;
+        assertEquals(numShards, metadata.shardCountBefore());
+        assertEquals(numShards * multiple, metadata.shardCountAfter());
         for (int i = 0; i < numShards; i++) {
-            assert split.getSourceShardState(i) == IndexReshardingState.Split.SourceShardState.SOURCE;
+            assertSame(IndexReshardingState.Split.SourceShardState.SOURCE, split.getSourceShardState(i));
+            assertFalse(split.isTargetShard(i));
         }
         for (int i = numShards; i < numShards * multiple; i++) {
-            assert split.getTargetShardState(i) == IndexReshardingState.Split.TargetShardState.CLONE;
+            assertSame(IndexReshardingState.Split.TargetShardState.CLONE, split.getTargetShardState(i));
+            assertTrue(split.isTargetShard(i));
         }
 
         // advance split state randomly and expect to terminate
@@ -62,10 +64,12 @@ public void testSplit() {
         }
 
         for (int i = 0; i < numShards; i++) {
-            assert split.getSourceShardState(i) == IndexReshardingState.Split.SourceShardState.DONE;
+            assertSame(IndexReshardingState.Split.SourceShardState.DONE, split.getSourceShardState(i));
+            assertFalse(split.isTargetShard(i));
         }
         for (int i = numShards; i < numShards * multiple; i++) {
-            assert split.getTargetShardState(i) == IndexReshardingState.Split.TargetShardState.DONE;
+            assertSame(IndexReshardingState.Split.TargetShardState.DONE, split.getTargetShardState(i));
+            assertTrue(split.isTargetShard(i));
         }
     }
 }