fix thread synchronization

Author: zhongkechen

sdk-integration-tests/pom.xml

@@ -41,6 +41,12 @@
             <artifactId>junit-jupiter</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.slf4j</groupId>
+            <artifactId>slf4j-simple</artifactId>
+            <version>${slf4j.version}</version>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>

sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ParallelIntegrationTest.java

@@ -109,16 +109,15 @@ void testParallelPartialFailure_failedBranchDoesNotPreventOthers() {
     void testParallelAllBranchesFail() {
         var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
             var config = ParallelConfig.builder().build();
-            var futures = new ArrayList<DurableFuture<String>>();
             var parallel = context.parallel("all-fail", config);
 
             try (parallel) {
-                futures.add(parallel.branch("branch-x", String.class, ctx -> {
+                parallel.branch("branch-x", String.class, ctx -> {
                     throw new RuntimeException("fail-x");
-                }));
-                futures.add(parallel.branch("branch-y", String.class, ctx -> {
+                });
+                parallel.branch("branch-y", String.class, ctx -> {
                     throw new RuntimeException("fail-y");
-                }));
+                });
             }
 
             var result = parallel.get();
@@ -298,7 +297,7 @@ void testStepBeforeAndAfterParallel() {
 
             var config = ParallelConfig.builder().build();
             var futures = new ArrayList<DurableFuture<String>>();
-            var parallel = context.parallel("middle-parallel", config);
+            ParallelDurableFuture parallel = context.parallel("middle-parallel", config);
 
             try (parallel) {
                 futures.add(parallel.branch("branch-a", String.class, ctx -> "A"));
@@ -451,19 +450,18 @@ void testParallelUnlimitedConcurrencyWithToleratedFailureCount() {
             var config = ParallelConfig.builder()
                     .completionConfig(CompletionConfig.toleratedFailureCount(1))
                     .build();
-            var futures = new ArrayList<DurableFuture<String>>();
-            var parallel = context.parallel("unlimited-tolerated", config);
+            ParallelDurableFuture parallel = context.parallel("unlimited-tolerated", config);
 
             try (parallel) {
-                futures.add(parallel.branch("branch-ok1", String.class, ctx -> "OK1"));
-                futures.add(parallel.branch("branch-fail1", String.class, ctx -> {
+                parallel.branch("branch-ok1", String.class, ctx -> "OK1");
+                parallel.branch("branch-fail1", String.class, ctx -> {
                     throw new RuntimeException("failed: fail1");
-                }));
-                futures.add(parallel.branch("branch-ok2", String.class, ctx -> "OK2"));
-                futures.add(parallel.branch("branch-fail2", String.class, ctx -> {
+                });
+                parallel.branch("branch-ok2", String.class, ctx -> "OK2");
+                parallel.branch("branch-fail2", String.class, ctx -> {
                     throw new RuntimeException("failed: fail2");
-                }));
-                futures.add(parallel.branch("branch-ok3", String.class, ctx -> "OK3"));
+                });
+                parallel.branch("branch-ok3", String.class, ctx -> "OK3");
             }
 
             var result = parallel.get();
@@ -508,7 +506,7 @@ void testParallelBranchesReturnDifferentTypes() {
     void testParallelResultSummary_succeededAndFailedCounts() {
         var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
             var config = ParallelConfig.builder().build();
-            var parallel = context.parallel("count-check", config);
+            ParallelDurableFuture parallel = context.parallel("count-check", config);
 
             try (parallel) {
                 parallel.branch("ok1", String.class, ctx -> "OK1");
@@ -595,15 +593,14 @@ void testParallel50BranchesWithWaitForCallback_maxConcurrency5() {
 
         var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
             var config = ParallelConfig.builder().maxConcurrency(5).build();
-            var futures = new ArrayList<DurableFuture<String>>();
             var parallel = context.parallel("50-callbacks-limited", config);
 
             try (parallel) {
                 for (int i = 0; i < branchCount; i++) {
                     var idx = i;
-                    futures.add(parallel.branch("branch-" + i, String.class, ctx -> {
+                    parallel.branch("branch-" + i, String.class, ctx -> {
                         return ctx.waitForCallback("cb-" + idx, String.class, (callbackId, stepCtx) -> {});
-                    }));
+                    });
                 }
             }
 
@@ -644,15 +641,14 @@ void testParallel50BranchesWithWaitForCallback_partialFailure() {
 
         var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
             var config = ParallelConfig.builder().build();
-            var futures = new ArrayList<DurableFuture<String>>();
             var parallel = context.parallel("50-callbacks-partial-fail", config);
 
             try (parallel) {
                 for (int i = 0; i < branchCount; i++) {
                     var idx = i;
-                    futures.add(parallel.branch("branch-" + i, String.class, ctx -> {
+                    parallel.branch("branch-" + i, String.class, ctx -> {
                         return ctx.waitForCallback("approval-" + idx, String.class, (callbackId, stepCtx) -> {});
-                    }));
+                    });
                 }
             }
 
@@ -697,18 +693,17 @@ void testParallel50BranchesWithWaitForCallback_stepsBeforeAndAfterCallback() {
 
         var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
             var config = ParallelConfig.builder().build();
-            var futures = new ArrayList<DurableFuture<String>>();
-            var parallel = context.parallel("50-callbacks-with-steps", config);
+            ParallelDurableFuture parallel = context.parallel("50-callbacks-with-steps", config);
 
             try (parallel) {
                 for (int i = 0; i < branchCount; i++) {
                     var idx = i;
-                    futures.add(parallel.branch("branch-" + i, String.class, ctx -> {
+                    parallel.branch("branch-" + i, String.class, ctx -> {
                         var before = ctx.step("prepare-" + idx, String.class, stepCtx -> "prepared-" + idx);
                         var approval =
                                 ctx.waitForCallback("approval-" + idx, String.class, (callbackId, stepCtx) -> {});
                         return ctx.step("finalize-" + idx, String.class, stepCtx -> before + ":" + approval + ":done");
-                    }));
+                    });
                 }
             }
 
@@ -890,20 +885,19 @@ void testParallel50BranchesMixed_callbackAndCondition() {
 
         var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
             var config = ParallelConfig.builder().build();
-            var futures = new ArrayList<DurableFuture<String>>();
             var parallel = context.parallel("50-mixed", config);
 
             try (parallel) {
                 for (int i = 0; i < branchCount; i++) {
                     var idx = i;
                     if (i % 2 == 0) {
                         // Even branches: waitForCallback
-                        futures.add(parallel.branch("branch-" + i, String.class, ctx -> {
+                        parallel.branch("branch-" + i, String.class, ctx -> {
                             return ctx.waitForCallback("cb-" + idx, String.class, (callbackId, stepCtx) -> {});
-                        }));
+                        });
                     } else {
                         // Odd branches: waitForCondition
-                        futures.add(parallel.branch("branch-" + i, String.class, ctx -> {
+                        parallel.branch("branch-" + i, String.class, ctx -> {
                             var strategy = WaitStrategies.<Integer>fixedDelay(5, Duration.ofSeconds(1));
                             var wfcConfig = WaitForConditionConfig.<Integer>builder()
                                     .waitStrategy(strategy)
@@ -916,7 +910,7 @@ void testParallel50BranchesMixed_callbackAndCondition() {
                                     wfcConfig);
 
                             return "polled-" + polled;
-                        }));
+                        });
                     }
                 }
             }

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

@@ -182,6 +182,7 @@ protected Operation waitForOperationCompletion() {
         validateCurrentThreadType();
 
         var threadContext = getCurrentThreadContext();
+        CompletableFuture<?> future = completionFuture;
 
         // It's important that we synchronize access to the future. Otherwise, a race condition could happen if the
         // completionFuture is completed by a user thread (a step or child context thread) when the execution here
@@ -201,7 +202,17 @@ protected Operation waitForOperationCompletion() {
                 // Add a completion stage to completionFuture so that when the completionFuture is completed,
                 // it will register the current Context thread synchronously to make sure it is always registered
                 // strictly before the execution thread (Step or child context) is deregistered.
-                completionFuture.thenRun(() -> registerActiveThread(threadContext.threadId()));
+                // chain them together
+                future = completionFuture.thenRun(() -> {
+                    logger.warn(
+                            "registering thread {} when operation {} ({}) completed ({})",
+                            threadContext.threadId(),
+                            getOperation(),
+                            getType(),
+                            completionFuture);
+
+                    registerActiveThread(threadContext.threadId());
+                });
 
                 // Deregister the current thread to allow suspension
                 executionManager.deregisterActiveThread(threadContext.threadId());
@@ -210,7 +221,7 @@ protected Operation waitForOperationCompletion() {
 
         // Block until operation completes. No-op if the future is already completed.
         try {
-            completionFuture.join();
+            future.join();
         } catch (Throwable throwable) {
             ExceptionHelper.sneakyThrow(ExceptionHelper.unwrapCompletableFuture(throwable));
         }
@@ -243,6 +254,7 @@ protected void runUserHandler(Runnable runnable, ThreadType threadType) {
             } finally {
                 if (operationId != null) {
                     try {
+                        logger.trace("deregistering thread {} after running user handler {}", operationId, getName());
                         // if this is a child context or a step context, we need to
                         // deregister the context's thread from the execution manager
                         executionManager.deregisterActiveThread(operationId);
@@ -271,6 +283,7 @@ protected void runUserHandler(Runnable runnable, ThreadType threadType) {
         // 2. setCurrentContext on the CHILD thread — sets the ThreadLocal so operations inside
         //    the child context know which context they belong to.
         // registerActiveThread is idempotent (no-op if already registered).
+        logger.trace("registering thread {} before running user handler {}", operationId, getName());
         registerActiveThread(operationId);
 
         runningUserHandler.set(CompletableFuture.runAsync(

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

@@ -255,7 +255,10 @@ private BaseDurableOperation waitForChildCompletion(
             future = CompletableFuture.anyOf(futures.toArray(CompletableFuture[]::new));
             // skip deregistering the current thread if there is more completed future to process
             if (!future.isDone()) {
-                future.thenRun(() -> registerActiveThread(threadContext.threadId()));
+                future = future.thenApply(o -> {
+                    registerActiveThread(threadContext.threadId());
+                    return o;
+                });
                 // Deregister the current thread to allow suspension
                 executionManager.deregisterActiveThread(threadContext.threadId());
             }