Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Backport 2.x] Fix a flaky test that relies on terminate_after for an exact count ma… #12343

Merged
merged 1 commit into from
Feb 16, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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 ≥ 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();
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
Original file line number Diff line number Diff line change
Expand Up @@ -307,6 +307,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
Loading