handle drop error policy to avoid failing entire polled batch

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

Author: varunbharadwaj

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

@@ -139,20 +139,20 @@ public void testCloseIndex() throws Exception {
     }
 
     public void testKafkaIngestionWithMultipleProcessorThreads() throws Exception {
-        for (int i = 1; i <= 100; i++) {
+        for (int i = 1; i <= 2000; i++) {
             produceData(String.valueOf(i), "name" + i, "25");
         }
         createIndexWithDefaultSettings(indexName, 1, 0, 3);
         RangeQueryBuilder query = new RangeQueryBuilder("age").gte(21);
         waitForState(() -> {
             refresh(indexName);
             SearchResponse response = client().prepareSearch(indexName).setQuery(query).get();
-            assertThat(response.getHits().getTotalHits().value(), is(100L));
+            assertThat(response.getHits().getTotalHits().value(), is(2000L));
             PollingIngestStats stats = client().admin().indices().prepareStats(indexName).get().getIndex(indexName).getShards()[0]
                 .getPollingIngestStats();
             assertNotNull(stats);
-            assertThat(stats.getMessageProcessorStats().getTotalProcessedCount(), is(100L));
-            assertThat(stats.getConsumerStats().getTotalPolledCount(), is(100L));
+            assertThat(stats.getMessageProcessorStats().getTotalProcessedCount(), is(2000L));
+            assertThat(stats.getConsumerStats().getTotalPolledCount(), is(2000L));
             return true;
         });
     }

server/src/main/java/org/opensearch/indices/pollingingest/DefaultStreamPoller.java

@@ -144,6 +144,8 @@ protected void startPoll() {
         }
         logger.info("Starting poller for shard {}", consumer.getShardId());
 
+        IngestionShardPointer lastProcessedPointer = null;
+        boolean encounteredError = false;
         while (true) {
             try {
                 if (closed) {
@@ -190,29 +192,27 @@ protected void startPoll() {
                 }
 
                 state = State.POLLING;
-
                 List<IngestionShardConsumer.ReadResult<? extends IngestionShardPointer, ? extends Message>> results;
 
-                if (includeBatchStartPointer) {
+                // todo: handle multi-writer scenarios to provide atleast once semantics
+                if (encounteredError && lastProcessedPointer != null) {
+                    results = consumer.readNext(lastProcessedPointer, false, MAX_POLL_SIZE, POLL_TIMEOUT);
+                } else if (includeBatchStartPointer) {
                     results = consumer.readNext(batchStartPointer, true, MAX_POLL_SIZE, POLL_TIMEOUT);
                 } else {
                     results = consumer.readNext(MAX_POLL_SIZE, POLL_TIMEOUT);
                 }
 
+                encounteredError = false;
                 if (results.isEmpty()) {
                     // no new records
                     continue;
                 }
 
                 state = State.PROCESSING;
                 // process the records
-                boolean firstInBatch = true;
                 for (IngestionShardConsumer.ReadResult<? extends IngestionShardPointer, ? extends Message> result : results) {
-                    if (firstInBatch) {
-                        // update the batch start pointer to the next batch
-                        batchStartPointer = result.getPointer();
-                        firstInBatch = false;
-                    }
+                    lastProcessedPointer = result.getPointer();
 
                     // check if the message is already processed
                     if (isProcessed(result.getPointer())) {
@@ -231,6 +231,7 @@ protected void startPoll() {
                 // for future reads, we do not need to include the batch start pointer, and read from the last successful pointer.
                 includeBatchStartPointer = false;
             } catch (Throwable e) {
+                encounteredError = true;
                 logger.error("Error in polling the shard {}: {}", consumer.getShardId(), e);
                 errorStrategy.handleError(e, IngestionErrorStrategy.ErrorStage.POLLING);
 

server/src/main/java/org/opensearch/indices/pollingingest/PartitionedBlockingQueueContainer.java

@@ -130,6 +130,7 @@ public void add(IngestionShardConsumer.ReadResult<? extends IngestionShardPointe
             id = (String) payloadMap.get(ID);
         } else {
             id = IngestionUtils.generateID();
+            payloadMap.put(ID, id);
             autoGeneratedIdTimestamp = System.currentTimeMillis();
         }