From 527d8be979113b7313d32b3d38366be7448ea7f8 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Oct 2026 00:37:11 +0000 Subject: [PATCH 01/48] fix(otel): scope root handler instrumentation to its execution --- otel-plugin/README.md | 13 ++ .../durable/otel/ExecutionOtelPlugin.java | 11 + .../durable/otel/InvocationOtelPlugin.java | 11 + .../otel/HandlerContextIntegrationTest.java | 219 ++++++++++++++++++ .../durable/execution/DurableExecutor.java | 3 +- .../plugin/DurableExecutionPlugin.java | 11 + .../lambda/durable/plugin/PluginRunner.java | 26 +++ .../durable/plugin/HandlerScopeTest.java | 78 +++++++ 8 files changed, 371 insertions(+), 1 deletion(-) create mode 100644 otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerContextIntegrationTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeTest.java diff --git a/otel-plugin/README.md b/otel-plugin/README.md index c51032aaa..e88c95499 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -13,6 +13,19 @@ OpenTelemetry instrumentation plugin for the AWS Lambda Durable Execution SDK fo - **ADOT Java Agent Integration**: `new InvocationOtelPlugin()` late-binds the ADOT Java agent's global provider with no handler-side OpenTelemetry initialization - **Lambda Layer Discovery**: `DURABLE_EXECUTION_PLUGINS` loads either OTel plugin from a JAR under a layer's `java/lib` directory +## Root handler context + +With a core that supports `DurableExecutionPlugin.openHandlerScope()`, user instrumentation in the root handler joins +its canonical execution trace. A valid same-trace ambient Lambda span stays current. If ambient context is absent or +belongs to another trace, invocation view activates `Invocation`; execution view activates the deterministic +`Workflow` context, including its unsampled non-recording form. Nested operations retain their existing contexts. + +The core opens the scope after invocation startup and closes it on the same handler thread on success, failure, or +suspension. Invocation-end hooks can execute on a different thread and do not own this scope. The additive hook uses +only a JDK type and defaults to no scope: old plugins on a new core and new plugin layers on an older core retain their +existing behavior. Root-handler fallback activation requires both the updated core and plugin; no provider API or +dependency floor changes are required. + ## Installation ```xml diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java index 09214055c..625cda26e 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java @@ -271,6 +271,17 @@ public void onInvocationStart(InvocationInfo info) { tracingEnabled = true; } + @Override + public AutoCloseable openHandlerScope() { + var trace = executionTrace; + if (!tracingEnabled || trace == null) return null; + var ambient = Span.current().getSpanContext(); + // Preserve a compatible ambient Lambda span. An absent or unrelated ambient span must not leave + // handler instrumentation outside the durable execution's canonical trace. + if (ambient.isValid() && trace.traceId().equals(ambient.getTraceId())) return null; + return Span.wrap(workflowSpanContext).makeCurrent(); + } + @Override public void onInvocationEnd(InvocationEndInfo info) { if (!tracingEnabled) { diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java index e89804bdc..de3b1c9a1 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java @@ -268,6 +268,17 @@ public void onInvocationStart(InvocationInfo info) { tracingEnabled = true; } + @Override + public AutoCloseable openHandlerScope() { + var trace = executionTrace; + if (!tracingEnabled || trace == null) return null; + var ambient = Span.current().getSpanContext(); + // Preserve a compatible ambient Lambda span. An absent or unrelated ambient span must not leave + // handler instrumentation outside the durable execution's canonical trace. + if (ambient.isValid() && trace.traceId().equals(ambient.getTraceId())) return null; + return invocationSpan.makeCurrent(); + } + @Override public void onInvocationEnd(InvocationEndInfo info) { if (!tracingEnabled) { diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerContextIntegrationTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerContextIntegrationTest.java new file mode 100644 index 000000000..e7665a3c2 --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerContextIntegrationTest.java @@ -0,0 +1,219 @@ +// 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 io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.api.trace.TraceFlags; +import io.opentelemetry.api.trace.TraceState; +import io.opentelemetry.context.Context; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; +import software.amazon.lambda.durable.plugin.InvocationInfo; +import software.amazon.lambda.durable.plugin.PluginRunner; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class HandlerContextIntegrationTest { + private static final String TRACE_ID = "12345678901234567890123456789012"; + + @ParameterizedTest + @CsvSource({"true,success", "false,success", "true,failure", "false,failure", "true,suspension", "false,suspension" + }) + void reusedHandlerWorkerHasNoLeakedScope(boolean executionView, String outcome) throws Exception { + var exporter = InMemorySpanExporter.create(); + var builder = SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)); + var pluginConfig = OtelPluginConfig.builder() + .enableMdc(false) + .contextExtractor( + () -> new ExtractedContext(TRACE_ID, "1234567890123456", ExtractedContext.Sampling.SAMPLED)) + .build(); + DurableExecutionPlugin plugin = executionView + ? new ExecutionOtelPlugin(builder, pluginConfig) + : new InvocationOtelPlugin(builder, pluginConfig); + var executor = Executors.newSingleThreadExecutor(); + var config = DurableConfig.builder() + .withExecutorService(executor) + .withPlugins(plugin) + .build(); + try { + var runner = LocalDurableTestRunner.create( + String.class, + (input, ctx) -> { + assertEquals(TRACE_ID, Span.current().getSpanContext().getTraceId()); + if (outcome.equals("failure")) throw new IllegalStateException("user failure"); + if (outcome.equals("suspension")) ctx.wait("resume", Duration.ofSeconds(1)); + return "done"; + }, + config); + var result = runner.run("input"); + assertEquals( + switch (outcome) { + case "failure" -> ExecutionStatus.FAILED; + case "suspension" -> ExecutionStatus.PENDING; + default -> ExecutionStatus.SUCCEEDED; + }, + result.getStatus()); + assertFalse( + executor.submit(() -> Span.current().getSpanContext().isValid()) + .get(5, TimeUnit.SECONDS), + "cleanup must run on the worker, even if invocation finalization ran on another thread"); + } finally { + executor.shutdownNow(); + } + } + + @ParameterizedTest + @CsvSource({"true,true", "false,true", "true,false", "false,false"}) + void preservesCompatibleAmbientAndRestoresUnrelatedAmbientAfterFailure(boolean executionView, boolean sameTrace) { + var exporter = InMemorySpanExporter.create(); + var builder = SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)); + var pluginConfig = OtelPluginConfig.builder() + .enableMdc(false) + .contextExtractor( + () -> new ExtractedContext(TRACE_ID, "1234567890123456", ExtractedContext.Sampling.SAMPLED)) + .build(); + DurableExecutionPlugin plugin = executionView + ? new ExecutionOtelPlugin(builder, pluginConfig) + : new InvocationOtelPlugin(builder, pluginConfig); + var ambient = SpanContext.create( + sameTrace ? TRACE_ID : "abcdefabcdefabcdefabcdefabcdefab", + "abcdefabcdefabcd", + TraceFlags.getSampled(), + TraceState.getDefault()); + var previous = Context.current(); + try (var ignored = Span.wrap(ambient).makeCurrent()) { + plugin.onInvocationStart(new InvocationInfo( + "req", + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/name/id", + true, + Instant.now())); + var error = new IllegalStateException("handler failure"); + var runner = new PluginRunner(List.of(plugin)); + assertSame( + error, + assertThrows( + IllegalStateException.class, + () -> runner.runHandler(() -> { + var active = Span.current().getSpanContext(); + assertEquals(TRACE_ID, active.getTraceId()); + if (sameTrace) assertEquals(ambient, active); + else assertNotEquals(ambient.getSpanId(), active.getSpanId()); + throw error; + }))); + assertEquals(ambient, Span.current().getSpanContext()); + } + assertSame(previous, Context.current()); + } + + @ParameterizedTest + @CsvSource({"true,true", "false,true", "true,false", "false,false"}) + void rootContextIsValidAndRestoredAcrossNestedWorkAndResume(boolean executionView, boolean sampled) { + var exporter = InMemorySpanExporter.create(); + var builder = SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)); + var pluginConfig = OtelPluginConfig.builder() + .enableMdc(false) + .contextExtractor(() -> new ExtractedContext( + TRACE_ID, + "1234567890123456", + sampled ? ExtractedContext.Sampling.SAMPLED : ExtractedContext.Sampling.NOT_SAMPLED)) + .build(); + DurableExecutionPlugin plugin = executionView + ? new ExecutionOtelPlugin(builder, pluginConfig) + : new InvocationOtelPlugin(builder, pluginConfig); + var executor = Executors.newCachedThreadPool(); + var config = DurableConfig.builder() + .withExecutorService(executor) + .withPlugins(plugin) + .build(); + var userProvider = SdkTracerProvider.builder() + .addSpanProcessor(SimpleSpanProcessor.create(exporter)) + .build(); + var userTracer = userProvider.get("user"); + var roots = new ArrayList(); + var bodyCalls = new AtomicInteger(); + try { + var runner = LocalDurableTestRunner.create( + String.class, + (input, ctx) -> { + var root = Span.current().getSpanContext(); + assertTrue(root.isValid(), "root handler must have a valid context without an agent"); + assertEquals(TRACE_ID, root.getTraceId()); + assertEquals(sampled, root.isSampled()); + roots.add(root); + userTracer.spanBuilder("user-handler").startSpan().end(); + ctx.step("before", String.class, step -> { + bodyCalls.incrementAndGet(); + assertNotEquals( + root.getSpanId(), + Span.current().getSpanContext().getSpanId()); + return "before"; + }); + assertEquals(root, Span.current().getSpanContext()); + ctx.runInChildContext("child", String.class, child -> { + assertNotEquals( + root.getSpanId(), + Span.current().getSpanContext().getSpanId()); + return "child"; + }); + assertEquals(root, Span.current().getSpanContext()); + userTracer + .spanBuilder("user-handler-restored") + .startSpan() + .end(); + ctx.wait("resume", Duration.ofSeconds(1)); + assertEquals(root, Span.current().getSpanContext()); + userTracer + .spanBuilder("user-handler-after-resume") + .startSpan() + .end(); + return "done"; + }, + config); + var first = runner.run("input"); + assertEquals(ExecutionStatus.PENDING, first.getStatus()); + runner.advanceTime(); + var last = runner.run("input"); + assertEquals(ExecutionStatus.SUCCEEDED, last.getStatus()); + assertEquals(2, roots.size()); + assertEquals(1, bodyCalls.get()); + assertFalse(Span.current().getSpanContext().isValid()); + if (sampled) { + var rootName = executionView ? "Workflow" : "Invocation"; + var spans = exporter.getFinishedSpanItems(); + var userSpans = spans.stream() + .filter(span -> span.getName().startsWith("user-handler")) + .toList(); + assertEquals(5, userSpans.size()); + for (var span : userSpans) { + assertEquals(TRACE_ID, span.getTraceId()); + assertTrue(roots.stream().anyMatch(root -> root.getSpanId().equals(span.getParentSpanId()))); + } + for (var root : roots) { + assertTrue(spans.stream() + .anyMatch(s -> rootName.equals(s.getName()) + && root.getSpanId().equals(s.getSpanId()))); + } + } else { + assertTrue(exporter.getFinishedSpanItems().isEmpty()); + } + } finally { + executor.shutdownNow(); + userProvider.close(); + } + } +} 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 d8db91326..5e8d1319e 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 @@ -109,7 +109,8 @@ public static DurableExecutionOutput execute( DurableContextImpl.setCurrentContext(context); // use a try-with-resources to clear logger properties try (var ignored = DurableLogger.attachContext()) { - return handler.apply(userInput, context); + var handlerInput = userInput; + return pluginRunner.runHandler(() -> handler.apply(handlerInput, context)); } }, config.getExecutorService()); // Get executor from config for running user code diff --git a/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java b/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java index e2f8a46df..fca92fe90 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java @@ -36,6 +36,17 @@ default void onInvocationStart(InvocationInfo info) {} */ default void onInvocationEnd(InvocationEndInfo info) {} + /** + * Opens an optional scope around the root handler body, after invocation startup. + * + *

The SDK opens and closes this scope on the handler thread, including when the handler fails or suspends. + * Scopes close in reverse plugin order when the handler exits. Return {@code null} for no scope. This additive + * capability requires a core that calls it; older cores retain their original hook behavior. + */ + default AutoCloseable openHandlerScope() { + return null; + } + // ─── Operation-level hooks ─────────────────────────────────────────── /** Called when an operation starts (including replay). Use for logging/metrics that want replay visibility. */ diff --git a/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java b/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java index e3a5707c4..b0ec6439a 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java @@ -2,9 +2,11 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.plugin; +import java.util.ArrayDeque; import java.util.Collections; import java.util.List; import java.util.function.Consumer; +import java.util.function.Supplier; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -55,6 +57,30 @@ private void run(Consumer hook) { } } + /** Runs the root handler with optional plugin scopes, closing them on the same thread in reverse order. */ + public T runHandler(Supplier handler) { + var scopes = new ArrayDeque(); + try { + for (var plugin : plugins) { + try { + var scope = plugin.openHandlerScope(); + if (scope != null) scopes.push(scope); + } catch (Exception e) { + logger.warn("Plugin handler scope threw exception", e); + } + } + return handler.get(); + } finally { + while (!scopes.isEmpty()) { + try { + scopes.pop().close(); + } catch (Exception e) { + logger.warn("Plugin handler scope cleanup threw exception", e); + } + } + } + } + public void onInvocationStart(InvocationInfo info) { run(p -> p.onInvocationStart(info)); } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeTest.java b/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeTest.java new file mode 100644 index 000000000..2ee586a57 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeTest.java @@ -0,0 +1,78 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.plugin; + +import static org.junit.jupiter.api.Assertions.*; + +import java.util.ArrayList; +import java.util.List; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import software.amazon.lambda.durable.execution.SuspendExecutionException; +import software.amazon.lambda.durable.util.ExceptionHelper; + +class HandlerScopeTest { + @ParameterizedTest + @ValueSource(strings = {"success", "failure", "suspension"}) + void scopesCloseInReverseOrderOnTheHandlerThread(String outcome) { + var calls = new ArrayList(); + var owner = Thread.currentThread(); + var runner = new PluginRunner(List.of(scope("a", calls, owner), scope("b", calls, owner))); + Throwable error = + outcome.equals("suspension") ? new SuspendExecutionException() : new IllegalStateException("user"); + if (outcome.equals("success")) { + assertEquals("ok", runner.runHandler(() -> { + calls.add("handler"); + return "ok"; + })); + } else { + assertSame( + error, + assertThrows( + Throwable.class, + () -> runner.runHandler(() -> { + calls.add("handler"); + ExceptionHelper.sneakyThrow(error); + return null; + }))); + } + assertEquals(List.of("open-a", "open-b", "handler", "close-b", "close-a"), calls); + } + + @Test + void ordinarySetupAndCleanupFailuresDoNotReplaceTheHandlerResult() { + var calls = new ArrayList(); + var brokenSetup = new DurableExecutionPlugin() { + @Override + public AutoCloseable openHandlerScope() { + throw new IllegalStateException("setup"); + } + }; + var brokenClose = new DurableExecutionPlugin() { + @Override + public AutoCloseable openHandlerScope() { + return () -> { + throw new IllegalStateException("cleanup"); + }; + } + }; + var runner = new PluginRunner(List.of(scope("first", calls, Thread.currentThread()), brokenSetup, brokenClose)); + assertEquals("ok", runner.runHandler(() -> "ok")); + assertEquals(List.of("open-first", "close-first"), calls); + } + + private static DurableExecutionPlugin scope(String name, List calls, Thread owner) { + return new DurableExecutionPlugin() { + @Override + public AutoCloseable openHandlerScope() { + assertSame(owner, Thread.currentThread()); + calls.add("open-" + name); + return () -> { + assertSame(owner, Thread.currentThread()); + calls.add("close-" + name); + }; + } + }; + } +} From 06727e76b49949909ab59b45ab2444920770e8a7 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Oct 2026 01:37:44 +0000 Subject: [PATCH 02/48] fix(otel): finalize handler scopes before invocation response --- otel-plugin/README.md | 9 +- .../durable/otel/ExecutionOtelPlugin.java | 2 +- .../durable/otel/InvocationOtelPlugin.java | 2 +- ...ndlerScopeFinalizationIntegrationTest.java | 111 ++++++++++++++++++ .../durable/execution/DurableExecutor.java | 32 ++++- .../plugin/DurableExecutionPlugin.java | 6 +- .../lambda/durable/plugin/PluginRunner.java | 28 +++-- .../durable/plugin/HandlerScopeTest.java | 73 ++++++++++++ 8 files changed, 246 insertions(+), 17 deletions(-) create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/HandlerScopeFinalizationIntegrationTest.java diff --git a/otel-plugin/README.md b/otel-plugin/README.md index e88c95499..f4734a226 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -21,10 +21,15 @@ belongs to another trace, invocation view activates `Invocation`; execution view `Workflow` context, including its unsampled non-recording form. Nested operations retain their existing contexts. The core opens the scope after invocation startup and closes it on the same handler thread on success, failure, or -suspension. Invocation-end hooks can execute on a different thread and do not own this scope. The additive hook uses +suspension. When a handler scope is opened, suspension/termination finalization waits asynchronously for the handler to unwind +and close it, before invocation-end hooks flush telemetry or a response is returned. A compatible ambient span uses +a no-op scope so it receives the same finalization ordering. Invocation-end hooks can execute on a different thread +and do not own this scope. The additive hook uses only a JDK type and defaults to no scope: old plugins on a new core and new plugin layers on an older core retain their existing behavior. Root-handler fallback activation requires both the updated core and plugin; no provider API or -dependency floor changes are required. +dependency floor changes are required. The new scope boundary isolates ordinary exceptions and nonfatal linkage +errors during open and close, continues earlier scope cleanup, and preserves the handler outcome. JVM fatal errors +remain outside that containment. ## Installation diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java index 625cda26e..0982b16bd 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java @@ -278,7 +278,7 @@ public AutoCloseable openHandlerScope() { var ambient = Span.current().getSpanContext(); // Preserve a compatible ambient Lambda span. An absent or unrelated ambient span must not leave // handler instrumentation outside the durable execution's canonical trace. - if (ambient.isValid() && trace.traceId().equals(ambient.getTraceId())) return null; + if (ambient.isValid() && trace.traceId().equals(ambient.getTraceId())) return Scope.noop(); return Span.wrap(workflowSpanContext).makeCurrent(); } diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java index de3b1c9a1..1531e0163 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java @@ -275,7 +275,7 @@ public AutoCloseable openHandlerScope() { var ambient = Span.current().getSpanContext(); // Preserve a compatible ambient Lambda span. An absent or unrelated ambient span must not leave // handler instrumentation outside the durable execution's canonical trace. - if (ambient.isValid() && trace.traceId().equals(ambient.getTraceId())) return null; + if (ambient.isValid() && trace.traceId().equals(ambient.getTraceId())) return Scope.noop(); return invocationSpan.makeCurrent(); } diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/HandlerScopeFinalizationIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/HandlerScopeFinalizationIntegrationTest.java new file mode 100644 index 000000000..1974bad3e --- /dev/null +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/HandlerScopeFinalizationIntegrationTest.java @@ -0,0 +1,111 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.*; + +import java.time.Duration; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import software.amazon.awssdk.services.lambda.model.ErrorObject; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +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.testing.LocalDurableTestRunner; + +class HandlerScopeFinalizationIntegrationTest { + @ParameterizedTest + @CsvSource({"true,false", "false,false", "true,true", "false,true"}) + void finalizationWaitsForAnOpenedScopeWithoutChangingLegacyNoScopeTiming(boolean hasScope, boolean terminate) + throws Exception { + var finallyEntered = new CountDownLatch(1); + var releaseFinally = new CountDownLatch(1); + var handlerExited = new CountDownLatch(1); + var scopeClosed = new AtomicBoolean(); + var endCalled = new AtomicBoolean(); + var scopeClosedAtEnd = new AtomicBoolean(); + var plugin = new DurableExecutionPlugin() { + @Override + public AutoCloseable openHandlerScope() { + if (!hasScope) return null; + var owner = Thread.currentThread(); + return () -> { + assertSame(owner, Thread.currentThread()); + scopeClosed.set(true); + }; + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + scopeClosedAtEnd.set(scopeClosed.get()); + endCalled.set(true); + } + }; + var config = DurableConfig.builder().withPlugins(plugin).build(); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + try { + if (terminate) { + context.step("terminate", String.class, step -> { + throw new UnrecoverableDurableExecutionException( + ErrorObject.builder() + .errorMessage("retry invocation") + .build(), + true); + }); + } else { + context.wait("pause", Duration.ofSeconds(1)); + } + return "done"; + } finally { + finallyEntered.countDown(); + try { + if (!releaseFinally.await(5, TimeUnit.SECONDS)) + throw new AssertionError("finally not released"); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } finally { + handlerExited.countDown(); + } + } + }, + config); + var caller = Executors.newSingleThreadExecutor(task -> { + var thread = new Thread(task, "test-invocation-caller"); + thread.setDaemon(true); + return thread; + }); + try { + var response = caller.submit(() -> runner.run("input")); + assertTrue(finallyEntered.await(5, TimeUnit.SECONDS)); + if (hasScope) { + assertThrows(TimeoutException.class, () -> response.get(150, TimeUnit.MILLISECONDS)); + assertFalse(endCalled.get(), "flush must wait for same-thread scope cleanup"); + releaseFinally.countDown(); + } + if (terminate) { + var failure = assertThrows(ExecutionException.class, () -> response.get(5, TimeUnit.SECONDS)); + assertInstanceOf(UnrecoverableDurableExecutionException.class, failure.getCause()); + } else { + assertEquals( + ExecutionStatus.PENDING, + response.get(5, TimeUnit.SECONDS).getStatus()); + } + assertTrue(endCalled.get()); + assertEquals(hasScope, scopeClosedAtEnd.get()); + } finally { + releaseFinally.countDown(); + assertTrue(handlerExited.await(5, TimeUnit.SECONDS)); + caller.shutdownNow(); + } + } +} 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 5e8d1319e..24f7e7041 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 @@ -7,6 +7,7 @@ import java.nio.charset.StandardCharsets; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; import org.slf4j.Logger; @@ -64,6 +65,7 @@ public static DurableExecutionOutput execute( executionManager.registerActiveThread(null); // Captured for onInvocationEnd, which runs outside the handler thread below. var pluginExecutionInput = new AtomicReference<>(); + var hasHandlerScope = new AtomicBoolean(); var handlerFuture = CompletableFuture.supplyAsync( () -> { executionManager.setCurrentThreadContext(new ThreadContext(null, ThreadType.CONTEXT)); @@ -110,7 +112,8 @@ public static DurableExecutionOutput execute( // use a try-with-resources to clear logger properties try (var ignored = DurableLogger.attachContext()) { var handlerInput = userInput; - return pluginRunner.runHandler(() -> handler.apply(handlerInput, context)); + return pluginRunner.runHandler( + () -> handler.apply(handlerInput, context), () -> hasHandlerScope.set(true)); } }, config.getExecutorService()); // Get executor from config for running user code @@ -119,8 +122,10 @@ public static DurableExecutionOutput execute( // will be returned. Otherwise, it will complete exceptionally with a SuspendExecutionException or a // failure. try { - return executionManager - .runUntilCompleteOrSuspend(handlerFuture) + return awaitHandlerScopes( + executionManager.runUntilCompleteOrSuspend(handlerFuture), + handlerFuture, + hasHandlerScope) .handle((result, ex) -> { if (ex != null) { // an exception thrown from handlerFuture or suspension/termination occurred @@ -200,6 +205,27 @@ public static DurableExecutionOutput execute( } } + private static CompletableFuture awaitHandlerScopes( + CompletableFuture executionFuture, CompletableFuture handlerFuture, AtomicBoolean hasHandlerScope) { + return executionFuture + .handle((result, failure) -> { + if (failure == null) return CompletableFuture.completedFuture(result); + if (!hasHandlerScope.get()) return CompletableFuture.failedFuture(failure); + // Suspension/termination can be signaled from inside the handler before its finally blocks unwind. + // A continuation releases that worker to close its scopes; joining it here would deadlock. + // Legacy plugins returning no scope retain their original finalization timing. + return handlerFuture.handle((ignored, handlerFailure) -> { + var cause = ExceptionHelper.unwrapCompletableFuture(handlerFailure); + ExceptionHelper.sneakyThrow( + cause != null && !(cause instanceof SuspendExecutionException) + ? handlerFailure + : failure); + return result; + }); + }) + .thenCompose(future -> future); + } + private static void fireOnInvocationEnd( PluginRunner pluginRunner, ExecutionManager executionManager, diff --git a/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java b/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java index fca92fe90..556a2741d 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java @@ -40,8 +40,10 @@ default void onInvocationEnd(InvocationEndInfo info) {} * Opens an optional scope around the root handler body, after invocation startup. * *

The SDK opens and closes this scope on the handler thread, including when the handler fails or suspends. - * Scopes close in reverse plugin order when the handler exits. Return {@code null} for no scope. This additive - * capability requires a core that calls it; older cores retain their original hook behavior. + * Scopes close in reverse plugin order when the handler exits. If any scope is returned, invocation finalization + * waits for the handler to unwind and close it, including after suspension/termination is signaled. Return + * {@code null} for no scope. This additive capability requires a core that calls it; older cores retain their + * original hook behavior. */ default AutoCloseable openHandlerScope() { return null; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java b/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java index b0ec6439a..d3422e1e0 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java @@ -59,24 +59,36 @@ private void run(Consumer hook) { /** Runs the root handler with optional plugin scopes, closing them on the same thread in reverse order. */ public T runHandler(Supplier handler) { + return runHandler(handler, () -> {}); + } + + /** Runs the handler and notifies the invocation when a scope requires same-thread finalization. */ + public T runHandler(Supplier handler, Runnable onScopeOpened) { var scopes = new ArrayDeque(); try { for (var plugin : plugins) { try { var scope = plugin.openHandlerScope(); - if (scope != null) scopes.push(scope); - } catch (Exception e) { + if (scope != null) { + scopes.push(scope); + onScopeOpened.run(); + } + } catch (Exception | LinkageError e) { logger.warn("Plugin handler scope threw exception", e); } } return handler.get(); } finally { - while (!scopes.isEmpty()) { - try { - scopes.pop().close(); - } catch (Exception e) { - logger.warn("Plugin handler scope cleanup threw exception", e); - } + closeHandlerScopes(scopes); + } + } + + private static void closeHandlerScopes(ArrayDeque scopes) { + while (!scopes.isEmpty()) { + try { + scopes.pop().close(); + } catch (Exception | LinkageError e) { + logger.warn("Plugin handler scope cleanup threw exception", e); } } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeTest.java b/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeTest.java index 2ee586a57..701236798 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeTest.java @@ -6,6 +6,7 @@ import java.util.ArrayList; import java.util.List; +import java.util.function.Supplier; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; @@ -62,6 +63,78 @@ public AutoCloseable openHandlerScope() { assertEquals(List.of("open-first", "close-first"), calls); } + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void linkageFailuresPreserveTheBodyOutcomeAndAllEarlierScopes(boolean bodyFails) { + var active = new ThreadLocal(); + var calls = new ArrayList(); + var healthy = new DurableExecutionPlugin() { + @Override + public AutoCloseable openHandlerScope() { + calls.add("open-healthy"); + active.set("healthy"); + return () -> { + calls.add("close-healthy"); + active.remove(); + }; + } + }; + var brokenOpen = new DurableExecutionPlugin() { + @Override + public AutoCloseable openHandlerScope() { + calls.add("open-broken"); + throw new NoSuchMethodError("optional API missing"); + } + }; + var brokenClose = new DurableExecutionPlugin() { + @Override + public AutoCloseable openHandlerScope() { + calls.add("open-last"); + return () -> { + calls.add("close-last"); + throw new NoClassDefFoundError("optional class missing"); + }; + } + }; + var runner = new PluginRunner(List.of(healthy, brokenOpen, brokenClose)); + var bodyFailure = new IllegalStateException("original body error"); + Supplier body = () -> { + calls.add("body"); + assertEquals("healthy", active.get()); + if (bodyFails) throw bodyFailure; + return "ok"; + }; + try { + if (bodyFails) + assertSame(bodyFailure, assertThrows(IllegalStateException.class, () -> runner.runHandler(body))); + else assertEquals("ok", runner.runHandler(body)); + assertNull(active.get(), "a later linkage failure must not prevent earlier context cleanup"); + assertEquals( + List.of("open-healthy", "open-broken", "open-last", "body", "close-last", "close-healthy"), calls); + } finally { + active.remove(); + } + } + + @SuppressWarnings("removal") + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void fatalScopeFailuresStillPropagate(boolean duringOpen) { + for (Error fatal : List.of(new InternalError("fatal VM failure"), new ThreadDeath())) { + var plugin = new DurableExecutionPlugin() { + @Override + public AutoCloseable openHandlerScope() { + if (duringOpen) throw fatal; + return () -> { + throw fatal; + }; + } + }; + assertSame( + fatal, assertThrows(Error.class, () -> new PluginRunner(List.of(plugin)).runHandler(() -> "ok"))); + } + } + private static DurableExecutionPlugin scope(String name, List calls, Thread owner) { return new DurableExecutionPlugin() { @Override From ba840d282e663e65b9f0b0bd934ef62065bc1545 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Oct 2026 02:01:21 +0000 Subject: [PATCH 03/48] fix(plugin): preserve wrapped fatal handler-scope failures --- otel-plugin/README.md | 2 +- .../lambda/durable/plugin/PluginRunner.java | 18 ++++++++++++++++-- .../durable/plugin/HandlerScopeTest.java | 17 ++++++++++++++--- 3 files changed, 31 insertions(+), 6 deletions(-) diff --git a/otel-plugin/README.md b/otel-plugin/README.md index f4734a226..2cc23e4bb 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -29,7 +29,7 @@ only a JDK type and defaults to no scope: old plugins on a new core and new plug existing behavior. Root-handler fallback activation requires both the updated core and plugin; no provider API or dependency floor changes are required. The new scope boundary isolates ordinary exceptions and nonfatal linkage errors during open and close, continues earlier scope cleanup, and preserves the handler outcome. JVM fatal errors -remain outside that containment. +remain outside that containment, including when wrapped by asynchronous completion/future exceptions. ## Installation diff --git a/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java b/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java index d3422e1e0..22c034634 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java @@ -5,6 +5,8 @@ import java.util.ArrayDeque; import java.util.Collections; import java.util.List; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ExecutionException; import java.util.function.Consumer; import java.util.function.Supplier; import org.slf4j.Logger; @@ -74,7 +76,7 @@ public T runHandler(Supplier handler, Runnable onScopeOpened) { onScopeOpened.run(); } } catch (Exception | LinkageError e) { - logger.warn("Plugin handler scope threw exception", e); + reportHandlerScopeFailure("Plugin handler scope threw exception", e); } } return handler.get(); @@ -88,11 +90,23 @@ private static void closeHandlerScopes(ArrayDeque scopes) { try { scopes.pop().close(); } catch (Exception | LinkageError e) { - logger.warn("Plugin handler scope cleanup threw exception", e); + reportHandlerScopeFailure("Plugin handler scope cleanup threw exception", e); } } } + @SuppressWarnings("removal") + private static void reportHandlerScopeFailure(String message, Throwable failure) { + var cause = failure; + while ((cause instanceof CompletionException || cause instanceof ExecutionException) + && cause.getCause() != null) { + cause = cause.getCause(); + } + if (cause instanceof VirtualMachineError fatal) throw fatal; + if (cause instanceof ThreadDeath fatal) throw fatal; + logger.warn(message, failure); + } + public void onInvocationStart(InvocationInfo info) { run(p -> p.onInvocationStart(info)); } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeTest.java b/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeTest.java index 701236798..814248b91 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeTest.java @@ -6,9 +6,12 @@ import java.util.ArrayList; import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; import java.util.function.Supplier; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; import org.junit.jupiter.params.provider.ValueSource; import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.util.ExceptionHelper; @@ -118,14 +121,22 @@ public AutoCloseable openHandlerScope() { @SuppressWarnings("removal") @ParameterizedTest - @ValueSource(booleans = {true, false}) - void fatalScopeFailuresStillPropagate(boolean duringOpen) { + @CsvSource({"true,false", "false,false", "true,true", "false,true"}) + void fatalScopeFailuresStillPropagate(boolean duringOpen, boolean wrapped) { for (Error fatal : List.of(new InternalError("fatal VM failure"), new ThreadDeath())) { var plugin = new DurableExecutionPlugin() { @Override public AutoCloseable openHandlerScope() { - if (duringOpen) throw fatal; + if (duringOpen) { + if (wrapped) + CompletableFuture.failedFuture(new ExecutionException(fatal)) + .join(); + throw fatal; + } return () -> { + if (wrapped) + CompletableFuture.failedFuture(new ExecutionException(fatal)) + .join(); throw fatal; }; } From 6269b00a615f3e051499442fd4e113286491aa83 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Oct 2026 02:53:18 +0000 Subject: [PATCH 04/48] fix(ci): keep PR review intake read-only --- .github/workflows/ai-pr-review-address.yml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/ai-pr-review-address.yml b/.github/workflows/ai-pr-review-address.yml index 2b6e25d79..5bd13defb 100644 --- a/.github/workflows/ai-pr-review-address.yml +++ b/.github/workflows/ai-pr-review-address.yml @@ -23,7 +23,10 @@ jobs: contents: read issues: read pull-requests: read - uses: aws/aws-durable-execution-ci/.github/workflows/ai-pr-review-address.yml@d6b017da14385908951d23e26c790b28a4e5f9f0 + uses: aws/aws-durable-execution-ci/.github/workflows/ai-work-item-resolver.yml@d6b017da14385908951d23e26c790b28a4e5f9f0 + with: + work-scope: review + upload-work-items: true address: if: >- From 3e769629986cce6c304573867277d49673f06a11 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Oct 2026 03:03:55 +0000 Subject: [PATCH 05/48] fix(execution): bound handler scope cleanup and preserve control flow --- otel-plugin/README.md | 11 +-- ...ndlerScopeFinalizationIntegrationTest.java | 68 +++++++++++++++++++ .../durable/execution/DurableExecutor.java | 48 +++++++------ .../plugin/DurableExecutionPlugin.java | 3 +- .../execution/HandlerScopeHandoffTest.java | 66 ++++++++++++++++++ 5 files changed, 170 insertions(+), 26 deletions(-) create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/execution/HandlerScopeHandoffTest.java diff --git a/otel-plugin/README.md b/otel-plugin/README.md index 2cc23e4bb..75d5ab56b 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -21,11 +21,12 @@ belongs to another trace, invocation view activates `Invocation`; execution view `Workflow` context, including its unsampled non-recording form. Nested operations retain their existing contexts. The core opens the scope after invocation startup and closes it on the same handler thread on success, failure, or -suspension. When a handler scope is opened, suspension/termination finalization waits asynchronously for the handler to unwind -and close it, before invocation-end hooks flush telemetry or a response is returned. A compatible ambient span uses -a no-op scope so it receives the same finalization ordering. Invocation-end hooks can execute on a different thread -and do not own this scope. The additive hook uses -only a JDK type and defaults to no scope: old plugins on a new core and new plugin layers on an older core retain their +suspension. On suspension or termination, an opened scope gets up to 500ms to unwind before finalization, capped by +remaining Lambda time with 50ms response headroom. The invocation caller waits; the owning handler worker remains free to exit. +If cleanup is blocked, the SDK logs the timeout and preserves the original PENDING/RETRYING outcome. The scope still +closes on its owning thread when the handler eventually exits; late cleanup telemetry is best effort. A compatible +ambient span uses a no-op scope and participates in the same bounded handoff. Invocation-end hooks can execute on a +different thread and do not own this scope. The additive hook uses only a JDK type and defaults to no scope: old plugins on a new core and new plugin layers on an older core retain their existing behavior. Root-handler fallback activation requires both the updated core and plugin; no provider API or dependency floor changes are required. The new scope boundary isolates ordinary exceptions and nonfatal linkage errors during open and close, continues earlier scope cleanup, and preserves the handler outcome. JVM fatal errors diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/HandlerScopeFinalizationIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/HandlerScopeFinalizationIntegrationTest.java index 1974bad3e..f74943571 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/HandlerScopeFinalizationIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/HandlerScopeFinalizationIntegrationTest.java @@ -13,6 +13,7 @@ import java.util.concurrent.atomic.AtomicBoolean; 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.ErrorObject; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.model.ExecutionStatus; @@ -21,6 +22,73 @@ import software.amazon.lambda.durable.testing.LocalDurableTestRunner; class HandlerScopeFinalizationIntegrationTest { + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void blockedCleanupReturnsTheOriginalOutcomeThenClosesOnItsOwner(boolean terminate) throws Exception { + var releaseFinally = new CountDownLatch(1); + var enteredFinally = new CountDownLatch(1); + var scopeClosed = new CountDownLatch(1); + var closedAtEnd = new AtomicBoolean(); + var plugin = new DurableExecutionPlugin() { + @Override + public AutoCloseable openHandlerScope() { + var owner = Thread.currentThread(); + return () -> { + assertSame(owner, Thread.currentThread()); + scopeClosed.countDown(); + }; + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + closedAtEnd.set(scopeClosed.getCount() == 0); + } + }; + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + try { + if (terminate) + context.step("terminate", String.class, step -> { + throw new UnrecoverableDurableExecutionException( + ErrorObject.builder() + .errorMessage("retry") + .build(), + true); + }); + else context.wait("pause", Duration.ofSeconds(1)); + return "done"; + } finally { + enteredFinally.countDown(); + try { + if (!releaseFinally.await(5, TimeUnit.SECONDS)) throw new AssertionError("not released"); + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + throw new AssertionError(error); + } + } + }, + DurableConfig.builder().withPlugins(plugin).build()); + var caller = Executors.newSingleThreadExecutor(); + try { + var response = caller.submit(() -> runner.run("input")); + assertTrue(enteredFinally.await(5, TimeUnit.SECONDS)); + if (terminate) { + var failure = assertThrows(ExecutionException.class, () -> response.get(2, TimeUnit.SECONDS)); + assertInstanceOf(UnrecoverableDurableExecutionException.class, failure.getCause()); + } else + assertEquals( + ExecutionStatus.PENDING, + response.get(2, TimeUnit.SECONDS).getStatus()); + assertFalse(closedAtEnd.get()); + assertEquals(1L, scopeClosed.getCount()); + } finally { + releaseFinally.countDown(); + assertTrue(scopeClosed.await(5, TimeUnit.SECONDS)); + caller.shutdownNow(); + } + } + @ParameterizedTest @CsvSource({"true,false", "false,false", "true,true", "false,true"}) void finalizationWaitsForAnOpenedScopeWithoutChangingLegacyNoScopeTiming(boolean hasScope, boolean terminate) 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 24f7e7041..b924ae7bb 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 @@ -7,6 +7,9 @@ import java.nio.charset.StandardCharsets; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; @@ -125,7 +128,8 @@ public static DurableExecutionOutput execute( return awaitHandlerScopes( executionManager.runUntilCompleteOrSuspend(handlerFuture), handlerFuture, - hasHandlerScope) + hasHandlerScope, + lambdaContext) .handle((result, ex) -> { if (ex != null) { // an exception thrown from handlerFuture or suspension/termination occurred @@ -205,25 +209,29 @@ public static DurableExecutionOutput execute( } } - private static CompletableFuture awaitHandlerScopes( - CompletableFuture executionFuture, CompletableFuture handlerFuture, AtomicBoolean hasHandlerScope) { - return executionFuture - .handle((result, failure) -> { - if (failure == null) return CompletableFuture.completedFuture(result); - if (!hasHandlerScope.get()) return CompletableFuture.failedFuture(failure); - // Suspension/termination can be signaled from inside the handler before its finally blocks unwind. - // A continuation releases that worker to close its scopes; joining it here would deadlock. - // Legacy plugins returning no scope retain their original finalization timing. - return handlerFuture.handle((ignored, handlerFailure) -> { - var cause = ExceptionHelper.unwrapCompletableFuture(handlerFailure); - ExceptionHelper.sneakyThrow( - cause != null && !(cause instanceof SuspendExecutionException) - ? handlerFailure - : failure); - return result; - }); - }) - .thenCompose(future -> future); + static CompletableFuture awaitHandlerScopes( + CompletableFuture executionFuture, + CompletableFuture handlerFuture, + AtomicBoolean hasHandlerScope, + Context lambdaContext) { + // This method runs on the invocation caller, never as a callback on the signaling handler worker. + // Observe the winning outcome before bounded cleanup, then return that same outcome unchanged. + var failure = executionFuture.handle((result, error) -> error).join(); + if (failure == null || !hasHandlerScope.get()) return executionFuture; + var budgetMillis = lambdaContext == null + ? 500L + : Math.min(500L, Math.max(0L, (long) lambdaContext.getRemainingTimeInMillis() - 50L)); + try { + handlerFuture.handle((result, error) -> null).get(budgetMillis, TimeUnit.MILLISECONDS); + } catch (TimeoutException e) { + logger.warn("Handler scope cleanup exceeded its handoff budget; cleanup continues on the handler thread"); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + logger.warn("Interrupted while awaiting handler scope cleanup; preserving the execution outcome", e); + } catch (ExecutionException e) { + logger.warn("Could not observe handler scope cleanup; preserving the execution outcome", e); + } + return executionFuture; } private static void fireOnInvocationEnd( diff --git a/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java b/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java index 556a2741d..7bf966199 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java @@ -41,7 +41,8 @@ default void onInvocationEnd(InvocationEndInfo info) {} * *

The SDK opens and closes this scope on the handler thread, including when the handler fails or suspends. * Scopes close in reverse plugin order when the handler exits. If any scope is returned, invocation finalization - * waits for the handler to unwind and close it, including after suspension/termination is signaled. Return + * gives the handler a bounded opportunity to unwind after suspension/termination. A blocked handler keeps ownership + * of its scope and closes it when it eventually exits; cleanup never replaces the winning execution outcome. Return * {@code null} for no scope. This additive capability requires a core that calls it; older cores retain their * original hook behavior. */ 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 new file mode 100644 index 000000000..8702de106 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/HandlerScopeHandoffTest.java @@ -0,0 +1,66 @@ +// 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 java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import software.amazon.awssdk.services.lambda.model.ErrorObject; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; + +class HandlerScopeHandoffTest { + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void throwingCleanupCannotReplaceTheWinningControlFlow(boolean retry) throws Exception { + var original = control(retry); + var execution = CompletableFuture.failedFuture(original); + var handler = CompletableFuture.failedFuture(new IllegalStateException("finally failed")); + var result = handoff(execution, handler); + var thrown = assertThrows(CompletionException.class, result::join); + assertSame(original, thrown.getCause()); + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void nonUnwindingCleanupDoesNotWaitForever(boolean retry) throws Exception { + var original = control(retry); + var execution = CompletableFuture.failedFuture(original); + var handler = new CompletableFuture(); + var caller = Executors.newSingleThreadExecutor(task -> { + var thread = new Thread(task, "handoff-caller"); + thread.setDaemon(true); + return thread; + }); + try { + var response = caller.submit(() -> { + var result = handoff(execution, handler); + var thrown = assertThrows(CompletionException.class, result::join); + assertSame(original, thrown.getCause()); + return null; + }); + response.get(2, TimeUnit.SECONDS); + assertFalse(handler.isDone(), "timeout must not complete or cancel the owner task"); + } finally { + handler.complete("released"); + caller.shutdownNow(); + } + } + + private static CompletableFuture handoff( + CompletableFuture execution, CompletableFuture handler) { + return DurableExecutor.awaitHandlerScopes(execution, handler, new AtomicBoolean(true), null); + } + + private static Throwable control(boolean retry) { + return retry + ? new UnrecoverableDurableExecutionException( + ErrorObject.builder().errorMessage("retry").build(), true) + : new SuspendExecutionException(); + } +} From 8132a08d61075a135e2bed8747e8c1da872a8033 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Oct 2026 04:09:49 +0000 Subject: [PATCH 06/48] test: verify earlier context cleanup after a scope fatal --- otel-plugin/README.md | 8 +- .../lambda/durable/otel/DelayedFlush.java | 57 +++ .../durable/otel/ExecutionOtelPluginTest.java | 5 +- .../otel/HandlerScopeFinalizationTest.java | 328 ++++++++++++++++++ .../durable/execution/DurableExecutor.java | 33 +- .../plugin/DurableExecutionPlugin.java | 8 +- .../lambda/durable/plugin/PluginRunner.java | 34 +- .../execution/HandlerScopeHandoffTest.java | 36 +- .../durable/plugin/HandlerScopeTest.java | 24 +- 9 files changed, 509 insertions(+), 24 deletions(-) create mode 100644 otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DelayedFlush.java create mode 100644 otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerScopeFinalizationTest.java diff --git a/otel-plugin/README.md b/otel-plugin/README.md index 75d5ab56b..e04d24471 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -22,8 +22,12 @@ belongs to another trace, invocation view activates `Invocation`; execution view The core opens the scope after invocation startup and closes it on the same handler thread on success, failure, or suspension. On suspension or termination, an opened scope gets up to 500ms to unwind before finalization, capped by -remaining Lambda time with 50ms response headroom. The invocation caller waits; the owning handler worker remains free to exit. -If cleanup is blocked, the SDK logs the timeout and preserves the original PENDING/RETRYING outcome. The scope still +remaining Lambda time after reserving five seconds per configured plugin and one second for shutdown/response. +This is a best-effort reserve, not a total deadline bound: existing plugin callbacks and checkpoint draining may exceed it. The invocation caller waits; the owning handler worker remains free to exit. +If cleanup is blocked, the SDK logs the timeout and preserves the original PENDING/RETRYING outcome. +An observed `VirtualMachineError` or `ThreadDeath` from the new scope callbacks instead escapes the invocation caller; +ordinary cleanup failures and legacy body failures retain the original outcome. Fatals reported only after the response +cannot retroactively change it. The scope still closes on its owning thread when the handler eventually exits; late cleanup telemetry is best effort. A compatible ambient span uses a no-op scope and participates in the same bounded handoff. Invocation-end hooks can execute on a different thread and do not own this scope. The additive hook uses only a JDK type and defaults to no scope: old plugins on a new core and new plugin layers on an older core retain their diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DelayedFlush.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DelayedFlush.java new file mode 100644 index 000000000..a762d23e4 --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DelayedFlush.java @@ -0,0 +1,57 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.otel; + +import io.opentelemetry.context.Context; +import io.opentelemetry.sdk.common.CompletableResultCode; +import io.opentelemetry.sdk.trace.ReadWriteSpan; +import io.opentelemetry.sdk.trace.ReadableSpan; +import io.opentelemetry.sdk.trace.SpanProcessor; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; + +class DelayedFlush implements SpanProcessor { + public void onStart(Context parent, ReadWriteSpan span) {} + + final HandlerScopeFinalizationTest.Deadline deadline; + final AtomicInteger calls = new AtomicInteger(); + final AtomicInteger remainingAtFlush = new AtomicInteger(); + final AtomicLong flushMillis = new AtomicLong(); + final ScheduledExecutorService clock; + + DelayedFlush(HandlerScopeFinalizationTest.Deadline deadline, ScheduledExecutorService clock) { + this.deadline = deadline; + this.clock = clock; + } + + public boolean isStartRequired() { + return false; + } + + public void onEnd(ReadableSpan span) {} + + public boolean isEndRequired() { + return false; + } + + public CompletableResultCode shutdown() { + return CompletableResultCode.ofSuccess(); + } + + public CompletableResultCode forceFlush() { + calls.incrementAndGet(); + remainingAtFlush.set(deadline.remaining()); + var result = new CompletableResultCode(); + long start = System.nanoTime(); + clock.schedule( + () -> { + flushMillis.set(TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start)); + result.succeed(); + }, + 250, + TimeUnit.MILLISECONDS); + return result; + } +} diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExecutionOtelPluginTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExecutionOtelPluginTest.java index 8b84efc10..9614759f6 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExecutionOtelPluginTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExecutionOtelPluginTest.java @@ -280,7 +280,8 @@ void workflowAndInvocationSpans_shareExecutionTrace_withoutAmbientContext() { void invocationStart_joinsAmbientTrace_whenAmbientIsOnExecutionTrace() { // Drive an invocation to learn the canonical execution trace ID, then start a fresh invocation with an ambient // span on that same trace: the Invocation span joins the ambient span directly. - plugin.onInvocationStart(new InvocationInfo("req-0", ARN, true, Instant.now())); + var executionStart = Instant.parse("2026-01-01T00:00:00Z"); + plugin.onInvocationStart(new InvocationInfo("req-0", ARN, true, executionStart)); plugin.onInvocationEnd(new InvocationEndInfo("req-0", ARN, true, InvocationStatus.SUCCEEDED, null)); var canonicalTraceId = spanByName(spanExporter.getFinishedSpanItems(), "Workflow").getTraceId(); @@ -290,7 +291,7 @@ void invocationStart_joinsAmbientTrace_whenAmbientIsOnExecutionTrace() { var ambient = SpanContext.create(canonicalTraceId, ambientSpanId, TraceFlags.getSampled(), TraceState.getDefault()); try (var ignored = Span.wrap(ambient).makeCurrent()) { - plugin.onInvocationStart(new InvocationInfo("req-1", ARN, false, Instant.now())); + plugin.onInvocationStart(new InvocationInfo("req-1", ARN, false, executionStart)); } plugin.onInvocationEnd(new InvocationEndInfo("req-1", ARN, false, InvocationStatus.SUCCEEDED, null)); diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerScopeFinalizationTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerScopeFinalizationTest.java new file mode 100644 index 000000000..2176540cc --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerScopeFinalizationTest.java @@ -0,0 +1,328 @@ +// 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 com.amazonaws.services.lambda.runtime.Context; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import java.lang.reflect.Proxy; +import java.time.Duration; +import java.time.Instant; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import org.junit.jupiter.api.Test; +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.TypeToken; +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.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.testing.local.LocalMemoryExecutionClient; + +class HandlerScopeFinalizationTest { + @ParameterizedTest + @CsvSource({ + "false,false,false", + "false,false,true", + "false,true,false", + "false,true,true", + "true,false,false", + "true,false,true", + "true,true,false", + "true,true,true" + }) + void preservesTimeForRealOtelFinalization(boolean executionView, boolean retry, boolean hasScope) throws Exception { + var deadline = new Deadline(); + var clock = Executors.newSingleThreadScheduledExecutor(r -> daemon(r, "flush-clock")); + var flush = new DelayedFlush(deadline, clock); + var builder = SdkTracerProvider.builder().addSpanProcessor(flush); + var settings = OtelPluginConfig.builder() + .enableMdc(false) + .contextExtractor(() -> new ExtractedContext( + "12345678901234567890123456789012", "1234567890123456", ExtractedContext.Sampling.SAMPLED)) + .build(); + DurableExecutionPlugin plugin = executionView + ? new ExecutionOtelPlugin(builder, settings) { + public AutoCloseable openHandlerScope() { + return hasScope ? super.openHandlerScope() : null; + } + } + : new InvocationOtelPlugin(builder, settings) { + public AutoCloseable openHandlerScope() { + return hasScope ? super.openHandlerScope() : null; + } + }; + var workers = Executors.newCachedThreadPool(r -> daemon(r, "handler-owner")); + var callers = Executors.newSingleThreadExecutor(r -> daemon(r, "invocation-caller")); + var enteredFinally = new CountDownLatch(1); + var releaseFinally = new CountDownLatch(1); + var original = retryError(); + var cfg = config(workers, plugin); + try { + var response = callers.submit(() -> DurableExecutor.execute( + input(), + deadline.context(), + TypeToken.get(String.class), + (value, ctx) -> { + deadline.arm(400); + try { + if (retry) + ctx.step("retry", String.class, step -> { + throw original; + }); + else ctx.wait("pause", Duration.ofSeconds(1)); + return "done"; + } finally { + enteredFinally.countDown(); + await(releaseFinally); + } + }, + cfg)); + check(enteredFinally.await(2, TimeUnit.SECONDS), "handler finally not reached"); + outcome(response, retry, original); + long total = deadline.elapsedMillis(); + check(flush.calls.get() == 1, "expected one actual bundled OTel forceFlush"); + check(flush.flushMillis.get() >= 240, "actual bundled OTel join did not await forceFlush result"); + check(total < 400, "unexpected response/deadline relation: " + total); + check(releaseFinally.getCount() == 1, "test cleanup was not blocked"); + + } finally { + releaseFinally.countDown(); + workers.shutdown(); + workers.awaitTermination(2, TimeUnit.SECONDS); + callers.shutdownNow(); + clock.shutdownNow(); + } + } + + @SuppressWarnings("removal") + @ParameterizedTest + @CsvSource({ + "false,false,false", + "false,false,true", + "false,true,false", + "false,true,true", + "true,false,false", + "true,false,true", + "true,true,false", + "true,true,true" + }) + void scopeFatal(boolean retry, boolean wrapped, boolean threadDeath) throws Exception { + scopeFatal(retry, wrapped, threadDeath ? new ThreadDeath() : new InternalError("scope fatal"), false); + } + + @Test + void losingBodyFinallyFatalPreservesTheOriginalWinner() throws Exception { + scopeFatal(false, false, new InternalError("body finally fatal"), true); + } + + @Test + void fatalAfterTheHandoffTimeoutCannotChangeAnAlreadyReturnedOutcome() throws Exception { + scopeFatal(false, false, new InternalError("late scope fatal"), false, true); + } + + static void scopeFatal(boolean retry, boolean wrapped, Error fatal, boolean fromBodyFinally) throws Exception { + scopeFatal(retry, wrapped, fatal, fromBodyFinally, false); + } + + static void scopeFatal(boolean retry, boolean wrapped, Error fatal, boolean fromBodyFinally, boolean afterTimeout) + throws Exception { + var deadline = new Deadline(); + var closeEntered = new CountDownLatch(1); + var releaseFatal = new CountDownLatch(1); + var fatalRaised = new CountDownLatch(1); + var fatalBeforeEnd = new AtomicBoolean(); + var plugin = new DurableExecutionPlugin() { + public AutoCloseable openHandlerScope() { + var owner = Thread.currentThread(); + return () -> { + check(owner == Thread.currentThread(), "wrong scope close thread"); + if (!fromBodyFinally) raiseFatal(closeEntered, releaseFatal, fatalRaised, wrapped, fatal); + }; + } + + public void onInvocationEnd(InvocationEndInfo info) { + fatalBeforeEnd.set(fatalRaised.getCount() == 0); + } + }; + var workers = Executors.newCachedThreadPool(r -> daemon(r, "fatal-owner")); + var callers = Executors.newSingleThreadExecutor(r -> daemon(r, "fatal-caller")); + var original = retryError(); + var cfg = config(workers, plugin); + try { + var response = callers.submit(() -> DurableExecutor.execute( + input(), + deadline.context(), + TypeToken.get(String.class), + (value, ctx) -> { + deadline.arm(30_000); + try { + if (retry) + ctx.step("retry", String.class, step -> { + throw original; + }); + else ctx.wait("pause", Duration.ofSeconds(1)); + return "done"; + } finally { + if (fromBodyFinally) raiseFatal(closeEntered, releaseFatal, fatalRaised, wrapped, fatal); + } + }, + cfg)); + check(closeEntered.await(2, TimeUnit.SECONDS), "fatal site not reached"); + check(deadline.budgetRead.await(2, TimeUnit.SECONDS), "winner not observed before fatal"); + if (afterTimeout) { + outcome(response, retry, original); + assertFalse(fatalBeforeEnd.get()); + releaseFatal.countDown(); + assertTrue(fatalRaised.await(2, TimeUnit.SECONDS)); + outcome(response, retry, original); + return; + } + releaseFatal.countDown(); + if (fromBodyFinally) { + outcome(response, retry, original); + assertTrue(fatalBeforeEnd.get()); + } else { + var thrown = assertThrows(ExecutionException.class, () -> response.get(3, TimeUnit.SECONDS)); + assertSame(fatal, thrown.getCause(), "scope-owned fatal must escape the invocation caller"); + } + + } finally { + releaseFatal.countDown(); + workers.shutdown(); + workers.awaitTermination(2, TimeUnit.SECONDS); + callers.shutdownNow(); + } + } + + @SuppressWarnings("removal") + @ParameterizedTest + @CsvSource({"false", "true"}) + void legacyBodyFatalRetainsFailedMapping(boolean threadDeath) { + Error fatal = threadDeath ? new ThreadDeath() : new InternalError("body fatal"); + var workers = Executors.newCachedThreadPool(r -> daemon(r, "body-fatal")); + try { + var output = DurableExecutor.execute( + input(), + null, + TypeToken.get(String.class), + (value, ctx) -> { + throw fatal; + }, + config(workers)); + check(output.status() == ExecutionStatus.FAILED, "legacy body fatal not FAILED"); + check(output.error().errorType().equals(fatal.getClass().getName()), "legacy body fatal error type lost"); + + } finally { + workers.shutdownNow(); + } + } + + static void raiseFatal( + CountDownLatch entered, CountDownLatch release, CountDownLatch raised, boolean wrapped, Error fatal) { + entered.countDown(); + await(release); + raised.countDown(); + if (wrapped) throw new CompletionException(new ExecutionException(fatal)); + throw fatal; + } + + static void outcome(Future response, boolean retry, Throwable original) throws Exception { + if (retry) { + try { + response.get(3, TimeUnit.SECONDS); + throw new AssertionError("retry returned output"); + } catch (ExecutionException e) { + check(e.getCause() == original, "retry failure changed to " + e.getCause()); + } + } else check(response.get(3, TimeUnit.SECONDS).status() == ExecutionStatus.PENDING, "suspension changed"); + } + + static DurableConfig config(ExecutorService workers, DurableExecutionPlugin... plugins) { + return DurableConfig.builder() + .withDurableExecutionClient(new LocalMemoryExecutionClient()) + .withExecutorService(workers) + .withCheckpointDelay(Duration.ZERO) + .withPlugins(plugins) + .build(); + } + + static UnrecoverableDurableExecutionException retryError() { + return new UnrecoverableDurableExecutionException( + ErrorObject.builder().errorMessage("original retry").build(), true); + } + + static DurableExecutionInput input() { + var operation = Operation.builder() + .id("id") + .name("test") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.now()) + .executionDetails( + ExecutionDetails.builder().inputPayload("\"input\"").build()) + .build(); + return new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/name/id", + "token", + CheckpointUpdatedExecutionState.builder().operations(operation).build()); + } + + static Thread daemon(Runnable runnable, String name) { + var thread = new Thread(runnable, name); + thread.setDaemon(true); + return thread; + } + + static void await(CountDownLatch latch) { + try { + check(latch.await(3, TimeUnit.SECONDS), "fixture latch timed out"); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError("owner interrupted", e); + } + } + + static void check(boolean condition, String message) { + if (!condition) throw new AssertionError(message); + } + + static class Deadline { + volatile long start; + volatile long end; + final CountDownLatch budgetRead = new CountDownLatch(1); + + void arm(int millis) { + start = System.nanoTime(); + end = start + TimeUnit.MILLISECONDS.toNanos(millis); + } + + long elapsedMillis() { + return TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start); + } + + int remaining() { + return end == 0 ? 30_000 : (int) Math.max(0L, TimeUnit.NANOSECONDS.toMillis(end - System.nanoTime())); + } + + Context context() { + return (Context) Proxy.newProxyInstance( + Context.class.getClassLoader(), new Class[] {Context.class}, (proxy, method, args) -> { + if (method.getName().equals("getRemainingTimeInMillis")) { + budgetRead.countDown(); + return remaining(); + } + if (method.getName().equals("getAwsRequestId")) return "test-request"; + if (method.getName().equals("getMemoryLimitInMB")) return 512; + return null; + }); + } + } +} 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 b924ae7bb..483b80803 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 @@ -51,6 +51,11 @@ public class DurableExecutor { // Lambda response size limit is 6MB minus small epsilon for envelope private static final int LAMBDA_RESPONSE_SIZE_LIMIT = 6 * 1024 * 1024 - 50; + // Best-effort allowance for each configured plugin's finalization (including the bundled OTel 5s join), + // plus shutdown/response headroom. Existing arbitrary callbacks and checkpoint draining can exceed it. + private static final long PLUGIN_FINALIZATION_RESERVE_MILLIS = 5_000; + private static final long SHUTDOWN_RESPONSE_RESERVE_MILLIS = 1_000; + private DurableExecutor() {} public static DurableExecutionOutput execute( @@ -69,6 +74,7 @@ public static DurableExecutionOutput execute( // Captured for onInvocationEnd, which runs outside the handler thread below. var pluginExecutionInput = new AtomicReference<>(); var hasHandlerScope = new AtomicBoolean(); + var scopeFatal = new AtomicReference(); var handlerFuture = CompletableFuture.supplyAsync( () -> { executionManager.setCurrentThreadContext(new ThreadContext(null, ThreadType.CONTEXT)); @@ -116,7 +122,9 @@ public static DurableExecutionOutput execute( try (var ignored = DurableLogger.attachContext()) { var handlerInput = userInput; return pluginRunner.runHandler( - () -> handler.apply(handlerInput, context), () -> hasHandlerScope.set(true)); + () -> handler.apply(handlerInput, context), + () -> hasHandlerScope.set(true), + fatal -> scopeFatal.compareAndSet(null, fatal)); } }, config.getExecutorService()); // Get executor from config for running user code @@ -129,7 +137,9 @@ public static DurableExecutionOutput execute( executionManager.runUntilCompleteOrSuspend(handlerFuture), handlerFuture, hasHandlerScope, - lambdaContext) + lambdaContext, + pluginRunner.getPlugins().size(), + scopeFatal) .handle((result, ex) -> { if (ex != null) { // an exception thrown from handlerFuture or suspension/termination occurred @@ -205,6 +215,9 @@ public static DurableExecutionOutput execute( // unwrap the CompletionException and rethrow the wrapped exception ExceptionHelper.sneakyThrow(ExceptionHelper.unwrapCompletableFuture(e)); return null; + } finally { + // Also observe a scope fatal reported during invocation-end hooks, without extending the wait. + throwIfScopeFatal(scopeFatal); } } } @@ -213,14 +226,18 @@ static CompletableFuture awaitHandlerScopes( CompletableFuture executionFuture, CompletableFuture handlerFuture, AtomicBoolean hasHandlerScope, - Context lambdaContext) { + Context lambdaContext, + int pluginCount, + AtomicReference scopeFatal) { // This method runs on the invocation caller, never as a callback on the signaling handler worker. - // Observe the winning outcome before bounded cleanup, then return that same outcome unchanged. + // Preserve the winning outcome except for an observed fatal error from the new scope callbacks. var failure = executionFuture.handle((result, error) -> error).join(); + throwIfScopeFatal(scopeFatal); if (failure == null || !hasHandlerScope.get()) return executionFuture; + var reserve = SHUTDOWN_RESPONSE_RESERVE_MILLIS + PLUGIN_FINALIZATION_RESERVE_MILLIS * pluginCount; var budgetMillis = lambdaContext == null ? 500L - : Math.min(500L, Math.max(0L, (long) lambdaContext.getRemainingTimeInMillis() - 50L)); + : Math.min(500L, Math.max(0L, (long) lambdaContext.getRemainingTimeInMillis() - reserve)); try { handlerFuture.handle((result, error) -> null).get(budgetMillis, TimeUnit.MILLISECONDS); } catch (TimeoutException e) { @@ -231,9 +248,15 @@ static CompletableFuture awaitHandlerScopes( } catch (ExecutionException e) { logger.warn("Could not observe handler scope cleanup; preserving the execution outcome", e); } + throwIfScopeFatal(scopeFatal); return executionFuture; } + private static void throwIfScopeFatal(AtomicReference scopeFatal) { + var fatal = scopeFatal.get(); + if (fatal != null) throw fatal; + } + private static void fireOnInvocationEnd( PluginRunner pluginRunner, ExecutionManager executionManager, diff --git a/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java b/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java index 7bf966199..a2cd13688 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java @@ -42,9 +42,11 @@ default void onInvocationEnd(InvocationEndInfo info) {} *

The SDK opens and closes this scope on the handler thread, including when the handler fails or suspends. * Scopes close in reverse plugin order when the handler exits. If any scope is returned, invocation finalization * gives the handler a bounded opportunity to unwind after suspension/termination. A blocked handler keeps ownership - * of its scope and closes it when it eventually exits; cleanup never replaces the winning execution outcome. Return - * {@code null} for no scope. This additive capability requires a core that calls it; older cores retain their - * original hook behavior. + * of its scope and closes it when it eventually exits. Ordinary cleanup failures preserve the winning outcome; an + * observed VirtualMachineError or ThreadDeath from scope callbacks escapes the invocation caller. Optional waiting + * reserves five seconds per configured plugin plus one second for shutdown/response; this is best effort because + * existing finalizers and checkpoint draining can exceed that allowance. Return {@code null} for no scope. This + * additive capability requires a core that calls it; older cores retain their original hook behavior. */ default AutoCloseable openHandlerScope() { return null; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java b/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java index 22c034634..0908b1a85 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java @@ -66,6 +66,12 @@ public T runHandler(Supplier handler) { /** Runs the handler and notifies the invocation when a scope requires same-thread finalization. */ public T runHandler(Supplier handler, Runnable onScopeOpened) { + return runHandler(handler, onScopeOpened, fatal -> {}); + } + + /** Reports fatal errors originating only in scope callbacks, before rethrowing on the owner thread. */ + @SuppressWarnings("removal") + public T runHandler(Supplier handler, Runnable onScopeOpened, Consumer onScopeFatal) { var scopes = new ArrayDeque(); try { for (var plugin : plugins) { @@ -75,35 +81,45 @@ public T runHandler(Supplier handler, Runnable onScopeOpened) { scopes.push(scope); onScopeOpened.run(); } - } catch (Exception | LinkageError e) { - reportHandlerScopeFailure("Plugin handler scope threw exception", e); + } catch (Exception | LinkageError | VirtualMachineError | ThreadDeath e) { + reportHandlerScopeFailure("Plugin handler scope threw exception", e, onScopeFatal); } } return handler.get(); } finally { - closeHandlerScopes(scopes); + closeHandlerScopes(scopes, onScopeFatal); } } - private static void closeHandlerScopes(ArrayDeque scopes) { + @SuppressWarnings("removal") + private static void closeHandlerScopes(ArrayDeque scopes, Consumer onScopeFatal) { + Error firstFatal = null; while (!scopes.isEmpty()) { try { scopes.pop().close(); - } catch (Exception | LinkageError e) { - reportHandlerScopeFailure("Plugin handler scope cleanup threw exception", e); + } catch (Exception | LinkageError | VirtualMachineError | ThreadDeath e) { + try { + reportHandlerScopeFailure("Plugin handler scope cleanup threw exception", e, onScopeFatal); + } catch (VirtualMachineError | ThreadDeath fatal) { + if (firstFatal == null) firstFatal = fatal; + } } } + if (firstFatal != null) throw firstFatal; } @SuppressWarnings("removal") - private static void reportHandlerScopeFailure(String message, Throwable failure) { + private static void reportHandlerScopeFailure(String message, Throwable failure, Consumer onScopeFatal) { var cause = failure; while ((cause instanceof CompletionException || cause instanceof ExecutionException) && cause.getCause() != null) { cause = cause.getCause(); } - if (cause instanceof VirtualMachineError fatal) throw fatal; - if (cause instanceof ThreadDeath fatal) throw fatal; + if (cause instanceof VirtualMachineError || cause instanceof ThreadDeath) { + var fatal = (Error) cause; + onScopeFatal.accept(fatal); + throw fatal; + } logger.warn(message, failure); } 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 8702de106..0f107f043 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 @@ -3,12 +3,17 @@ 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.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +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.ErrorObject; @@ -52,9 +57,38 @@ void nonUnwindingCleanupDoesNotWaitForever(boolean retry) throws Exception { } } + @Test + void reservesFinalizationForEveryConfiguredPluginWithoutCancellingTheOwner() { + var context = mock(Context.class); + when(context.getRemainingTimeInMillis()).thenReturn(10_000); + var execution = CompletableFuture.failedFuture(control(false)); + var handler = new CompletableFuture(); + assertTimeoutPreemptively( + Duration.ofMillis(200), + () -> assertSame( + execution, + DurableExecutor.awaitHandlerScopes( + execution, handler, new AtomicBoolean(true), context, 2, new AtomicReference<>()))); + assertFalse(handler.isDone()); + } + + @Test + void noScopeDoesNotConsultTheRemainingTimeOrWaitForTheHandler() { + var context = mock(Context.class); + var execution = CompletableFuture.failedFuture(control(false)); + var handler = new CompletableFuture(); + assertSame( + execution, + DurableExecutor.awaitHandlerScopes( + execution, handler, new AtomicBoolean(false), context, 1, new AtomicReference<>())); + verifyNoInteractions(context); + assertFalse(handler.isDone()); + } + private static CompletableFuture handoff( CompletableFuture execution, CompletableFuture handler) { - return DurableExecutor.awaitHandlerScopes(execution, handler, new AtomicBoolean(true), null); + return DurableExecutor.awaitHandlerScopes( + execution, handler, new AtomicBoolean(true), null, 1, new AtomicReference<>()); } private static Throwable control(boolean retry) { diff --git a/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeTest.java b/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeTest.java index 814248b91..0704fd639 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeTest.java @@ -8,6 +8,7 @@ import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.Supplier; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; @@ -141,8 +142,27 @@ public AutoCloseable openHandlerScope() { }; } }; - assertSame( - fatal, assertThrows(Error.class, () -> new PluginRunner(List.of(plugin)).runHandler(() -> "ok"))); + var reported = new AtomicReference(); + var calls = new ArrayList(); + var active = new ThreadLocal(); + var owner = Thread.currentThread(); + var healthy = new DurableExecutionPlugin() { + @Override + public AutoCloseable openHandlerScope() { + active.set("healthy"); + calls.add("open-healthy"); + return () -> { + assertSame(owner, Thread.currentThread()); + active.remove(); + calls.add("close-healthy"); + }; + } + }; + var runner = new PluginRunner(List.of(healthy, plugin)); + assertSame(fatal, assertThrows(Error.class, () -> runner.runHandler(() -> "ok", () -> {}, reported::set))); + assertSame(fatal, reported.get(), "only a scope callback reports this fatal"); + assertNull(active.get(), "earlier context must not leak into reuse of the owner thread"); + assertEquals(List.of("open-healthy", "close-healthy"), calls); } } From 4d81d8c8748da85bb95b6dc442dffb3d6c031c38 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Mon, 5 Oct 2026 22:24:26 -0700 Subject: [PATCH 07/48] test: cover fatal cleanup during invocation end dispatch --- .../otel/HandlerScopeFinalizationTest.java | 71 ++++++++++++------- .../durable/execution/DurableExecutor.java | 41 +++++++++-- 2 files changed, 82 insertions(+), 30 deletions(-) diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerScopeFinalizationTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerScopeFinalizationTest.java index 668046ed7..a549d8146 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerScopeFinalizationTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerScopeFinalizationTest.java @@ -27,6 +27,7 @@ import software.amazon.lambda.durable.plugin.HandlerScoped; import software.amazon.lambda.durable.plugin.InvocationEndInfo; import software.amazon.lambda.durable.plugin.InvocationInfo; +import software.amazon.lambda.durable.plugin.InvocationStatus; import software.amazon.lambda.durable.testing.local.LocalMemoryExecutionClient; class HandlerScopeFinalizationTest { @@ -139,13 +140,21 @@ void fatalAfterTheHandoffTimeoutCannotChangeAnAlreadyReturnedOutcome() throws Ex scopeFatal(false, false, new InternalError("late scope fatal"), false, true); } - @Test - void observesScopeFatalReportedDuringManagerShutdown() throws Exception { + @SuppressWarnings("removal") + @ParameterizedTest + @CsvSource({"false,false", "true,false", "false,true", "true,true"}) + void scopeFatalAfterEndDispatchPreservesCallerOutcomeAndEscapesOwner(boolean threadDeath, boolean duringEnd) + throws Exception { var releaseScope = new CountDownLatch(1); var ownerFinished = new CountDownLatch(1); var shutdownEntered = new AtomicBoolean(); var ended = new AtomicBoolean(); - var fatal = new InternalError("scope failure during manager shutdown"); + var endCalls = new AtomicInteger(); + var endInfo = new AtomicReference(); + var fatalObservedDuringEnd = new AtomicBoolean(); + Error fatal = threadDeath ? new ThreadDeath() : new InternalError("late handler scope failure"); + var ownerFatal = new AtomicReference(); + var fatalEscaped = new CountDownLatch(1); var plugin = new ScopedPlugin() { public AutoCloseable openHandlerScope() { return () -> { @@ -156,16 +165,24 @@ public AutoCloseable openHandlerScope() { public void onInvocationEnd(InvocationEndInfo info) { ended.set(true); + endCalls.incrementAndGet(); + endInfo.set(info); + if (duringEnd) { + releaseScope.countDown(); + await(fatalEscaped); + fatalObservedDuringEnd.set(ownerFatal.get() == fatal); + } } }; var workers = - new ThreadPoolExecutor( - 0, - Integer.MAX_VALUE, - 60, - TimeUnit.SECONDS, - new SynchronousQueue<>(), - task -> daemon(task, "shutdown-owner")) { + new ThreadPoolExecutor(0, Integer.MAX_VALUE, 60, TimeUnit.SECONDS, new SynchronousQueue<>(), task -> { + var owner = daemon(task, "shutdown-owner"); + owner.setUncaughtExceptionHandler((thread, failure) -> { + ownerFatal.set(failure); + fatalEscaped.countDown(); + }); + return owner; + }) { @Override public void execute(Runnable task) { super.execute(() -> { @@ -190,21 +207,27 @@ public int getActiveCount() { }; var deadline = new Deadline(); try { - assertSame( - fatal, - assertThrows( - InternalError.class, - () -> DurableExecutor.execute( - input(), - deadline.context(), - TypeToken.get(String.class), - (value, ctx) -> { - deadline.arm(0); - ctx.wait("pause", Duration.ofSeconds(1)); - return "done"; - }, - config(workers, plugin)))); + var output = assertDoesNotThrow(() -> DurableExecutor.execute( + input(), + deadline.context(), + TypeToken.get(String.class), + (value, ctx) -> { + deadline.arm(0); + ctx.wait("pause", Duration.ofSeconds(1)); + return "done"; + }, + config(workers, plugin))); + assertEquals(ExecutionStatus.PENDING, output.status(), "the already finalized caller outcome stays frozen"); + assertTrue(fatalEscaped.await(2, TimeUnit.SECONDS), "fatal must still escape the actual owner thread"); + assertSame(fatal, ownerFatal.get()); + assertEquals(duringEnd, fatalObservedDuringEnd.get(), "blocked end hook must observe the owner fatal"); assertTrue(shutdownEntered.get()); + assertEquals(1, endCalls.get(), "a late fatal must not replay already delivered end hooks"); + assertEquals( + InvocationStatus.PENDING, + endInfo.get().invocationStatus(), + "snapshot reflects outcome known at dispatch"); + assertNull(endInfo.get().executionError(), "the fatal is reported after the selected end snapshot"); } finally { releaseScope.countDown(); workers.shutdown(); 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 e26a9fcdd..428ebc8b6 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 @@ -8,6 +8,7 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; @@ -78,7 +79,7 @@ public static DurableExecutionOutput execute( // Captured for onInvocationEnd, which runs outside the handler thread below. var pluginExecutionInput = new AtomicReference<>(); var hasHandlerScope = new AtomicBoolean(); - var handlerFuture = CompletableFuture.supplyAsync( + var handlerFuture = supplyHandler( preservingMdc(pluginRunner, () -> { executionManager.setCurrentThreadContext(new ThreadContext(null, ThreadType.CONTEXT)); @@ -130,7 +131,9 @@ public static DurableExecutionOutput execute( fatal -> scopeFatal.compareAndSet(null, fatal)); } }), - config.getExecutorService()); // Get executor from config for running user code + config.getExecutorService(), + pluginRunner, + scopeFatal); // Get executor from config for running user code // Execute the handlerFuture in ExecutionManager. If it completes successfully, the output of user function // will be returned. Otherwise, it will complete exceptionally with a SuspendExecutionException or a @@ -151,6 +154,7 @@ public static DurableExecutionOutput execute( // return PENDING if it's SuspendExecutionException if (cause instanceof SuspendExecutionException) { fireOnInvocationEnd( + scopeFatal, pluginRunner, executionManager, requestId, @@ -170,6 +174,7 @@ public static DurableExecutionOutput execute( unrecoverableDurableExecutionException && unrecoverableDurableExecutionException.isRetryable()) { fireOnInvocationEnd( + scopeFatal, pluginRunner, executionManager, requestId, @@ -185,6 +190,7 @@ public static DurableExecutionOutput execute( // fail the execution otherwise logger.debug("Execution failed: {}", cause.getMessage()); fireOnInvocationEnd( + scopeFatal, pluginRunner, executionManager, requestId, @@ -202,6 +208,7 @@ public static DurableExecutionOutput execute( var output = DurableExecutionOutput.success(handleLargePayload(executionManager, outputPayload)); fireOnInvocationEnd( + scopeFatal, pluginRunner, executionManager, requestId, @@ -219,13 +226,31 @@ public static DurableExecutionOutput execute( ExceptionHelper.sneakyThrow(ExceptionHelper.unwrapCompletableFuture(e)); return null; } - } finally { - // Resource shutdown may wait for other work after the bounded handoff. Observe any scope fatal - // reported during that existing wait before the caller commits its response, without waiting again. - throwIfScopeFatal(scopeFatal); } } + private static CompletableFuture supplyHandler( + Supplier task, Executor executor, PluginRunner plugins, AtomicReference scopeFatal) { + if (plugins.isEmpty()) return CompletableFuture.supplyAsync(task, executor); + var result = new CompletableFuture(); + var caller = Thread.currentThread(); + Runnable work = (Runnable & CompletableFuture.AsynchronousCompletionTask) () -> { + try { + result.complete(task.get()); + } catch (Throwable failure) { + result.completeExceptionally(failure); + // Preserve legacy user-body Error handling. Only a positively reported new scope fatal also + // escapes its actual owner after its observation future settles. Direct callers finalize below. + var fatal = scopeFatal.get(); + if (Thread.currentThread() != caller + && fatal != null + && ExceptionHelper.unwrapCompletableFuture(failure) == fatal) throw fatal; + } + }; + executor.execute(work); + return result; + } + static CompletableFuture awaitHandlerScopes( CompletableFuture executionFuture, CompletableFuture handlerFuture, @@ -280,6 +305,7 @@ private static SafeCloseable restoreMdcOnClose() { } private static void fireOnInvocationEnd( + AtomicReference scopeFatal, PluginRunner pluginRunner, ExecutionManager executionManager, String requestId, @@ -292,6 +318,9 @@ private static void fireOnInvocationEnd( if (pluginRunner.isEmpty()) { return; } + // Freeze the selected end snapshot here. A scope fatal already observed still escapes before dispatch; + // a later scope fatal escapes its owner instead of rewriting this caller outcome after finalization. + throwIfScopeFatal(scopeFatal); // Finalization can run on the invocation caller rather than the handler worker. try (var ignored = restoreMdcOnClose()) { pluginRunner.onInvocationEnd(new InvocationEndInfo( From ddcd77b4f7656d92b15f27ecbd0534ca45e05eae Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Tue, 6 Oct 2026 10:26:38 -0700 Subject: [PATCH 08/48] fix: preserve invocation end completion-thread dispatch --- .../otel/InvocationEndCompatibilityTest.java | 211 ++++++++++++++++++ .../durable/execution/DurableExecutor.java | 56 ++++- 2 files changed, 257 insertions(+), 10 deletions(-) create mode 100644 otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndCompatibilityTest.java diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndCompatibilityTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndCompatibilityTest.java new file mode 100644 index 000000000..d36280d50 --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndCompatibilityTest.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.otel; + +import static org.junit.jupiter.api.Assertions.*; + +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.*; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.locks.LockSupport; +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.TypeToken; +import software.amazon.lambda.durable.execution.DurableExecutor; +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.InvocationInfo; +import software.amazon.lambda.durable.testing.local.LocalMemoryExecutionClient; + +/** Preserves legacy completion-thread dispatch without promising owner affinity for precompleted futures. */ +class InvocationEndCompatibilityTest { + @ParameterizedTest + @CsvSource({ + "false,false,false", "true,false,false", "false,true,false", "true,true,false", + "false,false,true", "true,false,true", "false,true,true", "true,true,true" + }) + void legacyEndThread(boolean mixed, boolean precompleted, boolean pending) throws Exception { + var local = new ThreadLocal(); + var observations = new ArrayList(); + var callerThread = new AtomicReference(); + var caller = Executors.newSingleThreadExecutor(task -> { + var thread = daemon(task, "legacy-probe-caller"); + callerThread.set(thread); + return thread; + }); + var workers = + new ThreadPoolExecutor( + 1, + 1, + 0, + TimeUnit.SECONDS, + new LinkedBlockingQueue<>(), + task -> daemon(task, "legacy-probe-worker")) { + @Override + public void execute(Runnable task) { + var finished = new CountDownLatch(1); + super.execute(() -> { + try { + task.run(); + } finally { + finished.countDown(); + } + }); + if (precompleted) await(finished); + } + }; + try { + for (int round = 1; round <= 2; round++) { + var marker = "inv" + round; + var entered = new CountDownLatch(1); + var release = new CountDownLatch(1); + var endValue = new AtomicReference(); + var startThread = new AtomicReference(); + var endThread = new AtomicReference(); + var previous = new AtomicReference(); + var endCalls = new AtomicInteger(); + var endOrder = new CopyOnWriteArrayList(); + var legacy = new DurableExecutionPlugin() { + @Override + public void onInvocationStart(InvocationInfo info) { + startThread.set(Thread.currentThread()); + previous.set(local.get()); + local.set(marker); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + endCalls.incrementAndGet(); + endOrder.add("legacy"); + endThread.set(Thread.currentThread()); + endValue.set(local.get()); + local.remove(); + } + }; + var plugins = new ArrayList(); + plugins.add(legacy); + if (mixed) + plugins.add( + 0, + new InvocationOtelPlugin( + SdkTracerProvider.builder(), + OtelPluginConfig.builder().enableMdc(false).build()) { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + endOrder.add("otel"); + super.onInvocationEnd(info); + } + }); + var config = DurableConfig.builder() + .withExecutorService(workers) + .withDurableExecutionClient(new LocalMemoryExecutionClient()) + .withCheckpointDelay(Duration.ZERO) + .withPlugins(plugins.toArray(DurableExecutionPlugin[]::new)) + .build(); + var result = caller.submit(() -> DurableExecutor.execute( + input(marker), + null, + TypeToken.get(String.class), + (value, context) -> { + entered.countDown(); + if (!precompleted) await(release); + if (pending) context.wait("pause", Duration.ofSeconds(1)); + return "done"; + }, + config)); + try { + assertTrue(entered.await(3, TimeUnit.SECONDS)); + if (!precompleted) { + awaitCallerJoin(callerThread.get()); + release.countDown(); + } + assertEquals( + pending ? ExecutionStatus.PENDING : ExecutionStatus.SUCCEEDED, + result.get(5, TimeUnit.SECONDS).status()); + var after = workers.submit(local::get).get(3, TimeUnit.SECONDS); + observations.add(new Observation( + startThread.get(), endThread.get(), endValue.get(), after, previous.get(), endCalls.get())); + assertEquals( + mixed ? List.of("otel", "legacy") : List.of("legacy"), + endOrder, + "scoped and legacy end hooks retain registration order and execute once"); + } finally { + release.countDown(); + } + } + for (int index = 0; index < observations.size(); index++) { + var observation = observations.get(index); + assertEquals(1, observation.endCalls()); + if (precompleted) { + assertSame(callerThread.get(), observation.endThread(), "completed futures keep caller dispatch"); + } else if (!pending) { + assertSame(observation.startThread(), observation.endThread()); + assertEquals("inv" + (index + 1), observation.endValue()); + assertNull(observation.workerAfter(), "legacy end hook clears its original worker value"); + assertNull(observation.previous(), "reused worker must not retain the prior invocation"); + } + // Suspension can complete from the handler or checkpoint processing; no owner-thread guarantee. + } + } finally { + caller.shutdownNow(); + workers.shutdownNow(); + assertTrue(caller.awaitTermination(3, TimeUnit.SECONDS)); + assertTrue(workers.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + private record Observation( + Thread startThread, Thread endThread, String endValue, String workerAfter, String previous, int endCalls) {} + + 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()) + .anyMatch(frame -> frame.getClassName().equals(CompletableFuture.class.getName()) + && frame.getMethodName().equals("join"))) return; + LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1)); + } + fail("invocation caller did not reach its future join before releasing the handler"); + } + + private static DurableExecutionInput input(String marker) { + var operation = Operation.builder() + .id("id") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails( + ExecutionDetails.builder().inputPayload("\"input\"").build()) + .build(); + return new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/" + marker + "/id", + "token", + CheckpointUpdatedExecutionState.builder().operations(operation).build()); + } + + private static Thread daemon(Runnable task, String name) { + var thread = new Thread(task, name); + thread.setDaemon(true); + return thread; + } + + private static void await(CountDownLatch latch) { + try { + assertTrue(latch.await(3, TimeUnit.SECONDS)); + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + throw new AssertionError(error); + } + } +} 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 428ebc8b6..490e6aa97 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 @@ -139,14 +139,14 @@ public static DurableExecutionOutput execute( // will be returned. Otherwise, it will complete exceptionally with a SuspendExecutionException or a // failure. try { - return awaitHandlerScopes( - executionManager.runUntilCompleteOrSuspend(handlerFuture), - handlerFuture, - hasHandlerScope, - lambdaContext, - pluginRunner.getPlugins().size(), - scopeFatal) - .handle((result, ex) -> { + return finalizeAfterHandlerScopes( + executionManager.runUntilCompleteOrSuspend(handlerFuture), + handlerFuture, + hasHandlerScope, + lambdaContext, + pluginRunner.getPlugins().size(), + scopeFatal, + (result, ex) -> { if (ex != null) { // an exception thrown from handlerFuture or suspension/termination occurred Throwable cause = ExceptionHelper.unwrapCompletableFuture(ex); @@ -219,8 +219,7 @@ public static DurableExecutionOutput execute( pluginExecutionInput.get(), result); return output; - }) - .join(); + }); } catch (CompletionException e) { // unwrap the CompletionException and rethrow the wrapped exception ExceptionHelper.sneakyThrow(ExceptionHelper.unwrapCompletableFuture(e)); @@ -251,6 +250,43 @@ private static CompletableFuture supplyHandler( return result; } + private static R finalizeAfterHandlerScopes( + CompletableFuture executionFuture, + CompletableFuture handlerFuture, + AtomicBoolean hasHandlerScope, + Context lambdaContext, + int pluginCount, + AtomicReference scopeFatal, + BiFunction finalizer) { + var started = new AtomicBoolean(); + var finalized = new CompletableFuture(); + // Attach before waiting so legacy hooks retain normal CompletableFuture completion-thread dispatch. + var ready = executionFuture.handle((value, failure) -> { + Runnable finish = () -> completeFinalization(started, finalized, () -> finalizer.apply(value, failure)); + if (!hasHandlerScope.get() || handlerFuture.isDone()) finish.run(); + else handlerFuture.whenComplete((ignored, ignoredFailure) -> finish.run()); + return finish; + }); + var finish = ready.join(); + if (!started.get()) { + // Only the caller waits: a signaling handler must be free to unwind and close its scopes. + awaitHandlerScopes(executionFuture, handlerFuture, hasHandlerScope, lambdaContext, pluginCount, scopeFatal); + finish.run(); + } + return finalized.join(); + } + + private static void completeFinalization( + AtomicBoolean started, CompletableFuture result, Supplier finalizer) { + if (!started.compareAndSet(false, true)) return; + try { + result.complete(finalizer.get()); + } catch (Throwable failure) { + // Match CompletableFuture.handle: transfer failures unchanged for the invocation caller to rethrow. + result.completeExceptionally(failure); + } + } + static CompletableFuture awaitHandlerScopes( CompletableFuture executionFuture, CompletableFuture handlerFuture, From de9ea32ae9628ad4ec64a6489838abffbe1abcfb Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Tue, 6 Oct 2026 12:57:30 -0700 Subject: [PATCH 09/48] ci: queue serial conformance runs without replacing pending work --- .github/workflows/conformance-tests.yml | 1 + .github/workflows/e2e-tests.yml | 1 + .github/workflows/otel-conformance-tests.yml | 5 +++++ 3 files changed, 7 insertions(+) diff --git a/.github/workflows/conformance-tests.yml b/.github/workflows/conformance-tests.yml index 313c6c589..71f459dc7 100644 --- a/.github/workflows/conformance-tests.yml +++ b/.github/workflows/conformance-tests.yml @@ -25,6 +25,7 @@ concurrency: # same stack -- mirrors e2e-tests.yml. group: conformance-tests cancel-in-progress: false + queue: max permissions: contents: read diff --git a/.github/workflows/e2e-tests.yml b/.github/workflows/e2e-tests.yml index 0b8a0c296..0058552b9 100644 --- a/.github/workflows/e2e-tests.yml +++ b/.github/workflows/e2e-tests.yml @@ -26,6 +26,7 @@ on: concurrency: group: e2e-tests cancel-in-progress: false + queue: max # permission can be added at job level or workflow level permissions: diff --git a/.github/workflows/otel-conformance-tests.yml b/.github/workflows/otel-conformance-tests.yml index 68b3bc3f2..43a3e693f 100644 --- a/.github/workflows/otel-conformance-tests.yml +++ b/.github/workflows/otel-conformance-tests.yml @@ -53,6 +53,11 @@ on: permissions: {} +concurrency: + group: otel-conformance-tests + cancel-in-progress: false + queue: max + jobs: opentelemetry: # 1. Run for non-PR events, such as scheduled runs and manual invocations From b114ab75aa220ea6c5c2bb5794930033c38fee5d Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Tue, 6 Oct 2026 22:51:44 +0000 Subject: [PATCH 10/48] Register the replay validation test caller as active --- .../amazon/lambda/durable/ReplayValidationTest.java | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/sdk/src/test/java/software/amazon/lambda/durable/ReplayValidationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/ReplayValidationTest.java index 98f505ad6..bf62d52fa 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/ReplayValidationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/ReplayValidationTest.java @@ -3,6 +3,7 @@ package software.amazon.lambda.durable; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -48,7 +49,9 @@ private DurableContext createTestContext(List initialOperations) { null); var context = DurableContextImpl.createRootContext( executionManager, DurableConfig.builder().build(), null); - executionManager.setCurrentThreadContext(new ThreadContext(EXECUTION_OP_ID + "-execution", ThreadType.CONTEXT)); + // Mirror DurableExecutor: the caller must stay active if a worker finishes before step().get() waits. + executionManager.registerActiveThread(null); + executionManager.setCurrentThreadContext(new ThreadContext(null, ThreadType.CONTEXT)); return context; } @@ -59,7 +62,7 @@ void shouldPassValidationWhenNoCheckpointExists() { var context = createTestContext(List.of()); // When & Then: Should not throw - assertDoesNotThrow(() -> context.step("test", String.class, stepCtx -> "result")); + assertEquals("result", context.step("test", String.class, stepCtx -> "result")); } @Test From ee73b45ea3e69aa07e068d31449fcf6035ccabf2 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 7 Oct 2026 01:34:41 +0000 Subject: [PATCH 11/48] ci: run OTel conformance on validated CodeBuild Java 21 --- .github/workflows/otel-conformance-tests.yml | 33 +++++++++++++++++--- 1 file changed, 29 insertions(+), 4 deletions(-) diff --git a/.github/workflows/otel-conformance-tests.yml b/.github/workflows/otel-conformance-tests.yml index 43a3e693f..4f948b298 100644 --- a/.github/workflows/otel-conformance-tests.yml +++ b/.github/workflows/otel-conformance-tests.yml @@ -67,8 +67,9 @@ jobs: actions: write contents: read id-token: write - uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@43872f8d5dd917fe3ee6d01a6c953c1ff67c5fe4 + uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@a66037abbbfa55fde97f714e30f0bc262edefd63 with: + runs_on: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.repo.full_name != github.repository && 'ubuntu-latest' || format('codebuild-github-actions-runner-{0}-{1}', github.run_id, github.run_attempt) }} language: java resource_prefix: j sdk_repository: aws/aws-durable-execution-sdk-java @@ -80,13 +81,37 @@ jobs: # checked out (.build/durable-sdk). examples_dir: .build/durable-sdk/conformance-tests-otel setup_command: | - if [ -z "${JAVA_HOME_21_X64:-}" ]; then - echo "The runner does not provide Java 21" + set -euo pipefail + if [ -n "${JAVA_HOME_21_X64:-}" ]; then + export JAVA_HOME="$JAVA_HOME_21_X64" + elif [ -x /usr/lib/jvm/java-21-amazon-corretto/bin/java ]; then + export JAVA_HOME=/usr/lib/jvm/java-21-amazon-corretto + elif [ -z "${JAVA_HOME:-}" ]; then + export JAVA_HOME="$(dirname "$(dirname "$(readlink -f "$(command -v java)")")")" + fi + if [ ! -x "$JAVA_HOME/bin/java" ] || [ ! -x "$JAVA_HOME/bin/javac" ]; then + echo "Selected JAVA_HOME=$JAVA_HOME is not a complete JDK" exit 1 fi - export JAVA_HOME="$JAVA_HOME_21_X64" export PATH="$JAVA_HOME/bin:$PATH" + java_spec=$(java -XshowSettings:properties -version 2>&1 | awk '$1 == "java.specification.version" {print $3}') + if [ "$java_spec" != "21" ]; then + echo "Java 21 is required; selected JAVA_HOME=$JAVA_HOME reports specification version $java_spec" + exit 1 + fi + javac_version=$(javac -version 2>&1) + if [[ ! "$javac_version" =~ ^javac[[:space:]]21([.]|[[:space:]]|$) ]]; then + echo "Java 21 javac is required; selected compiler reports $javac_version" + exit 1 + fi java -version + javac -version + if [ -n "${GITHUB_ENV:-}" ]; then + echo "JAVA_HOME=$JAVA_HOME" >> "$GITHUB_ENV" + fi + if [ -n "${GITHUB_PATH:-}" ]; then + echo "$JAVA_HOME/bin" >> "$GITHUB_PATH" + fi prepare_command: | JAVA_SDK_VERSION=$( mvn -B -q \ From 0e0daeb948637eb497f0ee4c589c101415ed476f Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 7 Oct 2026 03:18:05 +0000 Subject: [PATCH 12/48] fix: preserve end-hook completion ownership and MDC lifetime --- .../otel/HandlerMdcIntegrationTest.java | 98 +++++++++++++++ .../otel/InvocationEndCompatibilityTest.java | 112 ++++++++++++++++++ .../durable/execution/DurableExecutor.java | 40 +++---- .../execution/HandlerScopeHandoffTest.java | 80 +++++++++++++ 4 files changed, 307 insertions(+), 23 deletions(-) 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 c525de2ba..7fe9f18cf 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 @@ -6,11 +6,16 @@ import io.opentelemetry.sdk.trace.SdkTracerProvider; import java.time.Duration; +import java.util.Arrays; import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.locks.LockSupport; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.CsvSource; import org.junit.jupiter.params.provider.ValueSource; @@ -19,6 +24,8 @@ import software.amazon.lambda.durable.TypeToken; 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.serde.JacksonSerDes; import software.amazon.lambda.durable.serde.SerDes; import software.amazon.lambda.durable.testing.LocalDurableTestRunner; @@ -174,6 +181,97 @@ public T deserialize(String data, TypeToken type) { }; } + @ParameterizedTest + @ValueSource(strings = {"success", "failure", "inputFailure"}) + void endHookObservesTaskMdcBeforeReusedWorkerAmbientStateIsRestored(String outcome) throws Exception { + var worker = Executors.newSingleThreadExecutor(); + var callerThread = new AtomicReference(); + var caller = Executors.newSingleThreadExecutor(task -> { + var thread = new Thread(task, "mdc-end-caller"); + callerThread.set(thread); + return thread; + }); + var started = new CountDownLatch(1); + var release = new CountDownLatch(1); + var startThread = new AtomicReference(); + var endThread = new AtomicReference(); + var endMdc = new AtomicReference>(); + var ambient = Map.of("worker", "ambient"); + var plugin = new DurableExecutionPlugin() { + @Override + public void onInvocationStart(InvocationInfo info) { + startThread.set(Thread.currentThread()); + MDC.put("start-hook", "request"); + started.countDown(); + awaitMdcLatch(release); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + endThread.set(Thread.currentThread()); + var values = MDC.getCopyOfContextMap(); + endMdc.set(values == null ? Map.of() : values); + MDC.put("end-hook", "must-not-leak"); + } + }; + try { + worker.submit(() -> MDC.setContextMap(ambient)).get(3, TimeUnit.SECONDS); + var config = DurableConfig.builder() + .withExecutorService(worker) + .withPlugins(plugin) + .withSerDes(outcome.equals("inputFailure") ? failingInputSerDes() : new JacksonSerDes()) + .build(); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + if (outcome.equals("failure")) throw new IllegalStateException("body failure"); + return "done"; + }, + config); + var result = caller.submit(() -> runner.run("input")); + assertTrue(started.await(3, TimeUnit.SECONDS)); + awaitMdcCallerJoin(callerThread.get()); + release.countDown(); + assertEquals( + outcome.equals("success") ? ExecutionStatus.SUCCEEDED : ExecutionStatus.FAILED, + result.get(5, TimeUnit.SECONDS).getStatus()); + assertSame(startThread.get(), endThread.get()); + // Input failure skips DurableLogger; entered handlers retain its established MDC-clearing behavior. + assertEquals( + outcome.equals("inputFailure") ? Map.of("worker", "ambient", "start-hook", "request") : Map.of(), + endMdc.get(), + "End hooks must not see prematurely restored ambient worker MDC"); + assertEquals(ambient, worker.submit(MDC::getCopyOfContextMap).get(3, TimeUnit.SECONDS)); + } finally { + release.countDown(); + caller.shutdownNow(); + worker.shutdownNow(); + assertTrue(caller.awaitTermination(3, TimeUnit.SECONDS)); + assertTrue(worker.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + private static void awaitMdcLatch(CountDownLatch latch) { + try { + assertTrue(latch.await(3, TimeUnit.SECONDS)); + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + throw new AssertionError(error); + } + } + + 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()) + .anyMatch(frame -> frame.getClassName().equals(CompletableFuture.class.getName()) + && frame.getMethodName().equals("join"))) return; + LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1)); + } + fail("Caller did not attach its completion observer before releasing the handler"); + } + private static void assertRestored( ExecutorService executor, Map workerMdc, Map callerMdc) throws Exception { var callerAfter = MDC.getCopyOfContextMap(); diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndCompatibilityTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndCompatibilityTest.java index d36280d50..7104be813 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndCompatibilityTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndCompatibilityTest.java @@ -14,6 +14,8 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.locks.LockSupport; +import java.util.function.Function; +import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.CsvSource; import software.amazon.awssdk.services.lambda.model.*; @@ -23,6 +25,7 @@ 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.HandlerScoped; import software.amazon.lambda.durable.plugin.InvocationEndInfo; import software.amazon.lambda.durable.plugin.InvocationInfo; import software.amazon.lambda.durable.testing.local.LocalMemoryExecutionClient; @@ -164,6 +167,115 @@ public void onInvocationEnd(InvocationEndInfo info) { } } + @Test + void scopedSuspensionWithinBudgetKeepsLegacyCleanupOnTheReusedWorker() throws Exception { + var local = new ThreadLocal(); + var callerThread = new AtomicReference(); + var callers = Executors.newSingleThreadExecutor(task -> { + var thread = daemon(task, "scoped-handoff-caller"); + callerThread.set(thread); + return thread; + }); + var workers = Executors.newSingleThreadExecutor(task -> daemon(task, "scoped-handoff-worker")); + try { + for (var round = 0; round < 3; round++) { + var marker = "scoped-" + round; + var scoped = new PausingOtelPlugin(); + var previous = new AtomicReference(); + var observed = new AtomicReference(); + var owner = new AtomicReference(); + var endThread = new AtomicReference(); + var legacy = new DurableExecutionPlugin() { + @Override + public void onInvocationStart(InvocationInfo info) { + owner.set(Thread.currentThread()); + previous.set(local.get()); + local.set(marker); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + endThread.set(Thread.currentThread()); + observed.set(local.get()); + local.remove(); + } + }; + var config = DurableConfig.builder() + .withExecutorService(workers) + .withDurableExecutionClient(new LocalMemoryExecutionClient()) + .withCheckpointDelay(Duration.ZERO) + .withPlugins(scoped, legacy) + .build(); + var result = callers.submit(() -> DurableExecutor.execute( + input(marker), + null, + TypeToken.get(String.class), + (value, context) -> { + context.wait("pause", Duration.ofSeconds(1)); + return "done"; + }, + config)); + try { + assertTrue(scoped.closeEntered.await(3, TimeUnit.SECONDS)); + awaitCallerHandoff(callerThread.get()); + scoped.releaseClose.countDown(); + assertEquals( + ExecutionStatus.PENDING, + result.get(5, TimeUnit.SECONDS).status()); + assertSame(owner.get(), endThread.get()); + assertEquals(marker, observed.get()); + assertNull(previous.get(), "The prior invocation's end hook must clear the reused worker"); + assertNull(workers.submit(local::get).get(3, TimeUnit.SECONDS)); + } finally { + scoped.releaseClose.countDown(); + } + } + } finally { + callers.shutdownNow(); + workers.shutdownNow(); + assertTrue(callers.awaitTermination(3, TimeUnit.SECONDS)); + assertTrue(workers.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + private static void awaitCallerHandoff(Thread caller) { + var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(3); + while (System.nanoTime() < deadline) { + if (Arrays.stream(caller.getStackTrace()) + .anyMatch(frame -> frame.getClassName().equals(DurableExecutor.class.getName()) + && frame.getMethodName().equals("awaitHandlerScopes"))) return; + LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1)); + } + fail("The caller did not begin its bounded cleanup handoff"); + } + + @HandlerScoped(PausingOtelPlugin.Opener.class) + private static final class PausingOtelPlugin extends InvocationOtelPlugin { + private final CountDownLatch closeEntered = new CountDownLatch(1); + private final CountDownLatch releaseClose = new CountDownLatch(1); + + private PausingOtelPlugin() { + super( + SdkTracerProvider.builder(), + OtelPluginConfig.builder().enableMdc(false).build()); + } + + public static final class Opener implements Function { + @Override + public AutoCloseable apply(PausingOtelPlugin plugin) { + var actual = new InvocationOtelPlugin.HandlerScopeOpener().apply(plugin); + return () -> { + try { + if (actual != null) actual.close(); + } finally { + plugin.closeEntered.countDown(); + await(plugin.releaseClose); + } + }; + } + } + } + private record Observation( Thread startThread, Thread endThread, String endValue, String workerAfter, String previous, int endCalls) {} 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 490e6aa97..713066869 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,7 +80,7 @@ public static DurableExecutionOutput execute( var pluginExecutionInput = new AtomicReference<>(); var hasHandlerScope = new AtomicBoolean(); var handlerFuture = supplyHandler( - preservingMdc(pluginRunner, () -> { + () -> { executionManager.setCurrentThreadContext(new ThreadContext(null, ThreadType.CONTEXT)); // Deserialize once and share the value with the plugin hooks and the handler below. A second @@ -130,7 +130,7 @@ public static DurableExecutionOutput execute( () -> hasHandlerScope.set(true), fatal -> scopeFatal.compareAndSet(null, fatal)); } - }), + }, config.getExecutorService(), pluginRunner, scopeFatal); // Get executor from config for running user code @@ -234,16 +234,18 @@ private static CompletableFuture supplyHandler( var result = new CompletableFuture(); var caller = Thread.currentThread(); Runnable work = (Runnable & CompletableFuture.AsynchronousCompletionTask) () -> { - try { - result.complete(task.get()); - } catch (Throwable failure) { - result.completeExceptionally(failure); - // Preserve legacy user-body Error handling. Only a positively reported new scope fatal also - // escapes its actual owner after its observation future settles. Direct callers finalize below. - var fatal = scopeFatal.get(); - if (Thread.currentThread() != caller - && fatal != null - && ExceptionHelper.unwrapCompletableFuture(failure) == fatal) throw fatal; + try (var ignored = restoreMdcOnClose()) { + try { + result.complete(task.get()); + } catch (Throwable failure) { + result.completeExceptionally(failure); + // Preserve legacy user-body Error handling. Only a positively reported new scope fatal also + // escapes its actual owner after its observation future settles. Direct callers finalize below. + var fatal = scopeFatal.get(); + if (Thread.currentThread() != caller + && fatal != null + && ExceptionHelper.unwrapCompletableFuture(failure) == fatal) throw fatal; + } } }; executor.execute(work); @@ -271,7 +273,9 @@ private static R finalizeAfterHandlerScopes( if (!started.get()) { // Only the caller waits: a signaling handler must be free to unwind and close its scopes. awaitHandlerScopes(executionFuture, handlerFuture, hasHandlerScope, lambdaContext, pluginCount, scopeFatal); - finish.run(); + // Completion publishes the handler result before its callbacks necessarily run. Once cleanup completed, + // let its registered callback finalize on the owner instead of racing it from this waiter. + if (!handlerFuture.isDone()) finish.run(); } return finalized.join(); } @@ -322,16 +326,6 @@ private static void throwIfScopeFatal(AtomicReference scopeFatal) { if (fatal != null) throw fatal; } - /** Restores the worker even when plugin startup runs before a failing input deserialization. */ - private static Supplier preservingMdc(PluginRunner plugins, Supplier task) { - if (plugins.isEmpty()) return task; - return () -> { - try (var ignored = restoreMdcOnClose()) { - return task.get(); - } - }; - } - private static SafeCloseable restoreMdcOnClose() { var previous = MDC.getCopyOfContextMap(); return () -> { 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 0f107f043..15574a55b 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 @@ -9,10 +9,15 @@ import java.time.Duration; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executors; +import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BiConsumer; +import java.util.function.BiFunction; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; @@ -85,6 +90,81 @@ void noScopeDoesNotConsultTheRemainingTimeOrWaitForTheHandler() { assertFalse(handler.isDone()); } + @Test + void completedHandoffWaitsForTheRegisteredOwnerFinalizer() throws Exception { + var callbackRegistered = new CountDownLatch(1); + var callbacks = new LinkedBlockingQueue(); + var handler = new CompletableFuture() { + @Override + public CompletableFuture whenComplete(BiConsumer action) { + callbackRegistered.countDown(); + // A completed future may publish its result before its registered completion action executes. + return super.whenCompleteAsync(action, callbacks::add); + } + }; + var local = new ThreadLocal(); + var owner = new AtomicReference(); + var finalizerThread = new AtomicReference(); + var finalizerValue = new AtomicReference(); + var workers = Executors.newSingleThreadExecutor(); + var callers = Executors.newSingleThreadExecutor(); + var method = DurableExecutor.class.getDeclaredMethod( + "finalizeAfterHandlerScopes", + CompletableFuture.class, + CompletableFuture.class, + AtomicBoolean.class, + Context.class, + int.class, + AtomicReference.class, + BiFunction.class); + method.setAccessible(true); + BiFunction end = (value, failure) -> { + finalizerThread.set(Thread.currentThread()); + finalizerValue.set(local.get()); + local.remove(); + return "pending"; + }; + try { + var result = callers.submit(() -> method.invoke( + null, + CompletableFuture.failedFuture(control(false)), + handler, + new AtomicBoolean(true), + null, + 2, + new AtomicReference(), + end)); + assertTrue(callbackRegistered.await(3, TimeUnit.SECONDS)); + workers.submit(() -> { + owner.set(Thread.currentThread()); + local.set("invocation"); + handler.complete("handler finished"); + }) + .get(3, TimeUnit.SECONDS); + assertTrue(handler.isDone()); + assertThrows( + TimeoutException.class, + () -> result.get(100, TimeUnit.MILLISECONDS), + "The waiter must not steal a completed handoff's pending completion callback"); + var callback = callbacks.poll(3, TimeUnit.SECONDS); + assertNotNull(callback); + workers.submit(callback).get(3, TimeUnit.SECONDS); + assertEquals("pending", result.get(3, TimeUnit.SECONDS)); + assertSame(owner.get(), finalizerThread.get()); + assertEquals("invocation", finalizerValue.get()); + assertNull(workers.submit(local::get).get(3, TimeUnit.SECONDS), "The reused owner must be clean"); + } finally { + handler.complete("cleanup"); + Runnable callback; + while ((callback = callbacks.poll()) != null) + workers.submit(callback).get(3, TimeUnit.SECONDS); + callers.shutdownNow(); + workers.shutdownNow(); + assertTrue(callers.awaitTermination(3, TimeUnit.SECONDS)); + assertTrue(workers.awaitTermination(3, TimeUnit.SECONDS)); + } + } + private static CompletableFuture handoff( CompletableFuture execution, CompletableFuture handler) { return DurableExecutor.awaitHandlerScopes( From b5c688c10e6ba4961c4a4ce1fa8239bb9e552c99 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 7 Oct 2026 03:36:02 +0000 Subject: [PATCH 13/48] fix: settle handler future when MDC capture fails --- .../durable/execution/DurableExecutor.java | 10 ++- .../execution/HandlerScopeHandoffTest.java | 80 +++++++++++++++++++ 2 files changed, 89 insertions(+), 1 deletion(-) 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 713066869..8641d4444 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 @@ -234,7 +234,15 @@ private static CompletableFuture supplyHandler( var result = new CompletableFuture(); var caller = Thread.currentThread(); Runnable work = (Runnable & CompletableFuture.AsynchronousCompletionTask) () -> { - try (var ignored = restoreMdcOnClose()) { + SafeCloseable mdcRestore; + try { + mdcRestore = restoreMdcOnClose(); + } catch (Throwable failure) { + // Initialization can fail before task execution; the observer must still receive that failure. + result.completeExceptionally(failure); + return; + } + try (var ignored = mdcRestore) { try { result.complete(task.get()); } catch (Throwable failure) { 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 15574a55b..8db315f81 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 @@ -7,9 +7,13 @@ import com.amazonaws.services.lambda.runtime.Context; import java.time.Duration; +import java.util.List; +import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; import java.util.concurrent.Executors; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; @@ -18,11 +22,15 @@ import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiConsumer; import java.util.function.BiFunction; +import java.util.function.Supplier; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; +import org.slf4j.MDC; import software.amazon.awssdk.services.lambda.model.ErrorObject; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; +import software.amazon.lambda.durable.plugin.PluginRunner; class HandlerScopeHandoffTest { @ParameterizedTest @@ -165,6 +173,78 @@ public CompletableFuture whenComplete(BiConsumer worker.execute(() -> { + // Static mocks are thread-scoped: fail only the real asynchronous worker's MDC capture. + try (var mdc = mockStatic(MDC.class, CALLS_REAL_METHODS)) { + mdc.when(MDC::getCopyOfContextMap).thenThrow(failure); + task.run(); + } + }); + var supply = DurableExecutor.class.getDeclaredMethod( + "supplyHandler", Supplier.class, Executor.class, PluginRunner.class, AtomicReference.class); + supply.setAccessible(true); + try { + Supplier task = () -> { + taskCalled.set(true); + return "unexpected"; + }; + var result = (CompletableFuture) supply.invoke( + null, + task, + failingCaptureWorker, + new PluginRunner(List.of(new DurableExecutionPlugin() {})), + new AtomicReference()); + var thrown = assertThrows(ExecutionException.class, () -> result.get(2, TimeUnit.SECONDS)); + assertSame( + failure, thrown.getCause(), "Initialization failure must settle the observation future unchanged"); + assertFalse(taskCalled.get(), "Failed MDC capture must not start user/plugin work"); + } finally { + worker.shutdownNow(); + assertTrue(worker.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + @Test + void mdcRestorationFailureAfterCompletionStillEscapesItsOwner() throws Exception { + var failure = new IllegalStateException("MDC restoration failed"); + var ownerFailure = new CompletableFuture(); + var worker = Executors.newSingleThreadExecutor(); + Executor failingRestoreWorker = task -> worker.execute(() -> { + try (var mdc = mockStatic(MDC.class, CALLS_REAL_METHODS)) { + mdc.when(MDC::getCopyOfContextMap).thenReturn(Map.of("worker", "ambient")); + mdc.when(() -> MDC.setContextMap(Map.of("worker", "ambient"))).thenThrow(failure); + try { + task.run(); + } catch (Throwable thrown) { + ownerFailure.complete(thrown); + } + } + }); + var supply = DurableExecutor.class.getDeclaredMethod( + "supplyHandler", Supplier.class, Executor.class, PluginRunner.class, AtomicReference.class); + supply.setAccessible(true); + try { + Supplier task = () -> "completed"; + var result = (CompletableFuture) supply.invoke( + null, + task, + failingRestoreWorker, + new PluginRunner(List.of(new DurableExecutionPlugin() {})), + new AtomicReference()); + assertEquals("completed", result.get(2, TimeUnit.SECONDS)); + assertSame(failure, ownerFailure.get(2, TimeUnit.SECONDS)); + assertEquals("completed", result.join(), "Post-completion restoration must not rewrite the result"); + } finally { + worker.shutdownNow(); + assertTrue(worker.awaitTermination(3, TimeUnit.SECONDS)); + } + } + private static CompletableFuture handoff( CompletableFuture execution, CompletableFuture handler) { return DurableExecutor.awaitHandlerScopes( From d5bd3b8e12fc862737f6d2c50e261878f0c147bd Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Tue, 6 Oct 2026 21:34:05 -0700 Subject: [PATCH 14/48] fix: wake invocation caller on reported scope fatal --- .../durable/otel/ScopeFatalSignalTest.java | 93 +++++++++++++++++++ .../durable/execution/DurableExecutor.java | 30 +++++- 2 files changed, 119 insertions(+), 4 deletions(-) create mode 100644 otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ScopeFatalSignalTest.java diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ScopeFatalSignalTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ScopeFatalSignalTest.java new file mode 100644 index 000000000..faf6a5775 --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ScopeFatalSignalTest.java @@ -0,0 +1,93 @@ +// 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 java.util.concurrent.*; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; +import software.amazon.lambda.durable.plugin.HandlerScoped; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; +import software.amazon.lambda.durable.util.ExceptionHelper; + +class ScopeFatalSignalTest { + @ParameterizedTest + @CsvSource({"false,false", "true,false", "false,true", "true,true"}) + @SuppressWarnings("removal") + void fatalCloseWakesCallerAndStillAllowsBoundedEarlierCleanup(boolean threadDeath, boolean releaseInBudget) + throws Exception { + var earlierEntered = new CountDownLatch(1); + var release = new CountDownLatch(1); + var closed = new AtomicBoolean(); + var ownerFatal = new AtomicReference(); + var ownerObserved = new CountDownLatch(1); + Error fatal = threadDeath ? new ThreadDeath() : new InternalError("later scope close"); + var workers = Executors.newSingleThreadExecutor(task -> { + var thread = new Thread(task, "scope-fatal-owner"); + thread.setUncaughtExceptionHandler((owner, failure) -> { + ownerFatal.set(failure); + ownerObserved.countDown(); + }); + return thread; + }); + var callers = Executors.newSingleThreadExecutor(); + try { + var earlier = new Scoped(() -> { + earlierEntered.countDown(); + if (!release.await(5, TimeUnit.SECONDS)) throw new AssertionError("release timed out"); + closed.set(true); + }); + var later = new Scoped(() -> { + throw fatal; + }); + var config = DurableConfig.builder() + .withExecutorService(workers) + .withPlugins(earlier, later) + .build(); + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> input, config); + var result = callers.submit(() -> runner.run("input")); + assertTrue(earlierEntered.await(3, TimeUnit.SECONDS)); + assertThrows( + TimeoutException.class, + () -> result.get(100, TimeUnit.MILLISECONDS), + "earlier cleanup must receive its handoff budget"); + if (releaseInBudget) release.countDown(); + var failure = assertThrows( + ExecutionException.class, + () -> result.get(2, TimeUnit.SECONDS), + "reported fatal must wake the invocation without waiting forever for earlier cleanup"); + assertSame(fatal, ExceptionHelper.unwrapCompletableFuture(failure.getCause())); + assertEquals(releaseInBudget, closed.get()); + } finally { + release.countDown(); + callers.shutdown(); + workers.shutdown(); + assertTrue(callers.awaitTermination(3, TimeUnit.SECONDS)); + assertTrue(workers.awaitTermination(3, TimeUnit.SECONDS)); + } + assertTrue(ownerObserved.await(3, TimeUnit.SECONDS)); + assertSame(fatal, ownerFatal.get(), "the fatal still escapes its actual owner"); + } + + @HandlerScoped(Opener.class) + public static final class Scoped implements DurableExecutionPlugin { + private final AutoCloseable closer; + + Scoped(AutoCloseable closer) { + this.closer = closer; + } + } + + public static final class Opener implements Function { + @Override + public AutoCloseable apply(Object plugin) { + return ((Scoped) plugin).closer; + } + } +} 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 8641d4444..1f4865642 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 @@ -70,6 +70,7 @@ public static DurableExecutionOutput execute( DurableConfig config) { var pluginRunner = config.getPluginRunner(); var scopeFatal = new AtomicReference(); + var scopeFailure = new CompletableFuture(); try (var executionManager = new ExecutionManager(input, config, lambdaContext)) { var isFirstInvocation = !executionManager.isReplaying(); var requestId = lambdaContext != null ? lambdaContext.getAwsRequestId() : null; @@ -128,7 +129,10 @@ public static DurableExecutionOutput execute( return pluginRunner.runHandler( () -> handler.apply(handlerInput, context), () -> hasHandlerScope.set(true), - fatal -> scopeFatal.compareAndSet(null, fatal)); + fatal -> { + scopeFatal.compareAndSet(null, fatal); + scopeFailure.completeExceptionally(scopeFatal.get()); + }); } }, config.getExecutorService(), @@ -139,8 +143,9 @@ public static DurableExecutionOutput execute( // will be returned. Otherwise, it will complete exceptionally with a SuspendExecutionException or a // failure. try { + var executionFuture = executionManager.runUntilCompleteOrSuspend(handlerFuture); return finalizeAfterHandlerScopes( - executionManager.runUntilCompleteOrSuspend(handlerFuture), + pluginRunner.isEmpty() ? executionFuture : observeScopeFailure(executionFuture, scopeFailure), handlerFuture, hasHandlerScope, lambdaContext, @@ -260,6 +265,20 @@ private static CompletableFuture supplyHandler( return result; } + /** Wakes the caller when a reported scope fatal precedes completion of earlier scope cleanup. */ + private static CompletableFuture observeScopeFailure( + CompletableFuture execution, CompletableFuture scopeFailure) { + var result = new CompletableFuture(); + execution.whenComplete((value, failure) -> { + if (failure == null) result.complete(value); + else result.completeExceptionally(failure); + }); + scopeFailure.whenComplete((ignored, failure) -> { + if (failure != null) result.completeExceptionally(failure); + }); + return result; + } + private static R finalizeAfterHandlerScopes( CompletableFuture executionFuture, CompletableFuture handlerFuture, @@ -309,8 +328,11 @@ static CompletableFuture awaitHandlerScopes( // This method runs on the invocation caller, never as a callback on the signaling handler worker. // Preserve the winning outcome except for an observed fatal error from the new scope callbacks. var failure = executionFuture.handle((result, error) -> error).join(); - throwIfScopeFatal(scopeFatal); - if (failure == null || !hasHandlerScope.get()) return executionFuture; + // A reported fatal can precede completion of earlier scopes; retain their bounded cleanup budget. + if (!hasHandlerScope.get() || (failure == null && scopeFatal.get() == null)) { + throwIfScopeFatal(scopeFatal); + return executionFuture; + } var reserve = SHUTDOWN_RESPONSE_RESERVE_MILLIS + PLUGIN_FINALIZATION_RESERVE_MILLIS * pluginCount; var budgetMillis = lambdaContext == null ? 500L From 0d8d8604e9d03efd71d4a3429ca701dd415f7c27 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 7 Oct 2026 18:50:53 +0000 Subject: [PATCH 15/48] fix: preserve execution outcomes across end MDC failures --- .../durable/otel/MdcFailureBoundaryTest.java | 314 ++++++++++++++++++ .../durable/execution/DurableExecutor.java | 75 ++++- 2 files changed, 385 insertions(+), 4 deletions(-) create mode 100644 otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFailureBoundaryTest.java diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFailureBoundaryTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFailureBoundaryTest.java new file mode 100644 index 000000000..015b6a0f5 --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFailureBoundaryTest.java @@ -0,0 +1,314 @@ +// 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 java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Proxy; +import java.lang.reflect.UndeclaredThrowableException; +import java.time.Duration; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.concurrent.AbstractExecutorService; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.ValueSource; +import org.slf4j.MDC; +import org.slf4j.spi.MDCAdapter; +import software.amazon.lambda.durable.DurableConfig; +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.testing.LocalDurableTestRunner; + +class MdcFailureBoundaryTest { + @Test + void workerCaptureFailureDoesNotEndAnInvocationThatNeverStarted() throws Exception { + var original = MDC.getMDCAdapter(); + var injected = new AtomicBoolean(); + var failure = new IllegalStateException("worker MDC capture failed"); + var plugin = new RecordingPlugin(); + var bodyCalls = new AtomicInteger(); + var executor = new CompletingWorker(); + try { + replaceAdapter(proxy(original, (method) -> { + if (method.equals("getCopyOfContextMap") + && Thread.currentThread().getName().equals("mdc-initialization-worker") + && injected.compareAndSet(false, true)) throw failure; + return null; + })); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + bodyCalls.incrementAndGet(); + return "done"; + }, + DurableConfig.builder() + .withExecutorService(executor) + .withPlugins(plugin) + .build()); + var result = runner.run("input"); + assertTrue(injected.get()); + assertEquals(ExecutionStatus.FAILED, result.getStatus()); + assertEquals(failure.getMessage(), result.getError().orElseThrow().errorMessage()); + assertEquals(0, bodyCalls.get()); + assertEquals(0, plugin.starts.get()); + assertEquals(0, plugin.ends.get(), "no end hook may consume stale state before start dispatch"); + } finally { + try { + executor.shutdownNow(); + assertTrue(executor.awaitTermination(3, TimeUnit.SECONDS)); + } finally { + replaceAdapter(original); + } + } + } + + @ParameterizedTest + @CsvSource({ + "getCopyOfContextMap,false", + "clear,false", + "setContextMap,false", + "getCopyOfContextMap,true", + "clear,true", + "setContextMap,true" + }) + void nonfatalEndMdcFailuresPreserveTheSelectedOutcome(String method, boolean suspend) throws Exception { + runEndFailure(method, suspend, new IllegalStateException("end MDC failure")); + } + + @SuppressWarnings("removal") + @ParameterizedTest + @CsvSource({ + "getCopyOfContextMap,false", + "clear,false", + "setContextMap,false", + "getCopyOfContextMap,true", + "clear,true", + "setContextMap,true" + }) + void fatalEndMdcFailuresStillEscape(String method, boolean wrapped) throws Exception { + Error fatal = wrapped ? new ThreadDeath() : new VirtualMachineError("fatal end MDC failure") {}; + runEndFailure(method, false, wrapped ? new CompletionException(new ExecutionException(fatal)) : fatal, fatal); + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void unreadableOrdinaryCausePreservesSuccessAndSuspension(boolean suspend) throws Exception { + var reads = new AtomicInteger(); + var failure = new CompletionException("unreadable", null) { + @Override + public synchronized Throwable getCause() { + reads.incrementAndGet(); + throw new IllegalStateException("diagnostic unavailable"); + } + }; + runEndFailure("getCopyOfContextMap", suspend, failure); + assertEquals(1, reads.get()); + } + + @SuppressWarnings("removal") + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void fatalCauseAccessorEscapesWithOriginalIdentity(boolean threadDeath) throws Exception { + var reads = new AtomicInteger(); + Error fatal = threadDeath ? new ThreadDeath() : new VirtualMachineError("fatal accessor") {}; + var failure = new CompletionException("unreadable", null) { + @Override + public synchronized Throwable getCause() { + reads.incrementAndGet(); + throw fatal; + } + }; + runEndFailure("setContextMap", false, failure, fatal); + assertEquals(1, reads.get()); + } + + @SuppressWarnings("removal") + @ParameterizedTest + @CsvSource({"reflection,false", "reflection,true", "proxy,false", "proxy,true"}) + void jdkReflectiveWrappersKeepFatalIdentity(String wrapper, boolean threadDeath) throws Exception { + Error fatal = threadDeath ? new ThreadDeath() : new VirtualMachineError("wrapped fatal") {}; + Throwable failure = wrapper.equals("reflection") + ? new InvocationTargetException(fatal) + : new UndeclaredThrowableException(fatal); + runEndFailure("clear", false, failure, fatal); + } + + @Test + void cyclicCauseIsReadOnlyOnceAndPreservesOutcome() throws Exception { + var reads = new AtomicInteger(); + var cycle = new CompletionException("cycle", null) { + @Override + public synchronized Throwable getCause() { + reads.incrementAndGet(); + return this; + } + }; + runEndFailure("getCopyOfContextMap", false, cycle); + assertEquals(1, reads.get()); + } + + @Test + void ordinaryApplicationCauseIsNotTreatedAsATransportWrapper() throws Exception { + runEndFailure("clear", false, new IllegalStateException("ordinary diagnostic", new VirtualMachineError() {})); + } + + private static void runEndFailure(String method, boolean suspend, Throwable failure) throws Exception { + runEndFailure(method, suspend, failure, null); + } + + private static void runEndFailure(String method, boolean suspend, Throwable failure, Error expectedFatal) + throws Exception { + var original = MDC.getMDCAdapter(); + var callerBefore = MDC.getCopyOfContextMap(); + var injected = new AtomicBoolean(); + var plugin = new RecordingPlugin(); + var worker = Executors.newSingleThreadExecutor(); + try { + replaceAdapter(proxy(original, name -> { + var inEnd = Arrays.stream(Thread.currentThread().getStackTrace()) + .anyMatch(frame -> frame.getClassName().endsWith(".DurableExecutor") + && frame.getMethodName().equals("fireOnInvocationEnd")); + if (inEnd && name.equals(method) && injected.compareAndSet(false, true)) throw failure; + // Select both legitimate MDC restoration branches, independently of the finalizer's thread. + if (inEnd && name.equals("getCopyOfContextMap")) + return method.equals("setContextMap") ? Map.of("ambient", "saved") : NullSnapshot.INSTANCE; + return null; + })); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + if (suspend) context.wait("resume", Duration.ofSeconds(1)); + return "done"; + }, + DurableConfig.builder() + .withExecutorService(worker) + .withPlugins(plugin) + .build()); + if (expectedFatal != null) { + assertSame(expectedFatal, assertThrows(Error.class, () -> runner.run("input"))); + assertEquals(method.equals("getCopyOfContextMap") ? 0 : 1, plugin.ends.get()); + } else { + var result = runner.run("input"); + assertEquals(suspend ? ExecutionStatus.PENDING : ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals(1, plugin.ends.get()); + } + assertTrue(injected.get(), "the actual execute() end boundary must exercise the adapter failure"); + assertEquals(1, plugin.starts.get()); + } finally { + try { + worker.shutdownNow(); + assertTrue(worker.awaitTermination(3, TimeUnit.SECONDS)); + } finally { + replaceAdapter(original); + if (callerBefore == null) MDC.clear(); + else MDC.setContextMap(callerBefore); + } + } + } + + private enum NullSnapshot { + INSTANCE + } + + @FunctionalInterface + private interface Fault { + Object apply(String method) throws Throwable; + } + + private static MDCAdapter proxy(MDCAdapter delegate, Fault fault) { + return (MDCAdapter) Proxy.newProxyInstance( + MDCAdapter.class.getClassLoader(), new Class[] {MDCAdapter.class}, (proxy, method, args) -> { + var value = fault.apply(method.getName()); + if (value == NullSnapshot.INSTANCE) return null; + if (value != null) return value; + try { + return method.invoke(delegate, args); + } catch (InvocationTargetException ex) { + throw ex.getCause(); + } + }); + } + + private static void replaceAdapter(MDCAdapter adapter) throws Exception { + var setter = MDC.class.getDeclaredMethod("setMDCAdapter", MDCAdapter.class); + setter.setAccessible(true); + setter.invoke(null, adapter); + } + + private static final class RecordingPlugin implements DurableExecutionPlugin { + final AtomicInteger starts = new AtomicInteger(); + final AtomicInteger ends = new AtomicInteger(); + + @Override + public void onInvocationStart(InvocationInfo info) { + starts.incrementAndGet(); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + ends.incrementAndGet(); + } + } + /** Executes on a real worker, completing its first task before returning to the invocation caller. */ + private static final class CompletingWorker extends AbstractExecutorService { + private final ExecutorService delegate = + Executors.newSingleThreadExecutor(task -> new Thread(task, "mdc-initialization-worker")); + private final AtomicBoolean first = new AtomicBoolean(true); + + @Override + public void execute(Runnable task) { + if (!first.compareAndSet(true, false)) { + delegate.execute(task); + return; + } + var done = new CompletableFuture(); + delegate.execute(() -> { + try { + task.run(); + } finally { + done.complete(null); + } + }); + done.orTimeout(3, TimeUnit.SECONDS).join(); + } + + @Override + public void shutdown() { + delegate.shutdown(); + } + + @Override + public List shutdownNow() { + return delegate.shutdownNow(); + } + + @Override + public boolean isShutdown() { + return delegate.isShutdown(); + } + + @Override + public boolean isTerminated() { + return delegate.isTerminated(); + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { + return delegate.awaitTermination(timeout, unit); + } + } +} 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 1f4865642..2ec749c81 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 @@ -4,7 +4,11 @@ import com.amazonaws.services.lambda.runtime.Context; import com.amazonaws.services.lambda.runtime.RequestHandler; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.UndeclaredThrowableException; import java.nio.charset.StandardCharsets; +import java.util.Collections; +import java.util.IdentityHashMap; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; import java.util.concurrent.ExecutionException; @@ -79,6 +83,7 @@ public static DurableExecutionOutput execute( executionManager.registerActiveThread(null); // Captured for onInvocationEnd, which runs outside the handler thread below. var pluginExecutionInput = new AtomicReference<>(); + var invocationStarted = new AtomicBoolean(); var hasHandlerScope = new AtomicBoolean(); var handlerFuture = supplyHandler( () -> { @@ -104,7 +109,7 @@ public static DurableExecutionOutput execute( // inject ThreadLocal objects, update MDC, etc. // executionStartTime comes from the initial EXECUTION operation in the first backend event. if (!pluginRunner.isEmpty()) { - pluginRunner.onInvocationStart(new InvocationInfo( + var invocationInfo = new InvocationInfo( requestId, executionArn, isFirstInvocation, @@ -115,7 +120,9 @@ public static DurableExecutionOutput execute( executionManager.getInitialOperationIds()), PluginInfoConverter.toOperationItemMap( executionManager.getUpdatedOperationsSnapshot(), - executionManager.getInitialOperationIds()))); + executionManager.getInitialOperationIds())); + invocationStarted.set(true); + pluginRunner.onInvocationStart(invocationInfo); } if (inputFailure != null) { ExceptionHelper.sneakyThrow(inputFailure); @@ -159,6 +166,7 @@ public static DurableExecutionOutput execute( // return PENDING if it's SuspendExecutionException if (cause instanceof SuspendExecutionException) { fireOnInvocationEnd( + invocationStarted.get(), scopeFatal, pluginRunner, executionManager, @@ -179,6 +187,7 @@ public static DurableExecutionOutput execute( unrecoverableDurableExecutionException && unrecoverableDurableExecutionException.isRetryable()) { fireOnInvocationEnd( + invocationStarted.get(), scopeFatal, pluginRunner, executionManager, @@ -195,6 +204,7 @@ public static DurableExecutionOutput execute( // fail the execution otherwise logger.debug("Execution failed: {}", cause.getMessage()); fireOnInvocationEnd( + invocationStarted.get(), scopeFatal, pluginRunner, executionManager, @@ -213,6 +223,7 @@ public static DurableExecutionOutput execute( var output = DurableExecutionOutput.success(handleLargePayload(executionManager, outputPayload)); fireOnInvocationEnd( + invocationStarted.get(), scopeFatal, pluginRunner, executionManager, @@ -364,7 +375,63 @@ private static SafeCloseable restoreMdcOnClose() { }; } + /** MDC is ancillary to end dispatch; ordinary adapter failures must not replace the selected outcome. */ + private static SafeCloseable preserveEndMdc() { + SafeCloseable restore; + try { + restore = restoreMdcOnClose(); + } catch (Throwable failure) { + reportEndMdcFailure(failure); + return () -> {}; + } + return () -> { + try { + restore.close(); + } catch (Throwable failure) { + reportEndMdcFailure(failure); + } + }; + } + + private static void reportEndMdcFailure(Throwable failure) { + throwIfMdcFatal(failure); + try { + // Do not ask a malformed Throwable for its message or cause again while logging. + logger.warn( + "Could not preserve MDC around invocation-end dispatch; preserving execution outcome ({})", + failure.getClass().getName()); + } catch (Throwable loggingFailure) { + // An MDC-backed logger may encounter the same adapter failure while reporting it. + throwIfMdcFatal(loggingFailure); + } + } + + @SuppressWarnings("removal") + private static void throwIfMdcFatal(Throwable failure) { + if (failure instanceof VirtualMachineError fatal) throw fatal; + if (failure instanceof ThreadDeath fatal) throw fatal; + var seen = Collections.newSetFromMap(new IdentityHashMap()); + var cause = failure; + while (seen.add(cause)) { + if (cause instanceof VirtualMachineError fatal) throw fatal; + if (cause instanceof ThreadDeath fatal) throw fatal; + if (!(cause instanceof CompletionException + || cause instanceof ExecutionException + || cause instanceof InvocationTargetException + || cause instanceof UndeclaredThrowableException)) return; + try { + cause = cause.getCause(); + } catch (VirtualMachineError | ThreadDeath fatal) { + throw fatal; + } catch (Throwable unreadableCause) { + return; + } + if (cause == null) return; + } + } + private static void fireOnInvocationEnd( + boolean invocationStarted, AtomicReference scopeFatal, PluginRunner pluginRunner, ExecutionManager executionManager, @@ -375,14 +442,14 @@ private static void fireOnInvocationEnd( Throwable error, Object executionInput, Object executionResult) { - if (pluginRunner.isEmpty()) { + if (pluginRunner.isEmpty() || !invocationStarted) { return; } // Freeze the selected end snapshot here. A scope fatal already observed still escapes before dispatch; // a later scope fatal escapes its owner instead of rewriting this caller outcome after finalization. throwIfScopeFatal(scopeFatal); // Finalization can run on the invocation caller rather than the handler worker. - try (var ignored = restoreMdcOnClose()) { + try (var ignored = preserveEndMdc()) { pluginRunner.onInvocationEnd(new InvocationEndInfo( requestId, executionArn, From 5a7f5494c3f4c9a3644fdbec6f0f2658742a1049 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 7 Oct 2026 18:56:24 +0000 Subject: [PATCH 16/48] fix: accept concrete handler scope return types --- .../lambda/durable/plugin/HandlerScoped.java | 2 +- .../plugin/ConcreteHandlerScopeTest.java | 48 +++++++++++++++++++ 2 files changed, 49 insertions(+), 1 deletion(-) create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/plugin/ConcreteHandlerScopeTest.java diff --git a/sdk/src/main/java/software/amazon/lambda/durable/plugin/HandlerScoped.java b/sdk/src/main/java/software/amazon/lambda/durable/plugin/HandlerScoped.java index e4ba959ca..c630341be 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/plugin/HandlerScoped.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/plugin/HandlerScoped.java @@ -26,5 +26,5 @@ @Retention(RetentionPolicy.RUNTIME) @Target(ElementType.TYPE) public @interface HandlerScoped { - Class> value(); + Class> value(); } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/plugin/ConcreteHandlerScopeTest.java b/sdk/src/test/java/software/amazon/lambda/durable/plugin/ConcreteHandlerScopeTest.java new file mode 100644 index 000000000..47ceea8c4 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/plugin/ConcreteHandlerScopeTest.java @@ -0,0 +1,48 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.plugin; + +import static org.junit.jupiter.api.Assertions.*; + +import java.util.List; +import java.util.function.Function; +import org.junit.jupiter.api.Test; + +class ConcreteHandlerScopeTest { + @Test + void concreteScopeReturnTypeOpensAndClosesOnItsOwner() { + var plugin = new ConcretePlugin(); + var runner = new PluginRunner(List.of(plugin)); + assertEquals("result", runner.runHandler(() -> { + assertNotNull(plugin.scope); + assertFalse(plugin.scope.closed); + return "result"; + })); + assertTrue(plugin.scope.closed); + assertSame(Thread.currentThread(), plugin.scope.owner); + } + + @HandlerScoped(ConcreteOpener.class) + public static final class ConcretePlugin implements DurableExecutionPlugin { + ConcreteScope scope; + } + + public static final class ConcreteOpener implements Function { + @Override + public ConcreteScope apply(ConcretePlugin plugin) { + return plugin.scope = new ConcreteScope(); + } + } + + public static final class ConcreteScope implements AutoCloseable { + final Thread owner = Thread.currentThread(); + boolean closed; + + @Override + public void close() { + assertSame(owner, Thread.currentThread()); + assertFalse(closed); + closed = true; + } + } +} From a9d328532c6c1100d5b2dd08d2a3d4a3a70cc9fb Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 7 Oct 2026 20:02:57 +0000 Subject: [PATCH 17/48] fix: contain malformed handler scope diagnostics --- .../durable/otel/ScopeFatalSignalTest.java | 22 +- .../lambda/durable/plugin/PluginRunner.java | 45 +++- .../HandlerScopeFailureInspectionTest.java | 212 ++++++++++++++++++ 3 files changed, 268 insertions(+), 11 deletions(-) create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeFailureInspectionTest.java diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ScopeFatalSignalTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ScopeFatalSignalTest.java index faf6a5775..f27e66a59 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ScopeFatalSignalTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ScopeFatalSignalTest.java @@ -18,10 +18,19 @@ class ScopeFatalSignalTest { @ParameterizedTest - @CsvSource({"false,false", "true,false", "false,true", "true,true"}) + @CsvSource({ + "false,false,false", + "true,false,false", + "false,true,false", + "true,true,false", + "false,false,true", + "true,false,true", + "false,true,true", + "true,true,true" + }) @SuppressWarnings("removal") - void fatalCloseWakesCallerAndStillAllowsBoundedEarlierCleanup(boolean threadDeath, boolean releaseInBudget) - throws Exception { + void fatalCloseWakesCallerAndStillAllowsBoundedEarlierCleanup( + boolean threadDeath, boolean releaseInBudget, boolean fatalAccessor) throws Exception { var earlierEntered = new CountDownLatch(1); var release = new CountDownLatch(1); var closed = new AtomicBoolean(); @@ -44,6 +53,13 @@ void fatalCloseWakesCallerAndStillAllowsBoundedEarlierCleanup(boolean threadDeat closed.set(true); }); var later = new Scoped(() -> { + if (fatalAccessor) + throw new CompletionException("scope cause accessor", null) { + @Override + public synchronized Throwable getCause() { + throw fatal; + } + }; throw fatal; }); var config = DurableConfig.builder() diff --git a/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java b/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java index e3e2a120a..887b172d3 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java @@ -3,9 +3,11 @@ package software.amazon.lambda.durable.plugin; import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.UndeclaredThrowableException; import java.util.ArrayDeque; import java.util.Collections; import java.util.HashMap; +import java.util.IdentityHashMap; import java.util.LinkedHashSet; import java.util.List; import java.util.Set; @@ -167,19 +169,46 @@ private static void closeHandlerScopes(ArrayDeque scopes, Consume if (firstFatal != null) throw firstFatal; } - @SuppressWarnings("removal") private static void reportHandlerScopeFailure(String message, Throwable failure, Consumer onScopeFatal) { - var cause = failure; - while ((cause instanceof CompletionException || cause instanceof ExecutionException) - && cause.getCause() != null) { - cause = cause.getCause(); + reportHandlerScopeFatal(failure, onScopeFatal); + try { + // A plugin-controlled diagnostic must not be inspected again by the logger. + logger.warn("{} ({})", message, failure.getClass().getName()); + } catch (Throwable loggingFailure) { + reportHandlerScopeFatal(loggingFailure, onScopeFatal); } - if (cause instanceof VirtualMachineError || cause instanceof ThreadDeath) { - var fatal = (Error) cause; + } + + private static void reportHandlerScopeFatal(Throwable failure, Consumer onScopeFatal) { + var fatal = findHandlerScopeFatal(failure); + if (fatal != null) { + // The caller must wake even when the fatal originated in an exception's cause accessor. onScopeFatal.accept(fatal); throw fatal; } - logger.warn(message, failure); + } + + @SuppressWarnings("removal") + private static Error findHandlerScopeFatal(Throwable failure) { + if (failure instanceof VirtualMachineError || failure instanceof ThreadDeath) return (Error) failure; + var seen = Collections.newSetFromMap(new IdentityHashMap()); + var cause = failure; + while (seen.add(cause)) { + if (cause instanceof VirtualMachineError || cause instanceof ThreadDeath) return (Error) cause; + if (!(cause instanceof CompletionException + || cause instanceof ExecutionException + || cause instanceof InvocationTargetException + || cause instanceof UndeclaredThrowableException)) return null; + try { + cause = cause.getCause(); + } catch (VirtualMachineError | ThreadDeath fatal) { + return fatal; + } catch (Throwable unreadableCause) { + return null; + } + if (cause == null) return null; + } + return null; } public void onInvocationStart(InvocationInfo info) { diff --git a/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeFailureInspectionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeFailureInspectionTest.java new file mode 100644 index 000000000..fac872ad4 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeFailureInspectionTest.java @@ -0,0 +1,212 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.plugin; + +import static org.junit.jupiter.api.Assertions.*; + +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.UndeclaredThrowableException; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; +import java.util.function.Supplier; +import java.util.stream.Stream; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import software.amazon.lambda.durable.util.ExceptionHelper; + +class HandlerScopeFailureInspectionTest { + static Stream ordinaryFailures() { + return Stream.of("throwing", "cycle", "two-cycle", "changing", "null", "application-cause") + .flatMap(kind -> Stream.of(false, true) + .flatMap(open -> Stream.of(false, true).map(bodyFails -> Arguments.of(kind, open, bodyFails)))); + } + + @ParameterizedTest + @MethodSource("ordinaryFailures") + void malformedScopeFailuresKeepTheBodyOutcomeAndReverseCleanup(String kind, boolean failOpen, boolean bodyFails) { + var calls = new ArrayList(); + var reads = new AtomicInteger(); + var failure = ordinary(kind, reads); + var owner = Thread.currentThread(); + var runner = new PluginRunner(List.of( + healthy("first", calls, owner), broken(failOpen, failure, calls), healthy("last", calls, owner))); + var original = new IllegalStateException("body failure"); + var reported = new AtomicReference(); + var bodyCalls = new AtomicInteger(); + var body = (Supplier) () -> { + bodyCalls.incrementAndGet(); + calls.add("body"); + if (bodyFails) throw original; + return "done"; + }; + if (bodyFails) + assertSame( + original, + assertThrows(IllegalStateException.class, () -> runner.runHandler(body, () -> {}, reported::set))); + else assertEquals("done", runner.runHandler(body, () -> {}, reported::set)); + assertEquals(1, bodyCalls.get()); + assertNull(reported.get()); + assertEquals( + failOpen + ? List.of("open-first", "open-bad", "open-last", "body", "close-last", "close-first") + : List.of( + "open-first", + "open-bad", + "open-last", + "body", + "close-last", + "close-bad", + "close-first"), + calls); + assertEquals(kind.equals("application-cause") ? 0 : kind.equals("two-cycle") ? 2 : 1, reads.get()); + } + + static Stream fatalFailures() { + return Stream.of("direct", "completion", "execution", "reflection", "proxy", "accessor") + .flatMap(kind -> Stream.of(false, true) + .flatMap(open -> + Stream.of(false, true).map(threadDeath -> Arguments.of(kind, open, threadDeath)))); + } + + @SuppressWarnings("removal") + @ParameterizedTest + @MethodSource("fatalFailures") + void fatalInspectionReportsBeforeEarlierCleanupAndKeepsIdentity( + String kind, boolean failOpen, boolean threadDeath) { + Error fatal = threadDeath ? new ThreadDeath() : new InternalError("scope fatal"); + var reported = new AtomicReference(); + var closed = new AtomicInteger(); + var bodyCalls = new AtomicInteger(); + var owner = Thread.currentThread(); + var earlier = new Scoped(() -> () -> { + assertSame(owner, Thread.currentThread()); + assertSame(fatal, reported.get(), "notify the invocation before potentially blocking earlier cleanup"); + closed.incrementAndGet(); + }); + var runner = new PluginRunner(List.of(earlier, broken(failOpen, fatalWrapper(kind, fatal), new ArrayList<>()))); + assertSame( + fatal, + assertThrows( + Error.class, + () -> runner.runHandler( + () -> { + bodyCalls.incrementAndGet(); + return "done"; + }, + () -> {}, + reported::set))); + assertSame(fatal, reported.get()); + assertEquals(1, closed.get()); + assertEquals(failOpen ? 0 : 1, bodyCalls.get()); + } + + private static Throwable ordinary(String kind, AtomicInteger reads) { + if (kind.equals("application-cause")) + return new IllegalStateException("application") { + @Override + public synchronized Throwable getCause() { + reads.incrementAndGet(); + return new InternalError("not a transport wrapper"); + } + }; + var first = new AtomicReference(); + var second = new ExecutionException("second", null) { + @Override + public synchronized Throwable getCause() { + guardNegativeTraversal(reads); + return first.get(); + } + }; + var failure = new CompletionException("scope", null) { + @Override + public synchronized Throwable getCause() { + guardNegativeTraversal(reads); + return switch (kind) { + case "throwing" -> throw new IllegalArgumentException("unreadable diagnostic"); + case "cycle" -> this; + case "two-cycle" -> second; + case "changing" -> + reads.get() == 1 + ? new IllegalStateException("first cause") + : new InternalError("changed cause"); + case "null" -> null; + default -> throw new AssertionError(kind); + }; + } + }; + first.set(failure); + return failure; + } + + private static void guardNegativeTraversal(AtomicInteger reads) { + // A broken classifier must fail the negative control rather than leave a spinning test thread. + if (reads.incrementAndGet() > 8) throw new AssertionError("test safety bound: repeated cause traversal"); + } + + private static Throwable fatalWrapper(String kind, Error fatal) { + return switch (kind) { + case "direct" -> fatal; + case "completion" -> new CompletionException(fatal); + case "execution" -> new ExecutionException(fatal); + case "reflection" -> new InvocationTargetException(fatal); + case "proxy" -> new UndeclaredThrowableException(fatal); + case "accessor" -> + new CompletionException("fatal accessor", null) { + @Override + public synchronized Throwable getCause() { + throw fatal; + } + }; + default -> throw new AssertionError(kind); + }; + } + + private static Scoped healthy(String name, List calls, Thread owner) { + return new Scoped(() -> { + assertSame(owner, Thread.currentThread()); + calls.add("open-" + name); + return () -> { + assertSame(owner, Thread.currentThread()); + calls.add("close-" + name); + }; + }); + } + + private static Scoped broken(boolean failOpen, Throwable failure, List calls) { + return new Scoped(() -> { + calls.add("open-bad"); + if (failOpen) ExceptionHelper.sneakyThrow(failure); + return () -> { + calls.add("close-bad"); + ExceptionHelper.sneakyThrow(failure); + }; + }); + } + + @FunctionalInterface + private interface Open { + AutoCloseable open(); + } + + @HandlerScoped(Opener.class) + public static final class Scoped implements DurableExecutionPlugin { + private final Open open; + + private Scoped(Open open) { + this.open = open; + } + } + + public static final class Opener implements Function { + @Override + public AutoCloseable apply(Scoped plugin) { + return plugin.open.open(); + } + } +} From 923a9147a6cc11f45b6d54bcdb7ff162aa150edd Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Tue, 6 Oct 2026 18:16:27 -0700 Subject: [PATCH 18/48] test: return one checkpoint state per operation in unit mocks (cherry picked from commit 5e171cbab84a00c121535cd5771eed055ccaa437) --- .../amazon/lambda/durable/TestUtils.java | 9 ++-- .../amazon/lambda/durable/TestUtilsTest.java | 51 +++++++++++++++++++ 2 files changed, 56 insertions(+), 4 deletions(-) create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/TestUtilsTest.java diff --git a/sdk/src/test/java/software/amazon/lambda/durable/TestUtils.java b/sdk/src/test/java/software/amazon/lambda/durable/TestUtils.java index 1aaba5d4c..44f0078a9 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/TestUtils.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/TestUtils.java @@ -5,7 +5,7 @@ import static org.mockito.ArgumentMatchers.*; import static org.mockito.Mockito.*; -import java.util.ArrayList; +import java.util.LinkedHashMap; import java.util.List; import java.util.UUID; import software.amazon.awssdk.services.lambda.model.*; @@ -18,7 +18,8 @@ public static DurableExecutionClient createMockClient() { var client = mock(DurableExecutionClient.class); when(client.checkpoint(any(), any(), any())).thenAnswer(invocation -> { var updates = (List) invocation.getArgument(2); - var responseOperations = new ArrayList(); + // A checkpoint returns the latest state of each operation, not one entry per update. + var responseOperations = new LinkedHashMap(); if (updates != null) { for (var update : updates) { @@ -52,14 +53,14 @@ public static DurableExecutionClient createMockClient() { opBuilder.contextDetails(contexDetail.build()); } } - responseOperations.add(opBuilder.build()); + responseOperations.put(update.id(), opBuilder.build()); } } return CheckpointDurableExecutionResponse.builder() .checkpointToken("new-token") .newExecutionState(CheckpointUpdatedExecutionState.builder() - .operations(responseOperations) + .operations(responseOperations.values()) .build()) .build(); }); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/TestUtilsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/TestUtilsTest.java new file mode 100644 index 000000000..1a2864de9 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/TestUtilsTest.java @@ -0,0 +1,51 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.*; + +import java.util.List; +import java.util.Set; +import org.junit.jupiter.api.Test; +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.plugin.PluginInfoConverter; + +class TestUtilsTest { + @Test + void batchedUpdatesReturnLatestStatePerOperation() { + var client = TestUtils.createMockClient(); + var response = client.checkpoint( + "arn:execution", + "token", + List.of( + update("first", OperationAction.START, null), + update("second", OperationAction.START, null), + update("first", OperationAction.SUCCEED, "\"result\""))); + var operations = response.newExecutionState().operations(); + + // Duplicate IDs make the real SDK callback reject a successful checkpoint. + var change = assertDoesNotThrow(() -> PluginInfoConverter.toOperationChangeInfo( + "request", "arn:execution", operations, operations, Set.of())); + assertEquals(2, operations.size()); + assertEquals(Set.of("first", "second"), change.updatedOperations().keySet()); + assertEquals( + OperationStatus.SUCCEEDED, + change.updatedOperations().get("first").status()); + assertEquals("\"result\"", change.updatedOperations().get("first").result()); + assertEquals( + OperationStatus.STARTED, + change.updatedOperations().get("second").status()); + } + + private static OperationUpdate update(String id, OperationAction action, String payload) { + return OperationUpdate.builder() + .id(id) + .type(OperationType.STEP) + .action(action) + .payload(payload) + .build(); + } +} From e72b3f2065abddc925c89dddf1f5dd51e0cea17d Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 7 Oct 2026 20:25:45 +0000 Subject: [PATCH 19/48] test: include failed execution details in executor reuse checks --- .../durable/execution/DurableExecutionTest.java | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutionTest.java index e0773aaeb..df72229ec 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutionTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutionTest.java @@ -411,7 +411,12 @@ void testExecutorNotShutdownAfterMultipleHandlerInvocations() { (userInput, ctx) -> ctx.step("test1", String.class, stepCtx -> "Result 1: " + userInput), config); - assertEquals(ExecutionStatus.SUCCEEDED, output1.status()); + assertEquals( + ExecutionStatus.SUCCEEDED, + output1.status(), + () -> output1.error() == null + ? "No error payload" + : output1.error().errorType() + ": " + output1.error().errorMessage()); assertFalse(sharedExecutor.isShutdown(), "Executor should not be shutdown after first execution"); // Create second input with different execution operation @@ -440,7 +445,12 @@ void testExecutorNotShutdownAfterMultipleHandlerInvocations() { (userInput, ctx) -> ctx.step("test2", String.class, stepCtx -> "Result 2: " + userInput), config); - assertEquals(ExecutionStatus.SUCCEEDED, output2.status()); + assertEquals( + ExecutionStatus.SUCCEEDED, + output2.status(), + () -> output2.error() == null + ? "No error payload" + : output2.error().errorType() + ": " + output2.error().errorMessage()); assertFalse(sharedExecutor.isShutdown(), "Executor should not be shutdown after second execution"); // Verify both executions completed successfully and used the same executor From fac86515fec0aadc75b85e42faac446a27c05a75 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Wed, 7 Oct 2026 15:53:17 -0700 Subject: [PATCH 20/48] refactor: finalize invocation hooks on the handler thread Run invocation start and end on the root handler thread and wait for handler cleanup and end hooks before returning. Unwind end hooks in reverse order, retain the selected suspension/retry result, and remove the handler-scope opener API and bounded handoff. Activate and restore OpenTelemetry context directly in the existing lifecycle hooks. --- otel-plugin/README.md | 46 +- .../durable/otel/ExecutionOtelPlugin.java | 26 +- .../durable/otel/InvocationOtelPlugin.java | 26 +- .../lambda/durable/otel/DelayedFlush.java | 57 -- .../durable/otel/ExecutionOtelPluginTest.java | 11 +- .../otel/HandlerContextIntegrationTest.java | 101 +++- .../otel/HandlerScopeFinalizationTest.java | 446 --------------- .../otel/InvocationEndCompatibilityTest.java | 62 +- .../otel/InvocationEndFailureTest.java | 133 +++++ .../otel/InvocationEndFinalizationTest.java | 195 +++++++ .../otel/InvocationOtelPluginTest.java | 2 +- .../LegacySubclassScopeCompatibilityTest.java | 9 +- .../durable/otel/MdcFailureBoundaryTest.java | 170 ++---- .../durable/otel/ScopeFatalSignalTest.java | 109 ---- ...ndlerScopeFinalizationIntegrationTest.java | 192 ------- ...InvocationFinalizationIntegrationTest.java | 226 ++++++++ .../durable/execution/DurableExecutor.java | 535 ++++++------------ .../plugin/DurableExecutionPlugin.java | 15 +- .../lambda/durable/plugin/HandlerScoped.java | 30 - .../lambda/durable/plugin/PluginRunner.java | 151 +---- .../execution/HandlerScopeHandoffTest.java | 260 --------- .../execution/InvocationLifecycleTest.java | 414 ++++++++++++++ .../plugin/ConcreteHandlerScopeTest.java | 48 -- .../HandlerScopeFailureInspectionTest.java | 212 ------- .../durable/plugin/HandlerScopeTest.java | 255 --------- .../LegacyHandlerScopeCompatibilityTest.java | 84 --- .../plugin/PluginLinkageErrorTest.java | 5 +- .../durable/plugin/PluginRunnerTest.java | 32 +- 28 files changed, 1405 insertions(+), 2447 deletions(-) delete mode 100644 otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DelayedFlush.java delete mode 100644 otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerScopeFinalizationTest.java create mode 100644 otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndFailureTest.java create mode 100644 otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndFinalizationTest.java delete mode 100644 otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ScopeFatalSignalTest.java delete mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/HandlerScopeFinalizationIntegrationTest.java create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/InvocationFinalizationIntegrationTest.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/plugin/HandlerScoped.java delete mode 100644 sdk/src/test/java/software/amazon/lambda/durable/execution/HandlerScopeHandoffTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationLifecycleTest.java delete mode 100644 sdk/src/test/java/software/amazon/lambda/durable/plugin/ConcreteHandlerScopeTest.java delete mode 100644 sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeFailureInspectionTest.java delete mode 100644 sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeTest.java delete mode 100644 sdk/src/test/java/software/amazon/lambda/durable/plugin/LegacyHandlerScopeCompatibilityTest.java diff --git a/otel-plugin/README.md b/otel-plugin/README.md index 610397872..6cebe01d0 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -15,26 +15,22 @@ OpenTelemetry instrumentation plugin for the AWS Lambda Durable Execution SDK fo ## Root handler context -With a core that supports the explicit `@HandlerScoped` capability, user instrumentation in the root handler joins -its canonical execution trace. A valid same-trace ambient Lambda span stays current. If ambient context is absent or -belongs to another trace, invocation view activates `Invocation`; execution view activates the deterministic -`Workflow` context, including its unsampled non-recording form. Nested operations retain their existing contexts. - -The core opens the scope after invocation startup and closes it on the same handler thread on success, failure, or -suspension. On suspension or termination, an opened scope gets up to 500ms to unwind before finalization, capped by -remaining Lambda time after reserving five seconds per configured plugin and one second for shutdown/response. -This is a best-effort reserve, not a total deadline bound: existing plugin callbacks and checkpoint draining may exceed it. The invocation caller waits; the owning handler worker remains free to exit. -If cleanup is blocked, the SDK logs the timeout and preserves the original PENDING/RETRYING outcome. -An observed `VirtualMachineError` or `ThreadDeath` from the new scope callbacks instead escapes the invocation caller; -ordinary cleanup failures and legacy body failures retain the original outcome. Fatals reported only after the response -cannot retroactively change it. The scope still -closes on its owning thread when the handler eventually exits; late cleanup telemetry is best effort. A compatible -ambient span uses a no-op scope and participates in the same bounded handoff. Invocation-end hooks can execute on a -different thread and do not own this scope. The explicit opener uses only JDK types; unannotated plugins have no scope: old plugins on a new core and new plugin layers on an older core retain their -existing behavior. Root-handler fallback activation requires both the updated core and plugin; no provider API or -dependency floor changes are required. The new scope boundary isolates ordinary exceptions and nonfatal linkage -errors during open and close, continues earlier scope cleanup, and preserves the handler outcome. JVM fatal errors -remain outside that containment, including when wrapped by asynchronous completion/future exceptions. +User instrumentation in the root handler joins its canonical execution trace. A valid same-trace ambient Lambda +span stays current. If ambient context is absent or belongs to another trace, invocation view activates `Invocation`; +execution view activates the deterministic `Workflow` context, including its unsampled non-recording form. +Nested operations retain their existing contexts. + +The plugin activates the context in `onInvocationStart` and restores the previous context in `onInvocationEnd`. +The core calls both hooks on the root handler thread and waits for `onInvocationEnd` to finish before returning, +including when execution suspends or terminates. Handler `finally` blocks must finish before invocation-end hooks +can run; a blocked handler cleanup therefore also blocks the invocation response. Context restoration still runs +when span finalization or flushing fails. Ordinary exceptions and nonfatal linkage errors retain the existing +plugin-hook isolation behavior. Invocation-end hooks run in reverse registration order so nested scopes unwind +correctly; remaining end hooks run even when another hook raises an Error, then the first Error propagates. + +This plugin version requires a core version that provides that same-thread invocation-hook contract. Upgrade the +core together with the plugin layer. Older cores may call `onInvocationEnd` on another thread and are not supported +with this plugin version. ## Installation @@ -90,7 +86,7 @@ Configuration reads this explicit opt-in annotation from the entire superclass c that happen to be named `getExclusiveGroup`. Existing subclasses retain their own methods while inheriting the bundled view restriction. A subclass may add another group, but cannot replace a superclass's group; repeated group names in one class hierarchy are checked once. -Older cores ignore the optional annotation and retain their prior behavior; no provider-version floor is raised. +The exclusivity annotation does not change the provider registration API; the same-thread core requirement above still applies. ## Quick Start using X-Ray/CloudWatch Tracing (ADOT Java Agent) @@ -430,16 +426,10 @@ var otelPlugin = new InvocationOtelPlugin( ## Requirements - Java 17+ -- AWS Durable Execution SDK for Java 2.0.0+ +- AWS Durable Execution SDK for Java with same-thread invocation hooks (use the core shipped with this plugin release or newer) - OpenTelemetry SDK 1.65.0+ (only for custom TracerProvider path) - ADOT Lambda Layer `AWSOpenTelemetryDistroJava` (for the no-arg constructor path) ## License Apache-2.0 - -The inherited `@HandlerScoped` annotation names an explicit JDK `Function` opener. The core passes the plugin -instance to that opener; it never discovers a method by name on the plugin. Bundled openers call private SDK code -without dispatching to coincidentally named subclass methods. No `openHandlerScope` method is added to either the -existing lifecycle interface or the bundled plugin superclasses. Older cores ignore the optional annotation and -still load the plugin layer. Custom opener classes need a public no-argument constructor accessible to the core. diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java index fed02f0b7..d0c46ea1a 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java @@ -21,13 +21,11 @@ import java.time.Instant; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeUnit; -import java.util.function.Function; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.slf4j.MDC; import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; import software.amazon.lambda.durable.plugin.ExclusivePluginGroup; -import software.amazon.lambda.durable.plugin.HandlerScoped; import software.amazon.lambda.durable.plugin.InvocationEndInfo; import software.amazon.lambda.durable.plugin.InvocationInfo; import software.amazon.lambda.durable.plugin.OperationEndInfo; @@ -91,7 +89,6 @@ *

Thread-safe: uses {@link ConcurrentHashMap} for span/scope storage since the SDK runs user code on multiple * threads. */ -@HandlerScoped(ExecutionOtelPlugin.HandlerScopeOpener.class) @ExclusivePluginGroup("durable-otel-view") public class ExecutionOtelPlugin implements DurableExecutionPlugin { @@ -108,6 +105,8 @@ public class ExecutionOtelPlugin implements DurableExecutionPlugin { // Per-invocation state private volatile boolean tracingEnabled; private volatile Span invocationSpan; + // Opened and closed by the invocation hooks on the root handler thread. + private Scope handlerScope; private volatile String durableExecutionArn; // Trace ID and flags of the execution trace, published together as one snapshot so readers never pair a trace ID @@ -274,17 +273,10 @@ public void onInvocationStart(InvocationInfo info) { invocationSpan.getSpanContext().getTraceId()); } tracingEnabled = true; + handlerScope = activateHandlerContext(); } - /** JDK-only scope bridge; does not dispatch to coincidental subclass methods. */ - public static final class HandlerScopeOpener implements Function { - @Override - public AutoCloseable apply(ExecutionOtelPlugin plugin) { - return plugin.activateHandlerContext(); - } - } - - private AutoCloseable activateHandlerContext() { + private Scope activateHandlerContext() { var trace = executionTrace; if (!tracingEnabled || trace == null) return null; var ambient = Span.current().getSpanContext(); @@ -296,6 +288,16 @@ private AutoCloseable activateHandlerContext() { @Override public void onInvocationEnd(InvocationEndInfo info) { + try { + endInvocation(info); + } finally { + var scope = handlerScope; + handlerScope = null; + if (scope != null) scope.close(); + } + } + + private void endInvocation(InvocationEndInfo info) { if (!tracingEnabled) { return; } diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java index 7bb452ecc..7e2a5acc5 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java @@ -22,13 +22,11 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedDeque; import java.util.concurrent.TimeUnit; -import java.util.function.Function; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.slf4j.MDC; import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; import software.amazon.lambda.durable.plugin.ExclusivePluginGroup; -import software.amazon.lambda.durable.plugin.HandlerScoped; import software.amazon.lambda.durable.plugin.InvocationEndInfo; import software.amazon.lambda.durable.plugin.InvocationInfo; import software.amazon.lambda.durable.plugin.OperationEndInfo; @@ -78,7 +76,6 @@ *

Thread-safe: uses {@link ConcurrentHashMap} for span/scope storage since the SDK runs user code on multiple * threads. */ -@HandlerScoped(InvocationOtelPlugin.HandlerScopeOpener.class) @ExclusivePluginGroup("durable-otel-view") public class InvocationOtelPlugin implements DurableExecutionPlugin { @@ -95,6 +92,8 @@ public class InvocationOtelPlugin implements DurableExecutionPlugin { // Per-invocation state private volatile boolean tracingEnabled; private volatile Span invocationSpan; + // Opened and closed by the invocation hooks on the root handler thread. + private Scope handlerScope; private volatile String durableExecutionArn; // Trace ID and flags of the execution trace, published together as one snapshot so readers never pair a trace ID // with mismatched flags. @@ -271,17 +270,10 @@ public void onInvocationStart(InvocationInfo info) { invocationSpan.getSpanContext().getTraceId()); } tracingEnabled = true; + handlerScope = activateHandlerContext(); } - /** JDK-only scope bridge; does not dispatch to coincidental subclass methods. */ - public static final class HandlerScopeOpener implements Function { - @Override - public AutoCloseable apply(InvocationOtelPlugin plugin) { - return plugin.activateHandlerContext(); - } - } - - private AutoCloseable activateHandlerContext() { + private Scope activateHandlerContext() { var trace = executionTrace; if (!tracingEnabled || trace == null) return null; var ambient = Span.current().getSpanContext(); @@ -293,6 +285,16 @@ private AutoCloseable activateHandlerContext() { @Override public void onInvocationEnd(InvocationEndInfo info) { + try { + endInvocation(info); + } finally { + var scope = handlerScope; + handlerScope = null; + if (scope != null) scope.close(); + } + } + + private void endInvocation(InvocationEndInfo info) { if (!tracingEnabled) { return; } diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DelayedFlush.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DelayedFlush.java deleted file mode 100644 index a762d23e4..000000000 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DelayedFlush.java +++ /dev/null @@ -1,57 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.otel; - -import io.opentelemetry.context.Context; -import io.opentelemetry.sdk.common.CompletableResultCode; -import io.opentelemetry.sdk.trace.ReadWriteSpan; -import io.opentelemetry.sdk.trace.ReadableSpan; -import io.opentelemetry.sdk.trace.SpanProcessor; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicLong; - -class DelayedFlush implements SpanProcessor { - public void onStart(Context parent, ReadWriteSpan span) {} - - final HandlerScopeFinalizationTest.Deadline deadline; - final AtomicInteger calls = new AtomicInteger(); - final AtomicInteger remainingAtFlush = new AtomicInteger(); - final AtomicLong flushMillis = new AtomicLong(); - final ScheduledExecutorService clock; - - DelayedFlush(HandlerScopeFinalizationTest.Deadline deadline, ScheduledExecutorService clock) { - this.deadline = deadline; - this.clock = clock; - } - - public boolean isStartRequired() { - return false; - } - - public void onEnd(ReadableSpan span) {} - - public boolean isEndRequired() { - return false; - } - - public CompletableResultCode shutdown() { - return CompletableResultCode.ofSuccess(); - } - - public CompletableResultCode forceFlush() { - calls.incrementAndGet(); - remainingAtFlush.set(deadline.remaining()); - var result = new CompletableResultCode(); - long start = System.nanoTime(); - clock.schedule( - () -> { - flushMillis.set(TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start)); - result.succeed(); - }, - 250, - TimeUnit.MILLISECONDS); - return result; - } -} diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExecutionOtelPluginTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExecutionOtelPluginTest.java index 9614759f6..875ae6c57 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExecutionOtelPluginTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExecutionOtelPluginTest.java @@ -292,8 +292,8 @@ void invocationStart_joinsAmbientTrace_whenAmbientIsOnExecutionTrace() { SpanContext.create(canonicalTraceId, ambientSpanId, TraceFlags.getSampled(), TraceState.getDefault()); try (var ignored = Span.wrap(ambient).makeCurrent()) { plugin.onInvocationStart(new InvocationInfo("req-1", ARN, false, executionStart)); + plugin.onInvocationEnd(new InvocationEndInfo("req-1", ARN, false, InvocationStatus.SUCCEEDED, null)); } - plugin.onInvocationEnd(new InvocationEndInfo("req-1", ARN, false, InvocationStatus.SUCCEEDED, null)); var invocationSpan = spanByName(spanExporter.getFinishedSpanItems(), "Invocation"); assertEquals(canonicalTraceId, invocationSpan.getTraceId()); @@ -312,8 +312,8 @@ void invocationStart_staysOnExecutionTrace_withoutLinkingAmbientSpan() { SpanContext.create(ambientTraceId, ambientSpanId, TraceFlags.getSampled(), TraceState.getDefault()); try (var ignored = Span.wrap(ambient).makeCurrent()) { plugin.onInvocationStart(new InvocationInfo("req-1", ARN, true, Instant.now())); + plugin.onInvocationEnd(new InvocationEndInfo("req-1", ARN, true, InvocationStatus.SUCCEEDED, null)); } - plugin.onInvocationEnd(new InvocationEndInfo("req-1", ARN, true, InvocationStatus.SUCCEEDED, null)); var spans = spanExporter.getFinishedSpanItems(); var workflowSpan = spanByName(spans, "Workflow"); @@ -350,8 +350,9 @@ void contextExtractor_isInvokedEveryInvocation_evenWithAmbientSpan_andBackendCon TraceState.getDefault()); try (var ignored = Span.wrap(ambient).makeCurrent()) { extractorPlugin.onInvocationStart(new InvocationInfo("req-1", ARN, true, Instant.now())); + extractorPlugin.onInvocationEnd( + new InvocationEndInfo("req-1", ARN, true, InvocationStatus.SUCCEEDED, null)); } - extractorPlugin.onInvocationEnd(new InvocationEndInfo("req-1", ARN, true, InvocationStatus.SUCCEEDED, null)); assertEquals(1, extractCalls.get(), "Extractor is invoked even when a valid ambient span is active"); var spans = exporter.getFinishedSpanItems(); @@ -378,16 +379,16 @@ void executionTrace_isStableAcrossReinvocations_withDifferentAmbientTraces() { try (var ignored = Span.wrap(ambientA).makeCurrent()) { plugin.onInvocationStart(new InvocationInfo("req-1", ARN, true, startTime)); + plugin.onInvocationEnd(new InvocationEndInfo("req-1", ARN, true, InvocationStatus.PENDING, null)); } - plugin.onInvocationEnd(new InvocationEndInfo("req-1", ARN, true, InvocationStatus.PENDING, null)); var firstInvocationTrace = spanByName(spanExporter.getFinishedSpanItems(), "Invocation").getTraceId(); spanExporter.reset(); try (var ignored = Span.wrap(ambientB).makeCurrent()) { plugin.onInvocationStart(new InvocationInfo("req-2", ARN, false, startTime)); + plugin.onInvocationEnd(new InvocationEndInfo("req-2", ARN, false, InvocationStatus.SUCCEEDED, null)); } - plugin.onInvocationEnd(new InvocationEndInfo("req-2", ARN, false, InvocationStatus.SUCCEEDED, null)); var secondInvocationTrace = spanByName(spanExporter.getFinishedSpanItems(), "Invocation").getTraceId(); diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerContextIntegrationTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerContextIntegrationTest.java index e7665a3c2..faa242bc0 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerContextIntegrationTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerContextIntegrationTest.java @@ -9,6 +9,8 @@ import io.opentelemetry.api.trace.TraceFlags; import io.opentelemetry.api.trace.TraceState; import io.opentelemetry.context.Context; +import io.opentelemetry.context.ContextKey; +import io.opentelemetry.context.Scope; import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; import io.opentelemetry.sdk.trace.SdkTracerProvider; import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; @@ -24,8 +26,9 @@ import software.amazon.lambda.durable.DurableConfig; 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.PluginRunner; +import software.amazon.lambda.durable.plugin.InvocationStatus; import software.amazon.lambda.durable.testing.LocalDurableTestRunner; class HandlerContextIntegrationTest { @@ -71,7 +74,7 @@ void reusedHandlerWorkerHasNoLeakedScope(boolean executionView, String outcome) assertFalse( executor.submit(() -> Span.current().getSpanContext().isValid()) .get(5, TimeUnit.SECONDS), - "cleanup must run on the worker, even if invocation finalization ran on another thread"); + "invocation-end cleanup must restore the handler worker before the response returns"); } finally { executor.shutdownNow(); } @@ -102,19 +105,15 @@ void preservesCompatibleAmbientAndRestoresUnrelatedAmbientAfterFailure(boolean e "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/name/id", true, Instant.now())); - var error = new IllegalStateException("handler failure"); - var runner = new PluginRunner(List.of(plugin)); - assertSame( - error, - assertThrows( - IllegalStateException.class, - () -> runner.runHandler(() -> { - var active = Span.current().getSpanContext(); - assertEquals(TRACE_ID, active.getTraceId()); - if (sameTrace) assertEquals(ambient, active); - else assertNotEquals(ambient.getSpanId(), active.getSpanId()); - throw error; - }))); + try { + var active = Span.current().getSpanContext(); + assertEquals(TRACE_ID, active.getTraceId()); + if (sameTrace) assertEquals(ambient, active); + else assertNotEquals(ambient.getSpanId(), active.getSpanId()); + } finally { + plugin.onInvocationEnd(new InvocationEndInfo( + "req", "arn", true, InvocationStatus.FAILED, new IllegalStateException("handler failure"))); + } assertEquals(ambient, Span.current().getSpanContext()); } assertSame(previous, Context.current()); @@ -216,4 +215,76 @@ void rootContextIsValidAndRestoredAcrossNestedWorkAndResume(boolean executionVie userProvider.close(); } } + + @ParameterizedTest + @CsvSource({"true,true", "true,false", "false,true", "false,false"}) + void customPluginScopesUnwindBeforeOtelScopeOnTheReusedWorker(boolean executionView, boolean suspend) + throws Exception { + var key = ContextKey.named("application-context"); + var exporter = InMemorySpanExporter.create(); + var builder = SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)); + var settings = OtelPluginConfig.builder() + .enableMdc(false) + .contextExtractor( + () -> new ExtractedContext(TRACE_ID, "1234567890123456", ExtractedContext.Sampling.SAMPLED)) + .build(); + DurableExecutionPlugin delegate = executionView + ? new ExecutionOtelPlugin(builder, settings) + : new InvocationOtelPlugin(builder, settings); + var ends = new ArrayList(); + var otel = new DurableExecutionPlugin() { + @Override + public void onInvocationStart(InvocationInfo info) { + delegate.onInvocationStart(info); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + ends.add("otel"); + delegate.onInvocationEnd(info); + } + }; + var custom = new DurableExecutionPlugin() { + private Scope scope; + + @Override + public void onInvocationStart(InvocationInfo info) { + scope = Context.current().with(key, "active").makeCurrent(); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + ends.add("custom"); + assertEquals("active", Context.current().get(key)); + scope.close(); + } + }; + var workers = Executors.newSingleThreadExecutor(); + try { + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + assertEquals("active", Context.current().get(key)); + assertEquals(TRACE_ID, Span.current().getSpanContext().getTraceId()); + if (suspend) context.wait("pause", Duration.ofSeconds(1)); + return "done"; + }, + DurableConfig.builder() + .withExecutorService(workers) + .withPlugins(otel, custom) + .build()); + for (var invocation = 0; invocation < 2; invocation++) { + runner.run("input"); + assertEquals(List.of("custom", "otel"), ends); + ends.clear(); + assertFalse(workers.submit(() -> Span.current().getSpanContext().isValid()) + .get(3, TimeUnit.SECONDS)); + assertNull(workers.submit(() -> Context.current().get(key)).get(3, TimeUnit.SECONDS)); + runner.advanceTime(); + } + } finally { + workers.shutdownNow(); + assertTrue(workers.awaitTermination(3, TimeUnit.SECONDS)); + } + } } diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerScopeFinalizationTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerScopeFinalizationTest.java deleted file mode 100644 index a549d8146..000000000 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerScopeFinalizationTest.java +++ /dev/null @@ -1,446 +0,0 @@ -// 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 com.amazonaws.services.lambda.runtime.Context; -import io.opentelemetry.sdk.trace.SdkTracerProvider; -import java.lang.reflect.Proxy; -import java.time.Duration; -import java.time.Instant; -import java.util.concurrent.*; -import java.util.concurrent.atomic.*; -import java.util.function.Function; -import org.junit.jupiter.api.Test; -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.TypeToken; -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.model.DurableExecutionOutput; -import software.amazon.lambda.durable.model.ExecutionStatus; -import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; -import software.amazon.lambda.durable.plugin.HandlerScoped; -import software.amazon.lambda.durable.plugin.InvocationEndInfo; -import software.amazon.lambda.durable.plugin.InvocationInfo; -import software.amazon.lambda.durable.plugin.InvocationStatus; -import software.amazon.lambda.durable.testing.local.LocalMemoryExecutionClient; - -class HandlerScopeFinalizationTest { - @ParameterizedTest - @CsvSource({ - "false,false,false", - "false,false,true", - "false,true,false", - "false,true,true", - "true,false,false", - "true,false,true", - "true,true,false", - "true,true,true" - }) - void preservesTimeForRealOtelFinalization(boolean executionView, boolean retry, boolean hasScope) throws Exception { - var deadline = new Deadline(); - var clock = Executors.newSingleThreadScheduledExecutor(r -> daemon(r, "flush-clock")); - var flush = new DelayedFlush(deadline, clock); - var builder = SdkTracerProvider.builder().addSpanProcessor(flush); - var settings = OtelPluginConfig.builder() - .enableMdc(false) - .contextExtractor(() -> new ExtractedContext( - "12345678901234567890123456789012", "1234567890123456", ExtractedContext.Sampling.SAMPLED)) - .build(); - DurableExecutionPlugin delegate = executionView - ? new ExecutionOtelPlugin(builder, settings) - : new InvocationOtelPlugin(builder, settings); - var plugin = new ScopedPlugin() { - public AutoCloseable openHandlerScope() { - if (!hasScope) return null; - if (delegate instanceof InvocationOtelPlugin invocation) - return new InvocationOtelPlugin.HandlerScopeOpener().apply(invocation); - return new ExecutionOtelPlugin.HandlerScopeOpener().apply((ExecutionOtelPlugin) delegate); - } - - public void onInvocationStart(InvocationInfo info) { - delegate.onInvocationStart(info); - } - - public void onInvocationEnd(InvocationEndInfo info) { - delegate.onInvocationEnd(info); - } - }; - var workers = Executors.newCachedThreadPool(r -> daemon(r, "handler-owner")); - var callers = Executors.newSingleThreadExecutor(r -> daemon(r, "invocation-caller")); - var enteredFinally = new CountDownLatch(1); - var releaseFinally = new CountDownLatch(1); - var original = retryError(); - var cfg = config(workers, plugin); - try { - var response = callers.submit(() -> DurableExecutor.execute( - input(), - deadline.context(), - TypeToken.get(String.class), - (value, ctx) -> { - deadline.arm(400); - try { - if (retry) - ctx.step("retry", String.class, step -> { - throw original; - }); - else ctx.wait("pause", Duration.ofSeconds(1)); - return "done"; - } finally { - enteredFinally.countDown(); - await(releaseFinally); - } - }, - cfg)); - check(enteredFinally.await(2, TimeUnit.SECONDS), "handler finally not reached"); - outcome(response, retry, original); - long total = deadline.elapsedMillis(); - check(flush.calls.get() == 1, "expected one actual bundled OTel forceFlush"); - check(flush.flushMillis.get() >= 240, "actual bundled OTel join did not await forceFlush result"); - check(total < 400, "unexpected response/deadline relation: " + total); - check(releaseFinally.getCount() == 1, "test cleanup was not blocked"); - - } finally { - releaseFinally.countDown(); - workers.shutdown(); - workers.awaitTermination(2, TimeUnit.SECONDS); - callers.shutdownNow(); - clock.shutdownNow(); - } - } - - @SuppressWarnings("removal") - @ParameterizedTest - @CsvSource({ - "false,false,false", - "false,false,true", - "false,true,false", - "false,true,true", - "true,false,false", - "true,false,true", - "true,true,false", - "true,true,true" - }) - void scopeFatal(boolean retry, boolean wrapped, boolean threadDeath) throws Exception { - scopeFatal(retry, wrapped, threadDeath ? new ThreadDeath() : new InternalError("scope fatal"), false); - } - - @Test - void losingBodyFinallyFatalPreservesTheOriginalWinner() throws Exception { - scopeFatal(false, false, new InternalError("body finally fatal"), true); - } - - @Test - void fatalAfterTheHandoffTimeoutCannotChangeAnAlreadyReturnedOutcome() throws Exception { - scopeFatal(false, false, new InternalError("late scope fatal"), false, true); - } - - @SuppressWarnings("removal") - @ParameterizedTest - @CsvSource({"false,false", "true,false", "false,true", "true,true"}) - void scopeFatalAfterEndDispatchPreservesCallerOutcomeAndEscapesOwner(boolean threadDeath, boolean duringEnd) - throws Exception { - var releaseScope = new CountDownLatch(1); - var ownerFinished = new CountDownLatch(1); - var shutdownEntered = new AtomicBoolean(); - var ended = new AtomicBoolean(); - var endCalls = new AtomicInteger(); - var endInfo = new AtomicReference(); - var fatalObservedDuringEnd = new AtomicBoolean(); - Error fatal = threadDeath ? new ThreadDeath() : new InternalError("late handler scope failure"); - var ownerFatal = new AtomicReference(); - var fatalEscaped = new CountDownLatch(1); - var plugin = new ScopedPlugin() { - public AutoCloseable openHandlerScope() { - return () -> { - await(releaseScope); - throw fatal; - }; - } - - public void onInvocationEnd(InvocationEndInfo info) { - ended.set(true); - endCalls.incrementAndGet(); - endInfo.set(info); - if (duringEnd) { - releaseScope.countDown(); - await(fatalEscaped); - fatalObservedDuringEnd.set(ownerFatal.get() == fatal); - } - } - }; - var workers = - new ThreadPoolExecutor(0, Integer.MAX_VALUE, 60, TimeUnit.SECONDS, new SynchronousQueue<>(), task -> { - var owner = daemon(task, "shutdown-owner"); - owner.setUncaughtExceptionHandler((thread, failure) -> { - ownerFatal.set(failure); - fatalEscaped.countDown(); - }); - return owner; - }) { - @Override - public void execute(Runnable task) { - super.execute(() -> { - try { - task.run(); - } finally { - ownerFinished.countDown(); - } - }); - } - - @Override - public int getActiveCount() { - // ExecutionManager.close reaches this after invocation-end callbacks. Hold resource closure - // until the owner has reported the fatal, without altering the handler/control futures. - assertTrue(ended.get()); - shutdownEntered.set(true); - releaseScope.countDown(); - await(ownerFinished); - return super.getActiveCount(); - } - }; - var deadline = new Deadline(); - try { - var output = assertDoesNotThrow(() -> DurableExecutor.execute( - input(), - deadline.context(), - TypeToken.get(String.class), - (value, ctx) -> { - deadline.arm(0); - ctx.wait("pause", Duration.ofSeconds(1)); - return "done"; - }, - config(workers, plugin))); - assertEquals(ExecutionStatus.PENDING, output.status(), "the already finalized caller outcome stays frozen"); - assertTrue(fatalEscaped.await(2, TimeUnit.SECONDS), "fatal must still escape the actual owner thread"); - assertSame(fatal, ownerFatal.get()); - assertEquals(duringEnd, fatalObservedDuringEnd.get(), "blocked end hook must observe the owner fatal"); - assertTrue(shutdownEntered.get()); - assertEquals(1, endCalls.get(), "a late fatal must not replay already delivered end hooks"); - assertEquals( - InvocationStatus.PENDING, - endInfo.get().invocationStatus(), - "snapshot reflects outcome known at dispatch"); - assertNull(endInfo.get().executionError(), "the fatal is reported after the selected end snapshot"); - } finally { - releaseScope.countDown(); - workers.shutdown(); - assertTrue(workers.awaitTermination(2, TimeUnit.SECONDS)); - } - } - - static void scopeFatal(boolean retry, boolean wrapped, Error fatal, boolean fromBodyFinally) throws Exception { - scopeFatal(retry, wrapped, fatal, fromBodyFinally, false); - } - - static void scopeFatal(boolean retry, boolean wrapped, Error fatal, boolean fromBodyFinally, boolean afterTimeout) - throws Exception { - var deadline = new Deadline(); - var closeEntered = new CountDownLatch(1); - var releaseFatal = new CountDownLatch(1); - var fatalRaised = new CountDownLatch(1); - var fatalBeforeEnd = new AtomicBoolean(); - var plugin = new ScopedPlugin() { - public AutoCloseable openHandlerScope() { - var owner = Thread.currentThread(); - return () -> { - check(owner == Thread.currentThread(), "wrong scope close thread"); - if (!fromBodyFinally) raiseFatal(closeEntered, releaseFatal, fatalRaised, wrapped, fatal); - }; - } - - public void onInvocationEnd(InvocationEndInfo info) { - fatalBeforeEnd.set(fatalRaised.getCount() == 0); - } - }; - var workers = Executors.newCachedThreadPool(r -> daemon(r, "fatal-owner")); - var callers = Executors.newSingleThreadExecutor(r -> daemon(r, "fatal-caller")); - var original = retryError(); - var cfg = config(workers, plugin); - try { - var response = callers.submit(() -> DurableExecutor.execute( - input(), - deadline.context(), - TypeToken.get(String.class), - (value, ctx) -> { - deadline.arm(30_000); - try { - if (retry) - ctx.step("retry", String.class, step -> { - throw original; - }); - else ctx.wait("pause", Duration.ofSeconds(1)); - return "done"; - } finally { - if (fromBodyFinally) raiseFatal(closeEntered, releaseFatal, fatalRaised, wrapped, fatal); - } - }, - cfg)); - check(closeEntered.await(2, TimeUnit.SECONDS), "fatal site not reached"); - check(deadline.budgetRead.await(2, TimeUnit.SECONDS), "winner not observed before fatal"); - if (afterTimeout) { - outcome(response, retry, original); - assertFalse(fatalBeforeEnd.get()); - releaseFatal.countDown(); - assertTrue(fatalRaised.await(2, TimeUnit.SECONDS)); - outcome(response, retry, original); - return; - } - releaseFatal.countDown(); - if (fromBodyFinally) { - outcome(response, retry, original); - assertTrue(fatalBeforeEnd.get()); - } else { - var thrown = assertThrows(ExecutionException.class, () -> response.get(3, TimeUnit.SECONDS)); - assertSame(fatal, thrown.getCause(), "scope-owned fatal must escape the invocation caller"); - } - - } finally { - releaseFatal.countDown(); - workers.shutdown(); - workers.awaitTermination(2, TimeUnit.SECONDS); - callers.shutdownNow(); - } - } - - @SuppressWarnings("removal") - @ParameterizedTest - @CsvSource({"false", "true"}) - void legacyBodyFatalRetainsFailedMapping(boolean threadDeath) { - Error fatal = threadDeath ? new ThreadDeath() : new InternalError("body fatal"); - var workers = Executors.newCachedThreadPool(r -> daemon(r, "body-fatal")); - try { - var output = DurableExecutor.execute( - input(), - null, - TypeToken.get(String.class), - (value, ctx) -> { - throw fatal; - }, - config(workers)); - check(output.status() == ExecutionStatus.FAILED, "legacy body fatal not FAILED"); - check(output.error().errorType().equals(fatal.getClass().getName()), "legacy body fatal error type lost"); - - } finally { - workers.shutdownNow(); - } - } - - static void raiseFatal( - CountDownLatch entered, CountDownLatch release, CountDownLatch raised, boolean wrapped, Error fatal) { - entered.countDown(); - await(release); - raised.countDown(); - if (wrapped) throw new CompletionException(new ExecutionException(fatal)); - throw fatal; - } - - static void outcome(Future response, boolean retry, Throwable original) throws Exception { - if (retry) { - try { - response.get(3, TimeUnit.SECONDS); - throw new AssertionError("retry returned output"); - } catch (ExecutionException e) { - check(e.getCause() == original, "retry failure changed to " + e.getCause()); - } - } else check(response.get(3, TimeUnit.SECONDS).status() == ExecutionStatus.PENDING, "suspension changed"); - } - - static DurableConfig config(ExecutorService workers, DurableExecutionPlugin... plugins) { - return DurableConfig.builder() - .withDurableExecutionClient(new LocalMemoryExecutionClient()) - .withExecutorService(workers) - .withCheckpointDelay(Duration.ZERO) - .withPlugins(plugins) - .build(); - } - - static UnrecoverableDurableExecutionException retryError() { - return new UnrecoverableDurableExecutionException( - ErrorObject.builder().errorMessage("original retry").build(), true); - } - - static DurableExecutionInput input() { - var operation = Operation.builder() - .id("id") - .name("test") - .type(OperationType.EXECUTION) - .status(OperationStatus.STARTED) - .startTimestamp(Instant.now()) - .executionDetails( - ExecutionDetails.builder().inputPayload("\"input\"").build()) - .build(); - return new DurableExecutionInput( - "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/name/id", - "token", - CheckpointUpdatedExecutionState.builder().operations(operation).build()); - } - - static Thread daemon(Runnable runnable, String name) { - var thread = new Thread(runnable, name); - thread.setDaemon(true); - return thread; - } - - static void await(CountDownLatch latch) { - try { - check(latch.await(3, TimeUnit.SECONDS), "fixture latch timed out"); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new AssertionError("owner interrupted", e); - } - } - - static void check(boolean condition, String message) { - if (!condition) throw new AssertionError(message); - } - - static class Deadline { - volatile long start; - volatile long end; - final CountDownLatch budgetRead = new CountDownLatch(1); - - void arm(int millis) { - start = System.nanoTime(); - end = start + TimeUnit.MILLISECONDS.toNanos(millis); - } - - long elapsedMillis() { - return TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start); - } - - int remaining() { - return end == 0 ? 30_000 : (int) Math.max(0L, TimeUnit.NANOSECONDS.toMillis(end - System.nanoTime())); - } - - Context context() { - return (Context) Proxy.newProxyInstance( - Context.class.getClassLoader(), new Class[] {Context.class}, (proxy, method, args) -> { - if (method.getName().equals("getRemainingTimeInMillis")) { - budgetRead.countDown(); - return remaining(); - } - if (method.getName().equals("getAwsRequestId")) return "test-request"; - if (method.getName().equals("getMemoryLimitInMB")) return 512; - return null; - }); - } - } - - @HandlerScoped(ScopedPlugin.Opener.class) - private abstract static class ScopedPlugin implements DurableExecutionPlugin { - public abstract AutoCloseable openHandlerScope(); - - public static class Opener implements Function { - public AutoCloseable apply(ScopedPlugin plugin) { - return plugin.openHandlerScope(); - } - } - } -} diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndCompatibilityTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndCompatibilityTest.java index 7104be813..9cb19b13f 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndCompatibilityTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndCompatibilityTest.java @@ -14,7 +14,6 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.locks.LockSupport; -import java.util.function.Function; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.CsvSource; @@ -25,19 +24,18 @@ 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.HandlerScoped; import software.amazon.lambda.durable.plugin.InvocationEndInfo; import software.amazon.lambda.durable.plugin.InvocationInfo; import software.amazon.lambda.durable.testing.local.LocalMemoryExecutionClient; -/** Preserves legacy completion-thread dispatch without promising owner affinity for precompleted futures. */ +/** Verifies invocation-hook thread affinity, including executors that complete tasks before submission returns. */ class InvocationEndCompatibilityTest { @ParameterizedTest @CsvSource({ "false,false,false", "true,false,false", "false,true,false", "true,true,false", "false,false,true", "true,false,true", "false,true,true", "true,true,true" }) - void legacyEndThread(boolean mixed, boolean precompleted, boolean pending) throws Exception { + void invocationEndRunsOnHandlerThread(boolean mixed, boolean precompleted, boolean pending) throws Exception { var local = new ThreadLocal(); var observations = new ArrayList(); var callerThread = new AtomicReference(); @@ -139,9 +137,9 @@ public void onInvocationEnd(InvocationEndInfo info) { observations.add(new Observation( startThread.get(), endThread.get(), endValue.get(), after, previous.get(), endCalls.get())); assertEquals( - mixed ? List.of("otel", "legacy") : List.of("legacy"), + mixed ? List.of("legacy", "otel") : List.of("legacy"), endOrder, - "scoped and legacy end hooks retain registration order and execute once"); + "all end hooks unwind registration order and execute once"); } finally { release.countDown(); } @@ -149,15 +147,11 @@ public void onInvocationEnd(InvocationEndInfo info) { for (int index = 0; index < observations.size(); index++) { var observation = observations.get(index); assertEquals(1, observation.endCalls()); - if (precompleted) { - assertSame(callerThread.get(), observation.endThread(), "completed futures keep caller dispatch"); - } else if (!pending) { - assertSame(observation.startThread(), observation.endThread()); - assertEquals("inv" + (index + 1), observation.endValue()); - assertNull(observation.workerAfter(), "legacy end hook clears its original worker value"); - assertNull(observation.previous(), "reused worker must not retain the prior invocation"); - } - // Suspension can complete from the handler or checkpoint processing; no owner-thread guarantee. + assertSame(observation.startThread(), observation.endThread()); + assertNotSame(callerThread.get(), observation.endThread()); + assertEquals("inv" + (index + 1), observation.endValue()); + assertNull(observation.workerAfter(), "end hook clears its original worker value"); + assertNull(observation.previous(), "reused worker must not retain the prior invocation"); } } finally { caller.shutdownNow(); @@ -168,7 +162,7 @@ public void onInvocationEnd(InvocationEndInfo info) { } @Test - void scopedSuspensionWithinBudgetKeepsLegacyCleanupOnTheReusedWorker() throws Exception { + void suspendedInvocationWaitsForEndOnTheReusedWorker() throws Exception { var local = new ThreadLocal(); var callerThread = new AtomicReference(); var callers = Executors.newSingleThreadExecutor(task -> { @@ -217,7 +211,10 @@ public void onInvocationEnd(InvocationEndInfo info) { config)); try { assertTrue(scoped.closeEntered.await(3, TimeUnit.SECONDS)); - awaitCallerHandoff(callerThread.get()); + assertThrows( + TimeoutException.class, + () -> result.get(600, TimeUnit.MILLISECONDS), + "a blocked end hook has no fallback to the caller thread"); scoped.releaseClose.countDown(); assertEquals( ExecutionStatus.PENDING, @@ -238,18 +235,6 @@ public void onInvocationEnd(InvocationEndInfo info) { } } - private static void awaitCallerHandoff(Thread caller) { - var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(3); - while (System.nanoTime() < deadline) { - if (Arrays.stream(caller.getStackTrace()) - .anyMatch(frame -> frame.getClassName().equals(DurableExecutor.class.getName()) - && frame.getMethodName().equals("awaitHandlerScopes"))) return; - LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1)); - } - fail("The caller did not begin its bounded cleanup handoff"); - } - - @HandlerScoped(PausingOtelPlugin.Opener.class) private static final class PausingOtelPlugin extends InvocationOtelPlugin { private final CountDownLatch closeEntered = new CountDownLatch(1); private final CountDownLatch releaseClose = new CountDownLatch(1); @@ -260,18 +245,13 @@ private PausingOtelPlugin() { OtelPluginConfig.builder().enableMdc(false).build()); } - public static final class Opener implements Function { - @Override - public AutoCloseable apply(PausingOtelPlugin plugin) { - var actual = new InvocationOtelPlugin.HandlerScopeOpener().apply(plugin); - return () -> { - try { - if (actual != null) actual.close(); - } finally { - plugin.closeEntered.countDown(); - await(plugin.releaseClose); - } - }; + @Override + public void onInvocationEnd(InvocationEndInfo info) { + try { + super.onInvocationEnd(info); + } finally { + closeEntered.countDown(); + await(releaseClose); } } } diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndFailureTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndFailureTest.java new file mode 100644 index 000000000..33ea7aada --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndFailureTest.java @@ -0,0 +1,133 @@ +// 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 io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.api.trace.TraceFlags; +import io.opentelemetry.api.trace.TraceState; +import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; +import io.opentelemetry.sdk.common.CompletableResultCode; +import io.opentelemetry.sdk.trace.ReadWriteSpan; +import io.opentelemetry.sdk.trace.ReadableSpan; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.SpanProcessor; +import java.time.Instant; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Stream; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +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; +import software.amazon.lambda.durable.plugin.PluginRunner; + +class InvocationEndFailureTest { + @ParameterizedTest + @MethodSource("failures") + @SuppressWarnings("removal") + void restoresContextWhenInvocationEndFails(boolean executionView, String phase, String failureKind) + throws Exception { + Throwable failure = + switch (failureKind) { + case "exception" -> new IllegalStateException("telemetry failed"); + case "linkage" -> new NoClassDefFoundError("incompatible telemetry dependency"); + case "fatal" -> new InternalError("fatal telemetry failure"); + default -> new ThreadDeath(); + }; + var processor = new FailingProcessor(phase, failure); + var builder = SdkTracerProvider.builder().addSpanProcessor(processor); + var config = OtelPluginConfig.builder() + .enableMdc(false) + .contextExtractor(() -> new ExtractedContext( + "12345678901234567890123456789012", "1234567890123456", ExtractedContext.Sampling.SAMPLED)) + .build(); + DurableExecutionPlugin plugin = + executionView ? new ExecutionOtelPlugin(builder, config) : new InvocationOtelPlugin(builder, config); + var healthyEnds = new AtomicInteger(); + var healthy = new DurableExecutionPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + healthyEnds.incrementAndGet(); + } + }; + var runner = new PluginRunner(List.of(healthy, plugin)); + var ambient = Span.wrap(SpanContext.create( + "abcdefabcdefabcdefabcdefabcdefab", + "abcdefabcdefabcd", + TraceFlags.getSampled(), + TraceState.getDefault())); + try (var ignored = ambient.makeCurrent()) { + var original = Context.current(); + runner.onInvocationStart(new InvocationInfo("req", "arn", true, Instant.now())); + assertNotSame(original, Context.current()); + if (phase.equals("scope")) { + var field = plugin.getClass().getDeclaredField("handlerScope"); + field.setAccessible(true); + var scope = (Scope) field.get(plugin); + field.set(plugin, (Scope) () -> { + scope.close(); + raise(failure); + }); + } + var end = new InvocationEndInfo("req", "arn", true, InvocationStatus.SUCCEEDED, null); + var fatal = failureKind.equals("fatal") || failureKind.equals("thread-death"); + if (fatal) assertSame(failure, assertThrows(Error.class, () -> runner.onInvocationEnd(end))); + else assertDoesNotThrow(() -> runner.onInvocationEnd(end)); + assertSame(original, Context.current(), "telemetry errors must not leave the handler context attached"); + assertEquals(1, healthyEnds.get(), "remaining invocation-end hooks must always release their resources"); + // Cleanup state is consumed even if the scope's close implementation throws. + assertDoesNotThrow(() -> plugin.onInvocationEnd(end)); + assertSame(original, Context.current()); + } + } + + private static Stream failures() { + return Stream.of(false, true) + .flatMap(executionView -> Stream.of("span", "flush", "scope") + .flatMap(phase -> Stream.of("exception", "linkage", "fatal", "thread-death") + .map(kind -> Arguments.of(executionView, phase, kind)))); + } + + private static void raise(Throwable failure) { + if (failure instanceof RuntimeException exception) throw exception; + throw (Error) failure; + } + + private record FailingProcessor(String phase, Throwable failure) implements SpanProcessor { + @Override + public void onStart(Context parent, ReadWriteSpan span) {} + + @Override + public boolean isStartRequired() { + return false; + } + + @Override + public void onEnd(ReadableSpan span) { + if (phase.equals("span")) raise(failure); + } + + @Override + public boolean isEndRequired() { + return true; + } + + @Override + public CompletableResultCode forceFlush() { + if (phase.equals("flush")) raise(failure); + return CompletableResultCode.ofSuccess(); + } + + @Override + public CompletableResultCode shutdown() { + return CompletableResultCode.ofSuccess(); + } + } +} diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndFinalizationTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndFinalizationTest.java new file mode 100644 index 000000000..5da60b868 --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndFinalizationTest.java @@ -0,0 +1,195 @@ +// 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 io.opentelemetry.api.trace.Span; +import io.opentelemetry.context.Context; +import io.opentelemetry.sdk.common.CompletableResultCode; +import io.opentelemetry.sdk.trace.ReadWriteSpan; +import io.opentelemetry.sdk.trace.ReadableSpan; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.SpanProcessor; +import java.time.Duration; +import java.time.Instant; +import java.util.concurrent.*; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +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.TypeToken; +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.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.testing.local.LocalMemoryExecutionClient; + +class InvocationEndFinalizationTest { + @ParameterizedTest + @CsvSource({"false,false", "false,true", "true,false", "true,true"}) + void waitsForHandlerFinallyAndRealOtelFlushBeforeResponding(boolean executionView, boolean retry) throws Exception { + var flush = new BlockingFlush(); + var builder = SdkTracerProvider.builder().addSpanProcessor(flush); + var settings = OtelPluginConfig.builder() + .enableMdc(false) + .contextExtractor(() -> new ExtractedContext( + "12345678901234567890123456789012", "1234567890123456", ExtractedContext.Sampling.SAMPLED)) + .build(); + DurableExecutionPlugin delegate = executionView + ? new ExecutionOtelPlugin(builder, settings) + : new InvocationOtelPlugin(builder, settings); + var owner = new AtomicReference(); + var endOwner = new AtomicReference(); + var ends = new AtomicInteger(); + var plugin = new DurableExecutionPlugin() { + @Override + public void onInvocationStart(InvocationInfo info) { + owner.set(Thread.currentThread()); + delegate.onInvocationStart(info); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + endOwner.set(Thread.currentThread()); + delegate.onInvocationEnd(info); + assertFalse( + Span.current().getSpanContext().isValid(), "End must restore ambient context before returning"); + ends.incrementAndGet(); + } + }; + var workers = Executors.newCachedThreadPool(); + var callers = Executors.newSingleThreadExecutor(); + var enteredFinally = new CountDownLatch(1); + var releaseFinally = new CountDownLatch(1); + var original = new UnrecoverableDurableExecutionException( + ErrorObject.builder().errorMessage("original retry").build(), true); + var config = DurableConfig.builder() + .withDurableExecutionClient(new LocalMemoryExecutionClient()) + .withExecutorService(workers) + .withCheckpointDelay(Duration.ZERO) + .withPlugins(plugin) + .build(); + try { + var response = callers.submit(() -> DurableExecutor.execute( + input(), + null, + TypeToken.get(String.class), + (value, ctx) -> { + assertSame(owner.get(), Thread.currentThread()); + assertTrue(Span.current().getSpanContext().isValid()); + try { + if (retry) + ctx.step("retry", String.class, step -> { + throw original; + }); + else ctx.wait("pause", Duration.ofSeconds(1)); + return "done"; + } finally { + enteredFinally.countDown(); + await(releaseFinally); + } + }, + config)); + assertTrue(enteredFinally.await(3, TimeUnit.SECONDS)); + assertThrows( + TimeoutException.class, + () -> response.get(600, TimeUnit.MILLISECONDS), + "suspension or termination cannot bypass a blocked handler finally"); + assertEquals(0, flush.calls.get(), "End must wait for the handler to unwind"); + assertEquals(0, ends.get()); + releaseFinally.countDown(); + assertTrue(flush.entered.await(3, TimeUnit.SECONDS)); + assertSame(owner.get(), endOwner.get()); + assertSame(owner.get(), flush.owner.get()); + assertThrows( + TimeoutException.class, + () -> response.get(100, TimeUnit.MILLISECONDS), + "the invocation response must wait for the plugin's actual forceFlush result"); + flush.result.succeed(); + if (retry) { + var thrown = assertThrows(ExecutionException.class, () -> response.get(3, TimeUnit.SECONDS)); + assertSame(original, thrown.getCause()); + } else { + assertEquals( + ExecutionStatus.PENDING, + response.get(3, TimeUnit.SECONDS).status()); + } + assertEquals(1, flush.calls.get()); + assertEquals(1, ends.get()); + } finally { + releaseFinally.countDown(); + flush.result.succeed(); + workers.shutdown(); + callers.shutdown(); + assertTrue(workers.awaitTermination(3, TimeUnit.SECONDS)); + assertTrue(callers.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + private static DurableExecutionInput input() { + var operation = Operation.builder() + .id("id") + .name("test") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.now()) + .executionDetails( + ExecutionDetails.builder().inputPayload("\"input\"").build()) + .build(); + return new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/name/id", + "token", + CheckpointUpdatedExecutionState.builder().operations(operation).build()); + } + + private static void await(CountDownLatch latch) { + try { + assertTrue(latch.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + throw new AssertionError(error); + } + } + + private static final class BlockingFlush implements SpanProcessor { + final AtomicInteger calls = new AtomicInteger(); + final AtomicReference owner = new AtomicReference<>(); + final CountDownLatch entered = new CountDownLatch(1); + final CompletableResultCode result = new CompletableResultCode(); + + @Override + public CompletableResultCode forceFlush() { + calls.incrementAndGet(); + owner.set(Thread.currentThread()); + entered.countDown(); + return result; + } + + @Override + public void onStart(Context parent, ReadWriteSpan span) {} + + @Override + public boolean isStartRequired() { + return false; + } + + @Override + public void onEnd(ReadableSpan span) {} + + @Override + public boolean isEndRequired() { + return false; + } + + @Override + public CompletableResultCode shutdown() { + return CompletableResultCode.ofSuccess(); + } + } +} diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOtelPluginTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOtelPluginTest.java index d8d7e7f43..79926ac69 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOtelPluginTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOtelPluginTest.java @@ -308,8 +308,8 @@ void invocationStart_staysOnExecutionTrace_withoutLinkingAmbientSpan() { try (var ignored = Span.wrap(ambientSpanContext).makeCurrent()) { plugin.onInvocationStart(new InvocationInfo("req-1", "arn:exec1", true, Instant.now())); + plugin.onInvocationEnd(new InvocationEndInfo("req-1", "arn:exec1", true, InvocationStatus.SUCCEEDED, null)); } - plugin.onInvocationEnd(new InvocationEndInfo("req-1", "arn:exec1", true, InvocationStatus.SUCCEEDED, null)); var invocationSpan = spanByName("Invocation"); var workflowSpan = spanByName("Workflow"); diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/LegacySubclassScopeCompatibilityTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/LegacySubclassScopeCompatibilityTest.java index 5c89f884a..2926a3506 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/LegacySubclassScopeCompatibilityTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/LegacySubclassScopeCompatibilityTest.java @@ -8,13 +8,15 @@ import java.net.URLClassLoader; import java.nio.file.Files; import java.nio.file.Path; -import java.util.List; +import java.time.Instant; import javax.tools.ToolProvider; import org.junit.jupiter.api.io.TempDir; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.CsvSource; import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; -import software.amazon.lambda.durable.plugin.PluginRunner; +import software.amazon.lambda.durable.plugin.InvocationEndInfo; +import software.amazon.lambda.durable.plugin.InvocationInfo; +import software.amazon.lambda.durable.plugin.InvocationStatus; class LegacySubclassScopeCompatibilityTest { @TempDir @@ -64,7 +66,8 @@ public void originalCall() throws Exception { try (var scope = openHandlerScope( new URL[] {classes.toUri().toURL()}, getClass().getClassLoader())) { var type = loader.loadClass("LegacySubclass"); var plugin = (DurableExecutionPlugin) type.getConstructor().newInstance(); - assertEquals("body", new PluginRunner(List.of(plugin)).runHandler(() -> "body")); + plugin.onInvocationStart(new InvocationInfo("req", "arn", true, Instant.now())); + plugin.onInvocationEnd(new InvocationEndInfo("req", "arn", true, InvocationStatus.SUCCEEDED, null)); assertEquals(0, type.getField("opened").get(null), "SDK must not invoke an unrelated subclass resource"); type.getMethod("originalCall").invoke(plugin); assertEquals( diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFailureBoundaryTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFailureBoundaryTest.java index 015b6a0f5..116738821 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFailureBoundaryTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFailureBoundaryTest.java @@ -6,24 +6,20 @@ import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Proxy; -import java.lang.reflect.UndeclaredThrowableException; import java.time.Duration; -import java.util.Arrays; import java.util.List; -import java.util.Map; import java.util.concurrent.AbstractExecutorService; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.CompletionException; -import java.util.concurrent.ExecutionException; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; 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.params.ParameterizedTest; import org.junit.jupiter.params.provider.CsvSource; -import org.junit.jupiter.params.provider.ValueSource; import org.slf4j.MDC; import org.slf4j.spi.MDCAdapter; import software.amazon.lambda.durable.DurableConfig; @@ -78,114 +74,39 @@ void workerCaptureFailureDoesNotEndAnInvocationThatNeverStarted() throws Excepti @ParameterizedTest @CsvSource({ - "getCopyOfContextMap,false", - "clear,false", - "setContextMap,false", - "getCopyOfContextMap,true", - "clear,true", - "setContextMap,true" + "false,false,exception", "true,false,exception", "false,true,exception", "true,true,exception", + "false,false,fatal", "true,false,fatal", "false,true,fatal", "true,true,fatal" }) - void nonfatalEndMdcFailuresPreserveTheSelectedOutcome(String method, boolean suspend) throws Exception { - runEndFailure(method, suspend, new IllegalStateException("end MDC failure")); - } - - @SuppressWarnings("removal") - @ParameterizedTest - @CsvSource({ - "getCopyOfContextMap,false", - "clear,false", - "setContextMap,false", - "getCopyOfContextMap,true", - "clear,true", - "setContextMap,true" - }) - void fatalEndMdcFailuresStillEscape(String method, boolean wrapped) throws Exception { - Error fatal = wrapped ? new ThreadDeath() : new VirtualMachineError("fatal end MDC failure") {}; - runEndFailure(method, false, wrapped ? new CompletionException(new ExecutionException(fatal)) : fatal, fatal); - } - - @ParameterizedTest - @ValueSource(booleans = {false, true}) - void unreadableOrdinaryCausePreservesSuccessAndSuspension(boolean suspend) throws Exception { - var reads = new AtomicInteger(); - var failure = new CompletionException("unreadable", null) { - @Override - public synchronized Throwable getCause() { - reads.incrementAndGet(); - throw new IllegalStateException("diagnostic unavailable"); - } - }; - runEndFailure("getCopyOfContextMap", suspend, failure); - assertEquals(1, reads.get()); - } - - @SuppressWarnings("removal") - @ParameterizedTest - @ValueSource(booleans = {false, true}) - void fatalCauseAccessorEscapesWithOriginalIdentity(boolean threadDeath) throws Exception { - var reads = new AtomicInteger(); - Error fatal = threadDeath ? new ThreadDeath() : new VirtualMachineError("fatal accessor") {}; - var failure = new CompletionException("unreadable", null) { - @Override - public synchronized Throwable getCause() { - reads.incrementAndGet(); - throw fatal; - } - }; - runEndFailure("setContextMap", false, failure, fatal); - assertEquals(1, reads.get()); - } - - @SuppressWarnings("removal") - @ParameterizedTest - @CsvSource({"reflection,false", "reflection,true", "proxy,false", "proxy,true"}) - void jdkReflectiveWrappersKeepFatalIdentity(String wrapper, boolean threadDeath) throws Exception { - Error fatal = threadDeath ? new ThreadDeath() : new VirtualMachineError("wrapped fatal") {}; - Throwable failure = wrapper.equals("reflection") - ? new InvocationTargetException(fatal) - : new UndeclaredThrowableException(fatal); - runEndFailure("clear", false, failure, fatal); - } - - @Test - void cyclicCauseIsReadOnlyOnceAndPreservesOutcome() throws Exception { - var reads = new AtomicInteger(); - var cycle = new CompletionException("cycle", null) { - @Override - public synchronized Throwable getCause() { - reads.incrementAndGet(); - return this; - } - }; - runEndFailure("getCopyOfContextMap", false, cycle); - assertEquals(1, reads.get()); - } - - @Test - void ordinaryApplicationCauseIsNotTreatedAsATransportWrapper() throws Exception { - runEndFailure("clear", false, new IllegalStateException("ordinary diagnostic", new VirtualMachineError() {})); - } - - private static void runEndFailure(String method, boolean suspend, Throwable failure) throws Exception { - runEndFailure(method, suspend, failure, null); - } - - private static void runEndFailure(String method, boolean suspend, Throwable failure, Error expectedFatal) + void workerMdcRestorationFailureCannotStrandTheInvocation(boolean ambientMdc, boolean suspend, String failureKind) throws Exception { var original = MDC.getMDCAdapter(); - var callerBefore = MDC.getCopyOfContextMap(); var injected = new AtomicBoolean(); var plugin = new RecordingPlugin(); - var worker = Executors.newSingleThreadExecutor(); + Throwable failure = failureKind.equals("fatal") + ? new InternalError("worker MDC restoration failed") + : new IllegalStateException("worker MDC restoration failed"); + var ownerFailure = new AtomicReference(); + var observed = new CountDownLatch(1); + var workers = Executors.newSingleThreadExecutor(task -> { + var thread = new Thread( + () -> { + if (ambientMdc) MDC.put("ambient", "saved"); + else MDC.clear(); + task.run(); + }, + "mdc-restoration-worker"); + thread.setUncaughtExceptionHandler((owner, error) -> { + ownerFailure.set(error); + observed.countDown(); + }); + return thread; + }); try { - replaceAdapter(proxy(original, name -> { - var inEnd = Arrays.stream(Thread.currentThread().getStackTrace()) - .anyMatch(frame -> frame.getClassName().endsWith(".DurableExecutor") - && frame.getMethodName().equals("fireOnInvocationEnd")); - if (inEnd && name.equals(method) && injected.compareAndSet(false, true)) throw failure; - // Select both legitimate MDC restoration branches, independently of the finalizer's thread. - if (inEnd && name.equals("getCopyOfContextMap")) - return method.equals("setContextMap") ? Map.of("ambient", "saved") : NullSnapshot.INSTANCE; + replaceAdapter(proxy(original, method -> { + if (Thread.currentThread().getName().equals("mdc-restoration-worker") + && plugin.ends.get() == 1 + && method.equals(ambientMdc ? "setContextMap" : "clear") + && injected.compareAndSet(false, true)) throw failure; return null; })); var runner = LocalDurableTestRunner.create( @@ -195,35 +116,23 @@ private static void runEndFailure(String method, boolean suspend, Throwable fail return "done"; }, DurableConfig.builder() - .withExecutorService(worker) + .withExecutorService(workers) .withPlugins(plugin) .build()); - if (expectedFatal != null) { - assertSame(expectedFatal, assertThrows(Error.class, () -> runner.run("input"))); - assertEquals(method.equals("getCopyOfContextMap") ? 0 : 1, plugin.ends.get()); - } else { - var result = runner.run("input"); - assertEquals(suspend ? ExecutionStatus.PENDING : ExecutionStatus.SUCCEEDED, result.getStatus()); - assertEquals(1, plugin.ends.get()); - } - assertTrue(injected.get(), "the actual execute() end boundary must exercise the adapter failure"); + var result = runner.run("input"); + assertEquals(suspend ? ExecutionStatus.PENDING : ExecutionStatus.SUCCEEDED, result.getStatus()); assertEquals(1, plugin.starts.get()); + assertEquals(1, plugin.ends.get()); + assertTrue(injected.get()); + assertTrue(observed.await(3, TimeUnit.SECONDS)); + assertSame(failure, ownerFailure.get(), "restoration failure still escapes its worker after End"); } finally { - try { - worker.shutdownNow(); - assertTrue(worker.awaitTermination(3, TimeUnit.SECONDS)); - } finally { - replaceAdapter(original); - if (callerBefore == null) MDC.clear(); - else MDC.setContextMap(callerBefore); - } + workers.shutdownNow(); + assertTrue(workers.awaitTermination(3, TimeUnit.SECONDS)); + replaceAdapter(original); } } - private enum NullSnapshot { - INSTANCE - } - @FunctionalInterface private interface Fault { Object apply(String method) throws Throwable; @@ -233,7 +142,6 @@ private static MDCAdapter proxy(MDCAdapter delegate, Fault fault) { return (MDCAdapter) Proxy.newProxyInstance( MDCAdapter.class.getClassLoader(), new Class[] {MDCAdapter.class}, (proxy, method, args) -> { var value = fault.apply(method.getName()); - if (value == NullSnapshot.INSTANCE) return null; if (value != null) return value; try { return method.invoke(delegate, args); diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ScopeFatalSignalTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ScopeFatalSignalTest.java deleted file mode 100644 index f27e66a59..000000000 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ScopeFatalSignalTest.java +++ /dev/null @@ -1,109 +0,0 @@ -// 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 java.util.concurrent.*; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Function; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.CsvSource; -import software.amazon.lambda.durable.DurableConfig; -import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; -import software.amazon.lambda.durable.plugin.HandlerScoped; -import software.amazon.lambda.durable.testing.LocalDurableTestRunner; -import software.amazon.lambda.durable.util.ExceptionHelper; - -class ScopeFatalSignalTest { - @ParameterizedTest - @CsvSource({ - "false,false,false", - "true,false,false", - "false,true,false", - "true,true,false", - "false,false,true", - "true,false,true", - "false,true,true", - "true,true,true" - }) - @SuppressWarnings("removal") - void fatalCloseWakesCallerAndStillAllowsBoundedEarlierCleanup( - boolean threadDeath, boolean releaseInBudget, boolean fatalAccessor) throws Exception { - var earlierEntered = new CountDownLatch(1); - var release = new CountDownLatch(1); - var closed = new AtomicBoolean(); - var ownerFatal = new AtomicReference(); - var ownerObserved = new CountDownLatch(1); - Error fatal = threadDeath ? new ThreadDeath() : new InternalError("later scope close"); - var workers = Executors.newSingleThreadExecutor(task -> { - var thread = new Thread(task, "scope-fatal-owner"); - thread.setUncaughtExceptionHandler((owner, failure) -> { - ownerFatal.set(failure); - ownerObserved.countDown(); - }); - return thread; - }); - var callers = Executors.newSingleThreadExecutor(); - try { - var earlier = new Scoped(() -> { - earlierEntered.countDown(); - if (!release.await(5, TimeUnit.SECONDS)) throw new AssertionError("release timed out"); - closed.set(true); - }); - var later = new Scoped(() -> { - if (fatalAccessor) - throw new CompletionException("scope cause accessor", null) { - @Override - public synchronized Throwable getCause() { - throw fatal; - } - }; - throw fatal; - }); - var config = DurableConfig.builder() - .withExecutorService(workers) - .withPlugins(earlier, later) - .build(); - var runner = LocalDurableTestRunner.create(String.class, (input, context) -> input, config); - var result = callers.submit(() -> runner.run("input")); - assertTrue(earlierEntered.await(3, TimeUnit.SECONDS)); - assertThrows( - TimeoutException.class, - () -> result.get(100, TimeUnit.MILLISECONDS), - "earlier cleanup must receive its handoff budget"); - if (releaseInBudget) release.countDown(); - var failure = assertThrows( - ExecutionException.class, - () -> result.get(2, TimeUnit.SECONDS), - "reported fatal must wake the invocation without waiting forever for earlier cleanup"); - assertSame(fatal, ExceptionHelper.unwrapCompletableFuture(failure.getCause())); - assertEquals(releaseInBudget, closed.get()); - } finally { - release.countDown(); - callers.shutdown(); - workers.shutdown(); - assertTrue(callers.awaitTermination(3, TimeUnit.SECONDS)); - assertTrue(workers.awaitTermination(3, TimeUnit.SECONDS)); - } - assertTrue(ownerObserved.await(3, TimeUnit.SECONDS)); - assertSame(fatal, ownerFatal.get(), "the fatal still escapes its actual owner"); - } - - @HandlerScoped(Opener.class) - public static final class Scoped implements DurableExecutionPlugin { - private final AutoCloseable closer; - - Scoped(AutoCloseable closer) { - this.closer = closer; - } - } - - public static final class Opener implements Function { - @Override - public AutoCloseable apply(Object plugin) { - return ((Scoped) plugin).closer; - } - } -} diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/HandlerScopeFinalizationIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/HandlerScopeFinalizationIntegrationTest.java deleted file mode 100644 index e1b52e08e..000000000 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/HandlerScopeFinalizationIntegrationTest.java +++ /dev/null @@ -1,192 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; - -import static org.junit.jupiter.api.Assertions.*; - -import java.time.Duration; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.Executors; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.function.Function; -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.ErrorObject; -import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; -import software.amazon.lambda.durable.model.ExecutionStatus; -import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; -import software.amazon.lambda.durable.plugin.HandlerScoped; -import software.amazon.lambda.durable.plugin.InvocationEndInfo; -import software.amazon.lambda.durable.testing.LocalDurableTestRunner; - -class HandlerScopeFinalizationIntegrationTest { - @ParameterizedTest - @ValueSource(booleans = {false, true}) - void blockedCleanupReturnsTheOriginalOutcomeThenClosesOnItsOwner(boolean terminate) throws Exception { - var releaseFinally = new CountDownLatch(1); - var enteredFinally = new CountDownLatch(1); - var scopeClosed = new CountDownLatch(1); - var closedAtEnd = new AtomicBoolean(); - var plugin = new ScopedPlugin() { - @Override - public AutoCloseable openHandlerScope() { - var owner = Thread.currentThread(); - return () -> { - assertSame(owner, Thread.currentThread()); - scopeClosed.countDown(); - }; - } - - @Override - public void onInvocationEnd(InvocationEndInfo info) { - closedAtEnd.set(scopeClosed.getCount() == 0); - } - }; - var runner = LocalDurableTestRunner.create( - String.class, - (input, context) -> { - try { - if (terminate) - context.step("terminate", String.class, step -> { - throw new UnrecoverableDurableExecutionException( - ErrorObject.builder() - .errorMessage("retry") - .build(), - true); - }); - else context.wait("pause", Duration.ofSeconds(1)); - return "done"; - } finally { - enteredFinally.countDown(); - try { - if (!releaseFinally.await(5, TimeUnit.SECONDS)) throw new AssertionError("not released"); - } catch (InterruptedException error) { - Thread.currentThread().interrupt(); - throw new AssertionError(error); - } - } - }, - DurableConfig.builder().withPlugins(plugin).build()); - var caller = Executors.newSingleThreadExecutor(); - try { - var response = caller.submit(() -> runner.run("input")); - assertTrue(enteredFinally.await(5, TimeUnit.SECONDS)); - if (terminate) { - var failure = assertThrows(ExecutionException.class, () -> response.get(2, TimeUnit.SECONDS)); - assertInstanceOf(UnrecoverableDurableExecutionException.class, failure.getCause()); - } else - assertEquals( - ExecutionStatus.PENDING, - response.get(2, TimeUnit.SECONDS).getStatus()); - assertFalse(closedAtEnd.get()); - assertEquals(1L, scopeClosed.getCount()); - } finally { - releaseFinally.countDown(); - assertTrue(scopeClosed.await(5, TimeUnit.SECONDS)); - caller.shutdownNow(); - } - } - - @ParameterizedTest - @CsvSource({"true,false", "false,false", "true,true", "false,true"}) - void finalizationWaitsForAnOpenedScopeWithoutChangingLegacyNoScopeTiming(boolean hasScope, boolean terminate) - throws Exception { - var finallyEntered = new CountDownLatch(1); - var releaseFinally = new CountDownLatch(1); - var handlerExited = new CountDownLatch(1); - var scopeClosed = new AtomicBoolean(); - var endCalled = new AtomicBoolean(); - var scopeClosedAtEnd = new AtomicBoolean(); - var plugin = new ScopedPlugin() { - @Override - public AutoCloseable openHandlerScope() { - if (!hasScope) return null; - var owner = Thread.currentThread(); - return () -> { - assertSame(owner, Thread.currentThread()); - scopeClosed.set(true); - }; - } - - @Override - public void onInvocationEnd(InvocationEndInfo info) { - scopeClosedAtEnd.set(scopeClosed.get()); - endCalled.set(true); - } - }; - var config = DurableConfig.builder().withPlugins(plugin).build(); - var runner = LocalDurableTestRunner.create( - String.class, - (input, context) -> { - try { - if (terminate) { - context.step("terminate", String.class, step -> { - throw new UnrecoverableDurableExecutionException( - ErrorObject.builder() - .errorMessage("retry invocation") - .build(), - true); - }); - } else { - context.wait("pause", Duration.ofSeconds(1)); - } - return "done"; - } finally { - finallyEntered.countDown(); - try { - if (!releaseFinally.await(5, TimeUnit.SECONDS)) - throw new AssertionError("finally not released"); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new AssertionError(e); - } finally { - handlerExited.countDown(); - } - } - }, - config); - var caller = Executors.newSingleThreadExecutor(task -> { - var thread = new Thread(task, "test-invocation-caller"); - thread.setDaemon(true); - return thread; - }); - try { - var response = caller.submit(() -> runner.run("input")); - assertTrue(finallyEntered.await(5, TimeUnit.SECONDS)); - if (hasScope) { - assertThrows(TimeoutException.class, () -> response.get(150, TimeUnit.MILLISECONDS)); - assertFalse(endCalled.get(), "flush must wait for same-thread scope cleanup"); - releaseFinally.countDown(); - } - if (terminate) { - var failure = assertThrows(ExecutionException.class, () -> response.get(5, TimeUnit.SECONDS)); - assertInstanceOf(UnrecoverableDurableExecutionException.class, failure.getCause()); - } else { - assertEquals( - ExecutionStatus.PENDING, - response.get(5, TimeUnit.SECONDS).getStatus()); - } - assertTrue(endCalled.get()); - assertEquals(hasScope, scopeClosedAtEnd.get()); - } finally { - releaseFinally.countDown(); - assertTrue(handlerExited.await(5, TimeUnit.SECONDS)); - caller.shutdownNow(); - } - } - - @HandlerScoped(ScopedPlugin.Opener.class) - private abstract static class ScopedPlugin implements DurableExecutionPlugin { - public abstract AutoCloseable openHandlerScope(); - - public static class Opener implements Function { - public AutoCloseable apply(ScopedPlugin plugin) { - return plugin.openHandlerScope(); - } - } - } -} diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/InvocationFinalizationIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/InvocationFinalizationIntegrationTest.java new file mode 100644 index 000000000..f9bfc5200 --- /dev/null +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/InvocationFinalizationIntegrationTest.java @@ -0,0 +1,226 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.*; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; +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.ErrorObject; +import software.amazon.lambda.durable.config.StepConfig; +import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.execution.ExecutionManager; +import software.amazon.lambda.durable.execution.SuspendExecutionException; +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; +import software.amazon.lambda.durable.retry.RetryStrategies; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class InvocationFinalizationIntegrationTest { + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void suspensionAndRetryWaitForHandlerFinallyAndOwnerEnd(boolean retry) throws Exception { + var enteredFinally = new CountDownLatch(1); + var releaseFinally = new CountDownLatch(1); + var plugin = new RecordingPlugin(); + var original = retryError(); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + try { + if (retry) + context.step("retry", String.class, step -> { + throw original; + }); + else context.wait("pause", Duration.ofSeconds(1)); + return "done"; + } finally { + enteredFinally.countDown(); + await(releaseFinally); + } + }, + DurableConfig.builder().withPlugins(plugin).build()); + var caller = Executors.newSingleThreadExecutor(); + try { + var response = caller.submit(() -> runner.run("input")); + assertTrue(enteredFinally.await(5, TimeUnit.SECONDS)); + assertThrows(TimeoutException.class, () -> response.get(700, TimeUnit.MILLISECONDS)); + assertTrue(plugin.ends.isEmpty(), "End must follow the handler's finally block"); + releaseFinally.countDown(); + if (retry) { + var thrown = assertThrows(ExecutionException.class, () -> response.get(5, TimeUnit.SECONDS)); + assertSame(original, thrown.getCause()); + } else + assertEquals( + ExecutionStatus.PENDING, + response.get(5, TimeUnit.SECONDS).getStatus()); + assertPaired(plugin, retry ? InvocationStatus.RETRYING : InvocationStatus.PENDING); + } finally { + releaseFinally.countDown(); + stop(caller); + } + } + + @ParameterizedTest + @CsvSource({"false,false", "false,true", "true,false", "true,true"}) + void externallySelectedOutcomeSurvivesReturningOrThrowingFinally(boolean retry, boolean throwsInFinally) + throws Exception { + var enteredFinally = new CountDownLatch(1); + var releaseFinally = new CountDownLatch(1); + var manager = new AtomicReference(); + var plugin = new RecordingPlugin(); + var original = retryError(); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + manager.set(((DurableContextImpl) context).getExecutionManager()); + try { + return "handler completed"; + } finally { + enteredFinally.countDown(); + await(releaseFinally); + if (throwsInFinally) throw new IllegalStateException("finally failed"); + } + }, + DurableConfig.builder().withPlugins(plugin).build()); + var caller = Executors.newSingleThreadExecutor(); + try { + var response = caller.submit(() -> runner.run("input")); + assertTrue(enteredFinally.await(5, TimeUnit.SECONDS)); + if (retry) + assertSame( + original, + assertThrows( + UnrecoverableDurableExecutionException.class, + () -> manager.get().terminateExecution(original))); + else + assertThrows( + SuspendExecutionException.class, () -> manager.get().suspendExecution()); + assertThrows(TimeoutException.class, () -> response.get(700, TimeUnit.MILLISECONDS)); + releaseFinally.countDown(); + if (retry) { + var thrown = assertThrows(ExecutionException.class, () -> response.get(5, TimeUnit.SECONDS)); + assertSame(original, thrown.getCause()); + } else + assertEquals( + ExecutionStatus.PENDING, + response.get(5, TimeUnit.SECONDS).getStatus()); + assertPaired(plugin, retry ? InvocationStatus.RETRYING : InvocationStatus.PENDING); + assertSame(retry ? original : null, plugin.ends.get(0).executionError()); + assertNull(plugin.ends.get(0).executionResult()); + } finally { + releaseFinally.countDown(); + stop(caller); + } + } + + @Test + void hooksStayPairedAcrossStepRetryAndReplayWithoutRepeatingCompletedWork() { + var plugin = new RecordingPlugin(); + var sideEffects = new AtomicInteger(); + var attempts = new AtomicInteger(); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + var result = context.step("completed", Integer.class, step -> sideEffects.incrementAndGet()); + context.step( + "retry", + String.class, + step -> { + if (attempts.incrementAndGet() == 1) throw new IllegalStateException("try again"); + return "retried"; + }, + StepConfig.builder() + .retryStrategy(RetryStrategies.fixedDelay(2, Duration.ofSeconds(1))) + .build()); + context.wait("pause", Duration.ofSeconds(1)); + return result; + }, + DurableConfig.builder().withPlugins(plugin).build()); + + var output = runner.runUntilComplete("input"); + + assertEquals(ExecutionStatus.SUCCEEDED, output.getStatus()); + assertEquals(1, output.getResult(Integer.class)); + assertEquals(1, sideEffects.get(), "the completed step must replay without repeating user code"); + assertEquals(2, attempts.get()); + assertTrue(plugin.starts.size() >= 3, "retry and wait should each suspend before the successful invocation"); + assertEquals(plugin.starts.size(), plugin.ends.size()); + assertEquals(plugin.startThreads, plugin.endThreads); + assertTrue(plugin.endLocalValues.stream().allMatch("invocation"::equals)); + assertEquals(InvocationStatus.PENDING, plugin.ends.get(0).invocationStatus()); + assertEquals( + InvocationStatus.SUCCEEDED, + plugin.ends.get(plugin.ends.size() - 1).invocationStatus()); + assertTrue(plugin.starts.get(0).isFirstInvocation()); + assertFalse(plugin.starts.get(plugin.starts.size() - 1).isFirstInvocation()); + } + + private static void assertPaired(RecordingPlugin plugin, InvocationStatus status) { + assertEquals(1, plugin.starts.size()); + assertEquals(1, plugin.ends.size()); + assertEquals(status, plugin.ends.get(0).invocationStatus()); + assertEquals(plugin.startThreads, plugin.endThreads); + assertEquals(List.of("invocation"), plugin.endLocalValues); + } + + private static UnrecoverableDurableExecutionException retryError() { + return new UnrecoverableDurableExecutionException( + ErrorObject.builder().errorMessage("retry invocation").build(), true); + } + + private static void await(CountDownLatch latch) { + try { + if (!latch.await(5, TimeUnit.SECONDS)) throw new AssertionError("finally not released"); + } catch (InterruptedException failure) { + Thread.currentThread().interrupt(); + throw new AssertionError(failure); + } + } + + private static void stop(ExecutorService executor) throws InterruptedException { + executor.shutdownNow(); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } + + private static class RecordingPlugin implements DurableExecutionPlugin { + final ThreadLocal local = new ThreadLocal<>(); + final List starts = new ArrayList<>(); + final List ends = new ArrayList<>(); + final List startThreads = new ArrayList<>(); + final List endThreads = new ArrayList<>(); + final List endLocalValues = new ArrayList<>(); + + @Override + public void onInvocationStart(InvocationInfo info) { + starts.add(info); + startThreads.add(Thread.currentThread()); + local.set("invocation"); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + ends.add(info); + endThreads.add(Thread.currentThread()); + endLocalValues.add(local.get()); + local.remove(); + } + } +} 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 2ec749c81..5197cc716 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 @@ -4,20 +4,12 @@ import com.amazonaws.services.lambda.runtime.Context; import com.amazonaws.services.lambda.runtime.RequestHandler; -import java.lang.reflect.InvocationTargetException; -import java.lang.reflect.UndeclaredThrowableException; import java.nio.charset.StandardCharsets; -import java.util.Collections; -import java.util.IdentityHashMap; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; -import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; +import java.util.function.Function; import java.util.function.Supplier; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -59,11 +51,6 @@ public class DurableExecutor { // Lambda response size limit is 6MB minus small epsilon for envelope private static final int LAMBDA_RESPONSE_SIZE_LIMIT = 6 * 1024 * 1024 - 50; - // Best-effort allowance for each configured plugin's finalization (including the bundled OTel 5s join), - // plus shutdown/response headroom. Existing arbitrary callbacks and checkpoint draining can exceed it. - private static final long PLUGIN_FINALIZATION_RESERVE_MILLIS = 5_000; - private static final long SHUTDOWN_RESPONSE_RESERVE_MILLIS = 1_000; - private DurableExecutor() {} public static DurableExecutionOutput execute( @@ -72,396 +59,202 @@ public static DurableExecutionOutput execute( TypeToken inputType, BiFunction handler, DurableConfig config) { - var pluginRunner = config.getPluginRunner(); - var scopeFatal = new AtomicReference(); - var scopeFailure = new CompletableFuture(); - try (var executionManager = new ExecutionManager(input, config, lambdaContext)) { - var isFirstInvocation = !executionManager.isReplaying(); - var requestId = lambdaContext != null ? lambdaContext.getAwsRequestId() : null; - var executionArn = input.durableExecutionArn(); - - executionManager.registerActiveThread(null); - // Captured for onInvocationEnd, which runs outside the handler thread below. - var pluginExecutionInput = new AtomicReference<>(); - var invocationStarted = new AtomicBoolean(); - var hasHandlerScope = new AtomicBoolean(); - var handlerFuture = supplyHandler( - () -> { - executionManager.setCurrentThreadContext(new ThreadContext(null, ThreadType.CONTEXT)); - - // Deserialize once and share the value with the plugin hooks and the handler below. A second - // deserialization would double the cost, hand plugins a different object than the handler, and - // re-run any side effects in a stateful custom SerDes. A failure is captured rather than thrown - // so onInvocationStart still fires before it surfaces, keeping the start/end hooks paired. - // SerDes is a public extension point whose deserialize declares no checked exceptions, so an - // implementation may sneaky-throw one; capture every Throwable and rethrow it unchanged. - I userInput = null; - Throwable inputFailure = null; - try { - userInput = extractUserInput( - executionManager.getExecutionOperation(), config.getSerDes(), inputType); - } catch (Throwable t) { - inputFailure = t; - } - pluginExecutionInput.set(userInput); - - // onInvocationStart runs on the user thread so plugins can - // inject ThreadLocal objects, update MDC, etc. - // executionStartTime comes from the initial EXECUTION operation in the first backend event. - if (!pluginRunner.isEmpty()) { - var invocationInfo = new InvocationInfo( - requestId, - executionArn, - isFirstInvocation, - executionManager.getExecutionOperation().startTimestamp(), - userInput, - PluginInfoConverter.toOperationItemMap( - executionManager.getOperationsSnapshot(), - executionManager.getInitialOperationIds()), - PluginInfoConverter.toOperationItemMap( - executionManager.getUpdatedOperationsSnapshot(), - executionManager.getInitialOperationIds())); - invocationStarted.set(true); - pluginRunner.onInvocationStart(invocationInfo); - } - if (inputFailure != null) { - ExceptionHelper.sneakyThrow(inputFailure); - } - - var context = DurableContextImpl.createRootContext(executionManager, config, lambdaContext); - DurableContextImpl.setCurrentContext(context); - // use a try-with-resources to clear logger properties - try (var ignored = DurableLogger.attachContext()) { - var handlerInput = userInput; - return pluginRunner.runHandler( - () -> handler.apply(handlerInput, context), - () -> hasHandlerScope.set(true), - fatal -> { - scopeFatal.compareAndSet(null, fatal); - scopeFailure.completeExceptionally(scopeFatal.get()); - }); - } - }, - config.getExecutorService(), - pluginRunner, - scopeFatal); // Get executor from config for running user code - - // Execute the handlerFuture in ExecutionManager. If it completes successfully, the output of user function - // will be returned. Otherwise, it will complete exceptionally with a SuspendExecutionException or a - // failure. + try (var manager = new ExecutionManager(input, config, lambdaContext)) { + manager.registerActiveThread(ROOT_THREAD_ID); + var invocation = new Invocation<>(input, lambdaContext, inputType, handler, config, manager); try { - var executionFuture = executionManager.runUntilCompleteOrSuspend(handlerFuture); - return finalizeAfterHandlerScopes( - pluginRunner.isEmpty() ? executionFuture : observeScopeFailure(executionFuture, scopeFailure), - handlerFuture, - hasHandlerScope, - lambdaContext, - pluginRunner.getPlugins().size(), - scopeFatal, - (result, ex) -> { - if (ex != null) { - // an exception thrown from handlerFuture or suspension/termination occurred - Throwable cause = ExceptionHelper.unwrapCompletableFuture(ex); - - // return PENDING if it's SuspendExecutionException - if (cause instanceof SuspendExecutionException) { - fireOnInvocationEnd( - invocationStarted.get(), - scopeFatal, - pluginRunner, - executionManager, - requestId, - executionArn, - isFirstInvocation, - InvocationStatus.PENDING, - null, - pluginExecutionInput.get(), - null); - return DurableExecutionOutput.pending(); - } - - // let the backend retry the invocation if the exception is retryable - if (cause - instanceof - UnrecoverableDurableExecutionException - unrecoverableDurableExecutionException - && unrecoverableDurableExecutionException.isRetryable()) { - fireOnInvocationEnd( - invocationStarted.get(), - scopeFatal, - pluginRunner, - executionManager, - requestId, - executionArn, - isFirstInvocation, - InvocationStatus.RETRYING, - cause, - pluginExecutionInput.get(), - null); - throw unrecoverableDurableExecutionException; - } - - // fail the execution otherwise - logger.debug("Execution failed: {}", cause.getMessage()); - fireOnInvocationEnd( - invocationStarted.get(), - scopeFatal, - pluginRunner, - executionManager, - requestId, - executionArn, - isFirstInvocation, - InvocationStatus.FAILED, - cause, - pluginExecutionInput.get(), - null); - return DurableExecutionOutput.failure(buildErrorObject(cause, config.getSerDes())); - } - // user handler complete successfully - logger.debug("Execution completed"); - var outputPayload = config.getSerDes().serialize(result); - var output = - DurableExecutionOutput.success(handleLargePayload(executionManager, outputPayload)); - fireOnInvocationEnd( - invocationStarted.get(), - scopeFatal, - pluginRunner, - executionManager, - requestId, - executionArn, - isFirstInvocation, - InvocationStatus.SUCCEEDED, - null, - pluginExecutionInput.get(), - result); - return output; - }); - } catch (CompletionException e) { - // unwrap the CompletionException and rethrow the wrapped exception - ExceptionHelper.sneakyThrow(ExceptionHelper.unwrapCompletableFuture(e)); + return invocation.execute().join(); + } catch (CompletionException failure) { + ExceptionHelper.sneakyThrow(ExceptionHelper.unwrapCompletableFuture(failure)); return null; } } } - private static CompletableFuture supplyHandler( - Supplier task, Executor executor, PluginRunner plugins, AtomicReference scopeFatal) { - if (plugins.isEmpty()) return CompletableFuture.supplyAsync(task, executor); - var result = new CompletableFuture(); - var caller = Thread.currentThread(); - Runnable work = (Runnable & CompletableFuture.AsynchronousCompletionTask) () -> { - SafeCloseable mdcRestore; + /** Invocation-local state, accessed on the handler thread when plugins are present. */ + private static final class Invocation { + private final Context lambdaContext; + private final TypeToken inputType; + private final BiFunction handler; + private final DurableConfig config; + private final ExecutionManager manager; + private final PluginRunner plugins; + private final String requestId; + private final String executionArn; + private final boolean isFirstInvocation; + private I userInput; + private boolean started; + + private Invocation( + DurableExecutionInput input, + Context lambdaContext, + TypeToken inputType, + BiFunction handler, + DurableConfig config, + ExecutionManager manager) { + this.lambdaContext = lambdaContext; + this.inputType = inputType; + this.handler = handler; + this.config = config; + this.manager = manager; + plugins = config.getPluginRunner(); + requestId = lambdaContext != null ? lambdaContext.getAwsRequestId() : null; + executionArn = input.durableExecutionArn(); + isFirstInvocation = !manager.isReplaying(); + } + + private CompletableFuture execute() { + if (plugins.isEmpty()) { + var body = CompletableFuture.supplyAsync(this::invokeHandler, config.getExecutorService()); + return manager.runUntilCompleteOrSuspend(body).handle(this::finishInvocation); + } + 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); + return supplyHandler( + () -> { + Outcome.capture(this::invokeHandler).complete(body); + var selected = outcome.join(); + return finishInvocation(selected.value(), selected.failure()); + }, + failure -> finishInvocation(null, failure), + config.getExecutorService()); + } + + private O invokeHandler() { + manager.setCurrentThreadContext(new ThreadContext(ROOT_THREAD_ID, ThreadType.CONTEXT)); + Throwable inputFailure = null; try { - mdcRestore = restoreMdcOnClose(); + userInput = extractUserInput(manager.getExecutionOperation(), config.getSerDes(), inputType); } catch (Throwable failure) { - // Initialization can fail before task execution; the observer must still receive that failure. - result.completeExceptionally(failure); - return; + // Deserialize only once. Even failed input gets paired start/end hooks with a null input value. + inputFailure = failure; } - try (var ignored = mdcRestore) { - try { - result.complete(task.get()); - } catch (Throwable failure) { - result.completeExceptionally(failure); - // Preserve legacy user-body Error handling. Only a positively reported new scope fatal also - // escapes its actual owner after its observation future settles. Direct callers finalize below. - var fatal = scopeFatal.get(); - if (Thread.currentThread() != caller - && fatal != null - && ExceptionHelper.unwrapCompletableFuture(failure) == fatal) throw fatal; - } + fireOnInvocationStart(); + if (inputFailure != null) ExceptionHelper.sneakyThrow(inputFailure); + var context = DurableContextImpl.createRootContext(manager, config, lambdaContext); + DurableContextImpl.setCurrentContext(context); + try (var ignored = DurableLogger.attachContext()) { + return handler.apply(userInput, context); } - }; - executor.execute(work); - return result; - } - - /** Wakes the caller when a reported scope fatal precedes completion of earlier scope cleanup. */ - private static CompletableFuture observeScopeFailure( - CompletableFuture execution, CompletableFuture scopeFailure) { - var result = new CompletableFuture(); - execution.whenComplete((value, failure) -> { - if (failure == null) result.complete(value); - else result.completeExceptionally(failure); - }); - scopeFailure.whenComplete((ignored, failure) -> { - if (failure != null) result.completeExceptionally(failure); - }); - return result; - } - - private static R finalizeAfterHandlerScopes( - CompletableFuture executionFuture, - CompletableFuture handlerFuture, - AtomicBoolean hasHandlerScope, - Context lambdaContext, - int pluginCount, - AtomicReference scopeFatal, - BiFunction finalizer) { - var started = new AtomicBoolean(); - var finalized = new CompletableFuture(); - // Attach before waiting so legacy hooks retain normal CompletableFuture completion-thread dispatch. - var ready = executionFuture.handle((value, failure) -> { - Runnable finish = () -> completeFinalization(started, finalized, () -> finalizer.apply(value, failure)); - if (!hasHandlerScope.get() || handlerFuture.isDone()) finish.run(); - else handlerFuture.whenComplete((ignored, ignoredFailure) -> finish.run()); - return finish; - }); - var finish = ready.join(); - if (!started.get()) { - // Only the caller waits: a signaling handler must be free to unwind and close its scopes. - awaitHandlerScopes(executionFuture, handlerFuture, hasHandlerScope, lambdaContext, pluginCount, scopeFatal); - // Completion publishes the handler result before its callbacks necessarily run. Once cleanup completed, - // let its registered callback finalize on the owner instead of racing it from this waiter. - if (!handlerFuture.isDone()) finish.run(); } - return finalized.join(); - } - private static void completeFinalization( - AtomicBoolean started, CompletableFuture result, Supplier finalizer) { - if (!started.compareAndSet(false, true)) return; - try { - result.complete(finalizer.get()); - } catch (Throwable failure) { - // Match CompletableFuture.handle: transfer failures unchanged for the invocation caller to rethrow. - result.completeExceptionally(failure); + private void fireOnInvocationStart() { + if (plugins.isEmpty()) return; + var info = new InvocationInfo( + requestId, + executionArn, + isFirstInvocation, + manager.getExecutionOperation().startTimestamp(), + userInput, + PluginInfoConverter.toOperationItemMap( + manager.getOperationsSnapshot(), manager.getInitialOperationIds()), + PluginInfoConverter.toOperationItemMap( + manager.getUpdatedOperationsSnapshot(), manager.getInitialOperationIds())); + started = true; + plugins.onInvocationStart(info); } - } - static CompletableFuture awaitHandlerScopes( - CompletableFuture executionFuture, - CompletableFuture handlerFuture, - AtomicBoolean hasHandlerScope, - Context lambdaContext, - int pluginCount, - AtomicReference scopeFatal) { - // This method runs on the invocation caller, never as a callback on the signaling handler worker. - // Preserve the winning outcome except for an observed fatal error from the new scope callbacks. - var failure = executionFuture.handle((result, error) -> error).join(); - // A reported fatal can precede completion of earlier scopes; retain their bounded cleanup budget. - if (!hasHandlerScope.get() || (failure == null && scopeFatal.get() == null)) { - throwIfScopeFatal(scopeFatal); - return executionFuture; + private DurableExecutionOutput finishInvocation(O value, Throwable failure) { + if (failure != null) return finishFailure(ExceptionHelper.unwrapCompletableFuture(failure)); + DurableExecutionOutput output; + try { + var payload = config.getSerDes().serialize(value); + output = DurableExecutionOutput.success(handleLargePayload(manager, payload)); + } catch (Throwable deliveryFailure) { + // Result serialization/checkpointing also belongs to this invocation. Close plugin resources even + // when delivery fails, then retain the original exception for the Lambda caller. + var cause = ExceptionHelper.unwrapCompletableFuture(deliveryFailure); + fireOnInvocationEnd(failureStatus(cause), cause, null); + ExceptionHelper.sneakyThrow(deliveryFailure); + return null; + } + fireOnInvocationEnd(InvocationStatus.SUCCEEDED, null, value); + return output; } - var reserve = SHUTDOWN_RESPONSE_RESERVE_MILLIS + PLUGIN_FINALIZATION_RESERVE_MILLIS * pluginCount; - var budgetMillis = lambdaContext == null - ? 500L - : Math.min(500L, Math.max(0L, (long) lambdaContext.getRemainingTimeInMillis() - reserve)); - try { - handlerFuture.handle((result, error) -> null).get(budgetMillis, TimeUnit.MILLISECONDS); - } catch (TimeoutException e) { - logger.warn("Handler scope cleanup exceeded its handoff budget; cleanup continues on the handler thread"); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - logger.warn("Interrupted while awaiting handler scope cleanup; preserving the execution outcome", e); - } catch (ExecutionException e) { - logger.warn("Could not observe handler scope cleanup; preserving the execution outcome", e); + + private DurableExecutionOutput finishFailure(Throwable cause) { + var status = failureStatus(cause); + fireOnInvocationEnd(status, status == InvocationStatus.PENDING ? null : cause, null); + if (status == InvocationStatus.PENDING) return DurableExecutionOutput.pending(); + if (status == InvocationStatus.RETRYING) { + ExceptionHelper.sneakyThrow(cause); + return null; + } + return DurableExecutionOutput.failure(buildErrorObject(cause, config.getSerDes())); } - throwIfScopeFatal(scopeFatal); - return executionFuture; - } - private static void throwIfScopeFatal(AtomicReference scopeFatal) { - var fatal = scopeFatal.get(); - if (fatal != null) throw fatal; + private void fireOnInvocationEnd(InvocationStatus status, Throwable error, Object result) { + if (!started) return; + plugins.onInvocationEnd(new InvocationEndInfo( + requestId, + executionArn, + isFirstInvocation, + manager.getExecutionOperation().startTimestamp(), + PluginInfoConverter.toOperationItemMap( + manager.getOperationsSnapshot(), manager.getInitialOperationIds()), + status, + error, + userInput, + result)); + } } - private static SafeCloseable restoreMdcOnClose() { - var previous = MDC.getCopyOfContextMap(); - return () -> { - if (previous == null) MDC.clear(); - else MDC.setContextMap(previous); - }; + private static InvocationStatus failureStatus(Throwable failure) { + if (failure instanceof SuspendExecutionException) return InvocationStatus.PENDING; + if (failure instanceof UnrecoverableDurableExecutionException unrecoverable && unrecoverable.isRetryable()) { + return InvocationStatus.RETRYING; + } + return InvocationStatus.FAILED; } - /** MDC is ancillary to end dispatch; ordinary adapter failures must not replace the selected outcome. */ - private static SafeCloseable preserveEndMdc() { - SafeCloseable restore; - try { - restore = restoreMdcOnClose(); - } catch (Throwable failure) { - reportEndMdcFailure(failure); - return () -> {}; - } - return () -> { + /** Captures task completion without running lifecycle hooks on CompletableFuture completion threads. */ + private record Outcome(T value, Throwable failure) { + private static Outcome capture(Supplier task) { try { - restore.close(); + return new Outcome<>(task.get(), null); } catch (Throwable failure) { - reportEndMdcFailure(failure); + return new Outcome<>(null, failure); } - }; - } + } - private static void reportEndMdcFailure(Throwable failure) { - throwIfMdcFatal(failure); - try { - // Do not ask a malformed Throwable for its message or cause again while logging. - logger.warn( - "Could not preserve MDC around invocation-end dispatch; preserving execution outcome ({})", - failure.getClass().getName()); - } catch (Throwable loggingFailure) { - // An MDC-backed logger may encounter the same adapter failure while reporting it. - throwIfMdcFatal(loggingFailure); + private void complete(CompletableFuture future) { + if (failure == null) future.complete(value); + else future.completeExceptionally(failure); } } - @SuppressWarnings("removal") - private static void throwIfMdcFatal(Throwable failure) { - if (failure instanceof VirtualMachineError fatal) throw fatal; - if (failure instanceof ThreadDeath fatal) throw fatal; - var seen = Collections.newSetFromMap(new IdentityHashMap()); - var cause = failure; - while (seen.add(cause)) { - if (cause instanceof VirtualMachineError fatal) throw fatal; - if (cause instanceof ThreadDeath fatal) throw fatal; - if (!(cause instanceof CompletionException - || cause instanceof ExecutionException - || cause instanceof InvocationTargetException - || cause instanceof UndeclaredThrowableException)) return; + private static CompletableFuture supplyHandler( + Supplier task, Function initializationFailure, Executor executor) { + var result = new CompletableFuture(); + Runnable work = (Runnable & CompletableFuture.AsynchronousCompletionTask) () -> { + SafeCloseable restore; try { - cause = cause.getCause(); - } catch (VirtualMachineError | ThreadDeath fatal) { - throw fatal; - } catch (Throwable unreadableCause) { + restore = restoreMdcOnClose(); + } catch (Throwable failure) { + Outcome.capture(() -> initializationFailure.apply(failure)).complete(result); return; } - if (cause == null) return; - } + var outcome = Outcome.capture(task); + try { + restore.close(); + } finally { + // End and worker restoration have run before publishing. A restoration failure still escapes its + // owner, without changing the invocation outcome already delivered to the end hooks. + outcome.complete(result); + } + }; + executor.execute(work); + return result; } - private static void fireOnInvocationEnd( - boolean invocationStarted, - AtomicReference scopeFatal, - PluginRunner pluginRunner, - ExecutionManager executionManager, - String requestId, - String executionArn, - boolean isFirstInvocation, - InvocationStatus status, - Throwable error, - Object executionInput, - Object executionResult) { - if (pluginRunner.isEmpty() || !invocationStarted) { - return; - } - // Freeze the selected end snapshot here. A scope fatal already observed still escapes before dispatch; - // a later scope fatal escapes its owner instead of rewriting this caller outcome after finalization. - throwIfScopeFatal(scopeFatal); - // Finalization can run on the invocation caller rather than the handler worker. - try (var ignored = preserveEndMdc()) { - pluginRunner.onInvocationEnd(new InvocationEndInfo( - requestId, - executionArn, - isFirstInvocation, - executionManager.getExecutionOperation().startTimestamp(), - PluginInfoConverter.toOperationItemMap( - executionManager.getOperationsSnapshot(), executionManager.getInitialOperationIds()), - status, - error, - executionInput, - executionResult)); - } + private static SafeCloseable restoreMdcOnClose() { + var previous = MDC.getCopyOfContextMap(); + return () -> { + if (previous == null) MDC.clear(); + else MDC.setContextMap(previous); + }; } private static String handleLargePayload(ExecutionManager executionManager, String outputPayload) { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java b/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java index e2f8a46df..f91bc8d4d 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java @@ -18,7 +18,9 @@ public interface DurableExecutionPlugin { // ─── Invocation-level hooks ────────────────────────────────────────── /** - * Called at the start of each Lambda invocation. Use to set up per-invocation state (trace ID, invocation span). + * Called on the root handler thread at the start of each Lambda invocation. Use to set up per-invocation state + * (trace ID, invocation span) and thread-local context. The paired {@link #onInvocationEnd} runs on this same + * thread. * *

Check {@link InvocationInfo#isFirstInvocation()} to detect the first invocation of an execution (useful for * sampling decisions or execution-level span creation). @@ -26,10 +28,17 @@ public interface DurableExecutionPlugin { default void onInvocationStart(InvocationInfo info) {} /** - * Called at the end of each Lambda invocation. Use to flush spans/metrics before Lambda freezes. + * Called on the same root handler thread as {@link #onInvocationStart}, after the handler and its finally blocks + * have exited. Use to restore thread-local context and flush spans/metrics before Lambda freezes. * *

This hook is awaited — the SDK blocks until it returns. This is the only safe flush point before Lambda - * freezes the execution environment. + * freezes the execution environment. Suspension and termination also wait for the handler to unwind and this hook + * to finish; a blocked handler, finally block, or end hook delays the invocation response. No caller-thread + * fallback or cleanup timeout is applied. An earlier suspension/termination retains its selected outcome even if a + * later handler finally block returns or throws. + * + *

Start hooks run in registration order; end hooks run in reverse order to unwind nested thread-local scopes. If + * an end hook throws an Error that is not isolated, the remaining end hooks still run before it is rethrown. * *

Check {@link InvocationEndInfo#invocationStatus()} to detect if the execution reached a terminal state in this * invocation (useful for writing summary records or flushing final data). diff --git a/sdk/src/main/java/software/amazon/lambda/durable/plugin/HandlerScoped.java b/sdk/src/main/java/software/amazon/lambda/durable/plugin/HandlerScoped.java deleted file mode 100644 index c630341be..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/plugin/HandlerScoped.java +++ /dev/null @@ -1,30 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.plugin; - -import java.lang.annotation.ElementType; -import java.lang.annotation.Inherited; -import java.lang.annotation.Retention; -import java.lang.annotation.RetentionPolicy; -import java.lang.annotation.Target; -import java.util.function.Function; - -/** - * Explicitly opts a plugin into a root-handler scope using a JDK {@link Function} opener with a public no-argument - * constructor. The core passes the plugin instance to that opener on the handler thread; null means no scope. No method - * name is discovered on the plugin or added to the existing lifecycle interface. Inherited metadata uses the declaring - * plugin's explicit opener, so unrelated methods on old subclasses retain their behavior. Older cores ignore this - * optional metadata; plugin layer signatures require no new shared SDK type. - * - *

Scopes open and close on the handler thread, in reverse order on exit. Suspension/termination gives cleanup a - * bounded opportunity to unwind. The handler retains ownership even after a timeout. Ordinary cleanup preserves the - * winning outcome; observed VirtualMachineError/ThreadDeath from scope callbacks escapes the invocation caller. - * Optional waiting reserves five seconds per configured plugin plus one second for shutdown/response. This is best - * effort: existing finalizers and checkpoint draining can exceed it. - */ -@Inherited -@Retention(RetentionPolicy.RUNTIME) -@Target(ElementType.TYPE) -public @interface HandlerScoped { - Class> value(); -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java b/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java index 887b172d3..0ad9558af 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java @@ -2,29 +2,21 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.plugin; -import java.lang.reflect.InvocationTargetException; -import java.lang.reflect.UndeclaredThrowableException; -import java.util.ArrayDeque; import java.util.Collections; import java.util.HashMap; -import java.util.IdentityHashMap; import java.util.LinkedHashSet; import java.util.List; import java.util.Set; -import java.util.concurrent.CompletionException; -import java.util.concurrent.ExecutionException; import java.util.function.Consumer; -import java.util.function.Function; -import java.util.function.Supplier; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import software.amazon.lambda.durable.util.ExceptionHelper; /** * Composes multiple {@link DurableExecutionPlugin} instances into a single dispatcher. * - *

Event hooks call each plugin in order. Exceptions and nonfatal linkage failures are isolated; other errors, - * including fatal JVM failures, retain their existing propagation behavior. + *

Event hooks call each plugin in registration order, except invocation end, which unwinds in reverse order. + * Exceptions and nonfatal linkage failures are isolated; other errors propagate. Invocation end finishes the remaining + * cleanup hooks before propagating the first error. * *

{@code onInvocationEnd} is awaited (the SDK blocks until it returns) to allow plugins to flush data before Lambda * freezes. @@ -91,136 +83,39 @@ public List getPlugins() { /** Calls a void hook on all plugins, isolating exceptions and incompatible binary dependencies. */ private void run(Consumer hook) { - for (var plugin : plugins) { - try { - hook.accept(plugin); - } catch (Exception e) { - logger.warn("Plugin hook threw exception", e); - } catch (LinkageError e) { - logger.warn("Plugin hook could not link a dependency; check SDK/plugin dependency compatibility", e); - } - } - } - - /** Runs the root handler with optional plugin scopes, closing them on the same thread in reverse order. */ - public T runHandler(Supplier handler) { - return runHandler(handler, () -> {}); - } - - /** Runs the handler and notifies the invocation when a scope requires same-thread finalization. */ - public T runHandler(Supplier handler, Runnable onScopeOpened) { - return runHandler(handler, onScopeOpened, fatal -> {}); - } - - /** Reports fatal errors originating only in scope callbacks, before rethrowing on the owner thread. */ - @SuppressWarnings("removal") - public T runHandler(Supplier handler, Runnable onScopeOpened, Consumer onScopeFatal) { - var scopes = new ArrayDeque(); - try { - for (var plugin : plugins) { - try { - var scope = openHandlerScope(plugin); - if (scope != null) { - scopes.push(scope); - onScopeOpened.run(); - } - } catch (Throwable e) { - reportHandlerScopeFailure("Plugin handler scope threw exception", e, onScopeFatal); - } - } - return handler.get(); - } finally { - closeHandlerScopes(scopes, onScopeFatal); - } - } - - @SuppressWarnings("unchecked") - private static AutoCloseable openHandlerScope(DurableExecutionPlugin plugin) throws ReflectiveOperationException { - var metadata = plugin.getClass().getAnnotation(HandlerScoped.class); - if (metadata == null) return null; - var constructor = metadata.value().getConstructor(); - if (!constructor.canAccess(null) && !constructor.trySetAccessible()) { - throw new IllegalAccessException( - "Cannot access @HandlerScoped opener " + metadata.value().getName()); - } - try { - var opener = (Function) constructor.newInstance(); - return opener.apply(plugin); - } catch (InvocationTargetException failure) { - ExceptionHelper.sneakyThrow(failure.getCause()); - return null; - } + for (var plugin : plugins) runHook(plugin, hook); } - @SuppressWarnings("removal") - private static void closeHandlerScopes(ArrayDeque scopes, Consumer onScopeFatal) { - Error firstFatal = null; - while (!scopes.isEmpty()) { - try { - scopes.pop().close(); - } catch (Throwable e) { - try { - reportHandlerScopeFailure("Plugin handler scope cleanup threw exception", e, onScopeFatal); - } catch (VirtualMachineError | ThreadDeath fatal) { - if (firstFatal == null) firstFatal = fatal; - } - } - } - if (firstFatal != null) throw firstFatal; - } - - private static void reportHandlerScopeFailure(String message, Throwable failure, Consumer onScopeFatal) { - reportHandlerScopeFatal(failure, onScopeFatal); + private void runHook(DurableExecutionPlugin plugin, Consumer hook) { try { - // A plugin-controlled diagnostic must not be inspected again by the logger. - logger.warn("{} ({})", message, failure.getClass().getName()); - } catch (Throwable loggingFailure) { - reportHandlerScopeFatal(loggingFailure, onScopeFatal); - } - } - - private static void reportHandlerScopeFatal(Throwable failure, Consumer onScopeFatal) { - var fatal = findHandlerScopeFatal(failure); - if (fatal != null) { - // The caller must wake even when the fatal originated in an exception's cause accessor. - onScopeFatal.accept(fatal); - throw fatal; + hook.accept(plugin); + } catch (Exception e) { + logger.warn("Plugin hook threw exception", e); + } catch (LinkageError e) { + logger.warn("Plugin hook could not link a dependency; check SDK/plugin dependency compatibility", e); } } - @SuppressWarnings("removal") - private static Error findHandlerScopeFatal(Throwable failure) { - if (failure instanceof VirtualMachineError || failure instanceof ThreadDeath) return (Error) failure; - var seen = Collections.newSetFromMap(new IdentityHashMap()); - var cause = failure; - while (seen.add(cause)) { - if (cause instanceof VirtualMachineError || cause instanceof ThreadDeath) return (Error) cause; - if (!(cause instanceof CompletionException - || cause instanceof ExecutionException - || cause instanceof InvocationTargetException - || cause instanceof UndeclaredThrowableException)) return null; - try { - cause = cause.getCause(); - } catch (VirtualMachineError | ThreadDeath fatal) { - return fatal; - } catch (Throwable unreadableCause) { - return null; - } - if (cause == null) return null; - } - return null; - } - public void onInvocationStart(InvocationInfo info) { run(p -> p.onInvocationStart(info)); } /** - * Called at the end of each invocation. Awaited — the SDK blocks until all plugins return, allowing plugins to - * flush spans/metrics before Lambda freezes. + * Called in reverse registration order on the root handler thread after it unwinds. Awaited — the SDK blocks until + * all plugins return, including during suspension or termination, allowing plugins to flush spans/metrics before + * Lambda freezes. */ public void onInvocationEnd(InvocationEndInfo info) { - run(p -> p.onInvocationEnd(info)); + Error firstError = null; + for (var index = plugins.size() - 1; index >= 0; index--) { + try { + runHook(plugins.get(index), p -> p.onInvocationEnd(info)); + } catch (Error failure) { + // Finish unwinding earlier plugins' thread-local scopes before propagating an end-hook error. + if (firstError == null) firstError = failure; + } + } + if (firstError != null) throw firstError; } public void onOperationStart(OperationInfo info) { 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 deleted file mode 100644 index 8db315f81..000000000 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/HandlerScopeHandoffTest.java +++ /dev/null @@ -1,260 +0,0 @@ -// 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.util.List; -import java.util.Map; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.CompletionException; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.Executor; -import java.util.concurrent.Executors; -import java.util.concurrent.LinkedBlockingQueue; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.BiConsumer; -import java.util.function.BiFunction; -import java.util.function.Supplier; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.ValueSource; -import org.slf4j.MDC; -import software.amazon.awssdk.services.lambda.model.ErrorObject; -import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; -import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; -import software.amazon.lambda.durable.plugin.PluginRunner; - -class HandlerScopeHandoffTest { - @ParameterizedTest - @ValueSource(booleans = {false, true}) - void throwingCleanupCannotReplaceTheWinningControlFlow(boolean retry) throws Exception { - var original = control(retry); - var execution = CompletableFuture.failedFuture(original); - var handler = CompletableFuture.failedFuture(new IllegalStateException("finally failed")); - var result = handoff(execution, handler); - var thrown = assertThrows(CompletionException.class, result::join); - assertSame(original, thrown.getCause()); - } - - @ParameterizedTest - @ValueSource(booleans = {false, true}) - void nonUnwindingCleanupDoesNotWaitForever(boolean retry) throws Exception { - var original = control(retry); - var execution = CompletableFuture.failedFuture(original); - var handler = new CompletableFuture(); - var caller = Executors.newSingleThreadExecutor(task -> { - var thread = new Thread(task, "handoff-caller"); - thread.setDaemon(true); - return thread; - }); - try { - var response = caller.submit(() -> { - var result = handoff(execution, handler); - var thrown = assertThrows(CompletionException.class, result::join); - assertSame(original, thrown.getCause()); - return null; - }); - response.get(2, TimeUnit.SECONDS); - assertFalse(handler.isDone(), "timeout must not complete or cancel the owner task"); - } finally { - handler.complete("released"); - caller.shutdownNow(); - } - } - - @Test - void reservesFinalizationForEveryConfiguredPluginWithoutCancellingTheOwner() { - var context = mock(Context.class); - when(context.getRemainingTimeInMillis()).thenReturn(10_000); - var execution = CompletableFuture.failedFuture(control(false)); - var handler = new CompletableFuture(); - assertTimeoutPreemptively( - Duration.ofMillis(200), - () -> assertSame( - execution, - DurableExecutor.awaitHandlerScopes( - execution, handler, new AtomicBoolean(true), context, 2, new AtomicReference<>()))); - assertFalse(handler.isDone()); - } - - @Test - void noScopeDoesNotConsultTheRemainingTimeOrWaitForTheHandler() { - var context = mock(Context.class); - var execution = CompletableFuture.failedFuture(control(false)); - var handler = new CompletableFuture(); - assertSame( - execution, - DurableExecutor.awaitHandlerScopes( - execution, handler, new AtomicBoolean(false), context, 1, new AtomicReference<>())); - verifyNoInteractions(context); - assertFalse(handler.isDone()); - } - - @Test - void completedHandoffWaitsForTheRegisteredOwnerFinalizer() throws Exception { - var callbackRegistered = new CountDownLatch(1); - var callbacks = new LinkedBlockingQueue(); - var handler = new CompletableFuture() { - @Override - public CompletableFuture whenComplete(BiConsumer action) { - callbackRegistered.countDown(); - // A completed future may publish its result before its registered completion action executes. - return super.whenCompleteAsync(action, callbacks::add); - } - }; - var local = new ThreadLocal(); - var owner = new AtomicReference(); - var finalizerThread = new AtomicReference(); - var finalizerValue = new AtomicReference(); - var workers = Executors.newSingleThreadExecutor(); - var callers = Executors.newSingleThreadExecutor(); - var method = DurableExecutor.class.getDeclaredMethod( - "finalizeAfterHandlerScopes", - CompletableFuture.class, - CompletableFuture.class, - AtomicBoolean.class, - Context.class, - int.class, - AtomicReference.class, - BiFunction.class); - method.setAccessible(true); - BiFunction end = (value, failure) -> { - finalizerThread.set(Thread.currentThread()); - finalizerValue.set(local.get()); - local.remove(); - return "pending"; - }; - try { - var result = callers.submit(() -> method.invoke( - null, - CompletableFuture.failedFuture(control(false)), - handler, - new AtomicBoolean(true), - null, - 2, - new AtomicReference(), - end)); - assertTrue(callbackRegistered.await(3, TimeUnit.SECONDS)); - workers.submit(() -> { - owner.set(Thread.currentThread()); - local.set("invocation"); - handler.complete("handler finished"); - }) - .get(3, TimeUnit.SECONDS); - assertTrue(handler.isDone()); - assertThrows( - TimeoutException.class, - () -> result.get(100, TimeUnit.MILLISECONDS), - "The waiter must not steal a completed handoff's pending completion callback"); - var callback = callbacks.poll(3, TimeUnit.SECONDS); - assertNotNull(callback); - workers.submit(callback).get(3, TimeUnit.SECONDS); - assertEquals("pending", result.get(3, TimeUnit.SECONDS)); - assertSame(owner.get(), finalizerThread.get()); - assertEquals("invocation", finalizerValue.get()); - assertNull(workers.submit(local::get).get(3, TimeUnit.SECONDS), "The reused owner must be clean"); - } finally { - handler.complete("cleanup"); - Runnable callback; - while ((callback = callbacks.poll()) != null) - workers.submit(callback).get(3, TimeUnit.SECONDS); - callers.shutdownNow(); - workers.shutdownNow(); - assertTrue(callers.awaitTermination(3, TimeUnit.SECONDS)); - assertTrue(workers.awaitTermination(3, TimeUnit.SECONDS)); - } - } - - @Test - void mdcCaptureFailureSettlesAsyncHandlerFuture() throws Exception { - var failure = new IllegalStateException("MDC capture failed"); - var taskCalled = new AtomicBoolean(); - var worker = Executors.newSingleThreadExecutor(); - Executor failingCaptureWorker = task -> worker.execute(() -> { - // Static mocks are thread-scoped: fail only the real asynchronous worker's MDC capture. - try (var mdc = mockStatic(MDC.class, CALLS_REAL_METHODS)) { - mdc.when(MDC::getCopyOfContextMap).thenThrow(failure); - task.run(); - } - }); - var supply = DurableExecutor.class.getDeclaredMethod( - "supplyHandler", Supplier.class, Executor.class, PluginRunner.class, AtomicReference.class); - supply.setAccessible(true); - try { - Supplier task = () -> { - taskCalled.set(true); - return "unexpected"; - }; - var result = (CompletableFuture) supply.invoke( - null, - task, - failingCaptureWorker, - new PluginRunner(List.of(new DurableExecutionPlugin() {})), - new AtomicReference()); - var thrown = assertThrows(ExecutionException.class, () -> result.get(2, TimeUnit.SECONDS)); - assertSame( - failure, thrown.getCause(), "Initialization failure must settle the observation future unchanged"); - assertFalse(taskCalled.get(), "Failed MDC capture must not start user/plugin work"); - } finally { - worker.shutdownNow(); - assertTrue(worker.awaitTermination(3, TimeUnit.SECONDS)); - } - } - - @Test - void mdcRestorationFailureAfterCompletionStillEscapesItsOwner() throws Exception { - var failure = new IllegalStateException("MDC restoration failed"); - var ownerFailure = new CompletableFuture(); - var worker = Executors.newSingleThreadExecutor(); - Executor failingRestoreWorker = task -> worker.execute(() -> { - try (var mdc = mockStatic(MDC.class, CALLS_REAL_METHODS)) { - mdc.when(MDC::getCopyOfContextMap).thenReturn(Map.of("worker", "ambient")); - mdc.when(() -> MDC.setContextMap(Map.of("worker", "ambient"))).thenThrow(failure); - try { - task.run(); - } catch (Throwable thrown) { - ownerFailure.complete(thrown); - } - } - }); - var supply = DurableExecutor.class.getDeclaredMethod( - "supplyHandler", Supplier.class, Executor.class, PluginRunner.class, AtomicReference.class); - supply.setAccessible(true); - try { - Supplier task = () -> "completed"; - var result = (CompletableFuture) supply.invoke( - null, - task, - failingRestoreWorker, - new PluginRunner(List.of(new DurableExecutionPlugin() {})), - new AtomicReference()); - assertEquals("completed", result.get(2, TimeUnit.SECONDS)); - assertSame(failure, ownerFailure.get(2, TimeUnit.SECONDS)); - assertEquals("completed", result.join(), "Post-completion restoration must not rewrite the result"); - } finally { - worker.shutdownNow(); - assertTrue(worker.awaitTermination(3, TimeUnit.SECONDS)); - } - } - - private static CompletableFuture handoff( - CompletableFuture execution, CompletableFuture handler) { - return DurableExecutor.awaitHandlerScopes( - execution, handler, new AtomicBoolean(true), null, 1, new AtomicReference<>()); - } - - private static Throwable control(boolean retry) { - return retry - ? new UnrecoverableDurableExecutionException( - ErrorObject.builder().errorMessage("retry").build(), true) - : new SuspendExecutionException(); - } -} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationLifecycleTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationLifecycleTest.java new file mode 100644 index 000000000..c2fca5f4f --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationLifecycleTest.java @@ -0,0 +1,414 @@ +// 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 java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.slf4j.MDC; +import org.slf4j.helpers.BasicMDCAdapter; +import org.slf4j.spi.MDCAdapter; +import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState; +import software.amazon.awssdk.services.lambda.model.ErrorObject; +import software.amazon.awssdk.services.lambda.model.ExecutionDetails; +import software.amazon.awssdk.services.lambda.model.Operation; +import software.amazon.awssdk.services.lambda.model.OperationStatus; +import software.amazon.awssdk.services.lambda.model.OperationType; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.context.BaseContextImpl; +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.InvocationInfo; +import software.amazon.lambda.durable.plugin.InvocationStatus; +import software.amazon.lambda.durable.serde.JacksonSerDes; +import software.amazon.lambda.durable.serde.SerDes; + +class InvocationLifecycleTest { + private MDCAdapter originalMdcAdapter; + + @BeforeEach + void installMdcAdapter() throws ReflectiveOperationException { + originalMdcAdapter = MDC.getMDCAdapter(); + setMdcAdapter(new BasicMDCAdapter()); + } + + @AfterEach + void restoreMdcAdapter() throws ReflectiveOperationException { + MDC.clear(); + BaseContextImpl.setCurrentContext(null); + setMdcAdapter(originalMdcAdapter); + } + + private static void setMdcAdapter(MDCAdapter adapter) throws ReflectiveOperationException { + var field = MDC.class.getDeclaredField("MDC_ADAPTER"); + field.setAccessible(true); + field.set(null, adapter); + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void precompletedHandlerStillEndsOnItsOwner(boolean direct) throws Exception { + var workers = Executors.newSingleThreadExecutor(); + var plugin = new RecordingPlugin(); + var handlerThread = new AtomicReference(); + var executor = executor(task -> { + if (direct) task.run(); + else awaitTask(workers, task); + }); + try { + var output = DurableExecutor.execute( + input(), + null, + get(String.class), + (input, context) -> { + handlerThread.set(Thread.currentThread()); + return "done"; + }, + config(executor, plugin)); + assertEquals(ExecutionStatus.SUCCEEDED, output.status()); + assertEquals(1, plugin.starts.size()); + assertEquals(1, plugin.ends.size()); + assertSame(handlerThread.get(), plugin.startThreads.get(0)); + assertSame(handlerThread.get(), plugin.endThreads.get(0)); + assertEquals("invocation", plugin.endLocalValues.get(0)); + if (!direct) assertNotSame(Thread.currentThread(), handlerThread.get()); + } finally { + stop(workers); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void reusedWorkerRestoresAmbientMdcAfterBothHooks(boolean fails) throws Exception { + var workers = Executors.newSingleThreadExecutor(); + var plugin = new RecordingPlugin(); + var ambient = Map.of("worker", "ambient"); + var callerMdc = MDC.getCopyOfContextMap(); + MDC.put("caller", "untouched"); + try { + workers.submit(() -> MDC.setContextMap(ambient)).get(5, TimeUnit.SECONDS); + var config = config(workers, plugin); + for (int invocation = 0; invocation < 2; invocation++) { + var output = DurableExecutor.execute( + input(), + null, + get(String.class), + (input, context) -> { + if (fails) throw new IllegalStateException("handler failure"); + return "done"; + }, + config); + assertEquals(fails ? ExecutionStatus.FAILED : ExecutionStatus.SUCCEEDED, output.status()); + assertEquals(ambient, workers.submit(MDC::getCopyOfContextMap).get(5, TimeUnit.SECONDS)); + assertNull(workers.submit(plugin.local::get).get(5, TimeUnit.SECONDS)); + assertEquals("untouched", MDC.get("caller")); + } + assertEquals(2, plugin.ends.size()); + assertEquals(List.of("invocation", "invocation"), plugin.endLocalValues); + assertEquals(plugin.startThreads, plugin.endThreads); + assertEquals(List.of(ambient, ambient), plugin.startMdc); + } finally { + if (callerMdc == null) MDC.clear(); + else MDC.setContextMap(callerMdc); + stop(workers); + } + } + + @Test + void inputDeserializationFailurePairsHooksOnWorkerWithoutRunningHandler() throws Exception { + var workers = Executors.newSingleThreadExecutor(); + var plugin = new RecordingPlugin(); + var handlerCalled = new AtomicBoolean(); + var serDes = mock(SerDes.class); + when(serDes.deserialize(any(), any())).thenThrow(new IllegalArgumentException("invalid input")); + var config = config(workers, plugin).toBuilder().withSerDes(serDes).build(); + try { + var output = DurableExecutor.execute( + input(), + null, + get(String.class), + (input, context) -> { + handlerCalled.set(true); + return "unreachable"; + }, + config); + assertEquals(ExecutionStatus.FAILED, output.status()); + assertFalse(handlerCalled.get()); + assertEquals(1, plugin.starts.size()); + assertNull(plugin.starts.get(0).executionInput()); + assertEquals(1, plugin.ends.size()); + assertEquals(InvocationStatus.FAILED, plugin.ends.get(0).invocationStatus()); + assertSame(plugin.startThreads.get(0), plugin.endThreads.get(0)); + assertNotSame(Thread.currentThread(), plugin.endThreads.get(0)); + assertEquals(List.of("invocation"), plugin.endLocalValues); + } finally { + stop(workers); + } + } + + @Test + void mdcCaptureFailureBeforeStartDoesNotDispatchEitherHook() throws Exception { + var workers = Executors.newSingleThreadExecutor(); + var plugin = new RecordingPlugin(); + var handlerCalled = new AtomicBoolean(); + var executor = executor(task -> workers.execute(() -> { + try (var mdc = mockStatic(MDC.class, CALLS_REAL_METHODS)) { + mdc.when(MDC::getCopyOfContextMap).thenThrow(new IllegalStateException("capture failed")); + task.run(); + } + })); + try { + var output = DurableExecutor.execute( + input(), + null, + get(String.class), + (input, context) -> { + handlerCalled.set(true); + return "unreachable"; + }, + config(executor, plugin)); + assertEquals(ExecutionStatus.FAILED, output.status()); + assertFalse(handlerCalled.get()); + assertTrue(plugin.starts.isEmpty()); + assertTrue(plugin.ends.isEmpty()); + } finally { + stop(workers); + } + } + + @Test + void callerCannotReturnUntilInvocationEndCompletes() throws Exception { + var enteredEnd = new CountDownLatch(1); + var releaseEnd = new CountDownLatch(1); + var plugin = new RecordingPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + enteredEnd.countDown(); + await(releaseEnd); + super.onInvocationEnd(info); + } + }; + var callers = Executors.newSingleThreadExecutor(); + var workers = Executors.newSingleThreadExecutor(); + try { + var response = callers.submit(() -> DurableExecutor.execute( + input(), null, get(String.class), (input, context) -> "done", config(workers, plugin))); + assertTrue(enteredEnd.await(5, TimeUnit.SECONDS)); + assertThrows(TimeoutException.class, () -> response.get(700, TimeUnit.MILLISECONDS)); + releaseEnd.countDown(); + assertEquals( + ExecutionStatus.SUCCEEDED, response.get(5, TimeUnit.SECONDS).status()); + assertEquals(1, plugin.ends.size()); + assertEquals(plugin.startThreads, plugin.endThreads); + } finally { + releaseEnd.countDown(); + stop(callers); + stop(workers); + } + } + + @Test + void callerCannotReturnUntilAmbientMdcIsRestored() throws Exception { + var restoring = new CountDownLatch(1); + var releaseRestore = new CountDownLatch(1); + var restored = new AtomicBoolean(); + var ambient = Map.of("worker", "ambient"); + var workers = Executors.newSingleThreadExecutor(); + var callers = Executors.newSingleThreadExecutor(); + var plugin = new RecordingPlugin(); + var executor = executor(task -> workers.execute(() -> { + MDC.setContextMap(ambient); + try (var mdc = mockStatic(MDC.class, CALLS_REAL_METHODS)) { + mdc.when(() -> MDC.setContextMap(ambient)).thenAnswer(call -> { + restoring.countDown(); + await(releaseRestore); + call.callRealMethod(); + restored.set(true); + return null; + }); + task.run(); + } + })); + try { + var response = callers.submit(() -> DurableExecutor.execute( + input(), null, get(String.class), (input, context) -> "done", config(executor, plugin))); + assertTrue(restoring.await(5, TimeUnit.SECONDS)); + assertThrows(TimeoutException.class, () -> response.get(700, TimeUnit.MILLISECONDS)); + releaseRestore.countDown(); + assertEquals( + ExecutionStatus.SUCCEEDED, response.get(5, TimeUnit.SECONDS).status()); + assertTrue(restored.get()); + assertEquals(1, plugin.ends.size()); + } finally { + releaseRestore.countDown(); + stop(callers); + stop(workers); + } + } + + @Test + void serializationFailureStillEndsOnTheHandlerThread() throws Exception { + var failure = new IllegalStateException("output cannot serialize"); + var serDes = spy(new JacksonSerDes()); + when(serDes.serialize("done")).thenThrow(failure); + var plugin = new RecordingPlugin(); + var workers = Executors.newSingleThreadExecutor(); + var config = config(workers, plugin).toBuilder().withSerDes(serDes).build(); + try { + assertSame( + failure, + assertThrows( + IllegalStateException.class, + () -> DurableExecutor.execute( + input(), null, get(String.class), (input, context) -> "done", config))); + assertEquals(1, plugin.starts.size()); + assertEquals(1, plugin.ends.size()); + assertEquals(plugin.startThreads, plugin.endThreads); + assertEquals(InvocationStatus.FAILED, plugin.ends.get(0).invocationStatus()); + assertSame(failure, plugin.ends.get(0).executionError()); + assertNull(plugin.ends.get(0).executionResult()); + } finally { + stop(workers); + } + } + + @Test + void retryableLargeResultCheckpointFailureEndsWithOriginalRetryError() throws Exception { + var failure = new UnrecoverableDurableExecutionException( + ErrorObject.builder().errorMessage("retry delivery").build(), true); + var client = TestUtils.createMockClient(); + when(client.checkpoint(any(), any(), any())).thenThrow(failure); + var plugin = new RecordingPlugin(); + var workers = Executors.newSingleThreadExecutor(); + var config = config(workers, plugin).toBuilder() + .withDurableExecutionClient(client) + .build(); + try { + assertSame( + failure, + assertThrows( + UnrecoverableDurableExecutionException.class, + () -> DurableExecutor.execute( + input(), + null, + get(String.class), + (input, context) -> "x".repeat(7 * 1024 * 1024), + config))); + assertEquals(1, plugin.starts.size()); + assertEquals(1, plugin.ends.size()); + assertEquals(plugin.startThreads, plugin.endThreads); + assertEquals(InvocationStatus.RETRYING, plugin.ends.get(0).invocationStatus()); + assertSame(failure, plugin.ends.get(0).executionError()); + assertNull(plugin.ends.get(0).executionResult()); + } finally { + stop(workers); + } + } + + private static DurableConfig config(ExecutorService executor, DurableExecutionPlugin plugin) { + return DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withExecutorService(executor) + .withPlugins(plugin) + .build(); + } + + private static ExecutorService executor(Consumer dispatch) { + var executor = mock(ExecutorService.class); + doAnswer(call -> { + dispatch.accept(call.getArgument(0)); + return null; + }) + .when(executor) + .execute(any(Runnable.class)); + return executor; + } + + private static void awaitTask(ExecutorService workers, Runnable task) { + try { + workers.submit(task).get(5, TimeUnit.SECONDS); + } catch (Exception failure) { + throw new AssertionError(failure); + } + } + + private static void await(CountDownLatch latch) { + try { + if (!latch.await(5, TimeUnit.SECONDS)) throw new AssertionError("latch not released"); + } catch (InterruptedException failure) { + Thread.currentThread().interrupt(); + throw new AssertionError(failure); + } + } + + private static void stop(ExecutorService executor) throws InterruptedException { + executor.shutdownNow(); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } + + private static DurableExecutionInput input() { + var execution = Operation.builder() + .id("execution") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.parse("2026-08-15T00:00:00Z")) + .executionDetails( + ExecutionDetails.builder().inputPayload("\"input\"").build()) + .build(); + return new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/test/execution", + "token", + CheckpointUpdatedExecutionState.builder().operations(execution).build()); + } + + private static class RecordingPlugin implements DurableExecutionPlugin { + final ThreadLocal local = new ThreadLocal<>(); + final List starts = new ArrayList<>(); + final List ends = new ArrayList<>(); + final List startThreads = new ArrayList<>(); + final List endThreads = new ArrayList<>(); + final List endLocalValues = new ArrayList<>(); + final List> startMdc = new ArrayList<>(); + + @Override + public void onInvocationStart(InvocationInfo info) { + starts.add(info); + startThreads.add(Thread.currentThread()); + startMdc.add(MDC.getCopyOfContextMap()); + local.set("invocation"); + MDC.put("plugin", "invocation"); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + ends.add(info); + endThreads.add(Thread.currentThread()); + endLocalValues.add(local.get()); + local.remove(); + MDC.put("plugin", "ended"); + } + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/plugin/ConcreteHandlerScopeTest.java b/sdk/src/test/java/software/amazon/lambda/durable/plugin/ConcreteHandlerScopeTest.java deleted file mode 100644 index 47ceea8c4..000000000 --- a/sdk/src/test/java/software/amazon/lambda/durable/plugin/ConcreteHandlerScopeTest.java +++ /dev/null @@ -1,48 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.plugin; - -import static org.junit.jupiter.api.Assertions.*; - -import java.util.List; -import java.util.function.Function; -import org.junit.jupiter.api.Test; - -class ConcreteHandlerScopeTest { - @Test - void concreteScopeReturnTypeOpensAndClosesOnItsOwner() { - var plugin = new ConcretePlugin(); - var runner = new PluginRunner(List.of(plugin)); - assertEquals("result", runner.runHandler(() -> { - assertNotNull(plugin.scope); - assertFalse(plugin.scope.closed); - return "result"; - })); - assertTrue(plugin.scope.closed); - assertSame(Thread.currentThread(), plugin.scope.owner); - } - - @HandlerScoped(ConcreteOpener.class) - public static final class ConcretePlugin implements DurableExecutionPlugin { - ConcreteScope scope; - } - - public static final class ConcreteOpener implements Function { - @Override - public ConcreteScope apply(ConcretePlugin plugin) { - return plugin.scope = new ConcreteScope(); - } - } - - public static final class ConcreteScope implements AutoCloseable { - final Thread owner = Thread.currentThread(); - boolean closed; - - @Override - public void close() { - assertSame(owner, Thread.currentThread()); - assertFalse(closed); - closed = true; - } - } -} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeFailureInspectionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeFailureInspectionTest.java deleted file mode 100644 index fac872ad4..000000000 --- a/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeFailureInspectionTest.java +++ /dev/null @@ -1,212 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.plugin; - -import static org.junit.jupiter.api.Assertions.*; - -import java.lang.reflect.InvocationTargetException; -import java.lang.reflect.UndeclaredThrowableException; -import java.util.ArrayList; -import java.util.List; -import java.util.concurrent.CompletionException; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Function; -import java.util.function.Supplier; -import java.util.stream.Stream; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.Arguments; -import org.junit.jupiter.params.provider.MethodSource; -import software.amazon.lambda.durable.util.ExceptionHelper; - -class HandlerScopeFailureInspectionTest { - static Stream ordinaryFailures() { - return Stream.of("throwing", "cycle", "two-cycle", "changing", "null", "application-cause") - .flatMap(kind -> Stream.of(false, true) - .flatMap(open -> Stream.of(false, true).map(bodyFails -> Arguments.of(kind, open, bodyFails)))); - } - - @ParameterizedTest - @MethodSource("ordinaryFailures") - void malformedScopeFailuresKeepTheBodyOutcomeAndReverseCleanup(String kind, boolean failOpen, boolean bodyFails) { - var calls = new ArrayList(); - var reads = new AtomicInteger(); - var failure = ordinary(kind, reads); - var owner = Thread.currentThread(); - var runner = new PluginRunner(List.of( - healthy("first", calls, owner), broken(failOpen, failure, calls), healthy("last", calls, owner))); - var original = new IllegalStateException("body failure"); - var reported = new AtomicReference(); - var bodyCalls = new AtomicInteger(); - var body = (Supplier) () -> { - bodyCalls.incrementAndGet(); - calls.add("body"); - if (bodyFails) throw original; - return "done"; - }; - if (bodyFails) - assertSame( - original, - assertThrows(IllegalStateException.class, () -> runner.runHandler(body, () -> {}, reported::set))); - else assertEquals("done", runner.runHandler(body, () -> {}, reported::set)); - assertEquals(1, bodyCalls.get()); - assertNull(reported.get()); - assertEquals( - failOpen - ? List.of("open-first", "open-bad", "open-last", "body", "close-last", "close-first") - : List.of( - "open-first", - "open-bad", - "open-last", - "body", - "close-last", - "close-bad", - "close-first"), - calls); - assertEquals(kind.equals("application-cause") ? 0 : kind.equals("two-cycle") ? 2 : 1, reads.get()); - } - - static Stream fatalFailures() { - return Stream.of("direct", "completion", "execution", "reflection", "proxy", "accessor") - .flatMap(kind -> Stream.of(false, true) - .flatMap(open -> - Stream.of(false, true).map(threadDeath -> Arguments.of(kind, open, threadDeath)))); - } - - @SuppressWarnings("removal") - @ParameterizedTest - @MethodSource("fatalFailures") - void fatalInspectionReportsBeforeEarlierCleanupAndKeepsIdentity( - String kind, boolean failOpen, boolean threadDeath) { - Error fatal = threadDeath ? new ThreadDeath() : new InternalError("scope fatal"); - var reported = new AtomicReference(); - var closed = new AtomicInteger(); - var bodyCalls = new AtomicInteger(); - var owner = Thread.currentThread(); - var earlier = new Scoped(() -> () -> { - assertSame(owner, Thread.currentThread()); - assertSame(fatal, reported.get(), "notify the invocation before potentially blocking earlier cleanup"); - closed.incrementAndGet(); - }); - var runner = new PluginRunner(List.of(earlier, broken(failOpen, fatalWrapper(kind, fatal), new ArrayList<>()))); - assertSame( - fatal, - assertThrows( - Error.class, - () -> runner.runHandler( - () -> { - bodyCalls.incrementAndGet(); - return "done"; - }, - () -> {}, - reported::set))); - assertSame(fatal, reported.get()); - assertEquals(1, closed.get()); - assertEquals(failOpen ? 0 : 1, bodyCalls.get()); - } - - private static Throwable ordinary(String kind, AtomicInteger reads) { - if (kind.equals("application-cause")) - return new IllegalStateException("application") { - @Override - public synchronized Throwable getCause() { - reads.incrementAndGet(); - return new InternalError("not a transport wrapper"); - } - }; - var first = new AtomicReference(); - var second = new ExecutionException("second", null) { - @Override - public synchronized Throwable getCause() { - guardNegativeTraversal(reads); - return first.get(); - } - }; - var failure = new CompletionException("scope", null) { - @Override - public synchronized Throwable getCause() { - guardNegativeTraversal(reads); - return switch (kind) { - case "throwing" -> throw new IllegalArgumentException("unreadable diagnostic"); - case "cycle" -> this; - case "two-cycle" -> second; - case "changing" -> - reads.get() == 1 - ? new IllegalStateException("first cause") - : new InternalError("changed cause"); - case "null" -> null; - default -> throw new AssertionError(kind); - }; - } - }; - first.set(failure); - return failure; - } - - private static void guardNegativeTraversal(AtomicInteger reads) { - // A broken classifier must fail the negative control rather than leave a spinning test thread. - if (reads.incrementAndGet() > 8) throw new AssertionError("test safety bound: repeated cause traversal"); - } - - private static Throwable fatalWrapper(String kind, Error fatal) { - return switch (kind) { - case "direct" -> fatal; - case "completion" -> new CompletionException(fatal); - case "execution" -> new ExecutionException(fatal); - case "reflection" -> new InvocationTargetException(fatal); - case "proxy" -> new UndeclaredThrowableException(fatal); - case "accessor" -> - new CompletionException("fatal accessor", null) { - @Override - public synchronized Throwable getCause() { - throw fatal; - } - }; - default -> throw new AssertionError(kind); - }; - } - - private static Scoped healthy(String name, List calls, Thread owner) { - return new Scoped(() -> { - assertSame(owner, Thread.currentThread()); - calls.add("open-" + name); - return () -> { - assertSame(owner, Thread.currentThread()); - calls.add("close-" + name); - }; - }); - } - - private static Scoped broken(boolean failOpen, Throwable failure, List calls) { - return new Scoped(() -> { - calls.add("open-bad"); - if (failOpen) ExceptionHelper.sneakyThrow(failure); - return () -> { - calls.add("close-bad"); - ExceptionHelper.sneakyThrow(failure); - }; - }); - } - - @FunctionalInterface - private interface Open { - AutoCloseable open(); - } - - @HandlerScoped(Opener.class) - public static final class Scoped implements DurableExecutionPlugin { - private final Open open; - - private Scoped(Open open) { - this.open = open; - } - } - - public static final class Opener implements Function { - @Override - public AutoCloseable apply(Scoped plugin) { - return plugin.open.open(); - } - } -} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeTest.java b/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeTest.java deleted file mode 100644 index 6b19a126c..000000000 --- a/sdk/src/test/java/software/amazon/lambda/durable/plugin/HandlerScopeTest.java +++ /dev/null @@ -1,255 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.plugin; - -import static org.junit.jupiter.api.Assertions.*; - -import java.util.ArrayList; -import java.util.List; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Function; -import java.util.function.Supplier; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.CsvSource; -import org.junit.jupiter.params.provider.ValueSource; -import software.amazon.lambda.durable.execution.SuspendExecutionException; -import software.amazon.lambda.durable.util.ExceptionHelper; - -class HandlerScopeTest { - @ParameterizedTest - @ValueSource(strings = {"success", "failure", "suspension"}) - void scopesCloseInReverseOrderOnTheHandlerThread(String outcome) { - var calls = new ArrayList(); - var owner = Thread.currentThread(); - var runner = new PluginRunner(List.of(scope("a", calls, owner), scope("b", calls, owner))); - Throwable error = - outcome.equals("suspension") ? new SuspendExecutionException() : new IllegalStateException("user"); - if (outcome.equals("success")) { - assertEquals("ok", runner.runHandler(() -> { - calls.add("handler"); - return "ok"; - })); - } else { - assertSame( - error, - assertThrows( - Throwable.class, - () -> runner.runHandler(() -> { - calls.add("handler"); - ExceptionHelper.sneakyThrow(error); - return null; - }))); - } - assertEquals(List.of("open-a", "open-b", "handler", "close-b", "close-a"), calls); - } - - @Test - void ordinarySetupAndCleanupFailuresDoNotReplaceTheHandlerResult() { - var calls = new ArrayList(); - var brokenSetup = new ScopedPlugin() { - @Override - public AutoCloseable openHandlerScope() { - throw new IllegalStateException("setup"); - } - }; - var brokenClose = new ScopedPlugin() { - @Override - public AutoCloseable openHandlerScope() { - return () -> { - throw new IllegalStateException("cleanup"); - }; - } - }; - var runner = new PluginRunner(List.of(scope("first", calls, Thread.currentThread()), brokenSetup, brokenClose)); - assertEquals("ok", runner.runHandler(() -> "ok")); - assertEquals(List.of("open-first", "close-first"), calls); - } - - @ParameterizedTest - @CsvSource({ - "false,linkage", - "true,linkage", - "false,assert-open", - "true,assert-open", - "false,assert-close", - "true,assert-close" - }) - void nonFatalErrorsPreserveTheBodyOutcomeAndAllEarlierScopes(boolean bodyFails, String failureSite) { - var active = new ThreadLocal(); - var calls = new ArrayList(); - var healthy = new ScopedPlugin() { - @Override - public AutoCloseable openHandlerScope() { - calls.add("open-healthy"); - active.set("healthy"); - return () -> { - calls.add("close-healthy"); - active.remove(); - }; - } - }; - var brokenOpen = new ScopedPlugin() { - @Override - public AutoCloseable openHandlerScope() { - calls.add("open-broken"); - if (failureSite.equals("assert-open")) throw new AssertionError("optional setup assertion"); - throw new NoSuchMethodError("optional API missing"); - } - }; - var brokenClose = new ScopedPlugin() { - @Override - public AutoCloseable openHandlerScope() { - calls.add("open-last"); - return () -> { - calls.add("close-last"); - if (failureSite.equals("assert-close")) throw new AssertionError("optional cleanup assertion"); - throw new NoClassDefFoundError("optional class missing"); - }; - } - }; - var runner = new PluginRunner(List.of(healthy, brokenOpen, brokenClose)); - var bodyFailure = new IllegalStateException("original body error"); - Supplier body = () -> { - calls.add("body"); - assertEquals("healthy", active.get()); - if (bodyFails) throw bodyFailure; - return "ok"; - }; - try { - if (bodyFails) - assertSame(bodyFailure, assertThrows(IllegalStateException.class, () -> runner.runHandler(body))); - else assertEquals("ok", runner.runHandler(body)); - assertNull(active.get(), "a later optional-scope failure must not prevent earlier context cleanup"); - assertEquals( - List.of("open-healthy", "open-broken", "open-last", "body", "close-last", "close-healthy"), calls); - } finally { - active.remove(); - } - } - - @SuppressWarnings("removal") - @ParameterizedTest - @CsvSource({"true,false", "false,false", "true,true", "false,true"}) - void fatalScopeFailuresStillPropagate(boolean duringOpen, boolean wrapped) { - for (Error fatal : List.of(new InternalError("fatal VM failure"), new ThreadDeath())) { - var plugin = new ScopedPlugin() { - @Override - public AutoCloseable openHandlerScope() { - if (duringOpen) { - if (wrapped) - CompletableFuture.failedFuture(new ExecutionException(fatal)) - .join(); - throw fatal; - } - return () -> { - if (wrapped) - CompletableFuture.failedFuture(new ExecutionException(fatal)) - .join(); - throw fatal; - }; - } - }; - var reported = new AtomicReference(); - var calls = new ArrayList(); - var active = new ThreadLocal(); - var owner = Thread.currentThread(); - var healthy = new ScopedPlugin() { - @Override - public AutoCloseable openHandlerScope() { - active.set("healthy"); - calls.add("open-healthy"); - return () -> { - assertSame(owner, Thread.currentThread()); - active.remove(); - calls.add("close-healthy"); - }; - } - }; - var runner = new PluginRunner(List.of(healthy, plugin)); - assertSame(fatal, assertThrows(Error.class, () -> runner.runHandler(() -> "ok", () -> {}, reported::set))); - assertSame(fatal, reported.get(), "only a scope callback reports this fatal"); - assertNull(active.get(), "earlier context must not leak into reuse of the owner thread"); - assertEquals(List.of("open-healthy", "close-healthy"), calls); - } - } - - @Test - void openerConstructionLinkageFailurePreservesBodyAndEarlierCleanup() { - var calls = new ArrayList(); - var runner = - new PluginRunner(List.of(scope("healthy", calls, Thread.currentThread()), new BrokenOpenerPlugin())); - assertEquals("body", runner.runHandler(() -> "body")); - assertEquals(List.of("open-healthy", "close-healthy"), calls); - } - - @Test - void openerConstructionFatalRetainsIdentityAndEarlierCleanup() { - var calls = new ArrayList(); - var reported = new AtomicReference(); - var runner = - new PluginRunner(List.of(scope("healthy", calls, Thread.currentThread()), new FatalOpenerPlugin())); - assertSame( - FatalOpener.FAILURE, - assertThrows( - InternalError.class, - () -> runner.runHandler(() -> fail("handler must not run"), () -> {}, reported::set))); - assertSame(FatalOpener.FAILURE, reported.get()); - assertEquals(List.of("open-healthy", "close-healthy"), calls); - } - - @HandlerScoped(BrokenOpener.class) - private static class BrokenOpenerPlugin implements DurableExecutionPlugin {} - - public static class BrokenOpener implements Function { - public BrokenOpener() { - throw new NoSuchMethodError("opener dependency"); - } - - public AutoCloseable apply(BrokenOpenerPlugin plugin) { - return null; - } - } - - @HandlerScoped(FatalOpener.class) - private static class FatalOpenerPlugin implements DurableExecutionPlugin {} - - public static class FatalOpener implements Function { - static final InternalError FAILURE = new InternalError("opener fatal"); - - public FatalOpener() { - throw FAILURE; - } - - public AutoCloseable apply(FatalOpenerPlugin plugin) { - return null; - } - } - - private static DurableExecutionPlugin scope(String name, List calls, Thread owner) { - return new ScopedPlugin() { - @Override - public AutoCloseable openHandlerScope() { - assertSame(owner, Thread.currentThread()); - calls.add("open-" + name); - return () -> { - assertSame(owner, Thread.currentThread()); - calls.add("close-" + name); - }; - } - }; - } - - @HandlerScoped(ScopedPlugin.Opener.class) - private abstract static class ScopedPlugin implements DurableExecutionPlugin { - public abstract AutoCloseable openHandlerScope(); - - public static class Opener implements Function { - public AutoCloseable apply(ScopedPlugin plugin) { - return plugin.openHandlerScope(); - } - } - } -} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/plugin/LegacyHandlerScopeCompatibilityTest.java b/sdk/src/test/java/software/amazon/lambda/durable/plugin/LegacyHandlerScopeCompatibilityTest.java deleted file mode 100644 index 2be3713cb..000000000 --- a/sdk/src/test/java/software/amazon/lambda/durable/plugin/LegacyHandlerScopeCompatibilityTest.java +++ /dev/null @@ -1,84 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.plugin; - -import static org.junit.jupiter.api.Assertions.*; - -import java.net.URL; -import java.net.URLClassLoader; -import java.nio.file.Files; -import java.nio.file.Path; -import java.util.List; -import javax.tools.ToolProvider; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.io.TempDir; - -class LegacyHandlerScopeCompatibilityTest { - @TempDir - Path directory; - - @Test - void precompiledUnrelatedDefaultMethodStillLinksAndConfigures() throws Exception { - var baseline = directory.resolve("baseline"); - var classes = directory.resolve("classes"); - Files.createDirectories(baseline); - Files.createDirectories(classes); - var oldApi = directory.resolve("DurableExecutionPlugin.java"); - Files.writeString(oldApi, """ - package software.amazon.lambda.durable.plugin; - public interface DurableExecutionPlugin {} - """); - var compiler = ToolProvider.getSystemJavaCompiler(); - assertEquals( - 0, compiler.run(null, null, null, "--release", "17", "-d", baseline.toString(), oldApi.toString())); - var source = directory.resolve("LegacyPlugin.java"); - Files.writeString(source, """ - import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; - interface ExistingCapability { - default AutoCloseable openHandlerScope() { LegacyPlugin.opened++; return () -> {}; } - } - public class LegacyPlugin implements DurableExecutionPlugin, ExistingCapability { - public static int opened; - public String originalCall() throws Exception { - try (var scope = openHandlerScope()) { return "legacy application value"; } - } - } - """); - assertEquals( - 0, - compiler.run( - null, - null, - null, - "--release", - "17", - "-cp", - baseline.toString(), - "-d", - classes.toString(), - source.toString())); - try (var loader = new URLClassLoader( - new URL[] {classes.toUri().toURL()}, getClass().getClassLoader())) { - var type = loader.loadClass("LegacyPlugin"); - var plugin = (DurableExecutionPlugin) type.getConstructor().newInstance(); - assertEquals("body", new PluginRunner(List.of(plugin)).runHandler(() -> "body")); - assertEquals(0, type.getField("opened").get(null), "unannotated application capability is not SDK opt-in"); - assertEquals( - "legacy application value", type.getMethod("originalCall").invoke(plugin)); - } - // Source compatibility matters too: the same downstream source compiles against this core. - assertEquals( - 0, - compiler.run( - null, - null, - null, - "--release", - "17", - "-cp", - System.getProperty("java.class.path"), - "-d", - classes.toString(), - source.toString())); - } -} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginLinkageErrorTest.java b/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginLinkageErrorTest.java index 276e60738..5611215fb 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginLinkageErrorTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginLinkageErrorTest.java @@ -41,7 +41,10 @@ void otherErrorsRetainTheirExistingPropagation( plugin(healthyCalls::incrementAndGet))); assertSame(failure, assertThrows(Error.class, () -> dispatch.accept(runner))); - assertEquals(0, healthyCalls.get(), "Fatal and unrelated errors must not be blanket-caught"); + assertEquals( + hook.equals("invocation end") ? 1 : 0, + healthyCalls.get(), + "Invocation End must finish cleanup before propagating errors; other hooks stop immediately"); } static Stream linkageFailures() { diff --git a/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginRunnerTest.java b/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginRunnerTest.java index 82f10f927..e8953a217 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginRunnerTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginRunnerTest.java @@ -155,7 +155,7 @@ void fireAndForget_callsAllHookTypes() { // ─── Awaited hooks ─────────────────────────────────────────────────── @Test - void awaitedHooks_callAllPlugins() { + void invocationEnd_unwindsPluginsInReverseRegistrationOrder() { var calls = new ArrayList(); var plugin1 = new TestPlugin("p1", calls); var plugin2 = new TestPlugin("p2", calls); @@ -163,7 +163,33 @@ void awaitedHooks_callAllPlugins() { runner.onInvocationEnd(invocationEndInfo()); - assertEquals(List.of("p1:onInvocationEnd", "p2:onInvocationEnd"), calls); + assertEquals(List.of("p2:onInvocationEnd", "p1:onInvocationEnd"), calls); + } + + @SuppressWarnings("removal") + @Test + void invocationEnd_runsRemainingCleanupThenRethrowsTheFirstError() { + for (var firstFailure : + List.of(new InternalError("fatal"), new AssertionError("assertion"), new ThreadDeath())) { + var calls = new ArrayList(); + var inner = new TestPlugin("inner", calls) { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + super.onInvocationEnd(info); + throw firstFailure; + } + }; + var middle = new TestPlugin("middle", calls) { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + super.onInvocationEnd(info); + throw new AssertionError("later cleanup failure"); + } + }; + var runner = new PluginRunner(List.of(new TestPlugin("outer", calls), middle, inner)); + assertSame(firstFailure, assertThrows(Error.class, () -> runner.onInvocationEnd(invocationEndInfo()))); + assertEquals(List.of("inner:onInvocationEnd", "middle:onInvocationEnd", "outer:onInvocationEnd"), calls); + } } @Test @@ -171,7 +197,7 @@ void awaitedHooks_swallowExceptions_butCallRemainingPlugins() { var calls = new ArrayList(); var throwingPlugin = new ThrowingPlugin(); var normalPlugin = new TestPlugin("p2", calls); - var runner = new PluginRunner(List.of(throwingPlugin, normalPlugin)); + var runner = new PluginRunner(List.of(normalPlugin, throwingPlugin)); assertDoesNotThrow(() -> runner.onInvocationEnd(invocationEndInfo())); assertEquals(List.of("p2:onInvocationEnd"), calls); From c1a33104784831a87e5d1fd9cf4b37417db20919 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 7 Oct 2026 22:55:22 +0000 Subject: [PATCH 21/48] ci: pin OTel queue workflow to published squash commit --- .github/workflows/otel-conformance-tests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/otel-conformance-tests.yml b/.github/workflows/otel-conformance-tests.yml index 452db089d..cfb82bbe7 100644 --- a/.github/workflows/otel-conformance-tests.yml +++ b/.github/workflows/otel-conformance-tests.yml @@ -68,7 +68,7 @@ jobs: actions: write contents: read id-token: write - uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@a66037abbbfa55fde97f714e30f0bc262edefd63 + uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@f18bd0b5f28c5c90e288d0fb8bca08a849b51863 with: runs_on: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.repo.full_name != github.repository && 'ubuntu-latest' || format('codebuild-github-actions-runner-{0}-{1}', github.run_id, github.run_attempt) }} language: java From 8c7d121baf431c6e6ce9334d545ba4add75ae26e Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 7 Oct 2026 23:19:38 +0000 Subject: [PATCH 22/48] fix(plugin): enforce invocation lifecycle boundaries --- .../scripts/verify_otel_api_compatibility.py | 7 +- otel-plugin/README.md | 15 +- .../durable/otel/ExecutionOtelPlugin.java | 2 + .../durable/otel/InvocationOtelPlugin.java | 2 + .../durable/otel/OtelPluginSupport.java | 12 + .../compatibility/b1/InstalledApiProbe.java | 42 +- .../otel/InvocationOutcomeBoundaryTest.java | 398 ++++++++++++++++++ .../durable/execution/DurableExecutor.java | 77 +++- .../plugin/DurableExecutionPlugin.java | 4 +- .../execution/InvocationLifecycleTest.java | 2 +- 10 files changed, 538 insertions(+), 23 deletions(-) create mode 100644 otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOutcomeBoundaryTest.java diff --git a/.github/scripts/verify_otel_api_compatibility.py b/.github/scripts/verify_otel_api_compatibility.py index b0e88c7aa..2ff92fd69 100644 --- a/.github/scripts/verify_otel_api_compatibility.py +++ b/.github/scripts/verify_otel_api_compatibility.py @@ -142,11 +142,14 @@ def run_matrix(args: argparse.Namespace) -> int: for view in ("otel-invocation", "otel-execution"): name = f"{label}-api{version}-{view}" negative = label == "old-old" and version == "1.49.0" + incompatible_core = label == "old-new" case: dict[str, object] = {"name": name, "expected_negative_control": negative, + "expected_core_rejection": incompatible_core, "api": jar_facts(api), "context": jar_facts(context)} command = [args.java, "-cp", os.pathsep.join(map(str, [classes, core, plugin, *dependencies])), PROBE, str(core), str(plugin), str(api), str(context), view, - str(negative).lower(), str(version == "1.66.0").lower()] + str(negative).lower(), str(version == "1.66.0").lower(), + str(incompatible_core).lower()] try: log = output / f"{name}.log" execute(command, log, env=probe_environment(view), timeout=90) @@ -155,6 +158,8 @@ def run_matrix(args: argparse.Namespace) -> int: raise RuntimeError("Probe did not report successful completion") if negative and "NEGATIVE_CONTROL_REPRODUCED" not in contents: raise RuntimeError("Released negative control did not reproduce the reported failure") + if incompatible_core and "CORE_LIFECYCLE_REJECTION_CONFIRMED" not in contents: + raise RuntimeError("Old core must reject the new plugin before invocation startup") case["passed"] = True except RuntimeError as error: failures += 1 diff --git a/otel-plugin/README.md b/otel-plugin/README.md index 6cebe01d0..8cb340d4b 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -27,10 +27,17 @@ can run; a blocked handler cleanup therefore also blocks the invocation response when span finalization or flushing fails. Ordinary exceptions and nonfatal linkage errors retain the existing plugin-hook isolation behavior. Invocation-end hooks run in reverse registration order so nested scopes unwind correctly; remaining end hooks run even when another hook raises an Error, then the first Error propagates. - -This plugin version requires a core version that provides that same-thread invocation-hook contract. Upgrade the -core together with the plugin layer. Older cores may call `onInvocationEnd` on another thread and are not supported -with this plugin version. +Invocations without plugins also wait for handler cleanup before returning the selected suspension or retry outcome. +Output-delivery failures that escape for a Lambda retry report `RETRYING`, so they do not prematurely end the Workflow +span. An ordinary MDC-restoration failure on an inline executor does not replace an already selected outcome. + +This plugin version requires the core's `DurableExecutor.supportsSameThreadInvocationHooks()` capability, introduced +in the 2.2.2 lifecycle contract (currently `2.2.2-SNAPSHOT`). Upgrade the core together with the plugin layer. Every +plugin constructor checks this capability before building a tracer provider or activating context. A core without +it, including released 2.2.1, is rejected with an explicit configuration error. Older cores may call +`onInvocationEnd` on another thread and are not supported with this plugin version. Existing plugin binaries remain +usable with the updated core; their invocation-end hooks now follow the same-thread, reverse-registration-order +contract described above. ## Installation diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java index d0c46ea1a..0702d84b4 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java @@ -179,6 +179,7 @@ public ExecutionOtelPlugin() { * @param config the plugin configuration */ public ExecutionOtelPlugin(SdkTracerProviderBuilder tracerProviderBuilder, OtelPluginConfig config) { + OtelPluginSupport.requireSameThreadInvocationHooks(); this.idGenerator = DeterministicIdGenerator.installOn(tracerProviderBuilder); // Wrap the configured sampler so durable spans use the execution's single precomputed decision. DurableSampler.installOn(tracerProviderBuilder); @@ -200,6 +201,7 @@ public ExecutionOtelPlugin(SdkTracerProviderBuilder tracerProviderBuilder, OtelP * @param config the plugin configuration */ public ExecutionOtelPlugin(OtelPluginConfig config) { + OtelPluginSupport.requireSameThreadInvocationHooks(); this.contextExtractor = config.contextExtractor(); this.enableMdc = config.enableMdc(); this.workflowSpanName = config.workflowSpanName(); diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java index 7e2a5acc5..0498e7727 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java @@ -172,6 +172,7 @@ public InvocationOtelPlugin() { * @param config the plugin configuration */ public InvocationOtelPlugin(SdkTracerProviderBuilder tracerProviderBuilder, OtelPluginConfig config) { + OtelPluginSupport.requireSameThreadInvocationHooks(); this.idGenerator = DeterministicIdGenerator.installOn(tracerProviderBuilder); // Wrap the configured sampler so durable spans use the execution's single precomputed decision. DurableSampler.installOn(tracerProviderBuilder); @@ -193,6 +194,7 @@ public InvocationOtelPlugin(SdkTracerProviderBuilder tracerProviderBuilder, Otel * @param config the plugin configuration */ public InvocationOtelPlugin(OtelPluginConfig config) { + OtelPluginSupport.requireSameThreadInvocationHooks(); this.contextExtractor = config.contextExtractor(); this.enableMdc = config.enableMdc(); this.workflowSpanName = config.workflowSpanName(); diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java index dea21b447..cfad36746 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java @@ -16,6 +16,7 @@ import java.util.Collections; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.lambda.durable.execution.DurableExecutor; /** Shared utilities for OTel plugin default constructor support (ADOT Java agent SPI path). */ final class OtelPluginSupport { @@ -24,6 +25,17 @@ final class OtelPluginSupport { private OtelPluginSupport() {} + /** Rejects an older core before a plugin builds a provider or activates any thread-local context. */ + static void requireSameThreadInvocationHooks() { + var message = "This OpenTelemetry plugin requires same-thread invocation hooks from the Durable Execution" + + " core (2.2.2 lifecycle capability). Upgrade the core together with the plugin layer."; + try { + if (!DurableExecutor.supportsSameThreadInvocationHooks()) throw new IllegalStateException(message); + } catch (NoSuchMethodError missingCapability) { + throw new IllegalStateException(message, missingCapability); + } + } + /** Creates a new DeterministicIdGenerator for the application-side state bridge. */ static DeterministicIdGenerator createDefaultIdGenerator() { return new DeterministicIdGenerator(); diff --git a/otel-plugin/src/test/compatibility/b1/InstalledApiProbe.java b/otel-plugin/src/test/compatibility/b1/InstalledApiProbe.java index 9d31b1d58..122fca380 100644 --- a/otel-plugin/src/test/compatibility/b1/InstalledApiProbe.java +++ b/otel-plugin/src/test/compatibility/b1/InstalledApiProbe.java @@ -14,6 +14,7 @@ import io.opentelemetry.sdk.trace.SdkTracerProvider; import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; import java.nio.file.Path; +import java.util.concurrent.Callable; import java.time.Duration; import java.util.List; import java.util.ServiceLoader; @@ -43,7 +44,8 @@ public static void main(String[] args) throws Exception { logs.start(); root.addAppender(logs); try { - exercise(view, Boolean.parseBoolean(args[5]), Boolean.parseBoolean(args[6]), logs); + if (Boolean.parseBoolean(args[7])) assertOldCoreRejected(view); + else exercise(view, Boolean.parseBoolean(args[5]), Boolean.parseBoolean(args[6]), logs); System.out.println("COMPAT_PASS " + view + " negative=" + args[5] + " api=" + args[2] + " core=" + args[0] + " plugin=" + args[1]); } finally { @@ -66,6 +68,44 @@ private static void verifyArtifacts(String[] args) throws Exception { checkSource(provider.getPluginType(), Path.of(args[1])); } + private static void assertOldCoreRejected(String view) throws Exception { + var provider = ServiceLoader.load(DurableExecutionPluginProvider.class).stream() + .map(ServiceLoader.Provider::get).filter(value -> value.getName().equals(view)) + .findFirst().orElseThrow(); + expectCoreRejection(provider::createPlugin); + var constructors = provider.getPluginType().getConstructors(); + check(constructors.length == 4, "exercise every published plugin constructor"); + for (var constructor : constructors) { + var types = constructor.getParameterTypes(); + var parameters = new Object[types.length]; + for (int i = 0; i < types.length; i++) { + parameters[i] = types[i].getName().endsWith("OtelPluginConfig") + ? types[i].getMethod("defaults").invoke(null) : SdkTracerProvider.builder(); + } + expectCoreRejection(() -> constructor.newInstance(parameters)); + } + var handlerCalls = new AtomicInteger(); + expectCoreRejection(() -> createRunner(handlerCalls, new AtomicInteger())); + check(handlerCalls.get() == 0 && HEALTHY_STARTS.get() == 0 && HEALTHY_ENDS.get() == 0, + "configuration rejection must precede handler execution and lifecycle hooks"); + System.out.println("CORE_LIFECYCLE_REJECTION_CONFIRMED " + view); + } + + private static void expectCoreRejection(Callable construction) { + Throwable rejected = null; + try { + construction.call(); + } catch (Throwable failure) { + rejected = failure; + } + check(rejected != null, "new plugin silently accepted the released core without its lifecycle guarantee"); + for (var failure = rejected; failure != null; failure = failure.getCause()) { + if (failure instanceof IllegalStateException && failure.getMessage() != null + && failure.getMessage().contains("requires same-thread invocation hooks")) return; + } + throw new AssertionError("missing explicit core compatibility diagnostic", rejected); + } + private static void exercise(String view, boolean negative, boolean compatible, ListAppender logs) { GlobalOpenTelemetry.resetForTest(); // Reproduce #763's documented visible-API skew, not a claim of a deployed agent test. diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOutcomeBoundaryTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOutcomeBoundaryTest.java new file mode 100644 index 000000000..998c2e520 --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOutcomeBoundaryTest.java @@ -0,0 +1,398 @@ +// 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 io.opentelemetry.context.Context; +import io.opentelemetry.context.ContextKey; +import io.opentelemetry.context.Scope; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Proxy; +import java.time.Duration; +import java.util.List; +import java.util.concurrent.CompletionException; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executors; +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.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.ValueSource; +import org.slf4j.MDC; +import org.slf4j.spi.MDCAdapter; +import software.amazon.awssdk.services.lambda.model.ErrorObject; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +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; +import software.amazon.lambda.durable.serde.JacksonSerDes; +import software.amazon.lambda.durable.serde.SerDes; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class InvocationOutcomeBoundaryTest { + @ParameterizedTest + @CsvSource({"false,false", "false,true", "true,false", "true,true"}) + void cleanupWaitContractAlsoAppliesWithoutPlugins(boolean withPlugin, boolean retry) throws Exception { + var entered = new CountDownLatch(1); + var release = new CountDownLatch(1); + var original = retryError(); + var config = DurableConfig.builder(); + if (withPlugin) config.withPlugins(new DurableExecutionPlugin() {}); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + try { + if (retry) + context.step("retry", String.class, step -> { + throw original; + }); + else context.wait("pause", Duration.ofSeconds(1)); + return "done"; + } finally { + entered.countDown(); + await(release); + } + }, + config.build()); + var caller = Executors.newSingleThreadExecutor(); + try { + var response = caller.submit(() -> runner.run("input")); + assertTrue(entered.await(5, TimeUnit.SECONDS)); + assertThrows(TimeoutException.class, () -> response.get(700, TimeUnit.MILLISECONDS)); + release.countDown(); + if (retry) + assertSame( + original, + assertThrows(ExecutionException.class, () -> response.get(5, TimeUnit.SECONDS)) + .getCause()); + else + assertEquals( + ExecutionStatus.PENDING, + response.get(5, TimeUnit.SECONDS).getStatus()); + } finally { + release.countDown(); + caller.shutdownNow(); + assertTrue(caller.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @SuppressWarnings("removal") + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void invocationEndFatalEscapesItsActualWorkerAfterSettlingCaller(boolean death) throws Exception { + Error fatal = death ? new ThreadDeath() : new InternalError("fatal end cleanup"); + var ownerFailure = new AtomicReference(); + var escaped = new CountDownLatch(1); + var ends = new AtomicInteger(); + var workers = Executors.newSingleThreadExecutor(task -> { + var worker = new Thread(task, "end-fatal-owner"); + worker.setUncaughtExceptionHandler((thread, failure) -> { + ownerFailure.set(failure); + escaped.countDown(); + }); + return worker; + }); + try { + var runner = LocalDurableTestRunner.create( + String.class, + (input, ctx) -> "done", + DurableConfig.builder() + .withExecutorService(workers) + .withPlugins(new DurableExecutionPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + ends.incrementAndGet(); + throw fatal; + } + }) + .build()); + assertSame(fatal, assertThrows(Error.class, () -> runner.run("input"))); + assertTrue(escaped.await(2, TimeUnit.SECONDS), "fatal end cleanup must escape its actual worker"); + assertSame(fatal, ownerFailure.get()); + assertEquals(1, ends.get()); + } finally { + workers.shutdownNow(); + assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @ParameterizedTest + @CsvSource({"SUCCEEDED,false", "SUCCEEDED,true", "PENDING,false", "PENDING,true", "RETRYING,false", "RETRYING,true" + }) + void inlineNonfatalMdcRestorationPreservesSelectedOutcome(String outcome, boolean ambient) throws Exception { + runInlineMdcFailure(outcome, ambient, new IllegalStateException("restore failed"), null); + } + + @SuppressWarnings("removal") + @ParameterizedTest + @CsvSource({"false,false", "false,true", "true,false", "true,true"}) + void inlineFatalRestorationStillEscapesWithOriginalIdentity(boolean death, boolean wrapped) throws Exception { + Error fatal = death ? new ThreadDeath() : new InternalError("fatal restoration"); + runInlineMdcFailure( + "SUCCEEDED", false, wrapped ? new CompletionException(new ExecutionException(fatal)) : fatal, fatal); + } + + @ParameterizedTest + @ValueSource(strings = {"unreadable", "cycle", "fatal"}) + void lifecycleDiagnosticTraversalIsBoundedAndPreservesFatalIdentity(String kind) throws Exception { + var reads = new AtomicInteger(); + var fatal = new InternalError("fatal diagnostic"); + var failure = new CompletionException("diagnostic", null) { + @Override + public synchronized Throwable getCause() { + reads.incrementAndGet(); + if (kind.equals("cycle")) return this; + if (kind.equals("fatal")) throw fatal; + throw new IllegalStateException("unreadable diagnostic"); + } + }; + runInlineMdcFailure("SUCCEEDED", false, failure, kind.equals("fatal") ? fatal : null); + assertEquals(1, reads.get()); + } + + private static void runInlineMdcFailure( + String outcome, boolean ambient, Throwable restoreFailure, Error expectedFatal) throws Exception { + var adapter = MDC.getMDCAdapter(); + var before = MDC.getCopyOfContextMap(); + var ends = new CopyOnWriteArrayList(); + var injected = new AtomicBoolean(); + var original = retryError(); + var executor = new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS, new SynchronousQueue()) { + @Override + public void execute(Runnable task) { + task.run(); + } + }; + try { + MDC.clear(); + if (ambient) MDC.put("ambient", "saved"); + replaceAdapter((MDCAdapter) Proxy.newProxyInstance( + MDCAdapter.class.getClassLoader(), new Class[] {MDCAdapter.class}, (proxy, method, args) -> { + if (!ends.isEmpty() + && method.getName().equals(ambient ? "setContextMap" : "clear") + && injected.compareAndSet(false, true)) throw restoreFailure; + try { + return method.invoke(adapter, args); + } catch (InvocationTargetException failure) { + throw failure.getCause(); + } + })); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + if (outcome.equals("PENDING")) context.wait("pause", Duration.ofSeconds(1)); + if (outcome.equals("RETRYING")) + context.step("retry", String.class, step -> { + throw original; + }); + return "done"; + }, + DurableConfig.builder() + .withExecutorService(executor) + .withPlugins(new DurableExecutionPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + ends.add(info); + } + }) + .build()); + if (expectedFatal != null) assertSame(expectedFatal, assertThrows(Error.class, () -> runner.run("input"))); + else if (outcome.equals("RETRYING")) + assertSame( + original, + assertThrows(UnrecoverableDurableExecutionException.class, () -> runner.run("input"))); + else + assertEquals( + ExecutionStatus.valueOf(outcome), runner.run("input").getStatus()); + assertTrue(injected.get()); + assertEquals(1, ends.size()); + assertEquals(InvocationStatus.valueOf(outcome), ends.get(0).invocationStatus()); + } finally { + executor.shutdownNow(); + replaceAdapter(adapter); + if (before == null) MDC.clear(); + else MDC.setContextMap(before); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void deliveryFailureRemainsRetryingUntilAReplayActuallySucceeds(boolean executionView) { + var exporter = InMemorySpanExporter.create(); + var provider = SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)); + var options = OtelPluginConfig.builder() + .contextExtractor(() -> null) + .enableMdc(false) + .build(); + DurableExecutionPlugin otel = executionView + ? new ExecutionOtelPlugin(provider, options) + : new InvocationOtelPlugin(provider, options); + var ends = new CopyOnWriteArrayList(); + var failDelivery = new AtomicBoolean(true); + var original = new IllegalStateException("cannot deliver result yet"); + var completedSideEffects = new AtomicInteger(); + var serDes = new SerDes() { + private final JacksonSerDes delegate = new JacksonSerDes(); + + @Override + public String serialize(Object value) { + if ("done".equals(value) && failDelivery.compareAndSet(true, false)) throw original; + return delegate.serialize(value); + } + + @Override + public T deserialize(String value, TypeToken type) { + return delegate.deserialize(value, type); + } + }; + var runner = LocalDurableTestRunner.create( + String.class, + (input, ctx) -> { + ctx.step("saved", String.class, step -> { + completedSideEffects.incrementAndGet(); + return "checkpointed"; + }); + return "done"; + }, + DurableConfig.builder() + .withSerDes(serDes) + .withPlugins(otel, new DurableExecutionPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + ends.add(info); + } + }) + .build()); + assertSame(original, assertThrows(IllegalStateException.class, () -> runner.run("input"))); + assertEquals(InvocationStatus.RETRYING, ends.get(0).invocationStatus()); + assertSame(original, ends.get(0).executionError()); + assertEquals( + 0, + exporter.getFinishedSpanItems().stream() + .filter(s -> s.getName().equals("Workflow")) + .count()); + assertEquals(ExecutionStatus.SUCCEEDED, runner.runUntilComplete("input").getStatus()); + assertEquals( + List.of(InvocationStatus.RETRYING, InvocationStatus.SUCCEEDED), + ends.stream().map(InvocationEndInfo::invocationStatus).toList()); + assertEquals(1, completedSideEffects.get()); + assertEquals( + 1, + exporter.getFinishedSpanItems().stream() + .filter(s -> s.getName().equals("Workflow")) + .count()); + } + + @Test + void partialStartStillUnwindsAllConfiguredEndHooksAndRestoresBothThreads() throws Exception { + var key = ContextKey.named("partial-start-context"); + var order = new CopyOnWriteArrayList(); + var scopes = new Scope[2]; + var bodyCalls = new AtomicInteger(); + var workers = Executors.newSingleThreadExecutor(task -> new Thread( + () -> { + try (var ignored = + Context.root().with(key, "worker ambient").makeCurrent()) { + task.run(); + } + }, + "partial-start-owner")); + var first = new DurableExecutionPlugin() { + @Override + public void onInvocationStart(InvocationInfo info) { + order.add("start:first"); + scopes[0] = Context.current().with(key, "first").makeCurrent(); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + order.add("end:first"); + scopes[0].close(); + } + }; + var failing = new DurableExecutionPlugin() { + @Override + public void onInvocationStart(InvocationInfo info) { + order.add("start:failing"); + scopes[1] = Context.current().with(key, "failing").makeCurrent(); + throw new AssertionError("start failed after acquiring context"); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + order.add("end:failing"); + scopes[1].close(); + } + }; + var unstarted = new DurableExecutionPlugin() { + @Override + public void onInvocationStart(InvocationInfo info) { + order.add("unexpected start"); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + order.add("end:unstarted"); + } + }; + try (var ignored = Context.current().with(key, "caller ambient").makeCurrent()) { + var before = Context.current(); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + bodyCalls.incrementAndGet(); + return "done"; + }, + DurableConfig.builder() + .withExecutorService(workers) + .withPlugins(first, failing, unstarted) + .build()); + assertEquals(ExecutionStatus.FAILED, runner.run("input").getStatus()); + assertEquals(List.of("start:first", "start:failing", "end:unstarted", "end:failing", "end:first"), order); + assertEquals(0, bodyCalls.get()); + assertSame(before, Context.current()); + assertEquals( + "worker ambient", + workers.submit(() -> Context.current().get(key)).get(5, TimeUnit.SECONDS)); + } finally { + workers.shutdownNow(); + assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private static UnrecoverableDurableExecutionException retryError() { + return new UnrecoverableDurableExecutionException( + ErrorObject.builder().errorMessage("retry").build(), true); + } + + private static void await(CountDownLatch latch) { + try { + assertTrue(latch.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException failure) { + Thread.currentThread().interrupt(); + throw new AssertionError(failure); + } + } + + private static void replaceAdapter(MDCAdapter adapter) throws Exception { + var setter = MDC.class.getDeclaredMethod("setMDCAdapter", MDCAdapter.class); + setter.setAccessible(true); + setter.invoke(null, adapter); + } +} 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 5197cc716..417ff32a8 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 @@ -4,9 +4,14 @@ import com.amazonaws.services.lambda.runtime.Context; import com.amazonaws.services.lambda.runtime.RequestHandler; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.UndeclaredThrowableException; import java.nio.charset.StandardCharsets; +import java.util.Collections; +import java.util.IdentityHashMap; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; import java.util.function.BiFunction; import java.util.function.Function; @@ -53,6 +58,15 @@ public class DurableExecutor { private DurableExecutor() {} + /** + * Returns whether this core runs invocation start/end hooks on the root handler thread and awaits handler cleanup + * and end-hook completion before returning. Plugins that retain thread-local scopes may require this capability + * before construction. Introduced with the 2.2.2 lifecycle contract. + */ + public static boolean supportsSameThreadInvocationHooks() { + return true; + } + public static DurableExecutionOutput execute( DurableExecutionInput input, Context lambdaContext, @@ -104,22 +118,18 @@ private Invocation( } private CompletableFuture execute() { - if (plugins.isEmpty()) { - var body = CompletableFuture.supplyAsync(this::invokeHandler, config.getExecutorService()); - return manager.runUntilCompleteOrSuspend(body).handle(this::finishInvocation); - } 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); - return supplyHandler( - () -> { - Outcome.capture(this::invokeHandler).complete(body); - var selected = outcome.join(); - return finishInvocation(selected.value(), selected.failure()); - }, - failure -> finishInvocation(null, failure), - config.getExecutorService()); + Supplier task = () -> { + Outcome.capture(this::invokeHandler).complete(body); + var selected = outcome.join(); + return finishInvocation(selected.value(), selected.failure()); + }; + // Cleanup is awaited even without plugins. Keep that path free of plugin MDC handling. + if (plugins.isEmpty()) return CompletableFuture.supplyAsync(task, config.getExecutorService()); + return supplyHandler(task, failure -> finishInvocation(null, failure), config.getExecutorService()); } private O invokeHandler() { @@ -166,7 +176,8 @@ private DurableExecutionOutput finishInvocation(O value, Throwable failure) { // Result serialization/checkpointing also belongs to this invocation. Close plugin resources even // when delivery fails, then retain the original exception for the Lambda caller. var cause = ExceptionHelper.unwrapCompletableFuture(deliveryFailure); - fireOnInvocationEnd(failureStatus(cause), cause, null); + // No terminal output is delivered: this exception escapes for a Lambda invocation retry. + fireOnInvocationEnd(InvocationStatus.RETRYING, cause, null); ExceptionHelper.sneakyThrow(deliveryFailure); return null; } @@ -238,17 +249,53 @@ private static CompletableFuture supplyHandler( } var outcome = Outcome.capture(task); try { - restore.close(); + try (var ignored = restore) { + // End/delivery failures must also escape their actual worker. Handler-body failures have + // already been mapped to their selected durable outcome; this does not reclassify them. + rethrowLifecycleFatal(outcome.failure()); + } } finally { // End and worker restoration have run before publishing. A restoration failure still escapes its // owner, without changing the invocation outcome already delivered to the end hooks. outcome.complete(result); } }; - executor.execute(work); + try { + executor.execute(work); + } catch (RuntimeException dispatchFailure) { + if (!result.isDone()) throw dispatchFailure; + // Inline execution can throw from MDC restoration after settling the selected outcome. Preserve that + // outcome for an ordinary cleanup failure, while a JVM-fatal failure still reaches the caller. + rethrowLifecycleFatal(dispatchFailure); + } return result; } + /** Inspects only standard transport wrappers at this lifecycle boundary, without trusting diagnostics. */ + @SuppressWarnings("removal") + private static void rethrowLifecycleFatal(Throwable failure) { + if (failure instanceof VirtualMachineError fatal) throw fatal; + if (failure instanceof ThreadDeath fatal) throw fatal; + if (failure == null) return; + var visited = Collections.newSetFromMap(new IdentityHashMap()); + var cause = failure; + while (cause != null && visited.add(cause)) { + if (cause instanceof VirtualMachineError fatal) throw fatal; + if (cause instanceof ThreadDeath fatal) throw fatal; + if (!(cause instanceof CompletionException + || cause instanceof ExecutionException + || cause instanceof InvocationTargetException + || cause instanceof UndeclaredThrowableException)) return; + try { + cause = cause.getCause(); + } catch (VirtualMachineError | ThreadDeath fatal) { + throw fatal; + } catch (Throwable unreadableDiagnostic) { + return; + } + } + } + private static SafeCloseable restoreMdcOnClose() { var previous = MDC.getCopyOfContextMap(); return () -> { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java b/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java index f91bc8d4d..88bbe82ba 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java @@ -38,7 +38,9 @@ default void onInvocationStart(InvocationInfo info) {} * later handler finally block returns or throws. * *

Start hooks run in registration order; end hooks run in reverse order to unwind nested thread-local scopes. If - * an end hook throws an Error that is not isolated, the remaining end hooks still run before it is rethrown. + * an end hook throws an Error that is not isolated, the remaining end hooks still run before it is rethrown. If a + * start hook throws an Error that is not isolated, later start hooks are not called, but all configured end hooks + * still run in reverse order. An end hook must therefore tolerate a start hook that did not complete. * *

Check {@link InvocationEndInfo#invocationStatus()} to detect if the execution reached a terminal state in this * invocation (useful for writing summary records or flushing final data). diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationLifecycleTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationLifecycleTest.java index c2fca5f4f..cb9259472 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationLifecycleTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/InvocationLifecycleTest.java @@ -287,7 +287,7 @@ void serializationFailureStillEndsOnTheHandlerThread() throws Exception { assertEquals(1, plugin.starts.size()); assertEquals(1, plugin.ends.size()); assertEquals(plugin.startThreads, plugin.endThreads); - assertEquals(InvocationStatus.FAILED, plugin.ends.get(0).invocationStatus()); + assertEquals(InvocationStatus.RETRYING, plugin.ends.get(0).invocationStatus()); assertSame(failure, plugin.ends.get(0).executionError()); assertNull(plugin.ends.get(0).executionResult()); } finally { From cacd73056ff1cfcdb0f9e544531fa76f46ecd0a9 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 00:33:14 +0000 Subject: [PATCH 23/48] fix: propagate fatal MDC initialization failures before startup --- otel-plugin/README.md | 4 + .../otel/MdcFatalInitializationTest.java | 252 ++++++++++++++++++ .../durable/execution/DurableExecutor.java | 9 + 3 files changed, 265 insertions(+) create mode 100644 otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFatalInitializationTest.java diff --git a/otel-plugin/README.md b/otel-plugin/README.md index 8cb340d4b..2327d9345 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -30,6 +30,10 @@ correctly; remaining end hooks run even when another hook raises an Error, then Invocations without plugins also wait for handler cleanup before returning the selected suspension or retry outcome. Output-delivery failures that escape for a Lambda retry report `RETRYING`, so they do not prematurely end the Workflow span. An ordinary MDC-restoration failure on an inline executor does not replace an already selected outcome. +Before invocation startup, a JVM-fatal error from MDC capture (direct or inside a standard transport wrapper) +completes the observation future exceptionally with that same fatal before escaping the handler worker. No start, +body, or end hook runs, and no durable `FAILED` response is produced for that fatal. Ordinary initialization errors +and handler/body failure classification retain their existing behavior. This plugin version requires the core's `DurableExecutor.supportsSameThreadInvocationHooks()` capability, introduced in the 2.2.2 lifecycle contract (currently `2.2.2-SNAPSHOT`). Upgrade the core together with the plugin layer. Every diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFatalInitializationTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFatalInitializationTest.java new file mode 100644 index 000000000..8f835bebf --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFatalInitializationTest.java @@ -0,0 +1,252 @@ +// 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 java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Proxy; +import java.lang.reflect.UndeclaredThrowableException; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.AbstractExecutorService; +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.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.stream.Stream; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import org.slf4j.MDC; +import org.slf4j.spi.MDCAdapter; +import software.amazon.lambda.durable.DurableConfig; +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.testing.LocalDurableTestRunner; + +class MdcFatalInitializationTest { + @ParameterizedTest + @MethodSource("fatalCases") + @SuppressWarnings("removal") + void fatalCaptureSettlesCallerThenEscapesOwnerBeforeAnyHook(String mode, boolean death, String wrapper) + throws Exception { + Error fatal = death ? new ThreadDeath() : new InternalError("fatal capture"); + Throwable failure = + switch (wrapper) { + case "completion" -> new CompletionException(fatal); + case "nested" -> new CompletionException(new ExecutionException(fatal)); + case "reflection" -> new InvocationTargetException(fatal); + case "proxy" -> new UndeclaredThrowableException(fatal); + default -> fatal; + }; + runCapture(mode, failure, fatal); + } + + @ParameterizedTest + @MethodSource("ordinaryCases") + void nonfatalInitializationKeepsItsDurableFailurePolicy(String mode, String kind) throws Exception { + Throwable failure = + switch (kind) { + case "assertion" -> new AssertionError("ordinary capture"); + case "application-cause" -> + new IllegalStateException("ordinary capture", new InternalError("diagnostic")); + default -> new IllegalStateException("ordinary capture"); + }; + runCapture(mode, failure, null); + } + + private static Stream fatalCases() { + return Stream.of("direct", "async", "precompleted") + .flatMap(mode -> Stream.of(false, true) + .flatMap(death -> Stream.of("direct", "completion", "nested", "reflection", "proxy") + .map(wrapper -> Arguments.of(mode, death, wrapper)))); + } + + private static Stream ordinaryCases() { + return Stream.of("direct", "async", "precompleted") + .flatMap(mode -> + Stream.of("exception", "assertion", "application-cause").map(kind -> Arguments.of(mode, kind))); + } + + private static void runCapture(String mode, Throwable failure, Error expectedFatal) throws Exception { + var originalAdapter = MDC.getMDCAdapter(); + var injected = new AtomicBoolean(); + var starts = new AtomicInteger(); + var bodies = new AtomicInteger(); + var ends = new AtomicInteger(); + var callerThread = new AtomicReference(); + var workers = new ObservedExecutor(mode); + var caller = Executors.newSingleThreadExecutor(task -> daemon(task, "capture-caller")); + try { + replaceAdapter((MDCAdapter) Proxy.newProxyInstance( + MDCAdapter.class.getClassLoader(), new Class[] {MDCAdapter.class}, (proxy, method, args) -> { + var capture = method.getName().equals("getCopyOfContextMap") + && Arrays.stream(Thread.currentThread().getStackTrace()) + .anyMatch(frame -> frame.getClassName().endsWith(".DurableExecutor") + && frame.getMethodName().equals("restoreMdcOnClose")); + if (capture && injected.compareAndSet(false, true)) throw failure; + try { + return method.invoke(originalAdapter, args); + } catch (InvocationTargetException invocation) { + throw invocation.getCause(); + } + })); + var plugin = new DurableExecutionPlugin() { + @Override + public void onInvocationStart(InvocationInfo info) { + starts.incrementAndGet(); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + ends.incrementAndGet(); + } + }; + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + bodies.incrementAndGet(); + return "unreachable"; + }, + DurableConfig.builder() + .withExecutorService(workers) + .withPlugins(plugin) + .build()); + var response = caller.submit(() -> { + callerThread.set(Thread.currentThread()); + return runner.run("input"); + }); + if (expectedFatal != null) { + var thrown = assertThrows( + ExecutionException.class, + () -> response.get(3, TimeUnit.SECONDS), + "the observation future must settle rather than strand the invocation"); + assertSame(expectedFatal, thrown.getCause()); + assertTrue(workers.escaped.await(3, TimeUnit.SECONDS)); + assertSame(expectedFatal, workers.escape.get()); + if (mode.equals("direct")) assertSame(callerThread.get(), workers.owner.get()); + else { + assertNotSame(callerThread.get(), workers.owner.get()); + assertTrue(workers.uncaught.await(3, TimeUnit.SECONDS)); + assertSame(expectedFatal, workers.uncaughtFailure.get()); + } + } else { + var result = response.get(3, TimeUnit.SECONDS); + assertEquals(ExecutionStatus.FAILED, result.getStatus()); + assertEquals( + failure.getClass().getName(), + result.getError().orElseThrow().errorType()); + assertEquals("ordinary capture", result.getError().orElseThrow().errorMessage()); + assertNull(workers.escape.get(), "ordinary initialization policy must not be broadened"); + } + assertTrue(injected.get()); + assertEquals(0, starts.get()); + assertEquals(0, bodies.get()); + assertEquals(0, ends.get()); + } finally { + workers.shutdownNow(); + caller.shutdownNow(); + try { + assertTrue(workers.awaitTermination(3, TimeUnit.SECONDS)); + assertTrue(caller.awaitTermination(3, TimeUnit.SECONDS)); + } finally { + replaceAdapter(originalAdapter); + } + } + } + + private static Thread daemon(Runnable task, String name) { + var thread = new Thread(task, name); + thread.setDaemon(true); + return thread; + } + + private static void replaceAdapter(MDCAdapter adapter) throws Exception { + var setter = MDC.class.getDeclaredMethod("setMDCAdapter", MDCAdapter.class); + setter.setAccessible(true); + setter.invoke(null, adapter); + } + + private static final class ObservedExecutor extends AbstractExecutorService { + private final String mode; + private final AtomicReference escape = new AtomicReference<>(); + private final AtomicReference owner = new AtomicReference<>(); + private final CountDownLatch escaped = new CountDownLatch(1); + private final AtomicReference uncaughtFailure = new AtomicReference<>(); + private final CountDownLatch uncaught = new CountDownLatch(1); + private final ExecutorService delegate = Executors.newSingleThreadExecutor(task -> { + var thread = daemon(task, "capture-worker"); + thread.setUncaughtExceptionHandler((worker, failure) -> { + uncaughtFailure.set(failure); + uncaught.countDown(); + }); + return thread; + }); + + private ObservedExecutor(String mode) { + this.mode = mode; + } + + @Override + public void execute(Runnable task) { + var completed = new CountDownLatch(1); + Runnable observed = () -> { + owner.set(Thread.currentThread()); + try { + task.run(); + } catch (Error failure) { + escape.set(failure); + escaped.countDown(); + throw failure; + } finally { + completed.countDown(); + } + }; + if (mode.equals("direct")) observed.run(); + else { + delegate.execute(observed); + if (mode.equals("precompleted")) { + try { + assertTrue(completed.await(3, TimeUnit.SECONDS)); + } catch (InterruptedException failure) { + Thread.currentThread().interrupt(); + throw new AssertionError(failure); + } + } + } + } + + @Override + public void shutdown() { + delegate.shutdown(); + } + + @Override + public List shutdownNow() { + return delegate.shutdownNow(); + } + + @Override + public boolean isShutdown() { + return delegate.isShutdown(); + } + + @Override + public boolean isTerminated() { + return delegate.isTerminated(); + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { + return delegate.awaitTermination(timeout, unit); + } + } +} 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 417ff32a8..e6438d643 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 @@ -236,6 +236,7 @@ private void complete(CompletableFuture future) { } } + @SuppressWarnings("removal") private static CompletableFuture supplyHandler( Supplier task, Function initializationFailure, Executor executor) { var result = new CompletableFuture(); @@ -244,6 +245,14 @@ private static CompletableFuture supplyHandler( try { restore = restoreMdcOnClose(); } catch (Throwable failure) { + try { + rethrowLifecycleFatal(failure); + } catch (VirtualMachineError | ThreadDeath fatal) { + // Wake the invocation caller with the original fatal before it escapes the actual worker. + // Throwing first would leave the observation future incomplete on an asynchronous executor. + result.completeExceptionally(fatal); + throw fatal; + } Outcome.capture(() -> initializationFailure.apply(failure)).complete(result); return; } From b1fb4ebae103b8951cf400d9c50d92d5d39a4899 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 01:12:52 +0000 Subject: [PATCH 24/48] fix: preserve invocation end failures through scope cleanup --- otel-plugin/README.md | 5 +- .../durable/otel/ExecutionOtelPlugin.java | 10 +-- .../durable/otel/InvocationOtelPlugin.java | 10 +-- .../durable/otel/OtelPluginSupport.java | 32 +++++++ .../otel/InvocationEndFailureTest.java | 85 +++++++++++++++++++ .../otel/InvocationOutcomeBoundaryTest.java | 31 +++++-- .../plugin/DurableExecutionPlugin.java | 8 +- .../lambda/durable/plugin/PluginRunner.java | 16 +++- .../durable/plugin/PluginRunnerTest.java | 4 +- 9 files changed, 172 insertions(+), 29 deletions(-) diff --git a/otel-plugin/README.md b/otel-plugin/README.md index 2327d9345..1c5835e9a 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -26,7 +26,10 @@ including when execution suspends or terminates. Handler `finally` blocks must f can run; a blocked handler cleanup therefore also blocks the invocation response. Context restoration still runs when span finalization or flushing fails. Ordinary exceptions and nonfatal linkage errors retain the existing plugin-hook isolation behavior. Invocation-end hooks run in reverse registration order so nested scopes unwind -correctly; remaining end hooks run even when another hook raises an Error, then the first Error propagates. +correctly; remaining end hooks run even when another hook raises an Error. The first unisolated Error propagates, +unless a later `VirtualMachineError` or `ThreadDeath` takes precedence over a non-JVM-fatal Error; other distinct +end-hook Errors are retained as suppressed failures. Scope cleanup preserves finalization failures using the same +JVM-fatal precedence, so an ordinary cleanup exception cannot hide an earlier Error. Invocations without plugins also wait for handler cleanup before returning the selected suspension or retry outcome. Output-delivery failures that escape for a Lambda retry report `RETRYING`, so they do not prematurely end the Workflow span. An ordinary MDC-restoration failure on an inline executor does not replace an already selected outcome. diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java index 0702d84b4..a621e8900 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java @@ -290,13 +290,9 @@ private Scope activateHandlerContext() { @Override public void onInvocationEnd(InvocationEndInfo info) { - try { - endInvocation(info); - } finally { - var scope = handlerScope; - handlerScope = null; - if (scope != null) scope.close(); - } + var scope = handlerScope; + handlerScope = null; + OtelPluginSupport.runInvocationEnd(scope, () -> endInvocation(info)); } private void endInvocation(InvocationEndInfo info) { diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java index 0498e7727..3042c603d 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java @@ -287,13 +287,9 @@ private Scope activateHandlerContext() { @Override public void onInvocationEnd(InvocationEndInfo info) { - try { - endInvocation(info); - } finally { - var scope = handlerScope; - handlerScope = null; - if (scope != null) scope.close(); - } + var scope = handlerScope; + handlerScope = null; + OtelPluginSupport.runInvocationEnd(scope, () -> endInvocation(info)); } private void endInvocation(InvocationEndInfo info) { diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java index cfad36746..9dd9684f5 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java @@ -9,6 +9,7 @@ import io.opentelemetry.api.trace.Tracer; import io.opentelemetry.api.trace.TracerProvider; import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; import io.opentelemetry.sdk.trace.SdkTracerProvider; import io.opentelemetry.sdk.trace.samplers.SamplingResult; import java.nio.file.Files; @@ -25,6 +26,37 @@ final class OtelPluginSupport { private OtelPluginSupport() {} + /** Closes the owning thread's scope without hiding a finalization error or a later JVM-fatal cleanup failure. */ + static void runInvocationEnd(Scope scope, Runnable end) { + Throwable primary = null; + try { + end.run(); + } catch (Throwable failure) { + primary = failure; + throw failure; + } finally { + try { + if (scope != null) scope.close(); + } catch (Throwable cleanup) { + if (primary == null) throw cleanup; + if (primary != cleanup) { + if (endFailurePriority(cleanup) > endFailurePriority(primary)) { + cleanup.addSuppressed(primary); + throw cleanup; + } + primary.addSuppressed(cleanup); + } + } + } + } + + @SuppressWarnings("removal") + private static int endFailurePriority(Throwable failure) { + if (failure instanceof VirtualMachineError || failure instanceof ThreadDeath) return 3; + if (failure instanceof Error && !(failure instanceof LinkageError)) return 2; + return 1; + } + /** Rejects an older core before a plugin builds a provider or activates any thread-local context. */ static void requireSameThreadInvocationHooks() { var message = "This OpenTelemetry plugin requires same-thread invocation hooks from the Durable Execution" diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndFailureTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndFailureTest.java index 33ea7aada..cff654b99 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndFailureTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndFailureTest.java @@ -88,6 +88,91 @@ public void onInvocationEnd(InvocationEndInfo info) { } } + @ParameterizedTest + @MethodSource("combinedFailures") + void preservesCombinedFinalizationAndScopeFailures(boolean executionView, String phase, String pair) + throws Exception { + Throwable primary = + switch (pair.split("-")[0]) { + case "fatal" -> new InternalError("primary finalization"); + case "assert" -> new AssertionError("primary finalization"); + case "linkage" -> new NoClassDefFoundError("primary finalization"); + default -> new IllegalStateException("primary finalization"); + }; + Throwable cleanup = + switch (pair.split("-")[1]) { + case "fatal" -> new InternalError("scope cleanup"); + case "assert" -> new AssertionError("scope cleanup"); + case "same" -> primary; + default -> new IllegalStateException("scope cleanup"); + }; + var cleanupWins = List.of("runtime-fatal", "assert-fatal", "runtime-assert", "linkage-assert") + .contains(pair); + var expected = cleanupWins ? cleanup : primary; + var secondary = cleanupWins ? primary : cleanup; + var processor = new FailingProcessor(phase, primary); + var config = OtelPluginConfig.builder() + .enableMdc(false) + .contextExtractor(() -> new ExtractedContext( + "12345678901234567890123456789012", "1234567890123456", ExtractedContext.Sampling.SAMPLED)) + .build(); + var builder = SdkTracerProvider.builder().addSpanProcessor(processor); + DurableExecutionPlugin plugin = + executionView ? new ExecutionOtelPlugin(builder, config) : new InvocationOtelPlugin(builder, config); + var healthyEnds = new AtomicInteger(); + var closed = new AtomicInteger(); + var healthy = new DurableExecutionPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + healthyEnds.incrementAndGet(); + } + }; + var runner = new PluginRunner(List.of(healthy, plugin)); + var ambient = Span.wrap(SpanContext.create( + "abcdefabcdefabcdefabcdefabcdefab", + "abcdefabcdefabcd", + TraceFlags.getSampled(), + TraceState.getDefault())); + try (var ignored = ambient.makeCurrent()) { + var original = Context.current(); + runner.onInvocationStart(new InvocationInfo("req", "arn", true, Instant.now())); + var field = plugin.getClass().getDeclaredField("handlerScope"); + field.setAccessible(true); + var scope = (Scope) field.get(plugin); + field.set(plugin, (Scope) () -> { + closed.incrementAndGet(); + scope.close(); + raise(cleanup); + }); + var end = new InvocationEndInfo("req", "arn", true, InvocationStatus.SUCCEEDED, null); + if (expected instanceof RuntimeException || expected instanceof LinkageError) { + assertDoesNotThrow(() -> runner.onInvocationEnd(end)); + } else assertSame(expected, assertThrows(Error.class, () -> runner.onInvocationEnd(end))); + assertEquals(expected == secondary ? List.of() : List.of(secondary), List.of(expected.getSuppressed())); + assertSame(original, Context.current()); + assertEquals(1, healthyEnds.get()); + assertEquals(1, closed.get()); + assertDoesNotThrow(() -> plugin.onInvocationEnd(end)); + assertEquals(1, closed.get(), "scope cleanup is one-shot even when both phases throw"); + } + } + + private static Stream combinedFailures() { + return Stream.of(false, true) + .flatMap(view -> Stream.of("span", "flush") + .flatMap(phase -> Stream.of( + "fatal-runtime", + "assert-runtime", + "runtime-fatal", + "assert-fatal", + "fatal-fatal", + "fatal-same", + "runtime-assert", + "linkage-assert", + "runtime-runtime") + .map(pair -> Arguments.of(view, phase, pair)))); + } + private static Stream failures() { return Stream.of(false, true) .flatMap(executionView -> Stream.of("span", "flush", "scope") diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOutcomeBoundaryTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOutcomeBoundaryTest.java index 998c2e520..2948d1433 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOutcomeBoundaryTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOutcomeBoundaryTest.java @@ -94,12 +94,27 @@ void cleanupWaitContractAlsoAppliesWithoutPlugins(boolean withPlugin, boolean re @SuppressWarnings("removal") @ParameterizedTest - @ValueSource(booleans = {false, true}) - void invocationEndFatalEscapesItsActualWorkerAfterSettlingCaller(boolean death) throws Exception { + @CsvSource({"false,false", "false,true", "true,false", "true,true"}) + void invocationEndFatalEscapesItsActualWorkerAfterSettlingCaller(boolean death, boolean earlierAssertion) + throws Exception { Error fatal = death ? new ThreadDeath() : new InternalError("fatal end cleanup"); var ownerFailure = new AtomicReference(); var escaped = new CountDownLatch(1); var ends = new AtomicInteger(); + var assertion = new AssertionError("earlier nonfatal cleanup"); + var fatalPlugin = new DurableExecutionPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + ends.incrementAndGet(); + throw fatal; + } + }; + var assertionPlugin = new DurableExecutionPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + throw assertion; + } + }; var workers = Executors.newSingleThreadExecutor(task -> { var worker = new Thread(task, "end-fatal-owner"); worker.setUncaughtExceptionHandler((thread, failure) -> { @@ -114,17 +129,15 @@ void invocationEndFatalEscapesItsActualWorkerAfterSettlingCaller(boolean death) (input, ctx) -> "done", DurableConfig.builder() .withExecutorService(workers) - .withPlugins(new DurableExecutionPlugin() { - @Override - public void onInvocationEnd(InvocationEndInfo info) { - ends.incrementAndGet(); - throw fatal; - } - }) + .withPlugins( + earlierAssertion + ? new DurableExecutionPlugin[] {fatalPlugin, assertionPlugin} + : new DurableExecutionPlugin[] {fatalPlugin}) .build()); assertSame(fatal, assertThrows(Error.class, () -> runner.run("input"))); assertTrue(escaped.await(2, TimeUnit.SECONDS), "fatal end cleanup must escape its actual worker"); assertSame(fatal, ownerFailure.get()); + assertEquals(earlierAssertion ? List.of(assertion) : List.of(), List.of(fatal.getSuppressed())); assertEquals(1, ends.get()); } finally { workers.shutdownNow(); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java b/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java index 88bbe82ba..d305a1c54 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/plugin/DurableExecutionPlugin.java @@ -38,9 +38,11 @@ default void onInvocationStart(InvocationInfo info) {} * later handler finally block returns or throws. * *

Start hooks run in registration order; end hooks run in reverse order to unwind nested thread-local scopes. If - * an end hook throws an Error that is not isolated, the remaining end hooks still run before it is rethrown. If a - * start hook throws an Error that is not isolated, later start hooks are not called, but all configured end hooks - * still run in reverse order. An end hook must therefore tolerate a start hook that did not complete. + * an end hook throws an Error that is not isolated, the remaining end hooks still run before it is rethrown. The + * first such Error propagates, except that a later VirtualMachineError or ThreadDeath takes precedence over a + * non-JVM-fatal Error. Other distinct end-hook Errors are retained as suppressed failures. If a start hook throws + * an Error that is not isolated, later start hooks are not called, but all configured end hooks still run in + * reverse order. An end hook must therefore tolerate a start hook that did not complete. * *

Check {@link InvocationEndInfo#invocationStatus()} to detect if the execution reached a terminal state in this * invocation (useful for writing summary records or flushing final data). diff --git a/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java b/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java index 0ad9558af..5532eaeb0 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginRunner.java @@ -112,12 +112,26 @@ public void onInvocationEnd(InvocationEndInfo info) { runHook(plugins.get(index), p -> p.onInvocationEnd(info)); } catch (Error failure) { // Finish unwinding earlier plugins' thread-local scopes before propagating an end-hook error. - if (firstError == null) firstError = failure; + if (firstError == null) { + firstError = failure; + } else if (firstError != failure) { + if (isJvmFatal(failure) && !isJvmFatal(firstError)) { + failure.addSuppressed(firstError); + firstError = failure; + } else { + firstError.addSuppressed(failure); + } + } } } if (firstError != null) throw firstError; } + @SuppressWarnings("removal") + private static boolean isJvmFatal(Error failure) { + return failure instanceof VirtualMachineError || failure instanceof ThreadDeath; + } + public void onOperationStart(OperationInfo info) { run(p -> p.onOperationStart(info)); } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginRunnerTest.java b/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginRunnerTest.java index e8953a217..c006b6252 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginRunnerTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginRunnerTest.java @@ -179,15 +179,17 @@ public void onInvocationEnd(InvocationEndInfo info) { throw firstFailure; } }; + var laterFailure = new AssertionError("later cleanup failure"); var middle = new TestPlugin("middle", calls) { @Override public void onInvocationEnd(InvocationEndInfo info) { super.onInvocationEnd(info); - throw new AssertionError("later cleanup failure"); + throw laterFailure; } }; var runner = new PluginRunner(List.of(new TestPlugin("outer", calls), middle, inner)); assertSame(firstFailure, assertThrows(Error.class, () -> runner.onInvocationEnd(invocationEndInfo()))); + assertEquals(List.of(laterFailure), List.of(firstFailure.getSuppressed())); assertEquals(List.of("inner:onInvocationEnd", "middle:onInvocationEnd", "outer:onInvocationEnd"), calls); } } From e9a07b2a5e46931d1b75be3c1fc37c8c28bb69ff Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 01:20:09 +0000 Subject: [PATCH 25/48] test: assert stored parallel tolerance outcome across racing branches Reuse the previously validated scheduling-independent example test and controlled zero-success lane, adapted to the 2.x plugin registration API. Preserve unlimited concurrency, exact failures, checkpointed outcome and replay side-effect checks. --- .../ParallelFailureToleranceExampleTest.java | 124 ++++++++++++++++-- 1 file changed, 114 insertions(+), 10 deletions(-) diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/parallel/ParallelFailureToleranceExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/parallel/ParallelFailureToleranceExampleTest.java index f1518e2ac..cb23be650 100644 --- a/examples/src/test/java/software/amazon/lambda/durable/examples/parallel/ParallelFailureToleranceExampleTest.java +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/parallel/ParallelFailureToleranceExampleTest.java @@ -5,8 +5,27 @@ import static org.junit.jupiter.api.Assertions.*; import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; 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.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.model.ConcurrencyCompletionStatus; import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.model.ParallelResult; +import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; +import software.amazon.lambda.durable.plugin.OperationEndInfo; +import software.amazon.lambda.durable.plugin.UserFunctionEndInfo; +import software.amazon.lambda.durable.plugin.UserFunctionStartInfo; +import software.amazon.lambda.durable.serde.JacksonSerDes; import software.amazon.lambda.durable.testing.LocalDurableTestRunner; class ParallelFailureToleranceExampleTest { @@ -41,19 +60,104 @@ void succeedsWhenAllBranchesSucceed() { assertEquals(3, output.succeeded()); } - @Test - void failsWhenFailuresExceedTolerance() { + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void failsWhenFailuresExceedTolerance(boolean holdHealthyBranch) throws Exception { var handler = new ParallelFailureToleranceExample(); - var runner = LocalDurableTestRunner.create(ParallelFailureToleranceExample.Input.class, handler); + var probe = new CompletionProbe(holdHealthyBranch); + var config = DurableConfig.builder().withPlugins(probe.newPlugin()).build(); + var runner = LocalDurableTestRunner.create(ParallelFailureToleranceExample.Input.class, handler) + .withDurableConfig(config); + var caller = Executors.newSingleThreadExecutor(); + try { + var input = new ParallelFailureToleranceExample.Input(List.of("svc-a", "bad-svc-b", "bad-svc-c"), 1, 2); + var invocation = caller.submit(() -> runner.runUntilComplete(input)); + if (holdHealthyBranch) { + await(probe.parallelStored); + probe.releaseHealthy.countDown(); + } + var result = invocation.get(5, TimeUnit.SECONDS); + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + var output = result.getResult(ParallelFailureToleranceExample.Output.class); + assertEquals(2, output.failed()); - // 2 bad services, toleratedFailureCount=1 — second failure exceeds tolerance - var input = new ParallelFailureToleranceExample.Input(List.of("svc-a", "bad-svc-b", "bad-svc-c"), 1, 2); - var result = runner.runUntilComplete(input); + var operation = result.getOperation("call-services"); + assertEquals(OperationStatus.SUCCEEDED, operation.getStatus()); + var stored = new JacksonSerDes() + .deserialize(operation.getContextDetails().result(), TypeToken.get(ParallelResult.class)); + assertEquals(ConcurrencyCompletionStatus.FAILURE_TOLERANCE_EXCEEDED, stored.completionStatus()); + assertFalse(stored.completionStatus().isSucceeded()); + assertEquals(3, stored.size()); + assertEquals( + List.of(ParallelResult.Status.FAILED, ParallelResult.Status.FAILED), + stored.statuses().subList(1, 3)); + assertTrue(List.of(ParallelResult.Status.SUCCEEDED, ParallelResult.Status.SKIPPED) + .contains(stored.statuses().get(0))); + assertEquals(output.succeeded(), stored.succeeded()); + assertEquals(output.failed(), stored.failed()); + assertEquals(1 - output.succeeded(), stored.skipped()); + if (holdHealthyBranch) assertEquals(0, output.succeeded()); - assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + var completedCalls = result.getOperations().stream() + .filter(op -> op.getType() == OperationType.STEP && op.isCompleted()) + .collect(Collectors.toMap(op -> op.getName(), op -> probe.completedCalls.get(op.getName()))); + assertEquals(1, completedCalls.get("invoke-bad-svc-b")); + assertEquals(1, completedCalls.get("invoke-bad-svc-c")); + var replay = runner.run(input); + assertEquals(ExecutionStatus.SUCCEEDED, replay.getStatus()); + assertEquals(output, replay.getResult(ParallelFailureToleranceExample.Output.class)); + completedCalls.forEach((name, count) -> assertEquals( + count, + probe.completedCalls.get(name), + "Completed step bodies must not run again on replay: " + name)); + } finally { + probe.releaseHealthy.countDown(); + caller.shutdownNow(); + assertTrue(caller.awaitTermination(5, TimeUnit.SECONDS)); + } + } - var output = result.getResult(ParallelFailureToleranceExample.Output.class); - assertEquals(2, output.failed()); - assertEquals(1, output.succeeded()); + private static void await(CountDownLatch latch) { + try { + assertTrue(latch.await(5, TimeUnit.SECONDS), "Controlled branch scheduling was not released"); + } catch (InterruptedException failure) { + Thread.currentThread().interrupt(); + throw new AssertionError(failure); + } + } + + private static final class CompletionProbe { + final boolean holdHealthy; + final CountDownLatch healthyStarted = new CountDownLatch(1); + final CountDownLatch releaseHealthy = new CountDownLatch(1); + final CountDownLatch parallelStored = new CountDownLatch(1); + final Map completedCalls = new ConcurrentHashMap<>(); + + CompletionProbe(boolean holdHealthy) { + this.holdHealthy = holdHealthy; + } + + DurableExecutionPlugin newPlugin() { + return new DurableExecutionPlugin() { + @Override + public void onUserFunctionStart(UserFunctionStartInfo info) { + if (!holdHealthy || !"STEP".equals(info.type())) return; + if ("invoke-svc-a".equals(info.name())) { + healthyStarted.countDown(); + await(releaseHealthy); + } else if (info.name().startsWith("invoke-bad-")) await(healthyStarted); + } + + @Override + public void onUserFunctionEnd(UserFunctionEndInfo info) { + if ("STEP".equals(info.type())) completedCalls.merge(info.name(), 1, Integer::sum); + } + + @Override + public void onOperationEnd(OperationEndInfo info) { + if ("call-services".equals(info.name())) parallelStored.countDown(); + } + }; + } } } From cdb5c0c8755be5de60c142e2fde2dfb60a97c983 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 01:45:15 +0000 Subject: [PATCH 26/48] fix: preserve delivery and settled invocation outcomes --- otel-plugin/README.md | 6 +- .../otel/InvocationOutcomeBoundaryTest.java | 102 ++++++++++++++++++ .../durable/execution/DurableExecutor.java | 36 ++++--- 3 files changed, 130 insertions(+), 14 deletions(-) diff --git a/otel-plugin/README.md b/otel-plugin/README.md index 1c5835e9a..aa47d6510 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -31,8 +31,10 @@ unless a later `VirtualMachineError` or `ThreadDeath` takes precedence over a no end-hook Errors are retained as suppressed failures. Scope cleanup preserves finalization failures using the same JVM-fatal precedence, so an ordinary cleanup exception cannot hide an earlier Error. Invocations without plugins also wait for handler cleanup before returning the selected suspension or retry outcome. -Output-delivery failures that escape for a Lambda retry report `RETRYING`, so they do not prematurely end the Workflow -span. An ordinary MDC-restoration failure on an inline executor does not replace an already selected outcome. +Result and failure-response serialization finish before terminal invocation-end notification. Delivery failures that +escape for a Lambda retry report `RETRYING`, so they do not prematurely end the Workflow span. Non-JVM-fatal +MDC-restoration failures on an inline executor, including `AssertionError` and `LinkageError`, do not replace an +already selected outcome. Before invocation startup, a JVM-fatal error from MDC capture (direct or inside a standard transport wrapper) completes the observation future exceptionally with that same fatal before escaping the handler worker. No start, body, or end hook runs, and no durable `FAILED` response is produced for that fatal. Ordinary initialization errors diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOutcomeBoundaryTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOutcomeBoundaryTest.java index 2948d1433..14185265f 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOutcomeBoundaryTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOutcomeBoundaryTest.java @@ -3,7 +3,9 @@ package software.amazon.lambda.durable.otel; import static org.junit.jupiter.api.Assertions.*; +import static software.amazon.lambda.durable.otel.SpanAttributes.DURABLE_INVOCATION_STATUS; +import io.opentelemetry.api.trace.StatusCode; import io.opentelemetry.context.Context; import io.opentelemetry.context.ContextKey; import io.opentelemetry.context.Scope; @@ -152,6 +154,23 @@ void inlineNonfatalMdcRestorationPreservesSelectedOutcome(String outcome, boolea runInlineMdcFailure(outcome, ambient, new IllegalStateException("restore failed"), null); } + @ParameterizedTest + @CsvSource({ + "SUCCEEDED,false,assertion", "SUCCEEDED,true,assertion", + "PENDING,false,assertion", "PENDING,true,assertion", + "RETRYING,false,assertion", "RETRYING,true,assertion", + "SUCCEEDED,false,linkage", "SUCCEEDED,true,linkage", + "PENDING,false,linkage", "PENDING,true,linkage", + "RETRYING,false,linkage", "RETRYING,true,linkage" + }) + void inlineNonfatalErrorRestorationPreservesSelectedOutcome(String outcome, boolean ambient, String kind) + throws Exception { + Error failure = kind.equals("linkage") + ? new NoClassDefFoundError("restore failed") + : new AssertionError("restore failed"); + runInlineMdcFailure(outcome, ambient, failure, null); + } + @SuppressWarnings("removal") @ParameterizedTest @CsvSource({"false,false", "false,true", "true,false", "true,true"}) @@ -312,6 +331,89 @@ public void onInvocationEnd(InvocationEndInfo info) { .count()); } + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void failureSerializationRemainsRetryingUntilFailureCanBeDelivered(boolean executionView) { + var exporter = InMemorySpanExporter.create(); + var provider = SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)); + var options = OtelPluginConfig.builder() + .contextExtractor(() -> null) + .enableMdc(false) + .build(); + DurableExecutionPlugin otel = executionView + ? new ExecutionOtelPlugin(provider, options) + : new InvocationOtelPlugin(provider, options); + var ends = new CopyOnWriteArrayList(); + var failDelivery = new AtomicBoolean(true); + var deliveryFailure = new IllegalStateException("cannot serialize failure yet"); + var bodyFailure = new IllegalArgumentException("handler failed"); + var completedSideEffects = new AtomicInteger(); + var serDes = new SerDes() { + private final JacksonSerDes delegate = new JacksonSerDes(); + + @Override + public String serialize(Object value) { + if (value == bodyFailure && failDelivery.compareAndSet(true, false)) throw deliveryFailure; + return delegate.serialize(value); + } + + @Override + public T deserialize(String value, TypeToken type) { + return delegate.deserialize(value, type); + } + }; + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + context.step("saved-before-failure", String.class, step -> { + completedSideEffects.incrementAndGet(); + return "checkpointed"; + }); + throw bodyFailure; + }, + DurableConfig.builder() + .withSerDes(serDes) + .withPlugins(otel, new DurableExecutionPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + ends.add(info); + } + }) + .build()); + assertSame(deliveryFailure, assertThrows(IllegalStateException.class, () -> runner.run("input"))); + assertEquals(1, ends.size()); + assertEquals(InvocationStatus.RETRYING, ends.get(0).invocationStatus()); + assertSame(deliveryFailure, ends.get(0).executionError()); + assertEquals( + 0, + exporter.getFinishedSpanItems().stream() + .filter(span -> span.getName().equals("Workflow")) + .count()); + var firstInvocation = exporter.getFinishedSpanItems().stream() + .filter(span -> span.getName().equals("Invocation")) + .findFirst() + .orElseThrow(); + assertEquals("RETRYING", firstInvocation.getAttributes().get(DURABLE_INVOCATION_STATUS)); + assertEquals(StatusCode.UNSET, firstInvocation.getStatus().getStatusCode()); + + var result = runner.run("input"); + assertEquals(ExecutionStatus.FAILED, result.getStatus()); + assertEquals( + bodyFailure.getClass().getName(), + result.getError().orElseThrow().errorType()); + assertEquals(bodyFailure.getMessage(), result.getError().orElseThrow().errorMessage()); + assertEquals( + List.of(InvocationStatus.RETRYING, InvocationStatus.FAILED), + ends.stream().map(InvocationEndInfo::invocationStatus).toList()); + assertSame(bodyFailure, ends.get(1).executionError()); + assertEquals(1, completedSideEffects.get(), "the completed step body must not run again after delivery retry"); + assertEquals( + 1, + exporter.getFinishedSpanItems().stream() + .filter(span -> span.getName().equals("Workflow")) + .count()); + } + @Test void partialStartStillUnwindsAllConfiguredEndHooksAndRestoresBothThreads() throws Exception { var key = ContextKey.named("partial-start-context"); 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 e6438d643..49ecc86e8 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 @@ -173,13 +173,7 @@ private DurableExecutionOutput finishInvocation(O value, Throwable failure) { var payload = config.getSerDes().serialize(value); output = DurableExecutionOutput.success(handleLargePayload(manager, payload)); } catch (Throwable deliveryFailure) { - // Result serialization/checkpointing also belongs to this invocation. Close plugin resources even - // when delivery fails, then retain the original exception for the Lambda caller. - var cause = ExceptionHelper.unwrapCompletableFuture(deliveryFailure); - // No terminal output is delivered: this exception escapes for a Lambda invocation retry. - fireOnInvocationEnd(InvocationStatus.RETRYING, cause, null); - ExceptionHelper.sneakyThrow(deliveryFailure); - return null; + return failDelivery(deliveryFailure); } fireOnInvocationEnd(InvocationStatus.SUCCEEDED, null, value); return output; @@ -187,13 +181,31 @@ private DurableExecutionOutput finishInvocation(O value, Throwable failure) { private DurableExecutionOutput finishFailure(Throwable cause) { var status = failureStatus(cause); + if (status == InvocationStatus.FAILED) return finishTerminalFailure(cause); fireOnInvocationEnd(status, status == InvocationStatus.PENDING ? null : cause, null); if (status == InvocationStatus.PENDING) return DurableExecutionOutput.pending(); - if (status == InvocationStatus.RETRYING) { - ExceptionHelper.sneakyThrow(cause); - return null; + ExceptionHelper.sneakyThrow(cause); + return null; + } + + private DurableExecutionOutput finishTerminalFailure(Throwable cause) { + DurableExecutionOutput output; + try { + output = DurableExecutionOutput.failure(buildErrorObject(cause, config.getSerDes())); + } catch (Throwable deliveryFailure) { + return failDelivery(deliveryFailure); } - return DurableExecutionOutput.failure(buildErrorObject(cause, config.getSerDes())); + fireOnInvocationEnd(InvocationStatus.FAILED, cause, null); + return output; + } + + private DurableExecutionOutput failDelivery(Throwable deliveryFailure) { + // Serialization/checkpointing did not produce a terminal response. Close plugin resources with RETRYING + // before preserving the original delivery failure for the Lambda caller. + var cause = ExceptionHelper.unwrapCompletableFuture(deliveryFailure); + fireOnInvocationEnd(InvocationStatus.RETRYING, cause, null); + ExceptionHelper.sneakyThrow(deliveryFailure); + return null; } private void fireOnInvocationEnd(InvocationStatus status, Throwable error, Object result) { @@ -271,7 +283,7 @@ private static CompletableFuture supplyHandler( }; try { executor.execute(work); - } catch (RuntimeException dispatchFailure) { + } catch (Throwable dispatchFailure) { if (!result.isDone()) throw dispatchFailure; // Inline execution can throw from MDC restoration after settling the selected outcome. Preserve that // outcome for an ordinary cleanup failure, while a JVM-fatal failure still reaches the caller. From 1dc8136ac4a76ce56dd7027f0ea9214799c065d6 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 02:16:06 +0000 Subject: [PATCH 27/48] docs: define the SDK invocation-end response boundary --- otel-plugin/README.md | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/otel-plugin/README.md b/otel-plugin/README.md index aa47d6510..b3f04a83a 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -31,10 +31,13 @@ unless a later `VirtualMachineError` or `ThreadDeath` takes precedence over a no end-hook Errors are retained as suppressed failures. Scope cleanup preserves finalization failures using the same JVM-fatal precedence, so an ordinary cleanup exception cannot hide an earlier Error. Invocations without plugins also wait for handler cleanup before returning the selected suspension or retry outcome. -Result and failure-response serialization finish before terminal invocation-end notification. Delivery failures that -escape for a Lambda retry report `RETRYING`, so they do not prematurely end the Workflow span. Non-JVM-fatal -MDC-restoration failures on an inline executor, including `AssertionError` and `LinkageError`, do not replace an -already selected outcome. +SDK output preparation, including customer `SerDes` calls and durable large-result checkpointing, finishes before +terminal invocation-end notification. Failures in this preparation report `RETRYING` instead of ending the Workflow +span. End describes the SDK outcome at that point, not acknowledgment of a response by the Lambda service. +Caller-side execution-manager cleanup, response-envelope encoding and output-stream writes follow End; runtime +response transport follows the handler return. Failures at those later boundaries still propagate, without a second +End dispatch or changing its already reported outcome. Non-JVM-fatal MDC-restoration failures on an inline executor, +including `AssertionError` and `LinkageError`, likewise do not replace an already selected outcome. Before invocation startup, a JVM-fatal error from MDC capture (direct or inside a standard transport wrapper) completes the observation future exceptionally with that same fatal before escaping the handler worker. No start, body, or end hook runs, and no durable `FAILED` response is produced for that fatal. Ordinary initialization errors From 7c09c60dd61c2c06154ab307395d123cf71dfbf5 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 02:54:24 +0000 Subject: [PATCH 28/48] fix: retain ready condition work across checkpoint responses Reuse the readiness repair from cca1b684 and terminal-poll safety follow-up5c123250, preserving their operation-worker publication ownership. Keep terminal replay/get and invocation lifecycle behavior unchanged. Record payload-free cached history metadata if the concurrent cloud example fails again. --- .../examples/CloudBasedIntegrationTest.java | 40 +- ...tForConditionReadinessIntegrationTest.java | 445 ++++++++++++++++++ .../operation/WaitForConditionOperation.java | 129 +++-- .../WaitForConditionTerminalPollingTest.java | 119 +++++ 4 files changed, 687 insertions(+), 46 deletions(-) create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionReadinessIntegrationTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionTerminalPollingTest.java diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/CloudBasedIntegrationTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/CloudBasedIntegrationTest.java index 7490961a5..58b5cf8bd 100644 --- a/examples/src/test/java/software/amazon/lambda/durable/examples/CloudBasedIntegrationTest.java +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/CloudBasedIntegrationTest.java @@ -21,6 +21,7 @@ import software.amazon.awssdk.regions.Region; import software.amazon.awssdk.services.lambda.LambdaClient; import software.amazon.awssdk.services.lambda.model.ErrorObject; +import software.amazon.awssdk.services.lambda.model.Event; import software.amazon.awssdk.services.lambda.model.OperationStatus; import software.amazon.awssdk.services.sts.StsClient; import software.amazon.lambda.durable.TypeToken; @@ -818,7 +819,20 @@ void testConcurrentWaitForConditionExample() { lambdaClient); var result = runner.run(new ConcurrentWaitForConditionExample.Input(3, 100, 50)); - assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + try { + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + } catch (RuntimeException | AssertionError failure) { + try { + // Preserve the already-fetched service timeline without payloads, callback IDs, or checkpoint tokens. + var history = result.getHistoryEvents().stream() + .map(CloudBasedIntegrationTest::historyMetadata) + .toList(); + System.err.println("Concurrent wait-for-condition history: " + new JacksonSerDes().serialize(history)); + } catch (RuntimeException diagnosticsFailure) { + failure.addSuppressed(diagnosticsFailure); + } + throw failure; + } // Verify each operation finished with 3 attempts var allOperationsOutput = result.getResult(); @@ -840,6 +854,30 @@ void testConcurrentWaitForConditionExample() { } } + private static Map historyMetadata(Event event) { + var row = new HashMap(); + row.put("eventId", event.eventId()); + row.put("type", event.eventTypeAsString()); + row.put("operationId", event.id()); + row.put("parentId", event.parentId()); + row.put("name", event.name()); + row.put("at", String.valueOf(event.eventTimestamp())); + var retries = event.stepSucceededDetails() != null + ? event.stepSucceededDetails().retryDetails() + : event.stepFailedDetails() != null ? event.stepFailedDetails().retryDetails() : null; + if (retries != null) { + row.put("attempt", retries.currentAttempt()); + row.put("nextAttemptDelaySeconds", retries.nextAttemptDelaySeconds()); + } + var invocation = event.invocationCompletedDetails(); + if (invocation != null) { + row.put("invocationStart", String.valueOf(invocation.startTimestamp())); + row.put("invocationEnd", String.valueOf(invocation.endTimestamp())); + row.put("invocationFailed", invocation.error() != null); + } + return row; + } + @Test void testPluginExample() { var runner = diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionReadinessIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionReadinessIntegrationTest.java new file mode 100644 index 000000000..111c3c984 --- /dev/null +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionReadinessIntegrationTest.java @@ -0,0 +1,445 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.*; + +import java.lang.management.ManagementFactory; +import java.lang.reflect.Method; +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Deque; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import java.util.concurrent.locks.LockSupport; +import java.util.function.IntFunction; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.ValueSource; +import org.slf4j.MDC; +import org.slf4j.spi.MDCAdapter; +import software.amazon.awssdk.services.lambda.model.*; +import software.amazon.lambda.durable.config.WaitForConditionConfig; +import software.amazon.lambda.durable.execution.DurableExecutor; +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.WaitForConditionResult; +import software.amazon.lambda.durable.serde.JacksonSerDes; +import software.amazon.lambda.durable.serde.SerDes; +import software.amazon.lambda.durable.testing.local.LocalMemoryExecutionClient; + +class WaitForConditionReadinessIntegrationTest { + @ParameterizedTest + @ValueSource(ints = {2, 25}) + void readyInRetryResponseContinuesWithoutSuspensionAndReplaysStoredResult(int threshold) throws Exception { + try (var run = new Run( + new ImmediateReadyClient(), + Executors.newCachedThreadPool(), + value -> value >= threshold + ? WaitForConditionResult.stopPolling(value) + : WaitForConditionResult.continuePolling(value + 1))) { + var first = run.start().get(5, TimeUnit.SECONDS); + assertEquals(ExecutionStatus.SUCCEEDED, first.status(), "READY work must not produce an invalid PENDING"); + assertEquals("\"" + threshold + "\"", first.result()); + assertEquals(threshold, run.checks.get()); + var updates = run.client.getOperationUpdates().size(); + var replay = run.start().get(5, TimeUnit.SECONDS); + assertEquals(first, replay); + assertEquals(threshold, run.checks.get(), "Completed checks must not be repeated on replay"); + assertEquals(updates, run.client.getOperationUpdates().size()); + } + } + + @Test + void immediateReadyFailureRemainsCheckpointedOnReplay() throws Exception { + try (var run = new Run(new ImmediateReadyClient(), Executors.newCachedThreadPool(), value -> { + if (value == 2) throw new IllegalStateException("predicate failure"); + return WaitForConditionResult.continuePolling(value + 1); + })) { + var first = run.start().get(5, TimeUnit.SECONDS); + assertEquals(ExecutionStatus.FAILED, first.status()); + assertEquals("predicate failure", first.error().errorMessage()); + var updates = run.client.getOperationUpdates().size(); + var replay = run.start().get(5, TimeUnit.SECONDS); + assertEquals(ExecutionStatus.FAILED, replay.status()); + assertEquals(first.error().errorType(), replay.error().errorType()); + assertEquals(first.error().errorMessage(), replay.error().errorMessage()); + assertEquals(2, run.checks.get()); + assertEquals(updates, run.client.getOperationUpdates().size()); + } + } + + @Test + void pendingRetryStillSuspendsAndResumesWhenBackendBecomesReady() throws Exception { + try (var run = new Run( + new LocalMemoryExecutionClient(), + Executors.newCachedThreadPool(), + value -> value == 2 + ? WaitForConditionResult.stopPolling(value) + : WaitForConditionResult.continuePolling(value + 1))) { + var first = run.start().get(5, TimeUnit.SECONDS); + assertEquals(ExecutionStatus.PENDING, first.status()); + assertEquals( + OperationStatus.PENDING, + run.client.getOperationByName("condition").status()); + assertEquals(1, run.checks.get()); + assertTrue(run.client.advanceTime()); + var resumed = run.start().get(5, TimeUnit.SECONDS); + assertEquals(ExecutionStatus.SUCCEEDED, resumed.status()); + assertEquals("\"2\"", resumed.result()); + assertEquals(2, run.checks.get()); + assertEquals(resumed, run.start().get(5, TimeUnit.SECONDS)); + assertEquals(2, run.checks.get()); + } + } + + @ParameterizedTest + @CsvSource({"true,false", "false,true", "true,true"}) + void asynchronousReadyHandoffDoesNotDeadlockOrLoseItsActivityLease( + boolean holdWorkerExit, boolean holdExecutorReturn) throws Exception { + var gate = new HandoffGate(holdWorkerExit, holdExecutorReturn); + var client = new DelayedReadyClient(gate); + try (var mdc = new MdcClearGate(gate); + var run = new Run(client, gate, value -> { + if (value == 1) { + gate.firstWorker.set(Thread.currentThread()); + return WaitForConditionResult.continuePolling(2); + } + gate.nextCheckEntered.countDown(); + await(gate.releaseNextCheck); + return WaitForConditionResult.stopPolling(value); + })) { + try { + var result = run.start(); + await(gate.pollEntered); + if (holdWorkerExit) await(gate.workerExitEntered); + else await(gate.executorReturnEntered); + assertTrue(client.advanceTime()); + gate.releaseReadyResponse.countDown(); + await(gate.readyResponseReturned); + awaitCheckpointHandoff(gate.pollThread.get()); + assertThrows(TimeoutException.class, () -> result.get(100, TimeUnit.MILLISECONDS)); + assertEquals(1, run.checks.get()); + gate.releaseWorkerExit.countDown(); + if (holdExecutorReturn) await(gate.executorReturnEntered); + gate.releaseExecutorReturn.countDown(); + await(gate.nextCheckEntered); + assertThrows( + TimeoutException.class, + () -> result.get(100, TimeUnit.MILLISECONDS), + "The next check must stay active after the old worker deregisters and the checkpoint returns"); + gate.releaseNextCheck.countDown(); + assertEquals( + ExecutionStatus.SUCCEEDED, + result.get(5, TimeUnit.SECONDS).status()); + assertEquals(2, run.checks.get()); + assertEquals(4, run.stateReads.get(), "In-process continuation must reuse normalized state"); + assertEquals( + ExecutionStatus.SUCCEEDED, + run.start().get(5, TimeUnit.SECONDS).status()); + assertEquals(2, run.checks.get()); + assertEquals(5, run.stateReads.get(), "Completed replay only decodes the stored result"); + } finally { + gate.releaseAll(); + } + } + } + + private static void awaitCheckpointHandoff(Thread checkpoint) { + var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(3); + while (System.nanoTime() < deadline) { + var stack = checkpoint.getStackTrace(); + if (checkpoint.getState() == Thread.State.WAITING + && Arrays.stream(stack) + .anyMatch(frame -> frame.getClassName().endsWith("WaitForConditionOperation")) + && Arrays.stream(stack) + .anyMatch(frame -> frame.getClassName().equals(CompletableFuture.class.getName()) + && frame.getMethodName().equals("join"))) { + var bean = ManagementFactory.getThreadMXBean(); + if (bean.isObjectMonitorUsageSupported()) { + var info = bean.getThreadInfo(new long[] {checkpoint.getId()}, true, true)[0]; + System.out.println( + "READY_HANDOFF checkpoint monitors=" + Arrays.toString(info.getLockedMonitors())); + } + return; + } + LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1)); + } + fail("The READY callback did not reach the controlled old-worker handoff window"); + } + + private static void await(CountDownLatch latch) { + try { + assertTrue(latch.await(5, TimeUnit.SECONDS), "Controlled readiness gate was not released"); + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + throw new AssertionError(error); + } + } + + private static final class ImmediateReadyClient extends LocalMemoryExecutionClient { + @Override + public CheckpointDurableExecutionResponse checkpoint(String arn, String token, List updates) { + var applied = super.checkpoint(arn, token, updates); + // Advance the test backend's virtual time before materializing the response. The configured + // one-second delay is valid under Lambda's minimum; no zero-delay service request is assumed. + advanceTime(); + var ready = super.checkpoint(arn, applied.checkpointToken(), List.of()); + var states = new LinkedHashMap(); + applied.newExecutionState().operations().forEach(op -> states.put(op.id(), op)); + ready.newExecutionState().operations().forEach(op -> states.put(op.id(), op)); + return ready.toBuilder() + .newExecutionState(CheckpointUpdatedExecutionState.builder() + .operations(states.values()) + .build()) + .build(); + } + } + + private static final class DelayedReadyClient extends LocalMemoryExecutionClient { + private final HandoffGate gate; + private final AtomicBoolean firstPoll = new AtomicBoolean(true); + + private DelayedReadyClient(HandoffGate gate) { + this.gate = gate; + } + + @Override + public CheckpointDurableExecutionResponse checkpoint(String arn, String token, List updates) { + if (updates.isEmpty() && firstPoll.compareAndSet(true, false)) { + gate.pollThread.set(Thread.currentThread()); + gate.pollEntered.countDown(); + await(gate.releaseReadyResponse); + var response = super.checkpoint(arn, token, updates); + gate.readyResponseReturned.countDown(); + return response; + } + return super.checkpoint(arn, token, updates); + } + } + + private static final class HandoffGate extends AbstractExecutorService { + private final ExecutorService delegate = Executors.newCachedThreadPool(); + private final AtomicInteger submitted = new AtomicInteger(); + private final AtomicBoolean exitHeld = new AtomicBoolean(); + private final AtomicReference firstWorker = new AtomicReference<>(); + private final AtomicReference pollThread = new AtomicReference<>(); + private final CountDownLatch pollEntered = new CountDownLatch(1); + private final CountDownLatch releaseReadyResponse = new CountDownLatch(1); + private final CountDownLatch readyResponseReturned = new CountDownLatch(1); + private final CountDownLatch workerExitEntered = new CountDownLatch(1); + private final CountDownLatch releaseWorkerExit = new CountDownLatch(1); + private final CountDownLatch executorReturnEntered = new CountDownLatch(1); + private final CountDownLatch releaseExecutorReturn = new CountDownLatch(1); + private final CountDownLatch nextCheckEntered = new CountDownLatch(1); + private final CountDownLatch releaseNextCheck = new CountDownLatch(1); + private final boolean holdWorkerExit; + private final boolean holdExecutorReturn; + + private HandoffGate(boolean holdWorkerExit, boolean holdExecutorReturn) { + this.holdWorkerExit = holdWorkerExit; + this.holdExecutorReturn = holdExecutorReturn; + } + + @Override + public void execute(Runnable task) { + var number = submitted.incrementAndGet(); + var finished = new CountDownLatch(1); + delegate.execute(() -> { + try { + task.run(); + } finally { + finished.countDown(); + } + }); + // Completing work before execute() returns is permitted by Executor's contract. + if (holdExecutorReturn && number == 2) { + await(finished); + executorReturnEntered.countDown(); + await(releaseExecutorReturn); + } + } + + private void releaseAll() { + releaseReadyResponse.countDown(); + releaseWorkerExit.countDown(); + releaseExecutorReturn.countDown(); + releaseNextCheck.countDown(); + } + + @Override + public void shutdown() { + delegate.shutdown(); + } + + @Override + public List shutdownNow() { + releaseAll(); + return delegate.shutdownNow(); + } + + @Override + public boolean isShutdown() { + return delegate.isShutdown(); + } + + @Override + public boolean isTerminated() { + return delegate.isTerminated(); + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { + return delegate.awaitTermination(timeout, unit); + } + } + + private static final class MdcClearGate implements AutoCloseable { + private final MDCAdapter previous = MDC.getMDCAdapter(); + private final Method setter = MDC.class.getDeclaredMethod("setMDCAdapter", MDCAdapter.class); + + private MdcClearGate(HandoffGate gate) throws Exception { + setter.setAccessible(true); + setter.invoke(null, new MDCAdapter() { + public void clear() { + if (gate.holdWorkerExit + && Thread.currentThread() == gate.firstWorker.get() + && gate.exitHeld.compareAndSet(false, true)) { + gate.workerExitEntered.countDown(); + await(gate.releaseWorkerExit); + } + previous.clear(); + } + + public void put(String key, String value) { + previous.put(key, value); + } + + public String get(String key) { + return previous.get(key); + } + + public void remove(String key) { + previous.remove(key); + } + + public Map getCopyOfContextMap() { + return previous.getCopyOfContextMap(); + } + + public void setContextMap(Map values) { + previous.setContextMap(values); + } + + public void pushByKey(String key, String value) { + previous.pushByKey(key, value); + } + + public String popByKey(String key) { + return previous.popByKey(key); + } + + public Deque getCopyOfDequeByKey(String key) { + return previous.getCopyOfDequeByKey(key); + } + + public void clearDequeByKey(String key) { + previous.clearDequeByKey(key); + } + }); + } + + public void close() throws Exception { + setter.invoke(null, previous); + } + } + + private static final class Run implements AutoCloseable { + private final LocalMemoryExecutionClient client; + private final ExecutorService workers; + private final ExecutorService caller = Executors.newSingleThreadExecutor(); + private final AtomicInteger checks = new AtomicInteger(); + private final AtomicInteger stateReads = new AtomicInteger(); + private final SerDes stateSerDes = new SerDes() { + private final JacksonSerDes delegate = new JacksonSerDes(); + + public String serialize(Object value) { + return delegate.serialize(value); + } + + public T deserialize(String value, TypeToken type) { + stateReads.incrementAndGet(); + return delegate.deserialize(value, type); + } + }; + private final IntFunction> check; + private final DurableConfig config; + private final WaitForConditionConfig wait = WaitForConditionConfig.builder() + .initialState(1) + .serDes(stateSerDes) + .waitStrategy((value, attempt) -> Duration.ofSeconds(1)) + .build(); + + private Run( + LocalMemoryExecutionClient client, + ExecutorService workers, + IntFunction> check) { + this.client = client; + this.workers = workers; + this.check = check; + config = DurableConfig.builder() + .withExecutorService(workers) + .withDurableExecutionClient(client) + .withCheckpointDelay(Duration.ZERO) + .withPollingStrategy(attempt -> Duration.ZERO) + .build(); + } + + private Future start() { + var execution = Operation.builder() + .id("execution") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails( + ExecutionDetails.builder().inputPayload("\"input\"").build()) + .build(); + var history = new ArrayList<>(client.getAllOperations()); + history.add(0, execution); + var input = new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/ready/execution", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(history) + .build()); + return caller.submit(() -> DurableExecutor.execute( + input, + null, + TypeToken.get(String.class), + (value, ctx) -> String.valueOf(ctx.waitForCondition( + "condition", + Integer.class, + (state, step) -> { + checks.incrementAndGet(); + return check.apply(state); + }, + wait)), + config)); + } + + @Override + public void close() throws Exception { + caller.shutdownNow(); + workers.shutdownNow(); + assertTrue(caller.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS)); + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java index 72eb65b8e..afff3deb8 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java @@ -4,6 +4,8 @@ import java.time.Duration; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; import software.amazon.awssdk.services.lambda.model.Operation; import software.amazon.awssdk.services.lambda.model.OperationAction; @@ -18,6 +20,7 @@ import software.amazon.lambda.durable.exception.DurableOperationException; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.exception.WaitForConditionFailedException; +import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.execution.ThreadType; import software.amazon.lambda.durable.logging.DurableLogger; @@ -108,68 +111,104 @@ private void resumeCheckLoop(Operation existing) { } private CompletableFuture pollReadyAndResumeCheckLoop(Operation existing) { - return pollForOperationUpdates() - .thenCompose(op -> op.status() == OperationStatus.READY - ? CompletableFuture.completedFuture(op) - : pollForOperationUpdates()) - .thenAccept(this::resumeCheckLoop); + return pollUntilReady().thenAccept(op -> { + if (!isOperationCompleted() && op.status() == OperationStatus.READY) resumeCheckLoop(op); + }); + } + + private CompletableFuture pollUntilReady() { + var known = getOperation(); + if (isReadyOrTerminal(known)) return CompletableFuture.completedFuture(known); + var update = pollForOperationUpdates(); + // Register before re-reading: another checkpoint may already have delivered READY before this poll existed. + known = getOperation(); + if (isReadyOrTerminal(known)) { + update.complete(known); + return CompletableFuture.completedFuture(known); + } + return update.thenCompose( + op -> isReadyOrTerminal(op) ? CompletableFuture.completedFuture(op) : pollUntilReady()); + } + + private boolean isReadyOrTerminal(Operation operation) { + if (operation == null) return false; + var status = operation.status(); + if (status == null || status == OperationStatus.UNKNOWN_TO_SDK_VERSION) { + throw terminateExecutionWithIllegalDurableOperationException( + "Unexpected waitForCondition status: " + operation.statusAsString()); + } + return status == OperationStatus.READY || ExecutionManager.isTerminalStatus(status); } private void executeCheckLogic(T currentState, int attempt) { - Runnable userHandler = () -> { + var publishedWorker = new CompletableFuture>(); + runUserHandler(() -> runCheckLoop(currentState, attempt, publishedWorker), ThreadType.STEP); + publishedWorker.complete(getRunningUserHandler()); + } + + private void runCheckLoop(T currentState, int attempt, CompletableFuture> publishedWorker) { + while (!isOperationCompleted()) { var stepContext = getContext().createStepContext(getOperationId(), getName(), attempt); BaseContextImpl.setCurrentContext(stepContext); try (var ignored = DurableLogger.attachContext()) { try { - // Checkpoint START if not already started var existing = getOperation(); if (existing == null || existing.status() != OperationStatus.STARTED) { - var startUpdate = OperationUpdate.builder().action(OperationAction.START); - sendOperationUpdateAsync(startUpdate); + sendOperationUpdateAsync(OperationUpdate.builder().action(OperationAction.START)); } - - // Execute check function inside the plugin hook boundary so a failure is reported - // through onUserFunctionEnd; checkpoint/poll handling stays outside the boundary. - WaitForConditionResult result = - runUserFunction(attempt, () -> checkFunc.apply(currentState, stepContext)); - - // Normalize the value through SerDes so first execution matches replay. + var stateForCheck = currentState; + var result = runUserFunction(attempt, () -> checkFunc.apply(stateForCheck, stepContext)); var serializedState = serializeAndDeserializeResult(result.value()); - T deserializedValue = serializedState.deserialized(); - + var deserializedValue = serializedState.deserialized(); if (result.isDone()) { - // Condition met — checkpoint SUCCEED - var successUpdate = OperationUpdate.builder() + sendOperationUpdate(OperationUpdate.builder() .action(OperationAction.SUCCEED) - .payload(serializedState.serialized()); - sendOperationUpdate(successUpdate); - } else { - // Compute delay from strategy - Duration delay = config.waitStrategy().evaluate(deserializedValue, attempt); - - // Checkpoint RETRY with delay - var retryUpdate = OperationUpdate.builder() - .action(OperationAction.RETRY) - .payload(serializedState.serialized()) - .stepOptions(StepOptions.builder() - .nextAttemptDelaySeconds(Math.toIntExact(delay.toSeconds())) - .build()); - sendOperationUpdate(retryUpdate); - - // Poll for READY, then continue the loop - pollForOperationUpdates() - .thenCompose(op -> op.status() == OperationStatus.READY - ? CompletableFuture.completedFuture(op) - : pollForOperationUpdates()) - .thenRun(() -> executeCheckLogic(deserializedValue, attempt + 1)); + .payload(serializedState.serialized())); + return; } - } catch (Throwable e) { - handleCheckFailure(e); + Duration delay = config.waitStrategy().evaluate(deserializedValue, attempt); + sendOperationUpdate(OperationUpdate.builder() + .action(OperationAction.RETRY) + .payload(serializedState.serialized()) + .stepOptions(StepOptions.builder() + .nextAttemptDelaySeconds(Math.toIntExact(delay.toSeconds())) + .build())); + var inlineReady = + continueInlineOrAfterCurrentWorker(publishedWorker, deserializedValue, attempt + 1); + if (inlineReady == null || inlineReady.status() != OperationStatus.READY) return; + // READY can be present in the RETRY response. Keep this worker active and continue without + // recursively starting an overlapping handler or suspending executable work. + currentState = deserializedValue; + attempt++; + } catch (Throwable failure) { + handleCheckFailure(failure); + return; } } - }; + } + } - runUserHandler(userHandler, ThreadType.STEP); + private Operation continueInlineOrAfterCurrentWorker( + CompletableFuture> publishedWorker, T nextState, int nextAttempt) { + var owner = Thread.currentThread(); + var inline = new AtomicReference(); + var acceptingInline = new AtomicBoolean(true); + try { + pollUntilReady().thenAccept(op -> { + if (Thread.currentThread() == owner && acceptingInline.get()) { + inline.set(op); + } else { + // A checkpoint callback can observe READY before this attempt's worker exits. Its checkpoint + // processing lease stays active through this handoff, preventing a false quiescence window. + publishedWorker.join().join(); + if (!isOperationCompleted() && op.status() == OperationStatus.READY) + executeCheckLogic(nextState, nextAttempt); + } + }); + } finally { + acceptingInline.set(false); + } + return inline.get(); } private void handleCheckFailure(Throwable exception) { diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionTerminalPollingTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionTerminalPollingTest.java new file mode 100644 index 000000000..ecac9a10b --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionTerminalPollingTest.java @@ -0,0 +1,119 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.operation; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.*; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.NullAndEmptySource; +import org.junit.jupiter.params.provider.ValueSource; +import software.amazon.awssdk.services.lambda.model.Operation; +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.config.WaitForConditionConfig; +import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.exception.IllegalDurableOperationException; +import software.amazon.lambda.durable.execution.ExecutionManager; +import software.amazon.lambda.durable.model.OperationIdentifier; +import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.model.WaitForConditionResult; + +class WaitForConditionTerminalPollingTest { + @ParameterizedTest + @CsvSource({ + "CANCELLED,false", + "CANCELLED,true", + "TIMED_OUT,false", + "TIMED_OUT,true", + "STOPPED,false", + "STOPPED,true" + }) + void terminalObservationNeverRegistersAnotherPoll(String status, boolean alreadyKnown) throws Exception { + var fixture = new Fixture(alreadyKnown ? status : "PENDING"); + var result = fixture.poll(); + if (!alreadyKnown) fixture.deliver(status); + assertEquals(status, result.get(1, TimeUnit.SECONDS).statusAsString()); + assertEquals(alreadyKnown ? 0 : 1, fixture.polls.size()); + } + + @ParameterizedTest + @ValueSource(strings = {"NEW_UNKNOWN_STATUS"}) + @NullAndEmptySource + void malformedStatusTerminatesInsteadOfPollingForever(String status) throws Exception { + var fixture = new Fixture("PENDING"); + var result = fixture.poll(); + fixture.deliver(status); + var failure = assertThrows(ExecutionException.class, () -> result.get(1, TimeUnit.SECONDS)); + assertInstanceOf(IllegalDurableOperationException.class, failure.getCause()); + verify(fixture.manager).terminateExecution((IllegalDurableOperationException) failure.getCause()); + assertEquals(1, fixture.polls.size()); + } + + @Test + void nonterminalUpdatesStillWaitForReady() throws Exception { + var fixture = new Fixture("PENDING"); + var result = fixture.poll(); + fixture.deliver("STARTED"); + assertFalse(result.isDone()); + fixture.deliver("PENDING"); + assertFalse(result.isDone()); + fixture.deliver("READY"); + assertEquals("READY", result.get(1, TimeUnit.SECONDS).statusAsString()); + assertEquals(3, fixture.polls.size()); + } + + private static final class Fixture { + final ExecutionManager manager = mock(ExecutionManager.class); + final AtomicReference known; + final List> polls = new ArrayList<>(); + final WaitForConditionOperation operation; + + Fixture(String status) { + known = new AtomicReference<>(snapshot(status)); + var context = mock(DurableContextImpl.class); + when(context.getExecutionManager()).thenReturn(manager); + when(context.getDurableConfig()).thenReturn(DurableConfig.builder().build()); + when(manager.getOperationAndUpdateReplayState("condition")).thenAnswer(call -> known.get()); + when(manager.pollForOperationUpdates("condition")).thenAnswer(call -> { + var poll = new CompletableFuture(); + polls.add(poll); + return poll; + }); + operation = new WaitForConditionOperation<>( + OperationIdentifier.of("condition", "condition", OperationSubType.WAIT_FOR_CONDITION), + (state, contextIgnored) -> WaitForConditionResult.stopPolling(state), + TypeToken.get(Integer.class), + WaitForConditionConfig.builder().initialState(1).build(), + context); + } + + @SuppressWarnings("unchecked") + CompletableFuture poll() throws Exception { + var method = WaitForConditionOperation.class.getDeclaredMethod("pollUntilReady"); + method.setAccessible(true); + return (CompletableFuture) method.invoke(operation); + } + + void deliver(String status) { + var update = snapshot(status); + known.set(update); + polls.get(polls.size() - 1).complete(update); + } + + private static Operation snapshot(String status) { + return Operation.builder() + .id("condition") + .type(OperationType.STEP) + .status(status) + .build(); + } + } +} From 7a4ff80af1ba53e760288135e4682053a03af792 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 03:27:25 +0000 Subject: [PATCH 29/48] fix: inspect MDC initialization causes without cycling --- otel-plugin/README.md | 6 + .../otel/MdcFatalInitializationTest.java | 141 +++++++++++++++++- .../durable/execution/DurableExecutor.java | 45 +++++- 3 files changed, 184 insertions(+), 8 deletions(-) diff --git a/otel-plugin/README.md b/otel-plugin/README.md index b3f04a83a..b989b0dbf 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -43,6 +43,12 @@ completes the observation future exceptionally with that same fatal before escap body, or end hook runs, and no durable `FAILED` response is produced for that fatal. Ordinary initialization errors and handler/body failure classification retain their existing behavior. +SDK inspection of MDC-capture failures reads each visited standard transport cause once and detects identity cycles. +Cyclic, null, or unreadable leading `CompletionException` chains retain the original wrapper; ordinary initialization +still reports `FAILED` when its error response can be serialized. This provides cycle safety for finite cause graphs, +not a fixed depth or time limit. Arbitrary custom `getCause`, other `Throwable` accessors, and customer `SerDes` +behavior remain outside that guarantee. + This plugin version requires the core's `DurableExecutor.supportsSameThreadInvocationHooks()` capability, introduced in the 2.2.2 lifecycle contract (currently `2.2.2-SNAPSHOT`). Upgrade the core together with the plugin layer. Every plugin constructor checks this capability before building a tracer provider or activating context. A core without diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFatalInitializationTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFatalInitializationTest.java index 8f835bebf..615a4be86 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFatalInitializationTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFatalInitializationTest.java @@ -7,6 +7,7 @@ import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Proxy; import java.lang.reflect.UndeclaredThrowableException; +import java.util.ArrayList; import java.util.Arrays; import java.util.List; import java.util.concurrent.AbstractExecutorService; @@ -19,6 +20,7 @@ import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Supplier; import java.util.stream.Stream; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; @@ -26,10 +28,13 @@ import org.slf4j.MDC; import org.slf4j.spi.MDCAdapter; import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TypeToken; 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.serde.JacksonSerDes; +import software.amazon.lambda.durable.serde.SerDes; import software.amazon.lambda.durable.testing.LocalDurableTestRunner; class MdcFatalInitializationTest { @@ -63,6 +68,120 @@ void nonfatalInitializationKeepsItsDurableFailurePolicy(String mode, String kind runCapture(mode, failure, null); } + @ParameterizedTest + @MethodSource("diagnosticCases") + @SuppressWarnings("removal") + void captureDiagnosticsAreBoundedAndPreserveTheSelectedFailure(String mode, String kind) throws Exception { + var abort = new AtomicBoolean(); + var reads = new ArrayList(); + var leaf = new IllegalStateException("ordinary capture"); + Error fatal = kind.equals("getter-death") ? new ThreadDeath() : new InternalError("fatal capture diagnostic"); + Throwable failure; + Throwable expected; + if (kind.equals("self")) { + var link = new AtomicReference(); + failure = diagnostic(link::get, abort, reads); + link.set(failure); + expected = failure; + } else if (kind.equals("pair")) { + var link = new AtomicReference(); + var inner = diagnostic(link::get, abort, reads); + failure = diagnostic(() -> inner, abort, reads); + link.set(failure); + expected = failure; + } else if (kind.equals("null")) { + failure = diagnostic(() -> null, abort, reads); + expected = failure; + } else if (kind.equals("unreadable")) { + failure = diagnostic( + () -> { + throw new IllegalStateException("unreadable cause"); + }, + abort, + reads); + expected = failure; + } else if (kind.equals("nested")) { + var inner = diagnostic(() -> leaf, abort, reads); + failure = diagnostic(() -> inner, abort, reads); + expected = leaf; + } else if (kind.equals("non-completion")) { + var wrapper = new UndeclaredThrowableException(leaf, "ordinary capture"); + failure = diagnostic(() -> wrapper, abort, reads); + expected = wrapper; // Ordinary initialization only unwraps a leading CompletionException chain. + } else if (kind.equals("changing")) { + var calls = new AtomicInteger(); + failure = diagnostic( + () -> { + if (calls.incrementAndGet() == 1) return leaf; + throw fatal; + }, + abort, + reads); + expected = leaf; + } else { + failure = diagnostic( + () -> { + throw fatal; + }, + abort, + reads); + expected = failure; + } + var serialized = new AtomicReference(); + var safeSerDes = new SerDes() { + private final JacksonSerDes delegate = new JacksonSerDes(); + + @Override + public String serialize(Object value) { + if (value instanceof Throwable error) { + serialized.set(error); + // Observe identity without making a customer serializer traverse the malformed cause graph. + return "\"capture failure\""; + } + return delegate.serialize(value); + } + + @Override + public T deserialize(String value, TypeToken type) { + return delegate.deserialize(value, type); + } + }; + var getterFatal = kind.startsWith("getter-"); + runCapture(mode, failure, getterFatal ? fatal : null, expected, safeSerDes, () -> abort.set(true)); + if (getterFatal) assertNull(serialized.get()); + else assertSame(expected, serialized.get(), "initialization must preserve the selected failure identity"); + reads.forEach(count -> assertEquals(1, count.get(), "read each diagnostic cause once")); + } + + private static CompletionException diagnostic( + Supplier cause, AtomicBoolean abort, List reads) { + var count = new AtomicInteger(); + reads.add(count); + return new CompletionException("ordinary capture", null) { + @Override + public synchronized Throwable getCause() { + if (abort.get()) return null; // Bound negative-control cleanup even when old code loops forever. + count.incrementAndGet(); + return cause.get(); + } + }; + } + + private static Stream diagnosticCases() { + return Stream.of("direct", "async") + .flatMap(mode -> Stream.of( + "self", + "pair", + "null", + "unreadable", + "nested", + "non-completion", + "changing", + "getter-vm", + "getter-death") + .map(kind -> Arguments.of(mode, kind))); + } + private static Stream fatalCases() { return Stream.of("direct", "async", "precompleted") .flatMap(mode -> Stream.of(false, true) @@ -77,6 +196,17 @@ private static Stream ordinaryCases() { } private static void runCapture(String mode, Throwable failure, Error expectedFatal) throws Exception { + runCapture(mode, failure, expectedFatal, failure, null, () -> {}); + } + + private static void runCapture( + String mode, + Throwable failure, + Error expectedFatal, + Throwable expectedOrdinary, + SerDes overrideSerDes, + Runnable release) + throws Exception { var originalAdapter = MDC.getMDCAdapter(); var injected = new AtomicBoolean(); var starts = new AtomicInteger(); @@ -110,16 +240,16 @@ public void onInvocationEnd(InvocationEndInfo info) { ends.incrementAndGet(); } }; + var configuration = + DurableConfig.builder().withExecutorService(workers).withPlugins(plugin); + if (overrideSerDes != null) configuration.withSerDes(overrideSerDes); var runner = LocalDurableTestRunner.create( String.class, (input, context) -> { bodies.incrementAndGet(); return "unreachable"; }, - DurableConfig.builder() - .withExecutorService(workers) - .withPlugins(plugin) - .build()); + configuration.build()); var response = caller.submit(() -> { callerThread.set(Thread.currentThread()); return runner.run("input"); @@ -142,7 +272,7 @@ public void onInvocationEnd(InvocationEndInfo info) { var result = response.get(3, TimeUnit.SECONDS); assertEquals(ExecutionStatus.FAILED, result.getStatus()); assertEquals( - failure.getClass().getName(), + expectedOrdinary.getClass().getName(), result.getError().orElseThrow().errorType()); assertEquals("ordinary capture", result.getError().orElseThrow().errorMessage()); assertNull(workers.escape.get(), "ordinary initialization policy must not be broadened"); @@ -152,6 +282,7 @@ public void onInvocationEnd(InvocationEndInfo info) { assertEquals(0, bodies.get()); assertEquals(0, ends.get()); } finally { + release.run(); workers.shutdownNow(); caller.shutdownNow(); try { 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 49ecc86e8..a3803c16f 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 @@ -129,7 +129,7 @@ private CompletableFuture execute() { }; // Cleanup is awaited even without plugins. Keep that path free of plugin MDC handling. if (plugins.isEmpty()) return CompletableFuture.supplyAsync(task, config.getExecutorService()); - return supplyHandler(task, failure -> finishInvocation(null, failure), config.getExecutorService()); + return supplyHandler(task, this::finishFailure, config.getExecutorService()); } private O invokeHandler() { @@ -257,15 +257,17 @@ private static CompletableFuture supplyHandler( try { restore = restoreMdcOnClose(); } catch (Throwable failure) { + Throwable initializationCause; try { - rethrowLifecycleFatal(failure); + initializationCause = normalizeMdcInitializationFailure(failure); } catch (VirtualMachineError | ThreadDeath fatal) { // Wake the invocation caller with the original fatal before it escapes the actual worker. // Throwing first would leave the observation future incomplete on an asynchronous executor. result.completeExceptionally(fatal); throw fatal; } - Outcome.capture(() -> initializationFailure.apply(failure)).complete(result); + Outcome.capture(() -> initializationFailure.apply(initializationCause)) + .complete(result); return; } var outcome = Outcome.capture(task); @@ -292,6 +294,43 @@ private static CompletableFuture supplyHandler( return result; } + /** Classifies capture failures once, retaining the ordinary policy of unwrapping only a completion prefix. */ + private static Throwable normalizeMdcInitializationFailure(Throwable failure) { + rethrowDirectMdcFatal(failure); + var visited = Collections.newSetFromMap(new IdentityHashMap()); + var cause = failure; + var normalized = failure; + var completionPrefix = true; + while (visited.add(cause)) { + rethrowDirectMdcFatal(cause); + if (!(cause instanceof CompletionException)) completionPrefix = false; + if (!(cause instanceof CompletionException + || cause instanceof ExecutionException + || cause instanceof InvocationTargetException + || cause instanceof UndeclaredThrowableException)) return normalized; + var next = readMdcCause(cause); + if (next == null || visited.contains(next)) return completionPrefix ? failure : normalized; + if (completionPrefix) normalized = next; + cause = next; + } + return completionPrefix ? failure : normalized; + } + + private static Throwable readMdcCause(Throwable failure) { + try { + return failure.getCause(); + } catch (Throwable unreadableDiagnostic) { + rethrowDirectMdcFatal(unreadableDiagnostic); + return null; + } + } + + @SuppressWarnings("removal") + private static void rethrowDirectMdcFatal(Throwable failure) { + if (failure instanceof VirtualMachineError fatal) throw fatal; + if (failure instanceof ThreadDeath fatal) throw fatal; + } + /** Inspects only standard transport wrappers at this lifecycle boundary, without trusting diagnostics. */ @SuppressWarnings("removal") private static void rethrowLifecycleFatal(Throwable failure) { From 58ef5699d08e81903dff49d2c8bb2923df74c9b2 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 04:06:58 +0000 Subject: [PATCH 30/48] fix: preserve terminal wait-for-condition outcomes on replay --- ...onditionTerminalReplayIntegrationTest.java | 221 ++++++++++++++++++ .../operation/WaitForConditionOperation.java | 5 +- 2 files changed, 224 insertions(+), 2 deletions(-) create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionTerminalReplayIntegrationTest.java diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionTerminalReplayIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionTerminalReplayIntegrationTest.java new file mode 100644 index 000000000..ed11e1d60 --- /dev/null +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionTerminalReplayIntegrationTest.java @@ -0,0 +1,221 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.*; + +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.stream.Stream; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import software.amazon.awssdk.services.lambda.model.CheckpointDurableExecutionResponse; +import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState; +import software.amazon.awssdk.services.lambda.model.ErrorObject; +import software.amazon.awssdk.services.lambda.model.ExecutionDetails; +import software.amazon.awssdk.services.lambda.model.Operation; +import software.amazon.awssdk.services.lambda.model.OperationAction; +import software.amazon.awssdk.services.lambda.model.OperationStatus; +import software.amazon.awssdk.services.lambda.model.OperationType; +import software.amazon.awssdk.services.lambda.model.OperationUpdate; +import software.amazon.awssdk.services.lambda.model.StepDetails; +import software.amazon.lambda.durable.config.WaitForConditionConfig; +import software.amazon.lambda.durable.exception.WaitForConditionFailedException; +import software.amazon.lambda.durable.execution.DurableExecutor; +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.WaitForConditionResult; +import software.amazon.lambda.durable.serde.JacksonSerDes; +import software.amazon.lambda.durable.testing.local.LocalMemoryExecutionClient; + +class WaitForConditionTerminalReplayIntegrationTest { + private static final String CACHED_ERROR = "cached terminal error"; + + static Stream terminals() { + return Stream.of( + OperationStatus.FAILED, + OperationStatus.CANCELLED, + OperationStatus.TIMED_OUT, + OperationStatus.STOPPED) + .flatMap(status -> Stream.of("absent", "empty", "stored") + .flatMap(details -> Stream.of("replay", "retry-response", "poll-response") + .map(delivery -> Arguments.of(status, details, delivery)))); + } + + @ParameterizedTest + @MethodSource("terminals") + void terminalSnapshotsReplayWithoutRepeatingChecksOrCheckpoints( + OperationStatus status, String details, String delivery) throws Exception { + var deliverWhileRunning = !delivery.equals("replay"); + var retried = new AtomicBoolean(); + var delivered = new AtomicBoolean(); + var pollEntered = new CountDownLatch(1); + var releasePoll = new CountDownLatch(1); + var terminal = new AtomicReference(); + var client = new LocalMemoryExecutionClient() { + @Override + public CheckpointDurableExecutionResponse checkpoint( + String arn, String token, List updates) { + var response = super.checkpoint(arn, token, updates); + var retry = updates.stream().anyMatch(u -> u.action() == OperationAction.RETRY); + if (retry) retried.set(true); + var returnTerminal = delivery.equals("retry-response") && retry + || delivery.equals("poll-response") && retried.get() && updates.isEmpty(); + if (!returnTerminal || !delivered.compareAndSet(false, true)) return response; + if (delivery.equals("poll-response")) { + pollEntered.countDown(); + await(releasePoll); + } + // Exercise an actual checkpoint-completion boundary with the backend's terminal snapshot. + var condition = getOperationByName("condition"); + terminal.set(terminalSnapshot(condition, status, details)); + return response.toBuilder() + .newExecutionState(CheckpointUpdatedExecutionState.builder() + .operations(terminal.get()) + .build()) + .build(); + } + }; + var workers = Executors.newCachedThreadPool(); + var caller = Executors.newSingleThreadExecutor(); + var checks = new AtomicInteger(); + try { + var config = DurableConfig.builder() + .withExecutorService(workers) + .withDurableExecutionClient(client) + .withCheckpointDelay(Duration.ZERO) + .withPollingStrategy(attempt -> Duration.ZERO) + .build(); + var first = caller.submit(() -> execute(List.of(), config, checks, status, details, () -> { + if (delivery.equals("poll-response")) { + await(pollEntered); + releasePoll.countDown(); + } + })) + .get(5, TimeUnit.SECONDS); + if (deliverWhileRunning) assertHandled(first, status, details); + else { + assertEquals(ExecutionStatus.PENDING, first.status()); + terminal.set(terminalSnapshot(client.getOperationByName("condition"), status, details)); + } + assertEquals(1, checks.get()); + var updates = client.getOperationUpdates().size(); + var history = List.of(terminal.get()); + var replay = caller.submit(() -> execute(history, config, checks, status, details, () -> {})) + .get(5, TimeUnit.SECONDS); + assertHandled(replay, status, details); + var repeated = caller.submit(() -> execute(history, config, checks, status, details, () -> {})) + .get(5, TimeUnit.SECONDS); + assertEquals(replay, repeated); + assertEquals(1, checks.get(), "Neither cancelled polling nor completed replay may re-enter the predicate"); + assertEquals(updates, client.getOperationUpdates().size(), "Replay must not rewrite a terminal checkpoint"); + } finally { + releasePoll.countDown(); + caller.shutdownNow(); + workers.shutdownNow(); + assertTrue(caller.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private static void await(CountDownLatch latch) { + try { + assertTrue(latch.await(5, TimeUnit.SECONDS), "Controlled terminal poll was not released"); + } catch (InterruptedException failure) { + Thread.currentThread().interrupt(); + throw new AssertionError(failure); + } + } + + private static DurableExecutionOutput execute( + List history, + DurableConfig config, + AtomicInteger checks, + OperationStatus status, + String details, + Runnable beforeAwait) { + var operations = new ArrayList(); + operations.add(Operation.builder() + .id("execution") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails( + ExecutionDetails.builder().inputPayload("\"input\"").build()) + .build()); + operations.addAll(history); + var input = new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/terminal/execution", + "token", + CheckpointUpdatedExecutionState.builder().operations(operations).build()); + return DurableExecutor.execute( + input, + null, + TypeToken.get(String.class), + (value, context) -> { + try { + var future = context.waitForConditionAsync( + "condition", + Integer.class, + (state, step) -> { + checks.incrementAndGet(); + return WaitForConditionResult.continuePolling(1); + }, + WaitForConditionConfig.builder() + .initialState(0) + .waitStrategy((state, attempt) -> Duration.ofSeconds(1)) + .build()); + beforeAwait.run(); + future.get(); + return "unexpected success"; + } catch (WaitForConditionFailedException failure) { + assertNotEquals("stored", details); + assertEquals(status, failure.getOperationStatus()); + assertEquals("condition", failure.getOperation().name()); + return status + ":no error details"; + } catch (IllegalStateException failure) { + assertEquals("stored", details); + assertEquals(CACHED_ERROR, failure.getMessage()); + return CACHED_ERROR; + } + }, + config); + } + + private static void assertHandled(DurableExecutionOutput output, OperationStatus status, String details) { + assertEquals( + ExecutionStatus.SUCCEEDED, + output.status(), + () -> output.error() == null + ? "No error payload" + : output.error().errorType() + ": " + output.error().errorMessage()); + assertEquals( + new JacksonSerDes().serialize(details.equals("stored") ? CACHED_ERROR : status + ":no error details"), + output.result()); + } + + private static Operation terminalSnapshot(Operation pending, OperationStatus status, String details) { + StepDetails snapshot = null; + if (details.equals("empty")) snapshot = StepDetails.builder().attempt(1).build(); + else if (details.equals("stored")) + snapshot = StepDetails.builder() + .attempt(1) + .error(ErrorObject.builder() + .errorType(IllegalStateException.class.getName()) + .errorMessage(CACHED_ERROR) + .errorData(new JacksonSerDes().serialize(new IllegalStateException(CACHED_ERROR))) + .build()) + .build(); + return pending.toBuilder().status(status).stepDetails(snapshot).build(); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java index afff3deb8..aaed6dbbe 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java @@ -66,7 +66,7 @@ protected void start() { @Override protected void replay(Operation existing) { switch (existing.status()) { - case SUCCEEDED, FAILED -> markAlreadyCompleted(); // Check if already completed / failed + case SUCCEEDED, FAILED, CANCELLED, TIMED_OUT, STOPPED -> markAlreadyCompleted(); case PENDING -> pollReadyAndResumeCheckLoop(existing); // Check if pending retry case STARTED, READY -> resumeCheckLoop(existing); default -> @@ -84,7 +84,8 @@ public T get() { var result = (stepDetails != null) ? stepDetails.result() : null; return deserializeResult(result); } else { - var errorObject = op.stepDetails().error(); + var stepDetails = op.stepDetails(); + var errorObject = stepDetails != null ? stepDetails.error() : null; // Attempt to reconstruct and throw the original exception Throwable original = deserializeException(errorObject); From 565d696b8c016e7630d7ec7e2212d9170b7f26dc Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 04:06:58 +0000 Subject: [PATCH 31/48] ci: retry readonly layer lookup for the existing 20 OTel cases --- .github/workflows/otel-conformance-tests.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/otel-conformance-tests.yml b/.github/workflows/otel-conformance-tests.yml index cfb82bbe7..bd7a4a9f6 100644 --- a/.github/workflows/otel-conformance-tests.yml +++ b/.github/workflows/otel-conformance-tests.yml @@ -68,8 +68,9 @@ jobs: actions: write contents: read id-token: write - uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@f18bd0b5f28c5c90e288d0fb8bca08a849b51863 + uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@1768f32d61d958b974ce9e66748df72e23bc19b2 with: + case_count: 20 runs_on: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.repo.full_name != github.repository && 'ubuntu-latest' || format('codebuild-github-actions-runner-{0}-{1}', github.run_id, github.run_attempt) }} language: java resource_prefix: j From bcb9a09050be67748bcd760a5f7fc51e9d0e4bab Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 04:31:03 +0000 Subject: [PATCH 32/48] fix: resume condition polling outside the checkpoint batcher Port of 880c672b694dd6e9859596872447c03fb779d522 after independent public-API deadlock reproduction; preserve this branch workflow, OTel topology and invocation lifecycle. --- ...tForConditionReadinessIntegrationTest.java | 175 ++++++++++++++++-- .../durable/execution/ExecutionManager.java | 38 +++- .../operation/WaitForConditionOperation.java | 22 ++- .../execution/ExecutionManagerTest.java | 113 +++++++++++ .../WaitForConditionOperationTest.java | 2 + 5 files changed, 329 insertions(+), 21 deletions(-) diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionReadinessIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionReadinessIntegrationTest.java index 111c3c984..4bf216f38 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionReadinessIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionReadinessIntegrationTest.java @@ -6,6 +6,8 @@ import java.lang.management.ManagementFactory; import java.lang.reflect.Method; +import java.nio.file.Files; +import java.nio.file.Path; import java.time.Duration; import java.time.Instant; import java.util.ArrayList; @@ -19,6 +21,7 @@ import java.util.concurrent.locks.LockSupport; import java.util.function.IntFunction; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.CsvSource; import org.junit.jupiter.params.provider.ValueSource; @@ -36,6 +39,146 @@ import software.amazon.lambda.durable.testing.local.LocalMemoryExecutionClient; class WaitForConditionReadinessIntegrationTest { + @TempDir + Path subprocessLogs; + + @ParameterizedTest + @ValueSource(strings = {"async", "direct", "submit-and-wait"}) + void delayedReadyResumesWithSynchronousOperationExecutors(String mode) throws Exception { + var log = subprocessLogs.resolve(mode + ".log"); + var process = new ProcessBuilder( + Path.of(System.getProperty("java.home"), "bin", "java").toString(), + "-cp", + System.getProperty("java.class.path"), + SynchronousProbe.class.getName(), + mode) + .redirectErrorStream(true) + .redirectOutput(log.toFile()) + .start(); + try { + assertTrue(process.waitFor(15, TimeUnit.SECONDS), "The isolated READY probe exceeded its cleanup bound"); + var output = Files.readString(log); + System.out.println("SYNC_READY_PROBE mode=" + mode + " exit=" + process.exitValue() + "\n" + output); + assertEquals(0, process.exitValue(), output); + assertTrue(output.contains("SYNC_READY_SUCCESS checks=2 replayChecks=2"), output); + } finally { + if (process.isAlive()) process.destroyForcibly(); + assertTrue(process.waitFor(5, TimeUnit.SECONDS)); + } + } + + /** A separate process bounds cleanup of an old-code deadlock without altering the observed protocol. */ + public static final class SynchronousProbe { + public static void main(String[] args) throws Exception { + var gate = new HandoffGate(true, false); + var workers = new SynchronousContinuationExecutor(args[0]); + var client = new DelayedReadyClient(gate); + var run = new Run(client, workers, value -> { + if (value == 1) gate.firstWorker.set(Thread.currentThread()); + return value == 2 + ? WaitForConditionResult.stopPolling(value) + : WaitForConditionResult.continuePolling(2); + }); + try (var mdc = new MdcClearGate(gate)) { + var first = run.start(); + await(gate.pollEntered); + await(gate.workerExitEntered); + assertTrue(client.advanceTime()); + gate.releaseReadyResponse.countDown(); + await(gate.readyResponseReturned); + gate.releaseWorkerExit.countDown(); + assertEquals( + ExecutionStatus.SUCCEEDED, + first.get(3, TimeUnit.SECONDS).status()); + assertEquals(2, run.checks.get()); + assertEquals( + ExecutionStatus.SUCCEEDED, + run.start().get(3, TimeUnit.SECONDS).status()); + assertEquals(2, run.checks.get()); + assertTrue(workers.continuationDispatches.get() > 0); + System.out.println("SYNC_READY_SUCCESS checks=2 replayChecks=2 mode=" + args[0]); + run.close(); + gate.shutdownNow(); + } catch (Throwable failure) { + System.out.println("SYNC_READY_BLOCKED mode=" + args[0] + " failure=" + + failure.getClass().getName()); + failure.printStackTrace(System.out); + var bean = ManagementFactory.getThreadMXBean(); + Thread.getAllStackTraces().forEach((thread, stack) -> { + if (!thread.getName().startsWith("durable-sdk-internal") + && !thread.getName().startsWith("sync-ready-worker")) return; + var info = bean.getThreadInfo(new long[] {thread.getId()}, true, true)[0]; + System.out.println("THREAD " + thread.getName() + " " + thread.getState() + " monitors=" + + (info == null ? "[]" : Arrays.toString(info.getLockedMonitors()))); + Arrays.stream(stack).forEach(frame -> System.out.println(" " + frame)); + }); + // The probe owns no cloud resources; exiting the child JVM releases only this fixture's blocked + // threads. + System.exit(2); + } + } + } + + private static final class SynchronousContinuationExecutor extends AbstractExecutorService { + private final String mode; + private final AtomicInteger submissions = new AtomicInteger(); + private final AtomicInteger continuationDispatches = new AtomicInteger(); + private final ExecutorService delegate = Executors.newCachedThreadPool(task -> { + var thread = new Thread(task, "sync-ready-worker"); + thread.setDaemon(true); + return thread; + }); + + private SynchronousContinuationExecutor(String mode) { + this.mode = mode; + } + + @Override + public void execute(Runnable task) { + if (submissions.incrementAndGet() <= 2) { + delegate.execute(task); // Root handler and first check retain their normal asynchronous setup. + return; + } + continuationDispatches.incrementAndGet(); + if (mode.equals("direct")) task.run(); + else if (mode.equals("submit-and-wait")) { + try { + delegate.submit(task).get(); + } catch (InterruptedException failure) { + Thread.currentThread().interrupt(); + throw new AssertionError(failure); + } catch (ExecutionException failure) { + throw new AssertionError(failure.getCause()); + } + } else delegate.execute(task); + } + + @Override + public void shutdown() { + delegate.shutdown(); + } + + @Override + public List shutdownNow() { + return delegate.shutdownNow(); + } + + @Override + public boolean isShutdown() { + return delegate.isShutdown(); + } + + @Override + public boolean isTerminated() { + return delegate.isTerminated(); + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { + return delegate.awaitTermination(timeout, unit); + } + } + @ParameterizedTest @ValueSource(ints = {2, 25}) void readyInRetryResponseContinuesWithoutSuspensionAndReplaysStoredResult(int threshold) throws Exception { @@ -155,24 +298,34 @@ void asynchronousReadyHandoffDoesNotDeadlockOrLoseItsActivityLease( private static void awaitCheckpointHandoff(Thread checkpoint) { var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(3); while (System.nanoTime() < deadline) { - var stack = checkpoint.getStackTrace(); - if (checkpoint.getState() == Thread.State.WAITING - && Arrays.stream(stack) - .anyMatch(frame -> frame.getClassName().endsWith("WaitForConditionOperation")) - && Arrays.stream(stack) - .anyMatch(frame -> frame.getClassName().equals(CompletableFuture.class.getName()) - && frame.getMethodName().equals("join"))) { + for (var entry : Thread.getAllStackTraces().entrySet()) { + var continuation = entry.getKey(); + var stack = entry.getValue(); + if (continuation.getState() != Thread.State.WAITING + || !continuation.getName().startsWith("durable-sdk-internal-") + || Arrays.stream(stack) + .noneMatch(frame -> frame.getClassName().endsWith("WaitForConditionOperation")) + || Arrays.stream(stack) + .noneMatch(frame -> frame.getClassName().equals(CompletableFuture.class.getName()) + && frame.getMethodName().equals("join"))) continue; + assertNotSame(checkpoint, continuation, "Worker publication must not block the checkpoint callback"); + assertTrue(Arrays.stream(stack) + .noneMatch(frame -> frame.getClassName().endsWith("ApiRequestDelayedBatcher"))); + if (Arrays.stream(checkpoint.getStackTrace()) + .anyMatch(frame -> frame.getClassName().endsWith("CheckpointManager") + && frame.getMethodName().equals("checkpointBatch"))) continue; var bean = ManagementFactory.getThreadMXBean(); if (bean.isObjectMonitorUsageSupported()) { - var info = bean.getThreadInfo(new long[] {checkpoint.getId()}, true, true)[0]; - System.out.println( - "READY_HANDOFF checkpoint monitors=" + Arrays.toString(info.getLockedMonitors())); + var info = bean.getThreadInfo(new long[] {continuation.getId()}, true, true)[0]; + assertEquals(0, info.getLockedMonitors().length); + System.out.println("READY_HANDOFF continuation=" + continuation.getName() + " checkpoint=" + + checkpoint.getName() + " continuationMonitors=[] checkpointBatchReturned=true"); } return; } LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1)); } - fail("The READY callback did not reach the controlled old-worker handoff window"); + fail("The independent READY continuation did not reach the controlled old-worker handoff window"); } private static void await(CountDownLatch latch) { 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 4c7feeda4..9a61c0224 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 @@ -15,6 +15,7 @@ import java.util.concurrent.CancellationException; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Executor; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; @@ -75,7 +76,7 @@ public class ExecutionManager implements SafeCloseable { private final Set activeThreads = Collections.synchronizedSet(new HashSet<>()); private static final ThreadLocal currentThreadContext = new ThreadLocal<>(); private final CompletableFuture executionExceptionFuture = new CompletableFuture<>(); - // Guarded by activeThreads so starting a checkpoint request is atomic with the last-thread suspension decision. + // Requests and polling continuations; guarded by activeThreads so admission is atomic with last-thread suspension. private int checkpointRequestsInFlight; /** @@ -544,6 +545,41 @@ void finishCheckpointProcessing() { } } + /** + * Runs a polling continuation away from the serialized checkpoint batcher. Keeps checkpoint processing active from + * receiving the callback until the continuation has registered its next worker (or finished inline). Canceling the + * returned future affects observation only; queued or running work still owns its activity lease. + * + * @param continuation operation work to dispatch after receiving a polling update + * @return completion of the continuation, or an already-completed future when execution has stopped + */ + public CompletableFuture runCheckpointContinuation(Runnable continuation) { + return runCheckpointContinuation(continuation, InternalExecutor.INSTANCE); + } + + CompletableFuture runCheckpointContinuation(Runnable continuation, Executor coordinator) { + if (!tryStartCheckpointProcessing()) return CompletableFuture.completedFuture(null); + try { + var completion = new CompletableFuture(); + coordinator.execute(() -> { + try { + try { + continuation.run(); + } finally { + finishCheckpointProcessing(); + } + completion.complete(null); + } catch (Throwable failure) { + completion.completeExceptionally(failure); + } + }); + return completion; + } catch (RuntimeException | Error failure) { + finishCheckpointProcessing(); + throw failure; + } + } + private boolean shouldSuspendExecution() { return activeThreads.isEmpty() && checkpointRequestsInFlight == 0 && !executionExceptionFuture.isDone(); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java index aaed6dbbe..1dbf5ba5b 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java @@ -112,9 +112,10 @@ private void resumeCheckLoop(Operation existing) { } private CompletableFuture pollReadyAndResumeCheckLoop(Operation existing) { - return pollUntilReady().thenAccept(op -> { - if (!isOperationCompleted() && op.status() == OperationStatus.READY) resumeCheckLoop(op); - }); + return pollUntilReady() + .thenCompose(op -> executionManager.runCheckpointContinuation(() -> { + if (!isOperationCompleted() && op.status() == OperationStatus.READY) resumeCheckLoop(op); + })); } private CompletableFuture pollUntilReady() { @@ -195,15 +196,18 @@ private Operation continueInlineOrAfterCurrentWorker( var inline = new AtomicReference(); var acceptingInline = new AtomicBoolean(true); try { - pollUntilReady().thenAccept(op -> { + pollUntilReady().thenCompose(op -> { if (Thread.currentThread() == owner && acceptingInline.get()) { inline.set(op); + return CompletableFuture.completedFuture(null); } else { - // A checkpoint callback can observe READY before this attempt's worker exits. Its checkpoint - // processing lease stays active through this handoff, preventing a false quiescence window. - publishedWorker.join().join(); - if (!isOperationCompleted() && op.status() == OperationStatus.READY) - executeCheckLogic(nextState, nextAttempt); + // A configured executor may run or wait for the next check inline. Leave the serialized + // checkpoint callback before dispatch, retaining activity through worker publication/registration. + return executionManager.runCheckpointContinuation(() -> { + publishedWorker.join().join(); + if (!isOperationCompleted() && op.status() == OperationStatus.READY) + executeCheckLogic(nextState, nextAttempt); + }); } }); } finally { diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java index 1552d2b21..dfb67639e 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java @@ -14,10 +14,15 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; 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.GetDurableExecutionStateResponse; import software.amazon.awssdk.services.lambda.model.Operation; @@ -372,6 +377,114 @@ void deferredSuspensionOccursWhenCheckpointFinishesWithoutReactivatingThread() { assertTrue(manager.isExecutionCompletedExceptionally()); } + @ParameterizedTest + @ValueSource(strings = {"complete", "register-worker", "fail"}) + void pollingContinuationRetainsActivityUntilItFinishesDispatch(String outcome) throws Exception { + var manager = createManager(List.of(executionOp(), stepOp("step", OperationStatus.PENDING))); + var entered = new CountDownLatch(1); + var release = new CountDownLatch(1); + var caller = Thread.currentThread(); + var failure = new IllegalStateException("continuation failure"); + manager.registerActiveThread("root"); + var completion = manager.runCheckpointContinuation(() -> { + assertNotSame(caller, Thread.currentThread()); + entered.countDown(); + try { + assertTrue(release.await(3, TimeUnit.SECONDS)); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw new AssertionError(interrupted); + } + if (outcome.equals("register-worker")) manager.registerActiveThread("next"); + if (outcome.equals("fail")) throw failure; + }); + try { + assertTrue(entered.await(3, TimeUnit.SECONDS)); + manager.deregisterActiveThread("root"); + assertFalse(manager.isExecutionCompletedExceptionally()); + assertFalse(completion.isDone()); + } finally { + release.countDown(); + } + if (outcome.equals("fail")) { + assertSame( + failure, + assertThrows(ExecutionException.class, () -> completion.get(3, TimeUnit.SECONDS)) + .getCause()); + } else completion.get(3, TimeUnit.SECONDS); + if (outcome.equals("register-worker")) { + assertFalse(manager.isExecutionCompletedExceptionally()); + assertThrows(SuspendExecutionException.class, () -> manager.deregisterActiveThread("next")); + } + assertTrue(manager.isExecutionCompletedExceptionally()); + manager.runCheckpointContinuation(() -> fail("No continuation may start after suspension")) + .get(3, TimeUnit.SECONDS); + } + + @Test + void cancellingQueuedContinuationDoesNotSkipItsWorkOrLeakActivity() { + var manager = createManager(List.of(executionOp(), stepOp("step", OperationStatus.PENDING))); + var queued = new AtomicReference(); + var calls = new AtomicInteger(); + manager.registerActiveThread("root"); + var completion = manager.runCheckpointContinuation(calls::incrementAndGet, queued::set); + assertTrue(completion.cancel(false)); + manager.deregisterActiveThread("root"); + assertFalse(manager.isExecutionCompletedExceptionally(), "Queued work still owns the activity lease"); + queued.get().run(); + assertAll( + () -> assertEquals(1, calls.get()), + () -> assertTrue(manager.isExecutionCompletedExceptionally(), "The actual runnable releases its lease"), + () -> assertTrue(completion.isCancelled())); + } + + @Test + void cancellingRunningContinuationDoesNotReleaseActivityBeforeItsCleanup() throws Exception { + var manager = createManager(List.of(executionOp(), stepOp("step", OperationStatus.PENDING))); + var queued = new AtomicReference(); + var entered = new CountDownLatch(1); + var release = new CountDownLatch(1); + manager.registerActiveThread("root"); + var completion = manager.runCheckpointContinuation( + () -> { + entered.countDown(); + try { + assertTrue(release.await(3, TimeUnit.SECONDS)); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw new AssertionError(interrupted); + } + }, + queued::set); + var actualWorker = CompletableFuture.runAsync(queued.get()); + try { + assertTrue(entered.await(3, TimeUnit.SECONDS)); + assertTrue(completion.cancel(false)); + manager.deregisterActiveThread("root"); + assertFalse(manager.isExecutionCompletedExceptionally(), "Cancellation must not release running work"); + } finally { + release.countDown(); + } + actualWorker.get(3, TimeUnit.SECONDS); + assertTrue(manager.isExecutionCompletedExceptionally()); + assertTrue(completion.isCancelled()); + } + + @Test + void rejectedContinuationReleasesActivityAndPreservesTheRejection() { + var manager = createManager(List.of(executionOp(), stepOp("step", OperationStatus.PENDING))); + var rejection = new RejectedExecutionException("coordinator unavailable"); + manager.registerActiveThread("root"); + assertSame( + rejection, + assertThrows( + RejectedExecutionException.class, + () -> manager.runCheckpointContinuation(() -> fail("Rejected work must not run"), task -> { + throw rejection; + }))); + assertThrows(SuspendExecutionException.class, () -> manager.deregisterActiveThread("root")); + } + @Test void checkpointDeliveryIsAtomicWithOperationRegistration() throws Exception { var manager = createManager(List.of(executionOp(), stepOp("step", OperationStatus.PENDING))); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionOperationTest.java index 6c923994f..c4df8aac5 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionOperationTest.java @@ -284,6 +284,8 @@ void getFailedWithNullErrorDataThrowsStepFailedException() { @Test void replayPendingPollsAndResumesCheckLoop() throws Exception { + when(executionManager.runCheckpointContinuation(any())) + .thenAnswer(call -> CompletableFuture.runAsync(call.getArgument(0))); var pendingOp = Operation.builder() .id(OPERATION_ID) .name(OPERATION_NAME) From 4b048d97e85b92dcd04c37d235f88d69558f4a7a Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 04:59:29 +0000 Subject: [PATCH 33/48] fix: yield READY checks to queued user operations (cherry picked from commit 1e27f85a6c4497b8f3fe731e489184a5a9f1e3fa) --- ...tForConditionReadinessIntegrationTest.java | 96 ++++++++++++++++--- .../operation/WaitForConditionOperation.java | 44 +++------ 2 files changed, 96 insertions(+), 44 deletions(-) diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionReadinessIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionReadinessIntegrationTest.java index 4bf216f38..9ed8cd295 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionReadinessIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionReadinessIntegrationTest.java @@ -19,6 +19,7 @@ import java.util.concurrent.*; import java.util.concurrent.atomic.*; import java.util.concurrent.locks.LockSupport; +import java.util.function.BiFunction; import java.util.function.IntFunction; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; @@ -45,12 +46,22 @@ class WaitForConditionReadinessIntegrationTest { @ParameterizedTest @ValueSource(strings = {"async", "direct", "submit-and-wait"}) void delayedReadyResumesWithSynchronousOperationExecutors(String mode) throws Exception { - var log = subprocessLogs.resolve(mode + ".log"); + assertProbe(SynchronousProbe.class, mode, "SYNC_READY_SUCCESS checks=2 replayChecks=2"); + } + + @ParameterizedTest + @ValueSource(strings = {"cached", "bounded"}) + void immediateReadyYieldsToTheSiblingThatUnlocksItsCondition(String mode) throws Exception { + assertProbe(FairnessProbe.class, mode, "FAIR_READY_SUCCESS siblingCalls=1 replayStable=true"); + } + + private void assertProbe(Class probe, String mode, String successMarker) throws Exception { + var log = subprocessLogs.resolve(probe.getSimpleName() + "-" + mode + ".log"); var process = new ProcessBuilder( Path.of(System.getProperty("java.home"), "bin", "java").toString(), "-cp", System.getProperty("java.class.path"), - SynchronousProbe.class.getName(), + probe.getName(), mode) .redirectErrorStream(true) .redirectOutput(log.toFile()) @@ -60,13 +71,73 @@ void delayedReadyResumesWithSynchronousOperationExecutors(String mode) throws Ex var output = Files.readString(log); System.out.println("SYNC_READY_PROBE mode=" + mode + " exit=" + process.exitValue() + "\n" + output); assertEquals(0, process.exitValue(), output); - assertTrue(output.contains("SYNC_READY_SUCCESS checks=2 replayChecks=2"), output); + assertTrue(output.contains(successMarker), output); } finally { if (process.isAlive()) process.destroyForcibly(); assertTrue(process.waitFor(5, TimeUnit.SECONDS)); } } + public static final class FairnessProbe { + public static void main(String[] args) throws Exception { + var workers = args[0].equals("bounded") ? Executors.newFixedThreadPool(2) : Executors.newCachedThreadPool(); + var siblingQueued = new CountDownLatch(1); + var unlocked = new AtomicBoolean(); + var checks = new AtomicInteger(); + var siblings = new AtomicInteger(); + var run = new Run(new ImmediateReadyClient(), workers, value -> { + throw new AssertionError("The explicit public handler supplies its own condition"); + }); + run.customHandler = (input, context) -> { + var condition = context.waitForConditionAsync( + "condition", + Integer.class, + (state, step) -> { + checks.incrementAndGet(); + await(siblingQueued); + return unlocked.get() + ? WaitForConditionResult.stopPolling(state) + : WaitForConditionResult.continuePolling(state + 1); + }, + run.wait); + var sibling = context.stepAsync("unlock-condition", String.class, step -> { + siblings.incrementAndGet(); + unlocked.set(true); + return "unlocked"; + }); + siblingQueued.countDown(); + var result = condition.get(); + assertEquals("unlocked", sibling.get()); + return String.valueOf(result); + }; + try { + var first = run.start().get(3, TimeUnit.SECONDS); + assertEquals(ExecutionStatus.SUCCEEDED, first.status()); + assertEquals(1, siblings.get()); + var completedChecks = checks.get(); + var checkpoints = run.client.getOperationUpdates().size(); + assertEquals(first, run.start().get(3, TimeUnit.SECONDS)); + assertEquals(completedChecks, checks.get()); + assertEquals(1, siblings.get()); + assertEquals(checkpoints, run.client.getOperationUpdates().size()); + System.out.println("FAIR_READY_SUCCESS siblingCalls=1 replayStable=true checks=" + completedChecks); + run.close(); + } catch (Throwable failure) { + System.out.println("FAIR_READY_FAILURE mode=" + args[0] + " checks=" + checks.get() + " siblingCalls=" + + siblings.get() + " queued=" + + ((ThreadPoolExecutor) workers).getQueue().size()); + failure.printStackTrace(System.out); + Thread.getAllStackTraces().forEach((thread, stack) -> { + if (!thread.getName().startsWith("durable-sdk-internal") + && !thread.getName().startsWith("pool-")) return; + System.out.println("THREAD " + thread.getName() + " " + thread.getState()); + Arrays.stream(stack).forEach(frame -> System.out.println(" " + frame)); + }); + System.exit(2); + } + } + } + /** A separate process bounds cleanup of an old-code deadlock without altering the observed protocol. */ public static final class SynchronousProbe { public static void main(String[] args) throws Exception { @@ -533,6 +604,7 @@ public T deserialize(String value, TypeToken type) { } }; private final IntFunction> check; + private BiFunction customHandler; private final DurableConfig config; private final WaitForConditionConfig wait = WaitForConditionConfig.builder() .initialState(1) @@ -576,14 +648,16 @@ private Future start() { input, null, TypeToken.get(String.class), - (value, ctx) -> String.valueOf(ctx.waitForCondition( - "condition", - Integer.class, - (state, step) -> { - checks.incrementAndGet(); - return check.apply(state); - }, - wait)), + customHandler != null + ? customHandler + : (value, ctx) -> String.valueOf(ctx.waitForCondition( + "condition", + Integer.class, + (state, step) -> { + checks.incrementAndGet(); + return check.apply(state); + }, + wait)), config)); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java index 1dbf5ba5b..36e9e986f 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java @@ -4,8 +4,6 @@ import java.time.Duration; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; import software.amazon.awssdk.services.lambda.model.Operation; import software.amazon.awssdk.services.lambda.model.OperationAction; @@ -149,7 +147,7 @@ private void executeCheckLogic(T currentState, int attempt) { } private void runCheckLoop(T currentState, int attempt, CompletableFuture> publishedWorker) { - while (!isOperationCompleted()) { + if (!isOperationCompleted()) { var stepContext = getContext().createStepContext(getOperationId(), getName(), attempt); BaseContextImpl.setCurrentContext(stepContext); try (var ignored = DurableLogger.attachContext()) { @@ -175,13 +173,9 @@ private void runCheckLoop(T currentState, int attempt, CompletableFuture> publishedWorker, T nextState, int nextAttempt) { - var owner = Thread.currentThread(); - var inline = new AtomicReference(); - var acceptingInline = new AtomicBoolean(true); - try { - pollUntilReady().thenCompose(op -> { - if (Thread.currentThread() == owner && acceptingInline.get()) { - inline.set(op); - return CompletableFuture.completedFuture(null); - } else { - // A configured executor may run or wait for the next check inline. Leave the serialized - // checkpoint callback before dispatch, retaining activity through worker publication/registration. - return executionManager.runCheckpointContinuation(() -> { - publishedWorker.join().join(); - if (!isOperationCompleted() && op.status() == OperationStatus.READY) - executeCheckLogic(nextState, nextAttempt); - }); - } - }); - } finally { - acceptingInline.set(false); - } - return inline.get(); + pollUntilReady() + .thenCompose(op -> executionManager.runCheckpointContinuation(() -> { + publishedWorker.join().join(); + if (!isOperationCompleted() && op.status() == OperationStatus.READY) + executeCheckLogic(nextState, nextAttempt); + })); } private void handleCheckFailure(Throwable exception) { From 23abc3875ae661e046a72562efbba4f926a04edb Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 05:35:48 +0000 Subject: [PATCH 34/48] fix: consume durable sampling intent before processor callbacks Reuse the isolated one-shot consumption and sampler ownership subset validated in PR767. Preserve ordinary-span delegate decisions without introducing synthetic roots or changing invocation lifecycle. --- .../lambda/durable/otel/DurableSampler.java | 2 +- .../durable/otel/DurableSamplingDecision.java | 44 ++-- .../durable/otel/ExecutionOtelPlugin.java | 4 +- .../durable/otel/InvocationOtelPlugin.java | 4 +- .../durable/otel/OtelPluginSupport.java | 5 + .../durable/otel/DurableSamplerTest.java | 31 ++- .../otel/SamplingProcessorIsolationTest.java | 217 ++++++++++++++++++ 7 files changed, 288 insertions(+), 19 deletions(-) create mode 100644 otel-plugin/src/test/java/software/amazon/lambda/durable/otel/SamplingProcessorIsolationTest.java diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSampler.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSampler.java index 7ce304a54..34e222ab3 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSampler.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSampler.java @@ -113,7 +113,7 @@ public SamplingResult shouldSample( SpanKind spanKind, Attributes attributes, List parentLinks) { - var intent = DurableSamplingDecision.get(parentContext); + var intent = DurableSamplingDecision.consume(parentContext); if (intent == null) { // Not a durable span: the customer's sampler governs it unchanged. return delegate.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks); diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSamplingDecision.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSamplingDecision.java index c43d8eb02..2341f8bdb 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSamplingDecision.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSamplingDecision.java @@ -6,6 +6,7 @@ import io.opentelemetry.context.ContextKey; import io.opentelemetry.sdk.trace.samplers.SamplingDecision; import io.opentelemetry.sdk.trace.samplers.SamplingResult; +import java.util.concurrent.atomic.AtomicReference; /** * Carries the durable execution's sampling intent to {@link DurableSampler} for one durable span. @@ -28,19 +29,21 @@ * identity, so a key created in one loader is not equal to the key created in the other. To bridge this: * *

    - *
  1. Context key — used when both sides share a class loader (an application-owned provider). It preserves - * the full {@link SamplingResult}, including any attributes a custom sampler attached. + *
  2. Context key — used when both sides share a class loader (an application-owned provider). Agent-backed + * spans omit this carrier because another loader cannot consume its key. It preserves the full + * {@link SamplingResult}, including any attributes a custom sampler attached. *
  3. Thread-scoped system property — a cross-class-loader fallback modelled on * {@link DeterministicIdGenerator}'s scoped-ID bridge. The intent is published on the thread that creates the - * durable span for the synchronous duration of {@code startSpan()} (the sampler runs on that same thread), keyed - * by thread ID under a bootstrap-visible {@link System} property so both class loaders read the same value. A - * resolved decision bridges its {@link SamplingDecision} name (the three built-in decisions carry no attributes, - * so reconstructing them is faithful); a deferral bridges a sentinel plus the canonical trace ID. + * durable span until its sampler consumes the value (before synchronous span processors run), keyed by thread ID + * under a bootstrap-visible {@link System} property so both class loaders read the same value. A resolved + * decision bridges its {@link SamplingDecision} name (the three built-in decisions carry no attributes, so + * reconstructing them is faithful); a deferral bridges a sentinel plus the canonical trace ID. *
* *

The scope is opened immediately around each durable {@code startSpan()} call and closed right after, so the - * property never leaks beyond the span it applies to. Nothing is persisted across invocations; cross-invocation - * consistency comes from recomputing the intent from stable inputs, not from sharing state. + * property is consumed once by the sampler before callbacks can create unrelated spans. Nothing is persisted across + * invocations; cross-invocation consistency comes from recomputing the intent from stable inputs, not from sharing + * state. */ final class DurableSamplingDecision { @@ -65,7 +68,7 @@ boolean isDeferred() { } } - private static final ContextKey KEY = + private static final ContextKey> KEY = ContextKey.named("software.amazon.lambda.durable.otel.durable-sampling-decision"); private static final String SCOPED_PROPERTY_PREFIX = "software.amazon.lambda.durable.otel.scopedSamplingDecision."; @@ -74,9 +77,9 @@ boolean isDeferred() { private DurableSamplingDecision() {} - /** Returns a context carrying the durable sampling intent (same-class-loader carrier), derived from the given. */ + /** Returns a context carrying a one-shot intent for one span, preserving the full SamplingResult. */ static Context store(Context context, Intent intent) { - return context.with(KEY, intent); + return context.with(KEY, new AtomicReference<>(intent)); } /** @@ -98,17 +101,30 @@ static Scope openScope(Intent intent) { } /** - * Returns the durable sampling intent for a span, or {@code null} when none is present. Prefers the full-fidelity - * context key (same class loader) and falls back to the thread-scoped system property (cross class loader). + * Inspects the unconsumed durable sampling intent for a span, or {@code null} when none is present. Prefers the + * full-fidelity context key (same class loader) and falls back to the thread-scoped system property (cross class + * loader). */ static Intent get(Context context) { - var fromContext = context.get(KEY); + var holder = context.get(KEY); + var fromContext = holder != null ? holder.get() : null; if (fromContext != null) { return fromContext; } return fromScopedProperty(); } + /** Consumes both carriers before callbacks can reuse the parent context or cross-loader bridge. */ + static Intent consume(Context context) { + try { + var holder = context.get(KEY); + var fromContext = holder != null ? holder.getAndSet(null) : null; + return fromContext != null ? fromContext : fromScopedProperty(); + } finally { + System.clearProperty(scopedProperty()); + } + } + private static String encode(Intent intent) { return intent.isDeferred() ? DEFERRED_PREFIX + intent.deferredTraceId() diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java index a621e8900..79c448ff4 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java @@ -655,7 +655,9 @@ private Context resolveParentContext(String parentId) { */ private Context withDurableDecision(Context context) { var intent = samplingIntent; - return intent != null ? DurableSamplingDecision.store(context, intent) : context; + return intent != null && OtelPluginSupport.usesLocalDurableSampler(sdkTracerProvider) + ? DurableSamplingDecision.store(context, intent) + : context; } private TraceFlags effectiveTraceFlags() { diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java index 3042c603d..c00c07508 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java @@ -673,7 +673,9 @@ private Context resolveParentContext(String parentId) { */ private Context withDurableDecision(Context context) { var intent = samplingIntent; - return intent != null ? DurableSamplingDecision.store(context, intent) : context; + return intent != null && OtelPluginSupport.usesLocalDurableSampler(sdkTracerProvider) + ? DurableSamplingDecision.store(context, intent) + : context; } /** diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java index 9dd9684f5..f06b8b787 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java @@ -26,6 +26,11 @@ final class OtelPluginSupport { private OtelPluginSupport() {} + /** Only the same DurableSampler class copy can consume this loader's parent-context holder. */ + static boolean usesLocalDurableSampler(SdkTracerProvider provider) { + return provider != null && provider.getSampler() instanceof DurableSampler; + } + /** Closes the owning thread's scope without hiding a finalization error or a later JVM-fatal cleanup failure. */ static void runInvocationEnd(Scope scope, Runnable end) { Throwable primary = null; diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DurableSamplerTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DurableSamplerTest.java index 6fea5784e..f6126492e 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DurableSamplerTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DurableSamplerTest.java @@ -62,6 +62,32 @@ void tearDown() { OtelPluginAutoConfigurationState.resetInstalledForTest(); } + @Test + void consumingContextIntentAlsoClearsCrossLoaderFallback() { + var delegate = new CountingSampler(Sampler.alwaysOff()); + var sampler = DurableSampler.wrap(delegate); + var parent = DurableSamplingDecision.store( + Context.root(), DurableSamplingDecision.Intent.resolved(SamplingResult.recordOnly())); + try (var ignored = DurableSamplingDecision.openScope( + DurableSamplingDecision.Intent.resolved(SamplingResult.recordAndSample()))) { + assertEquals( + SamplingDecision.RECORD_ONLY, + sampler.shouldSample(parent, TRACE_ID, "durable", SpanKind.INTERNAL, Attributes.empty(), List.of()) + .getDecision()); + assertEquals( + SamplingDecision.DROP, + sampler.shouldSample( + Context.root(), + TRACE_ID, + "callback", + SpanKind.INTERNAL, + Attributes.empty(), + List.of()) + .getDecision()); + assertEquals(1, delegate.count(), "The unrelated callback must use its own sampler"); + } + } + // ─── Unit tests for the wrapper ────────────────────────────────────── @Test @@ -124,9 +150,10 @@ void deferredIntent_evaluatesDelegateOncePerExecution() { // A stateful/quota delegate must be consulted once per execution (trace ID), not per durable span. var delegate = new CountingSampler(Sampler.alwaysOn()); var sampler = DurableSampler.wrap(delegate); - var parent = DurableSamplingDecision.store(Context.root(), DurableSamplingDecision.Intent.deferred(TRACE_ID)); - for (var i = 0; i < 4; i++) { + // Each SDK-owned span receives its own one-shot carrier, while the resolved decision stays execution-wide. + var parent = + DurableSamplingDecision.store(Context.root(), DurableSamplingDecision.Intent.deferred(TRACE_ID)); sampler.shouldSample(parent, TRACE_ID, "op" + i, SpanKind.INTERNAL, Attributes.empty(), List.of()); } diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/SamplingProcessorIsolationTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/SamplingProcessorIsolationTest.java new file mode 100644 index 000000000..863dacb01 --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/SamplingProcessorIsolationTest.java @@ -0,0 +1,217 @@ +// 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.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; + +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.api.trace.TracerProvider; +import io.opentelemetry.context.Context; +import io.opentelemetry.context.propagation.ContextPropagators; +import io.opentelemetry.sdk.trace.IdGenerator; +import io.opentelemetry.sdk.trace.ReadWriteSpan; +import io.opentelemetry.sdk.trace.ReadableSpan; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.SpanProcessor; +import io.opentelemetry.sdk.trace.samplers.Sampler; +import java.net.URL; +import java.net.URLClassLoader; +import java.time.Instant; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +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; + +/** + * Verifies the durable sampling decision crosses the application/Java-agent class-loader boundary. + * + *

Under the documented ADOT setup the plugin JAR is loaded twice: the application class loader computes and stores + * the decision, and a separate Java-agent extension class loader installs and runs the sampler. Because a + * {@link io.opentelemetry.context.ContextKey} uses reference identity, the two loaders hold distinct keys and the + * context carrier alone cannot bridge them. This test reproduces that topology with two child-first class loaders that + * each load {@code DurableSamplingDecision} separately while sharing the OpenTelemetry API/SDK types with the parent, + * then asserts the thread-scoped system-property bridge carries the decision from one loader to the other. + */ +class SamplingProcessorIsolationTest { + + @AfterEach + void clearBridge() { + DurableSamplingDecision.clearSharedStateForTest(); + } + + @ParameterizedTest + @CsvSource({ + "InvocationOtelPlugin,true,false,Invocation", + "InvocationOtelPlugin,false,false,Invocation", + "InvocationOtelPlugin,false,true,Invocation", + "ExecutionOtelPlugin,true,false,Invocation", + "ExecutionOtelPlugin,false,false,Invocation", + "ExecutionOtelPlugin,false,true,Invocation" + }) + void agentProcessorForwardingParentCannotReuseApplicationSamplingIntent( + String pluginName, boolean hideProvider, boolean localSampler, String observedSpan) throws Exception { + var previousHeader = System.getProperty("com.amazonaws.xray.traceHeader"); + GlobalOpenTelemetry.resetForTest(); + OtelPluginAutoConfigurationState.markInstalled(); + System.setProperty("com.amazonaws.xray.traceHeader", "Root=1-6955b900-123456789012345678901234;Sampled=1"); + try (var appLoader = pluginClassLoader(); + var agentLoader = pluginClassLoader()) { + var appSamplerType = Class.forName(DurableSampler.class.getName(), true, appLoader); + var appWrap = appSamplerType.getDeclaredMethod("wrap", Sampler.class); + appWrap.setAccessible(true); + var agentSamplerType = Class.forName(DurableSampler.class.getName(), true, agentLoader); + var agentWrap = agentSamplerType.getDeclaredMethod("wrap", Sampler.class); + agentWrap.setAccessible(true); + var appDecision = Class.forName(DurableSamplingDecision.class.getName(), true, appLoader); + var appGet = appDecision.getDeclaredMethod("get", Context.class); + appGet.setAccessible(true); + var agentDecision = Class.forName(DurableSamplingDecision.class.getName(), true, agentLoader); + var agentGet = agentDecision.getDeclaredMethod("get", Context.class); + agentGet.setAccessible(true); + var agentIdType = Class.forName(DeterministicIdGenerator.class.getName(), true, agentLoader); + var callbacks = new AtomicInteger(); + var leakedSampling = new AtomicBoolean(); + try (var appProvider = SdkTracerProvider.builder() + .setSampler((Sampler) appWrap.invoke(null, Sampler.alwaysOff())) + .build(); + var agentProvider = SdkTracerProvider.builder() + .setSampler( + (Sampler) (localSampler ? appWrap : agentWrap).invoke(null, Sampler.alwaysOff())) + .setIdGenerator( + (IdGenerator) agentIdType.getConstructor().newInstance()) + .addSpanProcessor(new SpanProcessor() { + @Override + public void onStart(Context parent, ReadWriteSpan span) { + if (!span.getName().equals(observedSpan)) return; + var unrelated = appProvider + .get("processor") + .spanBuilder("unrelated-forwarded-parent") + .setParent(parent) + .startSpan(); + callbacks.incrementAndGet(); + leakedSampling.compareAndSet(false, unrelated.isRecording()); + unrelated.end(); + } + + @Override + public boolean isStartRequired() { + return true; + } + + @Override + public void onEnd(ReadableSpan span) {} + + @Override + public boolean isEndRequired() { + return false; + } + }) + .build()) { + assertEquals(localSampler, agentProvider.getSampler().getClass() == appSamplerType); + var hiddenAgentProvider = new TracerProvider() { + @Override + public Tracer get(String name) { + return agentProvider.get(name); + } + + @Override + public Tracer get(String name, String version) { + return agentProvider.get(name, version); + } + }; + GlobalOpenTelemetry.set(new OpenTelemetry() { + @Override + public TracerProvider getTracerProvider() { + return hideProvider ? hiddenAgentProvider : agentProvider; + } + + @Override + public ContextPropagators getPropagators() { + return ContextPropagators.noop(); + } + }); + var plugin = (DurableExecutionPlugin) + Class.forName("software.amazon.lambda.durable.otel." + pluginName, true, appLoader) + .getConstructor() + .newInstance(); + var arn = "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/test/id"; + var previousAmbient = Span.current(); + var ambient = appProvider + .get("ambient") + .spanBuilder("ambient-control") + .startSpan(); + assertFalse(ambient.isRecording(), "The app provider drops unrelated spans outside the callback"); + try (var ambientScope = ambient.makeCurrent()) { + for (var first : new boolean[] {true, false}) { + plugin.onInvocationStart(new InvocationInfo("request", arn, first, Instant.EPOCH)); + plugin.onInvocationEnd( + new InvocationEndInfo("request", arn, first, InvocationStatus.PENDING, null)); + assertSame(ambient, Span.current(), "Plugin cleanup must preserve the caller's ambient span"); + assertNull(appGet.invoke(null, Context.root()), "Application sampling intent must be cleared"); + assertNull(agentGet.invoke(null, Context.root()), "Agent sampling intent must be cleared"); + } + } finally { + ambient.end(); + } + assertSame(previousAmbient, Span.current()); + assertEquals(2, callbacks.get(), "The selected span in both invocations must reach the real processor"); + assertFalse( + leakedSampling.get(), "The supplied parent must not override the app provider's DROP policy"); + } + } finally { + GlobalOpenTelemetry.resetForTest(); + OtelPluginAutoConfigurationState.resetInstalledForTest(); + if (previousHeader == null) System.clearProperty("com.amazonaws.xray.traceHeader"); + else System.setProperty("com.amazonaws.xray.traceHeader", previousHeader); + } + } + + /** + * A child-first class loader that loads {@code software.amazon.lambda.durable.otel.*} itself (so each instance + * holds its own copies, mirroring the two plugin class loaders) while delegating OpenTelemetry and JDK classes to + * the parent so those types are shared and interoperable across loaders. + */ + private static URLClassLoader pluginClassLoader() { + var classesDir = SamplingProcessorIsolationTest.class + .getProtectionDomain() + .getCodeSource() + .getLocation(); + // target/test-classes -> the main classes live in target/classes alongside it. + URL mainClasses; + try { + mainClasses = new URL(classesDir.toString().replace("/test-classes/", "/classes/")); + } catch (Exception e) { + throw new IllegalStateException(e); + } + var parent = SamplingProcessorIsolationTest.class.getClassLoader(); + return new URLClassLoader(new URL[] {mainClasses}, parent) { + @Override + protected Class loadClass(String name, boolean resolve) throws ClassNotFoundException { + if (name.startsWith("software.amazon.lambda.durable.otel.")) { + synchronized (getClassLoadingLock(name)) { + var loaded = findLoadedClass(name); + if (loaded == null) { + loaded = findClass(name); + } + if (resolve) { + resolveClass(loaded); + } + return loaded; + } + } + return super.loadClass(name, resolve); + } + }; + } +} From 85a09eb34a4843b648c132cff42447b507502cce Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 05:53:04 +0000 Subject: [PATCH 35/48] fix: drain owned condition continuations before manager shutdown (cherry picked from commit a08bdb46f84f7338e91d458504c6fd162f05de8c) --- .../ContinuationShutdownIntegrationTest.java | 431 ++++++++++++++++++ .../durable/execution/ExecutionManager.java | 110 ++++- .../operation/WaitForConditionOperation.java | 4 +- .../execution/ExecutionManagerTest.java | 61 +++ .../WaitForConditionOperationTest.java | 4 +- 5 files changed, 590 insertions(+), 20 deletions(-) create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/execution/ContinuationShutdownIntegrationTest.java diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/execution/ContinuationShutdownIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/execution/ContinuationShutdownIntegrationTest.java new file mode 100644 index 000000000..94b3489f8 --- /dev/null +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/execution/ContinuationShutdownIntegrationTest.java @@ -0,0 +1,431 @@ +// 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 java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import org.junit.jupiter.api.AfterEach; +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.*; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.config.CompletionConfig; +import software.amazon.lambda.durable.config.ParallelConfig; +import software.amazon.lambda.durable.config.WaitForConditionConfig; +import software.amazon.lambda.durable.context.BaseContext; +import software.amazon.lambda.durable.context.BaseContextImpl; +import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.model.WaitForConditionResult; +import software.amazon.lambda.durable.operation.BaseDurableOperation; +import software.amazon.lambda.durable.testing.local.LocalMemoryExecutionClient; + +class ContinuationShutdownIntegrationTest { + private ControlledManager initialManager; + private BaseContext previousContext; + private ThreadContext previousThreadContext; + + private void rememberContext(ControlledManager manager) { + if (initialManager != null) return; + initialManager = manager; + previousContext = BaseContext.getCurrentContext(); + previousThreadContext = manager.getCurrentThreadContext(); + } + + @AfterEach + void restoreCallerContext() { + if (initialManager == null) return; + BaseContextImpl.setCurrentContext(previousContext); + initialManager.setCurrentThreadContext(previousThreadContext); + } + + @Test + void runningPredicateFinishesItsCheckpointAndReplaysWithoutRepeatingWork() throws Exception { + var closed = new AtomicBoolean(); + var late = new AtomicInteger(); + var client = readyClient(closed, late); + var queue = new LinkedBlockingQueue(); + var workers = Executors.newCachedThreadPool(); + var started = new CountDownLatch(1); + var release = new CountDownLatch(1); + var calls = new AtomicInteger(); + var setup = setup(client, workers, queue, closed); + try { + setup.context.waitForConditionAsync( + "condition", + Integer.class, + (state, step) -> { + calls.incrementAndGet(); + if (state == 1) return WaitForConditionResult.continuePolling(2); + started.countDown(); + await(release); + return WaitForConditionResult.stopPolling(state); + }, + waitConfig()); + var task = queue.poll(3, TimeUnit.SECONDS); + assertNotNull(task); + CompletableFuture.runAsync(task).get(3, TimeUnit.SECONDS); + assertTrue(started.await(3, TimeUnit.SECONDS)); + var closing = CompletableFuture.runAsync(setup.manager::close); + assertTrue(setup.manager.closeEntered.await(3, TimeUnit.SECONDS)); + assertThrows(TimeoutException.class, () -> closing.get(150, TimeUnit.MILLISECONDS)); + release.countDown(); + closing.get(3, TimeUnit.SECONDS); + assertEquals(2, calls.get()); + assertEquals( + OperationStatus.SUCCEEDED, + client.getOperationByName("condition").status()); + var checkpoints = client.getOperationUpdates().size(); + var replay = setup(client, workers, new LinkedBlockingQueue<>(), new AtomicBoolean()); + try { + assertEquals( + 2, + replay.context.waitForCondition( + "condition", + Integer.class, + (state, step) -> { + calls.incrementAndGet(); + throw new AssertionError("A completed predicate must not run on replay"); + }, + waitConfig())); + assertEquals(2, calls.get()); + assertEquals(checkpoints, client.getOperationUpdates().size()); + assertEquals(0, late.get()); + } finally { + replay.manager.close(); + } + setup.manager + .runCheckpointContinuation(() -> fail("New work must be rejected after close")) + .get(3, TimeUnit.SECONDS); + assertTrue(queue.isEmpty()); + } finally { + release.countDown(); + setup.manager.close(); + workers.shutdownNow(); + assertTrue(workers.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + @Test + void earlyParallelCompletionStopsOnlyItsQueuedConditionAndUnblocksTheWaitingBranch() throws Exception { + var closed = new AtomicBoolean(); + var late = new AtomicInteger(); + var client = readyClient(closed, late); + var queue = new LinkedBlockingQueue(); + var workers = Executors.newCachedThreadPool(); + var created = new CountDownLatch(1); + var calls = new AtomicInteger(); + var condition = new AtomicReference(); + var setup = setup(client, workers, queue, closed); + try { + var parallel = setup.context.parallel( + "early", + ParallelConfig.builder() + .completionConfig(CompletionConfig.minSuccessful(1)) + .build()); + var task = new AtomicReference(); + try (parallel) { + parallel.branch("waiting", Integer.class, branch -> { + var future = branch.waitForConditionAsync( + "condition", + Integer.class, + (state, step) -> { + calls.incrementAndGet(); + return WaitForConditionResult.continuePolling(state + 1); + }, + waitConfig()); + condition.set((BaseDurableOperation) future); + created.countDown(); + return future.get(); + }); + parallel.branch("winner", String.class, branch -> { + await(created); + try { + task.set(queue.poll(3, TimeUnit.SECONDS)); + assertNotNull(task.get()); + condition.get().getRunningUserHandler().get(3, TimeUnit.SECONDS); + } catch (Exception failure) { + throw new AssertionError(failure); + } + return "winner"; + }); + } + var result = parallel.get(); + assertEquals(1, result.succeeded()); + assertFalse(condition.get().getCompletionFuture().isDone()); + var checkpointCount = client.getOperationUpdates().size(); + var closing = CompletableFuture.runAsync(setup.manager::close); + assertTrue(setup.manager.closeEntered.await(3, TimeUnit.SECONDS)); + assertThrows(TimeoutException.class, () -> closing.get(150, TimeUnit.MILLISECONDS)); + CompletableFuture.runAsync(task.get()).get(3, TimeUnit.SECONDS); + closing.get(3, TimeUnit.SECONDS); + assertTrue(condition.get().getCompletionFuture().isCompletedExceptionally()); + assertEquals(1, calls.get()); + assertEquals(0, late.get()); + assertEquals(checkpointCount, client.getOperationUpdates().size()); + assertEquals(1, result.succeeded(), "Cleanup must not rewrite the stored early result"); + } finally { + assertThrows(SuspendExecutionException.class, setup.manager::suspendExecution); + Runnable pending; + while ((pending = queue.poll()) != null) pending.run(); + workers.shutdownNow(); + assertTrue(workers.awaitTermination(3, TimeUnit.SECONDS)); + setup.manager.close(); + } + } + + private record Setup(ControlledManager manager, DurableContextImpl context) {} + + private Setup setup( + LocalMemoryExecutionClient client, + ExecutorService workers, + BlockingQueue queue, + AtomicBoolean closed) { + var config = DurableConfig.builder() + .withDurableExecutionClient(client) + .withExecutorService(workers) + .withCheckpointDelay(Duration.ZERO) + .withPollingStrategy(attempt -> Duration.ZERO) + .build(); + var manager = new ControlledManager(input(client), config, queue, closed); + rememberContext(manager); + manager.registerActiveThread(null); + manager.setCurrentThreadContext(new ThreadContext(null, ThreadType.CONTEXT)); + var context = DurableContextImpl.createRootContext(manager, config, null); + BaseContextImpl.setCurrentContext(context); + return new Setup(manager, context); + } + + private static WaitForConditionConfig waitConfig() { + return WaitForConditionConfig.builder() + .initialState(1) + .waitStrategy((state, attempt) -> Duration.ofSeconds(1)) + .build(); + } + + private static void await(CountDownLatch latch) { + try { + assertTrue(latch.await(3, TimeUnit.SECONDS)); + } catch (InterruptedException failure) { + Thread.currentThread().interrupt(); + throw new AssertionError(failure); + } + } + + private static LocalMemoryExecutionClient readyClient(AtomicBoolean closed, AtomicInteger late) { + return new LocalMemoryExecutionClient() { + public CheckpointDurableExecutionResponse checkpoint( + String arn, String token, List updates) { + if (closed.get()) late.incrementAndGet(); + var applied = super.checkpoint(arn, token, updates); + advanceTime(); + var ready = super.checkpoint(arn, applied.checkpointToken(), List.of()); + var states = new LinkedHashMap(); + applied.newExecutionState().operations().forEach(op -> states.put(op.id(), op)); + ready.newExecutionState().operations().forEach(op -> states.put(op.id(), op)); + return ready.toBuilder() + .newExecutionState(CheckpointUpdatedExecutionState.builder() + .operations(states.values()) + .build()) + .build(); + } + }; + } + + @ParameterizedTest + @ValueSource(strings = {"queued", "cancelled-observation", "after-old-worker-join", "already-suspended"}) + void normalCloseQuiescesAdmittedReadyWork(String mode) throws Exception { + var closed = new AtomicBoolean(); + var lateBackend = new AtomicInteger(); + var predicates = new AtomicInteger(); + var queue = new LinkedBlockingQueue(); + var workers = new GateExecutor(mode.equals("after-old-worker-join")); + var client = new LocalMemoryExecutionClient() { + public CheckpointDurableExecutionResponse checkpoint( + String arn, String token, List updates) { + if (closed.get()) lateBackend.incrementAndGet(); + var applied = super.checkpoint(arn, token, updates); + advanceTime(); + var ready = super.checkpoint(arn, applied.checkpointToken(), List.of()); + var states = new LinkedHashMap(); + applied.newExecutionState().operations().forEach(op -> states.put(op.id(), op)); + ready.newExecutionState().operations().forEach(op -> states.put(op.id(), op)); + return ready.toBuilder() + .newExecutionState(CheckpointUpdatedExecutionState.builder() + .operations(states.values()) + .build()) + .build(); + } + }; + var config = DurableConfig.builder() + .withDurableExecutionClient(client) + .withExecutorService(workers) + .withCheckpointDelay(Duration.ZERO) + .withPollingStrategy(attempt -> Duration.ZERO) + .build(); + var manager = new ControlledManager(input(client), config, queue, closed); + rememberContext(manager); + manager.registerActiveThread(null); + manager.setCurrentThreadContext(new ThreadContext(null, ThreadType.CONTEXT)); + var context = DurableContextImpl.createRootContext(manager, config, null); + BaseContextImpl.setCurrentContext(context); + var body = new CompletableFuture(); + var selected = manager.runUntilCompleteOrSuspend(body); + try { + context.waitForConditionAsync( + "condition", + Integer.class, + (state, step) -> { + predicates.incrementAndGet(); + return state == 1 + ? WaitForConditionResult.continuePolling(2) + : WaitForConditionResult.stopPolling(state); + }, + WaitForConditionConfig.builder() + .initialState(1) + .waitStrategy((state, attempt) -> Duration.ofSeconds(1)) + .build()); + var task = queue.poll(3, TimeUnit.SECONDS); + assertNotNull(task); + assertTrue(workers.firstFinished.await(3, TimeUnit.SECONDS)); + if (mode.equals("cancelled-observation")) assertTrue(manager.observation.cancel(false)); + CompletableFuture taskRun = null; + if (mode.equals("after-old-worker-join")) { + taskRun = CompletableFuture.runAsync(task); + assertTrue(workers.nextDispatchEntered.await(3, TimeUnit.SECONDS)); + } + var stopped = mode.equals("already-suspended"); + if (stopped) assertThrows(SuspendExecutionException.class, manager::suspendExecution); + else { + body.complete("root-done"); + assertEquals("root-done", selected.get(3, TimeUnit.SECONDS)); + } + var closing = CompletableFuture.runAsync(manager::close); + assertTrue(manager.closeEntered.await(3, TimeUnit.SECONDS)); + boolean returnedWhileHeld; + try { + closing.get(150, TimeUnit.MILLISECONDS); + returnedWhileHeld = true; + } catch (TimeoutException expected) { + returnedWhileHeld = false; + } + workers.releaseNextDispatch.countDown(); + if (taskRun == null) taskRun = CompletableFuture.runAsync(task); + taskRun.get(3, TimeUnit.SECONDS); + closing.get(3, TimeUnit.SECONDS); + workers.shutdown(); + assertTrue(workers.awaitTermination(3, TimeUnit.SECONDS)); + var returnedEarly = returnedWhileHeld; + System.out.println("CONTINUATION_CLOSE mode=" + mode + " returnedWhileHeld=" + returnedEarly + + " predicateCalls=" + predicates.get() + " lateCheckpointAttempts=" + + manager.lateCheckpointAttempts.get() + " lateBackendCalls=" + lateBackend.get()); + assertAll( + () -> { + if (!stopped) assertFalse(returnedEarly, "Close must wait for the actual admitted task"); + }, + () -> assertEquals(1, predicates.get(), "No new predicate may run after teardown"), + () -> assertEquals(0, manager.lateCheckpointAttempts.get()), + () -> assertEquals(0, lateBackend.get())); + if (!stopped) assertEquals("root-done", selected.join()); + } finally { + workers.releaseNextDispatch.countDown(); + Runnable pending; + while ((pending = queue.poll()) != null) pending.run(); + workers.shutdownNow(); + assertTrue(workers.awaitTermination(3, TimeUnit.SECONDS)); + manager.close(); + } + } + + private static final class ControlledManager extends ExecutionManager { + private final BlockingQueue queue; + private final AtomicBoolean closed; + private final CountDownLatch closeEntered = new CountDownLatch(1); + private final AtomicInteger lateCheckpointAttempts = new AtomicInteger(); + private volatile CompletableFuture observation; + + private ControlledManager( + DurableExecutionInput input, + DurableConfig config, + BlockingQueue queue, + AtomicBoolean closed) { + super(input, config, null); + this.queue = queue; + this.closed = closed; + } + + public CompletableFuture runCheckpointContinuation(BaseDurableOperation owner, Runnable task) { + observation = super.runCheckpointContinuation(owner, task, queue::add); + return observation; + } + + public CompletableFuture sendOperationUpdate(OperationUpdate update) { + if (closed.get()) lateCheckpointAttempts.incrementAndGet(); + return super.sendOperationUpdate(update); + } + + public void close() { + closeEntered.countDown(); + super.close(); + closed.set(true); + } + } + + private static final class GateExecutor extends ThreadPoolExecutor { + private final AtomicInteger submitted = new AtomicInteger(); + private final AtomicInteger finished = new AtomicInteger(); + private final boolean hold; + private final CountDownLatch firstFinished = new CountDownLatch(1); + private final CountDownLatch nextDispatchEntered = new CountDownLatch(1); + private final CountDownLatch releaseNextDispatch = new CountDownLatch(1); + + private GateExecutor(boolean hold) { + super(1, 1, 0, TimeUnit.SECONDS, new LinkedBlockingQueue<>()); + this.hold = hold; + } + + public void execute(Runnable task) { + if (submitted.incrementAndGet() == 2 && hold) { + nextDispatchEntered.countDown(); + try { + assertTrue(releaseNextDispatch.await(3, TimeUnit.SECONDS)); + } catch (InterruptedException failure) { + Thread.currentThread().interrupt(); + throw new AssertionError(failure); + } + } + super.execute(task); + } + + protected void afterExecute(Runnable task, Throwable failure) { + if (finished.incrementAndGet() == 1) firstFinished.countDown(); + } + } + + private static DurableExecutionInput input(LocalMemoryExecutionClient client) { + var states = new ArrayList<>(client.getAllOperations()); + states.add( + 0, + Operation.builder() + .id("execution") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()); + return new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/close/execution", + "token", + CheckpointUpdatedExecutionState.builder().operations(states).build()); + } +} 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 9a61c0224..9f260a688 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 @@ -7,6 +7,7 @@ import java.util.ArrayList; import java.util.Collection; import java.util.Collections; +import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; @@ -78,6 +79,9 @@ public class ExecutionManager implements SafeCloseable { private final CompletableFuture executionExceptionFuture = new CompletableFuture<>(); // Requests and polling continuations; guarded by activeThreads so admission is atomic with last-thread suspension. private int checkpointRequestsInFlight; + // Guarded by activeThreads; an observation future can be cancelled independently of its actual task. + private final Map checkpointContinuations = new HashMap<>(); + private boolean closing; /** * Per-wait state used to coordinate the caller thread with the future completion callback. @@ -554,34 +558,85 @@ void finishCheckpointProcessing() { * @return completion of the continuation, or an already-completed future when execution has stopped */ public CompletableFuture runCheckpointContinuation(Runnable continuation) { - return runCheckpointContinuation(continuation, InternalExecutor.INSTANCE); + return runCheckpointContinuation(null, continuation, InternalExecutor.INSTANCE); + } + + /** Schedules an operation continuation whose unfinished waiters are stopped during manager cleanup. */ + public CompletableFuture runCheckpointContinuation(BaseDurableOperation owner, Runnable continuation) { + return runCheckpointContinuation(owner, continuation, InternalExecutor.INSTANCE); } CompletableFuture runCheckpointContinuation(Runnable continuation, Executor coordinator) { - if (!tryStartCheckpointProcessing()) return CompletableFuture.completedFuture(null); + return runCheckpointContinuation(null, continuation, coordinator); + } + + CompletableFuture runCheckpointContinuation( + BaseDurableOperation owner, Runnable continuation, Executor coordinator) { + var registration = registerCheckpointContinuation(owner); + if (registration == null) { + if (isClosing()) stopContinuationOwner(owner); + return CompletableFuture.completedFuture(null); + } try { var completion = new CompletableFuture(); - coordinator.execute(() -> { - try { - try { - continuation.run(); - } finally { - finishCheckpointProcessing(); - } - completion.complete(null); - } catch (Throwable failure) { - completion.completeExceptionally(failure); - } - }); + coordinator.execute(() -> completeCheckpointContinuation(registration, continuation, completion)); return completion; } catch (RuntimeException | Error failure) { - finishCheckpointProcessing(); + finishCheckpointContinuation(registration); throw failure; } } + private void completeCheckpointContinuation( + Object registration, Runnable continuation, CompletableFuture completion) { + try { + try { + if (!isClosing()) continuation.run(); + } finally { + finishCheckpointContinuation(registration); + } + completion.complete(null); + } catch (Throwable failure) { + completion.completeExceptionally(failure); + } + } + + private Object registerCheckpointContinuation(BaseDurableOperation owner) { + synchronized (activeThreads) { + if (closing || executionExceptionFuture.isDone()) return null; + var registration = new Object(); + checkpointContinuations.put(registration, owner); + checkpointRequestsInFlight++; + return registration; + } + } + + private void finishCheckpointContinuation(Object registration) { + synchronized (activeThreads) { + checkpointContinuations.remove(registration); + try { + finishCheckpointProcessing(); + } finally { + activeThreads.notifyAll(); + } + } + } + + private boolean isClosing() { + synchronized (activeThreads) { + return closing; + } + } + + private static void stopContinuationOwner(BaseDurableOperation owner) { + if (owner != null) owner.getCompletionFuture().completeExceptionally(new SuspendExecutionException()); + } + private boolean shouldSuspendExecution() { - return activeThreads.isEmpty() && checkpointRequestsInFlight == 0 && !executionExceptionFuture.isDone(); + return !closing + && activeThreads.isEmpty() + && checkpointRequestsInFlight == 0 + && !executionExceptionFuture.isDone(); } private void preSuspendCheck() { @@ -631,11 +686,34 @@ public CompletableFuture pollForOperationUpdates(String operationId, /** Shutdown the checkpoint batcher. */ @Override public void close() { + stopCheckpointContinuations(); validateRunningThreads(); checkpointManager.shutdown(); } + private void stopCheckpointContinuations() { + List owners; + synchronized (activeThreads) { + closing = true; + owners = new ArrayList<>(checkpointContinuations.values()); + } + owners.forEach(ExecutionManager::stopContinuationOwner); + synchronized (activeThreads) { + while (!checkpointContinuations.isEmpty()) { + try { + // wait releases this monitor so tasks can finish and release their registrations. No user work + // or future join is performed while holding the coordination monitor. + activeThreads.wait(); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw new IllegalStateException( + "Interrupted while waiting for checkpoint continuations", interrupted); + } + } + } + } + private void validateRunningThreads() { // This will detect stuck user thread and thread leaks in the thread pool for (BaseDurableOperation op : registeredOperations.values()) { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java index 36e9e986f..40eb517ec 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java @@ -111,7 +111,7 @@ private void resumeCheckLoop(Operation existing) { private CompletableFuture pollReadyAndResumeCheckLoop(Operation existing) { return pollUntilReady() - .thenCompose(op -> executionManager.runCheckpointContinuation(() -> { + .thenCompose(op -> executionManager.runCheckpointContinuation(this, () -> { if (!isOperationCompleted() && op.status() == OperationStatus.READY) resumeCheckLoop(op); })); } @@ -187,7 +187,7 @@ private void runCheckLoop(T currentState, int attempt, CompletableFuture> publishedWorker, T nextState, int nextAttempt) { pollUntilReady() - .thenCompose(op -> executionManager.runCheckpointContinuation(() -> { + .thenCompose(op -> executionManager.runCheckpointContinuation(this, () -> { publishedWorker.join().join(); if (!isOperationCompleted() && op.status() == OperationStatus.READY) executeCheckLogic(nextState, nextAttempt); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java index dfb67639e..65d961f36 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java @@ -485,6 +485,67 @@ void rejectedContinuationReleasesActivityAndPreservesTheRejection() { assertThrows(SuspendExecutionException.class, () -> manager.deregisterActiveThread("root")); } + @ParameterizedTest + @ValueSource(strings = {"queued", "cancelled-observation", "rejected"}) + void admissionRacingCloseDrainsActualWorkAndPreservesUnrelatedOperations(String mode) throws Exception { + var manager = createManager(List.of(executionOp(), stepOp("step", OperationStatus.PENDING))); + manager.registerActiveThread("root"); + var owner = mock(BaseDurableOperation.class); + var ownerCompletion = new CompletableFuture(); + when(owner.getCompletionFuture()).thenReturn(ownerCompletion); + var unrelated = mock(BaseDurableOperation.class); + var unrelatedCompletion = new CompletableFuture(); + when(unrelated.getOperationId()).thenReturn("unrelated"); + when(unrelated.getCompletionFuture()).thenReturn(unrelatedCompletion); + manager.registerOperation(unrelated); + var dispatchEntered = new CountDownLatch(1); + var releaseDispatch = new CountDownLatch(1); + var queued = new AtomicReference(); + var rejection = new RejectedExecutionException("rejected during close"); + var admission = CompletableFuture.supplyAsync(() -> + manager.runCheckpointContinuation(owner, () -> fail("Closing must stop this admitted body"), task -> { + dispatchEntered.countDown(); + try { + assertTrue(releaseDispatch.await(3, TimeUnit.SECONDS)); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw new AssertionError(interrupted); + } + if (mode.equals("rejected")) throw rejection; + queued.set(task); + })); + assertTrue(dispatchEntered.await(3, TimeUnit.SECONDS)); + var closing = CompletableFuture.runAsync(manager::close); + try { + assertInstanceOf( + SuspendExecutionException.class, + assertThrows(ExecutionException.class, () -> ownerCompletion.get(3, TimeUnit.SECONDS)) + .getCause()); + assertThrows(TimeoutException.class, () -> closing.get(100, TimeUnit.MILLISECONDS)); + } finally { + releaseDispatch.countDown(); + } + if (mode.equals("rejected")) { + assertSame( + rejection, + assertThrows(ExecutionException.class, () -> admission.get(3, TimeUnit.SECONDS)) + .getCause()); + } else { + var observation = admission.get(3, TimeUnit.SECONDS); + if (mode.equals("cancelled-observation")) assertTrue(observation.cancel(false)); + assertThrows(TimeoutException.class, () -> closing.get(100, TimeUnit.MILLISECONDS)); + queued.get().run(); + } + closing.get(3, TimeUnit.SECONDS); + assertFalse(unrelatedCompletion.isDone(), "Normal close must not apply global stopAllOperations"); + assertFalse(manager.isExecutionCompletedExceptionally(), "Cleanup must not replace the selected root outcome"); + manager.runCheckpointContinuation( + () -> fail("Post-close admission must be rejected"), + task -> fail("Post-close work must not reach the executor")) + .get(3, TimeUnit.SECONDS); + assertDoesNotThrow(() -> manager.deregisterActiveThread("root")); + } + @Test void checkpointDeliveryIsAtomicWithOperationRegistration() throws Exception { var manager = createManager(List.of(executionOp(), stepOp("step", OperationStatus.PENDING))); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionOperationTest.java index c4df8aac5..39ea05e8d 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionOperationTest.java @@ -284,8 +284,8 @@ void getFailedWithNullErrorDataThrowsStepFailedException() { @Test void replayPendingPollsAndResumesCheckLoop() throws Exception { - when(executionManager.runCheckpointContinuation(any())) - .thenAnswer(call -> CompletableFuture.runAsync(call.getArgument(0))); + when(executionManager.runCheckpointContinuation(any(), any(Runnable.class))) + .thenAnswer(call -> CompletableFuture.runAsync(call.getArgument(1))); var pendingOp = Operation.builder() .id(OPERATION_ID) .name(OPERATION_NAME) From c0d7e060cdac84d1d95afb7ec19c7a247845d294 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 06:41:39 +0000 Subject: [PATCH 36/48] fix: retain live operation parents for native span clocks Code-only port of 9760d9b5666794256da34e4b8a1b523b84a71571 after actual-source clock negatives; preserve cached context fallback and all workflow/lifecycle/metadata settings. --- .../durable/otel/InvocationOtelPlugin.java | 5 + .../otel/InvocationClockContainmentTest.java | 145 ++++++++++++++++++ 2 files changed, 150 insertions(+) create mode 100644 otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationClockContainmentTest.java diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java index c00c07508..1cba31fd9 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java @@ -654,6 +654,11 @@ private Context invocationParentContext(ExecutionTraceContext execCtx, String ca private Context resolveParentContext(String parentId) { if (parentId != null) { + var parentSpan = operationSpans.get(parentId); + if (parentSpan != null) { + // Retain the provider's live span, including its clock, while this parent is open. + return withDurableDecision(Context.current().with(parentSpan)); + } var parentSpanContext = operationContexts.get(parentId); if (parentSpanContext != null) { return withDurableDecision(Context.current().with(Span.wrap(parentSpanContext))); diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationClockContainmentTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationClockContainmentTest.java new file mode 100644 index 000000000..da378fd86 --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationClockContainmentTest.java @@ -0,0 +1,145 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.otel; + +import static org.junit.jupiter.api.Assertions.*; + +import io.opentelemetry.sdk.common.Clock; +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.atomic.AtomicLong; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.EnumSource; +import software.amazon.lambda.durable.plugin.InvocationEndInfo; +import software.amazon.lambda.durable.plugin.InvocationInfo; +import software.amazon.lambda.durable.plugin.InvocationStatus; +import software.amazon.lambda.durable.plugin.OperationEndInfo; +import software.amazon.lambda.durable.plugin.OperationInfo; + +class InvocationClockContainmentTest { + @ParameterizedTest + @EnumSource(InvocationStatus.class) + void openChildAndParentShareTheConfiguredSdkClock(InvocationStatus status) { + var wall = new AtomicLong(100_000_000_000L); + var monotonic = new AtomicLong(); + var clock = new Clock() { + public long now() { + return wall.addAndGet(1_000_000); + } + + public long nanoTime() { + return monotonic.addAndGet(1_000); + } + }; + try (var exporter = InMemorySpanExporter.create()) { + var plugin = new InvocationOtelPlugin( + SdkTracerProvider.builder().setClock(clock).addSpanProcessor(SimpleSpanProcessor.create(exporter)), + OtelPluginConfig.builder() + .contextExtractor(() -> null) + .enableMdc(false) + .build()); + var start = Instant.ofEpochSecond(10); + var arn = "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/clock/execution"; + plugin.onInvocationStart(new InvocationInfo("request", arn, true, start)); + plugin.onOperationStart(new OperationInfo( + "parent", "parent", "CONTEXT", "WaitForCallback", null, start, null, null, false)); + plugin.onOperationStart( + new OperationInfo("child", "child", "CALLBACK", "Callback", "parent", start, null, null, false)); + plugin.onInvocationEnd(new InvocationEndInfo( + "request", + arn, + true, + status, + status == InvocationStatus.FAILED ? new IllegalStateException("failure") : null)); + var spans = exporter.getFinishedSpanItems(); + var child = spans.stream() + .filter(s -> s.getName().equals("child")) + .findFirst() + .orElseThrow(); + var parent = spans.stream() + .filter(s -> s.getName().equals("parent")) + .findFirst() + .orElseThrow(); + var invocation = spans.stream() + .filter(s -> s.getName().equals("Invocation")) + .findFirst() + .orElseThrow(); + assertEquals(parent.getSpanId(), child.getParentSpanId()); + assertEquals( + parent.getSpanContext(), + child.getParentSpanContext(), + "Live parent keeps the cached trace ID, span ID, flags and trace state"); + assertEquals(invocation.getSpanId(), parent.getParentSpanId()); + assertTrue(spans.indexOf(child) < spans.indexOf(parent), "The plugin already drains children first"); + assertTrue( + child.getEndEpochNanos() <= parent.getEndEpochNanos(), + "Separate anchor clocks must not invert child-first end timestamps"); + assertTrue(parent.getEndEpochNanos() <= invocation.getEndEpochNanos()); + assertTrue(child.getStartEpochNanos() >= parent.getStartEpochNanos()); + assertTrue( + invocation.getEndEpochNanos() < 101_000_000_000L, + "Keep the configured clock; do not substitute system time or relax containment"); + } + } + + @Test + void endedParentRetainsItsCachedContextForAReplayedChildWithoutReexport() { + try (var exporter = InMemorySpanExporter.create()) { + var plugin = new InvocationOtelPlugin( + SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)), + OtelPluginConfig.builder() + .contextExtractor(() -> null) + .enableMdc(false) + .build()); + var start = Instant.now(); + var arn = "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/fallback/execution"; + plugin.onInvocationStart(new InvocationInfo("request", arn, false, start)); + plugin.onOperationStart(new OperationInfo( + "parent", "parent", "CONTEXT", "RunInChildContext", null, start, null, null, true)); + plugin.onOperationEnd(new OperationEndInfo( + "parent", + "parent", + "CONTEXT", + "RunInChildContext", + null, + start, + Instant.now(), + "SUCCEEDED", + null, + true, + null, + null)); + var parent = exporter.getFinishedSpanItems().stream() + .filter(s -> s.getName().equals("parent")) + .findFirst() + .orElseThrow(); + plugin.onOperationEnd(new OperationEndInfo( + "child", + "child", + "CALLBACK", + "Callback", + "parent", + start, + Instant.now(), + "SUCCEEDED", + null, + true, + null, + null)); + plugin.onInvocationEnd(new InvocationEndInfo("request", arn, false, InvocationStatus.PENDING, null)); + var spans = exporter.getFinishedSpanItems(); + var child = spans.stream() + .filter(s -> s.getName().equals("child")) + .findFirst() + .orElseThrow(); + assertEquals(parent.getSpanContext(), child.getParentSpanContext()); + assertEquals( + 1, spans.stream().filter(s -> s.getName().equals("parent")).count()); + assertEquals( + 1, spans.stream().filter(s -> s.getName().equals("child")).count()); + } + } +} From c842f5a599ee3b028ce0c18be3caf3a188dbfab8 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 16:43:30 +0000 Subject: [PATCH 37/48] fix: preserve preparation failures and isolate sampling carriers --- otel-plugin/README.md | 4 +- .../durable/otel/ExecutionOtelPlugin.java | 4 +- .../durable/otel/InvocationOtelPlugin.java | 4 +- .../durable/otel/OtelPluginSupport.java | 6 + .../compatibility/b1/InstalledApiProbe.java | 4 +- .../otel/InvocationOutcomeBoundaryTest.java | 108 ++++++++++++++++++ .../otel/SamplingProcessorIsolationTest.java | 13 ++- .../durable/execution/DurableExecutor.java | 23 +++- 8 files changed, 155 insertions(+), 11 deletions(-) diff --git a/otel-plugin/README.md b/otel-plugin/README.md index b989b0dbf..4ca331989 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -33,7 +33,9 @@ JVM-fatal precedence, so an ordinary cleanup exception cannot hide an earlier Er Invocations without plugins also wait for handler cleanup before returning the selected suspension or retry outcome. SDK output preparation, including customer `SerDes` calls and durable large-result checkpointing, finishes before terminal invocation-end notification. Failures in this preparation report `RETRYING` instead of ending the Workflow -span. End describes the SDK outcome at that point, not acknowledgment of a response by the Lambda service. +span. If End also raises an unisolated Error, the preparation failure remains primary with the cleanup error +suppressed, unless cleanup introduces the first JVM-fatal error. An original JVM-fatal preparation error retains its +identity. End describes the SDK outcome at that point, not acknowledgment of a response by the Lambda service. Caller-side execution-manager cleanup, response-envelope encoding and output-stream writes follow End; runtime response transport follows the handler return. Failures at those later boundaries still propagate, without a second End dispatch or changing its already reported outcome. Non-JVM-fatal MDC-restoration failures on an inline executor, diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java index 79c448ff4..6f9a736fa 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java @@ -682,7 +682,7 @@ private TraceState effectiveTraceState() { */ private Span startDurableSpan(SpanBuilder spanBuilder) { var intent = samplingIntent; - if (intent == null) { + if (intent == null || !OtelPluginSupport.usesDurableSamplingBridge(sdkTracerProvider)) { return spanBuilder.startSpan(); } try (var ignored = DurableSamplingDecision.openScope(intent)) { @@ -693,7 +693,7 @@ private Span startDurableSpan(SpanBuilder spanBuilder) { /** Starts a durable span with a forced span ID, publishing the sampling intent as in {@link #startDurableSpan}. */ private Span startDurableSpan(SpanBuilder spanBuilder, String traceId, String spanId) { var intent = samplingIntent; - if (intent == null) { + if (intent == null || !OtelPluginSupport.usesDurableSamplingBridge(sdkTracerProvider)) { return idGenerator.startSpan(spanBuilder, traceId, spanId); } try (var ignored = DurableSamplingDecision.openScope(intent)) { diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java index 1cba31fd9..4e2e3cd70 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java @@ -691,7 +691,7 @@ private Context withDurableDecision(Context context) { */ private Span startDurableSpan(SpanBuilder spanBuilder) { var intent = samplingIntent; - if (intent == null) { + if (intent == null || !OtelPluginSupport.usesDurableSamplingBridge(sdkTracerProvider)) { return spanBuilder.startSpan(); } try (var ignored = DurableSamplingDecision.openScope(intent)) { @@ -702,7 +702,7 @@ private Span startDurableSpan(SpanBuilder spanBuilder) { /** Starts a durable span with a forced span ID, publishing the sampling intent as in {@link #startDurableSpan}. */ private Span startDurableSpan(SpanBuilder spanBuilder, String traceId, String spanId) { var intent = samplingIntent; - if (intent == null) { + if (intent == null || !OtelPluginSupport.usesDurableSamplingBridge(sdkTracerProvider)) { return idGenerator.startSpan(spanBuilder, traceId, spanId); } try (var ignored = DurableSamplingDecision.openScope(intent)) { diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java index f06b8b787..75c83d7a5 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java @@ -31,6 +31,12 @@ static boolean usesLocalDurableSampler(SdkTracerProvider provider) { return provider != null && provider.getSampler() instanceof DurableSampler; } + /** A visible replacement sampler cannot consume our bridge; an opaque agent provider may still need it. */ + static boolean usesDurableSamplingBridge(SdkTracerProvider provider) { + // Class-name matching is only for the cross-loader wire carrier, never for the context-key ownership check. + return provider == null || provider.getSampler().getClass().getName().equals(DurableSampler.class.getName()); + } + /** Closes the owning thread's scope without hiding a finalization error or a later JVM-fatal cleanup failure. */ static void runInvocationEnd(Scope scope, Runnable end) { Throwable primary = null; diff --git a/otel-plugin/src/test/compatibility/b1/InstalledApiProbe.java b/otel-plugin/src/test/compatibility/b1/InstalledApiProbe.java index 122fca380..26721887b 100644 --- a/otel-plugin/src/test/compatibility/b1/InstalledApiProbe.java +++ b/otel-plugin/src/test/compatibility/b1/InstalledApiProbe.java @@ -2,8 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.otel; -import ch.qos.logback.classic.Logger; import ch.qos.logback.classic.Level; +import ch.qos.logback.classic.Logger; import ch.qos.logback.classic.spi.ILoggingEvent; import ch.qos.logback.core.read.ListAppender; import io.opentelemetry.api.GlobalOpenTelemetry; @@ -14,10 +14,10 @@ import io.opentelemetry.sdk.trace.SdkTracerProvider; import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; import java.nio.file.Path; -import java.util.concurrent.Callable; import java.time.Duration; import java.util.List; import java.util.ServiceLoader; +import java.util.concurrent.Callable; import java.util.concurrent.atomic.AtomicInteger; import org.slf4j.LoggerFactory; import software.amazon.lambda.durable.model.ExecutionStatus; diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOutcomeBoundaryTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOutcomeBoundaryTest.java index 14185265f..f663951ab 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOutcomeBoundaryTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOutcomeBoundaryTest.java @@ -46,6 +46,7 @@ import software.amazon.lambda.durable.serde.JacksonSerDes; import software.amazon.lambda.durable.serde.SerDes; import software.amazon.lambda.durable.testing.LocalDurableTestRunner; +import software.amazon.lambda.durable.util.ExceptionHelper; class InvocationOutcomeBoundaryTest { @ParameterizedTest @@ -414,6 +415,113 @@ public void onInvocationEnd(InvocationEndInfo info) { .count()); } + @SuppressWarnings("removal") + @ParameterizedTest + @CsvSource({ + "ordinary,assertion,false,false", "ordinary,assertion,true,false", + "vm,assertion,false,false", "vm,assertion,true,false", + "death,assertion,false,false", "death,assertion,true,false", + "ordinary,vm,false,false", "ordinary,vm,true,false", + "vm,death,false,false", "vm,death,true,false", + "wrapped,assertion,false,false", "wrapped,assertion,true,false", + "ordinary,assertion,false,true", "ordinary,assertion,true,true", + "vm,assertion,false,true", "vm,assertion,true,true", + "death,assertion,false,true", "death,assertion,true,true", + "ordinary,vm,false,true", "ordinary,vm,true,true", + "vm,death,false,true", "vm,death,true,true", + "wrapped,assertion,false,true", "wrapped,assertion,true,true" + }) + void preparationAndEndFailuresRetainPrimaryAndFatalIdentity( + String preparationKind, String endKind, boolean serializeFailure, boolean inline) throws Exception { + Throwable original = + switch (preparationKind) { + case "vm", "wrapped" -> new InternalError("preparation fatal"); + case "death" -> new ThreadDeath(); + default -> new IllegalStateException("preparation failure"); + }; + var delivery = preparationKind.equals("wrapped") ? new CompletionException(original) : original; + Error endFailure = + switch (endKind) { + case "vm" -> new InternalError("end fatal"); + case "death" -> new ThreadDeath(); + default -> new AssertionError("end failure"); + }; + var expected = preparationKind.equals("ordinary") && endKind.equals("vm") ? endFailure : original; + var suppressed = expected == endFailure ? original : endFailure; + var ends = new CopyOnWriteArrayList(); + var bodyFailure = new IllegalArgumentException("body failure"); + var escaped = new CountDownLatch(1); + var workerFailure = new AtomicReference(); + var callers = Executors.newSingleThreadExecutor(); + var workers = + new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS, new SynchronousQueue(), task -> { + var thread = new Thread(task, "combined-failure-owner"); + thread.setUncaughtExceptionHandler((owner, failure) -> { + workerFailure.set(failure); + escaped.countDown(); + }); + return thread; + }) { + @Override + public void execute(Runnable task) { + if (inline) task.run(); + else super.execute(task); + } + }; + var serDes = new SerDes() { + private final JacksonSerDes delegate = new JacksonSerDes(); + + @Override + public String serialize(Object value) { + if (serializeFailure ? value == bodyFailure : "done".equals(value)) { + ExceptionHelper.sneakyThrow(delivery); + } + return delegate.serialize(value); + } + + @Override + public T deserialize(String value, TypeToken type) { + return delegate.deserialize(value, type); + } + }; + try { + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + if (serializeFailure) throw bodyFailure; + return "done"; + }, + DurableConfig.builder() + .withExecutorService(workers) + .withSerDes(serDes) + .withPlugins(new DurableExecutionPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + ends.add(info); + throw endFailure; + } + }) + .build()); + var response = callers.submit(() -> runner.run("input")); + var observed = assertThrows(ExecutionException.class, () -> response.get(3, TimeUnit.SECONDS)) + .getCause(); + assertSame(expected, observed); + assertEquals(List.of(suppressed), List.of(observed.getSuppressed())); + assertEquals(1, ends.size()); + assertEquals(InvocationStatus.RETRYING, ends.get(0).invocationStatus()); + assertSame(original, ends.get(0).executionError()); + if (!inline && (expected instanceof VirtualMachineError || expected instanceof ThreadDeath)) { + assertTrue(escaped.await(3, TimeUnit.SECONDS)); + assertSame(expected, workerFailure.get()); + } + } finally { + callers.shutdownNow(); + workers.shutdownNow(); + assertTrue(callers.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS)); + } + } + @Test void partialStartStillUnwindsAllConfiguredEndHooksAndRestoresBothThreads() throws Exception { var key = ContextKey.named("partial-start-context"); diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/SamplingProcessorIsolationTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/SamplingProcessorIsolationTest.java index 863dacb01..79981a9b4 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/SamplingProcessorIsolationTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/SamplingProcessorIsolationTest.java @@ -57,10 +57,14 @@ void clearBridge() { "InvocationOtelPlugin,false,true,Invocation", "ExecutionOtelPlugin,true,false,Invocation", "ExecutionOtelPlugin,false,false,Invocation", - "ExecutionOtelPlugin,false,true,Invocation" + "ExecutionOtelPlugin,false,true,Invocation", + "InvocationOtelPlugin,false,false,PlainInvocation", + "ExecutionOtelPlugin,false,false,PlainInvocation" }) void agentProcessorForwardingParentCannotReuseApplicationSamplingIntent( String pluginName, boolean hideProvider, boolean localSampler, String observedSpan) throws Exception { + var plainSampler = observedSpan.equals("PlainInvocation"); + var spanName = plainSampler ? "Invocation" : observedSpan; var previousHeader = System.getProperty("com.amazonaws.xray.traceHeader"); GlobalOpenTelemetry.resetForTest(); OtelPluginAutoConfigurationState.markInstalled(); @@ -87,13 +91,16 @@ void agentProcessorForwardingParentCannotReuseApplicationSamplingIntent( .build(); var agentProvider = SdkTracerProvider.builder() .setSampler( - (Sampler) (localSampler ? appWrap : agentWrap).invoke(null, Sampler.alwaysOff())) + plainSampler + ? Sampler.alwaysOn() + : (Sampler) (localSampler ? appWrap : agentWrap) + .invoke(null, Sampler.alwaysOff())) .setIdGenerator( (IdGenerator) agentIdType.getConstructor().newInstance()) .addSpanProcessor(new SpanProcessor() { @Override public void onStart(Context parent, ReadWriteSpan span) { - if (!span.getName().equals(observedSpan)) return; + if (!span.getName().equals(spanName)) return; var unrelated = appProvider .get("processor") .spanBuilder("unrelated-forwarded-parent") 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 a3803c16f..3f054489d 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 @@ -203,7 +203,11 @@ private DurableExecutionOutput failDelivery(Throwable deliveryFailure) { // Serialization/checkpointing did not produce a terminal response. Close plugin resources with RETRYING // before preserving the original delivery failure for the Lambda caller. var cause = ExceptionHelper.unwrapCompletableFuture(deliveryFailure); - fireOnInvocationEnd(InvocationStatus.RETRYING, cause, null); + try { + fireOnInvocationEnd(InvocationStatus.RETRYING, cause, null); + } catch (Error endFailure) { + deliveryFailure = combinePreparationAndEndFailures(deliveryFailure, endFailure); + } ExceptionHelper.sneakyThrow(deliveryFailure); return null; } @@ -224,6 +228,23 @@ private void fireOnInvocationEnd(InvocationStatus status, Throwable error, Objec } } + /** Preserves preparation as primary unless cleanup introduces the first JVM-fatal failure. */ + @SuppressWarnings("removal") + private static Throwable combinePreparationAndEndFailures(Throwable preparation, Error cleanup) { + try { + rethrowLifecycleFatal(preparation); + } catch (VirtualMachineError | ThreadDeath fatal) { + if (fatal != cleanup) fatal.addSuppressed(cleanup); + return fatal; + } + if (cleanup instanceof VirtualMachineError || cleanup instanceof ThreadDeath) { + if (cleanup != preparation) cleanup.addSuppressed(preparation); + return cleanup; + } + if (preparation != cleanup) preparation.addSuppressed(cleanup); + return preparation; + } + private static InvocationStatus failureStatus(Throwable failure) { if (failure instanceof SuspendExecutionException) return InvocationStatus.PENDING; if (failure instanceof UnrecoverableDurableExecutionException unrecoverable && unrecoverable.isRetryable()) { From 478c183131454e220afd5752f5ff5c133a37202f Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 17:06:42 +0000 Subject: [PATCH 38/48] fix: preserve foreign sampling metadata and execution isolation --- otel-plugin/README.md | 2 + .../lambda/durable/otel/DurableSampler.java | 17 +- .../durable/otel/DurableSamplingDecision.java | 8 +- .../durable/otel/OtelPluginSupport.java | 7 +- ...urableSamplingDecisionClassLoaderTest.java | 147 ++++++++++++++++++ 5 files changed, 167 insertions(+), 14 deletions(-) diff --git a/otel-plugin/README.md b/otel-plugin/README.md index 4ca331989..ab452d553 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -319,6 +319,8 @@ For precise, provider-independent control, set an explicit `Sampled` value upstr | `durable.attempt.number` | 1-based attempt number | | `durable.attempt.outcome` | SUCCEEDED (span status `OK`), FAILED (`ERROR`), or INCOMPLETE (`UNSET`) | +Deferred sampler results are isolated by execution ARN and canonical trace ID in the existing 256-entry LRU cache. The delegate is reused while an entry is resident; eviction can cause another evaluation. The delegate still receives the canonical trace ID. + ## Log Correlation (MDC) When `enableMdc` is true (default), the plugin injects these fields into SLF4J MDC during user function execution: diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSampler.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSampler.java index 34e222ab3..78dc3d84e 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSampler.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSampler.java @@ -44,13 +44,15 @@ final class DurableSampler implements Sampler { private static final int MAX_CACHED_DEFERRED_DECISIONS = 256; private final Sampler delegate; - // Caches the delegate's decision for a deferred durable execution, keyed by canonical trace ID, so a stateful or - // quota-based delegate is consulted once per execution rather than per span. Access-ordered LRU, size-capped and - // synchronized (contention is low: at most one miss per execution). - private final Map deferredDecisions = + // Caches the delegate's decision for a deferred durable execution, keyed by execution ARN and canonical trace ID, + // so a stateful or quota-based delegate is reused while that execution entry remains resident. + // Access-ordered LRU, size-capped and synchronized; eviction permits another delegate evaluation. + private record ExecutionKey(String traceId, String executionArn) {} + + private final Map deferredDecisions = Collections.synchronizedMap(new LinkedHashMap<>(16, 0.75f, true) { @Override - protected boolean removeEldestEntry(Map.Entry eldest) { + protected boolean removeEldestEntry(Map.Entry eldest) { return size() > MAX_CACHED_DEFERRED_DECISIONS; } }); @@ -126,8 +128,9 @@ public SamplingResult shouldSample( // Deferred (agent path, real sampler not reproducible here): evaluate the actual delegate once per execution // and reuse it, so an installed drop/rate-limit policy is honored and consulted only once. return deferredDecisions.computeIfAbsent( - intent.deferredTraceId(), - key -> delegate.shouldSample(Context.root(), key, name, spanKind, attributes, Collections.emptyList())); + new ExecutionKey(intent.deferredTraceId(), attributes.get(SpanAttributes.DURABLE_EXECUTION_ARN)), + key -> delegate.shouldSample( + Context.root(), key.traceId(), name, spanKind, attributes, Collections.emptyList())); } @Override diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSamplingDecision.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSamplingDecision.java index 2341f8bdb..1d2de1217 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSamplingDecision.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSamplingDecision.java @@ -19,8 +19,8 @@ * preserving the full three-way decision (including {@code RECORD_ONLY}) and consulting no delegate; or *

  • a deferral marker (carrying the canonical trace ID) — used on the Java-agent path when the real sampler * (a remote/custom/file-only policy) cannot be reproduced here. {@link DurableSampler} then evaluates its actual - * delegate once per execution, caches the result by trace ID, and reuses it for the execution's remaining durable - * spans, so an installed drop/rate-limit policy is honored and consulted only once. + * delegate once per execution, caches the result by execution ARN and trace ID, and reuses it for the execution's + * remaining durable spans, so an installed drop/rate-limit policy is honored and consulted only once. * * *

    Two carriers, because the plugin runs across two class loaders. Under the documented ADOT setup the plugin @@ -49,10 +49,10 @@ final class DurableSamplingDecision { /** * The durable sampling intent for a span: either a resolved {@link SamplingResult}, or a deferral to the agent-side - * sampler's own delegate keyed by the canonical trace ID. + * sampler's own delegate. The sampler uses the execution ARN already attached to the span with this trace ID. * * @param resolved the resolved decision, or null when deferring - * @param deferredTraceId the canonical trace ID to key the agent-side delegate cache on, or null when resolved + * @param deferredTraceId the canonical trace ID passed to the agent-side delegate, or null when resolved */ record Intent(SamplingResult resolved, String deferredTraceId) { static Intent resolved(SamplingResult result) { diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java index 75c83d7a5..3f28220a6 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java @@ -111,7 +111,7 @@ static DeterministicIdGenerator createDefaultIdGenerator() { * pipeline finally installs, and another extension's customizer can wrap or replace a recognized configured * sampler, so a reconstruction could disagree with the real delegate. Deferring routes the decision to the * agent-installed {@link DurableSampler}, which consults its actual delegate once per execution, caches the - * result by trace ID, and reuses it for the execution's remaining durable spans (see + * result by execution ARN and trace ID, and reuses it for the execution's remaining durable spans (see * {@link DurableSampler#shouldSample}). The delegate's decision is honored in full — including a * {@code DROP}/rate-limited outcome — so durable spans are not force-sampled. * @@ -162,8 +162,9 @@ static SamplingResult resolveSamplingResult( } return ambientSpan.isRecording() ? SamplingResult.recordOnly() : SamplingResult.drop(); } - // 3. An application-owned provider exposes the real sampler: evaluate it once, preserving its full result. - if (sdkTracerProvider != null) { + // 3. Only a same-copy DurableSampler can consume the full resolved context carrier. Defer for a foreign + // sampler so its own delegate retains attributes and updated trace state, not just the bridge decision. + if (usesLocalDurableSampler(sdkTracerProvider)) { return sdkTracerProvider .getSampler() .shouldSample( diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DurableSamplingDecisionClassLoaderTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DurableSamplingDecisionClassLoaderTest.java index 121f4c8fb..59b5cb32d 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DurableSamplingDecisionClassLoaderTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DurableSamplingDecisionClassLoaderTest.java @@ -7,13 +7,37 @@ import static org.junit.jupiter.api.Assertions.assertNotSame; import static org.junit.jupiter.api.Assertions.assertNull; +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.api.trace.TraceState; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.api.trace.TracerProvider; import io.opentelemetry.context.Context; +import io.opentelemetry.context.propagation.ContextPropagators; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.IdGenerator; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.LinkData; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import io.opentelemetry.sdk.trace.samplers.Sampler; import io.opentelemetry.sdk.trace.samplers.SamplingDecision; import io.opentelemetry.sdk.trace.samplers.SamplingResult; import java.net.URL; import java.net.URLClassLoader; +import java.time.Instant; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +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; /** * Verifies the durable sampling decision crosses the application/Java-agent class-loader boundary. @@ -26,6 +50,129 @@ * then asserts the thread-scoped system-property bridge carries the decision from one loader to the other. */ class DurableSamplingDecisionClassLoaderTest { + @ParameterizedTest + @CsvSource({ + "InvocationOtelPlugin,local,false", "InvocationOtelPlugin,local,true", + "InvocationOtelPlugin,foreign,false", "InvocationOtelPlugin,foreign,true", + "InvocationOtelPlugin,opaque,false", "InvocationOtelPlugin,opaque,true", + "ExecutionOtelPlugin,local,false", "ExecutionOtelPlugin,local,true", + "ExecutionOtelPlugin,foreign,false", "ExecutionOtelPlugin,foreign,true", + "ExecutionOtelPlugin,opaque,false", "ExecutionOtelPlugin,opaque,true" + }) + void customSamplerMetadataSurvivesProviderOwnershipBoundaries( + String pluginName, String topology, boolean sharedTraceExecutions) throws Exception { + var previousHeader = System.getProperty("com.amazonaws.xray.traceHeader"); + GlobalOpenTelemetry.resetForTest(); + OtelPluginAutoConfigurationState.markInstalled(); + System.setProperty("com.amazonaws.xray.traceHeader", "Root=1-6955b900-123456789012345678901234"); + var evaluations = new AtomicInteger(); + var key = AttributeKey.stringKey("sampler.extra"); + var delegate = new Sampler() { + public SamplingResult shouldSample( + Context parent, + String traceId, + String name, + SpanKind kind, + Attributes attributes, + List links) { + evaluations.incrementAndGet(); + var second = + attributes.get(SpanAttributes.DURABLE_EXECUTION_ARN).contains("/second/"); + var metadata = sharedTraceExecutions ? (second ? "second" : "first") : "kept"; + return new SamplingResult() { + public SamplingDecision getDecision() { + return sharedTraceExecutions && !second + ? SamplingDecision.RECORD_ONLY + : SamplingDecision.RECORD_AND_SAMPLE; + } + + public Attributes getAttributes() { + return Attributes.of(key, metadata); + } + + public TraceState getUpdatedTraceState(TraceState parentState) { + return parentState.toBuilder().put("vendor", metadata).build(); + } + }; + } + + public String getDescription() { + return "custom-metadata"; + } + }; + try (var appLoader = pluginClassLoader(); + var agentLoader = pluginClassLoader(); + var exporter = InMemorySpanExporter.create()) { + var loader = topology.equals("local") ? appLoader : agentLoader; + var samplerType = Class.forName(DurableSampler.class.getName(), true, loader); + var wrap = samplerType.getDeclaredMethod("wrap", Sampler.class); + wrap.setAccessible(true); + var idType = Class.forName(DeterministicIdGenerator.class.getName(), true, loader); + try (var provider = SdkTracerProvider.builder() + .setSampler((Sampler) wrap.invoke(null, delegate)) + .setIdGenerator((IdGenerator) idType.getConstructor().newInstance()) + .addSpanProcessor(SimpleSpanProcessor.builder(exporter) + .setExportUnsampledSpans(true) + .build()) + .build()) { + var hidden = new TracerProvider() { + public Tracer get(String name) { + return provider.get(name); + } + + public Tracer get(String name, String version) { + return provider.get(name, version); + } + }; + GlobalOpenTelemetry.set(new OpenTelemetry() { + public TracerProvider getTracerProvider() { + return topology.equals("opaque") ? hidden : provider; + } + + public ContextPropagators getPropagators() { + return ContextPropagators.noop(); + } + }); + var plugin = (DurableExecutionPlugin) + Class.forName("software.amazon.lambda.durable.otel." + pluginName, true, appLoader) + .getConstructor() + .newInstance(); + var executionCount = sharedTraceExecutions ? 2 : 1; + for (var index = 0; index < executionCount; index++) { + var executionName = index == 0 ? "first" : "second"; + var arn = "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/" + executionName + + "/id"; + for (var first : new boolean[] {true, false}) { + plugin.onInvocationStart(new InvocationInfo("request", arn, first, Instant.ofEpochSecond(10))); + plugin.onInvocationEnd( + new InvocationEndInfo("request", arn, first, InvocationStatus.PENDING, null)); + } + } + var spans = exporter.getFinishedSpanItems().stream() + .filter(s -> s.getName().equals("Invocation")) + .toList(); + assertEquals(executionCount * 2, spans.size()); + for (var span : spans) { + var second = span.getAttributes() + .get(SpanAttributes.DURABLE_EXECUTION_ARN) + .contains("/second/"); + var expected = sharedTraceExecutions ? (second ? "second" : "first") : "kept"; + assertEquals(expected, span.getAttributes().get(key), span.getName()); + assertEquals(expected, span.getSpanContext().getTraceState().get("vendor"), span.getName()); + assertEquals( + !sharedTraceExecutions || second, + span.getSpanContext().isSampled()); + assertEquals("6955b900123456789012345678901234", span.getTraceId()); + } + assertEquals(executionCount * (topology.equals("local") ? 2 : 1), evaluations.get()); + } + } finally { + GlobalOpenTelemetry.resetForTest(); + OtelPluginAutoConfigurationState.resetInstalledForTest(); + if (previousHeader == null) System.clearProperty("com.amazonaws.xray.traceHeader"); + else System.setProperty("com.amazonaws.xray.traceHeader", previousHeader); + } + } @AfterEach void clearBridge() { From c96e11a533bd22dd80e0e087c179aa3c438636a1 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 17:20:22 +0000 Subject: [PATCH 39/48] fix: select manager control flow before waking handlers --- otel-plugin/README.md | 2 + ...InvocationFinalizationIntegrationTest.java | 189 ++++++++++++++++++ .../durable/execution/ExecutionManager.java | 5 +- .../execution/ExecutionManagerTest.java | 33 +++ 4 files changed, 227 insertions(+), 2 deletions(-) diff --git a/otel-plugin/README.md b/otel-plugin/README.md index ab452d553..4ef52c93f 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -31,6 +31,8 @@ unless a later `VirtualMachineError` or `ThreadDeath` takes precedence over a no end-hook Errors are retained as suppressed failures. Scope cleanup preserves finalization failures using the same JVM-fatal precedence, so an ordinary cleanup exception cannot hide an earlier Error. Invocations without plugins also wait for handler cleanup before returning the selected suspension or retry outcome. +The manager selects suspension or termination before waking operation waiters, so a later returning or throwing +handler `finally` cannot replace that selected outcome. A handler outcome that was already selected remains primary. SDK output preparation, including customer `SerDes` calls and durable large-result checkpointing, finishes before terminal invocation-end notification. Failures in this preparation report `RETRYING` instead of ending the Workflow span. If End also raises an unisolated Error, the preparation failure remains primary with the cleanup error diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/InvocationFinalizationIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/InvocationFinalizationIntegrationTest.java index f9bfc5200..0e41ada4b 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/InvocationFinalizationIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/InvocationFinalizationIntegrationTest.java @@ -6,26 +6,33 @@ import java.time.Duration; import java.util.ArrayList; +import java.util.Arrays; import java.util.List; +import java.util.concurrent.AbstractExecutorService; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.locks.LockSupport; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.CsvSource; import org.junit.jupiter.params.provider.ValueSource; import software.amazon.awssdk.services.lambda.model.ErrorObject; import software.amazon.lambda.durable.config.StepConfig; +import software.amazon.lambda.durable.config.WaitForConditionConfig; import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.model.WaitForConditionResult; +import software.amazon.lambda.durable.operation.BaseDurableOperation; import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; import software.amazon.lambda.durable.plugin.InvocationEndInfo; import software.amazon.lambda.durable.plugin.InvocationInfo; @@ -131,6 +138,188 @@ void externallySelectedOutcomeSurvivesReturningOrThrowingFinally(boolean retry, } } + @ParameterizedTest + @CsvSource({ + "step,async,true", + "step,async,false", + "condition,async,true", + "condition,async,false", + "termination,async,true", + "termination,async,false", + "step,inline-root,true", + "step,inline-root,false", + "condition,inline-root,true", + "condition,inline-root,false", + "termination,inline-root,true", + "termination,inline-root,false", + "step,submit-wait,true", + "step,submit-wait,false", + "condition,submit-wait,true", + "condition,submit-wait,false", + "termination,submit-wait,true", + "termination,submit-wait,false" + }) + void naturalManagerOutcomePrecedesFastThrowingFinally(String kind, String mode, boolean withPlugin) + throws Exception { + var workers = new OutcomeExecutor(mode); + var releaseOperation = new CountDownLatch(mode.equals("submit-wait") ? 0 : 1); + var aboutToWait = new CountDownLatch(1); + var completedStepCalls = new AtomicInteger(); + var targetCalls = new AtomicInteger(); + var throwFinally = new AtomicBoolean(true); + var original = retryError(); + var plugin = new RecordingPlugin(); + var builder = DurableConfig.builder().withExecutorService(workers); + if (withPlugin) builder.withPlugins(plugin); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + context.step("saved", String.class, step -> { + completedStepCalls.incrementAndGet(); + return "saved"; + }); + DurableFuture future; + if (kind.equals("condition")) { + future = context.waitForConditionAsync( + "condition", + Integer.class, + (state, step) -> { + await(releaseOperation); + return targetCalls.incrementAndGet() == 1 + ? WaitForConditionResult.continuePolling(state + 1) + : WaitForConditionResult.stopPolling(state); + }, + WaitForConditionConfig.builder() + .initialState(0) + .waitStrategy((state, attempt) -> Duration.ofSeconds(1)) + .build()); + } else { + future = context.stepAsync( + "target", + String.class, + step -> { + await(releaseOperation); + if (targetCalls.incrementAndGet() == 1) { + if (kind.equals("termination")) throw original; + throw new IllegalArgumentException("retry step"); + } + return "target"; + }, + StepConfig.builder() + .retryStrategy(RetryStrategies.fixedDelay(2, Duration.ofSeconds(1))) + .build()); + } + // This schedules the existing completion window; suspension/termination is triggered only by the + // real + // operation. The root's dependent wait is registered later and wakes before this earlier callback + // drains. + ((BaseDurableOperation) future).getCompletionFuture().whenComplete((value, failure) -> { + if (failure != null && !mode.equals("submit-wait") && throwFinally.get()) + await(workers.rootReturned); + }); + try { + aboutToWait.countDown(); + future.get(); + return "done"; + } finally { + if (throwFinally.get()) throw new IllegalStateException("fast handler finally"); + } + }, + builder.build()); + var caller = Executors.newSingleThreadExecutor(); + try { + var response = caller.submit(() -> runner.run("input")); + await(aboutToWait); + if (!mode.equals("submit-wait")) { + var until = System.nanoTime() + TimeUnit.SECONDS.toNanos(3); + while (System.nanoTime() < until) { + var root = workers.rootOwner.get(); + if (root != null + && root.getState() == Thread.State.WAITING + && Arrays.stream(root.getStackTrace()) + .anyMatch(frame -> frame.getMethodName().equals("waitForOperationCompletion"))) + break; + LockSupport.parkNanos(100_000); + } + assertEquals(Thread.State.WAITING, workers.rootOwner.get().getState()); + releaseOperation.countDown(); + } + if (kind.equals("termination")) { + assertSame( + original, + assertThrows(ExecutionException.class, () -> response.get(5, TimeUnit.SECONDS)) + .getCause()); + } else { + assertEquals( + ExecutionStatus.PENDING, + response.get(5, TimeUnit.SECONDS).getStatus()); + } + if (withPlugin) + assertPaired(plugin, kind.equals("termination") ? InvocationStatus.RETRYING : InvocationStatus.PENDING); + throwFinally.set(false); + assertEquals( + ExecutionStatus.SUCCEEDED, runner.runUntilComplete("input").getStatus()); + assertEquals(1, completedStepCalls.get(), "The completed step must replay without another side effect"); + } finally { + releaseOperation.countDown(); + workers.rootReturned.countDown(); + stop(caller); + stop(workers); + } + } + + private static final class OutcomeExecutor extends AbstractExecutorService { + private final String mode; + private final ExecutorService delegate = Executors.newCachedThreadPool(); + private final AtomicBoolean first = new AtomicBoolean(true); + private final AtomicReference rootOwner = new AtomicReference<>(); + private final CountDownLatch rootReturned = new CountDownLatch(1); + + private OutcomeExecutor(String mode) { + this.mode = mode; + } + + public void execute(Runnable task) { + var root = first.getAndSet(false); + Runnable wrapped = () -> { + if (root) rootOwner.set(Thread.currentThread()); + try { + task.run(); + } finally { + if (root) rootReturned.countDown(); + } + }; + if (root && mode.equals("inline-root")) wrapped.run(); + else if (mode.equals("submit-wait")) { + try { + delegate.submit(wrapped).get(); + } catch (Exception failure) { + throw new AssertionError(failure); + } + } else delegate.execute(wrapped); + } + + public void shutdown() { + delegate.shutdown(); + } + + public List shutdownNow() { + return delegate.shutdownNow(); + } + + public boolean isShutdown() { + return delegate.isShutdown(); + } + + public boolean isTerminated() { + return delegate.isTerminated(); + } + + public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { + return delegate.awaitTermination(timeout, unit); + } + } + @Test void hooksStayPairedAcrossStepRetryAndReplayWithoutRepeatingCompletedWork() { var plugin = new RecordingPlugin(); 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 9f260a688..0ebfd7727 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 @@ -759,8 +759,9 @@ public static boolean isTerminalStatus(OperationStatus status) { * @param exception the unrecoverable exception that caused termination */ public void terminateExecution(UnrecoverableDurableExecutionException exception) { - stopAllOperations(exception); + // Select control flow before waking a handler whose finally block can complete its body future. executionExceptionFuture.completeExceptionally(exception); + stopAllOperations(exception); throw exception; } @@ -771,8 +772,8 @@ public void suspendExecution() { private SuspendExecutionException signalSuspension() { var ex = new SuspendExecutionException(); - stopAllOperations(ex); executionExceptionFuture.completeExceptionally(ex); + stopAllOperations(ex); return ex; } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java index 65d961f36..db3660903 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java @@ -12,6 +12,7 @@ import java.util.List; import java.util.Map; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; @@ -24,6 +25,7 @@ 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.ErrorObject; import software.amazon.awssdk.services.lambda.model.GetDurableExecutionStateResponse; import software.amazon.awssdk.services.lambda.model.Operation; import software.amazon.awssdk.services.lambda.model.OperationStatus; @@ -32,6 +34,7 @@ import software.amazon.lambda.durable.TestUtils; import software.amazon.lambda.durable.client.DurableExecutionClient; import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.model.DurableExecutionInput; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; @@ -55,6 +58,36 @@ private ExecutionManager createManager(List operations) { null); } + @SuppressWarnings("removal") + @ParameterizedTest + @ValueSource(strings = {"success", "ordinary", "vm", "death"}) + void bodySelectedBeforeManagerTerminationRetainsItsOutcome(String kind) { + try (var manager = createManager(List.of(executionOp()))) { + var body = new CompletableFuture(); + var selected = manager.runUntilCompleteOrSuspend(body); + Throwable failure = + switch (kind) { + case "ordinary" -> new IllegalArgumentException("body failure"); + case "vm" -> new InternalError("body fatal"); + case "death" -> new ThreadDeath(); + default -> null; + }; + if (failure == null) body.complete("body success"); + else body.completeExceptionally(failure); + var later = new UnrecoverableDurableExecutionException( + ErrorObject.builder().errorMessage("later manager failure").build(), true); + assertSame( + later, + assertThrows( + UnrecoverableDurableExecutionException.class, () -> manager.terminateExecution(later))); + if (failure == null) assertEquals("body success", selected.join()); + else + assertSame( + failure, + assertThrows(CompletionException.class, selected::join).getCause()); + } + } + private Operation executionOp() { return Operation.builder() .id(EXECUTION_OP_ID) From 5f23e7dfa5a46500221185da5088afd90919717a Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 18:13:37 +0000 Subject: [PATCH 40/48] fix: retain visible replacement sampler root policy --- otel-plugin/README.md | 39 ++++-- .../durable/otel/OtelPluginSupport.java | 8 +- .../otel/ReplacementSamplerPolicyTest.java | 127 ++++++++++++++++++ 3 files changed, 161 insertions(+), 13 deletions(-) create mode 100644 otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ReplacementSamplerPolicyTest.java diff --git a/otel-plugin/README.md b/otel-plugin/README.md index 4ef52c93f..3fa849c87 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -279,14 +279,35 @@ Operation and attempt spans link to the Workflow span. `ExecutionOtelPlugin` rev ### Sampling -The plugin decides sampling once per invocation and applies that single decision to every durable span (Workflow, Invocation, operation, attempt), so the configured sampler is not re-invoked per span and the full decision — including `RECORD_ONLY` — is preserved. The decision follows this precedence, highest first: - -1. **Backend decision** — `Sampled=1` / `Sampled=0` in the propagated header is authoritative and always preserved, regardless of the configured sampler. -2. **Same-trace ambient span** — when the header carries no usable `Sampled` value but a valid ambient span (for example an auto-instrumentation Lambda handler span) is already on the execution's trace, the plugin follows that span's decision: sampled → sampled; unsampled but still recording → `RECORD_ONLY`; unsampled and not recording → dropped. -3. **Configured sampler (application-owned provider)** — when you pass a `SdkTracerProvider` to the plugin, its sampler is read directly and evaluated once with the trace ID, span name, and attributes. A trace-ID-ratio sampler therefore produces a stable decision across reinvocations (the trace ID is stable). -4. **Installed sampler (Java-agent path)** — when the agent owns the provider, it is behind a classloader boundary and its *effective* sampler (which another agent extension may have wrapped or replaced) cannot be reliably read at decision time. Rather than guess, the plugin **defers**: it installs a delegating sampler through the agent's autoconfiguration and lets that wrapper consult the agent's real sampler. The delegate's decision is honored in full — if your configured policy is `always_off`, a rate limiter, or a remote sampler (`xray`, `jaeger_remote`) that returns drop, the durable spans are dropped; they are **not** force-sampled. To avoid consuming a stateful or quota-based sampler once per span, the wrapper consults the delegate once per execution (keyed by trace ID) and reuses that decision for the execution's remaining durable spans within the invocation. - -For precise, provider-independent control, set an explicit `Sampled` value upstream (for example by enabling X-Ray active tracing) — that backend decision takes precedence over everything else. +The SDK-wide sampling guarantees below require `DurableSampler` itself to be the final installed sampler. The +[builder constructors](#configuration) and [ADOT extension setup](#1-adot-lambda-layer) install it automatically; +configure its delegate through these documented paths and retain `DurableSampler` as the final sampler. An +unrecognized outer wrapper is treated as a plain replacement. The SDK wrapper applies its sampling intent to every durable +span (Workflow, Invocation, operation and attempt), preserving the full result, including `RECORD_ONLY`, sampler attributes +and trace state. + +The supported SDK sampler follows this precedence, highest first: + +1. **Backend decision** — with `DurableSampler` installed, `Sampled=1` / `Sampled=0` in the propagated header is + authoritative, regardless of the wrapped delegate's policy. +2. **Same-trace ambient span** — without an explicit header decision, a valid ambient span already on the execution + trace supplies its decision: sampled → sampled; unsampled and recording → `RECORD_ONLY`; unsampled and + non-recording → dropped. +3. **Same-copy durable sampler** — the visible SDK sampler is evaluated against root context once per invocation + using the canonical trace ID, span name and attributes. Its complete result is carried in that loader's context. +4. **Foreign or opaque SDK sampler** — resolution is deferred to the installed wrapper's actual delegate. Its full + result is cached by execution ARN and canonical trace ID in a 256-entry LRU cache and reused while resident. + Eviction can cause another evaluation. The delegate still receives the canonical trace ID; drop and + `RECORD_ONLY` decisions, attributes and updated trace state are retained. + +A visible plain replacement sampler supplies its root policy for execution-ancestor flags, then keeps its normal +per-span sampling behavior. It receives no SDK sampling carrier and does not provide the SDK decision-reuse +or full-result override guarantees above. Opaque providers cannot expose a later sampler replacement to the +application, so retain the wrapper through the documented customization setup. + +For explicit execution-level control, keep `DurableSampler` as the final installed sampler and set `Sampled` upstream, +for example through X-Ray +active tracing. The SDK sampling wrapper preserves that decision. ## Span Attributes @@ -321,8 +342,6 @@ For precise, provider-independent control, set an explicit `Sampled` value upstr | `durable.attempt.number` | 1-based attempt number | | `durable.attempt.outcome` | SUCCEEDED (span status `OK`), FAILED (`ERROR`), or INCOMPLETE (`UNSET`) | -Deferred sampler results are isolated by execution ARN and canonical trace ID in the existing 256-entry LRU cache. The delegate is reused while an entry is resident; eviction can cause another evaluation. The delegate still receives the canonical trace ID. - ## Log Correlation (MDC) When `enableMdc` is true (default), the plugin injects these fields into SLF4J MDC during user function execution: diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java index 3f28220a6..e3af17bab 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java @@ -162,9 +162,11 @@ static SamplingResult resolveSamplingResult( } return ambientSpan.isRecording() ? SamplingResult.recordOnly() : SamplingResult.drop(); } - // 3. Only a same-copy DurableSampler can consume the full resolved context carrier. Defer for a foreign - // sampler so its own delegate retains attributes and updated trace state, not just the bridge decision. - if (usesLocalDurableSampler(sdkTracerProvider)) { + // 3. Resolve a local durable sampler's full result, or a visible replacement's root policy for ancestor flags. + // A foreign DurableSampler still resolves in its own loader to retain full attributes and trace state. + // A plain replacement receives no carrier; the provider keeps its normal per-span sampling behavior. + if (sdkTracerProvider != null + && (usesLocalDurableSampler(sdkTracerProvider) || !usesDurableSamplingBridge(sdkTracerProvider))) { return sdkTracerProvider .getSampler() .shouldSample( diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ReplacementSamplerPolicyTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ReplacementSamplerPolicyTest.java new file mode 100644 index 000000000..2d3fbe8a5 --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ReplacementSamplerPolicyTest.java @@ -0,0 +1,127 @@ +// 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 io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.api.trace.TracerProvider; +import io.opentelemetry.context.Context; +import io.opentelemetry.context.propagation.ContextPropagators; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.LinkData; +import io.opentelemetry.sdk.trace.samplers.Sampler; +import io.opentelemetry.sdk.trace.samplers.SamplingResult; +import java.time.Duration; +import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class ReplacementSamplerPolicyTest { + @ParameterizedTest + @CsvSource({ + "invocation,false,root,false,plain", "execution,false,root,false,plain", + "invocation,true,root,true,plain", "execution,true,root,true,plain", + "invocation,false,sampled,true,plain", "execution,false,sampled,true,plain", + "invocation,true,unsampled,false,plain", "execution,true,unsampled,false,plain", + "invocation,false,root,false,local", "execution,false,root,false,local", + "invocation,false,root,false,opaque", "execution,false,root,false,opaque" + }) + void visibleReplacementParentBasedSamplerKeepsItsRootPolicy( + String view, boolean rootSamples, String upstream, boolean expectedSampled, String topology) { + var oldHeader = System.getProperty("com.amazonaws.xray.traceHeader"); + var evaluations = new AtomicInteger(); + var sideEffects = new AtomicInteger(); + var decisions = new CopyOnWriteArrayList(); + var header = "Root=1-6955b900-123456789012345678901234"; + if (upstream.equals("sampled")) header += ";Sampled=1"; + else if (upstream.equals("unsampled")) header += ";Sampled=0"; + var root = new Sampler() { + public SamplingResult shouldSample( + Context parent, + String traceId, + String name, + SpanKind kind, + Attributes attributes, + List links) { + evaluations.incrementAndGet(); + assertFalse(Span.fromContext(parent).getSpanContext().isValid()); + return rootSamples ? SamplingResult.recordAndSample() : SamplingResult.drop(); + } + + public String getDescription() { + return "counted-root-policy"; + } + }; + GlobalOpenTelemetry.resetForTest(); + OtelPluginAutoConfigurationState.markInstalled(); + System.setProperty("com.amazonaws.xray.traceHeader", header); + var configured = Sampler.parentBased(root); + try (var provider = SdkTracerProvider.builder() + .setSampler(topology.equals("plain") ? configured : DurableSampler.wrap(configured)) + .setIdGenerator(new DeterministicIdGenerator()) + .build()) { + var opaque = new TracerProvider() { + public Tracer get(String name) { + return provider.get(name); + } + + public Tracer get(String name, String version) { + return provider.get(name, version); + } + }; + GlobalOpenTelemetry.set(new OpenTelemetry() { + public TracerProvider getTracerProvider() { + return topology.equals("opaque") ? opaque : provider; + } + + public ContextPropagators getPropagators() { + return ContextPropagators.noop(); + } + }); + DurableExecutionPlugin plugin = + view.equals("invocation") ? new InvocationOtelPlugin() : new ExecutionOtelPlugin(); + var tracer = provider.get("application"); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + var span = tracer.spanBuilder("application-work").startSpan(); + decisions.add(span.isRecording()); + span.end(); + context.step("saved", String.class, step -> { + sideEffects.incrementAndGet(); + return "saved"; + }); + context.wait("pause", Duration.ofSeconds(1)); + return "done"; + }, + DurableConfig.builder().withPlugins(plugin).build()); + assertEquals(ExecutionStatus.PENDING, runner.run("input").getStatus()); + assertEquals( + ExecutionStatus.SUCCEEDED, runner.runUntilComplete("input").getStatus()); + assertTrue(decisions.size() >= 2, "Initial invocation and actual replay must both execute the handler"); + assertTrue(decisions.stream().allMatch(decision -> decision == expectedSampled), decisions.toString()); + assertEquals( + upstream.equals("root") ? (topology.equals("opaque") ? 1 : decisions.size()) : 0, + evaluations.get()); + assertEquals(1, sideEffects.get()); + } finally { + GlobalOpenTelemetry.resetForTest(); + OtelPluginAutoConfigurationState.resetInstalledForTest(); + DurableSamplingDecision.clearSharedStateForTest(); + if (oldHeader == null) System.clearProperty("com.amazonaws.xray.traceHeader"); + else System.setProperty("com.amazonaws.xray.traceHeader", oldHeader); + } + } +} From 80110fc612416cc3162cb3c4a85f8439212aa8d8 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 20:53:44 +0000 Subject: [PATCH 41/48] fix: propagate coordinator and post-End fatal failures --- otel-plugin/README.md | 16 +- .../durable/otel/MdcFailureBoundaryTest.java | 219 +++++++++++++++++- .../ContinuationFatalPublicProbeTest.java | 200 ++++++++++++++++ .../amazon/lambda/durable/DurableConfig.java | 5 +- ...nrecoverableDurableExecutionException.java | 7 +- .../durable/execution/DurableExecutor.java | 25 +- .../durable/execution/ExecutionManager.java | 22 ++ .../execution/ExecutionManagerTest.java | 101 ++++++++ 8 files changed, 580 insertions(+), 15 deletions(-) create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ContinuationFatalPublicProbeTest.java diff --git a/otel-plugin/README.md b/otel-plugin/README.md index 3fa849c87..db05972b7 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -40,13 +40,25 @@ suppressed, unless cleanup introduces the first JVM-fatal error. An original JVM identity. End describes the SDK outcome at that point, not acknowledgment of a response by the Lambda service. Caller-side execution-manager cleanup, response-envelope encoding and output-stream writes follow End; runtime response transport follows the handler return. Failures at those later boundaries still propagate, without a second -End dispatch or changing its already reported outcome. Non-JVM-fatal MDC-restoration failures on an inline executor, -including `AssertionError` and `LinkageError`, likewise do not replace an already selected outcome. +End dispatch or changing its already reported outcome. A JVM-fatal MDC-restoration failure after End (direct or +inside a standard transport wrapper) completes the caller's observation exceptionally with the original fatal +before that fatal escapes the worker, for both asynchronous and inline executors. The End notification remains +unchanged and is not repeated. An earlier JVM-fatal End/preparation failure remains primary; later restoration +failures are suppressed. A restoration fatal takes precedence over an earlier non-JVM-fatal delivery/End failure, +which remains suppressed. Non-JVM-fatal restoration failures, including `AssertionError` and `LinkageError`, retain +the selected caller outcome. Before invocation startup, a JVM-fatal error from MDC capture (direct or inside a standard transport wrapper) completes the observation future exceptionally with that same fatal before escaping the handler worker. No start, body, or end hook runs, and no durable `FAILED` response is produced for that fatal. Ordinary initialization errors and handler/body failure classification retain their existing behavior. +A direct `VirtualMachineError` or `ThreadDeath` from an SDK checkpoint continuation selects a retryable invocation +control failure before the continuation releases its activity lease. The caller receives the original fatal as the +cause of `UnrecoverableDurableExecutionException`, and started hooks receive `RETRYING`. The continuation's +observation future settles with the original fatal after lease release and before the fatal escapes its coordinator +worker. Already selected outcomes retain first-completion precedence. This direct-fatal continuation boundary +does not add general wrapped-fatal classification or change ordinary continuation and handler/body failures. + SDK inspection of MDC-capture failures reads each visited standard transport cause once and detects identity cycles. Cyclic, null, or unreadable leading `CompletionException` chains retain the original wrapper; ordinary initialization still reports `FAILED` when its error response can be serialized. This provides cycle safety for finite cause graphs, diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFailureBoundaryTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFailureBoundaryTest.java index 116738821..14ccc4e8b 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFailureBoundaryTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFailureBoundaryTest.java @@ -6,20 +6,26 @@ import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Proxy; +import java.lang.reflect.UndeclaredThrowableException; import java.time.Duration; import java.util.List; import java.util.concurrent.AbstractExecutorService; import java.util.concurrent.CompletableFuture; +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.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import java.util.stream.Stream; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.MethodSource; import org.slf4j.MDC; import org.slf4j.spi.MDCAdapter; import software.amazon.lambda.durable.DurableConfig; @@ -27,6 +33,7 @@ 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; import software.amazon.lambda.durable.testing.LocalDurableTestRunner; class MdcFailureBoundaryTest { @@ -75,16 +82,22 @@ void workerCaptureFailureDoesNotEndAnInvocationThatNeverStarted() throws Excepti @ParameterizedTest @CsvSource({ "false,false,exception", "true,false,exception", "false,true,exception", "true,true,exception", - "false,false,fatal", "true,false,fatal", "false,true,fatal", "true,true,fatal" + "false,false,fatal", "true,false,fatal", "false,true,fatal", "true,true,fatal", + "false,false,assertion", "true,false,assertion", "false,true,assertion", "true,true,assertion", + "false,false,linkage", "true,false,linkage", "false,true,linkage", "true,true,linkage" }) void workerMdcRestorationFailureCannotStrandTheInvocation(boolean ambientMdc, boolean suspend, String failureKind) throws Exception { var original = MDC.getMDCAdapter(); var injected = new AtomicBoolean(); var plugin = new RecordingPlugin(); - Throwable failure = failureKind.equals("fatal") - ? new InternalError("worker MDC restoration failed") - : new IllegalStateException("worker MDC restoration failed"); + Throwable failure = + switch (failureKind) { + case "fatal" -> new InternalError("worker MDC restoration failed"); + case "assertion" -> new AssertionError("worker MDC restoration failed"); + case "linkage" -> new NoClassDefFoundError("worker MDC restoration failed"); + default -> new IllegalStateException("worker MDC restoration failed"); + }; var ownerFailure = new AtomicReference(); var observed = new CountDownLatch(1); var workers = Executors.newSingleThreadExecutor(task -> { @@ -119,8 +132,12 @@ void workerMdcRestorationFailureCannotStrandTheInvocation(boolean ambientMdc, bo .withExecutorService(workers) .withPlugins(plugin) .build()); - var result = runner.run("input"); - assertEquals(suspend ? ExecutionStatus.PENDING : ExecutionStatus.SUCCEEDED, result.getStatus()); + if (failureKind.equals("fatal")) { + assertSame(failure, assertThrows(Error.class, () -> runner.run("input"))); + } else { + var result = runner.run("input"); + assertEquals(suspend ? ExecutionStatus.PENDING : ExecutionStatus.SUCCEEDED, result.getStatus()); + } assertEquals(1, plugin.starts.get()); assertEquals(1, plugin.ends.get()); assertTrue(injected.get()); @@ -133,6 +150,192 @@ void workerMdcRestorationFailureCannotStrandTheInvocation(boolean ambientMdc, bo } } + static Stream postEndFatalCases() { + return Stream.of(false, true) + .flatMap(inline -> Stream.of(false, true) + .flatMap(suspend -> Stream.of(false, true) + .flatMap(death -> Stream.of("direct", "completion", "execution", "reflection", "proxy") + .map(wrapper -> Arguments.of(inline, suspend, death, wrapper))))); + } + + @SuppressWarnings("removal") + @ParameterizedTest + @MethodSource("postEndFatalCases") + void postEndFatalReachesCallerAndWorkerWithoutChangingEnd( + boolean inline, boolean suspend, boolean death, String wrapper) throws Exception { + Error fatal = death ? new ThreadDeath() : new InternalError("post-End restore fatal"); + Throwable failure = + switch (wrapper) { + case "completion" -> new CompletionException(fatal); + case "execution" -> new ExecutionException(fatal); + case "reflection" -> new InvocationTargetException(fatal); + case "proxy" -> new UndeclaredThrowableException(fatal); + default -> fatal; + }; + exercisePostEndFailure(inline, suspend, failure, fatal, null); + } + + @SuppressWarnings("removal") + @ParameterizedTest + @CsvSource({ + "false,false,false", + "false,false,true", + "false,true,false", + "false,true,true", + "true,false,false", + "true,false,true", + "true,true,false", + "true,true,true" + }) + void fatalDiagnosticRetainsIdentityAndIsReadOnce(boolean inline, boolean suspend, boolean death) throws Exception { + Error fatal = death ? new ThreadDeath() : new InternalError("fatal diagnostic"); + var reads = new AtomicInteger(); + var wrapper = new CompletionException("diagnostic", null) { + @Override + public synchronized Throwable getCause() { + reads.incrementAndGet(); + throw fatal; + } + }; + exercisePostEndFailure(inline, suspend, wrapper, fatal, null); + assertEquals(1, reads.get()); + } + + @SuppressWarnings("removal") + @ParameterizedTest + @CsvSource({"false,false", "false,true", "true,false", "true,true"}) + void earlierEndFatalRemainsPrimaryAndRetainsRestorationDiagnostic(boolean inline, boolean death) throws Exception { + Error primary = death ? new ThreadDeath() : new InternalError("End fatal"); + var restoration = new InternalError("later restoration fatal"); + exercisePostEndFailure(inline, false, restoration, primary, primary); + assertArrayEquals(new Throwable[] {restoration}, primary.getSuppressed()); + } + + @ParameterizedTest + @CsvSource({"false,false", "false,true", "true,false", "true,true"}) + void restorationFatalRetainsEarlierNonfatalEndFailure(boolean inline, boolean wrapped) throws Exception { + var earlier = new AssertionError("ordinary End failure"); + var fatal = new InternalError("later restoration fatal"); + exercisePostEndFailure(inline, false, wrapped ? new CompletionException(fatal) : fatal, fatal, earlier); + assertArrayEquals(new Throwable[] {earlier}, fatal.getSuppressed()); + } + + private static void exercisePostEndFailure( + boolean inline, boolean suspend, Throwable failure, Error expectedFatal, Error endFailure) + throws Exception { + var original = MDC.getMDCAdapter(); + var injected = new AtomicBoolean(); + var plugin = new RecordingPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + super.onInvocationEnd(info); + if (endFailure != null) throw endFailure; + } + }; + var ownerFailure = new AtomicReference(); + var owner = new AtomicReference(); + var escaped = new CountDownLatch(1); + var caller = Executors.newSingleThreadExecutor(task -> new Thread(task, "post-end-caller")); + var delegate = Executors.newSingleThreadExecutor(task -> { + var thread = new Thread(task, "post-end-worker"); + thread.setUncaughtExceptionHandler((worker, error) -> { + ownerFailure.set(error); + escaped.countDown(); + }); + return thread; + }); + var executor = new AbstractExecutorService() { + @Override + public void execute(Runnable task) { + Runnable owned = () -> { + owner.set(Thread.currentThread()); + if (suspend) MDC.put("ambient", "saved"); + else MDC.clear(); + task.run(); + }; + if (inline) { + try { + owned.run(); + } catch (Throwable error) { + ownerFailure.set(error); + escaped.countDown(); + throw error; + } + } else delegate.execute(owned); + } + + @Override + public void shutdown() { + delegate.shutdown(); + } + + @Override + public List shutdownNow() { + return delegate.shutdownNow(); + } + + @Override + public boolean isShutdown() { + return delegate.isShutdown(); + } + + @Override + public boolean isTerminated() { + return delegate.isTerminated(); + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { + return delegate.awaitTermination(timeout, unit); + } + }; + try { + replaceAdapter(proxy(original, method -> { + if (Thread.currentThread() == owner.get() + && plugin.ends.get() == 1 + && method.equals(suspend ? "setContextMap" : "clear") + && injected.compareAndSet(false, true)) throw failure; + return null; + })); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + if (suspend) context.wait("resume", Duration.ofSeconds(1)); + return "done"; + }, + DurableConfig.builder() + .withExecutorService(executor) + .withPlugins(plugin) + .build()); + var observation = caller.submit(() -> { + try { + return (Object) runner.run("input"); + } catch (Throwable error) { + return error; + } + }); + assertSame( + expectedFatal, + observation.get(3, TimeUnit.SECONDS), + "caller receives original fatal without hanging"); + assertTrue(escaped.await(3, TimeUnit.SECONDS), "fatal escapes actual worker after observation settlement"); + assertSame(expectedFatal, ownerFailure.get()); + assertTrue(injected.get()); + assertEquals(1, plugin.starts.get()); + assertEquals(1, plugin.ends.get()); + assertEquals(suspend ? InvocationStatus.PENDING : InvocationStatus.SUCCEEDED, plugin.status.get()); + } finally { + caller.shutdownNow(); + executor.shutdownNow(); + try { + assertTrue(caller.awaitTermination(3, TimeUnit.SECONDS)); + assertTrue(executor.awaitTermination(3, TimeUnit.SECONDS)); + } finally { + replaceAdapter(original); + } + } + } + @FunctionalInterface private interface Fault { Object apply(String method) throws Throwable; @@ -157,9 +360,10 @@ private static void replaceAdapter(MDCAdapter adapter) throws Exception { setter.invoke(null, adapter); } - private static final class RecordingPlugin implements DurableExecutionPlugin { + private static class RecordingPlugin implements DurableExecutionPlugin { final AtomicInteger starts = new AtomicInteger(); final AtomicInteger ends = new AtomicInteger(); + final AtomicReference status = new AtomicReference<>(); @Override public void onInvocationStart(InvocationInfo info) { @@ -168,6 +372,7 @@ public void onInvocationStart(InvocationInfo info) { @Override public void onInvocationEnd(InvocationEndInfo info) { + status.set(info.invocationStatus()); ends.incrementAndGet(); } } diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ContinuationFatalPublicProbeTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ContinuationFatalPublicProbeTest.java new file mode 100644 index 000000000..770f99579 --- /dev/null +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ContinuationFatalPublicProbeTest.java @@ -0,0 +1,200 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.*; + +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import java.util.function.BiFunction; +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.config.WaitForConditionConfig; +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.model.DurableExecutionOutput; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.model.WaitForConditionResult; +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; +import software.amazon.lambda.durable.serde.SerDes; +import software.amazon.lambda.durable.testing.local.LocalMemoryExecutionClient; + +class ContinuationFatalPublicProbeTest { + @SuppressWarnings("removal") + @ParameterizedTest + @CsvSource({"false,false", "false,true", "true,false", "true,true"}) + void readyResumeSerdeFatalDoesNotDisappearIntoPending(boolean death, boolean withPlugin) throws Exception { + Error fatal = death ? new ThreadDeath() : new InternalError("resume serde fatal"); + var armed = new AtomicBoolean(); + var injected = new CountDownLatch(1); + var escaped = new CountDownLatch(1); + var injectionOwner = new AtomicReference(); + var ownerFailure = new AtomicReference(); + var checks = new AtomicInteger(); + var ends = new CopyOnWriteArrayList(); + var pollEntered = new CountDownLatch(1); + var releaseReady = new CountDownLatch(1); + var firstArmedPoll = new AtomicBoolean(true); + var onCoordinator = new AtomicBoolean(); + var previousHandler = Thread.getDefaultUncaughtExceptionHandler(); + Thread.setDefaultUncaughtExceptionHandler((thread, failure) -> { + if (failure == fatal) { + ownerFailure.set(failure); + escaped.countDown(); + } else if (previousHandler != null) previousHandler.uncaughtException(thread, failure); + }); + var serde = new SerDes() { + private final JacksonSerDes delegate = new JacksonSerDes(); + + public String serialize(Object value) { + return delegate.serialize(value); + } + + public T deserialize(String value, TypeToken type) { + if (armed.get() && Thread.currentThread().getName().startsWith("durable-sdk-internal-")) { + injectionOwner.set(Thread.currentThread().getName()); + onCoordinator.set(Arrays.stream(Thread.currentThread().getStackTrace()) + .anyMatch(frame -> frame.getMethodName().equals("completeCheckpointContinuation"))); + injected.countDown(); + throw fatal; + } + return delegate.deserialize(value, type); + } + }; + var client = new LocalMemoryExecutionClient() { + @Override + public CheckpointDurableExecutionResponse checkpoint( + String arn, String token, List updates) { + if (armed.get() && updates.isEmpty() && firstArmedPoll.compareAndSet(true, false)) { + pollEntered.countDown(); + try { + assertTrue(releaseReady.await(3, TimeUnit.SECONDS)); + } catch (InterruptedException failure) { + throw new AssertionError(failure); + } + } + return super.checkpoint(arn, token, updates); + } + }; + var workers = Executors.newCachedThreadPool(); + var caller = Executors.newSingleThreadExecutor(); + var builder = DurableConfig.builder() + .withExecutorService(workers) + .withDurableExecutionClient(client) + .withCheckpointDelay(Duration.ZERO) + .withPollingStrategy(attempt -> Duration.ZERO); + if (withPlugin) + builder.withPlugins(new DurableExecutionPlugin() { + public void onInvocationEnd(InvocationEndInfo info) { + ends.add(info); + } + }); + var config = builder.build(); + var condition = WaitForConditionConfig.builder() + .initialState(1) + .serDes(serde) + .waitStrategy((state, attempt) -> Duration.ofSeconds(1)) + .build(); + BiFunction handler = (input, context) -> { + var waiting = context.waitForConditionAsync( + "condition", + Integer.class, + (state, step) -> { + checks.incrementAndGet(); + return state == 2 + ? WaitForConditionResult.stopPolling(state) + : WaitForConditionResult.continuePolling(2); + }, + condition); + if (armed.get()) { + try { + assertTrue(pollEntered.await(3, TimeUnit.SECONDS)); + } catch (InterruptedException failure) { + throw new AssertionError(failure); + } + } + return String.valueOf(waiting.get()); + }; + try { + var first = caller.submit(() -> DurableExecutor.execute( + input(List.of()), null, TypeToken.get(String.class), handler, config)) + .get(3, TimeUnit.SECONDS); + assertEquals(ExecutionStatus.PENDING, first.status()); + var cachedPending = new ArrayList<>(client.getAllOperations()); + assertEquals( + OperationStatus.PENDING, + client.getOperationByName("condition").status()); + assertTrue(client.advanceTime()); + armed.set(true); + var second = caller.submit(() -> + DurableExecutor.execute(input(cachedPending), null, TypeToken.get(String.class), handler, config)); + assertTrue(pollEntered.await(3, TimeUnit.SECONDS)); + releaseReady.countDown(); + assertTrue(injected.await(3, TimeUnit.SECONDS)); + assertTrue(onCoordinator.get(), "Fault must run in the actual WFC coordinator continuation"); + DurableExecutionOutput output = null; + Throwable callerFailure = null; + try { + output = second.get(3, TimeUnit.SECONDS); + } catch (ExecutionException failure) { + callerFailure = failure.getCause(); + } + var ownerEscaped = escaped.await(300, TimeUnit.MILLISECONDS); + System.out.println("CONTINUATION_FATAL_PUBLIC owner=" + injectionOwner.get() + + " caller=" + + (output != null + ? output.status() + : callerFailure.getClass().getName()) + + " backend=" + client.getOperationByName("condition").status() + + " checks=" + checks.get() + " workerEscaped=" + ownerEscaped); + assertEquals(1, checks.get(), "A failed resumed state read must not rerun the condition body"); + assertTrue(ownerEscaped, "The actual coordinator worker must observe its JVM fatal"); + assertSame(fatal, ownerFailure.get()); + assertNull(output, "A failed coordinator must not return a durable terminal or PENDING response"); + var retry = assertInstanceOf(UnrecoverableDurableExecutionException.class, callerFailure); + assertTrue(retry.isRetryable()); + assertSame(fatal, retry.getCause()); + if (withPlugin) { + assertEquals( + List.of(InvocationStatus.PENDING, InvocationStatus.RETRYING), + ends.stream().map(InvocationEndInfo::invocationStatus).toList()); + assertSame(retry, ends.get(1).executionError()); + } else assertTrue(ends.isEmpty()); + } finally { + releaseReady.countDown(); + Thread.setDefaultUncaughtExceptionHandler(previousHandler); + caller.shutdownNow(); + workers.shutdownNow(); + assertTrue(caller.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private static DurableExecutionInput input(List operations) { + var execution = Operation.builder() + .id("execution") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails( + ExecutionDetails.builder().inputPayload("\"input\"").build()) + .build(); + var history = new ArrayList(); + history.add(execution); + history.addAll(operations); + return new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/fatal/execution", + "token", + CheckpointUpdatedExecutionState.builder().operations(history).build()); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableConfig.java index 54476ec79..ee81fcd23 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableConfig.java @@ -497,8 +497,9 @@ public Builder withCheckpointEmptyMap(boolean checkpointEmptyMap) { * constructed again when its exact concrete type is already explicitly configured. Unrelated plugins retain * their existing registration behavior. * - *

    Calling this method replaces any previously registered plugins. Plugins are called in registration order. - * A fresh builder combines this explicit list with environment-selected plugins. On a builder returned by + *

    Calling this method replaces any previously registered plugins. Hooks use registration order, except + * invocation end, which uses reverse order under the 2.2.2 same-thread lifecycle contract. A fresh builder + * combines this explicit list with environment-selected plugins. On a builder returned by * {@link DurableConfig#toBuilder()}, this method replaces the complete resolved list and dynamic discovery * remains disabled, including when the replacement list is empty. * diff --git a/sdk/src/main/java/software/amazon/lambda/durable/exception/UnrecoverableDurableExecutionException.java b/sdk/src/main/java/software/amazon/lambda/durable/exception/UnrecoverableDurableExecutionException.java index acd2fd9fc..4315ed4ac 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/exception/UnrecoverableDurableExecutionException.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/exception/UnrecoverableDurableExecutionException.java @@ -10,7 +10,12 @@ public class UnrecoverableDurableExecutionException extends DurableExecutionExce private final boolean retryable; public UnrecoverableDurableExecutionException(ErrorObject errorObject, boolean retryable) { - super(errorObject.errorMessage()); + this(errorObject, retryable, null); + } + + /** Creates an invocation control failure while retaining its original in-process cause. */ + public UnrecoverableDurableExecutionException(ErrorObject errorObject, boolean retryable, Throwable cause) { + super(errorObject.errorMessage(), cause); this.errorObject = errorObject; this.retryable = retryable; } 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 3f054489d..0c701be71 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 @@ -13,6 +13,7 @@ import java.util.concurrent.CompletionException; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; import java.util.function.Function; import java.util.function.Supplier; @@ -273,6 +274,7 @@ private void complete(CompletableFuture future) { private static CompletableFuture supplyHandler( Supplier task, Function initializationFailure, Executor executor) { var result = new CompletableFuture(); + var classifiedWorkerFailure = new AtomicReference(); Runnable work = (Runnable & CompletableFuture.AsynchronousCompletionTask) () -> { SafeCloseable restore; try { @@ -292,15 +294,30 @@ private static CompletableFuture supplyHandler( return; } var outcome = Outcome.capture(task); + var restoringAfterNonfatalOutcome = false; try { try (var ignored = restore) { // End/delivery failures must also escape their actual worker. Handler-body failures have // already been mapped to their selected durable outcome; this does not reclassify them. rethrowLifecycleFatal(outcome.failure()); + restoringAfterNonfatalOutcome = true; } + } catch (Throwable workerFailure) { + try { + rethrowLifecycleFatal(workerFailure); + } catch (VirtualMachineError | ThreadDeath fatal) { + if (restoringAfterNonfatalOutcome && outcome.failure() != null && outcome.failure() != fatal) { + fatal.addSuppressed(outcome.failure()); + } + // End describes the already selected SDK outcome, not successful return to the runtime. + // Publish this fatal before throwing it from the worker; do not dispatch End again. + result.completeExceptionally(fatal); + throw fatal; + } + classifiedWorkerFailure.set(workerFailure); + throw workerFailure; } finally { - // End and worker restoration have run before publishing. A restoration failure still escapes its - // owner, without changing the invocation outcome already delivered to the end hooks. + // Ordinary restoration failures retain the selected outcome. A fatal has already settled result. outcome.complete(result); } }; @@ -310,7 +327,9 @@ private static CompletableFuture supplyHandler( if (!result.isDone()) throw dispatchFailure; // Inline execution can throw from MDC restoration after settling the selected outcome. Preserve that // outcome for an ordinary cleanup failure, while a JVM-fatal failure still reaches the caller. - rethrowLifecycleFatal(dispatchFailure); + // An inline task already inspected this exact ordinary failure. Do not re-read custom diagnostics; + // a different failure raised by the executor itself still receives the existing classification. + if (dispatchFailure != classifiedWorkerFailure.get()) rethrowLifecycleFatal(dispatchFailure); } return result; } 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 0ebfd7727..2aff9a2c5 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 @@ -22,6 +22,7 @@ import java.util.stream.Collectors; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.services.lambda.model.ErrorObject; import software.amazon.awssdk.services.lambda.model.Operation; import software.amazon.awssdk.services.lambda.model.OperationStatus; import software.amazon.awssdk.services.lambda.model.OperationType; @@ -587,20 +588,39 @@ CompletableFuture runCheckpointContinuation( } } + @SuppressWarnings("removal") private void completeCheckpointContinuation( Object registration, Runnable continuation, CompletableFuture completion) { try { try { if (!isClosing()) continuation.run(); + } catch (VirtualMachineError | ThreadDeath fatal) { + // Choose retry control flow before releasing the final lease could otherwise select PENDING. + signalContinuationFatal(fatal); + throw fatal; } finally { finishCheckpointContinuation(registration); } completion.complete(null); } catch (Throwable failure) { completion.completeExceptionally(failure); + // The observation future may be ignored by a polling chain. Never hide a JVM fatal from its worker. + if (failure instanceof VirtualMachineError fatal) throw fatal; + if (failure instanceof ThreadDeath fatal) throw fatal; } } + private void signalContinuationFatal(Error fatal) { + var control = new UnrecoverableDurableExecutionException( + ErrorObject.builder() + .errorType(fatal.getClass().getName()) + .errorMessage("Fatal error in SDK checkpoint continuation") + .build(), + true, + fatal); + if (executionExceptionFuture.completeExceptionally(control)) stopAllOperations(control); + } + private Object registerCheckpointContinuation(BaseDurableOperation owner) { synchronized (activeThreads) { if (closing || executionExceptionFuture.isDone()) return null; @@ -613,6 +633,8 @@ private Object registerCheckpointContinuation(BaseDurableOperation owner) { private void finishCheckpointContinuation(Object registration) { synchronized (activeThreads) { + // An inline coordinator may rethrow after its task already released this registration. + if (!checkpointContinuations.containsKey(registration)) return; checkpointContinuations.remove(registration); try { finishCheckpointProcessing(); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java index db3660903..1c31805ea 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java @@ -16,13 +16,16 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; +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 java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.Test; 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.ErrorObject; @@ -88,6 +91,104 @@ void bodySelectedBeforeManagerTerminationRetainsItsOutcome(String kind) { } } + @SuppressWarnings("removal") + @ParameterizedTest + @CsvSource({"false,false", "false,true", "true,false", "true,true"}) + void continuationFatalSettlesOriginalBeforeEscapingWorker(boolean death, boolean rootActive) throws Exception { + Error fatal = death ? new ThreadDeath() : new InternalError("coordinator fatal"); + var escaped = new CountDownLatch(1); + var release = new CountDownLatch(1); + var ownerFailure = new AtomicReference(); + var observation = new AtomicReference>(); + var settledAtEscape = new AtomicBoolean(); + var executor = Executors.newSingleThreadExecutor(task -> { + var thread = new Thread(task, "continuation-fatal-probe"); + thread.setUncaughtExceptionHandler((owner, failure) -> { + ownerFailure.set(failure); + settledAtEscape.set(observation.get().isDone()); + escaped.countDown(); + }); + return thread; + }); + try (var manager = createManager(List.of(executionOp(), stepOp("pending", OperationStatus.PENDING)))) { + if (rootActive) manager.registerActiveThread("root-probe"); + var future = manager.runCheckpointContinuation( + () -> { + try { + assertTrue(release.await(3, TimeUnit.SECONDS)); + } catch (InterruptedException interrupted) { + throw new AssertionError(interrupted); + } + throw fatal; + }, + executor); + observation.set(future); + var closedFromObserver = new AtomicBoolean(); + future.whenComplete((ignored, failure) -> { + manager.close(); // Must not wait for this callback's own continuation registration. + closedFromObserver.set(true); + }); + release.countDown(); + assertSame( + fatal, + assertThrows(ExecutionException.class, () -> future.get(3, TimeUnit.SECONDS)) + .getCause()); + assertTrue(escaped.await(2, TimeUnit.SECONDS), "Fatal must escape the actual coordinator worker"); + assertSame(fatal, ownerFailure.get()); + assertTrue(settledAtEscape.get(), "Observation must settle before worker escape"); + assertTrue(closedFromObserver.get(), "The activity lease must be released before observer callbacks"); + } finally { + release.countDown(); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @SuppressWarnings("removal") + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void inlineFatalContinuationReleasesItsRegistrationOnlyOnce(boolean death) { + Error fatal = death ? new ThreadDeath() : new InternalError("inline coordinator fatal"); + try (var manager = createManager(List.of(executionOp()))) { + manager.registerActiveThread("root"); + assertSame( + fatal, + assertThrows( + Error.class, + () -> manager.runCheckpointContinuation( + () -> { + throw fatal; + }, + Runnable::run))); + assertTrue(manager.isExecutionCompletedExceptionally()); + } + } + + @Test + void ordinaryContinuationFailureRetainsItsExistingObservationPolicy() throws Exception { + var original = new IllegalArgumentException("ordinary continuation"); + var worker = new AtomicReference(); + var executor = Executors.newSingleThreadExecutor(); + try (var manager = createManager(List.of(executionOp()))) { + manager.registerActiveThread("root"); + var result = manager.runCheckpointContinuation( + () -> { + worker.set(Thread.currentThread()); + throw original; + }, + executor); + assertSame( + original, + assertThrows(ExecutionException.class, () -> result.get(3, TimeUnit.SECONDS)) + .getCause()); + assertFalse(manager.isExecutionCompletedExceptionally()); + assertSame(worker.get(), executor.submit(Thread::currentThread).get(3, TimeUnit.SECONDS)); + } finally { + executor.shutdownNow(); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + private Operation executionOp() { return Operation.builder() .id(EXECUTION_OP_ID) From 007031a9b825543c1e1bae2b412df2f3d2159461 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 21:27:28 +0000 Subject: [PATCH 42/48] fix: preserve owned continuation failures and rejected admission --- otel-plugin/README.md | 15 +- .../ContinuationFailurePublicProbeTest.java | 241 ++++++++++++++++++ .../durable/execution/ExecutionManager.java | 34 ++- .../operation/BaseDurableOperation.java | 17 +- .../execution/ExecutionManagerTest.java | 57 +++++ 5 files changed, 345 insertions(+), 19 deletions(-) create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ContinuationFailurePublicProbeTest.java diff --git a/otel-plugin/README.md b/otel-plugin/README.md index db05972b7..e3ffc8a5a 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -52,12 +52,15 @@ completes the observation future exceptionally with that same fatal before escap body, or end hook runs, and no durable `FAILED` response is produced for that fatal. Ordinary initialization errors and handler/body failure classification retain their existing behavior. -A direct `VirtualMachineError` or `ThreadDeath` from an SDK checkpoint continuation selects a retryable invocation -control failure before the continuation releases its activity lease. The caller receives the original fatal as the -cause of `UnrecoverableDurableExecutionException`, and started hooks receive `RETRYING`. The continuation's -observation future settles with the original fatal after lease release and before the fatal escapes its coordinator -worker. Already selected outcomes retain first-completion precedence. This direct-fatal continuation boundary -does not add general wrapped-fatal classification or change ordinary continuation and handler/body failures. +An SDK checkpoint continuation owned by an operation reports resumption/deserialization and dispatch failures +through retryable manager control before releasing its activity lease. The caller receives the original failure as +the cause of `UnrecoverableDurableExecutionException`, and started End hooks receive `RETRYING`; persisted operation +state is unchanged for the next invocation. A rejected user-worker dispatch rolls back its activity registration. +Ordinary failures during normal manager closing do not replace the selected outcome or stop unrelated operations. +Unowned ordinary helper failures retain their observation-only behavior. Direct `VirtualMachineError` or +`ThreadDeath` failures also settle their observation future after lease release and before escaping the coordinator +worker. Observation cancellation cannot skip the actual task or discard its operation failure. This boundary does +not reclassify handler/predicate failures or add general wrapped-fatal classification. SDK inspection of MDC-capture failures reads each visited standard transport cause once and detects identity cycles. Cyclic, null, or unreadable leading `CompletionException` chains retain the original wrapper; ordinary initialization diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ContinuationFailurePublicProbeTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ContinuationFailurePublicProbeTest.java new file mode 100644 index 000000000..7837527b5 --- /dev/null +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ContinuationFailurePublicProbeTest.java @@ -0,0 +1,241 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.*; + +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Set; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import java.util.function.BiFunction; +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.config.WaitForConditionConfig; +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.DurableExecutionOutput; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.model.WaitForConditionResult; +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; +import software.amazon.lambda.durable.serde.SerDes; +import software.amazon.lambda.durable.testing.local.LocalMemoryExecutionClient; + +class ContinuationFailurePublicProbeTest { + @SuppressWarnings("removal") + @ParameterizedTest + @CsvSource({"false,false", "false,true", "true,false", "true,true"}) + void readyResumeFailureDoesNotDisappearOrHang(boolean rejection, boolean withPlugin) throws Exception { + RuntimeException failure = rejection + ? new RejectedExecutionException("resumed worker rejected") + : new IllegalArgumentException("resumed state cannot deserialize"); + var managerForCleanup = new AtomicReference(); + var armed = new AtomicBoolean(); + var injected = new CountDownLatch(1); + var injectionOwner = new AtomicReference(); + var checks = new AtomicInteger(); + var ends = new CopyOnWriteArrayList(); + var pollEntered = new CountDownLatch(1); + var releaseReady = new CountDownLatch(1); + var firstArmedPoll = new AtomicBoolean(true); + var onCoordinator = new AtomicBoolean(); + var pendingCall = new AtomicReference>(); + var serde = new SerDes() { + private final JacksonSerDes delegate = new JacksonSerDes(); + + public String serialize(Object value) { + return delegate.serialize(value); + } + + public T deserialize(String value, TypeToken type) { + if (!rejection + && armed.get() + && Thread.currentThread().getName().startsWith("durable-sdk-internal-")) { + injectionOwner.set(Thread.currentThread().getName()); + onCoordinator.set(Arrays.stream(Thread.currentThread().getStackTrace()) + .anyMatch(frame -> frame.getMethodName().equals("completeCheckpointContinuation"))); + injected.countDown(); + throw failure; + } + return delegate.deserialize(value, type); + } + }; + var client = new LocalMemoryExecutionClient() { + @Override + public CheckpointDurableExecutionResponse checkpoint( + String arn, String token, List updates) { + if (armed.get() && updates.isEmpty() && firstArmedPoll.compareAndSet(true, false)) { + pollEntered.countDown(); + try { + assertTrue(releaseReady.await(3, TimeUnit.SECONDS)); + } catch (InterruptedException interrupted) { + throw new AssertionError(interrupted); + } + } + return super.checkpoint(arn, token, updates); + } + }; + var workers = + new ThreadPoolExecutor(0, Integer.MAX_VALUE, 60, TimeUnit.SECONDS, new SynchronousQueue()) { + @Override + public void execute(Runnable task) { + if (rejection + && armed.get() + && Thread.currentThread().getName().startsWith("durable-sdk-internal-")) { + injectionOwner.set(Thread.currentThread().getName()); + onCoordinator.set(Arrays.stream( + Thread.currentThread().getStackTrace()) + .anyMatch(frame -> frame.getMethodName().equals("completeCheckpointContinuation"))); + injected.countDown(); + throw failure; + } + super.execute(task); + } + }; + var caller = Executors.newSingleThreadExecutor(); + var builder = DurableConfig.builder() + .withExecutorService(workers) + .withDurableExecutionClient(client) + .withCheckpointDelay(Duration.ZERO) + .withPollingStrategy(attempt -> Duration.ZERO); + if (withPlugin) + builder.withPlugins(new DurableExecutionPlugin() { + public void onInvocationEnd(InvocationEndInfo info) { + ends.add(info); + } + }); + var config = builder.build(); + var condition = WaitForConditionConfig.builder() + .initialState(1) + .serDes(serde) + .waitStrategy((state, attempt) -> Duration.ofSeconds(1)) + .build(); + BiFunction handler = (input, context) -> { + managerForCleanup.set( + ((software.amazon.lambda.durable.context.BaseContextImpl) context).getExecutionManager()); + var waiting = context.waitForConditionAsync( + "condition", + Integer.class, + (state, step) -> { + checks.incrementAndGet(); + return state == 2 + ? WaitForConditionResult.stopPolling(state) + : WaitForConditionResult.continuePolling(2); + }, + condition); + if (armed.get()) { + try { + assertTrue(pollEntered.await(3, TimeUnit.SECONDS)); + } catch (InterruptedException interrupted) { + throw new AssertionError(interrupted); + } + } + return String.valueOf(waiting.get()); + }; + try { + var first = caller.submit(() -> DurableExecutor.execute( + input(List.of()), null, TypeToken.get(String.class), handler, config)) + .get(3, TimeUnit.SECONDS); + assertEquals(ExecutionStatus.PENDING, first.status()); + var cachedPending = new ArrayList<>(client.getAllOperations()); + assertEquals( + OperationStatus.PENDING, + client.getOperationByName("condition").status()); + assertTrue(client.advanceTime()); + armed.set(true); + var second = caller.submit(() -> + DurableExecutor.execute(input(cachedPending), null, TypeToken.get(String.class), handler, config)); + pendingCall.set(second); + assertTrue(pollEntered.await(3, TimeUnit.SECONDS)); + releaseReady.countDown(); + assertTrue(injected.await(3, TimeUnit.SECONDS)); + assertTrue(onCoordinator.get(), "Fault must run in the actual WFC coordinator continuation"); + DurableExecutionOutput output = null; + Throwable callerFailure = null; + try { + output = second.get(3, TimeUnit.SECONDS); + } catch (ExecutionException observationFailure) { + callerFailure = observationFailure.getCause(); + } + System.out.println("CONTINUATION_ORDINARY_PUBLIC owner=" + injectionOwner.get() + + " caller=" + + (output != null + ? output.status() + : callerFailure.getClass().getName()) + + " backend=" + client.getOperationByName("condition").status() + + " checks=" + checks.get()); + assertEquals(1, checks.get(), "A failed resumed state read must not rerun the condition body"); + assertNull(output, "An SDK continuation failure must not produce a durable response"); + var retry = assertInstanceOf(UnrecoverableDurableExecutionException.class, callerFailure); + assertTrue(retry.isRetryable()); + assertSame(failure, retry.getCause()); + assertEquals( + OperationStatus.READY, + client.getOperationByName("condition").status()); + var activeField = ExecutionManager.class.getDeclaredField("activeThreads"); + activeField.setAccessible(true); + assertFalse( + ((Set) activeField.get(managerForCleanup.get())) + .contains(client.getOperationByName("condition").id()), + "A rejected operation worker must not leave a phantom activity registration"); + if (withPlugin) + assertEquals( + List.of(InvocationStatus.PENDING, InvocationStatus.RETRYING), + ends.stream().map(InvocationEndInfo::invocationStatus).toList()); + else assertTrue(ends.isEmpty()); + armed.set(false); + var third = caller.submit(() -> DurableExecutor.execute( + input(client.getAllOperations()), null, TypeToken.get(String.class), handler, config)); + pendingCall.set(third); + assertEquals( + ExecutionStatus.SUCCEEDED, third.get(3, TimeUnit.SECONDS).status()); + assertEquals(2, checks.get(), "Only the original and successful resumed predicates run"); + } finally { + releaseReady.countDown(); + // Cleanup-only escape hatch AFTER the real public invocation observation/timeout; never test behavior. + var manager = managerForCleanup.get(); + if (pendingCall.get() != null + && !pendingCall.get().isDone() + && manager != null + && !manager.isExecutionCompletedExceptionally()) { + try { + manager.terminateExecution(new UnrecoverableDurableExecutionException( + ErrorObject.builder().errorMessage("probe cleanup").build(), true)); + } catch (UnrecoverableDurableExecutionException expected) { + } + } + caller.shutdownNow(); + workers.shutdownNow(); + assertTrue(caller.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private static DurableExecutionInput input(List operations) { + var execution = Operation.builder() + .id("execution") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails( + ExecutionDetails.builder().inputPayload("\"input\"").build()) + .build(); + var history = new ArrayList(); + history.add(execution); + history.addAll(operations); + return new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/fatal/execution", + "token", + CheckpointUpdatedExecutionState.builder().operations(history).build()); + } +} 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 2aff9a2c5..427e1e979 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 @@ -580,24 +580,36 @@ CompletableFuture runCheckpointContinuation( } try { var completion = new CompletableFuture(); - coordinator.execute(() -> completeCheckpointContinuation(registration, continuation, completion)); + coordinator.execute(() -> completeCheckpointContinuation(owner, registration, continuation, completion)); return completion; } catch (RuntimeException | Error failure) { - finishCheckpointContinuation(registration); + try { + if ((owner != null && !isClosing()) + || failure instanceof VirtualMachineError + || failure instanceof ThreadDeath) signalContinuationFailure(failure); + } finally { + finishCheckpointContinuation(registration); + } throw failure; } } @SuppressWarnings("removal") private void completeCheckpointContinuation( - Object registration, Runnable continuation, CompletableFuture completion) { + BaseDurableOperation owner, + Object registration, + Runnable continuation, + CompletableFuture completion) { try { try { if (!isClosing()) continuation.run(); - } catch (VirtualMachineError | ThreadDeath fatal) { - // Choose retry control flow before releasing the final lease could otherwise select PENDING. - signalContinuationFatal(fatal); - throw fatal; + } catch (Throwable failure) { + // An operation polling chain can ignore observation. Select retry control before the final lease + // could select PENDING; unowned ordinary helper failures retain their observation-only contract. + if ((owner != null && !isClosing()) + || failure instanceof VirtualMachineError + || failure instanceof ThreadDeath) signalContinuationFailure(failure); + throw failure; } finally { finishCheckpointContinuation(registration); } @@ -610,14 +622,14 @@ private void completeCheckpointContinuation( } } - private void signalContinuationFatal(Error fatal) { + private void signalContinuationFailure(Throwable failure) { var control = new UnrecoverableDurableExecutionException( ErrorObject.builder() - .errorType(fatal.getClass().getName()) - .errorMessage("Fatal error in SDK checkpoint continuation") + .errorType(failure.getClass().getName()) + .errorMessage("Error in SDK checkpoint continuation") .build(), true, - fatal); + failure); if (executionExceptionFuture.completeExceptionally(control)) stopAllOperations(control); } 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 a1514e6d5..bdb7fd11a 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 @@ -321,8 +321,21 @@ protected void runUserHandler(Runnable runnable, ThreadType threadType) { // registerActiveThread is idempotent (no-op if already registered). registerActiveThread(operationId); - runningUserHandler.set(CompletableFuture.runAsync( - wrapped, getContext().getDurableConfig().getExecutorService())); + try { + runningUserHandler.set(CompletableFuture.runAsync( + wrapped, getContext().getDurableConfig().getExecutorService())); + } catch (RuntimeException | Error dispatchFailure) { + // Admission failed, so no worker will execute its deregistration finally block. The submitting + // coordinator need not have a user ThreadContext; undo the registration directly on the manager. + if (operationId != null) { + try { + executionManager.deregisterActiveThread(operationId); + } catch (SuspendExecutionException ignored) { + // Preserve the dispatch failure; any suspension was already signaled by deregistration. + } + } + throw dispatchFailure; + } } /** diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java index 1c31805ea..25ea5a3fa 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java @@ -604,6 +604,63 @@ void cancellingRunningContinuationDoesNotReleaseActivityBeforeItsCleanup() throw assertTrue(completion.isCancelled()); } + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void ownedFailureReachesExecutionDespiteCancelledObservation(boolean cancelled) throws Exception { + var manager = createManager(List.of(executionOp(), stepOp("step", OperationStatus.PENDING))); + manager.registerActiveThread("root"); + var owner = mock(BaseDurableOperation.class); + when(owner.getCompletionFuture()).thenReturn(new CompletableFuture<>()); + var queued = new AtomicReference(); + var failure = new IllegalArgumentException("owned continuation failure"); + var selected = manager.runUntilCompleteOrSuspend(new CompletableFuture<>()); + var observation = manager.runCheckpointContinuation( + owner, + () -> { + throw failure; + }, + queued::set); + if (cancelled) assertTrue(observation.cancel(false)); + queued.get().run(); + var control = assertInstanceOf( + UnrecoverableDurableExecutionException.class, + assertThrows(ExecutionException.class, () -> selected.get(3, TimeUnit.SECONDS)) + .getCause()); + assertTrue(control.isRetryable()); + assertSame(failure, control.getCause()); + if (cancelled) assertTrue(observation.isCancelled()); + else + assertSame( + failure, + assertThrows(ExecutionException.class, () -> observation.get(3, TimeUnit.SECONDS)) + .getCause()); + CompletableFuture.runAsync(manager::close).get(3, TimeUnit.SECONDS); + } + + @Test + void ownedCoordinatorRejectionSettlesExecutionAndReleasesAdmission() throws Exception { + var manager = createManager(List.of(executionOp(), stepOp("step", OperationStatus.PENDING))); + manager.registerActiveThread("root"); + var owner = mock(BaseDurableOperation.class); + when(owner.getCompletionFuture()).thenReturn(new CompletableFuture<>()); + var selected = manager.runUntilCompleteOrSuspend(new CompletableFuture<>()); + var rejection = new RejectedExecutionException("coordinator admission rejected"); + assertSame( + rejection, + assertThrows( + RejectedExecutionException.class, + () -> manager.runCheckpointContinuation(owner, () -> fail("not admitted"), task -> { + throw rejection; + }))); + var control = assertInstanceOf( + UnrecoverableDurableExecutionException.class, + assertThrows(ExecutionException.class, () -> selected.get(3, TimeUnit.SECONDS)) + .getCause()); + assertTrue(control.isRetryable()); + assertSame(rejection, control.getCause()); + CompletableFuture.runAsync(manager::close).get(3, TimeUnit.SECONDS); + } + @Test void rejectedContinuationReleasesActivityAndPreservesTheRejection() { var manager = createManager(List.of(executionOp(), stepOp("step", OperationStatus.PENDING))); From 6794136501d52a1dd2c79ea3ff5b7081b995aaa1 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 22:48:46 +0000 Subject: [PATCH 43/48] fix: retain retry diagnostics and fatal identity through End cleanup --- otel-plugin/README.md | 7 ++-- .../durable/otel/MdcFailureBoundaryTest.java | 19 ++++++++++ .../ContinuationFatalPublicProbeTest.java | 35 ++++++++++++++++++- .../durable/execution/DurableExecutor.java | 25 +++++++++++-- .../execution/ExecutionManagerTest.java | 26 ++++++++++++++ 5 files changed, 106 insertions(+), 6 deletions(-) diff --git a/otel-plugin/README.md b/otel-plugin/README.md index e3ffc8a5a..9f1107b5a 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -37,14 +37,17 @@ SDK output preparation, including customer `SerDes` calls and durable large-resu terminal invocation-end notification. Failures in this preparation report `RETRYING` instead of ending the Workflow span. If End also raises an unisolated Error, the preparation failure remains primary with the cleanup error suppressed, unless cleanup introduces the first JVM-fatal error. An original JVM-fatal preparation error retains its -identity. End describes the SDK outcome at that point, not acknowledgment of a response by the Lambda service. +identity. Retryable control failures follow the same End-error combination: an ordinary End Error is suppressed +under the retry control, while a direct JVM-fatal End error retains its existing precedence with the control and +its cause preserved as diagnostics. This does not add arbitrary cause unwrapping to handler failure classification. +End describes the SDK outcome at that point, not acknowledgment of a response by the Lambda service. Caller-side execution-manager cleanup, response-envelope encoding and output-stream writes follow End; runtime response transport follows the handler return. Failures at those later boundaries still propagate, without a second End dispatch or changing its already reported outcome. A JVM-fatal MDC-restoration failure after End (direct or inside a standard transport wrapper) completes the caller's observation exceptionally with the original fatal before that fatal escapes the worker, for both asynchronous and inline executors. The End notification remains unchanged and is not repeated. An earlier JVM-fatal End/preparation failure remains primary; later restoration -failures are suppressed. A restoration fatal takes precedence over an earlier non-JVM-fatal delivery/End failure, +failures are suppressed, with an identity guard when restoration throws the same fatal object. A restoration fatal takes precedence over an earlier non-JVM-fatal delivery/End failure, which remains suppressed. Non-JVM-fatal restoration failures, including `AssertionError` and `LinkageError`, retain the selected caller outcome. Before invocation startup, a JVM-fatal error from MDC capture (direct or inside a standard transport wrapper) diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFailureBoundaryTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFailureBoundaryTest.java index 14ccc4e8b..93cce8a8a 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFailureBoundaryTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcFailureBoundaryTest.java @@ -211,6 +211,25 @@ void earlierEndFatalRemainsPrimaryAndRetainsRestorationDiagnostic(boolean inline assertArrayEquals(new Throwable[] {restoration}, primary.getSuppressed()); } + @SuppressWarnings("removal") + @ParameterizedTest + @CsvSource({ + "false,false,false", + "false,false,true", + "false,true,false", + "false,true,true", + "true,false,false", + "true,false,true", + "true,true,false", + "true,true,true" + }) + void sameFatalAtEndAndRestorationRetainsOwnerIdentity(boolean inline, boolean suspend, boolean death) + throws Exception { + Error fatal = death ? new ThreadDeath() : new InternalError("same End/restore fatal"); + exercisePostEndFailure(inline, suspend, fatal, fatal, fatal); + assertEquals(0, fatal.getSuppressed().length, "Never suppress the fatal onto itself"); + } + @ParameterizedTest @CsvSource({"false,false", "false,true", "true,false", "true,true"}) void restorationFatalRetainsEarlierNonfatalEndFailure(boolean inline, boolean wrapped) throws Exception { diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ContinuationFatalPublicProbeTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ContinuationFatalPublicProbeTest.java index 770f99579..43e76b8cb 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ContinuationFatalPublicProbeTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ContinuationFatalPublicProbeTest.java @@ -34,6 +34,24 @@ class ContinuationFatalPublicProbeTest { @ParameterizedTest @CsvSource({"false,false", "false,true", "true,false", "true,true"}) void readyResumeSerdeFatalDoesNotDisappearIntoPending(boolean death, boolean withPlugin) throws Exception { + runFatalWithEndFailure(death, withPlugin, null); + } + + @ParameterizedTest + @CsvSource({"false,assertion", "true,assertion", "false,fatal", "true,fatal"}) + void continuationFatalKeepsDiagnosticsWhenEndAlsoFails(boolean death, String endKind) throws Exception { + runFatalWithEndFailure(death, true, endKind); + } + + @SuppressWarnings("removal") + private static void runFatalWithEndFailure(boolean death, boolean withPlugin, String endKind) throws Exception { + Error endFailure = endKind == null + ? null + : endKind.equals("fatal") + ? (death ? new InternalError("End fatal") : new ThreadDeath()) + : new AssertionError("End assertion"); + var endEscaped = new CountDownLatch(1); + var endOwnerFailure = new AtomicReference(); Error fatal = death ? new ThreadDeath() : new InternalError("resume serde fatal"); var armed = new AtomicBoolean(); var injected = new CountDownLatch(1); @@ -51,6 +69,9 @@ void readyResumeSerdeFatalDoesNotDisappearIntoPending(boolean death, boolean wit if (failure == fatal) { ownerFailure.set(failure); escaped.countDown(); + } else if (failure == endFailure) { + endOwnerFailure.set(failure); + endEscaped.countDown(); } else if (previousHandler != null) previousHandler.uncaughtException(thread, failure); }); var serde = new SerDes() { @@ -97,6 +118,7 @@ public CheckpointDurableExecutionResponse checkpoint( builder.withPlugins(new DurableExecutionPlugin() { public void onInvocationEnd(InvocationEndInfo info) { ends.add(info); + if (info.invocationStatus() == InvocationStatus.RETRYING && endFailure != null) throw endFailure; } }); var config = builder.build(); @@ -161,7 +183,18 @@ public void onInvocationEnd(InvocationEndInfo info) { assertTrue(ownerEscaped, "The actual coordinator worker must observe its JVM fatal"); assertSame(fatal, ownerFailure.get()); assertNull(output, "A failed coordinator must not return a durable terminal or PENDING response"); - var retry = assertInstanceOf(UnrecoverableDurableExecutionException.class, callerFailure); + UnrecoverableDurableExecutionException retry; + if ("fatal".equals(endKind)) { + assertSame(endFailure, callerFailure, "Existing Root End JVM-fatal precedence remains"); + assertTrue(endEscaped.await(3, TimeUnit.SECONDS)); + assertSame(endFailure, endOwnerFailure.get()); + assertEquals(1, endFailure.getSuppressed().length); + retry = assertInstanceOf( + UnrecoverableDurableExecutionException.class, endFailure.getSuppressed()[0]); + } else { + retry = assertInstanceOf(UnrecoverableDurableExecutionException.class, callerFailure); + if (endFailure != null) assertArrayEquals(new Throwable[] {endFailure}, retry.getSuppressed()); + } assertTrue(retry.isRetryable()); assertSame(fatal, retry.getCause()); if (withPlugin) { 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 0c701be71..9d2150e49 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 @@ -183,7 +183,13 @@ private DurableExecutionOutput finishInvocation(O value, Throwable failure) { private DurableExecutionOutput finishFailure(Throwable cause) { var status = failureStatus(cause); if (status == InvocationStatus.FAILED) return finishTerminalFailure(cause); - fireOnInvocationEnd(status, status == InvocationStatus.PENDING ? null : cause, null); + try { + fireOnInvocationEnd(status, status == InvocationStatus.PENDING ? null : cause, null); + } catch (Error endFailure) { + if (status == InvocationStatus.RETRYING) + ExceptionHelper.sneakyThrow(combinePreparationAndEndFailures(cause, endFailure)); + throw endFailure; + } if (status == InvocationStatus.PENDING) return DurableExecutionOutput.pending(); ExceptionHelper.sneakyThrow(cause); return null; @@ -296,12 +302,16 @@ private static CompletableFuture supplyHandler( var outcome = Outcome.capture(task); var restoringAfterNonfatalOutcome = false; try { - try (var ignored = restore) { + try { // End/delivery failures must also escape their actual worker. Handler-body failures have // already been mapped to their selected durable outcome; this does not reclassify them. rethrowLifecycleFatal(outcome.failure()); - restoringAfterNonfatalOutcome = true; + } catch (VirtualMachineError | ThreadDeath fatal) { + closeMdcAfterFatal(restore, fatal); + throw fatal; } + restoringAfterNonfatalOutcome = true; + restore.close(); } catch (Throwable workerFailure) { try { rethrowLifecycleFatal(workerFailure); @@ -334,6 +344,15 @@ private static CompletableFuture supplyHandler( return result; } + /** A repeated fatal object must not be replaced by try-with-resources self-suppression failure. */ + private static void closeMdcAfterFatal(SafeCloseable restore, Error fatal) { + try { + restore.close(); + } catch (Throwable restorationFailure) { + if (restorationFailure != fatal) fatal.addSuppressed(restorationFailure); + } + } + /** Classifies capture failures once, retaining the ordinary policy of unwrapping only a completion prefix. */ private static Throwable normalizeMdcInitializationFailure(Throwable failure) { rethrowDirectMdcFatal(failure); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java index 25ea5a3fa..872a69181 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java @@ -604,6 +604,32 @@ void cancellingRunningContinuationDoesNotReleaseActivityBeforeItsCleanup() throw assertTrue(completion.isCancelled()); } + @SuppressWarnings("removal") + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void coordinatorAdmissionFatalSelectsRetryBeforeLeaseRelease(boolean death) throws Exception { + Error fatal = death ? new ThreadDeath() : new InternalError("coordinator admission fatal"); + var manager = createManager(List.of(executionOp(), stepOp("step", OperationStatus.PENDING))); + manager.registerActiveThread("root"); + var owner = mock(BaseDurableOperation.class); + when(owner.getCompletionFuture()).thenReturn(new CompletableFuture<>()); + var selected = manager.runUntilCompleteOrSuspend(new CompletableFuture<>()); + assertSame( + fatal, + assertThrows( + Error.class, + () -> manager.runCheckpointContinuation(owner, () -> fail("No accepted task"), task -> { + throw fatal; + }))); + var retry = assertInstanceOf( + UnrecoverableDurableExecutionException.class, + assertThrows(ExecutionException.class, () -> selected.get(3, TimeUnit.SECONDS)) + .getCause()); + assertTrue(retry.isRetryable()); + assertSame(fatal, retry.getCause()); + CompletableFuture.runAsync(manager::close).get(3, TimeUnit.SECONDS); + } + @ParameterizedTest @ValueSource(booleans = {false, true}) void ownedFailureReachesExecutionDespiteCancelledObservation(boolean cancelled) throws Exception { From 78e3b48d4534c400862aabed308e4129cc4c9c1d Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 23:32:13 +0000 Subject: [PATCH 44/48] fix: clear inheritable MDC after ordinary restoration failure --- otel-plugin/README.md | 7 +- .../otel/MdcWorkerReuseBoundaryTest.java | 161 ++++++++++++++++++ .../durable/execution/DurableExecutor.java | 17 ++ 3 files changed, 184 insertions(+), 1 deletion(-) create mode 100644 otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcWorkerReuseBoundaryTest.java diff --git a/otel-plugin/README.md b/otel-plugin/README.md index 9f1107b5a..503b7a8a5 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -49,7 +49,12 @@ before that fatal escapes the worker, for both asynchronous and inline executors unchanged and is not repeated. An earlier JVM-fatal End/preparation failure remains primary; later restoration failures are suppressed, with an identity guard when restoration throws the same fatal object. A restoration fatal takes precedence over an earlier non-JVM-fatal delivery/End failure, which remains suppressed. Non-JVM-fatal restoration failures, including `AssertionError` and `LinkageError`, retain -the selected caller outcome. +the selected caller outcome. After an ordinary outer restoration failure, the SDK makes one guarded `MDC.clear()` +attempt before rethrowing the original worker failure. A successful clear prevents inheritable invocation state +from reaching a replacement pool thread. A failing clear is retained as a suppressed diagnostic, or its JVM fatal +uses the same caller-settlement-before-worker-throw rule. Clearing can discard ambient MDC that could not be +restored. If the adapter's clear also fails, clean replacement state is not guaranteed; the SDK does not promise +quarantine for an arbitrary caller-owned executor. Initialization and handler/body policies are unchanged. Before invocation startup, a JVM-fatal error from MDC capture (direct or inside a standard transport wrapper) completes the observation future exceptionally with that same fatal before escaping the handler worker. No start, body, or end hook runs, and no durable `FAILED` response is produced for that fatal. Ordinary initialization errors diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcWorkerReuseBoundaryTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcWorkerReuseBoundaryTest.java new file mode 100644 index 000000000..d11cb44f5 --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcWorkerReuseBoundaryTest.java @@ -0,0 +1,161 @@ +// 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 java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Proxy; +import java.time.Duration; +import java.util.List; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import java.util.stream.Stream; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; +import org.slf4j.MDC; +import org.slf4j.helpers.BasicMDCAdapter; +import org.slf4j.spi.MDCAdapter; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.plugin.*; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class MdcWorkerReuseBoundaryTest { + @ParameterizedTest + @ValueSource(strings = {"start", "end"}) + void replacementWorkerDoesNotInheritInvocationMdcWhenClearSucceeds(String markerPhase) throws Exception { + exercise(markerPhase, false, "ok"); + } + + static Stream fallbackCases() { + return Stream.of(false, true) + .flatMap(suspend -> Stream.of( + "ordinary", "assertion", "same", "vm", "death", "wrapped-vm", "wrapped-death") + .map(kind -> Arguments.of(suspend, kind))); + } + + @ParameterizedTest + @MethodSource("fallbackCases") + void failedFallbackRetainsFailurePolicyAndFatalIdentity(boolean suspend, String kind) throws Exception { + exercise("end", suspend, kind); + } + + @SuppressWarnings("removal") + private static void exercise(String markerPhase, boolean suspend, String kind) throws Exception { + var original = MDC.getMDCAdapter(); + var basic = new BasicMDCAdapter(); + var ends = new AtomicInteger(); + var endStatus = new AtomicReference(); + var failed = new AtomicBoolean(); + var clearTried = new AtomicBoolean(); + var restoreFailure = new IllegalStateException("restore failed"); + Error expectedFatal = + switch (kind) { + case "vm", "wrapped-vm" -> new InternalError("fallback fatal"); + case "death", "wrapped-death" -> new ThreadDeath(); + default -> null; + }; + Throwable clearFailure = + switch (kind) { + case "ordinary" -> new IllegalArgumentException("clear failed"); + case "assertion" -> new AssertionError("clear failed"); + case "same" -> restoreFailure; + case "wrapped-vm", "wrapped-death" -> + new CompletionException(new ExecutionException(expectedFatal)); + default -> expectedFatal; + }; + var escaped = new CountDownLatch(1); + var failedWorker = new AtomicReference(); + var ownerFailure = new AtomicReference(); + var dirtyAtFailure = new AtomicReference(); + var count = new AtomicInteger(); + var workers = Executors.newSingleThreadExecutor(task -> { + var thread = new Thread(task, "mdc-inheritable-worker-" + count.incrementAndGet()); + thread.setUncaughtExceptionHandler((owner, failure) -> { + failedWorker.set(owner); + ownerFailure.set(failure); + dirtyAtFailure.set(MDC.get("invocation-marker")); + escaped.countDown(); + }); + return thread; + }); + try { + var adapter = (MDCAdapter) Proxy.newProxyInstance( + MDCAdapter.class.getClassLoader(), new Class[] {MDCAdapter.class}, (proxy, method, args) -> { + if (ends.get() == 1 + && method.getName().equals("setContextMap") + && Thread.currentThread().getName().startsWith("mdc-inheritable-worker-") + && failed.compareAndSet(false, true)) throw restoreFailure; + if (failed.get() + && method.getName().equals("clear") + && Thread.currentThread().getName().startsWith("mdc-inheritable-worker-") + && clearTried.compareAndSet(false, true) + && clearFailure != null) throw clearFailure; + try { + return method.invoke(basic, args); + } catch (InvocationTargetException failure) { + throw failure.getCause(); + } + }); + replace(adapter); + MDC.put("ambient", "saved"); + var plugin = new DurableExecutionPlugin() { + public void onInvocationStart(InvocationInfo info) { + if (markerPhase.equals("start")) MDC.put("invocation-marker", "previous-invocation"); + } + + public void onInvocationEnd(InvocationEndInfo info) { + if (markerPhase.equals("end")) MDC.put("invocation-marker", "previous-invocation"); + endStatus.set(info.invocationStatus()); + ends.incrementAndGet(); + } + }; + var runner = LocalDurableTestRunner.create( + String.class, + (value, context) -> { + if (suspend) context.wait("resume", Duration.ofSeconds(1)); + return "done"; + }, + DurableConfig.builder() + .withExecutorService(workers) + .withPlugins(plugin) + .build()); + if (expectedFatal != null) assertSame(expectedFatal, assertThrows(Error.class, () -> runner.run("input"))); + else + assertEquals( + suspend ? ExecutionStatus.PENDING : ExecutionStatus.SUCCEEDED, + runner.run("input").getStatus()); + assertTrue(escaped.await(3, TimeUnit.SECONDS)); + assertSame(expectedFatal != null ? expectedFatal : restoreFailure, ownerFailure.get()); + assertTrue(clearTried.get()); + assertEquals(1, ends.get()); + assertEquals(suspend ? InvocationStatus.PENDING : InvocationStatus.SUCCEEDED, endStatus.get()); + if (expectedFatal != null) assertEquals(List.of(restoreFailure), List.of(expectedFatal.getSuppressed())); + else + assertEquals( + clearFailure != null && clearFailure != restoreFailure ? List.of(clearFailure) : List.of(), + List.of(restoreFailure.getSuppressed())); + var nextThread = workers.submit(Thread::currentThread).get(3, TimeUnit.SECONDS); + var inherited = workers.submit(() -> MDC.get("invocation-marker")).get(3, TimeUnit.SECONDS); + assertNotSame(failedWorker.get(), nextThread); + System.out.println("MDC_REUSE_PUBLIC phase=" + markerPhase + " kind=" + kind + " replaced=true oldMarker=" + + dirtyAtFailure.get() + " inheritedMarker=" + inherited); + if (clearFailure == null) assertNull(inherited, "Successful fallback clears inheritable invocation state"); + // An adapter whose clear also fails cannot promise clean replacement state; its failure is exposed. + } finally { + workers.shutdownNow(); + assertTrue(workers.awaitTermination(3, TimeUnit.SECONDS)); + basic.clear(); + replace(original); + } + } + + private static void replace(MDCAdapter adapter) throws Exception { + var setter = MDC.class.getDeclaredMethod("setMDCAdapter", MDCAdapter.class); + setter.setAccessible(true); + setter.invoke(null, adapter); + } +} 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 9d2150e49..094ac8a7a 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 @@ -315,6 +315,7 @@ private static CompletableFuture supplyHandler( } catch (Throwable workerFailure) { try { rethrowLifecycleFatal(workerFailure); + clearMdcAfterOrdinaryFailure(workerFailure); } catch (VirtualMachineError | ThreadDeath fatal) { if (restoringAfterNonfatalOutcome && outcome.failure() != null && outcome.failure() != fatal) { fatal.addSuppressed(outcome.failure()); @@ -344,6 +345,22 @@ private static CompletableFuture supplyHandler( return result; } + @SuppressWarnings("removal") + private static void clearMdcAfterOrdinaryFailure(Throwable restorationFailure) { + try { + MDC.clear(); + } catch (Throwable clearFailure) { + if (clearFailure == restorationFailure) return; // Already classified; do not re-read its diagnostics. + try { + rethrowLifecycleFatal(clearFailure); + } catch (VirtualMachineError | ThreadDeath fatal) { + if (fatal != restorationFailure) fatal.addSuppressed(restorationFailure); + throw fatal; + } + restorationFailure.addSuppressed(clearFailure); + } + } + /** A repeated fatal object must not be replaced by try-with-resources self-suppression failure. */ private static void closeMdcAfterFatal(SafeCloseable restore, Error fatal) { try { From 43b8c34c9c8317e7c6dd5cca23b0e839f2a2362a Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 9 Oct 2026 00:09:04 +0000 Subject: [PATCH 45/48] fix: coordinate step retries outside the checkpoint batcher --- .../durable/StepRetryContinuationTest.java | 483 ++++++++++++++++++ .../durable/operation/StepOperation.java | 24 +- 2 files changed, 500 insertions(+), 7 deletions(-) create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StepRetryContinuationTest.java diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StepRetryContinuationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StepRetryContinuationTest.java new file mode 100644 index 000000000..ec723205c --- /dev/null +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StepRetryContinuationTest.java @@ -0,0 +1,483 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.*; + +import java.lang.management.ManagementFactory; +import java.lang.reflect.Method; +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Deque; +import java.util.List; +import java.util.Map; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import java.util.concurrent.locks.LockSupport; +import java.util.function.BiFunction; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.ValueSource; +import org.slf4j.MDC; +import org.slf4j.spi.MDCAdapter; +import software.amazon.awssdk.services.lambda.model.*; +import software.amazon.lambda.durable.config.StepConfig; +import software.amazon.lambda.durable.context.BaseContextImpl; +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.DurableExecutionOutput; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.retry.RetryDecision; +import software.amazon.lambda.durable.testing.local.LocalMemoryExecutionClient; + +class StepRetryContinuationTest { + @ParameterizedTest + @ValueSource(strings = {"async", "reject", "direct", "submit-wait"}) + void realPendingStepRetryMustNotLoseDispatchFailureOrBlockItsCheckpointBatcher(String mode) throws Exception { + var armed = new AtomicBoolean(); + var calls = new AtomicInteger(); + var pollEntered = new CountDownLatch(1); + var releasePoll = new CountDownLatch(1); + var pollOnce = new AtomicBoolean(true); + var dispatchTimedOut = new AtomicBoolean(); + var dispatchThread = new AtomicReference(); + var dispatchFailure = new RejectedExecutionException("resumed step rejected"); + var managerRef = new AtomicReference(); + var pendingCall = new AtomicReference>(); + var client = new LocalMemoryExecutionClient() { + @Override + public CheckpointDurableExecutionResponse checkpoint( + String arn, String token, List updates) { + if (armed.get() && updates.isEmpty() && pollOnce.compareAndSet(true, false)) { + pollEntered.countDown(); + await(releasePoll); + } + return super.checkpoint(arn, token, updates); + } + }; + var workers = + new ThreadPoolExecutor(0, Integer.MAX_VALUE, 60, TimeUnit.SECONDS, new SynchronousQueue()) { + @Override + public void execute(Runnable task) { + boolean resume = armed.get() + && Arrays.stream(Thread.currentThread().getStackTrace()) + .anyMatch(frame -> frame.getClassName().endsWith(".StepOperation") + && frame.getMethodName().equals("executeStepLogic")); + if (resume) { + dispatchThread.set(Thread.currentThread().getName()); + if (mode.equals("reject")) throw dispatchFailure; + if (mode.equals("direct")) { + task.run(); + return; + } + if (mode.equals("submit-wait")) { + var done = new CompletableFuture(); + super.execute(() -> { + try { + task.run(); + } finally { + done.complete(null); + } + }); + try { + done.get(2, TimeUnit.SECONDS); + } catch (TimeoutException expected) { + dispatchTimedOut.set(true); + } catch (Exception failure) { + throw new AssertionError(failure); + } + return; + } + } + super.execute(task); + } + }; + var config = DurableConfig.builder() + .withExecutorService(workers) + .withDurableExecutionClient(client) + .withCheckpointDelay(Duration.ZERO) + .withPollingStrategy(attempt -> Duration.ZERO) + .build(); + var stepConfig = StepConfig.builder() + .retryStrategy((error, attempt) -> + attempt < 2 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail()) + .build(); + var caller = Executors.newSingleThreadExecutor(); + BiFunction handler = (value, context) -> { + managerRef.set(((BaseContextImpl) context).getExecutionManager()); + var future = context.stepAsync( + "retry-step", + String.class, + step -> { + if (calls.incrementAndGet() == 1) throw new IllegalStateException("first attempt"); + return "stored"; + }, + stepConfig); + if (armed.get()) await(pollEntered); + return future.get(); + }; + try { + var first = caller.submit(() -> + DurableExecutor.execute(input(List.of()), null, TypeToken.get(String.class), handler, config)); + pendingCall.set(first); + assertEquals(ExecutionStatus.PENDING, first.get(5, TimeUnit.SECONDS).status()); + var pending = new ArrayList<>(client.getAllOperations()); + assertEquals( + OperationStatus.PENDING, + client.getOperationByName("retry-step").status()); + assertTrue(client.advanceTime()); + armed.set(true); + var second = caller.submit( + () -> DurableExecutor.execute(input(pending), null, TypeToken.get(String.class), handler, config)); + pendingCall.set(second); + assertTrue(pollEntered.await(5, TimeUnit.SECONDS)); + releasePoll.countDown(); + DurableExecutionOutput output = null; + Throwable failure = null; + try { + output = second.get(6, TimeUnit.SECONDS); + } catch (ExecutionException observed) { + failure = observed.getCause(); + } + System.out.println("STEP_RETRY_PUBLIC mode=" + mode + " thread=" + dispatchThread.get() + " output=" + + (output == null ? failure : output.status()) + " backend=" + + client.getOperationByName("retry-step").status() + + " calls=" + calls.get() + " dispatchTimedOut=" + dispatchTimedOut.get()); + assertFalse( + dispatchTimedOut.get(), "Submit-and-wait dispatch must not block the batcher needed by its step"); + if (mode.equals("reject")) { + var retry = assertInstanceOf(UnrecoverableDurableExecutionException.class, failure); + assertTrue(retry.isRetryable()); + assertSame(dispatchFailure, retry.getCause()); + assertNull(output); + assertEquals(1, calls.get()); + assertEquals( + OperationStatus.READY, + client.getOperationByName("retry-step").status()); + } else assertEquals(ExecutionStatus.SUCCEEDED, output.status()); + armed.set(false); + var resumed = caller.submit(() -> DurableExecutor.execute( + input(client.getAllOperations()), null, TypeToken.get(String.class), handler, config)); + pendingCall.set(resumed); + assertEquals( + ExecutionStatus.SUCCEEDED, resumed.get(5, TimeUnit.SECONDS).status()); + assertEquals(2, calls.get(), "Completed step bodies do not repeat on replay"); + } finally { + releasePoll.countDown(); + var manager = managerRef.get(); + if (pendingCall.get() != null + && !pendingCall.get().isDone() + && manager != null + && !manager.isExecutionCompletedExceptionally()) { + try { + manager.terminateExecution(new UnrecoverableDurableExecutionException( + ErrorObject.builder().errorMessage("probe cleanup").build(), true)); + } catch (UnrecoverableDurableExecutionException expected) { + } + } + caller.shutdownNow(); + workers.shutdownNow(); + assertTrue(caller.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private static void await(CountDownLatch latch) { + try { + assertTrue(latch.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw new AssertionError(interrupted); + } + } + + private static DurableExecutionInput input(List operations) { + var all = new ArrayList(); + all.add(Operation.builder() + .id("execution") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails( + ExecutionDetails.builder().inputPayload("\"input\"").build()) + .build()); + all.addAll(operations); + return new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/step-retry/execution", + "token", + CheckpointUpdatedExecutionState.builder().operations(all).build()); + } + + @ParameterizedTest + @CsvSource({"true,false", "false,true", "true,true"}) + void liveRetryRetiresItsWorkerBeforeReplacementWithoutLosingActivity( + boolean holdWorkerExit, boolean holdExecutorReturn) throws Exception { + var gate = new HandoffGate(holdWorkerExit, holdExecutorReturn); + var client = new DelayedReadyClient(gate); + var calls = new AtomicInteger(); + var caller = Executors.newSingleThreadExecutor(); + var config = DurableConfig.builder() + .withExecutorService(gate) + .withDurableExecutionClient(client) + .withCheckpointDelay(Duration.ZERO) + .withPollingStrategy(attempt -> Duration.ZERO) + .build(); + var retry = StepConfig.builder() + .retryStrategy((error, attempt) -> + attempt == 1 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail()) + .build(); + BiFunction handler = (value, context) -> { + var step = context.stepAsync( + "retry-step", + String.class, + stepContext -> { + if (calls.incrementAndGet() == 1) { + gate.firstWorker.set(Thread.currentThread()); + throw new IllegalStateException("retry once"); + } + gate.nextCheckEntered.countDown(); + await(gate.releaseNextCheck); + return "stored"; + }, + retry); + // Keep the root active until the live retry has registered its backend poll. + await(gate.pollEntered); + return step.get(); + }; + try (var mdc = new MdcClearGate(gate)) { + try { + var result = caller.submit(() -> + DurableExecutor.execute(input(List.of()), null, TypeToken.get(String.class), handler, config)); + await(gate.pollEntered); + if (holdWorkerExit) await(gate.workerExitEntered); + else await(gate.executorReturnEntered); + assertTrue(client.advanceTime()); + gate.releaseReadyResponse.countDown(); + await(gate.readyResponseReturned); + awaitCheckpointHandoff(gate.pollThread.get()); + assertThrows(TimeoutException.class, () -> result.get(100, TimeUnit.MILLISECONDS)); + assertEquals(1, calls.get()); + gate.releaseWorkerExit.countDown(); + if (holdExecutorReturn) await(gate.executorReturnEntered); + gate.releaseExecutorReturn.countDown(); + await(gate.nextCheckEntered); + assertThrows( + TimeoutException.class, + () -> result.get(100, TimeUnit.MILLISECONDS), + "The replacement retains activity after its previous worker deregisters"); + gate.releaseNextCheck.countDown(); + var completed = result.get(5, TimeUnit.SECONDS); + assertEquals(ExecutionStatus.SUCCEEDED, completed.status()); + assertEquals("\"stored\"", completed.result()); + assertEquals( + OperationStatus.SUCCEEDED, + client.getOperationByName("retry-step").status()); + var replay = caller.submit(() -> DurableExecutor.execute( + input(client.getAllOperations()), null, TypeToken.get(String.class), handler, config)); + assertEquals(completed, replay.get(5, TimeUnit.SECONDS)); + assertEquals(2, calls.get(), "The completed retry body is skipped on replay"); + } finally { + gate.releaseAll(); + caller.shutdownNow(); + gate.shutdownNow(); + assertTrue(caller.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(gate.awaitTermination(5, TimeUnit.SECONDS)); + } + } + } + + private static void awaitCheckpointHandoff(Thread checkpoint) { + var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(3); + while (System.nanoTime() < deadline) { + for (var entry : Thread.getAllStackTraces().entrySet()) { + var continuation = entry.getKey(); + var stack = entry.getValue(); + if (continuation.getState() != Thread.State.WAITING + || !continuation.getName().startsWith("durable-sdk-internal-") + || Arrays.stream(stack) + .noneMatch(frame -> frame.getClassName().endsWith("StepOperation")) + || Arrays.stream(stack) + .noneMatch(frame -> frame.getClassName().equals(CompletableFuture.class.getName()) + && frame.getMethodName().equals("join"))) continue; + assertNotSame(checkpoint, continuation, "Worker publication must not block the checkpoint callback"); + assertTrue(Arrays.stream(stack) + .noneMatch(frame -> frame.getClassName().endsWith("ApiRequestDelayedBatcher"))); + if (Arrays.stream(checkpoint.getStackTrace()) + .anyMatch(frame -> frame.getClassName().endsWith("CheckpointManager") + && frame.getMethodName().equals("checkpointBatch"))) continue; + var bean = ManagementFactory.getThreadMXBean(); + if (bean.isObjectMonitorUsageSupported()) { + var info = bean.getThreadInfo(new long[] {continuation.getId()}, true, true)[0]; + assertEquals(0, info.getLockedMonitors().length); + System.out.println("READY_HANDOFF continuation=" + continuation.getName() + " checkpoint=" + + checkpoint.getName() + " continuationMonitors=[] checkpointBatchReturned=true"); + } + return; + } + LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1)); + } + fail("The independent READY continuation did not reach the controlled old-worker handoff window"); + } + + private static final class DelayedReadyClient extends LocalMemoryExecutionClient { + private final HandoffGate gate; + private final AtomicBoolean firstPoll = new AtomicBoolean(true); + + private DelayedReadyClient(HandoffGate gate) { + this.gate = gate; + } + + @Override + public CheckpointDurableExecutionResponse checkpoint(String arn, String token, List updates) { + if (updates.isEmpty() && firstPoll.compareAndSet(true, false)) { + gate.pollThread.set(Thread.currentThread()); + gate.pollEntered.countDown(); + await(gate.releaseReadyResponse); + var response = super.checkpoint(arn, token, updates); + gate.readyResponseReturned.countDown(); + return response; + } + return super.checkpoint(arn, token, updates); + } + } + + private static final class HandoffGate extends AbstractExecutorService { + private final ExecutorService delegate = Executors.newCachedThreadPool(); + private final AtomicInteger submitted = new AtomicInteger(); + private final AtomicBoolean exitHeld = new AtomicBoolean(); + private final AtomicReference firstWorker = new AtomicReference<>(); + private final AtomicReference pollThread = new AtomicReference<>(); + private final CountDownLatch pollEntered = new CountDownLatch(1); + private final CountDownLatch releaseReadyResponse = new CountDownLatch(1); + private final CountDownLatch readyResponseReturned = new CountDownLatch(1); + private final CountDownLatch workerExitEntered = new CountDownLatch(1); + private final CountDownLatch releaseWorkerExit = new CountDownLatch(1); + private final CountDownLatch executorReturnEntered = new CountDownLatch(1); + private final CountDownLatch releaseExecutorReturn = new CountDownLatch(1); + private final CountDownLatch nextCheckEntered = new CountDownLatch(1); + private final CountDownLatch releaseNextCheck = new CountDownLatch(1); + private final boolean holdWorkerExit; + private final boolean holdExecutorReturn; + + private HandoffGate(boolean holdWorkerExit, boolean holdExecutorReturn) { + this.holdWorkerExit = holdWorkerExit; + this.holdExecutorReturn = holdExecutorReturn; + } + + @Override + public void execute(Runnable task) { + var number = submitted.incrementAndGet(); + var finished = new CountDownLatch(1); + delegate.execute(() -> { + try { + task.run(); + } finally { + finished.countDown(); + } + }); + // Completing work before execute() returns is permitted by Executor's contract. + if (holdExecutorReturn && number == 2) { + await(finished); + executorReturnEntered.countDown(); + await(releaseExecutorReturn); + } + } + + private void releaseAll() { + releaseReadyResponse.countDown(); + releaseWorkerExit.countDown(); + releaseExecutorReturn.countDown(); + releaseNextCheck.countDown(); + } + + @Override + public void shutdown() { + delegate.shutdown(); + } + + @Override + public List shutdownNow() { + releaseAll(); + return delegate.shutdownNow(); + } + + @Override + public boolean isShutdown() { + return delegate.isShutdown(); + } + + @Override + public boolean isTerminated() { + return delegate.isTerminated(); + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { + return delegate.awaitTermination(timeout, unit); + } + } + + private static final class MdcClearGate implements AutoCloseable { + private final MDCAdapter previous = MDC.getMDCAdapter(); + private final Method setter = MDC.class.getDeclaredMethod("setMDCAdapter", MDCAdapter.class); + + private MdcClearGate(HandoffGate gate) throws Exception { + setter.setAccessible(true); + setter.invoke(null, new MDCAdapter() { + public void clear() { + if (gate.holdWorkerExit + && Thread.currentThread() == gate.firstWorker.get() + && gate.exitHeld.compareAndSet(false, true)) { + gate.workerExitEntered.countDown(); + await(gate.releaseWorkerExit); + } + previous.clear(); + } + + public void put(String key, String value) { + previous.put(key, value); + } + + public String get(String key) { + return previous.get(key); + } + + public void remove(String key) { + previous.remove(key); + } + + public Map getCopyOfContextMap() { + return previous.getCopyOfContextMap(); + } + + public void setContextMap(Map values) { + previous.setContextMap(values); + } + + public void pushByKey(String key, String value) { + previous.pushByKey(key, value); + } + + public String popByKey(String key) { + return previous.popByKey(key); + } + + public Deque getCopyOfDequeByKey(String key) { + return previous.getCopyOfDequeByKey(key); + } + + public void clearDequeByKey(String key) { + previous.clearDequeByKey(key); + } + }); + } + + public void close() throws Exception { + setter.invoke(null, previous); + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java index 467a87b94..19fafbdc2 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java @@ -70,7 +70,7 @@ protected void replay(Operation existing) { case STARTED -> { if (isAtMostOnce()) { // AT_MOST_ONCE: treat as interrupted, go through retry logic - handleStepFailure(new StepInterruptedException(existing), attempt); + handleStepFailure(new StepInterruptedException(existing), attempt, null); } else { // AT_LEAST_ONCE: re-execute the step executeStepLogic(attempt); @@ -79,7 +79,7 @@ protected void replay(Operation existing) { // Step is pending retry - Start polling for PENDING -> READY transition case PENDING -> { if (existing.stepDetails() != null && existing.stepDetails().nextAttemptTimestamp() != null) { - pollReadyAndExecuteStepLogic(existing.stepDetails().nextAttemptTimestamp(), attempt); + pollReadyAndExecuteStepLogic(existing.stepDetails().nextAttemptTimestamp(), attempt, null); } else { throw terminateExecutionWithIllegalDurableOperationException( "Unexpected PENDING step without nextAttemptTimestamp: " + getOperationId()); @@ -93,16 +93,24 @@ protected void replay(Operation existing) { } } - private void pollReadyAndExecuteStepLogic(Instant nextAttemptInstant, int attempt) { + private void pollReadyAndExecuteStepLogic( + Instant nextAttemptInstant, int attempt, CompletableFuture> previousWorker) { pollForOperationUpdates(nextAttemptInstant) .thenCompose(op -> op.status() == OperationStatus.READY ? CompletableFuture.completedFuture(op) : pollForOperationUpdates(nextAttemptInstant)) - .thenRun(() -> executeStepLogic(attempt)); + .thenCompose(ignored -> executionManager.runCheckpointContinuation(this, () -> { + // Leave the serialized checkpoint callback before dispatching a configured executor. The lease + // spans retirement of the old attempt and registration of its replacement. + if (previousWorker != null) previousWorker.join().join(); + if (!isOperationCompleted()) executeStepLogic(attempt); + })); } private void executeStepLogic(int attempt) { + var publishedWorker = new CompletableFuture>(); Runnable userHandler = () -> { + if (isOperationCompleted()) return; // use a try-with-resources to // - add thread id/type to thread local when the step starts // - clear logger properties when the step finishes @@ -119,13 +127,14 @@ private void executeStepLogic(int attempt) { handleStepSucceeded(result); } catch (Throwable e) { - handleStepFailure(e, attempt); + handleStepFailure(e, attempt, publishedWorker); } } }; // Execute user provided step code in user-configured executor runUserHandler(userHandler, ThreadType.STEP); + publishedWorker.complete(getRunningUserHandler()); } private void checkpointStarted() { @@ -156,7 +165,8 @@ private void handleStepSucceeded(T result) { sendOperationUpdate(successUpdate); } - private void handleStepFailure(Throwable exception, int attempt) { + private void handleStepFailure( + Throwable exception, int attempt, CompletableFuture> publishedWorker) { exception = ExceptionHelper.unwrapCompletableFuture(exception); if (exception instanceof SuspendExecutionException suspendExecutionException) { throw suspendExecutionException; @@ -189,7 +199,7 @@ private void handleStepFailure(Throwable exception, int attempt) { sendOperationUpdate(retryUpdate); // Poll for READY status and then execute the step again - pollReadyAndExecuteStepLogic(Instant.now().plusSeconds(retryDelayInSeconds), attempt + 1); + pollReadyAndExecuteStepLogic(Instant.now().plusSeconds(retryDelayInSeconds), attempt + 1, publishedWorker); } else { // Send FAIL - retries exhausted var failUpdate = From 0e55c1a11999121bdc5bae521f92975724822b43 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 9 Oct 2026 00:57:17 +0000 Subject: [PATCH 46/48] fix: retain fatal cleanup and preparation error diagnostics --- otel-plugin/README.md | 8 +- .../otel/MdcWorkerReuseBoundaryTest.java | 49 +++++++- .../PreparationFailureDiagnosticsTest.java | 111 ++++++++++++++++++ .../durable/execution/DurableExecutor.java | 14 ++- 4 files changed, 171 insertions(+), 11 deletions(-) create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PreparationFailureDiagnosticsTest.java diff --git a/otel-plugin/README.md b/otel-plugin/README.md index 503b7a8a5..9596ebf53 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -50,9 +50,11 @@ unchanged and is not repeated. An earlier JVM-fatal End/preparation failure rema failures are suppressed, with an identity guard when restoration throws the same fatal object. A restoration fatal takes precedence over an earlier non-JVM-fatal delivery/End failure, which remains suppressed. Non-JVM-fatal restoration failures, including `AssertionError` and `LinkageError`, retain the selected caller outcome. After an ordinary outer restoration failure, the SDK makes one guarded `MDC.clear()` -attempt before rethrowing the original worker failure. A successful clear prevents inheritable invocation state -from reaching a replacement pool thread. A failing clear is retained as a suppressed diagnostic, or its JVM fatal -uses the same caller-settlement-before-worker-throw rule. Clearing can discard ambient MDC that could not be +attempt before rethrowing the original worker failure. This also applies when an earlier End/preparation JVM fatal +is already primary: restoration and clear failures remain suppressed under that first fatal. A successful clear +prevents inheritable invocation state from reaching a replacement pool thread. With no earlier JVM fatal, a failing +clear is retained as a suppressed diagnostic, or its JVM fatal uses the same caller-settlement-before-worker-throw +rule. Clearing can discard ambient MDC that could not be restored. If the adapter's clear also fails, clean replacement state is not guaranteed; the SDK does not promise quarantine for an arbitrary caller-owned executor. Initialization and handler/body policies are unchanged. Before invocation startup, a JVM-fatal error from MDC capture (direct or inside a standard transport wrapper) diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcWorkerReuseBoundaryTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcWorkerReuseBoundaryTest.java index d11cb44f5..60d7224c6 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcWorkerReuseBoundaryTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcWorkerReuseBoundaryTest.java @@ -43,8 +43,36 @@ void failedFallbackRetainsFailurePolicyAndFatalIdentity(boolean suspend, String exercise("end", suspend, kind); } + static Stream fatalEndCases() { + return Stream.of(false, true) + .flatMap(suspend -> Stream.of(false, true) + .flatMap(death -> Stream.of( + "ok", + "ordinary", + "assertion", + "same", + "vm", + "death", + "wrapped-vm", + "wrapped-death") + .map(kind -> Arguments.of(suspend, death, kind)))); + } + + @ParameterizedTest + @MethodSource("fatalEndCases") @SuppressWarnings("removal") + void fatalEndStillClearsOrdinaryRestoreFailureBeforeWorkerReplacement(boolean suspend, boolean death, String kind) + throws Exception { + exercise("end", suspend, kind, death ? new ThreadDeath() : new InternalError("original End fatal")); + } + private static void exercise(String markerPhase, boolean suspend, String kind) throws Exception { + exercise(markerPhase, suspend, kind, null); + } + + @SuppressWarnings("removal") + private static void exercise(String markerPhase, boolean suspend, String kind, Error selectedFatal) + throws Exception { var original = MDC.getMDCAdapter(); var basic = new BasicMDCAdapter(); var ends = new AtomicInteger(); @@ -52,7 +80,7 @@ private static void exercise(String markerPhase, boolean suspend, String kind) t var failed = new AtomicBoolean(); var clearTried = new AtomicBoolean(); var restoreFailure = new IllegalStateException("restore failed"); - Error expectedFatal = + Error clearFatal = switch (kind) { case "vm", "wrapped-vm" -> new InternalError("fallback fatal"); case "death", "wrapped-death" -> new ThreadDeath(); @@ -63,10 +91,10 @@ private static void exercise(String markerPhase, boolean suspend, String kind) t case "ordinary" -> new IllegalArgumentException("clear failed"); case "assertion" -> new AssertionError("clear failed"); case "same" -> restoreFailure; - case "wrapped-vm", "wrapped-death" -> - new CompletionException(new ExecutionException(expectedFatal)); - default -> expectedFatal; + case "wrapped-vm", "wrapped-death" -> new CompletionException(new ExecutionException(clearFatal)); + default -> clearFatal; }; + Error expectedFatal = selectedFatal != null ? selectedFatal : clearFatal; var escaped = new CountDownLatch(1); var failedWorker = new AtomicReference(); var ownerFailure = new AtomicReference(); @@ -111,6 +139,7 @@ public void onInvocationEnd(InvocationEndInfo info) { if (markerPhase.equals("end")) MDC.put("invocation-marker", "previous-invocation"); endStatus.set(info.invocationStatus()); ends.incrementAndGet(); + if (selectedFatal != null) throw selectedFatal; } }; var runner = LocalDurableTestRunner.create( @@ -130,10 +159,17 @@ public void onInvocationEnd(InvocationEndInfo info) { runner.run("input").getStatus()); assertTrue(escaped.await(3, TimeUnit.SECONDS)); assertSame(expectedFatal != null ? expectedFatal : restoreFailure, ownerFailure.get()); - assertTrue(clearTried.get()); + assertEquals(1, ends.get()); assertEquals(suspend ? InvocationStatus.PENDING : InvocationStatus.SUCCEEDED, endStatus.get()); - if (expectedFatal != null) assertEquals(List.of(restoreFailure), List.of(expectedFatal.getSuppressed())); + if (selectedFatal != null) { + assertEquals( + clearFailure != null && clearFailure != restoreFailure + ? List.of(restoreFailure, clearFailure) + : List.of(restoreFailure), + List.of(selectedFatal.getSuppressed())); + } else if (expectedFatal != null) + assertEquals(List.of(restoreFailure), List.of(expectedFatal.getSuppressed())); else assertEquals( clearFailure != null && clearFailure != restoreFailure ? List.of(clearFailure) : List.of(), @@ -144,6 +180,7 @@ public void onInvocationEnd(InvocationEndInfo info) { System.out.println("MDC_REUSE_PUBLIC phase=" + markerPhase + " kind=" + kind + " replaced=true oldMarker=" + dirtyAtFailure.get() + " inheritedMarker=" + inherited); if (clearFailure == null) assertNull(inherited, "Successful fallback clears inheritable invocation state"); + assertTrue(clearTried.get()); // An adapter whose clear also fails cannot promise clean replacement state; its failure is exposed. } finally { workers.shutdownNow(); diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PreparationFailureDiagnosticsTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PreparationFailureDiagnosticsTest.java new file mode 100644 index 000000000..66528416d --- /dev/null +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PreparationFailureDiagnosticsTest.java @@ -0,0 +1,111 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.*; + +import java.time.Duration; +import java.time.Instant; +import java.util.List; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import software.amazon.awssdk.services.lambda.model.*; +import software.amazon.lambda.durable.execution.DurableExecutor; +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.serde.SerDes; +import software.amazon.lambda.durable.testing.local.LocalMemoryExecutionClient; + +class PreparationFailureDiagnosticsTest { + @ParameterizedTest + @ValueSource(strings = {"serialize", "checkpoint"}) + void preparationCauseDeliveredToCallerRetainsEndError(String mode) throws Exception { + var preparation = new IllegalStateException("output preparation failed"); + var endFailure = new AssertionError("End error"); + var end = new AtomicReference(); + var endCalls = new AtomicInteger(); + var checkpoints = new AtomicInteger(); + var workers = Executors.newCachedThreadPool(); + var caller = Executors.newSingleThreadExecutor(); + var client = new LocalMemoryExecutionClient() { + @Override + public CheckpointDurableExecutionResponse checkpoint( + String arn, String token, List updates) { + if (mode.equals("checkpoint") + && updates.stream().anyMatch(op -> op.type() == OperationType.EXECUTION)) { + checkpoints.incrementAndGet(); + throw preparation; + } + return super.checkpoint(arn, token, updates); + } + }; + var serDes = new SerDes() { + private final JacksonSerDes delegate = new JacksonSerDes(); + + public String serialize(Object value) { + if (mode.equals("serialize") && "result".equals(value)) throw preparation; + return delegate.serialize(value); + } + + public T deserialize(String value, TypeToken type) { + return delegate.deserialize(value, type); + } + }; + var config = DurableConfig.builder() + .withDurableExecutionClient(client) + .withExecutorService(workers) + .withSerDes(serDes) + .withCheckpointDelay(Duration.ZERO) + .withPlugins(new DurableExecutionPlugin() { + public void onInvocationEnd(InvocationEndInfo info) { + endCalls.incrementAndGet(); + end.set(info); + throw endFailure; + } + }) + .build(); + try { + var result = caller.submit(() -> DurableExecutor.execute( + input(), + null, + TypeToken.get(String.class), + (value, context) -> mode.equals("checkpoint") ? "x".repeat(6 * 1024 * 1024) : "result", + config)); + var observed = assertThrows(ExecutionException.class, () -> result.get(5, TimeUnit.SECONDS)) + .getCause(); + System.out.println("PREPARATION_DIAGNOSTICS mode=" + mode + " original=" + (observed == preparation) + + " suppressed=" + List.of(observed.getSuppressed()) + " End=" + endCalls.get()); + assertSame(preparation, observed); + assertEquals(List.of(endFailure), List.of(observed.getSuppressed())); + assertEquals(1, endCalls.get()); + assertEquals(InvocationStatus.RETRYING, end.get().invocationStatus()); + assertSame(preparation, end.get().executionError()); + assertEquals(mode.equals("checkpoint") ? 1 : 0, checkpoints.get()); + } finally { + caller.shutdownNow(); + workers.shutdownNow(); + assertTrue(caller.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private static DurableExecutionInput input() { + return new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/diagnostics/execution", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("execution") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + } +} 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 094ac8a7a..8b97b9e27 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 @@ -213,7 +213,7 @@ private DurableExecutionOutput failDelivery(Throwable deliveryFailure) { try { fireOnInvocationEnd(InvocationStatus.RETRYING, cause, null); } catch (Error endFailure) { - deliveryFailure = combinePreparationAndEndFailures(deliveryFailure, endFailure); + deliveryFailure = combinePreparationAndEndFailures(cause, endFailure); } ExceptionHelper.sneakyThrow(deliveryFailure); return null; @@ -366,7 +366,17 @@ private static void closeMdcAfterFatal(SafeCloseable restore, Error fatal) { try { restore.close(); } catch (Throwable restorationFailure) { - if (restorationFailure != fatal) fatal.addSuppressed(restorationFailure); + if (restorationFailure == fatal) return; + fatal.addSuppressed(restorationFailure); + try { + rethrowLifecycleFatal(restorationFailure); + // The first fatal still owns caller/worker propagation. Clear only after an ordinary restore failure + // so a ThreadPoolExecutor replacement does not inherit an End hook's invocation-local map. + MDC.clear(); + } catch (Throwable cleanupFailure) { + if (cleanupFailure != fatal && cleanupFailure != restorationFailure) + fatal.addSuppressed(cleanupFailure); + } } } From d1483798017a38d827ff59385684b1d52d591c0f Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 9 Oct 2026 01:32:55 +0000 Subject: [PATCH 47/48] fix: select continuation failures atomically with close --- .../durable/execution/ExecutionManager.java | 12 +- .../ContinuationFailureCloseRaceTest.java | 133 ++++++++++++++++++ 2 files changed, 144 insertions(+), 1 deletion(-) create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/execution/ContinuationFailureCloseRaceTest.java 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 427e1e979..1aa148c2a 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 @@ -83,6 +83,8 @@ public class ExecutionManager implements SafeCloseable { // Guarded by activeThreads; an observation future can be cancelled independently of its actual task. private final Map checkpointContinuations = new HashMap<>(); private boolean closing; + // Guarded by activeThreads; selected against normal close before publishing callbacks outside the monitor. + private UnrecoverableDurableExecutionException selectedContinuationFailure; /** * Per-wait state used to coordinate the caller thread with the future completion callback. @@ -630,7 +632,15 @@ private void signalContinuationFailure(Throwable failure) { .build(), true, failure); - if (executionExceptionFuture.completeExceptionally(control)) stopAllOperations(control); + UnrecoverableDurableExecutionException selected; + synchronized (activeThreads) { + if (closing && !(failure instanceof VirtualMachineError || failure instanceof ThreadDeath)) return; + if (selectedContinuationFailure == null) selectedContinuationFailure = control; + selected = selectedContinuationFailure; + } + // A later admitted publisher may drain the winner's callbacks. Keep that CompletableFuture behavior, + // but publish/wake using the selected cause, outside the monitor needed by worker cleanup. + if (executionExceptionFuture.completeExceptionally(selected)) stopAllOperations(selected); } private Object registerCheckpointContinuation(BaseDurableOperation owner) { diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/ContinuationFailureCloseRaceTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/ContinuationFailureCloseRaceTest.java new file mode 100644 index 000000000..2f80d921a --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/ContinuationFailureCloseRaceTest.java @@ -0,0 +1,133 @@ +// 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.util.concurrent.*; +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.model.DurableExecutionInput; +import software.amazon.lambda.durable.operation.BaseDurableOperation; + +class ContinuationFailureCloseRaceTest { + @ParameterizedTest + @CsvSource({"body,true", "reject,true", "body,false", "reject,false"}) + void failureSelectionIsOrderedWithCloseAndDoesNotRunCallbacksUnderCoordinationLock(String mode, boolean closeFirst) + throws Exception { + var input = new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/close-race/execution", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("execution") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .build()) + .build()); + var manager = new ExecutionManager( + input, + DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .build(), + null); + manager.registerActiveThread("root"); + var callbacks = new CountDownLatch(1); + var callback = manager.runUntilCompleteOrSuspend(new CompletableFuture()) + .handle((value, error) -> { + try { + CompletableFuture.runAsync(() -> manager.registerActiveThread("callback-lock-proof")) + .get(3, TimeUnit.SECONDS); + } catch (Exception failure) { + throw new AssertionError("Failure callbacks must not hold the activeThreads monitor", failure); + } + callbacks.countDown(); + return error; + }); + var owner = mock(BaseDurableOperation.class); + var ownerCompletion = new CompletableFuture(); + when(owner.getCompletionFuture()).thenReturn(ownerCompletion); + var unrelated = mock(BaseDurableOperation.class); + var unrelatedCompletion = new CompletableFuture(); + when(unrelated.getOperationId()).thenReturn("unrelated"); + when(unrelated.getCompletionFuture()).thenReturn(unrelatedCompletion); + manager.registerOperation(unrelated); + var buildingFailure = new CountDownLatch(1); + var releaseFailure = new CountDownLatch(1); + RuntimeException failure = mode.equals("reject") + ? new RejectedExecutionException("coordinator rejected") + : new IllegalArgumentException("continuation failed"); + var publisher = CompletableFuture.runAsync(() -> { + // A scheduling gate inside ordinary control construction, after the old caller's isClosing precheck. + // It does not replace manager state, futures, close, or stopAllOperations. + try (var model = mockStatic(ErrorObject.class, CALLS_REAL_METHODS)) { + model.when(ErrorObject::builder).thenAnswer(call -> { + buildingFailure.countDown(); + assertTrue(releaseFailure.await(3, TimeUnit.SECONDS)); + return call.callRealMethod(); + }); + if (mode.equals("reject")) { + assertSame( + failure, + assertThrows( + RejectedExecutionException.class, + () -> manager.runCheckpointContinuation( + owner, () -> fail("Rejected work"), task -> { + throw failure; + }))); + } else { + var observed = manager.runCheckpointContinuation( + owner, + () -> { + throw failure; + }, + Runnable::run); + assertSame( + failure, + assertThrows(CompletionException.class, observed::join) + .getCause()); + } + } + }); + CompletableFuture closing = null; + try { + assertTrue(buildingFailure.await(3, TimeUnit.SECONDS)); + if (closeFirst) { + closing = CompletableFuture.runAsync(manager::close); + assertInstanceOf( + SuspendExecutionException.class, + assertThrows(ExecutionException.class, () -> ownerCompletion.get(3, TimeUnit.SECONDS)) + .getCause()); + assertFalse(unrelatedCompletion.isDone()); + var pendingClose = closing; + assertThrows(TimeoutException.class, () -> pendingClose.get(100, TimeUnit.MILLISECONDS)); + } + releaseFailure.countDown(); + publisher.get(3, TimeUnit.SECONDS); + if (closing == null) closing = CompletableFuture.runAsync(manager::close); + closing.get(3, TimeUnit.SECONDS); + System.out.println("CONTINUATION_CLOSE_RACE mode=" + mode + " closeFirst=" + closeFirst + + " unrelatedStopped=" + unrelatedCompletion.isDone() + + " managerFailure=" + manager.isExecutionCompletedExceptionally()); + assertEquals( + !closeFirst, + unrelatedCompletion.isDone(), + "Only failure selected before close may stop unrelated operations"); + assertEquals(!closeFirst, manager.isExecutionCompletedExceptionally()); + if (!closeFirst) { + assertNotNull(callback.get(3, TimeUnit.SECONDS)); + assertEquals(0, callbacks.getCount()); + } else assertEquals(1, callbacks.getCount()); + assertDoesNotThrow(() -> manager.deregisterActiveThread("root")); + } finally { + releaseFailure.countDown(); + publisher.get(3, TimeUnit.SECONDS); + if (closing == null) closing = CompletableFuture.runAsync(manager::close); + closing.get(3, TimeUnit.SECONDS); + } + } +} From be44fd7d1447323978b318872275a4db5afef6f1 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 9 Oct 2026 02:11:57 +0000 Subject: [PATCH 48/48] fix: prune completed polls before checkpoint batching --- .../durable/execution/CheckpointManager.java | 5 + .../execution/CompletedPollBatchTest.java | 99 +++++++++++++++++++ 2 files changed, 104 insertions(+) create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/execution/CompletedPollBatchTest.java 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 34f44134c..0430d6ab1 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 @@ -198,6 +198,11 @@ List fetchAllPages(CheckpointUpdatedExecutionState checkpointUpdatedE private void checkpointBatch(List updates) { synchronized (pollingFutures) { + // A READY recheck can complete a poll before its scheduled batch. Keep only consumers still waiting. + pollingFutures.entrySet().removeIf(entry -> { + entry.getValue().removeIf(CompletableFuture::isDone); + return entry.getValue().isEmpty(); + }); // filter the null values from pollers var request = updates.stream().filter(Objects::nonNull).toList(); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/CompletedPollBatchTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/CompletedPollBatchTest.java new file mode 100644 index 000000000..0646aaa46 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/CompletedPollBatchTest.java @@ -0,0 +1,99 @@ +// 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.*; +import static org.mockito.Mockito.*; + +import java.time.Duration; +import java.util.List; +import java.util.concurrent.*; +import java.util.concurrent.atomic.AtomicInteger; +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.client.DurableExecutionClient; + +class CompletedPollBatchTest { + @ParameterizedTest + @ValueSource( + strings = {"ready", "terminal", "exceptional", "cancelled", "checkpoint", "same-id", "other-id", "pending"}) + void completedPollsDoNotCauseEmptyRpcButLivePollsAndUpdatesStillFlush(String mode) throws Exception { + var client = mock(DurableExecutionClient.class); + var starts = new AtomicInteger(); + var finishes = new AtomicInteger(); + var callbacks = new AtomicInteger(); + var config = DurableConfig.builder() + .withDurableExecutionClient(client) + .withCheckpointDelay(Duration.ofHours(1)) + .build(); + var manager = new CheckpointManager( + config, + "arn:test", + "token", + values -> callbacks.incrementAndGet(), + () -> { + starts.incrementAndGet(); + return true; + }, + finishes::incrementAndGet); + var ready = Operation.builder() + .id("watched") + .type(OperationType.STEP) + .status(OperationStatus.READY) + .build(); + var terminal = ready.toBuilder().status(OperationStatus.SUCCEEDED).build(); + var liveId = mode.equals("other-id") ? "other" : "watched"; + var response = ready.toBuilder().id(liveId).build(); + when(client.checkpoint(anyString(), anyString(), anyList())) + .thenReturn(CheckpointDurableExecutionResponse.builder() + .checkpointToken("next") + .newExecutionState(CheckpointUpdatedExecutionState.builder() + .operations(response) + .build()) + .build()); + var update = OperationUpdate.builder() + .id("write") + .type(OperationType.STEP) + .action(OperationAction.START) + .build(); + try { + var poll = manager.pollForUpdate("watched", attempt -> Duration.ofHours(1)); + CompletableFuture live = null; + if (mode.equals("same-id") || mode.equals("other-id")) + live = manager.pollForUpdate(liveId, attempt -> Duration.ofHours(1)); + switch (mode) { + case "exceptional" -> poll.completeExceptionally(new IllegalStateException("already failed")); + case "cancelled" -> poll.cancel(false); + case "terminal" -> poll.complete(terminal); + case "pending" -> {} + default -> poll.complete(ready); + } + var checkpoint = mode.equals("checkpoint") ? manager.checkpoint(update) : null; + // Force the real delayed batch now, without closing CheckpointManager (which would clear the poll map). + // This exercises the same pending batch as its timer, without a timing-based no-extra-call assertion. + var field = CheckpointManager.class.getDeclaredField("checkpointApiRequestDelayedBatcher"); + field.setAccessible(true); + var dispatcher = (ApiRequestDelayedBatcher) field.get(manager); + dispatcher.shutdown(); + boolean required = live != null || checkpoint != null || mode.equals("pending"); + System.out.println("COMPLETED_POLL mode=" + mode + " APIleases=" + starts.get() + " finished=" + + finishes.get() + " callbacks=" + callbacks.get()); + verify(client, times(required ? 1 : 0)) + .checkpoint(eq("arn:test"), eq("token"), eq(checkpoint != null ? List.of(update) : List.of())); + assertEquals(required ? 1 : 0, starts.get()); + assertEquals(starts.get(), finishes.get()); + assertEquals(required ? 1 : 0, callbacks.get()); + if (live != null) assertEquals(response, live.get(3, TimeUnit.SECONDS)); + if (checkpoint != null) checkpoint.get(3, TimeUnit.SECONDS); + if (mode.equals("pending")) assertEquals(response, poll.get(3, TimeUnit.SECONDS)); + else if (mode.equals("cancelled")) assertTrue(poll.isCancelled()); + else if (mode.equals("exceptional")) assertTrue(poll.isCompletedExceptionally()); + else assertSame(mode.equals("terminal") ? terminal : ready, poll.getNow(null)); + } finally { + manager.shutdown(); + } + } +}