prevent a race condition in future completion

Author: zhongkechen

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,49 +199,62 @@ 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 (this) {
+                if (!skipDeregistration.get()) {
+                    registerActiveThread(threadContext.threadId());
+                }
             }
+        });
 
-            // 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());
+        // 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.
+
+            // 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 (this) {
+                if (!skipDeregistration.get() && !future.isDone()) {
+                    // Deregister the current thread to allow suspension
+                    executionManager.deregisterActiveThread(threadContext.threadId());
+                } else {
+                    skipDeregistration.set(true);
+                }
             }
         }
         return future.thenApply(o -> (BaseDurableOperation) o).join();
@@ -273,6 +302,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 +330,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();
     }