Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use set for shard routings in batch check. #14533

Merged
merged 1 commit into from
Jul 11, 2024
Merged

Use set for shard routings in batch check. #14533

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions server/src/main/java/org/opensearch/gateway/PrimaryShardBatchAllocator.java
View file
Open in desktop
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)) {
ashking94 marked this conversation as resolved.
Show resolved Hide resolved
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
11 changes: 9 additions & 2 deletions server/src/main/java/org/opensearch/gateway/ReplicaShardBatchAllocator.java
View file
Open in desktop
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
36 changes: 34 additions & 2 deletions server/src/test/java/org/opensearch/gateway/PrimaryShardBatchAllocatorTests.java
View file
Open in desktop
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
Loading