diff --git a/server/src/main/java/org/opensearch/gateway/PrimaryShardBatchAllocator.java b/server/src/main/java/org/opensearch/gateway/PrimaryShardBatchAllocator.java index 27f9bedc4e495..e57ba5cfd74d7 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.debug("Starting shard allocation execution for unassigned primary shards: {}", shardRoutings.size()); HashMap ineligibleShardAllocationDecisions = new HashMap<>(); List eligibleShards = new ArrayList<>(); List inEligibleShards = new ArrayList<>(); @@ -99,12 +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)) { + if (batchShardRoutingSet.contains(unassignedShard)) { assert unassignedShard.primary(); if (ineligibleShardAllocationDecisions.containsKey(unassignedShard.shardId())) { allocationDecision = ineligibleShardAllocationDecisions.get(unassignedShard.shardId()); @@ -113,8 +117,10 @@ public void allocateUnassignedBatch(List shardRoutings, RoutingAll allocationDecision = getAllocationDecision(unassignedShard, allocation, nodeShardStates, logger); } executeDecision(unassignedShard, allocationDecision, allocation, iterator); + batchShardRoutingSet.remove(unassignedShard); } } + logger.debug("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..9bef7cddc87ef 100644 --- a/server/src/main/java/org/opensearch/gateway/ReplicaShardBatchAllocator.java +++ b/server/src/main/java/org/opensearch/gateway/ReplicaShardBatchAllocator.java @@ -28,8 +28,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; import java.util.function.Supplier; import java.util.stream.Collectors; @@ -117,6 +119,7 @@ public AllocateUnassignedDecision makeAllocationDecision(ShardRouting unassigned * @param allocation the allocation state container object */ public void allocateUnassignedBatch(List shardRoutings, RoutingAllocation allocation) { + logger.debug("Starting shard allocation execution for unassigned replica shards: {}", shardRoutings.size()); List eligibleShards = new ArrayList<>(); List ineligibleShards = new ArrayList<>(); Map ineligibleShardAllocationDecisions = new HashMap<>(); @@ -135,7 +138,7 @@ 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 = shardRoutings.stream().map(ShardRouting::shardId).collect(Collectors.toCollection(HashSet::new)); RoutingNodes.UnassignedShards.UnassignedIterator iterator = allocation.routingNodes().unassigned().iterator(); while (iterator.hasNext()) { ShardRouting unassignedShard = iterator.next(); @@ -159,6 +162,7 @@ public void allocateUnassignedBatch(List shardRoutings, RoutingAll executeDecision(unassignedShard, allocateUnassignedDecision, allocation, iterator); } } + logger.debug("Finished shard allocation execution for unassigned replica shards: {}", shardRoutings.size()); } private AllocateUnassignedDecision getUnassignedShardAllocationDecision(