Fix a flaky test that relies on terminate_after for an exact count match

release-notes/opensearch.release-notes-2.11.0.md

@@ -5,7 +5,7 @@
 ### Added
 - Add coordinator level stats for search latency ([#8386](https://github.com/opensearch-project/OpenSearch/issues/8386))
 - Add metrics for thread_pool task wait time ([#9681](https://github.com/opensearch-project/OpenSearch/pull/9681))
-- Add parallel file download support for remote store based replication ([#8596](https://github.com/opensearch-project/OpenSearch/pull/8596)) 
+- Add parallel file download support for remote store based replication ([#8596](https://github.com/opensearch-project/OpenSearch/pull/8596))
 - Async blob read support for S3 plugin ([#9694](https://github.com/opensearch-project/OpenSearch/pull/9694))
 - [Telemetry-Otel] Added support for OtlpGrpcSpanExporter exporter ([#9666](https://github.com/opensearch-project/OpenSearch/pull/9666))
 - Async blob read support for encrypted containers ([#10131](https://github.com/opensearch-project/OpenSearch/pull/10131))

server/src/internalClusterTest/java/org/opensearch/search/simple/SimpleSearchIT.java

@@ -305,7 +305,15 @@ public void dotestSimpleTerminateAfterCountWithSize(int size, int max) throws Ex
                 .setSize(size)
                 .setTrackTotalHits(true)
                 .get();
-            assertHitCount(searchResponse, i);
+
+            // Do not expect an exact match as an optimization introduced by https://issues.apache.org/jira/browse/LUCENE-10620
+            // can produce a total hit count > terminated_after, but this only kicks in
+            // when size = 0 which is when TotalHitCountCollector is used.
+            if (size == 0) {
+                assertHitCount(searchResponse, i, max);
+            } else {
+                assertHitCount(searchResponse, i);
+            }
             assertTrue(searchResponse.isTerminatedEarly());
             assertEquals(Math.min(i, size), searchResponse.getHits().getHits().length);
         }
@@ -319,7 +327,6 @@ public void dotestSimpleTerminateAfterCountWithSize(int size, int max) throws Ex
         assertFalse(searchResponse.isTerminatedEarly());
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/opensearch-project/OpenSearch/issues/10435")
     public void testSimpleTerminateAfterCountSize0() throws Exception {
         int max = randomIntBetween(3, 29);
         dotestSimpleTerminateAfterCountWithSize(0, max);
@@ -330,6 +337,24 @@ public void testSimpleTerminateAfterCountRandomSize() throws Exception {
         dotestSimpleTerminateAfterCountWithSize(randomIntBetween(1, max), max);
     }
 
+    /**
+     * Special cases when size = 0:
+     *
+     * If track_total_hits = true:
+     *   Weight#count optimization can cause totalHits in the response to be up to the total doc count regardless of terminate_after.
+     *   So, we will have to do a range check, not an equality check.
+     *
+     * If track_total_hits != true, but set to a value AND terminate_after is set:
+     *   Again, due to the optimization, any count can be returned.
+     *   Up to terminate_after, relation == EQUAL_TO.
+     *   But if track_total_hits_up_to >= terminate_after, relation can be EQ _or_ GTE.
+     *   This ambiguity is due to the fact that totalHits == track_total_hits_up_to
+     *   or totalHits > track_total_hits_up_to and SearchPhaseController sets totalHits = track_total_hits_up_to when returning results
+     *   in which case relation = GTE.
+     *
+     * @param size
+     * @throws Exception
+     */
     public void doTestSimpleTerminateAfterTrackTotalHitsUpTo(int size) throws Exception {
         prepareCreate("test").setSettings(Settings.builder().put(SETTING_NUMBER_OF_SHARDS, 1).put(SETTING_NUMBER_OF_REPLICAS, 0)).get();
         ensureGreen();
@@ -346,6 +371,7 @@ public void doTestSimpleTerminateAfterTrackTotalHitsUpTo(int size) throws Except
         refresh();
 
         SearchResponse searchResponse;
+
         searchResponse = client().prepareSearch("test")
             .setQuery(QueryBuilders.rangeQuery("field").gte(1).lte(numDocs))
             .setTerminateAfter(10)
@@ -356,25 +382,28 @@ public void doTestSimpleTerminateAfterTrackTotalHitsUpTo(int size) throws Except
         assertEquals(5, searchResponse.getHits().getTotalHits().value);
         assertEquals(GREATER_THAN_OR_EQUAL_TO, searchResponse.getHits().getTotalHits().relation);
 
-        searchResponse = client().prepareSearch("test")
-            .setQuery(QueryBuilders.rangeQuery("field").gte(1).lte(numDocs))
-            .setTerminateAfter(5)
-            .setSize(size)
-            .setTrackTotalHitsUpTo(10)
-            .get();
-        assertTrue(searchResponse.isTerminatedEarly());
-        assertEquals(5, searchResponse.getHits().getTotalHits().value);
-        assertEquals(EQUAL_TO, searchResponse.getHits().getTotalHits().relation);
+        // For size = 0, the following queries terminate early, but hits and relation can vary.
+        if (size > 0) {
+            searchResponse = client().prepareSearch("test")
+                .setQuery(QueryBuilders.rangeQuery("field").gte(1).lte(numDocs))
+                .setTerminateAfter(5)
+                .setSize(size)
+                .setTrackTotalHitsUpTo(10)
+                .get();
+            assertTrue(searchResponse.isTerminatedEarly());
+            assertEquals(5, searchResponse.getHits().getTotalHits().value);
+            assertEquals(EQUAL_TO, searchResponse.getHits().getTotalHits().relation);
 
-        searchResponse = client().prepareSearch("test")
-            .setQuery(QueryBuilders.rangeQuery("field").gte(1).lte(numDocs))
-            .setTerminateAfter(5)
-            .setSize(size)
-            .setTrackTotalHitsUpTo(5)
-            .get();
-        assertTrue(searchResponse.isTerminatedEarly());
-        assertEquals(5, searchResponse.getHits().getTotalHits().value);
-        assertEquals(EQUAL_TO, searchResponse.getHits().getTotalHits().relation);
+            searchResponse = client().prepareSearch("test")
+                .setQuery(QueryBuilders.rangeQuery("field").gte(1).lte(numDocs))
+                .setTerminateAfter(5)
+                .setSize(size)
+                .setTrackTotalHitsUpTo(5)
+                .get();
+            assertTrue(searchResponse.isTerminatedEarly());
+            assertEquals(5, searchResponse.getHits().getTotalHits().value);
+            assertEquals(EQUAL_TO, searchResponse.getHits().getTotalHits().relation);
+        }
 
         searchResponse = client().prepareSearch("test")
             .setQuery(QueryBuilders.rangeQuery("field").gte(1).lte(numDocs))
@@ -383,7 +412,12 @@ public void doTestSimpleTerminateAfterTrackTotalHitsUpTo(int size) throws Except
             .setTrackTotalHits(true)
             .get();
         assertTrue(searchResponse.isTerminatedEarly());
-        assertEquals(5, searchResponse.getHits().getTotalHits().value);
+        if (size == 0) {
+            // Since terminate_after < track_total_hits, we need to do a range check.
+            assertHitCount(searchResponse, 5, numDocs);
+        } else {
+            assertEquals(5, searchResponse.getHits().getTotalHits().value);
+        }
         assertEquals(EQUAL_TO, searchResponse.getHits().getTotalHits().relation);
 
         searchResponse = client().prepareSearch("test")
@@ -393,7 +427,12 @@ public void doTestSimpleTerminateAfterTrackTotalHitsUpTo(int size) throws Except
             .setTrackTotalHits(true)
             .get();
         assertFalse(searchResponse.isTerminatedEarly());
-        assertEquals(numDocs, searchResponse.getHits().getTotalHits().value);
+        if (size == 0) {
+            // Since terminate_after < track_total_hits, we need to do a range check.
+            assertHitCount(searchResponse, 5, numDocs);
+        } else {
+            assertEquals(5, searchResponse.getHits().getTotalHits().value);
+        }
         assertEquals(EQUAL_TO, searchResponse.getHits().getTotalHits().relation);
 
         searchResponse = client().prepareSearch("test")
@@ -405,12 +444,11 @@ public void doTestSimpleTerminateAfterTrackTotalHitsUpTo(int size) throws Except
         assertEquals(GREATER_THAN_OR_EQUAL_TO, searchResponse.getHits().getTotalHits().relation);
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/opensearch-project/OpenSearch/issues/10435")
-    public void testSimpleTerminateAfterTrackTotalHitsUpToRandomSize() throws Exception {
+    public void testSimpleTerminateAfterTrackTotalHitsUpToRandomSize0() throws Exception {
         doTestSimpleTerminateAfterTrackTotalHitsUpTo(0);
     }
 
-    public void testSimpleTerminateAfterTrackTotalHitsUpToSize0() throws Exception {
+    public void testSimpleTerminateAfterTrackTotalHitsUpToSize() throws Exception {
         doTestSimpleTerminateAfterTrackTotalHitsUpTo(randomIntBetween(1, 29));
     }
 

server/src/test/java/org/opensearch/action/search/SearchRequestTests.java

@@ -247,6 +247,12 @@ private SearchRequest mutate(SearchRequest searchRequest) {
         mutators.add(
             () -> mutation.setPhaseTook(searchRequest.isPhaseTook() == null ? randomBoolean() : searchRequest.isPhaseTook() == false)
         );
+        mutators.add(
+            () -> mutation.setPhaseTook(searchRequest.isPhaseTook() == null ? randomBoolean() : searchRequest.isPhaseTook() == false)
+        );
+        mutators.add(
+            () -> mutation.setPhaseTook(searchRequest.isPhaseTook() == null ? randomBoolean() : searchRequest.isPhaseTook() == false)
+        );
         mutators.add(
             () -> mutation.setCancelAfterTimeInterval(
                 searchRequest.getCancelAfterTimeInterval() != null

test/framework/src/main/java/org/opensearch/test/hamcrest/OpenSearchAssertions.java

@@ -304,6 +304,22 @@ public static void assertHitCount(SearchResponse countResponse, long expectedHit
         }
     }
 
+    public static void assertHitCount(SearchResponse countResponse, long minHitCount, long maxHitCount) {
+        final TotalHits totalHits = countResponse.getHits().getTotalHits();
+        if (!(totalHits.relation == TotalHits.Relation.EQUAL_TO && totalHits.value >= minHitCount && totalHits.value <= maxHitCount)) {
+            fail(
+                "Count is "
+                    + totalHits
+                    + " not between "
+                    + minHitCount
+                    + " and "
+                    + maxHitCount
+                    + " inclusive. "
+                    + formatShardStatus(countResponse)
+            );
+        }
+    }
+
     public static void assertExists(GetResponse response) {
         String message = String.format(Locale.ROOT, "Expected %s/%s to exist, but does not", response.getIndex(), response.getId());
         assertThat(message, response.isExists(), is(true));