diff --git a/docs/advanced/configuration.md b/docs/advanced/configuration.md index 9097bb46a..e4969f146 100644 --- a/docs/advanced/configuration.md +++ b/docs/advanced/configuration.md @@ -105,3 +105,34 @@ file retention. Custom serializers can override `serialize(value, SerDesContext)` and `deserialize(data, typeToken, SerDesContext)` to receive stable execution/entity identity. Existing implementations inherit defaults that call the original methods. + +## Invocation deadline cleanup + +The SDK captures the Lambda invocation's remaining time once and uses a monotonic +budget for its waits and cleanup. Deadline exhaustion is an invocation interruption +that can be retried; it is not checkpointed as a business-operation failure. +Completed operation results and failures remain available for replay. + +Normal finalization drains invocation-owned user work before result preparation +and invocation End. Cleanup first requests cooperative cancellation when work does +not settle within its allowance, then rejects new durable admission. Required +checkpoint waits and cleanup use the remaining budget. A terminal execution result +already acknowledged by the service remains authoritative, while a JVM fatal is +never replaced by a cleanup timeout or an acknowledged-result shortcut. + +Interruption does not skip the remaining cleanup attempts: the SDK records it, +continues best-effort joins and checkpoint shutdown with the remaining budget, and +restores the interrupt flag. If an uncooperative root exits after its caller has +already timed out, its original worker still runs the matching invocation End once +to restore Start-installed scopes. This eventual cleanup does not admit new durable +work or prepare another result. An End hook that has already started can itself +remain blocked; the SDK cannot forcibly stop arbitrary callbacks. + +This is not an unconditional bounded-return or worker-reclamation guarantee. +Already-running application I/O, synchronous initial-state reads, custom clients, +and plugin hooks cannot be forcibly terminated. In particular, an executor that +runs the root inline, including a saturated `CallerRunsPolicy`, can block the runtime +caller before it reaches its timeout-aware wait. Use finite downstream I/O timeouts +and executor configurations that let the runtime caller make progress. Cancellation +cannot undo a request already accepted by the service; interrupted side effects +still require the documented retry and idempotency strategy. diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerMdcIntegrationTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerMdcIntegrationTest.java index 1369254cd..8c57a87d1 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerMdcIntegrationTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerMdcIntegrationTest.java @@ -22,6 +22,7 @@ import org.slf4j.MDC; import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.model.ExecutionStatus; import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; import software.amazon.lambda.durable.plugin.InvocationEndInfo; @@ -262,10 +263,19 @@ private static void awaitMdcLatch(CountDownLatch latch) { private static void awaitMdcCallerJoin(Thread caller) { var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(3); while (System.nanoTime() < deadline) { - if (caller.getState() == Thread.State.WAITING - && Arrays.stream(caller.getStackTrace()) + var state = caller.getState(); + var frames = caller.getStackTrace(); + var legacyJoin = Arrays.stream(frames) + .anyMatch(frame -> frame.getClassName().equals(CompletableFuture.class.getName()) + && frame.getMethodName().equals("join")); + var budgetedWait = Arrays.stream(frames) + .anyMatch(frame -> frame.getClassName().equals(ExecutionManager.class.getName()) + && frame.getMethodName().equals("awaitInvocationOutcome")) + && Arrays.stream(frames) .anyMatch(frame -> frame.getClassName().equals(CompletableFuture.class.getName()) - && frame.getMethodName().equals("join"))) return; + && frame.getMethodName().equals("get")); + if ((state == Thread.State.WAITING || state == Thread.State.TIMED_WAITING) && (legacyJoin || budgetedWait)) + return; LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1)); } fail("Caller did not attach its completion observer before releasing the handler"); diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndCompatibilityTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndCompatibilityTest.java index 9cb19b13f..9bf5d790b 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndCompatibilityTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndCompatibilityTest.java @@ -21,6 +21,7 @@ import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.execution.DurableExecutor; +import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.model.DurableExecutionInput; import software.amazon.lambda.durable.model.ExecutionStatus; import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; @@ -262,10 +263,19 @@ private record Observation( private static void awaitCallerJoin(Thread caller) { var end = System.nanoTime() + TimeUnit.SECONDS.toNanos(3); while (System.nanoTime() < end) { - if (caller.getState() == Thread.State.WAITING - && Arrays.stream(caller.getStackTrace()) + var state = caller.getState(); + var frames = caller.getStackTrace(); + var legacyJoin = Arrays.stream(frames) + .anyMatch(frame -> frame.getClassName().equals(CompletableFuture.class.getName()) + && frame.getMethodName().equals("join")); + var budgetedWait = Arrays.stream(frames) + .anyMatch(frame -> frame.getClassName().equals(ExecutionManager.class.getName()) + && frame.getMethodName().equals("awaitInvocationOutcome")) + && Arrays.stream(frames) .anyMatch(frame -> frame.getClassName().equals(CompletableFuture.class.getName()) - && frame.getMethodName().equals("join"))) return; + && frame.getMethodName().equals("get")); + if ((state == Thread.State.WAITING || state == Thread.State.TIMED_WAITING) && (legacyJoin || budgetedWait)) + return; LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1)); } fail("invocation caller did not reach its future join before releasing the handler"); 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 1c61311f4..f935d2c4e 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 @@ -21,6 +21,7 @@ import software.amazon.lambda.durable.config.*; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.DurableExecutor; +import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.model.DurableExecutionInput; import software.amazon.lambda.durable.model.ExecutionStatus; import software.amazon.lambda.durable.model.WaitForConditionResult; @@ -273,7 +274,9 @@ 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; + .anyMatch(frame -> frame.getClassName().equals(ExecutionManager.class.getName()) + && (frame.getMethodName().equals("waitForRunningUserHandlers") + || frame.getMethodName().equals("waitForRunningUserHandlersUntil")))) return; LockSupport.parkNanos(100_000); } fail("Root did not begin draining the pending child"); diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ContinuationEndCleanupOrderTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ContinuationEndCleanupOrderTest.java index de8b5c5e0..fdb3c266f 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ContinuationEndCleanupOrderTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ContinuationEndCleanupOrderTest.java @@ -341,9 +341,17 @@ public void onInvocationEnd(InvocationEndInfo info) { private static void awaitCallerOutcomeJoin(Thread caller) { var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(3); while (System.nanoTime() < deadline) { - if (Arrays.stream(caller.getStackTrace()) + var frames = caller.getStackTrace(); + var legacyJoin = Arrays.stream(frames) .anyMatch(frame -> frame.getClassName().equals(CompletableFuture.class.getName()) - && frame.getMethodName().equals("join"))) return; + && frame.getMethodName().equals("join")); + var budgetedWait = Arrays.stream(frames) + .anyMatch(frame -> frame.getClassName().equals(ExecutionManager.class.getName()) + && frame.getMethodName().equals("awaitInvocationOutcome")) + && Arrays.stream(frames) + .anyMatch(frame -> frame.getClassName().equals(CompletableFuture.class.getName()) + && frame.getMethodName().equals("get")); + if (legacyJoin || budgetedWait) return; Thread.yield(); } fail("Invocation caller did not reach its installed outcome observer"); 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 ab0edd6a9..e179ce10d 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 @@ -8,7 +8,9 @@ import java.util.Objects; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.function.Consumer; import java.util.function.Function; @@ -23,7 +25,7 @@ * @param Request type */ public class ApiRequestDelayedBatcher { - private static final Duration MAX_DELAY = Duration.ofMinutes(60); + private static final long NO_FLUSH_DEADLINE = Long.MAX_VALUE; /** Maximum items allowed in a single batch */ private final int maxItemCount; @@ -108,6 +110,12 @@ void shutdown() { flushingQueueFuture.join(); } + /** Flushes pending requests without waiting beyond the caller's remaining cleanup budget. */ + void shutdown(Duration timeout) throws InterruptedException, ExecutionException, TimeoutException { + flush(); + flushingQueueFuture.get(Math.max(0, timeout.toNanos()), TimeUnit.NANOSECONDS); + } + /** Schedules pending requests immediately without waiting for the current batch. */ void flush() { synchronized (delayedBatch) { @@ -121,8 +129,8 @@ void flush() { /** clear the current batch and creates a new batch */ private void initializeDelayedBatch() { this.delayedBatch.clear(); - // MAX_DELAY is longer than a single Lambda invocation - this.delayedBatchFlushTime = System.nanoTime() + MAX_DELAY.toNanos(); + // No timer is scheduled until an item supplies its flush time. + this.delayedBatchFlushTime = NO_FLUSH_DEADLINE; // the timer future is created initially without a timeout until an item is added to the batch this.delayedBatchFlushTimer = new CompletableFuture<>(); 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 67a02f5c0..d00d3c4dc 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 @@ -11,6 +11,8 @@ import java.util.Objects; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeoutException; import java.util.function.BooleanSupplier; import java.util.function.Consumer; import org.slf4j.Logger; @@ -18,9 +20,11 @@ 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.OperationAction; 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.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.retry.PollingStrategies; import software.amazon.lambda.durable.retry.PollingStrategy; import software.amazon.lambda.durable.util.DurableApiErrorClassifier; @@ -47,6 +51,9 @@ class CheckpointManager { private final BooleanSupplier tryStartCheckpointProcessing; private final Runnable onCheckpointProcessingFinish; private final Runnable onTokenRevoked; + private final Runnable requestAdmission; + private final Object checkpointLock = new Object(); + private volatile boolean terminalExecutionAcknowledged; private volatile String checkpointToken; private boolean pollingStopped; // Guarded by pollingFutures. /** @@ -71,6 +78,27 @@ class CheckpointManager { BooleanSupplier tryStartCheckpointProcessing, Runnable onCheckpointProcessingFinish, Runnable onTokenRevoked) { + this( + config, + durableExecutionArn, + checkpointToken, + callback, + tryStartCheckpointProcessing, + onCheckpointProcessingFinish, + onTokenRevoked, + () -> {}); + } + + CheckpointManager( + DurableConfig config, + String durableExecutionArn, + String checkpointToken, + Consumer> callback, + BooleanSupplier tryStartCheckpointProcessing, + Runnable onCheckpointProcessingFinish, + Runnable onTokenRevoked, + Runnable requestAdmission) { + this.requestAdmission = requestAdmission; this.config = config; this.durableExecutionArn = durableExecutionArn; this.callback = callback; @@ -186,6 +214,15 @@ void shutdown() { checkpointApiRequestDelayedBatcher.shutdown(); } + void shutdown(Duration timeout) throws InterruptedException, ExecutionException, TimeoutException { + stopPolling(new IllegalStateException("CheckpointManager shutdown")); + checkpointApiRequestDelayedBatcher.shutdown(timeout); + } + + boolean isTerminalExecutionAcknowledged() { + return terminalExecutionAcknowledged; + } + private void stopPolling(Throwable cause) { List>> allFutures; synchronized (pollingFutures) { @@ -204,6 +241,11 @@ private void stopPolling(Throwable cause) { * nextMarker) */ List fetchAllPages(CheckpointUpdatedExecutionState checkpointUpdatedExecutionState) { + return fetchAllPages(checkpointUpdatedExecutionState, false); + } + + private List fetchAllPages( + CheckpointUpdatedExecutionState checkpointUpdatedExecutionState, boolean terminal) { List operations = new ArrayList<>(); if (checkpointUpdatedExecutionState == null) { return operations; @@ -213,6 +255,12 @@ List fetchAllPages(CheckpointUpdatedExecutionState checkpointUpdatedE } var nextMarker = checkpointUpdatedExecutionState.nextMarker(); while (nextMarker != null && !nextMarker.isEmpty()) { + try { + requestAdmission.run(); + } catch (UnrecoverableDurableExecutionException failure) { + if (!terminal || !ExecutionManager.isLifecycleFailure(failure)) throw failure; + break; // Preserve inline/already-read metadata without issuing a new request after terminal ACK. + } var startTime = System.nanoTime(); try { var response = config.getDurableExecutionClient() @@ -231,12 +279,15 @@ List fetchAllPages(CheckpointUpdatedExecutionState checkpointUpdatedE } private void checkpointBatch(List updates) { - synchronized (pollingFutures) { - // A READY recheck can complete a poll before its scheduled batch. Keep only consumers still waiting. - pollingFutures.entrySet().removeIf(entry -> { - entry.getValue().removeIf(CompletableFuture::isDone); - return entry.getValue().isEmpty(); - }); + synchronized (checkpointLock) { + synchronized (pollingFutures) { + // A READY recheck can complete a poll before its scheduled batch. Keep only consumers still waiting. + pollingFutures.entrySet().removeIf(entry -> { + entry.getValue().removeIf(CompletableFuture::isDone); + return entry.getValue().isEmpty(); + }); + } + // filter the null values from pollers var request = updates.stream().filter(Objects::nonNull).toList(); @@ -277,7 +328,14 @@ private void checkpointBatch(List updates) { private void checkpointRequest(List request) { var startTime = System.nanoTime(); logger.debug("Calling durable checkpoint API with {} updates: {}", request.size(), request); + requestAdmission.run(); var response = config.getDurableExecutionClient().checkpoint(durableExecutionArn, checkpointToken, request); + if (request.stream() + .anyMatch(update -> + update.type() == OperationType.EXECUTION && update.action() == OperationAction.SUCCEED)) + terminalExecutionAcknowledged = true; + // A late response is durable at the service, but must not resume old SDK work after invocation cleanup. + if (!terminalExecutionAcknowledged) requestAdmission.run(); logger.debug("Durable checkpoint API called (latency={}ns): {}.", System.nanoTime() - startTime, response); checkpointToken = response.checkpointToken(); @@ -288,7 +346,7 @@ private void checkpointRequest(List request) { var state = response.newExecutionState(); if (state != null) { - processUpdatedOperations(fetchAllPages(state)); + processUpdatedOperations(fetchAllPages(state, terminalExecutionAcknowledged)); } } @@ -322,7 +380,10 @@ private void processUpdatedOperations(List operations) { callback.accept(operations); for (var operation : operations) { - var pollers = pollingFutures.remove(operation.id()); + List> pollers; + synchronized (pollingFutures) { + pollers = pollingFutures.remove(operation.id()); + } if (pollers != null) { completedFutures += pollers.size(); pollers.forEach(poller -> poller.complete(operation)); 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 dad337112..7f3d81cbb 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 @@ -80,10 +80,32 @@ public static DurableExecutionOutput execute( manager.registerActiveThread(ROOT_THREAD_ID); var invocation = new Invocation<>(input, lambdaContext, inputType, handler, config, manager); try { - return invocation.execute().join(); + var completion = invocation.execute(); + var ready = invocation.outcome.applyToEither( + completion.handle((ignored, failure) -> new Outcome(null, failure)), Function.identity()); + var selected = manager.awaitInvocationOutcome(ready); + var originalFailure = selected.failure() != null + ? selected.failure() + : ExceptionHelper.unwrapCompletableFuture( + selected.result().failure()); + var finalFailure = manager.finishInvocation(originalFailure); + var completedFailure = completion.isDone() + ? completion + .handle((ignored, failure) -> ExceptionHelper.unwrapCompletableFuture(failure)) + .join() + : null; + if (ExecutionManager.isFatal(completedFailure)) ExceptionHelper.sneakyThrow(completedFailure); + if (ExecutionManager.isFatal(finalFailure)) ExceptionHelper.sneakyThrow(finalFailure); + if (ExecutionManager.isLifecycleFailure(finalFailure) && manager.isTerminalExecutionAcknowledged()) + return DurableExecutionOutput.success(""); + if (finalFailure != null && (selected.failure() != null || finalFailure != originalFailure)) + ExceptionHelper.sneakyThrow(finalFailure); + return completion.join(); } catch (CompletionException failure) { ExceptionHelper.sneakyThrow(ExceptionHelper.unwrapCompletableFuture(failure)); return null; + } finally { + manager.markInvocationReported(); } } } @@ -101,6 +123,8 @@ private static final class Invocation { private final boolean isFirstInvocation; private I userInput; private boolean started; + private boolean ended; + private CompletableFuture> outcome; private Invocation( DurableExecutionInput input, @@ -124,9 +148,11 @@ private CompletableFuture execute() { var body = new CompletableFuture(); // Select the invocation outcome before scheduling, including with an inline executor. A suspension or // termination can win while the handler is still unwinding; its finally block must not replace it. - var outcome = manager.runUntilCompleteOrSuspend(body).handle(Outcome::new); + outcome = manager.runUntilCompleteOrSuspend(body).handle(Outcome::new); Supplier task = () -> { - Outcome.capture(this::invokeHandler).complete(body); + var handlerOutcome = Outcome.capture(this::invokeHandler); + manager.markRootHandlerExited(); + handlerOutcome.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 @@ -138,6 +164,11 @@ private CompletableFuture execute() { // 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. + var selectedFailure = ExceptionHelper.unwrapCompletableFuture(selected.failure()); + if (ExecutionManager.isFatal(selectedFailure)) { + if (selectedFailure != drainFailure) selectedFailure.addSuppressed(drainFailure); + return failDelivery(selectedFailure); + } return failDelivery(drainFailure); } return finishInvocation(selected.value(), selected.failure()); @@ -156,7 +187,9 @@ private O invokeHandler() { // Deserialize only once. Even failed input gets paired start/end hooks with a null input value. inputFailure = failure; } + manager.checkUserFunctionExecutionAllowed(); fireOnInvocationStart(); + manager.checkUserFunctionExecutionAllowed(); if (inputFailure != null) ExceptionHelper.sneakyThrow(inputFailure); var context = DurableContextImpl.createRootContext(manager, config, lambdaContext); DurableContextImpl.setCurrentContext(context); @@ -182,6 +215,12 @@ private void fireOnInvocationStart() { } private DurableExecutionOutput finishInvocation(O value, Throwable failure) { + if (manager.isInvocationReported()) + return failDelivery( + manager.getLifecycleFailure() != null + ? manager.getLifecycleFailure() + : new IllegalStateException("Invocation has already returned")); + if (manager.getLifecycleFailure() != null) return failDelivery(manager.getLifecycleFailure()); // 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()) @@ -247,7 +286,19 @@ private DurableExecutionOutput failDelivery(Throwable deliveryFailure) { } private void fireOnInvocationEnd(InvocationStatus status, Throwable error, Object result) { - if (!started) return; + if (!started || ended) return; + ended = true; + // An uncooperative root may exit after the caller's timeout. Its owner still owes Start's scopes one + // matching End; this cleanup must not reopen durable work or prepare another execution result. + if (manager.isInvocationReported() + && !manager.isTerminalExecutionAcknowledged() + && !ExecutionManager.isFatal(error)) { + status = InvocationStatus.RETRYING; + error = manager.getLifecycleFailure() != null + ? manager.getLifecycleFailure() + : new IllegalStateException("Invocation already ended before owner cleanup"); + result = null; + } plugins.onInvocationEnd(new InvocationEndInfo( requestId, executionArn, @@ -488,14 +539,13 @@ private static String handleLargePayload(ExecutionManager executionManager, Stri LAMBDA_RESPONSE_SIZE_LIMIT); // Checkpoint the large result and wait for it to complete - executionManager - .sendExecutionResultUpdate(OperationUpdate.builder() + executionManager.awaitCheckpointCompletion( + executionManager.sendExecutionResultUpdate(OperationUpdate.builder() .type(OperationType.EXECUTION) .id(executionManager.getExecutionOperation().id()) .action(OperationAction.SUCCEED) .payload(outputPayload) - .build()) - .join(); + .build())); // Return empty result, we checkpointed the data manually logger.debug("Execution completed (large response checkpointed)"); 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 91cbc48a5..9a65682f3 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 @@ -18,6 +18,7 @@ import java.util.Set; import java.util.concurrent.CancellationException; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; @@ -26,6 +27,7 @@ import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Supplier; @@ -74,6 +76,20 @@ public class ExecutionManager implements SafeCloseable { private static final AtomicLong LOCAL_INVOCATION_SEQUENCE = new AtomicLong(); + // Stop waiting for the root early enough to cancel tasks and flush checkpoints before the platform deadline. + private static final Duration EXECUTION_DEADLINE_HEADROOM = Duration.ofMillis(500); + // Leave a final margin for response serialization and return through the Lambda runtime wrapper. + private static final Duration RESPONSE_HEADROOM = Duration.ofMillis(100); + // Lambda always supplies a deadline; this cap also keeps local/null-context cleanup bounded. + private static final Duration MAX_CLEANUP_DURATION = Duration.ofSeconds(30); + private static final Duration MAX_GRACEFUL_DRAIN_DURATION = Duration.ofSeconds(5); + private static final Duration MAX_CANCELLATION_DRAIN_DURATION = Duration.ofSeconds(1); + private static final Duration CHECKPOINT_CLEANUP_RESERVE = Duration.ofMillis(100); + private static final String LIFECYCLE_ERROR_TYPE = + "software.amazon.lambda.durable.execution.InvocationLifecycleException"; + + record InvocationOutcome(T result, Throwable failure) {} + enum LifecycleState { /** The invocation accepts new durable operations, executor tasks, checkpoints, and polls. */ OPEN, @@ -108,7 +124,13 @@ enum LifecycleState { private final Map> activeExecutorTasks = new ConcurrentHashMap<>(); // The current task is shared across managers on a thread; admission verifies the exact handle's owner. private static final ThreadLocal> currentExecutorTask = new ThreadLocal<>(); - private LifecycleState lifecycleState = LifecycleState.OPEN; + private volatile LifecycleState lifecycleState = LifecycleState.OPEN; + private volatile boolean rootHandlerExited; + private final CompletableFuture rootHandlerExitSignal = new CompletableFuture<>(); + private volatile boolean invocationReported; + private boolean cleanupFinished; + private final AtomicReference lifecycleFailure = new AtomicReference<>(); + private final ThreadLocal cleanupInterrupted = ThreadLocal.withInitial(() -> false); // Guarded by admissionLock. Counts enqueue calls, not the lifetime of their returned futures. private int checkpointAdmissionsInFlight; @@ -161,7 +183,8 @@ public ExecutionManager(DurableExecutionInput input, DurableConfig config, Conte this::onCheckpointComplete, this::tryStartCheckpointProcessing, this::finishCheckpointProcessing, - this::signalSuspension); + this::signalSuspension, + this::checkCheckpointRequestAllowed); this.operationStorage = checkpointManager.fetchAllPages(input.initialExecutionState()).stream() .collect(Collectors.toConcurrentMap(Operation::id, op -> op)); @@ -938,14 +961,83 @@ public CompletableFuture pollForOperationUpdates(String operationId, /** Reserved result delivery after operation admission has closed; the caller observes this checkpoint. */ CompletableFuture sendExecutionResultUpdate(OperationUpdate update) { + checkCheckpointRequestAllowed(); return checkpointManager.checkpoint(update); } + void awaitCheckpointCompletion(CompletableFuture future) { + var deadline = + deadlineNanos == null ? addToNow(MAX_CLEANUP_DURATION) : deadlineNanos - RESPONSE_HEADROOM.toNanos(); + try { + future.get(remainingNanos(deadline), TimeUnit.NANOSECONDS); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw lifecycleFailure("Interrupted while waiting for checkpoint completion", interrupted); + } catch (TimeoutException timeout) { + if (!checkpointManager.isTerminalExecutionAcknowledged()) + throw lifecycleFailure("Checkpoint completion exceeded the invocation deadline", timeout); + } catch (ExecutionException failure) { + ExceptionHelper.sneakyThrow(ExceptionHelper.unwrapCompletableFuture(failure.getCause())); + } + } + + void markRootHandlerExited() { + rootHandlerExited = true; + rootHandlerExitSignal.complete(null); + } + + void markInvocationReported() { + invocationReported = true; + } + + boolean isInvocationReported() { + return invocationReported; + } + + UnrecoverableDurableExecutionException getLifecycleFailure() { + return lifecycleFailure.get(); + } + + boolean isTerminalExecutionAcknowledged() { + return checkpointManager.isTerminalExecutionAcknowledged(); + } + + /** Fences actual user callbacks, including callbacks delayed inside a plugin start hook. */ + public void checkUserFunctionExecutionAllowed() { + if (executionExceptionFuture.isCompletedExceptionally()) { + try { + executionExceptionFuture.join(); + } catch (CompletionException stopped) { + ExceptionHelper.sneakyThrow(ExceptionHelper.unwrapCompletableFuture(stopped)); + } + } + checkInvocationAdmission(); + } + + private void checkInvocationAdmission() { + if (invocationReported || lifecycleFailure.get() != null) { + var failure = lifecycleFailure.get(); + throw failure != null ? failure : lifecycleFailure("Invocation has already returned", null); + } + if (deadlineNanos != null && remainingNanos(deadlineNanos - EXECUTION_DEADLINE_HEADROOM.toNanos()) == 0) + throw lifecycleFailure("Invocation deadline reached before user work", new TimeoutException()); + } + + private void checkCheckpointRequestAllowed() { + if (invocationReported || lifecycleFailure.get() != null) { + var failure = lifecycleFailure.get(); + throw failure != null ? failure : lifecycleFailure("Invocation has already returned", null); + } + if (deadlineNanos != null && remainingNanos(deadlineNanos - RESPONSE_HEADROOM.toNanos()) == 0) + throw lifecycleFailure("Invocation deadline reached before checkpoint I/O", new TimeoutException()); + } + // ===== Utilities ===== /** Shutdown the checkpoint batcher. */ @Override public void close() { - drainOperations(); + var failure = finishInvocation(null); + if (failure != null) ExceptionHelper.sneakyThrow(failure); } /** @@ -954,59 +1046,33 @@ public void close() { * 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. beginDraining(); - stopCheckpointContinuations(); - waitForRunningUserHandlers(); - drainExecutorTasks(); - // An interrupted wait leaves admitted work draining. The caller's subsequent close must still allow its - // completion checkpoints; closing here on failure would discard results and repeat completed side effects. - closeLifecycle(); - validateExecutorPool(); - checkpointManager.shutdown(); - } - - private void stopCheckpointContinuations() { - List owners; - synchronized (activeThreads) { - closing = true; - owners = new ArrayList<>(checkpointContinuations.values()); - } - owners.forEach(ExecutionManager::stopContinuationOwner); - synchronized (activeThreads) { - while (!checkpointContinuations.isEmpty()) { - try { - // wait releases this monitor so tasks can finish and release their registrations. No user work - // or future join is performed while holding the coordination monitor. - activeThreads.wait(); - } catch (InterruptedException interrupted) { - Thread.currentThread().interrupt(); - throw new IllegalStateException( - "Interrupted while waiting for checkpoint continuations", interrupted); - } - } + var hardDeadline = cleanupDeadlineNanos(); + Throwable failure = null; + if (Thread.interrupted()) cleanupInterrupted.set(true); + try { + if (cleanupInterrupted.get()) throw new InterruptedException("Handler was interrupted before draining"); + stopCheckpointContinuationsUntil(hardDeadline); + waitForRunningUserHandlersUntil(hardDeadline); + if (!awaitTasksUntil(hardDeadline, false)) + failure = lifecycleFailure("Operation tasks exceeded the cleanup deadline", new TimeoutException()); + } catch (InterruptedException interrupted) { + cleanupInterrupted.set(true); + failure = new IllegalStateException("Interrupted while draining invocation operations", interrupted); + // An interrupted owner drain is not itself an invocation deadline. Keep already-admitted work able to + // finish and checkpoint during the remaining budget; the caller still cancels if that budget expires. + failure = settleInterruptedCleanup(hardDeadline, false, failure); + } catch (RuntimeException problem) { + failure = problem; + } finally { + failure = stopCheckpointAdmission(hardDeadline, failure); + failure = shutdownCheckpoints(hardDeadline, failure); + closeLifecycle(); + restoreCleanupInterrupt(); } - } - - 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(); - if (userHandlerFuture != null && !userHandlerFuture.isDone()) { - // Some user threads can still be running because - // the operations that run them have never been waiting for and the execution has completed. - logger.info("Waiting for operation to complete before shutting down: {}", op.getOperationId()); - try { - userHandlerFuture.get(); - } catch (InterruptedException | CancellationException e) { - // if the user handler is stuck - throw new IllegalStateException( - "Stuck running user handler when shutting down: " + op.getOperationId()); - } catch (Exception e) { - // ok if the future completed exceptionally - } - } + if (failure != null) { + signalLifecycleFailure(failure); + ExceptionHelper.sneakyThrow(failure); } } @@ -1019,20 +1085,8 @@ void beginDraining() { } private void closeLifecycle() { - var interrupted = false; synchronized (admissionLock) { lifecycleState = LifecycleState.CLOSED; - // Admission must stop before the final flush, including calls admitted before CLOSED but not yet enqueued. - while (checkpointAdmissionsInFlight > 0) { - try { - admissionLock.wait(); - } catch (InterruptedException e) { - interrupted = true; - } - } - } - if (interrupted) { - Thread.currentThread().interrupt(); } } @@ -1087,6 +1141,7 @@ private void requireTaskAdmission(BaseDurableOperation operation) { } private void requireOpen(String workType) { + checkInvocationAdmission(); if (lifecycleState != LifecycleState.OPEN) { throw rejected(workType); } @@ -1097,38 +1152,6 @@ private RejectedExecutionException rejected(String workType) { "Invocation " + invocationId + " is " + lifecycleState + "; cannot admit new " + workType); } - private void drainExecutorTasks() { - while (true) { - var tasks = operationTasksToDrain(); - if (tasks.isEmpty()) return; - for (var task : tasks) awaitTaskExit(task); - } - } - - private List> operationTasksToDrain() { - synchronized (admissionLock) { - var tasks = activeExecutorTasks.values().stream() - .filter(task -> task.role() != ExecutorTaskHandle.Role.ROOT) - .toList(); - if (tasks.isEmpty()) { - // Close under the submission lock so no task can enter after the final empty snapshot. - lifecycleState = LifecycleState.CLOSED; - } - return tasks; - } - } - - private void awaitTaskExit(ExecutorTaskHandle task) { - try { - task.exit().get(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new IllegalStateException("Interrupted while draining operation task: " + task.operationId(), e); - } catch (ExecutionException e) { - throw new IllegalStateException("Could not drain operation task: " + task.operationId(), e); - } - } - private void validateExecutorPool() { // double check if the thread pool is empty if (durableConfig.getExecutorService() instanceof ThreadPoolExecutor threadPoolExecutor) { @@ -1211,4 +1234,285 @@ public CompletableFuture runUntilCompleteOrSuspend(CompletableFuture u return null; }); } + /** Called only by the invocation caller, so root exit and owned-executor cleanup never self-join. */ + Throwable finishInvocation(Throwable originalFailure) { + if (cleanupFinished) return originalFailure; + beginDraining(); + var hardDeadline = cleanupDeadlineNanos(); + Throwable failure = null; + if (Thread.interrupted()) cleanupInterrupted.set(true); + try { + if (cleanupInterrupted.get() && !getActiveExecutorTasks().isEmpty()) + throw new InterruptedException("Invocation caller was interrupted"); + stopCheckpointContinuationsUntil(hardDeadline); + failure = drainInvocationTasks(originalFailure, hardDeadline); + } catch (InterruptedException interrupted) { + cleanupInterrupted.set(true); + failure = lifecycleFailure("Interrupted while draining invocation tasks", interrupted); + signalLifecycleFailure(failure); + failure = settleInterruptedCleanup(hardDeadline, true, failure); + } catch (RuntimeException problem) { + failure = problem; + signalLifecycleFailure(failure); + } finally { + failure = stopCheckpointAdmission(hardDeadline, failure); + failure = shutdownCheckpoints(hardDeadline, failure); + closeLifecycle(); + cleanupFinished = true; + restoreCleanupInterrupt(); + } + var recordedFailure = lifecycleFailure.get(); + var selectedFailure = isFatal(originalFailure) || recordedFailure == null ? originalFailure : recordedFailure; + if (failure == null) return selectedFailure; + return isFatal(selectedFailure) + ? combineFailures(selectedFailure, failure) + : combineFailures(failure, selectedFailure); + } + + private Throwable drainInvocationTasks(Throwable originalFailure, long hardDeadline) throws InterruptedException { + var taskDeadline = Math.max(System.nanoTime(), hardDeadline - CHECKPOINT_CLEANUP_RESERVE.toNanos()); + var recordedFailure = lifecycleFailure.get(); + var abortFailure = isFatal(originalFailure) + ? originalFailure + : recordedFailure != null + ? recordedFailure + : isLifecycleFailure(originalFailure) ? originalFailure : null; + var retrying = abortFailure != null; + if (retrying) signalLifecycleFailure(abortFailure); + else if (!awaitTasksUntil(minDeadline(taskDeadline, MAX_GRACEFUL_DRAIN_DURATION), false)) { + var failure = lifecycleFailure("Invocation user work exceeded its graceful cleanup budget", null); + signalLifecycleFailure(failure); + awaitTasksUntil(minDeadline(taskDeadline, MAX_CANCELLATION_DRAIN_DURATION), true); + return failure; + } + // The root body and operation work have exited. Let response preparation and End finish on their owner, + // within the remaining overall cleanup budget rather than charging an exporter to the user-work grace cap. + if (awaitTasksUntil(retrying ? minDeadline(taskDeadline, MAX_CANCELLATION_DRAIN_DURATION) : taskDeadline, true)) + return null; + var failure = lifecycleFailure("Invocation lifecycle did not exit before the cleanup deadline", null); + signalLifecycleFailure(failure); + return failure; + } + + private boolean awaitTasksUntil(long deadline, boolean includeRootLifecycle) throws InterruptedException { + while (true) { + List> tasks; + synchronized (admissionLock) { + tasks = getActiveExecutorTasks().stream() + .filter(task -> includeRootLifecycle + || task.role() != ExecutorTaskHandle.Role.ROOT + || !rootHandlerExited) + .toList(); + if (tasks.isEmpty()) { + // The final empty snapshot and admission closure are one transition. The internal, awaited + // EXECUTION-result path remains available; ordinary operations cannot enter after this point. + lifecycleState = LifecycleState.CLOSED; + return true; + } + } + var remaining = remainingNanos(deadline); + if (remaining == 0) return false; + try { + var exited = CompletableFuture.allOf( + tasks.stream().map(ExecutorTaskHandle::exit).toArray(CompletableFuture[]::new)); + var observed = !includeRootLifecycle + && tasks.stream().anyMatch(task -> task.role() == ExecutorTaskHandle.Role.ROOT) + ? CompletableFuture.anyOf(exited, rootHandlerExitSignal) + : exited; + observed.get(remaining, TimeUnit.NANOSECONDS); + } catch (ExecutionException impossible) { + throw new IllegalStateException("Executor task exit signal failed", impossible.getCause()); + } catch (TimeoutException timeout) { + return false; + } + } + } + + private void stopCheckpointContinuationsUntil(long deadline) throws InterruptedException { + List owners; + synchronized (activeThreads) { + closing = true; + owners = new ArrayList<>(checkpointContinuations.values()); + } + owners.forEach(ExecutionManager::stopContinuationOwner); + synchronized (activeThreads) { + while (!checkpointContinuations.isEmpty()) { + var remaining = remainingNanos(deadline); + if (remaining == 0) + throw lifecycleFailure( + "Checkpoint continuations exceeded the cleanup deadline", new TimeoutException()); + TimeUnit.NANOSECONDS.timedWait(activeThreads, remaining); + } + } + } + + private void waitForRunningUserHandlersUntil(long deadline) throws InterruptedException { + for (var operation : registeredOperations.values()) { + var task = operation.getRunningUserHandler(); + if (task == null || task.isDone()) continue; + try { + task.get(remainingNanos(deadline), TimeUnit.NANOSECONDS); + } catch (CancellationException cancelled) { + if (!isExecutionCompletedExceptionally()) + throw lifecycleFailure("Operation observation cancelled while draining", cancelled); + // Logical cancellation is not exit; the tracked task exit signal is still awaited below. + } catch (ExecutionException ignored) { + // The operation's durable outcome is selected separately; cleanup observes physical task exit below. + } catch (TimeoutException timeout) { + throw lifecycleFailure("Operation handler exceeded the cleanup deadline", timeout); + } + } + } + + private Throwable stopCheckpointAdmission(long deadline, Throwable priorFailure) { + synchronized (admissionLock) { + lifecycleState = LifecycleState.CLOSED; + while (checkpointAdmissionsInFlight > 0) { + var remaining = remainingNanos(deadline); + if (remaining == 0) + return combineFailures( + priorFailure, + lifecycleFailure("Checkpoint admission exceeded cleanup budget", new TimeoutException())); + try { + TimeUnit.NANOSECONDS.timedWait(admissionLock, remaining); + } catch (InterruptedException interrupted) { + cleanupInterrupted.set(true); + } + } + } + return priorFailure; + } + + private Throwable shutdownCheckpoints(long deadline, Throwable priorFailure) { + try { + checkpointManager.shutdown(remainingDuration(deadline)); + } catch (InterruptedException interrupted) { + cleanupInterrupted.set(true); + return combineFailures( + priorFailure, lifecycleFailure("Interrupted while shutting down checkpoints", interrupted)); + } catch (TimeoutException failure) { + return combineFailures( + priorFailure, lifecycleFailure("Checkpoint shutdown exceeded cleanup budget", failure)); + } catch (ExecutionException failure) { + return combineFailures(priorFailure, ExceptionHelper.unwrapCompletableFuture(failure.getCause())); + } catch (RuntimeException failure) { + return combineFailures(priorFailure, lifecycleFailure("Checkpoint shutdown failed", failure)); + } + return priorFailure; + } + + private Throwable settleInterruptedCleanup(long deadline, boolean includeRoot, Throwable primary) { + while (remainingNanos(deadline) > 0) { + try { + stopCheckpointContinuationsUntil(deadline); + if (!awaitTasksUntil(deadline, includeRoot)) + return combineFailures( + primary, + lifecycleFailure( + "Interrupted cleanup exhausted its remaining budget", new TimeoutException())); + return primary; + } catch (InterruptedException interrupted) { + cleanupInterrupted.set(true); + // Keep the first interruption/failure, but continue best-effort joins with the remaining budget. + } catch (RuntimeException failure) { + return combineFailures(primary, failure); + } + } + return primary; + } + + private void restoreCleanupInterrupt() { + if (cleanupInterrupted.get()) Thread.currentThread().interrupt(); + cleanupInterrupted.remove(); + } + + private void signalLifecycleFailure(Throwable failure) { + if (failure instanceof UnrecoverableDurableExecutionException problem + && LIFECYCLE_ERROR_TYPE.equals(problem.getErrorObject().errorType())) + lifecycleFailure.compareAndSet(null, problem); + // Interruption must precede arbitrary inline completion observers. + var current = currentExecutorTask.get(); + getActiveExecutorTasks().stream().filter(task -> task != current).forEach(task -> task.cancel(true)); + executionExceptionFuture.completeExceptionally(failure); + stopAllOperations(failure); + } + + InvocationOutcome awaitInvocationOutcome(CompletableFuture future) { + try { + var value = deadlineNanos == null + ? future.get() + : future.get( + remainingNanos(deadlineNanos - EXECUTION_DEADLINE_HEADROOM.toNanos()), + TimeUnit.NANOSECONDS); + return new InvocationOutcome<>(value, null); + } catch (InterruptedException interrupted) { + cleanupInterrupted.set(true); + var failure = lifecycleFailure("Invocation runtime thread was interrupted", interrupted); + lifecycleFailure.compareAndSet(null, failure); + return new InvocationOutcome<>(null, failure); + } catch (TimeoutException timeout) { + var failure = lifecycleFailure("Invocation deadline reached before execution completed", timeout); + lifecycleFailure.compareAndSet(null, failure); + return new InvocationOutcome<>(null, failure); + } catch (ExecutionException failure) { + return new InvocationOutcome<>(null, ExceptionHelper.unwrapCompletableFuture(failure.getCause())); + } + } + + private long cleanupDeadlineNanos() { + var localDeadline = addToNow(MAX_CLEANUP_DURATION); + if (deadlineNanos == null) { + return localDeadline; + } + return Math.min(localDeadline, deadlineNanos - RESPONSE_HEADROOM.toNanos()); + } + + private static UnrecoverableDurableExecutionException lifecycleFailure(String message, Throwable cause) { + return new UnrecoverableDurableExecutionException( + ErrorObject.builder() + .errorType(LIFECYCLE_ERROR_TYPE) + .errorMessage(message) + .build(), + true, + cause); + } + + private static Throwable combineFailures(Throwable primary, Throwable additional) { + if (primary == null) return additional; + if (additional == null || primary == additional) return primary; + if (isFatal(additional) && !isFatal(primary)) { + additional.addSuppressed(primary); + return additional; + } + primary.addSuppressed(additional); + return primary; + } + + @SuppressWarnings("removal") + static boolean isFatal(Throwable failure) { + return failure instanceof VirtualMachineError || failure instanceof ThreadDeath; + } + + static boolean isLifecycleFailure(Throwable failure) { + return failure instanceof UnrecoverableDurableExecutionException control + && LIFECYCLE_ERROR_TYPE.equals(control.getErrorObject().errorType()); + } + + private static long minDeadline(long deadline, Duration maximumWait) { + return Math.min(deadline, addToNow(maximumWait)); + } + + private static long addToNow(Duration duration) { + var now = System.nanoTime(); + var nanos = duration.toNanos(); + return now > Long.MAX_VALUE - nanos ? Long.MAX_VALUE : now + nanos; + } + + private static long remainingNanos(long deadline) { + return Math.max(0, deadline - System.nanoTime()); + } + + private static Duration remainingDuration(long deadline) { + return Duration.ofNanos(remainingNanos(deadline)); + } } 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 4db69eb34..8af2e739e 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 @@ -356,6 +356,7 @@ protected void runUserHandler(Runnable runnable, ThreadType threadType) { * @return the user function's result */ protected T runUserFunction(Integer attempt, Supplier userFunction) { + executionManager.checkUserFunctionExecutionAllowed(); var pluginRunner = getPluginRunner(); var startInfo = PluginInfoConverter.toUserFunctionStartInfo( operationIdentifier, @@ -364,6 +365,7 @@ protected T runUserFunction(Integer attempt, Supplier userFunction) { attempt); pluginRunner.onUserFunctionStart(startInfo); try { + executionManager.checkUserFunctionExecutionAllowed(); T result = userFunction.get(); pluginRunner.onUserFunctionEnd( PluginInfoConverter.toUserFunctionEndInfo(startInfo, UserFunctionOutcome.SUCCEEDED, null)); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/exception/UnrecoverableDurableExecutionExceptionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/exception/UnrecoverableDurableExecutionExceptionTest.java index 67c2f0fa1..b3232e38a 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/exception/UnrecoverableDurableExecutionExceptionTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/exception/UnrecoverableDurableExecutionExceptionTest.java @@ -5,6 +5,7 @@ import static org.junit.jupiter.api.Assertions.*; import org.junit.jupiter.api.Test; +import software.amazon.awssdk.services.lambda.model.ErrorObject; class UnrecoverableDurableExecutionExceptionTest { @@ -36,4 +37,15 @@ void testIllegalDurableOperationException() { assertInstanceOf(RuntimeException.class, exception); assertInstanceOf(DurableExecutionException.class, exception); } + + @Test + void preservesCauseForRetryableInvocationFailures() { + var cause = new InterruptedException("interrupted"); + var error = ErrorObject.builder().errorMessage("cleanup failed").build(); + + var exception = new UnrecoverableDurableExecutionException(error, true, cause); + + assertSame(cause, exception.getCause()); + assertTrue(exception.isRetryable()); + } } 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 75397781c..76e6029fd 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 @@ -319,4 +319,29 @@ private void await(CountDownLatch latch) { throw new AssertionError(e); } } + + @Test + void shutdownHonorsTimeoutWhileBatchActionIsBlocked() throws Exception { + var entered = new CountDownLatch(1); + var release = new CountDownLatch(1); + var batcher = new ApiRequestDelayedBatcher( + MAX_BATCH_SIZE, MAX_BATCH_BINARY_SIZE_IN_BYTES, item -> 0, items -> { + entered.countDown(); + try { + release.await(); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + } + }); + var future = batcher.submit(input, Duration.ZERO); + assertTrue(entered.await(5, TimeUnit.SECONDS)); + + try { + assertThrows(TimeoutException.class, () -> batcher.shutdown(Duration.ofMillis(50))); + assertFalse(future.isDone()); + } finally { + release.countDown(); + } + future.get(5, TimeUnit.SECONDS); + } } 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 916460e13..45303837b 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 @@ -38,6 +38,8 @@ class CheckpointManagerTest { + private static final Duration SHUTDOWN_TIMEOUT = Duration.ofSeconds(5); + private DurableConfig config; private DurableExecutionClient client; private CheckpointManager batcher; @@ -192,11 +194,11 @@ void pollForUpdate_handlesMultiplePollers() throws Exception { } @Test - void shutdown_completesAllPendingPollersWithException() { + void shutdown_completesAllPendingPollersWithException() throws Exception { var future1 = batcher.pollForUpdate("op-1"); var future2 = batcher.pollForUpdate("op-2"); - batcher.shutdown(); + batcher.shutdown(SHUTDOWN_TIMEOUT); assertTrue(future1.isCompletedExceptionally()); assertTrue(future2.isCompletedExceptionally()); @@ -218,7 +220,7 @@ void shutdown_waitsForPendingCheckpoints() throws Exception { .type(OperationType.STEP) .build()); - batcher.shutdown(); + batcher.shutdown(SHUTDOWN_TIMEOUT); assertTrue(future.isDone()); verify(client, atLeastOnce()).checkpoint(anyString(), anyString(), anyList()); 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 index 851e780c0..3e7b63956 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutorLifecycleRegressionTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutorLifecycleRegressionTest.java @@ -132,13 +132,15 @@ public void onInvocationEnd(InvocationEndInfo info) { // The interrupt flag is sticky, so the drain observes it whether it lands before or during the wait. rootThread.get().interrupt(); + assertThrows(TimeoutException.class, () -> invocation.get(100, TimeUnit.MILLISECONDS)); + assertEquals(1, ended.getCount(), "Interruption must not move End ahead of still-running work"); + releaseStep.countDown(); 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(); + // The admitted step has exited and checkpointed before the interrupted owner reports its failure. var failure = assertThrows(ExecutionException.class, () -> invocation.get(5, TimeUnit.SECONDS)); assertInstanceOf(IllegalStateException.class, failure.getCause()); } finally { diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerAdmissionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerAdmissionTest.java index 76a584591..810f62c0d 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerAdmissionTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerAdmissionTest.java @@ -14,6 +14,7 @@ import static org.mockito.Mockito.spy; import static org.mockito.Mockito.when; +import com.amazonaws.services.lambda.runtime.Context; import java.time.Duration; import java.time.Instant; import java.util.List; @@ -23,6 +24,7 @@ import java.util.concurrent.Executors; import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -38,6 +40,7 @@ import software.amazon.awssdk.services.lambda.model.OperationUpdate; import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.client.DurableExecutionClient; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.model.DurableExecutionInput; import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.operation.BaseDurableOperation; @@ -57,11 +60,28 @@ void setUp() { .withCheckpointDelay(Duration.ofHours(1)) .withPollingStrategy(attempt -> Duration.ofHours(1)) .build()); - manager = createManager(); + manager = createManager(null); manager.registerActiveThread(null); when(client.checkpoint(any(), any(), any())).thenReturn(response()); } + private ExecutionManager createManager(Context context) { + var execution = Operation.builder() + .id("execution") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .build(); + return new ExecutionManager( + new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/test/execution", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(execution) + .build()), + config, + context); + } + @AfterEach void tearDown() throws Exception { manager.close(); @@ -145,6 +165,38 @@ void shutdownWaitsForAdmittedPollAndPreservesInterruption(boolean interrupted) t } } + @Test + void stuckCheckpointAdmissionCannotExceedCleanupDeadline() throws Exception { + manager.close(); + var context = mock(Context.class); + when(context.getRemainingTimeInMillis()).thenReturn(500); + manager = createManager(context); + var admissionEntered = new CountDownLatch(1); + var releaseAdmission = new CountDownLatch(1); + doAnswer(invocation -> { + admissionEntered.countDown(); + await(releaseAdmission); + return invocation.callRealMethod(); + }) + .when(config) + .getPollingStrategy(); + try { + var admission = callers.submit(() -> manager.pollForOperationUpdates("pending")); + await(admissionEntered); + var shutdown = callers.submit(() -> manager.finishInvocation(null)); + var failure = + assertInstanceOf(UnrecoverableDurableExecutionException.class, shutdown.get(5, TimeUnit.SECONDS)); + assertTrue(failure.isRetryable()); + assertInstanceOf(TimeoutException.class, failure.getCause()); + assertEquals(ExecutionManager.LifecycleState.CLOSED, manager.getLifecycleState()); + assertThrows(RejectedExecutionException.class, () -> manager.sendOperationUpdate(update())); + releaseAdmission.countDown(); + admission.get(5, TimeUnit.SECONDS).cancel(false); + } finally { + releaseAdmission.countDown(); + } + } + @Test void failedAdmissionDoesNotPreventShutdown() throws Exception { doThrow(new IllegalArgumentException("polling configuration failed")) @@ -283,20 +335,7 @@ void alreadyRunningRootCannotAdmitOperationsAfterDrainingStarts() throws Excepti } private ExecutionManager createManager() { - var execution = Operation.builder() - .id("execution") - .type(OperationType.EXECUTION) - .status(OperationStatus.STARTED) - .build(); - return new ExecutionManager( - new DurableExecutionInput( - "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/test/execution", - "token", - CheckpointUpdatedExecutionState.builder() - .operations(execution) - .build()), - config, - null); + return createManager(null); } private static BaseDurableOperation contextOperation(String id, OperationSubType subType) { diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerLifecycleTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerLifecycleTest.java index edd1cd741..6f029486e 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerLifecycleTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerLifecycleTest.java @@ -4,6 +4,7 @@ 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.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertSame; @@ -11,9 +12,11 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import static software.amazon.lambda.durable.model.ExecutionStatus.PENDING; import com.amazonaws.services.lambda.runtime.Context; import java.time.Duration; +import java.time.Instant; import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; @@ -25,11 +28,13 @@ import java.util.concurrent.SynchronousQueue; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; 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.OperationAction; import software.amazon.awssdk.services.lambda.model.OperationStatus; @@ -37,8 +42,12 @@ 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.TypeToken; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.model.DurableExecutionInput; import software.amazon.lambda.durable.operation.BaseDurableOperation; +import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; +import software.amazon.lambda.durable.plugin.InvocationEndInfo; class ExecutionManagerLifecycleTest { private static final String EXECUTION_ID = "execution"; @@ -360,25 +369,174 @@ void drainingRejectsNewWorkButAllowsExistingCheckpointCleanup() throws Exception } } + @Test + void interruptedDrainStillShutsDownCheckpointsAndRestoresInterrupt() throws Exception { + var executor = Executors.newSingleThreadExecutor(); + var taskEntered = new CountDownLatch(1); + var taskInterrupted = new CountDownLatch(1); + var manager = createManager(null, executor); + manager.submitRootTask(() -> { + taskEntered.countDown(); + try { + new CountDownLatch(1).await(); + } catch (InterruptedException expected) { + taskInterrupted.countDown(); + } + return "cancelled"; + }); + assertTrue(taskEntered.await(5, TimeUnit.SECONDS)); + var checkpoint = manager.sendOperationUpdate(OperationUpdate.builder() + .id("step") + .name("step") + .type(OperationType.STEP) + .subType("Step") + .action(OperationAction.START) + .build()); + + Thread.currentThread().interrupt(); + try { + var failure = manager.finishInvocation(null); + + var lifecycleFailure = assertInstanceOf(UnrecoverableDurableExecutionException.class, failure); + assertInstanceOf(InterruptedException.class, lifecycleFailure.getCause()); + assertEquals(ExecutionManager.LifecycleState.CLOSED, manager.getLifecycleState()); + assertTrue(checkpoint.isDone(), "checkpoint shutdown should still be attempted"); + assertTrue(Thread.currentThread().isInterrupted()); + Thread.interrupted(); + assertTrue(taskInterrupted.await(5, TimeUnit.SECONDS)); + } finally { + Thread.interrupted(); + stop(executor); + } + } + + @Test + void invocationDeadlineCancelsAStuckRootAndReturnsRetryableFailure() throws Exception { + var context = mock(Context.class); + when(context.getAwsRequestId()).thenReturn("deadline-request"); + when(context.getRemainingTimeInMillis()).thenReturn(700); + var executor = Executors.newSingleThreadExecutor(); + var entered = new CountDownLatch(1); + var interrupted = new CountDownLatch(1); + var manager = createManager(context, executor); + var root = manager.submitRootTask(() -> { + entered.countDown(); + try { + new CountDownLatch(1).await(); + } catch (InterruptedException expected) { + interrupted.countDown(); + } + return "cancelled"; + }); + assertTrue(entered.await(5, TimeUnit.SECONDS)); + + try { + var outcome = manager.awaitInvocationOutcome(root); + var deadlineFailure = assertInstanceOf(UnrecoverableDurableExecutionException.class, outcome.failure()); + assertTrue(deadlineFailure.isRetryable()); + + assertEquals(deadlineFailure, manager.finishInvocation(deadlineFailure)); + assertTrue(interrupted.await(5, TimeUnit.SECONDS)); + assertEquals(ExecutionManager.LifecycleState.CLOSED, manager.getLifecycleState()); + } finally { + stop(executor); + } + } + + @Test + void responseCriticalCheckpointWaitIsBoundedByInvocationDeadline() { + var context = mock(Context.class); + when(context.getAwsRequestId()).thenReturn("checkpoint-deadline-request"); + when(context.getRemainingTimeInMillis()).thenReturn(250); + var manager = createManager(context, null); + + var failure = assertThrows( + UnrecoverableDurableExecutionException.class, + () -> manager.awaitCheckpointCompletion(new CompletableFuture<>())); + + assertTrue(failure.isRetryable()); + assertInstanceOf(TimeoutException.class, failure.getCause()); + assertEquals(failure, manager.finishInvocation(failure)); + } + + @Test + void pendingResponseAndInvocationEndWaitForRootFinally() throws Exception { + var users = Executors.newCachedThreadPool(); + var runtime = Executors.newSingleThreadExecutor(); + var cleanupEntered = new CountDownLatch(1); + var releaseCleanup = new CountDownLatch(1); + var invocationEnded = new CountDownLatch(1); + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withExecutorService(users) + .withPlugins(new DurableExecutionPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + invocationEnded.countDown(); + } + }) + .build(); + + try { + var response = runtime.submit(() -> DurableExecutor.execute( + input(), + context(10_000), + TypeToken.get(String.class), + (value, durableContext) -> { + try { + durableContext.wait("wait", Duration.ofSeconds(5)); + return value; + } finally { + cleanupEntered.countDown(); + await(releaseCleanup); + } + }, + config)); + assertTrue(cleanupEntered.await(5, TimeUnit.SECONDS)); + + assertThrows(TimeoutException.class, () -> response.get(100, TimeUnit.MILLISECONDS)); + assertEquals(1, invocationEnded.getCount()); + + releaseCleanup.countDown(); + assertEquals(PENDING, response.get(5, TimeUnit.SECONDS).status()); + assertTrue(invocationEnded.await(5, TimeUnit.SECONDS)); + } finally { + releaseCleanup.countDown(); + stop(users); + stop(runtime); + } + } + private ExecutionManager createManager(Context context, ExecutorService executor) { var config = DurableConfig.builder().withDurableExecutionClient(TestUtils.createMockClient()); if (executor != null) { config.withExecutorService(executor); } + return new ExecutionManager(input(), config.build(), context); + } + + private DurableExecutionInput input() { var execution = Operation.builder() .id(EXECUTION_ID) .type(OperationType.EXECUTION) .status(OperationStatus.STARTED) + .startTimestamp(Instant.now()) + .executionDetails( + ExecutionDetails.builder().inputPayload("\"input\"").build()) .build(); - return new ExecutionManager( - new DurableExecutionInput( - EXECUTION_ARN, - "token", - CheckpointUpdatedExecutionState.builder() - .operations(List.of(execution)) - .build()), - config.build(), - context); + return new DurableExecutionInput( + EXECUTION_ARN, + "token", + CheckpointUpdatedExecutionState.builder() + .operations(List.of(execution)) + .build()); + } + + private Context context(int remainingMillis) { + var context = mock(Context.class); + when(context.getAwsRequestId()).thenReturn("request-id"); + when(context.getRemainingTimeInMillis()).thenReturn(remainingMillis); + return context; } private static void await(CountDownLatch latch) { diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationCleanupBoundaryTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationCleanupBoundaryTest.java new file mode 100644 index 000000000..d2cc4bca1 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationCleanupBoundaryTest.java @@ -0,0 +1,496 @@ +// 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.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.*; + +import com.amazonaws.services.lambda.runtime.Context; +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 org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import software.amazon.awssdk.services.lambda.model.*; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.client.DurableExecutionClient; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.operation.BaseDurableOperation; +import software.amazon.lambda.durable.plugin.*; + +@Timeout(15) +class InvocationCleanupBoundaryTest { + private static final String ARN = + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/cleanup/execution"; + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + @SuppressWarnings("removal") + void cleanupBudgetCannotReplaceSelectedJvmFatal(boolean threadDeath) throws Exception { + var users = Executors.newSingleThreadExecutor(); + var manager = manager(users, context(2_000)); + var entered = new CountDownLatch(1); + var release = new CountDownLatch(1); + Error fatal = threadDeath ? new ThreadDeath() : new InternalError("selected fatal"); + try { + manager.submitRootTask(() -> { + entered.countDown(); + awaitIgnoringInterrupt(release); + return "ignored"; + }); + await(entered); + var task = manager.getActiveExecutorTasks().get(0); + assertSame(fatal, manager.finishInvocation(fatal)); + assertFalse(task.exit().isDone(), "The fixture deliberately retains an uncooperative owner"); + assertTrue(Arrays.stream(fatal.getSuppressed()).anyMatch(ExecutionManager::isLifecycleFailure)); + assertEquals(ExecutionManager.LifecycleState.CLOSED, manager.getLifecycleState()); + } finally { + release.countDown(); + stop(users); + manager.close(); + } + } + + @Test + void interruptedCallerStillObservesCooperativeExitAndSettlesCheckpoints() throws Exception { + var users = Executors.newSingleThreadExecutor(); + var callers = Executors.newSingleThreadExecutor(); + var manager = manager(users, null); + var entered = new CountDownLatch(1); + var interrupted = new CountDownLatch(1); + var releaseExit = new CountDownLatch(1); + try { + manager.submitRootTask(() -> { + entered.countDown(); + try { + new CountDownLatch(1).await(); + } catch (InterruptedException expected) { + interrupted.countDown(); + } finally { + awaitIgnoringInterrupt(releaseExit); + } + return "cancelled"; + }); + await(entered); + var task = manager.getActiveExecutorTasks().get(0); + var checkpoint = manager.sendOperationUpdate(OperationUpdate.builder() + .id("queued") + .type(OperationType.STEP) + .action(OperationAction.START) + .build()); + var finished = callers.submit(() -> { + Thread.currentThread().interrupt(); + var failure = manager.finishInvocation(null); + assertTrue(Thread.interrupted(), "Cleanup restores the invocation caller's interrupt"); + return failure; + }); + await(interrupted); + assertFalse(finished.isDone(), "Interruption cannot skip the cooperative owner's actual exit"); + assertFalse(task.exit().isDone()); + releaseExit.countDown(); + var failure = + assertInstanceOf(UnrecoverableDurableExecutionException.class, finished.get(5, TimeUnit.SECONDS)); + assertInstanceOf(InterruptedException.class, failure.getCause()); + assertTrue(task.exit().isDone()); + assertTrue(checkpoint.isDone(), "Checkpoint shutdown still settles queued waiters"); + } finally { + releaseExit.countDown(); + stop(users); + stop(callers); + manager.close(); + } + } + + @Test + void recordedDeadlineFailureSurvivesCleanupWithoutAnOriginalFailure() { + var users = Executors.newSingleThreadExecutor(); + var manager = manager(users, context(700)); + try { + var recorded = manager.awaitInvocationOutcome(new CompletableFuture()) + .failure(); + assertInstanceOf(UnrecoverableDurableExecutionException.class, recorded); + assertSame(recorded, manager.finishInvocation(null)); + } finally { + users.shutdownNow(); + manager.close(); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void terminalAcknowledgmentStopsAdditionalMetadataPagesAtCutoff(boolean admitFirstPage) throws Exception { + for (boolean terminal : List.of(false, true)) { + var client = mock(DurableExecutionClient.class); + var expired = new AtomicBoolean(); + var published = new AtomicReference>(); + var inline = Operation.builder() + .id("inline") + .type(OperationType.STEP) + .status(OperationStatus.SUCCEEDED) + .build(); + var paged = Operation.builder() + .id("paged") + .type(OperationType.STEP) + .status(OperationStatus.SUCCEEDED) + .build(); + var cutoff = cutoff(); + when(client.checkpoint(any(), any(), any())).thenAnswer(call -> { + expired.set(!admitFirstPage); + return CheckpointDurableExecutionResponse.builder() + .checkpointToken("accepted-token") + .newExecutionState(CheckpointUpdatedExecutionState.builder() + .operations(inline) + .nextMarker("first") + .build()) + .build(); + }); + when(client.getExecutionState(any(), any(), any())).thenAnswer(call -> { + expired.set(true); + return GetDurableExecutionStateResponse.builder() + .operations(paged) + .nextMarker("must-not-fetch") + .build(); + }); + var config = DurableConfig.builder() + .withDurableExecutionClient(client) + .withCheckpointDelay(Duration.ZERO) + .build(); + var checkpoints = + new CheckpointManager(config, ARN, "token", published::set, () -> true, () -> {}, () -> {}, () -> { + if (expired.get()) throw cutoff; + }); + try { + var future = checkpoints.checkpoint(OperationUpdate.builder() + .id("update") + .type(terminal ? OperationType.EXECUTION : OperationType.STEP) + .action(OperationAction.SUCCEED) + .build()); + if (terminal) { + future.get(5, TimeUnit.SECONDS); + assertEquals(admitFirstPage ? List.of(inline, paged) : List.of(inline), published.get()); + assertTrue(checkpoints.isTerminalExecutionAcknowledged()); + } else { + assertSame( + cutoff, + assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)) + .getCause()); + assertNull(published.get()); + } + verify(client, times(admitFirstPage ? 1 : 0)).getExecutionState(any(), any(), any()); + } finally { + checkpoints.shutdown(Duration.ofSeconds(1)); + } + } + } + + @Test + void terminalAcknowledgmentDoesNotHideUnrelatedMetadataFailure() throws Exception { + var client = mock(DurableExecutionClient.class); + var problem = new IllegalStateException("metadata store failed"); + when(client.checkpoint(any(), any(), any())) + .thenReturn(CheckpointDurableExecutionResponse.builder() + .checkpointToken("accepted-token") + .newExecutionState(CheckpointUpdatedExecutionState.builder() + .nextMarker("first") + .build()) + .build()); + when(client.getExecutionState(any(), any(), any())).thenThrow(problem); + var checkpoints = new CheckpointManager( + DurableConfig.builder() + .withDurableExecutionClient(client) + .withCheckpointDelay(Duration.ZERO) + .build(), + ARN, + "token", + ignored -> {}); + try { + var future = checkpoints.checkpoint(OperationUpdate.builder() + .id("execution") + .type(OperationType.EXECUTION) + .action(OperationAction.SUCCEED) + .build()); + assertSame( + problem, + assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)) + .getCause()); + assertTrue(checkpoints.isTerminalExecutionAcknowledged()); + } finally { + checkpoints.shutdown(Duration.ofSeconds(1)); + } + } + + @Test + void lateOwnerEndRestoresScopeBeforeTheSameWorkerIsReused() throws Exception { + var users = Executors.newSingleThreadExecutor(); + var callers = Executors.newSingleThreadExecutor(); + var scope = new ThreadLocal(); + var prior = new ThreadLocal(); + var owner = new AtomicReference(); + var ended = new AtomicInteger(); + var endStatus = new AtomicReference(); + var entered = new CountDownLatch(1); + var release = new CountDownLatch(1); + var end = new CountDownLatch(1); + var lateBodies = new AtomicInteger(); + var plugin = new DurableExecutionPlugin() { + @Override + public void onInvocationStart(InvocationInfo info) { + owner.set(Thread.currentThread()); + prior.set(scope.get()); + scope.set("invocation"); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + assertSame(owner.get(), Thread.currentThread()); + assertEquals("invocation", scope.get()); + scope.set(prior.get()); + prior.remove(); + endStatus.set(info.invocationStatus()); + ended.incrementAndGet(); + end.countDown(); + } + }; + try { + users.submit(() -> scope.set("host")).get(5, TimeUnit.SECONDS); + var config = DurableConfig.builder() + .withExecutorService(users) + .withDurableExecutionClient(TestUtils.createMockClient()) + .withPlugins(plugin) + .build(); + var invocation = callers.submit(() -> DurableExecutor.execute( + input(), + context(2_000), + TypeToken.get(String.class), + (value, ctx) -> { + entered.countDown(); + awaitIgnoringInterrupt(release); + assertThrows( + UnrecoverableDurableExecutionException.class, + () -> ctx.step("late", String.class, step -> { + lateBodies.incrementAndGet(); + return "late"; + })); + return "late-result"; + }, + config)); + await(entered); + assertInstanceOf( + UnrecoverableDurableExecutionException.class, + assertThrows(ExecutionException.class, () -> invocation.get(5, TimeUnit.SECONDS)) + .getCause()); + assertEquals(0, ended.get(), "The caller cannot run End on a different thread to force cleanup"); + release.countDown(); + await(end); + assertEquals(InvocationStatus.RETRYING, endStatus.get()); + assertEquals(1, ended.get()); + assertEquals(0, lateBodies.get()); + assertSame(owner.get(), users.submit(Thread::currentThread).get(5, TimeUnit.SECONDS)); + assertEquals("host", users.submit(scope::get).get(5, TimeUnit.SECONDS)); + } finally { + release.countDown(); + stop(users); + stop(callers); + } + } + + @Test + void terminalAckDoesNotHideAvailableFatalDuringTimedCheckpointCleanup() throws Exception { + var users = Executors.newSingleThreadExecutor(); + var callers = Executors.newSingleThreadExecutor(); + var releaseChange = new CountDownLatch(1); + var changeEntered = new CountDownLatch(1); + var fatalDelivered = new CountDownLatch(1); + var fatal = new InternalError("fatal from End after terminal ACK"); + var client = mock(DurableExecutionClient.class); + when(client.checkpoint(any(), any(), any())) + .thenReturn(CheckpointDurableExecutionResponse.builder() + .checkpointToken("accepted") + .newExecutionState(CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("execution") + .type(OperationType.EXECUTION) + .status(OperationStatus.SUCCEEDED) + .build()) + .build()) + .build()); + var plugin = new DurableExecutionPlugin() { + @Override + public void onOperationChange(OperationChangeInfo info) { + changeEntered.countDown(); + awaitIgnoringInterrupt(releaseChange); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + fatalDelivered.countDown(); + throw fatal; + } + }; + try { + var config = DurableConfig.builder() + .withExecutorService(users) + .withDurableExecutionClient(client) + .withPlugins(plugin) + .withCheckpointDelay(Duration.ZERO) + .build(); + var invocation = callers.submit(() -> DurableExecutor.execute( + input(), + context(2_000), + TypeToken.get(String.class), + (value, ctx) -> "x".repeat(6 * 1024 * 1024 + 100), + config)); + await(changeEntered); + var thrown = assertThrows(ExecutionException.class, () -> invocation.get(5, TimeUnit.SECONDS)); + assertSame( + fatal, + thrown.getCause(), + "An accepted terminal result cannot hide a fatal already delivered by End"); + assertEquals(0, fatalDelivered.getCount()); + } finally { + releaseChange.countDown(); + stop(users); + stop(callers); + } + } + + @Test + void finalEmptyTaskSnapshotClosesAdmissionBeforeAQueuedSubmitterCanEnter() throws Exception { + var users = Executors.newSingleThreadExecutor(); + var callers = Executors.newFixedThreadPool(2); + var emptySnapshot = new CountDownLatch(1); + var releaseSnapshot = new CountDownLatch(1); + var submitted = new CountDownLatch(1); + var submitter = new AtomicReference(); + var intercept = new AtomicBoolean(true); + var config = DurableConfig.builder() + .withExecutorService(users) + .withDurableExecutionClient(TestUtils.createMockClient()) + .build(); + var manager = new ExecutionManager(input(), config, null) { + @Override + List> getActiveExecutorTasks() { + var tasks = super.getActiveExecutorTasks(); + if (tasks.isEmpty() && intercept.compareAndSet(true, false)) { + emptySnapshot.countDown(); + await(releaseSnapshot); + } + return tasks; + } + }; + var operation = mock(BaseDurableOperation.class); + when(operation.getOperationId()).thenReturn("already-registered"); + when(operation.getType()).thenReturn(OperationType.STEP); + manager.registerOperation(operation); + var bodies = new AtomicInteger(); + try { + var closing = callers.submit(manager::close); + await(emptySnapshot); + var late = callers.submit(() -> { + submitter.set(Thread.currentThread()); + submitted.countDown(); + assertThrows( + RejectedExecutionException.class, + () -> manager.submitOperationTask(operation, bodies::incrementAndGet)); + }); + await(submitted); + var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(2); + var observed = submitter.get().getState(); + while (observed != Thread.State.BLOCKED && System.nanoTime() < deadline) { + Thread.onSpinWait(); + observed = submitter.get().getState(); + } + // Keep the observed contention, rather than race a second thread-state snapshot. + assertEquals(Thread.State.BLOCKED, observed); + releaseSnapshot.countDown(); + closing.get(5, TimeUnit.SECONDS); + late.get(5, TimeUnit.SECONDS); + assertEquals(0, bodies.get()); + } finally { + releaseSnapshot.countDown(); + stop(users); + stop(callers); + manager.close(); + } + } + + private static ExecutionManager manager(ExecutorService users, Context context) { + var result = new ExecutionManager( + input(), + DurableConfig.builder() + .withExecutorService(users) + .withDurableExecutionClient(TestUtils.createMockClient()) + .withCheckpointDelay(Duration.ofHours(1)) + .build(), + context); + result.registerActiveThread(null); + return result; + } + + private static DurableExecutionInput input() { + return new DurableExecutionInput( + ARN, + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("execution") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.now()) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + } + + private static Context context(int millis) { + var context = mock(Context.class); + when(context.getAwsRequestId()).thenReturn("cleanup-request"); + when(context.getRemainingTimeInMillis()).thenReturn(millis); + return context; + } + + private static UnrecoverableDurableExecutionException cutoff() { + return new UnrecoverableDurableExecutionException( + ErrorObject.builder() + .errorType("software.amazon.lambda.durable.execution.InvocationLifecycleException") + .errorMessage("controlled deadline cutoff") + .build(), + true, + new TimeoutException()); + } + + private static void await(CountDownLatch latch) { + try { + assertTrue(latch.await(5, TimeUnit.SECONDS), "Fixture gate must be released"); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw new AssertionError(interrupted); + } + } + + private static void awaitIgnoringInterrupt(CountDownLatch latch) { + while (latch.getCount() > 0) { + try { + latch.await(); + } catch (InterruptedException ignored) { + // This bounded fixture is released in finally; it models an uncooperative customer callback. + } + } + } + + private static void stop(ExecutorService executor) throws Exception { + executor.shutdownNow(); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } +}