From 318d78b3c97372f4e688d4b26348c04355834c76 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Fri, 23 Oct 2026 14:11:57 -0700 Subject: [PATCH] fix: bound invocation lifecycle and preserve cancellation failures Stop admitting user callbacks and checkpoint requests when the invocation lifecycle expires, request cancellation of admitted work, and wait for its actual owners and checkpoint transport to finish before returning. Preserve acknowledged results, canonical lifecycle failures, JVM fatal priority, and original delivery-error identity through cleanup and End. Cover deadline, interruption, task-admission, serialization, and owner-drain races with unit and OpenTelemetry integration regressions. Consolidate PR #733's existing changes relative to its last merged main, excluding dependency history already incorporated upstream. --- docs/advanced/configuration.md | 33 ++ docs/design.md | 9 +- lmi-tests/README.md | 11 +- .../otel/HandlerMdcIntegrationTest.java | 16 +- ...vocationPreparationAdmissionProbeTest.java | 5 +- .../otel/LegacyEndThreadProbeTest.java | 16 +- ...zationCancellationOtelIntegrationTest.java | 144 +++++ .../execution/ApiRequestDelayedBatcher.java | 14 +- .../durable/execution/CheckpointManager.java | 80 ++- .../durable/execution/DurableExecutor.java | 103 +++- .../durable/execution/ExecutionManager.java | 474 +++++++++++++++- .../operation/BaseDurableOperation.java | 2 + ...overableDurableExecutionExceptionTest.java | 12 + .../AcknowledgedDeliveryFailureTest.java | 265 +++++++++ .../ApiRequestDelayedBatcherTest.java | 25 + .../CanonicalLifecycleFailureTest.java | 209 +++++++ .../execution/CheckpointManagerTest.java | 8 +- .../DrainingAdmissionCancellationTest.java | 300 ++++++++++ .../ExecutionManagerAdmissionTest.java | 90 ++- .../ExecutionManagerLifecycleTest.java | 176 +++++- .../execution/HandlerScopeHandoffTest.java | 21 +- .../execution/InFlightFatalAbortTest.java | 8 +- .../IndependentDeadlineBoundaryTest.java | 211 +++++++ .../InvocationAdmissionFailureTest.java | 168 ++++++ .../InvocationCleanupBoundaryTest.java | 515 ++++++++++++++++++ ...nvocationDeadlineReviewRegressionTest.java | 265 +++++++++ .../InvocationPreparationBoundaryTest.java | 240 ++++++++ .../OwnerDrainLifecycleFailureTest.java | 319 +++++++++++ .../execution/TaskAdmissionFenceTest.java | 126 +++++ 29 files changed, 3796 insertions(+), 69 deletions(-) create mode 100644 otel-plugin/src/test/java/software/amazon/lambda/durable/otel/SerializationCancellationOtelIntegrationTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/execution/AcknowledgedDeliveryFailureTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/execution/CanonicalLifecycleFailureTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/execution/DrainingAdmissionCancellationTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/execution/IndependentDeadlineBoundaryTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationAdmissionFailureTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationCleanupBoundaryTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationDeadlineReviewRegressionTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationPreparationBoundaryTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/execution/OwnerDrainLifecycleFailureTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/execution/TaskAdmissionFenceTest.java diff --git a/docs/advanced/configuration.md b/docs/advanced/configuration.md index ed702ab0c1..00f37c8372 100644 --- a/docs/advanced/configuration.md +++ b/docs/advanced/configuration.md @@ -95,3 +95,36 @@ 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 monotonic +cutoffs for admission and its wait policy. Deadline exhaustion selects a retryable +invocation interruption; it is not checkpointed as a business-operation failure. +Catching an admission rejection cannot turn the stopped invocation into a +successful fallback. Completed operation results and failures remain available for +replay. + +Cleanup allowances determine when the SDK requests cooperative cancellation and +stops ordinary admission. They do not make cancellation equivalent to actual task +exit. Finalization still waits for admitted user work, continuations and in-flight +checkpoint/poll transport, and runs the matching invocation End on its owner. +The runtime caller waits for that owner and +End to finish, including after a deadline or interruption. End is not moved to a +completion-observer thread or deferred until after the SDK call returns. + +Interruption is recorded, remaining required cleanup still runs, and the interrupt +flag is restored. A terminal execution result acknowledged by the +service remains authoritative for a handled deadline cutoff. An accepted result +does not suppress an otherwise-propagating delivery or metadata failure, and an +available JVM fatal keeps its original identity and priority. + +These cutoffs do not guarantee bounded caller return or worker reclamation. +Already-running application I/O, synchronous initial-state reads, custom clients, +plugin hooks and other code that ignores interruption can keep the caller waiting +for actual cleanup beyond the policy budget. An executor that runs the root inline, +including a saturated `CallerRunsPolicy`, can also block the runtime caller before +it reaches a timeout-aware wait. Use finite downstream I/O timeouts and executor +configurations that allow 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/docs/design.md b/docs/design.md index 197ee6a7c7..e19cbd2f0f 100644 --- a/docs/design.md +++ b/docs/design.md @@ -151,10 +151,17 @@ public class MyHandler extends DurableHandler { ### Thread Pool Architecture +Each invocation tracks logical task outcomes separately from actual task exit. +Deadline and cleanup cutoffs fence admission and request cooperative cancellation, +while finalization still waits for admitted work and in-flight transport to finish. +The matching invocation End runs on its owner, and the runtime caller waits for +that owner and End before returning. Code that ignores interruption can therefore +outlast the policy budget; see [invocation deadline cleanup](advanced/configuration.md#invocation-deadline-cleanup). + The SDK uses two separate thread pools with distinct responsibilities: **User Executor (`DurableConfig.executorService`):** -- Runs user-defined operations (the code passed to `ctx.step()` and `ctx.stepAsync()`) +- Runs the root handler, steps, child-context user code, and map/parallel coordination - Configurable via `DurableConfig.builder().withExecutorService()` - Default: cached daemon thread pool diff --git a/lmi-tests/README.md b/lmi-tests/README.md index a88c5b73f9..b17a845376 100644 --- a/lmi-tests/README.md +++ b/lmi-tests/README.md @@ -19,6 +19,15 @@ those original snapshots and retain the existing zero-work and +16-thread limits Newly discovered request identities invalidate the warm cohort instead of moving its original baseline or final measurement to another request. +The current deadline controls fence admission and request cooperative cancellation, +but the SDK still awaits actual owner cleanup, admitted transport and the matching +invocation End before returning. Non-cooperative work can therefore keep the +runtime caller blocked beyond its deadline. This remains a partial change for +[#726](https://github.com/aws/aws-durable-execution-sdk-java/issues/726), not proof +that every desired LMI lifecycle bound is satisfied. The acceptance limits below +remain unchanged; local tests do not replace a deployed cloud acceptance run, and +eventual cleanup cannot turn an earlier lifecycle violation into a pass. + ## Test design * Each fixture is a published durable function with LMI invocation concurrency @@ -219,7 +228,7 @@ barrier diagnostics never authorize a retry. Every attempted cohort remains in the escape ledger, so a later successful cohort cannot hide an earlier pool escape. `timeouts/*.json` distinguishes server timeout logs from an SDK deadline -cancellation that returns early with an invocation error. A client HTTP timeout +cancellation reported as an invocation error after owner cleanup. A client HTTP timeout is a collection error. Deadline victims use asynchronous service invocation so a longer durable retry does not consume the driver HTTP timeout; invocation outcomes come from request-correlated runtime logs. Invocation request failures 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 dc28b137d7..55ef846b63 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 @@ -24,6 +24,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.DurableExecutionPluginFactory; @@ -325,10 +326,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/InvocationPreparationAdmissionProbeTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationPreparationAdmissionProbeTest.java index 73a536a0f9..cc764e756f 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/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/LegacyEndThreadProbeTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/LegacyEndThreadProbeTest.java index dd74fccbbc..6e70d6503d 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/LegacyEndThreadProbeTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/LegacyEndThreadProbeTest.java @@ -19,6 +19,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; @@ -147,10 +148,19 @@ public void onInvocationEnd(InvocationEndInfo info) { 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/SerializationCancellationOtelIntegrationTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/SerializationCancellationOtelIntegrationTest.java new file mode 100644 index 0000000000..726f4ecb93 --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/SerializationCancellationOtelIntegrationTest.java @@ -0,0 +1,144 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.otel; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; +import static software.amazon.lambda.durable.otel.SpanAttributes.DURABLE_INVOCATION_STATUS; + +import com.amazonaws.services.lambda.runtime.Context; +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import java.time.Instant; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +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.TypeToken; +import software.amazon.lambda.durable.client.DurableExecutionClient; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.execution.DurableExecutor; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.serde.JacksonSerDes; + +@Timeout(20) +class SerializationCancellationOtelIntegrationTest { + @BeforeEach + @AfterEach + void clearState() { + GlobalOpenTelemetry.resetForTest(); + DeterministicIdGenerator.clearSharedStateForTest(); + DurableSamplingDecision.clearSharedStateForTest(); + OtelPluginAutoConfigurationState.resetInstalledForTest(); + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void cancellationDuringSmallResultSerializationDoesNotCompleteWorkflow(boolean executionView) throws Exception { + var exporter = InMemorySpanExporter.create(); + var provider = SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)); + var pluginConfig = OtelPluginConfig.builder() + .contextExtractor(() -> null) + .enableMdc(false) + .build(); + var factory = executionView + ? ExecutionOtelPlugin.factory(provider, pluginConfig) + : InvocationOtelPlugin.factory(provider, pluginConfig); + var serializationEntered = new CountDownLatch(1); + var serializationInterrupted = new CountDownLatch(1); + var releaseSerialization = new CountDownLatch(1); + var caller = new AtomicReference(); + var callerInterrupted = new AtomicBoolean(); + var workers = Executors.newSingleThreadExecutor(); + var callers = Executors.newSingleThreadExecutor(); + var serializer = new JacksonSerDes() { + @Override + public String serialize(Object value) { + if ("result".equals(value)) { + serializationEntered.countDown(); + var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (true) { + try { + assertTrue(releaseSerialization.await( + Math.max(0, deadline - System.nanoTime()), TimeUnit.NANOSECONDS)); + break; + } catch (InterruptedException ignored) { + serializationInterrupted.countDown(); + } + } + } + return super.serialize(value); + } + }; + var config = DurableConfig.builder() + .withExecutorService(workers) + .withDurableExecutionClient(mock(DurableExecutionClient.class)) + .withSerDes(serializer) + .withPlugins(factory) + .build(); + var runtime = mock(Context.class); + when(runtime.getRemainingTimeInMillis()).thenReturn(10_000); + when(runtime.getAwsRequestId()).thenReturn("serialization-cancellation"); + var input = new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/otel/exec", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("exec") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.now()) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + try { + var response = callers.submit(() -> { + caller.set(Thread.currentThread()); + try { + return DurableExecutor.execute( + input, runtime, TypeToken.get(String.class), (value, context) -> "result", config); + } finally { + callerInterrupted.set(Thread.interrupted()); + } + }); + assertTrue(serializationEntered.await(5, TimeUnit.SECONDS)); + caller.get().interrupt(); + assertTrue(serializationInterrupted.await(5, TimeUnit.SECONDS)); + assertThrows(TimeoutException.class, () -> response.get(100, TimeUnit.MILLISECONDS)); + releaseSerialization.countDown(); + var failure = assertInstanceOf( + UnrecoverableDurableExecutionException.class, + assertThrows(ExecutionException.class, () -> response.get(5, TimeUnit.SECONDS)) + .getCause()); + assertTrue(failure.isRetryable()); + assertTrue(callerInterrupted.get()); + var spans = exporter.getFinishedSpanItems(); + var invocations = spans.stream() + .filter(span -> span.getName().equals("Invocation")) + .toList(); + assertEquals(1, invocations.size(), "The real plugin must still export the completed invocation"); + assertAll( + () -> assertEquals( + "RETRYING", invocations.get(0).getAttributes().get(DURABLE_INVOCATION_STATUS)), + () -> assertTrue( + spans.stream().noneMatch(span -> span.getName().equals("Workflow")), + "A cancelled preparation must not emit a terminal successful Workflow span")); + } finally { + releaseSerialization.countDown(); + workers.shutdownNow(); + callers.shutdownNow(); + assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(callers.awaitTermination(5, TimeUnit.SECONDS)); + exporter.close(); + } + } +} 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 2046e82e03..6a0fbf1985 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 @@ -11,7 +11,9 @@ import java.util.concurrent.CompletionException; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; import java.util.function.Function; @@ -28,7 +30,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; @@ -134,6 +136,12 @@ void shutdown() { if (fatal != null) throw new CompletionException(fatal); } + /** 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) { @@ -145,8 +153,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 e3a0dd0b30..fb82a20c6e 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 java.util.function.Supplier; @@ -19,9 +21,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; @@ -48,7 +52,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. /** @@ -130,6 +136,50 @@ class CheckpointManager { Runnable onCheckpointProcessingFinish, Runnable onTokenRevoked, Supplier pluginFatal) { + this( + config, + durableExecutionArn, + checkpointToken, + callback, + tryStartCheckpointProcessing, + onCheckpointProcessingFinish, + onTokenRevoked, + pluginFatal, + () -> {}); + } + + CheckpointManager( + DurableConfig config, + String durableExecutionArn, + String checkpointToken, + Consumer> callback, + BooleanSupplier tryStartCheckpointProcessing, + Runnable onCheckpointProcessingFinish, + Runnable onTokenRevoked, + Runnable requestAdmission) { + this( + config, + durableExecutionArn, + checkpointToken, + callback, + tryStartCheckpointProcessing, + onCheckpointProcessingFinish, + onTokenRevoked, + () -> null, + requestAdmission); + } + + CheckpointManager( + DurableConfig config, + String durableExecutionArn, + String checkpointToken, + Consumer> callback, + BooleanSupplier tryStartCheckpointProcessing, + Runnable onCheckpointProcessingFinish, + Runnable onTokenRevoked, + Supplier pluginFatal, + Runnable requestAdmission) { + this.requestAdmission = requestAdmission; this.config = config; this.durableExecutionArn = durableExecutionArn; this.callback = callback; @@ -256,6 +306,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) { @@ -274,6 +333,11 @@ private void stopPolling(Throwable cause) { * nextMarker) */ List fetchAllPages(CheckpointUpdatedExecutionState checkpointUpdatedExecutionState) { + return fetchAllPages(checkpointUpdatedExecutionState, false); + } + + private List fetchAllPages( + CheckpointUpdatedExecutionState checkpointUpdatedExecutionState, boolean terminal) { rethrowFatalIfAborted(); List operations = new ArrayList<>(); if (checkpointUpdatedExecutionState == null) { @@ -285,6 +349,12 @@ List fetchAllPages(CheckpointUpdatedExecutionState checkpointUpdatedE var nextMarker = checkpointUpdatedExecutionState.nextMarker(); while (nextMarker != null && !nextMarker.isEmpty()) { rethrowFatalIfAborted(); + 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() @@ -312,6 +382,7 @@ private void checkpointBatch(List updates) { return entry.getValue().isEmpty(); }); } + // filter the null values from pollers var request = updates.stream().filter(Objects::nonNull).toList(); @@ -353,8 +424,15 @@ private void checkpointRequest(List request) { var startTime = System.nanoTime(); logger.debug("Calling durable checkpoint API with {} updates: {}", request.size(), request); rethrowFatalIfAborted(); + 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; rethrowFatalIfAborted(); + // 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); rethrowFatalIfAborted(); @@ -366,7 +444,7 @@ private void checkpointRequest(List request) { var state = response.newExecutionState(); if (state != null) { - processUpdatedOperations(fetchAllPages(state)); + processUpdatedOperations(fetchAllPages(state, terminalExecutionAcknowledged)); } } 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 64fe95c5a2..a20d3c9b3a 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,41 @@ public static DurableExecutionOutput execute( manager.registerActiveThread(ROOT_THREAD_ID); var invocation = new Invocation<>(input, lambdaContext, inputType, handler, config, manager); try { - return invocation.execute().join(); - } catch (CompletionException failure) { - ExceptionHelper.sneakyThrow(ExceptionHelper.unwrapCompletableFuture(failure)); + 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.unwrapInvocationFailure( + selected.result().failure()); + var finalFailure = manager.finishInvocation(originalFailure); + var completedFailure = completion.isDone() + ? completion + .handle((ignored, failure) -> ExceptionHelper.unwrapInvocationFailure(failure)) + .join() + : null; + if (ExecutionManager.isFatal(completedFailure)) ExceptionHelper.sneakyThrow(completedFailure); + if (ExecutionManager.isFatal(finalFailure)) ExceptionHelper.sneakyThrow(finalFailure); + if (ExecutionManager.isLifecycleFailure(finalFailure) && manager.isTerminalExecutionAcknowledged()) { + // ACK outranks a lifecycle cutoff, not an unrelated failure already reported by delivery. + if (completedFailure != null && !ExecutionManager.isLifecycleFailure(completedFailure)) + ExceptionHelper.sneakyThrow(completedFailure); + return DurableExecutionOutput.success(""); + } + if (finalFailure != null && (selected.failure() != null || finalFailure != originalFailure)) + ExceptionHelper.sneakyThrow(finalFailure); + // Preserve the stored failure after control precedence; newer JDKs wrap join() cancellation. + if (completedFailure != null) ExceptionHelper.sneakyThrow(completedFailure); + return completion.join(); + } catch (Throwable failure) { + var originalFailure = ExceptionHelper.unwrapInvocationFailure(failure); + // Finalize synchronous admission/dispatch failures before automatic close observes them again. + var finalFailure = manager.finishInvocation(originalFailure); + ExceptionHelper.sneakyThrow(finalFailure != null ? finalFailure : originalFailure); return null; + } finally { + manager.markInvocationReported(); } } } @@ -102,6 +133,8 @@ private static final class Invocation { private final String xRayTraceId; private I userInput; private boolean started; + private boolean ended; + private CompletableFuture> outcome; private final AtomicReference losingBodyFatal = new AtomicReference<>(); private Invocation( @@ -128,9 +161,10 @@ 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 = () -> { var bodyOutcome = Outcome.capture(this::invokeHandler); + manager.markRootHandlerExited(); bodyOutcome.complete(body); var selected = outcome.join(); var selectedCause = ExceptionHelper.unwrapInvocationFailure(selected.failure()); @@ -152,7 +186,13 @@ 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. - return failDelivery(drainFailure); + var selectedFailure = ExceptionHelper.unwrapInvocationFailure(selected.failure()); + if (ExecutionManager.isFatal(selectedFailure)) { + var diagnostic = ExecutionManager.ownerDrainDiagnostic(drainFailure); + if (selectedFailure != diagnostic) selectedFailure.addSuppressed(diagnostic); + return failDelivery(selectedFailure); + } + return failDelivery(ExecutionManager.selectOwnerDrainFailure(drainFailure)); } catch (Error fatal) { return failDelivery(fatal); } @@ -172,19 +212,39 @@ private O invokeHandler() { // Deserialize only once. Even failed input gets paired start/end hooks with a null input value. inputFailure = failure; } + checkInputAdmission(inputFailure); fireOnInvocationStart(); + checkInputAdmission(inputFailure); if (inputFailure != null) ExceptionHelper.sneakyThrow(inputFailure); var context = DurableContextImpl.createRootContext(manager, config, lambdaContext); DurableContextImpl.setCurrentContext(context); try (var ignored = DurableLogger.attachContext()) { return plugins.runHandler( - () -> handler.apply(userInput, context), + () -> { + manager.checkUserFunctionExecutionAllowed(); + return handler.apply(userInput, context); + }, () -> {}, manager::recordHandlerScopeFatal, manager::recordRootHandlerFatal); } } + private void checkInputAdmission(Throwable inputFailure) { + try { + manager.checkUserFunctionExecutionAllowed(); + } catch (Throwable admissionFailure) { + // Keep existing input Start/End behavior while admission remains open. If it closes during input + // preparation, the captured JVM fatal must win without starting any additional plugin or user work. + var inputCause = ExceptionHelper.unwrapInvocationFailure(inputFailure); + if (isFatal(inputCause)) { + manager.recordInvocationFatal((Error) inputCause); + throw manager.getInvocationFatal(); + } + ExceptionHelper.sneakyThrow(admissionFailure); + } + } + private void fireOnInvocationStart() { if (plugins.isEmpty()) return; var info = new InvocationInfo( @@ -205,6 +265,12 @@ private void fireOnInvocationStart() { private DurableExecutionOutput finishInvocation(O value, Throwable failure) { var fatal = manager.getInvocationFatal(); if (fatal != null) return finishFailure(fatal); + 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()) @@ -217,7 +283,7 @@ private DurableExecutionOutput finishInvocation(O value, Throwable failure) { } 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) { + if (ExceptionHelper.unwrapInvocationFailure(deliveryFailure) instanceof SuspendExecutionException) { return finishFailure(new SuspendExecutionException(REVOKED_CHECKPOINT_TOKEN_MESSAGE)); } return failDelivery(deliveryFailure); @@ -273,7 +339,21 @@ private DurableExecutionOutput failDelivery(Throwable deliveryFailure) { } private void fireOnInvocationEnd(InvocationStatus status, Throwable error, Object result) { - if (!started) return; + if (!started || ended) return; + ended = true; + // End stays on its actual owner, including after a cutoff requested cancellation. The defensive + // reported-state guard must never reopen work or prepare another execution result. + var lifecycleFailure = manager.getLifecycleFailure(); + var terminal = status == InvocationStatus.SUCCEEDED || status == InvocationStatus.FAILED; + if ((manager.isInvocationReported() || (terminal && lifecycleFailure != null)) + && !manager.isTerminalExecutionAcknowledged() + && !ExecutionManager.isFatal(error)) { + status = InvocationStatus.RETRYING; + error = lifecycleFailure != null + ? lifecycleFailure + : new IllegalStateException("Invocation already ended before owner cleanup"); + result = null; + } plugins.onInvocationEnd(new InvocationEndInfo( requestId, executionArn, @@ -546,14 +626,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 06502b772a..ad06c1e02e 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.AtomicBoolean; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; @@ -77,6 +79,21 @@ 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 supplies a deadline. This fallback limits the initial cancellation phase for local/null contexts; + // actual owner and transport exit are still awaited without a hard return bound. + 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, @@ -119,7 +136,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; @@ -183,7 +206,8 @@ public ExecutionManager(DurableExecutionInput input, DurableConfig config, Conte this::tryStartCheckpointProcessing, this::finishCheckpointProcessing, this::signalSuspension, - pluginFatal::get); + pluginFatal::get, + this::checkCheckpointRequestAllowed); this.operationStorage = checkpointManager.fetchAllPages(input.initialExecutionState()).stream() .collect(Collectors.toConcurrentMap(Operation::id, op -> op)); @@ -332,6 +356,14 @@ public void registerOperation(BaseDurableOperation operation, BaseDurableOperati operationCompletionLocks.put(operationId, completionLock); registeredOperations.put(operationId, operation); } + // Cancellation can publish outside admissionLock and miss an operation inserted after its broadcast. + // Recheck after publication and release this operation's waiters before rejecting its constructor. + try { + checkInvocationAdmission(); + } catch (RuntimeException | Error failure) { + operation.getCompletionFuture().completeExceptionally(failure); + throw failure; + } // 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()) { @@ -441,6 +473,9 @@ private ExecutorTaskHandle registerExecutorTask( } try { + // The lifecycle fence can race task insertion after the check under admissionLock. The catch below + // settles actual-exit tracking if cancellation's task snapshot did not contain this registration. + checkInvocationAdmission(); dispatchExecutorTask(task, action, executor, preserveWorkerFailure); return task; } catch (RuntimeException | Error failure) { @@ -1035,15 +1070,124 @@ 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 failure = observeCheckpointUntilDeadline(future); + // A cutoff may settle the logical observation early. Main's End contract still requires admitted + // transport and its publication callbacks to finish before the owner finalizes the invocation. + try { + checkpointManager.shutdown(); + } catch (Throwable transportFailure) { + failure = combineFailures(failure, ExceptionHelper.unwrapInvocationFailure(transportFailure)); + } + var completedFailure = future.isDone() + ? future.handle((ignored, problem) -> ExceptionHelper.unwrapInvocationFailure(problem)) + .join() + : null; + if (isFatal(completedFailure)) ExceptionHelper.sneakyThrow(completedFailure); + if (isFatal(failure)) ExceptionHelper.sneakyThrow(failure); + if (completedFailure != null && !isLifecycleFailure(completedFailure)) + ExceptionHelper.sneakyThrow(completedFailure); + if (checkpointManager.isTerminalExecutionAcknowledged()) return; + if (failure != null) ExceptionHelper.sneakyThrow(failure); + future.join(); + } + + private Throwable observeCheckpointUntilDeadline(CompletableFuture future) { + var deadline = + deadlineNanos == null ? addToNow(MAX_CLEANUP_DURATION) : deadlineNanos - RESPONSE_HEADROOM.toNanos(); + try { + future.get(remainingNanos(deadline), TimeUnit.NANOSECONDS); + return null; + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + return lifecycleFailure("Interrupted while waiting for checkpoint completion", interrupted); + } catch (TimeoutException timeout) { + return lifecycleFailure("Checkpoint completion exceeded the invocation deadline", timeout); + } catch (CancellationException cancelled) { + // Newer JDKs wrap get() cancellation. Preserve the future's actual failure and still drain transport. + return future.handle((ignored, failure) -> ExceptionHelper.unwrapInvocationFailure(failure)) + .join(); + } catch (ExecutionException failure) { + return ExceptionHelper.unwrapInvocationFailure(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() { + rethrowPluginFatalIfPresent(); + if (executionExceptionFuture.isCompletedExceptionally()) { + try { + executionExceptionFuture.join(); + } catch (CompletionException stopped) { + ExceptionHelper.sneakyThrow(ExceptionHelper.unwrapInvocationFailure(stopped)); + } + } + checkInvocationAdmission(); + } + + private void checkInvocationAdmission() { + rethrowPluginFatalIfPresent(); + 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) { + var failure = lifecycleFailure("Invocation deadline reached before user work", new TimeoutException()); + // An existing suspension/termination already owns the outcome; still reject this later admission. + if (executionExceptionFuture.isCompletedExceptionally()) throw failure; + // This check may hold admissionLock. Latch without running cancellation/completion observers here, + // so catching the rejection cannot turn the invocation into a successful fallback. + lifecycleFailure.compareAndSet(null, failure); + throw lifecycleFailure.get(); + } + } + + private void checkCheckpointRequestAllowed() { + rethrowPluginFatalIfPresent(); + 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) { + var failure = lifecycleFailure("Invocation deadline reached before checkpoint I/O", new TimeoutException()); + if (executionExceptionFuture.isCompletedExceptionally()) throw failure; + lifecycleFailure.compareAndSet(null, failure); + throw lifecycleFailure.get(); + } + } + // ===== Utilities ===== /** Shutdown the checkpoint batcher. */ @Override public void close() { try { - drainOperations(); + var failure = finishInvocation(null); + if (failure != null) ExceptionHelper.sneakyThrow(failure); } finally { pluginRunner.releasePlugins(); } @@ -1105,7 +1249,7 @@ private void awaitUserHandler(BaseDurableOperation operation, DrainFailure drain drain.interrupted( "Interrupted while waiting for user handler " + operation.getOperationId(), interrupted); } catch (CancellationException cancelled) { - drain.record(new IllegalStateException( + drain.record(ownerDrainFailure( "User handler observation cancelled: " + operation.getOperationId(), cancelled)); return; } catch (Exception completedExceptionally) { @@ -1114,11 +1258,52 @@ private void awaitUserHandler(BaseDurableOperation operation, DrainFailure drain } } + /** Tags only synthetic owner-drain failures created after an invocation cancellation was recorded. */ + private RuntimeException ownerDrainFailure(String message, Throwable cause) { + var diagnostic = new IllegalStateException(message, cause); + var lifecycleCause = lifecycleFailure.get(); + return lifecycleCause == null ? diagnostic : new OwnerDrainFailure(diagnostic, lifecycleCause); + } + + private static final class OwnerDrainFailure extends RuntimeException { + private final IllegalStateException diagnostic; + private final UnrecoverableDurableExecutionException lifecycleCause; + + private OwnerDrainFailure( + IllegalStateException diagnostic, UnrecoverableDurableExecutionException lifecycleCause) { + super(diagnostic.getMessage(), diagnostic, false, false); + this.diagnostic = diagnostic; + this.lifecycleCause = lifecycleCause; + } + } + + static Throwable ownerDrainDiagnostic(Throwable failure) { + return failure instanceof OwnerDrainFailure owner ? owner.diagnostic : failure; + } + + static Throwable selectOwnerDrainFailure(Throwable failure) { + if (!(failure instanceof OwnerDrainFailure owner)) return failure; + for (var existing : owner.lifecycleCause.getSuppressed()) { + if (existing == owner.diagnostic) return owner.lifecycleCause; + } + owner.lifecycleCause.addSuppressed(owner.diagnostic); + return owner.lifecycleCause; + } + /** Retains cleanup failures while every required stage gets its finalization opportunity. */ - private static final class DrainFailure { + private final class DrainFailure { private Throwable failure; private boolean fatal; private boolean interrupted; + private final boolean lifecycleInterruptions; + + private DrainFailure() { + this(false); + } + + private DrainFailure(boolean lifecycleInterruptions) { + this.lifecycleInterruptions = lifecycleInterruptions; + } void run(Runnable stage) { try { @@ -1129,7 +1314,11 @@ void run(Runnable stage) { } void interrupted(String message, InterruptedException problem) { - if (!interrupted) record(new IllegalStateException(message, problem)); + if (!interrupted) + record( + lifecycleInterruptions + ? lifecycleFailure(message, problem) + : ownerDrainFailure(message, problem)); interrupted = true; } @@ -1139,17 +1328,24 @@ void preserveInterrupt() { @SuppressWarnings("removal") void record(Throwable problem) { - var cause = ExceptionHelper.unwrapInvocationFailure(problem); + var nextDiagnostic = ownerDrainDiagnostic(problem); + // Inspect wrappers once for fatal priority and aliases; retain the original diagnostic object below. + var cause = ExceptionHelper.unwrapAsyncFailure(nextDiagnostic); boolean nextFatal = cause instanceof VirtualMachineError || cause instanceof ThreadDeath; // Keep a transport wrapper intact: close may run while the same fatal already escapes the body. // Throwing that fatal object again from try-with-resources would trigger self-suppression. var next = problem; + var previousDiagnostic = ownerDrainDiagnostic(failure); if (failure == null || nextFatal && !fatal) { - if (failure != null && failure != next) next.addSuppressed(failure); + if (previousDiagnostic != null && previousDiagnostic != nextDiagnostic) + nextDiagnostic.addSuppressed(previousDiagnostic); failure = next; fatal = nextFatal; - } else if (failure != next) { - failure.addSuppressed(next); + } else if (previousDiagnostic != nextDiagnostic) { + // A shutdown re-observation of the same cancellation is not a new diagnostic. Keeping it here + // would create lifecycle -> drain -> lifecycle when the owner reports the canonical failure. + if (failure instanceof OwnerDrainFailure owner && cause == owner.lifecycleCause) return; + previousDiagnostic.addSuppressed(nextDiagnostic); } } @@ -1210,6 +1406,7 @@ private T admitCheckpoint(Supplier request, Supplier whenClosed) { } private void requireOperationAdmission() { + checkInvocationAdmission(); var task = currentExecutorTask.get(); if (lifecycleState == LifecycleState.DRAINING && task != null @@ -1224,6 +1421,7 @@ private void requireOperationAdmission() { } private void requireTaskAdmission(BaseDurableOperation operation) { + checkInvocationAdmission(); if (lifecycleState == LifecycleState.DRAINING && operation != null && registeredOperations.get(operation.getOperationId()) == operation) { @@ -1413,4 +1611,260 @@ 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 = signalLifecycleFailure( + lifecycleFailure("Interrupted while draining invocation tasks", interrupted)); + failure = settleInterruptedCleanup(hardDeadline, true, failure); + } catch (RuntimeException problem) { + failure = signalLifecycleFailure(problem); + } finally { + failure = finishActualInvocationCleanup(failure); + 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(); + // Fatal publication already stops operation peers. Retain the root/scope owner's cleanup allowance; + // only lifecycle cutoffs or interruption require immediate registry cancellation. + var abortFailure = 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 = signalLifecycleFailure( + lifecycleFailure("Invocation user work exceeded its graceful cleanup budget", null)); + 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; + return signalLifecycleFailure( + lifecycleFailure("Invocation lifecycle did not exit before the cleanup deadline", null)); + } + + 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 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 Throwable signalLifecycleFailure(Throwable failure) { + if (failure instanceof UnrecoverableDurableExecutionException problem + && LIFECYCLE_ERROR_TYPE.equals(problem.getErrorObject().errorType())) { + lifecycleFailure.compareAndSet(null, problem); + // End and the caller share the first selected lifecycle failure; keep later cleanup diagnostics. + failure = combineFailures(lifecycleFailure.get(), failure); + } + // 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); + return 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.unwrapInvocationFailure(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)); + } + + /** Cutoffs request cancellation; actual owner, continuation, and transport completion remain mandatory. */ + private Throwable finishActualInvocationCleanup(Throwable priorFailure) { + var drain = new DrainFailure(true); + drain.run(() -> stopCheckpointContinuations(drain)); + drain.run(() -> awaitActualInvocationTasks(drain)); + drain.run(() -> closeLifecycle(drain)); + drain.run(checkpointManager::shutdown); + try { + drain.finish(); + } catch (Throwable failure) { + // Keep shutdown's transport wrapper: close can run while its fatal cause already escapes the body. + return combineFailures(priorFailure, failure); + } + return priorFailure; + } + + private void awaitActualInvocationTasks(DrainFailure drain) { + while (true) { + List> tasks; + synchronized (admissionLock) { + tasks = getActiveExecutorTasks(); + if (tasks.isEmpty()) { + lifecycleState = LifecycleState.CLOSED; + return; + } + } + for (var task : tasks) awaitTaskExit(task, drain); + } + } } 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 cd884796e9..13d9edc22e 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 @@ -399,6 +399,7 @@ private void clearHandlerOwner(CompletableFuture completion) { * @return the user function's result */ protected T runUserFunction(Integer attempt, Supplier userFunction) { + executionManager.checkUserFunctionExecutionAllowed(); var pluginRunner = getPluginRunner().captureUserFunctionHooks(); var startInfo = PluginInfoConverter.toUserFunctionStartInfo( operationIdentifier, @@ -409,6 +410,7 @@ protected T runUserFunction(Integer attempt, Supplier userFunction) { T result; try { executionManager.rethrowPluginFatalIfPresent(); + executionManager.checkUserFunctionExecutionAllowed(); result = userFunction.get(); executionManager.rethrowPluginFatalIfPresent(); } catch (Throwable e) { 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 67c2f0fa15..b3232e38a4 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/AcknowledgedDeliveryFailureTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/AcknowledgedDeliveryFailureTest.java new file mode 100644 index 0000000000..e35731bf37 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/AcknowledgedDeliveryFailureTest.java @@ -0,0 +1,265 @@ +// 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 java.time.Duration; +import java.time.Instant; +import java.util.List; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletionException; +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.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +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.CheckpointDurableExecutionResponse; +import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState; +import software.amazon.awssdk.services.lambda.model.ExecutionDetails; +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.TypeToken; +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.model.ExecutionStatus; +import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; +import software.amazon.lambda.durable.plugin.InvocationEndInfo; +import software.amazon.lambda.durable.plugin.InvocationInfo; +import software.amazon.lambda.durable.plugin.InvocationStatus; + +@Timeout(15) +class AcknowledgedDeliveryFailureTest { + private static final String ARN = + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/ack/execution"; + private static final String RESULT = "x".repeat(6 * 1024 * 1024 + 100); + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void metadataFailureSurvivesCallerInterruptionDuringEnd(boolean cancelledMetadata) throws Exception { + RuntimeException failure = cancelledMetadata + ? new CancellationException("metadata request cancelled") + : new IllegalStateException("metadata store failed"); + var observed = run(failure, true); + assertSame(failure, observed.failure(), "Terminal ACK must not hide an unrelated delivery failure"); + assertNull(observed.output()); + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void causelessMetadataFailureRetainsItsIdentity(boolean interruptCaller) throws Exception { + var failure = new CompletionException("metadata failure without a cause", null); + var observed = run(failure, interruptCaller); + assertSame(failure, observed.failure(), "A non-null delivery error must not become successful ACK delivery"); + assertNull(observed.output()); + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void applicationExecutionExceptionRemainsTheDeliveryFailure(boolean interruptCaller) throws Exception { + var applicationFailure = new ExecutionException("application wrapper", new IllegalStateException("metadata")); + try (var fixture = new Fixture(new CompletionException(applicationFailure), applicationFailure)) { + var observed = fixture.observe(interruptCaller); + assertSame(applicationFailure, observed.failure()); + assertNull(observed.output()); + } + } + + @Test + void metadataFailureWithoutCallerInterruptionStillEscapes() throws Exception { + var failure = new IllegalStateException("metadata store failed"); + var observed = run(failure, false); + assertSame(failure, observed.failure()); + assertNull(observed.output()); + } + + @Test + void metadataCancellationWithoutCallerInterruptionRetainsIdentity() throws Exception { + var failure = new CancellationException("metadata request cancelled"); + var observed = run(failure, false); + assertSame(failure, observed.failure()); + assertNull(observed.output()); + } + + @Test + void acknowledgedResultStillWinsOverCallerInterruptionAlone() throws Exception { + var observed = run(null, true); + assertNull(observed.failure()); + assertEquals(ExecutionStatus.SUCCEEDED, observed.output().status()); + assertEquals("", observed.output().result(), "The acknowledged large payload is not returned twice"); + assertNull(observed.output().error()); + } + + private static Observed run(RuntimeException metadataFailure, boolean interruptCaller) throws Exception { + try (var fixture = new Fixture(metadataFailure)) { + return fixture.observe(interruptCaller); + } + } + + private static final class Fixture implements AutoCloseable, DurableExecutionPlugin { + private final ExecutorService users = Executors.newSingleThreadExecutor(); + private final ExecutorService callers = Executors.newSingleThreadExecutor(); + private final AtomicReference caller = new AtomicReference<>(); + private final AtomicReference owner = new AtomicReference<>(); + private final AtomicReference endInfo = new AtomicReference<>(); + private final CountDownLatch enteredEnd = new CountDownLatch(1); + private final CountDownLatch releaseEnd = new CountDownLatch(1); + private final CountDownLatch endFinished = new CountDownLatch(1); + private final AtomicBoolean endInterrupted = new AtomicBoolean(); + private final AtomicBoolean callerInterrupted = new AtomicBoolean(); + private final AtomicInteger endCalls = new AtomicInteger(); + private final AtomicReference> updates = new AtomicReference<>(); + private final Throwable expectedFailure; + private final DurableExecutionClient client; + private final DurableConfig config; + + private Fixture(RuntimeException metadataFailure) { + this(metadataFailure, metadataFailure); + } + + private Fixture(RuntimeException metadataFailure, Throwable expectedFailure) { + this.expectedFailure = expectedFailure; + client = client(metadataFailure, updates); + config = DurableConfig.builder() + .withExecutorService(users) + .withDurableExecutionClient(client) + .withCheckpointDelay(Duration.ZERO) + .withPlugins(ignoredInfo -> this) + .build(); + } + + @Override + public void onInvocationStart(InvocationInfo info) { + owner.set(Thread.currentThread()); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + assertSame(owner.get(), Thread.currentThread(), "End remains on its Start owner"); + endInfo.set(info); + endCalls.incrementAndGet(); + enteredEnd.countDown(); + try { + assertTrue(releaseEnd.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException expected) { + endInterrupted.set(true); + Thread.currentThread().interrupt(); + } finally { + endFinished.countDown(); + } + } + + private Future invoke() { + return callers.submit(() -> { + caller.set(Thread.currentThread()); + try { + return DurableExecutor.execute( + input(), null, TypeToken.get(String.class), (value, context) -> RESULT, config); + } finally { + callerInterrupted.set(Thread.interrupted()); + } + }); + } + + private Observed observe(boolean interruptCaller) throws Exception { + var invocation = invoke(); + assertTrue(enteredEnd.await(5, TimeUnit.SECONDS), "The controlled End hook must start"); + var status = expectedFailure == null ? InvocationStatus.SUCCEEDED : InvocationStatus.RETRYING; + assertEquals(status, endInfo.get().invocationStatus()); + assertSame(expectedFailure, endInfo.get().executionError()); + if (interruptCaller) caller.get().interrupt(); + else releaseEnd.countDown(); + var observed = resultOf(invocation); + assertLifecycle(interruptCaller); + assertTransport(); + return observed; + } + + private void assertLifecycle(boolean interrupted) { + assertEquals(0, endFinished.getCount(), "The cooperative End owner exits before caller completion"); + assertEquals(1, endCalls.get()); + assertEquals(interrupted, endInterrupted.get()); + assertEquals(interrupted, callerInterrupted.get(), "Runtime caller interruption is restored"); + } + + private void assertTransport() { + verify(client, times(1)).checkpoint(any(), any(), any()); + verify(client, times(1)).getExecutionState(ARN, "accepted-token", "metadata"); + assertEquals(1, updates.get().size()); + var update = updates.get().get(0); + assertEquals(OperationType.EXECUTION, update.type()); + assertEquals(OperationAction.SUCCEED, update.action()); + assertEquals("execution", update.id()); + assertEquals('"' + RESULT + '"', update.payload()); + } + + @Override + public void close() throws InterruptedException { + releaseEnd.countDown(); + users.shutdownNow(); + callers.shutdownNow(); + assertTrue(users.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(callers.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private static Observed resultOf(Future invocation) throws Exception { + try { + return new Observed(invocation.get(5, TimeUnit.SECONDS), null); + } catch (ExecutionException failed) { + return new Observed(null, failed.getCause()); + } + } + + private static DurableExecutionClient client( + RuntimeException metadataFailure, AtomicReference> updates) { + var client = mock(DurableExecutionClient.class); + when(client.checkpoint(any(), any(), any())).thenAnswer(call -> { + updates.set(List.copyOf(call.getArgument(2))); + return CheckpointDurableExecutionResponse.builder() + .checkpointToken("accepted-token") + .newExecutionState(CheckpointUpdatedExecutionState.builder() + .nextMarker("metadata") + .build()) + .build(); + }); + when(client.getExecutionState(any(), any(), any())).thenAnswer(call -> { + if (metadataFailure != null) throw metadataFailure; + return GetDurableExecutionStateResponse.builder().build(); + }); + return client; + } + + private static DurableExecutionInput input() { + var execution = Operation.builder() + .id("execution") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails( + ExecutionDetails.builder().inputPayload("\"input\"").build()) + .build(); + return new DurableExecutionInput( + ARN, + "token", + CheckpointUpdatedExecutionState.builder().operations(execution).build()); + } + + private record Observed(DurableExecutionOutput output, Throwable failure) {} +} 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 91d7d135ad..afdcbda6c2 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 @@ -437,4 +437,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/CanonicalLifecycleFailureTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/CanonicalLifecycleFailureTest.java new file mode 100644 index 0000000000..18997b685e --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/CanonicalLifecycleFailureTest.java @@ -0,0 +1,209 @@ +// 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.Mockito.*; + +import com.amazonaws.services.lambda.runtime.Context; +import java.time.Instant; +import java.util.Arrays; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import org.junit.jupiter.api.Timeout; +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.*; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.TypeToken; +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.plugin.*; + +@Timeout(20) +class CanonicalLifecycleFailureTest { + enum Stop { + INTERRUPT, + TIMEOUT + } + + @ParameterizedTest + @CsvSource({"false,INTERRUPT", "true,INTERRUPT", "false,TIMEOUT", "true,TIMEOUT"}) + void firstLifecycleFailureRemainsCanonicalThroughCleanup(boolean plugins, Stop stop) throws Exception { + try (var fixture = new Fixture(plugins)) { + var invocation = fixture.invoke(stop == Stop.TIMEOUT ? 3000 : 60_000); + assertTrue(fixture.bodyEntered.await(5, TimeUnit.SECONDS)); + if (stop == Stop.INTERRUPT) fixture.caller.get().interrupt(); + assertTrue(fixture.ownerInterrupted.await(5, TimeUnit.SECONDS)); + var first = fixture.manager.get().getLifecycleFailure(); + assertNotNull(first); + assertEquals( + stop == Stop.INTERRUPT + ? "Invocation runtime thread was interrupted" + : "Invocation deadline reached before execution completed", + first.getMessage()); + if (stop == Stop.TIMEOUT) awaitActualCleanup(fixture.caller.get()); + fixture.releaseOwner.countDown(); + var observed = assertThrows(ExecutionException.class, () -> invocation.get(5, TimeUnit.SECONDS)) + .getCause(); + assertSame(first, observed, "Cleanup must return the same failure already selected for End"); + assertTrue(first.isRetryable()); + if (stop == Stop.INTERRUPT) assertInstanceOf(InterruptedException.class, first.getCause()); + else assertInstanceOf(TimeoutException.class, first.getCause()); + var laterMessage = stop == Stop.INTERRUPT + ? "Interrupted while draining invocation tasks" + : "Invocation lifecycle did not exit before the cleanup deadline"; + assertTrue( + Arrays.stream(first.getSuppressed()) + .anyMatch(later -> later instanceof UnrecoverableDurableExecutionException + && laterMessage.equals(later.getMessage())), + "The later cleanup diagnostic must remain attached to the first lifecycle failure"); + fixture.assertEnded(first, plugins ? InvocationStatus.RETRYING : null); + assertEquals(stop == Stop.INTERRUPT, fixture.callerInterrupted.get()); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void normalCompletionDoesNotSelectALifecycleFailure(boolean plugins) throws Exception { + try (var fixture = new Fixture(plugins)) { + var invocation = fixture.invoke(60_000); + assertTrue(fixture.bodyEntered.await(5, TimeUnit.SECONDS)); + fixture.releaseOwner.countDown(); + assertEquals( + ExecutionStatus.SUCCEEDED, + invocation.get(5, TimeUnit.SECONDS).status()); + assertNull(fixture.manager.get().getLifecycleFailure()); + fixture.assertEnded(null, plugins ? InvocationStatus.SUCCEEDED : null); + assertFalse(fixture.callerInterrupted.get()); + } + } + + private static void awaitActualCleanup(Thread caller) { + var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (Arrays.stream(caller.getStackTrace()) + .noneMatch(frame -> frame.getClassName().equals(ExecutionManager.class.getName()) + && frame.getMethodName().equals("finishActualInvocationCleanup"))) { + assertTrue( + System.nanoTime() < deadline, "The cancellation allowance did not expire while the owner was held"); + Thread.yield(); + } + } + + private static final class Fixture implements AutoCloseable { + private final ExecutorService users = Executors.newSingleThreadExecutor(); + private final ExecutorService callers = Executors.newSingleThreadExecutor(); + private final CountDownLatch bodyEntered = new CountDownLatch(1); + private final CountDownLatch ownerInterrupted = new CountDownLatch(1); + private final CountDownLatch releaseOwner = new CountDownLatch(1); + private final AtomicReference caller = new AtomicReference<>(); + private final AtomicReference owner = new AtomicReference<>(); + private final AtomicReference manager = new AtomicReference<>(); + private final AtomicReference end = new AtomicReference<>(); + private final AtomicReference endThread = new AtomicReference<>(); + private final AtomicBoolean callerInterrupted = new AtomicBoolean(); + private final AtomicBoolean bodyExited = new AtomicBoolean(); + private final AtomicInteger endCalls = new AtomicInteger(); + private final DurableConfig config; + + private Fixture(boolean plugins) { + var builder = DurableConfig.builder() + .withExecutorService(users) + .withDurableExecutionClient(TestUtils.createMockClient()); + if (plugins) + builder.withPlugins(ignored -> new DurableExecutionPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + end.set(info); + endThread.set(Thread.currentThread()); + endCalls.incrementAndGet(); + } + }); + config = builder.build(); + } + + private Future invoke(int remainingMillis) { + var runtime = mock(Context.class); + when(runtime.getRemainingTimeInMillis()).thenReturn(remainingMillis); + when(runtime.getAwsRequestId()).thenReturn("canonical-life-request"); + return callers.submit(() -> { + caller.set(Thread.currentThread()); + try { + return DurableExecutor.execute( + input(), + runtime, + TypeToken.get(String.class), + (value, context) -> { + manager.set(((DurableContextImpl) context).getExecutionManager()); + owner.set(Thread.currentThread()); + bodyEntered.countDown(); + try { + awaitOwnerRelease(); + return "result"; + } finally { + bodyExited.set(true); + } + }, + config); + } finally { + callerInterrupted.set(Thread.interrupted()); + } + }); + } + + private void awaitOwnerRelease() { + var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10); + while (true) { + try { + assertTrue(releaseOwner.await(Math.max(0, deadline - System.nanoTime()), TimeUnit.NANOSECONDS)); + return; + } catch (InterruptedException ignored) { + ownerInterrupted.countDown(); + } + } + } + + private void assertEnded(Throwable failure, InvocationStatus status) { + assertTrue(bodyExited.get()); + if (status == null) { + assertEquals(0, endCalls.get()); + return; + } + assertEquals(1, endCalls.get()); + assertSame(owner.get(), endThread.get()); + assertEquals(status, end.get().invocationStatus()); + assertSame(failure, end.get().executionError()); + } + + @Override + public void close() throws InterruptedException { + releaseOwner.countDown(); + users.shutdownNow(); + callers.shutdownNow(); + assertTrue(users.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(callers.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private static DurableExecutionInput input() { + return new DurableExecutionInput( + "arn:execution/canonical", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("canonical") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + } +} 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 9ea4a5dba6..90975bd0f1 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 @@ -42,6 +42,8 @@ class CheckpointManagerTest { + private static final Duration SHUTDOWN_TIMEOUT = Duration.ofSeconds(5); + private DurableConfig config; private DurableExecutionClient client; private CheckpointManager batcher; @@ -271,11 +273,11 @@ void detachedPollersSettleAfterFirstObserverAborts(boolean earlierPollFailure, b } @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()); @@ -297,7 +299,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/DrainingAdmissionCancellationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/DrainingAdmissionCancellationTest.java new file mode 100644 index 0000000000..fcedd9e88a --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/DrainingAdmissionCancellationTest.java @@ -0,0 +1,300 @@ +// 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.Mockito.*; + +import com.amazonaws.services.lambda.runtime.Context; +import java.time.Duration; +import java.time.Instant; +import java.util.Arrays; +import java.util.Map; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import java.util.function.BooleanSupplier; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +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.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.model.OperationSubType; +import software.amazon.lambda.durable.operation.BaseDurableOperation; +import software.amazon.lambda.durable.plugin.*; + +@Timeout(20) +class DrainingAdmissionCancellationTest { + enum Timing { + AFTER_FINAL_BROADCAST, + DURING_INSERTION, + GRACEFUL + } + + @ParameterizedTest + @CsvSource({ + "AFTER_FINAL_BROADCAST,false", + "AFTER_FINAL_BROADCAST,true", + "DURING_INSERTION,false", + "DURING_INSERTION,true", + "GRACEFUL,false", + "GRACEFUL,true" + }) + void replayedWaitCannotMissCancellationButGracefulDescendantsRemainAllowed(Timing timing, boolean plugins) + throws Exception { + var childId = new OperationIdGenerator(null).nextOperationId(); + var waitId = new OperationIdGenerator(childId).nextOperationId(); + var childEntered = new CountDownLatch(1); + var releaseChild = new CountDownLatch(1); + var childExited = new CountDownLatch(1); + var interruptedChild = new CountDownLatch(1); + var beforeInsertion = new CountDownLatch(1); + var releaseInsertion = new CountDownLatch(1); + var manager = new AtomicReference(); + var caller = new AtomicReference(); + var rootOwner = new AtomicReference(); + var endOwner = new AtomicReference(); + var end = new AtomicReference(); + var ends = new AtomicInteger(); + var stepBodies = new AtomicInteger(); + var failureSeen = new AtomicReference(); + var lateWait = new AtomicReference(); + var callerInterrupted = new AtomicBoolean(); + var workers = Executors.newCachedThreadPool(); + var callers = Executors.newSingleThreadExecutor(); + var client = TestUtils.createMockClient(); + var config = DurableConfig.builder() + .withExecutorService(workers) + .withDurableExecutionClient(client) + .withCheckpointDelay(Duration.ZERO); + if (plugins) + config.withPlugins(info -> new DurableExecutionPlugin() { + public void onInvocationEnd(InvocationEndInfo info) { + endOwner.set(Thread.currentThread()); + end.set(info); + ends.incrementAndGet(); + } + }); + var runtime = mock(Context.class); + when(runtime.getRemainingTimeInMillis()).thenReturn(timing == Timing.AFTER_FINAL_BROADCAST ? 3000 : 10_000); + when(runtime.getAwsRequestId()).thenReturn("draining-admission"); + var pending = timing != Timing.GRACEFUL; + var history = CheckpointUpdatedExecutionState.builder() + .operations( + Operation.builder() + .id("exec") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build(), + Operation.builder() + .id(childId) + .name("child") + .type(OperationType.CONTEXT) + .subType(OperationSubType.RUN_IN_CHILD_CONTEXT.getValue()) + .status(OperationStatus.STARTED) + .build(), + Operation.builder() + .id(waitId) + .parentId(childId) + .name("old-wait") + .type(OperationType.WAIT) + .subType(OperationSubType.WAIT.getValue()) + .status(pending ? OperationStatus.PENDING : OperationStatus.SUCCEEDED) + .waitDetails(WaitDetails.builder() + .scheduledEndTimestamp(Instant.now().plusSeconds(3600)) + .build()) + .build()) + .build(); + var input = new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/draining/exec", + "token", + history); + Future response = null; + try { + response = callers.submit(() -> { + caller.set(Thread.currentThread()); + try { + return DurableExecutor.execute( + input, + runtime, + TypeToken.get(String.class), + (value, root) -> { + rootOwner.set(Thread.currentThread()); + var execution = ((DurableContextImpl) root).getExecutionManager(); + manager.set(execution); + if (timing == Timing.DURING_INSERTION) + installInsertionGate( + execution, waitId, beforeInsertion, releaseInsertion, lateWait); + root.runInChildContextAsync("child", String.class, child -> { + childEntered.countDown(); + awaitIgnoringInterrupt(releaseChild, interruptedChild); + try { + var wait = child.waitAsync("old-wait", Duration.ofHours(1)); + lateWait.set((BaseDurableOperation) wait); + wait.get(); + if (timing == Timing.GRACEFUL) + child.step("new-descendant", String.class, step -> { + stepBodies.incrementAndGet(); + return "step"; + }); + } catch (Throwable failure) { + failureSeen.set(failure); + } finally { + childExited.countDown(); + } + return "child"; + }); + await(childEntered); + return "root"; + }, + config.build()); + } finally { + callerInterrupted.set(Thread.interrupted()); + } + }); + await(childEntered); + awaitCondition(() -> inMethod(caller.get(), "awaitTasksUntil")); + if (timing == Timing.DURING_INSERTION) { + releaseChild.countDown(); + await(beforeInsertion); + caller.get().interrupt(); + awaitCondition(() -> manager.get().getLifecycleFailure() != null + && caller.get().getState() == Thread.State.BLOCKED + && inMethod(caller.get(), "awaitTasksUntil")); + assertFalse( + inMethod(caller.get(), "signalLifecycleFailure"), "Cancellation broadcast must already finish"); + releaseInsertion.countDown(); + } else if (timing == Timing.AFTER_FINAL_BROADCAST) { + caller.get().interrupt(); + await(interruptedChild); + awaitCondition(() -> inMethod(caller.get(), "awaitActualInvocationTasks")); + releaseChild.countDown(); + } else { + assertEquals( + ExecutionManager.LifecycleState.DRAINING, manager.get().getLifecycleState()); + assertNull(manager.get().getLifecycleFailure()); + releaseChild.countDown(); + } + assertTrue( + childExited.await(700, TimeUnit.MILLISECONDS), + "A cancelled descendant must reject the old PENDING wait, not block beyond the final broadcast"); + if (pending) { + var rejected = assertInstanceOf(UnrecoverableDurableExecutionException.class, failureSeen.get()); + assertSame(manager.get().getLifecycleFailure(), rejected); + var observedResponse = response; + var callerFailure = assertThrows(ExecutionException.class, () -> responseGet(observedResponse)); + assertSame(rejected, callerFailure.getCause()); + assertTrue(callerInterrupted.get()); + assertEquals(0, stepBodies.get()); + verify(client, never()).checkpoint(any(), any(), any()); + } else { + var result = (DurableExecutionOutput) response.get(5, TimeUnit.SECONDS); + assertEquals(ExecutionStatus.SUCCEEDED, result.status()); + assertNull(failureSeen.get()); + assertEquals(1, stepBodies.get()); + assertFalse(callerInterrupted.get()); + } + assertEquals(plugins ? 1 : 0, ends.get()); + if (plugins) { + assertSame(rootOwner.get(), endOwner.get()); + assertEquals( + pending ? InvocationStatus.RETRYING : InvocationStatus.SUCCEEDED, + end.get().invocationStatus()); + } + } finally { + releaseChild.countDown(); + releaseInsertion.countDown(); + var waiting = lateWait.get(); + if (waiting != null && !waiting.getCompletionFuture().isDone()) + waiting.getCompletionFuture() + .completeExceptionally(new IllegalStateException("test teardown releases blocked wait")); + if (response != null) { + try { + response.get(5, TimeUnit.SECONDS); + } catch (ExecutionException expected) { + } + } + workers.shutdownNow(); + callers.shutdownNow(); + assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(callers.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private static Object responseGet(Future response) throws Exception { + return response.get(5, TimeUnit.SECONDS); + } + + @SuppressWarnings("unchecked") + private static void installInsertionGate( + ExecutionManager manager, + String target, + CountDownLatch entered, + CountDownLatch release, + AtomicReference rescue) { + try { + var field = ExecutionManager.class.getDeclaredField("registeredOperations"); + field.setAccessible(true); + var original = (Map) field.get(manager); + var gated = new ConcurrentHashMap(original) { + @Override + public BaseDurableOperation put(String id, BaseDurableOperation operation) { + if (id.equals(target)) { + rescue.set(operation); + entered.countDown(); + awaitIgnoringInterrupt(release, new CountDownLatch(0)); + } + return super.put(id, operation); + } + }; + field.set(manager, gated); + } catch (ReflectiveOperationException failure) { + throw new AssertionError(failure); + } + } + + private static boolean inMethod(Thread thread, String method) { + return thread != null + && Arrays.stream(thread.getStackTrace()) + .anyMatch(frame -> frame.getMethodName().equals(method)); + } + + private static void awaitCondition(BooleanSupplier condition) throws InterruptedException { + var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(6); + while (!condition.getAsBoolean()) { + assertTrue(System.nanoTime() < deadline, "Fixture did not reach the required admission/broadcast phase"); + TimeUnit.MILLISECONDS.sleep(5); + } + } + + private static void await(CountDownLatch gate) { + try { + assertTrue(gate.await(8, TimeUnit.SECONDS)); + } catch (InterruptedException failure) { + Thread.currentThread().interrupt(); + throw new AssertionError(failure); + } + } + + private static void awaitIgnoringInterrupt(CountDownLatch gate, CountDownLatch interrupted) { + var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(12); + while (true) { + try { + assertTrue(gate.await(Math.max(0, deadline - System.nanoTime()), TimeUnit.NANOSECONDS)); + return; + } catch (InterruptedException ignored) { + interrupted.countDown(); + } + } + } +} 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 76a5845914..976d7355be 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 @@ -5,24 +5,31 @@ 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.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.doAnswer; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; 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; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; 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.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -38,6 +45,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 +65,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 +170,54 @@ void shutdownWaitsForAdmittedPollAndPreservesInterruption(boolean interrupted) t } } + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void stuckCheckpointAdmissionWaitsPastCleanupDeadlineAfterSelectingRetry(boolean interrupted) 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 restoredInterrupt = new AtomicBoolean(); + var shutdown = callers.submit(() -> { + var outcome = manager.awaitInvocationOutcome(new CompletableFuture<>()); + if (interrupted) Thread.currentThread().interrupt(); + try { + return manager.finishInvocation(outcome.failure()); + } finally { + restoredInterrupt.set(Thread.interrupted()); + } + }); + callers.submit(this::awaitClosedAdmission).get(5, TimeUnit.SECONDS); + var failure = assertInstanceOf(UnrecoverableDurableExecutionException.class, manager.getLifecycleFailure()); + assertTrue(failure.isRetryable()); + assertInstanceOf(TimeoutException.class, failure.getCause()); + // The cutoff selects RETRYING and fences admission, but cannot skip an already-admitted enqueue. + assertThrows(TimeoutException.class, () -> shutdown.get(600, TimeUnit.MILLISECONDS)); + assertThrows(RejectedExecutionException.class, () -> manager.sendOperationUpdate(update())); + verify(client, never()).checkpoint(any(), any(), any()); + releaseAdmission.countDown(); + var poll = admission.get(5, TimeUnit.SECONDS); + assertSame(failure, shutdown.get(5, TimeUnit.SECONDS)); + assertEquals(interrupted, restoredInterrupt.get()); + assertTrue(poll.isCompletedExceptionally()); + verify(client, never()).checkpoint(any(), any(), any()); + } finally { + releaseAdmission.countDown(); + } + } + @Test void failedAdmissionDoesNotPreventShutdown() throws Exception { doThrow(new IllegalArgumentException("polling configuration failed")) @@ -283,20 +356,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 edd1cd741a..5e26e95482 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(ignoredInfo -> 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/HandlerScopeHandoffTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/HandlerScopeHandoffTest.java index 3eef976a30..d874896332 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/HandlerScopeHandoffTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/HandlerScopeHandoffTest.java @@ -25,11 +25,11 @@ import software.amazon.lambda.durable.plugin.*; import software.amazon.lambda.durable.util.ExceptionHelper; -/** The old caller handoff is replaced by handler-owned completion, including exhausted runtime budgets. */ +/** The old caller handoff is replaced by handler-owned completion, after the execution cutoff. */ class HandlerScopeHandoffTest { @ParameterizedTest @ValueSource(booleans = {false, true}) - void controlWaitsForScopeOwnerAndEndDespiteExpiredHandoffBudget(boolean retry) throws Exception { + void selectedControlWaitsForScopeOwnerAndEndAfterExecutionCutoff(boolean retry) throws Exception { var scopeEntered = new CountDownLatch(1); var releaseScope = new CountDownLatch(1); var endEntered = new CountDownLatch(1); @@ -37,6 +37,7 @@ void controlWaitsForScopeOwnerAndEndDespiteExpiredHandoffBudget(boolean retry) t var owner = new AtomicReference(); var endThread = new AtomicReference(); var end = new AtomicReference(); + var executionManager = new AtomicReference(); var original = new UnrecoverableDurableExecutionException( ErrorObject.builder().errorMessage("retry").build(), true); var plugin = new ScopedPlugin( @@ -57,7 +58,7 @@ void controlWaitsForScopeOwnerAndEndDespiteExpiredHandoffBudget(boolean retry) t var workers = Executors.newSingleThreadExecutor(); var callers = Executors.newSingleThreadExecutor(); var runtime = mock(Context.class); - when(runtime.getRemainingTimeInMillis()).thenReturn(0); + when(runtime.getRemainingTimeInMillis()).thenReturn(3000); try { var config = DurableConfig.builder() .withExecutorService(workers) @@ -70,6 +71,7 @@ void controlWaitsForScopeOwnerAndEndDespiteExpiredHandoffBudget(boolean retry) t TypeToken.get(String.class), (value, context) -> { var manager = ((DurableContextImpl) context).getExecutionManager(); + executionManager.set(manager); try { if (retry) manager.terminateExecution(original); manager.suspendExecution(); @@ -80,12 +82,21 @@ void controlWaitsForScopeOwnerAndEndDespiteExpiredHandoffBudget(boolean retry) t }, config)); assertTrue(scopeEntered.await(3, TimeUnit.SECONDS)); - assertThrows(TimeoutException.class, () -> response.get(600, TimeUnit.MILLISECONDS)); + // Admission must precede the deadline. The selected control may then wait for its real scope + // owner past the execution cutoff, without consuming the separate cleanup budget. + var manager = executionManager.get(); + var untilCutoff = + manager.getRemainingInvocationTime().orElseThrow().toNanos() - TimeUnit.MILLISECONDS.toNanos(500); + assertTrue(untilCutoff > 0, "The scope must open before the execution cutoff"); + assertThrows( + TimeoutException.class, + () -> response.get(untilCutoff + TimeUnit.MILLISECONDS.toNanos(1), TimeUnit.NANOSECONDS)); + assertTrue(manager.getRemainingInvocationTime().orElseThrow().toMillis() < 500); assertEquals(1, endEntered.getCount(), "No End before the real scope owner finishes"); releaseScope.countDown(); assertTrue(endEntered.await(3, TimeUnit.SECONDS)); assertSame(owner.get(), endThread.get(), "End belongs to the handler thread"); - assertThrows(TimeoutException.class, () -> response.get(100, TimeUnit.MILLISECONDS)); + assertThrows(TimeoutException.class, () -> response.get(50, TimeUnit.MILLISECONDS)); releaseEnd.countDown(); if (retry) { var failure = assertThrows(ExecutionException.class, () -> response.get(3, TimeUnit.SECONDS)); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/InFlightFatalAbortTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/InFlightFatalAbortTest.java index 38f86793e5..51e02f98a5 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/InFlightFatalAbortTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/InFlightFatalAbortTest.java @@ -32,7 +32,9 @@ void activeRequestObservesFatalButShutdownWaitsForBackendReturn(boolean threadDe var request = batcher.submit("START", Duration.ZERO); try { assertTrue(entered.await(3, TimeUnit.SECONDS)); - var shutdown = closing.submit(batcher::shutdown); + var shutdown = closing.submit(() -> { + batcher.shutdown(); + }); assertThrows(TimeoutException.class, () -> shutdown.get(100, TimeUnit.MILLISECONDS)); batcher.abortPending(fatal); assertFatal(request, fatal); @@ -76,7 +78,9 @@ void checkpointShutdownWaitsForAdmittedTransportAfterFatal(boolean threadDeath) .build()); try { assertTrue(entered.await(3, TimeUnit.SECONDS)); - var shutdown = closing.submit(manager::shutdown); + var shutdown = closing.submit(() -> { + manager.shutdown(); + }); assertThrows(TimeoutException.class, () -> shutdown.get(100, TimeUnit.MILLISECONDS)); signal.set(fatal); manager.abortPending(fatal); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/IndependentDeadlineBoundaryTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/IndependentDeadlineBoundaryTest.java new file mode 100644 index 0000000000..bce33b17bd --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/IndependentDeadlineBoundaryTest.java @@ -0,0 +1,211 @@ +// 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.Instant; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import org.junit.jupiter.api.Timeout; +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.*; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.InvokeConfig; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.plugin.*; +import software.amazon.lambda.durable.serde.JacksonSerDes; +import software.amazon.lambda.durable.util.ExceptionHelper; + +@Timeout(15) +class IndependentDeadlineBoundaryTest { + @ParameterizedTest + @CsvSource({"false,false", "true,false", "false,true", "true,true"}) + void cancelledQueuedRootCannotHangOrRunAfterCallerReturns(boolean plugin, boolean interrupted) throws Exception { + var occupied = new CountDownLatch(1); + var release = new CountDownLatch(1); + var users = new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>()); + var callers = Executors.newSingleThreadExecutor(); + var caller = new AtomicReference(); + var interruptRestored = new AtomicBoolean(); + var bodies = new AtomicInteger(); + var starts = new AtomicInteger(); + var ends = new AtomicInteger(); + var client = TestUtils.createMockClient(); + var builder = DurableConfig.builder().withExecutorService(users).withDurableExecutionClient(client); + if (plugin) + builder.withPlugins(ignoredInvocation -> new DurableExecutionPlugin() { + @Override + public void onInvocationStart(InvocationInfo info) { + starts.incrementAndGet(); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + ends.incrementAndGet(); + } + }); + try { + users.execute(() -> { + occupied.countDown(); + try { + release.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + assertTrue(occupied.await(3, TimeUnit.SECONDS)); + var config = builder.build(); + var result = callers.submit(() -> { + caller.set(Thread.currentThread()); + try { + return DurableExecutor.execute( + input(), + runtime(interrupted ? 10_000 : 1_500), + TypeToken.get(String.class), + (value, context) -> { + bodies.incrementAndGet(); + return "unexpected"; + }, + config); + } finally { + interruptRestored.set(Thread.interrupted()); + } + }); + var limit = System.nanoTime() + TimeUnit.SECONDS.toNanos(3); + while (users.getQueue().isEmpty() && System.nanoTime() < limit) Thread.yield(); + assertEquals(1, users.getQueue().size(), "Root must actually be queued before cancellation"); + if (interrupted) caller.get().interrupt(); + var cause = assertThrows(ExecutionException.class, () -> result.get(3, TimeUnit.SECONDS)) + .getCause(); + var failure = assertInstanceOf(UnrecoverableDurableExecutionException.class, cause); + assertTrue(failure.isRetryable()); + assertEquals(interrupted, interruptRestored.get()); + release.countDown(); + users.submit(() -> {}).get(3, TimeUnit.SECONDS); + assertEquals(0, bodies.get()); + assertEquals(0, starts.get()); + assertEquals(0, ends.get()); + verify(client, never()).checkpoint(any(), any(), any()); + } finally { + release.countDown(); + callers.shutdownNow(); + users.shutdownNow(); + assertTrue(callers.awaitTermination(3, TimeUnit.SECONDS)); + assertTrue(users.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void caughtLateInvokeCheckpointCannotPublishFallback(boolean plugin) { + var caught = new AtomicReference(); + var serializations = new AtomicInteger(); + var fallbackSerializations = new AtomicInteger(); + var ends = new AtomicReference(); + var users = mock(ExecutorService.class); + doAnswer(call -> { + ((Runnable) call.getArgument(0)).run(); + return null; + }) + .when(users) + .execute(any()); + var client = TestUtils.createMockClient(); + var builder = DurableConfig.builder() + .withExecutorService(users) + .withDurableExecutionClient(client) + .withSerDes(new JacksonSerDes() { + @Override + public String serialize(Object value) { + if ("fallback".equals(value)) fallbackSerializations.incrementAndGet(); + return super.serialize(value); + } + }); + if (plugin) + builder.withPlugins(ignoredInvocation -> new DurableExecutionPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + ends.set(info); + } + }); + var payloadSerdes = new JacksonSerDes() { + @Override + public String serialize(Object value) { + serializations.incrementAndGet(); + try { + Thread.sleep(2_100); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return super.serialize(value); + } + }; + var config = builder.build(); + var returned = new AtomicReference(); + var failure = assertThrows( + UnrecoverableDurableExecutionException.class, + () -> returned.set(DurableExecutor.execute( + input(), + runtime(2_000), + TypeToken.get(String.class), + (value, context) -> { + try { + context.invoke( + "late-invoke", + "callee:1", + "payload", + String.class, + InvokeConfig.builder() + .payloadSerDes(payloadSerdes) + .build()); + } catch (Throwable error) { + caught.set(error); + } + return "fallback"; + }, + config)), + () -> "Observed output=" + returned.get() + "; caught=" + caught.get()); + assertEquals(1, serializations.get()); + assertNotNull(caught.get(), "The public handler must actually catch the late operation rejection"); + assertSame(failure, ExceptionHelper.unwrapCompletableFuture(caught.get())); + assertEquals(0, fallbackSerializations.get(), "A rejected checkpoint must not prepare a fallback success"); + assertTrue(failure.isRetryable()); + verify(client, never()).checkpoint(any(), any(), any()); + if (plugin) { + assertEquals(InvocationStatus.RETRYING, ends.get().invocationStatus()); + assertSame(failure, ends.get().executionError()); + } + } + + private static Context runtime(int millis) { + var context = mock(Context.class); + when(context.getRemainingTimeInMillis()).thenReturn(millis); + when(context.getAwsRequestId()).thenReturn("independent-request"); + return context; + } + + private static DurableExecutionInput input() { + return new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/independent/exec", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("exec") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationAdmissionFailureTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationAdmissionFailureTest.java new file mode 100644 index 0000000000..5fdf9e591d --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationAdmissionFailureTest.java @@ -0,0 +1,168 @@ +// 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.Mockito.*; + +import com.amazonaws.services.lambda.runtime.Context; +import java.time.Instant; +import java.util.concurrent.LinkedBlockingQueue; +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.AtomicInteger; +import org.junit.jupiter.api.Timeout; +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.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.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.model.DurableExecutionOutput; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; +import software.amazon.lambda.durable.plugin.InvocationEndInfo; +import software.amazon.lambda.durable.plugin.InvocationInfo; + +@Timeout(15) +class InvocationAdmissionFailureTest { + @ParameterizedTest + @CsvSource({"0,false", "0,true", "1,false", "1,true", "499,false", "499,true", "500,false", "500,true"}) + void expiredRootAdmissionRetainsRetryableFailureWithoutDispatch(int remainingMillis, boolean plugins) + throws Exception { + try (var fixture = new Fixture(plugins, null)) { + var failure = + assertThrows(UnrecoverableDurableExecutionException.class, () -> fixture.invoke(remainingMillis)); + assertTrue(failure.isRetryable()); + assertEquals( + "software.amazon.lambda.durable.execution.InvocationLifecycleException", + failure.getErrorObject().errorType()); + assertEquals("Invocation deadline reached before user work", failure.getMessage()); + assertInstanceOf(TimeoutException.class, failure.getCause()); + fixture.assertNoUserWork(); + assertEquals(0, fixture.dispatches.get()); + assertFalse(fixture.users.isShutdown(), "Shared executors remain caller-owned after rejected admission"); + verifyNoInteractions(fixture.client); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void usableDeadlineStillRunsTheHandler(boolean plugins) throws Exception { + try (var fixture = new Fixture(plugins, null)) { + var response = fixture.invoke(10_000); + assertEquals(ExecutionStatus.SUCCEEDED, response.status()); + assertEquals("\"input\"", response.result()); + assertEquals(1, fixture.dispatches.get()); + assertEquals(1, fixture.bodies.get()); + assertEquals(plugins ? 1 : 0, fixture.starts.get()); + assertEquals(plugins ? 1 : 0, fixture.ends.get()); + assertFalse(fixture.users.isShutdown()); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void synchronousExecutorRejectionRetainsItsOriginalIdentity(boolean plugins) throws Exception { + var rejection = new RejectedExecutionException("executor rejected the root"); + try (var fixture = new Fixture(plugins, rejection)) { + assertSame(rejection, assertThrows(RejectedExecutionException.class, () -> fixture.invoke(10_000))); + fixture.assertNoUserWork(); + assertEquals(1, fixture.dispatches.get()); + assertFalse(fixture.users.isShutdown()); + } + } + + private static final class Fixture implements AutoCloseable { + private final AtomicInteger dispatches = new AtomicInteger(); + private final AtomicInteger bodies = new AtomicInteger(); + private final AtomicInteger factories = new AtomicInteger(); + private final AtomicInteger starts = new AtomicInteger(); + private final AtomicInteger ends = new AtomicInteger(); + private final DurableExecutionClient client = mock(DurableExecutionClient.class); + private final ThreadPoolExecutor users; + private final DurableConfig config; + + private Fixture(boolean plugins, RejectedExecutionException rejection) { + users = new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>()) { + @Override + public void execute(Runnable command) { + dispatches.incrementAndGet(); + if (rejection != null) throw rejection; + super.execute(command); + } + }; + var builder = DurableConfig.builder().withExecutorService(users).withDurableExecutionClient(client); + if (plugins) + builder.withPlugins(ignored -> { + factories.incrementAndGet(); + return new DurableExecutionPlugin() { + @Override + public void onInvocationStart(InvocationInfo info) { + starts.incrementAndGet(); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + ends.incrementAndGet(); + } + }; + }); + config = builder.build(); + } + + private DurableExecutionOutput invoke(int remainingMillis) { + var runtime = mock(Context.class); + when(runtime.getRemainingTimeInMillis()).thenReturn(remainingMillis); + when(runtime.getAwsRequestId()).thenReturn("admission-request"); + return DurableExecutor.execute( + input(), + runtime, + TypeToken.get(String.class), + (value, context) -> { + bodies.incrementAndGet(); + return value; + }, + config); + } + + private void assertNoUserWork() { + assertEquals(0, bodies.get()); + assertEquals(0, factories.get()); + assertEquals(0, starts.get()); + assertEquals(0, ends.get()); + } + + @Override + public void close() throws InterruptedException { + users.shutdownNow(); + assertTrue(users.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + private static DurableExecutionInput input() { + return new DurableExecutionInput( + "arn:execution/admission", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("admission") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + } +} 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 0000000000..7fa150e6a2 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationCleanupBoundaryTest.java @@ -0,0 +1,515 @@ +// 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 callers = Executors.newSingleThreadExecutor(); + var manager = manager(users, context(2_000)); + var entered = new CountDownLatch(1); + var cancelled = new CountDownLatch(1); + var release = new CountDownLatch(1); + Error fatal = threadDeath ? new ThreadDeath() : new InternalError("selected fatal"); + try { + manager.submitRootTask(() -> { + entered.countDown(); + awaitIgnoringInterrupt(release, cancelled); + return "ignored"; + }); + await(entered); + var task = manager.getActiveExecutorTasks().get(0); + var finished = callers.submit(() -> manager.finishInvocation(fatal)); + await(cancelled); + assertThrows(TimeoutException.class, () -> finished.get(3, TimeUnit.SECONDS)); + assertFalse(task.exit().isDone(), "The cleanup budget cannot stand in for the owner's actual exit"); + release.countDown(); + assertSame(fatal, finished.get(5, TimeUnit.SECONDS)); + assertTrue(task.exit().isDone()); + assertTrue(Arrays.stream(fatal.getSuppressed()).anyMatch(ExecutionManager::isLifecycleFailure)); + assertEquals(ExecutionManager.LifecycleState.CLOSED, manager.getLifecycleState()); + } finally { + release.countDown(); + stop(users); + stop(callers); + 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 ownerEndFinishesBeforeCallerReturnAndRestoresScopeBeforeWorkerReuse() 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 cancelled = new CountDownLatch(1); + 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(ignoredInvocation -> plugin) + .build(); + var invocation = callers.submit(() -> DurableExecutor.execute( + input(), + context(2_000), + TypeToken.get(String.class), + (value, ctx) -> { + entered.countDown(); + awaitIgnoringInterrupt(release, cancelled); + assertThrows( + UnrecoverableDurableExecutionException.class, + () -> ctx.step("late", String.class, step -> { + lateBodies.incrementAndGet(); + return "late"; + })); + return "late-result"; + }, + config)); + await(entered); + await(cancelled); + assertThrows(TimeoutException.class, () -> invocation.get(3, TimeUnit.SECONDS)); + assertEquals(0, ended.get(), "The caller cannot run End on a different thread to force cleanup"); + release.countDown(); + assertInstanceOf( + UnrecoverableDurableExecutionException.class, + assertThrows(ExecutionException.class, () -> invocation.get(5, TimeUnit.SECONDS)) + .getCause()); + assertEquals(0, end.getCount(), "The owner's End must finish before the SDK caller returns"); + 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 terminalAckWaitsForAdmittedPublicationThenPropagatesEndFatal() 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(ignoredInvocation -> 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); + assertThrows(TimeoutException.class, () -> invocation.get(3, TimeUnit.SECONDS)); + assertEquals(1, fatalDelivered.getCount(), "End cannot overtake admitted checkpoint publication"); + releaseChange.countDown(); + 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) { + awaitIgnoringInterrupt(latch, new CountDownLatch(0)); + } + + private static void awaitIgnoringInterrupt(CountDownLatch latch, CountDownLatch interrupted) { + while (latch.getCount() > 0) { + try { + latch.await(); + } catch (InterruptedException ignored) { + interrupted.countDown(); + // 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)); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationDeadlineReviewRegressionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationDeadlineReviewRegressionTest.java new file mode 100644 index 0000000000..6769e45cc8 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationDeadlineReviewRegressionTest.java @@ -0,0 +1,265 @@ +// 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 static software.amazon.lambda.durable.TypeToken.get; + +import com.amazonaws.services.lambda.runtime.Context; +import java.time.Duration; +import java.time.Instant; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import org.junit.jupiter.api.Timeout; +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.*; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.model.ExecutionStatus; +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(20) +class InvocationDeadlineReviewRegressionTest { + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void caughtDeadlineAdmissionCannotPublishSuccessfulFallback(boolean plugins) { + var caught = new AtomicReference(); + var bodyEntered = new AtomicBoolean(); + var effects = new AtomicInteger(); + var fallbackSerializations = new AtomicInteger(); + var end = new AtomicReference(); + var config = config( + inlineExecutor(), + new JacksonSerDes() { + @Override + public String serialize(Object value) { + if ("fallback".equals(value)) fallbackSerializations.incrementAndGet(); + return super.serialize(value); + } + }, + plugins, + end); + var failure = assertThrows( + UnrecoverableDurableExecutionException.class, + () -> DurableExecutor.execute( + input(), + runtime(2_000), + get(String.class), + (value, context) -> { + bodyEntered.set(true); + try { + TimeUnit.MILLISECONDS.sleep(1_600); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw new AssertionError(interrupted); + } + try { + context.step("too-late", String.class, step -> { + effects.incrementAndGet(); + return "side-effect"; + }); + } catch (UnrecoverableDurableExecutionException interrupted) { + caught.set(interrupted); + } + return "fallback"; + }, + config)); + assertTrue(bodyEntered.get(), "The handler must reach the caught-admission path"); + assertSame(caught.get(), failure, "Swallowing a deadline rejection must not erase its selected failure"); + assertTrue(failure.isRetryable()); + assertEquals(0, effects.get()); + assertEquals(0, fallbackSerializations.get(), "An expired invocation must not prepare a fallback success"); + if (plugins) { + assertEquals(InvocationStatus.RETRYING, end.get().invocationStatus()); + assertSame(failure, end.get().executionError()); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void inlineSuccessBeforeTheCutoffIsStillAllowed(boolean plugins) { + var end = new AtomicReference(); + var config = config(inlineExecutor(), new JacksonSerDes(), plugins, end); + var output = DurableExecutor.execute( + input(), runtime(10_000), get(String.class), (value, context) -> "result", config); + assertEquals(ExecutionStatus.SUCCEEDED, output.status()); + assertEquals("\"result\"", output.result()); + if (plugins) assertEquals(InvocationStatus.SUCCEEDED, end.get().invocationStatus()); + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void anAlreadySelectedSuspensionSurvivesLaterCaughtAdmission(boolean plugins) { + var end = new AtomicReference(); + var effects = new AtomicInteger(); + var config = config(inlineExecutor(), new JacksonSerDes(), plugins, end); + var output = DurableExecutor.execute( + input(), + runtime(2_000), + get(String.class), + (value, context) -> { + try { + context.wait("pause", Duration.ofSeconds(5)); + throw new AssertionError("The wait should select suspension"); + } catch (SuspendExecutionException suspended) { + try { + TimeUnit.MILLISECONDS.sleep(1_600); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw new AssertionError(interrupted); + } + try { + context.step("too-late", String.class, step -> { + effects.incrementAndGet(); + return "side-effect"; + }); + } catch (RuntimeException stopped) { + // A rejected operation must not replace the suspension already selected by the durable + // wait. + } + return "fallback"; + } + }, + config); + assertEquals(ExecutionStatus.PENDING, output.status()); + assertEquals(0, effects.get()); + if (plugins) assertEquals(InvocationStatus.PENDING, end.get().invocationStatus()); + } + + @ParameterizedTest + @CsvSource({"false,true", "true,true", "false,false", "true,false"}) + void cancellationDuringSerializationRetainsFatalPriority(boolean plugins, boolean fatal) throws Exception { + var serializationEntered = new CountDownLatch(1); + var serializerInterrupted = new CountDownLatch(1); + var release = new CountDownLatch(1); + var uncaught = new CountDownLatch(1); + var uncaughtFailure = new AtomicReference(); + var runtimeThread = new AtomicReference(); + var end = new AtomicReference(); + Throwable serializationFailure = fatal + ? new InternalError("serializer fatal") + : new IllegalStateException("ordinary serializer failure"); + var workers = Executors.newSingleThreadExecutor(action -> { + var worker = new Thread(action, "serialization-owner"); + worker.setUncaughtExceptionHandler((thread, error) -> { + uncaughtFailure.set(error); + uncaught.countDown(); + }); + return worker; + }); + var callers = Executors.newSingleThreadExecutor(); + var serializer = new JacksonSerDes() { + @Override + public String serialize(Object value) { + if ("result".equals(value)) { + serializationEntered.countDown(); + try { + release.await(); + } catch (InterruptedException interrupted) { + serializerInterrupted.countDown(); + Thread.currentThread().interrupt(); + if (serializationFailure instanceof Error error) throw error; + throw (RuntimeException) serializationFailure; + } + } + return super.serialize(value); + } + }; + try { + var config = config(workers, serializer, plugins, end); + var invocation = callers.submit(() -> { + runtimeThread.set(Thread.currentThread()); + return DurableExecutor.execute( + input(), runtime(10_000), get(String.class), (value, context) -> "result", config); + }); + assertTrue(serializationEntered.await(5, TimeUnit.SECONDS)); + runtimeThread.get().interrupt(); + assertTrue(serializerInterrupted.await(5, TimeUnit.SECONDS), "Cleanup must physically interrupt the owner"); + var actual = assertThrows(ExecutionException.class, () -> invocation.get(5, TimeUnit.SECONDS)) + .getCause(); + if (fatal) { + assertSame(serializationFailure, actual, "An available JVM fatal must beat lifecycle cancellation"); + assertTrue(uncaught.await(5, TimeUnit.SECONDS), "The actual worker must not absorb a JVM fatal"); + assertSame(serializationFailure, uncaughtFailure.get()); + } else { + assertInstanceOf(UnrecoverableDurableExecutionException.class, actual); + assertNotSame( + serializationFailure, actual, "An ordinary late error must not replace caller interruption"); + } + if (plugins) { + assertEquals(InvocationStatus.RETRYING, end.get().invocationStatus()); + assertSame(serializationFailure, end.get().executionError()); + } + } finally { + release.countDown(); + callers.shutdownNow(); + workers.shutdownNow(); + assertTrue(callers.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private static DurableConfig config( + ExecutorService executor, + JacksonSerDes serializer, + boolean plugins, + AtomicReference end) { + var builder = DurableConfig.builder() + .withExecutorService(executor) + .withDurableExecutionClient(TestUtils.createMockClient()) + .withCheckpointDelay(Duration.ZERO) + .withSerDes(serializer); + if (plugins) + builder.withPlugins(ignoredInfo -> new DurableExecutionPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + end.set(info); + } + }); + return builder.build(); + } + + private static ExecutorService inlineExecutor() { + var executor = mock(ExecutorService.class); + doAnswer(call -> { + ((Runnable) call.getArgument(0)).run(); + return null; + }) + .when(executor) + .execute(any()); + return executor; + } + + private static Context runtime(int millis) { + var runtime = mock(Context.class); + when(runtime.getRemainingTimeInMillis()).thenReturn(millis); + when(runtime.getAwsRequestId()).thenReturn("runtime-request"); + return runtime; + } + + private static DurableExecutionInput input() { + return new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/review/exec", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("exec") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationPreparationBoundaryTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationPreparationBoundaryTest.java new file mode 100644 index 0000000000..3602bcb9b5 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationPreparationBoundaryTest.java @@ -0,0 +1,240 @@ +// 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.Mockito.*; + +import com.amazonaws.services.lambda.runtime.Context; +import java.time.Instant; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +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.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.plugin.*; +import software.amazon.lambda.durable.serde.JacksonSerDes; +import software.amazon.lambda.durable.util.ExceptionHelper; + +@Timeout(20) +class InvocationPreparationBoundaryTest { + @SuppressWarnings("removal") + @ParameterizedTest + @CsvSource({ + "false,true,false,false", "true,true,false,false", + "false,true,true,false", "true,true,true,false", + "false,false,false,false", "true,false,false,false", + "false,false,true,false", "true,false,true,false", + "false,true,false,true", "true,true,false,true" + }) + void capturedInputFatalSurvivesAdmissionCutoff(boolean plugins, boolean expired, boolean death, boolean wrapped) { + Error fatal = death ? new ThreadDeath() : new InternalError("input deserializer fatal"); + Throwable inputFailure = wrapped ? new CompletionException(fatal) : fatal; + var reads = new AtomicInteger(); + var bodies = new AtomicInteger(); + var starts = new AtomicInteger(); + var ends = new AtomicInteger(); + var end = new AtomicReference(); + var executor = mock(ExecutorService.class); + doAnswer(call -> { + ((Runnable) call.getArgument(0)).run(); + return null; + }) + .when(executor) + .execute(any()); + var serializer = new JacksonSerDes() { + @Override + public T deserialize(String data, TypeToken type) { + reads.incrementAndGet(); + if (expired) { + try { + assertFalse(new CountDownLatch(1).await(1600, TimeUnit.MILLISECONDS)); + } catch (InterruptedException interrupted) { + throw new AssertionError(interrupted); + } + } + ExceptionHelper.sneakyThrow(inputFailure); + return null; + } + }; + var client = TestUtils.createMockClient(); + var builder = DurableConfig.builder() + .withExecutorService(executor) + .withDurableExecutionClient(client) + .withSerDes(serializer); + if (plugins) + builder.withPlugins(info -> new DurableExecutionPlugin() { + public void onInvocationStart(InvocationInfo info) { + starts.incrementAndGet(); + assertNull(info.executionInput()); + } + + public void onInvocationEnd(InvocationEndInfo info) { + ends.incrementAndGet(); + end.set(info); + } + }); + var observed = assertThrows( + Throwable.class, + () -> DurableExecutor.execute( + input(), + runtime(expired ? 2000 : 10_000), + TypeToken.get(String.class), + (value, context) -> { + bodies.incrementAndGet(); + return value; + }, + builder.build())); + assertSame(fatal, observed, "A captured input JVM fatal must outrank a later admission rejection"); + assertEquals(1, reads.get()); + assertEquals(0, bodies.get()); + assertEquals(plugins && !expired ? 1 : 0, starts.get(), "Expired input cannot start new plugin work"); + assertEquals(starts.get(), ends.get()); + if (end.get() != null) { + assertEquals(InvocationStatus.RETRYING, end.get().invocationStatus()); + assertSame(fatal, end.get().executionError()); + } + verify(client, never()).checkpoint(any(), any(), any()); + } + + @ParameterizedTest + @CsvSource({"false,false", "false,true", "true,false", "true,true"}) + void terminalEndReflectsCancellationDuringNormalSerialization(boolean handlerFails, boolean cancel) + throws Exception { + var serializationEntered = new CountDownLatch(1); + var serializationInterrupted = new CountDownLatch(1); + var releaseSerialization = new CountDownLatch(1); + var endCount = new AtomicInteger(); + var end = new AtomicReference(); + var owner = new AtomicReference(); + var caller = new AtomicReference(); + var callerInterrupted = new AtomicBoolean(); + var applicationFailure = new IllegalStateException("handler failed"); + var workers = Executors.newSingleThreadExecutor(); + var callers = Executors.newSingleThreadExecutor(); + var serializer = new JacksonSerDes() { + @Override + public String serialize(Object value) { + if ("result".equals(value) || value == applicationFailure) { + assertSame(owner.get(), Thread.currentThread()); + serializationEntered.countDown(); + awaitIgnoringInterruptions(releaseSerialization, serializationInterrupted); + } + return super.serialize(value); + } + }; + var client = TestUtils.createMockClient(); + var config = DurableConfig.builder() + .withExecutorService(workers) + .withDurableExecutionClient(client) + .withSerDes(serializer) + .withPlugins(info -> new DurableExecutionPlugin() { + public void onInvocationStart(InvocationInfo info) { + owner.set(Thread.currentThread()); + } + + public void onInvocationEnd(InvocationEndInfo info) { + assertSame(owner.get(), Thread.currentThread()); + endCount.incrementAndGet(); + end.set(info); + } + }) + .build(); + try { + var response = callers.submit(() -> { + caller.set(Thread.currentThread()); + try { + return DurableExecutor.execute( + input(), + runtime(10_000), + TypeToken.get(String.class), + (value, context) -> { + if (handlerFails) throw applicationFailure; + return "result"; + }, + config); + } finally { + callerInterrupted.set(Thread.interrupted()); + } + }); + assertTrue(serializationEntered.await(5, TimeUnit.SECONDS)); + if (cancel) { + caller.get().interrupt(); + assertTrue(serializationInterrupted.await(5, TimeUnit.SECONDS)); + assertThrows(TimeoutException.class, () -> response.get(100, TimeUnit.MILLISECONDS)); + assertEquals(0, endCount.get(), "End must await the actual serializer owner"); + } + releaseSerialization.countDown(); + if (cancel) { + var failure = assertInstanceOf( + UnrecoverableDurableExecutionException.class, + assertThrows(ExecutionException.class, () -> response.get(5, TimeUnit.SECONDS)) + .getCause()); + assertTrue(failure.isRetryable()); + assertEquals(InvocationStatus.RETRYING, end.get().invocationStatus()); + assertSame(failure, end.get().executionError()); + assertNull(end.get().executionResult()); + } else { + assertEquals( + handlerFails ? ExecutionStatus.FAILED : ExecutionStatus.SUCCEEDED, + response.get(5, TimeUnit.SECONDS).status()); + assertEquals( + handlerFails ? InvocationStatus.FAILED : InvocationStatus.SUCCEEDED, + end.get().invocationStatus()); + assertSame(handlerFails ? applicationFailure : null, end.get().executionError()); + } + assertEquals(1, endCount.get()); + assertEquals(cancel, callerInterrupted.get()); + verify(client, never()).checkpoint(any(), any(), any()); + } finally { + releaseSerialization.countDown(); + workers.shutdownNow(); + callers.shutdownNow(); + assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(callers.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private static void awaitIgnoringInterruptions(CountDownLatch release, CountDownLatch interrupted) { + var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (true) { + try { + assertTrue(release.await(Math.max(0, deadline - System.nanoTime()), TimeUnit.NANOSECONDS)); + return; + } catch (InterruptedException ignored) { + interrupted.countDown(); + } + } + } + + private static Context runtime(int millis) { + var context = mock(Context.class); + when(context.getRemainingTimeInMillis()).thenReturn(millis); + when(context.getAwsRequestId()).thenReturn("preparation-request"); + return context; + } + + private static DurableExecutionInput input() { + return new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/review/exec", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("exec") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnerDrainLifecycleFailureTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnerDrainLifecycleFailureTest.java new file mode 100644 index 0000000000..2582eb24d0 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnerDrainLifecycleFailureTest.java @@ -0,0 +1,319 @@ +// 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.Mockito.*; + +import com.amazonaws.services.lambda.runtime.Context; +import java.time.Instant; +import java.util.*; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import java.util.function.BooleanSupplier; +import java.util.function.Supplier; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +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.context.DurableContextImpl; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.plugin.*; + +@Timeout(20) +class OwnerDrainLifecycleFailureTest { + enum Stop { + CALLER, + EXPIRY, + QUEUED, + QUEUED_OWNER, + OWNER, + NORMAL, + ALIAS, + TRANSPORT, + FATAL + } + + @ParameterizedTest + @CsvSource({ + "CALLER,true", + "CALLER,false", + "EXPIRY,true", + "QUEUED,true", + "QUEUED_OWNER,true", + "OWNER,true", + "NORMAL,true", + "ALIAS,true", + "TRANSPORT,true", + "FATAL,true" + }) + void ownerDrainKeepsTheRightFailureAndActualExitBoundary(Stop stop, boolean plugins) throws Exception { + var rootOwner = new AtomicReference(); + var runtimeCaller = new AtomicReference(); + var manager = new AtomicReference(); + var end = new AtomicReference(); + var endCalls = new AtomicInteger(); + var stepBodies = new AtomicInteger(); + var stepExited = new AtomicBoolean(); + var stepEntered = new CountDownLatch(1); + var stepInterrupted = new CountDownLatch(1); + var releaseStep = new CountDownLatch(1); + var rootReturning = new CountDownLatch(1); + var ignoredRootInterrupt = new CountDownLatch(1); + var uncaught = new AtomicReference(); + var fatalEscaped = new CountDownLatch(1); + var callerInterrupted = new AtomicBoolean(); + var suppressRootInterrupt = new AtomicBoolean(stop == Stop.QUEUED || stop == Stop.TRANSPORT); + var queued = stop == Stop.QUEUED || stop == Stop.QUEUED_OWNER; + var transport = new IllegalStateException( + "transport cleanup interrupted", new InterruptedException("transport interruption")); + var fatal = new InternalError("fatal transport cleanup"); + ThreadFactory threadFactory = command -> { + var thread = new Thread(command, "owner-drain-worker") { + @Override + public void interrupt() { + if (suppressRootInterrupt.get() && this == rootOwner.get()) { + ignoredRootInterrupt.countDown(); + return; + } + super.interrupt(); + } + }; + thread.setUncaughtExceptionHandler((owner, failure) -> { + uncaught.set(failure); + fatalEscaped.countDown(); + }); + return thread; + }; + var users = queued + ? new ThreadPoolExecutor( + 1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue(), threadFactory) + : new ThreadPoolExecutor( + 0, Integer.MAX_VALUE, 60, TimeUnit.SECONDS, new SynchronousQueue(), threadFactory); + var callers = Executors.newSingleThreadExecutor(); + var builder = DurableConfig.builder() + .withExecutorService(users) + .withDurableExecutionClient(TestUtils.createMockClient()); + if (plugins) + builder.withPlugins(info -> new DurableExecutionPlugin() { + public void onInvocationEnd(InvocationEndInfo info) { + assertSame(rootOwner.get(), Thread.currentThread()); + if (!queued) assertTrue(stepExited.get(), "End cannot precede actual step exit"); + end.set(info); + endCalls.incrementAndGet(); + } + }); + var runtime = mock(Context.class); + when(runtime.getRemainingTimeInMillis()).thenReturn(stop == Stop.EXPIRY ? 2500 : 20_000); + var input = new DurableExecutionInput( + "arn:execution/test", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("test") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + try { + var response = callers.submit(() -> { + runtimeCaller.set(Thread.currentThread()); + try { + return DurableExecutor.execute( + input, + runtime, + TypeToken.get(String.class), + (value, root) -> { + rootOwner.set(Thread.currentThread()); + var execution = ((DurableContextImpl) root).getExecutionManager(); + manager.set(execution); + if (stop == Stop.ALIAS || stop == Stop.TRANSPORT || stop == Stop.FATAL) + installShutdownFault( + execution, + rootOwner, + () -> stop == Stop.ALIAS + ? new CompletionException(execution.getLifecycleFailure()) + : stop == Stop.FATAL ? fatal : transport); + root.stepAsync("held", String.class, step -> { + stepBodies.incrementAndGet(); + stepEntered.countDown(); + try { + awaitIgnoringInterrupt(releaseStep, stepInterrupted); + return "step"; + } finally { + stepExited.set(true); + } + }); + if (!queued) await(stepEntered); + rootReturning.countDown(); + return "root-result"; + }, + builder.build()); + } finally { + callerInterrupted.set(Thread.interrupted()); + } + }); + await(rootReturning); + awaitCondition(() -> inMethod(rootOwner.get(), "awaitUserHandler")); + if (stop == Stop.NORMAL) { + assertEquals(0, endCalls.get()); + releaseStep.countDown(); + assertEquals( + ExecutionStatus.SUCCEEDED, + response.get(5, TimeUnit.SECONDS).status()); + assertEquals(InvocationStatus.SUCCEEDED, end.get().invocationStatus()); + assertNull(end.get().executionError()); + } else { + if (stop == Stop.OWNER) { + rootOwner.get().interrupt(); + awaitCondition( + () -> !rootOwner.get().isInterrupted() && inMethod(rootOwner.get(), "awaitUserHandler")); + } else if (stop == Stop.QUEUED_OWNER) { + var stepTask = manager.get().getActiveExecutorTasks().stream() + .filter(task -> task.role() == ExecutorTaskHandle.Role.STEP) + .findFirst() + .orElseThrow(); + assertTrue(stepTask.cancel(false)); + } else { + if (stop != Stop.EXPIRY) runtimeCaller.get().interrupt(); + if (stop == Stop.QUEUED) await(ignoredRootInterrupt); + else await(stepInterrupted); + } + if (!queued) { + assertThrows(TimeoutException.class, () -> response.get(100, TimeUnit.MILLISECONDS)); + assertEquals(0, endCalls.get()); + releaseStep.countDown(); + } + var observed = assertThrows(ExecutionException.class, () -> response.get(5, TimeUnit.SECONDS)) + .getCause(); + var lifecycle = manager.get().getLifecycleFailure(); + if (stop == Stop.FATAL) { + assertSame(fatal, observed); + assertSame(fatal, end.get().executionError()); + await(fatalEscaped); + assertSame(fatal, uncaught.get()); + } else if (stop == Stop.TRANSPORT) { + assertSame(lifecycle, observed); + assertSame( + transport, + end.get().executionError(), + "An unrelated transport failure is not a synthetic owner cancellation"); + } else if (stop == Stop.OWNER || stop == Stop.QUEUED_OWNER) { + assertNull(lifecycle); + assertEquals(IllegalStateException.class, observed.getClass()); + assertSame(observed, end.get().executionError()); + assertTrue( + stop == Stop.OWNER + ? observed.getCause() instanceof InterruptedException + : observed.getCause() instanceof CancellationException); + } else { + assertNotNull(lifecycle); + assertSame(lifecycle, observed); + if (plugins) + assertSame( + lifecycle, end.get().executionError(), "End must retain the original lifecycle cause"); + var diagnostic = Arrays.stream(lifecycle.getSuppressed()) + .filter(failure -> failure.getClass() == IllegalStateException.class + && (failure.getCause() instanceof InterruptedException + || failure.getCause() instanceof CancellationException)) + .toList(); + assertEquals(1, diagnostic.size(), "Keep the original synthetic drain error exactly once"); + assertFalse( + references(diagnostic.get(0), lifecycle), + "The diagnostic must not reference its canonical lifecycle owner"); + } + if (plugins) assertEquals(InvocationStatus.RETRYING, end.get().invocationStatus()); + } + assertEquals(plugins ? 1 : 0, endCalls.get()); + assertEquals(queued ? 0 : 1, stepBodies.get()); + assertEquals( + stop != Stop.NORMAL && stop != Stop.OWNER && stop != Stop.QUEUED_OWNER && stop != Stop.EXPIRY, + callerInterrupted.get()); + } finally { + releaseStep.countDown(); + suppressRootInterrupt.set(false); + users.shutdownNow(); + callers.shutdownNow(); + assertTrue(users.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(callers.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private static void installShutdownFault( + ExecutionManager manager, AtomicReference owner, Supplier fault) { + try { + var field = ExecutionManager.class.getDeclaredField("checkpointManager"); + field.setAccessible(true); + var checkpoint = spy((CheckpointManager) field.get(manager)); + doAnswer(call -> { + if (Thread.currentThread() != owner.get()) return call.callRealMethod(); + try { + call.callRealMethod(); + } finally { + throw fault.get(); + } + }) + .when(checkpoint) + .shutdown(); + field.set(manager, checkpoint); + } catch (ReflectiveOperationException failure) { + throw new AssertionError(failure); + } + } + + private static boolean references(Throwable root, Throwable target) { + var pending = new ArrayDeque(); + var seen = Collections.newSetFromMap(new IdentityHashMap()); + pending.add(root); + while (!pending.isEmpty()) { + var next = pending.remove(); + if (next == target) return true; + if (!seen.add(next)) continue; + if (next.getCause() != null) pending.add(next.getCause()); + pending.addAll(Arrays.asList(next.getSuppressed())); + } + return false; + } + + private static boolean inMethod(Thread thread, String name) { + return thread != null + && Arrays.stream(thread.getStackTrace()) + .anyMatch(frame -> frame.getMethodName().equals(name)); + } + + private static void awaitCondition(BooleanSupplier condition) throws InterruptedException { + var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (!condition.getAsBoolean()) { + assertTrue(System.nanoTime() < deadline); + TimeUnit.MILLISECONDS.sleep(5); + } + } + + private static void await(CountDownLatch gate) { + try { + assertTrue(gate.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException failure) { + throw new AssertionError(failure); + } + } + + private static void awaitIgnoringInterrupt(CountDownLatch gate, CountDownLatch interrupted) { + var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10); + while (true) { + try { + assertTrue(gate.await(Math.max(0, deadline - System.nanoTime()), TimeUnit.NANOSECONDS)); + return; + } catch (InterruptedException ignored) { + interrupted.countDown(); + } + } + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/TaskAdmissionFenceTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/TaskAdmissionFenceTest.java new file mode 100644 index 0000000000..caa545c930 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/TaskAdmissionFenceTest.java @@ -0,0 +1,126 @@ +// 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.Mockito.*; + +import java.time.Instant; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +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.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.operation.BaseDurableOperation; + +@Timeout(15) +class TaskAdmissionFenceTest { + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void registeredOperationCannotDispatchAfterLifecycleFence(boolean duringRegistration) throws Exception { + var actionCalls = new AtomicInteger(); + var dispatches = new AtomicInteger(); + var lookups = new AtomicInteger(); + var submitter = new AtomicReference(); + var lookupEntered = new CountDownLatch(1); + var releaseLookup = new CountDownLatch(1); + var users = new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue()) { + @Override + public void execute(Runnable action) { + dispatches.incrementAndGet(); + super.execute(action); + } + }; + var callers = Executors.newSingleThreadExecutor(); + var input = new DurableExecutionInput( + "arn:execution/test", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("test") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .build()) + .build()); + var manager = new ExecutionManager( + input, + DurableConfig.builder() + .withExecutorService(users) + .withDurableExecutionClient(TestUtils.createMockClient()) + .build(), + null); + var operation = mock(BaseDurableOperation.class); + when(operation.getOperationId()).thenAnswer(ignored -> { + if (duringRegistration && Thread.currentThread() == submitter.get() && lookups.incrementAndGet() == 2) { + lookupEntered.countDown(); + try { + assertTrue(releaseLookup.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException failure) { + throw new AssertionError(failure); + } + } + return "registered-step"; + }); + when(operation.getType()).thenReturn(OperationType.STEP); + when(operation.getCompletionFuture()).thenReturn(new CompletableFuture<>()); + manager.registerOperation(operation); + manager.beginDraining(); + Future submission = null; + UnrecoverableDurableExecutionException fence = null; + try { + if (!duringRegistration) fence = selectInterruption(manager); + submission = callers.submit(() -> { + submitter.set(Thread.currentThread()); + manager.submitOperationTask(operation, actionCalls::incrementAndGet) + .join(); + }); + if (duringRegistration) { + assertTrue(lookupEntered.await(5, TimeUnit.SECONDS)); + fence = selectInterruption(manager); + releaseLookup.countDown(); + } + var observed = submission; + var failure = assertThrows(ExecutionException.class, () -> observed.get(5, TimeUnit.SECONDS)); + assertSame(fence, failure.getCause()); + assertEquals(0, dispatches.get(), "A task fenced during registration must not be offered to the executor"); + assertEquals(0, actionCalls.get()); + assertTrue( + manager.getActiveExecutorTasks().isEmpty(), + "Rejected registration must release actual-exit tracking"); + } finally { + releaseLookup.countDown(); + if (submission != null) { + try { + submission.get(5, TimeUnit.SECONDS); + } catch (ExecutionException expected) { + } + } + try { + manager.close(); + } catch (UnrecoverableDurableExecutionException expected) { + } finally { + Thread.interrupted(); + } + users.shutdownNow(); + callers.shutdownNow(); + assertTrue(users.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(callers.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private static UnrecoverableDurableExecutionException selectInterruption(ExecutionManager manager) { + Thread.currentThread().interrupt(); + var selected = + manager.awaitInvocationOutcome(new CompletableFuture()).failure(); + assertFalse( + Thread.currentThread().isInterrupted(), + "The blocking observation consumes the interrupt before cleanup"); + return assertInstanceOf(UnrecoverableDurableExecutionException.class, selected); + } +}