@@ -116,6 +116,7 @@ public void tearDown() throws Exception {
116116 }
117117
118118 private static ContextIndexSearcher newContextSearcher (IndexReader reader ) throws IOException {
119+ // note that no executor is provided, as this test requires sequential execution
119120 return new ContextIndexSearcher (
120121 reader ,
121122 IndexSearcher .getDefaultSimilarity (),
@@ -125,27 +126,35 @@ private static ContextIndexSearcher newContextSearcher(IndexReader reader) throw
125126 );
126127 }
127128
128- public void testScorerTimeoutTerms () throws IOException {
129+ /**
130+ * Test that a timeout is appropriately handled when the (exitable) directory reader raises it while loading terms enum
131+ * as the scorer supplier is requested.
132+ */
133+ public void testScorerSupplierTimeoutTerms () throws IOException {
129134 assumeTrue ("Test requires more than one segment" , reader .leaves ().size () > 1 );
130135 int size = randomBoolean () ? 0 : randomIntBetween (100 , 500 );
131- scorerTimeoutTest (size , context -> {
136+ scorerSupplierTimeoutTest (size , context -> {
132137 final TermsEnum termsEnum = context .reader ().terms ("field" ).iterator ();
133138 termsEnum .next ();
134139 });
135140 }
136141
137- public void testScorerTimeoutPoints () throws IOException {
142+ /**
143+ * Test that a timeout is appropriately handled when the (exitable) directory reader raises it while loading points
144+ * as the scorer supplier is requested.
145+ */
146+ public void testScorerSupplierTimeoutPoints () throws IOException {
138147 assumeTrue ("Test requires more than one segment" , reader .leaves ().size () > 1 );
139148 int size = randomBoolean () ? 0 : randomIntBetween (100 , 500 );
140- scorerTimeoutTest (size , context -> {
149+ scorerSupplierTimeoutTest (size , context -> {
141150 PointValues pointValues = context .reader ().getPointValues ("long" );
142151 pointValues .size ();
143152 });
144153 }
145154
146- private void scorerTimeoutTest (int size , CheckedConsumer <LeafReaderContext , IOException > timeoutTrigger ) throws IOException {
155+ private void scorerSupplierTimeoutTest (int size , CheckedConsumer <LeafReaderContext , IOException > timeoutTrigger ) throws IOException {
147156 {
148- TimeoutQuery query = newMatchAllScorerTimeoutQuery (timeoutTrigger , false );
157+ TimeoutQuery query = newMatchAllScorerSupplierTimeoutQuery (timeoutTrigger , false );
149158 try (SearchContext context = createSearchContext (query , size )) {
150159 QueryPhase .executeQuery (context );
151160 assertFalse (context .queryResult ().searchTimedOut ());
@@ -154,18 +163,20 @@ private void scorerTimeoutTest(int size, CheckedConsumer<LeafReaderContext, IOEx
154163 }
155164 }
156165 {
157- TimeoutQuery query = newMatchAllScorerTimeoutQuery (timeoutTrigger , true );
166+ TimeoutQuery query = newMatchAllScorerSupplierTimeoutQuery (timeoutTrigger , true );
158167 try (SearchContext context = createSearchContextWithTimeout (query , size )) {
159168 QueryPhase .executeQuery (context );
160169 assertTrue (context .queryResult ().searchTimedOut ());
161170 int firstSegmentMaxDoc = reader .leaves ().get (0 ).reader ().maxDoc ();
162- assertEquals (Math .min (2048 , firstSegmentMaxDoc ), context .queryResult ().topDocs ().topDocs .totalHits .value ());
171+ // we are artificially raising the timeout when pulling the scorer supplier.
172+ // We score the entire first segment, then trigger timeout.
173+ assertEquals (firstSegmentMaxDoc , context .queryResult ().topDocs ().topDocs .totalHits .value ());
163174 assertEquals (Math .min (size , firstSegmentMaxDoc ), context .queryResult ().topDocs ().topDocs .scoreDocs .length );
164175 }
165176 }
166177 }
167178
168- private static TimeoutQuery newMatchAllScorerTimeoutQuery (
179+ private static TimeoutQuery newMatchAllScorerSupplierTimeoutQuery (
169180 CheckedConsumer <LeafReaderContext , IOException > timeoutTrigger ,
170181 boolean isTimeoutExpected
171182 ) {
@@ -177,6 +188,7 @@ public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float bo
177188
178189 @ Override
179190 public ScorerSupplier scorerSupplier (LeafReaderContext context ) throws IOException {
191+ // trigger the timeout as soon as the scorer supplier is request for the second segment
180192 if (firstSegment == false && isTimeoutExpected ) {
181193 shouldTimeout = true ;
182194 }
@@ -190,6 +202,96 @@ public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOExcepti
190202 };
191203 }
192204
205+ /**
206+ * Test that a timeout is appropriately handled when the (exitable) directory reader raises it while loading terms enum
207+ * as the scorer is retrieved from the scorer supplier
208+ */
209+ public void testScorerGetTimeoutTerms () throws IOException {
210+ assumeTrue ("Test requires more than one segment" , reader .leaves ().size () > 1 );
211+ int size = randomBoolean () ? 0 : randomIntBetween (100 , 500 );
212+ scorerGetTimeoutTest (size , context -> {
213+ final TermsEnum termsEnum = context .reader ().terms ("field" ).iterator ();
214+ termsEnum .next ();
215+ });
216+ }
217+
218+ /**
219+ * Test that a timeout is appropriately handled when the (exitable) directory reader raises it while loading points
220+ * as the scorer is retrieved from the scorer supplier
221+ */
222+ public void testScorerGetTimeoutPoints () throws IOException {
223+ assumeTrue ("Test requires more than one segment" , reader .leaves ().size () > 1 );
224+ int size = randomBoolean () ? 0 : randomIntBetween (100 , 500 );
225+ scorerGetTimeoutTest (size , context -> {
226+ PointValues pointValues = context .reader ().getPointValues ("long" );
227+ pointValues .size ();
228+ });
229+ }
230+
231+ private void scorerGetTimeoutTest (int size , CheckedConsumer <LeafReaderContext , IOException > timeoutTrigger ) throws IOException {
232+ {
233+ TimeoutQuery query = newMatchAllScorerGetTimeoutQuery (timeoutTrigger , false );
234+ try (SearchContext context = createSearchContext (query , size )) {
235+ QueryPhase .executeQuery (context );
236+ assertFalse (context .queryResult ().searchTimedOut ());
237+ assertEquals (numDocs , context .queryResult ().topDocs ().topDocs .totalHits .value ());
238+ assertEquals (size , context .queryResult ().topDocs ().topDocs .scoreDocs .length );
239+ }
240+ }
241+ {
242+ TimeoutQuery query = newMatchAllScorerGetTimeoutQuery (timeoutTrigger , true );
243+ try (SearchContext context = createSearchContextWithTimeout (query , size )) {
244+ QueryPhase .executeQuery (context );
245+ assertTrue (context .queryResult ().searchTimedOut ());
246+ int firstSegmentMaxDoc = reader .leaves ().get (0 ).reader ().maxDoc ();
247+ // we are artificially raising the timeout when pulling the scorer supplier.
248+ // We score the entire first segment, then trigger timeout.
249+ assertEquals (firstSegmentMaxDoc , context .queryResult ().topDocs ().topDocs .totalHits .value ());
250+ assertEquals (Math .min (size , firstSegmentMaxDoc ), context .queryResult ().topDocs ().topDocs .scoreDocs .length );
251+ }
252+ }
253+ }
254+
255+ private static TimeoutQuery newMatchAllScorerGetTimeoutQuery (
256+ CheckedConsumer <LeafReaderContext , IOException > timeoutTrigger ,
257+ boolean isTimeoutExpected
258+ ) {
259+ return new TimeoutQuery () {
260+ @ Override
261+ public Weight createWeight (IndexSearcher searcher , ScoreMode scoreMode , float boost ) {
262+ return new MatchAllWeight (this , boost , scoreMode ) {
263+ boolean firstSegment = true ;
264+
265+ @ Override
266+ public ScorerSupplier scorerSupplier (LeafReaderContext context ) throws IOException {
267+ ScorerSupplier scorerSupplier = super .scorerSupplier (context );
268+ return new ScorerSupplier () {
269+ @ Override
270+ public Scorer get (long leadCost ) throws IOException {
271+ // trigger the timeout as soon as the scorer is requested for the second segment
272+ if (firstSegment == false && isTimeoutExpected ) {
273+ shouldTimeout = true ;
274+ }
275+ timeoutTrigger .accept (context );
276+ assert shouldTimeout == false : "should have already timed out" ;
277+ firstSegment = false ;
278+ return scorerSupplier .get (leadCost );
279+ }
280+
281+ @ Override
282+ public long cost () {
283+ return scorerSupplier .cost ();
284+ }
285+ };
286+ }
287+ };
288+ }
289+ };
290+ }
291+
292+ /**
293+ * Test that a timeout is appropriately handled while bulk scoring, via cancellable bulk scorer
294+ */
193295 public void testBulkScorerTimeout () throws IOException {
194296 int size = randomBoolean () ? 0 : randomIntBetween (100 , 500 );
195297 {
@@ -207,6 +309,8 @@ public void testBulkScorerTimeout() throws IOException {
207309 QueryPhase .executeQuery (context );
208310 assertTrue (context .queryResult ().searchTimedOut ());
209311 int firstSegmentMaxDoc = reader .leaves ().get (0 ).reader ().maxDoc ();
312+ // See CancellableBulkScorer#INITIAL_INTERVAL for the source of 2048: we always score the first
313+ // batch of up to 2048 docs, and only then raise the timeout
210314 assertEquals (Math .min (2048 , firstSegmentMaxDoc ), context .queryResult ().topDocs ().topDocs .totalHits .value ());
211315 assertEquals (Math .min (size , firstSegmentMaxDoc ), context .queryResult ().topDocs ().topDocs .scoreDocs .length );
212316 }
@@ -233,7 +337,7 @@ public long cost() {
233337 }
234338
235339 @ Override
236- public BulkScorer bulkScorer () throws IOException {
340+ public BulkScorer bulkScorer () {
237341 final float score = score ();
238342 final int maxDoc = context .reader ().maxDoc ();
239343 return new BulkScorer () {
@@ -251,7 +355,7 @@ public int score(LeafCollector collector, Bits acceptDocs, int min, int max) thr
251355 }
252356 if (timeoutExpected ) {
253357 // timeout after collecting the first batch of documents from the 1st segment, or the entire 1st
254- // segment
358+ // segment if max > firstSegment.maxDoc()
255359 shouldTimeout = true ;
256360 }
257361 return max == maxDoc ? DocIdSetIterator .NO_MORE_DOCS : max ;
@@ -274,6 +378,9 @@ private TestSearchContext createSearchContextWithTimeout(TimeoutQuery query, int
274378 TestSearchContext context = new TestSearchContext (null , indexShard , newContextSearcher (reader )) {
275379 @ Override
276380 public long getRelativeTimeInMillis () {
381+ // this controls whether a timeout is raised or not. We abstract time away by pretending that the clock stops
382+ // when a timeout is not expected. The tiniest increment to relative time in millis triggers a timeout.
383+ // See QueryPhase#getTimeoutCheck
277384 return query .shouldTimeout ? 1L : 0L ;
278385 }
279386 };
0 commit comments