Added in-flight fetches for batch mode

Signed-off-by: Shivansh Arora <[email protected]>
shiv0408 · Oct 17, 2023 · f65ed4d · f65ed4d
1 parent d5a95b8
commit f65ed4d
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/server/src/main/java/org/opensearch/gateway/GatewayAllocator.java b/server/src/main/java/org/opensearch/gateway/GatewayAllocator.java
@@ -116,6 +116,14 @@ protected GatewayAllocator() {
     @Override
     public int getNumberOfInFlightFetches() {
         int count = 0;
+        // If fetching is done in non batched-mode then maps to maintain batches will be empty and vice versa for batch-mode
+        for (ShardsBatch batch : batchIdToStartedShardBatch.values()) {
+            count += (batch.getNumberOfInFlightFetches() * batch.getBatchedShards().size());
+        }
+        for (ShardsBatch batch : batchIdToStoreShardBatch.values()) {
+            count += (batch.getNumberOfInFlightFetches() * batch.getBatchedShards().size());
+        }
+
         for (AsyncShardFetch<TransportNodesListGatewayStartedShards.NodeGatewayStartedShards> fetch : asyncFetchStarted.values()) {
             count += fetch.getNumberOfInFlightFetches();
         }