update ingestion state on all-active shards on receiving cluster state update (opensearch-project#19407)

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

Author: varunbharadwaj

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

@@ -16,6 +16,7 @@
 import org.opensearch.action.admin.indices.stats.IndexStats;
 import org.opensearch.action.admin.indices.stats.ShardStats;
 import org.opensearch.action.admin.indices.streamingingestion.pause.PauseIngestionResponse;
+import org.opensearch.action.admin.indices.streamingingestion.resume.ResumeIngestionRequest;
 import org.opensearch.action.admin.indices.streamingingestion.resume.ResumeIngestionResponse;
 import org.opensearch.action.admin.indices.streamingingestion.state.GetIngestionStateResponse;
 import org.opensearch.action.search.SearchResponse;
@@ -254,8 +255,7 @@ public void testMultiThreadedWrites() throws Exception {
     }
 
     public void testAllActiveIngestion() throws Exception {
-        // Create pull-based index in default replication mode (docrep) and publish some messages
-
+        // Create all-active pull-based index
         internalCluster().startClusterManagerOnlyNode();
         final String nodeA = internalCluster().startDataOnlyNode();
         for (int i = 0; i < 10; i++) {
@@ -374,7 +374,7 @@ public void testAllActiveIngestion() throws Exception {
     }
 
     public void testReplicaPromotionOnAllActiveIngestion() throws Exception {
-        // Create pull-based index in default replication mode (docrep) and publish some messages
+        // Create all-active pull-based index
         internalCluster().startClusterManagerOnlyNode();
         final String nodeA = internalCluster().startDataOnlyNode();
         for (int i = 0; i < 10; i++) {
@@ -425,7 +425,7 @@ public void testReplicaPromotionOnAllActiveIngestion() throws Exception {
     }
 
     public void testSnapshotRestoreOnAllActiveIngestion() throws Exception {
-        // Create pull-based index in default replication mode (docrep) and publish some messages
+        // Create all-active pull-based index
         internalCluster().startClusterManagerOnlyNode();
         final String nodeA = internalCluster().startDataOnlyNode();
         final String nodeB = internalCluster().startDataOnlyNode();
@@ -505,6 +505,68 @@ public void testSnapshotRestoreOnAllActiveIngestion() throws Exception {
         assertThat(shardTypeToStats.get("replica").getConsumerStats().totalPollerMessageFailureCount(), is(0L));
     }
 
+    public void testResetPollerInAllActiveIngestion() throws Exception {
+        // Create all-active pull-based index
+        internalCluster().startClusterManagerOnlyNode();
+        final String nodeA = internalCluster().startDataOnlyNode();
+        final String nodeB = internalCluster().startDataOnlyNode();
+        for (int i = 0; i < 10; i++) {
+            produceData(Integer.toString(i), "name" + i, "30");
+        }
+
+        createIndex(
+            indexName,
+            Settings.builder()
+                .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
+                .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1)
+                .put("ingestion_source.type", "kafka")
+                .put("ingestion_source.param.topic", topicName)
+                .put("ingestion_source.param.bootstrap_servers", kafka.getBootstrapServers())
+                .put("ingestion_source.pointer.init.reset", "earliest")
+                .put("ingestion_source.all_active", true)
+                .build(),
+            "{\"properties\":{\"name\":{\"type\": \"text\"},\"age\":{\"type\": \"integer\"}}}}"
+        );
+
+        ensureGreen(indexName);
+        waitForSearchableDocs(10, List.of(nodeA, nodeB));
+
+        // pause ingestion
+        PauseIngestionResponse pauseResponse = pauseIngestion(indexName);
+        assertTrue(pauseResponse.isAcknowledged());
+        assertTrue(pauseResponse.isShardsAcknowledged());
+        waitForState(() -> {
+            GetIngestionStateResponse ingestionState = getIngestionState(indexName);
+            return ingestionState.getShardStates().length == 2
+                && ingestionState.getFailedShards() == 0
+                && Arrays.stream(ingestionState.getShardStates())
+                    .allMatch(state -> state.isPollerPaused() && state.getPollerState().equalsIgnoreCase("paused"));
+        });
+
+        // reset to offset=2 and resume ingestion
+        ResumeIngestionResponse resumeResponse = resumeIngestion(indexName, 0, ResumeIngestionRequest.ResetSettings.ResetMode.OFFSET, "2");
+        assertTrue(resumeResponse.isAcknowledged());
+        assertTrue(resumeResponse.isShardsAcknowledged());
+        waitForState(() -> {
+            GetIngestionStateResponse ingestionState = getIngestionState(indexName);
+            return ingestionState.getShardStates().length == 2
+                && Arrays.stream(ingestionState.getShardStates())
+                    .allMatch(
+                        state -> state.isPollerPaused() == false
+                            && (state.getPollerState().equalsIgnoreCase("polling") || state.getPollerState().equalsIgnoreCase("processing"))
+                    );
+        });
+
+        // validate there are 8 duplicate messages encountered after reset
+        waitForState(() -> {
+            Map<String, PollingIngestStats> shardTypeToStats = getPollingIngestStatsForPrimaryAndReplica(indexName);
+            assertNotNull(shardTypeToStats.get("primary"));
+            assertNotNull(shardTypeToStats.get("replica"));
+            return shardTypeToStats.get("primary").getConsumerStats().totalDuplicateMessageSkippedCount() == 8
+                && shardTypeToStats.get("replica").getConsumerStats().totalDuplicateMessageSkippedCount() == 8;
+        });
+    }
+
     // returns PollingIngestStats for single primary and single replica
     private Map<String, PollingIngestStats> getPollingIngestStatsForPrimaryAndReplica(String indexName) {
         IndexStats indexStats = client().admin().indices().prepareStats(indexName).get().getIndex(indexName);

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

@@ -1223,6 +1223,10 @@ public boolean useIngestionSource() {
         return ingestionSourceType != null && !(NONE_INGESTION_SOURCE_TYPE.equals(ingestionSourceType));
     }
 
+    public boolean isAllActiveIngestionEnabled() {
+        return INGESTION_SOURCE_ALL_ACTIVE_INGESTION_SETTING.get(settings);
+    }
+
     public IngestionStatus getIngestionStatus() {
         return ingestionStatus;
     }

server/src/main/java/org/opensearch/indices/cluster/IndicesClusterStateService.java

@@ -782,7 +782,8 @@ private void updateShard(
      */
     private void updateShardIngestionState(Shard shard, IndexMetadata indexMetadata, ShardRouting shardRouting) {
         try {
-            if (indexMetadata.useIngestionSource() && shardRouting.primary()) {
+            boolean isPrimaryOrAllActiveShard = shardRouting.primary() || indexMetadata.isAllActiveIngestionEnabled();
+            if (indexMetadata.useIngestionSource() && isPrimaryOrAllActiveShard) {
                 shard.updateShardIngestionState(indexMetadata);
             }
         } catch (Exception e) {