Fix how failures are handled for shard partitions in DDB source

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

@@ -10,11 +10,13 @@
 
 package org.opensearch.dataprepper.plugins.source.dynamodb.stream;
 
+import com.google.common.annotations.VisibleForTesting;
 import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
 import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager;
 import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourceCoordinator;
 import org.opensearch.dataprepper.common.concurrent.BackgroundThreadFactory;
 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.coordination.partition.StreamPartition;
 import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.state.StreamProgressState;
@@ -39,14 +41,16 @@
 
 import org.opensearch.dataprepper.plugins.source.dynamodb.model.ShardCheckpointStatus;
 
+import static org.opensearch.dataprepper.logging.DataPrepperMarkers.NOISY;
+
 public class ShardAcknowledgementManager {
     private static final Logger LOG = LoggerFactory.getLogger(ShardAcknowledgementManager.class);
 
     private static final String NULL_SEQUENCE_NUMBER = "null";
 
     private static final long WAIT_FOR_ACKNOWLEDGMENTS_TIMEOUT = 10L;
 
-    static final Duration CHECKPOINT_INTERVAL = Duration.ofMinutes(2);
+    static final Duration CHECKPOINT_INTERVAL = Duration.ofMinutes(3);
 
     private final DynamoDBSourceConfig dynamoDBSourceConfig;
     private final Map<StreamPartition, ConcurrentLinkedQueue<ShardCheckpointStatus>> checkpoints = new ConcurrentHashMap<>();
@@ -79,6 +83,20 @@ public ShardAcknowledgementManager(final AcknowledgementSetManager acknowledgeme
         executorService.submit(() -> monitorAcknowledgments(stopWorkerConsumer));
     }
 
+    @VisibleForTesting
+    public ShardAcknowledgementManager(final AcknowledgementSetManager acknowledgementSetManager,
+    final EnhancedSourceCoordinator sourceCoordinator,
+    final DynamoDBSourceConfig dynamoDBSourceConfig,
+                                       final ExecutorService executorService) {
+        this.executorService = executorService;
+        this.acknowledgementSetManager = acknowledgementSetManager;
+        this.sourceCoordinator = sourceCoordinator;
+        this.dynamoDBSourceConfig = dynamoDBSourceConfig;
+        this.partitionsToRemove = Collections.synchronizedList(new ArrayList<>());
+        this.partitionsToGiveUp = Collections.synchronizedList(new ArrayList<>());
+        this.lastCheckpointTime = Instant.now();
+    }
+
     void monitorAcknowledgments(final Consumer<StreamPartition> stopWorkerConsumer) {
         while (!Thread.currentThread().isInterrupted()) {
             boolean exit = runMonitorAcknowledgmentLoop(stopWorkerConsumer);
@@ -153,13 +171,9 @@ boolean runMonitorAcknowledgmentLoop(final Consumer<StreamPartition> stopWorkerC
                     sourceCoordinator.saveProgressStateForPartition(streamPartition, dynamoDBSourceConfig.getShardAcknowledgmentTimeout());
                     LOG.debug("Checkpointed shard {} with latest sequence number acknowledged {}", streamPartition.getShardId(), latestCheckpointForShard.getSequenceNumber());
                 }
-                if (partitionsToGiveUp.contains(streamPartition)) {
-                    partitionsToRemove.add(streamPartition);
-                    sourceCoordinator.giveUpPartition(streamPartition);
-                }
-
             } catch (final Exception e) {
-                LOG.error("Received exception while monitoring acknowledgments for stream partition {}", streamPartition.getPartitionKey(), e);
+                LOG.error(NOISY, "Received exception while monitoring acknowledgments for stream partition {}, stop processing shard", streamPartition.getPartitionKey(), e);
+                markPartitionForRemoval(streamPartition);
             }
         }
 
@@ -172,6 +186,13 @@ public AcknowledgementSet createAcknowledgmentSet(
         final boolean isFinalSetForPartition) {
         final String sequenceNumberNoNull = sequenceNumber == null ? NULL_SEQUENCE_NUMBER : sequenceNumber;
         final ShardCheckpointStatus shardCheckpointStatus = new ShardCheckpointStatus(sequenceNumber, Instant.now().toEpochMilli(), isFinalSetForPartition);
+
+        // 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)) {
+            throw new ShardNotTrackedException("The shard {} is not being tracked anymore, stop reading from shard");
+        }
         checkpoints.computeIfAbsent(streamPartition, segment -> new ConcurrentLinkedQueue<>()).add(shardCheckpointStatus);
         ackStatuses.computeIfAbsent(streamPartition, segment -> new ConcurrentHashMap<>());
         ackStatuses.get(streamPartition).put(sequenceNumberNoNull, shardCheckpointStatus);
@@ -199,7 +220,12 @@ void updateOwnershipForAllShardPartitions() {
         if (Duration.between(lastCheckpointTime, Instant.now()).compareTo(CHECKPOINT_INTERVAL) > 0) {
             for (final StreamPartition streamPartition : checkpoints.keySet()) {
                 if (!partitionsToRemove.contains(streamPartition)) {
-                    sourceCoordinator.saveProgressStateForPartition(streamPartition, dynamoDBSourceConfig.getShardAcknowledgmentTimeout());
+                    try {
+                        sourceCoordinator.saveProgressStateForPartition(streamPartition, dynamoDBSourceConfig.getShardAcknowledgmentTimeout());
+                    } catch (final PartitionUpdateException e) {
+                        LOG.warn(NOISY, "Failed to update progress state for shard {}, will stop tracking this shard as someone else owns it", streamPartition.getShardId());
+                        markPartitionForRemoval(streamPartition);
+                    }
                 }
             }
 
@@ -214,9 +240,8 @@ private void handleFailure(final StreamPartition streamPartition,
             streamProgressState.setSequenceNumber(latestCheckpointForShard.getSequenceNumber());
             sourceCoordinator.saveProgressStateForPartition(streamPartition, dynamoDBSourceConfig.getShardAcknowledgmentTimeout());
         }
-        partitionsToRemove.add(streamPartition);
-        sourceCoordinator.giveUpPartition(streamPartition);
-        partitionsToGiveUp.remove(streamPartition);
+
+        markPartitionForRemoval(streamPartition);
     }
 
     private void handleCompletedShard(final StreamPartition streamPartition) {
@@ -245,13 +270,17 @@ private void removePartitions() {
             ackStatuses.remove(streamPartition);
         });
 
+        partitionsToGiveUp.forEach(sourceCoordinator::giveUpPartition);
+
         partitionsToRemove.clear();
+        partitionsToGiveUp.clear();
     }
 
     public void giveUpPartition(final StreamPartition streamPartition) {
-        if (!partitionsToGiveUp.contains(streamPartition)) {
+        if (isStillTrackingShard(streamPartition)) {
             LOG.debug("Adding partition {} to give up list", streamPartition.getPartitionKey());
             partitionsToGiveUp.add(streamPartition);
+            partitionsToRemove.add(streamPartition);
         }
     }
 
@@ -260,7 +289,24 @@ public boolean isExportDone(StreamPartition streamPartition) {
         return globalPartition.isPresent();
     }
 
-    public void startUpdatingOwnershipForShard(final StreamPartition streamPartition) {
+    void startUpdatingOwnershipForShard(final StreamPartition streamPartition) {
         checkpoints.computeIfAbsent(streamPartition, segment -> new ConcurrentLinkedQueue<>());
     }
+
+    boolean isStillTrackingShard(final StreamPartition streamPartition) {
+        return isStillTrackingShardInternal(streamPartition);
+    }
+
+    private boolean isStillTrackingShardInternal(final StreamPartition streamPartition) {
+        LOG.info("partitions to remove: {}", partitionsToRemove);
+        LOG.info("Checkpoints: {}", checkpoints);
+        return !partitionsToRemove.contains(streamPartition) && checkpoints.containsKey(streamPartition);
+    }
+
+    private void markPartitionForRemoval(final StreamPartition streamPartition) {
+        if (!partitionsToRemove.contains(streamPartition)) {
+            partitionsToRemove.add(streamPartition);
+            partitionsToGiveUp.add(streamPartition);
+        }
+    }
 }

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

@@ -272,6 +272,11 @@ public void run() {
                 lastCheckpointTime = System.currentTimeMillis();
             }
 
+            if (shardAcknowledgementManager != null && !shardAcknowledgementManager.isStillTrackingShard(streamPartition)) {
+                LOG.warn("Shard {} is no longer being tracked by the acknowledgment manager, exiting", streamPartition.getShardId());
+                break;
+            }
+
             GetRecordsResponse response = callGetRecords(shardIterator);
             shardIterator = response.nextShardIterator();
             if (!response.records().isEmpty()) {

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

@@ -0,0 +1,12 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.plugins.source.dynamodb.stream;
+
+public class ShardNotTrackedException extends RuntimeException {
+    public ShardNotTrackedException(String message) {
+      super(message);
+    }
+}