Fix possible missing file count in data file loader (#6639)

Signed-off-by: Hai Yan <oeyh@amazon.com>

Author: oeyh

data-prepper-plugins/rds-source/src/main/java/org/opensearch/dataprepper/plugins/source/rds/export/DataFileScheduler.java

@@ -51,12 +51,11 @@ public class DataFileScheduler implements Runnable {
 
     private static final Duration DEFAULT_UPDATE_LOAD_STATUS_TIMEOUT = Duration.ofMinutes(30);
 
-
+    static final int CREATE_PARTITION_MAX_RETRIES = 3;
     static final String EXPORT_S3_OBJECTS_PROCESSED_COUNT = "exportS3ObjectsProcessed";
     static final String EXPORT_S3_OBJECTS_ERROR_COUNT = "exportS3ObjectsErrors";
     static final String ACTIVE_EXPORT_S3_OBJECT_CONSUMERS_GAUGE = "activeExportS3ObjectConsumers";
 
-
     private final EnhancedSourceCoordinator sourceCoordinator;
     private final ExecutorService executor;
     private final RdsSourceConfig sourceConfig;
@@ -176,35 +175,47 @@ private void processDataFilePartition(DataFilePartition dataFilePartition) {
         numOfWorkers.incrementAndGet();
     }
 
-    private void updateLoadStatus(String exportTaskId, Duration timeout) {
+    void updateLoadStatus(String exportTaskId, Duration timeout) {
 
         Instant endTime = Instant.now().plus(timeout);
+        LoadStatus loadStatus = null;
+        boolean savedSuccessfully = false;
         // Keep retrying in case update fails due to conflicts until timed out
         while (Instant.now().isBefore(endTime)) {
             Optional<EnhancedSourcePartition> globalStatePartition = sourceCoordinator.getPartition(exportTaskId);
             if (globalStatePartition.isEmpty()) {
                 LOG.error("Failed to get data file load status for {}", exportTaskId);
-                return;
+                // Transient read failures should be retried
+                continue;
             }
 
             GlobalState globalState = (GlobalState) globalStatePartition.get();
-            LoadStatus loadStatus = LoadStatus.fromMap(globalState.getProgressState().get());
+            loadStatus = LoadStatus.fromMap(globalState.getProgressState().get());
             loadStatus.setLoadedFiles(loadStatus.getLoadedFiles() + 1);
             LOG.info("Current data file load status: total {} loaded {}", loadStatus.getTotalFiles(), loadStatus.getLoadedFiles());
 
             globalState.setProgressState(loadStatus.toMap());
 
             try {
                 sourceCoordinator.saveProgressStateForPartition(globalState, null);
-                if (sourceConfig.isStreamEnabled() && loadStatus.getLoadedFiles() == loadStatus.getTotalFiles()) {
-                    LOG.info("All exports are done, streaming can continue...");
-                    sourceCoordinator.createPartition(new GlobalState("stream-for-" + sourceConfig.getDbIdentifier(), null));
-                }
+                savedSuccessfully = true;
                 break;
             } catch (Exception e) {
                 LOG.error("Failed to update the global status, looks like the status was out of date, will retry..");
             }
         }
+
+        if (savedSuccessfully && sourceConfig.isStreamEnabled() && loadStatus.getLoadedFiles() == loadStatus.getTotalFiles()) {
+            LOG.info("All exports are done, streaming can continue...");
+            for (int attempt = 0; attempt < CREATE_PARTITION_MAX_RETRIES; attempt++) {
+                try {
+                    sourceCoordinator.createPartition(new GlobalState("stream-for-" + sourceConfig.getDbIdentifier(), null));
+                    break;
+                } catch (Exception e) {
+                    LOG.error("Failed to create stream trigger partition for {} (attempt {}), will retry", sourceConfig.getDbIdentifier(), attempt + 1, e);
+                }
+            }
+        }
     }
 
     private BiConsumer<Void, Throwable> completeDataLoader(DataFilePartition dataFilePartition) {

data-prepper-plugins/rds-source/src/test/java/org/opensearch/dataprepper/plugins/source/rds/export/DataFileSchedulerTest.java

@@ -43,11 +43,13 @@
 import static org.awaitility.Awaitility.await;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.isNull;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.mockStatic;
 import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoInteractions;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
@@ -187,6 +189,69 @@ void test_data_file_loader_throws_exception_then_give_up_partition() {
         verify(sourceCoordinator).giveUpPartition(dataFilePartition);
     }
 
+    @Test
+    void test_when_getPartition_initially_returns_empty_then_retries_and_updates_load_status() {
+        final String exportTaskId = UUID.randomUUID().toString();
+
+        final GlobalState loadStatusGlobalState = mock(GlobalState.class);
+        final int totalFiles = 5;
+        final Map<String, Object> loadStatusMap = new LoadStatus(totalFiles, totalFiles - 2).toMap();
+        when(loadStatusGlobalState.getProgressState()).thenReturn(Optional.of(loadStatusMap));
+
+        // First call returns empty (transient failure), second call returns valid state
+        when(sourceCoordinator.getPartition(exportTaskId))
+                .thenReturn(Optional.empty())
+                .thenReturn(Optional.of(loadStatusGlobalState));
+
+        createObjectUnderTest().updateLoadStatus(exportTaskId, Duration.ofSeconds(2));
+
+        verify(sourceCoordinator).saveProgressStateForPartition(any(GlobalState.class), isNull());
+    }
+
+    @Test
+    void test_when_last_file_loaded_with_stream_enabled_then_creates_stream_trigger_partition() {
+        final String exportTaskId = UUID.randomUUID().toString();
+        final String dbIdentifier = UUID.randomUUID().toString();
+        when(sourceConfig.isStreamEnabled()).thenReturn(true);
+        when(sourceConfig.getDbIdentifier()).thenReturn(dbIdentifier);
+
+        final GlobalState loadStatusGlobalState = mock(GlobalState.class);
+        final int totalFiles = 3;
+        final Map<String, Object> loadStatusMap = new LoadStatus(totalFiles, totalFiles - 1).toMap();
+        when(loadStatusGlobalState.getProgressState()).thenReturn(Optional.of(loadStatusMap));
+        when(sourceCoordinator.getPartition(exportTaskId)).thenReturn(Optional.of(loadStatusGlobalState));
+
+        createObjectUnderTest().updateLoadStatus(exportTaskId, Duration.ofSeconds(2));
+
+        verify(sourceCoordinator).createPartition(any(GlobalState.class));
+    }
+
+    @Test
+    void test_when_createPartition_throws_transiently_then_retries_and_partition_is_completed() {
+        final String exportTaskId = UUID.randomUUID().toString();
+        final String dbIdentifier = UUID.randomUUID().toString();
+        when(sourceConfig.isStreamEnabled()).thenReturn(true);
+        when(sourceConfig.getDbIdentifier()).thenReturn(dbIdentifier);
+
+        final GlobalState loadStatusGlobalState = mock(GlobalState.class);
+        final int totalFiles = 3;
+        final Map<String, Object> loadStatusMap = new LoadStatus(totalFiles, totalFiles - 1).toMap();
+        when(loadStatusGlobalState.getProgressState()).thenReturn(Optional.of(loadStatusMap));
+        when(sourceCoordinator.getPartition(exportTaskId)).thenReturn(Optional.of(loadStatusGlobalState));
+
+        // createPartition throws once then succeeds — verifies retry behavior
+        when(sourceCoordinator.createPartition(any(GlobalState.class)))
+                .thenThrow(new RuntimeException("DynamoDB error"))
+                .thenReturn(true);
+
+        createObjectUnderTest().updateLoadStatus(exportTaskId, Duration.ofSeconds(2));
+
+        // saveProgressStateForPartition must be called exactly once — no retry/overshoot
+        verify(sourceCoordinator, times(1)).saveProgressStateForPartition(any(GlobalState.class), isNull());
+        // createPartition retried: first threw, second succeeded
+        verify(sourceCoordinator, times(2)).createPartition(any(GlobalState.class));
+    }
+
     @Disabled("Flaky test, needs to be fixed")
     @Test
     void test_shutdown() throws InterruptedException {