LUCENE-10151: Adding Timeout Support to IndexSearcher (#927)

Authored-by: Deepika Sharma <dpshrma@amazon.com>

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

@@ -37,6 +37,7 @@ import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.IndexReaderContext;
 import org.apache.lucene.index.IndexWriter;
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.QueryTimeout;
 import org.apache.lucene.index.ReaderUtil;
 import org.apache.lucene.index.StoredFieldVisitor;
 import org.apache.lucene.index.Term;
@@ -83,6 +84,8 @@ public class IndexSearcher {
   static int maxClauseCount = 1024;
   private static QueryCache DEFAULT_QUERY_CACHE;
   private static QueryCachingPolicy DEFAULT_CACHING_POLICY = new UsageTrackingQueryCachingPolicy();
+  private QueryTimeout queryTimeout = null;
+  private boolean partialResult = false;
 
   static {
     final int maxCachedQueries = 1000;
@@ -484,6 +487,10 @@ public class IndexSearcher {
     return search(query, manager);
   }
 
+  public void setTimeout(QueryTimeout queryTimeout) throws IOException {
+    this.queryTimeout = queryTimeout;
+  }
+
   /**
    * Finds the top <code>n</code> hits for <code>query</code>.
    *
@@ -507,6 +514,9 @@ public class IndexSearcher {
     search(leafContexts, createWeight(query, results.scoreMode(), 1), results);
   }
 
+  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
@@ -720,18 +730,29 @@ public class IndexSearcher {
       }
       BulkScorer scorer = weight.bulkScorer(ctx);
       if (scorer != null) {
-        try {
-          scorer.score(leafCollector, ctx.reader().getLiveDocs());
-        } catch (
-            @SuppressWarnings("unused")
-            CollectionTerminatedException e) {
-          // collection was terminated prematurely
-          // continue with the following leaf
+        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
+          }
         }
       }
     }
   }
-
   /**
    * Expert: called to re-write queries into primitive queries.
    *

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

@@ -0,0 +1,75 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.lucene.search;
+
+import java.io.IOException;
+import org.apache.lucene.index.QueryTimeout;
+import org.apache.lucene.util.Bits;
+
+/**
+ * The {@link TimeLimitingBulkScorer} is used to timeout search requests that take longer than the
+ * maximum allowed search time limit. After this time is exceeded, the search thread is stopped by
+ * throwing a {@link TimeLimitingBulkScorer.TimeExceededException}.
+ *
+ * @see org.apache.lucene.index.ExitableDirectoryReader
+ */
+public 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 {
+
+    private TimeExceededException() {
+      super("TimeLimit Exceeded");
+    }
+  }
+
+  private BulkScorer in;
+  private QueryTimeout queryTimeout;
+  /**
+   * Create a TimeLimitingBulkScorer wrapper over another {@link BulkScorer} with a specified
+   * timeout.
+   *
+   * @param bulkScorer the wrapped {@link BulkScorer}
+   * @param queryTimeout max time allowed for collecting hits after which {@link
+   *     TimeLimitingBulkScorer.TimeExceededException} is thrown
+   */
+  public TimeLimitingBulkScorer(BulkScorer bulkScorer, QueryTimeout queryTimeout) {
+    this.in = bulkScorer;
+    this.queryTimeout = queryTimeout;
+  }
+
+  @Override
+  public int score(LeafCollector collector, Bits acceptDocs, int min, int max) throws IOException {
+    while (min < max) {
+      final int newMax = (int) Math.min((long) min + INTERVAL, max);
+      if (queryTimeout.shouldExit() == true) {
+        throw new TimeLimitingBulkScorer.TimeExceededException();
+      }
+      min = in.score(collector, acceptDocs, min, newMax); // in is the wrapped bulk scorer
+    }
+    return min;
+  }
+
+  @Override
+  public long cost() {
+    return in.cost();
+  }
+}

lucene/core/src/test/org/apache/lucene/search/TestTimeLimitingBulkScorer.java

@@ -0,0 +1,84 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.lucene.search;
+
+import org.apache.lucene.document.Document;
+import org.apache.lucene.document.Field;
+import org.apache.lucene.index.*;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.tests.analysis.MockAnalyzer;
+import org.apache.lucene.tests.util.LuceneTestCase;
+
+/** Tests the {@link TimeLimitingBulkScorer}. */
+@LuceneTestCase.SuppressSysoutChecks(
+    bugUrl = "http://test.is.timing.sensitive.so.it.prints.instead.of.failing")
+public class TestTimeLimitingBulkScorer extends LuceneTestCase {
+
+  public void testTimeLimitingBulkScorer() throws Exception {
+    Directory directory = newDirectory();
+    IndexWriter writer =
+        new IndexWriter(directory, newIndexWriterConfig(new MockAnalyzer(random())));
+    int n = 10000;
+    for (int i = 0; i < n; i++) {
+      Document d = new Document();
+      d.add(newTextField("default", "ones ", Field.Store.YES));
+      writer.addDocument(d);
+    }
+    writer.forceMerge(1);
+    writer.commit();
+    writer.close();
+
+    DirectoryReader directoryReader;
+    IndexSearcher searcher;
+    TopDocs top;
+    ScoreDoc[] hits = null;
+
+    Query query = new TermQuery(new Term("default", "ones"));
+    directoryReader = DirectoryReader.open(directory);
+    searcher = new IndexSearcher(directoryReader);
+    searcher.setTimeout(CountingQueryTimeout(10));
+    top = searcher.search(query, n);
+    hits = top.scoreDocs;
+    assertTrue(
+        "Partial result and is aborted is true",
+        hits.length > 0 && hits.length < n && searcher.timedOut());
+    directoryReader.close();
+    directory.close();
+  }
+
+  public static QueryTimeout CountingQueryTimeout(int timeallowed) {
+
+    return new QueryTimeout() {
+      public static int counter = 0;
+
+      @Override
+      public boolean shouldExit() {
+        counter++;
+        if (counter == timeallowed) {
+          return true;
+        }
+        return false;
+      }
+
+      @Override
+      public boolean isTimeoutEnabled() {
+        return true;
+      }
+    };
+  }
+}