Skip to content

Commit

Permalink
Use set of shard routing for shard in unassigned shard batch check. (o…
Browse files Browse the repository at this point in the history
…pensearch-project#14533)

Signed-off-by: Swetha Guptha <[email protected]>
  • Loading branch information
SwethaGuptha authored Jul 11, 2024
1 parent dfb8449 commit 125b773
Show file tree
Hide file tree
Showing 3 changed files with 49 additions and 6 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -23,8 +23,10 @@
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
* PrimaryShardBatchAllocator is similar to {@link org.opensearch.gateway.PrimaryShardAllocator} only difference is
Expand Down Expand Up @@ -82,6 +84,7 @@ public AllocateUnassignedDecision makeAllocationDecision(ShardRouting unassigned
* @param allocation the allocation state container object
*/
public void allocateUnassignedBatch(List<ShardRouting> shardRoutings, RoutingAllocation allocation) {
logger.trace("Starting shard allocation execution for unassigned primary shards: {}", shardRoutings.size());
HashMap<ShardId, AllocateUnassignedDecision> ineligibleShardAllocationDecisions = new HashMap<>();
List<ShardRouting> eligibleShards = new ArrayList<>();
List<ShardRouting> inEligibleShards = new ArrayList<>();
Expand All @@ -99,13 +102,13 @@ public void allocateUnassignedBatch(List<ShardRouting> shardRoutings, RoutingAll
// only fetch data for eligible shards
final FetchResult<NodeGatewayStartedShardsBatch> shardsState = fetchData(eligibleShards, inEligibleShards, allocation);

Set<ShardRouting> batchShardRoutingSet = new HashSet<>(shardRoutings);
RoutingNodes.UnassignedShards.UnassignedIterator iterator = allocation.routingNodes().unassigned().iterator();
while (iterator.hasNext()) {
ShardRouting unassignedShard = iterator.next();
AllocateUnassignedDecision allocationDecision;

if (shardRoutings.contains(unassignedShard)) {
assert unassignedShard.primary();
if (unassignedShard.primary() && batchShardRoutingSet.contains(unassignedShard)) {
if (ineligibleShardAllocationDecisions.containsKey(unassignedShard.shardId())) {
allocationDecision = ineligibleShardAllocationDecisions.get(unassignedShard.shardId());
} else {
Expand All @@ -115,6 +118,7 @@ public void allocateUnassignedBatch(List<ShardRouting> shardRoutings, RoutingAll
executeDecision(unassignedShard, allocationDecision, allocation, iterator);
}
}
logger.trace("Finished shard allocation execution for unassigned primary shards: {}", shardRoutings.size());
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,10 +28,11 @@
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Supplier;
import java.util.stream.Collectors;

/**
* Allocates replica shards in a batch mode
Expand Down Expand Up @@ -117,6 +118,7 @@ public AllocateUnassignedDecision makeAllocationDecision(ShardRouting unassigned
* @param allocation the allocation state container object
*/
public void allocateUnassignedBatch(List<ShardRouting> shardRoutings, RoutingAllocation allocation) {
logger.trace("Starting shard allocation execution for unassigned replica shards: {}", shardRoutings.size());
List<ShardRouting> eligibleShards = new ArrayList<>();
List<ShardRouting> ineligibleShards = new ArrayList<>();
Map<ShardRouting, AllocateUnassignedDecision> ineligibleShardAllocationDecisions = new HashMap<>();
Expand All @@ -135,7 +137,11 @@ public void allocateUnassignedBatch(List<ShardRouting> shardRoutings, RoutingAll
// only fetch data for eligible shards
final FetchResult<NodeStoreFilesMetadataBatch> shardsState = fetchData(eligibleShards, ineligibleShards, allocation);

List<ShardId> shardIdsFromBatch = shardRoutings.stream().map(shardRouting -> shardRouting.shardId()).collect(Collectors.toList());
Set<ShardId> shardIdsFromBatch = new HashSet<>();
for (ShardRouting shardRouting : shardRoutings) {
ShardId shardId = shardRouting.shardId();
shardIdsFromBatch.add(shardId);
}
RoutingNodes.UnassignedShards.UnassignedIterator iterator = allocation.routingNodes().unassigned().iterator();
while (iterator.hasNext()) {
ShardRouting unassignedShard = iterator.next();
Expand All @@ -159,6 +165,7 @@ public void allocateUnassignedBatch(List<ShardRouting> shardRoutings, RoutingAll
executeDecision(unassignedShard, allocateUnassignedDecision, allocation, iterator);
}
}
logger.trace("Finished shard allocation execution for unassigned replica shards: {}", shardRoutings.size());
}

private AllocateUnassignedDecision getUnassignedShardAllocationDecision(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,10 @@ private void allocateAllUnassignedBatch(final RoutingAllocation allocation) {
final RoutingNodes.UnassignedShards.UnassignedIterator iterator = allocation.routingNodes().unassigned().iterator();
List<ShardRouting> shardsToBatch = new ArrayList<>();
while (iterator.hasNext()) {
shardsToBatch.add(iterator.next());
ShardRouting unassignedShardRouting = iterator.next();
if (unassignedShardRouting.primary()) {
shardsToBatch.add(unassignedShardRouting);
}
}
batchAllocator.allocateUnassignedBatch(shardsToBatch, allocation);
}
Expand Down Expand Up @@ -180,6 +183,35 @@ public void testInitializePrimaryShards() {
assertEquals(2, routingAllocation.routingNodes().getInitialPrimariesIncomingRecoveries(node1.getId()));
}

public void testInitializeOnlyPrimaryUnassignedShardsIgnoreReplicaShards() {
ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
AllocationDeciders allocationDeciders = randomAllocationDeciders(Settings.builder().build(), clusterSettings, random());
setUpShards(1);
final RoutingAllocation routingAllocation = routingAllocationWithOnePrimary(allocationDeciders, CLUSTER_RECOVERED, "allocId-0");

for (ShardId shardId : shardsInBatch) {
batchAllocator.addShardData(
node1,
"allocId-0",
shardId,
true,
new ReplicationCheckpoint(shardId, 20, 101, 1, Codec.getDefault().getName()),
null
);
}

allocateAllUnassignedBatch(routingAllocation);

List<ShardRouting> initializingShards = routingAllocation.routingNodes().shardsWithState(ShardRoutingState.INITIALIZING);
assertEquals(1, initializingShards.size());
assertTrue(shardsInBatch.contains(initializingShards.get(0).shardId()));
assertTrue(initializingShards.get(0).primary());
assertEquals(1, routingAllocation.routingNodes().getInitialPrimariesIncomingRecoveries(node1.getId()));
List<ShardRouting> unassignedShards = routingAllocation.routingNodes().shardsWithState(ShardRoutingState.UNASSIGNED);
assertEquals(1, unassignedShards.size());
assertTrue(!unassignedShards.get(0).primary());
}

public void testAllocateUnassignedBatchThrottlingAllocationDeciderIsHonoured() {
ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
AllocationDeciders allocationDeciders = randomAllocationDeciders(
Expand Down Expand Up @@ -258,7 +290,7 @@ private RoutingAllocation routingAllocationWithOnePrimary(
.routingTable(routingTableBuilder.build())
.nodes(DiscoveryNodes.builder().add(node1).add(node2).add(node3))
.build();
return new RoutingAllocation(deciders, new RoutingNodes(state, false), state, null, null, System.nanoTime());
return new RoutingAllocation(deciders, new RoutingNodes(state, false), state, ClusterInfo.EMPTY, null, System.nanoTime());
}

private RoutingAllocation routingAllocationWithMultiplePrimaries(
Expand Down

0 comments on commit 125b773

Please sign in to comment.