diff --git a/server/src/test/java/org/elasticsearch/search/query/QueryPhaseTimeoutTests.java b/server/src/test/java/org/elasticsearch/search/query/QueryPhaseTimeoutTests.java index aa9ca26974bad..b0d1c6b679892 100644 --- a/server/src/test/java/org/elasticsearch/search/query/QueryPhaseTimeoutTests.java +++ b/server/src/test/java/org/elasticsearch/search/query/QueryPhaseTimeoutTests.java @@ -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(), @@ -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 timeoutTrigger) throws IOException { + private void scorerSupplierTimeoutTest(int size, CheckedConsumer 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()); @@ -154,18 +163,20 @@ private void scorerTimeoutTest(int size, CheckedConsumer timeoutTrigger, boolean isTimeoutExpected ) { @@ -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; } @@ -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 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 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); { @@ -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); } @@ -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() { @@ -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; @@ -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; } };