Skip to content

Commit

Permalink
Fix a flaky test that relies on terminate_after for an exact count ma… (
Browse files Browse the repository at this point in the history 
#12179)

* Fix a flaky test that relies on terminate_after for an exact count match.  (Issue # 10435)

Signed-off-by: Austin Lee <austin@aryn.ai>

* Make comments HTML Javadoc friendly.

Signed-off-by: Austin Lee <austin@aryn.ai>

* Fix javadoc syntax

Signed-off-by: Andrew Ross <andrross@amazon.com>

---------

Signed-off-by: Austin Lee <austin@aryn.ai>
Signed-off-by: Andrew Ross <andrross@amazon.com>
Co-authored-by: Andrew Ross <andrross@amazon.com>
  • Loading branch information
@austintlee @andrross
austintlee and andrross authored Feb 15, 2024
1 parent 03e415e commit 8fe15e6
Show file tree
Hide file tree
Showing 2 changed files with 73 additions and 24 deletions.
81 changes: 57 additions & 24 deletions server/src/internalClusterTest/java/org/opensearch/search/simple/SimpleSearchIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -299,7 +299,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);
}
Expand All @@ -313,7 +321,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);
Expand All @@ -324,6 +331,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 &ge; terminate_after, relation can be EQ _or_ GTE.
* This ambiguity is due to the fact that totalHits == track_total_hits_up_to
* or totalHits &gt; 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();
Expand All @@ -340,6 +365,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)
Expand All @@ -350,25 +376,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))
Expand All @@ -377,7 +406,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")
Expand All @@ -399,12 +433,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));
}

Expand Down
16 changes: 16 additions & 0 deletions test/framework/src/main/java/org/opensearch/test/hamcrest/OpenSearchAssertions.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand Down

0 comments on commit 8fe15e6

Please sign in to comment.