prevent a race condition in concurrent operation (aws#250)

Author: zhongkechen

sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutionManager.java

@@ -191,12 +191,13 @@ public ThreadContext getCurrentThreadContext() {
      * @see ThreadContext
      */
     public void registerActiveThread(String threadId) {
-        if (activeThreads.contains(threadId)) {
-            logger.trace("Thread '{}' already registered as active", threadId);
-            return;
+        synchronized (activeThreads) {
+            if (activeThreads.add(threadId)) {
+                logger.trace("Registered thread '{}' as active. Active threads: {}", threadId, activeThreads.size());
+            } else {
+                logger.warn("Thread '{}' already registered as active", threadId);
+            }
         }
-        activeThreads.add(threadId);
-        logger.trace("Registered thread '{}' as active. Active threads: {}", threadId, activeThreads.size());
     }
 
     /**
@@ -210,16 +211,20 @@ public void deregisterActiveThread(String threadId) {
             return;
         }
 
-        boolean removed = activeThreads.remove(threadId);
-        if (removed) {
-            logger.trace("Deregistered thread '{}' Active threads: {}", threadId, activeThreads.size());
-        } else {
-            logger.warn("Thread '{}' not active, cannot deregister", threadId);
-        }
+        // Add synchronized block to avoid remove then check race condition and make sure that
+        // the suspendExecution is called only once
+        synchronized (activeThreads) {
+            boolean removed = activeThreads.remove(threadId);
+            if (removed) {
+                logger.trace("Deregistered thread '{}' Active threads: {}", threadId, activeThreads.size());
+            } else {
+                logger.warn("Thread '{}' not active, cannot deregister", threadId);
+            }
 
-        if (activeThreads.isEmpty()) {
-            logger.info("No active threads remaining - suspending execution");
-            suspendExecution();
+            if (activeThreads.isEmpty()) {
+                logger.info("No active threads remaining - suspending execution");
+                suspendExecution();
+            }
         }
     }
 

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

@@ -47,6 +47,7 @@ public abstract class BaseDurableOperation {
     private final OperationIdentifier operationIdentifier;
     protected final ExecutionManager executionManager;
     protected final CompletableFuture<BaseDurableOperation> completionFuture;
+    protected final BaseDurableOperation parentOperation;
     private final DurableContextImpl durableContext;
     private final AtomicReference<CompletableFuture<Void>> runningUserHandler = new AtomicReference<>(null);
 
@@ -55,9 +56,14 @@ public abstract class BaseDurableOperation {
      *
      * @param operationIdentifier the unique identifier for this operation
      * @param durableContext the parent context this operation belongs to
+     * @param parentOperation the parent operation if this is a branch/iteration of a ConcurrencyOperation
      */
-    protected BaseDurableOperation(OperationIdentifier operationIdentifier, DurableContextImpl durableContext) {
+    protected BaseDurableOperation(
+            OperationIdentifier operationIdentifier,
+            DurableContextImpl durableContext,
+            BaseDurableOperation parentOperation) {
         this.operationIdentifier = operationIdentifier;
+        this.parentOperation = parentOperation;
         this.durableContext = durableContext;
         this.executionManager = durableContext.getExecutionManager();
 
@@ -179,7 +185,9 @@ 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 this operation is a branch/iteration of a ConcurrencyOperation (map or parallel), the branches/iterations
+        // must be completed sequentially to avoid race conditions.
+        synchronized (parentOperation == null ? completionFuture : parentOperation) {
             if (!isOperationCompleted()) {
                 // Operation not done yet
                 logger.trace(
@@ -282,7 +290,7 @@ 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) {
+        synchronized (parentOperation == null ? completionFuture : parentOperation) {
             // 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

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

@@ -46,7 +46,6 @@ public class ChildContextOperation<T> extends SerializableDurableOperation<T> {
     private static final int LARGE_RESULT_THRESHOLD = 256 * 1024;
 
     private final Function<DurableContext, T> function;
-    private final ConcurrencyOperation<?> parentOperation;
     private final AtomicBoolean replayChildren = new AtomicBoolean(false);
     private T reconstructedResult;
 
@@ -66,9 +65,8 @@ public ChildContextOperation(
             RunInChildContextConfig config,
             DurableContextImpl durableContext,
             ConcurrencyOperation<?> parentOperation) {
-        super(operationIdentifier, resultTypeToken, config.serDes(), durableContext);
+        super(operationIdentifier, resultTypeToken, config.serDes(), durableContext, parentOperation);
         this.function = function;
-        this.parentOperation = parentOperation;
     }
 
     /** Starts the operation. */

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,22 +157,40 @@ 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;
                 }
                 var completionStatus = canComplete(succeededCount, failedCount, runningChildren);
                 if (completionStatus != null) {
-                    handleComplete(completionStatus);
+                    handleSuccess(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();
                 }
+
+                // If consumerThreadListener has been completed when processing above, waitForChildCompletion will
+                // immediately return null and repeat the above again
                 var child = waitForChildCompletion(succeededCount, failedCount, runningChildren);
-                // child may be null if the consumer thread is woken up due to a new item being added
+
+                // 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,12 +199,6 @@ 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
@@ -273,22 +283,15 @@ 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;
         }
 
         return null;
     }
 
-    private void handleComplete(ConcurrencyCompletionStatus status) {
-        synchronized (this) {
-            if (isOperationCompleted()) {
-                return;
-            }
-            handleSuccess(status);
-        }
-    }
-
     /**
      * Blocks the calling thread until the concurrency operation reaches a terminal state. Validates item count, handles
      * zero-branch case, then delegates to {@code waitForOperationCompletion()} from BaseDurableOperation.
@@ -299,8 +302,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();
     }
 

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

@@ -52,7 +52,16 @@ protected SerializableDurableOperation(
             TypeToken<T> resultTypeToken,
             SerDes resultSerDes,
             DurableContextImpl durableContext) {
-        super(operationIdentifier, durableContext);
+        this(operationIdentifier, resultTypeToken, resultSerDes, durableContext, null);
+    }
+
+    protected SerializableDurableOperation(
+            OperationIdentifier operationIdentifier,
+            TypeToken<T> resultTypeToken,
+            SerDes resultSerDes,
+            DurableContextImpl durableContext,
+            BaseDurableOperation parentOperation) {
+        super(operationIdentifier, durableContext, parentOperation);
         this.resultTypeToken = resultTypeToken;
         this.resultSerDes = resultSerDes;
     }

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

@@ -29,7 +29,7 @@ public class WaitOperation extends BaseDurableOperation implements DurableFuture
 
     public WaitOperation(
             OperationIdentifier operationIdentifier, Duration duration, DurableContextImpl durableContext) {
-        super(operationIdentifier, durableContext);
+        super(operationIdentifier, durableContext, null);
         this.duration = duration;
     }