Make sure processPostCollection is called in case of early termination (opensearch-project#14208)

Signed-off-by: Jay Deng <[email protected]>

Author: jed326

server/src/internalClusterTest/java/org/opensearch/search/aggregations/AggregationsIntegrationIT.java

@@ -187,4 +187,27 @@ public void testAggsOnEmptyShards() {
         // Validate non-global agg does not throw an exception
         assertSearchResponse(client().prepareSearch("idx").addAggregation(stats("value_stats").field("score")).get());
     }
+
+    public void testAggsWithTerminateAfter() throws InterruptedException {
+        assertAcked(
+            prepareCreate(
+                "terminate_index",
+                Settings.builder().put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1).put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0)
+            ).setMapping("f", "type=keyword").get()
+        );
+        List<IndexRequestBuilder> docs = new ArrayList<>();
+        for (int i = 0; i < randomIntBetween(5, 20); ++i) {
+            docs.add(client().prepareIndex("terminate_index").setSource("f", Integer.toString(i / 3)));
+        }
+        indexRandom(true, docs);
+
+        SearchResponse response = client().prepareSearch("terminate_index")
+            .setSize(2)
+            .setTerminateAfter(1)
+            .addAggregation(terms("f").field("f"))
+            .get();
+        assertSearchResponse(response);
+        assertTrue(response.isTerminatedEarly());
+        assertEquals(response.getHits().getHits().length, 1);
+    }
 }

server/src/main/java/org/opensearch/search/query/QueryPhase.java

@@ -354,6 +354,9 @@ private static boolean searchWithCollector(
         try {
             searcher.search(query, queryCollector);
         } catch (EarlyTerminatingCollector.EarlyTerminationException e) {
+            // EarlyTerminationException is not caught in ContextIndexSearcher to allow force termination of collection. Postcollection
+            // still needs to be processed for Aggregations when early termination takes place.
+            searchContext.bucketCollectorProcessor().processPostCollection(queryCollector);
             queryResult.terminatedEarly(true);
         }
         if (searchContext.isSearchTimedOut()) {