Fix SearchProgressActionListenerIT (#120888) (#121019)

This fixes SearchProgressActionListenerIT by notifying the search progress
listener before search operation terminates.

(cherry picked from commit cbb62c2)
Signed-off-by: Andrei Dan <[email protected]>

Author: andreidan

muted-tests.yml

@@ -438,9 +438,6 @@ tests:
 - class: org.elasticsearch.xpack.test.rest.XPackRestIT
   method: test {p0=transform/transforms_start_stop/Test schedule_now on an already started transform}
   issue: https://github.com/elastic/elasticsearch/issues/120720
-- class: org.elasticsearch.action.search.SearchProgressActionListenerIT
-  method: testSearchProgressWithHitsAndAggs
-  issue: https://github.com/elastic/elasticsearch/issues/120583
 - class: org.elasticsearch.xpack.esql.action.CrossClusterAsyncEnrichStopIT
   method: testEnrichAfterStop
   issue: https://github.com/elastic/elasticsearch/issues/120757

server/src/main/java/org/elasticsearch/action/search/FetchSearchPhase.java

@@ -155,15 +155,21 @@ private void innerRunFetch(ScoreDoc[] scoreDocs, int numShards, SearchPhaseContr
         );
         for (int i = 0; i < docIdsToLoad.length; i++) {
             List<Integer> entry = docIdsToLoad[i];
+            SearchPhaseResult shardPhaseResult = searchPhaseShardResults.get(i);
             if (entry == null) { // no results for this shard ID
                 // if we got some hits from this shard we have to release the context
                 // we do this below after sending out the fetch requests relevant to the search to give priority to those requests
                 // that contribute to the final search response
                 // in any case we count down this result since we don't talk to this shard anymore
+                if (shardPhaseResult != null) {
+                    // notifying the listener here as otherwise the search operation might finish before we
+                    // get a chance to notify the progress listener for some fetch results
+                    progressListener.notifyFetchResult(i);
+                }
                 counter.countDown();
             } else {
                 executeFetch(
-                    searchPhaseShardResults.get(i),
+                    shardPhaseResult,
                     counter,
                     entry,
                     rankDocsPerShard == null || rankDocsPerShard.get(i).isEmpty() ? null : new RankDocShardInfo(rankDocsPerShard.get(i)),
@@ -176,7 +182,6 @@ private void innerRunFetch(ScoreDoc[] scoreDocs, int numShards, SearchPhaseContr
                 SearchPhaseResult shardPhaseResult = searchPhaseShardResults.get(i);
                 if (shardPhaseResult != null) {
                     releaseIrrelevantSearchContext(shardPhaseResult, context);
-                    progressListener.notifyFetchResult(i);
                 }
             }
         }