Add ownership renewal during S3 partition creation

Prevents timeout during long-running operations by renewing global state ownership every 3 minutes

Signed-off-by: Siqi Ding <dingdd@amazon.com>

Author: Davidding4718

data-prepper-api/src/main/java/org/opensearch/dataprepper/model/source/coordinator/SourceCoordinator.java

@@ -164,4 +164,9 @@ public interface SourceCoordinator<T> {
     void renewPartitionOwnership(final String partitionKey);
 
     void deletePartition(final String partitionKey);
+
+    /**
+     * Renews the global state ownership for partition creation to prevent timeout during long-running operations
+     */
+    void renewGlobalStateOwnershipForPartitionCreation();
 }

data-prepper-core/src/main/java/org/opensearch/dataprepper/core/sourcecoordination/LeaseBasedSourceCoordinator.java

@@ -61,6 +61,7 @@ public class LeaseBasedSourceCoordinator<T> implements SourceCoordinator<T> {
     static final String PARTITION_UPDATE_ERROR_COUNT = "PartitionUpdateErrors";
     static final Duration DEFAULT_LEASE_TIMEOUT = Duration.ofMinutes(10);
     static final Duration DEFAULT_RENEW_TIMEOUT = Duration.ofMinutes(3);
+    static final Duration OWNERSHIP_RENEWAL_INTERVAL = Duration.ofMinutes(3);
 
     private static final String hostName;
     static final String PARTITION_TYPE = "PARTITION";
@@ -208,7 +209,7 @@ public void createPartitions(final List<PartitionIdentifier> partitionIdentifier
         Instant lastOwnershipRenewal = Instant.now();
 
         for (final PartitionIdentifier partitionIdentifier : partitionIdentifiers) {
-            if (Instant.now().isAfter(lastOwnershipRenewal.plus(DEFAULT_RENEW_TIMEOUT))) {
+            if (Instant.now().isAfter(lastOwnershipRenewal.plus(OWNERSHIP_RENEWAL_INTERVAL))) {
                 LOG.info("Renewing global state ownership for partition creation");
                 renewGlobalStateOwnershipForPartitionCreation();
                 lastOwnershipRenewal = Instant.now();
@@ -526,7 +527,8 @@ private Optional<SourcePartitionStoreItem> acquireGlobalStateForPartitionCreatio
         return globalStateItemForPartitionCreation;
     }
 
-    private void renewGlobalStateOwnershipForPartitionCreation() {
+    @Override
+    public void renewGlobalStateOwnershipForPartitionCreation() {
         try {
             final Optional<SourcePartitionStoreItem> globalStateItem = sourceCoordinationStore.getSourcePartitionItem(
                     sourceIdentifierWithGlobalStateType, GLOBAL_STATE_SOURCE_PARTITION_KEY_FOR_CREATING_PARTITIONS);

data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/ScanObjectWorker.java

@@ -58,6 +58,7 @@ public class ScanObjectWorker implements Runnable {
     private static final Integer MAX_OBJECTS_PER_ACKNOWLEDGMENT_SET = 1;
     static final Integer MAX_RETRIES = 5;
     static final Duration CHECKPOINT_OWNERSHIP_INTERVAL = Duration.ofMinutes(2);
+    static final Duration OWNERSHIP_RENEWAL_INTERVAL = Duration.ofMinutes(3);
 
     static final Duration NO_OBJECTS_FOUND_BEFORE_PARTITION_DELETION_DURATION = Duration.ofHours(1);
     private static final int RETRY_BACKOFF_ON_EXCEPTION_MILLIS = 5_000;
@@ -305,6 +306,7 @@ private void processFolderPartition(final SourcePartition<S3SourceProgressState>
         final String bucket = folderPartition.getPartitionKey().split("\\|")[0];
         final String s3Prefix = folderPartition.getPartitionKey().split("\\|")[1];
 
+        sourceCoordinator.renewGlobalStateOwnershipForPartitionCreation();
         final List<S3ObjectReference> objectsToProcess = getObjectsForPrefix(bucket, s3Prefix);
 
         Optional<S3SourceProgressState> folderPartitionState = folderPartition.getPartitionState();
@@ -337,12 +339,19 @@ private void processFolderPartition(final SourcePartition<S3SourceProgressState>
     private List<S3ObjectReference> getObjectsForPrefix(final String bucket, final String s3Prefix) {
         ListObjectsV2Response listObjectsV2Response = null;
         final List<S3ObjectReference> objectsToProcess = new ArrayList<>();
+        Instant lastOwnershipRenewal = Instant.now();
 
         final ListObjectsV2Request.Builder listObjectsV2Request = ListObjectsV2Request.builder()
                 .bucket(bucket)
                 .prefix(s3Prefix);
 
         do {
+            // Renew ownership every 3 minutes during long listing operations
+            if (Duration.between(lastOwnershipRenewal, Instant.now()).compareTo(OWNERSHIP_RENEWAL_INTERVAL) >= 0) {
+                sourceCoordinator.renewGlobalStateOwnershipForPartitionCreation();
+                lastOwnershipRenewal = Instant.now();
+            }
+
             listObjectsV2Response = s3Client.listObjectsV2(listObjectsV2Request
                     .fetchOwner(true)
                     .continuationToken(Objects.nonNull(listObjectsV2Response) ? listObjectsV2Response.nextContinuationToken() : null)
@@ -379,8 +388,13 @@ private void processObjectsForFolderPartition(final List<S3ObjectReference> obje
         int objectIndex = 0;
         String activeAcknowledgmentSetId = null;
         AcknowledgementSet acknowledgementSet = null;
+        Instant lastOwnershipRenewal = Instant.now();
 
         while (objectIndex < objectsToProcess.size() && objectsProcessed < folderPartitioningOptions.getMaxObjectsPerOwnership()) {
+            if (Duration.between(lastOwnershipRenewal, Instant.now()).compareTo(OWNERSHIP_RENEWAL_INTERVAL) >= 0) {
+                sourceCoordinator.renewGlobalStateOwnershipForPartitionCreation();
+                lastOwnershipRenewal = Instant.now();
+            }
             final S3ObjectReference s3ObjectReference = objectsToProcess.get(objectIndex);
             if (objectsProcessed % MAX_OBJECTS_PER_ACKNOWLEDGMENT_SET == 0) {
                 if (acknowledgementSet != null) {