Fix SearchProgressActionListenerIT (#120888)

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

Author: andreidan

muted-tests.yml

@@ -213,9 +213,6 @@ tests:
   issue: https://github.com/elastic/elasticsearch/issues/120482
 - class: org.elasticsearch.xpack.sql.expression.function.scalar.datetime.DateTimeToCharProcessorTests
   issue: https://github.com/elastic/elasticsearch/issues/120575
-- class: org.elasticsearch.action.search.SearchProgressActionListenerIT
-  method: testSearchProgressWithHitsAndAggs
-  issue: https://github.com/elastic/elasticsearch/issues/120583
 - class: org.elasticsearch.index.reindex.BulkByScrollUsesAllScrollDocumentsAfterConflictsIntegTests
   method: testReindex
   issue: https://github.com/elastic/elasticsearch/issues/120605
@@ -232,9 +229,6 @@ tests:
   issue: https://github.com/elastic/elasticsearch/issues/120767
 - class: org.elasticsearch.search.fieldcaps.FieldCapabilitiesIT
   issue: https://github.com/elastic/elasticsearch/issues/120772
-- class: org.elasticsearch.action.search.SearchProgressActionListenerIT
-  method: testSearchProgressWithHits
-  issue: https://github.com/elastic/elasticsearch/issues/120671
 - class: org.elasticsearch.xpack.test.rest.XPackRestIT
   method: test {p0=ml/3rd_party_deployment/Test start deployment fails while model download in progress}
   issue: https://github.com/elastic/elasticsearch/issues/120810

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

@@ -158,15 +158,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)),
@@ -179,7 +185,6 @@ private void innerRunFetch(ScoreDoc[] scoreDocs, int numShards, SearchPhaseContr
                 SearchPhaseResult shardPhaseResult = searchPhaseShardResults.get(i);
                 if (shardPhaseResult != null) {
                     releaseIrrelevantSearchContext(shardPhaseResult, context);
-                    progressListener.notifyFetchResult(i);
                 }
             }
         }