Use set for shard routings in batch check.

Signed-off-by: Swetha Guptha <gupthasg@amazon.com>
opensearch-project · Jun 25, 2024 · cf04569 · cf04569
1 parent 9729a92
commit cf04569
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git 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<ShardRouting> shardRoutings, RoutingAllocation allocation) {
+        logger.debug("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<>();
@@ -99,12 +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)) {
+            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<ShardRouting> 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
@@ -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<ShardRouting> shardRoutings, RoutingAllocation allocation) {
+        logger.debug("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<>();
@@ -135,7 +138,7 @@ 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 = 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<ShardRouting> shardRoutings, RoutingAll
                 executeDecision(unassignedShard, allocateUnassignedDecision, allocation, iterator);
             }
         }
+        logger.debug("Finished shard allocation execution for unassigned replica shards: {}", shardRoutings.size());
     }
 
     private AllocateUnassignedDecision getUnassignedShardAllocationDecision(