Fix issues with acknowledgment timeouts and data loss with checkpointing acknowledgments

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

Author: graytaylor0

data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/coordination/partition/StreamPartition.java

@@ -9,6 +9,7 @@
 import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourcePartition;
 import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.state.StreamProgressState;
 
+import java.util.Objects;
 import java.util.Optional;
 
 public class StreamPartition extends EnhancedSourcePartition<StreamProgressState> {
@@ -61,4 +62,16 @@ public String getStreamArn() {
     public String getShardId() {
         return shardId;
     }
+
+    @Override
+    public boolean equals(Object o) {
+        if (o == null || getClass() != o.getClass()) return false;
+        StreamPartition that = (StreamPartition) o;
+        return Objects.equals(shardId, that.shardId);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hashCode(shardId);
+    }
 }

data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardAcknowledgementManager.java

@@ -85,6 +85,12 @@ void monitorAcknowledgments(final Consumer<StreamPartition> stopWorkerConsumer)
             if (exit) {
                 break;
             }
+
+            try {
+                Thread.sleep(2000);
+            } catch (InterruptedException e) {
+                throw new RuntimeException(e);
+            }
         }
 
         LOG.info("Exiting acknowledgment manager");
@@ -111,6 +117,9 @@ boolean runMonitorAcknowledgmentLoop(final Consumer<StreamPartition> stopWorkerC
 
                     if (checkpointStatuses.peek().isPositiveAcknowledgement()) {
                         latestCheckpointForShard = checkpointStatuses.poll();
+                        if (latestCheckpointForShard != null) {
+                            ackStatuses.get(streamPartition).remove(latestCheckpointForShard.getSequenceNumber());
+                        }
                     } else if (checkpointStatuses.peek().isNegativeAcknowledgement()
                         || checkpointStatuses.peek().isExpired(dynamoDBSourceConfig.getShardAcknowledgmentTimeout())) {
                         handleFailure(streamPartition, streamProgressState, latestCheckpointForShard);

data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumer.java

@@ -257,7 +257,7 @@ public void run() {
             if (shardIterator == null) {
                 // End of Shard
                 if (shardAcknowledgementManager != null && !createdFinalAcknowledgmentSetForShard) {
-                    final AcknowledgementSet finalAcknowledgmentSet = shardAcknowledgementManager.createAcknowledgmentSet(streamPartition, sequenceNumber, true);
+                    final AcknowledgementSet finalAcknowledgmentSet = shardAcknowledgementManager.createAcknowledgmentSet(streamPartition, "END_OF_SHARD", true);
                     finalAcknowledgmentSet.complete();
                 }
                 LOG.debug("Reached end of shard");

data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactory.java

@@ -142,11 +142,10 @@ public String getShardIterator(String streamArn, String shardId, String sequence
 
         if (sequenceNumber != null && !sequenceNumber.isEmpty()) {
             LOG.debug("Get Shard Iterator at {}", sequenceNumber);
-            // There may be an overlap for 1 record
             getShardIteratorRequest = GetShardIteratorRequest.builder()
                     .shardId(shardId)
                     .streamArn(streamArn)
-                    .shardIteratorType(ShardIteratorType.AT_SEQUENCE_NUMBER)
+                    .shardIteratorType(ShardIteratorType.AFTER_SEQUENCE_NUMBER)
                     .sequenceNumber(sequenceNumber)
                     .build();
         } else {

data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactoryTest.java

@@ -13,7 +13,6 @@
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.opensearch.dataprepper.metrics.PluginMetrics;
-import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
 import org.opensearch.dataprepper.model.buffer.Buffer;
 import org.opensearch.dataprepper.model.event.Event;
 import org.opensearch.dataprepper.model.record.Record;
@@ -131,8 +130,6 @@ public void test_create_shardConsumer_correctly() {
 
     @Test
     public void test_create_shardConsumer_correctly_with_is_disable_checkpointing_enabled_starts_from_trim_horizon() {
-
-        final AcknowledgementSet acknowledgementSet = mock(AcknowledgementSet.class);
         when(streamConfig.isDisableCheckpointing()).thenReturn(true);
         StreamProgressState state = new StreamProgressState();
         state.setWaitForExport(false);
@@ -141,7 +138,7 @@ public void test_create_shardConsumer_correctly_with_is_disable_checkpointing_en
         streamPartition = new StreamPartition(streamArn, shardId, Optional.of(state));
 
         ShardConsumerFactory consumerFactory = new ShardConsumerFactory(coordinator, dynamoDbStreamsClient, pluginMetrics, dynamoDBSourceAggregateMetrics, buffer, streamConfig);
-        Runnable consumer = consumerFactory.createConsumer(streamPartition, acknowledgementSet, null);
+        Runnable consumer = consumerFactory.createConsumer(streamPartition, Duration.ofMinutes(1), mock(ShardAcknowledgementManager.class));
         assertThat(consumer, notNullValue());
 
         final ArgumentCaptor<GetShardIteratorRequest> captor = ArgumentCaptor.forClass(GetShardIteratorRequest.class);