Make acknowledgment_timeout configurable for s3 scan source, configure timeout to 10 minutes for MongoDB L2 transform

graytaylor0 · graytaylor0 · commit 300afcac28f6 · 2024-09-30T15:36:17.000-05:00
Signed-off-by: Taylor Gray &lt;tylgry@amazon.com&gt;
diff --git a/data-prepper-plugins/mongodb/src/main/resources/org/opensearch/dataprepper/transforms/templates/documentdb-template.yaml b/data-prepper-plugins/mongodb/src/main/resources/org/opensearch/dataprepper/transforms/templates/documentdb-template.yaml
@@ -66,6 +66,7 @@
       delete_s3_objects_on_read: true
       disable_s3_metadata_in_event: true
       scan:
+        acknowledgment_timeout: "PT10M"
         folder_partitions:
           depth: "<<FUNCTION_NAME:calculateDepth,PARAMETER:$.<<pipeline-name>>.source.documentdb.s3_prefix>>"
           max_objects_per_ownership: 50
@@ -75,7 +76,7 @@
               filter:
                 include_prefix: ["<<FUNCTION_NAME:getSourceCoordinationIdentifierEnvVariable,PARAMETER:$.<<pipeline-name>>.source.documentdb.s3_prefix>>"]
         scheduling:
-          interval: "60s"
+          interval: "20s"
   processor: "<<$.<<pipeline-name>>.processor>>"
   sink: "<<$.<<pipeline-name>>.sink>>"
   routes: "<<$.<<pipeline-name>>.routes>>" # In placeholder, routes or route (defined as alias) will be transformed to route in json as route will be primarily picked in pipelineModel.
diff --git a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanPartitionCreationSupplier.java b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanPartitionCreationSupplier.java
@@ -90,6 +90,7 @@ public List<PartitionIdentifier> apply(final Map<String, Object> globalStateMap)
             final Instant updatedScanTime = Instant.now();
             if (Objects.nonNull(s3ScanKeyPathOption) && Objects.nonNull(s3ScanKeyPathOption.getS3scanIncludePrefixOptions()))
                 s3ScanKeyPathOption.getS3scanIncludePrefixOptions().forEach(includePath -> {
+                    LOG.info("Include path: {}", includePath);
                     listObjectsV2Request.prefix(includePath);
                     objectsToProcess.addAll(listFilteredS3ObjectsForBucket(excludeItems, listObjectsV2Request,
                             scanOptions.getBucketOption().getName(), scanOptions.getUseStartDateTime(), scanOptions.getUseEndDateTime(), globalStateMap));
diff --git a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/ScanObjectWorker.java b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/ScanObjectWorker.java
@@ -56,8 +56,6 @@ public class ScanObjectWorker implements Runnable {
 
     static final Duration NO_OBJECTS_FOUND_BEFORE_PARTITION_DELETION_DURATION = Duration.ofHours(1);
     private static final int RETRY_BACKOFF_ON_EXCEPTION_MILLIS = 5_000;
-
-    static final Duration ACKNOWLEDGEMENT_SET_TIMEOUT = Duration.ofHours(2);
     static final String ACKNOWLEDGEMENT_SET_CALLBACK_METRIC_NAME = "acknowledgementSetCallbackCounter";
 
     static final String NO_OBJECTS_FOUND_FOR_FOLDER_PARTITION = "folderPartitionNoObjectsFound";
@@ -99,6 +97,8 @@ public class ScanObjectWorker implements Runnable {
 
     private final Map<String, AtomicInteger> acknowledgmentsRemainingForPartitions;
 
+    private final Duration acknowledgmentSetTimeout;
+
     public ScanObjectWorker(final S3Client s3Client,
                             final List<ScanOptions> scanOptionsBuilderList,
                             final S3ObjectHandler s3ObjectHandler,
@@ -127,6 +127,7 @@ public ScanObjectWorker(final S3Client s3Client,
         this.sourceCoordinator.initialize();
         this.partitionKeys = new ArrayList<>();
         this.folderPartitioningOptions = s3SourceConfig.getS3ScanScanOptions().getPartitioningOptions();
+        this.acknowledgmentSetTimeout = s3SourceConfig.getS3ScanScanOptions().getAcknowledgmentTimeout();
 
         this.partitionCreationSupplier = new S3ScanPartitionCreationSupplier(s3Client, bucketOwnerProvider, scanOptionsBuilderList, s3ScanSchedulingOptions, s3SourceConfig.getS3ScanScanOptions().getPartitioningOptions());
         this.acknowledgmentsRemainingForPartitions = new ConcurrentHashMap<>();
@@ -214,7 +215,7 @@ private void startProcessingObject(final long waitTimeMillis) {
                         sourceCoordinator.giveUpPartition(objectToProcess.get().getPartitionKey());
                     }
                     partitionKeys.remove(objectToProcess.get().getPartitionKey());
-                }, ACKNOWLEDGEMENT_SET_TIMEOUT);
+                }, acknowledgmentSetTimeout);
 
                 addProgressCheck(acknowledgementSet, objectToProcess.get());
             }
@@ -226,7 +227,7 @@ private void startProcessingObject(final long waitTimeMillis) {
             if (endToEndAcknowledgementsEnabled) {
                 deleteObjectRequest.ifPresent(deleteRequest -> objectsToDeleteForAcknowledgmentSets.put(objectToProcess.get().getPartitionKey(), Set.of(deleteRequest)));
                 try {
-                    sourceCoordinator.updatePartitionForAcknowledgmentWait(objectToProcess.get().getPartitionKey(), ACKNOWLEDGEMENT_SET_TIMEOUT);
+                    sourceCoordinator.updatePartitionForAcknowledgmentWait(objectToProcess.get().getPartitionKey(), acknowledgmentSetTimeout);
                 } catch (final PartitionUpdateException e) {
                     LOG.debug("Failed to update the partition for the acknowledgment wait.");
                 }
@@ -375,7 +376,7 @@ private void processObjectsForFolderPartition(final List<S3ObjectReference> obje
             objectIndex++;
         }
 
-        sourceCoordinator.updatePartitionForAcknowledgmentWait(folderPartition.getPartitionKey(), ACKNOWLEDGEMENT_SET_TIMEOUT);
+        sourceCoordinator.updatePartitionForAcknowledgmentWait(folderPartition.getPartitionKey(), acknowledgmentSetTimeout);
 
         if (acknowledgementSet != null) {
             acknowledgementSet.complete();
@@ -402,7 +403,7 @@ private AcknowledgementSet createAcknowledgmentSetForFolderPartition(final Sourc
                 LOG.info("Received all acknowledgments for folder partition {}, giving up this partition", folderPartition.getPartitionKey());
                 sourceCoordinator.giveUpPartition(folderPartition.getPartitionKey(), Instant.now());
             }
-        }, ACKNOWLEDGEMENT_SET_TIMEOUT);
+        }, acknowledgmentSetTimeout);
     }
 
     private void addProgressCheck(final AcknowledgementSet acknowledgementSet, final SourcePartition<S3SourceProgressState> objectToProcess) {
diff --git a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/configuration/S3ScanScanOptions.java b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/configuration/S3ScanScanOptions.java
@@ -22,6 +22,9 @@
  */
 public class S3ScanScanOptions {
 
+    @JsonProperty("acknowledgment_timeout")
+    private Duration acknowledgmentTimeout = Duration.ofHours(2);
+
     @JsonProperty("folder_partitions")
     @Valid
     private FolderPartitioningOptions folderPartitioningOptions;
@@ -84,4 +87,6 @@ public S3ScanSchedulingOptions getSchedulingOptions() {
     }
 
     public FolderPartitioningOptions getPartitioningOptions() { return folderPartitioningOptions; }
+
+    public Duration getAcknowledgmentTimeout() { return acknowledgmentTimeout; }
 }
diff --git a/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanObjectWorkerTest.java b/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanObjectWorkerTest.java
@@ -69,7 +69,6 @@
 import static org.mockito.Mockito.when;
 import static org.opensearch.dataprepper.model.source.s3.S3ScanEnvironmentVariables.STOP_S3_SCAN_PROCESSING_PROPERTY;
 import static org.opensearch.dataprepper.plugins.source.s3.ScanObjectWorker.ACKNOWLEDGEMENT_SET_CALLBACK_METRIC_NAME;
-import static org.opensearch.dataprepper.plugins.source.s3.ScanObjectWorker.ACKNOWLEDGEMENT_SET_TIMEOUT;
 import static org.opensearch.dataprepper.plugins.source.s3.ScanObjectWorker.CHECKPOINT_OWNERSHIP_INTERVAL;
 import static org.opensearch.dataprepper.plugins.source.s3.ScanObjectWorker.NO_OBJECTS_FOUND_BEFORE_PARTITION_DELETION_DURATION;
 import static org.opensearch.dataprepper.plugins.source.s3.ScanObjectWorker.NO_OBJECTS_FOUND_FOR_FOLDER_PARTITION;
@@ -125,10 +124,14 @@ class S3ScanObjectWorkerTest {
 
     private List<ScanOptions> scanOptionsList;
 
+    @Mock
+    private Duration acknowledgmentSetTimeout;
+
     @BeforeEach
     void setup() {
         scanOptionsList = new ArrayList<>();
         when(s3ScanScanOptions.getPartitioningOptions()).thenReturn(null);
+        when(s3ScanScanOptions.getAcknowledgmentTimeout()).thenReturn(acknowledgmentSetTimeout);
     }
 
     private ScanObjectWorker createObjectUnderTest() {
@@ -232,7 +235,7 @@ void buildDeleteObjectRequest_should_be_invoked_after_processing_when_deleteS3Ob
 
         final InOrder inOrder = inOrder(sourceCoordinator, acknowledgementSet, s3ObjectDeleteWorker);
         inOrder.verify(s3ObjectDeleteWorker).buildDeleteObjectRequest(bucket, objectKey);
-        inOrder.verify(sourceCoordinator).updatePartitionForAcknowledgmentWait(partitionKey, ACKNOWLEDGEMENT_SET_TIMEOUT);
+        inOrder.verify(sourceCoordinator).updatePartitionForAcknowledgmentWait(partitionKey, acknowledgmentSetTimeout);
         inOrder.verify(acknowledgementSet).complete();
         inOrder.verify(sourceCoordinator).renewPartitionOwnership(partitionKey);
         inOrder.verify(sourceCoordinator).completePartition(partitionKey, true);
@@ -289,7 +292,7 @@ void acknowledgment_progress_check_increments_ownership_error_metric_when_partit
 
         final InOrder inOrder = inOrder(sourceCoordinator, acknowledgementSet, s3ObjectDeleteWorker);
         inOrder.verify(s3ObjectDeleteWorker).buildDeleteObjectRequest(bucket, objectKey);
-        inOrder.verify(sourceCoordinator).updatePartitionForAcknowledgmentWait(partitionKey, ACKNOWLEDGEMENT_SET_TIMEOUT);
+        inOrder.verify(sourceCoordinator).updatePartitionForAcknowledgmentWait(partitionKey, acknowledgmentSetTimeout);
         inOrder.verify(acknowledgementSet).complete();
         inOrder.verify(sourceCoordinator).renewPartitionOwnership(partitionKey);
         inOrder.verify(sourceCoordinator).completePartition(partitionKey, true);
@@ -536,7 +539,7 @@ 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(sourceCoordinator).updatePartitionForAcknowledgmentWait(partitionKey, ACKNOWLEDGEMENT_SET_TIMEOUT);
+        inOrder.verify(sourceCoordinator).updatePartitionForAcknowledgmentWait(partitionKey, acknowledgmentSetTimeout);
         inOrder.verify(acknowledgementSet2).complete();
 
         final Consumer<Boolean> firstAckCallback = ackCallbacks.get(0);
@@ -616,7 +619,7 @@ 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(sourceCoordinator).updatePartitionForAcknowledgmentWait(partitionKey, ACKNOWLEDGEMENT_SET_TIMEOUT);
+        inOrder.verify(sourceCoordinator).updatePartitionForAcknowledgmentWait(partitionKey, acknowledgmentSetTimeout);
         inOrder.verify(acknowledgementSet1).complete();
 
         final Consumer<Boolean> ackCallback = consumerArgumentCaptor.getValue();

Original file line number	Diff line number	Diff line change
`@@ -22,6 +22,9 @@`
`22`	`22`	`*/`
`23`	`23`	`public class S3ScanScanOptions {`
`24`	`24`
	`25`	`+ @JsonProperty("acknowledgment_timeout")`
	`26`	`+ private Duration acknowledgmentTimeout = Duration.ofHours(2);`
	`27`	`+`
`25`	`28`	`@JsonProperty("folder_partitions")`
`26`	`29`	`@Valid`
`27`	`30`	`private FolderPartitioningOptions folderPartitioningOptions;`
`@@ -84,4 +87,6 @@ public S3ScanSchedulingOptions getSchedulingOptions() {`
`84`	`87`	`}`
`85`	`88`
`86`	`89`	`public FolderPartitioningOptions getPartitioningOptions() { return folderPartitioningOptions; }`
	`90`	`+`
	`91`	`+ public Duration getAcknowledgmentTimeout() { return acknowledgmentTimeout; }`
`87`	`92`	`}`