Change to not create queue in createAcknowledgmentSet to prevent race condition

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

Author: graytaylor0

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

@@ -190,10 +190,12 @@ public AcknowledgementSet createAcknowledgmentSet(
         // Shard should already be in checkpoints map from call to startUpdatingOwnershipForShard, if it is not in the map
         // that means that ShardAcknowledgmentManager stopped tracking it due to some error, and another worker will pick it up
         // We throw an error in this case to have the ShardConsumer exit and stop reading data from the shard
-        if (!isStillTrackingShard(streamPartition)) {
+        ConcurrentLinkedQueue<ShardCheckpointStatus> queue = checkpoints.get(streamPartition);
+        if (queue == null) {
             throw new ShardNotTrackedException("The shard {} is not being tracked anymore, stop reading from shard");
         }
-        checkpoints.computeIfAbsent(streamPartition, segment -> new ConcurrentLinkedQueue<>()).add(shardCheckpointStatus);
+        queue.add(shardCheckpointStatus);
+
         ackStatuses.computeIfAbsent(streamPartition, segment -> new ConcurrentHashMap<>());
         ackStatuses.get(streamPartition).put(sequenceNumberNoNull, shardCheckpointStatus);
 
@@ -217,6 +219,7 @@ public AcknowledgementSet createAcknowledgmentSet(
     }
 
     void updateOwnershipForAllShardPartitions() {
+
         if (Duration.between(lastCheckpointTime, Instant.now()).compareTo(CHECKPOINT_INTERVAL) > 0) {
             for (final StreamPartition streamPartition : checkpoints.keySet()) {
                 if (!partitionsToRemove.contains(streamPartition)) {

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

@@ -31,6 +31,8 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.equalTo;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.ArgumentMatchers.same;
+import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.lenient;
 import static org.mockito.Mockito.mock;
@@ -282,7 +284,7 @@ void create_acknowledgment_set_throws_ShardNotTrackedException_if_shard_is_not_b
     }
 
     @Test
-    void failure_to_update_state_for_one_partition_stops_tracking_that_partition_and_continues_on_others() {
+    void failure_to_update_state_for_one_partition_stops_tracking_that_partition_and_continues_on_others() throws InterruptedException {
 
         final ShardAcknowledgementManager shardAcknowledgementManager = createObjectUnderTest();
         final StreamPartition secondPartition = mock(StreamPartition.class);
@@ -311,14 +313,13 @@ void failure_to_update_state_for_one_partition_stops_tracking_that_partition_and
         callback1.accept(true);  // Positive ack
         callback2.accept(true);
 
-        doThrow(PartitionUpdateException.class).when(sourceCoordinator).saveProgressStateForPartition(eq(streamPartition),
-                        any(Duration.class));
+//        doThrow(PartitionUpdateException.class).when(sourceCoordinator).saveProgressStateForPartition(eq(streamPartition),
+//                        any(Duration.class));
+        doThrow(new PartitionUpdateException("Failed update", mock(Throwable.class))).when(sourceCoordinator)
+                .saveProgressStateForPartition(same(streamPartition), any(Duration.class));
 
-
-        // Run the acknowledgment loop
-        assertThat(shardAcknowledgementManager.isStillTrackingShard(streamPartition), equalTo(true));
-        assertThat(shardAcknowledgementManager.isStillTrackingShard(secondPartition), equalTo(false));
-        shardAcknowledgementManager.runMonitorAcknowledgmentLoop(stopWorkerConsumer);
+        doNothing().when(sourceCoordinator)
+                        .saveProgressStateForPartition(same(secondPartition), any(Duration.class));
 
         shardAcknowledgementManager.startUpdatingOwnershipForShard(secondPartition);
         // Create first acknowledgment set and capture its callback
@@ -331,15 +332,21 @@ void failure_to_update_state_for_one_partition_stops_tracking_that_partition_and
         callback1_2.accept(true);
         callback2_2.accept(true);
 
+        Thread.sleep(1000);
+
+        // Run the acknowledgment loop
+        assertThat(shardAcknowledgementManager.isStillTrackingShard(streamPartition), equalTo(true));
+        assertThat(shardAcknowledgementManager.isStillTrackingShard(secondPartition), equalTo(true));
+        shardAcknowledgementManager.runMonitorAcknowledgmentLoop(stopWorkerConsumer);
+
         assertThat(shardAcknowledgementManager.isStillTrackingShard(streamPartition), equalTo(false));
         assertThat(shardAcknowledgementManager.isStillTrackingShard(secondPartition), equalTo(true));
         shardAcknowledgementManager.runMonitorAcknowledgmentLoop(stopWorkerConsumer);
         assertThat(shardAcknowledgementManager.isStillTrackingShard(secondPartition), equalTo(true));
 
-        final InOrder inOrder = Mockito.inOrder(sourceCoordinator, secondStreamProgressState, streamProgressState);
-        inOrder.verify(streamProgressState).setSequenceNumber("seq2");
-        inOrder.verify(sourceCoordinator).giveUpPartition(streamPartition);
-        inOrder.verify(secondStreamProgressState).setSequenceNumber("seq124");
-        inOrder.verify(sourceCoordinator).saveProgressStateForPartition(secondPartition, dynamoDBSourceConfig.getShardAcknowledgmentTimeout());
+        verify(streamProgressState).setSequenceNumber("seq2");
+        verify(sourceCoordinator).giveUpPartition(streamPartition);
+        verify(secondStreamProgressState).setSequenceNumber("seq124");
+        verify(sourceCoordinator).saveProgressStateForPartition(secondPartition, dynamoDBSourceConfig.getShardAcknowledgmentTimeout());
     }
 }