Signed-off-by: Divyansh Bokadia <dbokadia@amazon.com> (opensearch-project#5621)

Adding retries and backoff to handle delayed sync between GSI and primary table when DDB is used Source Coordination Store

Author: divbok

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

@@ -130,6 +130,17 @@ public interface SourceCoordinator<T> {
     void giveUpPartition(final String partitionKey, final Instant priorityTimestamp);
 
 
+    /**
+     * Should be called by the source when it is shutting down to indicate that it will no longer be able to perform work on partitions,
+     * or can be called to give up ownership of its partitions in order to pick up new ones with {@link #getNextPartition(Function)} ()}.
+     * @param partitionKey - Key used as the partition key.
+     * @param priorityTimestamp - A timestamp that will determine the order that UNASSIGNED partitions are acquired after they are given up.
+     * @param maxRetries - The number of times to retry giving up the partition before throwing an exception
+     * @throws org.opensearch.dataprepper.model.source.coordinator.exceptions.PartitionUpdateException if the partition could not be given up due to some failure
+     * @since 2.12
+     */
+    void giveUpPartition(final String partitionKey, final Instant priorityTimestamp,final Integer maxRetries);
+
     /**
      * Should be called by the source after when acknowledgments are enabled to keep ownership of the partition for acknowledgmentTimeout amount of time
      * before another instance of Data Prepper can pick it up for processing. Allows the source to acquire another partition immediately for processing

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

@@ -351,6 +351,31 @@ public void giveUpPartition(final String partitionKey, final Instant priorityTim
         giveUpPartitionInternal(partitionKey, priorityTimestamp);
     }
 
+    @Override
+    public void giveUpPartition(final String partitionKey, final Instant priorityTimestamp,final Integer maxRetries) {
+        final long backoffMillis = 200;
+        for (int attempt = 1; attempt <= maxRetries; attempt++) {
+            try {
+                giveUpPartitionInternal(partitionKey, priorityTimestamp);
+                return;
+            } catch (PartitionNotOwnedException e) {
+                if (attempt == maxRetries) {
+                    throw e;
+                }
+                LOG.warn("Partition {} not owned on attempt {}/{}. Will retry giving up the partition", partitionKey,
+                        attempt, maxRetries);
+            }
+
+            try {
+                Thread.sleep(backoffMillis);
+            } catch (InterruptedException ie) {
+                Thread.currentThread().interrupt();
+                LOG.warn("Interrupted while retrying giveUpPartition for partition {} with message {}", partitionKey, ie);
+                throw new RuntimeException(ie);
+            }
+        }
+    }
+
     private void giveUpPartitionInternal(final String partitionKey, final Instant priorityTimestamp) {
         if (!initialized) {
             return;

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

@@ -56,7 +56,7 @@ public class ScanObjectWorker implements Runnable {
 
     private static final Logger LOG = LoggerFactory.getLogger(ScanObjectWorker.class);
     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 NO_OBJECTS_FOUND_BEFORE_PARTITION_DELETION_DURATION = Duration.ofHours(1);
@@ -215,7 +215,7 @@ private void startProcessingObject(final long waitTimeMillis) {
                 processFolderPartition(objectToProcess.get());
             } catch (final Exception e) {
                 LOG.error("An exception occurred while processing folder partition {}, giving up this partition", objectToProcess.get().getPartitionKey(), e);
-                sourceCoordinator.giveUpPartition(objectToProcess.get().getPartitionKey(), Instant.now());
+                sourceCoordinator.giveUpPartition(objectToProcess.get().getPartitionKey(), Instant.now(), MAX_RETRIES);
                 partitionKeys.remove(objectToProcess.get().getPartitionKey());
             }
             return;
@@ -323,7 +323,7 @@ private void processFolderPartition(final SourcePartition<S3SourceProgressState>
                 return;
             }
             LOG.debug("No objects to process, giving up partition");
-            sourceCoordinator.giveUpPartition(folderPartition.getPartitionKey(), Instant.now());
+            sourceCoordinator.giveUpPartition(folderPartition.getPartitionKey(), Instant.now(), MAX_RETRIES);
             return;
         }
 
@@ -441,7 +441,7 @@ private AcknowledgementSet createAcknowledgmentSetForFolderPartition(final Sourc
                 objectsToDeleteForAcknowledgmentSets.remove(acknowledgmentSetId);
                 partitionKeys.remove(folderPartition.getPartitionKey());
                 LOG.info("Received all acknowledgments for folder partition {}, giving up this partition", folderPartition.getPartitionKey());
-                sourceCoordinator.giveUpPartition(folderPartition.getPartitionKey(), Instant.now());
+                sourceCoordinator.giveUpPartition(folderPartition.getPartitionKey(), Instant.now(), MAX_RETRIES);
             }
         }, acknowledgmentSetTimeout);
     }

data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanObjectWorkerTest.java

@@ -74,6 +74,7 @@
 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.CHECKPOINT_OWNERSHIP_INTERVAL;
+import static org.opensearch.dataprepper.plugins.source.s3.ScanObjectWorker.MAX_RETRIES;
 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;
 import static org.opensearch.dataprepper.plugins.source.s3.ScanObjectWorker.PARTITION_OWNERSHIP_UPDATE_ERRORS;
@@ -523,14 +524,15 @@ void processing_with_folder_partitions_with_no_objects_gives_up_that_partition()
         when(listObjectsV2Response.contents()).thenReturn(Collections.emptyList());
 
         when(s3Client.listObjectsV2(any(ListObjectsV2Request.class))).thenReturn(listObjectsV2Response);
-        doNothing().when(sourceCoordinator).giveUpPartition(eq(partitionKey), any(Instant.class));
+        doNothing().when(sourceCoordinator).giveUpPartition(eq(partitionKey), any(Instant.class),
+                eq(MAX_RETRIES));
         doNothing().when(sourceCoordinator).saveProgressStateForPartition(eq(partitionKey), any(S3SourceProgressState.class));
 
         final ScanObjectWorker scanObjectWorker = createObjectUnderTest();
         scanObjectWorker.runWithoutInfiniteLoop();
 
         verify(sourceCoordinator).saveProgressStateForPartition(eq(partitionKey), any(S3SourceProgressState.class));
-        verify(sourceCoordinator).giveUpPartition(eq(partitionKey), any(Instant.class));
+        verify(sourceCoordinator).giveUpPartition(eq(partitionKey), any(Instant.class), eq(MAX_RETRIES));
 
         final ArgumentCaptor<ListObjectsV2Request> listObjectsV2RequestArgumentCaptor = ArgumentCaptor.forClass(ListObjectsV2Request.class);
 
@@ -653,7 +655,7 @@ void processing_with_folder_partition_processes_objects_in_folder_and_deletes_th
         secondAckCallback.accept(true);
 
         inOrder.verify(s3ObjectDeleteWorker).deleteS3Object(secondObjectDeleteRequest);
-        inOrder.verify(sourceCoordinator).giveUpPartition(eq(partitionKey), any(Instant.class));
+        inOrder.verify(sourceCoordinator).giveUpPartition(eq(partitionKey), any(Instant.class), eq(MAX_RETRIES));
     }
 
     @Test
@@ -727,7 +729,7 @@ void processing_with_folder_partition_processes_objects_in_folder_until_max_obje
         ackCallback.accept(true);
 
         inOrder.verify(s3ObjectDeleteWorker).deleteS3Object(firstObjectDeleteRequest);
-        inOrder.verify(sourceCoordinator).giveUpPartition(eq(partitionKey), any(Instant.class));
+        inOrder.verify(sourceCoordinator).giveUpPartition(eq(partitionKey), any(Instant.class), any(Integer.class));
     }
 
     @Test