LUCENE-8026: ExitableDirectoryReader does not instrument points

Closes #497

Signed-off-by: Adrien Grand <jpountz@gmail.com>

lucene/CHANGES.txt

@@ -198,6 +198,12 @@ Build:
 
 * LUCENE-8537: ant test command fails under lucene/tools (Peter Somogyi)
 
+New Features
+
+* LUCENE-8026: ExitableDirectoryReader may now time out queries that run on
+  points such as range queries or geo queries.
+  (Christophe Bismuth via Adrien Grand)
+
 ======================= Lucene 7.6.0 =======================
 
 Build:

lucene/core/src/java/org/apache/lucene/index/ExitableDirectoryReader.java

@@ -77,6 +77,15 @@ public class ExitableDirectoryReader extends FilterDirectoryReader {
       this.queryTimeout = queryTimeout;
     }
 
+    @Override
+    public PointValues getPointValues(String field) throws IOException {
+      final PointValues pointValues = in.getPointValues(field);
+      if (pointValues == null) {
+        return null;
+      }
+      return (queryTimeout.isTimeoutEnabled()) ? new ExitablePointValues(pointValues, queryTimeout) : pointValues;
+    }
+
     @Override
     public Terms terms(String field) throws IOException {
       Terms terms = in.terms(field);
@@ -100,13 +109,156 @@ public class ExitableDirectoryReader extends FilterDirectoryReader {
 
   }
 
+  /**
+   * Wrapper class for another PointValues implementation that is used by ExitableFields.
+   */
+  private static class ExitablePointValues extends PointValues {
+
+    private final PointValues in;
+    private final QueryTimeout queryTimeout;
+
+    private ExitablePointValues(PointValues in, QueryTimeout queryTimeout) {
+      this.in = in;
+      this.queryTimeout = queryTimeout;
+      checkAndThrow();
+    }
+
+    /**
+     * Throws {@link ExitingReaderException} if {@link QueryTimeout#shouldExit()} returns true,
+     * or if {@link Thread#interrupted()} returns true.
+     */
+    private void checkAndThrow() {
+      if (queryTimeout.shouldExit()) {
+        throw new ExitingReaderException("The request took too long to iterate over point values. Timeout: "
+            + queryTimeout.toString()
+            + ", PointValues=" + in
+        );
+      } else if (Thread.interrupted()) {
+        throw new ExitingReaderException("Interrupted while iterating over point values. PointValues=" + in);
+      }
+    }
+
+    @Override
+    public void intersect(IntersectVisitor visitor) throws IOException {
+      checkAndThrow();
+      in.intersect(new ExitableIntersectVisitor(visitor, queryTimeout));
+    }
+
+    @Override
+    public long estimatePointCount(IntersectVisitor visitor) {
+      checkAndThrow();
+      return in.estimatePointCount(visitor);
+    }
+
+    @Override
+    public byte[] getMinPackedValue() throws IOException {
+      checkAndThrow();
+      return in.getMinPackedValue();
+    }
+
+    @Override
+    public byte[] getMaxPackedValue() throws IOException {
+      checkAndThrow();
+      return in.getMaxPackedValue();
+    }
+
+    @Override
+    public int getNumDataDimensions() throws IOException {
+      checkAndThrow();
+      return in.getNumDataDimensions();
+    }
+
+    @Override
+    public int getNumIndexDimensions() throws IOException {
+      checkAndThrow();
+      return in.getNumIndexDimensions();
+    }
+
+    @Override
+    public int getBytesPerDimension() throws IOException {
+      checkAndThrow();
+      return in.getBytesPerDimension();
+    }
+
+    @Override
+    public long size() {
+      checkAndThrow();
+      return in.size();
+    }
+
+    @Override
+    public int getDocCount() {
+      checkAndThrow();
+      return in.getDocCount();
+    }
+  }
+
+  private static class ExitableIntersectVisitor implements PointValues.IntersectVisitor {
+
+    private static final int MAX_CALLS_BEFORE_QUERY_TIMEOUT_CHECK = 10;
+
+    private final PointValues.IntersectVisitor in;
+    private final QueryTimeout queryTimeout;
+    private int calls;
+
+    private ExitableIntersectVisitor(PointValues.IntersectVisitor in, QueryTimeout queryTimeout) {
+      this.in = in;
+      this.queryTimeout = queryTimeout;
+    }
+
+    /**
+     * Throws {@link ExitingReaderException} if {@link QueryTimeout#shouldExit()} returns true,
+     * or if {@link Thread#interrupted()} returns true.
+     */
+    private void checkAndThrowWithSampling() {
+      if (calls++ % MAX_CALLS_BEFORE_QUERY_TIMEOUT_CHECK == 0) {
+        checkAndThrow();
+      }
+    }
+
+    private void checkAndThrow() {
+      if (queryTimeout.shouldExit()) {
+        throw new ExitingReaderException("The request took too long to intersect point values. Timeout: "
+            + queryTimeout.toString()
+            + ", PointValues=" + in
+        );
+      } else if (Thread.interrupted()) {
+        throw new ExitingReaderException("Interrupted while intersecting point values. PointValues=" + in);
+      }
+    }
+
+    @Override
+    public void visit(int docID) throws IOException {
+      checkAndThrowWithSampling();
+      in.visit(docID);
+    }
+
+    @Override
+    public void visit(int docID, byte[] packedValue) throws IOException {
+      checkAndThrowWithSampling();
+      in.visit(docID, packedValue);
+    }
+
+    @Override
+    public PointValues.Relation compare(byte[] minPackedValue, byte[] maxPackedValue) {
+      checkAndThrow();
+      return in.compare(minPackedValue, maxPackedValue);
+    }
+
+    @Override
+    public void grow(int count) {
+      checkAndThrow();
+      in.grow(count);
+    }
+  }
+
   /**
    * Wrapper class for another Terms implementation that is used by ExitableFields.
    */
   public static class ExitableTerms extends FilterTerms {
 
     private QueryTimeout queryTimeout;
-    
+
     /** Constructor **/
     public ExitableTerms(Terms terms, QueryTimeout queryTimeout) {
       super(terms);

lucene/core/src/test/org/apache/lucene/index/TestExitableDirectoryReader.java

@@ -21,6 +21,7 @@ import java.io.IOException;
 import org.apache.lucene.analysis.MockAnalyzer;
 import org.apache.lucene.document.Document;
 import org.apache.lucene.document.Field;
+import org.apache.lucene.document.IntPoint;
 import org.apache.lucene.index.ExitableDirectoryReader.ExitingReaderException;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.PrefixQuery;
@@ -28,7 +29,6 @@ import org.apache.lucene.search.Query;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.LuceneTestCase;
-import org.junit.Ignore;
 
 /**
  * Test that uses a default/lucene Implementation of {@link QueryTimeout}
@@ -92,8 +92,7 @@ public class TestExitableDirectoryReader extends LuceneTestCase {
    * Tests timing out of TermsEnum iterations
    * @throws Exception on error
    */
-  @Ignore("this test relies on wall clock time and sometimes false fails")
-  public void testExitableFilterIndexReader() throws Exception {
+  public void testExitableFilterTermsIndexReader() throws Exception {
     Directory directory = newDirectory();
     IndexWriter writer = new IndexWriter(directory, newIndexWriterConfig(new MockAnalyzer(random())));
 
@@ -108,8 +107,8 @@ public class TestExitableDirectoryReader extends LuceneTestCase {
     Document d3 = new Document();
     d3.add(newTextField("default", "ones two four", Field.Store.YES));
     writer.addDocument(d3);
-    writer.forceMerge(1);
 
+    writer.forceMerge(1);
     writer.commit();
     writer.close();
 
@@ -120,19 +119,18 @@ public class TestExitableDirectoryReader extends LuceneTestCase {
 
     Query query = new PrefixQuery(new Term("default", "o"));
 
-    // Set a fairly high timeout value (1 second) and expect the query to complete in that time frame.
+    // Set a fairly high timeout value (infinite) and expect the query to complete in that time frame.
     // Not checking the validity of the result, all we are bothered about in this test is the timing out.
     directoryReader = DirectoryReader.open(directory);
-    exitableDirectoryReader = new ExitableDirectoryReader(directoryReader, new QueryTimeoutImpl(1000));
+    exitableDirectoryReader = new ExitableDirectoryReader(directoryReader, infiniteQueryTimeout());
     reader = new TestReader(getOnlyLeafReader(exitableDirectoryReader));
     searcher = new IndexSearcher(reader);
     searcher.search(query, 10);
     reader.close();
 
-
-    // Set a really low timeout value (1 millisecond) and expect an Exception
+    // Set a really low timeout value (immediate) and expect an Exception
     directoryReader = DirectoryReader.open(directory);
-    exitableDirectoryReader = new ExitableDirectoryReader(directoryReader, new QueryTimeoutImpl(1));
+    exitableDirectoryReader = new ExitableDirectoryReader(directoryReader, immediateQueryTimeout());
     reader = new TestReader(getOnlyLeafReader(exitableDirectoryReader));
     IndexSearcher slowSearcher = new IndexSearcher(reader);
     expectThrows(ExitingReaderException.class, () -> {
@@ -143,7 +141,7 @@ public class TestExitableDirectoryReader extends LuceneTestCase {
     // Set maximum time out and expect the query to complete. 
     // Not checking the validity of the result, all we are bothered about in this test is the timing out.
     directoryReader = DirectoryReader.open(directory);
-    exitableDirectoryReader = new ExitableDirectoryReader(directoryReader, new QueryTimeoutImpl(Long.MAX_VALUE));
+    exitableDirectoryReader = new ExitableDirectoryReader(directoryReader, infiniteQueryTimeout());
     reader = new TestReader(getOnlyLeafReader(exitableDirectoryReader));
     searcher = new IndexSearcher(reader);
     searcher.search(query, 10);
@@ -152,7 +150,7 @@ public class TestExitableDirectoryReader extends LuceneTestCase {
     // Set a negative time allowed and expect the query to complete (should disable timeouts)
     // Not checking the validity of the result, all we are bothered about in this test is the timing out.
     directoryReader = DirectoryReader.open(directory);
-    exitableDirectoryReader = new ExitableDirectoryReader(directoryReader, new QueryTimeoutImpl(-189034L));
+    exitableDirectoryReader = new ExitableDirectoryReader(directoryReader, disabledQueryTimeout());
     reader = new TestReader(getOnlyLeafReader(exitableDirectoryReader));
     searcher = new IndexSearcher(reader);
     searcher.search(query, 10);
@@ -160,5 +158,122 @@ public class TestExitableDirectoryReader extends LuceneTestCase {
 
     directory.close();
   }
+
+  /**
+   * Tests timing out of PointValues queries
+   *
+   * @throws Exception on error
+   */
+  public void testExitablePointValuesIndexReader() throws Exception {
+    Directory directory = newDirectory();
+    IndexWriter writer = new IndexWriter(directory, newIndexWriterConfig(new MockAnalyzer(random())));
+
+    Document d1 = new Document();
+    d1.add(new IntPoint("default", 10));
+    writer.addDocument(d1);
+
+    Document d2 = new Document();
+    d2.add(new IntPoint("default", 100));
+    writer.addDocument(d2);
+
+    Document d3 = new Document();
+    d3.add(new IntPoint("default", 1000));
+    writer.addDocument(d3);
+
+    writer.forceMerge(1);
+    writer.commit();
+    writer.close();
+
+    DirectoryReader directoryReader;
+    DirectoryReader exitableDirectoryReader;
+    IndexReader reader;
+    IndexSearcher searcher;
+
+    Query query = IntPoint.newRangeQuery("default", 10, 20);
+
+    // Set a fairly high timeout value (infinite) and expect the query to complete in that time frame.
+    // Not checking the validity of the result, all we are bothered about in this test is the timing out.
+    directoryReader = DirectoryReader.open(directory);
+    exitableDirectoryReader = new ExitableDirectoryReader(directoryReader, infiniteQueryTimeout());
+    reader = new TestReader(getOnlyLeafReader(exitableDirectoryReader));
+    searcher = new IndexSearcher(reader);
+    searcher.search(query, 10);
+    reader.close();
+
+    // Set a really low timeout value (immediate) and expect an Exception
+    directoryReader = DirectoryReader.open(directory);
+    exitableDirectoryReader = new ExitableDirectoryReader(directoryReader, immediateQueryTimeout());
+    reader = new TestReader(getOnlyLeafReader(exitableDirectoryReader));
+    IndexSearcher slowSearcher = new IndexSearcher(reader);
+    expectThrows(ExitingReaderException.class, () -> {
+      slowSearcher.search(query, 10);
+    });
+    reader.close();
+
+    // Set maximum time out and expect the query to complete.
+    // Not checking the validity of the result, all we are bothered about in this test is the timing out.
+    directoryReader = DirectoryReader.open(directory);
+    exitableDirectoryReader = new ExitableDirectoryReader(directoryReader, infiniteQueryTimeout());
+    reader = new TestReader(getOnlyLeafReader(exitableDirectoryReader));
+    searcher = new IndexSearcher(reader);
+    searcher.search(query, 10);
+    reader.close();
+
+    // Set a negative time allowed and expect the query to complete (should disable timeouts)
+    // Not checking the validity of the result, all we are bothered about in this test is the timing out.
+    directoryReader = DirectoryReader.open(directory);
+    exitableDirectoryReader = new ExitableDirectoryReader(directoryReader, disabledQueryTimeout());
+    reader = new TestReader(getOnlyLeafReader(exitableDirectoryReader));
+    searcher = new IndexSearcher(reader);
+    searcher.search(query, 10);
+    reader.close();
+
+    directory.close();
+  }
+
+  private static QueryTimeout disabledQueryTimeout() {
+    return new QueryTimeout() {
+
+      @Override
+      public boolean shouldExit() {
+        return false;
+      }
+
+      @Override
+      public boolean isTimeoutEnabled() {
+        return false;
+      }
+    };
+  }
+
+  private static QueryTimeout infiniteQueryTimeout() {
+    return new QueryTimeout() {
+
+      @Override
+      public boolean shouldExit() {
+        return false;
+      }
+
+      @Override
+      public boolean isTimeoutEnabled() {
+        return true;
+      }
+    };
+  }
+
+  private static QueryTimeout immediateQueryTimeout() {
+    return new QueryTimeout() {
+
+      @Override
+      public boolean shouldExit() {
+        return true;
+      }
+
+      @Override
+      public boolean isTimeoutEnabled() {
+        return true;
+      }
+    };
+  }
 }