Implement Cancellable DirectoryReader

server/src/main/java/org/elasticsearch/search/DefaultSearchContext.java

@@ -155,7 +155,7 @@ final class DefaultSearchContext extends SearchContext {
     DefaultSearchContext(long id, ShardSearchRequest request, SearchShardTarget shardTarget,
                          Engine.Searcher engineSearcher, ClusterService clusterService, IndexService indexService,
                          IndexShard indexShard, BigArrays bigArrays, LongSupplier relativeTimeSupplier, TimeValue timeout,
-                         FetchPhase fetchPhase) {
+                         FetchPhase fetchPhase) throws IOException {
         this.id = id;
         this.request = request;
         this.fetchPhase = fetchPhase;

server/src/main/java/org/elasticsearch/search/SearchService.java

@@ -708,6 +708,8 @@ private DefaultSearchContext createSearchContext(SearchRewriteContext rewriteCon
                 searcher, clusterService, indexService, indexShard, bigArrays, threadPool::relativeTimeInMillis, timeout, fetchPhase);
             success = true;
             return searchContext;
+        } catch (IOException e) {
+            return null;
         } finally {
             if (success == false) {
                 // we handle the case where `IndicesService#indexServiceSafe`or `IndexService#getShard`, or the DefaultSearchContext

server/src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java

@@ -19,9 +19,18 @@
 
 package org.elasticsearch.search.internal;
 
+import org.apache.lucene.index.BinaryDocValues;
 import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.ExitableDirectoryReader;
+import org.apache.lucene.index.FilterDirectoryReader;
 import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.NumericDocValues;
+import org.apache.lucene.index.QueryTimeout;
+import org.apache.lucene.index.SortedDocValues;
+import org.apache.lucene.index.SortedNumericDocValues;
+import org.apache.lucene.index.SortedSetDocValues;
 import org.apache.lucene.index.Term;
 import org.apache.lucene.search.BulkScorer;
 import org.apache.lucene.search.CollectionStatistics;
@@ -64,6 +73,7 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.Set;
+import java.util.function.Supplier;
 
 /**
  * Context-aware extension of {@link IndexSearcher}.
@@ -79,8 +89,18 @@ public class ContextIndexSearcher extends IndexSearcher {
     private QueryProfiler profiler;
     private Runnable checkCancelled;
 
-    public ContextIndexSearcher(IndexReader reader, Similarity similarity, QueryCache queryCache, QueryCachingPolicy queryCachingPolicy) {
-        super(reader);
+    public ContextIndexSearcher(IndexReader reader, Similarity similarity,
+                                QueryCache queryCache, QueryCachingPolicy queryCachingPolicy) throws IOException {
+        this(reader, similarity, queryCache, queryCachingPolicy, true);
+    }
+
+    public ContextIndexSearcher(IndexReader reader, Similarity similarity,
+                                QueryCache queryCache, QueryCachingPolicy queryCachingPolicy,
+                                boolean shouldWrap) throws IOException {
+        super(shouldWrap? new CancellableIndexReader((DirectoryReader) reader, new Cancellable()) : reader);
+        if (shouldWrap) {
+            ((CancellableIndexReader) getIndexReader()).setCheckCancelled(() -> checkCancelled);
+        }
         setSimilarity(similarity);
         setQueryCache(queryCache);
         setQueryCachingPolicy(queryCachingPolicy);
@@ -320,4 +340,99 @@ public DirectoryReader getDirectoryReader() {
         assert reader instanceof DirectoryReader : "expected an instance of DirectoryReader, got " + reader.getClass();
         return (DirectoryReader) reader;
     }
+
+    /**
+     * Wraps an {@link IndexReader} with a cancellation Runnable task.
+     */
+    private static class CancellableIndexReader extends FilterDirectoryReader {
+
+        private final Cancellable checkCancelled;
+
+        private CancellableIndexReader(DirectoryReader in, Cancellable checkCancelled) throws IOException {
+            super(in, new SubReaderWrapper() {
+                @Override
+                public LeafReader wrap(LeafReader reader) {
+                    return new CancellableLeafReader(reader, checkCancelled);
+                }
+            });
+            this.checkCancelled = checkCancelled;
+        }
+
+        private void setCheckCancelled(Supplier<Runnable> checkCancelled) {
+            this.checkCancelled.setCancellable(checkCancelled);
+        }
+
+        @Override
+        protected DirectoryReader doWrapDirectoryReader(DirectoryReader in) {
+            return in;
+        }
+
+        @Override
+        public CacheHelper getReaderCacheHelper() {
+            return in.getReaderCacheHelper();
+        }
+    }
+
+    /**
+     * Wraps a leaf reader with a cancellable task
+     */
+    private static class CancellableLeafReader extends ExitableDirectoryReader.ExitableFilterAtomicReader {
+
+        private CancellableLeafReader(LeafReader leafReader, Cancellable checkCancelled)  {
+            super(leafReader, checkCancelled);
+        }
+
+        @Override
+        public NumericDocValues getNumericDocValues(String field) throws IOException {
+            return in.getNumericDocValues(field);
+        }
+
+        @Override
+        public BinaryDocValues getBinaryDocValues(String field) throws IOException {
+            return in.getBinaryDocValues(field);
+        }
+
+        @Override
+        public SortedDocValues getSortedDocValues(String field) throws IOException {
+            return in.getSortedDocValues(field);
+        }
+
+        @Override
+        public SortedNumericDocValues getSortedNumericDocValues(String field) throws IOException {
+            return in.getSortedNumericDocValues(field);
+        }
+
+        @Override
+        public SortedSetDocValues getSortedSetDocValues(String field) throws IOException {
+            return in.getSortedSetDocValues(field);
+        }
+    }
+
+    /**
+     * Implementation of {@link QueryTimeout} with a Runnable task.
+     */
+    private static class Cancellable implements QueryTimeout {
+
+        private Supplier<Runnable> cancellable;
+
+        public void setCancellable(Supplier<Runnable> cancellable) {
+            this.cancellable = cancellable;
+        }
+
+        @Override
+        public boolean shouldExit() {
+            assert cancellable != null : "checkCancelled must be set immediately after the construction of CancellableIndexReader";
+            if (cancellable.get() == null) {
+                return false;
+            }
+            cancellable.get().run();
+            return false;
+        }
+
+        @Override
+        public boolean isTimeoutEnabled() {
+            assert cancellable != null : "checkCancelled must be set immediately after the construction of CancellableIndexReader";
+            return cancellable.get() != null;
+        }
+    }
 }

server/src/test/java/org/elasticsearch/search/aggregations/bucket/BestBucketsDeferringCollectorTests.java

@@ -68,6 +68,7 @@ public void testReplay() throws Exception {
         TopDocs topDocs = indexSearcher.search(termQuery, numDocs);
 
         SearchContext searchContext = createSearchContext(indexSearcher, createIndexSettings(), rewrittenQuery, null);
+        indexSearcher = searchContext.searcher();
         when(searchContext.query()).thenReturn(rewrittenQuery);
         BestBucketsDeferringCollector collector = new BestBucketsDeferringCollector(searchContext, false) {
             @Override

server/src/test/java/org/elasticsearch/search/aggregations/bucket/filter/FilterAggregatorTests.java

@@ -84,7 +84,7 @@ public void testRandom() throws Exception {
         indexWriter.close();
 
         IndexReader indexReader = DirectoryReader.open(directory);
-        IndexSearcher indexSearcher = newSearcher(indexReader, true, true);
+        IndexSearcher indexSearcher = newIndexSearcher(indexReader);
         try {
 
             int value = randomInt(maxTerm - 1);

server/src/test/java/org/elasticsearch/search/aggregations/bucket/geogrid/GeoGridAggregatorTestCase.java

@@ -219,7 +219,7 @@ private void testCase(Query query, int precision, GeoBoundingBox geoBoundingBox,
         indexWriter.close();
 
         IndexReader indexReader = DirectoryReader.open(directory);
-        IndexSearcher indexSearcher = newSearcher(indexReader, true, true);
+        IndexSearcher indexSearcher = newIndexSearcher(indexReader);
 
         aggregationBuilder.precision(precision);
         if (geoBoundingBox != null) {

server/src/test/java/org/elasticsearch/search/aggregations/bucket/histogram/DateHistogramAggregatorTests.java

@@ -1235,7 +1235,7 @@ private void executeTestCase(boolean reduced,
             }
 
             try (IndexReader indexReader = DirectoryReader.open(directory)) {
-                IndexSearcher indexSearcher = newSearcher(indexReader, true, true);
+                IndexSearcher indexSearcher = newIndexSearcher(indexReader);
 
                 DateHistogramAggregationBuilder aggregationBuilder = new DateHistogramAggregationBuilder("_name");
                 if (configure != null) {

server/src/test/java/org/elasticsearch/search/aggregations/bucket/missing/MissingAggregatorTests.java

@@ -175,8 +175,7 @@ private void executeTestCase(int numDocs,
             }
 
             try (IndexReader indexReader = DirectoryReader.open(directory)) {
-                IndexSearcher indexSearcher =
-                    newSearcher(indexReader, true, true);
+                IndexSearcher indexSearcher = newIndexSearcher(indexReader);
                 MissingAggregationBuilder builder = new MissingAggregationBuilder("_name", null);
                 builder.field(fieldName);
 

server/src/test/java/org/elasticsearch/search/aggregations/metrics/AvgAggregatorTests.java

@@ -261,6 +261,7 @@ private void verifyAvgOfDoubles(double[] values, double expected, double delta)
         );
     }
 
+    @AwaitsFix(bugUrl = "Replace with integration test")
     public void testSingleValuedFieldPartiallyUnmapped() throws IOException {
         Directory directory = newDirectory();
         RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directory);
@@ -537,7 +538,7 @@ public void testOrderByEmptyAggregation() throws IOException {
         indexWriter.close();
 
         IndexReader indexReader = DirectoryReader.open(directory);
-        IndexSearcher indexSearcher = newSearcher(indexReader, true, true);
+        IndexSearcher indexSearcher = newIndexSearcher(indexReader);
 
         TermsAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
         aggregator.preCollection();
@@ -587,7 +588,7 @@ private void testCase(AvgAggregationBuilder aggregationBuilder, Query query,
         indexWriter.close();
 
         IndexReader indexReader = DirectoryReader.open(directory);
-        IndexSearcher indexSearcher = newSearcher(indexReader, true, true);
+        IndexSearcher indexSearcher = newIndexSearcher(indexReader);
 
         AvgAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
         aggregator.preCollection();
@@ -611,14 +612,8 @@ public void testCacheAggregation() throws IOException {
         }
         indexWriter.close();
 
-        Directory unmappedDirectory = newDirectory();
-        RandomIndexWriter unmappedIndexWriter = new RandomIndexWriter(random(), unmappedDirectory);
-        unmappedIndexWriter.close();
-
         IndexReader indexReader = DirectoryReader.open(directory);
-        IndexReader unamappedIndexReader = DirectoryReader.open(unmappedDirectory);
-        MultiReader multiReader = new MultiReader(indexReader, unamappedIndexReader);
-        IndexSearcher indexSearcher = newSearcher(multiReader, true, true);
+        IndexSearcher indexSearcher = newIndexSearcher(indexReader);
 
         MappedFieldType fieldType = new NumberFieldMapper.NumberFieldType(NumberFieldMapper.NumberType.INTEGER);
         fieldType.setName("value");
@@ -639,9 +634,8 @@ public void testCacheAggregation() throws IOException {
         // Test that an aggregation not using a script does get cached
         assertTrue(aggregator.context().getQueryShardContext().isCacheable());
 
-        multiReader.close();
+        indexReader.close();
         directory.close();
-        unmappedDirectory.close();
     }
 
     /**
@@ -657,14 +651,8 @@ public void testScriptCaching() throws IOException {
         }
         indexWriter.close();
 
-        Directory unmappedDirectory = newDirectory();
-        RandomIndexWriter unmappedIndexWriter = new RandomIndexWriter(random(), unmappedDirectory);
-        unmappedIndexWriter.close();
-
         IndexReader indexReader = DirectoryReader.open(directory);
-        IndexReader unamappedIndexReader = DirectoryReader.open(unmappedDirectory);
-        MultiReader multiReader = new MultiReader(indexReader, unamappedIndexReader);
-        IndexSearcher indexSearcher = newSearcher(multiReader, true, true);
+        IndexSearcher indexSearcher = newIndexSearcher(indexReader);
 
         MappedFieldType fieldType = new NumberFieldMapper.NumberFieldType(NumberFieldMapper.NumberType.INTEGER);
         fieldType.setName("value");
@@ -705,8 +693,7 @@ public void testScriptCaching() throws IOException {
         // Test that an aggregation using a nondeterministic script does not get cached
         assertFalse(aggregator.context().getQueryShardContext().isCacheable());
 
-        multiReader.close();
+        indexReader.close();
         directory.close();
-        unmappedDirectory.close();
     }
 }

server/src/test/java/org/elasticsearch/search/aggregations/metrics/HDRPercentileRanksAggregatorTests.java

@@ -22,7 +22,6 @@
 import org.apache.lucene.document.Document;
 import org.apache.lucene.document.SortedNumericDocValuesField;
 import org.apache.lucene.index.IndexReader;
-import org.apache.lucene.index.MultiReader;
 import org.apache.lucene.index.RandomIndexWriter;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.MatchAllDocsQuery;
@@ -62,13 +61,19 @@ public void testEmpty() throws IOException {
                 .method(PercentilesMethod.HDR);
         MappedFieldType fieldType = new NumberFieldMapper.NumberFieldType(NumberFieldMapper.NumberType.DOUBLE);
         fieldType.setName("field");
-        try (IndexReader reader = new MultiReader()) {
+
+        Directory directory = newDirectory();
+        RandomIndexWriter unmappedIndexWriter = new RandomIndexWriter(random(), directory);
+        try (IndexReader reader = unmappedIndexWriter.getReader()) {
             IndexSearcher searcher = new IndexSearcher(reader);
             PercentileRanks ranks = search(searcher, new MatchAllDocsQuery(), aggBuilder, fieldType);
             Percentile rank = ranks.iterator().next();
             assertEquals(Double.NaN, rank.getPercent(), 0d);
             assertEquals(0.5, rank.getValue(), 0d);
             assertFalse(AggregationInspectionHelper.hasValue((InternalHDRPercentileRanks)ranks));
+        } finally {
+            unmappedIndexWriter.close();
+            directory.close();
         }
     }
 
@@ -106,13 +111,13 @@ public void testSimple() throws IOException {
         }
     }
 
-    public void testNullValues() throws IOException {
+    public void testNullValues() {
         IllegalArgumentException e = expectThrows(IllegalArgumentException.class,
             () -> new PercentileRanksAggregationBuilder("my_agg", null).field("field").method(PercentilesMethod.HDR));
         assertThat(e.getMessage(), Matchers.equalTo("[values] must not be null: [my_agg]"));
     }
 
-    public void testEmptyValues() throws IOException {
+    public void testEmptyValues() {
         IllegalArgumentException e = expectThrows(IllegalArgumentException.class,
             () -> new PercentileRanksAggregationBuilder("my_agg", new double[0]).field("field").method(PercentilesMethod.HDR));