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.

Sign up for GitHub

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

[9.0] Address test issue in QueryPhaseTimeoutTests (#124327) #124515

Merged
merged 1 commit into from
Mar 10, 2025
Merged

[9.0] Address test issue in QueryPhaseTimeoutTests (#124327) #124515

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
129 changes: 118 additions & 11 deletions server/src/test/java/org/elasticsearch/search/query/QueryPhaseTimeoutTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,7 @@ public void tearDown() throws Exception {
}

private static ContextIndexSearcher newContextSearcher(IndexReader reader) throws IOException {
// note that no executor is provided, as this test requires sequential execution
return new ContextIndexSearcher(
reader,
IndexSearcher.getDefaultSimilarity(),
Expand All @@ -125,27 +126,35 @@ private static ContextIndexSearcher newContextSearcher(IndexReader reader) throw
);
}

public void testScorerTimeoutTerms() throws IOException {
/**
* Test that a timeout is appropriately handled when the (exitable) directory reader raises it while loading terms enum
* as the scorer supplier is requested.
*/
public void testScorerSupplierTimeoutTerms() throws IOException {
assumeTrue("Test requires more than one segment", reader.leaves().size() > 1);
int size = randomBoolean() ? 0 : randomIntBetween(100, 500);
scorerTimeoutTest(size, context -> {
scorerSupplierTimeoutTest(size, context -> {
final TermsEnum termsEnum = context.reader().terms("field").iterator();
termsEnum.next();
});
}

public void testScorerTimeoutPoints() throws IOException {
/**
* Test that a timeout is appropriately handled when the (exitable) directory reader raises it while loading points
* as the scorer supplier is requested.
*/
public void testScorerSupplierTimeoutPoints() throws IOException {
assumeTrue("Test requires more than one segment", reader.leaves().size() > 1);
int size = randomBoolean() ? 0 : randomIntBetween(100, 500);
scorerTimeoutTest(size, context -> {
scorerSupplierTimeoutTest(size, context -> {
PointValues pointValues = context.reader().getPointValues("long");
pointValues.size();
});
}

private void scorerTimeoutTest(int size, CheckedConsumer<LeafReaderContext, IOException> timeoutTrigger) throws IOException {
private void scorerSupplierTimeoutTest(int size, CheckedConsumer<LeafReaderContext, IOException> timeoutTrigger) throws IOException {
{
TimeoutQuery query = newMatchAllScorerTimeoutQuery(timeoutTrigger, false);
TimeoutQuery query = newMatchAllScorerSupplierTimeoutQuery(timeoutTrigger, false);
try (SearchContext context = createSearchContext(query, size)) {
QueryPhase.executeQuery(context);
assertFalse(context.queryResult().searchTimedOut());
Expand All @@ -154,18 +163,20 @@ private void scorerTimeoutTest(int size, CheckedConsumer<LeafReaderContext, IOEx
}
}
{
TimeoutQuery query = newMatchAllScorerTimeoutQuery(timeoutTrigger, true);
TimeoutQuery query = newMatchAllScorerSupplierTimeoutQuery(timeoutTrigger, true);
try (SearchContext context = createSearchContextWithTimeout(query, size)) {
QueryPhase.executeQuery(context);
assertTrue(context.queryResult().searchTimedOut());
int firstSegmentMaxDoc = reader.leaves().get(0).reader().maxDoc();
assertEquals(Math.min(2048, firstSegmentMaxDoc), context.queryResult().topDocs().topDocs.totalHits.value());
// we are artificially raising the timeout when pulling the scorer supplier.
// We score the entire first segment, then trigger timeout.
assertEquals(firstSegmentMaxDoc, context.queryResult().topDocs().topDocs.totalHits.value());
assertEquals(Math.min(size, firstSegmentMaxDoc), context.queryResult().topDocs().topDocs.scoreDocs.length);
}
}
}

private static TimeoutQuery newMatchAllScorerTimeoutQuery(
private static TimeoutQuery newMatchAllScorerSupplierTimeoutQuery(
CheckedConsumer<LeafReaderContext, IOException> timeoutTrigger,
boolean isTimeoutExpected
) {
Expand All @@ -177,6 +188,7 @@ public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float bo

@Override
public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOException {
// trigger the timeout as soon as the scorer supplier is request for the second segment
if (firstSegment == false && isTimeoutExpected) {
shouldTimeout = true;
}
Expand All @@ -190,6 +202,96 @@ public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOExcepti
};
}

/**
* Test that a timeout is appropriately handled when the (exitable) directory reader raises it while loading terms enum
* as the scorer is retrieved from the scorer supplier
*/
public void testScorerGetTimeoutTerms() throws IOException {
assumeTrue("Test requires more than one segment", reader.leaves().size() > 1);
int size = randomBoolean() ? 0 : randomIntBetween(100, 500);
scorerGetTimeoutTest(size, context -> {
final TermsEnum termsEnum = context.reader().terms("field").iterator();
termsEnum.next();
});
}

/**
* Test that a timeout is appropriately handled when the (exitable) directory reader raises it while loading points
* as the scorer is retrieved from the scorer supplier
*/
public void testScorerGetTimeoutPoints() throws IOException {
assumeTrue("Test requires more than one segment", reader.leaves().size() > 1);
int size = randomBoolean() ? 0 : randomIntBetween(100, 500);
scorerGetTimeoutTest(size, context -> {
PointValues pointValues = context.reader().getPointValues("long");
pointValues.size();
});
}

private void scorerGetTimeoutTest(int size, CheckedConsumer<LeafReaderContext, IOException> timeoutTrigger) throws IOException {
{
TimeoutQuery query = newMatchAllScorerGetTimeoutQuery(timeoutTrigger, false);
try (SearchContext context = createSearchContext(query, size)) {
QueryPhase.executeQuery(context);
assertFalse(context.queryResult().searchTimedOut());
assertEquals(numDocs, context.queryResult().topDocs().topDocs.totalHits.value());
assertEquals(size, context.queryResult().topDocs().topDocs.scoreDocs.length);
}
}
{
TimeoutQuery query = newMatchAllScorerGetTimeoutQuery(timeoutTrigger, true);
try (SearchContext context = createSearchContextWithTimeout(query, size)) {
QueryPhase.executeQuery(context);
assertTrue(context.queryResult().searchTimedOut());
int firstSegmentMaxDoc = reader.leaves().get(0).reader().maxDoc();
// we are artificially raising the timeout when pulling the scorer supplier.
// We score the entire first segment, then trigger timeout.
assertEquals(firstSegmentMaxDoc, context.queryResult().topDocs().topDocs.totalHits.value());
assertEquals(Math.min(size, firstSegmentMaxDoc), context.queryResult().topDocs().topDocs.scoreDocs.length);
}
}
}

private static TimeoutQuery newMatchAllScorerGetTimeoutQuery(
CheckedConsumer<LeafReaderContext, IOException> timeoutTrigger,
boolean isTimeoutExpected
) {
return new TimeoutQuery() {
@Override
public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) {
return new MatchAllWeight(this, boost, scoreMode) {
boolean firstSegment = true;

@Override
public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOException {
ScorerSupplier scorerSupplier = super.scorerSupplier(context);
return new ScorerSupplier() {
@Override
public Scorer get(long leadCost) throws IOException {
// trigger the timeout as soon as the scorer is requested for the second segment
if (firstSegment == false && isTimeoutExpected) {
shouldTimeout = true;
}
timeoutTrigger.accept(context);
assert shouldTimeout == false : "should have already timed out";
firstSegment = false;
return scorerSupplier.get(leadCost);
}

@Override
public long cost() {
return scorerSupplier.cost();
}
};
}
};
}
};
}

/**
* Test that a timeout is appropriately handled while bulk scoring, via cancellable bulk scorer
*/
public void testBulkScorerTimeout() throws IOException {
int size = randomBoolean() ? 0 : randomIntBetween(100, 500);
{
Expand All @@ -207,6 +309,8 @@ public void testBulkScorerTimeout() throws IOException {
QueryPhase.executeQuery(context);
assertTrue(context.queryResult().searchTimedOut());
int firstSegmentMaxDoc = reader.leaves().get(0).reader().maxDoc();
// See CancellableBulkScorer#INITIAL_INTERVAL for the source of 2048: we always score the first
// batch of up to 2048 docs, and only then raise the timeout
assertEquals(Math.min(2048, firstSegmentMaxDoc), context.queryResult().topDocs().topDocs.totalHits.value());
assertEquals(Math.min(size, firstSegmentMaxDoc), context.queryResult().topDocs().topDocs.scoreDocs.length);
}
Expand All @@ -233,7 +337,7 @@ public long cost() {
}

@Override
public BulkScorer bulkScorer() throws IOException {
public BulkScorer bulkScorer() {
final float score = score();
final int maxDoc = context.reader().maxDoc();
return new BulkScorer() {
Expand All @@ -251,7 +355,7 @@ public int score(LeafCollector collector, Bits acceptDocs, int min, int max) thr
}
if (timeoutExpected) {
// timeout after collecting the first batch of documents from the 1st segment, or the entire 1st
// segment
// segment if max > firstSegment.maxDoc()
shouldTimeout = true;
}
return max == maxDoc ? DocIdSetIterator.NO_MORE_DOCS : max;
Expand All @@ -274,6 +378,9 @@ private TestSearchContext createSearchContextWithTimeout(TimeoutQuery query, int
TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader)) {
@Override
public long getRelativeTimeInMillis() {
// this controls whether a timeout is raised or not. We abstract time away by pretending that the clock stops
// when a timeout is not expected. The tiniest increment to relative time in millis triggers a timeout.
// See QueryPhase#getTimeoutCheck
return query.shouldTimeout ? 1L : 0L;
}
};
Expand Down