From f65ed4d4399120e786dea13b0b17a8733ea00419 Mon Sep 17 00:00:00 2001
From: Shivansh Arora <shivansh.arora@protonmail.com>
Date: Wed, 30 Aug 2023 13:18:04 +0530
Subject: [PATCH] Added in-flight fetches for batch mode

Signed-off-by: Shivansh Arora <shivansh.arora@protonmail.com>
---
 .../java/org/opensearch/gateway/GatewayAllocator.java     | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/server/src/main/java/org/opensearch/gateway/GatewayAllocator.java b/server/src/main/java/org/opensearch/gateway/GatewayAllocator.java
index b5a00b1a47523..1baea091fc7f2 100644
--- 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();
         }