Remove the each time update in ScanObjectWorker; Add the lastGlobalOwnershipRenewal related renewGlobalOwnershipIfNeeded to ensures ownership is maintained both before starting and throughout the entire partition creation process

Author: Davidding4718

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

@@ -96,6 +96,7 @@ public class LeaseBasedSourceCoordinator<T> implements SourceCoordinator<T> {
     private final ReentrantLock lock;
 
     private Instant lastSupplierRunTime;
+    private Instant lastGlobalOwnershipRenewal;
 
     static final Duration FORCE_SUPPLIER_AFTER_DURATION = Duration.ofMinutes(5);
 
@@ -138,6 +139,7 @@ public LeaseBasedSourceCoordinator(final Class<T> partitionProgressStateClass,
         this.partitionsDeleted = pluginMetrics.counter(PARTITIONS_DELETED);
         this.lock = new ReentrantLock();
         this.lastSupplierRunTime = Instant.now();
+        this.lastGlobalOwnershipRenewal = Instant.now();
     }
 
     @Override
@@ -206,16 +208,10 @@ private Optional<SourcePartition<T>> getNextPartitionInternal(final Function<Map
 
     @Override
     public void createPartitions(final List<PartitionIdentifier> partitionIdentifiers) {
-        Instant lastOwnershipRenewal = Instant.now();
+        renewGlobalOwnershipIfNeeded();
 
         for (final PartitionIdentifier partitionIdentifier : partitionIdentifiers) {
-            if (Instant.now().isAfter(lastOwnershipRenewal.plus(OWNERSHIP_RENEWAL_INTERVAL))) {
-                LOG.info("Renewing global state ownership for partition creation");
-                renewGlobalStateOwnershipForPartitionCreation();
-                lastOwnershipRenewal = Instant.now();
-            }
-            
-
+            renewGlobalOwnershipIfNeeded();
 
             final Optional<SourcePartitionStoreItem> optionalPartitionItem = sourceCoordinationStore.getSourcePartitionItem(sourceIdentifierWithPartitionType, partitionIdentifier.getPartitionKey());
 
@@ -542,6 +538,14 @@ public void renewGlobalStateOwnershipForPartitionCreation() {
         }
     }
 
+    private void renewGlobalOwnershipIfNeeded() {
+        if (Instant.now().isAfter(lastGlobalOwnershipRenewal.plus(OWNERSHIP_RENEWAL_INTERVAL))) {
+            LOG.info("Renewing global state ownership for partition creation");
+            renewGlobalStateOwnershipForPartitionCreation();
+            lastGlobalOwnershipRenewal = Instant.now();
+        }
+    }
+
     private void giveUpAndSaveGlobalStateForPartitionCreation(final SourcePartitionStoreItem globalStateItemForPartitionCreation, final Map<String, Object> globalStateMap) {
 
         globalStateItemForPartitionCreation.setSourcePartitionStatus(SourcePartitionStatus.UNASSIGNED);

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

@@ -58,7 +58,6 @@ 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;
@@ -306,7 +305,6 @@ 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();
@@ -339,19 +337,12 @@ 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)
@@ -388,13 +379,8 @@ 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) {