Make sure processPostCollection is called in case of early termination (

opensearch-project#14208) Signed-off-by: Jay Deng <jayd0104@gmail.com>
harshavamsi · Jul 12, 2024 · f8bd939 · f8bd939
1 parent 94ae591
commit f8bd939
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 0 deletions.
diff --git a/...nternalClusterTest/java/org/opensearch/search/aggregations/AggregationsIntegrationIT.java b/...nternalClusterTest/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);
+    }
 }
diff --git a/server/src/main/java/org/opensearch/search/query/QueryPhase.java b/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()) {