Use set of shard routing for shard in unassigned shard batch check.

Signed-off-by: Swetha Guptha <[email protected]>
SwethaGuptha · Jul 3, 2024 · 475da7b · 475da7b
1 parent 9729a92
commit 475da7b
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 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.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<>();
@@ -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 (unassignedShard.primary() && batchShardRoutingSet.contains(unassignedShard)) {
                 assert unassignedShard.primary();
                 if (ineligibleShardAllocationDecisions.containsKey(unassignedShard.shardId())) {
                     allocationDecision = ineligibleShardAllocationDecisions.get(unassignedShard.shardId());
@@ -115,6 +119,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());
     }
 
     /**

diff --git 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<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<>();
@@ -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();
@@ -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(