feat(map): Fix map replay, wait-inside-map, and concurrency race conditions by ayushiahjolia · Pull Request #229 · aws/aws-durable-execution-sdk-java

ayushiahjolia · 2026-03-18T19:45:31Z

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

Issue Link, if available

#39

Description

Fix BaseConcurrentOperation to properly handle suspension when branches call wait(), add synchronized concurrency control, and move parent checkpoint from get() into the last-completing-branch callback.
Convert MapResult/MapResultItem to records with ErrorObject for checkpoint-replay durability. Remove throws Exception from MapFunction.
Added more integration tests.

Demo/Screenshots

Checklist

I have filled out every section of the PR template
I have thoroughly tested this change

Testing

Unit Tests

Have unit tests been written for these changes? N/A

Integration Tests

Have integration tests been written for these changes? Yes

Examples

Has a new example been added for the change? (if applicable) Yes

…itions

zhongkechen · 2026-03-18T23:41:43Z

unit tests failed:

[ERROR] Errors: 
[ERROR]   ParallelOperationTest.branchCreation_multipleBranchesAllCreated:106 » IllegalState Cannot add items to a completed operation

This is due to the race condition introduced in the following change:

-        // Block until operation completes. No-op if the future is already completed.
-        completionFuture.join();
+        // Block until operation completes or execution suspends.
+        // Using runUntilCompleteOrSuspend races completionFuture against executionExceptionFuture,
+        // so when all active threads suspend (e.g., wait inside map branches), the
+        // SuspendExecutionException propagates and this thread is freed — preventing thread leaks
+        // on shared executor pools across invocations.
+        executionManager.runUntilCompleteOrSuspend(completionFuture).join();

Removing this change should fix the tests

ayushiahjolia · 2026-03-19T04:06:37Z

When all map branches call wait(), every branch thread deregisters and suspendExecution() fires. But the parent thread sitting on completionFuture.join() in BaseConcurrentOperation.get() is stuck forever - no branch will ever complete to trigger onChildContextComplete and finalize the map's completionFuture. runUntilCompleteOrSuspend races completionFuture against executionExceptionFuture, so when suspension fires, the parent thread gets freed via SuspendExecutionException instead of blocking indefinitely.

zhongkechen · 2026-03-19T04:15:11Z

sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java

+        // so when all active threads suspend (e.g., wait inside map branches), the
+        // SuspendExecutionException propagates and this thread is freed — preventing thread leaks
+        // on shared executor pools across invocations.
+        executionManager.runUntilCompleteOrSuspend(completionFuture).join();


This is not the correct behavior. This will unblock all the threads waiting for completion of this operation when suspension.

…itions

zhongkechen · 2026-03-19T05:16:09Z

sdk/src/main/java/software/amazon/lambda/durable/operation/BaseConcurrentOperation.java

    public R get() {
-        var op = waitForOperationCompletion();
+        var executionManager = getContext().getExecutionManager();
+        var threadContext = getCurrentThreadContext();


why don't we just call waitForOpeartionCompletion here?

zhongkechen · 2026-03-19T05:17:24Z

sdk/src/main/java/software/amazon/lambda/durable/serde/JacksonSerDes.java

    public JacksonSerDes() {
        this(new ObjectMapper()
                .registerModule(new JavaTimeModule())
+                .registerModule(new AwsSdkV2Module())


Do not add this module to user's serDes. Why do we need this?

zhongkechen · 2026-03-19T05:20:25Z

sdk/src/main/java/software/amazon/lambda/durable/operation/BaseConcurrentOperation.java

        sendOperationUpdateAsync(
                OperationUpdate.builder().action(OperationAction.START).subType(subType.getValue()));
+        var parentThreadContext = getCurrentThreadContext();
        this.rootContext = getContext().createChildContext(getOperationId(), getName());


Why do we initialize here?

zhongkechen · 2026-03-19T05:21:30Z

sdk/src/main/java/software/amazon/lambda/durable/operation/BaseConcurrentOperation.java


    protected CompletionReason getCompletionReason() {
-        return completionReason;
+        synchronized (lock) {


Why do we need a lock to access a field? Same below

zhongkechen · 2026-03-19T05:24:03Z

sdk/src/main/java/software/amazon/lambda/durable/operation/BaseConcurrentOperation.java

+        } else if (op.status() == OperationStatus.STARTED) {
+            // All branches suspended (e.g., wait inside map branches) — propagate suspension.
+            // onChildContextSuspended completed completionFuture when activeBranches hit 0.
+            throw new SuspendExecutionException();


When completionFuture is completed, the status can't be STARTED

ayushiahjolia self-assigned this Mar 18, 2026

ayushiahjolia force-pushed the map_bug_fixes branch from cae0595 to 00210be Compare March 18, 2026 19:57

feat(map): Fix map replay, wait-inside-map, and concurrency race cond…

14a215e

…itions

ayushiahjolia force-pushed the map_bug_fixes branch from 00210be to 14a215e Compare March 18, 2026 21:22

Merge branch 'main' into map_bug_fixes

fde6a1b

ayushiahjolia marked this pull request as ready for review March 18, 2026 22:08

ayushiahjolia requested a review from a team March 18, 2026 22:08

zhongkechen force-pushed the map_bug_fixes branch from fde6a1b to 3c9974e Compare March 18, 2026 23:41

ayushiahjolia force-pushed the map_bug_fixes branch from 3c9974e to bb907d9 Compare March 19, 2026 03:35

ayushiahjolia force-pushed the map_bug_fixes branch from bb907d9 to 33e5548 Compare March 19, 2026 04:07

zhongkechen requested changes Mar 19, 2026

View reviewed changes

feat(map): Fix map replay, wait-inside-map, and concurrency race cond…

d7339bc

…itions

ayushiahjolia force-pushed the map_bug_fixes branch from 72e4167 to d7339bc Compare March 19, 2026 04:42

Merge branch 'main' into map_bug_fixes

977217c

zhongkechen requested changes Mar 19, 2026

View reviewed changes

zhongkechen reviewed Mar 19, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat(map): Fix map replay, wait-inside-map, and concurrency race conditions#229

feat(map): Fix map replay, wait-inside-map, and concurrency race conditions#229
ayushiahjolia wants to merge 4 commits intomainfrom
map_bug_fixes

ayushiahjolia commented Mar 18, 2026 •

edited

Loading

Uh oh!

zhongkechen commented Mar 18, 2026 •

edited

Loading

Uh oh!

ayushiahjolia commented Mar 19, 2026 •

edited

Loading

Uh oh!

zhongkechen Mar 19, 2026

Uh oh!

zhongkechen Mar 19, 2026

Uh oh!

zhongkechen Mar 19, 2026

Uh oh!

zhongkechen Mar 19, 2026

Uh oh!

zhongkechen Mar 19, 2026

Uh oh!

zhongkechen Mar 19, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

ayushiahjolia commented Mar 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Issue Link, if available

Description

Demo/Screenshots

Checklist

Testing

Unit Tests

Integration Tests

Examples

Uh oh!

zhongkechen commented Mar 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ayushiahjolia commented Mar 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

zhongkechen Mar 19, 2026

Choose a reason for hiding this comment

Uh oh!

zhongkechen Mar 19, 2026

Choose a reason for hiding this comment

Uh oh!

zhongkechen Mar 19, 2026

Choose a reason for hiding this comment

Uh oh!

zhongkechen Mar 19, 2026

Choose a reason for hiding this comment

Uh oh!

zhongkechen Mar 19, 2026

Choose a reason for hiding this comment

Uh oh!

zhongkechen Mar 19, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

ayushiahjolia commented Mar 18, 2026 •

edited

Loading

zhongkechen commented Mar 18, 2026 •

edited

Loading

ayushiahjolia commented Mar 19, 2026 •

edited

Loading