diff --git a/CHANGELOG.md b/CHANGELOG.md index 276ba96567449..84b598fd7a9f8 100644 --- a/CHANGELOG.md +++ b/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)) diff --git a/plugins/ingestion-kafka/src/internalClusterTest/java/org/opensearch/plugin/kafka/IngestFromKafkaIT.java b/plugins/ingestion-kafka/src/internalClusterTest/java/org/opensearch/plugin/kafka/IngestFromKafkaIT.java index 814021ddf0b32..e0958b32f0bbc 100644 --- a/plugins/ingestion-kafka/src/internalClusterTest/java/org/opensearch/plugin/kafka/IngestFromKafkaIT.java +++ b/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()); + } } diff --git a/plugins/ingestion-kafka/src/test/java/org/opensearch/plugin/kafka/KafkaSingleNodeTests.java b/plugins/ingestion-kafka/src/test/java/org/opensearch/plugin/kafka/KafkaSingleNodeTests.java index 7a14782d0a1a8..2aa0c81731698 100644 --- a/plugins/ingestion-kafka/src/test/java/org/opensearch/plugin/kafka/KafkaSingleNodeTests.java +++ b/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, diff --git a/server/src/main/java/org/opensearch/cluster/metadata/IndexMetadata.java b/server/src/main/java/org/opensearch/cluster/metadata/IndexMetadata.java index e43df0ea35b80..eb528115bbc02 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/IndexMetadata.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/IndexMetadata.java @@ -1031,6 +1031,32 @@ public Iterator> 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 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 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 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; diff --git a/server/src/main/java/org/opensearch/cluster/metadata/IngestionSource.java b/server/src/main/java/org/opensearch/cluster/metadata/IngestionSource.java index ff0bda28cc88e..6b1400c305e08 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/IngestionSource.java +++ b/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 mapperSettings; + private final WarmupConfig warmupConfig; private IngestionSource( String type, @@ -58,7 +61,8 @@ private IngestionSource( boolean allActiveIngestion, TimeValue pointerBasedLagUpdateInterval, IngestionMessageMapper.MapperType mapperType, - Map mapperSettings + Map 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 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 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 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 ); } diff --git a/server/src/main/java/org/opensearch/common/settings/IndexScopedSettings.java b/server/src/main/java/org/opensearch/common/settings/IndexScopedSettings.java index ce63a66a4b27c..e242304e446b5 100644 --- a/server/src/main/java/org/opensearch/common/settings/IndexScopedSettings.java +++ b/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, diff --git a/server/src/main/java/org/opensearch/index/engine/IngestionEngine.java b/server/src/main/java/org/opensearch/index/engine/IngestionEngine.java index ef6df80f2a47f..ebcbd776135cd 100644 --- a/server/src/main/java/org/opensearch/index/engine/IngestionEngine.java +++ b/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); + } + } + } diff --git a/server/src/main/java/org/opensearch/index/shard/IndexShard.java b/server/src/main/java/org/opensearch/index/shard/IndexShard.java index a15dd0b1030b8..1c155c897acba 100644 --- a/server/src/main/java/org/opensearch/index/shard/IndexShard.java +++ b/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 */ diff --git a/server/src/main/java/org/opensearch/indices/pollingingest/DefaultStreamPoller.java b/server/src/main/java/org/opensearch/indices/pollingingest/DefaultStreamPoller.java index 2571f7b7ac29a..4dc51e0539d83 100644 --- a/server/src/main/java/org/opensearch/indices/pollingingest/DefaultStreamPoller.java +++ b/server/src/main/java/org/opensearch/indices/pollingingest/DefaultStreamPoller.java @@ -16,6 +16,7 @@ import org.opensearch.cluster.metadata.IngestionSource; import org.opensearch.common.Nullable; import org.opensearch.common.metrics.CounterMetric; +import org.opensearch.common.unit.TimeValue; import org.opensearch.index.IndexSettings; import org.opensearch.index.IngestionConsumerFactory; import org.opensearch.index.IngestionShardConsumer; @@ -30,8 +31,10 @@ import java.util.Locale; import java.util.Map; import java.util.Objects; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; /** * Default implementation of {@link StreamPoller} @@ -56,9 +59,15 @@ public class DefaultStreamPoller implements StreamPoller { private volatile boolean reinitializeConsumer; private volatile long lastPolledMessageTimestamp = 0; - private volatile long cachedPointerBasedLag = 0; + private volatile long cachedPointerBasedLag = -1; // -1 indicates poller has not consumed any message yet private volatile long lastPointerBasedLagUpdateTime = 0; + // Warmup configuration and state + private final IngestionSource.WarmupConfig warmupConfig; + private volatile boolean warmupComplete = false; + private volatile long warmupStartTime = 0; + private final CountDownLatch warmupLatch = new CountDownLatch(1); + @Nullable private IngestionShardConsumer consumer; private IngestionConsumerFactory consumerFactory; @@ -107,7 +116,8 @@ private DefaultStreamPoller( int blockingQueueSize, long pointerBasedLagUpdateIntervalMs, IngestionMessageMapper.MapperType mapperType, - Map mapperSettings + Map mapperSettings, + IngestionSource.WarmupConfig warmupConfig ) { this( startPointer, @@ -123,7 +133,8 @@ private DefaultStreamPoller( pollTimeout, pointerBasedLagUpdateIntervalMs, ingestionEngine.config().getIndexSettings(), - IngestionMessageMapper.create(mapperType.getName(), shardId, mapperSettings) + IngestionMessageMapper.create(mapperType.getName(), shardId, mapperSettings), + warmupConfig ); } @@ -144,7 +155,8 @@ private DefaultStreamPoller( int pollTimeout, long pointerBasedLagUpdateIntervalMs, IndexSettings indexSettings, - IngestionMessageMapper messageMapper + IngestionMessageMapper messageMapper, + IngestionSource.WarmupConfig warmupConfig ) { this.consumerFactory = Objects.requireNonNull(consumerFactory); this.consumerClientId = Objects.requireNonNull(consumerClientId); @@ -163,9 +175,14 @@ private DefaultStreamPoller( this.errorStrategy = errorStrategy; this.indexName = indexSettings.getIndex().getName(); this.messageMapper = Objects.requireNonNull(messageMapper); + this.warmupConfig = Objects.requireNonNull(warmupConfig); // handle initial poller states this.paused = initialState == State.PAUSED; + // If warmup is disabled, mark as complete immediately + if (!warmupConfig.isEnabled()) { + this.warmupComplete = true; + } } @Override @@ -195,10 +212,17 @@ protected void startPoll() { } logger.info("Starting poller for shard {}", shardId); + // Initialize warmup if enabled + if (warmupConfig.isEnabled() && !warmupComplete) { + warmupStartTime = System.currentTimeMillis(); + state = State.WARMING_UP; + logger.info("Starting warmup phase for index {} shard {}, waiting for lag to catch up", indexName, shardId); + } + while (true) { try { if (closed) { - state = State.CLOSED; + setStateWithWarmupAwareness(State.CLOSED); closeConsumer(); break; } @@ -212,8 +236,13 @@ protected void startPoll() { // Update lag periodically. Lag is updated even if the poller is paused. updatePointerBasedLagIfNeeded(); + // Check warmup status if not yet complete + if (!warmupComplete && warmupConfig.isEnabled()) { + updateWarmupStatus(); + } + if (paused || isWriteBlockEnabled) { - state = State.PAUSED; + setStateWithWarmupAwareness(State.PAUSED); try { Thread.sleep(DEFAULT_POLLER_SLEEP_PERIOD_MS); } catch (Throwable e) { @@ -222,7 +251,7 @@ protected void startPoll() { continue; } - state = State.POLLING; + setStateWithWarmupAwareness(State.POLLING); List> results; // Force the consumer to start from forcedShardPointer if available @@ -240,7 +269,7 @@ protected void startPoll() { continue; } - state = State.PROCESSING; + setStateWithWarmupAwareness(State.PROCESSING); // processRecords returns failed shard pointers. Update forcedShardPointer to the failed pointer to retry on next iteration // in case of failures forcedShardPointer = processRecords(results); @@ -355,6 +384,101 @@ public boolean isClosed() { return closed; } + @Override + public boolean isWarmupComplete() { + return warmupComplete || !warmupConfig.isEnabled(); + } + + /** + * Sets the poller state with warmup-aware logic. + * During warmup, POLLING and PROCESSING states are reported as WARMING_UP + * to allow monitoring via the ingestion state API. + * + * @param newState the desired state to set + */ + private void setStateWithWarmupAwareness(State newState) { + // CLOSED and PAUSED always take effect + if (newState == State.CLOSED || newState == State.PAUSED) { + this.state = newState; + return; + } + + // During warmup, stay in WARMING_UP instead of POLLING/PROCESSING + if (!isWarmupComplete()) { + this.state = State.WARMING_UP; + return; + } + + this.state = newState; + } + + @Override + public boolean awaitWarmupComplete(long timeoutMs) throws InterruptedException { + if (!warmupConfig.isEnabled() || isWarmupComplete()) { + return true; + } + + boolean completed = warmupLatch.await(timeoutMs, TimeUnit.MILLISECONDS); + if (!completed) { + logger.warn( + "Warmup timeout for index {} shard {} - proceeding with current lag (warmupComplete={})", + indexName, + shardId, + isWarmupComplete() + ); + } + return completed; + } + + /** + * Check if warmup conditions are met and mark warmup as complete if so. + * + * Warmup uses offset-based lag (cachedPointerBasedLag) which tracks the difference between + * the current consumer position and the end of the stream. This is the preferred mode for + * Kafka and other sources that support offset-based lag calculation. + * Note: cachedPointerBasedLag is -1 by default (indicating no messages consumed yet) and is only updated after updatePointerBasedLagIfNeeded() + * is called. + */ + private void updateWarmupStatus() { + // Skip warmup if poller is paused + if (paused) { + warmupComplete = true; + warmupLatch.countDown(); + logger.info("Warmup skipped for index {} shard {} - poller is paused", indexName, shardId); + return; + } + + long currentLag = cachedPointerBasedLag; + long threshold = warmupConfig.lagThreshold(); + + long elapsedTime = System.currentTimeMillis() - warmupStartTime; + boolean lagBelowThreshold = currentLag >= 0 && currentLag <= threshold; + boolean timeoutReached = elapsedTime >= warmupConfig.timeout().millis(); + + if (lagBelowThreshold) { + warmupComplete = true; + warmupLatch.countDown(); + logger.info( + "Warmup complete for index {} shard {} - lag {} is at or below threshold {}", + indexName, + shardId, + currentLag, + threshold + ); + } else if (timeoutReached) { + warmupComplete = true; + warmupLatch.countDown(); + logger.warn( + "Warmup timeout for index {} shard {} after {}ms - proceeding with lag {} (threshold was {})", + indexName, + shardId, + elapsedTime, + currentLag, + threshold + ); + } + } + /** * Returns the batch start pointer from where the poller can resume in case of shard recovery. The poller and * processor are decoupled in this implementation, and hence the latest pointer tracked by the processor acts as the @@ -608,6 +732,8 @@ public static class Builder { private long pointerBasedLagUpdateIntervalMs = 10000; private IngestionMessageMapper.MapperType mapperType = IngestionMessageMapper.MapperType.DEFAULT; private Map mapperSettings = Collections.emptyMap(); + // Warmup configuration - default matches IndexMetadata settings + private IngestionSource.WarmupConfig warmupConfig = new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(-1), 100L); /** * Initialize the builder with mandatory parameters @@ -715,6 +841,15 @@ public Builder mapperSettings(Map mapperSettings) { return this; } + /** + * Set warmup enabled + * Set warmup configuration + */ + public Builder warmupConfig(IngestionSource.WarmupConfig warmupConfig) { + this.warmupConfig = Objects.requireNonNull(warmupConfig); + return this; + } + /** * Build the DefaultStreamPoller instance */ @@ -735,7 +870,8 @@ public DefaultStreamPoller build() { blockingQueueSize, pointerBasedLagUpdateIntervalMs, mapperType, - mapperSettings + mapperSettings, + warmupConfig ); } } diff --git a/server/src/main/java/org/opensearch/indices/pollingingest/StreamPoller.java b/server/src/main/java/org/opensearch/indices/pollingingest/StreamPoller.java index dee5bd3384389..1d9249cb3f57c 100644 --- a/server/src/main/java/org/opensearch/indices/pollingingest/StreamPoller.java +++ b/server/src/main/java/org/opensearch/indices/pollingingest/StreamPoller.java @@ -83,11 +83,25 @@ public interface StreamPoller extends Closeable, ClusterStateListener { */ void requestConsumerReinitialization(IngestionSource updatedIngestionSource); + /** + * @return true if the warmup phase is complete and the shard is ready to serve + */ + boolean isWarmupComplete(); + + /** + * Block until warmup is complete or timeout occurs. + * @param timeoutMs maximum time to wait in milliseconds + * @return true if warmup completed, false if timeout + * @throws InterruptedException if the thread is interrupted while waiting + */ + boolean awaitWarmupComplete(long timeoutMs) throws InterruptedException; + /** * A state to indicate the current state of the poller */ enum State { NONE, + WARMING_UP, CLOSED, PAUSED, POLLING, diff --git a/server/src/test/java/org/opensearch/cluster/metadata/IngestionSourceTests.java b/server/src/test/java/org/opensearch/cluster/metadata/IngestionSourceTests.java index 1c3f74ead393e..f9a1ec9bb3f7d 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/IngestionSourceTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/IngestionSourceTests.java @@ -108,7 +108,7 @@ public void testToString() { .setErrorStrategy(DROP) .build(); String expected = - "IngestionSource{type='type',pointer_init_reset='PointerInitReset{type='RESET_BY_OFFSET', value=1000}',error_strategy='DROP', params={key=value}, maxPollSize=1000, pollTimeout=1000, numProcessorThreads=1, blockingQueueSize=100, allActiveIngestion=false, pointerBasedLagUpdateInterval=10s, mapperType='DEFAULT', mapperSettings={}}"; + "IngestionSource{type='type',pointer_init_reset='PointerInitReset{type='RESET_BY_OFFSET', value=1000}',error_strategy='DROP', params={key=value}, maxPollSize=1000, pollTimeout=1000, numProcessorThreads=1, blockingQueueSize=100, allActiveIngestion=false, pointerBasedLagUpdateInterval=10s, mapperType='DEFAULT', mapperSettings={}, warmupConfig=WarmupConfig[timeout=-1, lagThreshold=100]}"; assertEquals(expected, source.toString()); } @@ -177,4 +177,101 @@ public void testMapperSettings() { .build(); assertTrue(sourceNoMapperSettings.getMapperSettings().isEmpty()); } + + public void testWarmupConfigurationConstructorAndGetters() { + Map params = new HashMap<>(); + params.put("key", "value"); + + // Test with warmup configuration (timeout >= 0 means enabled) + IngestionSource source = new IngestionSource.Builder("type").setParams(params) + .setPointerInitReset(pointerInitReset) + .setErrorStrategy(DROP) + .setWarmupTimeout(TimeValue.timeValueMinutes(10)) + .setWarmupLagThreshold(100) + .build(); + + assertTrue("Warmup should be enabled", source.getWarmupConfig().isEnabled()); + assertEquals(TimeValue.timeValueMinutes(10), source.getWarmupConfig().timeout()); + assertEquals(100, source.getWarmupConfig().lagThreshold()); + } + + public void testWarmupConfigurationDefaults() { + // Test default warmup values (timeout=-1 means disabled) + IngestionSource source = new IngestionSource.Builder("type").build(); + + assertFalse("Warmup should be disabled by default", source.getWarmupConfig().isEnabled()); + assertEquals(TimeValue.timeValueMillis(-1), source.getWarmupConfig().timeout()); + assertEquals(100, source.getWarmupConfig().lagThreshold()); + } + + public void testWarmupConfigurationEquality() { + Map params = new HashMap<>(); + params.put("key", "value"); + + IngestionSource source1 = new IngestionSource.Builder("type").setParams(params) + .setWarmupTimeout(TimeValue.timeValueMinutes(10)) + .setWarmupLagThreshold(100) + .build(); + + IngestionSource source2 = new IngestionSource.Builder("type").setParams(params) + .setWarmupTimeout(TimeValue.timeValueMinutes(10)) + .setWarmupLagThreshold(100) + .build(); + + assertEquals(source1, source2); + assertEquals(source1.hashCode(), source2.hashCode()); + + // Test inequality with different warmup settings (disabled vs enabled) + IngestionSource source3 = new IngestionSource.Builder("type").setParams(params) + .setWarmupTimeout(TimeValue.timeValueMillis(-1)) + .setWarmupLagThreshold(100) + .build(); + + assertNotEquals(source1, source3); + } + + public void testWarmupConfigurationCopiedByBuilder() { + IngestionSource original = new IngestionSource.Builder("type").setWarmupTimeout(TimeValue.timeValueMinutes(10)) + .setWarmupLagThreshold(500) + .build(); + + // Create a copy using the copy constructor + IngestionSource copy = new IngestionSource.Builder(original).build(); + + assertEquals(original.getWarmupConfig().timeout(), copy.getWarmupConfig().timeout()); + assertEquals(original.getWarmupConfig().lagThreshold(), copy.getWarmupConfig().lagThreshold()); + } + + public void testWarmupConfigClass() { + IngestionSource.WarmupConfig config1 = new IngestionSource.WarmupConfig(TimeValue.timeValueMinutes(10), 100); + + assertTrue(config1.isEnabled()); + assertEquals(TimeValue.timeValueMinutes(10), config1.timeout()); + assertEquals(100, config1.lagThreshold()); + + // Test equals and hashCode + IngestionSource.WarmupConfig config2 = new IngestionSource.WarmupConfig(TimeValue.timeValueMinutes(10), 100); + assertEquals(config1, config2); + assertEquals(config1.hashCode(), config2.hashCode()); + + // Test inequality + IngestionSource.WarmupConfig config3 = new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(-1), 100); + assertNotEquals(config1, config3); + assertFalse(config3.isEnabled()); + + // Test toString + String expectedToString = "WarmupConfig[timeout=10m, lagThreshold=100]"; + assertEquals(expectedToString, config1.toString()); + } + + public void testSetWarmupConfig() { + IngestionSource.WarmupConfig warmupConfig = new IngestionSource.WarmupConfig(TimeValue.timeValueMinutes(15), 200); + + IngestionSource source = new IngestionSource.Builder("type").setWarmupConfig(warmupConfig).build(); + + assertEquals(warmupConfig, source.getWarmupConfig()); + assertTrue(source.getWarmupConfig().isEnabled()); + assertEquals(TimeValue.timeValueMinutes(15), source.getWarmupConfig().timeout()); + assertEquals(200, source.getWarmupConfig().lagThreshold()); + } } diff --git a/server/src/test/java/org/opensearch/indices/pollingingest/DefaultStreamPollerTests.java b/server/src/test/java/org/opensearch/indices/pollingingest/DefaultStreamPollerTests.java index ff36de1e525d3..7979c2c86dec1 100644 --- a/server/src/test/java/org/opensearch/indices/pollingingest/DefaultStreamPollerTests.java +++ b/server/src/test/java/org/opensearch/indices/pollingingest/DefaultStreamPollerTests.java @@ -16,6 +16,7 @@ import org.opensearch.cluster.block.ClusterBlocks; import org.opensearch.cluster.metadata.IngestionSource; import org.opensearch.common.settings.Settings; +import org.opensearch.common.unit.TimeValue; import org.opensearch.core.rest.RestStatus; import org.opensearch.index.IndexSettings; import org.opensearch.index.IngestionConsumerFactory; @@ -95,7 +96,8 @@ public void setUp() throws Exception { 1000, 10000, indexSettings, - new DefaultIngestionMessageMapper() + new DefaultIngestionMessageMapper(), + new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(-1), 0) ); partitionedBlockingQueueContainer.startProcessorThreads(); } @@ -169,7 +171,8 @@ public void testResetStateEarliest() throws InterruptedException { 1000, 10000, indexSettings, - new DefaultIngestionMessageMapper() + new DefaultIngestionMessageMapper(), + new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(-1), 0) ); CountDownLatch latch = new CountDownLatch(2); doAnswer(invocation -> { @@ -203,7 +206,8 @@ public void testResetStateLatest() throws InterruptedException { 1000, 10000, indexSettings, - new DefaultIngestionMessageMapper() + new DefaultIngestionMessageMapper(), + new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(-1), 0) ); // Set up latch to wait for 2 messages to be processed @@ -245,7 +249,8 @@ public void testResetStateRewindByOffset() throws InterruptedException { 1000, 10000, indexSettings, - new DefaultIngestionMessageMapper() + new DefaultIngestionMessageMapper(), + new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(-1), 0) ); CountDownLatch latch = new CountDownLatch(1); doAnswer(invocation -> { @@ -331,7 +336,8 @@ public void testDropErrorIngestionStrategy() throws TimeoutException, Interrupte 1000, 10000, indexSettings, - new DefaultIngestionMessageMapper() + new DefaultIngestionMessageMapper(), + new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(-1), 0) ); poller.start(); Thread.sleep(sleepTime); @@ -394,7 +400,8 @@ public void testBlockErrorIngestionStrategy() throws TimeoutException, Interrupt 1000, 10000, indexSettings, - new DefaultIngestionMessageMapper() + new DefaultIngestionMessageMapper(), + new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(-1), 0) ); poller.start(); Thread.sleep(sleepTime); @@ -431,7 +438,8 @@ public void testProcessingErrorWithBlockErrorIngestionStrategy() throws TimeoutE 1000, 10000, indexSettings, - new DefaultIngestionMessageMapper() + new DefaultIngestionMessageMapper(), + new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(-1), 0) ); poller.start(); Thread.sleep(sleepTime); @@ -504,7 +512,8 @@ public void testPersistedBatchStartPointer() throws TimeoutException, Interrupte 1000, 10000, indexSettings, - new DefaultIngestionMessageMapper() + new DefaultIngestionMessageMapper(), + new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(-1), 0) ); poller.start(); Thread.sleep(sleepTime); @@ -574,7 +583,8 @@ public void testConsumerInitializationRetry() throws Exception { 1000, 10000, indexSettings, - new DefaultIngestionMessageMapper() + new DefaultIngestionMessageMapper(), + new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(-1), 0) ); poller.start(); @@ -617,7 +627,8 @@ public void testConsumerReinitializationAfterProcessingMessages() throws Excepti 1000, 10000, indexSettings, - new DefaultIngestionMessageMapper() + new DefaultIngestionMessageMapper(), + new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(-1), 0) ); // Start and wait for 2 messages to be processed @@ -673,7 +684,8 @@ public void testConsumerReinitializationWithNoInitialMessages() throws Exception 1000, 10000, indexSettings, - new DefaultIngestionMessageMapper() + new DefaultIngestionMessageMapper(), + new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(-1), 0) ); // Start poller @@ -717,10 +729,251 @@ public void testGetBatchStartPointerWithNullInitialPointer() { 1000, 10000, indexSettings, - new DefaultIngestionMessageMapper() + new DefaultIngestionMessageMapper(), + new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(-1), 0) ); // When all queues return null and initialBatchStartPointer is null, getBatchStartPointer should return null assertNull(poller.getBatchStartPointer()); } + + // ==================== Warmup Tests ==================== + + public void testWarmupDisabledBehavior() { + // When warmup is disabled, isWarmupComplete should return true immediately + // The default poller in setUp has warmup disabled (timeout=-1) + assertTrue(poller.isWarmupComplete()); // Warmup disabled means it's considered complete + + // Create another poller with warmup explicitly disabled + DefaultStreamPoller warmupDisabledPoller = new DefaultStreamPoller( + new FakeIngestionSource.FakeIngestionShardPointer(0), + fakeConsumerFactory, + "", + 0, + partitionedBlockingQueueContainer, + StreamPoller.ResetState.NONE, + "", + errorStrategy, + StreamPoller.State.NONE, + 1000, + 1000, + 10000, + indexSettings, + new DefaultIngestionMessageMapper(), + new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(-1), 0) + ); + + // Warmup should be considered complete when disabled + assertTrue(warmupDisabledPoller.isWarmupComplete()); + warmupDisabledPoller.close(); + } + + public void testWarmupCompletesImmediatelyWhenLagIsZero() throws InterruptedException { + // Create a poller with warmup enabled and lag threshold of 0 + // FakeIngestionSource returns 0 for lag by default + DefaultStreamPoller warmupPoller = new DefaultStreamPoller( + new FakeIngestionSource.FakeIngestionShardPointer(0), + fakeConsumerFactory, + "", + 0, + partitionedBlockingQueueContainer, + StreamPoller.ResetState.NONE, + "", + errorStrategy, + StreamPoller.State.NONE, + 1000, + 1000, + 10000, + indexSettings, + new DefaultIngestionMessageMapper(), + new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(300000), 0) + ); + + // Initially warmup is not complete + assertFalse(warmupPoller.isWarmupComplete()); + + // Start the poller - it should complete warmup quickly since lag is 0 + warmupPoller.start(); + + // Wait for warmup with timeout + boolean completed = warmupPoller.awaitWarmupComplete(5000); + assertTrue("Warmup should complete when lag is at threshold", completed); + assertTrue(warmupPoller.isWarmupComplete()); + + warmupPoller.close(); + } + + public void testWarmupTimesOutAndProceeds() throws InterruptedException, TimeoutException { + // Create a mock consumer factory that always reports high lag + IngestionConsumerFactory mockFactory = mock(IngestionConsumerFactory.class); + IngestionShardConsumer mockConsumer = mock(IngestionShardConsumer.class); + when(mockFactory.createShardConsumer(anyString(), anyInt())).thenReturn(mockConsumer); + when(mockConsumer.getPointerBasedLag(any())).thenReturn(1000L); // High lag + when(mockConsumer.readNext(anyLong(), anyInt())).thenReturn(Collections.emptyList()); + + DefaultStreamPoller warmupPoller = new DefaultStreamPoller( + new FakeIngestionSource.FakeIngestionShardPointer(0), + mockFactory, + "", + 0, + partitionedBlockingQueueContainer, + StreamPoller.ResetState.NONE, + "", + errorStrategy, + StreamPoller.State.NONE, + 1000, + 1000, + 10000, + indexSettings, + new DefaultIngestionMessageMapper(), + new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(500), 0) + ); + + warmupPoller.start(); + + // Wait for timeout - should eventually complete due to timeout + boolean completed = warmupPoller.awaitWarmupComplete(2000); + // Even if lag never reaches threshold, warmup completes on timeout + assertTrue("Warmup should complete after timeout", warmupPoller.isWarmupComplete()); + + warmupPoller.close(); + } + + public void testWarmupStateTransitions() throws InterruptedException { + DefaultStreamPoller warmupPoller = new DefaultStreamPoller( + new FakeIngestionSource.FakeIngestionShardPointer(0), + fakeConsumerFactory, + "", + 0, + partitionedBlockingQueueContainer, + StreamPoller.ResetState.NONE, + "", + errorStrategy, + StreamPoller.State.NONE, + 1000, + 1000, + 10000, + indexSettings, + new DefaultIngestionMessageMapper(), + new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(300000), 0) + ); + + // Initial state should be NONE + assertEquals(StreamPoller.State.NONE, warmupPoller.getState()); + + warmupPoller.start(); + + // Give it time to start + Thread.sleep(200); + + // State should eventually transition from WARMING_UP to POLLING/PROCESSING + // (since lag should be 0 from fake source) + warmupPoller.awaitWarmupComplete(5000); + assertTrue(warmupPoller.isWarmupComplete()); + + warmupPoller.close(); + } + + public void testWarmupCompletesViaTimeoutWhenPointerLagNegative() throws InterruptedException, TimeoutException { + // Create a mock consumer factory that returns negative pointer-based lag (unsupported) + IngestionConsumerFactory mockFactory = mock(IngestionConsumerFactory.class); + IngestionShardConsumer mockConsumer = mock(IngestionShardConsumer.class); + when(mockFactory.createShardConsumer(anyString(), anyInt())).thenReturn(mockConsumer); + when(mockConsumer.getPointerBasedLag(any())).thenReturn(-1L); // Negative means unsupported + when(mockConsumer.readNext(anyLong(), anyInt())).thenReturn(Collections.emptyList()); + + DefaultStreamPoller warmupPoller = new DefaultStreamPoller( + new FakeIngestionSource.FakeIngestionShardPointer(0), + mockFactory, + "", + 0, + partitionedBlockingQueueContainer, + StreamPoller.ResetState.NONE, + "", + errorStrategy, + StreamPoller.State.NONE, + 1000, + 1000, + 10000, + indexSettings, + new DefaultIngestionMessageMapper(), + new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(500), 0) + ); + + warmupPoller.start(); + + // Should eventually complete via timeout since negative lag won't satisfy threshold + boolean completed = warmupPoller.awaitWarmupComplete(2000); + assertTrue("Warmup should complete after timeout when pointer lag is unsupported", warmupPoller.isWarmupComplete()); + + warmupPoller.close(); + } + + public void testWarmupWithNonZeroLagThreshold() throws InterruptedException, TimeoutException { + // Create a mock consumer factory that reports lag slightly above threshold then drops to threshold + IngestionConsumerFactory mockFactory = mock(IngestionConsumerFactory.class); + IngestionShardConsumer mockConsumer = mock(IngestionShardConsumer.class); + when(mockFactory.createShardConsumer(anyString(), anyInt())).thenReturn(mockConsumer); + // Return lag of 50 which is below threshold of 100 + when(mockConsumer.getPointerBasedLag(any())).thenReturn(50L); + when(mockConsumer.readNext(anyLong(), anyInt())).thenReturn(Collections.emptyList()); + + DefaultStreamPoller warmupPoller = new DefaultStreamPoller( + new FakeIngestionSource.FakeIngestionShardPointer(0), + mockFactory, + "", + 0, + partitionedBlockingQueueContainer, + StreamPoller.ResetState.NONE, + "", + errorStrategy, + StreamPoller.State.NONE, + 1000, + 1000, + 10000, + indexSettings, + new DefaultIngestionMessageMapper(), + new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(30000), 100) + ); + + warmupPoller.start(); + + // Should complete quickly since lag (50) is below threshold (100) + boolean completed = warmupPoller.awaitWarmupComplete(5000); + assertTrue("Warmup should complete when lag is below threshold", completed); + assertTrue(warmupPoller.isWarmupComplete()); + + warmupPoller.close(); + } + + public void testWarmupAwaitReturnsImmediatelyWhenAlreadyComplete() throws InterruptedException { + // Create a poller with warmup disabled (which means warmup is considered complete) + DefaultStreamPoller warmupPoller = new DefaultStreamPoller( + new FakeIngestionSource.FakeIngestionShardPointer(0), + fakeConsumerFactory, + "", + 0, + partitionedBlockingQueueContainer, + StreamPoller.ResetState.NONE, + "", + errorStrategy, + StreamPoller.State.NONE, + 1000, + 1000, + 10000, + indexSettings, + new DefaultIngestionMessageMapper(), + new IngestionSource.WarmupConfig(TimeValue.timeValueMillis(-1), 0) + ); + + // Should return immediately without blocking since warmup is disabled + long startTime = System.currentTimeMillis(); + boolean completed = warmupPoller.awaitWarmupComplete(60000); + long elapsed = System.currentTimeMillis() - startTime; + + assertTrue("Should return true when warmup is disabled", completed); + assertTrue("Should return immediately (< 1000ms)", elapsed < 1000); + + warmupPoller.close(); + } }