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,9 +27,10 @@
 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> {
+public class LeaderOnlyTokenCrawler implements Crawler<PaginationCrawlerWorkerProgressState> {
     private static final Logger log = LoggerFactory.getLogger(LeaderOnlyTokenCrawler.class);
     private static final Duration NO_ACK_TIME_OUT_SECONDS = Duration.ofSeconds(900);
     private static final Duration CHECKPOINT_INTERVAL = Duration.ofMinutes(1);
@@ -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;
@@ -121,8 +122,8 @@ public Instant crawl(LeaderPartition leaderPartition,
     }
 
     @Override
-    public void executePartition(SaasWorkerProgressState state, Buffer buffer, AcknowledgementSet acknowledgementSet) {
-
+    public void executePartition(PaginationCrawlerWorkerProgressState state, Buffer buffer, AcknowledgementSet acknowledgementSet) {
+        client.executePartition(state, buffer, acknowledgementSet);
     }
 
     private List<ItemInfo> collectBatch(Iterator<ItemInfo> iterator) {
@@ -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

@@ -1,5 +1,8 @@
 package org.opensearch.dataprepper.plugins.source.source_crawler.base;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -14,6 +17,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.opensearch.dataprepper.plugins.source.source_crawler.model.TestItemInfo;
@@ -29,12 +34,14 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.verify;
 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)
@@ -144,29 +151,79 @@ void testProgressStateUpdate() {
     }
 
     @Test
-    void testNegativeAcknowledgment() {
-        List<ItemInfo> items = createTestItems(BATCH_SIZE + 1);
-        when(client.listItems(INITIAL_TOKEN)).thenReturn(items.iterator());
-        when(acknowledgementSetManager.create(any(), eq(TEST_TIMEOUT)))
-                .thenReturn(acknowledgementSet);
+    void testRetryPartitionStateSerializationAndDeserialization() {
+        List<ItemInfo> batch = createTestItems( 1);
 
-        ArgumentCaptor<Consumer<Boolean>> callbackCaptor = ArgumentCaptor.forClass(Consumer.class);
+        // Mock client to return our test items
+        when(client.listItems(INITIAL_TOKEN)).thenReturn(batch.iterator());
 
+        // Capture the created partition
+        ArgumentCaptor<SaasSourcePartition> partitionCaptor = ArgumentCaptor.forClass(SaasSourcePartition.class);
+
+        // Simulate negative acknowledgment
+        when(acknowledgementSetManager.create(any(), any())).thenAnswer(inv -> {
+            Consumer<Boolean> callback = inv.getArgument(0);
+            callback.accept(false);
+            return acknowledgementSet;
+        });
+
+        // Create retry partition
         crawler.setAcknowledgementsEnabled(true);
         crawler.crawl(leaderPartition, coordinator);
 
-        verify(acknowledgementSetManager).create(callbackCaptor.capture(), eq(TEST_TIMEOUT));
+        // Verify partition creation and get original state
+        verify(coordinator).createPartition(partitionCaptor.capture());
+        SaasSourcePartition originalPartition = partitionCaptor.getValue();
+        PaginationCrawlerWorkerProgressState originalState =
+                (PaginationCrawlerWorkerProgressState) originalPartition.getProgressState().get();
+
+        // Test serialization/deserialization
+        ObjectMapper objectMapper = new ObjectMapper()
+                .registerModule(new JavaTimeModule());  // For Instant serialization
+
+        try {
+            // Serialize state
+            String serializedState = objectMapper.writeValueAsString(originalState);
+
+            // Deserialize state
+            PaginationCrawlerWorkerProgressState deserializedState =
+                    objectMapper.readValue(serializedState, PaginationCrawlerWorkerProgressState.class);
+
+            // Verify deserialized state matches original
+            assertEquals(originalState.getItemIds(), deserializedState.getItemIds());
+            assertEquals(originalState.getKeyAttributes(), deserializedState.getKeyAttributes());
+            assertEquals(originalState.getExportStartTime(), deserializedState.getExportStartTime());
+            assertEquals(originalState.getLoadedItems(), deserializedState.getLoadedItems());
+        } catch (JsonProcessingException e) {
+            fail("Serialization/deserialization failed", e);
+        }
+    }
 
-        // Simulate negative acknowledgment
-        callbackCaptor.getValue().accept(false);
+    @Test
+    void testNegativeAcknowledgment() {
+        List<ItemInfo> items = createTestItems(1);
+        when(client.listItems(INITIAL_TOKEN)).thenReturn(items.iterator());
+
+        // 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 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 +235,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());
     }