diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ParallelIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ParallelIntegrationTest.java index 2e6c86eab..72f7f6628 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ParallelIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ParallelIntegrationTest.java @@ -7,8 +7,10 @@ import java.time.Duration; import java.util.ArrayList; import java.util.List; +import java.util.concurrent.ForkJoinPool; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.CsvSource; import software.amazon.lambda.durable.config.CompletionConfig; @@ -25,6 +27,34 @@ class ParallelIntegrationTest { + @Test + void singleWorkerForkJoinPoolDoesNotStarveCoordinator() { + var executor = new ForkJoinPool(1); + try { + var config = DurableConfig.builder().withExecutorService(executor).build(); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + var futures = new ArrayList>(); + var parallel = context.parallel("process-items"); + try (parallel) { + futures.add(parallel.branch("branch-a", String.class, ctx -> "A")); + futures.add(parallel.branch("branch-b", String.class, ctx -> "B")); + } + parallel.get(); + return String.join( + ",", futures.stream().map(DurableFuture::get).toList()); + }, + config); + + var result = assertTimeoutPreemptively(Duration.ofSeconds(5), () -> runner.runUntilComplete("test")); + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("A,B", result.getResult(String.class)); + } finally { + executor.shutdownNow(); + } + } + @ParameterizedTest @CsvSource({"FLAT, 2", "NESTED, 8"}) void testSimpleParallel(NestingType nestingType, int events) { 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 index 321894bd3..646b0685a 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ConcurrencyOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/ConcurrencyOperation.java @@ -9,11 +9,10 @@ import java.util.Objects; import java.util.Queue; import java.util.Set; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ConcurrentLinkedDeque; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -41,10 +40,11 @@ *

Key design points: * *

* * @param the result type of this operation @@ -53,6 +53,26 @@ public abstract class ConcurrencyOperation extends SerializableDurableOperati protected record ExpectedCompletionStatus(int completed, CompletionConfig.CompletionDecision completionDecision) {} + private record CoordinatorEvent(ChildContextOperation completedChild, Throwable failure) { + private static CoordinatorEvent stateChanged() { + return new CoordinatorEvent(null, null); + } + + private static CoordinatorEvent childCompleted(ChildContextOperation child, Throwable failure) { + return new CoordinatorEvent(child, failure); + } + + private static CoordinatorEvent failed(Throwable failure) { + return new CoordinatorEvent(null, failure); + } + } + + private static final class CoordinatorState { + private final Set> runningChildren = new HashSet<>(); + private int succeededCount; + private int failedCount; + } + private static final Logger logger = LoggerFactory.getLogger(ConcurrencyOperation.class); private final int maxConcurrency; @@ -61,18 +81,22 @@ protected record ExpectedCompletionStatus(int completed, CompletionConfig.Comple private final DurableContextImpl rootContext; private final NestingType nestingType; - // access by context thread only + // added by the context thread and read by the coordinator and result aggregation private final List> branches = Collections.synchronizedList(new ArrayList<>()); - // put only by context thread and consume only by consumer thread - private final Queue> pendingQueue = new ConcurrentLinkedDeque<>(); + // produced by the context thread and consumed by the coordinator + private final Queue> pendingQueue = new ConcurrentLinkedQueue<>(); + + // workers publish events; only the coordinator consumes them and mutates scheduling state + private final BlockingQueue coordinatorEvents = new LinkedBlockingQueue<>(); + + // guarded by completionFuture + private boolean stateChangedQueued; + private boolean coordinatorWaiting; // set by context thread and used by consumer thread protected final AtomicBoolean isJoined = new AtomicBoolean(false); - // used to wake up consumer thread for either new items or checking completion condition (isJoined changed) - private final AtomicReference> consumerThreadListener; - protected ConcurrencyOperation( OperationIdentifier operationIdentifier, TypeToken resultTypeToken, @@ -87,8 +111,12 @@ protected ConcurrencyOperation( this.operationIdGenerator = new OperationIdGenerator(getOperationId()); // root context of the concurrency operation is always non-virtual this.rootContext = durableContext.createChildContext(getOperationId(), getName(), false); - this.consumerThreadListener = new AtomicReference<>(new CompletableFuture<>()); this.nestingType = nestingType; + completionFuture.whenComplete((ignored, failure) -> { + if (failure != null) { + publishCoordinatorEvent(CoordinatorEvent.failed(failure)); + } + }); } // ========== Template methods for subclasses ========== @@ -147,14 +175,32 @@ protected ChildContextOperation enqueueItem( logger.debug("Item enqueued {}", name); pendingQueue.add(childOp); } - // notify the consumer thread a new item is available - notifyConsumerThread(); + notifyCoordinatorStateChanged(); return childOp; } - private void notifyConsumerThread() { + private void notifyCoordinatorStateChanged() { synchronized (completionFuture) { - consumerThreadListener.get().complete(null); + if (!stateChangedQueued) { + stateChangedQueued = true; + publishCoordinatorEventLocked(CoordinatorEvent.stateChanged()); + } + } + } + + private void publishCoordinatorEvent(CoordinatorEvent event) { + synchronized (completionFuture) { + publishCoordinatorEventLocked(event); + } + } + + private void publishCoordinatorEventLocked(CoordinatorEvent event) { + coordinatorEvents.add(event); + if (coordinatorWaiting) { + coordinatorWaiting = false; + if (event.failure() == null) { + registerActiveThread(getOperationId()); + } } } @@ -165,65 +211,75 @@ protected void executeItems() { /** Starts execution of all enqueued items until the expectedCompletionStatus is met. */ protected void executeItems(ExpectedCompletionStatus expectedCompletionStatus) { - // variables accessed only by the consumer thread. Put them here to avoid accidentally used by other threads - Set runningChildren = new HashSet<>(); - AtomicInteger succeededCount = new AtomicInteger(0); - AtomicInteger failedCount = new AtomicInteger(0); - - Runnable consumer = () -> { - try { - 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 (completionFuture) { - 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 completionDecision = canComplete(succeededCount, failedCount, expectedCompletionStatus); - if (completionDecision != null) { - handleCompletion(completionDecision); - 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, expectedCompletionStatus); - - // 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); - onItemComplete(succeededCount, failedCount, (ChildContextOperation) child); - } else { - throw new IllegalStateException("Unexpected completion: " + child); - } - } + // run consumer in the user thread pool, although it's not a real user thread + runUserHandler(() -> runCoordinator(expectedCompletionStatus), ThreadType.CONTEXT); + } + + private void runCoordinator(ExpectedCompletionStatus expectedCompletionStatus) { + var state = new CoordinatorState(); + + try { + while (!isOperationCompleted()) { + var completionDecision = canComplete(state.succeededCount, state.failedCount, expectedCompletionStatus); + if (completionDecision != null) { + handleCompletion(completionDecision); + return; } - } catch (Throwable ex) { - handleException(ex); + startPendingItems(state); + processCoordinatorEvent(waitForCoordinatorEvent(), state); } - }; - // run consumer in the user thread pool, although it's not a real user thread - runUserHandler(consumer, ThreadType.CONTEXT); + } catch (Throwable ex) { + handleException(ex); + } + } + + private void startPendingItems(CoordinatorState state) { + ChildContextOperation next; + while (state.runningChildren.size() < maxConcurrency && (next = pendingQueue.poll()) != null) { + var child = next; + state.runningChildren.add(child); + child.getCompletionFuture() + .whenComplete((ignored, failure) -> + publishCoordinatorEvent(CoordinatorEvent.childCompleted(child, failure))); + logger.debug("Executing operation {}", child.getName()); + child.execute(); + } + } + + private void processCoordinatorEvent(CoordinatorEvent event, CoordinatorState state) { + if (event.failure() != null) { + ExceptionHelper.sneakyThrow(ExceptionHelper.unwrapCompletableFuture(event.failure())); + } + var child = event.completedChild(); + if (child == null) { + synchronized (completionFuture) { + stateChangedQueued = false; + } + return; + } + if (!state.runningChildren.remove(child)) { + throw new IllegalStateException("Unexpected completion: " + child); + } + onItemComplete(state, child); + } + + private CoordinatorEvent waitForCoordinatorEvent() { + var threadContext = getCurrentThreadContext(); + synchronized (completionFuture) { + var event = coordinatorEvents.poll(); + if (event != null || isOperationCompleted()) { + return event != null ? event : CoordinatorEvent.stateChanged(); + } + coordinatorWaiting = true; + deregisterActiveThread(threadContext.threadId()); + } + + try { + return coordinatorEvents.take(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("Concurrency coordinator interrupted", e); + } } private void handleException(Throwable ex) { @@ -240,70 +296,22 @@ private void handleException(Throwable ex) { String.format("Unexpected exception in concurrency operation: %s", throwable)); } - private BaseDurableOperation waitForChildCompletion( - AtomicInteger succeededCount, - AtomicInteger failedCount, - Set runningChildren, - ExpectedCompletionStatus expectedCompletionStatus) { - var threadContext = getCurrentThreadContext(); - CompletableFuture future; - - synchronized (completionFuture) { - // check again in synchronized block to prevent race conditions - if (isOperationCompleted()) { - return null; - } - var completionDecision = canComplete(succeededCount, failedCount, expectedCompletionStatus); - if (completionDecision != null) { - return null; - } - ArrayList> 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()); - } - - // 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 = future.thenApply(o -> { - registerActiveThread(threadContext.threadId()); - return o; - }); - // Deregister the current thread to allow suspension - deregisterActiveThread(threadContext.threadId()); - } - } - try { - return future.thenApply(o -> (BaseDurableOperation) o).join(); - } catch (Throwable throwable) { - ExceptionHelper.sneakyThrow(ExceptionHelper.unwrapCompletableFuture(throwable)); - throw throwable; - } - } - /** * Called by a ChildContextOperation BEFORE it closes its child context. Updates counters, checks completion * criteria, and either triggers the next queued item or completes the operation. * * @param child the child operation that completed */ - private void onItemComplete( - AtomicInteger succeededCount, AtomicInteger failedCount, ChildContextOperation child) { + private void onItemComplete(CoordinatorState state, ChildContextOperation child) { // Evaluate child result outside the lock — child.get() may block waiting for a checkpoint response. logger.debug("OnItemComplete called by {}, Id: {}", child.getName(), child.getOperationId()); try { child.get(); logger.debug("Result succeeded - {}", child.getName()); - succeededCount.incrementAndGet(); + state.succeededCount++; } catch (Throwable e) { logger.debug("Child operation {} failed: {}", child.getOperationId(), e.getMessage()); - failedCount.incrementAndGet(); + state.failedCount++; } } @@ -314,12 +322,7 @@ private void onItemComplete( * @return the completion status if the operation is complete, or null if it should continue */ private CompletionConfig.CompletionDecision canComplete( - AtomicInteger succeededCount, - AtomicInteger failedCount, - ExpectedCompletionStatus expectedCompletionStatus) { - int succeeded = succeededCount.get(); - int failed = failedCount.get(); - + int succeeded, int failed, ExpectedCompletionStatus expectedCompletionStatus) { if (expectedCompletionStatus != null) { if (succeeded + failed >= expectedCompletionStatus.completed) { return expectedCompletionStatus.completionDecision; @@ -351,9 +354,7 @@ private boolean allItemsRegistered() { protected void join() { isJoined.set(true); - // Notify the consumer thread this concurrency operation is joined. Consumer thread need to check the - // completion condition again. - notifyConsumerThread(); + notifyCoordinatorStateChanged(); waitForOperationCompletion(); } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/ConcurrencyOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/ConcurrencyOperationTest.java index e51bd13da..b6488139f 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/ConcurrencyOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/ConcurrencyOperationTest.java @@ -9,10 +9,16 @@ import java.lang.reflect.Field; import java.lang.reflect.InvocationTargetException; +import java.util.ArrayList; +import java.util.List; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Function; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import software.amazon.awssdk.services.lambda.model.ContextDetails; @@ -20,13 +26,16 @@ import software.amazon.awssdk.services.lambda.model.OperationStatus; import software.amazon.awssdk.services.lambda.model.OperationType; import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.TestUtils; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.CompletionConfig; import software.amazon.lambda.durable.config.NestingType; +import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.execution.OperationIdGenerator; +import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.execution.ThreadContext; import software.amazon.lambda.durable.execution.ThreadType; import software.amazon.lambda.durable.model.OperationIdentifier; @@ -103,14 +112,10 @@ private void setOperationIdGenerator(ConcurrencyOperation op, OperationIdGene private CompletionConfig.CompletionDecision canComplete( ConcurrencyOperation op, int succeededCount, int failedCount) throws Exception { var method = ConcurrencyOperation.class.getDeclaredMethod( - "canComplete", - AtomicInteger.class, - AtomicInteger.class, - ConcurrencyOperation.ExpectedCompletionStatus.class); + "canComplete", int.class, int.class, ConcurrencyOperation.ExpectedCompletionStatus.class); method.setAccessible(true); try { - return (CompletionConfig.CompletionDecision) - method.invoke(op, new AtomicInteger(succeededCount), new AtomicInteger(failedCount), null); + return (CompletionConfig.CompletionDecision) method.invoke(op, succeededCount, failedCount, null); } catch (InvocationTargetException e) { var cause = e.getCause(); if (cause instanceof RuntimeException runtimeException) { @@ -284,6 +289,68 @@ void canComplete_whenShouldCompleteReturnsNull_shouldThrow() throws Exception { assertEquals("shouldComplete must return a completion decision", exception.getMessage()); } + @Test + void coordinatorStartsNextQueuedChildAfterCompletionEvent() throws Exception { + var activeCount = new AtomicInteger(0); + var peakCount = new AtomicInteger(0); + var startOrder = new CopyOnWriteArrayList(); + var op = new QueueTestConcurrencyOperation( + OperationIdentifier.of(OPERATION_ID, "test-concurrency", OperationSubType.PARALLEL), + durableContext, + childContext, + activeCount, + peakCount, + startOrder); + + op.execute(); + op.enqueueItem( + "branch-1", + ctx -> "result-1", + TypeToken.get(String.class), + SER_DES, + OperationSubType.PARALLEL_BRANCH, + false); + op.enqueueItem( + "branch-2", + ctx -> "result-2", + TypeToken.get(String.class), + SER_DES, + OperationSubType.PARALLEL_BRANCH, + false); + + var first = op.getControlledChild(0); + var second = op.getControlledChild(1); + assertTrue(first.awaitStarted()); + assertFalse(second.hasStarted()); + + first.completeSuccessfully(); + assertTrue(second.awaitStarted()); + second.completeSuccessfully(); + op.exposedJoin(); + + assertTrue(op.isSuccessHandled()); + assertEquals(1, peakCount.get()); + assertEquals(List.of("branch-1", "branch-2"), startOrder); + } + + @Test + void exceptionalCompletionWakesWaitingCoordinator() throws Exception { + var op = new QueueTestConcurrencyOperation( + OperationIdentifier.of(OPERATION_ID, "test-concurrency", OperationSubType.PARALLEL), + durableContext, + childContext, + new AtomicInteger(), + new AtomicInteger(), + new CopyOnWriteArrayList<>()); + + op.execute(); + verify(executionManager, timeout(5_000)).deregisterActiveThread("Root"); + + op.suspend(); + + assertTrue(op.awaitCoordinatorStopped()); + } + // ===== Test subclass ===== static class TestConcurrencyOperation extends ConcurrencyOperation { @@ -356,4 +423,156 @@ DurableContextImpl getLastParentContext() { return lastParentContext; } } + + static class QueueTestConcurrencyOperation extends ConcurrencyOperation { + + private final DurableContextImpl childContext; + private final AtomicInteger activeCount; + private final AtomicInteger peakCount; + private final List startOrder; + private final List> controlledChildren = new ArrayList<>(); + private volatile boolean successHandled; + + QueueTestConcurrencyOperation( + OperationIdentifier operationIdentifier, + DurableContextImpl durableContext, + DurableContextImpl childContext, + AtomicInteger activeCount, + AtomicInteger peakCount, + List startOrder) { + super( + operationIdentifier, + RESULT_TYPE, + SER_DES, + durableContext, + 1, + CompletionConfig.allSuccessful().completionDecisionFunction(), + NestingType.NESTED); + this.childContext = childContext; + this.activeCount = activeCount; + this.peakCount = peakCount; + this.startOrder = startOrder; + } + + @Override + protected ChildContextOperation createItem( + String operationId, + String name, + Function function, + TypeToken resultType, + SerDes serDes, + OperationSubType branchSubType) { + var child = new ControlledChildOperation<>( + OperationIdentifier.of(operationId, name, branchSubType), + resultType, + serDes, + childContext, + this, + activeCount, + peakCount, + startOrder); + controlledChildren.add(child); + return child; + } + + @Override + protected void handleCompletion(CompletionConfig.CompletionDecision completionDecision) { + successHandled = true; + onCheckpointComplete(Operation.builder() + .id(getOperationId()) + .status(OperationStatus.SUCCEEDED) + .build()); + } + + @Override + protected void start() { + executeItems(); + } + + @Override + protected void replay(Operation existing) { + executeItems(); + } + + @Override + public Void get() { + return null; + } + + ControlledChildOperation getControlledChild(int index) { + return controlledChildren.get(index); + } + + void exposedJoin() { + join(); + } + + boolean isSuccessHandled() { + return successHandled; + } + + void suspend() { + completionFuture.completeExceptionally(new SuspendExecutionException()); + } + + boolean awaitCoordinatorStopped() throws Exception { + getRunningUserHandler().get(5, TimeUnit.SECONDS); + return getRunningUserHandler().isDone(); + } + } + + static class ControlledChildOperation extends ChildContextOperation { + + private final AtomicInteger activeCount; + private final AtomicInteger peakCount; + private final List startOrder; + private final CountDownLatch started = new CountDownLatch(1); + + ControlledChildOperation( + OperationIdentifier operationIdentifier, + TypeToken resultType, + SerDes serDes, + DurableContextImpl childContext, + ConcurrencyOperation parent, + AtomicInteger activeCount, + AtomicInteger peakCount, + List startOrder) { + super( + operationIdentifier, + ctx -> null, + resultType, + RunInChildContextConfig.builder().serDes(serDes).build(), + childContext, + parent); + this.activeCount = activeCount; + this.peakCount = peakCount; + this.startOrder = startOrder; + } + + @Override + public void execute() { + startOrder.add(getName()); + var current = activeCount.incrementAndGet(); + peakCount.updateAndGet(peak -> Math.max(peak, current)); + started.countDown(); + } + + @Override + public R get() { + return null; + } + + boolean awaitStarted() throws InterruptedException { + return started.await(5, TimeUnit.SECONDS); + } + + boolean hasStarted() { + return started.getCount() == 0; + } + + void completeSuccessfully() { + activeCount.decrementAndGet(); + markAlreadyCompleted(); + } + } }