diff --git a/server/src/main/java/org/elasticsearch/cluster/metadata/IndexReshardingMetadata.java b/server/src/main/java/org/elasticsearch/cluster/metadata/IndexReshardingMetadata.java index 3fd00dd7449bd..4ec325c42f3f3 100644 --- a/server/src/main/java/org/elasticsearch/cluster/metadata/IndexReshardingMetadata.java +++ b/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()); } /** diff --git a/server/src/main/java/org/elasticsearch/cluster/metadata/IndexReshardingState.java b/server/src/main/java/org/elasticsearch/cluster/metadata/IndexReshardingState.java index 3dbadff3a1c58..76f8f58d0b3e4 100644 --- a/server/src/main/java/org/elasticsearch/cluster/metadata/IndexReshardingState.java +++ b/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 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 diff --git a/server/src/test/java/org/elasticsearch/cluster/metadata/IndexReshardingMetadataTests.java b/server/src/test/java/org/elasticsearch/cluster/metadata/IndexReshardingMetadataTests.java index 8ec3b6e67860d..462fe26313682 100644 --- a/server/src/test/java/org/elasticsearch/cluster/metadata/IndexReshardingMetadataTests.java +++ b/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)); } } }