From 482ab3edbceb02e383f9956a6f147181362b3906 Mon Sep 17 00:00:00 2001 From: hippi345 <13539685+hippi345@users.noreply.github.com> Date: Sat, 17 Oct 2026 16:43:33 +0000 Subject: [PATCH 1/4] fix(sdk): throw when getting a skipped parallel branch get() on a parallel branch recorded as SKIPPED now throws ParallelBranchSkippedException instead of suspending with no pending operations. This covers branches that never started and in-flight branches skipped at completion, on the first invocation and on replay. --- docs/core/parallel.md | 26 ++++- .../durable/ParallelIntegrationTest.java | 106 +++++++++++++++++- .../ParallelBranchSkippedException.java | 32 ++++++ .../operation/BaseDurableOperation.java | 10 +- .../operation/ChildContextOperation.java | 19 ++++ .../durable/operation/ParallelOperation.java | 37 ++++-- .../ParallelBranchSkippedExceptionTest.java | 25 +++++ .../operation/ParallelOperationTest.java | 7 ++ 8 files changed, 246 insertions(+), 16 deletions(-) create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/exception/ParallelBranchSkippedException.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/exception/ParallelBranchSkippedExceptionTest.java diff --git a/docs/core/parallel.md b/docs/core/parallel.md index 99bdb26f86..d56723ffc8 100644 --- a/docs/core/parallel.md +++ b/docs/core/parallel.md @@ -39,8 +39,12 @@ try (var parallel = ctx.parallel("work")) { | `size()` | Total number of registered branches | | `succeeded()` | Number of branches that succeeded | | `failed()` | Number of branches that failed | +| `skipped()` | Number of branches recorded as `SKIPPED` | +| `statuses()` | Per-branch status: `SUCCEEDED`, `FAILED`, or `SKIPPED` | | `completionStatus()` | Why the operation completed (`ALL_COMPLETED`, `MIN_SUCCESSFUL_REACHED`, `FAILURE_TOLERANCE_EXCEEDED`) | +When a completion config finishes the parallel operation early, branches that never started — and branches that were still running at completion — are recorded as `SKIPPED`. `get()` on a skipped branch throws `ParallelBranchSkippedException` on the first invocation and on every replay. The SDK does not start or resume that branch later. + ### ParallelConfig Configure concurrency limits and completion criteria: @@ -110,9 +114,29 @@ try { } ``` +A skipped branch is different: the parallel result recorded it as `SKIPPED`, so `get()` throws `ParallelBranchSkippedException` instead of waiting: + +```java +var parallel = ctx.parallel("first-successful", ParallelConfig.builder() + .maxConcurrency(1) + .completionConfig(CompletionConfig.firstSuccessful()) + .build()); +var a = parallel.branch("branch-a", String.class, branchCtx -> "A"); +var b = parallel.branch("branch-b", String.class, branchCtx -> "B"); +parallel.get(); // statuses [SUCCEEDED, SKIPPED] +a.get(); // "A" +try { + b.get(); // throws on the first invocation and on every replay +} catch (ParallelBranchSkippedException e) { + // Branch never started, or was still running when the parallel completed. + // e.getBranchName() is the name passed to branch(). +} +``` + | Exception | When Thrown | |-----------|-------------| | `ParallelBranchFailedException` | Branch failed and the original exception could not be reconstructed | +| `ParallelBranchSkippedException` | Branch was recorded as `SKIPPED` in the parallel result | | User's exception | Branch threw a reconstructable exception — propagated through `get()` | ### Checkpoint-and-Replay @@ -121,4 +145,4 @@ Parallel operations are fully durable. On replay after interruption: - Completed branches return cached results without re-execution - Incomplete branches resume from their last checkpoint -- Branches that never started execute fresh +- Branches recorded as `SKIPPED` are not started; `get()` on those branch futures throws `ParallelBranchSkippedException` 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 8ec7e20578..e83c261a95 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,7 +7,9 @@ import java.time.Duration; import java.util.ArrayList; import java.util.List; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ForkJoinPool; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.Test; @@ -19,7 +21,7 @@ import software.amazon.lambda.durable.config.NestingType; import software.amazon.lambda.durable.config.ParallelConfig; import software.amazon.lambda.durable.config.WaitForConditionConfig; -import software.amazon.lambda.durable.exception.IllegalDurableOperationException; +import software.amazon.lambda.durable.exception.ParallelBranchSkippedException; import software.amazon.lambda.durable.model.ConcurrencyCompletionStatus; import software.amazon.lambda.durable.model.ExecutionStatus; import software.amazon.lambda.durable.model.ParallelResult; @@ -1387,12 +1389,104 @@ void testInvalidSuspensionWhenWaitingOnSkippedBranch() { } parallel.get(); - return futures.get(1).get(); + var skipped = assertThrows( + ParallelBranchSkippedException.class, () -> futures.get(1).get()); + assertEquals("branch-2", skipped.getBranchName()); + return "done"; }); - var ex = assertThrows(IllegalDurableOperationException.class, () -> runner.run("test")); - assertTrue(ex.isRetryable()); - assertTrue(ex.getMessage().contains("Invalid suspension. No operation is pending.")); - assertTrue(ex.getMessage().contains("branch-2")); + var result = assertTimeoutPreemptively(Duration.ofSeconds(10), () -> runner.runUntilComplete("test")); + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("done", result.getResult(String.class)); + } + + @Test + void getOnSkippedBranchThatNeverStartedThrowsOnFirstRunAndReplay() { + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { + var config = ParallelConfig.builder() + .maxConcurrency(1) + .completionConfig(CompletionConfig.firstSuccessful()) + .build(); + var parallel = context.parallel("first-successful", config); + var a = parallel.branch("branch-a", String.class, ctx -> "A"); + var b = parallel.branch("branch-b", String.class, ctx -> "B"); + var result = parallel.get(); + assertEquals(ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED, result.completionStatus()); + assertEquals(List.of(ParallelResult.Status.SUCCEEDED, ParallelResult.Status.SKIPPED), result.statuses()); + assertEquals("A", a.get()); + var skipped = assertThrows(ParallelBranchSkippedException.class, b::get); + assertEquals("branch-b", skipped.getBranchName()); + assertEquals("Parallel branch 'branch-b' was skipped", skipped.getMessage()); + return "done"; + }); + + var first = assertTimeoutPreemptively(Duration.ofSeconds(10), () -> runner.runUntilComplete("input")); + assertEquals(ExecutionStatus.SUCCEEDED, first.getStatus()); + assertEquals("done", first.getResult(String.class)); + + var replay = assertTimeoutPreemptively(Duration.ofSeconds(10), () -> runner.run("input")); + assertEquals(ExecutionStatus.SUCCEEDED, replay.getStatus()); + assertEquals("done", replay.getResult(String.class)); + } + + @Test + void getOnSkippedInFlightBranchThrowsOnFirstRunAndReplay() throws Exception { + var branchBStarted = new CountDownLatch(1); + var releaseBranchB = new CountDownLatch(1); + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { + var config = ParallelConfig.builder() + .completionConfig(CompletionConfig.firstSuccessful()) + .build(); + var parallel = context.parallel("first-successful", config); + var b = parallel.branch("branch-b", String.class, ctx -> { + branchBStarted.countDown(); + try { + if (!releaseBranchB.await(10, TimeUnit.SECONDS)) { + throw new IllegalStateException("timed out waiting to release branch-b"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + return "B"; + }); + var a = parallel.branch("branch-a", String.class, ctx -> { + try { + if (!branchBStarted.await(10, TimeUnit.SECONDS)) { + throw new IllegalStateException("timed out waiting for branch-b to start"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + return "A"; + }); + var result = parallel.get(); + assertEquals(ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED, result.completionStatus()); + assertEquals(List.of(ParallelResult.Status.SKIPPED, ParallelResult.Status.SUCCEEDED), result.statuses()); + try { + var skipped = assertThrows(ParallelBranchSkippedException.class, b::get); + assertEquals("branch-b", skipped.getBranchName()); + assertEquals("A", a.get()); + context.wait("pause", Duration.ofSeconds(1)); + var skippedOnReplay = assertThrows(ParallelBranchSkippedException.class, b::get); + assertEquals("branch-b", skippedOnReplay.getBranchName()); + return "done"; + } finally { + releaseBranchB.countDown(); + } + }); + + try { + var first = assertTimeoutPreemptively(Duration.ofSeconds(15), () -> runner.run("input")); + assertEquals(ExecutionStatus.PENDING, first.getStatus()); + + runner.advanceTime(); + var replay = assertTimeoutPreemptively(Duration.ofSeconds(15), () -> runner.run("input")); + assertEquals(ExecutionStatus.SUCCEEDED, replay.getStatus()); + assertEquals("done", replay.getResult(String.class)); + } finally { + releaseBranchB.countDown(); + } } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/exception/ParallelBranchSkippedException.java b/sdk/src/main/java/software/amazon/lambda/durable/exception/ParallelBranchSkippedException.java new file mode 100644 index 0000000000..9bee4766d0 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/exception/ParallelBranchSkippedException.java @@ -0,0 +1,32 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.exception; + +/** + * Thrown when {@code get()} is called on a parallel branch that the checkpointed parallel result recorded as + * {@code SKIPPED}. + * + *

This happens when a completion config such as {@code firstSuccessful()} finishes the parallel operation before + * every branch has started, or while a branch is still running. The SDK does not start or resume that branch later, so + * {@code get()} throws this exception on the first invocation and on every replay. + */ +public class ParallelBranchSkippedException extends DurableExecutionException { + private final String branchName; + + public ParallelBranchSkippedException(String branchName) { + super(formatMessage(branchName)); + this.branchName = branchName; + } + + /** Returns the name of the skipped branch. */ + public String getBranchName() { + return branchName; + } + + private static String formatMessage(String branchName) { + if (branchName == null) { + return "Parallel branch was skipped"; + } + return String.format("Parallel branch '%s' was skipped", branchName); + } +} 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 index 31c9926156..efe7da8151 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java @@ -243,8 +243,16 @@ private static ExecutionManager getExecutionManager(List extends SerializableDurableOperation { private final Function function; private final AtomicBoolean replayChildren = new AtomicBoolean(false); + private final AtomicBoolean skipped = new AtomicBoolean(false); private final AtomicReference> cachedOperationResult = new AtomicReference<>(null); // child context for RunInChildContext @@ -223,8 +225,25 @@ private void handleChildContextFailure(Throwable exception) { OperationUpdate.builder().action(OperationAction.FAIL).error(errorObject)); } + /** + * Completes this branch as skipped so {@link #get()} throws {@link ParallelBranchSkippedException} instead of + * waiting for an operation that will never be started or checkpointed. + */ + void completeAsSkipped() { + skipped.set(true); + markAlreadyCompleted(); + } + + boolean isSkipped() { + return skipped.get(); + } + @Override public T get() { + waitForCompletionFuture(); + if (skipped.get()) { + throw new ParallelBranchSkippedException(getName()); + } var op = waitForOperationCompletion(); if (cachedOperationResult.get() != null) { // we have a result, just return it directly diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java index 41a687e6b6..e2dbd90d35 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java @@ -17,6 +17,7 @@ import software.amazon.lambda.durable.config.ParallelBranchConfig; import software.amazon.lambda.durable.config.ParallelConfig; import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.exception.ParallelBranchSkippedException; import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; @@ -85,6 +86,7 @@ protected void handleCompletion(CompletionConfig.CompletionDecision completionDe statuses); var serializedResult = serializeAndDeserializeResult(cachedResult); cachedResult = serializedResult.deserialized(); + completeSkippedBranches(items, statuses); // Branches added after checkpoint will not exist in the checkpointed result, but they'll be in the returned // value from get() method. @@ -95,13 +97,23 @@ protected void handleCompletion(CompletionConfig.CompletionDecision completionDe .contextOptions(ContextOptions.builder().replayChildren(true).build())); } + private void completeSkippedBranches(List> items, List statuses) { + for (int i = 0; i < items.size(); i++) { + if (statuses.get(i) == ParallelResult.Status.SKIPPED) { + items.get(i).completeAsSkipped(); + } + } + } + private ParallelResult.Status getParallelItemStatus(ChildContextOperation childContextOperation) { - if (!childContextOperation.isOperationCompleted()) { + if (!childContextOperation.isOperationCompleted() || childContextOperation.isSkipped()) { return ParallelResult.Status.SKIPPED; } try { childContextOperation.get(); return ParallelResult.Status.SUCCEEDED; + } catch (ParallelBranchSkippedException e) { + return ParallelResult.Status.SKIPPED; } catch (Throwable t) { return ParallelResult.Status.FAILED; } @@ -180,13 +192,22 @@ public DurableFuture branch( var nextBranchIndex = getBranches().size(); - // ConcurrencyOperation will skip this branch if skip=true: - // 1. if the parallel operation is already completed (partialResult is not null) - // 2. if the branch is already skipped in the partialResult or nonexistent in the partialResult - var skip = partialResult != null - && (partialResult.statuses().size() <= nextBranchIndex - || partialResult.statuses().get(nextBranchIndex) == ParallelResult.Status.SKIPPED); + // Skip when the checkpointed or in-memory parallel result already completed without this branch. + var skip = shouldSkipBranch(nextBranchIndex); var serDes = config.serDes() == null ? getContext().getDurableConfig().getSerDes() : config.serDes(); - return enqueueItem(name, func, resultType, serDes, OperationSubType.PARALLEL_BRANCH, skip); + var branch = enqueueItem(name, func, resultType, serDes, OperationSubType.PARALLEL_BRANCH, skip); + if (skip) { + branch.completeAsSkipped(); + } + return branch; + } + + private boolean shouldSkipBranch(int nextBranchIndex) { + var result = partialResult != null ? partialResult : cachedResult; + if (result == null) { + return false; + } + return result.statuses().size() <= nextBranchIndex + || result.statuses().get(nextBranchIndex) == ParallelResult.Status.SKIPPED; } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/exception/ParallelBranchSkippedExceptionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/exception/ParallelBranchSkippedExceptionTest.java new file mode 100644 index 0000000000..a8b535174b --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/exception/ParallelBranchSkippedExceptionTest.java @@ -0,0 +1,25 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.exception; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; + +import org.junit.jupiter.api.Test; + +class ParallelBranchSkippedExceptionTest { + + @Test + void namesTheSkippedBranchInTheMessage() { + var exception = new ParallelBranchSkippedException("branch-b"); + assertEquals("branch-b", exception.getBranchName()); + assertEquals("Parallel branch 'branch-b' was skipped", exception.getMessage()); + } + + @Test + void formatsAFallbackMessageWhenTheBranchNameIsNull() { + var exception = new ParallelBranchSkippedException(null); + assertNull(exception.getBranchName()); + assertEquals("Parallel branch was skipped", exception.getMessage()); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/ParallelOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/ParallelOperationTest.java index 1a29032513..99519fe113 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/ParallelOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/ParallelOperationTest.java @@ -9,6 +9,7 @@ import static org.mockito.Mockito.*; import java.lang.reflect.Field; +import java.time.Duration; import java.util.ArrayList; import java.util.List; import java.util.concurrent.CompletableFuture; @@ -33,6 +34,7 @@ import software.amazon.lambda.durable.config.ParallelBranchConfig; import software.amazon.lambda.durable.config.ParallelConfig; import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.exception.ParallelBranchSkippedException; import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.execution.ThreadContext; import software.amazon.lambda.durable.execution.ThreadType; @@ -364,6 +366,11 @@ void minSuccessful_notExecuteSkippedBranchWhenReplay() { assertEquals(1, result.skipped()); assertEquals(ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED, result.completionStatus()); assertTrue(result.completionStatus().isSucceeded()); + + var skipped = op.getBranches().get(0); + var thrown = assertTimeoutPreemptively( + Duration.ofSeconds(2), () -> assertThrows(ParallelBranchSkippedException.class, skipped::get)); + assertEquals("branch-1", thrown.getBranchName()); } // ===== Context hierarchy ===== From 0b805a42f39188ddf40253ce73df66b5c177c1b8 Mon Sep 17 00:00:00 2001 From: hippi345 <13539685+hippi345@users.noreply.github.com> Date: Sat, 17 Oct 2026 16:59:20 +0000 Subject: [PATCH 2/4] docs(sdk): restore waitForOperationCompletion javadoc Move the existing waitForOperationCompletion javadoc back onto that method and document waitForCompletionFuture separately. --- .../durable/operation/BaseDurableOperation.java | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) 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 index efe7da8151..551915794c 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java @@ -238,19 +238,20 @@ private static ExecutionManager getExecutionManager(ListSkipped parallel branches that never started can settle this way because they have no checkpointed operation. */ protected void waitForCompletionFuture() { executionManager.waitForOperationCompletion(this); } + /** + * Waits for the operation to complete. Deregisters the current thread to allow Lambda suspension if the operation + * is still in progress, then re-registers when the operation completes. + * + * @return the completed operation + */ protected Operation waitForOperationCompletion() { waitForCompletionFuture(); From aed55d3cf1aba09b735a490fd0c92d1fa414c041 Mon Sep 17 00:00:00 2001 From: hippi345 <13539685+hippi345@users.noreply.github.com> Date: Mon, 26 Oct 2026 15:56:16 +0000 Subject: [PATCH 3/4] fix(sdk): keep skipped parallel branches stable for anyOf and get anyOf ignores a branch that settled only because it was skipped, so replay returns the same winner. The parallel snapshot marks an incomplete branch skipped under the same lock that publishes its result, and a nested skip exception stays a failure of the branch that propagated it. --- docs/core/parallel.md | 4 +- .../ParallelSkippedBranchReviewTest.java | 205 ++++++++++++++++++ .../amazon/lambda/durable/DurableFuture.java | 4 + .../durable/execution/ExecutionManager.java | 21 ++ .../operation/BaseDurableOperation.java | 11 + .../operation/ChildContextOperation.java | 63 +++++- .../durable/operation/ParallelOperation.java | 28 +-- ...seDurableOperationFirstCompletionTest.java | 73 +++++++ 8 files changed, 381 insertions(+), 28 deletions(-) create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ParallelSkippedBranchReviewTest.java diff --git a/docs/core/parallel.md b/docs/core/parallel.md index d56723ffc8..c02d11451d 100644 --- a/docs/core/parallel.md +++ b/docs/core/parallel.md @@ -43,7 +43,9 @@ try (var parallel = ctx.parallel("work")) { | `statuses()` | Per-branch status: `SUCCEEDED`, `FAILED`, or `SKIPPED` | | `completionStatus()` | Why the operation completed (`ALL_COMPLETED`, `MIN_SUCCESSFUL_REACHED`, `FAILURE_TOLERANCE_EXCEEDED`) | -When a completion config finishes the parallel operation early, branches that never started — and branches that were still running at completion — are recorded as `SKIPPED`. `get()` on a skipped branch throws `ParallelBranchSkippedException` on the first invocation and on every replay. The SDK does not start or resume that branch later. +When a completion config finishes the parallel operation early, branches that never started — and branches that were still running at completion — are recorded as `SKIPPED`. The snapshot and the skip happen together with that branch's own completion, so `get()` does not return a value and later throw for the same call sequence. `get()` on a skipped branch throws `ParallelBranchSkippedException` on the first invocation and on every replay. The SDK does not start or resume that branch later. + +`DurableFuture.anyOf` does not select a skipped branch while another future can still settle, so the winner from the first invocation is the winner on replay. If every future passed to `anyOf` is skipped, it throws `ParallelBranchSkippedException`. ### ParallelConfig diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ParallelSkippedBranchReviewTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ParallelSkippedBranchReviewTest.java new file mode 100644 index 0000000000..a8ee9bfc7a --- /dev/null +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ParallelSkippedBranchReviewTest.java @@ -0,0 +1,205 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.time.Duration; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.services.lambda.model.OperationStatus; +import software.amazon.lambda.durable.config.CompletionConfig; +import software.amazon.lambda.durable.config.ParallelConfig; +import software.amazon.lambda.durable.exception.ParallelBranchSkippedException; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.model.ParallelResult; +import software.amazon.lambda.durable.serde.JacksonSerDes; +import software.amazon.lambda.durable.serde.SerDes; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +/** + * Repros for the three review findings on skipped parallel branches: {@code anyOf} replay, the completion-snapshot + * race, and nested skip exceptions classified as failures. + */ +class ParallelSkippedBranchReviewTest { + + private static final Duration DEADLINE = Duration.ofSeconds(20); + + @Test + void anyOfKeepsTheRunningWinnerWhenTheOtherBranchIsSkippedOnReplay() { + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { + var handlerThread = Thread.currentThread(); + var parallel = context.parallel( + "first-successful", + ParallelConfig.builder() + .maxConcurrency(1) + .completionConfig(CompletionConfig.firstSuccessful()) + .build()); + var winner = parallel.branch("winner", String.class, ctx -> { + waitUntilParked(handlerThread); + return "winner"; + }); + var skipped = parallel.branch("skipped", String.class, ctx -> "skipped"); + var first = String.valueOf(DurableFuture.anyOf(skipped, winner)); + context.wait("pause", Duration.ofSeconds(1)); + var second = String.valueOf(DurableFuture.anyOf(skipped, winner)); + return first + "," + second; + }); + + var pending = assertTimeoutPreemptively(DEADLINE, () -> runner.run("input")); + assertEquals(ExecutionStatus.PENDING, pending.getStatus()); + + runner.advanceTime(); + var replay = assertTimeoutPreemptively(DEADLINE, () -> runner.run("input")); + assertEquals(ExecutionStatus.SUCCEEDED, replay.getStatus()); + assertEquals("winner,winner", replay.getResult(String.class)); + } + + @Test + void branchGetKeepsTheSameOutcomeWhenItFinishesDuringResultSerialization() { + var branchRunning = new CountDownLatch(1); + var allowBranch = new CountDownLatch(1); + var branch = new AtomicReference>(); + var duringSerialization = new AtomicReference(); + var serDes = new ParallelResultSerializationGate(branch, allowBranch, duringSerialization); + var config = DurableConfig.builder().withSerDes(serDes).build(); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + var parallel = context.parallel( + "first-successful", + ParallelConfig.builder() + .completionConfig(CompletionConfig.firstSuccessful()) + .build()); + var slow = parallel.branch("branch-b", String.class, ctx -> { + branchRunning.countDown(); + await(allowBranch, "release branch-b"); + return "B"; + }); + branch.set(slow); + parallel.branch("branch-a", String.class, ctx -> { + await(branchRunning, "branch-b to start"); + return "A"; + }); + var result = parallel.get(); + return duringSerialization.get() + "|" + observe(slow) + "|" + result.statuses(); + }, + config); + + try { + var execution = assertTimeoutPreemptively(DEADLINE, () -> runner.runUntilComplete("input")); + assertEquals(ExecutionStatus.SUCCEEDED, execution.getStatus()); + assertEquals("SKIPPED:branch-b|SKIPPED:branch-b|[SKIPPED, SUCCEEDED]", execution.getResult(String.class)); + } finally { + allowBranch.countDown(); + } + } + + @Test + void nestedSkipExceptionStaysAFailedOuterBranch() { + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { + var outer = context.parallel("outer", ParallelConfig.builder().build()); + var outerBranch = outer.branch("outer-branch", String.class, ctx -> { + var inner = ctx.parallel( + "inner", + ParallelConfig.builder() + .maxConcurrency(1) + .completionConfig(CompletionConfig.firstSuccessful()) + .build()); + inner.branch("inner-winner", String.class, innerCtx -> "ok"); + var skipped = inner.branch("inner-skipped", String.class, innerCtx -> "no"); + inner.get(); + skipped.get(); + return "unreachable"; + }); + + var result = outer.get(); + assertEquals(List.of(ParallelResult.Status.FAILED), result.statuses()); + assertEquals(1, result.failed()); + assertEquals(0, result.skipped()); + var thrown = assertThrows(ParallelBranchSkippedException.class, outerBranch::get); + assertEquals("inner-skipped", thrown.getBranchName()); + return "done"; + }); + + var first = assertTimeoutPreemptively(DEADLINE, () -> runner.runUntilComplete("input")); + assertEquals(ExecutionStatus.SUCCEEDED, first.getStatus()); + assertEquals("done", first.getResult(String.class)); + var checkpoint = runner.getOperation("outer-branch"); + assertEquals(OperationStatus.FAILED, checkpoint.getStatus()); + assertTrue(checkpoint.getContextDetails().error().errorMessage().contains("inner-skipped")); + + var replay = assertTimeoutPreemptively(DEADLINE, () -> runner.run("input")); + assertEquals(ExecutionStatus.SUCCEEDED, replay.getStatus()); + assertEquals("done", replay.getResult(String.class)); + } + + private static void waitUntilParked(Thread handlerThread) { + var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10); + while (handlerThread.getState() != Thread.State.WAITING + && handlerThread.getState() != Thread.State.TIMED_WAITING) { + if (System.nanoTime() >= deadline) { + throw new IllegalStateException("timed out waiting for anyOf to park"); + } + Thread.onSpinWait(); + } + } + + private static void await(CountDownLatch latch, String description) { + try { + if (!latch.await(10, TimeUnit.SECONDS)) { + throw new IllegalStateException("timed out waiting for " + description); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("interrupted while waiting for " + description, e); + } + } + + private static String observe(DurableFuture future) { + try { + return String.valueOf(future.get()); + } catch (ParallelBranchSkippedException e) { + return "SKIPPED:" + e.getBranchName(); + } + } + + /** Blocks inside {@link ParallelResult} serialization so a branch can finish in that window. */ + private static final class ParallelResultSerializationGate implements SerDes { + private final JacksonSerDes delegate = new JacksonSerDes(); + private final AtomicReference> branch; + private final CountDownLatch allowBranch; + private final AtomicReference duringSerialization; + private final AtomicBoolean gated = new AtomicBoolean(); + + private ParallelResultSerializationGate( + AtomicReference> branch, + CountDownLatch allowBranch, + AtomicReference duringSerialization) { + this.branch = branch; + this.allowBranch = allowBranch; + this.duringSerialization = duringSerialization; + } + + @Override + public String serialize(Object value) { + if (value instanceof ParallelResult && gated.compareAndSet(false, true)) { + allowBranch.countDown(); + duringSerialization.set(observe(branch.get())); + } + return delegate.serialize(value); + } + + @Override + public T deserialize(String data, TypeToken typeToken) { + return delegate.deserialize(data, typeToken); + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableFuture.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableFuture.java index 91f3351ae2..8ad0b51b7a 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableFuture.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableFuture.java @@ -58,6 +58,10 @@ static List allOf(List> futures) { /** * Waits for any of the provided futures to settle and returns its result. * + *

A parallel branch that settled only because it was skipped does not win while another future can still settle. + * That keeps the winner stable when the skipped branch is completed immediately on replay. If every future is + * skipped, this throws {@code ParallelBranchSkippedException} from {@link #get()}. + * * @param futures the futures to wait for * @return the result of the first future to settle */ 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 index 06502b772a..d5381d33d5 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutionManager.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutionManager.java @@ -691,6 +691,10 @@ public void waitForOperationCompletion(BaseDurableOperation operation) { *

This is the multi-operation equivalent of {@code BaseDurableOperation.waitForOperationCompletion()}: it builds * a combined future so this manager can suspend while waiting for the first operation to settle. * + *

Operations that settled only because they were skipped are not selected while another operation can still + * settle. That preserves the {@code anyOf} winner when a skipped parallel branch is completed immediately on + * replay. + * * @param operations operations to wait on * @return the first durable future to settle */ @@ -701,6 +705,23 @@ public DurableFuture waitForFirstOperationCompletion(List(operations); + while (true) { + var eligible = pending.stream().filter(op -> !op.isSkipSettlement()).toList(); + var candidates = eligible.isEmpty() ? List.copyOf(pending) : eligible; + var settled = awaitFirstSettledOperation(candidates); + if (eligible.isEmpty() || !isSkipSettlement(settled)) { + return settled; + } + pending.remove(settled); + } + } + + private static boolean isSkipSettlement(DurableFuture settled) { + return settled instanceof BaseDurableOperation operation && operation.isSkipSettlement(); + } + + private DurableFuture awaitFirstSettledOperation(List operations) { var threadId = getCurrentThreadContext().threadId(); var wait = new CompletionWait(); var anyOfFuture = CompletableFuture.anyOf(operations.stream() 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 index 8cefe2de35..d1c4d908bc 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java @@ -207,6 +207,17 @@ protected boolean isOperationCompleted() { return completionFuture.isDone(); } + /** + * Whether this operation settled because it was skipped, rather than because its own work finished. + * + *

{@link software.amazon.lambda.durable.DurableFuture#anyOf} ignores a skipped operation while another future + * can still settle. A branch marked skipped on replay therefore does not replace the winner observed on the first + * invocation. + */ + public boolean isSkipSettlement() { + return false; + } + /** * Waits for the first of the given operations to complete. * diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java index 11afc9294c..371875098a 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java @@ -36,6 +36,7 @@ import software.amazon.lambda.durable.logging.DurableLogger; import software.amazon.lambda.durable.model.DeserializedOperationResult; import software.amazon.lambda.durable.model.OperationIdentifier; +import software.amazon.lambda.durable.model.ParallelResult; import software.amazon.lambda.durable.util.ExceptionHelper; /** @@ -55,6 +56,7 @@ public class ChildContextOperation extends SerializableDurableOperation { private final Function function; private final AtomicBoolean replayChildren = new AtomicBoolean(false); private final AtomicBoolean skipped = new AtomicBoolean(false); + private final Object settlementLock = new Object(); private final AtomicReference> cachedOperationResult = new AtomicReference<>(null); // child context for RunInChildContext @@ -154,7 +156,17 @@ private void executeChildContext() { private void handleChildContextSuccess(T result) { var serializedResult = serializeAndDeserializeResult(result); + synchronized (settlementLock) { + // A parallel snapshot may already have recorded this branch as skipped. Do not publish a result after that + // decision, or get() can return a value and later throw. + if (skipped.get()) { + return; + } + publishSuccess(serializedResult); + } + } + private void publishSuccess(SerializedResult serializedResult) { if (replayChildren.get() || isVirtual || parentOperation != null && parentOperation.isOperationCompleted()) { // Skip checkpointing if // - parent ConcurrencyOperation has already completed, preventing race conditions where a child finishes @@ -167,9 +179,9 @@ private void handleChildContextSuccess(T result) { fireOnOperationEnd(null, null, false); } markAlreadyCompleted(); - } else { - checkpointSuccess(serializedResult.deserialized(), serializedResult.serialized()); + return; } + checkpointSuccess(serializedResult.deserialized(), serializedResult.serialized()); } private void checkpointSuccess(T result, String serialized) { @@ -207,6 +219,15 @@ private void handleChildContextFailure(Throwable exception) { errorObject = serializeException(exception); } + synchronized (settlementLock) { + if (skipped.get()) { + return; + } + publishFailure(exception, errorObject); + } + } + + private void publishFailure(Throwable exception, ErrorObject errorObject) { var op = createVirtualOperation(errorObject); cachedOperationResult.set(DeserializedOperationResult.failed(translateException(op, errorObject))); @@ -229,8 +250,41 @@ private void handleChildContextFailure(Throwable exception) { /** * Completes this branch as skipped so {@link #get()} throws {@link ParallelBranchSkippedException} instead of * waiting for an operation that will never be started or checkpointed. + * + *

A branch that has already published its own result keeps that result. The parallel snapshot calls this only + * while holding the same lock the result is published under. */ void completeAsSkipped() { + synchronized (settlementLock) { + if (isOperationCompleted() && !skipped.get()) { + return; + } + markSkipped(); + } + } + + /** + * Records this branch in a parallel completion snapshot. An incomplete branch is marked skipped before the lock is + * released, so a later completion cannot make {@code get()} succeed and then throw. + */ + ParallelResult.Status settleForParallelSnapshot() { + synchronized (settlementLock) { + if (skipped.get() || !isOperationCompleted()) { + markSkipped(); + return ParallelResult.Status.SKIPPED; + } + try { + get(); + return ParallelResult.Status.SUCCEEDED; + } catch (Throwable ignored) { + // A skip exception that propagated from the branch body is a failure of this branch. Only the explicit + // skipped flag above means the parallel operation itself skipped the branch. + return ParallelResult.Status.FAILED; + } + } + } + + private void markSkipped() { skipped.set(true); markAlreadyCompleted(); } @@ -239,6 +293,11 @@ boolean isSkipped() { return skipped.get(); } + @Override + public boolean isSkipSettlement() { + return skipped.get(); + } + @Override public T get() { waitForCompletionFuture(); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java index e2dbd90d35..d5b0681079 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java @@ -17,7 +17,6 @@ import software.amazon.lambda.durable.config.ParallelBranchConfig; import software.amazon.lambda.durable.config.ParallelConfig; import software.amazon.lambda.durable.context.DurableContextImpl; -import software.amazon.lambda.durable.exception.ParallelBranchSkippedException; import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; @@ -70,7 +69,9 @@ public ParallelOperation( protected void handleCompletion(CompletionConfig.CompletionDecision completionDecision) { var items = List.copyOf(getBranches()); - var statuses = items.stream().map(this::getParallelItemStatus).toList(); + var statuses = items.stream() + .map(ChildContextOperation::settleForParallelSnapshot) + .toList(); int succeededCount = Math.toIntExact(statuses.stream() .filter(s -> s == ParallelResult.Status.SUCCEEDED) .count()); @@ -86,7 +87,6 @@ protected void handleCompletion(CompletionConfig.CompletionDecision completionDe statuses); var serializedResult = serializeAndDeserializeResult(cachedResult); cachedResult = serializedResult.deserialized(); - completeSkippedBranches(items, statuses); // Branches added after checkpoint will not exist in the checkpointed result, but they'll be in the returned // value from get() method. @@ -97,28 +97,6 @@ protected void handleCompletion(CompletionConfig.CompletionDecision completionDe .contextOptions(ContextOptions.builder().replayChildren(true).build())); } - private void completeSkippedBranches(List> items, List statuses) { - for (int i = 0; i < items.size(); i++) { - if (statuses.get(i) == ParallelResult.Status.SKIPPED) { - items.get(i).completeAsSkipped(); - } - } - } - - private ParallelResult.Status getParallelItemStatus(ChildContextOperation childContextOperation) { - if (!childContextOperation.isOperationCompleted() || childContextOperation.isSkipped()) { - return ParallelResult.Status.SKIPPED; - } - try { - childContextOperation.get(); - return ParallelResult.Status.SUCCEEDED; - } catch (ParallelBranchSkippedException e) { - return ParallelResult.Status.SKIPPED; - } catch (Throwable t) { - return ParallelResult.Status.FAILED; - } - } - private ParallelResult rebuildParallelResult() { var statuses = new ArrayList<>(cachedResult.statuses()); while (statuses.size() < getBranches().size()) { diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/BaseDurableOperationFirstCompletionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/BaseDurableOperationFirstCompletionTest.java index d2b0c08758..4efaa36dfb 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/BaseDurableOperationFirstCompletionTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/BaseDurableOperationFirstCompletionTest.java @@ -17,6 +17,7 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import software.amazon.awssdk.services.lambda.model.CheckpointDurableExecutionResponse; @@ -138,6 +139,66 @@ void returnsAlreadySettledOperationWithoutParkingTheCaller() { assertFalse(executionManager.isExecutionCompletedExceptionally(), "caller must not have triggered suspension"); } + @Test + void returnsAlreadySettledOperationAheadOfAnEarlierSkippedFuture() { + var skipped = operation(FIRST_OP_ID); + var winner = operation(SECOND_OP_ID); + skipped.skip(); + winner.settle(); + + var operation = executionManager.waitForFirstOperationCompletion(List.of(skipped, winner)); + + assertEquals(winner, operation); + } + + @Test + void returnsLaterNonSkippedOperationWhenAnEarlierFutureWasSkipped() { + executionManager.registerActiveThread(OTHER_THREAD); + var skipped = operation(FIRST_OP_ID); + var winner = operation(SECOND_OP_ID); + skipped.skip(); + CompletableFuture.runAsync(winner::settle, CompletableFuture.delayedExecutor(50, MILLISECONDS)); + + var operation = awaitFirstSettled(List.of(skipped, winner)); + + assertEquals(winner, operation); + } + + @Test + void ignoresASkipThatSettlesWhileAnotherOperationIsStillRunning() { + executionManager.registerActiveThread(OTHER_THREAD); + var skipping = operation(FIRST_OP_ID); + var winner = operation(SECOND_OP_ID); + CompletableFuture.runAsync(skipping::skip, CompletableFuture.delayedExecutor(20, MILLISECONDS)); + CompletableFuture.runAsync(winner::settle, CompletableFuture.delayedExecutor(80, MILLISECONDS)); + + var operation = awaitFirstSettled(List.of(skipping, winner)); + + assertEquals(winner, operation); + } + + @Test + void suspendsWhenOnlySkippedFuturesHaveSettled() { + initExecution(List.of(pendingStepOperation(SECOND_OP_ID))); + var skipped = operation(FIRST_OP_ID); + var pending = operation(SECOND_OP_ID); + skipped.skip(); + + assertThrows( + SuspendExecutionException.class, + () -> executionManager.waitForFirstOperationCompletion(List.of(skipped, pending))); + } + + @Test + void returnsSkippedFutureWhenEveryFutureIsSkipped() { + var skipped = operation(FIRST_OP_ID); + skipped.skip(); + + var operation = executionManager.waitForFirstOperationCompletion(List.of(skipped)); + + assertEquals(skipped, operation); + } + @Test void suspendsWhenNoOperationCanSettle() { initExecution(List.of(pendingStepOperation(FIRST_OP_ID), pendingStepOperation(SECOND_OP_ID))); @@ -303,6 +364,8 @@ private TestOperation operation(String operationId) { /** Minimal operation that settles on demand, so the waiting behaviour can be driven directly. */ private static final class TestOperation extends BaseDurableOperation implements DurableFuture { + private final AtomicBoolean skipSettlement = new AtomicBoolean(); + private TestOperation(String operationId, DurableContextImpl durableContext) { super(OperationIdentifier.of(operationId, operationId, OperationSubType.STEP), durableContext, null); } @@ -313,6 +376,11 @@ protected void start() {} @Override protected void replay(Operation existing) {} + @Override + public boolean isSkipSettlement() { + return skipSettlement.get(); + } + @Override public Void get() { return null; @@ -321,6 +389,11 @@ public Void get() { private void settle() { markAlreadyCompleted(); } + + private void skip() { + skipSettlement.set(true); + settle(); + } } private static final class BlockingReactivationExecutionManager extends ExecutionManager { From a7d38acb4c20f35f16bc74965287bd79685540f3 Mon Sep 17 00:00:00 2001 From: hippi345 <13539685+hippi345@users.noreply.github.com> Date: Mon, 26 Oct 2026 20:26:56 +0000 Subject: [PATCH 4/4] fix(sdk): settle parallel branches registered after early completion firstSuccessful() can checkpoint a parallel operation before a later branch is registered. get() on that branch then suspends with no pending operation. Decide the skip and publish the completion snapshot under one lock, and settle any branch still missing from that result. --- .../ParallelBranchRegistrationRaceTest.java | 64 +++++++++++++ .../operation/ChildContextOperation.java | 11 ++- .../durable/operation/ParallelOperation.java | 91 ++++++++++++------- 3 files changed, 130 insertions(+), 36 deletions(-) create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/operation/ParallelBranchRegistrationRaceTest.java diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/operation/ParallelBranchRegistrationRaceTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/operation/ParallelBranchRegistrationRaceTest.java new file mode 100644 index 0000000000..5098960447 --- /dev/null +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/operation/ParallelBranchRegistrationRaceTest.java @@ -0,0 +1,64 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.operation; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.config.CompletionConfig; +import software.amazon.lambda.durable.config.ParallelConfig; +import software.amazon.lambda.durable.exception.ParallelBranchSkippedException; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class ParallelBranchRegistrationRaceTest { + + @AfterEach + void clearRegistrationPause() { + ParallelOperation.beforeBranchRegistration(null); + } + + @Test + void branchRegisteredAfterSnapshotIsSkippedInsteadOfSuspending() { + var calls = new AtomicInteger(); + ParallelOperation.beforeBranchRegistration(() -> { + if (calls.incrementAndGet() != 2) { + return; + } + try { + Thread.sleep(100); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(e); + } + }); + + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { + var config = ParallelConfig.builder() + .completionConfig(CompletionConfig.firstSuccessful()) + .maxConcurrency(1) + .build(); + var futures = new ArrayList>(); + var parallel = context.parallel("first-success-parallel", config); + try (parallel) { + futures.add(parallel.branch("branch-1", String.class, ctx -> "success-1")); + futures.add(parallel.branch("branch-2", String.class, ctx -> "skipped-2")); + } + var skipped = assertThrows( + ParallelBranchSkippedException.class, () -> futures.get(1).get()); + assertEquals("branch-2", skipped.getBranchName()); + return "done"; + }); + + var result = assertTimeoutPreemptively(Duration.ofSeconds(10), () -> runner.runUntilComplete("test")); + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("done", result.getResult(String.class)); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java index 371875098a..7527d4de76 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java @@ -300,10 +300,9 @@ public boolean isSkipSettlement() { @Override public T get() { + throwIfSkipped(); waitForCompletionFuture(); - if (skipped.get()) { - throw new ParallelBranchSkippedException(getName()); - } + throwIfSkipped(); var op = waitForOperationCompletion(); if (cachedOperationResult.get() != null) { // we have a result, just return it directly @@ -323,6 +322,12 @@ public T get() { } } + private void throwIfSkipped() { + if (skipped.get()) { + throw new ParallelBranchSkippedException(getName()); + } + } + private Throwable translateException(Operation op, ErrorObject errorObject) { // Attempt to reconstruct and throw the original exception Throwable original = deserializeException(errorObject); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java index d5b0681079..fa4252653e 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java @@ -46,9 +46,16 @@ */ public class ParallelOperation extends ConcurrencyOperation implements ParallelDurableFuture { + private static volatile Runnable beforeBranchRegistration; + // this field could be written and read in different threads private volatile ParallelResult cachedResult; private volatile ParallelResult partialResult; + private final Object branchMembership = new Object(); + + static void beforeBranchRegistration(Runnable action) { + beforeBranchRegistration = action; + } public ParallelOperation( OperationIdentifier operationIdentifier, @@ -67,26 +74,11 @@ public ParallelOperation( @Override protected void handleCompletion(CompletionConfig.CompletionDecision completionDecision) { - - var items = List.copyOf(getBranches()); - var statuses = items.stream() - .map(ChildContextOperation::settleForParallelSnapshot) - .toList(); - int succeededCount = Math.toIntExact(statuses.stream() - .filter(s -> s == ParallelResult.Status.SUCCEEDED) - .count()); - int failedCount = Math.toIntExact( - statuses.stream().filter(s -> s == ParallelResult.Status.FAILED).count()); - int skippedCount = items.size() - succeededCount - failedCount; - cachedResult = new ParallelResult( - items.size(), - succeededCount, - failedCount, - skippedCount, - completionDecision.completionStatus(), - statuses); + snapshotBranches(completionDecision); var serializedResult = serializeAndDeserializeResult(cachedResult); - cachedResult = serializedResult.deserialized(); + synchronized (branchMembership) { + cachedResult = serializedResult.deserialized(); + } // Branches added after checkpoint will not exist in the checkpointed result, but they'll be in the returned // value from get() method. @@ -97,12 +89,27 @@ protected void handleCompletion(CompletionConfig.CompletionDecision completionDe .contextOptions(ContextOptions.builder().replayChildren(true).build())); } - private ParallelResult rebuildParallelResult() { - var statuses = new ArrayList<>(cachedResult.statuses()); - while (statuses.size() < getBranches().size()) { - statuses.add(ParallelResult.Status.SKIPPED); + private void snapshotBranches(CompletionConfig.CompletionDecision completionDecision) { + synchronized (branchMembership) { + cachedResult = parallelResult(completionDecision, settleCurrentBranches()); } + } + private List settleCurrentBranches() { + var statuses = new ArrayList(); + var settled = 0; + do { + var items = List.copyOf(getBranches()); + for (var i = settled; i < items.size(); i++) { + statuses.add(items.get(i).settleForParallelSnapshot()); + } + settled = items.size(); + } while (getBranches().size() != settled); + return statuses; + } + + private static ParallelResult parallelResult( + CompletionConfig.CompletionDecision completionDecision, List statuses) { int succeededCount = Math.toIntExact(statuses.stream() .filter(status -> status == ParallelResult.Status.SUCCEEDED) .count()); @@ -114,10 +121,23 @@ private ParallelResult rebuildParallelResult() { succeededCount, failedCount, statuses.size() - succeededCount - failedCount, - cachedResult.completionStatus(), + completionDecision.completionStatus(), List.copyOf(statuses)); } + private ParallelResult rebuildParallelResult() { + synchronized (branchMembership) { + var branches = List.copyOf(getBranches()); + var statuses = new ArrayList<>(cachedResult.statuses()); + while (statuses.size() < branches.size()) { + branches.get(statuses.size()).completeAsSkipped(); + statuses.add(ParallelResult.Status.SKIPPED); + } + return parallelResult( + CompletionConfig.CompletionDecision.complete(cachedResult.completionStatus()), statuses); + } + } + @Override protected void start() { sendOperationUpdateAsync(OperationUpdate.builder() @@ -168,16 +188,21 @@ public DurableFuture branch( throw new IllegalStateException("Cannot add branches after join() has been called"); } - var nextBranchIndex = getBranches().size(); - - // Skip when the checkpointed or in-memory parallel result already completed without this branch. - var skip = shouldSkipBranch(nextBranchIndex); - var serDes = config.serDes() == null ? getContext().getDurableConfig().getSerDes() : config.serDes(); - var branch = enqueueItem(name, func, resultType, serDes, OperationSubType.PARALLEL_BRANCH, skip); - if (skip) { - branch.completeAsSkipped(); + var pause = beforeBranchRegistration; + if (pause != null) { + pause.run(); + } + synchronized (branchMembership) { + var nextBranchIndex = getBranches().size(); + var skip = shouldSkipBranch(nextBranchIndex); + var serDes = + config.serDes() == null ? getContext().getDurableConfig().getSerDes() : config.serDes(); + var branch = enqueueItem(name, func, resultType, serDes, OperationSubType.PARALLEL_BRANCH, skip); + if (skip) { + branch.completeAsSkipped(); + } + return branch; } - return branch; } private boolean shouldSkipBranch(int nextBranchIndex) {