Introducing MergedSegmentWarmerFactory to support the extension of IndexWriter.IndexReaderWarmer

CHANGELOG.md

@@ -36,6 +36,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add update and delete support in pull-based ingestion ([#17822](https://github.com/opensearch-project/OpenSearch/pull/17822))
 - Allow maxPollSize and pollTimeout in IngestionSource to be configurable ([#17863](https://github.com/opensearch-project/OpenSearch/pull/17863))
 - [Star Tree] [Search] Add query changes to support unsigned-long in star tree ([#17275](https://github.com/opensearch-project/OpenSearch/pull/17275))
+- Introduce MergedSegmentWarmerFactory for Local/Remote merged segment pre copy ([#17881](https://github.com/opensearch-project/OpenSearch/pull/17881))
 
 ### Changed
 - Migrate BC libs to their FIPS counterparts ([#14912](https://github.com/opensearch-project/OpenSearch/pull/14912))

server/src/internalClusterTest/java/org/opensearch/index/shard/IndexShardIT.java

@@ -74,6 +74,7 @@
 import org.opensearch.index.VersionType;
 import org.opensearch.index.engine.CommitStats;
 import org.opensearch.index.engine.Engine;
+import org.opensearch.index.engine.MergedSegmentWarmerFactory;
 import org.opensearch.index.engine.NoOpEngine;
 import org.opensearch.index.flush.FlushStats;
 import org.opensearch.index.mapper.MapperService;
@@ -719,7 +720,8 @@ public static final IndexShard newIndexShard(
             DefaultRemoteStoreSettings.INSTANCE,
             false,
             IndexShardTestUtils.getFakeDiscoveryNodes(initializingShardRouting),
-            mock(Function.class)
+            mock(Function.class),
+            new MergedSegmentWarmerFactory(null, null, null)
         );
     }
 

server/src/main/java/org/opensearch/common/settings/FeatureFlagSettings.java

@@ -39,6 +39,7 @@ protected FeatureFlagSettings(
         FeatureFlags.APPLICATION_BASED_CONFIGURATION_TEMPLATES_SETTING,
         FeatureFlags.READER_WRITER_SPLIT_EXPERIMENTAL_SETTING,
         FeatureFlags.TERM_VERSION_PRECOMMIT_ENABLE_SETTING,
-        FeatureFlags.ARROW_STREAMS_SETTING
+        FeatureFlags.ARROW_STREAMS_SETTING,
+        FeatureFlags.MERGED_SEGMENT_WARMER_EXPERIMENTAL_SETTING
     );
 }

server/src/main/java/org/opensearch/common/util/FeatureFlags.java

@@ -77,6 +77,12 @@ public class FeatureFlags {
 
     public static final String READER_WRITER_SPLIT_EXPERIMENTAL = FEATURE_FLAG_PREFIX + "read.write.split.enabled";
 
+    /**
+     * Gates the functionality of merged segment warmer in local/remote segment replication.
+     * Once the feature is ready for release, this feature flag can be removed.
+     */
+    public static final String MERGED_SEGMENT_WARMER_EXPERIMENTAL_FLAG = "opensearch.experimental.feature.merged_segment_warmer.enabled";
+
     public static final Setting<Boolean> REMOTE_STORE_MIGRATION_EXPERIMENTAL_SETTING = Setting.boolSetting(
         REMOTE_STORE_MIGRATION_EXPERIMENTAL,
         false,
@@ -105,6 +111,12 @@ public class FeatureFlags {
         Property.NodeScope
     );
 
+    public static final Setting<Boolean> MERGED_SEGMENT_WARMER_EXPERIMENTAL_SETTING = Setting.boolSetting(
+        MERGED_SEGMENT_WARMER_EXPERIMENTAL_FLAG,
+        false,
+        Property.NodeScope
+    );
+
     /**
      * Gates the functionality of star tree index, which improves the performance of search
      * aggregations.
@@ -162,6 +174,7 @@ static class FeatureFlagsImpl {
                     SEARCHABLE_SNAPSHOT_EXTENDED_COMPATIBILITY_SETTING,
                     SEARCHABLE_SNAPSHOT_EXTENDED_COMPATIBILITY_SETTING.getDefault(Settings.EMPTY)
                 );
+                put(MERGED_SEGMENT_WARMER_EXPERIMENTAL_SETTING, MERGED_SEGMENT_WARMER_EXPERIMENTAL_SETTING.getDefault(Settings.EMPTY));
             }
         };
 

server/src/main/java/org/opensearch/index/IndexService.java

@@ -76,6 +76,7 @@
 import org.opensearch.index.engine.Engine;
 import org.opensearch.index.engine.EngineConfigFactory;
 import org.opensearch.index.engine.EngineFactory;
+import org.opensearch.index.engine.MergedSegmentWarmerFactory;
 import org.opensearch.index.fielddata.IndexFieldDataCache;
 import org.opensearch.index.fielddata.IndexFieldDataService;
 import org.opensearch.index.mapper.MapperService;
@@ -577,7 +578,8 @@ public synchronized IndexShard createShard(
         final RepositoriesService repositoriesService,
         final DiscoveryNode targetNode,
         @Nullable DiscoveryNode sourceNode,
-        DiscoveryNodes discoveryNodes
+        DiscoveryNodes discoveryNodes,
+        MergedSegmentWarmerFactory mergedSegmentWarmerFactory
     ) throws IOException {
         Objects.requireNonNull(retentionLeaseSyncer);
         /*
@@ -704,7 +706,8 @@ protected void closeInternal() {
                 remoteStoreSettings,
                 seedRemote,
                 discoveryNodes,
-                segmentReplicationStatsProvider
+                segmentReplicationStatsProvider,
+                mergedSegmentWarmerFactory
             );
             eventListener.indexShardStateChanged(indexShard, null, indexShard.state(), "shard created");
             eventListener.afterIndexShardCreated(indexShard);

server/src/main/java/org/opensearch/index/engine/EngineConfig.java

@@ -33,6 +33,7 @@
 
 import org.apache.lucene.analysis.Analyzer;
 import org.apache.lucene.codecs.Codec;
+import org.apache.lucene.index.IndexWriter;
 import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.MergePolicy;
 import org.apache.lucene.search.QueryCache;
@@ -248,6 +249,8 @@
 
     private final TranslogFactory translogFactory;
 
+    private final IndexWriter.IndexReaderWarmer indexReaderWarmer;
+
     /**
      * Creates a new {@link org.opensearch.index.engine.EngineConfig}
      */
@@ -299,6 +302,7 @@
         this.translogFactory = builder.translogFactory;
         this.leafSorter = builder.leafSorter;
         this.documentMapperForTypeSupplier = builder.documentMapperForTypeSupplier;
+        this.indexReaderWarmer = builder.indexReaderWarmer;
     }
 
     /**
@@ -523,6 +527,14 @@
         return translogFactory;
     }
 
+    /**
+     * Returns the underlying indexReaderWarmer
+     * @return the indexReaderWarmer
+     */
+    public IndexWriter.IndexReaderWarmer getIndexReaderWarmer() {
+        return indexReaderWarmer;
+    }
+
     /**
      * A supplier supplies tombstone documents which will be used in soft-update methods.
      * The returned document consists only _uid, _seqno, _term and _version fields; other metadata fields are excluded.
@@ -598,6 +610,7 @@
         private TranslogFactory translogFactory = new InternalTranslogFactory();
         private Supplier<DocumentMapperForType> documentMapperForTypeSupplier;
         Comparator<LeafReader> leafSorter;
+        private IndexWriter.IndexReaderWarmer indexReaderWarmer;
 
         public Builder shardId(ShardId shardId) {
             this.shardId = shardId;
@@ -739,6 +752,11 @@
             return this;
         }
 
+        public Builder indexReaderWarmer(IndexWriter.IndexReaderWarmer indexReaderWarmer) {
+            this.indexReaderWarmer = indexReaderWarmer;
+            return this;
+        }
+
         public EngineConfig build() {
             return new EngineConfig(this);
         }

server/src/main/java/org/opensearch/index/engine/EngineConfigFactory.java

@@ -10,6 +10,7 @@
 
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.analysis.Analyzer;
+import org.apache.lucene.index.IndexWriter;
 import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.MergePolicy;
 import org.apache.lucene.search.QueryCache;
@@ -156,7 +157,8 @@ public EngineConfig newEngineConfig(
         BooleanSupplier startedPrimarySupplier,
         TranslogFactory translogFactory,
         Comparator<LeafReader> leafSorter,
-        Supplier<DocumentMapperForType> documentMapperForTypeSupplier
+        Supplier<DocumentMapperForType> documentMapperForTypeSupplier,
+        IndexWriter.IndexReaderWarmer indexReaderWarmer
     ) {
         CodecService codecServiceToUse = codecService;
         if (codecService == null && this.codecServiceFactory != null) {
@@ -191,6 +193,7 @@ public EngineConfig newEngineConfig(
             .translogFactory(translogFactory)
             .leafSorter(leafSorter)
             .documentMapperForTypeSupplier(documentMapperForTypeSupplier)
+            .indexReaderWarmer(indexReaderWarmer)
             .build();
     }
 

server/src/main/java/org/opensearch/index/engine/InternalEngine.java

@@ -81,6 +81,7 @@
 import org.opensearch.common.lucene.uid.VersionsAndSeqNoResolver.DocIdAndSeqNo;
 import org.opensearch.common.metrics.CounterMetric;
 import org.opensearch.common.unit.TimeValue;
+import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.common.util.concurrent.AbstractRunnable;
 import org.opensearch.common.util.concurrent.KeyedLock;
 import org.opensearch.common.util.concurrent.ReleasableLock;
@@ -2382,6 +2383,10 @@
         if (config().getLeafSorter() != null) {
             iwc.setLeafSorter(config().getLeafSorter()); // The default segment search order
         }
+        if (FeatureFlags.isEnabled(FeatureFlags.MERGED_SEGMENT_WARMER_EXPERIMENTAL_SETTING)
+            && config().getIndexSettings().isSegRepEnabledOrRemoteNode()) {
+            iwc.setMergedSegmentWarmer(config().getIndexReaderWarmer());
+        }
         return iwc;
     }
 

server/src/main/java/org/opensearch/index/engine/LocalMergedSegmentWarmer.java

@@ -0,0 +1,63 @@
+/*
+ * 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.
+ */
+
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+package org.opensearch.index.engine;
+
+import org.apache.lucene.index.IndexWriter;
+import org.apache.lucene.index.LeafReader;
+import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.indices.recovery.RecoverySettings;
+import org.opensearch.transport.TransportService;
+
+import java.io.IOException;
+
+/**
+ * Implementation of a {@link IndexWriter.IndexReaderWarmer} when local on-disk segment replication is enabled.
+ *
+ * @opensearch.internal
+ */
+public class LocalMergedSegmentWarmer implements IndexWriter.IndexReaderWarmer {
+    private final TransportService transportService;
+    private final RecoverySettings recoverySettings;
+    private final ClusterService clusterService;
+
+    public LocalMergedSegmentWarmer(TransportService transportService, RecoverySettings recoverySettings, ClusterService clusterService) {
+        this.transportService = transportService;
+        this.recoverySettings = recoverySettings;
+        this.clusterService = clusterService;
+    }
+
+    @Override
+    public void warm(LeafReader leafReader) throws IOException {
+        // TODO: node-node segment replication merged segment warmer
+    }
+}

server/src/main/java/org/opensearch/index/engine/MergedSegmentWarmerFactory.java

@@ -0,0 +1,69 @@
+/*
+ * 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.
+ */
+
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+package org.opensearch.index.engine;
+
+import org.apache.lucene.index.IndexWriter;
+import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.annotation.ExperimentalApi;
+import org.opensearch.index.shard.IndexShard;
+import org.opensearch.indices.recovery.RecoverySettings;
+import org.opensearch.transport.TransportService;
+
+/**
+ * MergedSegmentWarmerFactory to enable creation of various local on-disk
+ * and remote store flavors of {@link IndexWriter.IndexReaderWarmer}
+ *
+ * @opensearch.internal
+ */
+@ExperimentalApi
+public class MergedSegmentWarmerFactory {
+    private final TransportService transportService;
+    private final RecoverySettings recoverySettings;
+    private final ClusterService clusterService;
+
+    public MergedSegmentWarmerFactory(TransportService transportService, RecoverySettings recoverySettings, ClusterService clusterService) {
+        this.transportService = transportService;
+        this.recoverySettings = recoverySettings;
+        this.clusterService = clusterService;
+    }
+
+    public IndexWriter.IndexReaderWarmer get(IndexShard shard) {
+        if (shard.indexSettings().isAssignedOnRemoteNode()) {
+            return new RemoteStoreMergedSegmentWarmer(transportService, recoverySettings, clusterService);
+        } else if (shard.indexSettings().isSegRepLocalEnabled()) {
+            return new LocalMergedSegmentWarmer(transportService, recoverySettings, clusterService);
+        } else {
+            return null;
+        }
+    }
+}