Skip second completePartition when WorkerScheduler is processing partition (#6066)

Signed-off-by: Brendan Benner <bbenner@amazon.com>
Co-authored-by: Brendan Benner <bbenner@amazon.com>

Author: bbenner7635

data-prepper-plugins/saas-source-plugins/source-crawler/src/main/java/org/opensearch/dataprepper/plugins/source/source_crawler/coordination/scheduler/WorkerScheduler.java

@@ -13,6 +13,7 @@
 import org.opensearch.dataprepper.plugins.source.source_crawler.base.CrawlerSourceConfig;
 import org.opensearch.dataprepper.plugins.source.source_crawler.base.SaasWorkerProgressState;
 import org.opensearch.dataprepper.plugins.source.source_crawler.coordination.partition.SaasSourcePartition;
+import com.google.common.annotations.VisibleForTesting;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -111,13 +112,21 @@ private void processPartition(EnhancedSourcePartition partition, Buffer<Record<E
             AcknowledgementSet acknowledgementSet = null;
             if (sourceConfig.isAcknowledgments()) {
                 acknowledgementSet = createAcknowledgementSet(partition);
+                // When acknowledgments are enabled, partition completion is handled in the acknowledgment callback
+                crawler.executePartition((SaasWorkerProgressState) partition.getProgressState().get(), buffer, acknowledgementSet);
+            } else {
+                // When acknowledgments are disabled, complete the partition immediately after execution
+                crawler.executePartition((SaasWorkerProgressState) partition.getProgressState().get(), buffer, acknowledgementSet);
+                sourceCoordinator.completePartition(partition);
             }
-            crawler.executePartition((SaasWorkerProgressState) partition.getProgressState().get(), buffer, acknowledgementSet);
+        } else {
+            // If no progress state, complete the partition immediately
+            sourceCoordinator.completePartition(partition);
         }
-        sourceCoordinator.completePartition(partition);
     }
 
-    private AcknowledgementSet createAcknowledgementSet(EnhancedSourcePartition partition) {
+    @VisibleForTesting
+    AcknowledgementSet createAcknowledgementSet(EnhancedSourcePartition partition) {
         return acknowledgementSetManager.create((result) -> {
             if (result) {
                 acknowledgementSetSuccesses.increment();

data-prepper-plugins/saas-source-plugins/source-crawler/src/test/java/org/opensearch/dataprepper/plugins/source/source_crawler/coordination/scheduler/WorkerSchedulerTest.java

@@ -30,9 +30,11 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.BDDMockito.given;
 import static org.mockito.Mockito.atLeast;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoInteractions;
 import static org.mockito.Mockito.when;
@@ -216,4 +218,45 @@ void testRetryBackOffTriggeredWhenExceptionOccurred() throws InterruptedExceptio
         // Crawler shouldn't be invoked in this case
         verifyNoInteractions(crawler);
     }
+
+    @Test
+    void testCompletePartitionWithAcknowledgements() throws InterruptedException {
+        WorkerScheduler workerScheduler = new WorkerScheduler(pluginName, buffer,
+                coordinator, sourceConfig, crawler, pluginMetrics, acknowledgementSetManager);
+        WorkerScheduler spyWorkerScheduler = org.mockito.Mockito.spy(workerScheduler);
+        
+        when(sourceConfig.isAcknowledgments()).thenReturn(true);
+        
+        PaginationCrawlerWorkerProgressState mockProgressState = new PaginationCrawlerWorkerProgressState();
+        
+        EnhancedSourcePartition mockPartition = org.mockito.Mockito.mock(EnhancedSourcePartition.class);
+        when(mockPartition.getProgressState()).thenReturn(Optional.of(mockProgressState));
+        
+        org.mockito.Mockito.doAnswer(invocation -> {
+            coordinator.completePartition(mockPartition);
+            return acknowledgementSet;
+        }).when(spyWorkerScheduler).createAcknowledgementSet(any(EnhancedSourcePartition.class));
+        
+        given(coordinator.acquireAvailablePartition(SaasSourcePartition.PARTITION_TYPE))
+            .willReturn(Optional.of(mockPartition))
+            .willReturn(Optional.empty());
+
+        ExecutorService executorService = Executors.newSingleThreadExecutor();
+        executorService.submit(spyWorkerScheduler);
+
+        Thread.sleep(500);
+        executorService.shutdownNow();
+        
+        verify(mockPartition, atLeast(1)).getProgressState();
+        verify(sourceConfig, atLeast(1)).isAcknowledgments();
+        
+        // Verify that createAcknowledgementSet was called
+        verify(spyWorkerScheduler, times(1)).createAcknowledgementSet(eq(mockPartition));
+        
+        // Verify that crawler.executePartition was called with acknowledgement set
+        verify(crawler, times(1)).executePartition(any(SaasWorkerProgressState.class), eq(buffer), eq(acknowledgementSet));
+        
+        // Verify that coordinator.completePartition was called (from the acknowledgement callback with true)
+        verify(coordinator, times(1)).completePartition(eq(mockPartition));
+    }
 }