Add tiered-storage module with stored fields prefetch support

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))

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'
+}

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<Setting<?>> 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<TieredStoragePrefetchSettings> getPrefetchSettingsSupplier() {
+        return () -> this.tieredStoragePrefetchSettings;
+    }
+
+    @Override
+    public Collection<Object> createComponents(
+        Client client,
+        ClusterService clusterService,
+        ThreadPool threadPool,
+        ResourceWatcherService resourceWatcherService,
+        ScriptService scriptService,
+        NamedXContentRegistry xContentRegistry,
+        Environment environment,
+        NodeEnvironment nodeEnvironment,
+        NamedWriteableRegistry namedWriteableRegistry,
+        IndexNameExpressionResolver indexNameExpressionResolver,
+        Supplier<RepositoriesService> 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()));
+        }
+    }
+}

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;

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<TieredStoragePrefetchSettings> tieredStoragePrefetchSettingsSupplier;
+
+    /**
+     * Creates a new StoredFieldsPrefetch instance.
+     * @param tieredStoragePrefetchSettingsSupplier supplier for prefetch settings
+     */
+    public StoredFieldsPrefetch(Supplier<TieredStoragePrefetchSettings> 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();
+    }
+}

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<Integer> 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<Boolean> 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<String> READ_AHEAD_ENABLE_FILE_FORMATS = List.of(DVD_FILE_SUFFIX);
+    private int readAheadBlockCount;
+    private final List<String> 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<String> getReadAheadEnableFileFormats() {
+        return this.readAheadEnableFileFormats;
+    }
+}

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;