diff --git a/docs/advanced/configuration.md b/docs/advanced/configuration.md index 00f37c8372..e1b9f6389a 100644 --- a/docs/advanced/configuration.md +++ b/docs/advanced/configuration.md @@ -34,11 +34,22 @@ public class OrderProcessor extends DurableHandler { | `withLambdaClientBuilder()` | Custom AWS Lambda client | Auto-configured Lambda client | | `withSerDes()` | Serializer for step results | Jackson with default settings | | `withExecutorService()` | Thread pool for user-defined operations | Cached daemon thread pool | +| `withInvocationExecutorFactory()` | Factory for isolated per-invocation user executors | Disabled | | `withLoggerConfig()` | Logger behavior configuration | Suppress logs during replay | | `withPollingStrategy()` | Backend polling strategy | Exponential backoff: 1s base, 2x rate, FULL jitter, 10s max | | `withCheckpointDelay()` | How often the SDK checkpoints updates | `Duration.ofSeconds(0)` (as soon as possible) | -The `withExecutorService()` option configures the thread pool used for running user-defined operations. Internal SDK coordination (checkpoint batching, polling) runs on an SDK-managed thread pool. +The `withExecutorService()` option configures a caller-owned executor shared by every invocation that uses the configuration. The SDK never shuts it down. To isolate concurrent Lambda Managed Instances invocations, configure a factory instead: + +```java +DurableConfig.builder() + .withInvocationExecutorFactory(() -> Executors.newFixedThreadPool(4)) + .build(); +``` + +The factory is called once per invocation and may be called concurrently on Lambda Managed Instances. It must return a fresh executor. The invocation caller waits for actual admitted work and End before shutting that executor down, including after cancellation. Code that ignores interruption can delay this shutdown beyond the policy budget; executor termination is then observed with any cleanup time that remains. Returning an executor already leased to an active invocation is rejected. The factory and `withExecutorService()` are mutually exclusive. Both options run all user code, including plugin-enabled root handlers, steps, and child contexts. Internal SDK coordination such as checkpoint batching and polling remains on the SDK-managed executor. `toBuilder()` and the local test runner preserve the factory, including across replay/resume. + +A bounded executor supplied through `withExecutorService()` remains shared. On Lambda Managed Instances, concurrent roots can occupy all of its workers while waiting for queued steps. Use the invocation factory when capacity must be isolated between native invocations. Isolation does not define the capacity required by arbitrarily nested synchronous child contexts; size or design those separately. ### Dynamic plugin loading diff --git a/docs/design.md b/docs/design.md index e19cbd2f0f..1fc7f8aa3e 100644 --- a/docs/design.md +++ b/docs/design.md @@ -145,6 +145,7 @@ public class MyHandler extends DurableHandler { | `lambdaClientBuilder` | Auto-created `LambdaClient` for current region, primed for performance (see [`DurableConfig.java`](../sdk/src/main/java/software/amazon/lambda/durable/DurableConfig.java)) | | `serDes` | `JacksonSerDes` | | `executorService` | `Executors.newCachedThreadPool()` (for user-defined operations only) | +| `invocationExecutorFactory` | Disabled; when configured, creates an SDK-owned executor for each invocation | | `loggerConfig` | `LoggerConfig.defaults()` (suppress replay logs) | | `pollingStrategy` | Exponential backoff: 1s base, 2x rate, FULL jitter, 10s max | | `checkpointDelay` | `Duration.ofSeconds(0)` (checkpoint as soon as possible) | @@ -163,8 +164,15 @@ The SDK uses two separate thread pools with distinct responsibilities: **User Executor (`DurableConfig.executorService`):** - Runs the root handler, steps, child-context user code, and map/parallel coordination - Configurable via `DurableConfig.builder().withExecutorService()` +- Caller-owned and shared across invocations; the SDK never shuts it down - Default: cached daemon thread pool +**Invocation Executor Factory (`DurableConfig.invocationExecutorFactory`):** +- Optional alternative to the shared user executor +- Creates one executor per invocation for all root, step, child-context user code, and map/parallel coordination +- SDK-owned: the invocation caller waits for actual admitted work and End before requesting executor shutdown, including after cancellation. Non-cooperative work can delay that shutdown beyond the policy budget. Executor termination is observed with any cleanup time that remains. +- Must return a fresh executor and cannot be combined with `withExecutorService()` + **Internal Executor (`InternalExecutor.INSTANCE`):** - Runs SDK coordination tasks: checkpoint batching, polling for wait completion - Dedicated cached thread pool with daemon threads named `durable-sdk-internal-*` diff --git a/lmi-tests/README.md b/lmi-tests/README.md index b17a845376..97f7e50e7b 100644 --- a/lmi-tests/README.md +++ b/lmi-tests/README.md @@ -39,7 +39,7 @@ eventual cleanup cannot turn an earlier lifecycle violation into a pass. SDK's minimum JDK prevents fixture-only changes from introducing newer JDK APIs into the shared Maven reactor. On manual dispatch, a separate serialized cloud job builds the selected commit, deploys all five - fixture functions, runs all 13 cases, then collects evidence. A persistent + fixture functions, runs all 12 gating cases, then collects evidence. A persistent CloudFormation stack owns all five functions and log groups. The stack, functions and bucket remain after both successful and failed test runs; later deployments update the same resources. Each uses 2 GiB / 1 vCPU and a single @@ -74,17 +74,14 @@ eventual cleanup cannot turn an earlier lifecycle violation into a pass. Each real checkpoint/poll call has a request-local call ID paired with its `finally` exit. Missing or mismatched exits prevent cleanup from reusing the fixture, and an exit after wrapper return remains a lifecycle failure. -* A bounded same-JVM root barrier establishes the fixed-pool reproduction. - `fixed2` uses a shared two-thread pool so two roots starve their first steps. - `nested2` uses four threads and a separate cohort-scoped child barrier. Both - child handlers must enter before either schedules nested map/parallel work. - The child barrier has an independent 8-second setup budget; the nested progress - budget starts when it releases, not when the roots passed their earlier barrier. - A broken child barrier fails setup without growing the pool; an escape after - successful child admission still fails the starvation assertion. - Other cases hold steps with private S3 control objects. The driver requires - distinct request IDs active in the same JVM. Placement has its own deadline - and failure category. Environment replacement is never worker recovery. +* A bounded same-JVM root barrier verifies that `fixed2` gives each invocation + its own two-thread executor. `nested2` is retained in the persistent stack for + baseline diagnostics but does not produce a gating LMI case. It uses four + shared threads and a separate cohort-scoped child barrier before scheduling + nested map/parallel work. Other cases hold steps with private S3 control + objects. The driver requires distinct request IDs active in the same JVM. + Placement has its own deadline and failure category. Environment replacement + is never worker recovery. * Timeout victims start before healthy peers, leaving the peers time to hold the other runtime slots during recovery. Diagnostics capture the real context deadline; no fake clock, context, checkpoint backend, or time-skipping runner @@ -179,8 +176,9 @@ python3 lmi-tests/cloud_suite.py collect ``` The deploy command creates or updates and retains `default1`, `default2`, `default8`, -`fixed2`, and `nested2` together. The test command requires all five and runs the -full suite. CI publishes one combined artifact, named `lmi-e2e-RUN-ATTEMPT`, with +`fixed2`, and `nested2` together. The test command requires all five and runs 12 +gating cases; `nested2` is retained but excluded. CI publishes one combined +artifact, named `lmi-e2e-RUN-ATTEMPT`, with the single `template.json` and per-function configuration snapshots. The CloudFormation execution role needs permission to manage the function scaling configuration; the driver calls `lambda:GetFunctionScalingConfig` to verify it @@ -260,9 +258,10 @@ a normally ordered retry cannot replace an admitted participant. After stopping durable executions, cleanup allows any pre-stop stubborn retry its configured residual hold (capped at 120 seconds), plus the existing 30-second collection margin and 20-second quiet period. This cleanup allowance does not alter -the original invocation deadline + 5s or progress/recovery acceptance limits. No virtual-thread executor variant is deployed -until its executor contract is defined; default cached and shared fixed pools -are covered separately. +the original invocation deadline + 5s or progress/recovery acceptance limits. +No virtual-thread executor variant is deployed until its executor contract is defined. +The cloud suite gates default cached and invocation-local fixed pools; the shared +fixed nesting fixture remains available as a non-gating baseline diagnostic. CloudFormation's native [FunctionScalingConfig](https://docs.aws.amazon.com/AWSCloudFormation/latest/TemplateReference/aws-properties-lambda-function-functionscalingconfig.html) sets the limits in the resource declaration. LMI provisioning and version behavior diff --git a/lmi-tests/cloud_suite.py b/lmi-tests/cloud_suite.py index 95d51e25bc..b1666623d9 100644 --- a/lmi-tests/cloud_suite.py +++ b/lmi-tests/cloud_suite.py @@ -25,7 +25,7 @@ OWNER = "java-sdk-lmi-e2e" DEFAULT_STACK = "java-lmi-e2e" FIXTURES = {"default1": (1, "default"), "default2": (2, "default"), - "default8": (8, "default"), "fixed2": (2, "fixed"), "nested2": (2, "nested")} + "default8": (8, "default"), "fixed2": (2, "invocation-fixed"), "nested2": (2, "nested")} PEER_LOG_OBSERVATION_SECONDS = 20 PEER_PLACEMENT_SECONDS = 60 @@ -201,7 +201,7 @@ def deploy_fixtures(manifest, tags, seconds=1800): wait_stack(manifest["stack"], expected, remaining_budget(deadline, seconds)) stack = existing_stack(manifest["stack"]) outputs = {output["OutputKey"]: output["OutputValue"] for output in stack["Outputs"]} - require(set(outputs) == set(FIXTURES), "The stack must expose all five LMI fixtures") + require(set(outputs) == set(FIXTURES), "The stack must expose all required LMI fixtures") for fixture in FIXTURES: arn = outputs[fixture] record_fixture(manifest, fixture, arn) @@ -893,10 +893,10 @@ def collect_warm_snapshots(cloud, items, seconds=30): def cases_for_fixture(cloud, fixture): if fixture == "default1": return [("baseline-concurrency1", lambda: replay_case(cloud, fixture, "baseline"))] - if fixture in {"fixed2", "nested2"}: - scenario = "fixed" if fixture == "fixed2" else "nested" - return [("fixed-two-roots" if fixture == "fixed2" else "fixed-nested-map-parallel", - lambda: fixed_case(cloud, fixture, scenario))] + if fixture == "fixed2": + return [("isolated-fixed-two-roots", lambda: fixed_case(cloud, fixture, "fixed"))] + if fixture == "nested2": + return [] cases = [(fixture + "-isolation", lambda: overlap_case(cloud, fixture)), (fixture + "-suspend-cleanup-replay", lambda: replay_case(cloud, fixture, "suspend")), (fixture + "-timeout-recovery", lambda: timeout_case(cloud, fixture))] @@ -1022,7 +1022,7 @@ def revalidate_cases(cloud, records): def run_tests(): manifest = json.loads(MANIFEST.read_text()) if set(manifest["functions"]) != set(FIXTURES): - raise PreconditionError("All five functions must be deployed before running the cloud suite") + raise PreconditionError("All required functions must be deployed before running the cloud suite") cloud = Cloud(manifest, ARTIFACTS) suite = ET.Element("testsuite", name="LMI cloud lifecycle") records = [] diff --git a/lmi-tests/cloud_support.py b/lmi-tests/cloud_support.py index e7a18b9aac..0ed28528be 100644 --- a/lmi-tests/cloud_support.py +++ b/lmi-tests/cloud_support.py @@ -213,7 +213,7 @@ def assert_fixed(events, markers, progress_start="BARRIER_PASSED"): require(len({e["requestId"] for e in entered}) == len(markers), "Missing distinct runtime invocations") require(max(e["sequence"] for e in entered) < min(e["sequence"] for e in passed), "Root invocations did not overlap at the barrier") - require(not selected(events, "ESCAPE"), "Shared fixed executor starved its own queued work") + require(not selected(events, "ESCAPE"), "Fixed executor starved its own queued work") # Child admission is one cohort-wide phase. A participant scheduled later must # not get a fresh progress budget after its peer has already left the barrier. cohort_start = min((e["nanos"] for e in selected(events, progress_start)), default=None) @@ -222,7 +222,7 @@ def assert_fixed(events, markers, progress_start="BARRIER_PASSED"): start = selected(events, progress_start, marker) started = cohort_start if progress_start == "CHILD_BARRIER_PASSED" else (start[0]["nanos"] if start else None) require(start and progress and 0 <= (progress[-1]["nanos"] - started) < 8_000_000_000, - "Shared executor did not progress within budget") + "Fixed executor did not progress within budget") def assert_nested_reached(events, markers): diff --git a/lmi-tests/src/main/java/software/amazon/lambda/durable/lmi/LifecycleHandler.java b/lmi-tests/src/main/java/software/amazon/lambda/durable/lmi/LifecycleHandler.java index 4db3041e24..b8ce3845fb 100644 --- a/lmi-tests/src/main/java/software/amazon/lambda/durable/lmi/LifecycleHandler.java +++ b/lmi-tests/src/main/java/software/amazon/lambda/durable/lmi/LifecycleHandler.java @@ -25,6 +25,7 @@ import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicReference; import software.amazon.awssdk.services.lambda.model.OperationType; import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.DurableContext; @@ -75,27 +76,42 @@ public void handleRequest(InputStream source, OutputStream destination, Context var input = WIRE.deserialize( new String(source.readAllBytes(), StandardCharsets.UTF_8), TypeToken.get(DurableExecutionInput.class)); var trace = new InvocationTrace(readUserInput(input), runtime, input.durableExecutionArn()); + var invocationExecutor = new AtomicReference(); trace.wrapperEnter(); var status = "THREW"; try { - var config = DurableConfig.builder() - .withExecutorService(shared.getExecutorService()) - .withDurableExecutionClient(new ObservedClient(shared.getDurableExecutionClient(), trace)) - .build(); + var configBuilder = DurableConfig.builder() + .withDurableExecutionClient(new ObservedClient(shared.getDurableExecutionClient(), trace)); + if ("invocation-fixed".equals(System.getenv("LMI_EXECUTOR"))) { + configBuilder.withInvocationExecutorFactory(() -> { + var executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(2); + invocationExecutor.set(executor); + return executor; + }); + } else { + configBuilder.withExecutorService(shared.getExecutorService()); + } + var config = configBuilder.build(); var response = DurableExecutor.execute( input, runtime, TypeToken.get(FixtureInput.class), - (value, context) -> handle(value, context, trace), + (value, context) -> handle(value, context, trace, executor(invocationExecutor)), config); destination.write(WIRE.serialize(response).getBytes(StandardCharsets.UTF_8)); status = response.status().name(); } finally { - trace.snapshot((ThreadPoolExecutor) shared.getExecutorService()); + trace.snapshot(executor(invocationExecutor)); trace.wrapperExit(status); } } + private ThreadPoolExecutor executor(AtomicReference invocationExecutor) { + return invocationExecutor.get() != null + ? invocationExecutor.get() + : (ThreadPoolExecutor) shared.getExecutorService(); + } + private static FixtureInput readUserInput(DurableExecutionInput input) { var operation = input.initialExecutionState().operations().stream() .filter(op -> op.type() == OperationType.EXECUTION) @@ -104,7 +120,8 @@ private static FixtureInput readUserInput(DurableExecutionInput input) { return JSON.deserialize(operation.executionDetails().inputPayload(), TypeToken.get(FixtureInput.class)); } - String handle(FixtureInput input, DurableContext context, InvocationTrace trace) { + String handle( + FixtureInput input, DurableContext context, InvocationTrace trace, ThreadPoolExecutor invocationExecutor) { trace.rootEnter(); try { if (!context.isReplaying() @@ -113,7 +130,7 @@ String handle(FixtureInput input, DurableContext context, InvocationTrace trace) trace.event("PLACEMENT_MISS"); return "PLACEMENT_MISS"; } - return runScenario(input, context, trace); + return runScenario(input, context, trace, invocationExecutor); } finally { if ("suspend".equals(input.scenario())) { trace.event("CLEANUP_ENTER"); @@ -124,13 +141,14 @@ String handle(FixtureInput input, DurableContext context, InvocationTrace trace) } } - private String runScenario(FixtureInput input, DurableContext context, InvocationTrace trace) { + private String runScenario( + FixtureInput input, DurableContext context, InvocationTrace trace, ThreadPoolExecutor invocationExecutor) { return switch (input.scenario()) { case "baseline", "replay", "suspend" -> replay(input, context, trace); case "hold", "probe" -> context.step("held-step", String.class, step -> hold(input, trace), NO_RETRY); case "timeout", "failure-inflight", "return-inflight" -> inFlight(input, context, trace); case "stubborn" -> stubborn(input, context, trace); - case "fixed", "nested" -> fixed(input, context, trace); + case "fixed", "nested" -> fixed(input, context, trace, invocationExecutor); case "success" -> context.step("success", String.class, step -> body(trace, "success", input.marker())); case "failure" -> context.step("failure", String.class, step -> fail(trace, input.marker()), NO_RETRY); default -> throw new IllegalArgumentException("Unknown fixture scenario"); @@ -226,8 +244,7 @@ private String stubborn(FixtureInput input, DurableContext context, InvocationTr return input.marker(); } - private String fixed(FixtureInput input, DurableContext context, InvocationTrace trace) { - var pool = (ThreadPoolExecutor) shared.getExecutorService(); + private String fixed(FixtureInput input, DurableContext context, InvocationTrace trace, ThreadPoolExecutor pool) { awaitRootBarrier(input, trace); if ("nested".equals(input.scenario())) { var value = context.step("success", String.class, step -> body(trace, "success", input.marker())); diff --git a/lmi-tests/tests/test_evidence.py b/lmi-tests/tests/test_evidence.py index d7a814e01a..cb0973bfbc 100644 --- a/lmi-tests/tests/test_evidence.py +++ b/lmi-tests/tests/test_evidence.py @@ -186,14 +186,15 @@ def test_artifacts_redact_control_credentials_in_nested_payloads(self): value = {"controlUrl": "secret", "InputPayload": '{"controlUrl":"https://bucket/key?X-Amz-Signature=secret"}'} self.assertNotIn("secret", json.dumps(scrub(value))) - def test_one_stack_declares_all_five_functions_and_their_scaling_limits(self): + def test_one_stack_declares_all_functions_and_their_scaling_limits(self): manifest = {"stack": "test", "role": "role", "bucket": "bucket", "provider": "provider", "commit": "sha", "invocationTimeout": 60, "codeSha256": "digest", "codeKey": "code/digest.jar", "runId": "run-1"} spec = template(manifest) functions = [r["Properties"] for r in spec["Resources"].values() if r["Type"] == "AWS::Lambda::Function"] - self.assertEqual(5, len(functions)) + self.assertEqual(len(FIXTURES), len(functions)) self.assertEqual(set(FIXTURES), set(spec["Outputs"])) - self.assertEqual(5, sum(r["Type"] == "AWS::Logs::LogGroup" for r in spec["Resources"].values())) + self.assertEqual( + len(FIXTURES), sum(r["Type"] == "AWS::Logs::LogGroup" for r in spec["Resources"].values())) concurrencies = set() for function in functions: capacity = function["CapacityProviderConfig"]["LambdaManagedInstancesCapacityProviderConfig"] @@ -213,7 +214,7 @@ def test_one_stack_declares_all_five_functions_and_their_scaling_limits(self): self.assertEqual({1, 2, 8}, concurrencies) fixed_executor = spec["Resources"]["fixed2Function"]["Properties"]["Environment"]["Variables"] nested_executor = spec["Resources"]["nested2Function"]["Properties"]["Environment"]["Variables"] - self.assertEqual("fixed", fixed_executor["LMI_EXECUTOR"]) + self.assertEqual("invocation-fixed", fixed_executor["LMI_EXECUTOR"]) self.assertEqual("nested", nested_executor["LMI_EXECUTOR"]) self.assertGreater(PEER_LOG_OBSERVATION_SECONDS, 9) @@ -226,7 +227,10 @@ def test_full_suite_preserves_all_cases_on_their_own_fixture(self): try: for fixture in FIXTURES: cases = cases_for_fixture(cloud, fixture) - self.assertTrue(cases) + if fixture == "nested2": + self.assertEqual([], cases) + else: + self.assertTrue(cases) names.extend(name for name, _ in cases) for _, run in cases: run() @@ -237,15 +241,15 @@ def test_full_suite_preserves_all_cases_on_their_own_fixture(self): finally: for p in patches: p.stop() - self.assertEqual(13, len(names)) - self.assertEqual(13, len(set(names))) + self.assertEqual(12, len(names)) + self.assertEqual(12, len(set(names))) @patch("cloud_suite.Cloud") - def test_cloud_suite_requires_all_five_functions(self, cloud): + def test_cloud_suite_requires_all_functions(self, cloud): with TemporaryDirectory() as directory: path = Path(directory) / "manifest.json" path.write_text(json.dumps({"functions": {"default1": {}}})) - with patch("cloud_suite.MANIFEST", path), self.assertRaisesRegex(PreconditionError, "All five"): + with patch("cloud_suite.MANIFEST", path), self.assertRaisesRegex(PreconditionError, "All required"): run_tests() cloud.assert_not_called() diff --git a/lmi-tests/tests/test_persistence.py b/lmi-tests/tests/test_persistence.py index bafc949e31..50f56aaa7b 100644 --- a/lmi-tests/tests/test_persistence.py +++ b/lmi-tests/tests/test_persistence.py @@ -23,6 +23,16 @@ def stack(self, status="CREATE_COMPLETE"): "Outputs": [{"OutputKey": key, "OutputValue": "function:" + key + ":$LATEST.PUBLISHED"} for key in FIXTURES]} + def test_fixed_fixture_uses_invocation_factory(self): + resources = template(self.manifest())["Resources"] + + self.assertEqual( + "invocation-fixed", + resources["fixed2Function"]["Properties"]["Environment"]["Variables"]["LMI_EXECUTOR"]) + self.assertEqual( + "nested", + resources["nested2Function"]["Properties"]["Environment"]["Variables"]["LMI_EXECUTOR"]) + @patch("cloud_suite.save") @patch("cloud_suite.wait_stack") @patch("cloud_suite.wait_for_idle_functions") @@ -51,7 +61,7 @@ def record_one(manifest, fixture, arn): self.assertEqual("run-2", after["Environment"]["Variables"]["LMI_TEST_RUN_ID"]) self.assertNotIn("TimeoutInMinutes", api.call_args_list[1].args[2]) self.assertEqual(["CREATE_COMPLETE", "UPDATE_COMPLETE"], [call.args[1] for call in wait.call_args_list]) - self.assertEqual(10, verify.call_count) + self.assertEqual(2 * len(FIXTURES), verify.call_count) idle.assert_called_once() @patch("cloud_suite.save") @@ -66,7 +76,7 @@ def test_unchanged_stack_still_gets_readback_without_recreation(self, api, exist deploy_fixtures(self.manifest(), []) self.assertEqual("update-stack", api.call_args.args[1]) wait.assert_not_called() - self.assertEqual(5, verify.call_count) + self.assertEqual(len(FIXTURES), verify.call_count) @patch("cloud_suite.save") @patch("cloud_suite.existing_stack") diff --git a/sdk-testing/src/test/java/software/amazon/lambda/durable/testing/LocalDurableTestRunnerTest.java b/sdk-testing/src/test/java/software/amazon/lambda/durable/testing/LocalDurableTestRunnerTest.java index 8967ad2e01..58a54e4ae8 100644 --- a/sdk-testing/src/test/java/software/amazon/lambda/durable/testing/LocalDurableTestRunnerTest.java +++ b/sdk-testing/src/test/java/software/amazon/lambda/durable/testing/LocalDurableTestRunnerTest.java @@ -10,6 +10,9 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.Test; import software.amazon.awssdk.services.lambda.model.Operation; @@ -204,4 +207,36 @@ void await() { assertTrue(ex.getMessage().contains("Invalid suspension. No operation is pending.")); assertTrue(ex.getMessage().contains("unregistered-op")); } + + @Test + void preservesInvocationExecutorFactoryAcrossReinvocations() { + var executors = new CopyOnWriteArrayList(); + var config = DurableConfig.builder() + .withInvocationExecutorFactory(() -> { + var executor = Executors.newFixedThreadPool(2); + executors.add(executor); + return executor; + }) + .build(); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + context.wait("wait", Duration.ofMinutes(1)); + return input; + }, + config); + + try { + var first = runner.run("test"); + runner.advanceTime(); + var second = runner.run("test"); + + assertEquals(ExecutionStatus.PENDING, first.getStatus()); + assertEquals(ExecutionStatus.SUCCEEDED, second.getStatus()); + assertEquals(2, executors.size()); + assertTrue(executors.stream().allMatch(ExecutorService::isTerminated)); + } finally { + executors.forEach(ExecutorService::shutdownNow); + } + } } 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 6b9bda1071..2640f3ab29 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableConfig.java @@ -8,9 +8,11 @@ import java.util.ArrayList; import java.util.List; import java.util.Objects; +import java.util.Optional; import java.util.Properties; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.function.Supplier; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.awssdk.auth.credentials.EnvironmentVariableCredentialsProvider; @@ -95,6 +97,7 @@ public final class DurableConfig { private final DurableExecutionClient durableExecutionClient; private final SerDes serDes; private final ExecutorService executorService; + private final Supplier invocationExecutorFactory; private final LoggerConfig loggerConfig; private final PollingStrategy pollingStrategy; private final Duration checkpointDelay; @@ -103,6 +106,10 @@ public final class DurableConfig { private final List pluginFactories; private DurableConfig(Builder builder) { + if (builder.executorService != null && builder.invocationExecutorFactory != null) { + throw new IllegalArgumentException( + "ExecutorService and invocation executor factory cannot both be configured"); + } this.pluginFactories = builder.loadDynamicPlugins ? DynamicPluginLoader.loadConfiguredPluginFactories(builder.pluginFactories) : List.copyOf(builder.pluginFactories); @@ -112,6 +119,7 @@ private DurableConfig(Builder builder) { this.serDes = Objects.requireNonNullElseGet(builder.serDes, JacksonSerDes::new); this.executorService = Objects.requireNonNullElseGet(builder.executorService, DurableConfig::createDefaultExecutor); + this.invocationExecutorFactory = builder.invocationExecutorFactory; this.loggerConfig = Objects.requireNonNullElseGet(builder.loggerConfig, LoggerConfig::defaults); this.pollingStrategy = Objects.requireNonNullElse(builder.pollingStrategy, PollingStrategies.Presets.DEFAULT); this.checkpointDelay = Objects.requireNonNullElseGet(builder.checkpointDelay, () -> Duration.ofSeconds(0)); @@ -166,6 +174,19 @@ public ExecutorService getExecutorService() { return executorService; } + /** + * Gets the optional factory for invocation-owned executors. + * + *

When present, {@code ExecutionManager} calls the factory once per Lambda invocation and owns the returned + * executor's shutdown. {@link #getExecutorService()} remains the shared executor used when no factory is + * configured. + * + * @return configured invocation executor factory, if any + */ + public Optional> getInvocationExecutorFactory() { + return Optional.ofNullable(invocationExecutorFactory); + } + /** * Gets the configured LoggerConfig. * @@ -239,12 +260,16 @@ public Builder toBuilder() { var builder = new Builder() .withDurableExecutionClient(durableExecutionClient) .withSerDes(serDes) - .withExecutorService(executorService) .withLoggerConfig(loggerConfig) .withPollingStrategy(pollingStrategy) .withCheckpointDelay(checkpointDelay) .withDeserializeAfterSerialization(deserializeAfterSerialization) .withCheckpointEmptyMap(checkpointEmptyMap); + if (invocationExecutorFactory != null) { + builder.withInvocationExecutorFactory(invocationExecutorFactory); + } else { + builder.withExecutorService(executorService); + } builder.pluginFactories = new ArrayList<>(pluginFactories); builder.loadDynamicPlugins = false; return builder; @@ -341,6 +366,7 @@ public static final class Builder { private DurableExecutionClient durableExecutionClient; private SerDes serDes; private ExecutorService executorService; + private Supplier invocationExecutorFactory; private LoggerConfig loggerConfig; private PollingStrategy pollingStrategy; private Duration checkpointDelay; @@ -422,6 +448,26 @@ public Builder withExecutorService(ExecutorService executorService) { return this; } + /** + * Sets a factory that creates an isolated executor for each Lambda invocation. + * + *

The factory is called once by each {@code DurableExecutor.execute()} call and may be called concurrently + * on Lambda Managed Instances. It must return a fresh, non-null {@link ExecutorService}; the SDK shuts that + * executor down only after actual admitted work and End finish, then observes termination within the remaining + * cleanup budget. Running code that ignores interruption may delay shutdown beyond that budget. Do not return + * an executor shared with another invocation or application component. This option cannot be combined with + * {@link #withExecutorService(ExecutorService)}. + * + * @param invocationExecutorFactory factory for invocation-owned executors + * @return This builder + * @throws NullPointerException if invocationExecutorFactory is null + */ + public Builder withInvocationExecutorFactory(Supplier invocationExecutorFactory) { + this.invocationExecutorFactory = + Objects.requireNonNull(invocationExecutorFactory, "Invocation executor factory cannot be null"); + return this; + } + /** * Sets a custom LoggerConfig. If not set, defaults to suppressing replay logs. * 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 a20d3c9b3a..2ceecec9c1 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 @@ -77,9 +77,9 @@ public static DurableExecutionOutput execute( BiFunction handler, DurableConfig config) { try (var manager = new ExecutionManager(input, config, lambdaContext)) { - manager.registerActiveThread(ROOT_THREAD_ID); - var invocation = new Invocation<>(input, lambdaContext, inputType, handler, config, manager); try { + manager.registerActiveThread(ROOT_THREAD_ID); + var invocation = new Invocation<>(input, lambdaContext, inputType, handler, config, manager); var completion = invocation.execute(); var ready = invocation.outcome.applyToEither( completion.handle((ignored, failure) -> new Outcome(null, failure)), Function.identity()); @@ -88,14 +88,15 @@ public static DurableExecutionOutput execute( ? selected.failure() : ExceptionHelper.unwrapInvocationFailure( selected.result().failure()); - var finalFailure = manager.finishInvocation(originalFailure); + var finalFailure = manager.finishInvocation(originalFailure, completion); var completedFailure = completion.isDone() ? completion .handle((ignored, failure) -> ExceptionHelper.unwrapInvocationFailure(failure)) .join() : null; if (ExecutionManager.isFatal(completedFailure)) ExceptionHelper.sneakyThrow(completedFailure); - if (ExecutionManager.isFatal(finalFailure)) ExceptionHelper.sneakyThrow(finalFailure); + if (ExecutionManager.isFatal(finalFailure)) + ExceptionHelper.sneakyThrow(retainCompletionDiagnostic(finalFailure, completedFailure)); if (ExecutionManager.isLifecycleFailure(finalFailure) && manager.isTerminalExecutionAcknowledged()) { // ACK outranks a lifecycle cutoff, not an unrelated failure already reported by delivery. if (completedFailure != null && !ExecutionManager.isLifecycleFailure(completedFailure)) @@ -103,13 +104,13 @@ public static DurableExecutionOutput execute( return DurableExecutionOutput.success(""); } if (finalFailure != null && (selected.failure() != null || finalFailure != originalFailure)) - ExceptionHelper.sneakyThrow(finalFailure); + ExceptionHelper.sneakyThrow(retainCompletionDiagnostic(finalFailure, completedFailure)); // Preserve the stored failure after control precedence; newer JDKs wrap join() cancellation. if (completedFailure != null) ExceptionHelper.sneakyThrow(completedFailure); return completion.join(); } catch (Throwable failure) { - var originalFailure = ExceptionHelper.unwrapInvocationFailure(failure); - // Finalize synchronous admission/dispatch failures before automatic close observes them again. + var originalFailure = ExecutionManager.normalizeInvocationBoundaryFailure(failure); + // Finalize synchronous setup/admission/dispatch failures before automatic close observes them again. var finalFailure = manager.finishInvocation(originalFailure); ExceptionHelper.sneakyThrow(finalFailure != null ? finalFailure : originalFailure); return null; @@ -119,6 +120,23 @@ public static DurableExecutionOutput execute( } } + private static Throwable retainCompletionDiagnostic(Throwable primary, Throwable completion) { + if (completion == null || completion == primary) return primary; + // Re-observing the same delivery error as an SDK cleanup cause does not add a new diagnostic. + if (hasSdkLifecycleCause(primary, completion) || hasSdkLifecycleCause(completion, primary)) return primary; + for (var diagnostic : primary.getSuppressed()) { + if (diagnostic == completion || hasSdkLifecycleCause(diagnostic, completion)) return primary; + } + primary.addSuppressed(completion); + return primary; + } + + private static boolean hasSdkLifecycleCause(Throwable failure, Throwable cause) { + return failure.getClass() == UnrecoverableDurableExecutionException.class + && ExecutionManager.isLifecycleFailure(failure) + && failure.getCause() == cause; + } + /** Invocation-local state, accessed on the handler thread when plugins are present. */ private static final class Invocation { private final Context lambdaContext; 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 ad06c1e02e..bb6f5a9530 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 @@ -11,6 +11,7 @@ import java.util.Comparator; import java.util.HashMap; import java.util.HashSet; +import java.util.IdentityHashMap; import java.util.List; import java.util.Map; import java.util.Objects; @@ -78,6 +79,8 @@ public class ExecutionManager implements SafeCloseable { private static final Logger logger = LoggerFactory.getLogger(ExecutionManager.class); private static final AtomicLong LOCAL_INVOCATION_SEQUENCE = new AtomicLong(); + private static final Set ACTIVE_INVOCATION_EXECUTORS = + Collections.synchronizedSet(Collections.newSetFromMap(new IdentityHashMap<>())); // Stop waiting for the root early enough to cancel tasks and flush checkpoints before the platform deadline. private static final Duration EXECUTION_DEADLINE_HEADROOM = Duration.ofMillis(500); @@ -88,6 +91,7 @@ public class ExecutionManager implements SafeCloseable { private static final Duration MAX_CLEANUP_DURATION = Duration.ofSeconds(30); private static final Duration MAX_GRACEFUL_DRAIN_DURATION = Duration.ofSeconds(5); private static final Duration MAX_CANCELLATION_DRAIN_DURATION = Duration.ofSeconds(1); + private static final Duration EXECUTOR_SHUTDOWN_RESERVE = Duration.ofMillis(100); private static final Duration CHECKPOINT_CLEANUP_RESERVE = Duration.ofMillis(100); private static final String LIFECYCLE_ERROR_TYPE = "software.amazon.lambda.durable.execution.InvocationLifecycleException"; @@ -117,6 +121,8 @@ enum LifecycleState { private final Context lambdaContext; private final AtomicReference executionMode; private final DurableConfig durableConfig; + private final ExecutorService executorService; + private final boolean ownsExecutorService; private final Set updatedOperationIdsSinceLastInvocation; private final Set initialOperationIds; @@ -230,10 +236,80 @@ public ExecutionManager(DurableExecutionInput input, DurableConfig config, Conte if (executionOp == null) { throw new IllegalStateException("EXECUTION operation not found"); } - logger.debug("DurableExecution.execute() called"); - logger.debug("DurableExecutionArn: {}", durableExecutionArn); - logger.debug("Initial operations count: {}", operationStorage.size()); - logger.debug("EXECUTION operation found: {}", executionOp.id()); + var invocationExecutorFactory = config.getInvocationExecutorFactory(); + if (invocationExecutorFactory.isPresent()) { + this.executorService = Objects.requireNonNull( + invocationExecutorFactory.get().get(), "Invocation executor factory returned null"); + this.ownsExecutorService = true; + } else { + this.executorService = config.getExecutorService(); + this.ownsExecutorService = false; + } + var leaseAcquired = ownsExecutorService && ACTIVE_INVOCATION_EXECUTORS.add(executorService); + if (ownsExecutorService && !leaseAcquired) { + throw new IllegalStateException("Invocation executor factory returned an executor already in use"); + } + try { + if (ownsExecutorService && executorService.isShutdown()) { + throw new IllegalStateException("Invocation executor factory returned a shutdown executor"); + } + logger.debug("DurableExecution.execute() called"); + logger.debug("DurableExecutionArn: {}", durableExecutionArn); + logger.debug("Initial operations count: {}", operationStorage.size()); + logger.debug("EXECUTION operation found: {}", executionOp.id()); + } catch (Throwable initializationFailure) { + if (leaseAcquired) initializationFailure = cleanupFailedExecutorInitialization(initializationFailure); + ExceptionHelper.sneakyThrow(initializationFailure); + } + } + + private Throwable cleanupFailedExecutorInitialization(Throwable initializationFailure) { + var initial = inspectOwnedExecutorFailure(initializationFailure); + var initialCause = normalizeOwnedExecutorFailure(initial, null); + var initialFatal = isFatal(initialCause) ? (Error) initialCause : null; + var deadline = cleanupDeadlineNanos(); + try { + Throwable cleanupFailure; + try { + // No SDK work has been dispatched. Avoid repeating the state callback that failed validation. + cleanupFailure = attemptOwnedExecutorShutdown(deadline, null, false, initialFatal); + } catch (Throwable failure) { + cleanupFailure = failure; + } + var selected = combineOwnedExecutorFailures(initializationFailure, initial, cleanupFailure); + return cleanupFailure == null ? selected : forceStopOwnedExecutor(deadline, selected); + } finally { + ACTIVE_INVOCATION_EXECUTORS.remove(executorService); + restoreCleanupInterrupt(); + } + } + + private Throwable forceStopOwnedExecutor(long deadline, Throwable priorFailure) { + var failure = priorFailure; + try { + executorService.shutdownNow(); + } catch (Throwable shutdownFailure) { + failure = combineOwnedExecutorFailures(failure, shutdownFailure); + } + try { + if (!awaitOwnedExecutorTermination(deadline)) + failure = combineOwnedExecutorFailures( + failure, lifecycleFailure("Invocation executor did not terminate during forced cleanup", null)); + } catch (InterruptedException interrupted) { + cleanupInterrupted.set(true); + failure = combineOwnedExecutorFailures( + failure, lifecycleFailure("Interrupted while forcing invocation executor shutdown", interrupted)); + } catch (Throwable terminationFailure) { + failure = combineOwnedExecutorFailures(failure, terminationFailure); + } + return failure; + } + + private boolean awaitOwnedExecutorTermination(long deadline) throws InterruptedException { + var terminated = executorService.awaitTermination(remainingNanos(deadline), TimeUnit.NANOSECONDS); + // A later state-query failure must not retain an executor whose termination was already confirmed. + if (terminated) ACTIVE_INVOCATION_EXECUTORS.remove(executorService); + return terminated; } private static String resolveInvocationId(Context lambdaContext, String durableExecutionArn) { @@ -373,7 +449,7 @@ public void registerOperation(BaseDurableOperation operation, BaseDurableOperati /** Submits the invocation's root handler and records its executor task separately from its logical result. */ CompletableFuture submitRootTask(Supplier action) { - return submitExecutorTask(ExecutorTaskHandle.Role.ROOT, null, durableConfig.getExecutorService(), action); + return submitExecutorTask(ExecutorTaskHandle.Role.ROOT, null, executorService, action); } /** Submits an operation handler and records its executor task separately from its logical result. */ @@ -388,7 +464,7 @@ public void executeOperationTask( ExecutorTaskHandle task = registerExecutorTask( operationRole(operation), operation, - durableConfig.getExecutorService(), + executorService, () -> { started.set(true); command.run(); @@ -409,7 +485,7 @@ private CompletableFuture submitOperationTask( return submitExecutorTask( operationRole(operation), operation, - durableConfig.getExecutorService(), + executorService, () -> { action.run(); return null; @@ -439,7 +515,7 @@ void executeRootTask(Runnable command) { submitExecutorTask( ExecutorTaskHandle.Role.ROOT, null, - durableConfig.getExecutorService(), + executorService, () -> { command.run(); return null; @@ -1186,7 +1262,7 @@ private void checkCheckpointRequestAllowed() { @Override public void close() { try { - var failure = finishInvocation(null); + var failure = finishInvocation(null, true); if (failure != null) ExceptionHelper.sneakyThrow(failure); } finally { pluginRunner.releasePlugins(); @@ -1478,7 +1554,7 @@ private void awaitTaskExit(ExecutorTaskHandle task, DrainFailure drain) { private void validateExecutorPool() { // double check if the thread pool is empty - if (durableConfig.getExecutorService() instanceof ThreadPoolExecutor threadPoolExecutor) { + if (executorService instanceof ThreadPoolExecutor threadPoolExecutor) { var threadCount = threadPoolExecutor.getActiveCount(); // This may or may not be a problem because getActiveCount doesn't return an accurate number if (threadCount > 0) { @@ -1613,10 +1689,24 @@ public CompletableFuture runUntilCompleteOrSuspend(CompletableFuture u } /** Called only by the invocation caller, so root exit and owned-executor cleanup never self-join. */ Throwable finishInvocation(Throwable originalFailure) { + return finishInvocation(originalFailure, false); + } + + Throwable finishInvocation(Throwable originalFailure, CompletableFuture completedRoot) { + return finishInvocation(originalFailure, false, completedRoot); + } + + private Throwable finishInvocation(Throwable originalFailure, boolean closeObservation) { + return finishInvocation(originalFailure, closeObservation, null); + } + + private Throwable finishInvocation( + Throwable originalFailure, boolean closeObservation, CompletableFuture completedRoot) { if (cleanupFinished) return originalFailure; beginDraining(); var hardDeadline = cleanupDeadlineNanos(); Throwable failure = null; + Throwable completionFatal = null; if (Thread.interrupted()) cleanupInterrupted.set(true); try { if (cleanupInterrupted.get() && !getActiveExecutorTasks().isEmpty()) @@ -1631,20 +1721,232 @@ Throwable finishInvocation(Throwable originalFailure) { } catch (RuntimeException problem) { failure = signalLifecycleFailure(problem); } finally { - failure = finishActualInvocationCleanup(failure); + failure = finishActualInvocationCleanup(failure, ownsExecutorService && !closeObservation); + if (ownsExecutorService && !closeObservation && isFatal(originalFailure)) + failure = combineCompletionFailure(originalFailure, failure); + if (ownsExecutorService && !closeObservation && failure != null && failure == getInvocationFatal()) + completionFatal = failure; + if (completedRoot != null && completedRoot.isDone() && getInvocationFatal() == null) { + var rootFailure = + completedRoot.handle((ignored, problem) -> problem).join(); + // MDC restoration runs after End. Keep its fatal local so cleanup can retain later diagnostics + // without changing the selected End outcome or publishing new checkpoint/cancellation control. + if (isFatal(rootFailure)) { + completionFatal = rootFailure; + failure = combineCompletionFailure(rootFailure, failure); + } + } + try { + failure = shutdownOwnedExecutor(hardDeadline, failure); + } catch (Throwable shutdownFailure) { + // Executor callbacks, including fallback shutdown and state queries, must not skip finalization. + shutdownFailure = normalizeOwnedExecutorFailure( + inspectOwnedExecutorFailure(shutdownFailure), isFatal(failure) ? (Error) failure : null); + failure = combineFailures(failure, shutdownFailure); + var firstFatal = getInvocationFatal(); + if (firstFatal != null) failure = combineFailures(firstFatal, failure); + // Inline execution can already be throwing this fatal through try-with-resources close. + // Normal completion keeps the raw fatal ahead of lifecycle/ACK handling. + if (closeObservation && isFatal(failure)) failure = new CompletionException(failure); + } cleanupFinished = true; restoreCleanupInterrupt(); } var recordedFailure = lifecycleFailure.get(); var selectedFailure = isFatal(originalFailure) || recordedFailure == null ? originalFailure : recordedFailure; if (failure == null) return selectedFailure; - return isFatal(selectedFailure) - ? combineFailures(selectedFailure, failure) + // Cleanup can return a later lifecycle failure without signaling it; retain the latched invocation failure. + if (isFatal(selectedFailure) || (recordedFailure != null && isLifecycleFailure(failure))) + return combineFailures(selectedFailure, failure); + return failure == completionFatal + ? combineCompletionFailure(failure, selectedFailure) : combineFailures(failure, selectedFailure); } + private static Throwable combineCompletionFailure(Throwable completionFatal, Throwable additional) { + // The completed root may already carry this same delivery failure from MDC restoration handling. + for (var diagnostic : completionFatal.getSuppressed()) { + if (diagnostic == additional) return completionFatal; + } + return combineFailures(completionFatal, additional); + } + + private Throwable shutdownOwnedExecutor(long hardDeadline, Throwable priorFailure) { + if (!ownsExecutorService) return priorFailure; + Throwable result = priorFailure; + Throwable thrown = null; + try { + result = attemptOwnedExecutorShutdown( + hardDeadline, priorFailure, true, isFatal(priorFailure) ? (Error) priorFailure : null); + } catch (Throwable failure) { + thrown = isFatal(priorFailure) ? combineOwnedExecutorFailures(priorFailure, failure) : failure; + } + try { + if (executorService.isShutdown()) ACTIVE_INVOCATION_EXECUTORS.remove(executorService); + } catch (Throwable stateFailure) { + // The lease query must not replace the shutdown attempt's fatal or erase its diagnostics. + thrown = combineOwnedExecutorFailures(thrown != null ? thrown : result, stateFailure); + } + // Preserve the outer finalizer's first-fatal and explicit-close observation handling. + if (thrown != null) ExceptionHelper.sneakyThrow(thrown); + return result; + } + + private static Throwable combineOwnedExecutorFailures(Throwable earlier, Throwable later) { + return combineOwnedExecutorFailures(earlier, inspectOwnedExecutorFailure(earlier), later); + } + + private static Throwable combineOwnedExecutorFailures(Throwable earlier, OwnedFailure inspected, Throwable later) { + var earlierCause = normalizeOwnedExecutorFailure(inspected, null); + var laterCause = normalizeOwnedExecutorFailure( + inspectOwnedExecutorFailure(later), isFatal(earlierCause) ? (Error) earlierCause : null); + return combineFailures( + isFatal(earlierCause) ? earlierCause : earlier, isFatal(laterCause) ? laterCause : later); + } + + static Throwable normalizeInvocationBoundaryFailure(Throwable failure) { + var inspected = inspectOwnedExecutorFailure(failure); + var discarded = new ArrayList(); + // A retained application ExecutionException and its cause graph remain intact, including their diagnostics. + for (var wrapper : inspected.wrappers()) { + if (wrapper == inspected.selected()) break; + discarded.add(wrapper); + } + if (!discarded.isEmpty()) retainOwnedWrapperDiagnostics(inspected.selected(), inspected.selected(), discarded); + return inspected.selected(); + } + + private record OwnedFailure(Throwable selected, List wrappers) {} + + @SuppressWarnings("removal") + private static OwnedFailure inspectOwnedExecutorFailure(Throwable failure) { + // Match invocation unwrapping, retaining its observed wrappers without reading a cause accessor twice. + var wrappers = new ArrayList(); + var visited = Collections.newSetFromMap(new IdentityHashMap()); + Throwable applicationWrapper = null; + while (failure instanceof CompletionException || failure instanceof ExecutionException) { + if (!visited.add(failure)) break; + wrappers.add(failure); + if (applicationWrapper == null && failure instanceof ExecutionException) applicationWrapper = failure; + Throwable cause; + try { + cause = failure.getCause(); + } catch (VirtualMachineError | ThreadDeath fatal) { + failure = fatal; + break; + } catch (Throwable unreadableCause) { + break; + } + if (cause == null) break; + failure = cause; + } + return new OwnedFailure( + applicationWrapper != null && !isFatal(failure) ? applicationWrapper : failure, wrappers); + } + + private static Throwable normalizeOwnedExecutorFailure(OwnedFailure failure, Error preferredFatal) { + if (isFatal(failure.selected()) && !failure.wrappers().isEmpty()) { + var fatal = (Error) failure.selected(); + retainOwnedWrapperDiagnostics(preferredFatal != null ? preferredFatal : fatal, fatal, failure.wrappers()); + } + return failure.selected(); + } + + private static void retainOwnedWrapperDiagnostics(Throwable primary, Throwable cause, List wrappers) { + var omitted = Collections.newSetFromMap(new IdentityHashMap()); + omitted.addAll(wrappers); + var pending = new ArrayList(); + wrappers.forEach(wrapper -> Collections.addAll(pending, wrapper.getSuppressed())); + var seen = Collections.newSetFromMap(new IdentityHashMap()); + for (var index = 0; index < pending.size(); index++) { + var diagnostic = pending.get(index); + if (diagnostic == primary || diagnostic == cause || omitted.contains(diagnostic) || !seen.add(diagnostic)) + continue; + var aliases = ownedAliasWrappers(diagnostic, primary, cause, omitted); + if (aliases != null) { + omitted.addAll(aliases); + aliases.forEach(alias -> Collections.addAll(pending, alias.getSuppressed())); + } else { + retainWrapperDiagnostic(primary, diagnostic); + } + } + } + + private static void retainWrapperDiagnostic(Throwable primary, Throwable diagnostic) { + // This transfers diagnostics from a discarded wrapper; it must not select a different failure. + for (var retained : primary.getSuppressed()) { + if (retained == diagnostic) return; + } + primary.addSuppressed(diagnostic); + } + + private static List ownedAliasWrappers( + Throwable failure, Throwable primary, Throwable cause, Set omitted) { + var aliases = new ArrayList(); + var seen = Collections.newSetFromMap(new IdentityHashMap()); + while (failure != primary && failure != cause && !omitted.contains(failure)) { + if (failure == null || !seen.add(failure)) return null; + var type = failure.getClass(); + if (type != CompletionException.class && type != ExecutionException.class) return null; + aliases.add(failure); + try { + failure = failure.getCause(); + } catch (Throwable unreadableCause) { + // An unreadable diagnostic link must not replace the selected fatal. + return null; + } + } + return aliases; + } + + private void requestOwnedExecutorShutdown(long deadline, boolean forceOnFailure, Error preferredFatal) { + try { + executorService.shutdown(); + } catch (Throwable shutdownFailure) { + if (preferredFatal != null) shutdownFailure = combineOwnedExecutorFailures(preferredFatal, shutdownFailure); + if (forceOnFailure) shutdownFailure = forceStopOwnedExecutor(deadline, shutdownFailure); + // The existing attempt and outer finalizer still classify the original failure and any later fatal. + ExceptionHelper.sneakyThrow(shutdownFailure); + } + } + + private Throwable attemptOwnedExecutorShutdown( + long hardDeadline, Throwable priorFailure, boolean forceOnShutdownFailure, Error preferredFatal) { + var executorDeadline = Math.max(System.nanoTime(), hardDeadline - CHECKPOINT_CLEANUP_RESERVE.toNanos()); + try { + // Constructor cleanup already recovers any failed attempt; normal finalization recovers this callback. + requestOwnedExecutorShutdown(executorDeadline, forceOnShutdownFailure, preferredFatal); + if (!getActiveExecutorTasks().isEmpty()) { + executorService.shutdownNow(); + } + if (awaitOwnedExecutorTermination(executorDeadline)) { + return priorFailure; + } + executorService.shutdownNow(); + if (awaitOwnedExecutorTermination(executorDeadline)) { + return priorFailure; + } + return combineFailures( + priorFailure, + lifecycleFailure("Invocation executor did not terminate before cleanup deadline", null)); + } catch (InterruptedException interrupted) { + cleanupInterrupted.set(true); + executorService.shutdownNow(); + return combineFailures( + priorFailure, lifecycleFailure("Interrupted while shutting down invocation executor", interrupted)); + } catch (RuntimeException failure) { + var cause = normalizeOwnedExecutorFailure(inspectOwnedExecutorFailure(failure), preferredFatal); + if (isFatal(cause)) ExceptionHelper.sneakyThrow(cause); + return combineFailures(priorFailure, lifecycleFailure("Invocation executor shutdown failed", failure)); + } + } + private Throwable drainInvocationTasks(Throwable originalFailure, long hardDeadline) throws InterruptedException { - var taskDeadline = Math.max(System.nanoTime(), hardDeadline - CHECKPOINT_CLEANUP_RESERVE.toNanos()); + var taskDeadline = Math.max( + System.nanoTime(), + hardDeadline + - CHECKPOINT_CLEANUP_RESERVE.toNanos() + - (ownsExecutorService ? EXECUTOR_SHUTDOWN_RESERVE.toNanos() : 0)); var recordedFailure = lifecycleFailure.get(); // Fatal publication already stops operation peers. Retain the root/scope owner's cleanup allowance; // only lifecycle cutoffs or interruption require immediate registry cancellation. @@ -1839,7 +2141,7 @@ private static Duration remainingDuration(long deadline) { } /** Cutoffs request cancellation; actual owner, continuation, and transport completion remain mandatory. */ - private Throwable finishActualInvocationCleanup(Throwable priorFailure) { + private Throwable finishActualInvocationCleanup(Throwable priorFailure, boolean retainRecordedFatal) { var drain = new DrainFailure(true); drain.run(() -> stopCheckpointContinuations(drain)); drain.run(() -> awaitActualInvocationTasks(drain)); @@ -1848,12 +2150,38 @@ private Throwable finishActualInvocationCleanup(Throwable priorFailure) { try { drain.finish(); } catch (Throwable failure) { + var recordedFatal = retainRecordedFatal ? getInvocationFatal() : null; + if (recordedFatal != null && isRecordedFatalAlias(failure, recordedFatal)) + return combineRecordedFatal(recordedFatal, priorFailure, failure); // Keep shutdown's transport wrapper: close can run while its fatal cause already escapes the body. return combineFailures(priorFailure, failure); } return priorFailure; } + private static boolean isRecordedFatalAlias(Throwable failure, Error fatal) { + var type = failure.getClass(); + return (type == CompletionException.class || type == ExecutionException.class) && failure.getCause() == fatal; + } + + private static Throwable combineRecordedFatal(Error fatal, Throwable priorFailure, Throwable observation) { + var pending = new ArrayList(); + if (priorFailure != null) pending.add(priorFailure); + pending.add(observation); + var seen = Collections.newSetFromMap(new IdentityHashMap()); + for (var index = 0; index < pending.size(); index++) { + var diagnostic = pending.get(index); + if (diagnostic == fatal || !seen.add(diagnostic)) continue; + if (isRecordedFatalAlias(diagnostic, fatal)) { + // Omit the observation carrier, which points back to this fatal, but keep its diagnostics. + Collections.addAll(pending, diagnostic.getSuppressed()); + } else { + combineCompletionFailure(fatal, diagnostic); + } + } + return fatal; + } + private void awaitActualInvocationTasks(DrainFailure drain) { while (true) { List> tasks; diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableConfigTest.java index a79b8c515d..ad2d8fbed5 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableConfigTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableConfigTest.java @@ -18,6 +18,7 @@ import java.util.List; import java.util.Map; import java.util.concurrent.ExecutorService; +import java.util.function.Supplier; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import software.amazon.awssdk.core.client.config.SdkAdvancedClientOption; @@ -92,6 +93,40 @@ void testBuilder_WithCustomExecutorService() { assertNotNull(config.getSerDes()); } + @Test + void testBuilder_WithInvocationExecutorFactory() { + Supplier factory = () -> mockExecutor; + + var config = DurableConfig.builder() + .withDurableExecutionClient(mockClient) + .withInvocationExecutorFactory(factory) + .build(); + + assertSame(factory, config.getInvocationExecutorFactory().orElseThrow()); + } + + @Test + void testBuilder_SharedExecutorAndInvocationFactoryAreMutuallyExclusive() { + var builder = DurableConfig.builder() + .withDurableExecutionClient(mockClient) + .withExecutorService(mockExecutor) + .withInvocationExecutorFactory(() -> mock(ExecutorService.class)); + + var exception = assertThrows(IllegalArgumentException.class, builder::build); + + assertEquals( + "ExecutorService and invocation executor factory cannot both be configured", exception.getMessage()); + } + + @Test + void testBuilder_NullInvocationExecutorFactoryThrowsException() { + var builder = DurableConfig.builder(); + + var exception = assertThrows(NullPointerException.class, () -> builder.withInvocationExecutorFactory(null)); + + assertEquals("Invocation executor factory cannot be null", exception.getMessage()); + } + @Test void testBuilder_DeserializeAfterSerializationDefaultsToTrue() { var config = @@ -174,6 +209,7 @@ void testBuilder_FluentAPI() { assertSame(builder, builder.withDurableExecutionClient(mockClient)); assertSame(builder, builder.withSerDes(mockSerDes)); assertSame(builder, builder.withExecutorService(mockExecutor)); + assertSame(builder, builder.withInvocationExecutorFactory(() -> mock(ExecutorService.class))); assertSame(builder, builder.withDeserializeAfterSerialization(false)); } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/ContinuationFatalBudgetTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/ContinuationFatalBudgetTest.java index 38bdb53eda..d78fab8925 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/ContinuationFatalBudgetTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/ContinuationFatalBudgetTest.java @@ -18,7 +18,7 @@ class ContinuationFatalBudgetTest { @ParameterizedTest @ValueSource(booleans = {false, true}) @Timeout(10) - void admittedContinuationMustSettleBeforeFatalClose(boolean reportWhileClosing) throws Exception { + void admittedContinuationMustExitBeforeFatalClose(boolean reportWhileClosing) throws Exception { var input = new DurableExecutionInput( "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/budget/execution", "token", @@ -40,6 +40,7 @@ void admittedContinuationMustSettleBeforeFatalClose(boolean reportWhileClosing) var caller = Executors.newSingleThreadExecutor(); var entered = new CountDownLatch(1); var release = new CountDownLatch(1); + var bodyExited = new CountDownLatch(1); var fatal = new InternalError("simulated invocation fatal"); var selected = manager.runUntilCompleteOrSuspend(new CompletableFuture()); var observed = manager.runCheckpointContinuation( @@ -50,6 +51,8 @@ void admittedContinuationMustSettleBeforeFatalClose(boolean reportWhileClosing) } catch (InterruptedException failure) { Thread.currentThread().interrupt(); throw new AssertionError(failure); + } finally { + bodyExited.countDown(); } }, worker); @@ -69,14 +72,15 @@ void admittedContinuationMustSettleBeforeFatalClose(boolean reportWhileClosing) manager.recordInvocationFatal(fatal); } assertThrows(TimeoutException.class, () -> closing.get(700, TimeUnit.MILLISECONDS)); + assertEquals(1, bodyExited.getCount(), "The continuation body is still held"); assertFalse(observed.isDone(), "The actual continuation is still owned while held"); release.countDown(); var closeFailure = closing.get(2, TimeUnit.SECONDS); if (closeFailure != null) assertSame(fatal, ExceptionHelper.unwrapAsyncFailure(closeFailure)); assertSame( fatal, ExceptionHelper.unwrapAsyncFailure(assertThrows(CompletionException.class, selected::join))); - assertTrue(observed.isDone(), "Close must await actual continuation settlement"); - release.countDown(); + assertEquals(0, bodyExited.getCount(), "Close must await actual continuation exit"); + // Observer publication follows lease release so its inline callbacks can safely call close again. observed.get(2, TimeUnit.SECONDS); } finally { release.countDown(); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerAdmissionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerAdmissionTest.java index 976d7355be..cdc5938c82 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerAdmissionTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerAdmissionTest.java @@ -266,7 +266,10 @@ void drainingAllowsOnlyTasksOfPreviouslyRegisteredOperations() throws Exception @ValueSource(booleans = {false, true}) void admittedContextTasksCanRegisterDescendantsWithoutLeakingAdmission(boolean coordinator) throws Exception { var executor = Executors.newSingleThreadExecutor(); + manager.close(); when(config.getExecutorService()).thenReturn(executor); + manager = createManager(); + manager.registerActiveThread(null); var parent = contextOperation("parent", coordinator ? OperationSubType.MAP : OperationSubType.RUN_IN_CHILD_CONTEXT); var descendant = step("descendant"); @@ -311,7 +314,10 @@ void inlineNestedInvocationCannotUseItsCallersDrainingPermission() throws Except }) .when(inline) .execute(any(Runnable.class)); + manager.close(); when(config.getExecutorService()).thenReturn(inline); + manager = createManager(); + manager.registerActiveThread(null); var other = createManager(); var parent = contextOperation("parent", OperationSubType.RUN_IN_CHILD_CONTEXT); var nested = contextOperation("nested", OperationSubType.RUN_IN_CHILD_CONTEXT); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerLifecycleTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerLifecycleTest.java index 5e26e95482..98321b071f 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerLifecycleTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerLifecycleTest.java @@ -7,18 +7,21 @@ import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNotSame; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import static software.amazon.lambda.durable.model.ExecutionStatus.PENDING; +import static software.amazon.lambda.durable.model.ExecutionStatus.SUCCEEDED; import com.amazonaws.services.lambda.runtime.Context; import java.time.Duration; import java.time.Instant; import java.util.List; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; @@ -30,6 +33,7 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.IntStream; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; @@ -51,8 +55,6 @@ class ExecutionManagerLifecycleTest { private static final String EXECUTION_ID = "execution"; - private static final String EXECUTION_ARN = - "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/test/" + EXECUTION_ID; @Test void capturesRequestIdAndInvocationDeadline() { @@ -507,6 +509,110 @@ public void onInvocationEnd(InvocationEndInfo info) { } } + @Test + void invocationExecutorFactoryIsolatesConcurrentRootsAndOwnsShutdown() throws Exception { + var executors = new CopyOnWriteArrayList(); + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withInvocationExecutorFactory(() -> { + var executor = Executors.newFixedThreadPool(2); + executors.add(executor); + return executor; + }) + .build(); + var runtime = Executors.newFixedThreadPool(2); + var rootsEntered = new CountDownLatch(2); + + try { + var responses = IntStream.range(0, 2) + .mapToObj(index -> runtime.submit(() -> DurableExecutor.execute( + input("isolated-" + index), + context(10_000), + TypeToken.get(String.class), + (value, durableContext) -> { + rootsEntered.countDown(); + await(rootsEntered); + return durableContext.step("step", String.class, step -> value); + }, + config))) + .toList(); + + for (var response : responses) { + assertEquals(SUCCEEDED, response.get(5, TimeUnit.SECONDS).status()); + } + assertEquals(2, executors.size()); + assertNotSame(executors.get(0), executors.get(1)); + assertTrue(executors.stream().allMatch(ExecutorService::isShutdown)); + assertTrue(executors.stream().allMatch(ExecutorService::isTerminated)); + } finally { + stop(runtime); + for (var executor : executors) { + stop(executor); + } + } + } + + @Test + void sharedExecutorRemainsCallerOwned() throws Exception { + var sharedExecutor = Executors.newCachedThreadPool(); + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withExecutorService(sharedExecutor) + .build(); + + try { + var response = DurableExecutor.execute( + input("shared"), + context(10_000), + TypeToken.get(String.class), + (value, durableContext) -> value, + config); + + assertEquals(SUCCEEDED, response.status()); + assertFalse(sharedExecutor.isShutdown()); + } finally { + stop(sharedExecutor); + } + } + + @Test + void invocationExecutorFactoryMustReturnUsableExecutor() { + var nullConfig = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withInvocationExecutorFactory(() -> null) + .build(); + assertThrows(NullPointerException.class, () -> new ExecutionManager(input("null-executor"), nullConfig, null)); + + var shutdownExecutor = Executors.newSingleThreadExecutor(); + shutdownExecutor.shutdown(); + var shutdownConfig = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withInvocationExecutorFactory(() -> shutdownExecutor) + .build(); + assertThrows( + IllegalStateException.class, + () -> new ExecutionManager(input("shutdown-executor"), shutdownConfig, null)); + } + + @Test + void invocationExecutorFactoryRejectsExecutorAlreadyLeasedToAnotherInvocation() { + var executor = Executors.newFixedThreadPool(2); + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withInvocationExecutorFactory(() -> executor) + .build(); + var first = new ExecutionManager(input("first"), config, null); + + try { + var exception = assertThrows( + IllegalStateException.class, () -> new ExecutionManager(input("second"), config, null)); + assertEquals("Invocation executor factory returned an executor already in use", exception.getMessage()); + } finally { + first.close(); + } + assertTrue(executor.isTerminated()); + } + private ExecutionManager createManager(Context context, ExecutorService executor) { var config = DurableConfig.builder().withDurableExecutionClient(TestUtils.createMockClient()); if (executor != null) { @@ -516,8 +622,12 @@ private ExecutionManager createManager(Context context, ExecutorService executor } private DurableExecutionInput input() { + return input(EXECUTION_ID); + } + + private DurableExecutionInput input(String executionId) { var execution = Operation.builder() - .id(EXECUTION_ID) + .id(executionId) .type(OperationType.EXECUTION) .status(OperationStatus.STARTED) .startTimestamp(Instant.now()) @@ -525,7 +635,7 @@ private DurableExecutionInput input() { ExecutionDetails.builder().inputPayload("\"input\"").build()) .build(); return new DurableExecutionInput( - EXECUTION_ARN, + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/test/" + executionId, "token", CheckpointUpdatedExecutionState.builder() .operations(List.of(execution)) diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/LmiLifecycleRegressionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/LmiLifecycleRegressionTest.java index 6d53fc0184..5b5d816df9 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/LmiLifecycleRegressionTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/LmiLifecycleRegressionTest.java @@ -10,17 +10,15 @@ import com.amazonaws.services.lambda.runtime.Context; import java.time.Duration; import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; -import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.stream.IntStream; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.condition.EnabledIfSystemProperty; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.ValueSource; import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState; import software.amazon.awssdk.services.lambda.model.ExecutionDetails; import software.amazon.awssdk.services.lambda.model.Operation; @@ -102,14 +100,20 @@ void deadlineMustCancelAnInvocationOwnedTaskAndBoundCleanup() throws Exception { } } - @ParameterizedTest - @ValueSource(booleans = {false, true}) - void sharedFixedExecutorMustProgress(boolean nested) throws Exception { + @Test + void invocationExecutorFactoryIsolatesConcurrentRoots() throws Exception { var entered = new CountDownLatch(2); - var users = (ThreadPoolExecutor) Executors.newFixedThreadPool(2); + var invocationExecutors = new CopyOnWriteArrayList(); var runtime = Executors.newFixedThreadPool(2); try { - var config = config(users); + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withInvocationExecutorFactory(() -> { + var executor = Executors.newFixedThreadPool(2); + invocationExecutors.add(executor); + return executor; + }) + .build(); var responses = IntStream.range(0, 2) .mapToObj(index -> runtime.submit(() -> DurableExecutor.execute( input("fixed-" + index), @@ -118,16 +122,6 @@ void sharedFixedExecutorMustProgress(boolean nested) throws Exception { (value, ctx) -> { entered.countDown(); await(entered); - if (nested) { - return ctx.runInChildContext( - "child", - String.class, - child -> child.runInChildContext( - "grandchild", - String.class, - grandchild -> - grandchild.step("step", String.class, step -> value))); - } return ctx.step("step", String.class, step -> value); }, config))) @@ -142,10 +136,12 @@ void sharedFixedExecutorMustProgress(boolean nested) throws Exception { .status(), "Blocking orchestration must not starve the work it awaits"); } + assertEquals(2, invocationExecutors.size()); + assertTrue(invocationExecutors.stream().allMatch(ExecutorService::isTerminated)); } finally { - users.setMaximumPoolSize(16); - users.setCorePoolSize(16); // Escape only after the assertion, so the test process cannot hang. - stop(users, runtime); + invocationExecutors.forEach(ExecutorService::shutdownNow); + runtime.shutdownNow(); + assertTrue(runtime.awaitTermination(10, TimeUnit.SECONDS)); } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/MdcFinalizationAdjacentTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/MdcFinalizationAdjacentTest.java new file mode 100644 index 0000000000..b2fdc34c3c --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/MdcFinalizationAdjacentTest.java @@ -0,0 +1,196 @@ +// 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.lang.reflect.InvocationTargetException; +import java.lang.reflect.Proxy; +import java.time.Instant; +import java.util.Arrays; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.slf4j.MDC; +import org.slf4j.spi.MDCAdapter; +import software.amazon.awssdk.services.lambda.model.*; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestContext; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.plugin.*; + +@Timeout(15) +class MdcFinalizationAdjacentTest { + @ParameterizedTest + @CsvSource({"true,false", "false,true", "true,true"}) + void restorationFatalRetainsTaskAndOwnedDiagnosticsOnce(boolean fallbackClear, boolean retrying) throws Exception { + var original = MDC.getMDCAdapter(); + var before = MDC.getCopyOfContextMap(); + var first = new InternalError("MDC restoration fatal"); + var cleanup = new InternalError("owned shutdown fatal"); + var ordinaryRestore = new IllegalStateException("ordinary restore failure"); + var retry = new UnrecoverableDurableExecutionException( + ErrorObject.builder() + .errorType("test.retry") + .errorMessage("retryable handler outcome") + .build(), + true); + var owner = new AtomicReference(); + var workerError = new AtomicReference(); + var escaped = new CountDownLatch(1); + var outerInjected = new AtomicBoolean(); + var fallbackInjected = new AtomicBoolean(); + var ends = new AtomicInteger(); + var end = new AtomicReference(); + var shutdowns = new AtomicInteger(); + var users = + new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue(), task -> { + var thread = new Thread(task, "adjacent-mdc-owner"); + thread.setUncaughtExceptionHandler((ignored, failure) -> { + workerError.set(failure); + escaped.countDown(); + }); + return thread; + }) { + @Override + public void shutdown() { + shutdowns.incrementAndGet(); + super.shutdown(); + throw cleanup; + } + }; + var callers = Executors.newSingleThreadExecutor(); + var proxy = (MDCAdapter) Proxy.newProxyInstance( + MDCAdapter.class.getClassLoader(), new Class[] {MDCAdapter.class}, (unused, method, arguments) -> { + if (Thread.currentThread() == owner.get()) { + if ((method.getName().equals("clear") + || method.getName().equals("setContextMap")) + && outerRestore() + && outerInjected.compareAndSet(false, true)) + throw fallbackClear ? ordinaryRestore : first; + if (fallbackClear + && method.getName().equals("clear") + && inMethod("clearMdcAfterOrdinaryFailure") + && fallbackInjected.compareAndSet(false, true)) throw first; + } + try { + return method.invoke(original, arguments); + } catch (InvocationTargetException failure) { + throw failure.getCause(); + } + }); + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withInvocationExecutorFactory(() -> users) + .withPlugins(ignored -> new DurableExecutionPlugin() { + public void onInvocationStart(InvocationInfo info) { + owner.set(Thread.currentThread()); + } + + public void onInvocationEnd(InvocationEndInfo info) { + assertSame(owner.get(), Thread.currentThread()); + end.set(info); + ends.incrementAndGet(); + } + }) + .build(); + try { + replaceAdapter(proxy); + var response = callers.submit(() -> DurableExecutor.execute( + input(), + new TestContext("mdc-adjacent"), + TypeToken.get(String.class), + (value, context) -> { + if (retrying) throw retry; + return "result"; + }, + config)); + var observed = assertThrows(ExecutionException.class, () -> response.get(5, TimeUnit.SECONDS)) + .getCause(); + assertSame(first, observed); + assertTrue(escaped.await(5, TimeUnit.SECONDS)); + assertSame(first, workerError.get()); + assertTrue(outerInjected.get()); + assertEquals(fallbackClear, fallbackInjected.get()); + assertEquals(1, ends.get()); + assertEquals( + retrying ? InvocationStatus.RETRYING : InvocationStatus.SUCCEEDED, + end.get().invocationStatus()); + assertSame( + retrying ? retry : null, + end.get().executionError(), + "Later MDC failure must not rewrite or re-fire End"); + assertEquals( + 1, + Arrays.stream(first.getSuppressed()) + .filter(error -> error == cleanup) + .count(), + "Retain later owned fatal once"); + if (retrying) + assertEquals( + 1, + Arrays.stream(first.getSuppressed()) + .filter(error -> error == retry) + .count(), + "The pre-existing retry diagnostic must not be duplicated during finalization"); + if (fallbackClear) + assertEquals( + 1, + Arrays.stream(first.getSuppressed()) + .filter(error -> error == ordinaryRestore) + .count()); + assertEquals(1, shutdowns.get()); + } finally { + try { + users.shutdownNow(); + callers.shutdownNow(); + assertTrue(users.awaitTermination(3, TimeUnit.SECONDS)); + assertTrue(callers.awaitTermination(3, TimeUnit.SECONDS)); + } finally { + replaceAdapter(original); + if (before == null) MDC.clear(); + else MDC.setContextMap(before); + } + } + } + + private static boolean outerRestore() { + return Arrays.stream(Thread.currentThread().getStackTrace()) + .anyMatch(frame -> frame.getClassName().equals(DurableExecutor.class.getName()) + && frame.getMethodName().startsWith("lambda$restoreMdcOnClose")); + } + + private static boolean inMethod(String method) { + return Arrays.stream(Thread.currentThread().getStackTrace()) + .anyMatch(frame -> frame.getClassName().equals(DurableExecutor.class.getName()) + && frame.getMethodName().equals(method)); + } + + private static void replaceAdapter(MDCAdapter adapter) throws Exception { + var method = MDC.class.getDeclaredMethod("setMDCAdapter", MDCAdapter.class); + method.setAccessible(true); + method.invoke(null, adapter); + } + + private static DurableExecutionInput input() { + return new DurableExecutionInput( + "arn:execution/test", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("test") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedCallbackFailureAggregationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedCallbackFailureAggregationTest.java new file mode 100644 index 0000000000..456a636776 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedCallbackFailureAggregationTest.java @@ -0,0 +1,303 @@ +// 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.Instant; +import java.util.ArrayDeque; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Set; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import java.util.stream.Stream; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; +import software.amazon.awssdk.services.lambda.model.*; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestContext; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.plugin.*; +import software.amazon.lambda.durable.serde.JacksonSerDes; +import software.amazon.lambda.durable.serde.SerDes; +import software.amazon.lambda.durable.util.ExceptionHelper; + +@Timeout(15) +class OwnedCallbackFailureAggregationTest { + enum Stage { + SHUTDOWN, + AWAIT, + STATE, + FORCED, + INITIAL, + SETUP + } + + enum Wrapper { + RAW, + COMPLETION, + EXECUTION, + CUSTOM + } + + record Scenario( + Stage stage, + Wrapper wrapper, + boolean recorded, + boolean repeatAtState, + boolean wrappedState, + boolean death) {} + + static Stream cases() { + var triple = Stream.of(Wrapper.values()) + .flatMap(wrapper -> Stream.of(false, true) + .flatMap(wrappedState -> Stream.of(false, true) + .map(death -> new Scenario(Stage.SHUTDOWN, wrapper, true, true, wrappedState, death)))); + var wrappers = Stream.of(Stage.values()) + .filter(stage -> stage != Stage.SETUP) + .flatMap(stage -> Stream.of(Wrapper.COMPLETION, Wrapper.EXECUTION) + .flatMap(wrapper -> (stage == Stage.INITIAL ? Stream.of(false) : Stream.of(false, true)) + .map(recorded -> new Scenario(stage, wrapper, recorded, false, false, false)))); + return Stream.concat( + Stream.concat(triple, wrappers), + Stream.of( + new Scenario(Stage.SHUTDOWN, Wrapper.RAW, false, true, false, false), + new Scenario(Stage.INITIAL, Wrapper.CUSTOM, false, false, false, false), + new Scenario(Stage.SETUP, Wrapper.COMPLETION, true, true, true, false))); + } + + @SuppressWarnings("removal") + @ParameterizedTest + @MethodSource("cases") + void callbackFailuresKeepTheFirstFatalAndAllNonAliasDiagnostics(Scenario scenario) throws Exception { + Error first = scenario.death() ? new ThreadDeath() : new InternalError("serializer fatal"); + var later = new InternalError("executor callback fatal"); + var diagnostic = new IllegalArgumentException("wrapper cleanup diagnostic"); + var nestedDiagnostic = new IllegalStateException("nested wrapper diagnostic"); + var queryDiagnostic = new IllegalArgumentException("query wrapper diagnostic"); + var ordinaryShutdown = new IllegalStateException("graceful shutdown failed before delegation"); + var causeReads = new AtomicInteger(); + Throwable callback = switch (scenario.wrapper()) { + case RAW -> later; + case COMPLETION -> new CompletionException(later); + case EXECUTION -> new ExecutionException(later); + case CUSTOM -> + new CompletionException(later) { + @Override + public synchronized Throwable getCause() { + assertEquals(1, causeReads.incrementAndGet(), "Read a custom wrapper cause only once"); + return super.getCause(); + } + }; + }; + if (callback != later) { + callback.addSuppressed(diagnostic); + var alias = new CompletionException(later); + alias.addSuppressed(nestedDiagnostic); + callback.addSuppressed(alias); + if (scenario.recorded()) callback.addSuppressed(first); + } + Throwable repeated = scenario.wrappedState() ? new CompletionException(first) : first; + if (repeated != first) repeated.addSuppressed(queryDiagnostic); + var started = new AtomicBoolean(); + var inject = new AtomicBoolean(true); + var firstAwait = new AtomicBoolean(true); + var firstInitial = new AtomicBoolean(true); + var shutdowns = new AtomicInteger(); + var forced = new AtomicInteger(); + var stateFailures = new AtomicInteger(); + var bodies = new AtomicInteger(); + var serializations = new AtomicInteger(); + var manager = new AtomicReference(); + var fatalAtShutdown = new AtomicReference(); + var ends = new AtomicInteger(); + var end = new AtomicReference(); + var workerFatal = new AtomicReference(); + var workerEscaped = new CountDownLatch(1); + var pool = + new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue(), task -> { + var thread = new Thread(task, "callback-failure-owner"); + thread.setUncaughtExceptionHandler((ignored, failure) -> { + workerFatal.set(failure); + workerEscaped.countDown(); + }); + return thread; + }) { + @Override + public void shutdown() { + shutdowns.incrementAndGet(); + started.set(true); + if (manager.get() != null) + fatalAtShutdown.set(manager.get().getInvocationFatal()); + if (inject.get() && scenario.stage() == Stage.FORCED) throw ordinaryShutdown; + super.shutdown(); + if (inject.get() && (scenario.stage() == Stage.SHUTDOWN || scenario.stage() == Stage.SETUP)) + ExceptionHelper.sneakyThrow(callback); + } + + @Override + public List shutdownNow() { + var tasks = super.shutdownNow(); + if (inject.get()) { + forced.incrementAndGet(); + if (scenario.stage() == Stage.FORCED) ExceptionHelper.sneakyThrow(callback); + } + return tasks; + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { + if (inject.get() && scenario.stage() == Stage.AWAIT && firstAwait.compareAndSet(true, false)) + ExceptionHelper.sneakyThrow(callback); + return super.awaitTermination(timeout, unit); + } + + @Override + public boolean isShutdown() { + if (inject.get()) { + if (scenario.stage() == Stage.INITIAL && firstInitial.compareAndSet(true, false)) { + stateFailures.incrementAndGet(); + ExceptionHelper.sneakyThrow(callback); + } + if (started.get() && (scenario.repeatAtState() || scenario.stage() == Stage.STATE)) { + stateFailures.incrementAndGet(); + ExceptionHelper.sneakyThrow(scenario.repeatAtState() ? repeated : callback); + } + } + return super.isShutdown(); + } + }; + pool.prestartAllCoreThreads(); + var serde = new JacksonSerDes(); + var config = DurableConfig.builder() + .withInvocationExecutorFactory(() -> pool) + .withDurableExecutionClient(TestUtils.createMockClient()) + .withSerDes(new SerDes() { + public String serialize(Object value) { + serializations.incrementAndGet(); + if (scenario.recorded()) throw first; + return serde.serialize(value); + } + + public T deserialize(String value, TypeToken type) { + return serde.deserialize(value, type); + } + }) + .withPlugins(ignored -> new DurableExecutionPlugin() { + public void onInvocationEnd(InvocationEndInfo info) { + ends.incrementAndGet(); + end.set(info); + } + }) + .build(); + var requestReads = new AtomicInteger(); + var runtime = scenario.stage() == Stage.SETUP + ? new TestContext("setup-callback") { + @Override + public String getAwsRequestId() { + if (requestReads.incrementAndGet() == 2) throw first; + return super.getAwsRequestId(); + } + } + : null; + try { + var observed = assertThrows( + Throwable.class, + () -> DurableExecutor.execute( + input(), + runtime, + TypeToken.get(String.class), + (value, context) -> { + manager.set(((DurableContextImpl) context).getExecutionManager()); + bodies.incrementAndGet(); + return "output"; + }, + config)); + assertSame(scenario.recorded() ? first : later, observed); + if (scenario.recorded() && scenario.stage() != Stage.SETUP) { + assertTrue(workerEscaped.await(3, TimeUnit.SECONDS)); + assertSame(first, workerFatal.get()); + assertSame(first, fatalAtShutdown.get()); + } + assertTrue(contains(observed, later), "Retain the later executor callback fatal"); + if (scenario.wrapper() != Wrapper.RAW) { + assertTrue(contains(observed, diagnostic), "Retain discarded wrapper diagnostics"); + assertTrue(contains(observed, nestedDiagnostic), "Retain diagnostics on wrapper aliases"); + } + if (scenario.wrappedState()) assertTrue(contains(observed, queryDiagnostic)); + if (scenario.stage() == Stage.FORCED) assertTrue(contains(observed, ordinaryShutdown)); + if (scenario.repeatAtState()) { + assertEquals(1, stateFailures.get()); + assertTrue(contains(observed, first)); + } + assertEquals(1, shutdowns.get()); + if (scenario.stage() == Stage.INITIAL || scenario.stage() == Stage.SETUP) { + assertEquals(0, bodies.get()); + assertEquals(0, serializations.get()); + assertEquals(0, ends.get()); + } else { + assertEquals(1, bodies.get()); + assertEquals(1, serializations.get()); + assertEquals(1, ends.get()); + assertEquals( + scenario.recorded() ? InvocationStatus.RETRYING : InvocationStatus.SUCCEEDED, + end.get().invocationStatus()); + assertSame(scenario.recorded() ? first : null, end.get().executionError()); + } + if (scenario.wrapper() == Wrapper.CUSTOM) assertEquals(1, causeReads.get()); + if (scenario.stage() == Stage.SETUP) { + assertEquals(2, requestReads.get()); + assertNull(fatalAtShutdown.get(), "Setup fatal is carried as originalFailure, without plugin state"); + } + assertAcyclic(observed, Collections.newSetFromMap(new IdentityHashMap())); + } finally { + inject.set(false); + pool.shutdownNow(); + assertTrue(pool.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + private static boolean contains(Throwable root, Throwable target) { + var pending = new ArrayDeque(); + var seen = Collections.newSetFromMap(new IdentityHashMap()); + pending.add(root); + while (!pending.isEmpty()) { + var next = pending.remove(); + if (next == target) return true; + if (!seen.add(next)) continue; + if (next.getCause() != null) pending.add(next.getCause()); + Collections.addAll(pending, next.getSuppressed()); + } + return false; + } + + private static void assertAcyclic(Throwable failure, Set path) { + assertTrue(path.add(failure), "Executor callback diagnostics must not point back to the first fatal"); + if (failure.getCause() != null) assertAcyclic(failure.getCause(), path); + for (var diagnostic : failure.getSuppressed()) assertAcyclic(diagnostic, path); + path.remove(failure); + } + + private static DurableExecutionInput input() { + return new DurableExecutionInput( + "arn:execution/callback-fatal", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("callback-fatal") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedCompletionAdjacentTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedCompletionAdjacentTest.java new file mode 100644 index 0000000000..77b56e4c19 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedCompletionAdjacentTest.java @@ -0,0 +1,231 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.execution; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import com.amazonaws.services.lambda.runtime.Context; +import java.time.Duration; +import java.time.Instant; +import java.util.*; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import software.amazon.awssdk.services.lambda.model.*; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.client.DurableExecutionClient; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.plugin.*; +import software.amazon.lambda.durable.serde.JacksonSerDes; +import software.amazon.lambda.durable.util.ExceptionHelper; + +@Timeout(15) +class OwnedCompletionAdjacentTest { + @Test + void acknowledgedDeliveryFailureIsRetainedByLaterOwnedFatal() throws Exception { + run(true, true); + } + + @Test + void acknowledgedDeliveryFailureStillOutranksOrdinaryOwnedFailure() throws Exception { + run(true, false); + } + + @Test + void callerCutoffRetainsLaterSerializerFailureAndOwnedFailure() throws Exception { + run(false, false); + } + + @Test + void callerCutoffDoesNotDuplicateTheSerializerUsedAsShutdownFailure() throws Exception { + run(false, false, true); + } + + private static void run(boolean acknowledged, boolean cleanupFatal) throws Exception { + run(acknowledged, cleanupFatal, false); + } + + private static void run(boolean acknowledged, boolean cleanupFatal, boolean sameError) throws Exception { + var delivery = new IllegalStateException("late delivery failure"); + Throwable shutdown = sameError + ? delivery + : cleanupFatal + ? new InternalError("later owned shutdown fatal") + : new IllegalStateException("later owned shutdown failure"); + var injected = new AtomicBoolean(true); + var shutdownCalls = new AtomicInteger(); + var users = new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>()) { + @Override + public void shutdown() { + shutdownCalls.incrementAndGet(); + super.shutdown(); + if (injected.get()) ExceptionHelper.sneakyThrow(shutdown); + } + }; + var callers = Executors.newSingleThreadExecutor(); + var caller = new AtomicReference(); + var callerInterrupted = new AtomicBoolean(); + var serializationEntered = new CountDownLatch(1); + var ownerInterrupted = new CountDownLatch(1); + var releaseSerialization = new CountDownLatch(1); + var owner = new AtomicReference(); + var ends = new AtomicInteger(); + var end = new AtomicReference(); + var updates = new AtomicReference>(); + DurableExecutionClient client; + if (acknowledged) { + client = mock(DurableExecutionClient.class); + when(client.checkpoint(any(), any(), any())).thenAnswer(call -> { + updates.set(List.copyOf(call.getArgument(2))); + return CheckpointDurableExecutionResponse.builder() + .checkpointToken("accepted-token") + .newExecutionState(CheckpointUpdatedExecutionState.builder() + .nextMarker("metadata") + .build()) + .build(); + }); + when(client.getExecutionState(any(), any(), any())).thenThrow(delivery); + } else { + client = TestUtils.createMockClient(); + } + var result = acknowledged ? "x".repeat(6 * 1024 * 1024 + 100) : "result"; + var serializer = new JacksonSerDes() { + @Override + public String serialize(Object value) { + if (!acknowledged && result.equals(value)) { + serializationEntered.countDown(); + boolean interrupted = false; + try { + while (true) { + try { + assertTrue(releaseSerialization.await(5, TimeUnit.SECONDS)); + break; + } catch (InterruptedException observed) { + interrupted = true; + ownerInterrupted.countDown(); + } + } + } finally { + if (interrupted) Thread.currentThread().interrupt(); + } + throw delivery; + } + return super.serialize(value); + } + }; + var config = DurableConfig.builder() + .withInvocationExecutorFactory(() -> users) + .withCheckpointDelay(Duration.ZERO) + .withDurableExecutionClient(client) + .withSerDes(serializer) + .withPlugins(ignored -> new DurableExecutionPlugin() { + @Override + public void onInvocationStart(InvocationInfo info) { + owner.set(Thread.currentThread()); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + assertSame(owner.get(), Thread.currentThread()); + ends.incrementAndGet(); + end.set(info); + } + }) + .build(); + try { + var invocation = callers.submit(() -> { + caller.set(Thread.currentThread()); + try { + return DurableExecutor.execute( + input(), runtime(), TypeToken.get(String.class), (value, context) -> result, config); + } finally { + callerInterrupted.set(Thread.interrupted()); + } + }); + if (!acknowledged) { + assertTrue(serializationEntered.await(5, TimeUnit.SECONDS)); + caller.get().interrupt(); + assertTrue(ownerInterrupted.await(5, TimeUnit.SECONDS)); + assertFalse(invocation.isDone(), "Caller must still wait for the actual serializer owner"); + releaseSerialization.countDown(); + } + var observed = assertThrows(ExecutionException.class, () -> invocation.get(5, TimeUnit.SECONDS)) + .getCause(); + assertAcyclic(observed, Collections.newSetFromMap(new IdentityHashMap<>())); + if (cleanupFatal) { + assertSame(shutdown, observed); + assertEquals(1, countIdentity(observed, delivery)); + } else if (acknowledged) { + assertSame(delivery, observed, "ACK must not erase the unrelated metadata failure"); + } else { + assertTrue( + ExecutionManager.isLifecycleFailure(observed), "The caller's recorded cutoff remains primary"); + assertEquals(1, countIdentity(observed, delivery)); + assertEquals(1, countIdentity(observed, shutdown)); + } + assertEquals(!acknowledged, callerInterrupted.get()); + assertEquals(1, ends.get()); + assertEquals(InvocationStatus.RETRYING, end.get().invocationStatus()); + assertSame(delivery, end.get().executionError()); + assertEquals(1, shutdownCalls.get()); + if (acknowledged) { + verify(client, times(1)).checkpoint(any(), any(), any()); + verify(client, times(1)).getExecutionState(any(), eq("accepted-token"), eq("metadata")); + assertEquals(1, updates.get().size()); + assertEquals(OperationAction.SUCCEED, updates.get().get(0).action()); + assertEquals('"' + result + '"', updates.get().get(0).payload()); + } + } finally { + releaseSerialization.countDown(); + injected.set(false); + users.shutdownNow(); + callers.shutdownNow(); + assertTrue(users.awaitTermination(3, TimeUnit.SECONDS)); + assertTrue(callers.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + private static int countIdentity(Throwable node, Throwable expected) { + if (node == null) return 0; + int count = node == expected ? 1 : 0; + count += countIdentity(node.getCause(), expected); + for (var suppressed : node.getSuppressed()) count += countIdentity(suppressed, expected); + return count; + } + + private static void assertAcyclic(Throwable node, Set path) { + if (node == null) return; + assertTrue(path.add(node), "SDK aggregation must not create a cycle"); + assertAcyclic(node.getCause(), path); + for (var suppressed : node.getSuppressed()) assertAcyclic(suppressed, path); + path.remove(node); + } + + private static Context runtime() { + var context = mock(Context.class); + when(context.getRemainingTimeInMillis()).thenReturn(10_000); + when(context.getAwsRequestId()).thenReturn("completion-adjacent-request"); + return context; + } + + private static DurableExecutionInput input() { + return new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/adjacent/exec", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("exec") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorAdmissionFailureTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorAdmissionFailureTest.java new file mode 100644 index 0000000000..9b9a54563a --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorAdmissionFailureTest.java @@ -0,0 +1,102 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.execution; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import com.amazonaws.services.lambda.runtime.Context; +import java.time.Instant; +import java.util.concurrent.*; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import software.amazon.awssdk.services.lambda.model.*; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.plugin.*; + +@Timeout(10) +class OwnedExecutorAdmissionFailureTest { + @ParameterizedTest + @CsvSource({"0,false", "0,true", "500,false", "500,true"}) + void cutoffBeforeDispatchFinalizesOwnedExecutorOnce(int remainingMillis, boolean plugins) throws Exception { + var dispatches = new AtomicInteger(); + var shutdowns = new AtomicInteger(); + var bodies = new AtomicInteger(); + var starts = new AtomicInteger(); + var ends = new AtomicInteger(); + var owned = new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue()) { + @Override + public void execute(Runnable command) { + dispatches.incrementAndGet(); + super.execute(command); + } + + @Override + public void shutdown() { + shutdowns.incrementAndGet(); + super.shutdown(); + } + }; + var runtime = mock(Context.class); + when(runtime.getRemainingTimeInMillis()).thenReturn(remainingMillis); + var builder = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withInvocationExecutorFactory(() -> owned); + if (plugins) + builder.withPlugins(ignored -> new DurableExecutionPlugin() { + public void onInvocationStart(InvocationInfo info) { + starts.incrementAndGet(); + } + + public void onInvocationEnd(InvocationEndInfo info) { + ends.incrementAndGet(); + } + }); + var input = new DurableExecutionInput( + "arn:execution/test", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("test") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + try { + var failure = assertThrows( + Throwable.class, + () -> DurableExecutor.execute( + input, + runtime, + TypeToken.get(String.class), + (value, context) -> { + bodies.incrementAndGet(); + return value; + }, + builder.build())); + assertAll( + () -> assertInstanceOf(UnrecoverableDurableExecutionException.class, failure), + () -> assertEquals(0, dispatches.get()), + () -> assertEquals(0, bodies.get()), + () -> assertEquals(0, starts.get()), + () -> assertEquals(0, ends.get()), + () -> assertEquals( + 1, shutdowns.get(), "Synchronous failure must finalize owned resources only once"), + () -> assertTrue(owned.isTerminated())); + assertTrue(((UnrecoverableDurableExecutionException) failure).isRetryable()); + } finally { + owned.shutdownNow(); + assertTrue(owned.awaitTermination(3, TimeUnit.SECONDS)); + } + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorDeliveryFailureTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorDeliveryFailureTest.java new file mode 100644 index 0000000000..d5c1cbae6e --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorDeliveryFailureTest.java @@ -0,0 +1,156 @@ +// 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.Instant; +import java.util.ArrayDeque; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import java.util.stream.Stream; +import org.junit.jupiter.api.Timeout; +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.*; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.plugin.*; +import software.amazon.lambda.durable.serde.SerDes; +import software.amazon.lambda.durable.util.ExceptionHelper; + +@Timeout(15) +class OwnedExecutorDeliveryFailureTest { + enum Cleanup { + NONE, + ORDINARY, + FATAL, + SAME_ERROR + } + + static Stream cases() { + return Stream.of(Cleanup.values()) + .flatMap(cleanup -> Stream.of(false, true).map(plugins -> Arguments.of(cleanup, plugins))); + } + + @ParameterizedTest + @MethodSource("cases") + void ownedCleanupRetainsSerializerFailureWithoutChangingWinner(Cleanup cleanup, boolean plugins) throws Exception { + var delivery = new IllegalStateException("output serialization failed"); + var ordinary = new IllegalArgumentException("owned shutdown failed"); + var fatal = new InternalError("owned shutdown fatal"); + Throwable shutdownFailure = switch (cleanup) { + case NONE -> null; + case ORDINARY -> ordinary; + case FATAL -> fatal; + case SAME_ERROR -> delivery; + }; + var bodies = new AtomicInteger(); + var serializations = new AtomicInteger(); + var shutdowns = new AtomicInteger(); + var ends = new AtomicInteger(); + var end = new AtomicReference(); + var pool = new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue()) { + @Override + public void shutdown() { + shutdowns.incrementAndGet(); + super.shutdown(); + if (shutdownFailure != null) ExceptionHelper.sneakyThrow(shutdownFailure); + } + }; + pool.prestartAllCoreThreads(); + var builder = DurableConfig.builder() + .withInvocationExecutorFactory(() -> pool) + .withDurableExecutionClient(TestUtils.createMockClient()) + .withSerDes(new SerDes() { + public String serialize(Object value) { + serializations.incrementAndGet(); + throw delivery; + } + + @SuppressWarnings("unchecked") + public T deserialize(String value, TypeToken type) { + return (T) "input"; + } + }); + if (plugins) + builder.withPlugins(ignored -> new DurableExecutionPlugin() { + public void onInvocationEnd(InvocationEndInfo info) { + ends.incrementAndGet(); + end.set(info); + } + }); + try { + var observed = assertThrows( + Throwable.class, + () -> DurableExecutor.execute( + input(), + null, + TypeToken.get(String.class), + (value, context) -> { + bodies.incrementAndGet(); + return "output"; + }, + builder.build())); + if (cleanup == Cleanup.NONE) assertSame(delivery, observed); + else if (cleanup == Cleanup.FATAL) assertSame(fatal, observed); + else { + var winner = assertInstanceOf(UnrecoverableDurableExecutionException.class, observed); + assertEquals("Invocation executor shutdown failed", winner.getMessage()); + assertSame(shutdownFailure, winner.getCause()); + } + assertTrue(contains(observed, delivery), "The completion-only serializer diagnostic must be retained"); + assertEquals(1, bodies.get()); + assertEquals(1, serializations.get()); + assertEquals(1, shutdowns.get()); + assertTrue(pool.isTerminated()); + assertEquals(plugins ? 1 : 0, ends.get()); + if (plugins) { + assertEquals(InvocationStatus.RETRYING, end.get().invocationStatus()); + assertSame(delivery, end.get().executionError()); + } + if (cleanup == Cleanup.SAME_ERROR) + assertEquals(0, observed.getSuppressed().length, "Do not duplicate the known lifecycle cause"); + } finally { + pool.shutdownNow(); + assertTrue(pool.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + private static boolean contains(Throwable root, Throwable expected) { + var pending = new ArrayDeque(); + var seen = Collections.newSetFromMap(new IdentityHashMap()); + pending.add(root); + while (!pending.isEmpty()) { + var next = pending.remove(); + if (next == expected) return true; + if (!seen.add(next)) continue; + if (next.getCause() != null) pending.add(next.getCause()); + Collections.addAll(pending, next.getSuppressed()); + } + return false; + } + + private static DurableExecutionInput input() { + return new DurableExecutionInput( + "arn:execution/delivery", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("delivery") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorFailureOrderingTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorFailureOrderingTest.java new file mode 100644 index 0000000000..3f9448e891 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorFailureOrderingTest.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.execution; + +import static org.junit.jupiter.api.Assertions.*; + +import java.time.Instant; +import java.util.ArrayDeque; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import java.util.stream.Stream; +import org.junit.jupiter.api.Timeout; +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.*; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.plugin.*; +import software.amazon.lambda.durable.util.ExceptionHelper; + +@Timeout(15) +class OwnedExecutorFailureOrderingTest { + enum Stage { + SHUTDOWN, + AWAIT + } + + enum Failures { + FATAL_THEN_ORDINARY, + DEATH_THEN_FATAL, + ORDINARY_THEN_FATAL, + COMPLETION_FATAL, + EXECUTION_FATAL + } + + static Stream failureOrders() { + return Stream.of(Stage.values()) + .flatMap(stage -> Stream.of(Failures.values()) + .flatMap(order -> Stream.of(false, true).map(prior -> Arguments.of(stage, order, prior)))); + } + + @SuppressWarnings("removal") + @ParameterizedTest + @MethodSource("failureOrders") + void stateQueryCannotReplaceEarlierFailures(Stage stage, Failures order, boolean priorHandlerFatal) + throws Exception { + var handlerFatal = new InternalError("earlier handler fatal"); + Error shutdownFatal = + order == Failures.DEATH_THEN_FATAL ? new ThreadDeath() : new InternalError("shutdown fatal"); + var ordinaryShutdown = new IllegalStateException("shutdown failed"); + Throwable shutdownFailure = switch (order) { + case ORDINARY_THEN_FATAL -> ordinaryShutdown; + case COMPLETION_FATAL -> new CompletionException(shutdownFatal); + case EXECUTION_FATAL -> new ExecutionException(shutdownFatal); + default -> shutdownFatal; + }; + Throwable queryFailure = order == Failures.DEATH_THEN_FATAL || order == Failures.ORDINARY_THEN_FATAL + ? new InternalError("later state-query fatal") + : new IllegalArgumentException("later state-query failure"); + var expected = + priorHandlerFatal ? handlerFatal : order == Failures.ORDINARY_THEN_FATAL ? queryFailure : shutdownFatal; + var inject = new AtomicBoolean(true); + var shutdownStarted = new AtomicBoolean(); + var bodyExited = new AtomicBoolean(); + var endBeforeShutdown = new AtomicBoolean(); + var shutdownCalls = new AtomicInteger(); + var stateCalls = new AtomicInteger(); + var events = new CopyOnWriteArrayList(); + var owner = new AtomicReference(); + var endThread = new AtomicReference(); + var end = new AtomicReference(); + var ends = new AtomicInteger(); + var workerFailure = new AtomicReference(); + var workerEscaped = new CountDownLatch(1); + var users = + new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>(), action -> { + var thread = new Thread(action, "ordered-shutdown-owner"); + thread.setUncaughtExceptionHandler((ignored, failure) -> { + workerFailure.set(failure); + workerEscaped.countDown(); + }); + return thread; + }) { + @Override + public void shutdown() { + if (inject.get()) { + events.add("shutdown"); + shutdownCalls.incrementAndGet(); + endBeforeShutdown.set(bodyExited.get() && ends.get() == 1); + } + shutdownStarted.set(true); + super.shutdown(); + if (inject.get() && stage == Stage.SHUTDOWN) ExceptionHelper.sneakyThrow(shutdownFailure); + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { + if (inject.get() && stage == Stage.AWAIT) { + events.add("await"); + ExceptionHelper.sneakyThrow(shutdownFailure); + } + return super.awaitTermination(timeout, unit); + } + + @Override + public boolean isShutdown() { + var stopped = super.isShutdown(); + if (inject.get() && shutdownStarted.get()) { + events.add("state"); + stateCalls.incrementAndGet(); + ExceptionHelper.sneakyThrow(queryFailure); + } + return stopped; + } + }; + var callers = Executors.newSingleThreadExecutor(); + var config = DurableConfig.builder() + .withInvocationExecutorFactory(() -> users) + .withDurableExecutionClient(TestUtils.createMockClient()) + .withPlugins(ignored -> new DurableExecutionPlugin() { + public void onInvocationStart(InvocationInfo info) { + owner.set(Thread.currentThread()); + } + + public void onInvocationEnd(InvocationEndInfo info) { + endThread.set(Thread.currentThread()); + end.set(info); + ends.incrementAndGet(); + events.add("end"); + } + }) + .build(); + try { + var result = callers.submit(() -> DurableExecutor.execute( + input(), + null, + TypeToken.get(String.class), + (value, context) -> { + try { + if (priorHandlerFatal) throw handlerFatal; + return "result"; + } finally { + bodyExited.set(true); + } + }, + config)); + var observed = assertThrows(ExecutionException.class, () -> result.get(5, TimeUnit.SECONDS)) + .getCause(); + assertSame(expected, observed, "The first recognized fatal must survive the later state query"); + assertTrue( + contains(observed, order == Failures.ORDINARY_THEN_FATAL ? ordinaryShutdown : shutdownFatal), + "The shutdown attempt's original failure remains reachable"); + assertTrue(contains(observed, queryFailure), "The later state-query diagnostic remains reachable"); + assertEquals( + stage == Stage.SHUTDOWN + ? List.of("end", "shutdown", "state") + : List.of("end", "shutdown", "await", "state"), + events); + assertTrue(endBeforeShutdown.get(), "Owned cleanup follows actual body and End completion"); + assertEquals(1, shutdownCalls.get(), "Automatic close cannot repeat finalized shutdown callbacks"); + assertEquals(1, stateCalls.get()); + assertEquals(1, ends.get()); + assertSame(owner.get(), endThread.get()); + assertEquals( + priorHandlerFatal ? InvocationStatus.RETRYING : InvocationStatus.SUCCEEDED, + end.get().invocationStatus()); + assertSame(priorHandlerFatal ? handlerFatal : null, end.get().executionError()); + if (priorHandlerFatal) { + assertTrue(workerEscaped.await(3, TimeUnit.SECONDS)); + assertSame(handlerFatal, workerFailure.get()); + } else assertNull(workerFailure.get()); + } finally { + inject.set(false); + users.shutdownNow(); + callers.shutdownNow(); + assertTrue(users.awaitTermination(3, TimeUnit.SECONDS)); + assertTrue(callers.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + private static boolean contains(Throwable root, Throwable expected) { + var seen = Collections.newSetFromMap(new IdentityHashMap()); + var pending = new ArrayDeque(); + pending.add(root); + while (!pending.isEmpty()) { + var failure = pending.removeFirst(); + if (failure == expected) return true; + if (!seen.add(failure)) continue; + if (failure.getCause() != null) pending.add(failure.getCause()); + Collections.addAll(pending, failure.getSuppressed()); + } + return false; + } + + private static DurableExecutionInput input() { + return new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/ordering/exec", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("exec") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorForcedShutdownTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorForcedShutdownTest.java new file mode 100644 index 0000000000..a44027ad6b --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorForcedShutdownTest.java @@ -0,0 +1,323 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.execution; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import java.time.Instant; +import java.util.ArrayDeque; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import java.util.stream.Stream; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +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.*; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.plugin.*; +import software.amazon.lambda.durable.util.ExceptionHelper; + +@Timeout(20) +class OwnedExecutorForcedShutdownTest { + enum First { + ORDINARY, + FATAL, + WRAPPED_FATAL + } + + enum Later { + NONE, + FORCED_FATAL, + STATE_FATAL, + FORCED_ORDINARY_STATE_FATAL + } + + static Stream cases() { + var normal = Stream.of(First.values()) + .flatMap(first -> Stream.of(Later.values()) + .flatMap(later -> + Stream.of(false, true).map(prior -> Arguments.of(first, later, prior, false)))); + return Stream.concat( + normal, + Stream.of( + Arguments.of(First.ORDINARY, Later.NONE, false, true), + Arguments.of(First.ORDINARY, Later.FORCED_FATAL, false, true))); + } + + @ParameterizedTest + @MethodSource("cases") + void aThrowingGracefulShutdownStillStopsItsOwnedWorkers( + First first, Later later, boolean handlerFatal, boolean acknowledged) throws Exception { + var earlierHandlerFatal = new InternalError("earlier handler fatal"); + var shutdownFatal = new InternalError("graceful shutdown fatal"); + var ordinaryShutdown = new IllegalStateException("graceful shutdown failed before delegation"); + Throwable shutdownFailure = switch (first) { + case ORDINARY -> ordinaryShutdown; + case FATAL -> shutdownFatal; + case WRAPPED_FATAL -> new CompletionException(shutdownFatal); + }; + Throwable forcedFailure = later == Later.FORCED_FATAL + ? new InternalError("forced stop fatal") + : later == Later.FORCED_ORDINARY_STATE_FATAL + ? new IllegalArgumentException("forced stop diagnostic") + : null; + Throwable stateFailure = later == Later.STATE_FATAL || later == Later.FORCED_ORDINARY_STATE_FATAL + ? new InternalError("state query fatal") + : null; + Throwable expectedFatal = handlerFatal + ? earlierHandlerFatal + : first != First.ORDINARY ? shutdownFatal : later == Later.FORCED_FATAL ? forcedFailure : stateFailure; + var end = new AtomicReference(); + var ends = new AtomicInteger(); + var owner = new AtomicReference(); + var endOwner = new AtomicReference(); + var caller = new AtomicReference(); + var callerInterrupted = new AtomicBoolean(); + var manager = new AtomicReference(); + var endEntered = new CountDownLatch(1); + var releaseEnd = new CountDownLatch(acknowledged ? 1 : 0); + var workerFailure = new AtomicReference(); + var workerEscaped = new CountDownLatch(1); + var users = new FailingPool(shutdownFailure, forcedFailure, stateFailure, workerFailure, workerEscaped); + assertEquals(1, users.prestartAllCoreThreads()); + var callers = Executors.newSingleThreadExecutor(); + var client = TestUtils.createMockClient(); + var value = acknowledged ? "x".repeat(6 * 1024 * 1024 + 100) : "result"; + var config = DurableConfig.builder() + .withInvocationExecutorFactory(() -> users) + .withDurableExecutionClient(client) + .withPlugins(ignored -> new DurableExecutionPlugin() { + public void onInvocationStart(InvocationInfo info) { + owner.set(Thread.currentThread()); + } + + public void onInvocationEnd(InvocationEndInfo info) { + endOwner.set(Thread.currentThread()); + end.set(info); + ends.incrementAndGet(); + endEntered.countDown(); + try { + assertTrue(releaseEnd.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + } + } + }) + .build(); + try { + var invocation = callers.submit(() -> { + caller.set(Thread.currentThread()); + try { + return DurableExecutor.execute( + input(), + null, + TypeToken.get(String.class), + (input, context) -> { + manager.set(((DurableContextImpl) context).getExecutionManager()); + if (handlerFatal) throw earlierHandlerFatal; + return value; + }, + config); + } finally { + callerInterrupted.set(Thread.interrupted()); + } + }); + assertTrue(endEntered.await(5, TimeUnit.SECONDS)); + if (acknowledged) caller.get().interrupt(); + Throwable observed; + if (acknowledged && expectedFatal == null) { + var output = invocation.get(5, TimeUnit.SECONDS); + assertEquals(ExecutionStatus.SUCCEEDED, output.status()); + assertEquals("", output.result()); + observed = manager.get().getLifecycleFailure(); + } else { + observed = assertThrows(ExecutionException.class, () -> invocation.get(5, TimeUnit.SECONDS)) + .getCause(); + if (expectedFatal != null) assertSame(expectedFatal, observed); + else { + var failure = assertInstanceOf(UnrecoverableDurableExecutionException.class, observed); + assertEquals("Invocation executor shutdown failed", failure.getMessage()); + assertSame(ordinaryShutdown, failure.getCause()); + } + } + assertTrue(users.isTerminated(), "A graceful callback failure must not retain live owned workers"); + assertEquals(1, users.gracefulCalls.get()); + assertEquals(1, users.forcedCalls.get(), "Normal finalization forces cleanup once, without close retry"); + assertTrue(contains(observed, first == First.ORDINARY ? ordinaryShutdown : shutdownFatal)); + if (forcedFailure != null) assertTrue(contains(observed, forcedFailure)); + if (stateFailure != null) assertTrue(contains(observed, stateFailure)); + assertEquals(acknowledged, callerInterrupted.get()); + assertEquals(1, ends.get()); + assertSame(owner.get(), endOwner.get()); + assertEquals( + handlerFatal ? InvocationStatus.RETRYING : InvocationStatus.SUCCEEDED, + end.get().invocationStatus()); + assertSame(handlerFatal ? earlierHandlerFatal : null, end.get().executionError()); + if (handlerFatal) { + assertTrue(workerEscaped.await(3, TimeUnit.SECONDS)); + assertSame(earlierHandlerFatal, workerFailure.get()); + } else assertNull(workerFailure.get()); + if (acknowledged) + verify(client, times(1)) + .checkpoint( + any(), + any(), + argThat(updates -> updates.size() == 1 + && updates.get(0).type() == OperationType.EXECUTION + && updates.get(0).action() == OperationAction.SUCCEED)); + if (stateFailure == null) { + users.inject.set(false); + var rejected = assertThrows( + IllegalStateException.class, + () -> DurableExecutor.execute( + input(), + null, + TypeToken.get(String.class), + (input, context) -> "must not run", + config)); + assertEquals( + "Invocation executor factory returned a shutdown executor", + rejected.getMessage(), + "The stopped executor's normal lease is released"); + } + } finally { + releaseEnd.countDown(); + users.dispose(); + callers.shutdownNow(); + assertTrue(callers.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + @Test + void sharedExecutorsDoNotEnterOwnedShutdownRecovery() throws Exception { + var users = new FailingPool( + new IllegalStateException("must not shut down shared executor"), + null, + null, + new AtomicReference<>(), + new CountDownLatch(1)); + users.prestartAllCoreThreads(); + try { + var config = DurableConfig.builder() + .withExecutorService(users) + .withDurableExecutionClient(TestUtils.createMockClient()) + .build(); + assertEquals( + ExecutionStatus.SUCCEEDED, + DurableExecutor.execute( + input(), null, TypeToken.get(String.class), (value, context) -> value, config) + .status()); + assertEquals(0, users.gracefulCalls.get()); + assertEquals(0, users.forcedCalls.get()); + assertFalse(users.isShutdown()); + } finally { + users.dispose(); + } + } + + private static final class FailingPool extends ThreadPoolExecutor { + private final Throwable first; + private final Throwable forced; + private final Throwable state; + private final AtomicBoolean inject = new AtomicBoolean(true); + private final AtomicBoolean firstCall = new AtomicBoolean(true); + private final AtomicBoolean shutdownAttempted = new AtomicBoolean(); + private final AtomicInteger gracefulCalls = new AtomicInteger(); + private final AtomicInteger forcedCalls = new AtomicInteger(); + + private FailingPool( + Throwable first, + Throwable forced, + Throwable state, + AtomicReference workerFailure, + CountDownLatch escaped) { + super(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>(), command -> { + var thread = new Thread(command, "forced-shutdown-owner"); + thread.setUncaughtExceptionHandler((ignored, failure) -> { + workerFailure.set(failure); + escaped.countDown(); + }); + return thread; + }); + this.first = first; + this.forced = forced; + this.state = state; + } + + @Override + public void shutdown() { + if (inject.get()) { + gracefulCalls.incrementAndGet(); + shutdownAttempted.set(true); + if (firstCall.compareAndSet(true, false)) ExceptionHelper.sneakyThrow(first); + } + super.shutdown(); + } + + @Override + public List shutdownNow() { + var result = super.shutdownNow(); + if (inject.get()) { + forcedCalls.incrementAndGet(); + if (forced != null) ExceptionHelper.sneakyThrow(forced); + } + return result; + } + + @Override + public boolean isShutdown() { + if (inject.get() && shutdownAttempted.get() && state != null) ExceptionHelper.sneakyThrow(state); + return super.isShutdown(); + } + + private void dispose() throws InterruptedException { + inject.set(false); + super.shutdownNow(); + assertTrue(super.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + private static boolean contains(Throwable root, Throwable target) { + if (root == null) return false; + var pending = new ArrayDeque(); + var seen = Collections.newSetFromMap(new IdentityHashMap()); + pending.add(root); + while (!pending.isEmpty()) { + var next = pending.removeFirst(); + if (next == target) return true; + if (!seen.add(next)) continue; + if (next.getCause() != null) pending.add(next.getCause()); + Collections.addAll(pending, next.getSuppressed()); + } + return false; + } + + private static DurableExecutionInput input() { + return new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/forced/exec", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("exec") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorInitializationFailureTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorInitializationFailureTest.java new file mode 100644 index 0000000000..db4b4e0e4b --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorInitializationFailureTest.java @@ -0,0 +1,263 @@ +// 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.Instant; +import java.util.ArrayDeque; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import java.util.stream.Stream; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +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 software.amazon.awssdk.services.lambda.model.*; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.plugin.*; +import software.amazon.lambda.durable.util.ExceptionHelper; + +@Timeout(15) +class OwnedExecutorInitializationFailureTest { + enum Kind { + ORDINARY, + FATAL, + WRAPPED_FATAL + } + + enum Cleanup { + NONE, + ORDINARY, + FATAL + } + + static Stream failures() { + return Stream.of(Kind.values()) + .flatMap(kind -> Stream.of(Cleanup.values()).map(cleanup -> Arguments.of(kind, cleanup))); + } + + @ParameterizedTest + @MethodSource("failures") + void failedStateValidationDisposesOnlyTheClaimedPool(Kind kind, Cleanup cleanup) throws Exception { + var initialFatal = new InternalError("initial state query fatal"); + Throwable initial = switch (kind) { + case ORDINARY -> new IllegalArgumentException("initial state query failed"); + case FATAL -> initialFatal; + case WRAPPED_FATAL -> new CompletionException(initialFatal); + }; + Throwable cleanupFailure = switch (cleanup) { + case NONE -> null; + case ORDINARY -> new IllegalStateException("graceful shutdown failed"); + case FATAL -> new InternalError("later shutdown fatal"); + }; + var expected = kind != Kind.ORDINARY ? initialFatal : cleanup == Cleanup.FATAL ? cleanupFailure : initial; + var bodies = new AtomicInteger(); + var starts = new AtomicInteger(); + var ends = new AtomicInteger(); + var pool = new QueryFailingPool(initial, cleanupFailure); + assertEquals(1, pool.prestartAllCoreThreads()); + var config = config(pool, starts, ends); + try { + var observed = assertThrows(Throwable.class, () -> invoke(config, bodies)); + assertSame(expected, observed, "Preserve the original failure unless cleanup introduces the first fatal"); + assertTrue(pool.isTerminated(), "Prestarted workers must be disposed before failed setup returns"); + assertTrue(pool.shutdownCalls.get() > 0 || pool.forcedCalls.get() > 0); + assertEquals(1, pool.queryCalls.get(), "Cleanup does not need to repeat the failed state query"); + assertEquals(0, bodies.get()); + assertEquals(0, starts.get()); + assertEquals(0, ends.get()); + if (cleanupFailure != null) { + assertTrue(contains(observed, cleanupFailure), "Retain later cleanup diagnostics"); + assertTrue(contains(observed, kind == Kind.ORDINARY ? initial : initialFatal)); + assertTrue(pool.forcedCalls.get() > 0, "Failed graceful shutdown still gets a best-effort forced stop"); + } + pool.inject.set(false); + var rejected = assertThrows(IllegalStateException.class, () -> invoke(config, bodies)); + assertEquals( + "Invocation executor factory returned a shutdown executor", + rejected.getMessage(), + "Failed initialization must release its lease, so closed-state validation remains authoritative"); + assertEquals(0, bodies.get()); + } finally { + pool.dispose(); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void initializationCleanupPreservesCallerInterrupt(boolean interrupted) throws Exception { + var initial = new IllegalArgumentException("initial query failed"); + var pool = new QueryFailingPool(initial, null); + pool.prestartAllCoreThreads(); + try { + if (interrupted) Thread.currentThread().interrupt(); + var observed = assertThrows( + Throwable.class, + () -> invoke(config(pool, new AtomicInteger(), new AtomicInteger()), new AtomicInteger())); + assertSame(initial, observed); + assertEquals(interrupted, Thread.interrupted()); + assertTrue(pool.isTerminated()); + } finally { + Thread.interrupted(); + pool.dispose(); + } + } + + @Test + void anAlreadyLeasedPoolIsNeitherQueriedNorStopped() throws Exception { + var pool = new QueryFailingPool(new IllegalArgumentException("foreign state query must not run"), null); + pool.inject.set(false); + pool.prestartAllCoreThreads(); + var entered = new CountDownLatch(1); + var release = new CountDownLatch(1); + var callers = Executors.newSingleThreadExecutor(); + var config = config(pool, new AtomicInteger(), new AtomicInteger()); + try { + var first = callers.submit(() -> DurableExecutor.execute( + input(), + null, + TypeToken.get(String.class), + (value, context) -> { + entered.countDown(); + try { + assertTrue(release.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException interrupted) { + throw new AssertionError(interrupted); + } + return value; + }, + config)); + assertTrue(entered.await(5, TimeUnit.SECONDS)); + var queriesBeforeSecond = pool.queryCalls.get(); + pool.inject.set(true); + var secondBodies = new AtomicInteger(); + var rejected = assertThrows(IllegalStateException.class, () -> invoke(config, secondBodies)); + assertEquals("Invocation executor factory returned an executor already in use", rejected.getMessage()); + assertEquals(queriesBeforeSecond, pool.queryCalls.get(), "Do not call foreign executor state callbacks"); + assertEquals(0, pool.shutdownCalls.get()); + assertEquals(0, pool.forcedCalls.get()); + assertEquals(0, secondBodies.get()); + assertFalse(first.isDone()); + pool.inject.set(false); + release.countDown(); + assertEquals( + ExecutionStatus.SUCCEEDED, first.get(5, TimeUnit.SECONDS).status()); + } finally { + pool.inject.set(false); + release.countDown(); + callers.shutdownNow(); + assertTrue(callers.awaitTermination(5, TimeUnit.SECONDS)); + pool.dispose(); + } + } + + private static DurableConfig config(QueryFailingPool pool, AtomicInteger starts, AtomicInteger ends) { + return DurableConfig.builder() + .withInvocationExecutorFactory(() -> pool) + .withDurableExecutionClient(TestUtils.createMockClient()) + .withPlugins(ignored -> new DurableExecutionPlugin() { + public void onInvocationStart(InvocationInfo info) { + starts.incrementAndGet(); + } + + public void onInvocationEnd(InvocationEndInfo info) { + ends.incrementAndGet(); + } + }) + .build(); + } + + private static Object invoke(DurableConfig config, AtomicInteger bodies) { + return DurableExecutor.execute( + input(), + null, + TypeToken.get(String.class), + (value, context) -> { + bodies.incrementAndGet(); + return value; + }, + config); + } + + private static final class QueryFailingPool extends ThreadPoolExecutor { + private final Throwable initial; + private final Throwable cleanup; + private final AtomicBoolean inject = new AtomicBoolean(true); + private final AtomicInteger queryCalls = new AtomicInteger(); + private final AtomicInteger shutdownCalls = new AtomicInteger(); + private final AtomicInteger forcedCalls = new AtomicInteger(); + + private QueryFailingPool(Throwable initial, Throwable cleanup) { + super(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>()); + this.initial = initial; + this.cleanup = cleanup; + } + + @Override + public boolean isShutdown() { + queryCalls.incrementAndGet(); + if (inject.get()) ExceptionHelper.sneakyThrow(initial); + return super.isShutdown(); + } + + @Override + public void shutdown() { + shutdownCalls.incrementAndGet(); + if (inject.get() && cleanup != null) ExceptionHelper.sneakyThrow(cleanup); + super.shutdown(); + } + + @Override + public List shutdownNow() { + forcedCalls.incrementAndGet(); + return super.shutdownNow(); + } + + private void dispose() throws InterruptedException { + inject.set(false); + super.shutdownNow(); + assertTrue(super.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + private static boolean contains(Throwable root, Throwable target) { + var pending = new ArrayDeque(); + var seen = Collections.newSetFromMap(new IdentityHashMap()); + pending.add(root); + while (!pending.isEmpty()) { + var next = pending.removeFirst(); + if (next == target) return true; + if (!seen.add(next)) continue; + if (next.getCause() != null) pending.add(next.getCause()); + Collections.addAll(pending, next.getSuppressed()); + } + return false; + } + + private static DurableExecutionInput input() { + return new DurableExecutionInput( + "arn:execution/init", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("init") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorLifecycleIdentityTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorLifecycleIdentityTest.java new file mode 100644 index 0000000000..3ae04d556b --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorLifecycleIdentityTest.java @@ -0,0 +1,141 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.execution; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import com.amazonaws.services.lambda.runtime.Context; +import java.time.Instant; +import java.util.Arrays; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import software.amazon.awssdk.services.lambda.model.*; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.plugin.*; + +@Timeout(15) +class OwnedExecutorLifecycleIdentityTest { + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void ownedTerminationTimeoutDoesNotReplacePreviouslySelectedCutoff(boolean plugins) throws Exception { + var releaseWorkerTail = new CountDownLatch(1); + var workerTailEntered = new CountDownLatch(1); + var first = new AtomicReference(); + var end = new AtomicReference(); + var stepCalls = new AtomicInteger(); + var owned = new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue()) { + @Override + public void execute(Runnable command) { + var commandReturned = new CountDownLatch(1); + super.execute(() -> { + try { + command.run(); + } finally { + commandReturned.countDown(); + } + }); + // Execute may be synchronous, as with CallerRunsPolicy. Runtime observation starts after the owner. + awaitIgnoringInterrupt(commandReturned); + } + + @Override + protected void afterExecute(Runnable command, Throwable failure) { + super.afterExecute(command, failure); + workerTailEntered.countDown(); + // Hold a real executor termination boundary after the tracked SDK task has actually exited. + awaitIgnoringInterrupt(releaseWorkerTail); + } + }; + var runtime = mock(Context.class); + when(runtime.getRemainingTimeInMillis()).thenReturn(2000); + var builder = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withInvocationExecutorFactory(() -> owned); + if (plugins) + builder.withPlugins(info -> new DurableExecutionPlugin() { + public void onInvocationEnd(InvocationEndInfo info) { + end.set(info); + } + }); + var input = new DurableExecutionInput( + "arn:execution/test", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("test") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + try { + var observed = assertThrows( + UnrecoverableDurableExecutionException.class, + () -> DurableExecutor.execute( + input, + runtime, + TypeToken.get(String.class), + (value, context) -> { + try { + assertFalse(new CountDownLatch(1).await(1600, TimeUnit.MILLISECONDS)); + } catch (InterruptedException interrupted) { + throw new AssertionError(interrupted); + } + try { + context.step("expired", String.class, step -> { + stepCalls.incrementAndGet(); + return "late"; + }); + } catch (UnrecoverableDurableExecutionException selected) { + first.set(selected); + } + return "fallback"; + }, + builder.build())); + assertNotNull(first.get()); + assertEquals( + "Invocation deadline reached before user work", first.get().getMessage()); + assertSame( + first.get(), + observed, + "Owned termination is a later lifecycle diagnostic, not a replacement outcome"); + assertTrue(Arrays.stream(first.get().getSuppressed()) + .anyMatch(later -> later instanceof UnrecoverableDurableExecutionException + && later.getMessage() + .equals("Invocation executor did not terminate before cleanup deadline"))); + assertTrue(workerTailEntered.await(2, TimeUnit.SECONDS)); + assertFalse(owned.isTerminated(), "The fixture must exercise a real termination timeout"); + assertEquals(0, stepCalls.get()); + if (plugins) { + assertEquals(InvocationStatus.RETRYING, end.get().invocationStatus()); + assertSame(first.get(), end.get().executionError()); + } + } finally { + releaseWorkerTail.countDown(); + owned.shutdownNow(); + assertTrue(owned.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + private static void awaitIgnoringInterrupt(CountDownLatch gate) { + var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(8); + while (true) { + try { + assertTrue(gate.await(Math.max(0, deadline - System.nanoTime()), TimeUnit.NANOSECONDS)); + return; + } catch (InterruptedException ignored) { + } + } + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorMdcFailureTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorMdcFailureTest.java new file mode 100644 index 0000000000..4a0e5893e3 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorMdcFailureTest.java @@ -0,0 +1,203 @@ +// 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.Instant; +import java.util.ArrayDeque; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import java.util.stream.Stream; +import org.junit.jupiter.api.Timeout; +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.helpers.BasicMDCAdapter; +import software.amazon.awssdk.services.lambda.model.*; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.plugin.*; +import software.amazon.lambda.durable.util.ExceptionHelper; + +@Timeout(15) +class OwnedExecutorMdcFailureTest { + enum Cleanup { + NONE, + ORDINARY, + FATAL, + WRAPPED_FATAL, + SAME_FATAL, + WRAPPED_SAME_FATAL + } + + static Stream cases() { + return Stream.concat( + Stream.of(false, true) + .flatMap(death -> + Stream.of(Cleanup.values()).map(cleanup -> Arguments.of(death, cleanup, false))), + Stream.of(false, true).map(death -> Arguments.of(death, Cleanup.NONE, true))); + } + + @SuppressWarnings("removal") + @ParameterizedTest + @MethodSource("cases") + void asynchronousMdcFatalRetainsLaterOwnedCleanup(boolean death, Cleanup cleanup, boolean shared) throws Exception { + Error first = death ? new ThreadDeath() : new InternalError("MDC restoration fatal"); + var laterFatal = new InternalError("later owned shutdown fatal"); + var ordinary = new IllegalStateException("later owned shutdown diagnostic"); + Throwable cleanupFailure = switch (cleanup) { + case NONE -> null; + case ORDINARY -> ordinary; + case FATAL -> laterFatal; + case WRAPPED_FATAL -> new CompletionException(laterFatal); + case SAME_FATAL -> first; + case WRAPPED_SAME_FATAL -> new CompletionException(first); + }; + var end = new AtomicReference(); + var owner = new AtomicReference(); + var ends = new AtomicInteger(); + var restored = new AtomicInteger(); + var shutdowns = new AtomicInteger(); + var manager = new AtomicReference(); + var fatalSeenByShutdown = new AtomicReference(); + var workerFailure = new AtomicReference(); + var workerEscaped = new CountDownLatch(1); + var originalAdapter = MDC.getMDCAdapter(); + var adapterField = MDC.class.getDeclaredField("MDC_ADAPTER"); + adapterField.setAccessible(true); + adapterField.set(null, new BasicMDCAdapter() { + @Override + public void setContextMap(Map values) { + if (Thread.currentThread() == owner.get() + && end.get() != null + && values.equals(Map.of("seed", "before-invocation"))) { + restored.incrementAndGet(); + throw first; + } + super.setContextMap(values); + } + }); + var pool = + new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue(), task -> { + var thread = new Thread( + () -> { + MDC.put("seed", "before-invocation"); + task.run(); + }, + "mdc-restoration-owner"); + thread.setUncaughtExceptionHandler((ignored, failure) -> { + workerFailure.set(failure); + workerEscaped.countDown(); + }); + return thread; + }) { + @Override + public void shutdown() { + shutdowns.incrementAndGet(); + fatalSeenByShutdown.set(manager.get().getInvocationFatal()); + super.shutdown(); + if (cleanupFailure != null) ExceptionHelper.sneakyThrow(cleanupFailure); + } + }; + pool.prestartAllCoreThreads(); + var builder = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withPlugins(ignored -> new DurableExecutionPlugin() { + public void onInvocationStart(InvocationInfo info) { + owner.set(Thread.currentThread()); + } + + public void onInvocationEnd(InvocationEndInfo info) { + ends.incrementAndGet(); + end.set(info); + } + }); + if (shared) builder.withExecutorService(pool); + else builder.withInvocationExecutorFactory(() -> pool); + try { + var observed = assertThrows( + Throwable.class, + () -> DurableExecutor.execute( + input(), + null, + TypeToken.get(String.class), + (value, context) -> { + manager.set(((DurableContextImpl) context).getExecutionManager()); + return value; + }, + builder.build())); + assertSame(first, observed, "The first root completion fatal must reach the caller unchanged"); + assertTrue(workerEscaped.await(3, TimeUnit.SECONDS)); + assertSame(first, workerFailure.get()); + assertEquals(1, restored.get(), "Only final handler MDC restoration injects the fatal"); + assertEquals(1, ends.get()); + assertEquals(InvocationStatus.SUCCEEDED, end.get().invocationStatus()); + assertNull(end.get().executionError(), "End must keep its already-selected successful body outcome"); + assertEquals(shared ? 0 : 1, shutdowns.get()); + if (shared) assertFalse(pool.isShutdown()); + else assertTrue(pool.isTerminated()); + assertNull(fatalSeenByShutdown.get(), "This regression reaches cleanup through the root future alone"); + if (cleanupFailure != null) { + var expected = cleanup == Cleanup.ORDINARY + ? ordinary + : cleanup == Cleanup.SAME_FATAL || cleanup == Cleanup.WRAPPED_SAME_FATAL ? first : laterFatal; + assertTrue(contains(observed, expected), "Retain the later owned cleanup diagnostic"); + } + assertNoCycle(observed, Collections.newSetFromMap(new IdentityHashMap())); + } finally { + try { + pool.shutdownNow(); + assertTrue(pool.awaitTermination(3, TimeUnit.SECONDS)); + } finally { + adapterField.set(null, originalAdapter); + } + } + } + + private static boolean contains(Throwable root, Throwable expected) { + var pending = new ArrayDeque(); + var seen = Collections.newSetFromMap(new IdentityHashMap()); + pending.add(root); + while (!pending.isEmpty()) { + var failure = pending.remove(); + if (failure == expected) return true; + if (!seen.add(failure)) continue; + if (failure.getCause() != null) pending.add(failure.getCause()); + Collections.addAll(pending, failure.getSuppressed()); + } + return false; + } + + private static void assertNoCycle(Throwable failure, Set path) { + assertTrue(path.add(failure), "Diagnostics must not contain a cause/suppression cycle"); + if (failure.getCause() != null) assertNoCycle(failure.getCause(), path); + for (var suppressed : failure.getSuppressed()) assertNoCycle(suppressed, path); + path.remove(failure); + } + + private static DurableExecutionInput input() { + return new DurableExecutionInput( + "arn:execution/mdc", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("mdc") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorSetupFailureTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorSetupFailureTest.java new file mode 100644 index 0000000000..276510c743 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorSetupFailureTest.java @@ -0,0 +1,452 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.execution; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import com.amazonaws.services.lambda.runtime.Context; +import java.time.Instant; +import java.util.ArrayDeque; +import java.util.Arrays; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.concurrent.*; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.ValueSource; +import software.amazon.awssdk.services.lambda.model.*; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestContext; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.model.DurableExecutionInput; +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.util.ExceptionHelper; + +class OwnedExecutorSetupFailureTest { + @SuppressWarnings("removal") + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void headerFatalRetainsIdentityWhenOwnedCloseReportsTheSameFatal(boolean death) throws Exception { + Error failure = death ? new ThreadDeath() : new OutOfMemoryError("simulated reused VM fatal"); + var shutdownCalls = new AtomicInteger(); + var bodies = new AtomicInteger(); + var owned = new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue()) { + @Override + public void shutdown() { + shutdownCalls.incrementAndGet(); + super.shutdown(); + throw failure; + } + }; + var runtime = mock(RuntimeContext.class); + when(runtime.getRemainingTimeInMillis()).thenReturn(10_000); + when(runtime.getXrayTraceId()).thenThrow(failure); + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withInvocationExecutorFactory(() -> owned) + .withPlugins(ignored -> new DurableExecutionPlugin() {}) + .build(); + var input = new DurableExecutionInput( + "arn:execution/test", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("test") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + try { + var observed = assertThrows( + Throwable.class, + () -> DurableExecutor.execute( + input, + runtime, + TypeToken.get(String.class), + (value, context) -> { + bodies.incrementAndGet(); + return value; + }, + config)); + assertSame(failure, observed, "Implicit close must not self-suppress a fatal raised before dispatch"); + assertEquals(0, bodies.get()); + assertEquals(1, shutdownCalls.get()); + assertTrue(owned.isTerminated()); + } finally { + owned.shutdownNow(); + assertTrue(owned.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + enum Setup { + ORDINARY, + FATAL, + WRAPPED_FATAL, + SAME_FATAL, + NESTED_WRAPPED_FATAL, + CHANGING_FATAL_CAUSE, + THROWING_FATAL_CAUSE, + APPLICATION_WRAPPER, + WRAPPED_ORDINARY, + NESTED_WRAPPED_ORDINARY, + CHANGING_ORDINARY_CAUSE, + COMPLETED_APPLICATION_WRAPPER, + WRAPPED_ORDINARY_FATAL_DIAGNOSTIC + } + + enum Cleanup { + NONE, + ORDINARY, + FATAL, + THREAD_DEATH, + WRAPPED_FATAL + } + + @SuppressWarnings("removal") + @ParameterizedTest + @CsvSource({ + "ORDINARY,FATAL,false", "ORDINARY,FATAL,true", + "ORDINARY,THREAD_DEATH,false", "ORDINARY,THREAD_DEATH,true", + "ORDINARY,WRAPPED_FATAL,false", "ORDINARY,WRAPPED_FATAL,true", + "WRAPPED_FATAL,FATAL,false", "WRAPPED_FATAL,FATAL,true", + "FATAL,FATAL,false", "FATAL,FATAL,true", + "SAME_FATAL,FATAL,false", "SAME_FATAL,FATAL,true", + "ORDINARY,NONE,false", "ORDINARY,NONE,true", + "ORDINARY,ORDINARY,false", "ORDINARY,ORDINARY,true", + "APPLICATION_WRAPPER,NONE,false", "APPLICATION_WRAPPER,NONE,true", + "APPLICATION_WRAPPER,FATAL,false", "APPLICATION_WRAPPER,FATAL,true", + "WRAPPED_FATAL,NONE,false", "WRAPPED_FATAL,NONE,true", + "NESTED_WRAPPED_FATAL,NONE,false", "NESTED_WRAPPED_FATAL,NONE,true", + "NESTED_WRAPPED_FATAL,FATAL,false", "NESTED_WRAPPED_FATAL,FATAL,true", + "CHANGING_FATAL_CAUSE,NONE,false", "CHANGING_FATAL_CAUSE,NONE,true", + "THROWING_FATAL_CAUSE,NONE,false", "THROWING_FATAL_CAUSE,NONE,true", + "WRAPPED_ORDINARY,NONE,false", "WRAPPED_ORDINARY,NONE,true", + "WRAPPED_ORDINARY,ORDINARY,false", "WRAPPED_ORDINARY,ORDINARY,true", + "WRAPPED_ORDINARY,FATAL,false", "WRAPPED_ORDINARY,FATAL,true", + "NESTED_WRAPPED_ORDINARY,NONE,false", "NESTED_WRAPPED_ORDINARY,NONE,true", + "NESTED_WRAPPED_ORDINARY,FATAL,false", "NESTED_WRAPPED_ORDINARY,FATAL,true", + "CHANGING_ORDINARY_CAUSE,NONE,false", "CHANGING_ORDINARY_CAUSE,NONE,true", + "COMPLETED_APPLICATION_WRAPPER,NONE,false", "COMPLETED_APPLICATION_WRAPPER,NONE,true", + "COMPLETED_APPLICATION_WRAPPER,FATAL,false", "COMPLETED_APPLICATION_WRAPPER,FATAL,true", + "WRAPPED_ORDINARY_FATAL_DIAGNOSTIC,NONE,false", "WRAPPED_ORDINARY_FATAL_DIAGNOSTIC,NONE,true", + "WRAPPED_ORDINARY_FATAL_DIAGNOSTIC,FATAL,false", "WRAPPED_ORDINARY_FATAL_DIAGNOSTIC,FATAL,true" + }) + void postManagerSetupFailureIsExplicitlyFinalized(Setup setup, Cleanup cleanup, boolean plugins) throws Exception { + var setupFatal = new InternalError("first setup fatal"); + var setupOrdinary = new IllegalStateException("nonfatal setup cause"); + var applicationCause = new CompletionException(new IllegalStateException("application cause")); + var innerApplicationDiagnostic = new IllegalArgumentException("retained cause diagnostic"); + applicationCause.addSuppressed(innerApplicationDiagnostic); + var applicationWrapper = new ExecutionException("application setup failure", applicationCause); + var applicationDiagnostic = new IllegalStateException("retained application diagnostic"); + applicationWrapper.addSuppressed(applicationDiagnostic); + Throwable wrapperDiagnostic = setup == Setup.WRAPPED_ORDINARY_FATAL_DIAGNOSTIC + ? new InternalError("suppressed wrapper diagnostic") + : new IllegalArgumentException("setup wrapper diagnostic"); + var nestedDiagnostic = new IllegalStateException("nested setup wrapper diagnostic"); + var aliasDiagnostic = new IllegalArgumentException("alias wrapper diagnostic"); + var causeReads = new AtomicInteger(); + Throwable setupFailure = switch (setup) { + case ORDINARY -> new IllegalStateException("second request ID read failed"); + case FATAL, SAME_FATAL -> setupFatal; + case WRAPPED_FATAL -> new CompletionException(setupFatal); + case NESTED_WRAPPED_FATAL -> { + var inner = new ExecutionException(setupFatal); + inner.addSuppressed(nestedDiagnostic); + var alias = new CompletionException(setupFatal); + alias.addSuppressed(aliasDiagnostic); + inner.addSuppressed(alias); + yield new CompletionException(inner); + } + case CHANGING_FATAL_CAUSE -> + new CompletionException("changing cause", null) { + @Override + public synchronized Throwable getCause() { + return causeReads.incrementAndGet() == 1 + ? setupFatal + : new IllegalStateException("second read"); + } + }; + case THROWING_FATAL_CAUSE -> + new CompletionException("fatal cause accessor", null) { + @Override + public synchronized Throwable getCause() { + causeReads.incrementAndGet(); + throw setupFatal; + } + }; + case APPLICATION_WRAPPER -> applicationWrapper; + case COMPLETED_APPLICATION_WRAPPER -> new CompletionException(applicationWrapper); + case WRAPPED_ORDINARY, WRAPPED_ORDINARY_FATAL_DIAGNOSTIC -> new CompletionException(setupOrdinary); + case NESTED_WRAPPED_ORDINARY -> { + var inner = new CompletionException(setupOrdinary); + inner.addSuppressed(nestedDiagnostic); + var alias = new CompletionException(setupOrdinary); + alias.addSuppressed(aliasDiagnostic); + inner.addSuppressed(alias); + yield new CompletionException(inner); + } + case CHANGING_ORDINARY_CAUSE -> + new CompletionException("changing nonfatal cause", null) { + @Override + public synchronized Throwable getCause() { + return causeReads.incrementAndGet() == 1 + ? setupOrdinary + : new IllegalStateException("second read"); + } + }; + }; + var wrappedFatal = setup == Setup.WRAPPED_FATAL + || setup == Setup.NESTED_WRAPPED_FATAL + || setup == Setup.CHANGING_FATAL_CAUSE + || setup == Setup.THROWING_FATAL_CAUSE; + var wrappedOrdinary = setup == Setup.WRAPPED_ORDINARY + || setup == Setup.WRAPPED_ORDINARY_FATAL_DIAGNOSTIC + || setup == Setup.NESTED_WRAPPED_ORDINARY + || setup == Setup.CHANGING_ORDINARY_CAUSE + || setup == Setup.COMPLETED_APPLICATION_WRAPPER; + var setupIsFatal = setup == Setup.FATAL || setup == Setup.SAME_FATAL || wrappedFatal; + var selectedSetup = setupIsFatal + ? setupFatal + : setup == Setup.COMPLETED_APPLICATION_WRAPPER + ? applicationWrapper + : wrappedOrdinary ? setupOrdinary : setupFailure; + if (wrappedFatal || wrappedOrdinary) { + setupFailure.addSuppressed(wrapperDiagnostic); + // Raw and observed-wrapper aliases must not create a cycle when their diagnostics transfer. + setupFailure.addSuppressed(selectedSetup); + if (setup == Setup.NESTED_WRAPPED_FATAL || setup == Setup.NESTED_WRAPPED_ORDINARY) { + setupFailure.getCause().addSuppressed(setupFailure); + setupFailure.addSuppressed(nestedDiagnostic); + } + } + Error cleanupFatal = setup == Setup.SAME_FATAL + ? setupFatal + : cleanup == Cleanup.THREAD_DEATH ? new ThreadDeath() : new InternalError("owned cleanup fatal"); + Throwable cleanupFailure = switch (cleanup) { + case NONE -> null; + case ORDINARY -> new IllegalArgumentException("ordinary owned cleanup failure"); + case FATAL, THREAD_DEATH -> cleanupFatal; + case WRAPPED_FATAL -> new CompletionException(cleanupFatal); + }; + var requests = new AtomicInteger(); + var starts = new AtomicInteger(); + var ends = new AtomicInteger(); + var bodies = new AtomicInteger(); + var factoryCalls = new AtomicInteger(); + var shutdowns = new AtomicInteger(); + var owned = new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue()) { + @Override + public void shutdown() { + shutdowns.incrementAndGet(); + super.shutdown(); + if (cleanupFailure != null) ExceptionHelper.sneakyThrow(cleanupFailure); + } + }; + assertEquals(1, owned.prestartAllCoreThreads()); + var runtime = new TestContext("setup-request") { + @Override + public String getAwsRequestId() { + if (requests.incrementAndGet() == 2) ExceptionHelper.sneakyThrow(setupFailure); + return super.getAwsRequestId(); + } + }; + var builder = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withInvocationExecutorFactory(() -> { + factoryCalls.incrementAndGet(); + return owned; + }); + if (plugins) + builder.withPlugins(ignored -> new DurableExecutionPlugin() { + public void onInvocationStart(InvocationInfo info) { + starts.incrementAndGet(); + } + + public void onInvocationEnd(InvocationEndInfo info) { + ends.incrementAndGet(); + } + }); + try { + var observed = assertThrows( + Throwable.class, + () -> DurableExecutor.execute( + setupInput(), + runtime, + TypeToken.get(String.class), + (value, context) -> { + bodies.incrementAndGet(); + return value; + }, + builder.build())); + if (setupIsFatal) { + assertSame(setupFatal, observed, "The first setup JVM fatal must remain primary"); + } else if (cleanup == Cleanup.FATAL + || cleanup == Cleanup.THREAD_DEATH + || cleanup == Cleanup.WRAPPED_FATAL) { + assertSame( + cleanupFatal, observed, "A cleanup JVM fatal must outrank an earlier nonfatal setup failure"); + } else if (cleanup == Cleanup.NONE) { + assertSame(selectedSetup, observed); + } + assertTrue(contains(observed, selectedSetup), "Retain setup diagnostics"); + if (cleanupFailure != null) + assertTrue( + contains(observed, cleanup == Cleanup.ORDINARY ? cleanupFailure : cleanupFatal), + "Retain later cleanup diagnostics"); + if (setup == Setup.APPLICATION_WRAPPER || setup == Setup.COMPLETED_APPLICATION_WRAPPER) { + assertSame(applicationCause, applicationWrapper.getCause()); + assertArrayEquals(new Throwable[] {innerApplicationDiagnostic}, applicationCause.getSuppressed()); + assertTrue(contains(observed, applicationDiagnostic)); + assertTrue(contains(observed, innerApplicationDiagnostic)); + assertEquals(setup == Setup.APPLICATION_WRAPPER ? 1 : 2, applicationWrapper.getSuppressed().length); + assertSame(applicationDiagnostic, applicationWrapper.getSuppressed()[0]); + assertFalse( + Arrays.asList(applicationWrapper.getSuppressed()).contains(innerApplicationDiagnostic), + "Do not promote diagnostics from the still-reachable application cause"); + } + if (wrappedFatal || wrappedOrdinary) { + assertTrue(contains(observed, wrapperDiagnostic), "The final winner must retain setup diagnostics"); + assertEquals( + 1, + Arrays.stream(selectedSetup.getSuppressed()) + .filter(e -> e == wrapperDiagnostic) + .count()); + assertFalse( + contains(observed, setupFailure), + "Do not attach a wrapper that points back to the selected cause"); + if (setup == Setup.NESTED_WRAPPED_FATAL || setup == Setup.NESTED_WRAPPED_ORDINARY) { + assertEquals( + 1, + Arrays.stream(selectedSetup.getSuppressed()) + .filter(e -> e == nestedDiagnostic) + .count()); + assertEquals( + 1, + Arrays.stream(selectedSetup.getSuppressed()) + .filter(e -> e == aliasDiagnostic) + .count()); + } + if (setup == Setup.CHANGING_FATAL_CAUSE + || setup == Setup.THROWING_FATAL_CAUSE + || setup == Setup.CHANGING_ORDINARY_CAUSE) + assertEquals(1, causeReads.get(), "Inspect each recognized setup wrapper cause only once"); + } + if (setup == Setup.WRAPPED_ORDINARY_FATAL_DIAGNOSTIC) + assertFalse( + contains(wrapperDiagnostic, selectedSetup), + "Diagnostic transfer must not reverse priority or create a cycle"); + assertEquals(2, requests.get(), "The failure must occur after the manager claimed its executor"); + assertEquals(1, factoryCalls.get()); + assertEquals(1, shutdowns.get(), "Automatic close must not repeat explicit finalization"); + assertEquals(0, bodies.get()); + assertEquals(0, starts.get()); + assertEquals(0, ends.get()); + assertTrue(owned.awaitTermination(3, TimeUnit.SECONDS)); + } finally { + owned.shutdownNow(); + assertTrue(owned.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void nonfatalPostManagerSetupStillLeavesSharedExecutorCallerOwned(boolean plugins) throws Exception { + var failure = new IllegalStateException("second request ID read failed"); + var requests = new AtomicInteger(); + var bodies = new AtomicInteger(); + var hooks = new AtomicInteger(); + var shared = new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue()); + shared.prestartAllCoreThreads(); + var runtime = new TestContext("setup-request") { + @Override + public String getAwsRequestId() { + if (requests.incrementAndGet() == 2) throw failure; + return super.getAwsRequestId(); + } + }; + var builder = DurableConfig.builder() + .withExecutorService(shared) + .withDurableExecutionClient(TestUtils.createMockClient()); + if (plugins) + builder.withPlugins(ignored -> new DurableExecutionPlugin() { + public void onInvocationStart(InvocationInfo info) { + hooks.incrementAndGet(); + } + + public void onInvocationEnd(InvocationEndInfo info) { + hooks.incrementAndGet(); + } + }); + try { + assertSame( + failure, + assertThrows( + Throwable.class, + () -> DurableExecutor.execute( + setupInput(), + runtime, + TypeToken.get(String.class), + (value, context) -> { + bodies.incrementAndGet(); + return value; + }, + builder.build()))); + assertEquals(2, requests.get()); + assertEquals(0, bodies.get()); + assertEquals(0, hooks.get()); + assertFalse(shared.isShutdown()); + } finally { + shared.shutdownNow(); + assertTrue(shared.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + private static DurableExecutionInput setupInput() { + return new DurableExecutionInput( + "arn:execution/setup", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("setup") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + } + + private static boolean contains(Throwable root, Throwable expected) { + var queue = new ArrayDeque(); + var seen = Collections.newSetFromMap(new IdentityHashMap()); + queue.add(root); + while (!queue.isEmpty()) { + var next = queue.remove(); + if (next == expected) return true; + if (!seen.add(next)) continue; + var cause = next.getCause(); + if (cause != null) queue.add(cause); + queue.addAll(Arrays.asList(next.getSuppressed())); + } + return false; + } + + private abstract static class RuntimeContext implements Context { + @Override + public String getXrayTraceId() { + return null; + } + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorShutdownFailureTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorShutdownFailureTest.java new file mode 100644 index 0000000000..2103cf2787 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorShutdownFailureTest.java @@ -0,0 +1,262 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.execution; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import com.amazonaws.services.lambda.runtime.Context; +import java.time.Instant; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import software.amazon.awssdk.services.lambda.model.*; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.plugin.*; +import software.amazon.lambda.durable.serde.JacksonSerDes; + +@Timeout(15) +class OwnedExecutorShutdownFailureTest { + enum ShutdownStage { + SHUTDOWN, + INTERRUPTED_FALLBACK, + STATE_QUERY + } + + @SuppressWarnings("removal") + @ParameterizedTest + @CsvSource({ + "false, SHUTDOWN, false", "true, SHUTDOWN, false", + "false, INTERRUPTED_FALLBACK, false", "true, INTERRUPTED_FALLBACK, false", + "false, STATE_QUERY, false", "true, STATE_QUERY, false", + "false, SHUTDOWN, true", "true, SHUTDOWN, true", + "false, INTERRUPTED_FALLBACK, true", "true, INTERRUPTED_FALLBACK, true", + "false, STATE_QUERY, true", "true, STATE_QUERY, true" + }) + void serializationFatalRemainsPrimaryWhenOwnedShutdownIntroducesAnotherFatal( + boolean deathFirst, ShutdownStage stage, boolean inline) throws Exception { + Error first = deathFirst ? new ThreadDeath() : new InternalError("first serializer fatal"); + Error second = deathFirst ? new InternalError("later owned shutdown fatal") : new ThreadDeath(); + var workerFailure = new AtomicReference(); + var escaped = new CountDownLatch(1); + var ends = new AtomicInteger(); + var shutdownCalls = new AtomicInteger(); + var injectFailure = new AtomicBoolean(true); + var callerInterrupted = new AtomicBoolean(); + var owner = new AtomicReference(); + var users = + new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>(), action -> { + var thread = new Thread(action, "owned-fatal-owner"); + thread.setUncaughtExceptionHandler((ignored, failure) -> { + workerFailure.set(failure); + escaped.countDown(); + }); + return thread; + }) { + @Override + public void execute(Runnable command) { + if (inline) command.run(); + else super.execute(command); + } + + @Override + public void shutdown() { + shutdownCalls.incrementAndGet(); + super.shutdown(); + if (injectFailure.get() && stage == ShutdownStage.SHUTDOWN) throw second; + } + + @Override + public List shutdownNow() { + var queued = super.shutdownNow(); + if (injectFailure.get() && stage == ShutdownStage.INTERRUPTED_FALLBACK) throw second; + return queued; + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { + if (injectFailure.get() && stage == ShutdownStage.INTERRUPTED_FALLBACK) { + Thread.currentThread().interrupt(); + new CountDownLatch(1) + .await(); // Observe the interruption even if an inline pool terminated. + } + return super.awaitTermination(timeout, unit); + } + + @Override + public boolean isShutdown() { + var stopped = super.isShutdown(); + if (injectFailure.get() && stage == ShutdownStage.STATE_QUERY && stopped) throw second; + return stopped; + } + }; + var callers = Executors.newSingleThreadExecutor(); + var serializer = new JacksonSerDes() { + @Override + public String serialize(Object value) { + if ("result".equals(value)) throw first; + return super.serialize(value); + } + }; + try { + var config = DurableConfig.builder() + .withInvocationExecutorFactory(() -> users) + .withDurableExecutionClient(TestUtils.createMockClient()) + .withSerDes(serializer) + .withPlugins(ignoredInvocation -> new DurableExecutionPlugin() { + @Override + public void onInvocationStart(InvocationInfo info) { + owner.set(Thread.currentThread()); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + assertSame(owner.get(), Thread.currentThread()); + assertEquals(InvocationStatus.RETRYING, info.invocationStatus()); + assertSame(first, info.executionError()); + ends.incrementAndGet(); + } + }) + .build(); + var invocation = callers.submit(() -> { + try { + return DurableExecutor.execute( + input(), + runtime(10_000), + TypeToken.get(String.class), + (value, context) -> "result", + config); + } finally { + callerInterrupted.set(Thread.interrupted()); + } + }); + assertSame( + first, + assertThrows(ExecutionException.class, () -> invocation.get(5, TimeUnit.SECONDS)) + .getCause()); + if (inline) { + assertNull(workerFailure.get(), "The inline fatal belongs to the invocation caller"); + } else { + assertTrue(escaped.await(3, TimeUnit.SECONDS)); + assertSame(first, workerFailure.get()); + } + assertEquals(1, ends.get()); + assertEquals(1, shutdownCalls.get(), "Implicit close must not retry an already-finalized invocation"); + assertEquals(stage == ShutdownStage.INTERRUPTED_FALLBACK, callerInterrupted.get()); + assertTrue(Arrays.asList(first.getSuppressed()).contains(second)); + } finally { + injectFailure.set(false); + callers.shutdownNow(); + users.shutdownNow(); + assertTrue(callers.awaitTermination(3, TimeUnit.SECONDS)); + assertTrue(users.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + @Test + void ownedShutdownFatalOutranksAcknowledgedSuccessAndCallerInterruption() throws Exception { + var fatal = new InternalError("owned shutdown fatal after ACK"); + var injectFailure = new AtomicBoolean(true); + var shutdownCalls = new AtomicInteger(); + var users = new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>()) { + @Override + public void shutdown() { + shutdownCalls.incrementAndGet(); + super.shutdown(); + if (injectFailure.get()) throw fatal; + } + }; + var callers = Executors.newSingleThreadExecutor(); + var caller = new AtomicReference(); + var callerInterrupted = new AtomicBoolean(); + var endInfo = new AtomicReference(); + var endEntered = new CountDownLatch(1); + var releaseEnd = new CountDownLatch(1); + var client = TestUtils.createMockClient(); + var result = "x".repeat(6 * 1024 * 1024 + 100); // Force the explicit EXECUTION SUCCEED checkpoint. + var config = DurableConfig.builder() + .withInvocationExecutorFactory(() -> users) + .withDurableExecutionClient(client) + .withPlugins(ignored -> new DurableExecutionPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + endInfo.set(info); + endEntered.countDown(); + try { + assertTrue(releaseEnd.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + } + } + }) + .build(); + try { + var invocation = callers.submit(() -> { + caller.set(Thread.currentThread()); + try { + return DurableExecutor.execute( + input(), runtime(10_000), TypeToken.get(String.class), (value, context) -> result, config); + } finally { + callerInterrupted.set(Thread.interrupted()); + } + }); + assertTrue(endEntered.await(5, TimeUnit.SECONDS)); + assertEquals(InvocationStatus.SUCCEEDED, endInfo.get().invocationStatus()); + assertNull(endInfo.get().executionError()); + caller.get().interrupt(); + assertSame( + fatal, + assertThrows(ExecutionException.class, () -> invocation.get(5, TimeUnit.SECONDS)) + .getCause()); + assertTrue(callerInterrupted.get()); + assertEquals(1, shutdownCalls.get()); + verify(client, times(1)) + .checkpoint( + any(), + any(), + argThat(updates -> updates.size() == 1 + && updates.get(0).type() == OperationType.EXECUTION + && updates.get(0).action() == OperationAction.SUCCEED + && updates.get(0).payload().equals('"' + result + '"'))); + } finally { + releaseEnd.countDown(); + injectFailure.set(false); + users.shutdownNow(); + callers.shutdownNow(); + assertTrue(users.awaitTermination(3, TimeUnit.SECONDS)); + assertTrue(callers.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + private static Context runtime(int millis) { + var context = mock(Context.class); + when(context.getRemainingTimeInMillis()).thenReturn(millis); + when(context.getAwsRequestId()).thenReturn("independent-request"); + return context; + } + + private static DurableExecutionInput input() { + return new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/independent/exec", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("exec") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorTerminationLeaseTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorTerminationLeaseTest.java new file mode 100644 index 0000000000..fd555989a7 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorTerminationLeaseTest.java @@ -0,0 +1,271 @@ +// 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.Instant; +import java.util.ArrayDeque; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import java.util.stream.Stream; +import org.junit.jupiter.api.Timeout; +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 software.amazon.awssdk.services.lambda.model.*; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.util.ExceptionHelper; + +@Timeout(15) +class OwnedExecutorTerminationLeaseTest { + enum Path { + GRACEFUL, + SECOND_WAIT, + FORCED + } + + static Stream confirmedCases() { + return Stream.of(Path.values()) + .flatMap(path -> Stream.of(false, true) + .flatMap(queryFatal -> + Stream.of(false, true).map(priorFatal -> Arguments.of(path, queryFatal, priorFatal)))); + } + + @ParameterizedTest + @MethodSource("confirmedCases") + void confirmedTerminationReleasesLeaseEvenWhenLaterStateQueryFails( + Path path, boolean queryFatal, boolean priorFatal) throws Exception { + var firstFatal = new InternalError("first handler fatal"); + Throwable queryFailure = queryFatal + ? new InternalError("state query fatal") + : new IllegalStateException("state query diagnostic"); + var shutdownFailure = new IllegalArgumentException("graceful shutdown failed before delegation"); + var queryEnabled = new AtomicBoolean(true); + var confirmed = new AtomicBoolean(); + var awaits = new AtomicInteger(); + var forced = new AtomicInteger(); + var queriesAfterTermination = new AtomicInteger(); + var tailEntered = new CountDownLatch(1); + var releaseTail = new CountDownLatch(path == Path.SECOND_WAIT ? 1 : 0); + var workerFatal = new AtomicReference(); + var workerEscaped = new CountDownLatch(1); + var pool = + new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue(), task -> { + var thread = new Thread(task, "termination-lease-owner"); + thread.setUncaughtExceptionHandler((ignored, failure) -> { + workerFatal.set(failure); + workerEscaped.countDown(); + }); + return thread; + }) { + @Override + public void shutdown() { + if (path == Path.FORCED && queryEnabled.get()) throw shutdownFailure; + super.shutdown(); + } + + @Override + public List shutdownNow() { + forced.incrementAndGet(); + return super.shutdownNow(); + } + + @Override + protected void afterExecute(Runnable task, Throwable failure) { + tailEntered.countDown(); + try { + assertTrue(releaseTail.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + } + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { + if (path == Path.SECOND_WAIT && awaits.incrementAndGet() == 1) { + assertTrue(tailEntered.await(3, TimeUnit.SECONDS)); + var stopped = super.awaitTermination(0, TimeUnit.NANOSECONDS); + assertFalse(stopped, "The actual worker tail is still held at the first wait"); + releaseTail.countDown(); + return stopped; + } + var stopped = super.awaitTermination(timeout, unit); + if (stopped) confirmed.set(true); + return stopped; + } + + @Override + public boolean isShutdown() { + if (queryEnabled.get() && confirmed.get()) { + queriesAfterTermination.incrementAndGet(); + ExceptionHelper.sneakyThrow(queryFailure); + } + return super.isShutdown(); + } + }; + assertEquals(1, pool.prestartAllCoreThreads()); + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withInvocationExecutorFactory(() -> pool) + .build(); + try { + var observed = assertThrows( + Throwable.class, + () -> DurableExecutor.execute( + input(), + null, + TypeToken.get(String.class), + (value, context) -> { + if (priorFatal) throw firstFatal; + return value; + }, + config)); + if (priorFatal) assertSame(firstFatal, observed); + else if (queryFatal || path != Path.FORCED) assertSame(queryFailure, observed); + else + assertSame( + shutdownFailure, + assertInstanceOf(UnrecoverableDurableExecutionException.class, observed) + .getCause()); + assertTrue(contains(observed, queryFailure), "The state-query diagnostic must remain visible"); + if (path == Path.FORCED) assertTrue(contains(observed, shutdownFailure)); + assertTrue(confirmed.get(), "Termination must have been confirmed by a real pool wait"); + assertTrue(pool.isTerminated()); + assertEquals(1, queriesAfterTermination.get(), "The later diagnostic query still runs"); + assertEquals(path == Path.GRACEFUL ? 0 : 1, forced.get()); + if (priorFatal) { + assertTrue(workerEscaped.await(3, TimeUnit.SECONDS)); + assertSame(firstFatal, workerFatal.get()); + } + + queryEnabled.set(false); + var rejected = assertThrows( + IllegalStateException.class, + () -> DurableExecutor.execute( + input(), + null, + TypeToken.get(String.class), + (value, context) -> fail("must not dispatch"), + config)); + assertEquals( + "Invocation executor factory returned a shutdown executor", + rejected.getMessage(), + "Released lease allows closed-state validation instead of the stale already-in-use rejection"); + } finally { + queryEnabled.set(false); + releaseTail.countDown(); + pool.shutdownNow(); + assertTrue(pool.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void unconfirmedOrActiveExecutorKeepsItsLease(boolean stillActive) throws Exception { + var injected = new AtomicBoolean(true); + var attempted = new AtomicBoolean(); + var confirms = new AtomicInteger(); + var queries = new AtomicInteger(); + var awaitFailure = new IllegalArgumentException("termination not observed"); + var queryFailure = new IllegalStateException("state unknown"); + var pool = new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue()) { + @Override + public void shutdown() { + attempted.set(true); + if (!injected.get() || !stillActive) super.shutdown(); + } + + @Override + public List shutdownNow() { + return injected.get() && stillActive ? List.of() : super.shutdownNow(); + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { + if (injected.get() && !stillActive) throw awaitFailure; + var stopped = super.awaitTermination(injected.get() ? 0 : timeout, unit); + if (stopped) confirms.incrementAndGet(); + return stopped; + } + + @Override + public boolean isShutdown() { + if (injected.get() && attempted.get()) { + queries.incrementAndGet(); + throw queryFailure; + } + return super.isShutdown(); + } + }; + pool.prestartAllCoreThreads(); + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withInvocationExecutorFactory(() -> pool) + .build(); + try { + var observed = assertThrows( + Throwable.class, + () -> DurableExecutor.execute( + input(), null, TypeToken.get(String.class), (value, context) -> value, config)); + assertTrue(contains(observed, queryFailure)); + if (!stillActive) assertTrue(contains(observed, awaitFailure)); + assertEquals(0, confirms.get()); + if (stillActive) assertFalse(pool.isTerminated()); + var queryCount = queries.get(); + var rejected = assertThrows( + IllegalStateException.class, + () -> DurableExecutor.execute( + input(), + null, + TypeToken.get(String.class), + (value, context) -> fail("must not dispatch"), + config)); + assertEquals("Invocation executor factory returned an executor already in use", rejected.getMessage()); + assertEquals(queryCount, queries.get(), "An existing lease must be rejected before any state callback"); + } finally { + injected.set(false); + pool.shutdownNow(); + assertTrue(pool.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + private static boolean contains(Throwable root, Throwable target) { + var pending = new ArrayDeque(); + var seen = Collections.newSetFromMap(new IdentityHashMap()); + pending.add(root); + while (!pending.isEmpty()) { + var next = pending.remove(); + if (next == target) return true; + if (!seen.add(next)) continue; + if (next.getCause() != null) pending.add(next.getCause()); + Collections.addAll(pending, next.getSuppressed()); + } + return false; + } + + private static DurableExecutionInput input() { + return new DurableExecutionInput( + "arn:execution/lease", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("lease") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorWrappedFailureTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorWrappedFailureTest.java new file mode 100644 index 0000000000..a8fadaafe9 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedExecutorWrappedFailureTest.java @@ -0,0 +1,130 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.execution; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import com.amazonaws.services.lambda.runtime.Context; +import java.time.Instant; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import software.amazon.awssdk.services.lambda.model.*; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.plugin.*; +import software.amazon.lambda.durable.util.ExceptionHelper; + +@Timeout(15) +class OwnedExecutorWrappedFailureTest { + @ParameterizedTest + @CsvSource({"false, false, true", "true, false, true", "false, true, true", "true, true, true", "true, false, false" + }) + void wrappedOwnedShutdownFailureRetainsIdentity(boolean executionWrapper, boolean acknowledged, boolean fatalCause) + throws Exception { + Throwable cause = fatalCause + ? new InternalError("wrapped owned shutdown fatal") + : new IllegalStateException("ordinary application failure"); + Throwable wrapped = executionWrapper ? new ExecutionException(cause) : new CompletionException(cause); + var injectFailure = new AtomicBoolean(true); + var shutdownCalls = new AtomicInteger(); + var users = new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>()) { + @Override + public void shutdown() { + shutdownCalls.incrementAndGet(); + super.shutdown(); + if (injectFailure.get()) ExceptionHelper.sneakyThrow(wrapped); + } + }; + var callers = Executors.newSingleThreadExecutor(); + var caller = new AtomicReference(); + var callerInterrupted = new AtomicBoolean(); + var endInfo = new AtomicReference(); + var endEntered = new CountDownLatch(1); + var releaseEnd = new CountDownLatch(acknowledged ? 1 : 0); + var client = TestUtils.createMockClient(); + var result = acknowledged ? "x".repeat(6 * 1024 * 1024 + 100) : "result"; + var config = DurableConfig.builder() + .withInvocationExecutorFactory(() -> users) + .withDurableExecutionClient(client) + .withPlugins(ignored -> new DurableExecutionPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + endInfo.set(info); + endEntered.countDown(); + try { + assertTrue(releaseEnd.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + } + } + }) + .build(); + try { + var invocation = callers.submit(() -> { + caller.set(Thread.currentThread()); + try { + return DurableExecutor.execute( + input(), runtime(10_000), TypeToken.get(String.class), (value, context) -> result, config); + } finally { + callerInterrupted.set(Thread.interrupted()); + } + }); + assertTrue(endEntered.await(5, TimeUnit.SECONDS)); + assertEquals(InvocationStatus.SUCCEEDED, endInfo.get().invocationStatus()); + assertNull(endInfo.get().executionError()); + if (acknowledged) caller.get().interrupt(); + assertSame( + fatalCause ? cause : wrapped, + assertThrows(ExecutionException.class, () -> invocation.get(5, TimeUnit.SECONDS)) + .getCause()); + assertEquals(acknowledged, callerInterrupted.get()); + assertEquals(1, shutdownCalls.get()); + if (acknowledged) + verify(client, times(1)) + .checkpoint( + any(), + any(), + argThat(updates -> updates.size() == 1 + && updates.get(0).type() == OperationType.EXECUTION + && updates.get(0).action() == OperationAction.SUCCEED + && updates.get(0).payload().equals('"' + result + '"'))); + } finally { + releaseEnd.countDown(); + injectFailure.set(false); + users.shutdownNow(); + callers.shutdownNow(); + assertTrue(users.awaitTermination(3, TimeUnit.SECONDS)); + assertTrue(callers.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + private static Context runtime(int millis) { + var context = mock(Context.class); + when(context.getRemainingTimeInMillis()).thenReturn(millis); + when(context.getAwsRequestId()).thenReturn("independent-request"); + return context; + } + + private static DurableExecutionInput input() { + return new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/independent/exec", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("exec") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedInvocationExecutorLifecycleTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedInvocationExecutorLifecycleTest.java new file mode 100644 index 0000000000..5a511a1f4c --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/OwnedInvocationExecutorLifecycleTest.java @@ -0,0 +1,286 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.execution; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.amazonaws.services.lambda.runtime.Context; +import java.time.Instant; +import java.util.List; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.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.CheckpointUpdatedExecutionState; +import software.amazon.awssdk.services.lambda.model.ExecutionDetails; +import software.amazon.awssdk.services.lambda.model.Operation; +import software.amazon.awssdk.services.lambda.model.OperationStatus; +import software.amazon.awssdk.services.lambda.model.OperationType; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; +import software.amazon.lambda.durable.plugin.InvocationEndInfo; +import software.amazon.lambda.durable.plugin.InvocationInfo; + +class OwnedInvocationExecutorLifecycleTest { + @Test + void callerShutsDownOwnedExecutorAfterHandlerEndReturns() throws Exception { + var ownedThreads = ConcurrentHashMap.newKeySet(); + var owned = Executors.newFixedThreadPool(2, action -> { + var thread = new Thread(action, "invocation-owned"); + ownedThreads.add(thread); + return thread; + }); + var runtime = Executors.newSingleThreadExecutor(); + var startThread = new AtomicReference(); + var stepThread = new AtomicReference(); + var endThread = new AtomicReference(); + var callerThread = new AtomicReference(); + var endEntered = new CountDownLatch(1); + var releaseEnd = new CountDownLatch(1); + var plugin = new DurableExecutionPlugin() { + @Override + public void onInvocationStart(InvocationInfo info) { + startThread.set(Thread.currentThread()); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + endThread.set(Thread.currentThread()); + endEntered.countDown(); + await(releaseEnd); + } + }; + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withInvocationExecutorFactory(() -> owned) + .withPlugins(ignoredInfo -> plugin) + .build(); + try { + var response = runtime.submit(() -> { + callerThread.set(Thread.currentThread()); + return DurableExecutor.execute( + input(), + context(), + TypeToken.get(String.class), + (value, durable) -> durable.step("step", String.class, step -> { + stepThread.set(Thread.currentThread()); + return value; + }), + config); + }); + assertTrue(endEntered.await(5, TimeUnit.SECONDS)); + assertSame(startThread.get(), endThread.get()); + assertTrue(ownedThreads.contains(startThread.get()), "Plugin-enabled roots must use the owned executor"); + assertTrue(ownedThreads.contains(stepThread.get()), "Operation wrappers must use the owned executor"); + assertNotSame(callerThread.get(), endThread.get()); + assertFalse(owned.isShutdown(), "The handler-owned End hook must run before executor shutdown"); + assertFalse(response.isDone(), "The runtime caller must wait for End"); + releaseEnd.countDown(); + assertEquals( + ExecutionStatus.SUCCEEDED, response.get(5, TimeUnit.SECONDS).status()); + assertTrue(owned.isTerminated(), "Executor shutdown must not wait for its own root worker"); + } finally { + releaseEnd.countDown(); + stop(runtime); + stop(owned); + } + } + + @Test + void interruptedCallerStillWaitsForEndBeforeShuttingDownOwnedExecutor() throws Exception { + var owned = Executors.newFixedThreadPool(2); + var runtime = Executors.newSingleThreadExecutor(); + var callerThread = new AtomicReference(); + var callerInterrupted = new AtomicBoolean(); + var endEntered = new CountDownLatch(1); + var endInterrupted = new CountDownLatch(1); + var releaseEnd = new CountDownLatch(1); + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withInvocationExecutorFactory(() -> owned) + .withPlugins(ignoredInfo -> new DurableExecutionPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + endEntered.countDown(); + awaitIgnoringInterruptions(releaseEnd, endInterrupted); + } + }) + .build(); + try { + var response = runtime.submit(() -> { + callerThread.set(Thread.currentThread()); + try { + return DurableExecutor.execute( + input(), context(), TypeToken.get(String.class), (value, durable) -> value, config); + } finally { + callerInterrupted.set(Thread.interrupted()); + } + }); + assertTrue(endEntered.await(5, TimeUnit.SECONDS)); + callerThread.get().interrupt(); + assertTrue(endInterrupted.await(5, TimeUnit.SECONDS)); + assertThrows(TimeoutException.class, () -> response.get(1200, TimeUnit.MILLISECONDS)); + assertFalse(owned.isShutdown(), "Cancellation cannot move executor shutdown ahead of actual End exit"); + releaseEnd.countDown(); + var failure = assertInstanceOf( + UnrecoverableDurableExecutionException.class, + assertThrows(ExecutionException.class, () -> response.get(5, TimeUnit.SECONDS)) + .getCause()); + assertTrue(failure.isRetryable()); + assertInstanceOf(InterruptedException.class, failure.getCause()); + assertTrue(callerInterrupted.get()); + assertTrue(owned.isTerminated()); + } finally { + releaseEnd.countDown(); + stop(runtime); + stop(owned); + } + } + + private static void awaitIgnoringInterruptions(CountDownLatch release, CountDownLatch interrupted) { + var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10); + while (true) { + try { + assertTrue(release.await(Math.max(0, deadline - System.nanoTime()), TimeUnit.NANOSECONDS)); + return; + } catch (InterruptedException ignored) { + interrupted.countDown(); + } + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void rootStepAndNestedChildUseTheInvocationSelectedExecutor(boolean pluginsEnabled) throws Exception { + var ownedThreads = ConcurrentHashMap.newKeySet(); + var observedThreads = new CopyOnWriteArrayList(); + // Root, child and nested step can wait on each other; give this scoped routing test enough capacity. + var owned = Executors.newFixedThreadPool(4, action -> { + var thread = new Thread(action, "routing-owned"); + ownedThreads.add(thread); + return thread; + }); + var builder = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withInvocationExecutorFactory(() -> owned); + if (pluginsEnabled) builder.withPlugins(ignoredInfo -> new DurableExecutionPlugin() {}); + try { + var response = DurableExecutor.execute( + input(), + context(), + TypeToken.get(String.class), + (value, durable) -> { + observedThreads.add(Thread.currentThread()); + durable.step("root-step", String.class, step -> { + observedThreads.add(Thread.currentThread()); + return value; + }); + return durable.runInChildContext("child", String.class, child -> { + observedThreads.add(Thread.currentThread()); + return child.step("nested-step", String.class, step -> { + observedThreads.add(Thread.currentThread()); + return value; + }); + }); + }, + builder.build()); + assertEquals(ExecutionStatus.SUCCEEDED, response.status()); + assertEquals("\"input\"", response.result()); + assertEquals(4, observedThreads.size()); + assertTrue( + ownedThreads.containsAll(observedThreads), + "Root, step, child and nested operation code must all use the invocation-selected executor"); + assertTrue(owned.isTerminated()); + } finally { + stop(owned); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void rejectedRootStillShutsDownInvocationOwnedExecutor(boolean pluginsEnabled) throws Exception { + var owned = new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>()) { + @Override + public void execute(Runnable command) { + throw new RejectedExecutionException("test root submission rejection"); + } + }; + var builder = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withInvocationExecutorFactory(() -> owned); + if (pluginsEnabled) builder.withPlugins(ignoredInfo -> new DurableExecutionPlugin() {}); + var config = builder.build(); + try { + assertThrows( + RuntimeException.class, + () -> DurableExecutor.execute( + input(), context(), TypeToken.get(String.class), (value, durable) -> value, config)); + assertTrue(owned.isTerminated()); + } finally { + stop(owned); + } + } + + private static DurableExecutionInput input() { + var operation = Operation.builder() + .id("execution") + .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/test/execution", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(List.of(operation)) + .build()); + } + + private static Context context() { + var context = mock(Context.class); + when(context.getAwsRequestId()).thenReturn("owned-executor-request"); + when(context.getRemainingTimeInMillis()).thenReturn(20_000); + return context; + } + + private static void await(CountDownLatch latch) { + try { + if (!latch.await(5, TimeUnit.SECONDS)) throw new IllegalStateException("Test gate timed out"); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(interrupted); + } + } + + private static void stop(ExecutorService executor) throws InterruptedException { + executor.shutdownNow(); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/RecordedFatalCleanupTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/RecordedFatalCleanupTest.java new file mode 100644 index 0000000000..f2d475855f --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/RecordedFatalCleanupTest.java @@ -0,0 +1,220 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.execution; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import java.time.Instant; +import java.util.ArrayDeque; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.Set; +import java.util.concurrent.*; +import java.util.concurrent.atomic.*; +import java.util.stream.Stream; +import org.junit.jupiter.api.Timeout; +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.*; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.plugin.*; +import software.amazon.lambda.durable.serde.SerDes; +import software.amazon.lambda.durable.util.ExceptionHelper; + +@Timeout(15) +class RecordedFatalCleanupTest { + enum Cleanup { + NONE, + ORDINARY, + FATAL, + WRAPPED_FATAL, + SAME_FATAL, + WRAPPED_SAME_FATAL, + WRAPPER_DIAGNOSTICS + } + + static Stream cases() { + return Stream.of(Cleanup.values()) + .flatMap(cleanup -> Stream.of(false, true) + .flatMap( + plugins -> Stream.of(false, true).map(death -> Arguments.of(cleanup, plugins, death)))); + } + + @SuppressWarnings("removal") + @ParameterizedTest + @MethodSource("cases") + void recordedSerializerFatalOwnsCleanupDiagnosticsWithoutCycles(Cleanup cleanup, boolean plugins, boolean death) + throws Exception { + Error first = death ? new ThreadDeath() : new InternalError("fatal output serialization"); + var ordinary = new IllegalStateException("ordinary owned shutdown failure"); + var laterFatal = new InternalError("later owned shutdown fatal"); + var wrapperDiagnostic = new IllegalArgumentException("transport-wrapper diagnostic"); + var nestedDiagnostic = new IllegalArgumentException("nested observation diagnostic"); + Throwable shutdownFailure = switch (cleanup) { + case NONE -> null; + case ORDINARY, WRAPPER_DIAGNOSTICS -> ordinary; + case FATAL -> laterFatal; + case WRAPPED_FATAL -> new CompletionException(laterFatal); + case SAME_FATAL -> first; + case WRAPPED_SAME_FATAL -> new CompletionException(first); + }; + var bodies = new AtomicInteger(); + var serialized = new AtomicInteger(); + var shutdowns = new AtomicInteger(); + var ends = new AtomicInteger(); + var end = new AtomicReference(); + var manager = new AtomicReference(); + var fatalAtShutdown = new AtomicReference(); + var workerFatal = new AtomicReference(); + var workerEscaped = new CountDownLatch(1); + var pool = + new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue(), task -> { + var worker = new Thread(task, "recorded-fatal-owner"); + worker.setUncaughtExceptionHandler((ignored, failure) -> { + workerFatal.set(failure); + workerEscaped.countDown(); + }); + return worker; + }) { + @Override + public void shutdown() { + shutdowns.incrementAndGet(); + fatalAtShutdown.set(manager.get().getInvocationFatal()); + super.shutdown(); + if (shutdownFailure != null) ExceptionHelper.sneakyThrow(shutdownFailure); + } + }; + pool.prestartAllCoreThreads(); + var builder = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withInvocationExecutorFactory(() -> pool) + .withSerDes(new SerDes() { + public String serialize(Object value) { + serialized.incrementAndGet(); + throw first; + } + + @SuppressWarnings("unchecked") + public T deserialize(String value, TypeToken type) { + return (T) "input"; + } + }); + if (plugins) + builder.withPlugins(ignored -> new DurableExecutionPlugin() { + public void onInvocationEnd(InvocationEndInfo info) { + ends.incrementAndGet(); + end.set(info); + } + }); + try { + var observed = assertThrows( + Throwable.class, + () -> DurableExecutor.execute( + input(), + null, + TypeToken.get(String.class), + (value, context) -> { + manager.set(((DurableContextImpl) context).getExecutionManager()); + if (cleanup == Cleanup.WRAPPER_DIAGNOSTICS) + assertDoesNotThrow(() -> decorateShutdownObservation( + manager.get(), first, wrapperDiagnostic, nestedDiagnostic)); + bodies.incrementAndGet(); + return "output"; + }, + builder.build())); + assertSame(first, observed, "The first recorded fatal must remain the caller outcome"); + assertTrue(workerEscaped.await(3, TimeUnit.SECONDS)); + assertSame(first, workerFatal.get()); + assertSame(first, fatalAtShutdown.get(), "This fatal was recorded before owned executor cleanup"); + assertEquals(1, bodies.get()); + assertEquals(1, serialized.get()); + assertEquals(1, shutdowns.get()); + assertTrue(pool.isTerminated()); + assertEquals(plugins ? 1 : 0, ends.get()); + if (plugins) { + assertEquals(InvocationStatus.RETRYING, end.get().invocationStatus()); + assertSame(first, end.get().executionError()); + } + if (cleanup == Cleanup.ORDINARY || cleanup == Cleanup.WRAPPER_DIAGNOSTICS) + assertTrue(contains(first, ordinary), "Retain the ordinary shutdown diagnostic on the recorded fatal"); + if (cleanup == Cleanup.WRAPPER_DIAGNOSTICS) { + assertTrue(contains(first, wrapperDiagnostic)); + assertTrue(contains(first, nestedDiagnostic)); + } + if (cleanup == Cleanup.FATAL || cleanup == Cleanup.WRAPPED_FATAL) + assertTrue(contains(first, laterFatal), "Retain the later cleanup fatal"); + assertAcyclic(first, Collections.newSetFromMap(new IdentityHashMap())); + } finally { + pool.shutdownNow(); + assertTrue(pool.awaitTermination(3, TimeUnit.SECONDS)); + } + } + + private static void decorateShutdownObservation( + ExecutionManager manager, Error fatal, Throwable diagnostic, Throwable nestedDiagnostic) + throws ReflectiveOperationException { + var field = ExecutionManager.class.getDeclaredField("checkpointManager"); + field.setAccessible(true); + var checkpointManager = spy((CheckpointManager) field.get(manager)); + doAnswer(call -> { + try { + call.callRealMethod(); + return null; + } catch (CompletionException observation) { + assertSame(fatal, observation.getCause(), "Decorate the actual shutdown observation"); + observation.addSuppressed(diagnostic); + var nested = new CompletionException(fatal); + nested.addSuppressed(nestedDiagnostic); + observation.addSuppressed(nested); + throw observation; + } + }) + .when(checkpointManager) + .shutdown(); + field.set(manager, checkpointManager); + } + + private static boolean contains(Throwable root, Throwable expected) { + var queue = new ArrayDeque(); + var seen = Collections.newSetFromMap(new IdentityHashMap()); + queue.add(root); + while (!queue.isEmpty()) { + var next = queue.remove(); + if (next == expected) return true; + if (!seen.add(next)) continue; + if (next.getCause() != null) queue.add(next.getCause()); + Collections.addAll(queue, next.getSuppressed()); + } + return false; + } + + private static void assertAcyclic(Throwable failure, Set path) { + assertTrue(path.add(failure), "Cleanup diagnostics must not point back to the selected fatal"); + if (failure.getCause() != null) assertAcyclic(failure.getCause(), path); + for (var additional : failure.getSuppressed()) assertAcyclic(additional, path); + path.remove(failure); + } + + private static DurableExecutionInput input() { + return new DurableExecutionInput( + "arn:execution/recorded-fatal", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(Operation.builder() + .id("recorded-fatal") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .startTimestamp(Instant.EPOCH) + .executionDetails(ExecutionDetails.builder() + .inputPayload("\"input\"") + .build()) + .build()) + .build()); + } +}