Add thread-safe synchronization to startUpdatingOwnershipForShard

divbok · divbok · commit f35146091b8f · 2026-01-23T12:06:39.000-06:00
Signed-off-by: Divyansh Bokadia &lt;dbokadia@amazon.com&gt;
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardAcknowledgementManager.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardAcknowledgementManager.java
@@ -320,7 +320,12 @@ public boolean isExportDone(StreamPartition streamPartition) {
     }
 
     void startUpdatingOwnershipForShard(final StreamPartition streamPartition) {
-        checkpoints.computeIfAbsent(streamPartition, segment -> new ConcurrentLinkedQueue<>());
+        lock.lock();
+        try {
+            checkpoints.computeIfAbsent(streamPartition, segment -> new ConcurrentLinkedQueue<>());
+        } finally {
+            lock.unlock();
+        }
     }
 
     boolean isStillTrackingShard(final StreamPartition streamPartition) {

Original file line number	Diff line number	Diff line change
`@@ -320,7 +320,12 @@ public boolean isExportDone(StreamPartition streamPartition) {`
`320`	`320`	`}`
`321`	`321`
`322`	`322`	`void startUpdatingOwnershipForShard(final StreamPartition streamPartition) {`
`323`		`- checkpoints.computeIfAbsent(streamPartition, segment -> new ConcurrentLinkedQueue<>());`
	`323`	`+ lock.lock();`
	`324`	`+ try {`
	`325`	`+ checkpoints.computeIfAbsent(streamPartition, segment -> new ConcurrentLinkedQueue<>());`
	`326`	`+ } finally {`
	`327`	`+ lock.unlock();`
	`328`	`+ }`
`324`	`329`	`}`
`325`	`330`
`326`	`331`	`boolean isStillTrackingShard(final StreamPartition streamPartition) {`