prevent a race condition in future completion

zhongkechen · zhongkechen · commit c6214b34a527 · 2026-03-23T14:52:37.000-07:00
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutionManager.java b/sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutionManager.java
@@ -191,12 +191,11 @@ public ThreadContext getCurrentThreadContext() {
      * @see ThreadContext
      */
     public void registerActiveThread(String threadId) {
-        if (activeThreads.contains(threadId)) {
+        if (activeThreads.add(threadId)) {
+            logger.trace("Registered thread '{}' as active. Active threads: {}", threadId, activeThreads.size());
+        } else {
             logger.trace("Thread '{}' already registered as active", threadId);
-            return;
         }
-        activeThreads.add(threadId);
-        logger.trace("Registered thread '{}' as active. Active threads: {}", threadId, activeThreads.size());
     }
 
     /**
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java
@@ -6,6 +6,7 @@
 import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -179,23 +180,45 @@ protected Operation waitForOperationCompletion() {
         // 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
         // is between `isOperationCompleted` and `thenRun`.
-        synchronized (completionFuture) {
-            if (!isOperationCompleted()) {
-                // Operation not done yet
-                logger.trace(
-                        "deregistering thread {} when waiting for operation {} ({}) to complete ({})",
-                        threadContext.threadId(),
-                        getOperation(),
-                        getType(),
-                        completionFuture);
-
-                // 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()));
-
-                // Deregister the current thread to allow suspension
-                executionManager.deregisterActiveThread(threadContext.threadId());
+        AtomicBoolean skipDeregistration = new AtomicBoolean(false);
+        long callerThreadId = Thread.currentThread().getId();
+
+        // 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(() -> {
+            // at this point, future.isDone = true
+            if (Thread.currentThread().getId() == callerThreadId) {
+                // If the completion thread is the same as the consumer thread (immediately completed),
+                // we don't want to deregister and register the thread.
+                skipDeregistration.set(true);
+                return;
+            }
+
+            synchronized (skipDeregistration) {
+                if (!skipDeregistration.get()) {
+                    registerActiveThread(threadContext.threadId());
+                }
+            }
+        });
+
+        // skip deregistration if future is always completed or skipDeregistration is true
+        if (!isOperationCompleted() && !completionFuture.isDone() && !skipDeregistration.get()) {
+            synchronized (skipDeregistration) {
+                if (!skipDeregistration.get() && !completionFuture.isDone()) {
+                    // Operation not done yet
+                    logger.trace(
+                            "deregistering thread {} when waiting for operation {} ({}) to complete ({})",
+                            threadContext.threadId(),
+                            getOperation(),
+                            getType(),
+                            completionFuture);
+
+                    // Deregister the current thread to allow suspension
+                    executionManager.deregisterActiveThread(threadContext.threadId());
+                } else {
+                    skipDeregistration.set(true);
+                }
             }
         }
 
@@ -282,13 +305,11 @@ protected void markAlreadyCompleted() {
     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);
-        }
+        // 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);
     }
 
     /**
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ConcurrencyOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ConcurrencyOperation.java
@@ -82,7 +82,7 @@ protected ConcurrencyOperation(
         this.toleratedFailureCount = toleratedFailureCount;
         this.operationIdGenerator = new OperationIdGenerator(getOperationId());
         this.rootContext = durableContext.createChildContext(getOperationId(), getName());
-        this.consumerThreadListener = new AtomicReference<>(null);
+        this.consumerThreadListener = new AtomicReference<>(new CompletableFuture<>());
     }
 
     // ========== Template methods for subclasses ==========
@@ -138,15 +138,13 @@ protected <R> ChildContextOperation<R> enqueueItem(
         pendingQueue.add(childOp);
         logger.debug("Item enqueued {}", name);
         // notify the consumer thread a new item is available
-        completeVacancyListenerIfSet();
+        notifyConsumerThread();
         return childOp;
     }
 
-    private void completeVacancyListenerIfSet() {
-        synchronized (this) {
-            if (consumerThreadListener.get() != null) {
-                consumerThreadListener.get().complete(null);
-            }
+    private void notifyConsumerThread() {
+        synchronized (consumerThreadListener) {
+            consumerThreadListener.get().complete(null);
         }
     }
 
@@ -159,6 +157,17 @@ protected void executeItems() {
 
         Runnable consumer = () -> {
             while (true) {
+                // Set a new future if it's completed so that it will be able to receive a notification of
+                // new items when the thread is checking completion condition and processing
+                // the queued items below.
+                synchronized (consumerThreadListener) {
+                    if (consumerThreadListener.get() != null
+                            && consumerThreadListener.get().isDone()) {
+                        consumerThreadListener.set(new CompletableFuture<>());
+                    }
+                }
+
+                // Process completion condition. Quit the loop if the condition is met.
                 if (isOperationCompleted()) {
                     return;
                 }
@@ -167,14 +176,21 @@ protected void executeItems() {
                     handleComplete(completionStatus);
                     return;
                 }
+
+                // process new items in the queue
                 while (runningChildren.size() < maxConcurrency && !pendingQueue.isEmpty()) {
                     var next = pendingQueue.poll();
                     runningChildren.add(next);
                     logger.debug("Executing operation {}", next.getName());
                     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 consumerThreadListener has been completed when processing above, waitForChildCompletion will
+                // immediately return null and repeat the above again
+                var child = waitForChildCompletion(runningChildren);
+
+                // child may be null if the consumer thread is woken up due to new items added or completion condition
+                // changed
                 if (child != null) {
                     if (runningChildren.contains(child)) {
                         runningChildren.remove(child);
@@ -183,49 +199,67 @@ protected void executeItems() {
                         throw new IllegalStateException("Unexpected completion: " + child);
                     }
                 }
-                synchronized (this) {
-                    if (consumerThreadListener.get() != null
-                            && consumerThreadListener.get().isDone()) {
-                        consumerThreadListener.set(null);
-                    }
-                }
             }
         };
         // run consumer in the user thread pool, although it's not a real user thread
         runUserHandler(consumer, getOperationId(), ThreadType.CONTEXT);
     }
 
-    private BaseDurableOperation waitForChildCompletion(
-            AtomicInteger succeededCount, AtomicInteger failedCount, Set<BaseDurableOperation> runningChildren) {
+    private BaseDurableOperation waitForChildCompletion(Set<BaseDurableOperation> runningChildren) {
         var threadContext = getCurrentThreadContext();
-        CompletableFuture<Object> future;
+        ArrayList<CompletableFuture<BaseDurableOperation>> futures = new ArrayList<>(runningChildren.stream()
+                .map(BaseDurableOperation::getCompletionFuture)
+                .toList());
 
-        synchronized (this) {
-            // check again in synchronized block to prevent race conditions
-            if (isOperationCompleted()) {
-                return null;
-            }
-            var completionStatus = canComplete(succeededCount, failedCount, runningChildren);
-            if (completionStatus != null) {
-                return null;
+        // always add the future to listen to the new items or condition changes. This might have been
+        // completed during the period of consuming items from the queue.
+        futures.add(consumerThreadListener.get());
+
+        // future will be completed immediately if any future of the list is already completed
+        CompletableFuture<Object> future = CompletableFuture.anyOf(futures.toArray(CompletableFuture[]::new));
+
+        // skip deregistering the current thread if it is already completed
+        AtomicBoolean skipDeregistration = new AtomicBoolean(false);
+        long callerThreadId = Thread.currentThread().getId();
+
+        // attach the stage for registering thread without checking future status to avoid race condition
+        future.thenRun(() -> {
+            // at this point, future.isDone = true
+            if (Thread.currentThread().getId() == callerThreadId) {
+                // If the completion thread is the same as the consumer thread (immediately completed),
+                // we don't want to deregister and register the thread.
+                skipDeregistration.set(true);
+                return;
             }
-            ArrayList<CompletableFuture<BaseDurableOperation>> futures;
-            futures = new ArrayList<>(runningChildren.stream()
-                    .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());
+
+            synchronized (skipDeregistration) {
+                if (!skipDeregistration.get()) {
+                    registerActiveThread(threadContext.threadId());
+                }
             }
+        });
+
+        // skip deregistration if future is always completed or skipDeregistration is true
+        if (!future.isDone() && !skipDeregistration.get()) {
+            // We tried best to skip deregistration/registration using thread id check above. However,
+            // there is still a small window where the completionFuture is completed right after the check
+            // and before the deregisterActiveThread call.
 
-            // future will be completed immediately if any future of the list is already completed
-            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()));
-                // Deregister the current thread to allow suspension
-                executionManager.deregisterActiveThread(threadContext.threadId());
+            // Registering and deregistering threads might be competing for this lock.
+            // If this deregistration thread skips deregistering, skipDeregistration will be set to true
+            // so that the registration thread will skip registration
+            synchronized (skipDeregistration) {
+                if (!skipDeregistration.get() && !future.isDone()) {
+                    // Operation not done yet
+                    logger.trace(
+                            "deregistering consumer thread {} when waiting for {} branches to complete",
+                            threadContext.threadId(),
+                            getType());
+                    // Deregister the current thread to allow suspension
+                    executionManager.deregisterActiveThread(threadContext.threadId());
+                } else {
+                    skipDeregistration.set(true);
+                }
             }
         }
         return future.thenApply(o -> (BaseDurableOperation) o).join();
@@ -273,6 +307,8 @@ private ConcurrencyCompletionStatus canComplete(
         }
 
         // All items finished — complete
+        // This condition relies on isJoined, so the consumer will wake up and check this again when
+        // isJoined is set to true.
         if (isJoined.get() && pendingQueue.isEmpty() && runningChildren.isEmpty()) {
             return ConcurrencyCompletionStatus.ALL_COMPLETED;
         }
@@ -281,12 +317,10 @@ private ConcurrencyCompletionStatus canComplete(
     }
 
     private void handleComplete(ConcurrencyCompletionStatus status) {
-        synchronized (this) {
-            if (isOperationCompleted()) {
-                return;
-            }
-            handleSuccess(status);
+        if (isOperationCompleted()) {
+            return;
         }
+        handleSuccess(status);
     }
 
     /**
@@ -299,8 +333,10 @@ protected void join() {
                     + ") exceeds the number of registered items (" + branches.size() + ")");
         }
         isJoined.set(true);
-        // notify the execution thread this concurrency operation is joined
-        completeVacancyListenerIfSet();
+
+        // Notify the consumer thread this concurrency operation is joined. Consumer thread need to check the
+        // completion condition again.
+        notifyConsumerThread();
         waitForOperationCompletion();
     }
 
