Fix IndexSearcherWrapper visibility (#39071) (#40145)

This change adds a wrapper for IndexSearcher that makes IndexSearcher#search(List, Weight, Collector) visible by
sub-classes. The wrapper is used by the ContextIndexSearcher to call this protected method on a searcher created by a plugin.
This ensures that an override of the protected method in an IndexSearcherWrapper plugin is called when a search is executed.

Closes #30758

server/src/main/java/org/apache/lucene/search/XIndexSearcher.java

@@ -0,0 +1,46 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch 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.index.LeafReaderContext;
+
+import java.io.IOException;
+import java.util.List;
+
+/**
+ * A wrapper for {@link IndexSearcher} that makes {@link IndexSearcher#search(List, Weight, Collector)}
+ * visible by sub-classes.
+ */
+public class XIndexSearcher extends IndexSearcher {
+    private final IndexSearcher in;
+
+    public XIndexSearcher(IndexSearcher in) {
+        super(in.getIndexReader());
+        this.in = in;
+        setSimilarity(in.getSimilarity());
+        setQueryCache(in.getQueryCache());
+        setQueryCachingPolicy(in.getQueryCachingPolicy());
+    }
+
+    @Override
+    public void search(List<LeafReaderContext> leaves, Weight weight, Collector collector) throws IOException {
+        in.search(leaves, weight, collector);
+    }
+}

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

@@ -35,6 +35,7 @@ import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.Scorer;
 import org.apache.lucene.search.TermStatistics;
 import org.apache.lucene.search.Weight;
+import org.apache.lucene.search.XIndexSearcher;
 import org.elasticsearch.common.lease.Releasable;
 import org.elasticsearch.index.engine.Engine;
 import org.elasticsearch.search.dfs.AggregatedDfs;
@@ -56,7 +57,7 @@ public class ContextIndexSearcher extends IndexSearcher implements Releasable {
     /** The wrapped {@link IndexSearcher}. The reason why we sometimes prefer delegating to this searcher instead of {@code super} is that
      *  this instance may have more assertions, for example if it comes from MockInternalEngine which wraps the IndexSearcher into an
      *  AssertingIndexSearcher. */
-    private final IndexSearcher in;
+    private final XIndexSearcher in;
 
     private AggregatedDfs aggregatedDfs;
 
@@ -67,11 +68,10 @@ public class ContextIndexSearcher extends IndexSearcher implements Releasable {
 
     private Runnable checkCancelled;
 
-    public ContextIndexSearcher(Engine.Searcher searcher,
+    public ContextIndexSearcher(Engine.Searcher searcher, QueryCache queryCache, QueryCachingPolicy queryCachingPolicy) {
-            QueryCache queryCache, QueryCachingPolicy queryCachingPolicy) {
         super(searcher.reader());
-        in = searcher.searcher();
         engineSearcher = searcher;
+        in = new XIndexSearcher(searcher.searcher());
         setSimilarity(searcher.searcher().getSimilarity());
         setQueryCache(queryCache);
         setQueryCachingPolicy(queryCachingPolicy);
@@ -174,7 +174,7 @@ public class ContextIndexSearcher extends IndexSearcher implements Releasable {
         } else {
             cancellableWeight = weight;
         }
-        super.search(leaves, cancellableWeight, collector);
+        in.search(leaves, cancellableWeight, collector);
     }
 
     @Override

server/src/test/java/org/elasticsearch/index/shard/IndexSearcherWrapperTests.java

@@ -28,23 +28,31 @@ import org.apache.lucene.index.FilterDirectoryReader;
 import org.apache.lucene.index.IndexWriter;
 import org.apache.lucene.index.IndexWriterConfig;
 import org.apache.lucene.index.LeafReader;
+import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.Term;
+import org.apache.lucene.search.Collector;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.TermQuery;
 import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.search.TotalHitCountCollector;
+import org.apache.lucene.search.Weight;
 import org.apache.lucene.store.Directory;
 import org.elasticsearch.core.internal.io.IOUtils;
 import org.elasticsearch.common.lucene.index.ElasticsearchDirectoryReader;
 import org.elasticsearch.index.engine.Engine;
 import org.elasticsearch.index.engine.EngineException;
+import org.elasticsearch.search.internal.ContextIndexSearcher;
 import org.elasticsearch.test.ESTestCase;
 
 import java.io.IOException;
 import java.util.Collections;
+import java.util.List;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import static org.hamcrest.Matchers.equalTo;
+
 public class IndexSearcherWrapperTests extends ESTestCase {
 
     public void testReaderCloseListenerIsCalled() throws IOException {
@@ -159,6 +167,50 @@ public class IndexSearcherWrapperTests extends ESTestCase {
         IOUtils.close(writer, dir);
     }
 
+    public void testWrapVisibility() throws IOException {
+        Directory dir = newDirectory();
+        IndexWriterConfig iwc = newIndexWriterConfig();
+        IndexWriter writer = new IndexWriter(dir, iwc);
+        Document doc = new Document();
+        doc.add(new StringField("id", "1", random().nextBoolean() ? Field.Store.YES : Field.Store.NO));
+        doc.add(new TextField("field", "doc", random().nextBoolean() ? Field.Store.YES : Field.Store.NO));
+        writer.addDocument(doc);
+        DirectoryReader open = ElasticsearchDirectoryReader.wrap(DirectoryReader.open(writer), new ShardId("foo", "_na_", 1));
+        IndexSearcher searcher = new IndexSearcher(open);
+        assertEquals(1, searcher.search(new TermQuery(new Term("field", "doc")), 1).totalHits.value);
+        IndexSearcherWrapper wrapper = new IndexSearcherWrapper() {
+            @Override
+            public DirectoryReader wrap(DirectoryReader reader) throws IOException {
+                return reader;
+            }
+
+            @Override
+            public IndexSearcher wrap(IndexSearcher searcher) throws EngineException {
+                return new IndexSearcher(searcher.getIndexReader()) {
+                    @Override
+                    protected void search(List<LeafReaderContext> leaves, Weight weight, Collector collector) throws IOException {
+                        throw new IllegalStateException("boum");
+                    }
+                };
+            }
+
+        };
+        final AtomicBoolean closeCalled = new AtomicBoolean(false);
+        final Engine.Searcher wrap =  wrapper.wrap(new Engine.Searcher("foo", searcher, () -> closeCalled.set(true)));
+        assertEquals(1, wrap.reader().getRefCount());
+        ContextIndexSearcher contextSearcher = new ContextIndexSearcher(wrap, wrap.searcher().getQueryCache(),
+            wrap.searcher().getQueryCachingPolicy());
+        IllegalStateException exc = expectThrows(IllegalStateException.class,
+            () -> contextSearcher.search(new TermQuery(new Term("field", "doc")), new TotalHitCountCollector()));
+        assertThat(exc.getMessage(), equalTo("boum"));
+        wrap.close();
+        assertFalse("wrapped reader is closed", wrap.reader().tryIncRef());
+        assertTrue(closeCalled.get());
+
+        IOUtils.close(open, writer, dir);
+        assertEquals(0, open.getRefCount());
+    }
+
     private static class FieldMaskingReader extends FilterDirectoryReader {
         private final String field;
         private final AtomicInteger closeCalls;

test/framework/src/main/java/org/elasticsearch/search/aggregations/AggregatorTestCase.java

@@ -32,6 +32,7 @@ import org.apache.lucene.search.QueryCache;
 import org.apache.lucene.search.QueryCachingPolicy;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.Weight;
+import org.apache.lucene.search.XIndexSearcher;
 import org.elasticsearch.Version;
 import org.elasticsearch.cluster.metadata.IndexMetaData;
 import org.elasticsearch.common.lease.Releasable;
@@ -447,7 +448,16 @@ public abstract class AggregatorTestCase extends ESTestCase {
      */
     protected static IndexSearcher newIndexSearcher(IndexReader indexReader) {
         if (randomBoolean()) {
-            return new AssertingIndexSearcher(random(), indexReader);
+            final IndexSearcher delegate = new IndexSearcher(indexReader);
+            final XIndexSearcher wrappedSearcher = new XIndexSearcher(delegate);
+            // this executes basic query checks and asserts that weights are normalized only once etc.
+            return new AssertingIndexSearcher(random(), indexReader) {
+                @Override
+                protected void search(List<LeafReaderContext> leaves, Weight weight, Collector collector) throws IOException {
+                    // we cannot use the asserting searcher because the weight is created by the ContextIndexSearcher
+                    wrappedSearcher.search(leaves, weight, collector);
+                }
+             };
         } else {
             return new IndexSearcher(indexReader);
         }

test/framework/src/main/java/org/elasticsearch/test/engine/MockEngineSupport.java

@@ -24,9 +24,14 @@ import org.apache.lucene.index.AssertingDirectoryReader;
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.FilterDirectoryReader;
 import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.AssertingIndexSearcher;
+import org.apache.lucene.search.Collector;
+import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.QueryCache;
 import org.apache.lucene.search.QueryCachingPolicy;
+import org.apache.lucene.search.Weight;
+import org.apache.lucene.search.XIndexSearcher;
 import org.apache.lucene.util.LuceneTestCase;
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.common.settings.Setting;
@@ -42,6 +47,7 @@ import java.io.Closeable;
 import java.io.IOException;
 import java.lang.reflect.Constructor;
 import java.util.IdentityHashMap;
+import java.util.List;
 import java.util.Random;
 import java.util.concurrent.atomic.AtomicBoolean;
 
@@ -145,8 +151,19 @@ public final class MockEngineSupport {
         if (reader instanceof DirectoryReader && mockContext.wrapReader) {
             wrappedReader = wrapReader((DirectoryReader) reader);
         }
+        final IndexSearcher delegate = new IndexSearcher(wrappedReader);
+        delegate.setSimilarity(searcher.searcher().getSimilarity());
+        delegate.setQueryCache(filterCache);
+        delegate.setQueryCachingPolicy(filterCachingPolicy);
+        final XIndexSearcher wrappedSearcher = new XIndexSearcher(delegate);
         // this executes basic query checks and asserts that weights are normalized only once etc.
-        final AssertingIndexSearcher assertingIndexSearcher = new AssertingIndexSearcher(mockContext.random, wrappedReader);
+        final AssertingIndexSearcher assertingIndexSearcher = new AssertingIndexSearcher(mockContext.random, wrappedReader) {
+            @Override
+            protected void search(List<LeafReaderContext> leaves, Weight weight, Collector collector) throws IOException {
+                // we cannot use the asserting searcher because the weight is created by the ContextIndexSearcher
+                wrappedSearcher.search(leaves, weight, collector);
+            }
+        };
         assertingIndexSearcher.setSimilarity(searcher.searcher().getSimilarity());
         assertingIndexSearcher.setQueryCache(filterCache);
         assertingIndexSearcher.setQueryCachingPolicy(filterCachingPolicy);