From 2b497b97e390acebfc341397a0b02093a2a5e81d Mon Sep 17 00:00:00 2001
From: Kaituo Li <kaituo@amazon.com>
Date: Wed, 23 Oct 2024 18:07:19 -0700
Subject: [PATCH] make sure increment before schedule

Signed-off-by: Kaituo Li <kaituo@amazon.com>
---
 .../opensearch/timeseries/transport/ResultProcessor.java    | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/main/java/org/opensearch/timeseries/transport/ResultProcessor.java b/src/main/java/org/opensearch/timeseries/transport/ResultProcessor.java
index 7b3424d69..2b7ebca65 100644
--- a/src/main/java/org/opensearch/timeseries/transport/ResultProcessor.java
+++ b/src/main/java/org/opensearch/timeseries/transport/ResultProcessor.java
@@ -229,6 +229,9 @@ class PageListener implements ActionListener<CompositeRetriever.Page> {
 
         @Override
         public void onResponse(CompositeRetriever.Page entityFeatures) {
+            // Increment pagesInFlight to track the processing of this page
+            pagesInFlight.incrementAndGet();
+
             // start processing next page after sending out features for previous page
             if (pageIterator.hasNext()) {
                 pageIterator.next(this);
@@ -236,9 +239,6 @@ public void onResponse(CompositeRetriever.Page entityFeatures) {
                 scheduleImputeHCTask();
             }
 
-            // Increment pagesInFlight to track the processing of this page
-            pagesInFlight.incrementAndGet();
-
             if (entityFeatures != null && false == entityFeatures.isEmpty()) {
                 LOG
                     .info(