prevent a race condition in future completion

Author: zhongkechen

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,21 +180,31 @@ 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()));
+        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.
+            AtomicBoolean alreadyCompleted = new AtomicBoolean(false);
+            long callerThreadId = Thread.currentThread().getId();
+            completionFuture.thenRun(() -> {
+                if (Thread.currentThread().getId() == callerThreadId) {
+                    // If the calling thread is same as the thread that initiated the wait, meaning the future is
+                    // already completed, we should skip deregister/register the thread.
+                    alreadyCompleted.set(true);
+                    return;
+                }
+                registerActiveThread(threadContext.threadId());
+            });
 
+            if (!alreadyCompleted.get()) {
                 // Deregister the current thread to allow suspension
                 executionManager.deregisterActiveThread(threadContext.threadId());
             }
@@ -266,8 +277,11 @@ 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);
-
-            markCompletionFutureCompleted();
+            // 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);
         }
     }
 
@@ -276,19 +290,11 @@ 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);
-        }
+        // 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/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() {
+    private void notifyConsumerThread() {
         synchronized (this) {
-            if (consumerThreadListener.get() != null) {
-                consumerThreadListener.get().complete(null);
-            }
+            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 (this) {
+                    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,47 +199,39 @@ 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;
-            }
-            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());
-            }
+        // 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));
 
-            // 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()));
+        // skip deregistering the current thread if there is more completed future to process
+        AtomicBoolean futureCompletedImmediately = new AtomicBoolean(false);
+        long callerThreadId = Thread.currentThread().getId();
+        if (!future.isDone()) {
+            future.thenRun(() -> {
+                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.
+                    futureCompletedImmediately.set(true);
+                    return;
+                }
+                registerActiveThread(threadContext.threadId());
+            });
+            if (!futureCompletedImmediately.get()) {
                 // Deregister the current thread to allow suspension
                 executionManager.deregisterActiveThread(threadContext.threadId());
             }
@@ -273,6 +281,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;
         }
@@ -299,8 +309,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();
     }