avoid marking branch done when parent is done

Author: zhongkechen

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

@@ -266,15 +266,8 @@ public void onCheckpointComplete(Operation operation) {
             // This method handles only terminal status updates. Override this method if a DurableOperation needs to
             // handle other updates.
             logger.trace("In onCheckpointComplete, completing operation {} ({})", getOperationId(), completionFuture);
-            // It's important that we synchronize access to the future, otherwise the processing could happen
-            // on someone else's thread and cause a race condition.
-            synchronized (completionFuture) {
-                // Completing the future here will also run any other completion stages that have been attached
-                // to the future. In our case, other contexts may have attached a function to reactivate themselves,
-                // so they will definitely have a chance to reactivate before we finish completing and deactivating
-                // whatever operations were just checkpointed.
-                completionFuture.complete(this);
-            }
+
+            markCompletionFutureCompleted();
         }
     }
 
@@ -283,10 +276,17 @@ protected void markAlreadyCompleted() {
         // When the operation is already completed in a replay, we complete completionFuture immediately
         // so that the `get` method will be unblocked and the context thread will be registered
         logger.trace("In markAlreadyCompleted, completing operation: {} ({}).", getOperationId(), completionFuture);
+        markCompletionFutureCompleted();
+    }
 
+    private void markCompletionFutureCompleted() {
         // It's important that we synchronize access to the future, otherwise the processing could happen
         // on someone else's thread and cause a race condition.
         synchronized (completionFuture) {
+            // Completing the future here will also run any other completion stages that have been attached
+            // to the future. In our case, other contexts may have attached a function to reactivate themselves,
+            // so they will definitely have a chance to reactivate before we finish completing and deactivating
+            // whatever operations were just checkpointed.
             completionFuture.complete(this);
         }
     }

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

@@ -5,6 +5,7 @@
 import static software.amazon.lambda.durable.execution.ExecutionManager.isTerminalStatus;
 
 import java.nio.charset.StandardCharsets;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
 import software.amazon.awssdk.services.lambda.model.ContextOptions;
 import software.amazon.awssdk.services.lambda.model.ErrorObject;
@@ -46,7 +47,7 @@ public class ChildContextOperation<T> extends SerializableDurableOperation<T> {
 
     private final Function<DurableContext, T> function;
     private final ConcurrencyOperation<?> parentOperation;
-    private boolean replayChildContext;
+    private final AtomicBoolean replayChildren = new AtomicBoolean(false);
     private T reconstructedResult;
 
     public ChildContextOperation(
@@ -86,7 +87,7 @@ protected void replay(Operation existing) {
                 if (existing.contextDetails() != null
                         && Boolean.TRUE.equals(existing.contextDetails().replayChildren())) {
                     // Large result: re-execute child context to reconstruct result
-                    replayChildContext = true;
+                    replayChildren.set(true);
                     executeChildContext();
                 } else {
                     markAlreadyCompleted();
@@ -100,11 +101,6 @@ protected void replay(Operation existing) {
         }
     }
 
-    @Override
-    protected void markAlreadyCompleted() {
-        super.markAlreadyCompleted();
-    }
-
     private void executeChildContext() {
         // The operationId is already globally unique (prefixed by parent context path via
         // DurableContext.nextOperationId), so we use it directly as the contextId.
@@ -137,7 +133,7 @@ private void executeChildContext() {
     }
 
     private void handleChildContextSuccess(T result) {
-        if (replayChildContext) {
+        if (replayChildren.get()) {
             // Replaying a SUCCEEDED child with replayChildren=true — skip checkpointing.
             // Mark the completableFuture completed so get() doesn't block waiting for a checkpoint response.
             this.reconstructedResult = result;
@@ -151,8 +147,6 @@ private void checkpointSuccess(T result) {
         // Skip checkpointing if parent ConcurrencyOperation has already completed —
         // prevents race conditions where a child finishes after the parent has already completed.
         if (parentOperation != null && parentOperation.isOperationCompleted()) {
-            this.reconstructedResult = result;
-            markAlreadyCompleted();
             return;
         }
 
@@ -187,7 +181,6 @@ private void handleChildContextFailure(Throwable exception) {
         // Skip checkpointing if parent ConcurrencyOperation has already completed —
         // prevents race conditions where a child finishes after the parent has already succeeded.
         if (parentOperation != null && parentOperation.isOperationCompleted()) {
-            markAlreadyCompleted();
             return;
         }
 

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

@@ -174,6 +174,7 @@ protected void executeItems() {
                     next.execute();
                 }
                 var child = waitForChildCompletion(succeededCount, failedCount, runningChildren);
+                // child may be null if the consumer thread is woken up due to a new item being added
                 if (child != null) {
                     if (runningChildren.contains(child)) {
                         runningChildren.remove(child);
@@ -213,6 +214,7 @@ private BaseDurableOperation waitForChildCompletion(
                     .map(BaseDurableOperation::getCompletionFuture)
                     .toList());
             if (futures.size() < maxConcurrency) {
+                // add a future to listen to the new items if there is a vacancy
                 consumerThreadListener.compareAndSet(null, new CompletableFuture<>());
                 futures.add(consumerThreadListener.get());
             }