Infrastructure for resharding search filters

server/src/internalClusterTest/java/org/elasticsearch/get/GetActionIT.java

@@ -87,7 +87,7 @@ public static class SearcherWrapperPlugin extends Plugin {
         public void onIndexModule(IndexModule indexModule) {
             super.onIndexModule(indexModule);
             if (enabled) {
-                indexModule.setReaderWrapper(indexService -> {
+                indexModule.addReaderWrapper(indexService -> {
                     CheckedFunction<DirectoryReader, DirectoryReader, IOException> wrapper = EngineTestCase.randomReaderWrapper();
                     return reader -> {
                         calls.incrementAndGet();

server/src/main/java/org/elasticsearch/index/IndexModule.java

@@ -166,8 +166,8 @@ public interface DirectoryWrapper {
     private final AnalysisRegistry analysisRegistry;
     private final EngineFactory engineFactory;
     private final SetOnce<DirectoryWrapper> indexDirectoryWrapper = new SetOnce<>();
-    private final SetOnce<Function<IndexService, CheckedFunction<DirectoryReader, DirectoryReader, IOException>>> indexReaderWrapper =
-        new SetOnce<>();
+    private final List<Function<IndexService, CheckedFunction<DirectoryReader, DirectoryReader, IOException>>> indexReaderWrappers =
+        new ArrayList<>();
     private final Set<IndexEventListener> indexEventListeners = new HashSet<>();
     private final Map<String, TriFunction<Settings, IndexVersion, ScriptService, Similarity>> similarities = new HashMap<>();
     private final Map<String, IndexStorePlugin.DirectoryFactory> directoryFactories;
@@ -368,9 +368,9 @@ public void addSimilarity(String name, TriFunction<Settings, IndexVersion, Scrip
     }
 
     /**
-     * Sets the factory for creating new {@link DirectoryReader} wrapper instances.
+     * Adds a new instance of factory creating new {@link DirectoryReader} wrapper instances.
      * The factory ({@link Function}) is called once the IndexService is fully constructed.
-     * NOTE: this method can only be called once per index. Multiple wrappers are not supported.
+     * The order of execution of wrappers is not guaranteed.
      * <p>
      * The {@link CheckedFunction} is invoked each time a {@link Engine.Searcher} is requested to do an operation,
      * for example search, and must return a new directory reader wrapping the provided directory reader or if no
@@ -384,11 +384,11 @@ public void addSimilarity(String name, TriFunction<Settings, IndexVersion, Scrip
      * The returned reader is closed once it goes out of scope.
      * </p>
      */
-    public void setReaderWrapper(
+    public void addReaderWrapper(
         Function<IndexService, CheckedFunction<DirectoryReader, DirectoryReader, IOException>> indexReaderWrapperFactory
     ) {
         ensureNotFrozen();
-        this.indexReaderWrapper.set(indexReaderWrapperFactory);
+        this.indexReaderWrappers.add(indexReaderWrapperFactory);
     }
 
     /**
@@ -499,8 +499,26 @@ public IndexService newIndexService(
         QueryRewriteInterceptor queryRewriteInterceptor
     ) throws IOException {
         final IndexEventListener eventListener = freeze();
-        Function<IndexService, CheckedFunction<DirectoryReader, DirectoryReader, IOException>> readerWrapperFactory = indexReaderWrapper
-            .get() == null ? (shard) -> null : indexReaderWrapper.get();
+        Function<IndexService, CheckedFunction<DirectoryReader, DirectoryReader, IOException>> readerWrapperFactory =
+            switch (indexReaderWrappers.size()) {
+                case 0 -> indexService -> null;
+                case 1 -> indexReaderWrappers.get(0);
+                default -> indexService -> {
+                    // Call factories only one time when creating index service to avoid duplicate work.
+                    var wrappers = new ArrayList<CheckedFunction<DirectoryReader, DirectoryReader, IOException>>();
+                    for (var indexReaderWrapper : indexReaderWrappers) {
+                        wrappers.add(indexReaderWrapper.apply(indexService));
+                    }
+
+                    return directoryReader -> {
+                        var wrapped = wrappers.get(0).apply(directoryReader);
+                        for (int i = 1; i < wrappers.size(); i++) {
+                            wrapped = wrappers.get(i).apply(wrapped);
+                        }
+                        return wrapped;
+                    };
+                };
+            };
         eventListener.beforeIndexCreated(indexSettings.getIndex(), indexSettings.getSettings());
         final IndexStorePlugin.DirectoryFactory directoryFactory = getDirectoryFactory(indexSettings, directoryFactories);
         final IndexStorePlugin.RecoveryStateFactory recoveryStateFactory = getRecoveryStateFactory(indexSettings, recoveryStateFactories);

server/src/test/java/org/elasticsearch/index/IndexModuleTests.java

@@ -12,7 +12,10 @@
 import org.apache.lucene.analysis.standard.StandardTokenizer;
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.FieldInvertState;
+import org.apache.lucene.index.FilterDirectoryReader;
 import org.apache.lucene.index.IndexCommit;
+import org.apache.lucene.index.IndexWriter;
+import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.search.CollectionStatistics;
 import org.apache.lucene.search.QueryCachingPolicy;
 import org.apache.lucene.search.TermStatistics;
@@ -267,14 +270,78 @@ public void testWrapperIsBound() throws IOException {
             new SearchStatsSettings(ClusterSettings.createBuiltInClusterSettings()),
             MergeMetrics.NOOP
         );
-        module.setReaderWrapper(s -> new Wrapper());
+        module.addReaderWrapper(s -> new Wrapper());
 
         IndexService indexService = newIndexService(module);
         assertTrue(indexService.getReaderWrapper() instanceof Wrapper);
         assertSame(indexService.getEngineFactory(), module.getEngineFactory());
         closeIndexService(indexService);
     }
 
+    public void testMultipleReaderWrappers() throws IOException {
+        final MockEngineFactory engineFactory = new MockEngineFactory(AssertingDirectoryReader.class);
+        IndexModule module = new IndexModule(
+            indexSettings,
+            emptyAnalysisRegistry,
+            engineFactory,
+            Collections.emptyMap(),
+            () -> true,
+            indexNameExpressionResolver,
+            Collections.emptyMap(),
+            mock(SlowLogFieldProvider.class),
+            MapperMetrics.NOOP,
+            emptyList(),
+            new IndexingStatsSettings(ClusterSettings.createBuiltInClusterSettings()),
+            new SearchStatsSettings(ClusterSettings.createBuiltInClusterSettings()),
+            MergeMetrics.NOOP
+        );
+
+        class Wrapper extends FilterDirectoryReader {
+            final String name;
+            final DirectoryReader wrappedReader;
+
+            Wrapper(String name, DirectoryReader in) throws IOException {
+                super(in, new SubReaderWrapper() {
+                    @Override
+                    public LeafReader wrap(LeafReader reader) {
+                        return reader;
+                    }
+                });
+                this.name = name;
+                this.wrappedReader = in;
+            }
+
+            @Override
+            protected DirectoryReader doWrapDirectoryReader(DirectoryReader in) throws IOException {
+                return null;
+            }
+
+            @Override
+            public CacheHelper getReaderCacheHelper() {
+                return null;
+            }
+        }
+
+        module.addReaderWrapper(s -> reader -> new Wrapper("A", reader));
+        module.addReaderWrapper(s -> reader -> new Wrapper("B", reader));
+
+        IndexService indexService = newIndexService(module);
+
+        var wrapper = indexService.getReaderWrapper();
+
+        try (var directory = newDirectory(); var reader = new DummyDirectoryReader(directory)) {
+            var wrapped = wrapper.apply(reader);
+            // B is the outermost wrapper since it was added last.
+            assertTrue(wrapped instanceof Wrapper w && w.name.equals("B"));
+            var secondLevel = ((Wrapper) wrapped).wrappedReader;
+            assertTrue(secondLevel instanceof Wrapper w && w.name.equals("A"));
+            var thirdLevel = ((Wrapper) secondLevel).wrappedReader;
+            assertTrue(thirdLevel instanceof DummyDirectoryReader);
+        }
+
+        closeIndexService(indexService);
+    }
+
     public void testRegisterIndexStore() throws IOException {
         final Settings settings = Settings.builder()
             .put(IndexMetadata.SETTING_VERSION_CREATED, IndexVersion.current())
@@ -489,7 +556,7 @@ public void testFrozen() {
         assertEquals(msg, expectThrows(IllegalStateException.class, () -> module.addIndexEventListener(null)).getMessage());
         assertEquals(msg, expectThrows(IllegalStateException.class, () -> module.addIndexOperationListener(null)).getMessage());
         assertEquals(msg, expectThrows(IllegalStateException.class, () -> module.addSimilarity(null, null)).getMessage());
-        assertEquals(msg, expectThrows(IllegalStateException.class, () -> module.setReaderWrapper(null)).getMessage());
+        assertEquals(msg, expectThrows(IllegalStateException.class, () -> module.addReaderWrapper(null)).getMessage());
         assertEquals(msg, expectThrows(IllegalStateException.class, () -> module.forceQueryCacheProvider(null)).getMessage());
         assertEquals(msg, expectThrows(IllegalStateException.class, () -> module.setDirectoryWrapper(null)).getMessage());
         assertEquals(msg, expectThrows(IllegalStateException.class, () -> module.setIndexCommitListener(null)).getMessage());
@@ -894,4 +961,50 @@ protected WrappedDirectory(Directory in, ShardRouting shardRouting) {
             this.shardRouting = shardRouting;
         }
     }
+
+    private static class DummyDirectoryReader extends DirectoryReader {
+        DummyDirectoryReader(Directory directory) throws IOException {
+            super(directory, new LeafReader[0], null);
+        }
+
+        @Override
+        protected DirectoryReader doOpenIfChanged() throws IOException {
+            return null;
+        }
+
+        @Override
+        protected DirectoryReader doOpenIfChanged(IndexCommit commit) throws IOException {
+            return null;
+        }
+
+        @Override
+        protected DirectoryReader doOpenIfChanged(IndexWriter writer, boolean applyAllDeletes) throws IOException {
+            return null;
+        }
+
+        @Override
+        public long getVersion() {
+            return 0;
+        }
+
+        @Override
+        public boolean isCurrent() throws IOException {
+            return false;
+        }
+
+        @Override
+        public IndexCommit getIndexCommit() throws IOException {
+            return null;
+        }
+
+        @Override
+        protected void doClose() throws IOException {
+
+        }
+
+        @Override
+        public CacheHelper getReaderCacheHelper() {
+            return null;
+        }
+    }
 }

server/src/test/java/org/elasticsearch/search/SearchServiceSingleNodeTests.java

@@ -199,7 +199,7 @@ protected Collection<Class<? extends Plugin>> getPlugins() {
     public static class ReaderWrapperCountPlugin extends Plugin {
         @Override
         public void onIndexModule(IndexModule indexModule) {
-            indexModule.setReaderWrapper(service -> SearchServiceSingleNodeTests::apply);
+            indexModule.addReaderWrapper(service -> SearchServiceSingleNodeTests::apply);
         }
     }
 

x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/Security.java

@@ -1598,7 +1598,7 @@ public void onIndexModule(IndexModule module) {
             assert getLicenseState() != null;
             if (XPackSettings.DLS_FLS_ENABLED.get(settings)) {
                 assert dlsBitsetCache.get() != null;
-                module.setReaderWrapper(
+                module.addReaderWrapper(
                     indexService -> new SecurityIndexReaderWrapper(
                         shardId -> indexService.newSearchExecutionContext(
                             shardId.id(),

x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/SecurityTests.java

@@ -477,7 +477,7 @@ public void testOnIndexModuleIsNoOpWithSecurityDisabled() throws Exception {
         );
         security.onIndexModule(indexModule);
         // indexReaderWrapper is a SetOnce so if Security#onIndexModule had already set an ReaderWrapper we would get an exception here
-        indexModule.setReaderWrapper(null);
+        indexModule.addReaderWrapper(null);
     }
 
     public void testFilteredSettings() throws Exception {