refactor poller and handle race conditions on consumer reinitializations

varunbharadwaj · varunbharadwaj · commit 9d56ce66df9b · 2025-11-19T00:04:19.000-08:00
Signed-off-by: Varun Bharadwaj &lt;varunbharadwaj1995@gmail.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -30,7 +30,6 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Implement GRPC FunctionScoreQuery ([#19888](https://github.com/opensearch-project/OpenSearch/pull/19888))
 - Implement error_trace parameter for bulk requests ([#19985](https://github.com/opensearch-project/OpenSearch/pull/19985))
 - Allow the truncate filter in normalizers ([#19778](https://github.com/opensearch-project/OpenSearch/issues/19778))
-- Support pull-based ingestion message mappers and raw payload support ([#19765](https://github.com/opensearch-project/OpenSearch/pull/19765)]
 - Support pull-based ingestion message mappers and raw payload support ([#19765](https://github.com/opensearch-project/OpenSearch/pull/19765))
 - Support dynamic consumer configuration update in pull-based ingestion ([#19963](https://github.com/opensearch-project/OpenSearch/pull/19963))
 
diff --git a/server/src/main/java/org/opensearch/index/engine/IngestionEngine.java b/server/src/main/java/org/opensearch/index/engine/IngestionEngine.java
@@ -534,25 +534,16 @@ private void updateIngestionSourceParams(Map<String, Object> updatedParams) {
             return;
         }
 
-        try {
-            logger.info("Ingestion source params updated, reinitializing consumer");
-
-            // Get current ingestion source with updated params from index metadata
-            IndexMetadata indexMetadata = engineConfig.getIndexSettings().getIndexMetadata();
-            assert indexMetadata != null;
-            IngestionSource updatedIngestionSource = Objects.requireNonNull(indexMetadata.getIngestionSource());
+        logger.info("Ingestion source params updated, reinitializing consumer");
 
-            // Initialize the factory with updated params
-            ingestionConsumerFactory.initialize(updatedIngestionSource);
-
-            // Request consumer reinitialization in the poller
-            streamPoller.requestConsumerReinitialization();
+        // Get current ingestion source with updated params from index metadata
+        IndexMetadata indexMetadata = engineConfig.getIndexSettings().getIndexMetadata();
+        assert indexMetadata != null;
+        IngestionSource updatedIngestionSource = Objects.requireNonNull(indexMetadata.getIngestionSource());
 
-            logger.info("Successfully processed ingestion source params update");
-        } catch (Exception e) {
-            logger.error("Failed to update ingestion source params", e);
-            throw new OpenSearchException("Failed to update ingestion source params", e);
-        }
+        // Request consumer reinitialization in the poller
+        streamPoller.requestConsumerReinitialization(updatedIngestionSource);
+        logger.info("Successfully processed ingestion source params update");
     }
 
     /**
diff --git a/server/src/main/java/org/opensearch/indices/pollingingest/DefaultStreamPoller.java b/server/src/main/java/org/opensearch/indices/pollingingest/DefaultStreamPoller.java
@@ -13,6 +13,7 @@
 import org.opensearch.cluster.ClusterChangedEvent;
 import org.opensearch.cluster.ClusterState;
 import org.opensearch.cluster.block.ClusterBlockLevel;
+import org.opensearch.cluster.metadata.IngestionSource;
 import org.opensearch.common.Nullable;
 import org.opensearch.common.metrics.CounterMetric;
 import org.opensearch.index.IndexSettings;
@@ -66,7 +67,6 @@ public class DefaultStreamPoller implements StreamPoller {
 
     // start of the batch, inclusive
     private IngestionShardPointer initialBatchStartPointer;
-    private boolean includeBatchStartPointer = false;
 
     private ResetState resetState;
     private final String resetValue;
@@ -86,6 +86,9 @@ public class DefaultStreamPoller implements StreamPoller {
 
     private PartitionedBlockingQueueContainer blockingQueueContainer;
 
+    // Force the consumer to start reading from this pointer. This is used in case of failures, or during reinitialization.
+    private IngestionShardPointer forcedShardPointer = null;
+
     private DefaultStreamPoller(
         IngestionShardPointer startPointer,
         IngestionConsumerFactory consumerFactory,
@@ -173,8 +176,6 @@ public void start() {
         }
 
         started = true;
-        // when we start, we need to include the batch start pointer in the read for the first read
-        includeBatchStartPointer = true;
         consumerThread.submit(this::startPoll);
         blockingQueueContainer.startProcessorThreads();
     }
@@ -191,33 +192,20 @@ protected void startPoll() {
         }
         logger.info("Starting poller for shard {}", shardId);
 
-        // Force the consumer to start reading from this pointer. This is used in case of failures, or during reinitialization.
-        IngestionShardPointer forcedShardPointer = null;
         while (true) {
             try {
                 if (closed) {
                     state = State.CLOSED;
+                    closeConsumer();
                     break;
                 }
 
-                // Initialize consumer if not already initialized or if reinitialization is requested
+                // Initialize/reinitialization consumer
                 if (this.consumer == null || reinitializeConsumer) {
-                    handleConsumerInitialization(CONSUMER_INIT_RETRY_INTERVAL_MS);
-
-                    // If consumer reinitialization is requested, update the new consumer's start offset to the latest batch start pointer.
-                    // First clear the blocking queue partitions, and then retrieve the latest batch start pointer. This
-                    // will ensure we resume from the earliest offset possible without too much duplicate processing.
-                    if (reinitializeConsumer && includeBatchStartPointer == false) {
-                        blockingQueueContainer.clearAllQueues();
-                        forcedShardPointer = getBatchStartPointer();
-                    }
-                    reinitializeConsumer = false;
+                    handleConsumerInitialization();
                     continue;
                 }
 
-                // reset the consumer offset
-                handleResetState();
-
                 // Update lag periodically. Lag is updated even if the poller is paused.
                 updatePointerBasedLagIfNeeded();
 
@@ -234,10 +222,8 @@ protected void startPoll() {
                 state = State.POLLING;
                 List<IngestionShardConsumer.ReadResult<? extends IngestionShardPointer, ? extends Message>> results;
 
-                if (includeBatchStartPointer) {
-                    results = consumer.readNext(initialBatchStartPointer, true, maxPollSize, pollTimeout);
-                    includeBatchStartPointer = false;
-                } else if (forcedShardPointer != null) {
+                // Force the consumer to start from forcedShardPointer if available
+                if (forcedShardPointer != null) {
                     results = consumer.readNext(forcedShardPointer, true, maxPollSize, pollTimeout);
                     forcedShardPointer = null;
                 } else {
@@ -335,7 +321,6 @@ public void close() {
         closed = true;
         if (!started) {
             logger.info("consumer thread not started");
-            closeConsumer();
             return;
         }
         long startTime = System.currentTimeMillis(); // Record the start time
@@ -354,7 +339,6 @@ public void close() {
         }
         consumerThread.shutdown();
         blockingQueueContainer.close();
-        closeConsumer();
         logger.info("closed the poller of shard {}", shardId);
     }
 
@@ -478,15 +462,18 @@ public IngestionShardConsumer getConsumer() {
 
     /**
      * Mark the poller's consumer for reinitialization. A new consumer will be initialized and start consuming from the
-     * latest batchStartPointer.
+     * latest batchStartPointer. This method also reinitializes the consumer factory with the updated ingestion source.
+     * @param updatedIngestionSource the updated ingestion source with new configuration parameters
      */
     @Override
-    public void requestConsumerReinitialization() {
+    public synchronized void requestConsumerReinitialization(IngestionSource updatedIngestionSource) {
         if (closed) {
             logger.warn("Cannot reinitialize consumer for closed poller of shard {}", shardId);
             return;
         }
 
+        // Reinitialize the consumer factory with updated configuration
+        consumerFactory.initialize(updatedIngestionSource);
         logger.info("Configuration parameters updated for index {} shard {}, requesting consumer reinitialization", indexName, shardId);
         reinitializeConsumer = true;
     }
@@ -508,63 +495,75 @@ public void clusterChanged(ClusterChangedEvent event) {
     }
 
     /**
-     * Handles the reset state logic, updating the initialBatchStartPointer based on the reset state.
+     * Handles the reset state logic.
+     * Returns the resulting IngestionShardPointer for reset or null if no reset required.
      */
-    private void handleResetState() {
-        if (resetState != ResetState.NONE) {
+    private IngestionShardPointer getResetShardPointer() {
+        IngestionShardPointer resetPointer = null;
+        if (resetState != ResetState.NONE && consumer != null) {
             switch (resetState) {
                 case EARLIEST:
-                    initialBatchStartPointer = consumer.earliestPointer();
-                    logger.info("Resetting pointer by seeking to earliest pointer {}", initialBatchStartPointer.asString());
+                    resetPointer = consumer.earliestPointer();
+                    logger.info("Resetting pointer by seeking to earliest pointer {}", resetPointer.asString());
                     break;
                 case LATEST:
-                    initialBatchStartPointer = consumer.latestPointer();
-                    logger.info("Resetting pointer by seeking to latest pointer {}", initialBatchStartPointer.asString());
+                    resetPointer = consumer.latestPointer();
+                    logger.info("Resetting pointer by seeking to latest pointer {}", resetPointer.asString());
                     break;
                 case RESET_BY_OFFSET:
-                    initialBatchStartPointer = consumer.pointerFromOffset(resetValue);
-                    logger.info("Resetting pointer by seeking to pointer {}", initialBatchStartPointer.asString());
+                    resetPointer = consumer.pointerFromOffset(resetValue);
+                    logger.info("Resetting pointer by seeking to pointer {}", resetPointer.asString());
                     break;
                 case RESET_BY_TIMESTAMP:
-                    initialBatchStartPointer = consumer.pointerFromTimestampMillis(Long.parseLong(resetValue));
+                    resetPointer = consumer.pointerFromTimestampMillis(Long.parseLong(resetValue));
                     logger.info(
                         "Resetting pointer by seeking to timestamp {}, corresponding pointer {}",
                         resetValue,
-                        initialBatchStartPointer.asString()
+                        resetPointer.asString()
                     );
                     break;
             }
             resetState = ResetState.NONE;
         }
+
+        return resetPointer;
     }
 
     /**
-     * Handles consumer initialization and reinitialization logic.
-     * If reinitializing, closes the existing consumer before creating a new one.
-     *
-     * @param sleepDurationOnError duration to sleep if initialization fails
+     * Handles consumer initialization and reinitialization logic. Closes existing consumer if available and clears the
+     * blocking queues before initializing a new consumer. Also forces the consumer to start reading from the initial
+     * batchStartPointer if first time initialization, or from the latest available batchStartPointer on reinitialization.
      */
-    private void handleConsumerInitialization(int sleepDurationOnError) {
-        if (reinitializeConsumer) {
-            logger.info("Reinitializing consumer for index {} shard {} due to configuration changes", indexName, shardId);
-            closeConsumer();
+    private void handleConsumerInitialization() {
+        closeConsumer();
+        blockingQueueContainer.clearAllQueues();
+        initializeConsumer();
+
+        // Handle consumer offset reset the first time an index is created. The reset offset takes precedence if available.
+        IngestionShardPointer resetShardPointer = getResetShardPointer();
+        if (resetShardPointer != null) {
+            initialBatchStartPointer = resetShardPointer;
         }
-        initializeConsumer(sleepDurationOnError);
+
+        // Force the consumer to start from the batchStartPointer. This will be the initialBatchStartPointer for first
+        // time initialization, or the latest batchStartPointer based on processed messages.
+        forcedShardPointer = getBatchStartPointer();
     }
 
     /**
      * Initializes the consumer using the provided consumerFactory. If an error is encountered during initialization,
      * the poller thread sleeps for the provided duration before retrying/proceeding with the polling loop.
      */
-    private void initializeConsumer(int sleepDurationOnError) {
+    private synchronized void initializeConsumer() {
         try {
+            reinitializeConsumer = false;
             this.consumer = consumerFactory.createShardConsumer(consumerClientId, shardId);
             logger.info("Successfully initialized consumer for shard {}", shardId);
         } catch (Exception e) {
             logger.warn("Failed to create consumer for shard {}: {}", shardId, e.getMessage());
             totalConsumerErrorCount.inc();
             try {
-                Thread.sleep(sleepDurationOnError);
+                Thread.sleep(CONSUMER_INIT_RETRY_INTERVAL_MS);
             } catch (InterruptedException ie) {
                 Thread.currentThread().interrupt();
             }
diff --git a/server/src/main/java/org/opensearch/indices/pollingingest/StreamPoller.java b/server/src/main/java/org/opensearch/indices/pollingingest/StreamPoller.java
@@ -9,6 +9,7 @@
 package org.opensearch.indices.pollingingest;
 
 import org.opensearch.cluster.ClusterStateListener;
+import org.opensearch.cluster.metadata.IngestionSource;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.index.IngestionShardConsumer;
 import org.opensearch.index.IngestionShardPointer;
@@ -76,10 +77,11 @@ public interface StreamPoller extends Closeable, ClusterStateListener {
     IngestionShardConsumer getConsumer();
 
     /**
-     * Requests the poller to reinitialize the consumer.
+     * Requests the poller to reinitialize the consumer with updated ingestion source configuration.
      * This is called when ingestion source params are dynamically updated.
+     * @param updatedIngestionSource the updated ingestion source with new configuration parameters
      */
-    void requestConsumerReinitialization();
+    void requestConsumerReinitialization(IngestionSource updatedIngestionSource);
 
     /**
      * A state to indicate the current state of the poller
diff --git a/server/src/test/java/org/opensearch/indices/pollingingest/DefaultStreamPollerTests.java b/server/src/test/java/org/opensearch/indices/pollingingest/DefaultStreamPollerTests.java
@@ -14,6 +14,7 @@
 import org.opensearch.cluster.block.ClusterBlock;
 import org.opensearch.cluster.block.ClusterBlockLevel;
 import org.opensearch.cluster.block.ClusterBlocks;
+import org.opensearch.cluster.metadata.IngestionSource;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.index.IndexSettings;
@@ -29,6 +30,7 @@
 
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.EnumSet;
 import java.util.List;
@@ -183,6 +185,10 @@ public void testResetStateEarliest() throws InterruptedException {
     }
 
     public void testResetStateLatest() throws InterruptedException {
+        // Clear messages first and add them after poller starts
+        // This ensures latestPointer() returns the correct value at initialization time
+        messages.clear();
+
         poller = new DefaultStreamPoller(
             new FakeIngestionSource.FakeIngestionShardPointer(0),
             fakeConsumerFactory,
@@ -200,12 +206,28 @@ public void testResetStateLatest() throws InterruptedException {
             new DefaultIngestionMessageMapper()
         );
 
+        // Set up latch to wait for 2 messages to be processed
+        CountDownLatch latch = new CountDownLatch(2);
+        doAnswer(invocation -> {
+            latch.countDown();
+            return null;
+        }).when(processor).process(any(), any());
+
         poller.start();
         waitUntil(() -> poller.getState() == DefaultStreamPoller.State.POLLING, awaitTime, TimeUnit.MILLISECONDS);
-        // no messages processed
-        verify(processor, never()).process(any(), any());
-        // reset to the latest
-        assertEquals(new FakeIngestionSource.FakeIngestionShardPointer(2), poller.getBatchStartPointer());
+
+        // Verify batch start pointer was set to latest (which is 0 since messages list was empty)
+        assertEquals(new FakeIngestionSource.FakeIngestionShardPointer(0), poller.getBatchStartPointer());
+
+        // Now add messages after poller has started with LATEST reset
+        messages.add("{\"_id\":\"1\",\"_source\":{\"name\":\"bob\", \"age\": 24}}".getBytes(StandardCharsets.UTF_8));
+        messages.add("{\"_id\":\"2\",\"_source\":{\"name\":\"alice\", \"age\": 21}}".getBytes(StandardCharsets.UTF_8));
+
+        // Wait for messages to be processed
+        latch.await();
+
+        // Verify that the messages added after starting from latest are processed
+        verify(processor, times(2)).process(any(), any());
     }
 
     public void testResetStateRewindByOffset() throws InterruptedException {
@@ -610,7 +632,9 @@ public void testConsumerReinitializationAfterProcessingMessages() throws Excepti
             return null;
         }).when(processor).process(any(), any());
 
-        poller.requestConsumerReinitialization();
+        // Create a mock ingestion source for reinitialization
+        IngestionSource mockIngestionSource = new IngestionSource.Builder("test").build();
+        poller.requestConsumerReinitialization(mockIngestionSource);
 
         // Add a 3rd message
         messages.add("{\"_id\":\"3\",\"_source\":{\"name\":\"charlie\", \"age\": 30}}".getBytes(StandardCharsets.UTF_8));
@@ -653,7 +677,8 @@ public void testConsumerReinitializationWithNoInitialMessages() throws Exception
         verify(processor, never()).process(any(), any());
 
         // Request consumer reinitialization
-        poller.requestConsumerReinitialization();
+        IngestionSource mockIngestionSource = new IngestionSource.Builder("test").build();
+        poller.requestConsumerReinitialization(mockIngestionSource);
 
         // Add a message
         messages.add("{\"_id\":\"1\",\"_source\":{\"name\":\"bob\", \"age\": 24}}".getBytes(StandardCharsets.UTF_8));
@@ -670,4 +695,31 @@ public void testConsumerReinitializationWithNoInitialMessages() throws Exception
         // Verify 1 message was processed
         verify(processor, times(1)).process(any(), any());
     }
+
+    public void testGetBatchStartPointerWithNullInitialPointer() {
+        // Create a mock blocking queue container that returns null pointers
+        PartitionedBlockingQueueContainer mockContainer = mock(PartitionedBlockingQueueContainer.class);
+        when(mockContainer.getCurrentShardPointers()).thenReturn(Arrays.asList(null, null, null));
+
+        // Create poller with null initial batch start pointer
+        poller = new DefaultStreamPoller(
+            null,
+            fakeConsumerFactory,
+            "",
+            0,
+            mockContainer,
+            StreamPoller.ResetState.NONE,
+            "",
+            errorStrategy,
+            StreamPoller.State.NONE,
+            1000,
+            1000,
+            10000,
+            indexSettings,
+            new DefaultIngestionMessageMapper()
+        );
+
+        // When all queues return null and initialBatchStartPointer is null, getBatchStartPointer should return null
+        assertNull(poller.getBatchStartPointer());
+    }
 }
