From baa07b7788d81e4cab40ce5ec3a2a216ab7752a8 Mon Sep 17 00:00:00 2001 From: Swetha Guptha Date: Tue, 25 Jun 2024 20:28:45 +0530 Subject: [PATCH] Use set of shard routing for shard in unassigned shard batch check. Signed-off-by: Swetha Guptha --- .../gateway/PrimaryShardBatchAllocator.java | 8 +++-- .../gateway/ReplicaShardBatchAllocator.java | 11 ++++-- .../PrimaryShardBatchAllocatorTests.java | 36 +++++++++++++++++-- 3 files changed, 49 insertions(+), 6 deletions(-) diff --git a/server/src/main/java/org/opensearch/gateway/PrimaryShardBatchAllocator.java b/server/src/main/java/org/opensearch/gateway/PrimaryShardBatchAllocator.java index 27f9bedc4e495..c493bf717c97f 100644 --- a/server/src/main/java/org/opensearch/gateway/PrimaryShardBatchAllocator.java +++ b/server/src/main/java/org/opensearch/gateway/PrimaryShardBatchAllocator.java @@ -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 @@ -82,6 +84,7 @@ public AllocateUnassignedDecision makeAllocationDecision(ShardRouting unassigned * @param allocation the allocation state container object */ public void allocateUnassignedBatch(List shardRoutings, RoutingAllocation allocation) { + logger.trace("Starting shard allocation execution for unassigned primary shards: {}", shardRoutings.size()); HashMap ineligibleShardAllocationDecisions = new HashMap<>(); List eligibleShards = new ArrayList<>(); List inEligibleShards = new ArrayList<>(); @@ -99,13 +102,13 @@ public void allocateUnassignedBatch(List shardRoutings, RoutingAll // only fetch data for eligible shards final FetchResult shardsState = fetchData(eligibleShards, inEligibleShards, allocation); + Set 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 { @@ -115,6 +118,7 @@ public void allocateUnassignedBatch(List shardRoutings, RoutingAll executeDecision(unassignedShard, allocationDecision, allocation, iterator); } } + logger.trace("Finished shard allocation execution for unassigned primary shards: {}", shardRoutings.size()); } /** diff --git a/server/src/main/java/org/opensearch/gateway/ReplicaShardBatchAllocator.java b/server/src/main/java/org/opensearch/gateway/ReplicaShardBatchAllocator.java index f2cb3d053440d..7c75f2a5d1a8f 100644 --- a/server/src/main/java/org/opensearch/gateway/ReplicaShardBatchAllocator.java +++ b/server/src/main/java/org/opensearch/gateway/ReplicaShardBatchAllocator.java @@ -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 @@ -117,6 +118,7 @@ public AllocateUnassignedDecision makeAllocationDecision(ShardRouting unassigned * @param allocation the allocation state container object */ public void allocateUnassignedBatch(List shardRoutings, RoutingAllocation allocation) { + logger.trace("Starting shard allocation execution for unassigned replica shards: {}", shardRoutings.size()); List eligibleShards = new ArrayList<>(); List ineligibleShards = new ArrayList<>(); Map ineligibleShardAllocationDecisions = new HashMap<>(); @@ -135,7 +137,11 @@ public void allocateUnassignedBatch(List shardRoutings, RoutingAll // only fetch data for eligible shards final FetchResult shardsState = fetchData(eligibleShards, ineligibleShards, allocation); - List shardIdsFromBatch = shardRoutings.stream().map(shardRouting -> shardRouting.shardId()).collect(Collectors.toList()); + Set 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(); @@ -159,6 +165,7 @@ public void allocateUnassignedBatch(List shardRoutings, RoutingAll executeDecision(unassignedShard, allocateUnassignedDecision, allocation, iterator); } } + logger.trace("Finished shard allocation execution for unassigned replica shards: {}", shardRoutings.size()); } private AllocateUnassignedDecision getUnassignedShardAllocationDecision( diff --git a/server/src/test/java/org/opensearch/gateway/PrimaryShardBatchAllocatorTests.java b/server/src/test/java/org/opensearch/gateway/PrimaryShardBatchAllocatorTests.java index e90850de3fe33..8ad8bcda95f40 100644 --- a/server/src/test/java/org/opensearch/gateway/PrimaryShardBatchAllocatorTests.java +++ b/server/src/test/java/org/opensearch/gateway/PrimaryShardBatchAllocatorTests.java @@ -85,7 +85,10 @@ private void allocateAllUnassignedBatch(final RoutingAllocation allocation) { final RoutingNodes.UnassignedShards.UnassignedIterator iterator = allocation.routingNodes().unassigned().iterator(); List shardsToBatch = new ArrayList<>(); while (iterator.hasNext()) { - shardsToBatch.add(iterator.next()); + ShardRouting unassignedShardRouting = iterator.next(); + if (unassignedShardRouting.primary()) { + shardsToBatch.add(unassignedShardRouting); + } } batchAllocator.allocateUnassignedBatch(shardsToBatch, allocation); } @@ -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 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 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( @@ -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(