Skip to content

Commit

Permalink
Added tests for routing table diff
Browse files Browse the repository at this point in the history
Signed-off-by: Bukhtawar Khan <bukhtawa@amazon.com>
  • Loading branch information
Bukhtawar committed Aug 4, 2024
1 parent 4607c5d commit 9b5c271
Show file tree
Hide file tree
Showing 2 changed files with 114 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,44 @@ public void testRoutingTableUpserts() {
}
}

public void testRoutingTableDeletes() {
assertThat(clusterState.routingTable().shardsWithState(ShardRoutingState.UNASSIGNED).size(), is(this.totalNumberOfShards));
initPrimaries();
int expectedUnassignedShardCount = this.totalNumberOfShards - 2 * this.numberOfShards;
assertThat(
clusterState.routingTable().shardsWithState(ShardRoutingState.UNASSIGNED).size(),
is(expectedUnassignedShardCount)
);
assertThat(clusterState.routingTable().shardsWithState(ShardRoutingState.INITIALIZING).size(), is(2 * this.numberOfShards));
Metadata metadata = Metadata.builder().put(createIndexMetadata(TEST_INDEX_1)).put(createIndexMetadata(TEST_INDEX_2)).build();
ClusterState oldClusterState = clusterState;

//delete index routing table for TEST_INDEX_1
metadata = Metadata.builder(metadata).put(createIndexMetadata(TEST_INDEX_3)).build();
RoutingTable testRoutingTable = new RoutingTable.Builder(clusterState.routingTable())
.remove(TEST_INDEX_1)
.build();
this.clusterState = ClusterState.builder(org.opensearch.cluster.ClusterName.CLUSTER_NAME_SETTING.getDefault(Settings.EMPTY))
.metadata(metadata)
.routingTable(testRoutingTable)
.build();
this.totalNumberOfShards = this.shardsPerIndex;
assertThat(clusterState.routingTable().shardsWithState(ShardRoutingState.UNASSIGNED).size(),
is(expectedUnassignedShardCount - this.numberOfShards * this.numberOfReplicas));
Diff<RoutingTable> fullDiff = clusterState.routingTable().diff(oldClusterState.getRoutingTable());
Diff<RoutingTable> incrementalDiff = clusterState.routingTable().incrementalDiff(oldClusterState.getRoutingTable());
RoutingTable newRoutingTable = incrementalDiff.apply(oldClusterState.getRoutingTable());
for (IndexRoutingTable indexRoutingTable : clusterState.routingTable()) {
assertEquals(clusterState.routingTable().version(), newRoutingTable.version());
assertEquals(indexRoutingTable, newRoutingTable.index(indexRoutingTable.getIndex()));
}
RoutingTable newRoutingTableWithFullDiff = fullDiff.apply(oldClusterState.getRoutingTable());
for (IndexRoutingTable indexRoutingTable : clusterState.routingTable()) {
assertEquals(clusterState.routingTable().version(), newRoutingTableWithFullDiff.version());
assertEquals(indexRoutingTable, newRoutingTableWithFullDiff.index(indexRoutingTable.getIndex()));
}
}

public void testRoutingTableUpsertsWithDiff() {
assertThat(clusterState.routingTable().shardsWithState(ShardRoutingState.UNASSIGNED).size(), is(this.totalNumberOfShards));
initPrimaries();
Expand All @@ -162,7 +200,7 @@ public void testRoutingTableUpsertsWithDiff() {
assertThat(clusterState.routingTable().shardsWithState(ShardRoutingState.UNASSIGNED).size(),
is(expectedUnassignedShardCount + this.shardsPerIndex));
initPrimaries();
clusterState = startRandomInitializingShard(clusterState, ALLOCATION_SERVICE);
clusterState = startRandomInitializingShard(clusterState, ALLOCATION_SERVICE, TEST_INDEX_2);
//assertThat(clusterState.routingTable().shardsWithState(ShardRoutingState.INITIALIZING).size(), is(2 * this.numberOfShards + 1));
Diff<RoutingTable> fullDiff = clusterState.routingTable().diff(oldClusterState.getRoutingTable());
Diff<RoutingTable> incrementalDiff = clusterState.routingTable().incrementalDiff(oldClusterState.getRoutingTable());
Expand All @@ -178,6 +216,72 @@ public void testRoutingTableUpsertsWithDiff() {
}
}

public void testRoutingTableDiffWithReplicaAdded() {
assertThat(clusterState.routingTable().shardsWithState(ShardRoutingState.UNASSIGNED).size(), is(this.totalNumberOfShards));
initPrimaries();
int expectedUnassignedShardCount = this.totalNumberOfShards - 2 * this.numberOfShards;
assertThat(
clusterState.routingTable().shardsWithState(ShardRoutingState.UNASSIGNED).size(),
is(expectedUnassignedShardCount)
);
assertThat(clusterState.routingTable().shardsWithState(ShardRoutingState.INITIALIZING).size(), is(2 * this.numberOfShards));
ClusterState oldClusterState = clusterState;

//update replica count for TEST_INDEX_1
RoutingTable updatedRoutingTable = RoutingTable.builder(clusterState.routingTable())
.updateNumberOfReplicas(this.numberOfReplicas + 1, new String[]{TEST_INDEX_1})
.build();
Metadata metadata = Metadata.builder(clusterState.metadata()).updateNumberOfReplicas(this.numberOfReplicas + 1, new String[]{TEST_INDEX_1}).build();
clusterState = ClusterState.builder(clusterState).routingTable(updatedRoutingTable).metadata(metadata).build();
assertThat(clusterState.routingTable().shardsWithState(ShardRoutingState.UNASSIGNED).size(),
is(expectedUnassignedShardCount + this.numberOfShards));
Diff<RoutingTable> fullDiff = clusterState.routingTable().diff(oldClusterState.getRoutingTable());
Diff<RoutingTable> incrementalDiff = clusterState.routingTable().incrementalDiff(oldClusterState.getRoutingTable());
RoutingTable newRoutingTable = incrementalDiff.apply(oldClusterState.getRoutingTable());
for (IndexRoutingTable indexRoutingTable : clusterState.routingTable()) {
assertEquals(clusterState.routingTable().version(), newRoutingTable.version());
assertEquals(indexRoutingTable, newRoutingTable.index(indexRoutingTable.getIndex()));
}
RoutingTable newRoutingTableWithFullDiff = fullDiff.apply(oldClusterState.getRoutingTable());
for (IndexRoutingTable indexRoutingTable : clusterState.routingTable()) {
assertEquals(clusterState.routingTable().version(), newRoutingTableWithFullDiff.version());
assertEquals(indexRoutingTable, newRoutingTableWithFullDiff.index(indexRoutingTable.getIndex()));
}
}

public void testRoutingTableDiffWithReplicaRemoved() {
assertThat(clusterState.routingTable().shardsWithState(ShardRoutingState.UNASSIGNED).size(), is(this.totalNumberOfShards));
initPrimaries();
int expectedUnassignedShardCount = this.totalNumberOfShards - 2 * this.numberOfShards;
assertThat(
clusterState.routingTable().shardsWithState(ShardRoutingState.UNASSIGNED).size(),
is(expectedUnassignedShardCount)
);
assertThat(clusterState.routingTable().shardsWithState(ShardRoutingState.INITIALIZING).size(), is(2 * this.numberOfShards));
ClusterState oldClusterState = clusterState;

//update replica count for TEST_INDEX_1
RoutingTable updatedRoutingTable = RoutingTable.builder(clusterState.routingTable())
.updateNumberOfReplicas(this.numberOfReplicas - 1, new String[]{TEST_INDEX_1})
.build();
Metadata metadata = Metadata.builder(clusterState.metadata()).updateNumberOfReplicas(this.numberOfReplicas - 1, new String[]{TEST_INDEX_1}).build();
clusterState = ClusterState.builder(clusterState).routingTable(updatedRoutingTable).metadata(metadata).build();
assertThat(clusterState.routingTable().shardsWithState(ShardRoutingState.UNASSIGNED).size(),
is(expectedUnassignedShardCount - this.numberOfShards));
Diff<RoutingTable> fullDiff = clusterState.routingTable().diff(oldClusterState.getRoutingTable());
Diff<RoutingTable> incrementalDiff = clusterState.routingTable().incrementalDiff(oldClusterState.getRoutingTable());
RoutingTable newRoutingTable = incrementalDiff.apply(oldClusterState.getRoutingTable());
for (IndexRoutingTable indexRoutingTable : clusterState.routingTable()) {
assertEquals(clusterState.routingTable().version(), newRoutingTable.version());
assertEquals(indexRoutingTable, newRoutingTable.index(indexRoutingTable.getIndex()));
}
RoutingTable newRoutingTableWithFullDiff = fullDiff.apply(oldClusterState.getRoutingTable());
for (IndexRoutingTable indexRoutingTable : clusterState.routingTable()) {
assertEquals(clusterState.routingTable().version(), newRoutingTableWithFullDiff.version());
assertEquals(indexRoutingTable, newRoutingTableWithFullDiff.index(indexRoutingTable.getIndex()));
}
}

public void testRoutingTableDiffsWithStartedState() {
assertThat(clusterState.routingTable().shardsWithState(ShardRoutingState.UNASSIGNED).size(), is(this.totalNumberOfShards));
initPrimaries();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -179,6 +179,15 @@ protected static DiscoveryNode newNode(String nodeId, Version version) {

protected static ClusterState startRandomInitializingShard(ClusterState clusterState, AllocationService strategy) {
List<ShardRouting> initializingShards = clusterState.getRoutingNodes().shardsWithState(INITIALIZING);
return startInitialisingShardsAndReroute(strategy, clusterState, initializingShards);
}

protected static ClusterState startRandomInitializingShard(ClusterState clusterState, AllocationService strategy, String index) {
List<ShardRouting> initializingShards = clusterState.getRoutingNodes().shardsWithState(index, INITIALIZING);
return startInitialisingShardsAndReroute(strategy, clusterState, initializingShards);
}

private static ClusterState startInitialisingShardsAndReroute(AllocationService strategy, ClusterState clusterState, List<ShardRouting> initializingShards) {
if (initializingShards.isEmpty()) {
return clusterState;
}
Expand Down

0 comments on commit 9b5c271

Please sign in to comment.