diff --git a/CHANGELOG.md b/CHANGELOG.md index f766abb8cbfdf..63f032c028b17 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,6 +27,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Add warmup phase to wait for lag to catch up in pull-based ingestion before serving ([#20526](https://github.com/opensearch-project/OpenSearch/pull/20526)) - Add a new static method to IndicesOptions API to expose `STRICT_EXPAND_OPEN_HIDDEN_FORBID_CLOSED` index option ([#20980](https://github.com/opensearch-project/OpenSearch/pull/20980)) +- Add tiered-storage module with stored fields prefetch support ([#20962](https://github.com/opensearch-project/OpenSearch/pull/20962)) ### Changed - Make telemetry `Tags` immutable ([#20788](https://github.com/opensearch-project/OpenSearch/pull/20788)) diff --git a/modules/tiered-storage/build.gradle b/modules/tiered-storage/build.gradle new file mode 100644 index 0000000000000..78339c1eb8ab5 --- /dev/null +++ b/modules/tiered-storage/build.gradle @@ -0,0 +1,21 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + * + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +opensearchplugin { + description = 'Module for tiered storage and writable warm index support' + classname = 'org.opensearch.storage.TieredStoragePlugin' +} + +test { + include '**/*Tests.class' + include '**/*Test.class' +} diff --git a/modules/tiered-storage/src/main/java/org/opensearch/storage/TieredStoragePlugin.java b/modules/tiered-storage/src/main/java/org/opensearch/storage/TieredStoragePlugin.java new file mode 100644 index 0000000000000..2802c2d755592 --- /dev/null +++ b/modules/tiered-storage/src/main/java/org/opensearch/storage/TieredStoragePlugin.java @@ -0,0 +1,86 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.storage; + +import org.opensearch.cluster.metadata.IndexNameExpressionResolver; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.util.FeatureFlags; +import org.opensearch.core.common.io.stream.NamedWriteableRegistry; +import org.opensearch.core.xcontent.NamedXContentRegistry; +import org.opensearch.env.Environment; +import org.opensearch.env.NodeEnvironment; +import org.opensearch.index.IndexModule; +import org.opensearch.plugins.Plugin; +import org.opensearch.repositories.RepositoriesService; +import org.opensearch.script.ScriptService; +import org.opensearch.storage.prefetch.StoredFieldsPrefetch; +import org.opensearch.storage.prefetch.TieredStoragePrefetchSettings; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.client.Client; +import org.opensearch.watcher.ResourceWatcherService; + +import java.util.Collection; +import java.util.Collections; +import java.util.List; +import java.util.function.Supplier; + +/** + * Plugin to support writable warm index and other related features. + */ +public class TieredStoragePlugin extends Plugin { + + /** + * Default constructor. + */ + public TieredStoragePlugin() {} + + private TieredStoragePrefetchSettings tieredStoragePrefetchSettings; + + @Override + public List> getSettings() { + return List.of( + TieredStoragePrefetchSettings.READ_AHEAD_BLOCK_COUNT, + TieredStoragePrefetchSettings.STORED_FIELDS_PREFETCH_ENABLED_SETTING + ); + } + + /** + * Returns a supplier for the tiered storage prefetch settings. + * @return supplier of {@link TieredStoragePrefetchSettings} + */ + public Supplier getPrefetchSettingsSupplier() { + return () -> this.tieredStoragePrefetchSettings; + } + + @Override + public Collection createComponents( + Client client, + ClusterService clusterService, + ThreadPool threadPool, + ResourceWatcherService resourceWatcherService, + ScriptService scriptService, + NamedXContentRegistry xContentRegistry, + Environment environment, + NodeEnvironment nodeEnvironment, + NamedWriteableRegistry namedWriteableRegistry, + IndexNameExpressionResolver indexNameExpressionResolver, + Supplier repositoriesServiceSupplier + ) { + this.tieredStoragePrefetchSettings = new TieredStoragePrefetchSettings(clusterService.getClusterSettings()); + return Collections.emptyList(); + } + + @Override + public void onIndexModule(IndexModule indexModule) { + if (FeatureFlags.isEnabled(FeatureFlags.WRITABLE_WARM_INDEX_EXPERIMENTAL_FLAG)) { + indexModule.addSearchOperationListener(new StoredFieldsPrefetch(getPrefetchSettingsSupplier())); + } + } +} diff --git a/modules/tiered-storage/src/main/java/org/opensearch/storage/package-info.java b/modules/tiered-storage/src/main/java/org/opensearch/storage/package-info.java new file mode 100644 index 0000000000000..056f41f7ef772 --- /dev/null +++ b/modules/tiered-storage/src/main/java/org/opensearch/storage/package-info.java @@ -0,0 +1,12 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/** + * Tiered storage plugin for writable warm index support. + */ +package org.opensearch.storage; diff --git a/modules/tiered-storage/src/main/java/org/opensearch/storage/prefetch/StoredFieldsPrefetch.java b/modules/tiered-storage/src/main/java/org/opensearch/storage/prefetch/StoredFieldsPrefetch.java new file mode 100644 index 0000000000000..6ee446d583589 --- /dev/null +++ b/modules/tiered-storage/src/main/java/org/opensearch/storage/prefetch/StoredFieldsPrefetch.java @@ -0,0 +1,116 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.storage.prefetch; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.apache.lucene.index.FilterLeafReader; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.ReaderUtil; +import org.apache.lucene.index.SegmentReader; +import org.apache.lucene.index.StoredFields; +import org.apache.lucene.util.BitSet; +import org.opensearch.ExceptionsHelper; +import org.opensearch.common.lucene.search.Queries; +import org.opensearch.index.shard.SearchOperationListener; +import org.opensearch.search.internal.SearchContext; + +import java.io.IOException; +import java.util.function.Supplier; + +/** + * Search operation listener that prefetches stored fields for tiered storage indices. + */ +public class StoredFieldsPrefetch implements SearchOperationListener { + + private static final Logger log = LogManager.getLogger(StoredFieldsPrefetch.class); + private final Supplier tieredStoragePrefetchSettingsSupplier; + + /** + * Creates a new StoredFieldsPrefetch instance. + * @param tieredStoragePrefetchSettingsSupplier supplier for prefetch settings + */ + public StoredFieldsPrefetch(Supplier tieredStoragePrefetchSettingsSupplier) { + this.tieredStoragePrefetchSettingsSupplier = tieredStoragePrefetchSettingsSupplier; + } + + @Override + public void onPreFetchPhase(SearchContext searchContext) { + if (checkIfStoredFieldsPrefetchEnabled()) { + executePrefetch(searchContext); + } + } + + private void executePrefetch(SearchContext context) { + int currentReaderIndex = -1; + LeafReaderContext currentReaderContext = null; + StoredFields currentReader = null; + log.debug("Stored Field Execute prefetch was triggered: {}", context.docIdsToLoadSize()); + for (int index = 0; index < context.docIdsToLoadSize(); index++) { + int docId = context.docIdsToLoad()[context.docIdsToLoadFrom() + index]; + try { + int readerIndex = ReaderUtil.subIndex(docId, context.searcher().getIndexReader().leaves()); + if (currentReaderIndex != readerIndex) { + currentReaderContext = context.searcher().getIndexReader().leaves().get(readerIndex); + currentReaderIndex = readerIndex; + + // Unwrap the reader here + LeafReader innerLeafReader = currentReaderContext.reader(); + while (innerLeafReader instanceof FilterLeafReader) { + innerLeafReader = ((FilterLeafReader) innerLeafReader).getDelegate(); + } + // never be the case, just sanity check + if (!(innerLeafReader instanceof SegmentReader)) { + // disable prefetch on stored fields for this segment + log.warn("Unexpected reader type [{}], skipping stored fields prefetch", innerLeafReader.getClass().getName()); + currentReader = null; + continue; + } + currentReader = innerLeafReader.storedFields(); + } + assert currentReaderContext != null; + if (currentReader == null) { + continue; + } + log.debug( + "Prefetching stored fields for index shard: {}, docId: {}, readerIndex: {}", + context.indexShard().shardId(), + docId, + readerIndex + ); + + // nested docs logic + final int subDocId = docId - currentReaderContext.docBase; + final int rootDocId = findRootDocumentIfNested(context, currentReaderContext, subDocId); + if (rootDocId != -1) { + currentReader.prefetch(rootDocId); + } + currentReader.prefetch(subDocId); + } catch (Exception e) { + throw ExceptionsHelper.convertToOpenSearchException(e); + } + } + } + + private int findRootDocumentIfNested(SearchContext context, LeafReaderContext subReaderContext, int subDocId) throws IOException { + if (context.mapperService().hasNested()) { + BitSet bits = context.bitsetFilterCache().getBitSetProducer(Queries.newNonNestedFilter()).getBitSet(subReaderContext); + if (bits != null && !bits.get(subDocId)) { + return bits.nextSetBit(subDocId); + } + } + return -1; + } + + private boolean checkIfStoredFieldsPrefetchEnabled() { + TieredStoragePrefetchSettings settings = tieredStoragePrefetchSettingsSupplier.get(); + return settings != null && settings.isStoredFieldsPrefetchEnabled(); + } +} diff --git a/modules/tiered-storage/src/main/java/org/opensearch/storage/prefetch/TieredStoragePrefetchSettings.java b/modules/tiered-storage/src/main/java/org/opensearch/storage/prefetch/TieredStoragePrefetchSettings.java new file mode 100644 index 0000000000000..cdc515a4f215d --- /dev/null +++ b/modules/tiered-storage/src/main/java/org/opensearch/storage/prefetch/TieredStoragePrefetchSettings.java @@ -0,0 +1,101 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.storage.prefetch; + +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; + +import java.util.List; + +/** + * Settings for tiered storage prefetch behavior. + */ +public class TieredStoragePrefetchSettings { + + /** Default number of blocks to read ahead during prefetch. */ + public static final int DEFAULT_READ_AHEAD_BLOCK_COUNT = 4; + /** File suffix for DVD format files. */ + public static final String DVD_FILE_SUFFIX = "dvd"; + /** File suffix for CFS format files. */ + public static final String CFS_FILE_SUFFIX = "cfs"; + /** Setting for the number of blocks to read ahead. */ + public static final Setting READ_AHEAD_BLOCK_COUNT = Setting.intSetting( + "tiering.service.prefetch.read_ahead.block_count", + DEFAULT_READ_AHEAD_BLOCK_COUNT, + 0, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + + /** Setting to enable or disable stored fields prefetch. */ + public static final Setting STORED_FIELDS_PREFETCH_ENABLED_SETTING = Setting.boolSetting( + "tiering.service.prefetch.stored_fields.enabled", + true, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + + /** List of file formats for which read-ahead is enabled. */ + public static final List READ_AHEAD_ENABLE_FILE_FORMATS = List.of(DVD_FILE_SUFFIX); + private int readAheadBlockCount; + private final List readAheadEnableFileFormats; + private boolean storedFieldsPrefetchEnabled; + + /** + * Creates a new TieredStoragePrefetchSettings instance. + * @param clusterSettings the cluster settings to read prefetch configuration from + */ + public TieredStoragePrefetchSettings(ClusterSettings clusterSettings) { + this.readAheadBlockCount = clusterSettings.get(READ_AHEAD_BLOCK_COUNT); + clusterSettings.addSettingsUpdateConsumer(READ_AHEAD_BLOCK_COUNT, this::setReadAheadBlockCount); + this.readAheadEnableFileFormats = READ_AHEAD_ENABLE_FILE_FORMATS; + this.storedFieldsPrefetchEnabled = clusterSettings.get(STORED_FIELDS_PREFETCH_ENABLED_SETTING); + clusterSettings.addSettingsUpdateConsumer(STORED_FIELDS_PREFETCH_ENABLED_SETTING, this::setStoredFieldsPrefetchEnabled); + } + + /** + * Sets the read-ahead block count. + * @param readAheadBlockCount the number of blocks to read ahead + */ + public void setReadAheadBlockCount(int readAheadBlockCount) { + this.readAheadBlockCount = readAheadBlockCount; + } + + /** + * Sets whether stored fields prefetch is enabled. + * @param storedFieldsPrefetchEnabled true to enable, false to disable + */ + public void setStoredFieldsPrefetchEnabled(boolean storedFieldsPrefetchEnabled) { + this.storedFieldsPrefetchEnabled = storedFieldsPrefetchEnabled; + } + + /** + * Returns whether stored fields prefetch is enabled. + * @return true if enabled + */ + public boolean isStoredFieldsPrefetchEnabled() { + return storedFieldsPrefetchEnabled; + } + + /** + * Returns the read-ahead block count. + * @return the number of blocks to read ahead + */ + public int getReadAheadBlockCount() { + return this.readAheadBlockCount; + } + + /** + * Returns the list of file formats for which read-ahead is enabled. + * @return list of file format suffixes + */ + public List getReadAheadEnableFileFormats() { + return this.readAheadEnableFileFormats; + } +} diff --git a/modules/tiered-storage/src/main/java/org/opensearch/storage/prefetch/package-info.java b/modules/tiered-storage/src/main/java/org/opensearch/storage/prefetch/package-info.java new file mode 100644 index 0000000000000..85a72b15fd068 --- /dev/null +++ b/modules/tiered-storage/src/main/java/org/opensearch/storage/prefetch/package-info.java @@ -0,0 +1,12 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/** + * Prefetch support for tiered storage, including stored fields prefetching. + */ +package org.opensearch.storage.prefetch; diff --git a/modules/tiered-storage/src/test/java/org/opensearch/storage/TieredStoragePluginTests.java b/modules/tiered-storage/src/test/java/org/opensearch/storage/TieredStoragePluginTests.java new file mode 100644 index 0000000000000..2518fb1aa459b --- /dev/null +++ b/modules/tiered-storage/src/test/java/org/opensearch/storage/TieredStoragePluginTests.java @@ -0,0 +1,155 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.storage; + +import org.opensearch.cluster.metadata.IndexNameExpressionResolver; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.FeatureFlags; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.core.common.io.stream.NamedWriteableRegistry; +import org.opensearch.core.xcontent.NamedXContentRegistry; +import org.opensearch.env.Environment; +import org.opensearch.index.IndexModule; +import org.opensearch.index.IndexSettings; +import org.opensearch.index.engine.EngineConfigFactory; +import org.opensearch.index.engine.InternalEngineFactory; +import org.opensearch.repositories.RepositoriesService; +import org.opensearch.script.ScriptService; +import org.opensearch.storage.prefetch.TieredStoragePrefetchSettings; +import org.opensearch.test.ClusterServiceUtils; +import org.opensearch.test.IndexSettingsModule; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.threadpool.TestThreadPool; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.client.Client; +import org.opensearch.watcher.ResourceWatcherService; +import org.junit.After; +import org.junit.Before; + +import java.util.Collection; +import java.util.Collections; +import java.util.HashSet; +import java.util.Set; +import java.util.function.Supplier; + +import static org.opensearch.common.settings.ClusterSettings.BUILT_IN_CLUSTER_SETTINGS; +import static org.mockito.Mockito.mock; + +public class TieredStoragePluginTests extends OpenSearchTestCase { + + private ThreadPool threadPool; + private ClusterService clusterService; + + @Before + public void setUp() throws Exception { + super.setUp(); + Set> clusterSettingsToAdd = new HashSet<>(BUILT_IN_CLUSTER_SETTINGS); + clusterSettingsToAdd.add(TieredStoragePrefetchSettings.READ_AHEAD_BLOCK_COUNT); + clusterSettingsToAdd.add(TieredStoragePrefetchSettings.STORED_FIELDS_PREFETCH_ENABLED_SETTING); + ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, clusterSettingsToAdd); + threadPool = new TestThreadPool("TieredStoragePluginTests"); + clusterService = ClusterServiceUtils.createClusterService(Settings.EMPTY, clusterSettings, threadPool); + } + + @After + public void tearDown() throws Exception { + super.tearDown(); + threadPool.shutdownNow(); + } + + public void testConstructor() { + TieredStoragePlugin plugin = new TieredStoragePlugin(); + assertNotNull(plugin); + } + + public void testGetPrefetchSettingsSupplier_BeforeCreateComponents() { + TieredStoragePlugin plugin = new TieredStoragePlugin(); + Supplier supplier = plugin.getPrefetchSettingsSupplier(); + assertNotNull(supplier); + assertNull(supplier.get()); + } + + public void testCreateComponents() { + TieredStoragePlugin plugin = new TieredStoragePlugin(); + Collection components = plugin.createComponents( + mock(Client.class), + clusterService, + threadPool, + mock(ResourceWatcherService.class), + mock(ScriptService.class), + mock(NamedXContentRegistry.class), + mock(Environment.class), + null, + mock(NamedWriteableRegistry.class), + mock(IndexNameExpressionResolver.class), + () -> mock(RepositoriesService.class) + ); + assertNotNull(components); + assertTrue(components.isEmpty()); + + Supplier supplier = plugin.getPrefetchSettingsSupplier(); + assertNotNull(supplier.get()); + } + + public void testGetPrefetchSettingsSupplier_AfterCreateComponents() { + TieredStoragePlugin plugin = new TieredStoragePlugin(); + plugin.createComponents( + mock(Client.class), + clusterService, + threadPool, + mock(ResourceWatcherService.class), + mock(ScriptService.class), + mock(NamedXContentRegistry.class), + mock(Environment.class), + null, + mock(NamedWriteableRegistry.class), + mock(IndexNameExpressionResolver.class), + () -> mock(RepositoriesService.class) + ); + TieredStoragePrefetchSettings settings = plugin.getPrefetchSettingsSupplier().get(); + assertNotNull(settings); + assertTrue(settings.isStoredFieldsPrefetchEnabled()); + assertEquals(TieredStoragePrefetchSettings.DEFAULT_READ_AHEAD_BLOCK_COUNT, settings.getReadAheadBlockCount()); + } + + public void testOnIndexModule_WhenFeatureFlagDisabled() { + TieredStoragePlugin plugin = new TieredStoragePlugin(); + IndexSettings indexSettings = IndexSettingsModule.newIndexSettings("test-index", Settings.EMPTY); + IndexModule indexModule = new IndexModule( + indexSettings, + null, + new InternalEngineFactory(), + new EngineConfigFactory(indexSettings), + Collections.emptyMap(), + () -> true, + new IndexNameExpressionResolver(new ThreadContext(Settings.EMPTY)), + Collections.emptyMap() + ); + plugin.onIndexModule(indexModule); + } + + public void testOnIndexModule_WhenFeatureFlagEnabled() throws Exception { + TieredStoragePlugin plugin = new TieredStoragePlugin(); + IndexSettings indexSettings = IndexSettingsModule.newIndexSettings("test-index", Settings.EMPTY); + IndexModule indexModule = new IndexModule( + indexSettings, + null, + new InternalEngineFactory(), + new EngineConfigFactory(indexSettings), + Collections.emptyMap(), + () -> true, + new IndexNameExpressionResolver(new ThreadContext(Settings.EMPTY)), + Collections.emptyMap() + ); + FeatureFlags.TestUtils.with(FeatureFlags.WRITABLE_WARM_INDEX_EXPERIMENTAL_FLAG, () -> { plugin.onIndexModule(indexModule); }); + } +} diff --git a/modules/tiered-storage/src/test/java/org/opensearch/storage/prefetch/StoredFieldsPrefetchTests.java b/modules/tiered-storage/src/test/java/org/opensearch/storage/prefetch/StoredFieldsPrefetchTests.java new file mode 100644 index 0000000000000..12c91a2ba9a4d --- /dev/null +++ b/modules/tiered-storage/src/test/java/org/opensearch/storage/prefetch/StoredFieldsPrefetchTests.java @@ -0,0 +1,359 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.storage.prefetch; + +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.NumericDocValuesField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FilterDirectoryReader; +import org.apache.lucene.index.FilterLeafReader; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.store.Directory; +import org.opensearch.Version; +import org.opensearch.cluster.metadata.IndexMetadata; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; +import org.opensearch.core.index.shard.ShardId; +import org.opensearch.index.IndexSettings; +import org.opensearch.index.cache.bitset.BitsetFilterCache; +import org.opensearch.index.mapper.MapperService; +import org.opensearch.index.mapper.SeqNoFieldMapper; +import org.opensearch.index.shard.IndexShard; +import org.opensearch.search.internal.ContextIndexSearcher; +import org.opensearch.search.internal.SearchContext; +import org.opensearch.test.ClusterServiceUtils; +import org.opensearch.test.IndexSettingsModule; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.threadpool.TestThreadPool; +import org.opensearch.threadpool.ThreadPool; +import org.junit.After; +import org.junit.Before; + +import java.io.IOException; +import java.util.HashSet; +import java.util.Set; +import java.util.function.Supplier; + +import static org.opensearch.common.settings.ClusterSettings.BUILT_IN_CLUSTER_SETTINGS; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class StoredFieldsPrefetchTests extends OpenSearchTestCase { + + private ClusterService clusterService; + private ThreadPool threadPool; + private SearchContext searchContext; + private TieredStoragePrefetchSettings tieredStoragePrefetchSettings; + private StoredFieldsPrefetch storedFieldsPrefetch; + private Directory directory; + private IndexReader indexReader; + + @Before + public void setUp() throws Exception { + super.setUp(); + Set> clusterSettingsToAdd = new HashSet<>(BUILT_IN_CLUSTER_SETTINGS); + clusterSettingsToAdd.add(TieredStoragePrefetchSettings.READ_AHEAD_BLOCK_COUNT); + clusterSettingsToAdd.add(TieredStoragePrefetchSettings.STORED_FIELDS_PREFETCH_ENABLED_SETTING); + ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, clusterSettingsToAdd); + threadPool = new TestThreadPool("TieredStoragePrefetchSettingsTests"); + clusterService = ClusterServiceUtils.createClusterService(Settings.EMPTY, clusterSettings, threadPool); + this.tieredStoragePrefetchSettings = new TieredStoragePrefetchSettings(clusterService.getClusterSettings()); + searchContext = mock(SearchContext.class); + storedFieldsPrefetch = new StoredFieldsPrefetch(getPrefetchSettingsSupplier()); + + directory = newDirectory(); + IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig()); + Document doc = new Document(); + doc.add(new StringField("id", "1", Field.Store.YES)); + writer.addDocument(doc); + doc = new Document(); + doc.add(new StringField("id", "2", Field.Store.YES)); + writer.addDocument(doc); + doc = new Document(); + doc.add(new StringField("id", "3", Field.Store.YES)); + writer.addDocument(doc); + writer.commit(); + writer.close(); + indexReader = DirectoryReader.open(directory); + } + + @After + public void tearDown() throws Exception { + super.tearDown(); + indexReader.close(); + directory.close(); + threadPool.shutdownNow(); + } + + public Supplier getPrefetchSettingsSupplier() { + return () -> this.tieredStoragePrefetchSettings; + } + + /** + * Helper to set up SearchContext mock with the real index reader and given doc IDs. + */ + private void setupSearchContext(int[] docIds, boolean hasNested) { + when(searchContext.docIdsToLoadSize()).thenReturn(docIds.length); + when(searchContext.docIdsToLoad()).thenReturn(docIds); + when(searchContext.docIdsToLoadFrom()).thenReturn(0); + + ContextIndexSearcher searcher = mock(ContextIndexSearcher.class); + when(searchContext.searcher()).thenReturn(searcher); + when(searcher.getIndexReader()).thenReturn(indexReader); + + IndexShard indexShard = mock(IndexShard.class); + when(searchContext.indexShard()).thenReturn(indexShard); + when(indexShard.shardId()).thenReturn(new ShardId("test-index", "uuid", 0)); + + MapperService mapperService = mock(MapperService.class); + when(searchContext.mapperService()).thenReturn(mapperService); + when(mapperService.hasNested()).thenReturn(hasNested); + + if (hasNested) { + IndexSettings indexSettings = IndexSettingsModule.newIndexSettings( + "test-index", + Settings.builder().put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT).build() + ); + BitsetFilterCache bitsetFilterCache = new BitsetFilterCache(indexSettings, mock(BitsetFilterCache.Listener.class)); + when(searchContext.bitsetFilterCache()).thenReturn(bitsetFilterCache); + } + } + + public void testOnPreFetchPhase_WhenPrefetchDisabled() { + Settings settings = Settings.builder() + .put(TieredStoragePrefetchSettings.STORED_FIELDS_PREFETCH_ENABLED_SETTING.getKey(), false) + .build(); + clusterService.getClusterSettings().applySettings(settings); + storedFieldsPrefetch.onPreFetchPhase(searchContext); + verify(searchContext, never()).docIdsToLoadSize(); + } + + public void testOnPreFetchPhase_WhenSettingsSupplierReturnsNull() { + StoredFieldsPrefetch prefetchWithNull = new StoredFieldsPrefetch(() -> null); + prefetchWithNull.onPreFetchPhase(searchContext); + verify(searchContext, never()).docIdsToLoadSize(); + } + + public void testOnPreFetchPhase_WhenPrefetchEnabled_WithSegmentReader() throws IOException { + setupSearchContext(new int[] { 0 }, false); + storedFieldsPrefetch.onPreFetchPhase(searchContext); + verify(searchContext, atLeastOnce()).docIdsToLoadSize(); + } + + public void testOnPreFetchPhase_WithNonSegmentReader_SkipsPrefetch() throws IOException { + setupSearchContext(new int[] { 0 }, false); + storedFieldsPrefetch.onPreFetchPhase(searchContext); + verify(searchContext, atLeastOnce()).docIdsToLoadSize(); + } + + public void testOnPreFetchPhase_MultipleDocsInSameSegment() throws IOException { + setupSearchContext(new int[] { 0, 1, 2 }, false); + storedFieldsPrefetch.onPreFetchPhase(searchContext); + verify(searchContext, atLeastOnce()).docIdsToLoadSize(); + } + + public void testOnPreFetchPhase_WithNestedMapping_NonNestedDoc() throws IOException { + setupSearchContext(new int[] { 0 }, true); + expectThrows(Exception.class, () -> storedFieldsPrefetch.onPreFetchPhase(searchContext)); + } + + public void testOnPreFetchPhase_WithNoNestedMapping() throws IOException { + setupSearchContext(new int[] { 0 }, false); + storedFieldsPrefetch.onPreFetchPhase(searchContext); + verify(searchContext, atLeastOnce()).docIdsToLoadSize(); + } + + public void testOnPreFetchPhase_ExceptionWrappedAsOpenSearchException() throws IOException { + when(searchContext.docIdsToLoadSize()).thenReturn(1); + when(searchContext.docIdsToLoad()).thenReturn(new int[] { 0 }); + when(searchContext.docIdsToLoadFrom()).thenReturn(0); + + ContextIndexSearcher searcher = mock(ContextIndexSearcher.class); + when(searchContext.searcher()).thenReturn(searcher); + when(searcher.getIndexReader()).thenReturn(indexReader); + + IndexShard indexShard = mock(IndexShard.class); + when(searchContext.indexShard()).thenReturn(indexShard); + when(indexShard.shardId()).thenReturn(new ShardId("test-index", "uuid", 0)); + + MapperService mapperService = mock(MapperService.class); + when(searchContext.mapperService()).thenReturn(mapperService); + when(mapperService.hasNested()).thenThrow(new RuntimeException("simulated failure")); + + expectThrows(Exception.class, () -> storedFieldsPrefetch.onPreFetchPhase(searchContext)); + } + + public void testOnPreFetchPhase_NullBitSetFilterCache_ThrowsException() throws IOException { + when(searchContext.docIdsToLoadSize()).thenReturn(1); + when(searchContext.docIdsToLoad()).thenReturn(new int[] { 0 }); + when(searchContext.docIdsToLoadFrom()).thenReturn(0); + + ContextIndexSearcher searcher = mock(ContextIndexSearcher.class); + when(searchContext.searcher()).thenReturn(searcher); + when(searcher.getIndexReader()).thenReturn(indexReader); + + IndexShard indexShard = mock(IndexShard.class); + when(searchContext.indexShard()).thenReturn(indexShard); + when(indexShard.shardId()).thenReturn(new ShardId("test-index", "uuid", 0)); + + MapperService mapperService = mock(MapperService.class); + when(searchContext.mapperService()).thenReturn(mapperService); + when(mapperService.hasNested()).thenReturn(true); + + expectThrows(Exception.class, () -> storedFieldsPrefetch.onPreFetchPhase(searchContext)); + } + + public void testOnPreFetchPhase_WithNonSegmentReaderViaFilterDirectoryReader() throws IOException { + DirectoryReader realReader = DirectoryReader.open(directory); + DirectoryReader wrappedReader = new NonSegmentReaderDirectoryReader(realReader); + try { + when(searchContext.docIdsToLoadSize()).thenReturn(1); + when(searchContext.docIdsToLoad()).thenReturn(new int[] { 0 }); + when(searchContext.docIdsToLoadFrom()).thenReturn(0); + + ContextIndexSearcher searcher = mock(ContextIndexSearcher.class); + when(searchContext.searcher()).thenReturn(searcher); + when(searcher.getIndexReader()).thenReturn(wrappedReader); + + IndexShard indexShard = mock(IndexShard.class); + when(searchContext.indexShard()).thenReturn(indexShard); + when(indexShard.shardId()).thenReturn(new ShardId("test-index", "uuid", 0)); + + MapperService mapperService = mock(MapperService.class); + when(searchContext.mapperService()).thenReturn(mapperService); + when(mapperService.hasNested()).thenReturn(false); + + storedFieldsPrefetch.onPreFetchPhase(searchContext); + + verify(searchContext, atLeastOnce()).docIdsToLoadSize(); + } finally { + wrappedReader.close(); + } + } + + public void testOnPreFetchPhase_NullCurrentReaderContinuesForSubsequentDocs() throws IOException { + DirectoryReader realReader = DirectoryReader.open(directory); + DirectoryReader wrappedReader = new NonSegmentReaderDirectoryReader(realReader); + try { + when(searchContext.docIdsToLoadSize()).thenReturn(2); + when(searchContext.docIdsToLoad()).thenReturn(new int[] { 0, 1 }); + when(searchContext.docIdsToLoadFrom()).thenReturn(0); + + ContextIndexSearcher searcher = mock(ContextIndexSearcher.class); + when(searchContext.searcher()).thenReturn(searcher); + when(searcher.getIndexReader()).thenReturn(wrappedReader); + + IndexShard indexShard = mock(IndexShard.class); + when(searchContext.indexShard()).thenReturn(indexShard); + when(indexShard.shardId()).thenReturn(new ShardId("test-index", "uuid", 0)); + + MapperService mapperService = mock(MapperService.class); + when(searchContext.mapperService()).thenReturn(mapperService); + when(mapperService.hasNested()).thenReturn(false); + + storedFieldsPrefetch.onPreFetchPhase(searchContext); + + verify(searchContext, atLeastOnce()).docIdsToLoadSize(); + } finally { + wrappedReader.close(); + } + } + + public void testOnPreFetchPhase_NestedChildDoc_PrefetchesRootDoc() throws IOException { + Directory nestedDir = newDirectory(); + IndexWriter writer = new IndexWriter(nestedDir, new IndexWriterConfig()); + Document childDoc = new Document(); + childDoc.add(new StringField("nested_field", "child_value", Field.Store.YES)); + Document rootDoc = new Document(); + rootDoc.add(new StringField("id", "1", Field.Store.YES)); + rootDoc.add(new NumericDocValuesField(SeqNoFieldMapper.PRIMARY_TERM_NAME, 1)); + writer.addDocuments(java.util.List.of(childDoc, rootDoc)); + writer.commit(); + writer.close(); + + IndexReader nestedReader = DirectoryReader.open(nestedDir); + try { + when(searchContext.docIdsToLoadSize()).thenReturn(1); + when(searchContext.docIdsToLoad()).thenReturn(new int[] { 0 }); + when(searchContext.docIdsToLoadFrom()).thenReturn(0); + + ContextIndexSearcher searcher = mock(ContextIndexSearcher.class); + when(searchContext.searcher()).thenReturn(searcher); + when(searcher.getIndexReader()).thenReturn(nestedReader); + + IndexShard indexShard = mock(IndexShard.class); + when(searchContext.indexShard()).thenReturn(indexShard); + when(indexShard.shardId()).thenReturn(new ShardId("test-index", "uuid", 0)); + + MapperService mapperService = mock(MapperService.class); + when(searchContext.mapperService()).thenReturn(mapperService); + when(mapperService.hasNested()).thenReturn(true); + + IndexSettings indexSettings = IndexSettingsModule.newIndexSettings( + "test-index", + Settings.builder().put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT).build() + ); + BitsetFilterCache bitsetFilterCache = new BitsetFilterCache(indexSettings, mock(BitsetFilterCache.Listener.class)); + when(searchContext.bitsetFilterCache()).thenReturn(bitsetFilterCache); + + expectThrows(Exception.class, () -> storedFieldsPrefetch.onPreFetchPhase(searchContext)); + } finally { + nestedReader.close(); + nestedDir.close(); + } + } + + private static class NonSegmentReaderDirectoryReader extends FilterDirectoryReader { + NonSegmentReaderDirectoryReader(DirectoryReader in) throws IOException { + super(in, new SubReaderWrapper() { + @Override + public LeafReader wrap(LeafReader reader) { + return new FilterLeafReader(reader) { + private final LeafReader fakeDelegate = mock(LeafReader.class); + + @Override + public LeafReader getDelegate() { + return fakeDelegate; + } + + @Override + public CacheHelper getCoreCacheHelper() { + return reader.getCoreCacheHelper(); + } + + @Override + public CacheHelper getReaderCacheHelper() { + return reader.getReaderCacheHelper(); + } + }; + } + }); + } + + @Override + protected DirectoryReader doWrapDirectoryReader(DirectoryReader in) throws IOException { + return new NonSegmentReaderDirectoryReader(in); + } + + @Override + public CacheHelper getReaderCacheHelper() { + return in.getReaderCacheHelper(); + } + } +} diff --git a/modules/tiered-storage/src/test/java/org/opensearch/storage/prefetch/TieredStoragePrefetchSettingsTests.java b/modules/tiered-storage/src/test/java/org/opensearch/storage/prefetch/TieredStoragePrefetchSettingsTests.java new file mode 100644 index 0000000000000..68151bb384a31 --- /dev/null +++ b/modules/tiered-storage/src/test/java/org/opensearch/storage/prefetch/TieredStoragePrefetchSettingsTests.java @@ -0,0 +1,68 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.storage.prefetch; + +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; +import org.opensearch.test.ClusterServiceUtils; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.threadpool.TestThreadPool; +import org.opensearch.threadpool.ThreadPool; +import org.junit.Before; + +import java.util.HashSet; +import java.util.Set; + +import static org.opensearch.common.settings.ClusterSettings.BUILT_IN_CLUSTER_SETTINGS; + +public class TieredStoragePrefetchSettingsTests extends OpenSearchTestCase { + + private ClusterService clusterService; + private ThreadPool threadPool; + + @Before + public void setUp() throws Exception { + super.setUp(); + threadPool = new TestThreadPool("TieredStoragePrefetchSettingsTests"); + Set> clusterSettingsToAdd = new HashSet<>(BUILT_IN_CLUSTER_SETTINGS); + clusterSettingsToAdd.add(TieredStoragePrefetchSettings.READ_AHEAD_BLOCK_COUNT); + clusterSettingsToAdd.add(TieredStoragePrefetchSettings.STORED_FIELDS_PREFETCH_ENABLED_SETTING); + clusterService = ClusterServiceUtils.createClusterService( + Settings.EMPTY, + new ClusterSettings(Settings.EMPTY, clusterSettingsToAdd), + threadPool + ); + } + + @Override + public void tearDown() throws Exception { + super.tearDown(); + threadPool.shutdownNow(); + } + + public void testDefaultSettings() { + TieredStoragePrefetchSettings settings = new TieredStoragePrefetchSettings(clusterService.getClusterSettings()); + assertEquals(TieredStoragePrefetchSettings.DEFAULT_READ_AHEAD_BLOCK_COUNT, settings.getReadAheadBlockCount()); + assertEquals(TieredStoragePrefetchSettings.READ_AHEAD_ENABLE_FILE_FORMATS, settings.getReadAheadEnableFileFormats()); + assertEquals(true, settings.isStoredFieldsPrefetchEnabled()); + } + + public void testUpdateAfterGetDefaultSettings() { + TieredStoragePrefetchSettings tieringServicePrefetchSettings = new TieredStoragePrefetchSettings( + clusterService.getClusterSettings() + ); + assertEquals(tieringServicePrefetchSettings.getReadAheadBlockCount(), TieredStoragePrefetchSettings.DEFAULT_READ_AHEAD_BLOCK_COUNT); + assertEquals(tieringServicePrefetchSettings.isStoredFieldsPrefetchEnabled(), true); + Settings settings = Settings.builder().put(TieredStoragePrefetchSettings.READ_AHEAD_BLOCK_COUNT.getKey(), 10).build(); + clusterService.getClusterSettings().applySettings(settings); + assertEquals(tieringServicePrefetchSettings.getReadAheadBlockCount(), 10); + } +}