Fix no acknowledgments for closed shard issue (opensearch-project#3651)

Signed-off-by: Aiden Dai <[email protected]>
asifsmohammed · Nov 14, 2023 · 3322c8a · 3322c8a
1 parent 8922a0f
commit 3322c8a
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/.../main/java/org/opensearch/dataprepper/plugins/source/dynamodb/leader/LeaderScheduler.java b/.../main/java/org/opensearch/dataprepper/plugins/source/dynamodb/leader/LeaderScheduler.java
@@ -201,7 +201,7 @@ private void compareAndCreateChildrenPartitions(List<EnhancedSourcePartition> so
         sourcePartitions.forEach(sourcePartition -> {
             StreamPartition streamPartition = (StreamPartition) sourcePartition;
             List<String> childShardIds = shardManager.findChildShardIds(streamPartition.getStreamArn(), streamPartition.getShardId());
-            if (!childShardIds.isEmpty()) {
+            if (childShardIds != null && !childShardIds.isEmpty()) {
                 childShardIds.forEach(
                         shardId -> {
                             if (!completedShardIds.contains(shardId)) {

diff --git a/...rc/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumer.java b/...rc/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumer.java
@@ -198,7 +198,13 @@ public ShardConsumer build() {
     public void run() {
         LOG.debug("Shard Consumer start to run...");
         // Check should skip processing or not.
-        if (shouldSkip()) return;
+        if (shouldSkip()) {
+            if (acknowledgementSet != null) {
+                checkpointer.updateShardForAcknowledgmentWait(shardAcknowledgmentTimeout);
+                acknowledgementSet.complete();
+            }
+            return;
+        }
 
         long lastCheckpointTime = System.currentTimeMillis();
         String sequenceNumber = "";