Add WarmerRefreshListerner to NRTReplicationEngine to warm replica shards (opensearch-project#20650)

Kartik031299 · kbansal2 · Deepti24 · commit b49875e26698 · 2026-03-10T14:17:02.000+05:30
Signed-off-by: Kartik Bansal &lt;kbansal2@atlassian.com&gt;
Co-authored-by: Kartik Bansal &lt;kbansal2@atlassian.com&gt;
Signed-off-by: Deepti24 &lt;chauhan.deepti24@gmail.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -37,6 +37,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Introduce AdditionalCodecs and EnginePlugin::getAdditionalCodecs hook to allow additional Codec registration ([#20411](https://github.com/opensearch-project/OpenSearch/pull/20411))
 - Introduced strategy planner interfaces for indexing and deletion ([#20585](https://github.com/opensearch-project/OpenSearch/pull/20585))
 - Implement FieldMappingIngestionMessageMapper for pull-based ingestion ([#20729](https://github.com/opensearch-project/OpenSearch/pull/20729))
+- Added support of WarmerRefreshListener in NRTReplicationEngine to trigger warmer after replication on replica shards ([#20650](https://github.com/opensearch-project/OpenSearch/pull/20650))
 
 ### Changed
 - Move Randomness from server to libs/common ([#20570](https://github.com/opensearch-project/OpenSearch/pull/20570))
diff --git a/server/src/internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationReplicaIndexWarmerIT.java b/server/src/internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationReplicaIndexWarmerIT.java
@@ -0,0 +1,141 @@
+/*
+ * 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.indices.replication;
+
+import org.opensearch.action.admin.indices.stats.IndicesStatsResponse;
+import org.opensearch.action.admin.indices.stats.ShardStats;
+import org.opensearch.index.warmer.WarmerStats;
+import org.opensearch.test.OpenSearchIntegTestCase;
+import org.junit.Before;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked;
+import static org.hamcrest.Matchers.greaterThan;
+
+/**
+ * Integration tests that verify index warming (e.g., eager global ordinals loading) works correctly
+ * on NRT replica shards during segment replication. This validates the end-to-end flow where
+ * index warmer gets triggered when new segments arrive on replicas via segment replication.
+ *
+ * <p>Uses a keyword field with {@code eager_global_ordinals: true}
+ * to exercise the warming path during segment replication.</p>
+ */
+@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, numDataNodes = 0)
+public class SegmentReplicationReplicaIndexWarmerIT extends SegmentReplicationBaseIT {
+    @Before
+    public void setup() {
+        internalCluster().startClusterManagerOnlyNode();
+    }
+
+    /**
+     * Verifies that eager global ordinals are loaded on both primary and replica shards
+     * after segment replication, by checking warmer invocation metrics from the index stats API.
+     *
+     * <p>This test ensures that the {@code WarmerRefreshListener} in {@code NRTReplicationEngine}
+     * correctly invokes the {@code IndexWarmer} chain on replica shards, which in turn loads
+     * global ordinals for the keyword field with {@code eager_global_ordinals: true}.</p>
+     */
+    public void testEagerGlobalOrdinalsLoadedOnReplicaAfterSegmentReplication() throws Exception {
+        final String primaryNode = internalCluster().startDataOnlyNode();
+        assertAcked(
+            prepareCreate(INDEX_NAME).setSettings(indexSettings()).setMapping("category", "type=keyword,eager_global_ordinals=true")
+        );
+        ensureYellow(INDEX_NAME);
+        final String replicaNode = internalCluster().startDataOnlyNode();
+        ensureGreen(INDEX_NAME);
+
+        final int docCount = randomIntBetween(10, 20);
+        indexTestDocuments(docCount, -1);
+        final Map<String, Long> warmerTotalBeforeRefresh = getWarmerTotalPerShardType(INDEX_NAME);
+        refresh(INDEX_NAME);
+        waitForSearchableDocs(docCount, primaryNode, replicaNode);
+
+        // Assert warmer invoked metrics from index stats API for both primary and replica shards
+        assertBusy(() -> {
+            Map<String, Long> warmerTotalsAfterRefresh = getWarmerTotalPerShardType(INDEX_NAME);
+            compareWarmerTotals(warmerTotalBeforeRefresh, warmerTotalsAfterRefresh);
+        });
+    }
+
+    /**
+     * Verifies that warmer invocations continue on replica shards after a force merge on the primary,
+     * followed by segment replication. This ensures warming is triggered on subsequent segment updates.
+     */
+    public void testWarmerInvokedOnReplicaAfterForceMerge() throws Exception {
+        final String primaryNode = internalCluster().startDataOnlyNode();
+        assertAcked(
+            prepareCreate(INDEX_NAME).setSettings(indexSettings()).setMapping("category", "type=keyword,eager_global_ordinals=true")
+        );
+        ensureYellow(INDEX_NAME);
+        final String replicaNode = internalCluster().startDataOnlyNode();
+        ensureGreen(INDEX_NAME);
+
+        // Index documents in batches to create multiple segments
+        final int docCount = randomIntBetween(10, 20);
+        indexTestDocuments(docCount, 3);
+        waitForSearchableDocs(docCount, primaryNode, replicaNode);
+
+        // Capture warmer stats before force merge
+        final Map<String, Long> warmerTotalBeforeForceMerge = getWarmerTotalPerShardType(INDEX_NAME);
+
+        // Force merge and wait for segment replication to complete
+        client().admin().indices().prepareForceMerge(INDEX_NAME).setMaxNumSegments(1).setFlush(true).get();
+
+        // Verify warmer was invoked again on replica after force merge replication
+        assertBusy(() -> {
+            final Map<String, Long> warmerTotalAfterForceMerge = getWarmerTotalPerShardType(INDEX_NAME);
+            compareWarmerTotals(warmerTotalBeforeForceMerge, warmerTotalAfterForceMerge, true);
+        });
+    }
+
+    private Map<String, Long> getWarmerTotalPerShardType(String indexName) {
+        IndicesStatsResponse shardStatsArray = client().admin().indices().prepareStats(indexName).clear().setWarmer(true).get();
+        Map<String, Long> warmerTotals = new HashMap<>();
+        warmerTotals.put("primary", 0L);
+        warmerTotals.put("replica", 0L);
+        for (ShardStats shardStats : shardStatsArray.getShards()) {
+            WarmerStats warmerStats = shardStats.getStats().getWarmer();
+            if (warmerStats != null) {
+                String key = shardStats.getShardRouting().primary() ? "primary" : "replica";
+                warmerTotals.merge(key, warmerStats.total(), Long::sum);
+            }
+        }
+        return warmerTotals;
+    }
+
+    private void indexTestDocuments(int docCount, int refreshAfterDocs) {
+        for (int i = 0; i < docCount; i++) {
+            client().prepareIndex(INDEX_NAME).setId(Integer.toString(i)).setSource("category", "value-" + i).get();
+            if (refreshAfterDocs != -1 && i % refreshAfterDocs == 0) {
+                refresh(INDEX_NAME);
+            }
+        }
+    }
+
+    private void compareWarmerTotals(Map<String, Long> before, Map<String, Long> after) {
+        compareWarmerTotals(before, after, false);
+    }
+
+    private void compareWarmerTotals(Map<String, Long> before, Map<String, Long> after, boolean onReplicaOnly) {
+        for (String shardType : before.keySet()) {
+            if (onReplicaOnly && shardType.equals("primary")) {
+                continue; // Skip primary if we're only checking replica warmer totals
+            }
+            long beforeTotal = before.get(shardType);
+            long afterTotal = after.get(shardType);
+            assertThat(
+                "Warmer total for " + shardType + " should increase after segment replication",
+                afterTotal,
+                greaterThan(beforeTotal)
+            );
+        }
+    }
+}
diff --git a/server/src/main/java/org/opensearch/index/engine/NRTReplicationEngine.java b/server/src/main/java/org/opensearch/index/engine/NRTReplicationEngine.java
@@ -8,6 +8,7 @@
 
 package org.opensearch.index.engine;
 
+import org.apache.logging.log4j.Logger;
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.IndexCommit;
 import org.apache.lucene.index.SegmentInfos;
@@ -43,6 +44,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.BiFunction;
@@ -100,6 +102,9 @@ public NRTReplicationEngine(EngineConfig engineConfig) {
             for (ReferenceManager.RefreshListener listener : engineConfig.getInternalRefreshListener()) {
                 this.readerManager.addListener(listener);
             }
+            // Wire up a warmer listener to trigger index warming
+            // when new segments arrive via segment replication
+            this.readerManager.addListener(new WarmerRefreshListener(logger, isClosed, engineConfig, this.readerManager));
             final Map<String, String> userData = this.lastCommittedSegmentInfos.getUserData();
             final String translogUUID = Objects.requireNonNull(userData.get(Translog.TRANSLOG_UUID_KEY));
             translogManagerRef = new WriteOnlyTranslogManager(
@@ -560,4 +565,56 @@ private DirectoryReader getDirectoryReader() throws IOException {
             Lucene.SOFT_DELETES_FIELD
         );
     }
+
+    /**
+     * A {@link ReferenceManager.RefreshListener} that warms new segments when the reader is refreshed
+     * during segment replication. This ensures index warming (e.g., loading global ordinals via
+     * {@link Engine.Warmer}) occurs on NRT replica shards, consistent with the warming behavior
+     * in {@link InternalEngine} used by NRT primary shards.
+     *
+     * @opensearch.internal
+     */
+    static final class WarmerRefreshListener implements ReferenceManager.RefreshListener {
+        private final Engine.Warmer warmer;
+        private final Logger logger;
+        private final AtomicBoolean isEngineClosed;
+        private final NRTReplicationReaderManager readerManager;
+
+        WarmerRefreshListener(
+            Logger logger,
+            AtomicBoolean isEngineClosed,
+            EngineConfig engineConfig,
+            NRTReplicationReaderManager readerManager
+        ) {
+            this.warmer = engineConfig.getWarmer();
+            this.logger = logger;
+            this.isEngineClosed = isEngineClosed;
+            this.readerManager = readerManager;
+        }
+
+        @Override
+        public void beforeRefresh() throws IOException {}
+
+        @Override
+        public void afterRefresh(boolean didRefresh) {
+            if (didRefresh && warmer != null) {
+                try {
+                    OpenSearchDirectoryReader reader = readerManager.acquire();
+                    try {
+                        warmer.warm(reader);
+                    } catch (Exception e) {
+                        if (isEngineClosed.get() == false) {
+                            logger.warn("failed to warm reader replica", e);
+                        }
+                    } finally {
+                        readerManager.release(reader);
+                    }
+                } catch (IOException e) {
+                    if (isEngineClosed.get() == false) {
+                        logger.warn("failed to acquire reader for warming on replica", e);
+                    }
+                }
+            }
+        }
+    }
 }
diff --git a/server/src/test/java/org/opensearch/index/engine/InternalEngineTests.java b/server/src/test/java/org/opensearch/index/engine/InternalEngineTests.java
@@ -3477,7 +3477,8 @@ public void onFailedEngine(String reason, Exception e) {
                     globalCheckpoint::get,
                     retentionLeasesHolder::get,
                     new NoneCircuitBreakerService(),
-                    eventListener
+                    eventListener,
+                    null
                 )
             );
 
@@ -3578,7 +3579,8 @@ public void onFailedEngine(String reason, Exception e) {
                     globalCheckpoint::get,
                     retentionLeasesHolder::get,
                     new NoneCircuitBreakerService(),
-                    eventListener
+                    eventListener,
+                    null
                 )
             );
 
@@ -3672,7 +3674,8 @@ public void onFailedEngine(String reason, Exception e) {
                     globalCheckpoint::get,
                     retentionLeasesHolder::get,
                     new NoneCircuitBreakerService(),
-                    eventListener
+                    eventListener,
+                    null
                 )
             );
 
@@ -3769,7 +3772,8 @@ public void onFailedEngine(String reason, Exception e) {
                     globalCheckpoint::get,
                     retentionLeasesHolder::get,
                     new NoneCircuitBreakerService(),
-                    eventListener
+                    eventListener,
+                    null
                 )
             );
 
@@ -7216,7 +7220,18 @@ public void testKeepMinRetainedSeqNoByMergePolicy() throws IOException {
         Set<Long> existingSeqNos = new HashSet<>();
         store = createStore();
         engine = createEngine(
-            config(indexSettings, store, createTempDir(), newMergePolicy(), null, null, globalCheckpoint::get, retentionLeasesHolder::get)
+            config(
+                indexSettings,
+                store,
+                createTempDir(),
+                newMergePolicy(),
+                null,
+                null,
+                null,
+                globalCheckpoint::get,
+                retentionLeasesHolder::get,
+                new NoneCircuitBreakerService()
+            )
         );
         assertThat(engine.getMinRetainedSeqNo(), equalTo(0L));
         long lastMinRetainedSeqNo = engine.getMinRetainedSeqNo();
diff --git a/server/src/test/java/org/opensearch/index/engine/NRTReplicationEngineTests.java b/server/src/test/java/org/opensearch/index/engine/NRTReplicationEngineTests.java
diff --git a/test/framework/src/main/java/org/opensearch/index/engine/EngineTestCase.java b/test/framework/src/main/java/org/opensearch/index/engine/EngineTestCase.java
