Add worker retry mechanism for failed batch processing in LeaderOnlyTokenCrawler

Signed-off-by: Alekhya Parisha <aparisha@amazon.com>

Author: Alekhya Parisha

data-prepper-plugins/saas-source-plugins/source-crawler/src/main/java/org/opensearch/dataprepper/plugins/source/source_crawler/base/LeaderOnlyTokenCrawler.java

@@ -12,6 +12,8 @@
 import org.opensearch.dataprepper.model.record.Record;
 import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourceCoordinator;
 import org.opensearch.dataprepper.plugins.source.source_crawler.coordination.partition.LeaderPartition;
+import org.opensearch.dataprepper.plugins.source.source_crawler.coordination.partition.SaasSourcePartition;
+import org.opensearch.dataprepper.plugins.source.source_crawler.coordination.state.PaginationCrawlerWorkerProgressState;
 import org.opensearch.dataprepper.plugins.source.source_crawler.coordination.state.TokenPaginationCrawlerLeaderProgressState;
 import org.opensearch.dataprepper.plugins.source.source_crawler.model.ItemInfo;
 import org.slf4j.Logger;
@@ -25,6 +27,7 @@
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.stream.Collectors;
 
 @Named
 public class LeaderOnlyTokenCrawler implements Crawler<SaasWorkerProgressState> {
@@ -54,7 +57,6 @@ public class LeaderOnlyTokenCrawler implements Crawler<SaasWorkerProgressState>
     private final Timer bufferWriteTimer;
 
     private String lastToken;
-    private boolean shouldStopCrawl = false;
     private Duration noAckTimeout;
 
     public LeaderOnlyTokenCrawler(
@@ -73,7 +75,6 @@ public LeaderOnlyTokenCrawler(
     @Override
     public Instant crawl(LeaderPartition leaderPartition,
                          EnhancedSourceCoordinator coordinator) {
-        shouldStopCrawl = false;
         long startTime = System.currentTimeMillis();
         Instant lastCheckpointTime = Instant.now();
         TokenPaginationCrawlerLeaderProgressState leaderProgressState =
@@ -84,7 +85,7 @@ public Instant crawl(LeaderPartition leaderPartition,
 
         Iterator<ItemInfo> itemIterator = ((LeaderOnlyTokenCrawlerClient) client).listItems(lastToken);
 
-        while (itemIterator.hasNext() && !shouldStopCrawl) {
+        while (itemIterator.hasNext()) {
             List<ItemInfo> batch = collectBatch(itemIterator);
             if (batch.isEmpty()) {
                 continue;
@@ -148,12 +149,12 @@ private void processBatch(List<ItemInfo> batch,
                         if (success) {
                             // On success: update checkpoint
                             acknowledgementSetSuccesses.increment();
-                            updateLeaderProgressState(leaderPartition, lastToken, coordinator);
                         } else {
-                            // On failure: Stop the crawl
+                            // On failure: Create a retry partition
                             acknowledgementSetFailures.increment();
-                            log.warn("Batch processing received negative acknowledgment for token: {}. Stopping current crawl.", lastToken);
-                            shouldStopCrawl = true;
+                            log.warn("Batch processing received negative acknowledgment for token: {}. Creating retry " +
+                                    "partition.", lastToken);
+                            createRetryPartition(batch, coordinator);
                         }
                     },
                     noAckTimeout
@@ -172,17 +173,19 @@ private void processBatch(List<ItemInfo> batch,
 
                         if (!ackWaitDuration.minus(noAckTimeout).isNegative()) {
                             // No ack received within NO_ACK_TIME_OUT_SECONDS
-                            log.warn("Acknowledgment not received for batch with token {} past wait time. Stopping current crawl.", lastToken);
-                            shouldStopCrawl = true;
+                            log.warn("No acknowledgment received for batch with token: {}. Creating retry partition.", lastToken);
+                            createRetryPartition(batch, coordinator);
                             break;
                         }
                     }
+                    updateLeaderProgressState(leaderPartition, lastToken, coordinator);
                 } catch (InterruptedException e) {
                     Thread.currentThread().interrupt();
                     throw new RuntimeException("Interrupted while waiting for acknowledgment", e);
                 } catch (Exception e) {
                     log.error("Failed to process batch ending with token {}", lastToken, e);
                     acknowledgementSet.complete();
+                    createRetryPartition(batch, coordinator);
                     throw e;
                 }
             });
@@ -201,6 +204,23 @@ private void processBatch(List<ItemInfo> batch,
         }
     }
 
+    private void createRetryPartition(List<ItemInfo> itemInfoList, EnhancedSourceCoordinator coordinator) {
+        if (itemInfoList.isEmpty()) {
+            return;
+        }
+        ItemInfo itemInfo = itemInfoList.get(0);
+        String partitionKey = itemInfo.getPartitionKey();
+        List<String> itemIds = itemInfoList.stream().map(ItemInfo::getId).collect(Collectors.toList());
+        PaginationCrawlerWorkerProgressState state = new PaginationCrawlerWorkerProgressState();
+        state.setKeyAttributes(itemInfo.getKeyAttributes());
+        state.setItemIds(itemIds);
+        state.setExportStartTime(Instant.now());
+        state.setLoadedItems(itemInfoList.size());
+        SaasSourcePartition sourcePartition = new SaasSourcePartition(state, partitionKey);
+        coordinator.createPartition(sourcePartition);
+    }
+
+
     private void updateLeaderProgressState(LeaderPartition leaderPartition,
                                            String updatedToken,
                                            EnhancedSourceCoordinator coordinator) {

data-prepper-plugins/saas-source-plugins/source-crawler/src/test/java/org/opensearch/dataprepper/plugins/source/source_crawler/base/LeaderOnlyTokenCrawlerTest.java

@@ -35,6 +35,7 @@
 import static org.mockito.Mockito.when;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.internal.verification.VerificationModeFactory.times;
 
 @ExtendWith(MockitoExtension.class)
@@ -145,28 +146,29 @@ void testProgressStateUpdate() {
 
     @Test
     void testNegativeAcknowledgment() {
-        List<ItemInfo> items = createTestItems(BATCH_SIZE + 1);
+        List<ItemInfo> items = createTestItems(1);
         when(client.listItems(INITIAL_TOKEN)).thenReturn(items.iterator());
-        when(acknowledgementSetManager.create(any(), eq(TEST_TIMEOUT)))
-                .thenReturn(acknowledgementSet);
 
-        ArgumentCaptor<Consumer<Boolean>> callbackCaptor = ArgumentCaptor.forClass(Consumer.class);
+        // Setup immediate negative acknowledgment
+        doAnswer(invocation -> {
+            Consumer<Boolean> callback = invocation.getArgument(0);
+            callback.accept(false);  // Trigger negative ack immediately
+            return acknowledgementSet;
+        }).when(acknowledgementSetManager).create(any(), eq(TEST_TIMEOUT));
 
         crawler.setAcknowledgementsEnabled(true);
         crawler.crawl(leaderPartition, coordinator);
 
-        verify(acknowledgementSetManager).create(callbackCaptor.capture(), eq(TEST_TIMEOUT));
-
-        // Simulate negative acknowledgment
-        callbackCaptor.getValue().accept(false);
-
+        // Verify behavior
         verify(client, times(1)).writeBatchToBuffer(any(), any(), any());
-        verify(coordinator, never()).saveProgressStateForPartition(eq(leaderPartition), any(Duration.class));
+        verify(coordinator, times(1)).createPartition(any());
+        verify(acknowledgementSet, times(1)).complete();
     }
 
+
     @Test
     void testAcknowledgmentTimeout() {
-        List<ItemInfo> items = createTestItems(BATCH_SIZE + 1);
+        List<ItemInfo> items = createTestItems( 1);
         when(client.listItems(INITIAL_TOKEN)).thenReturn(items.iterator());
         when(acknowledgementSetManager.create(any(), eq(TEST_TIMEOUT)))
                 .thenReturn(acknowledgementSet);
@@ -178,13 +180,10 @@ void testAcknowledgmentTimeout() {
 
         verify(acknowledgementSetManager).create(callbackCaptor.capture(), eq(TEST_TIMEOUT));
 
-        // Verify:
-        // 1. Only first batch was processed
+        // Verify timeout behavior
         verify(client, times(1)).writeBatchToBuffer(any(), any(), any());
-        // 2. No checkpoint update happened
-        verify(coordinator, never()).saveProgressStateForPartition(eq(leaderPartition), any(Duration.class));
-        // 3. Acknowledgment set was completed
         verify(acknowledgementSet).complete();
+        verify(coordinator).createPartition(any());
     }