Add warmup phase for pull-based ingestion

CHANGELOG.md

@@ -24,6 +24,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add ref_path support for package-based hunspell dictionary loading ([#20840](https://github.com/opensearch-project/OpenSearch/pull/20840))
 - Add support for enabling pluggable data formats, starting with phase-1 of decoupling shard from engine, and introducing basic abstractions ([#20675](https://github.com/opensearch-project/OpenSearch/pull/20675))
 
+- 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))
 ### Changed
 - Make telemetry `Tags` immutable ([#20788](https://github.com/opensearch-project/OpenSearch/pull/20788))
 - Move Randomness from server to libs/common ([#20570](https://github.com/opensearch-project/OpenSearch/pull/20570))

plugins/ingestion-kafka/src/internalClusterTest/java/org/opensearch/plugin/kafka/IngestFromKafkaIT.java

@@ -1338,4 +1338,55 @@ public void testKafkaIngestionWithFieldMappingMapper_VariousConfigurations() thr
                 && !docs.get("jkl").containsKey("expired");
         });
     }
+
+    public void testWarmupPhase() throws Exception {
+        // Step 1: Publish 10 messages before creating the index
+        for (int i = 0; i < 10; i++) {
+            produceData(Integer.toString(i), "name" + i, "25");
+        }
+
+        // Step 2: Start cluster
+        internalCluster().startClusterManagerOnlyNode();
+        final String nodeA = internalCluster().startDataOnlyNode();
+
+        // Step 3: Create index with warmup enabled, lag_threshold=0, and long timeout
+        createIndex(
+            indexName,
+            Settings.builder()
+                .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
+                .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0)
+                .put("ingestion_source.type", "kafka")
+                .put("ingestion_source.pointer.init.reset", "earliest")
+                .put("ingestion_source.param.topic", topicName)
+                .put("ingestion_source.param.bootstrap_servers", kafka.getBootstrapServers())
+                .put("ingestion_source.warmup.lag_threshold", 0)
+                .put("ingestion_source.warmup.timeout", "10m")
+                .put("ingestion_source.all_active", true)
+                .build(),
+            "{\"properties\":{\"name\":{\"type\": \"text\"},\"age\":{\"type\": \"integer\"}}}}"
+        );
+
+        ensureGreen(indexName);
+
+        // Step 4: Wait for poller to enter POLLING state (warmup complete)
+        waitForState(() -> {
+            GetIngestionStateResponse ingestionState = getIngestionState(indexName);
+            return ingestionState.getShardStates().length == 1
+                && ingestionState.getShardStates()[0].getPollerState().equalsIgnoreCase("polling");
+        });
+
+        // Step 5: Validate all 10 documents are searchable after warmup
+        // Refresh to make all docs visible, then assert directly (not waitForSearchableDocs)
+        // since all 10 docs should already be indexed during the warmup phase
+        client(nodeA).admin().indices().prepareRefresh(indexName).get();
+        long docCount = client(nodeA).prepareSearch(indexName).setSize(0).get().getHits().getTotalHits().value();
+        assertEquals("All 10 documents should be searchable immediately after warmup completes", 10L, docCount);
+
+        // Step 6: Verify stats
+        PollingIngestStats stats = client(nodeA).admin().indices().prepareStats(indexName).get().getIndex(indexName).getShards()[0]
+            .getPollingIngestStats();
+        assertNotNull(stats);
+        assertEquals(10L, stats.getMessageProcessorStats().totalProcessedCount());
+        assertEquals(10L, stats.getConsumerStats().totalPolledCount());
+    }
 }

plugins/ingestion-kafka/src/test/java/org/opensearch/plugin/kafka/KafkaSingleNodeTests.java

@@ -140,6 +140,7 @@ public void testPauseAndResumeAPIs() throws Exception {
     }
 
     // This test validates shard initialization does not fail due to kafka connection errors.
+    // Warmup is disabled since we're testing error handling, not warmup behavior.
     public void testShardInitializationUsingUnknownTopic() throws Exception {
         createIndexWithMappingSource(
             indexName,

server/src/main/java/org/opensearch/cluster/metadata/IndexMetadata.java

@@ -1031,6 +1031,32 @@ public Iterator<Setting<?>> settings() {
         }, Property.IndexScope, Property.Final)
     );
 
+    /**
+     * Defines the maximum time to wait for lag to catch up during warmup phase.
+     * A value of -1 means warmup is disabled (the default). A value >= 0 enables warmup with that timeout.
+     */
+    public static final String SETTING_INGESTION_SOURCE_WARMUP_TIMEOUT = "index.ingestion_source.warmup.timeout";
+    public static final Setting<TimeValue> INGESTION_SOURCE_WARMUP_TIMEOUT_SETTING = Setting.timeSetting(
+        SETTING_INGESTION_SOURCE_WARMUP_TIMEOUT,
+        TimeValue.timeValueMillis(-1),
+        TimeValue.timeValueMillis(-1),
+        Property.IndexScope,
+        Property.Final
+    );
+
+    /**
+     * Defines the acceptable pointer-based lag threshold. Warmup completes when lag is at or below this value.
+     * A value of 0 means fully caught up (no lag).
+     */
+    public static final String SETTING_INGESTION_SOURCE_WARMUP_LAG_THRESHOLD = "index.ingestion_source.warmup.lag_threshold";
+    public static final Setting<Long> INGESTION_SOURCE_WARMUP_LAG_THRESHOLD_SETTING = Setting.longSetting(
+        SETTING_INGESTION_SOURCE_WARMUP_LAG_THRESHOLD,
+        100L,
+        0L,
+        Property.IndexScope,
+        Property.Final
+    );
+
     /**
      * an internal index format description, allowing us to find out if this index is upgraded or needs upgrading
      */
@@ -1301,6 +1327,12 @@ public IngestionSource getIngestionSource() {
             final IngestionMessageMapper.MapperType mapperType = INGESTION_SOURCE_MAPPER_TYPE_SETTING.get(settings);
             final Map<String, Object> mapperSettings = INGESTION_SOURCE_MAPPER_SETTINGS.getAsMap(settings);
 
+            // Warmup settings
+            final IngestionSource.WarmupConfig warmupConfig = new IngestionSource.WarmupConfig(
+                INGESTION_SOURCE_WARMUP_TIMEOUT_SETTING.get(settings),
+                INGESTION_SOURCE_WARMUP_LAG_THRESHOLD_SETTING.get(settings)
+            );
+
             return new IngestionSource.Builder(ingestionSourceType).setParams(ingestionSourceParams)
                 .setPointerInitReset(pointerInitReset)
                 .setErrorStrategy(errorStrategy)
@@ -1312,6 +1344,7 @@ public IngestionSource getIngestionSource() {
                 .setPointerBasedLagUpdateInterval(pointerBasedLagUpdateInterval)
                 .setMapperType(mapperType)
                 .setMapperSettings(mapperSettings)
+                .setWarmupConfig(warmupConfig)
                 .build();
         }
         return null;

server/src/main/java/org/opensearch/cluster/metadata/IngestionSource.java

@@ -27,6 +27,8 @@
 import static org.opensearch.cluster.metadata.IndexMetadata.INGESTION_SOURCE_NUM_PROCESSOR_THREADS_SETTING;
 import static org.opensearch.cluster.metadata.IndexMetadata.INGESTION_SOURCE_POINTER_BASED_LAG_UPDATE_INTERVAL_SETTING;
 import static org.opensearch.cluster.metadata.IndexMetadata.INGESTION_SOURCE_POLL_TIMEOUT;
+import static org.opensearch.cluster.metadata.IndexMetadata.INGESTION_SOURCE_WARMUP_LAG_THRESHOLD_SETTING;
+import static org.opensearch.cluster.metadata.IndexMetadata.INGESTION_SOURCE_WARMUP_TIMEOUT_SETTING;
 
 /**
  * Class encapsulating the configuration of an ingestion source.
@@ -45,6 +47,7 @@ public class IngestionSource {
     private final TimeValue pointerBasedLagUpdateInterval;
     private final IngestionMessageMapper.MapperType mapperType;
     private final Map<String, Object> mapperSettings;
+    private final WarmupConfig warmupConfig;
 
     private IngestionSource(
         String type,
@@ -58,7 +61,8 @@ private IngestionSource(
         boolean allActiveIngestion,
         TimeValue pointerBasedLagUpdateInterval,
         IngestionMessageMapper.MapperType mapperType,
-        Map<String, Object> mapperSettings
+        Map<String, Object> mapperSettings,
+        WarmupConfig warmupConfig
     ) {
         this.type = type;
         this.pointerInitReset = pointerInitReset;
@@ -72,6 +76,7 @@ private IngestionSource(
         this.pointerBasedLagUpdateInterval = pointerBasedLagUpdateInterval;
         this.mapperType = mapperType;
         this.mapperSettings = mapperSettings != null ? Collections.unmodifiableMap(mapperSettings) : Collections.emptyMap();
+        this.warmupConfig = warmupConfig;
     }
 
     public String getType() {
@@ -122,6 +127,10 @@ public Map<String, Object> getMapperSettings() {
         return mapperSettings;
     }
 
+    public WarmupConfig getWarmupConfig() {
+        return warmupConfig;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
@@ -138,7 +147,8 @@ public boolean equals(Object o) {
             && Objects.equals(allActiveIngestion, ingestionSource.allActiveIngestion)
             && Objects.equals(pointerBasedLagUpdateInterval, ingestionSource.pointerBasedLagUpdateInterval)
             && Objects.equals(mapperType, ingestionSource.mapperType)
-            && Objects.equals(mapperSettings, ingestionSource.mapperSettings);
+            && Objects.equals(mapperSettings, ingestionSource.mapperSettings)
+            && Objects.equals(warmupConfig, ingestionSource.warmupConfig);
     }
 
     @Override
@@ -155,7 +165,8 @@ public int hashCode() {
             allActiveIngestion,
             pointerBasedLagUpdateInterval,
             mapperType,
-            mapperSettings
+            mapperSettings,
+            warmupConfig
         );
     }
 
@@ -190,6 +201,8 @@ public String toString() {
             + '\''
             + ", mapperSettings="
             + mapperSettings
+            + ", warmupConfig="
+            + warmupConfig
             + '}';
     }
 
@@ -233,6 +246,21 @@ public String toString() {
         }
     }
 
+    /**
+     * Record encapsulating the warmup configuration for pull-based ingestion.
+     * When warmup is enabled (timeout >= 0), shards will wait for lag to catch up before serving queries
+     * after node restart or shard relocation. A timeout of -1 means warmup is disabled.
+     */
+    @PublicApi(since = "3.6.0")
+    public record WarmupConfig(TimeValue timeout, long lagThreshold) {
+        /**
+         * Returns true if warmup is enabled (timeout >= 0).
+         */
+        public boolean isEnabled() {
+            return timeout.millis() >= 0;
+        }
+    }
+
     /**
      * Builder for {@link IngestionSource}.
      *
@@ -253,6 +281,9 @@ public static class Builder {
         );
         private IngestionMessageMapper.MapperType mapperType = INGESTION_SOURCE_MAPPER_TYPE_SETTING.getDefault(Settings.EMPTY);
         private Map<String, Object> mapperSettings = new HashMap<>();
+        // Warmup configuration
+        private TimeValue warmupTimeout = INGESTION_SOURCE_WARMUP_TIMEOUT_SETTING.getDefault(Settings.EMPTY);
+        private long warmupLagThreshold = INGESTION_SOURCE_WARMUP_LAG_THRESHOLD_SETTING.getDefault(Settings.EMPTY);
 
         public Builder(String type) {
             this.type = type;
@@ -269,6 +300,10 @@ public Builder(IngestionSource ingestionSource) {
             this.pointerBasedLagUpdateInterval = ingestionSource.pointerBasedLagUpdateInterval;
             this.mapperType = ingestionSource.mapperType;
             this.mapperSettings = new HashMap<>(ingestionSource.mapperSettings);
+            // Copy warmup config
+            WarmupConfig wc = ingestionSource.warmupConfig;
+            this.warmupTimeout = wc.timeout();
+            this.warmupLagThreshold = wc.lagThreshold();
         }
 
         public Builder setPointerInitReset(PointerInitReset pointerInitReset) {
@@ -331,7 +366,24 @@ public Builder setMapperSettings(Map<String, Object> mapperSettings) {
             return this;
         }
 
+        public Builder setWarmupTimeout(TimeValue warmupTimeout) {
+            this.warmupTimeout = warmupTimeout;
+            return this;
+        }
+
+        public Builder setWarmupLagThreshold(long warmupLagThreshold) {
+            this.warmupLagThreshold = warmupLagThreshold;
+            return this;
+        }
+
+        public Builder setWarmupConfig(WarmupConfig warmupConfig) {
+            this.warmupTimeout = warmupConfig.timeout();
+            this.warmupLagThreshold = warmupConfig.lagThreshold();
+            return this;
+        }
+
         public IngestionSource build() {
+            WarmupConfig warmupConfig = new WarmupConfig(warmupTimeout, warmupLagThreshold);
             return new IngestionSource(
                 type,
                 pointerInitReset,
@@ -344,7 +396,8 @@ public IngestionSource build() {
                 allActiveIngestion,
                 pointerBasedLagUpdateInterval,
                 mapperType,
-                mapperSettings
+                mapperSettings,
+                warmupConfig
             );
         }
 

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

@@ -292,6 +292,8 @@ public final class IndexScopedSettings extends AbstractScopedSettings {
                 IndexMetadata.INGESTION_SOURCE_POINTER_BASED_LAG_UPDATE_INTERVAL_SETTING,
                 IndexMetadata.INGESTION_SOURCE_MAPPER_TYPE_SETTING,
                 IndexMetadata.INGESTION_SOURCE_MAPPER_SETTINGS,
+                IndexMetadata.INGESTION_SOURCE_WARMUP_TIMEOUT_SETTING,
+                IndexMetadata.INGESTION_SOURCE_WARMUP_LAG_THRESHOLD_SETTING,
 
                 // Settings for search replica
                 IndexMetadata.INDEX_NUMBER_OF_SEARCH_REPLICAS_SETTING,

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

@@ -155,6 +155,7 @@ private void initializeStreamPoller(
             .pointerBasedLagUpdateInterval(ingestionSource.getPointerBasedLagUpdateInterval().millis())
             .mapperType(ingestionSource.getMapperType())
             .mapperSettings(ingestionSource.getMapperSettings())
+            .warmupConfig(ingestionSource.getWarmupConfig())
             .build();
         registerStreamPollerListener();
 
@@ -662,4 +663,25 @@ public ShardIngestionState getIngestionState() {
             shardPointer != null ? shardPointer.toString() : ""
         );
     }
+
+    /**
+     * Block until warmup is complete or timeout occurs.
+     * This method handles all warmup logic internally. On timeout, always logs a warning and proceeds.
+     *
+     * @throws InterruptedException if the thread is interrupted while waiting
+     */
+    public void awaitWarmupComplete() throws InterruptedException {
+        IngestionSource ingestionSource = engineConfig.getIndexSettings().getIndexMetadata().getIngestionSource();
+        if (ingestionSource == null || !ingestionSource.getWarmupConfig().isEnabled() || streamPoller.isPaused()) {
+            return;
+        }
+
+        long timeoutMs = ingestionSource.getWarmupConfig().timeout().millis();
+        boolean completed = streamPoller.awaitWarmupComplete(timeoutMs);
+
+        if (!completed) {
+            logger.warn("Ingestion warmup timed out for shard after {}ms, proceeding with potentially stale data.", timeoutMs);
+        }
+    }
+
 }

server/src/main/java/org/opensearch/index/shard/IndexShard.java

@@ -2487,6 +2487,10 @@ public void postRecovery(String reason) throws IndexShardStartedException, Index
             // responded to in addRefreshListener. The refresh must happen under the same mutex used in addRefreshListener
             // and before moving this shard to POST_RECOVERY state (i.e., allow to read from this shard).
             getIndexer().refresh("post_recovery");
+
+            // Wait for ingestion warmup if enabled (pull-based ingestion only)
+            handlePullBasedIngestionWarmup(getIndexer());
+
             synchronized (mutex) {
                 if (state == IndexShardState.CLOSED) {
                     throw new IndexShardClosedException(shardId);
@@ -2500,6 +2504,29 @@ public void postRecovery(String reason) throws IndexShardStartedException, Index
         }
     }
 
+    /**
+     * Handles warmup for pull-based ingestion (PBI) engines.
+     * When warmup is enabled, this method blocks until the shard has caught up with the streaming source
+     * or until the configured timeout is reached.
+     *
+     * @param indexer the indexer to check for warmup
+     */
+    private void handlePullBasedIngestionWarmup(Indexer indexer) {
+        if (!(indexer instanceof EngineBackedIndexer)) {
+            return;
+        }
+        Engine engine = ((EngineBackedIndexer) indexer).getEngine();
+        if (engine instanceof IngestionEngine) {
+            IngestionEngine ingestionEngine = (IngestionEngine) engine;
+            try {
+                ingestionEngine.awaitWarmupComplete();
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+                throw new OpenSearchException("Interrupted waiting for ingestion warmup", e);
+            }
+        }
+    }
+
     /**
      * called before starting to copy index files over
      */