diff --git a/server/src/internalClusterTest/java/org/elasticsearch/get/GetActionIT.java b/server/src/internalClusterTest/java/org/elasticsearch/get/GetActionIT.java index f06810377771b..34a225f19498a 100644 --- a/server/src/internalClusterTest/java/org/elasticsearch/get/GetActionIT.java +++ b/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 wrapper = EngineTestCase.randomReaderWrapper(); return reader -> { calls.incrementAndGet(); diff --git a/server/src/main/java/org/elasticsearch/index/IndexModule.java b/server/src/main/java/org/elasticsearch/index/IndexModule.java index 42ab9ae362509..3823e6cd3b756 100644 --- a/server/src/main/java/org/elasticsearch/index/IndexModule.java +++ b/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 indexDirectoryWrapper = new SetOnce<>(); - private final SetOnce>> indexReaderWrapper = - new SetOnce<>(); + private final List>> indexReaderWrappers = + new ArrayList<>(); private final Set indexEventListeners = new HashSet<>(); private final Map> similarities = new HashMap<>(); private final Map directoryFactories; @@ -368,9 +368,9 @@ public void addSimilarity(String name, TriFunction * 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 */ - public void setReaderWrapper( + public void addReaderWrapper( Function> 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> readerWrapperFactory = indexReaderWrapper - .get() == null ? (shard) -> null : indexReaderWrapper.get(); + Function> 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>(); + 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); diff --git a/server/src/test/java/org/elasticsearch/index/IndexModuleTests.java b/server/src/test/java/org/elasticsearch/index/IndexModuleTests.java index 088c748bde5f6..d40fbbef1e31c 100644 --- a/server/src/test/java/org/elasticsearch/index/IndexModuleTests.java +++ b/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,7 +270,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); @@ -275,6 +278,70 @@ public void testWrapperIsBound() throws IOException { 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; + } + } } diff --git a/server/src/test/java/org/elasticsearch/search/SearchServiceSingleNodeTests.java b/server/src/test/java/org/elasticsearch/search/SearchServiceSingleNodeTests.java index 0cd60823a22cc..8bc4075dbd86a 100644 --- a/server/src/test/java/org/elasticsearch/search/SearchServiceSingleNodeTests.java +++ b/server/src/test/java/org/elasticsearch/search/SearchServiceSingleNodeTests.java @@ -199,7 +199,7 @@ protected Collection> getPlugins() { public static class ReaderWrapperCountPlugin extends Plugin { @Override public void onIndexModule(IndexModule indexModule) { - indexModule.setReaderWrapper(service -> SearchServiceSingleNodeTests::apply); + indexModule.addReaderWrapper(service -> SearchServiceSingleNodeTests::apply); } } diff --git a/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/Security.java b/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/Security.java index c9f0771e93068..0a48f3b499100 100644 --- a/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/Security.java +++ b/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(), diff --git a/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/SecurityTests.java b/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/SecurityTests.java index 93d1e15aebd3a..fba6ff1474430 100644 --- a/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/SecurityTests.java +++ b/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 {