Keep initializing shards in batch

Signed-off-by: Shivansh Arora <[email protected]>
Gaurav614 · Apr 22, 2024 · 27c9b98 · 27c9b98
1 parent 90659b5
commit 27c9b98
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/server/src/main/java/org/opensearch/cluster/routing/RoutingNode.java b/server/src/main/java/org/opensearch/cluster/routing/RoutingNode.java
@@ -204,6 +204,10 @@ public int size() {
         return shards.size();
     }
 
+    public Collection<ShardRouting> getInitializingShards() {
+        return initializingShards;
+    }
+
     /**
      * Add a new shard to this node
      * @param shard Shard to create on this Node

diff --git a/server/src/main/java/org/opensearch/gateway/ShardsBatchGatewayAllocator.java b/server/src/main/java/org/opensearch/gateway/ShardsBatchGatewayAllocator.java
@@ -251,6 +251,16 @@ else if (shardRouting.primary() == primary) {
             }
         });
 
+        allocation.routingNodes().forEach(
+            routingNode ->
+            routingNode.getInitializingShards().forEach(
+                shardRouting -> {
+                    if (currentBatchedShards.containsKey(shardRouting.shardId()) && shardRouting.primary() == primary) {
+                        batchedShardsToAssign.add(shardRouting.shardId());
+                    }
+                })
+        );
+
         refreshShardBatches(currentBatches, batchedShardsToAssign, primary);
 
         Iterator<ShardRouting> iterator = newShardsToBatch.iterator();