diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndAdmissionTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndAdmissionTest.java index 9781b334f8..863ca67b5b 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndAdmissionTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndAdmissionTest.java @@ -25,9 +25,11 @@ import software.amazon.lambda.durable.config.ParallelConfig; import software.amazon.lambda.durable.config.WaitForConditionConfig; import software.amazon.lambda.durable.execution.DurableExecutor; +import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.model.DurableExecutionInput; import software.amazon.lambda.durable.model.ExecutionStatus; import software.amazon.lambda.durable.model.WaitForConditionResult; +import software.amazon.lambda.durable.operation.BaseDurableOperation; import software.amazon.lambda.durable.plugin.*; import software.amazon.lambda.durable.testing.local.LocalMemoryExecutionClient; @@ -38,7 +40,7 @@ void alreadyAdmittedWorkersRetainTheBaseCleanupBoundary(String mode) throws Exce boolean awaitResult = mode.equals("awaited"); boolean alreadyRunning = mode.equals("already-running"); var runningEntered = new CountDownLatch(1); - var resumedPredicate = new CountDownLatch(1); + var conditionStopped = new CountDownLatch(1); var releaseRunning = new CountDownLatch(1); var nextDispatch = new CountDownLatch(1); var releaseDispatch = new CountDownLatch(1); @@ -117,8 +119,8 @@ public void onOperationEnd(OperationEndInfo info) { } }; // The merged lifecycle unwinds End in reverse order, so OTel must be registered after the gate. - // This fixture parks inside Executor.execute AFTER SDK admission. The remaining worker-start boundary - // is tracked in #787; new registration during preparation is covered separately without weakening that gate. + // This fixture pauses a continuation after admission. Draining must stop its owner and join the + // continuation before End. A predicate that already started must finish its checkpoint first. var config = DurableConfig.builder() .withExecutorService(workers) .withDurableExecutionClient(client) @@ -133,25 +135,31 @@ public void onOperationEnd(OperationEndInfo info) { null, TypeToken.get(String.class), (value, context) -> { - Function> startCondition = - branch -> branch.waitForConditionAsync( - "condition", - Integer.class, - (state, step) -> { - predicates.incrementAndGet(); - if (state == 2) resumedPredicate.countDown(); - if (alreadyRunning && state == 2) { - runningEntered.countDown(); - await(releaseRunning); - } - return state == 1 - ? WaitForConditionResult.continuePolling(2) - : WaitForConditionResult.stopPolling(state); - }, - WaitForConditionConfig.builder() - .initialState(1) - .waitStrategy((state, attempt) -> Duration.ofSeconds(1)) - .build()); + Function> startCondition = branch -> { + var condition = branch.waitForConditionAsync( + "condition", + Integer.class, + (state, step) -> { + predicates.incrementAndGet(); + if (alreadyRunning && state == 2) { + runningEntered.countDown(); + await(releaseRunning); + } + return state == 1 + ? WaitForConditionResult.continuePolling(2) + : WaitForConditionResult.stopPolling(state); + }, + WaitForConditionConfig.builder() + .initialState(1) + .waitStrategy((state, attempt) -> Duration.ofSeconds(1)) + .build()); + ((BaseDurableOperation) condition) + .getCompletionFuture() + .whenComplete((result, failure) -> { + if (failure instanceof SuspendExecutionException) conditionStopped.countDown(); + }); + return condition; + }; if (mode.equals("early-parallel")) { var parallel = context.parallel( "early", @@ -182,17 +190,19 @@ public void onOperationEnd(OperationEndInfo info) { if (alreadyRunning) { releaseDispatch.countDown(); assertTrue(runningEntered.await(3, TimeUnit.SECONDS)); - assertTrue(endEntered.await(3, TimeUnit.SECONDS), "End does not join already-running predicates"); - releaseEnd.countDown(); + assertEquals(1, endEntered.getCount(), "End must wait for the running predicate"); assertThrows( TimeoutException.class, () -> call.get(100, TimeUnit.MILLISECONDS), - "Normal close still waits for the real handler"); + "Draining waits for the real handler before End"); releaseRunning.countDown(); - } else if (!awaitResult) assertTrue(endEntered.await(3, TimeUnit.SECONDS)); + } else if (!awaitResult) { + assertThrows(TimeoutException.class, () -> call.get(100, TimeUnit.MILLISECONDS)); + assertEquals(1, endEntered.getCount(), "End must join the admitted continuation"); + assertTrue(conditionStopped.await(3, TimeUnit.SECONDS), "Draining must release the condition owner"); + } releaseDispatch.countDown(); - if (!awaitResult && !alreadyRunning) assertTrue(resumedPredicate.await(3, TimeUnit.SECONDS)); - assertTrue(terminal.await(3, TimeUnit.SECONDS)); + if (awaitResult || alreadyRunning) assertTrue(terminal.await(3, TimeUnit.SECONDS)); assertTrue(endEntered.await(3, TimeUnit.SECONDS)); releaseEnd.countDown(); assertEquals( @@ -205,23 +215,14 @@ public void onOperationEnd(OperationEndInfo info) { System.out.println("END_SNAPSHOT_PUBLIC mode=" + mode + " snapshot=" + snapshotStatus.get() + " backend=" + client.getOperationByName("condition").status() + " predicates=" + predicates.get() + " lateHooks=" + lateHooks.get() + " lateStarts=" + lateStarts.get() + " operationSpans=" + spans); + assertEquals(0, lateStarts.get(), "Drained continuations must not start user functions after End"); + assertEquals(0, lateHooks.get()); + assertEquals(awaitResult || alreadyRunning ? 2 : 1, predicates.get()); assertEquals( - awaitResult || alreadyRunning ? 0 : 1, - lateStarts.get(), - "Already-admitted workers may start after End; this is not a new-registration guarantee"); - assertEquals(2, predicates.get()); - if (alreadyRunning) { - assertNotEquals( - OperationStatus.SUCCEEDED, snapshotStatus.get(), "The existing snapshot boundary is retained"); - assertEquals( - OperationStatus.SUCCEEDED, - client.getOperationByName("condition").status()); - } else { - assertEquals(awaitResult ? OperationStatus.SUCCEEDED : OperationStatus.READY, snapshotStatus.get()); - assertEquals( - OperationStatus.SUCCEEDED, - client.getOperationByName("condition").status()); - } + awaitResult || alreadyRunning ? OperationStatus.SUCCEEDED : OperationStatus.READY, + snapshotStatus.get()); + assertEquals( + snapshotStatus.get(), client.getOperationByName("condition").status()); } finally { releaseRunning.countDown(); releaseDispatch.countDown(); diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationLateFlushProbeTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationLateFlushProbeTest.java index 108756b60f..b6bb682aaa 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationLateFlushProbeTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationLateFlushProbeTest.java @@ -9,6 +9,7 @@ import io.opentelemetry.sdk.trace.export.BatchSpanProcessor; import java.time.Duration; import java.time.Instant; +import java.util.List; import java.util.concurrent.*; import java.util.concurrent.atomic.*; import org.junit.jupiter.params.ParameterizedTest; @@ -24,11 +25,11 @@ import software.amazon.lambda.durable.plugin.*; import software.amazon.lambda.durable.testing.local.LocalMemoryExecutionClient; -/** Diagnostic for the pre-existing End/close flush boundary; no proposed production policy in this class. */ +/** Verifies that draining operations precedes invocation End and its telemetry flush. */ class InvocationLateFlushProbeTest { @ParameterizedTest @ValueSource(booleans = {false, true}) - void alreadyRunningStepCompletesAfterEnd(boolean earlyParallel) throws Exception { + void alreadyRunningStepCompletesBeforeEnd(boolean earlyParallel) throws Exception { var running = new CountDownLatch(1); var release = new CountDownLatch(1); var ended = new CountDownLatch(1); @@ -101,9 +102,11 @@ public void onOperationEnd(OperationEndInfo info) { return "root-success"; }, config)); - assertTrue(ended.await(5, TimeUnit.SECONDS)); + assertTrue(running.await(5, TimeUnit.SECONDS)); assertThrows(TimeoutException.class, () -> call.get(100, TimeUnit.MILLISECONDS)); + assertEquals(1, ended.getCount(), "End must wait for the running step"); release.countDown(); + assertTrue(ended.await(5, TimeUnit.SECONDS)); assertEquals( ExecutionStatus.SUCCEEDED, call.get(5, TimeUnit.SECONDS).status()); assertEquals( @@ -122,8 +125,9 @@ public void onOperationEnd(OperationEndInfo info) { System.out.println("LATE_FLUSH_PUBLIC earlyParallel=" + earlyParallel + " backend=SUCCEEDED bodyCalls=" + bodyCalls.get() + " lateOperationEnds=" + lateOperationEnds.get() + " beforeManualFlush=" + before + " afterManualFlush=" + after); - assertTrue(lateOperationEnds.get() > 0, "Existing public operation-end occurs after invocation End"); - // This diagnostic records actual export behavior; it does not redefine which outcome wins. + assertEquals(0, lateOperationEnds.get(), "Operation End must precede invocation End"); + assertEquals(List.of("SUCCEEDED"), before, "Invocation End must flush the completed step"); + assertEquals(before, after, "Manual flush must not reveal an unexported step"); } finally { release.countDown(); caller.shutdownNow(); diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationPreparationAdmissionProbeTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationPreparationAdmissionProbeTest.java index 665bca59e5..1c61311f4e 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationPreparationAdmissionProbeTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationPreparationAdmissionProbeTest.java @@ -6,9 +6,11 @@ import java.time.Duration; import java.time.Instant; +import java.util.Arrays; import java.util.List; import java.util.concurrent.*; import java.util.concurrent.atomic.*; +import java.util.concurrent.locks.LockSupport; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.CsvSource; @@ -48,6 +50,8 @@ void preparationAdmissionWindow(String mode, boolean rejectResume) throws Except var preparing = new AtomicBoolean(); var predicates = new AtomicInteger(); var preparationDispatches = new AtomicInteger(); + var resumedDispatches = new AtomicInteger(); + var rootThread = new AtomicReference(); var childFailure = new AtomicReference(); var ends = new CopyOnWriteArrayList(); var rejection = new RejectedExecutionException("controlled READY admission rejection"); @@ -56,6 +60,7 @@ void preparationAdmissionWindow(String mode, boolean rejectResume) throws Except @Override public void execute(Runnable task) { if (backendReleased.get()) { + resumedDispatches.incrementAndGet(); if (preparing.get()) preparationDispatches.incrementAndGet(); if (rejectResume) throw rejection; } @@ -107,6 +112,7 @@ public void onInvocationEnd(InvocationEndInfo info) { null, TypeToken.get(String.class), (value, root) -> { + rootThread.set(Thread.currentThread()); if (mode.equals("awaited")) { root.runInChildContextAsync("keep-active", String.class, child -> { keeperEntered.countDown(); @@ -151,13 +157,14 @@ public void onInvocationEnd(InvocationEndInfo info) { if (call.isDone()) System.out.println("PROBE_EARLY_RESULT " + call.get()); fail("Real backend did not reach PENDING"); } - if (!mode.equals("awaited")) - assertTrue( - preparationEntered.await(5, TimeUnit.SECONDS), - "Root success selected before real output serializer"); + // Draining now precedes serialization. Release READY only after the root reaches that boundary. + // A child waiting on the condition must settle before output preparation can start. + if (!mode.equals("awaited")) awaitDrain(rootThread.get()); backendReleased.set(true); client.advanceTime(); assertTrue(conditionFinished.await(5, TimeUnit.SECONDS), "Real durable waiter settled after backend READY"); + if (!mode.equals("awaited")) + assertTrue(preparationEntered.await(5, TimeUnit.SECONDS), "Preparation starts after the child drains"); releasePreparation.countDown(); releaseKeeper.countDown(); if (mode.equals("awaited") && rejectResume) { @@ -182,8 +189,8 @@ public void onInvocationEnd(InvocationEndInfo info) { + " backend=" + client.getOperationByName("condition").status() + " End=" + ends.get(0).invocationStatus()); if (!mode.equals("awaited")) { - // Proposed admission guarantee, distinct from replacing the selected SUCCEEDED outcome. - assertEquals(0, preparationDispatches.get(), "No new READY user-worker admission after root selection"); + assertEquals(0, resumedDispatches.get(), "No new READY user-worker admission while draining"); + assertEquals(0, preparationDispatches.get(), "No new READY user-worker admission during preparation"); assertEquals(1, predicates.get()); } else assertEquals(rejectResume ? 1 : 2, predicates.get()); } finally { @@ -262,6 +269,16 @@ private static int condition( } } + private static void awaitDrain(Thread root) { + var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(3); + while (System.nanoTime() < deadline) { + if (Arrays.stream(root.getStackTrace()) + .anyMatch(frame -> frame.getMethodName().equals("waitForRunningUserHandlers"))) return; + LockSupport.parkNanos(100_000); + } + fail("Root did not begin draining the pending child"); + } + private static void await(CountDownLatch latch) { try { assertTrue(latch.await(8, TimeUnit.SECONDS), "Probe gate released within bounded cleanup"); diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationWorkerContextCleanupProbeTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationWorkerContextCleanupProbeTest.java index baddca6832..9530e75bdb 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationWorkerContextCleanupProbeTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationWorkerContextCleanupProbeTest.java @@ -30,7 +30,7 @@ import software.amazon.lambda.durable.plugin.*; import software.amazon.lambda.durable.testing.local.LocalMemoryExecutionClient; -/** Diagnostic for the pre-existing End/close flush boundary; no proposed production policy in this class. */ +/** Verifies worker context restoration and export when operations drain before invocation End. */ class InvocationWorkerContextCleanupProbeTest { @ParameterizedTest @CsvSource({ @@ -41,7 +41,7 @@ class InvocationWorkerContextCleanupProbeTest { "early-parallel,true", "awaited,true" }) - void alreadyRunningStepCompletesAfterEnd(String mode, boolean executionView) throws Exception { + void alreadyRunningStepCompletesBeforeEnd(String mode, boolean executionView) throws Exception { boolean earlyParallel = mode.equals("early-parallel"); boolean awaited = mode.equals("awaited"); var running = new CountDownLatch(1); @@ -138,13 +138,11 @@ public void onOperationEnd(OperationEndInfo info) { return "root-success"; }, config)); - if (awaited) { - assertTrue(running.await(5, TimeUnit.SECONDS)); - release.countDown(); - } - assertTrue(ended.await(5, TimeUnit.SECONDS)); - if (!awaited) assertThrows(TimeoutException.class, () -> call.get(100, TimeUnit.MILLISECONDS)); + assertTrue(running.await(5, TimeUnit.SECONDS)); + assertThrows(TimeoutException.class, () -> call.get(100, TimeUnit.MILLISECONDS)); + assertEquals(1, ended.getCount(), "End must wait for the running step"); release.countDown(); + assertTrue(ended.await(5, TimeUnit.SECONDS)); assertEquals( ExecutionStatus.SUCCEEDED, call.get(5, TimeUnit.SECONDS).status()); assertEquals( @@ -164,9 +162,9 @@ public void onOperationEnd(OperationEndInfo info) { System.out.println("LATE_FLUSH_PUBLIC earlyParallel=" + earlyParallel + " backend=SUCCEEDED bodyCalls=" + bodyCalls.get() + " lateOperationEnds=" + lateOperationEnds.get() + " beforeManualFlush=" + before + " afterManualFlush=" + after); - if (!awaited) - assertTrue( - lateOperationEnds.get() > 0, "Existing public operation-end occurs after invocation End"); + assertEquals(0, lateOperationEnds.get(), "Operation End must precede invocation End"); + assertEquals(List.of("SUCCEEDED"), before, "Invocation End must flush the completed step"); + assertEquals(before, after, "Manual flush must not reveal an unexported step"); System.out.println("WORKER_CONTEXT_PUBLIC executionView=" + executionView + " mode=" + mode + " restored=" + (initial.get(owner.get()) == workerContext.get()) + " remainingSpan=" + Span.fromContext(workerContext.get()).getSpanContext()); @@ -174,7 +172,6 @@ public void onOperationEnd(OperationEndInfo info) { initial.get(owner.get()), workerContext.get(), "Real worker ambient Context restored after its task finishes"); - // This diagnostic records actual export behavior; it does not redefine which outcome wins. } finally { release.countDown(); caller.shutdownNow(); @@ -297,8 +294,11 @@ public void onInvocationEnd(InvocationEndInfo info) { return "root-success"; }, config)); - assertTrue(ended.await(5, TimeUnit.SECONDS)); + assertTrue(running.await(5, TimeUnit.SECONDS)); + assertThrows(TimeoutException.class, () -> call.get(100, TimeUnit.MILLISECONDS)); + assertEquals(1, ended.getCount(), "End must wait for the running step"); release.countDown(); + assertTrue(ended.await(5, TimeUnit.SECONDS)); assertEquals( ExecutionStatus.SUCCEEDED, call.get(5, TimeUnit.SECONDS).status()); diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/CheckpointTokenRevocationIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/CheckpointTokenRevocationIntegrationTest.java new file mode 100644 index 0000000000..29a487efcd --- /dev/null +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/CheckpointTokenRevocationIntegrationTest.java @@ -0,0 +1,145 @@ +// 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.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static software.amazon.lambda.durable.config.StepSemantics.AT_MOST_ONCE_PER_RETRY; +import static software.amazon.lambda.durable.retry.RetryStrategies.Presets.NO_RETRY; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import software.amazon.awssdk.services.lambda.model.OperationStatus; +import software.amazon.lambda.durable.config.StepConfig; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; +import software.amazon.lambda.durable.plugin.InvocationInfo; +import software.amazon.lambda.durable.plugin.OperationEndInfo; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class CheckpointTokenRevocationIntegrationTest { + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void acceptedStepOutcomeReplaysAfterResumeWithoutRepeatingUserCode(boolean fail) { + var calls = new AtomicInteger(); + var runner = createPausedRunner(fail, calls); + + assertEquals(ExecutionStatus.PENDING, runner.runUntilComplete("input").getStatus()); + assertEquals(1, calls.get()); + runner.resumeExecution(); + var resumed = runner.runUntilComplete("input"); + + assertEquals(ExecutionStatus.SUCCEEDED, resumed.getStatus()); + assertEquals(fail ? "stored-error" : "stored-result", resumed.getResult(String.class)); + assertEquals(1, calls.get()); + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void acceptedStepOutcomeIsReportedToPluginsAfterResume(boolean fail) { + var calls = new AtomicInteger(); + var plugin = new RecordingPlugin(); + var config = DurableConfig.builder().withPlugins(plugin).build(); + var runner = createPausedRunner(fail, calls, config); + + assertEquals(ExecutionStatus.PENDING, runner.runUntilComplete("input").getStatus()); + assertTrue(plugin.operationEnds.isEmpty()); + runner.resumeExecution(); + assertEquals(ExecutionStatus.SUCCEEDED, runner.runUntilComplete("input").getStatus()); + + assertPluginObservedAcceptedOutcome(plugin, fail); + } + + private LocalDurableTestRunner createPausedRunner(boolean fail, AtomicInteger calls) { + return createPausedRunner(fail, calls, DurableConfig.builder().build()); + } + + private LocalDurableTestRunner createPausedRunner( + boolean fail, AtomicInteger calls, DurableConfig config) { + var runnerRef = new AtomicReference>(); + var runner = LocalDurableTestRunner.create( + String.class, + (input, ctx) -> { + try { + return ctx.step( + "once", + String.class, + stepCtx -> { + calls.incrementAndGet(); + runnerRef.get().pauseExecution(); + if (fail) { + throw new IllegalArgumentException("stored-error"); + } + return "stored-result"; + }, + noRetryConfig()); + } catch (IllegalArgumentException e) { + return e.getMessage(); + } + }, + config); + runnerRef.set(runner); + return runner; + } + + private void assertPluginObservedAcceptedOutcome(RecordingPlugin plugin, boolean fail) { + assertEquals(2, plugin.invocationStarts.size()); + var updatedOperation = plugin.invocationStarts.get(1).updatedOperations().values().stream() + .filter(operation -> "once".equals(operation.name())) + .findFirst() + .orElseThrow(); + var expectedStatus = fail ? OperationStatus.FAILED : OperationStatus.SUCCEEDED; + assertEquals(expectedStatus, updatedOperation.status()); + assertTrue(updatedOperation.isReplay()); + assertStoredOutcome(updatedOperation.result(), updatedOperation.error(), fail); + + var operationEnds = plugin.operationEnds.stream() + .filter(operation -> "once".equals(operation.name())) + .toList(); + assertEquals(1, operationEnds.size()); + var operationEnd = operationEnds.get(0); + assertEquals(expectedStatus.toString(), operationEnd.status()); + assertTrue(operationEnd.isReplay()); + assertStoredOutcome(operationEnd.result(), operationEnd.error(), fail); + } + + private void assertStoredOutcome(String result, Throwable error, boolean fail) { + if (fail) { + assertNull(result); + assertNotNull(error); + assertTrue(error.getMessage().contains("stored-error")); + } else { + assertEquals("\"stored-result\"", result); + assertNull(error); + } + } + + private StepConfig noRetryConfig() { + return StepConfig.builder() + .semanticsPerRetry(AT_MOST_ONCE_PER_RETRY) + .retryStrategy(NO_RETRY) + .build(); + } + + private static class RecordingPlugin implements DurableExecutionPlugin { + private final List invocationStarts = Collections.synchronizedList(new ArrayList<>()); + private final List operationEnds = Collections.synchronizedList(new ArrayList<>()); + + @Override + public void onInvocationStart(InvocationInfo info) { + invocationStarts.add(info); + } + + @Override + public void onOperationEnd(OperationEndInfo info) { + operationEnds.add(info); + } + } +} diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/InvocationFinalizationIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/InvocationFinalizationIntegrationTest.java index 0e41ada4b3..0471a27953 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/InvocationFinalizationIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/InvocationFinalizationIntegrationTest.java @@ -164,6 +164,8 @@ void naturalManagerOutcomePrecedesFastThrowingFinally(String kind, String mode, var workers = new OutcomeExecutor(mode); var releaseOperation = new CountDownLatch(mode.equals("submit-wait") ? 0 : 1); var aboutToWait = new CountDownLatch(1); + var handlerFinished = new CountDownLatch(1); + var controlSelectedBeforeCallback = new AtomicBoolean(); var completedStepCalls = new AtomicInteger(); var targetCalls = new AtomicInteger(); var throwFinally = new AtomicBoolean(true); @@ -209,19 +211,22 @@ void naturalManagerOutcomePrecedesFastThrowingFinally(String kind, String mode, .retryStrategy(RetryStrategies.fixedDelay(2, Duration.ofSeconds(1))) .build()); } - // This schedules the existing completion window; suspension/termination is triggered only by the - // real - // operation. The root's dependent wait is registered later and wakes before this earlier callback - // drains. + var manager = ((DurableContextImpl) context).getExecutionManager(); + // The manager must select suspension or termination before waking operation waiters. + // Observe that ordering directly, before the handler's finally block can change the outcome. + // Wait for the handler body, not its SDK task: the SDK task drains this callback before returning. ((BaseDurableOperation) future).getCompletionFuture().whenComplete((value, failure) -> { - if (failure != null && !mode.equals("submit-wait") && throwFinally.get()) - await(workers.rootReturned); + if (failure != null && throwFinally.get()) { + controlSelectedBeforeCallback.set(manager.isExecutionCompletedExceptionally()); + if (!mode.equals("submit-wait")) await(handlerFinished); + } }); try { aboutToWait.countDown(); future.get(); return "done"; } finally { + handlerFinished.countDown(); if (throwFinally.get()) throw new IllegalStateException("fast handler finally"); } }, @@ -254,6 +259,7 @@ void naturalManagerOutcomePrecedesFastThrowingFinally(String kind, String mode, ExecutionStatus.PENDING, response.get(5, TimeUnit.SECONDS).getStatus()); } + assertTrue(controlSelectedBeforeCallback.get(), "Manager control must be selected before waking waiters"); if (withPlugin) assertPaired(plugin, kind.equals("termination") ? InvocationStatus.RETRYING : InvocationStatus.PENDING); throwFinally.set(false); @@ -262,7 +268,7 @@ void naturalManagerOutcomePrecedesFastThrowingFinally(String kind, String mode, assertEquals(1, completedStepCalls.get(), "The completed step must replay without another side effect"); } finally { releaseOperation.countDown(); - workers.rootReturned.countDown(); + handlerFinished.countDown(); stop(caller); stop(workers); } @@ -273,7 +279,6 @@ private static final class OutcomeExecutor extends AbstractExecutorService { private final ExecutorService delegate = Executors.newCachedThreadPool(); private final AtomicBoolean first = new AtomicBoolean(true); private final AtomicReference rootOwner = new AtomicReference<>(); - private final CountDownLatch rootReturned = new CountDownLatch(1); private OutcomeExecutor(String mode) { this.mode = mode; @@ -283,11 +288,7 @@ public void execute(Runnable task) { var root = first.getAndSet(false); Runnable wrapped = () -> { if (root) rootOwner.set(Thread.currentThread()); - try { - task.run(); - } finally { - if (root) rootReturned.countDown(); - } + task.run(); }; if (root && mode.equals("inline-root")) wrapped.run(); else if (mode.equals("submit-wait")) { diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/execution/ContinuationShutdownIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/execution/ContinuationShutdownIntegrationTest.java index 94b3489f8d..800d94c7f9 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/execution/ContinuationShutdownIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/execution/ContinuationShutdownIntegrationTest.java @@ -114,8 +114,10 @@ void runningPredicateFinishesItsCheckpointAndReplaysWithoutRepeatingWork() throw } } - @Test - void earlyParallelCompletionStopsOnlyItsQueuedConditionAndUnblocksTheWaitingBranch() throws Exception { + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void earlyParallelCompletionStopsOnlyItsQueuedConditionAndUnblocksTheWaitingBranch(boolean beforeEnd) + throws Exception { var closed = new AtomicBoolean(); var late = new AtomicInteger(); var client = readyClient(closed, late); @@ -162,8 +164,11 @@ void earlyParallelCompletionStopsOnlyItsQueuedConditionAndUnblocksTheWaitingBran assertEquals(1, result.succeeded()); assertFalse(condition.get().getCompletionFuture().isDone()); var checkpointCount = client.getOperationUpdates().size(); - var closing = CompletableFuture.runAsync(setup.manager::close); - assertTrue(setup.manager.closeEntered.await(3, TimeUnit.SECONDS)); + var closing = CompletableFuture.runAsync(() -> { + if (beforeEnd) setup.manager.drainOperations(); + else setup.manager.close(); + }); + if (!beforeEnd) assertTrue(setup.manager.closeEntered.await(3, TimeUnit.SECONDS)); assertThrows(TimeoutException.class, () -> closing.get(150, TimeUnit.MILLISECONDS)); CompletableFuture.runAsync(task.get()).get(3, TimeUnit.SECONDS); closing.get(3, TimeUnit.SECONDS); diff --git a/sdk-testing/src/main/java/software/amazon/lambda/durable/testing/CloudDurableTestRunner.java b/sdk-testing/src/main/java/software/amazon/lambda/durable/testing/CloudDurableTestRunner.java index f268c82706..fa30098efe 100644 --- a/sdk-testing/src/main/java/software/amazon/lambda/durable/testing/CloudDurableTestRunner.java +++ b/sdk-testing/src/main/java/software/amazon/lambda/durable/testing/CloudDurableTestRunner.java @@ -10,6 +10,7 @@ import software.amazon.awssdk.services.lambda.model.InvocationType; import software.amazon.awssdk.services.lambda.model.InvokeRequest; import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.annotations.Experimental; import software.amazon.lambda.durable.serde.JacksonSerDes; import software.amazon.lambda.durable.serde.SerDes; import software.amazon.lambda.durable.testing.cloud.HistoryEventProcessor; @@ -216,4 +217,21 @@ public TestOperation getOperation(String name) { } return lastResult.getOperation(name); } + + /** + * Always throws. Withholding a checkpoint token is a behaviour of the in-memory checkpoint client, and a cloud + * execution checkpoints against the service instead. + */ + @Experimental + public void pauseExecution() { + throw new UnsupportedOperationException( + "pauseExecution is supported by LocalDurableTestRunner only, not by CloudDurableTestRunner"); + } + + /** Always throws, for the same reason as {@link #pauseExecution()}. */ + @Experimental + public void resumeExecution() { + throw new UnsupportedOperationException( + "resumeExecution is supported by LocalDurableTestRunner only, not by CloudDurableTestRunner"); + } } diff --git a/sdk-testing/src/main/java/software/amazon/lambda/durable/testing/LocalDurableTestRunner.java b/sdk-testing/src/main/java/software/amazon/lambda/durable/testing/LocalDurableTestRunner.java index 5be36c45eb..0eaf52c11b 100644 --- a/sdk-testing/src/main/java/software/amazon/lambda/durable/testing/LocalDurableTestRunner.java +++ b/sdk-testing/src/main/java/software/amazon/lambda/durable/testing/LocalDurableTestRunner.java @@ -19,6 +19,7 @@ import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.annotations.Experimental; import software.amazon.lambda.durable.execution.DurableExecutor; import software.amazon.lambda.durable.model.DurableExecutionInput; import software.amazon.lambda.durable.model.ExecutionStatus; @@ -281,9 +282,10 @@ public TestResult runUntilComplete(I input) { for (int i = 0; i < MAX_INVOCATIONS; i++) { result = run(input); - if (result.getStatus() != ExecutionStatus.PENDING || !storage.advanceTime()) { + if (result.getStatus() != ExecutionStatus.PENDING || storage.isPaused() || !storage.advanceTime()) { // break the loop if // - Return SUCCEEDED or FAILED - we're done + // - Paused: no checkpoint can be recorded, so every further invocation would report PENDING again // - Return PENDING and let test manually advance operations if no operations can be auto advanced break; } @@ -291,6 +293,30 @@ public TestResult runUntilComplete(I input) { return result; } + /** + * Answers every later checkpoint without a checkpoint token, until {@link #resumeExecution()}. + * + *

The invocation running at the time stops checkpointing and returns PENDING. The checkpoint that received the + * token-less response was recorded, so its updates survive; work the SDK had not yet sent is abandoned and replays + * after {@link #resumeExecution()}. {@link #runUntilComplete(Object)} returns as soon as it sees that PENDING, + * because no further invocation could record anything either. + * + *

Advance waits and complete callbacks while paused as usual; the next invocation observes them only once the + * execution is resumed, since this runner invokes the handler only when a test asks it to. + * + *

Idempotent. + */ + @Experimental + public void pauseExecution() { + storage.pause(); + } + + /** Answers checkpoints with a checkpoint token again, so the next {@code run} replays the abandoned work. */ + @Experimental + public void resumeExecution() { + storage.resume(); + } + /** Resets a named step operation to STARTED status, simulating a checkpoint failure. */ public void resetCheckpointToStarted(String stepName) { storage.resetCheckpointToStarted(stepName); diff --git a/sdk-testing/src/main/java/software/amazon/lambda/durable/testing/local/LocalMemoryExecutionClient.java b/sdk-testing/src/main/java/software/amazon/lambda/durable/testing/local/LocalMemoryExecutionClient.java index 4a2f0c91a2..11532d60eb 100644 --- a/sdk-testing/src/main/java/software/amazon/lambda/durable/testing/local/LocalMemoryExecutionClient.java +++ b/sdk-testing/src/main/java/software/amazon/lambda/durable/testing/local/LocalMemoryExecutionClient.java @@ -38,29 +38,51 @@ public class LocalMemoryExecutionClient implements DurableExecutionClient { private final List operationUpdates = new CopyOnWriteArrayList<>(); private final Map updatedOperations = new HashMap<>(); private final Set operationIdsUpdatedSinceLastInvocation = new HashSet<>(); + private volatile boolean paused; @Override public CheckpointDurableExecutionResponse checkpoint(String arn, String token, List updates) { operationUpdates.addAll(updates); updates.forEach(update -> applyUpdate(update, true)); - var newToken = UUID.randomUUID().toString(); - CheckpointDurableExecutionResponse response; synchronized (updatedOperations) { - response = CheckpointDurableExecutionResponse.builder() - .checkpointToken(newToken) + var responseBuilder = CheckpointDurableExecutionResponse.builder() .newExecutionState(CheckpointUpdatedExecutionState.builder() .operations(updatedOperations.values()) - .build()) - .build(); + .build()); + if (!paused) { + responseBuilder.checkpointToken(UUID.randomUUID().toString()); + } + response = responseBuilder.build(); // updatedOperations was copied into response, so clearing it is safe here updatedOperations.clear(); } + if (response.checkpointToken() == null) { + updates.forEach(update -> operationIdsUpdatedSinceLastInvocation.add(update.id())); + } return response; } + /** + * Answers every later checkpoint without a checkpoint token, until {@link #resume()}. Those checkpoints are still + * recorded, so the next invocation replays them. + */ + public void pause() { + paused = true; + } + + /** Answers checkpoints with a checkpoint token again. */ + public void resume() { + paused = false; + } + + /** Returns {@code true} while checkpoints are answered without a checkpoint token. */ + public boolean isPaused() { + return paused; + } + @Override public GetDurableExecutionStateResponse getExecutionState(String arn, String checkpointToken, String marker) { // local runner doesn't use this API at all @@ -257,9 +279,9 @@ private void updateOperation(OperationUpdate update, Operation op, boolean withi synchronized (updatedOperations) { updatedOperations.put(op.id(), op); } - // Only track operations updated outside of a checkpoint call (i.e., between invocations) - // for the updatedOperationIds field. Operations updated during a checkpoint are already - // visible to the SDK via the checkpoint response. + // Only track operations updated outside of a checkpoint call (i.e., between invocations) here. + // Checkpoint updates are tracked by checkpoint() when its tokenless response ends the invocation + // before the SDK can observe the returned state. if (!withinCheckpoint) { operationIdsUpdatedSinceLastInvocation.add(op.id()); } diff --git a/sdk-testing/src/test/java/software/amazon/lambda/durable/testing/CloudDurableTestRunnerTest.java b/sdk-testing/src/test/java/software/amazon/lambda/durable/testing/CloudDurableTestRunnerTest.java index b1fde42e93..bf1962ca18 100644 --- a/sdk-testing/src/test/java/software/amazon/lambda/durable/testing/CloudDurableTestRunnerTest.java +++ b/sdk-testing/src/test/java/software/amazon/lambda/durable/testing/CloudDurableTestRunnerTest.java @@ -31,4 +31,15 @@ void testPlaceholderMethods() { assertThrows(IllegalStateException.class, () -> runner.getOperation("test")); } + + @Test + void pauseAndResumeExecutionAreNotSupported() { + var mockClient = mock(LambdaClient.class); + var runner = CloudDurableTestRunner.create( + "arn:aws:lambda:us-east-2:123:function:test", String.class, String.class, mockClient); + + assertThrows(UnsupportedOperationException.class, runner::pauseExecution); + assertThrows(UnsupportedOperationException.class, runner::resumeExecution); + verifyNoInteractions(mockClient); + } } diff --git a/sdk-testing/src/test/java/software/amazon/lambda/durable/testing/LocalDurableTestRunnerTest.java b/sdk-testing/src/test/java/software/amazon/lambda/durable/testing/LocalDurableTestRunnerTest.java index 3060f4651b..0b41339645 100644 --- a/sdk-testing/src/test/java/software/amazon/lambda/durable/testing/LocalDurableTestRunnerTest.java +++ b/sdk-testing/src/test/java/software/amazon/lambda/durable/testing/LocalDurableTestRunnerTest.java @@ -10,6 +10,7 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.Test; import software.amazon.awssdk.services.lambda.model.Operation; import software.amazon.lambda.durable.DurableConfig; @@ -121,6 +122,56 @@ public void onInvocationStart(InvocationInfo info) { assertEquals(executionStartTimes.get(0), executionStartTimes.get(1)); } + @Test + void pausedExecutionReportsPendingAndFinishesAfterResume() { + var runner = LocalDurableTestRunner.create(String.class, (input, ctx) -> { + ctx.step("only-step", String.class, stepCtx -> "stepped"); + return "done"; + }); + + runner.pauseExecution(); + var paused = runner.runUntilComplete("input"); + + assertEquals(ExecutionStatus.PENDING, paused.getStatus()); + + runner.resumeExecution(); + var resumed = runner.runUntilComplete("input"); + + assertEquals(ExecutionStatus.SUCCEEDED, resumed.getStatus()); + assertEquals("done", resumed.getResult(String.class)); + assertEquals("stepped", runner.getOperation("only-step").getStepResult(String.class)); + } + + @Test + void runUntilCompleteStopsWhilePausedAndFinishesAfterResume() { + var stepRuns = new AtomicInteger(); + var runner = LocalDurableTestRunner.create(String.class, (input, ctx) -> { + ctx.wait("hold", Duration.ofMinutes(5)); + return ctx.step("after-wait", String.class, stepCtx -> { + stepRuns.incrementAndGet(); + return "done"; + }); + }); + var waiting = runner.run("input"); + + assertEquals(ExecutionStatus.PENDING, waiting.getStatus()); + + runner.pauseExecution(); + var paused = runner.runUntilComplete("input"); + + assertEquals(ExecutionStatus.PENDING, paused.getStatus()); + // The wait could normally be auto-advanced by runUntilComplete. While paused, runUntilComplete must stop + // instead of re-invoking and running after-wait, because any new checkpoint from that work will not be able + // to be recorded. + assertEquals(0, stepRuns.get()); + + runner.resumeExecution(); + var resumed = runner.runUntilComplete("input"); + + assertEquals(ExecutionStatus.SUCCEEDED, resumed.getStatus()); + assertEquals("done", resumed.getResult(String.class)); + } + @Test void testInvalidSuspensionThrowsIllegalDurableOperationException() { var runner = LocalDurableTestRunner.create(String.class, (input, ctx) -> { diff --git a/sdk-testing/src/test/java/software/amazon/lambda/durable/testing/local/LocalMemoryExecutionClientTest.java b/sdk-testing/src/test/java/software/amazon/lambda/durable/testing/local/LocalMemoryExecutionClientTest.java index 233f7b8a86..6d8aa971b2 100644 --- a/sdk-testing/src/test/java/software/amazon/lambda/durable/testing/local/LocalMemoryExecutionClientTest.java +++ b/sdk-testing/src/test/java/software/amazon/lambda/durable/testing/local/LocalMemoryExecutionClientTest.java @@ -74,4 +74,36 @@ void testCompleteInvocationThrowsWhenOperationNotFound() { "nonexistent-invoke", new OperationResult(OperationStatus.TIMED_OUT, null, null))); assertEquals("Operation not found: nonexistent-invoke", exception.getMessage()); } + + @Test + void pausedCheckpointOmitsTheTokenAndStillRecordsTheUpdates() { + var client = new LocalMemoryExecutionClient(); + var update = OperationUpdate.builder() + .id("1") + .name("step-1") + .type(OperationType.STEP) + .action(OperationAction.SUCCEED) + .payload("\"result1\"") + .build(); + + client.pause(); + var response = client.checkpoint("arn", "token", List.of(update)); + + assertTrue(client.isPaused()); + assertNull(response.checkpointToken()); + assertNotNull(client.getOperationByName("step-1")); + assertEquals(List.of("1"), client.getUpdatedOperationIdsSinceLastInvocation()); + } + + @Test + void resumeAnswersCheckpointsWithATokenAgain() { + var client = new LocalMemoryExecutionClient(); + client.pause(); + assertNull(client.checkpoint("arn", "token", List.of()).checkpointToken()); + + client.resume(); + + assertFalse(client.isPaused()); + assertNotNull(client.checkpoint("arn", "token", List.of()).checkpointToken()); + } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/execution/ApiRequestDelayedBatcher.java b/sdk/src/main/java/software/amazon/lambda/durable/execution/ApiRequestDelayedBatcher.java index 3b0def81fb..ab0edd6a9a 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/execution/ApiRequestDelayedBatcher.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/execution/ApiRequestDelayedBatcher.java @@ -104,15 +104,18 @@ CompletableFuture submit(T request, Duration flushDelay) { /** Flushes pending batch and waits for completion */ void shutdown() { + flush(); + flushingQueueFuture.join(); + } + + /** Schedules pending requests immediately without waiting for the current batch. */ + void flush() { synchronized (delayedBatch) { // cancel the flush timer if it has not been triggered this.delayedBatchFlushTimer.cancel(false); // execute the current batch now flushDelayedBatch(); } - - // wait for previous batches to be flushed - flushingQueueFuture.join(); } /** clear the current batch and creates a new batch */ @@ -146,6 +149,7 @@ private void flushDelayedBatch() { // the new future will just do nothing. flushingQueueFuture = flushingQueueFuture.thenRunAsync(this::flushQueue, InternalExecutor.INSTANCE); } + /** Call checkpoint API with items in the flushing queue */ private void flushQueue() { // There could be more items to flush because diff --git a/sdk/src/main/java/software/amazon/lambda/durable/execution/CheckpointManager.java b/sdk/src/main/java/software/amazon/lambda/durable/execution/CheckpointManager.java index 5f1747444d..67a02f5c0c 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/execution/CheckpointManager.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/execution/CheckpointManager.java @@ -18,6 +18,7 @@ import software.amazon.awssdk.awscore.exception.AwsServiceException; import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState; import software.amazon.awssdk.services.lambda.model.Operation; +import software.amazon.awssdk.services.lambda.model.OperationType; import software.amazon.awssdk.services.lambda.model.OperationUpdate; import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.retry.PollingStrategies; @@ -34,6 +35,8 @@ class CheckpointManager { private static final int MAX_BATCH_SIZE_BYTES = 750 * 1024; // 750KB private static final int MAX_ITEM_COUNT = 200; // max updates in one batch private static final int FIRST_ATTEMPT = 1; + private static final String REVOKED_CHECKPOINT_TOKEN_MESSAGE = + "The checkpoint response did not include a checkpoint token"; private static final Logger logger = LoggerFactory.getLogger(CheckpointManager.class); private final Consumer> callback; @@ -42,20 +45,22 @@ class CheckpointManager { private final ApiRequestDelayedBatcher checkpointApiRequestDelayedBatcher; private final DurableConfig config; private final BooleanSupplier tryStartCheckpointProcessing; - private final Runnable finishCheckpointProcessing; - private String checkpointToken; + private final Runnable onCheckpointProcessingFinish; + private final Runnable onTokenRevoked; + private volatile String checkpointToken; + private boolean pollingStopped; // Guarded by pollingFutures. + /** + * Set when a nonterminal checkpoint response omits its token. This invocation cannot checkpoint again, so + * revocation is never cleared. + */ private volatile boolean checkpointTokenRevoked; - boolean isCheckpointTokenRevoked() { - return checkpointTokenRevoked; - } - CheckpointManager( DurableConfig config, String durableExecutionArn, String checkpointToken, Consumer> callback) { - this(config, durableExecutionArn, checkpointToken, callback, () -> true, () -> {}); + this(config, durableExecutionArn, checkpointToken, callback, () -> true, () -> {}, () -> {}); } CheckpointManager( @@ -64,25 +69,42 @@ boolean isCheckpointTokenRevoked() { String checkpointToken, Consumer> callback, BooleanSupplier tryStartCheckpointProcessing, - Runnable finishCheckpointProcessing) { + Runnable onCheckpointProcessingFinish, + Runnable onTokenRevoked) { this.config = config; this.durableExecutionArn = durableExecutionArn; this.callback = callback; this.checkpointToken = checkpointToken; this.tryStartCheckpointProcessing = tryStartCheckpointProcessing; - this.finishCheckpointProcessing = finishCheckpointProcessing; + this.onCheckpointProcessingFinish = onCheckpointProcessingFinish; + this.onTokenRevoked = onTokenRevoked; this.checkpointApiRequestDelayedBatcher = new ApiRequestDelayedBatcher<>( MAX_ITEM_COUNT, MAX_BATCH_SIZE_BYTES, CheckpointManager::estimateSize, this::checkpointBatch); } + /** Returns {@code true} once the service has withdrawn this invocation's checkpoint token. */ + boolean isCheckpointTokenRevoked() { + return checkpointTokenRevoked; + } + /** * Queues a checkpoint request for batched execution * * @return a future that completes when the checkpoint request is executed */ CompletableFuture checkpoint(OperationUpdate update) { + if (checkpointToken == null) { + return CompletableFuture.failedFuture(new SuspendExecutionException(REVOKED_CHECKPOINT_TOKEN_MESSAGE)); + } + logger.debug("Checkpoint request received: Action {}", update.action()); - return checkpointApiRequestDelayedBatcher.submit(update, config.getCheckpointDelay()); + var future = checkpointApiRequestDelayedBatcher.submit(update, config.getCheckpointDelay()); + // Withdrawal can flush before this request is submitted, so check again to avoid waiting for its timer. + if (checkpointToken == null) { + checkpointApiRequestDelayedBatcher.flush(); + } + + return future; } /** @@ -114,6 +136,9 @@ CompletableFuture pollForUpdate(String operationId, PollingStrategy p logger.debug("Polling request received: operation id {}", operationId); var future = new CompletableFuture(); synchronized (pollingFutures) { + if (pollingStopped) { + return CompletableFuture.failedFuture(new SuspendExecutionException()); + } // register the future in pollingFutures, which will be completed by the polling thread pollingFutures .computeIfAbsent(operationId, k -> Collections.synchronizedList(new ArrayList<>())) @@ -156,18 +181,22 @@ private CompletableFuture pollForUpdateInternal( /** Cancels all polling futures and waits for all pending checkpoint requests to complete */ void shutdown() { // complete all polling futures with an exception + stopPolling(new IllegalStateException("CheckpointManager shutdown")); + // wait for all non-polling checkpoint requests to complete + checkpointApiRequestDelayedBatcher.shutdown(); + } + + private void stopPolling(Throwable cause) { List>> allFutures; synchronized (pollingFutures) { + pollingStopped = true; allFutures = new ArrayList<>(pollingFutures.values()); pollingFutures.clear(); } for (var futures : allFutures) { - futures.forEach(f -> f.completeExceptionally(new IllegalStateException("CheckpointManager shutdown"))); + futures.forEach(future -> future.completeExceptionally(cause)); } - - // wait for all non-polling checkpoint requests to complete - checkpointApiRequestDelayedBatcher.shutdown(); } /** @@ -216,6 +245,12 @@ private void checkpointBatch(List updates) { return; } + // An earlier batch already came back without a checkpoint token, so this invocation must not issue + // further checkpoints. + if (checkpointToken == null) { + throw new SuspendExecutionException(REVOKED_CHECKPOINT_TOKEN_MESSAGE); + } + // Starting the backend request is coordinated with the last-thread suspension decision. Once suspension // wins that race, no later poll/checkpoint may advance backend state behind the PENDING response. if (!tryStartCheckpointProcessing.getAsBoolean()) { @@ -226,53 +261,78 @@ private void checkpointBatch(List updates) { request.size()); } logger.debug("Skipping checkpoint API call because execution has already completed"); - return; + throw new SuspendExecutionException(); } - var startTime = System.nanoTime(); - logger.debug("Calling durable checkpoint API with {} updates: {}", updates.size(), request); try { - var response = - config.getDurableExecutionClient().checkpoint(durableExecutionArn, checkpointToken, request); - logger.debug( - "Durable checkpoint API called (latency={}ns): {}.", System.nanoTime() - startTime, response); - - // Notify callback of completion - checkpointToken = response.checkpointToken(); - if (checkpointToken == null) { - checkpointTokenRevoked = true; - } - if (response.newExecutionState() != null) { - // fetch all pages of operations - var operations = fetchAllPages(response.newExecutionState()); - - var processStartTime = System.nanoTime(); - int completedFutures = 0; - logger.debug( - "Processing {} operations. ({} pending pollers)", operations.size(), pollingFutures.size()); - // call the callback - callback.accept(operations); - - // complete the registered pollingFutures - for (var operation : operations) { - var pollers = pollingFutures.remove(operation.id()); - if (pollers != null) { - completedFutures += pollers.size(); - pollers.forEach(poller -> poller.complete(operation)); - } - } - logger.debug( - "{} operations processed and {} pollers completed (latency={}ns). ", - operations.size(), - completedFutures, - System.nanoTime() - processStartTime); - } + checkpointRequest(request); } catch (AwsServiceException e) { throw DurableApiErrorClassifier.classifyException(e); } finally { - finishCheckpointProcessing.run(); + onCheckpointProcessingFinish.run(); + } + } + } + + private void checkpointRequest(List request) { + var startTime = System.nanoTime(); + logger.debug("Calling durable checkpoint API with {} updates: {}", request.size(), request); + var response = config.getDurableExecutionClient().checkpoint(durableExecutionArn, checkpointToken, request); + logger.debug("Durable checkpoint API called (latency={}ns): {}.", System.nanoTime() - startTime, response); + + checkpointToken = response.checkpointToken(); + if (checkpointToken == null) { + handleMissingToken(request); + return; // The token is gone, so this invocation cannot use the response state or fetch more pages. + } + + var state = response.newExecutionState(); + if (state != null) { + processUpdatedOperations(fetchAllPages(state)); + } + } + + private void handleMissingToken(List updates) { + // Publish before completing pollers, whose dependent stages can run inline. + // A terminal execution update was accepted, so only nonterminal requests require suspension. + checkpointTokenRevoked = updates.stream().noneMatch(update -> update.type() == OperationType.EXECUTION); + + var cause = new SuspendExecutionException(REVOKED_CHECKPOINT_TOKEN_MESSAGE); + stopPolling(cause); + checkpointApiRequestDelayedBatcher.flush(); + + if (checkpointTokenRevoked) { + logger.warn( + "The checkpoint response for durable execution {} did not include a checkpoint token. The SDK will" + + " stop checkpointing and report the invocation as PENDING.", + durableExecutionArn); + onTokenRevoked.run(); + throw cause; + } + + logger.info( + "Checkpoint token withheld on durable execution {}'s terminal update; execution is finished.", + durableExecutionArn); + } + + private void processUpdatedOperations(List operations) { + var processStartTime = System.nanoTime(); + int completedFutures = 0; + logger.debug("Processing {} operations. ({} pending pollers)", operations.size(), pollingFutures.size()); + callback.accept(operations); + + for (var operation : operations) { + var pollers = pollingFutures.remove(operation.id()); + if (pollers != null) { + completedFutures += pollers.size(); + pollers.forEach(poller -> poller.complete(operation)); } } + logger.debug( + "{} operations processed and {} pollers completed (latency={}ns). ", + operations.size(), + completedFutures, + System.nanoTime() - processStartTime); } private static int estimateSize(OperationUpdate update) { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/execution/DurableExecutor.java b/sdk/src/main/java/software/amazon/lambda/durable/execution/DurableExecutor.java index da8e0e6ab3..082fb50aa6 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/execution/DurableExecutor.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/execution/DurableExecutor.java @@ -52,6 +52,8 @@ */ public class DurableExecutor { private static final String ROOT_THREAD_ID = null; + private static final String REVOKED_CHECKPOINT_TOKEN_MESSAGE = + "The checkpoint response did not include a checkpoint token"; private static final Logger logger = LoggerFactory.getLogger(DurableExecutor.class); // Lambda response size limit is 6MB minus small epsilon for envelope @@ -126,6 +128,18 @@ private CompletableFuture execute() { Supplier task = () -> { Outcome.capture(this::invokeHandler).complete(body); var selected = outcome.join(); + // The handler has settled, but its operations may still be running or still have checkpoints queued. + // Drain them here so every checkpoint response, including one that withdraws the token, is observed + // before the terminal decision below. Later polls therefore cannot change that decision. + try { + manager.drainOperations(); + } catch (RuntimeException drainFailure) { + // The drain throws only when this thread is interrupted or an operation future is cancelled. + // No terminal response exists yet, so the service must retry this invocation. Plugins still + // need their end hook. Classifying this as a handler failure would mark the execution FAILED, + // so it is reported as a delivery failure: RETRYING, then the original exception is rethrown. + return failDelivery(drainFailure); + } return finishInvocation(selected.value(), selected.failure()); }; // Cleanup is awaited even without plugins. Keep that path free of plugin MDC handling. @@ -168,13 +182,21 @@ private void fireOnInvocationStart() { } private DurableExecutionOutput finishInvocation(O value, Throwable failure) { + // A withdrawn token means this invocation must stop and report PENDING. That overrides a handler success, + // a handler failure, and a retryable failure, so it is checked before any of them. + if (manager.isCheckpointTokenRevoked()) + return finishFailure(new SuspendExecutionException(REVOKED_CHECKPOINT_TOKEN_MESSAGE)); if (failure != null) return finishFailure(ExceptionHelper.unwrapCompletableFuture(failure)); - manager.closeCheckpointContinuationAdmission(); DurableExecutionOutput output; try { var payload = config.getSerDes().serialize(value); output = DurableExecutionOutput.success(handleLargePayload(manager, payload)); } catch (Throwable deliveryFailure) { + // Checkpointing a large result fails with SuspendExecutionException when the token is already gone. + // That is a suspension, not a delivery error, so it reports PENDING rather than RETRYING. + if (ExceptionHelper.unwrapCompletableFuture(deliveryFailure) instanceof SuspendExecutionException) { + return finishFailure(new SuspendExecutionException(REVOKED_CHECKPOINT_TOKEN_MESSAGE)); + } return failDelivery(deliveryFailure); } fireOnInvocationEnd(InvocationStatus.SUCCEEDED, null, value); 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 2526c74109..007df86658 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 @@ -118,7 +118,8 @@ public ExecutionManager(DurableExecutionInput input, DurableConfig config, Conte input.checkpointToken(), this::onCheckpointComplete, this::tryStartCheckpointProcessing, - this::finishCheckpointProcessing); + this::finishCheckpointProcessing, + this::signalSuspension); this.operationStorage = checkpointManager.fetchAllPages(input.initialExecutionState()).stream() .collect(Collectors.toConcurrentMap(Operation::id, op -> op)); @@ -230,6 +231,11 @@ public void registerOperation(BaseDurableOperation operation, BaseDurableOperati : completionLockFor(completionLockParent); operationCompletionLocks.put(operationId, completionLock); registeredOperations.put(operationId, operation); + // Revocation is published before existing operations are suspended. Check after insertion so an + // operation registered too late for that notification still releases its waiters. + if (isCheckpointTokenRevoked()) { + operation.getCompletionFuture().completeExceptionally(new SuspendExecutionException()); + } } // ===== Checkpoint Completion Handler ===== @@ -741,17 +747,21 @@ public CompletableFuture pollForOperationUpdates(String operationId, /** Shutdown the checkpoint batcher. */ @Override public void close() { - stopCheckpointContinuations(); - validateRunningThreads(); - - checkpointManager.shutdown(); + drainOperations(); } - /** Closes new checkpoint-continuation admission without signaling owners or joining their cleanup. */ - void closeCheckpointContinuationAdmission() { - synchronized (activeThreads) { - closing = true; - } + /** + * Settles submitted operation handlers and stops polling before flushing checkpoints. Terminal decisions therefore + * observe token withdrawal without allowing later polls to change that decision. Safe to call more than once: a + * second call finds no running handlers and an already-stopped batcher. + */ + void drainOperations() { + // Closing admission can skip queued continuations. Stop their owners before joining user handlers, + // or a child waiting on a skipped continuation would prevent the drain from completing. + stopCheckpointContinuations(); + waitForRunningUserHandlers(); + validateExecutorPool(); + checkpointManager.shutdown(); } private void stopCheckpointContinuations() { @@ -776,7 +786,7 @@ private void stopCheckpointContinuations() { } } - private void validateRunningThreads() { + private void waitForRunningUserHandlers() { // This will detect stuck user thread and thread leaks in the thread pool for (BaseDurableOperation op : registeredOperations.values()) { var userHandlerFuture = op.getRunningUserHandler(); @@ -795,7 +805,9 @@ private void validateRunningThreads() { } } } + } + private void validateExecutorPool() { // double check if the thread pool is empty if (durableConfig.getExecutorService() instanceof ThreadPoolExecutor threadPoolExecutor) { var threadCount = threadPoolExecutor.getActiveCount(); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/execution/SuspendExecutionException.java b/sdk/src/main/java/software/amazon/lambda/durable/execution/SuspendExecutionException.java index 468dfc195e..6524f364d1 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/execution/SuspendExecutionException.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/execution/SuspendExecutionException.java @@ -7,4 +7,9 @@ public class SuspendExecutionException extends DurableExecutionError { public SuspendExecutionException() { super("Execution suspended for wait operation"); } + + /** @param message identifies why the invocation is being suspended, for log correlation. */ + public SuspendExecutionException(String message) { + super(message); + } } 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 bdb7fd11a9..31c9926156 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 @@ -277,6 +277,9 @@ protected void runUserHandler(Runnable runnable, ThreadType threadType) { executionManager.setCurrentThreadContext(new ThreadContext(operationId, threadType)); try { + if (executionManager.isCheckpointTokenRevoked()) { + throw new SuspendExecutionException(); + } runnable.run(); } catch (Throwable throwable) { // Operations wrap the user function and handle all outcomes except for SuspendExecutionException. diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/ApiRequestDelayedBatcherTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/ApiRequestDelayedBatcherTest.java index 6aec8f7c14..75397781c0 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/ApiRequestDelayedBatcherTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/ApiRequestDelayedBatcherTest.java @@ -18,11 +18,15 @@ import java.util.ArrayList; import java.util.List; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Consumer; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; class ApiRequestDelayedBatcherTest { private static final Duration SHORT_DELAY = Duration.ofMillis(5); @@ -190,6 +194,77 @@ void whenBatchActionThrowsRuntimeException_futuresReceiveOriginalCause() { CompletableFuture.allOf(resultFuture1, resultFuture2, resultFuture3).join(); } + @Test + void whenFlushCalled_longDelayedItemsCompletePromptly() throws Exception { + var future1 = cut.submit(input, Duration.ofHours(1)); + var future2 = cut.submit(input, Duration.ofHours(1)); + + cut.flush(); + + CompletableFuture.allOf(future1, future2).get(1, TimeUnit.SECONDS); + verify(doBatchAction).accept(argThat(list -> list.size() == 2)); + cut.shutdown(); + } + + @Test + @Timeout(10) + void whenFlushCalledDuringActiveBatch_queuedBatchesWaitWithoutBlockingFlush() throws Exception { + var firstBatchStarted = new CountDownLatch(1); + var releaseFirstBatch = new CountDownLatch(1); + var executions = new AtomicInteger(); + var singleItemCut = blockingBatcher(firstBatchStarted, releaseFirstBatch, executions); + try { + var firstFuture = singleItemCut.submit(input, Duration.ofHours(1)); + singleItemCut.flush(); + await(firstBatchStarted); + var queuedFuture1 = singleItemCut.submit(input, Duration.ofHours(1)); + var queuedFuture2 = singleItemCut.submit(input, Duration.ofHours(1)); + singleItemCut.flush(); + + assertFalse(firstFuture.isDone()); + assertFalse(queuedFuture1.isDone()); + assertFalse(queuedFuture2.isDone()); + assertEquals(1, executions.get()); + + releaseFirstBatch.countDown(); + CompletableFuture.allOf(firstFuture, queuedFuture1, queuedFuture2).get(1, TimeUnit.SECONDS); + assertEquals(3, executions.get()); + } finally { + releaseFirstBatch.countDown(); + singleItemCut.shutdown(); + } + } + + @Test + @Timeout(10) + void whenShutdownCalledDuringActiveBatch_itWaitsForCompletion() throws Exception { + var batchStarted = new CountDownLatch(1); + var releaseBatch = new CountDownLatch(1); + var shutdownStarted = new CountDownLatch(1); + var blockingCut = blockingBatcher(batchStarted, releaseBatch, new AtomicInteger()); + var caller = Executors.newSingleThreadExecutor(); + try { + var future = blockingCut.submit(input, Duration.ofHours(1)); + blockingCut.flush(); + await(batchStarted); + var shutdown = caller.submit(() -> { + shutdownStarted.countDown(); + blockingCut.shutdown(); + }); + await(shutdownStarted); + assertThrows(TimeoutException.class, () -> shutdown.get(100, TimeUnit.MILLISECONDS)); + assertFalse(future.isDone()); + + releaseBatch.countDown(); + shutdown.get(1, TimeUnit.SECONDS); + assertTrue(future.isDone()); + } finally { + releaseBatch.countDown(); + blockingCut.shutdown(); + caller.shutdownNow(); + } + } + @Test void whenShutdownCalled_pendingItemsAreFlushedImmediately() { var future = cut.submit(input, LONG_DELAY); @@ -226,4 +301,22 @@ void whenMultipleBatchesFlushedViaShutdown_allFuturesComplete() { cut.shutdown(); assertTrue(future2.isDone()); } + + private ApiRequestDelayedBatcher blockingBatcher( + CountDownLatch started, CountDownLatch release, AtomicInteger executions) { + return new ApiRequestDelayedBatcher<>(1, MAX_BATCH_BINARY_SIZE_IN_BYTES, item -> 0, requests -> { + executions.incrementAndGet(); + started.countDown(); + await(release); + }); + } + + private void await(CountDownLatch latch) { + try { + assertTrue(latch.await(5, TimeUnit.SECONDS), "Timed out waiting for the test latch"); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/CheckpointManagerTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/CheckpointManagerTest.java index d6de81482b..916460e138 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/CheckpointManagerTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/CheckpointManagerTest.java @@ -10,12 +10,15 @@ import java.time.Instant; import java.util.ArrayList; import java.util.List; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import software.amazon.awssdk.awscore.exception.AwsErrorDetails; import software.amazon.awssdk.awscore.exception.AwsServiceException; import software.amazon.awssdk.http.SdkHttpResponse; @@ -80,14 +83,22 @@ void checkpoint_sendsUpdateAndReturnsCompletedFuture() throws Exception { void checkpoint_skipsApiCallWhenExecutionAlreadyCompleted() throws Exception { var finishCheckpointProcessing = mock(Runnable.class); var guardedBatcher = new CheckpointManager( - config, "arn:test", "token-1", callbackOperations::addAll, () -> false, finishCheckpointProcessing); + config, + "arn:test", + "token-1", + callbackOperations::addAll, + () -> false, + finishCheckpointProcessing, + () -> {}); var update = OperationUpdate.builder() .id("op-1") .type(OperationType.STEP) .action(OperationAction.START) .build(); - guardedBatcher.checkpoint(update).get(200, TimeUnit.MILLISECONDS); + var error = assertThrows( + Exception.class, () -> guardedBatcher.checkpoint(update).get(200, TimeUnit.MILLISECONDS)); + assertInstanceOf(SuspendExecutionException.class, error.getCause()); verifyNoInteractions(client); verify(finishCheckpointProcessing, never()).run(); @@ -672,6 +683,139 @@ void checkpointBatch_nonRetryableError_throwsUnrecoverable() { assertInstanceOf(UnrecoverableDurableExecutionException.class, ex.getCause()); } + // --- Checkpoint token revocation tests (a checkpoint response with no checkpoint token) --- + + @Test + void checkpointBatch_nonTerminalTokenWithdrawal_suspendsAndRejectsFurtherRequests() { + var signalSuspension = mock(Runnable.class); + var revocableBatcher = new CheckpointManager( + config, "arn:test", "token-1", callbackOperations::addAll, () -> true, () -> {}, signalSuspension); + var update = stepUpdate("op-1"); + when(client.checkpoint(any(), any(), any())).thenReturn(missingTokenResponse(update)); + var poller = revocableBatcher.pollForUpdate("op-1", Instant.now().plusSeconds(60)); + try { + assertSuspended(revocableBatcher.checkpoint(update)); + assertSuspended(poller); + assertTrue(revocableBatcher.isCheckpointTokenRevoked()); + verify(signalSuspension).run(); + assertWithdrawalStopsFurtherRequests(revocableBatcher, update); + } finally { + revocableBatcher.shutdown(); + } + } + + @Test + void checkpointBatch_nonTerminalTokenWithdrawal_publishesRevocationBeforeCompletingPollers() throws Exception { + var update = stepUpdate("op-1"); + + var poller = batcher.pollForUpdate("op-1", Instant.now().plusSeconds(60)); + var revocationObservedByPoller = poller.handle((operation, error) -> batcher.isCheckpointTokenRevoked()); + when(client.checkpoint(any(), any(), any())).thenReturn(missingTokenResponse(update)); + + try { + assertSuspended(batcher.checkpoint(update)); + assertSuspended(poller); + assertTrue(revocationObservedByPoller.get(1, TimeUnit.SECONDS)); + } finally { + batcher.shutdown(); + } + } + + @Test + void checkpointBatch_terminalTokenWithdrawal_completesAndRejectsFurtherRequests() throws Exception { + var signalSuspension = mock(Runnable.class); + var revocableBatcher = new CheckpointManager( + config, "arn:test", "token-1", callbackOperations::addAll, () -> true, () -> {}, signalSuspension); + var update = OperationUpdate.builder() + .id("exec-op") + .type(OperationType.EXECUTION) + .action(OperationAction.SUCCEED) + .payload("\"result\"") + .build(); + when(client.checkpoint(any(), any(), any())).thenReturn(missingTokenResponse(update)); + var poller = revocableBatcher.pollForUpdate("exec-op", Instant.now().plusSeconds(60)); + try { + revocableBatcher.checkpoint(update).get(1, TimeUnit.SECONDS); + assertSuspended(poller); + assertFalse(revocableBatcher.isCheckpointTokenRevoked()); + verify(signalSuspension, never()).run(); + assertWithdrawalStopsFurtherRequests(revocableBatcher, update); + } finally { + revocableBatcher.shutdown(); + } + } + + @Test + void shutdown_rejectsLaterPollRegistration() throws Exception { + batcher.shutdown(); + + var poller = batcher.pollForUpdate("late", Instant.now().plusSeconds(60)); + assertTrue(poller.isCompletedExceptionally()); + var error = assertThrows(Exception.class, () -> poller.get(1, TimeUnit.SECONDS)); + assertInstanceOf(SuspendExecutionException.class, error.getCause()); + batcher.shutdown(); + verifyNoInteractions(client); + } + + @Test + @Timeout(10) + void checkpointBatch_revokedToken_flushesLongDelayedCheckpointsAcrossMultipleBatches() throws Exception { + var revocationStarted = new CountDownLatch(1); + var returnRevokedResponse = new CountDownLatch(1); + var delayedConfig = longDelayedConfig(); + var revocableBatcher = new CheckpointManager(delayedConfig, "arn:test", "token-1", callbackOperations::addAll); + revokeAfterRelease(revocationStarted, returnRevokedResponse); + try { + var poller = revocableBatcher.pollForUpdate("poll", Instant.now()); + await(revocationStarted); + var queued = new ArrayList>(); + for (int i = 0; i < 401; i++) { + queued.add(revocableBatcher.checkpoint(stepUpdate("op-" + i))); + } + returnRevokedResponse.countDown(); + assertSuspended(poller); + for (var future : queued) { + assertSuspended(future); + } + verify(client, times(1)).checkpoint(anyString(), anyString(), anyList()); + verify(client, never()).getExecutionState(any(), any(), any()); + } finally { + returnRevokedResponse.countDown(); + revocableBatcher.shutdown(); + } + } + + @Test + @Timeout(10) + void checkpoint_withdrawalBetweenTokenCheckAndSubmission_flushesLateRequest() throws Exception { + var requestStarted = new CountDownLatch(1); + var returnRevokedResponse = new CountDownLatch(1); + var submissionPaused = new CountDownLatch(1); + var resumeSubmission = new CountDownLatch(1); + var delayedConfig = blockedCheckpointConfig(submissionPaused, resumeSubmission); + var revocableBatcher = new CheckpointManager(delayedConfig, "arn:test", "token-1", callbackOperations::addAll); + var caller = Executors.newSingleThreadExecutor(); + revokeAfterRelease(requestStarted, returnRevokedResponse); + try { + var poller = revocableBatcher.pollForUpdate("poll", Instant.now()); + await(requestStarted); + var submission = caller.submit(() -> revocableBatcher.checkpoint(stepUpdate("late"))); + await(submissionPaused); + returnRevokedResponse.countDown(); + assertSuspended(poller); + revocableBatcher.shutdown(); + resumeSubmission.countDown(); + assertSuspended(submission.get(1, TimeUnit.SECONDS)); + verify(client, times(1)).checkpoint(anyString(), anyString(), anyList()); + verify(client, never()).getExecutionState(any(), any(), any()); + } finally { + returnRevokedResponse.countDown(); + resumeSubmission.countDown(); + caller.shutdownNow(); + revocableBatcher.shutdown(); + } + } + @Test void fetchAllPages_nonRetryableError_throwsUnrecoverable() { when(client.getExecutionState(eq("arn:test"), eq("token-1"), eq("marker-1"))) @@ -694,4 +838,79 @@ void fetchAllPages_nonRetryableError_throwsUnrecoverable() { assertThrows(UnrecoverableDurableExecutionException.class, () -> batcher.fetchAllPages(state)); } + + private CheckpointDurableExecutionResponse missingTokenResponse(OperationUpdate update) { + var operation = Operation.builder() + .id(update.id()) + .type(update.type()) + .status(OperationStatus.SUCCEEDED) + .build(); + return CheckpointDurableExecutionResponse.builder() + .newExecutionState(CheckpointUpdatedExecutionState.builder() + .operations(operation) + .nextMarker("unused") + .build()) + .build(); + } + + private void assertWithdrawalStopsFurtherRequests(CheckpointManager manager, OperationUpdate update) { + assertTrue(callbackOperations.isEmpty()); + var lateCheckpoint = manager.checkpoint(update); + var latePoller = manager.pollForUpdate("late", Instant.now().plusSeconds(60)); + assertTrue(lateCheckpoint.isCompletedExceptionally()); + assertTrue(latePoller.isCompletedExceptionally()); + assertSuspended(lateCheckpoint); + assertSuspended(latePoller); + verify(client, times(1)).checkpoint(any(), any(), any()); + verify(client, never()).getExecutionState(any(), any(), any()); + } + + private DurableConfig longDelayedConfig() { + return DurableConfig.builder() + .withDurableExecutionClient(client) + .withCheckpointDelay(Duration.ofHours(1)) + .build(); + } + + private DurableConfig blockedCheckpointConfig(CountDownLatch paused, CountDownLatch resume) { + var blockedConfig = spy(longDelayedConfig()); + doAnswer(invocation -> { + paused.countDown(); + await(resume); + return Duration.ofHours(1); + }) + .when(blockedConfig) + .getCheckpointDelay(); + return blockedConfig; + } + + private void revokeAfterRelease(CountDownLatch started, CountDownLatch release) { + when(client.checkpoint(anyString(), anyString(), anyList())).thenAnswer(invocation -> { + started.countDown(); + await(release); + return CheckpointDurableExecutionResponse.builder().build(); + }); + } + + private OperationUpdate stepUpdate(String id) { + return OperationUpdate.builder() + .id(id) + .type(OperationType.STEP) + .action(OperationAction.START) + .build(); + } + + private void assertSuspended(CompletableFuture future) { + var error = assertThrows(Exception.class, () -> future.get(1, TimeUnit.SECONDS)); + assertInstanceOf(SuspendExecutionException.class, error.getCause()); + } + + private void await(CountDownLatch latch) { + try { + assertTrue(latch.await(5, TimeUnit.SECONDS), "Timed out waiting for the test latch"); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/CompletedPollBatchTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/CompletedPollBatchTest.java index 0646aaa46f..100bdd30bd 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/CompletedPollBatchTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/CompletedPollBatchTest.java @@ -38,7 +38,8 @@ void completedPollsDoNotCauseEmptyRpcButLivePollsAndUpdatesStillFlush(String mod starts.incrementAndGet(); return true; }, - finishes::incrementAndGet); + finishes::incrementAndGet, + () -> {}); var ready = Operation.builder() .id("watched") .type(OperationType.STEP) diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutionTest.java index d8f5a849c4..1597efa916 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutionTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutionTest.java @@ -4,24 +4,37 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static software.amazon.lambda.durable.TypeToken.get; +import java.time.Duration; import java.time.Instant; import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BiFunction; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.ValueSource; import software.amazon.awssdk.services.lambda.model.CheckpointDurableExecutionResponse; import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState; import software.amazon.awssdk.services.lambda.model.ErrorObject; @@ -33,12 +46,14 @@ import software.amazon.awssdk.services.lambda.model.OperationUpdate; import software.amazon.awssdk.services.lambda.model.StepDetails; import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.TestUtils; import software.amazon.lambda.durable.client.DurableExecutionClient; import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.exception.IllegalDurableOperationException; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.model.DurableExecutionOutput; import software.amazon.lambda.durable.model.ExecutionStatus; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; @@ -119,7 +134,7 @@ void testExecutePending() { get(String.class), (userInput, ctx) -> { ctx.step("step1", String.class, stepCtx -> "Done"); - ctx.wait(null, java.time.Duration.ofSeconds(60)); + ctx.wait(null, Duration.ofSeconds(60)); return "Should not reach here"; }, configWithMockClient()); @@ -471,6 +486,125 @@ void testExecutorNotShutdownAfterMultipleHandlerInvocations() { assertTrue(output2.result().contains("Result 2: test-input-2")); } + private enum OperationKind { + STEP, + WAIT, + CALLBACK, + MAP, + PARALLEL + } + + @ParameterizedTest + @EnumSource(OperationKind.class) + @Timeout(10) + void checkpointTokenRevoked_duringOperation_suspendsExecutionAsPending(OperationKind kind) { + var client = mock(DurableExecutionClient.class); + var revoked = new AtomicBoolean(); + when(client.checkpoint(any(), any(), any())).thenAnswer(invocation -> { + List updates = invocation.getArgument(2); + var revoke = updates.stream().anyMatch(update -> "revoke".equals(update.name())); + if (revoke) { + revoked.set(true); + } + return CheckpointDurableExecutionResponse.builder() + .checkpointToken(revoke ? null : "next-token") + .build(); + }); + + var config = DurableConfig.builder().withDurableExecutionClient(client).build(); + var output = execute(config, (userInput, ctx) -> { + executeRevokingOperation(kind, ctx); + return "unreachable"; + }); + + assertTrue(revoked.get(), "The selected operation must receive a token-less checkpoint response"); + assertEquals(ExecutionStatus.PENDING, output.status()); + assertNull(output.result()); + assertNull(output.error()); + verify(client, never()).getExecutionState(any(), any(), any()); + } + + private void executeRevokingOperation(OperationKind kind, DurableContext ctx) { + switch (kind) { + case STEP -> ctx.step("revoke", String.class, stepCtx -> "done"); + case WAIT -> ctx.wait("revoke", Duration.ofSeconds(60)); + case CALLBACK -> ctx.createCallback("revoke", String.class).get(); + case MAP -> + ctx.map( + "map", + List.of("item"), + String.class, + (item, index, child) -> child.step("revoke", String.class, stepCtx -> "done")); + case PARALLEL -> { + try (var parallel = ctx.parallel("parallel")) { + parallel.branch( + "branch", String.class, child -> child.step("revoke", String.class, stepCtx -> "done")); + parallel.get(); + } + } + } + } + + @ParameterizedTest(name = "largeResult={0}") + @ValueSource(booleans = {false, true}) + @Timeout(10) + void checkpointTokenRevoked_beforeHandlerReturns_reportsPending(boolean largeResult) throws Exception { + var client = mock(DurableExecutionClient.class); + when(client.checkpoint(any(), any(), any())) + .thenReturn(CheckpointDurableExecutionResponse.builder().build()); + var plugin = mock(DurableExecutionPlugin.class); + var config = DurableConfig.builder() + .withDurableExecutionClient(client) + .withPlugins(plugin) + .build(); + + var handlerReadyToReturn = new CountDownLatch(1); + var output = execute(config, handlerAfterRevocation(largeResult, handlerReadyToReturn)); + + assertTrue(handlerReadyToReturn.await(5, TimeUnit.SECONDS), "Handler must observe suspension before returning"); + assertEquals(ExecutionStatus.PENDING, output.status()); + assertNull(output.result()); + assertNull(output.error()); + verify(client, times(1)).checkpoint(any(), any(), any()); + verify(plugin, times(1)) + .onInvocationEnd(argThat(info -> info.invocationStatus() == InvocationStatus.PENDING + && info.executionError() == null + && info.executionResult() == null)); + } + + private BiFunction handlerAfterRevocation( + boolean largeResult, CountDownLatch handlerReadyToReturn) { + return (userInput, ctx) -> { + var manager = ((DurableContextImpl) ctx).getExecutionManager(); + var checkpoint = manager.sendOperationUpdate(OperationUpdate.builder() + .id("bg-op") + .type(OperationType.STEP) + .action(OperationAction.START) + .build()); + assertSuspended(checkpoint); + assertTrue(manager.isCheckpointTokenRevoked()); + assertTrue(manager.isExecutionCompletedExceptionally()); + var result = largeResult ? "x".repeat(7 * 1024 * 1024) : "small-result"; + handlerReadyToReturn.countDown(); + return result; + }; + } + + private DurableExecutionOutput execute(DurableConfig config, BiFunction handler) { + var input = new DurableExecutionInput( + EXECUTION_ARN, + "token1", + CheckpointUpdatedExecutionState.builder() + .operations(executionOp()) + .build()); + return DurableExecutor.execute(input, null, get(String.class), handler, config); + } + + private void assertSuspended(CompletableFuture future) { + var error = assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)); + assertInstanceOf(SuspendExecutionException.class, error.getCause()); + } + @Test void invalidSuspensionThrowsIllegalDurableOperationExceptionAndNotifiesPluginsAsRetrying() { var input = new DurableExecutionInput( diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutorLifecycleRegressionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutorLifecycleRegressionTest.java new file mode 100644 index 0000000000..ea50f1f2af --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutorLifecycleRegressionTest.java @@ -0,0 +1,274 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.execution; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static software.amazon.lambda.durable.TypeToken.get; + +import java.time.Instant; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BiFunction; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState; +import software.amazon.awssdk.services.lambda.model.ExecutionDetails; +import software.amazon.awssdk.services.lambda.model.Operation; +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.client.DurableExecutionClient; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.model.DurableExecutionOutput; +import software.amazon.lambda.durable.operation.BaseDurableOperation; +import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; +import software.amazon.lambda.durable.plugin.InvocationEndInfo; +import software.amazon.lambda.durable.plugin.InvocationStatus; +import software.amazon.lambda.durable.serde.JacksonSerDes; + +@Timeout(30) +class DurableExecutorLifecycleRegressionTest { + @Test + void revocationMustWaitForRootHandlerBeforeReturning() throws Exception { + try (var scenario = new RootHandlerScenario(false)) { + scenario.awaitFirstCheckpoint(); + assertThrows( + TimeoutException.class, + () -> scenario.invocation.get(500, TimeUnit.MILLISECONDS), + "Invocation must not return while its root handler is still running"); + + scenario.releaseRoot.countDown(); + scenario.invocation.get(5, TimeUnit.SECONDS); + assertEquals(0, scenario.rootFinished.getCount()); + } + } + + @Test + void laterStepMustUnwindAfterTokenWithdrawal() throws Exception { + try (var scenario = new RootHandlerScenario(true)) { + scenario.awaitFirstCheckpoint(); + scenario.releaseRoot.countDown(); + assertTrue( + scenario.rootFinished.await(5, TimeUnit.SECONDS), + "Root remains blocked on a later operation after token withdrawal"); + scenario.invocation.get(5, TimeUnit.SECONDS); + } + } + + @Test + void failureSerializationStillFiresInvocationEnd() { + var ended = new AtomicInteger(); + var serializationFailure = new IllegalStateException("Cannot serialize failure"); + var config = DurableConfig.builder() + .withDurableExecutionClient(mock(DurableExecutionClient.class)) + .withSerDes(failingErrorSerializer(serializationFailure)) + .withPlugins(recordingPlugin(ended)) + .build(); + + var actual = assertThrows( + IllegalStateException.class, + () -> execute(config, (value, ctx) -> { + throw new IllegalArgumentException("handler failed"); + })); + + assertSame(serializationFailure, actual); + assertEquals(1, ended.get(), "Error serialization must not suppress the invocation-end hook"); + } + + @Test + void interruptedDrainStillFiresInvocationEnd() throws Exception { + var releaseStep = new CountDownLatch(1); + var handlerReturning = new CountDownLatch(1); + var ended = new CountDownLatch(1); + var endInfo = new AtomicReference(); + var rootThread = new AtomicReference(); + var userExecutor = Executors.newFixedThreadPool(2); + var caller = Executors.newSingleThreadExecutor(); + try { + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withExecutorService(userExecutor) + .withPlugins(new DurableExecutionPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + endInfo.set(info); + ended.countDown(); + } + }) + .build(); + + var invocation = caller.submit(() -> execute(config, (value, context) -> { + // Never awaited, so the root reaches drainOperations() while this step is still running. + context.stepAsync("blocked", String.class, step -> { + await(releaseStep); + return "late"; + }); + rootThread.set(Thread.currentThread()); + handlerReturning.countDown(); + return "root-result"; + })); + + await(handlerReturning); + // The interrupt flag is sticky, so the drain observes it whether it lands before or during the wait. + rootThread.get().interrupt(); + + assertTrue( + ended.await(5, TimeUnit.SECONDS), "An interrupted drain must still fire the invocation-end hook"); + assertEquals(InvocationStatus.RETRYING, endInfo.get().invocationStatus()); + assertInstanceOf(IllegalStateException.class, endInfo.get().executionError()); + + // Only now let the step finish, so the caller-side drain in ExecutionManager.close() can complete. + releaseStep.countDown(); + var failure = assertThrows(ExecutionException.class, () -> invocation.get(5, TimeUnit.SECONDS)); + assertInstanceOf(IllegalStateException.class, failure.getCause()); + } finally { + releaseStep.countDown(); + caller.shutdownNow(); + userExecutor.shutdownNow(); + assertTrue(caller.awaitTermination(5, TimeUnit.SECONDS), "Invocation executor did not terminate"); + assertTrue(userExecutor.awaitTermination(5, TimeUnit.SECONDS), "User executor did not terminate"); + } + } + + private static JacksonSerDes failingErrorSerializer(RuntimeException failure) { + return new JacksonSerDes() { + @Override + public String serialize(Object value) { + if (value instanceof Throwable) { + throw failure; + } + return super.serialize(value); + } + }; + } + + private static DurableExecutionPlugin recordingPlugin(AtomicInteger ended) { + return new DurableExecutionPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + ended.incrementAndGet(); + } + }; + } + + private static DurableExecutionOutput execute( + DurableConfig config, BiFunction handler) { + var execution = Operation.builder() + .id("exec") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails( + ExecutionDetails.builder().inputPayload("\"input\"").build()) + .build(); + var input = new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/lifecycle/exec", + "token", + CheckpointUpdatedExecutionState.builder().operations(execution).build()); + return DurableExecutor.execute(input, null, get(String.class), handler, config); + } + + private static class RootHandlerScenario implements AutoCloseable { + private final CountDownLatch rootPaused = new CountDownLatch(1); + private final CountDownLatch releaseRoot = new CountDownLatch(1); + private final CountDownLatch continuationReached = new CountDownLatch(1); + private final CountDownLatch rootFinished = new CountDownLatch(1); + private final AtomicReference firstOperation = new AtomicReference<>(); + private final AtomicReference lateOperation = new AtomicReference<>(); + private final ExecutorService userExecutor = Executors.newFixedThreadPool(2); + private final ExecutorService caller = Executors.newSingleThreadExecutor(); + private final Future invocation; + + private RootHandlerScenario(boolean startLaterStep) { + var config = DurableConfig.builder() + .withDurableExecutionClient(revokingClient()) + .withExecutorService(userExecutor) + .build(); + invocation = caller.submit(() -> execute(config, (value, context) -> runHandler(context, startLaterStep))); + } + + private DurableExecutionClient revokingClient() { + var delegate = TestUtils.createMockClient(); + var client = mock(DurableExecutionClient.class); + var checkpoints = new AtomicInteger(); + when(client.checkpoint(any(), any(), any())).thenAnswer(call -> { + await(rootPaused); + var response = delegate.checkpoint(call.getArgument(0), call.getArgument(1), call.getArgument(2)); + // Preserve accepted operation state so the same test can run against the pre-change implementation. + return checkpoints.incrementAndGet() == 1 + ? response.toBuilder().checkpointToken(null).build() + : response; + }); + return client; + } + + private String runHandler(DurableContext context, boolean startLaterStep) { + try { + firstOperation.set((BaseDurableOperation) context.stepAsync("async", String.class, step -> "done")); + rootPaused.countDown(); + // Hold the root between durable calls to make token withdrawal during orchestration reproducible. + await(releaseRoot); + if (startLaterStep) { + var next = context.stepAsync("next", String.class, step -> "next"); + lateOperation.set((BaseDurableOperation) next); + continuationReached.countDown(); + return next.get(); + } + return "root-result"; + } finally { + // Immediate suspension from stepAsync is also a valid way to unwind the root. + continuationReached.countDown(); + rootFinished.countDown(); + } + } + + private void awaitFirstCheckpoint() throws Exception { + await(rootPaused); + firstOperation.get().getRunningUserHandler().get(5, TimeUnit.SECONDS); + } + + @Override + public void close() throws InterruptedException { + releaseRoot.countDown(); + try { + await(continuationReached); + var operation = lateOperation.get(); + if (operation != null) { + // Cleanup only: release the stranded future after the assertion so a failing test cannot leak it. + operation.getCompletionFuture().completeExceptionally(new SuspendExecutionException()); + } + assertTrue(rootFinished.await(5, TimeUnit.SECONDS), "Root handler did not exit during cleanup"); + } finally { + caller.shutdownNow(); + userExecutor.shutdownNow(); + assertTrue(caller.awaitTermination(5, TimeUnit.SECONDS), "Invocation executor did not terminate"); + assertTrue(userExecutor.awaitTermination(5, TimeUnit.SECONDS), "User executor did not terminate"); + } + } + } + + private static void await(CountDownLatch latch) { + try { + assertTrue(latch.await(10, TimeUnit.SECONDS), "Timed out waiting for test coordination"); + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + throw new AssertionError(error); + } + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutorLifecycleTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutorLifecycleTest.java new file mode 100644 index 0000000000..d70291e2ab --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutorLifecycleTest.java @@ -0,0 +1,362 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.execution; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.argThat; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static software.amazon.lambda.durable.TypeToken.get; + +import java.time.Duration; +import java.time.Instant; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BiFunction; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.EnumSource; +import software.amazon.awssdk.services.lambda.model.CheckpointDurableExecutionResponse; +import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState; +import software.amazon.awssdk.services.lambda.model.ErrorObject; +import software.amazon.awssdk.services.lambda.model.ExecutionDetails; +import software.amazon.awssdk.services.lambda.model.Operation; +import software.amazon.awssdk.services.lambda.model.OperationAction; +import software.amazon.awssdk.services.lambda.model.OperationStatus; +import software.amazon.awssdk.services.lambda.model.OperationType; +import software.amazon.awssdk.services.lambda.model.OperationUpdate; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.client.DurableExecutionClient; +import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.model.DurableExecutionOutput; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.operation.BaseDurableOperation; +import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; +import software.amazon.lambda.durable.plugin.InvocationStatus; +import software.amazon.lambda.durable.serde.JacksonSerDes; +import software.amazon.lambda.durable.util.ExceptionHelper; + +@Timeout(10) +class DurableExecutorLifecycleTest { + private enum HandlerOutcome { + SUCCESS, + FAILURE, + RETRY, + LARGE_RESULT + } + + private final DurableExecutionClient client = mock(DurableExecutionClient.class); + private final DurableExecutionPlugin plugin = mock(DurableExecutionPlugin.class); + + @ParameterizedTest + @EnumSource(HandlerOutcome.class) + void queuedRevocationOverridesHandlerOutcomeAndClearsPluginError(HandlerOutcome outcome) { + when(client.checkpoint(any(), any(), any())) + .thenReturn(CheckpointDurableExecutionResponse.builder().build()); + var handlerError = switch (outcome) { + case FAILURE -> new IllegalArgumentException("handler failed"); + case RETRY -> + new UnrecoverableDurableExecutionException( + ErrorObject.builder().errorMessage("retry").build(), true); + default -> null; + }; + var output = execute( + config().withCheckpointDelay(Duration.ofHours(1)).build(), + handlerWithQueuedCheckpoint(outcome, handlerError)); + + assertEquals(ExecutionStatus.PENDING, output.status()); + assertNull(output.result()); + assertNull(output.error()); + verify(client, times(1)).checkpoint(any(), any(), any()); + verify(plugin, times(1)) + .onInvocationEnd(argThat(info -> info.invocationStatus() == InvocationStatus.PENDING + && info.executionError() == null + && info.executionResult() == null)); + } + + @Test + void queuedStepDrainsOnBoundedExecutor() throws Exception { + // The invocation drains operations on the root handler thread, so a queued step needs a second user thread. + // Single-thread executor support is outside this test's scope. + var userExecutor = Executors.newFixedThreadPool(2); + var caller = Executors.newSingleThreadExecutor(); + var handlerReady = new CountDownLatch(1); + var releaseHandler = new CountDownLatch(1); + var ran = new AtomicBoolean(); + try { + var config = config().withDurableExecutionClient(TestUtils.createMockClient()) + .withExecutorService(userExecutor) + .build(); + var invocation = + caller.submit(() -> execute(config, handlerWithQueuedStep(ran, handlerReady, releaseHandler))); + await(handlerReady); + assertThrows(TimeoutException.class, () -> invocation.get(100, TimeUnit.MILLISECONDS)); + releaseHandler.countDown(); + + assertEquals( + ExecutionStatus.SUCCEEDED, + invocation.get(5, TimeUnit.SECONDS).status()); + assertTrue(ran.get()); + verify(plugin, times(1)) + .onInvocationEnd(argThat(info -> info.invocationStatus() == InvocationStatus.SUCCEEDED)); + } finally { + releaseHandler.countDown(); + caller.shutdownNow(); + userExecutor.shutdownNow(); + } + } + + @Test + void queuedStepSkipsUserCodeAfterRevocation() { + // The root handler thread drains on itself, so the queued step needs a second thread to unwind on. A blocker + // occupies that thread until the handler has observed the revocation, so the step still starts after it. + var userExecutor = Executors.newFixedThreadPool(2); + var releaseQueuedStep = new CountDownLatch(1); + var ran = new AtomicBoolean(); + when(client.checkpoint(any(), any(), any())) + .thenReturn(CheckpointDurableExecutionResponse.builder().build()); + try { + userExecutor.execute(() -> await(releaseQueuedStep)); + var output = execute(config().withExecutorService(userExecutor).build(), (value, ctx) -> { + ctx.stepAsync("queued", String.class, stepCtx -> { + ran.set(true); + return "done"; + }); + try { + ctx.wait("revoke", Duration.ofSeconds(1)); + } finally { + releaseQueuedStep.countDown(); + } + return "unreachable"; + }); + + assertEquals(ExecutionStatus.PENDING, output.status()); + assertFalse(ran.get()); + verify(client, times(1)).checkpoint(any(), any(), any()); + } finally { + releaseQueuedStep.countDown(); + userExecutor.shutdownNow(); + } + } + + @Test + void runningStepFinishesBeforePendingPluginEvent() throws Exception { + var started = new CountDownLatch(1); + var release = new CountDownLatch(1); + var revoked = new CountDownLatch(1); + var finished = new AtomicBoolean(); + when(client.checkpoint(any(), any(), any())).thenAnswer(call -> { + await(started); + revoked.countDown(); + return CheckpointDurableExecutionResponse.builder().build(); + }); + var caller = Executors.newSingleThreadExecutor(); + try { + var invocation = + caller.submit(() -> execute(config().build(), handlerWithRunningStep(started, release, finished))); + await(revoked); + assertThrows(TimeoutException.class, () -> invocation.get(100, TimeUnit.MILLISECONDS)); + verify(plugin, never()).onInvocationEnd(any()); + release.countDown(); + + assertEquals( + ExecutionStatus.PENDING, invocation.get(5, TimeUnit.SECONDS).status()); + assertTrue(finished.get()); + verify(plugin, times(1)) + .onInvocationEnd(argThat(info -> + info.invocationStatus() == InvocationStatus.PENDING && info.executionError() == null)); + } finally { + release.countDown(); + caller.shutdownNow(); + } + } + + @Test + void continuationAdmissionClosesBeforeDrainingHandlers() { + var attempted = new AtomicBoolean(); + var dispatched = new AtomicBoolean(); + var output = execute( + config().withDurableExecutionClient(TestUtils.createMockClient()) + .build(), + (value, ctx) -> { + var manager = ((DurableContextImpl) ctx).getExecutionManager(); + var operation = mock(BaseDurableOperation.class); + when(operation.getOperationId()).thenReturn("drain-probe"); + when(operation.getRunningUserHandler()).thenAnswer(call -> { + attempted.set(true); + manager.runCheckpointContinuation(() -> {}, task -> { + dispatched.set(true); + task.run(); + }); + return CompletableFuture.completedFuture(null); + }); + manager.registerOperation(operation); + return "done"; + }); + + assertEquals(ExecutionStatus.SUCCEEDED, output.status()); + assertTrue(attempted.get()); + assertFalse(dispatched.get(), "Draining must not admit another checkpoint continuation"); + } + + @Test + void pollingStopsBeforeOutputSerialization() throws Exception { + var serializing = new CountDownLatch(1); + var release = new CountDownLatch(1); + var manager = new AtomicReference(); + var caller = Executors.newSingleThreadExecutor(); + try { + var config = config().withDurableExecutionClient(TestUtils.createMockClient()) + .withSerDes(blockedSerializer(serializing, release)) + .build(); + var invocation = caller.submit(() -> execute(config, (value, ctx) -> { + manager.set(((DurableContextImpl) ctx).getExecutionManager()); + ctx.waitAsync("unawaited", Duration.ofHours(1)); + return "done"; + })); + await(serializing); + var latePoller = + manager.get().pollForOperationUpdates("late", Instant.now().plusSeconds(60)); + assertTrue(latePoller.isCompletedExceptionally()); + var error = assertThrows(CompletionException.class, latePoller::join); + assertInstanceOf(SuspendExecutionException.class, error.getCause()); + release.countDown(); + + assertEquals( + ExecutionStatus.SUCCEEDED, + invocation.get(5, TimeUnit.SECONDS).status()); + } finally { + release.countDown(); + caller.shutdownNow(); + } + } + + @Test + void terminalCheckpointWithoutTokenRetainsPreviouslyObservedPluginStateWithoutPagination() { + var terminalOperation = executionOperation().toBuilder() + .status(OperationStatus.SUCCEEDED) + .build(); + when(client.checkpoint(any(), any(), any())) + .thenReturn(CheckpointDurableExecutionResponse.builder() + .newExecutionState(CheckpointUpdatedExecutionState.builder() + .operations(terminalOperation) + .nextMarker("unused") + .build()) + .build()); + + var output = execute(config().build(), (value, ctx) -> "x".repeat(6 * 1024 * 1024)); + + assertEquals(ExecutionStatus.SUCCEEDED, output.status()); + assertEquals("", output.result()); + verify(client, times(1)).checkpoint(any(), any(), any()); + verify(client, never()).getExecutionState(any(), any(), any()); + verify(plugin, times(1)) + .onInvocationEnd(argThat(info -> info.invocationStatus() == InvocationStatus.SUCCEEDED + && info.operations().get("exec").status() == OperationStatus.STARTED)); + } + + private DurableConfig.Builder config() { + return DurableConfig.builder().withDurableExecutionClient(client).withPlugins(plugin); + } + + private DurableExecutionOutput execute(DurableConfig config, BiFunction handler) { + var input = new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/lifecycle/exec", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(executionOperation()) + .build()); + return DurableExecutor.execute(input, null, get(String.class), handler, config); + } + + private Operation executionOperation() { + return Operation.builder() + .id("exec") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails( + ExecutionDetails.builder().inputPayload("\"input\"").build()) + .build(); + } + + private BiFunction handlerWithQueuedCheckpoint( + HandlerOutcome outcome, Throwable error) { + return (value, ctx) -> { + ((DurableContextImpl) ctx) + .getExecutionManager() + .sendOperationUpdate(OperationUpdate.builder() + .id("queued") + .type(OperationType.STEP) + .action(OperationAction.START) + .build()); + if (error != null) { + ExceptionHelper.sneakyThrow(error); + } + return outcome == HandlerOutcome.LARGE_RESULT ? "x".repeat(6 * 1024 * 1024) : "done"; + }; + } + + private BiFunction handlerWithQueuedStep( + AtomicBoolean ran, CountDownLatch handlerReady, CountDownLatch releaseHandler) { + return (value, ctx) -> { + ctx.stepAsync("queued", String.class, stepCtx -> { + ran.set(true); + return "done"; + }); + handlerReady.countDown(); + await(releaseHandler); + return "done"; + }; + } + + private BiFunction handlerWithRunningStep( + CountDownLatch started, CountDownLatch release, AtomicBoolean finished) { + return (value, ctx) -> ctx.step("running", String.class, stepCtx -> { + started.countDown(); + await(release); + finished.set(true); + return "done"; + }); + } + + private JacksonSerDes blockedSerializer(CountDownLatch serializing, CountDownLatch release) { + return new JacksonSerDes() { + @Override + public String serialize(Object value) { + serializing.countDown(); + await(release); + return super.serialize(value); + } + }; + } + + private void await(CountDownLatch latch) { + try { + assertTrue(latch.await(5, TimeUnit.SECONDS), "Timed out waiting for the test latch"); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java index afecc8ddb9..7bb872911c 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java @@ -7,6 +7,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import java.time.Instant; import java.util.ArrayList; import java.util.Collections; import java.util.List; @@ -27,12 +28,15 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.CsvSource; import org.junit.jupiter.params.provider.ValueSource; +import software.amazon.awssdk.services.lambda.model.CheckpointDurableExecutionResponse; import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState; import software.amazon.awssdk.services.lambda.model.ErrorObject; import software.amazon.awssdk.services.lambda.model.GetDurableExecutionStateResponse; import software.amazon.awssdk.services.lambda.model.Operation; +import software.amazon.awssdk.services.lambda.model.OperationAction; import software.amazon.awssdk.services.lambda.model.OperationStatus; import software.amazon.awssdk.services.lambda.model.OperationType; +import software.amazon.awssdk.services.lambda.model.OperationUpdate; import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.TestUtils; import software.amazon.lambda.durable.client.DurableExecutionClient; @@ -215,6 +219,106 @@ private Operation namedStepOp(String id, OperationStatus status) { .build(); } + @Test + void operationRegisteredBeforeTokenWithdrawalIsSuspended() throws Exception { + try (var manager = createManager(List.of(executionOp()))) { + var operation = mockRegisteredOperation("step"); + manager.registerOperation(operation); + assertFalse(operation.getCompletionFuture().isDone()); + + withdrawToken(manager); + + assertSuspended(operation); + } + } + + @Test + void operationRegisteredAfterTokenWithdrawalIsSuspended() throws Exception { + try (var manager = createManager(List.of(executionOp()))) { + withdrawToken(manager); + var operation = mockRegisteredOperation("late"); + + manager.registerOperation(operation); + + assertSuspended(operation); + } + } + + @Test + void registrationDuringPollerCompletionObservesWithdrawalBeforeSuspensionSignal() throws Exception { + try (var manager = createManager(List.of(executionOp()))) { + var poller = manager.pollForOperationUpdates("poll", Instant.now().plusSeconds(3600)); + var registration = poller.handle((ignored, failure) -> { + // Pollers complete after revocation is published, but before the invocation is signaled. + assertTrue(manager.isCheckpointTokenRevoked()); + assertFalse(manager.isExecutionCompletedExceptionally()); + assertInstanceOf(SuspendExecutionException.class, failure); + + var operation = mockRegisteredOperation("late"); + manager.registerOperation(operation); + assertSuspended(operation); + + return null; + }); + + withdrawToken(manager); + + registration.get(5, TimeUnit.SECONDS); + } + } + + @Test + void registrationFinishingAfterTokenWithdrawalIsSuspended() throws Exception { + try (var manager = createManager(List.of(executionOp()))) { + var registrationStarted = new CountDownLatch(1); + var releaseRegistration = new CountDownLatch(1); + + var operation = mockRegisteredOperation("late"); + when(operation.getOperationId()).thenAnswer(call -> { + registrationStarted.countDown(); + assertTrue(releaseRegistration.await(5, TimeUnit.SECONDS)); + return "late"; + }); + + var registration = CompletableFuture.runAsync(() -> manager.registerOperation(operation)); + try { + assertTrue(registrationStarted.await(5, TimeUnit.SECONDS)); + withdrawToken(manager); + } finally { + releaseRegistration.countDown(); + } + registration.get(5, TimeUnit.SECONDS); + + assertSuspended(operation); + } + } + + private BaseDurableOperation mockRegisteredOperation(String id) { + var operation = mock(BaseDurableOperation.class); + when(operation.getOperationId()).thenReturn(id); + when(operation.getCompletionFuture()).thenReturn(new CompletableFuture<>()); + return operation; + } + + private void withdrawToken(ExecutionManager manager) throws Exception { + when(client.checkpoint(any(), any(), any())) + .thenReturn(CheckpointDurableExecutionResponse.builder().build()); + var checkpoint = manager.sendOperationUpdate(OperationUpdate.builder() + .id("revoke") + .type(OperationType.STEP) + .action(OperationAction.START) + .build()); + var error = assertThrows(ExecutionException.class, () -> checkpoint.get(5, TimeUnit.SECONDS)); + assertInstanceOf(SuspendExecutionException.class, error.getCause()); + } + + private void assertSuspended(BaseDurableOperation operation) { + assertTrue(operation.getCompletionFuture().isCompletedExceptionally()); + var error = assertThrows( + ExecutionException.class, () -> operation.getCompletionFuture().get(5, TimeUnit.SECONDS)); + assertInstanceOf(SuspendExecutionException.class, error.getCause()); + } + @Test void startsInReplayModeWhenOperationsExist() { var manager = createManager(List.of(executionOp(), stepOp("1", OperationStatus.SUCCEEDED))); @@ -494,7 +598,9 @@ protected void replay(Operation existing) {} waiter.get(5, TimeUnit.SECONDS); checkpoint.get(5, TimeUnit.SECONDS); assertDoesNotThrow(() -> manager.deregisterActiveThread(otherThread)); - assertThrows(SuspendExecutionException.class, () -> manager.deregisterActiveThread(callerThread)); + var suspension = + assertThrows(SuspendExecutionException.class, () -> manager.deregisterActiveThread(callerThread)); + assertEquals("Execution suspended for wait operation", suspension.getMessage()); } @Test