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); + } +}