Move consumer initialization to the poller to prevent engine failure

Signed-off-by: Varun Bharadwaj <[email protected]>

Author: varunbharadwaj

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

@@ -37,6 +37,7 @@
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 
+import static org.hamcrest.Matchers.greaterThan;
 import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked;
 import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertHitCount;
 import static org.hamcrest.Matchers.is;
@@ -749,6 +750,47 @@ public void testIndexRelocation() throws Exception {
         waitForSearchableDocs(4, List.of(nodeB));
     }
 
+    public void testKafkaConnectionLost() throws Exception {
+        // Step 1: Create 2 nodes
+        internalCluster().startClusterManagerOnlyNode();
+        final String nodeA = internalCluster().startDataOnlyNode();
+        final String nodeB = internalCluster().startDataOnlyNode();
+
+        // Step 2: Create index
+        createIndexWithDefaultSettings(1, 0);
+        ensureGreen(indexName);
+        assertTrue(nodeA.equals(primaryNodeName(indexName)));
+
+        // Step 3: Write documents and verify
+        produceData("1", "name1", "24");
+        produceData("2", "name2", "20");
+        refresh(indexName);
+        waitForSearchableDocs(2, List.of(nodeA));
+        flush(indexName);
+
+        // Step 4: Stop kafka and relocate index to nodeB
+        kafka.stop();
+        assertAcked(
+            client().admin()
+                .indices()
+                .prepareUpdateSettings(indexName)
+                .setSettings(Settings.builder().put("index.routing.allocation.require._name", nodeB))
+                .get()
+        );
+
+        // Step 5: Wait for relocation to complete
+        waitForState(() -> nodeB.equals(primaryNodeName(indexName)));
+
+        // Step 6: Ensure index is searchable on nodeB even though kafka is down
+        ensureGreen(indexName);
+        waitForSearchableDocs(2, List.of(nodeB));
+        waitForState(() -> {
+            PollingIngestStats stats = client().admin().indices().prepareStats(indexName).get().getIndex(indexName).getShards()[0]
+                .getPollingIngestStats();
+            return stats.getConsumerStats().totalConsumerErrorCount() > 0;
+        });
+    }
+
     private void verifyRemoteStoreEnabled(String node) {
         GetSettingsResponse settingsResponse = client(node).admin().indices().prepareGetSettings(indexName).get();
         String remoteStoreEnabled = settingsResponse.getIndexToSettings().get(indexName).get("index.remote_store.enabled");

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

@@ -133,6 +133,24 @@ public void testPauseAndResumeAPIs() throws Exception {
         });
     }
 
+    // This test validates shard initialization does not fail due to kafka connection errors.
+    public void testShardInitializationUsingUnknownTopic() throws Exception {
+        createIndexWithMappingSource(
+            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", "unknownTopic")
+                .put("ingestion_source.param.bootstrap_servers", kafka.getBootstrapServers())
+                .put("index.replication.type", "SEGMENT")
+                .build(),
+            mappings
+        );
+        ensureGreen(indexName);
+    }
+
     private void setupKafka() {
         kafka = new KafkaContainer(DockerImageName.parse("confluentinc/cp-kafka:6.2.1"))
             // disable topic auto creation

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

@@ -98,11 +98,7 @@ private void initializeStreamPoller(
             + engineConfig.getIndexSettings().getIndex().getName()
             + "-"
             + engineConfig.getShardId().getId();
-        IngestionShardConsumer ingestionShardConsumer = this.ingestionConsumerFactory.createShardConsumer(
-            clientId,
-            engineConfig.getShardId().getId()
-        );
-        logger.info("created ingestion consumer for shard [{}]", engineConfig.getShardId());
+
         Map<String, String> commitData = commitDataAsMap(indexWriter);
         StreamPoller.ResetState resetState = ingestionSource.getPointerInitReset().getType();
         String resetValue = ingestionSource.getPointerInitReset().getValue();
@@ -146,20 +142,25 @@ private void initializeStreamPoller(
         StreamPoller.State initialPollerState = indexMetadata.getIngestionStatus().isPaused()
             ? StreamPoller.State.PAUSED
             : StreamPoller.State.NONE;
-        streamPoller = new DefaultStreamPoller(
+
+        // initialize the stream poller
+        DefaultStreamPoller.Builder streamPollerBuilder = new DefaultStreamPoller.Builder(
             startPointer,
             persistedPointers,
-            ingestionShardConsumer,
-            this,
-            resetState,
-            resetValue,
-            ingestionErrorStrategy,
-            initialPollerState,
-            ingestionSource.getMaxPollSize(),
-            ingestionSource.getPollTimeout(),
-            ingestionSource.getNumProcessorThreads(),
-            ingestionSource.getBlockingQueueSize()
-        );
+            ingestionConsumerFactory,
+            clientId,
+            engineConfig.getShardId().getId(),
+            this);
+        streamPoller = streamPollerBuilder
+            .resetState(resetState)
+            .resetValue(resetValue)
+            .errorStrategy(ingestionErrorStrategy)
+            .initialState(initialPollerState)
+            .maxPollSize(ingestionSource.getMaxPollSize())
+            .pollTimeout(ingestionSource.getPollTimeout())
+            .numProcessorThreads(ingestionSource.getNumProcessorThreads())
+            .blockingQueueSize(ingestionSource.getBlockingQueueSize())
+            .build();
         registerStreamPollerListener();
 
         // start the polling loop
@@ -575,6 +576,10 @@ private void resetStreamPoller(StreamPoller.ResetState resetState, String resetV
             throw new IllegalStateException("Cannot reset consumer when poller is not paused");
         }
 
+        if (streamPoller.getConsumer() == null) {
+            throw new OpenSearchException("Consumer is not yet initialized");
+        }
+
         try {
             // refresh is needed for persisted pointers to be visible
             refresh("reset poller", SearcherScope.INTERNAL, true);