Handle DynamoDB source leader exceptions correctly by attempting to reacquire partition (#6195)

Signed-off-by: Taylor Gray <tylgry@amazon.com>

Author: graytaylor0

data-prepper-plugins/dynamodb-source-coordination-store/src/main/java/org/opensearch/dataprepper/plugins/sourcecoordinator/dynamodb/DynamoDbClientWrapper.java

@@ -183,8 +183,8 @@ private void tryUpdateItem(final DynamoDbSourcePartitionItem dynamoDbSourceParti
         } catch (final ConditionalCheckFailedException e) {
             final String message = String.format(
                     "ConditionalCheckFailed while updating partition %s. This partition item was either deleted from the dynamo table, " +
-                            "or another instance of Data Prepper has modified it.",
-                    dynamoDbSourcePartitionItem.getSourcePartitionKey());
+                            "or another instance of Data Prepper has modified it. Expected version: %s",
+                    dynamoDbSourcePartitionItem.getSourcePartitionKey(), dynamoDbSourcePartitionItem.getVersion() - 1L);
             throw new PartitionUpdateException(message, e);
         } catch (final Exception e) {
             final String errorMessage = String.format("An exception occurred while attempting to update a DynamoDb partition item %s",

data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/leader/LeaderScheduler.java

@@ -129,7 +129,8 @@ public void run() {
                     try {
                         coordinator.saveProgressStateForPartition(leaderPartition, Duration.ofMinutes(DEFAULT_EXTEND_LEASE_MINUTES));
                     } catch (final Exception e) {
-                        LOG.error("Failed to update ownership for leader partition. Retrying...");
+                        LOG.error("Failed to update ownership for leader partition. Will attempt to reacquire this partition...");
+                        leaderPartition = null;
                     }
                 }
                 try {

data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/leader/LeaderSchedulerTest.java

@@ -8,6 +8,7 @@
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourceCoordinator;
 import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourcePartition;
+import org.opensearch.dataprepper.model.source.coordinator.exceptions.PartitionUpdateException;
 import org.opensearch.dataprepper.plugins.source.dynamodb.DynamoDBSourceConfig;
 import org.opensearch.dataprepper.plugins.source.dynamodb.configuration.ExportConfig;
 import org.opensearch.dataprepper.plugins.source.dynamodb.configuration.StreamConfig;
@@ -16,6 +17,7 @@
 import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.LeaderPartition;
 import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.StreamPartition;
 import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.state.StreamProgressState;
+import org.opensearch.dataprepper.test.helper.ReflectivelySetField;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.dynamodb.model.ContinuousBackupsDescription;
 import software.amazon.awssdk.services.dynamodb.model.ContinuousBackupsStatus;
@@ -55,6 +57,7 @@
 import static org.mockito.ArgumentMatchers.isNull;
 import static org.mockito.BDDMockito.given;
 import static org.mockito.Mockito.atLeast;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.lenient;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
@@ -330,6 +333,35 @@ void run_without_acquiring_leader_partition_does_not_save_null_state() {
         verify(coordinator, never()).saveProgressStateForPartition(isNull(), any(Duration.class));
     }
 
+    @Test
+    void test_shardDiscovery_with_failure_to_save_partition_state_reacquires_partition() throws InterruptedException, NoSuchFieldException, IllegalAccessException {
+        leaderScheduler = new LeaderScheduler(coordinator, dynamoDbClient, shardManager, List.of(tableConfig));
+        leaderPartition = new LeaderPartition();
+        leaderPartition.getProgressState().get().setInitialized(true);
+        leaderPartition.getProgressState().get().setStreamArns(List.of(streamArn));
+        given(coordinator.acquireAvailablePartition(LeaderPartition.PARTITION_TYPE)).willReturn(Optional.of(leaderPartition));
+        doThrow(PartitionUpdateException.class).when(coordinator).saveProgressStateForPartition(eq(leaderPartition), any(Duration.class));
+
+        ReflectivelySetField.setField(LeaderScheduler.class, leaderScheduler, "leaseInterval", Duration.ofMillis(40));
+
+        ExecutorService executorService = Executors.newSingleThreadExecutor();
+        executorService.submit(() -> leaderScheduler.run());
+
+        Thread.sleep(100);
+        executorService.shutdownNow();
+        // Already init
+        verifyNoInteractions(dynamoDbClient);
+
+        // Should check the completed partitions
+        verify(coordinator, atLeast(2)).queryCompletedPartitions(eq(StreamPartition.PARTITION_TYPE), any(Instant.class));
+
+        // Should create 3 stream partitions for child shards found
+        verify(coordinator, atLeast(3)).createPartition(any(EnhancedSourcePartition.class));
+
+        verify(coordinator, atLeast(2)).saveProgressStateForPartition(eq(leaderPartition), any(Duration.class));
+        verify(coordinator, atLeast(2)).acquireAvailablePartition(LeaderPartition.PARTITION_TYPE);
+    }
+
 
     /**
      * Helper function to mock DescribeContinuousBackupsResponse