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

Signed-off-by: Swetha Guptha <gupthasg@amazon.com>

Author: Swetha Guptha

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,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)) {
                 if (ineligibleShardAllocationDecisions.containsKey(unassignedShard.shardId())) {
                     allocationDecision = ineligibleShardAllocationDecisions.get(unassignedShard.shardId());
                 } else {
@@ -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());
     }
 
     /**

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(

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