Fix bug where race condition on ack callback could cause S3 folder pa…

…rtition to not be given up (#4835) Signed-off-by: Taylor Gray <tylgry@amazon.com>
opensearch-project · Aug 19, 2024 · 9a82590 · 9a82590
1 parent 04098b3
commit 9a82590
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/...3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/ScanObjectWorker.java b/...3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/ScanObjectWorker.java
@@ -280,8 +280,6 @@ private void processFolderPartition(final SourcePartition<S3SourceProgressState>
         sourceCoordinator.saveProgressStateForPartition(folderPartition.getPartitionKey(), folderPartitionState.get());
 
         processObjectsForFolderPartition(objectsToProcess, folderPartition);
-
-        sourceCoordinator.updatePartitionForAcknowledgmentWait(folderPartition.getPartitionKey(), ACKNOWLEDGEMENT_SET_TIMEOUT);
     }
 
     private List<S3ObjectReference> getObjectsForPrefix(final String bucket, final String s3Prefix) {
@@ -364,7 +362,8 @@ private void processObjectsForFolderPartition(final List<S3ObjectReference> obje
             objectIndex++;
         }
 
-        // Complete the final acknowledgment set
+        sourceCoordinator.updatePartitionForAcknowledgmentWait(folderPartition.getPartitionKey(), ACKNOWLEDGEMENT_SET_TIMEOUT);
+
         if (acknowledgementSet != null) {
             acknowledgementSet.complete();
         }

diff --git a/...ce/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanObjectWorkerTest.java b/...ce/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanObjectWorkerTest.java
@@ -460,8 +460,8 @@ void processing_with_folder_partition_processes_objects_in_folder_and_deletes_th
         inOrder.verify(s3ObjectDeleteWorker).buildDeleteObjectRequest(bucket, firstObject.key());
         inOrder.verify(acknowledgementSet1).complete();
         inOrder.verify(s3ObjectDeleteWorker).buildDeleteObjectRequest(bucket, secondObject.key());
-        inOrder.verify(acknowledgementSet2).complete();
         inOrder.verify(sourceCoordinator).updatePartitionForAcknowledgmentWait(partitionKey, ACKNOWLEDGEMENT_SET_TIMEOUT);
+        inOrder.verify(acknowledgementSet2).complete();
 
         final Consumer<Boolean> firstAckCallback = ackCallbacks.get(0);
         firstAckCallback.accept(true);
@@ -532,8 +532,8 @@ void processing_with_folder_partition_processes_objects_in_folder_until_max_obje
         final InOrder inOrder = inOrder(sourceCoordinator, acknowledgementSet1, s3ObjectDeleteWorker);
 
         inOrder.verify(s3ObjectDeleteWorker).buildDeleteObjectRequest(bucket, firstObject.key());
-        inOrder.verify(acknowledgementSet1).complete();
         inOrder.verify(sourceCoordinator).updatePartitionForAcknowledgmentWait(partitionKey, ACKNOWLEDGEMENT_SET_TIMEOUT);
+        inOrder.verify(acknowledgementSet1).complete();
 
         final Consumer<Boolean> ackCallback = consumerArgumentCaptor.getValue();
         ackCallback.accept(true);