Address test issue in QueryPhaseTimeoutTests (elastic#124327) (elastic#124515)

The test that is verifying timeout handling when pulling a scorer supplier is
going to always score the entire segment. The test needs to be adjusted
accordingly. While at it, I added docs and clarified the tests a bit, as well
as added a couple new tests that cover for timeout handling when retrieving a
scorer from the scorer supplier.

Closes elastic#124140

Author: javanna

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<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());
@@ -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
     ) {
@@ -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<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);
         {
@@ -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;
             }
         };