Added in-flight fetches for batchMode

opensearch-project · Aug 31, 2023 · f0cdd68 · f0cdd68
1 parent 18a5d9c
commit f0cdd68
Showing 1 changed file with 19 additions and 4 deletions.
diff --git a/server/src/main/java/org/opensearch/gateway/GatewayAllocator.java b/server/src/main/java/org/opensearch/gateway/GatewayAllocator.java
@@ -146,12 +146,23 @@ protected GatewayAllocator() {
 
     @Override
     public int getNumberOfInFlightFetches() {
+        boolean batchModeEnabled = true;
         int count = 0;
-        for (AsyncShardFetch<TransportNodesListGatewayStartedShards.NodeGatewayStartedShards> fetch : asyncFetchStarted.values()) {
-            count += fetch.getNumberOfInFlightFetches();
+        if (batchModeEnabled) {
+            for (ShardsBatch batch : batchIdToStartedShardBatch.values()) {
+                count += batch.getNumberOfInFlightFetches();
+            }
+            for (ShardsBatch batch : batchIdToStoreShardBatch.values()) {
+                count += batch.getNumberOfInFlightFetches();
+            }
         }
-        for (AsyncShardFetch<TransportNodesListShardStoreMetadata.NodeStoreFilesMetadata> fetch : asyncFetchStore.values()) {
-            count += fetch.getNumberOfInFlightFetches();
+        else {
+            for (AsyncShardFetch<TransportNodesListGatewayStartedShards.NodeGatewayStartedShards> fetch : asyncFetchStarted.values()) {
+                count += fetch.getNumberOfInFlightFetches();
+            }
+            for (AsyncShardFetch<TransportNodesListShardStoreMetadata.NodeStoreFilesMetadata> fetch : asyncFetchStore.values()) {
+                count += fetch.getNumberOfInFlightFetches();
+            }
         }
         return count;
     }
@@ -596,6 +607,10 @@ AsyncBatchShardFetch<? extends BaseNodeResponse> getAsyncFetcher() {
             return asyncBatch;
         }
 
+        public int getNumberOfInFlightFetches() {
+            return asyncBatch.getNumberOfInFlightFetches();
+        }
+
         @Override
         public boolean equals(Object o) {
             if (this == o) {