LUCENE-10151: Some fixes to query timeouts. (#996)

I noticed some minor bugs in the original PR #927 that this PR should fix:
 - When a timeout is set, we would no longer catch
   `CollectionTerminatedException`.
 - I added randomization to `LuceneTestCase` to randomly set a timeout, it
   would have caught the above bug.
 - Fixed visibility of `TimeLimitingBulkScorer`.

lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java

@@ -85,7 +85,11 @@ public class IndexSearcher {
   private static QueryCache DEFAULT_QUERY_CACHE;
   private static QueryCachingPolicy DEFAULT_CACHING_POLICY = new UsageTrackingQueryCachingPolicy();
   private QueryTimeout queryTimeout = null;
-  private boolean partialResult = false;
+  // partialResult may be set on one of the threads of the executor. It may be correct to not make
+  // this variable volatile since joining these threads should ensure a happens-before relationship
+  // that guarantees that writes become visible on the main thread, but making the variable volatile
+  // shouldn't hurt either.
+  private volatile boolean partialResult = false;
 
   static {
     final int maxCachedQueries = 1000;
@@ -487,7 +491,8 @@ public class IndexSearcher {
     return search(query, manager);
   }
 
-  public void setTimeout(QueryTimeout queryTimeout) throws IOException {
+  /** Set a {@link QueryTimeout} for all searches that run through this {@link IndexSearcher}. */
+  public void setTimeout(QueryTimeout queryTimeout) {
     this.queryTimeout = queryTimeout;
   }
 
@@ -514,9 +519,11 @@ public class IndexSearcher {
     search(leafContexts, createWeight(query, results.scoreMode(), 1), results);
   }
 
+  /** Returns true if any search hit the {@link #setTimeout(QueryTimeout) timeout}. */
   public boolean timedOut() {
     return partialResult;
   }
+
   /**
    * Search implementation with arbitrary sorting, plus control over whether hit scores and max
    * score should be computed. Finds the top <code>n</code> hits for <code>query</code>, and sorting
@@ -730,29 +737,25 @@ public class IndexSearcher {
       }
       BulkScorer scorer = weight.bulkScorer(ctx);
       if (scorer != null) {
-        if (queryTimeout != null) {
-          TimeLimitingBulkScorer timeLimitingBulkScorer =
-              new TimeLimitingBulkScorer(scorer, queryTimeout);
-          try {
-            timeLimitingBulkScorer.score(leafCollector, ctx.reader().getLiveDocs());
-          } catch (
-              @SuppressWarnings("unused")
-              TimeLimitingBulkScorer.TimeExceededException e) {
-            partialResult = true;
-          }
-        } else {
-          try {
-            scorer.score(leafCollector, ctx.reader().getLiveDocs());
-          } catch (
-              @SuppressWarnings("unused")
-              CollectionTerminatedException e) {
-            // collection was terminated prematurely
-            // continue with the following leaf
-          }
+        if (queryTimeout != null && queryTimeout.isTimeoutEnabled()) {
+          scorer = new TimeLimitingBulkScorer(scorer, queryTimeout);
+        }
+        try {
+          scorer.score(leafCollector, ctx.reader().getLiveDocs());
+        } catch (
+            @SuppressWarnings("unused")
+            CollectionTerminatedException e) {
+          // collection was terminated prematurely
+          // continue with the following leaf
+        } catch (
+            @SuppressWarnings("unused")
+            TimeLimitingBulkScorer.TimeExceededException e) {
+          partialResult = true;
         }
       }
     }
   }
+
   /**
    * Expert: called to re-write queries into primitive queries.
    *

lucene/core/src/java/org/apache/lucene/search/TimeLimitingBulkScorer.java

@@ -28,10 +28,11 @@ import org.apache.lucene.util.Bits;
  *
  * @see org.apache.lucene.index.ExitableDirectoryReader
  */
-public class TimeLimitingBulkScorer extends BulkScorer {
+final class TimeLimitingBulkScorer extends BulkScorer {
   // We score chunks of documents at a time so as to avoid the cost of checking the timeout for
   // every document we score.
   static final int INTERVAL = 100;
+
   /** Thrown when elapsed search time exceeds allowed search time. */
   @SuppressWarnings("serial")
   static class TimeExceededException extends RuntimeException {
@@ -41,8 +42,9 @@ public class TimeLimitingBulkScorer extends BulkScorer {
     }
   }
 
-  private BulkScorer in;
-  private QueryTimeout queryTimeout;
+  private final BulkScorer in;
+  private final QueryTimeout queryTimeout;
+
   /**
    * Create a TimeLimitingBulkScorer wrapper over another {@link BulkScorer} with a specified
    * timeout.

lucene/facet/src/test/org/apache/lucene/facet/TestDrillSideways.java

@@ -132,7 +132,11 @@ public class TestDrillSideways extends FacetTestCase {
   private IndexSearcher getNewSearcher(IndexReader reader) {
     // Do not wrap with an asserting searcher, since DrillSidewaysQuery doesn't
     // implement all the required components like Weight#scorer.
-    return newSearcher(reader, true, false, random().nextBoolean());
+    IndexSearcher searcher = newSearcher(reader, true, false, random().nextBoolean());
+    // DrillSideways requires the entire range of docs to be scored at once, so it doesn't support
+    // timeouts whose implementation scores one window of doc IDs at a time.
+    searcher.setTimeout(null);
+    return searcher;
   }
 
   // See LUCENE-10060:

lucene/facet/src/test/org/apache/lucene/facet/range/TestRangeFacetCounts.java

@@ -458,6 +458,9 @@ public class TestRangeFacetCounts extends FacetTestCase {
     final TaxonomyReader tr = new DirectoryTaxonomyReader(tw);
 
     IndexSearcher s = newSearcher(r, false, false);
+    // DrillSideways requires the entire range of docs to be scored at once, so it doesn't support
+    // timeouts whose implementation scores one window of doc IDs at a time.
+    s.setTimeout(null);
 
     if (VERBOSE) {
       System.out.println("TEST: searcher=" + s);
@@ -1555,6 +1558,9 @@ public class TestRangeFacetCounts extends FacetTestCase {
     IndexReader r = writer.getReader();
 
     IndexSearcher s = newSearcher(r, false, false);
+    // DrillSideways requires the entire range of docs to be scored at once, so it doesn't support
+    // timeouts whose implementation scores one window of doc IDs at a time.
+    s.setTimeout(null);
     FacetsCollector fc = s.search(new MatchAllDocsQuery(), new FacetsCollectorManager());
 
     final DoubleRange[] ranges =

lucene/test-framework/src/java/org/apache/lucene/tests/util/LuceneTestCase.java

@@ -138,6 +138,7 @@ import org.apache.lucene.index.ParallelCompositeReader;
 import org.apache.lucene.index.ParallelLeafReader;
 import org.apache.lucene.index.PointValues;
 import org.apache.lucene.index.PostingsEnum;
+import org.apache.lucene.index.QueryTimeout;
 import org.apache.lucene.index.SerialMergeScheduler;
 import org.apache.lucene.index.SimpleMergedSegmentWarmer;
 import org.apache.lucene.index.SortedDocValues;
@@ -1993,6 +1994,15 @@ public abstract class LuceneTestCase extends Assert {
       }
       ret.setSimilarity(classEnvRule.similarity);
       ret.setQueryCachingPolicy(MAYBE_CACHE_POLICY);
+      if (random().nextBoolean()) {
+        ret.setTimeout(
+            new QueryTimeout() {
+              @Override
+              public boolean shouldExit() {
+                return false;
+              }
+            });
+      }
       return ret;
     }
   }