Ensure shards are completed when last getRecords call has no records and no shardIterator (#5958)

Signed-off-by: Jonah Calvo <caljonah@amazon.com>

Author: Jonah Calvo

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

@@ -252,10 +252,14 @@ public void run() {
         String sequenceNumber = "";
         int interval;
         List<software.amazon.awssdk.services.dynamodb.model.Record> records;
-
+        boolean createdFinalAcknowledgmentSetForShard = false;
         while (!shouldStop) {
             if (shardIterator == null) {
                 // End of Shard
+                if (shardAcknowledgementManager != null && !createdFinalAcknowledgmentSetForShard) {
+                    final AcknowledgementSet finalAcknowledgmentSet = shardAcknowledgementManager.createAcknowledgmentSet(streamPartition, sequenceNumber, true);
+                    finalAcknowledgmentSet.complete();
+                }
                 LOG.debug("Reached end of shard");
                 break;
             }
@@ -287,6 +291,9 @@ public void run() {
                 AcknowledgementSet acknowledgementSet = null;
                 if (shardAcknowledgementManager != null) {
                     acknowledgementSet = shardAcknowledgementManager.createAcknowledgmentSet(streamPartition, sequenceNumber, shardIterator == null);
+                    if (shardIterator == null) {
+                        createdFinalAcknowledgmentSetForShard = true;
+                    }
                 }
 
                 records = response.records().stream()

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

@@ -49,6 +49,7 @@
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.BDDMockito.given;
 import static org.mockito.Mockito.lenient;
 import static org.mockito.Mockito.mock;
@@ -347,6 +348,9 @@ void test_run_shardConsumer_catches_4xx_exception_and_increments_metric() {
 
     @Test
     void test_run_shardConsumer_calls_startUpdatingOwnershipForShard() throws Exception {
+        final AcknowledgementSet finalAcknowledgementSet = mock(AcknowledgementSet.class);
+        when(shardAcknowledgementManager.createAcknowledgmentSet(any(StreamPartition.class), any(String.class), any(Boolean.class)))
+            .thenReturn(finalAcknowledgementSet);
         try (final MockedStatic<BufferAccumulator> bufferAccumulatorMockedStatic = mockStatic(BufferAccumulator.class)) {
             bufferAccumulatorMockedStatic.when(() -> BufferAccumulator.create(buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT)).thenReturn(bufferAccumulator);
             ShardConsumer shardConsumer = ShardConsumer.builder(dynamoDbStreamsClient, pluginMetrics, aggregateMetrics, buffer, streamConfig)
@@ -364,6 +368,94 @@ void test_run_shardConsumer_calls_startUpdatingOwnershipForShard() throws Except
         verify(shardAcknowledgementManager).startUpdatingOwnershipForShard(streamPartition);
     }
 
+    @Test
+    void test_shard_has_records_null_iterator() throws Exception {
+        final AcknowledgementSet finalAcknowledgementSet = mock(AcknowledgementSet.class);
+        when(shardAcknowledgementManager.createAcknowledgmentSet(any(StreamPartition.class), any(String.class), any(Boolean.class)))
+                .thenReturn(finalAcknowledgementSet);
+
+        // Set up response with null nextShardIterator to trigger end of shard
+        GetRecordsResponse response = GetRecordsResponse.builder()
+                .records(buildRecords(1))
+                .nextShardIterator(null)
+                .build();
+        when(dynamoDbStreamsClient.getRecords(any(GetRecordsRequest.class))).thenReturn(response);
+
+        try (MockedStatic<ShardConsumer> shardConsumerMockedStatic = mockStatic(ShardConsumer.class, invocation -> {
+            if (invocation.getMethod().getName().equals("stopAll")) {
+                return null;
+            } else if (invocation.getMethod().getName().equals("shouldStop")) {
+                return false;
+            }
+            return invocation.callRealMethod();
+        })) {
+            ShardConsumer shardConsumer;
+            try (final MockedStatic<BufferAccumulator> bufferAccumulatorMockedStatic = mockStatic(BufferAccumulator.class)) {
+                bufferAccumulatorMockedStatic.when(() -> BufferAccumulator.create(buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT)).thenReturn(bufferAccumulator);
+                shardConsumer = ShardConsumer.builder(dynamoDbStreamsClient, pluginMetrics, aggregateMetrics, buffer, streamConfig)
+                        .shardIterator(shardIterator)
+                        .shardAcknowledgementManager(shardAcknowledgementManager)
+                        .streamPartition(streamPartition)
+                        .tableInfo(tableInfo)
+                        .startTime(null)
+                        .waitForExport(false)
+                        .build();
+            }
+
+            shardConsumer.run();
+
+            // Verify acknowledgment set created for records with shardIterator == null (true)
+            verify(shardAcknowledgementManager).createAcknowledgmentSet(eq(streamPartition), any(String.class), eq(true));
+            // Verify final acknowledgment set created and completed when shardIterator is null
+            verify(finalAcknowledgementSet).complete();
+
+        }
+    }
+
+
+    @Test
+    void test_shard_has_no_records_null_iterator() throws Exception {
+        final AcknowledgementSet finalAcknowledgementSet = mock(AcknowledgementSet.class);
+        when(shardAcknowledgementManager.createAcknowledgmentSet(any(StreamPartition.class), any(String.class), any(Boolean.class)))
+            .thenReturn(finalAcknowledgementSet);
+
+        // Set up response with null nextShardIterator to trigger end of shard
+        GetRecordsResponse response = GetRecordsResponse.builder()
+            .records(List.of())
+            .nextShardIterator(null)
+            .build();
+        when(dynamoDbStreamsClient.getRecords(any(GetRecordsRequest.class))).thenReturn(response);
+
+        try (MockedStatic<ShardConsumer> shardConsumerMockedStatic = mockStatic(ShardConsumer.class, invocation -> {
+            if (invocation.getMethod().getName().equals("stopAll")) {
+                return null;
+            } else if (invocation.getMethod().getName().equals("shouldStop")) {
+                return false;
+            }
+            return invocation.callRealMethod();
+        })) {
+            ShardConsumer shardConsumer;
+            try (final MockedStatic<BufferAccumulator> bufferAccumulatorMockedStatic = mockStatic(BufferAccumulator.class)) {
+                bufferAccumulatorMockedStatic.when(() -> BufferAccumulator.create(buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT)).thenReturn(bufferAccumulator);
+                shardConsumer = ShardConsumer.builder(dynamoDbStreamsClient, pluginMetrics, aggregateMetrics, buffer, streamConfig)
+                    .shardIterator(shardIterator)
+                    .shardAcknowledgementManager(shardAcknowledgementManager)
+                    .streamPartition(streamPartition)
+                    .tableInfo(tableInfo)
+                    .startTime(null)
+                    .waitForExport(false)
+                    .build();
+            }
+
+            shardConsumer.run();
+
+            // Verify acknowledgment set created for records with shardIterator == null (true)
+            verify(shardAcknowledgementManager).createAcknowledgmentSet(eq(streamPartition), any(String.class), eq(true));
+            // Verify final acknowledgment set created and completed when shardIterator is null
+            verify(finalAcknowledgementSet).complete();
+
+        }
+    }
     private List<Record> buildRecords(int count) {
         List<Record> records = new ArrayList<>();
         for (int i = 0; i < count; i++) {