Infrastructure for resharding search filters

Author: lkts

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.
+     * All added wrappers are applied in the order they have been added.
      * <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,14 @@ 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 = indexReaderWrappers
+            .isEmpty() ? indexService -> null : indexService -> (directoryReader) -> {
+                var wrapped = indexReaderWrappers.get(0).apply(indexService).apply(directoryReader);
+                for (int i = 1; i < indexReaderWrappers.size(); i++) {
+                    wrapped = indexReaderWrappers.get(i).apply(indexService).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

@@ -267,7 +267,7 @@ 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);
@@ -489,7 +489,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());

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

@@ -1585,7 +1585,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

@@ -475,7 +475,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 {