From 17ba43bc0f5f9234bd0c82283116d7bf8e048a0e Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 00:15:40 +0000 Subject: [PATCH 01/40] docs: design custom extension operations API --- ...8-10-custom-extension-operations-design.md | 278 ++++++++++++++++++ 1 file changed, 278 insertions(+) create mode 100644 docs/superpowers/specs/2026-08-10-custom-extension-operations-design.md diff --git a/docs/superpowers/specs/2026-08-10-custom-extension-operations-design.md b/docs/superpowers/specs/2026-08-10-custom-extension-operations-design.md new file mode 100644 index 000000000..e745491fb --- /dev/null +++ b/docs/superpowers/specs/2026-08-10-custom-extension-operations-design.md @@ -0,0 +1,278 @@ +# Custom Extension Operations Design + +## Objective + +Provide a supported public API for implementing reusable durable extension operations in a separate Maven module. +Extensions compose SDK-owned primitive operations without adding backend operation types, sending raw checkpoint +updates, or depending on SDK implementation classes. + +Application code calls ordinary static extension methods: + +```java +import static software.amazon.lambda.durable.dag.DagOperations.dag; + +var result = dag("etl", definition -> { + // Define DAG nodes. +}); +``` + +The public `DurableContext` interface remains unchanged. Its existing instance methods continue to work for backward +compatibility. + +## Operation Categories + +### Core operations + +Core operations correspond to SDK-owned primitive behavior: + +- `step` +- `wait` +- chained `invoke` +- `createCallback` +- `runInChildContext` + +`DurableOperations` exposes context-free static facades for these operations. The facades obtain the active +`DurableContext` from SDK-managed current-context storage and delegate to the existing instance methods. New static +step APIs use `StepContext` functions and do not reproduce the deprecated `Supplier` step overloads. + +### Extension operations + +Extension operations compose core operations: + +- `waitForCallback` +- `waitForCondition` +- `withRetry` +- `map` +- `parallel` +- third-party operations such as DAG + +`DurableExtensionOperations` exposes the built-in extension operations through the same static-import style. +Existing `DurableContext` instance methods and their behavior remain unchanged. + +An extension does not automatically create a child context. Each extension chooses its scope: + +- Replay-safe value helpers can create a step directly in the current context. +- Recursive invocation helpers can delegate directly to `invoke`. +- `map`, `parallel`, or `withRetry` can explicitly create child contexts when isolation is part of their semantics. +- DAG can reserve primitive identities in the current scope or explicitly create a child context if the DAG contract + requires one. + +## Extension Authoring Contract + +There is no universal `DurableExtensions.run` or `runAsync` method. An extension is an ordinary public static method +that uses the static operation facades and, when it needs stable deferred identities, the active `ExtensionContext`. + +```java +public interface ExtensionContext extends BaseContext { + static ExtensionContext getCurrentContext() { + var context = BaseContext.getCurrentContext(); + if (context instanceof ExtensionContext extensionContext) { + return extensionContext; + } + throw new IllegalStateException( + "ExtensionContext is only available from a durable handler or child-context thread"); + } + + boolean isReplaying(); + + ExtensionOperation reserve(String name); +} +``` + +SDK-managed handler and child contexts implement both `DurableContext` and `ExtensionContext`. Step contexts do not. +`ExtensionContext.getCurrentContext()` therefore succeeds only on supported handler and child-context threads. + +`ExtensionContext` exposes metadata through `BaseContext`, replay state, and deterministic primitive reservations. It +does not expose execution managers, checkpoint models, backend operation types, operation updates, or raw operation +IDs. + +## Primitive Reservations + +`ExtensionContext.reserve(name)` immediately consumes the next sequential operation ID in the active durable scope and +returns an opaque, one-shot `ExtensionOperation`. The ID remains hidden from extension code. + +```java +public interface ExtensionOperation { + DurableFuture stepAsync( + TypeToken resultType, + Function function, + StepConfig config); + + DurableFuture waitAsync(Duration duration); + + DurableFuture invokeAsync( + String functionName, + U payload, + TypeToken resultType, + InvokeConfig config); + + DurableCallbackFuture createCallback( + TypeToken resultType, + CallbackConfig config); + + DurableFuture runInChildContextAsync( + TypeToken resultType, + Function function, + RunInChildContextConfig config); + + // Class, synchronous, and default-configuration overloads are default methods. +} +``` + +The operation name is bound by `reserve` and is not repeated when selecting the primitive. A reservation can execute +exactly one primitive operation. Reuse fails with `IllegalStateException`. + +The SDK implements reservations by allocating an ID through the current context's normal `OperationIdGenerator`. +Package-private explicit-ID variants of primitive creation methods consume the reserved ID. These internal methods are +not part of the extension API. + +## DAG Usage + +A DAG module uses only public SDK contracts: + +```java +public static DagResult dag(String name, Consumer register) { + var extension = ExtensionContext.getCurrentContext(); + var dag = new DagContext(name, extension); + register.accept(dag); + return dag.execute(); +} +``` + +During its deterministic definition phase, the DAG reserves primitive positions: + +```java +var extract = extension.reserve("extract"); +var transform = extension.reserve("transform"); +var load = extension.reserve("load"); +``` + +The scheduler can later execute those reservations in any dependency-valid order: + +```java +var transformFuture = transform.stepAsync(String.class, step -> transformData()); +var extractFuture = extract.runInChildContextAsync( + ExtractResult.class, + child -> executeExtraction()); +``` + +Registration order determines IDs; launch order does not. This supports graph scheduling without name-derived IDs or +public explicit-ID APIs. + +The production DAG module is outside this issue. A small extension fixture in a separate repository Maven module +proves that an external module can compile and execute using only the supported contracts. + +## Current Context + +The SDK binds and restores current context around: + +- the durable handler +- child-context functions +- step functions +- wait-for-condition check functions + +`DurableContext.getCurrentContext()` continues to use its existing signature. Its failure behavior is clarified: + +- Handler or child-context thread: returns the active durable context. +- Step thread: throws `IllegalStateException` directing callers to `StepContext`. +- Unsupported or application-created thread: throws `IllegalStateException` explaining that no durable context is + active. + +`ExtensionContext.getCurrentContext()` returns the active extension-capable handler or child context. It throws a +clear `IllegalStateException` from step threads and unsupported threads. + +Current context is not propagated to application-created threads. Extensions must create durable primitives on +SDK-managed durable context threads. + +## Static Operation Facades + +`DurableOperations` contains only core operations. Each method obtains the current durable context internally. Its +child-context static methods use callbacks that do not require callers to receive a `DurableContext`; code inside the +callback can use static operations or `ExtensionContext.getCurrentContext()`. + +`DurableExtensionOperations` contains built-in composed operations. The initial implementation delegates to the +existing `DurableContext` methods to preserve behavior. The classification and static API do not require rewriting +each established operation implementation in this change. + +Both facade classes are stateless utility classes. Calling either facade outside a supported durable context produces +the same clear failure as `DurableContext.getCurrentContext()`. + +## Durable Futures + +Asynchronous extension methods may return SDK operation futures or custom composed `DurableFuture` implementations. +`DurableFuture` therefore exposes a public non-mutating completion signal: + +```java +default CompletableFuture completionFuture() { + throw new UnsupportedOperationException( + "This DurableFuture does not expose a completion signal"); +} +``` + +SDK operation implementations return a derived completion future whose completion or cancellation cannot mutate the +durable operation. `DurableFuture.anyOf` uses this public contract instead of downcasting to +`BaseDurableOperation`. Custom futures that support `anyOf` override `completionFuture()`. + +## Replay and Compatibility + +Reservations must be created in the same deterministic order on every replay. Reordering, inserting, or removing +reservations can associate existing checkpoints with different logical primitives and is a workflow compatibility +change. After registration, executing reservations in a different order is supported. + +Direct static core calls allocate IDs when invoked, matching existing `DurableContext` semantics. They are appropriate +when call order is deterministic. + +Nested extension calls execute in the active scope unless an extension explicitly creates a child context. No +extension-specific recursion limit is introduced. + +Public compatibility guarantees apply to: + +- `DurableOperations` +- `DurableExtensionOperations` +- `ExtensionContext` +- `ExtensionOperation` +- `DurableFuture.completionFuture()` + +Compatible SDK releases may add new default overloads or new primitive capabilities. Existing reservation ordering, +one-shot behavior, and primitive semantics change only in a breaking release. + +## Plugins and Failures + +There is no automatic extension lifecycle boundary because an extension is an ordinary composition method. Plugins +observe every primitive created by the extension. If the extension explicitly creates a child context, plugins also +observe that child-context operation. + +Primitive serialization, exception, suspension, cancellation, retry, and checkpoint behavior remain owned by the +existing primitive implementation. Extension code cannot send checkpoint updates or define backend operation +subtypes. + +Invalid names and null arguments use the existing SDK validators. Reusing a reservation or requesting current context +from an unsupported thread fails before creating a primitive. + +## Verification + +Unit tests cover: + +- current durable, step, and extension context lookup +- restoration of nested current-context bindings +- deterministic reservation allocation +- out-of-order reservation execution +- one-shot reservation enforcement +- each reserved primitive delegation path +- custom `DurableFuture` participation in `anyOf` +- static facade failure outside a durable context + +Integration tests in `sdk-integration-tests` cover: + +- an extension fixture compiled in a separate Maven module +- initial execution and replay +- suspension and resume +- reservations launched in different orders across replays +- nested extensions in the same scope +- extensions that explicitly create child contexts +- static core and built-in extension facades +- primitive plugin lifecycle events + +Formatting runs through `mvn spotless:apply`. Verification starts with focused SDK and integration tests, then expands +to the full reactor because the change affects public APIs, execution context propagation, replay identity, and +durable futures. From 3cb3728a1ebdd87693c59223289e533fb736a183 Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 00:21:46 +0000 Subject: [PATCH 02/40] docs: split extension operation facades --- ...8-10-custom-extension-operations-design.md | 46 +++++++++++++------ 1 file changed, 33 insertions(+), 13 deletions(-) diff --git a/docs/superpowers/specs/2026-08-10-custom-extension-operations-design.md b/docs/superpowers/specs/2026-08-10-custom-extension-operations-design.md index e745491fb..1c138abff 100644 --- a/docs/superpowers/specs/2026-08-10-custom-extension-operations-design.md +++ b/docs/superpowers/specs/2026-08-10-custom-extension-operations-design.md @@ -31,7 +31,7 @@ Core operations correspond to SDK-owned primitive behavior: - `createCallback` - `runInChildContext` -`DurableOperations` exposes context-free static facades for these operations. The facades obtain the active +`DurableCoreOperations` exposes context-free static facades for these operations. The facades obtain the active `DurableContext` from SDK-managed current-context storage and delegate to the existing instance methods. New static step APIs use `StepContext` functions and do not reproduce the deprecated `Supplier` step overloads. @@ -46,8 +46,16 @@ Extension operations compose core operations: - `parallel` - third-party operations such as DAG -`DurableExtensionOperations` exposes the built-in extension operations through the same static-import style. -Existing `DurableContext` instance methods and their behavior remain unchanged. +Each built-in extension has an independently maintained static facade: + +- `DurableMapOperations` +- `DurableParallelOperations` +- `DurableWaitForCallbackOperations` +- `DurableWaitForConditionOperations` +- `DurableWithRetryOperations` + +Each class owns only its operation's overloads, tests, and documentation. Existing `DurableContext` instance methods +and their behavior remain unchanged. An extension does not automatically create a child context. Each extension chooses its scope: @@ -186,16 +194,24 @@ SDK-managed durable context threads. ## Static Operation Facades -`DurableOperations` contains only core operations. Each method obtains the current durable context internally. Its -child-context static methods use callbacks that do not require callers to receive a `DurableContext`; code inside the -callback can use static operations or `ExtensionContext.getCurrentContext()`. +`DurableCoreOperations` contains only core operations. Each method obtains the current durable context internally. +Its child-context static methods use callbacks that do not require callers to receive a `DurableContext`; code inside +the callback can use static operations or `ExtensionContext.getCurrentContext()`. + +Each built-in extension facade contains only one operation family: -`DurableExtensionOperations` contains built-in composed operations. The initial implementation delegates to the -existing `DurableContext` methods to preserve behavior. The classification and static API do not require rewriting -each established operation implementation in this change. +| Facade | Methods | +| --- | --- | +| `DurableMapOperations` | `map`, `mapAsync` | +| `DurableParallelOperations` | `parallel` and its branch-building API | +| `DurableWaitForCallbackOperations` | `waitForCallback`, `waitForCallbackAsync` | +| `DurableWaitForConditionOperations` | `waitForCondition`, `waitForConditionAsync` | +| `DurableWithRetryOperations` | `withRetry`, `withRetryAsync` | -Both facade classes are stateless utility classes. Calling either facade outside a supported durable context produces -the same clear failure as `DurableContext.getCurrentContext()`. +The initial implementations delegate to the existing `DurableContext` methods to preserve behavior. The +classification and static API do not require rewriting each established operation implementation in this change. +Each facade is a stateless utility class. Calling a facade outside a supported durable context produces the same clear +failure as `DurableContext.getCurrentContext()`. ## Durable Futures @@ -227,8 +243,12 @@ extension-specific recursion limit is introduced. Public compatibility guarantees apply to: -- `DurableOperations` -- `DurableExtensionOperations` +- `DurableCoreOperations` +- `DurableMapOperations` +- `DurableParallelOperations` +- `DurableWaitForCallbackOperations` +- `DurableWaitForConditionOperations` +- `DurableWithRetryOperations` - `ExtensionContext` - `ExtensionOperation` - `DurableFuture.completionFuture()` From 8c4b3eef41cf0d929e6579d9ce52352ce1f515df Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 00:41:08 +0000 Subject: [PATCH 03/40] docs: define TLS-only extension callbacks --- ...8-10-custom-extension-operations-design.md | 135 ++++++++++++++++-- 1 file changed, 123 insertions(+), 12 deletions(-) diff --git a/docs/superpowers/specs/2026-08-10-custom-extension-operations-design.md b/docs/superpowers/specs/2026-08-10-custom-extension-operations-design.md index 1c138abff..7d0291a0b 100644 --- a/docs/superpowers/specs/2026-08-10-custom-extension-operations-design.md +++ b/docs/superpowers/specs/2026-08-10-custom-extension-operations-design.md @@ -11,8 +11,9 @@ Application code calls ordinary static extension methods: ```java import static software.amazon.lambda.durable.dag.DagOperations.dag; -var result = dag("etl", definition -> { - // Define DAG nodes. +var result = dag("etl", () -> { + var dag = DagContext.getCurrentContext(); + // Define DAG nodes through the current DAG context. }); ``` @@ -32,8 +33,9 @@ Core operations correspond to SDK-owned primitive behavior: - `runInChildContext` `DurableCoreOperations` exposes context-free static facades for these operations. The facades obtain the active -`DurableContext` from SDK-managed current-context storage and delegate to the existing instance methods. New static -step APIs use `StepContext` functions and do not reproduce the deprecated `Supplier` step overloads. +`DurableContext` from SDK-managed current-context storage and delegate to the existing instance methods. User +functions in the new APIs do not receive SDK context objects. For example, static step methods accept `Supplier`; +step code obtains `StepContext` through `StepContext.getCurrentContext()`. ### Extension operations @@ -103,7 +105,7 @@ returns an opaque, one-shot `ExtensionOperation`. The ID remains hidden from ext public interface ExtensionOperation { DurableFuture stepAsync( TypeToken resultType, - Function function, + Supplier function, StepConfig config); DurableFuture waitAsync(Duration duration); @@ -120,7 +122,7 @@ public interface ExtensionOperation { DurableFuture runInChildContextAsync( TypeToken resultType, - Function function, + Supplier function, RunInChildContextConfig config); // Class, synchronous, and default-configuration overloads are default methods. @@ -139,10 +141,12 @@ not part of the extension API. A DAG module uses only public SDK contracts: ```java -public static DagResult dag(String name, Consumer register) { +public static DagResult dag(String name, Runnable register) { var extension = ExtensionContext.getCurrentContext(); var dag = new DagContext(name, extension); - register.accept(dag); + try (var ignored = DagContext.attach(dag)) { + register.run(); + } return dag.execute(); } ``` @@ -158,10 +162,10 @@ var load = extension.reserve("load"); The scheduler can later execute those reservations in any dependency-valid order: ```java -var transformFuture = transform.stepAsync(String.class, step -> transformData()); +var transformFuture = transform.stepAsync(String.class, () -> transformData()); var extractFuture = extract.runInChildContextAsync( ExtractResult.class, - child -> executeExtraction()); + () -> executeExtraction()); ``` Registration order determines IDs; launch order does not. This supports graph scheduling without name-derived IDs or @@ -178,6 +182,9 @@ The SDK binds and restores current context around: - child-context functions - step functions - wait-for-condition check functions +- map item functions +- wait-for-callback submitters +- with-retry bodies `DurableContext.getCurrentContext()` continues to use its existing signature. Its failure behavior is clarified: @@ -192,11 +199,110 @@ clear `IllegalStateException` from step threads and unsupported threads. Current context is not propagated to application-created threads. Extensions must create durable primitives on SDK-managed durable context threads. +## User Function Signatures + +New static APIs and extension reservations never pass SDK-created context or metadata values as user-function +arguments. User functions receive only values supplied by the application or values from the application's durable +data flow. + +Examples: + +```java +var result = DurableCoreOperations.step( + "process", + Result.class, + () -> { + var step = StepContext.getCurrentContext(); + return process(step.getAttempt()); + }); +``` + +```java +var result = DurableMapOperations.map( + "process", + items, + Result.class, + item -> { + var mapItem = MapItemContext.getCurrentContext(); + return process(item, mapItem.getIndex()); + }); +``` + +```java +var result = DurableWaitForCallbackOperations.waitForCallback( + "approval", + Approval.class, + () -> { + var callback = WaitForCallbackContext.getCurrentContext(); + submitApproval(callback.getCallbackId()); + }); +``` + +```java +var result = DurableWithRetryOperations.withRetry( + "transaction", + () -> { + var retry = WithRetryContext.getCurrentContext(); + return executeAttempt(retry.getAttempt()); + }); +``` + +The new callback shapes are: + +- step and child-context functions: `Supplier` +- map item functions: `Function`; `MapItemContext` exposes the item index +- parallel branch functions: `Supplier` +- wait-for-callback submitters: `Runnable`; `WaitForCallbackContext` exposes the callback ID +- wait-for-condition checks: receive only the durable state value; attempt metadata is available from + `StepContext.getCurrentContext()` +- with-retry bodies: `Supplier`; `WithRetryContext` exposes the attempt number +- extension child-context reservations: `Supplier`; `ExtensionContext` is obtained through TLS + +`MapItemContext`, `WaitForCallbackContext`, `WithRetryContext`, and any equivalent operation-specific context provide +`getCurrentContext()` static accessors. Each accessor fails clearly outside its matching user-function scope. These +operation-specific contexts are bound in addition to the base durable or step context so static core operations +continue to resolve the active `DurableContext` or `StepContext`. + +The initial operation-specific metadata contracts are: + +```java +public interface MapItemContext { + static MapItemContext getCurrentContext() { + return OperationContextStorage.get(MapItemContext.class); + } + + int getIndex(); +} + +public interface WaitForCallbackContext { + static WaitForCallbackContext getCurrentContext() { + return OperationContextStorage.get(WaitForCallbackContext.class); + } + + String getCallbackId(); +} + +public interface WithRetryContext { + static WithRetryContext getCurrentContext() { + return OperationContextStorage.get(WithRetryContext.class); + } + + int getAttempt(); +} +``` + +Each context uses a scoped SDK-managed `ThreadLocal`. Entering a nested operation stores the previous value, and +closing the scope restores it. The thread-local value is removed when no previous value exists. +`OperationContextStorage` is a package-private SDK implementation detail. + +Existing context-accepting functions on `DurableContext`, including `Function`, +`Function`, and existing map/retry callback types, remain unchanged for backward compatibility. + ## Static Operation Facades `DurableCoreOperations` contains only core operations. Each method obtains the current durable context internally. -Its child-context static methods use callbacks that do not require callers to receive a `DurableContext`; code inside -the callback can use static operations or `ExtensionContext.getCurrentContext()`. +Its step and child-context static methods accept context-free suppliers. Code inside those callbacks uses typed +current-context accessors when it needs SDK metadata. Each built-in extension facade contains only one operation family: @@ -251,6 +357,9 @@ Public compatibility guarantees apply to: - `DurableWithRetryOperations` - `ExtensionContext` - `ExtensionOperation` +- `MapItemContext` +- `WaitForCallbackContext` +- `WithRetryContext` - `DurableFuture.completionFuture()` Compatible SDK releases may add new default overloads or new primitive capabilities. Existing reservation ordering, @@ -274,6 +383,7 @@ from an unsupported thread fails before creating a primitive. Unit tests cover: - current durable, step, and extension context lookup +- operation-specific context lookup and scope validation - restoration of nested current-context bindings - deterministic reservation allocation - out-of-order reservation execution @@ -291,6 +401,7 @@ Integration tests in `sdk-integration-tests` cover: - nested extensions in the same scope - extensions that explicitly create child contexts - static core and built-in extension facades +- context-free user functions with TLS-based metadata access - primitive plugin lifecycle events Formatting runs through `mvn spotless:apply`. Verification starts with focused SDK and integration tests, then expands From 0aaaacee94272e95ef48439d349ad5f385e65684 Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 00:48:53 +0000 Subject: [PATCH 04/40] docs: add custom extension operations ADR --- docs/adr/006-custom-extension-operations.md | 280 ++++++++++++++++++++ 1 file changed, 280 insertions(+) create mode 100644 docs/adr/006-custom-extension-operations.md diff --git a/docs/adr/006-custom-extension-operations.md b/docs/adr/006-custom-extension-operations.md new file mode 100644 index 000000000..32c1186ac --- /dev/null +++ b/docs/adr/006-custom-extension-operations.md @@ -0,0 +1,280 @@ +# ADR-006: Public API for Custom Extension Operations + +**Status:** Proposed +**Date:** 2026-08-10 + +## Context + +Issue [#571](https://github.com/aws/aws-durable-execution-sdk-java/issues/571) requests a supported way to +implement reusable durable operations in separate Maven modules without changing or rebuilding the core SDK. + +The SDK currently exposes all operations as instance methods on `DurableContext`. This creates several constraints: + +- Adding an optional operation requires changing the core `DurableContext` interface and `DurableContextImpl`. +- Extension libraries cannot create primitive operations with stable identities when registration order and execution + order differ. +- Extension code would need SDK implementation classes or public explicit-operation-ID methods to implement schedulers + such as DAG. +- Existing user-function APIs receive SDK-created context and metadata parameters, coupling new APIs to callback + signatures instead of the SDK-managed current context. +- A single facade containing every built-in extension would couple unrelated operation families and make independent + maintenance difficult. + +Extension operations do not share one execution scope. Some extensions are direct primitive wrappers in the current +scope, while others deliberately create child contexts. The extension mechanism must not impose a child-context +boundary. + +The SDK must continue to own primitive operation IDs, checkpointing, replay, suspension, serialization, failures, and +backend communication. This decision does not make backend operation types extensible. + +The detailed API specification is in +[Custom Extension Operations Design](../superpowers/specs/2026-08-10-custom-extension-operations-design.md). + +## Decision + +### Preserve DurableContext + +Keep the public `DurableContext` interface unchanged. Its existing instance methods and context-accepting callback +types remain supported for backward compatibility. + +SDK-managed handler and child contexts additionally implement a new public `ExtensionContext` interface. Step +contexts do not implement it. + +```java +public interface ExtensionContext extends BaseContext { + static ExtensionContext getCurrentContext() { + var context = BaseContext.getCurrentContext(); + if (context instanceof ExtensionContext extensionContext) { + return extensionContext; + } + throw new IllegalStateException( + "ExtensionContext is only available from a durable handler or child-context thread"); + } + + boolean isReplaying(); + + ExtensionOperation reserve(String name); +} +``` + +Extension libraries expose ordinary static methods. There is no required registration mechanism and no universal +`DurableExtensions.run` method. + +```java +public final class DagOperations { + public static DagResult dag(String name, Runnable definition) { + var extension = ExtensionContext.getCurrentContext(); + return executeDag(name, extension, definition); + } +} +``` + +An extension decides whether to execute in the current scope or explicitly create a child context. + +### Separate Core and Extension Facades + +Expose primitive operations through `DurableCoreOperations`: + +- `step` +- `wait` +- chained `invoke` +- `createCallback` +- `runInChildContext` + +Expose each built-in extension family through an independently maintained class: + +| Facade | Operation family | +| --- | --- | +| `DurableMapOperations` | `map`, `mapAsync` | +| `DurableParallelOperations` | `parallel` and branch construction | +| `DurableWaitForCallbackOperations` | `waitForCallback`, `waitForCallbackAsync` | +| `DurableWaitForConditionOperations` | `waitForCondition`, `waitForConditionAsync` | +| `DurableWithRetryOperations` | `withRetry`, `withRetryAsync` | + +These classes obtain the active durable context from SDK-managed thread-local storage and delegate to existing +operation implementations. The facade split does not require rewriting the established operation implementations. + +### Use TLS for SDK Context and Metadata + +User functions in the new static APIs receive only application-provided values or values from the application's +durable data flow. They do not receive `DurableContext`, `StepContext`, `ExtensionContext`, or SDK-generated metadata +as callback parameters. + +Examples: + +- Step and child-context functions use `Supplier`. +- Map functions receive the item; `MapItemContext.getCurrentContext()` provides the item index. +- Parallel branches use `Supplier`. +- Wait-for-callback submitters use `Runnable`; + `WaitForCallbackContext.getCurrentContext()` provides the callback ID. +- Wait-for-condition checks receive the durable state; + `StepContext.getCurrentContext()` provides attempt metadata. +- With-retry bodies use `Supplier`; + `WithRetryContext.getCurrentContext()` provides the attempt number. + +Operation-specific contexts use scoped SDK-managed thread-local storage. Nested scopes restore the previous value, and +the value is removed when no previous scope exists. Operation-specific metadata TLS is bound in addition to the base +durable or step context, allowing static core operations to resolve the correct active context. + +Current context is available only on SDK-managed user-code threads. It is not propagated to application-created +threads. + +### Reserve Primitive Identities + +`ExtensionContext.reserve(name)` immediately allocates the next sequential operation ID in the active durable scope +and returns an opaque, one-shot `ExtensionOperation`. + +```java +public interface ExtensionOperation { + DurableFuture stepAsync( + TypeToken resultType, + Supplier function, + StepConfig config); + + DurableFuture waitAsync(Duration duration); + + DurableFuture invokeAsync( + String functionName, + U payload, + TypeToken resultType, + InvokeConfig config); + + DurableCallbackFuture createCallback( + TypeToken resultType, + CallbackConfig config); + + DurableFuture runInChildContextAsync( + TypeToken resultType, + Supplier function, + RunInChildContextConfig config); +} +``` + +The SDK binds the operation name and hidden ID to the reservation. The reservation can create exactly one primitive; +reusing it throws `IllegalStateException`. + +Extensions with deterministic invocation order can call `DurableCoreOperations` directly. Extensions such as DAG +reserve identities during deterministic definition, then execute the reservations in any dependency-valid order. + +The implementation may add package-private explicit-ID primitive constructors. Extension code cannot access those +methods or raw IDs. + +### Support Composed Durable Futures + +Add a public, non-mutating completion signal to `DurableFuture`: + +```java +default CompletableFuture completionFuture() { + throw new UnsupportedOperationException( + "This DurableFuture does not expose a completion signal"); +} +``` + +SDK operations return a derived completion future that cannot mutate the underlying durable operation. +`DurableFuture.anyOf` uses this public contract instead of downcasting to `BaseDurableOperation`. Custom composed +futures override the method when they support `anyOf`. + +### Preserve Primitive Plugin Semantics + +Extensions do not create an automatic lifecycle or checkpoint boundary. Plugins observe the primitives created by an +extension. If an extension explicitly creates a child context, plugins also observe that context operation. + +No extension-specific backend operation type or raw checkpoint API is added. + +## Alternatives Considered + +### Add extension methods to DurableContext + +Add `runExtensionAsync` or operation-specific methods to `DurableContext`. + +**Rejected because:** + +- It changes the interface that this decision must preserve. +- Optional extension families would continue to expand the core API. +- It makes extension execution appear to require a special runtime boundary. + +### Require every extension to run in a child context + +Provide a universal `DurableExtensions.run` method that creates a child context. + +**Rejected because:** + +- Direct primitive wrappers do not need a child context. +- Child-context checkpoint and replay behavior would be imposed even when it is not part of the extension semantics. +- Extensions such as map or retry must remain responsible for selecting their own isolation strategy. + +### Compose directly through DurableContext only + +Let extension implementations retrieve `DurableContext` and invoke its existing methods. + +**Rejected because:** + +- It exposes the entire legacy operation surface instead of a stable extension contract. +- Operations receive IDs when executed, so schedulers cannot register identities before varying launch order. +- Extension implementations remain coupled to context-accepting legacy callbacks. + +### Expose raw or name-derived operation IDs + +Allow extension libraries to supply operation IDs or derive them from operation names. + +**Rejected because:** + +- The SDK must retain ownership of global uniqueness and backend identity rules. +- Name-derived IDs introduce collision, normalization, and compatibility requirements. +- Public explicit-ID methods expose checkpoint protocol details. + +### Pass contexts and generated metadata as callback arguments + +Mirror the existing `DurableContext` callback signatures in the new static APIs. + +**Rejected because:** + +- The new API uses SDK-managed current context consistently across core and extension operations. +- Generated values such as map index, retry attempt, and callback ID belong to typed operation contexts. +- Context-free callbacks make extension methods compose without threading SDK objects through application code. + +### Use one built-in extension facade + +Place map, parallel, callback, condition, and retry methods in one `DurableExtensionOperations` class. + +**Rejected because:** + +- Unrelated overload sets, tests, and documentation would change together. +- Large operation families such as map and parallel need independent ownership. +- Separate classes align the public API with independently maintained extension implementations. + +## Consequences + +**Positive:** + +- Third-party Maven modules can publish static durable operations using supported public contracts. +- Application call sites do not pass or qualify `DurableContext`. +- `DurableContext` remains source- and binary-compatible. +- Extensions choose their own scope instead of inheriting a mandatory child-context boundary. +- Deterministic reservations support replay-safe schedulers whose launch order can vary. +- Primitive IDs, checkpointing, replay, and backend communication remain SDK-owned. +- Built-in extension families can evolve independently. +- New callback APIs consistently use TLS for SDK context and generated metadata. +- Custom composed futures work with public future combinators without internal downcasts. + +**Negative:** + +- The SDK must manage multiple scoped thread-local context types and restore them correctly across nested calls. +- Static APIs depend on execution from SDK-managed threads and fail from application-created threads. +- Reservations add package-private explicit-ID paths that must remain consistent with ordinary primitive creation. +- The static facade API duplicates overloads that remain on `DurableContext` for compatibility. +- Extension authors must understand that reservation order is part of workflow replay compatibility. + +**Compatibility requirements:** + +- Reservations must be created in the same deterministic order on every replay. +- Reordering, inserting, or removing reservations can rebind existing checkpoints and is a workflow compatibility + change. +- Launching already reserved operations in a different order is supported. +- Existing `DurableContext` methods and callback signatures remain unchanged. + +**Deferred:** + +- A production DAG extension module. +- Reimplementing every built-in extension through the new public reservation contract. +- Propagating current context to application-created threads. From c03d6dd9633db96746bd823527dcf6ce8c6a29df Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 00:57:17 +0000 Subject: [PATCH 05/40] docs: plan custom extension operations implementation --- .../2026-08-10-custom-extension-operations.md | 608 ++++++++++++++++++ 1 file changed, 608 insertions(+) create mode 100644 docs/superpowers/plans/2026-08-10-custom-extension-operations.md diff --git a/docs/superpowers/plans/2026-08-10-custom-extension-operations.md b/docs/superpowers/plans/2026-08-10-custom-extension-operations.md new file mode 100644 index 000000000..12dfc111f --- /dev/null +++ b/docs/superpowers/plans/2026-08-10-custom-extension-operations.md @@ -0,0 +1,608 @@ +# Custom Extension Operations Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add a supported public API for context-free static durable operations and replay-safe custom extensions while leaving the existing `DurableContext` interface surface unchanged. + +**Architecture:** SDK-managed handler and child contexts implement a minimal `ExtensionContext` that can reserve opaque, one-shot primitive identities. `DurableCoreOperations` and one facade per built-in extension family adapt context-free user callbacks to the existing `DurableContext` implementation, while typed scoped TLS contexts expose SDK-generated metadata. Existing primitive operation classes continue to own checkpointing, replay, suspension, serialization, and plugin events. + +**Tech Stack:** Java 17, Maven reactor, JUnit 6, Mockito 5, `LocalDurableTestRunner`, Palantir Java Format through Spotless. + +## Global Constraints + +- Keep all existing `DurableContext` methods and callback signatures unchanged for backward compatibility. +- Do not add `runExtensionAsync` or any other extension entry point to `DurableContext`. +- New user callbacks receive only application-provided values; SDK contexts and generated metadata are retrieved through TLS. +- Keep primitive operation IDs opaque and SDK-owned. +- A reservation is one-shot and allocates its sequential ID when `reserve` is called, not when the primitive is launched. +- Extensions do not automatically create child contexts. +- Do not add dependencies. +- Run `mvn spotless:apply` after Java changes. + +--- + +### Task 1: Scoped Current Contexts + +**Files:** +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/BaseContextImpl.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/execution/DurableExecutor.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java` +- Create: `sdk/src/main/java/software/amazon/lambda/durable/OperationContextStorage.java` +- Create: `sdk/src/main/java/software/amazon/lambda/durable/MapItemContext.java` +- Create: `sdk/src/main/java/software/amazon/lambda/durable/WaitForCallbackContext.java` +- Create: `sdk/src/main/java/software/amazon/lambda/durable/WithRetryContext.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/CurrentContextTest.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/OperationContextStorageTest.java` + +**Interfaces:** +- Consumes: Existing `BaseContext.getCurrentContext()`, `DurableContext.getCurrentContext()`, and `StepContext.getCurrentContext()`. +- Produces: `MapItemContext.getCurrentContext().getIndex()`, `WaitForCallbackContext.getCurrentContext().getCallbackId()`, and `WithRetryContext.getCurrentContext().getAttempt()`. + +- [ ] **Step 1: Write failing lookup and restoration tests** + +Add tests that establish these behaviors: + +```java +@Test +void mapItemContextRestoresNestedScope() { + assertThrows(IllegalStateException.class, MapItemContext::getCurrentContext); + try (var outer = MapItemContext.attach(2)) { + assertEquals(2, MapItemContext.getCurrentContext().getIndex()); + try (var inner = MapItemContext.attach(7)) { + assertEquals(7, MapItemContext.getCurrentContext().getIndex()); + } + assertEquals(2, MapItemContext.getCurrentContext().getIndex()); + } + assertThrows(IllegalStateException.class, MapItemContext::getCurrentContext); +} +``` + +Add equivalent outside-scope and nested-restoration assertions for callback IDs and retry attempts. In +`CurrentContextTest`, assert that handler/child contexts resolve as `DurableContext`, step scopes reject +`DurableContext` with guidance to use `StepContext`, and all scopes restore the preceding base context. + +- [ ] **Step 2: Run the focused tests and verify RED** + +Run: + +```bash +JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ +/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ +-Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ +-Djacoco.skip=true \ +-DargLine=-javaagent:/home/czk/.codex-tmp/m2/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar \ +-pl sdk -Dtest=CurrentContextTest,OperationContextStorageTest test +``` + +Expected: compilation fails because the three operation-specific context classes and their scoped attachment methods +do not exist. + +- [ ] **Step 3: Implement scoped storage and SDK binding** + +Implement a package-private generic storage helper: + +```java +final class OperationContextStorage { + private final String contextName; + private final ThreadLocal current = new ThreadLocal<>(); + + T getCurrentContext() { + var context = current.get(); + if (context == null) { + throw new IllegalStateException(contextName + " is not active on the current thread"); + } + return context; + } + + SafeCloseable attach(T context) { + var previous = current.get(); + current.set(context); + return () -> { + if (previous == null) { + current.remove(); + } else { + current.set(previous); + } + }; + } +} +``` + +Each public final metadata context owns a private static storage, a private immutable value, a public static lookup, +a public getter, and a package-private `attach` used by same-package facades. Preserve the existing +`DurableContext` signatures while clarifying its current-context failure behavior. Bind base contexts with +try-with-resources around handler, child, step, and condition user functions so nested calls restore rather than +blindly clear TLS. + +- [ ] **Step 4: Run focused tests and verify GREEN** + +Run the command from Step 2. Expected: all current-context and operation-context tests pass. + +- [ ] **Step 5: Commit** + +```bash +git add sdk/src/main/java/software/amazon/lambda/durable \ + sdk/src/test/java/software/amazon/lambda/durable/CurrentContextTest.java \ + sdk/src/test/java/software/amazon/lambda/durable/OperationContextStorageTest.java +git commit -m "feat: add scoped durable operation contexts" +``` + +### Task 2: Opaque Primitive Reservations + +**Files:** +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionContext.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java` +- Delete: `sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionContextImpl.java` +- Delete: `sdk/src/main/java/software/amazon/lambda/durable/DurableExtensions.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java` +- Delete: `sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionContextImplTest.java` + +**Interfaces:** +- Consumes: `DurableContextImpl` primitive construction and `OperationIdGenerator`. +- Produces: + +```java +public interface ExtensionContext extends BaseContext { + static ExtensionContext getCurrentContext(); + boolean isReplaying(); + ExtensionOperation reserve(String name); +} +``` + +```java +public interface ExtensionOperation { + DurableFuture stepAsync( + TypeToken resultType, Supplier function, StepConfig config); + DurableFuture waitAsync(Duration duration); + DurableFuture invokeAsync( + String functionName, U payload, TypeToken resultType, InvokeConfig config); + DurableCallbackFuture createCallback( + TypeToken resultType, CallbackConfig config); + DurableFuture runInChildContextAsync( + TypeToken resultType, Supplier function, RunInChildContextConfig config); +} +``` + +- [ ] **Step 1: Write failing reservation tests** + +Create tests that use a real `DurableContextImpl` with mocked dependencies to prove: + +```java +var first = context.reserve("first"); +var second = context.reserve("second"); + +second.stepAsync(String.class, () -> "second"); +first.stepAsync(String.class, () -> "first"); + +verify(operationFactory).createStepOperation("2", "second", ...); +verify(operationFactory).createStepOperation("1", "first", ...); +``` + +Also test every reserved primitive path, `ExtensionContext.getCurrentContext()` on handler/child versus step/no scope, +and a second use of the same reservation throwing: + +```java +assertThrows(IllegalStateException.class, () -> reservation.waitAsync(Duration.ZERO)); +``` + +- [ ] **Step 2: Run focused tests and verify RED** + +Run: + +```bash +JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ +/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ +-Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ +-Djacoco.skip=true \ +-DargLine=-javaagent:/home/czk/.codex-tmp/m2/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar \ +-pl sdk -Dtest=ExtensionOperationImplTest,CurrentContextTest test +``` + +Expected: compilation fails because reservations still accept context-bearing functions and +`DurableContextImpl` does not directly implement the final `ExtensionContext` contract. + +- [ ] **Step 3: Implement the minimal reservation path** + +Make `DurableContextImpl` implement `ExtensionContext`, with: + +```java +@Override +public ExtensionOperation reserve(String name) { + return new ExtensionOperationImpl(this, operationIdGenerator.next(), name); +} +``` + +Retain package-private explicit-ID helpers for step, wait, invoke, callback, and child context. Adapt `Supplier` to +the existing primitive callbacks inside `ExtensionOperationImpl`; current TLS is already attached when the user +supplier executes. Guard all primitive selectors with a single `AtomicBoolean.compareAndSet(false, true)`. + +Remove `runExtensionAsync` from `DurableContext`, remove the universal `DurableExtensions` facade, remove the wrapper +`ExtensionContextImpl`, and keep `ExtensionContext` limited to current lookup, replay state, and reservation. + +- [ ] **Step 4: Run focused tests and verify GREEN** + +Run the command from Step 2. Expected: all reservation and current-context tests pass. + +- [ ] **Step 5: Commit** + +```bash +git add sdk/src/main/java/software/amazon/lambda/durable \ + sdk/src/test/java/software/amazon/lambda/durable/context +git commit -m "feat: add opaque extension operation reservations" +``` + +### Task 3: Context-Free Core Static Facade + +**Files:** +- Create: `sdk/src/main/java/software/amazon/lambda/durable/DurableCoreOperations.java` +- Delete: `sdk/src/main/java/software/amazon/lambda/durable/DurableOperations.java` +- Create: `sdk/src/test/java/software/amazon/lambda/durable/DurableCoreOperationsTest.java` + +**Interfaces:** +- Consumes: `DurableContext.getCurrentContext()` and all existing primitive instance methods. +- Produces: sync/async, `Class`/`TypeToken`, default/custom config overloads for `step`, `wait`, chained + `invoke`, `createCallback`, and `runInChildContext`. + +- [ ] **Step 1: Write failing facade tests** + +Write tests against a TLS-bound mocked `DurableContext` proving the core overloads delegate and strip SDK context +parameters: + +```java +var result = DurableCoreOperations.step("step", String.class, () -> { + assertSame(stepContext, StepContext.getCurrentContext()); + return "done"; +}); +assertEquals("done", result); +``` + +For child contexts, verify a zero-argument supplier can obtain both `DurableContext` and `ExtensionContext` from TLS. +Also assert every facade family throws `IllegalStateException` outside a durable context. + +- [ ] **Step 2: Run the focused test and verify RED** + +Run: + +```bash +JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ +/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ +-Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ +-Djacoco.skip=true \ +-DargLine=-javaagent:/home/czk/.codex-tmp/m2/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar \ +-pl sdk -Dtest=DurableCoreOperationsTest test +``` + +Expected: compilation fails because `DurableCoreOperations` does not exist. + +- [ ] **Step 3: Implement only primitive facade overloads** + +Create a stateless final utility class whose step methods accept `Supplier` and delegate with +`ignored -> function.get()`. Child-context methods also accept `Supplier` and delegate with +`ignored -> function.get()`. Invoke, wait, and callback methods delegate values unchanged. Do not include map, +parallel, callback composition, condition polling, or retry methods. + +- [ ] **Step 4: Run focused tests and verify GREEN** + +Run the command from Step 2. Expected: all core facade tests pass. + +- [ ] **Step 5: Commit** + +```bash +git add sdk/src/main/java/software/amazon/lambda/durable/DurableCoreOperations.java \ + sdk/src/main/java/software/amazon/lambda/durable/DurableOperations.java \ + sdk/src/test/java/software/amazon/lambda/durable/DurableCoreOperationsTest.java +git commit -m "feat: add context-free core operation facade" +``` + +### Task 4: Independently Maintained Extension Facades + +**Files:** +- Create: `sdk/src/main/java/software/amazon/lambda/durable/DurableMapOperations.java` +- Create: `sdk/src/main/java/software/amazon/lambda/durable/DurableParallelOperations.java` +- Create: `sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForCallbackOperations.java` +- Create: `sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForConditionOperations.java` +- Create: `sdk/src/main/java/software/amazon/lambda/durable/DurableWithRetryOperations.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java` +- Create: `sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationsTest.java` +- Create: `sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationsTest.java` +- Create: `sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationsTest.java` +- Create: `sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForConditionOperationsTest.java` +- Create: `sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationsTest.java` + +**Interfaces:** +- Consumes: Existing `DurableContext` map, parallel, wait-for-callback, wait-for-condition, and with-retry methods. +- Produces: + - map callbacks as `Function` with index in `MapItemContext` + - parallel branch callbacks as `Supplier` + - callback submitters as `Runnable` with ID in `WaitForCallbackContext` + - condition checks as `Function>` + - retry bodies as `Supplier` with attempt in `WithRetryContext` + +- [ ] **Step 1: Write one failing metadata test per facade** + +Tests must invoke the adapted legacy callback and assert the public callback sees TLS metadata: + +```java +DurableMapOperations.map("map", List.of("a"), String.class, item -> { + assertEquals(3, MapItemContext.getCurrentContext().getIndex()); + return item.toUpperCase(); +}); +``` + +Use the mocked legacy callback to supply index `3`; repeat for callback ID `"cb-1"` and retry attempt `2`. For +wait-for-condition, bind a `StepContext` and prove the check function receives only the state value. For parallel, +compile and execute `parallel.branch("branch", String.class, () -> "done")`. + +- [ ] **Step 2: Run the five focused tests and verify RED** + +Run: + +```bash +JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ +/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ +-Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ +-Djacoco.skip=true \ +-DargLine=-javaagent:/home/czk/.codex-tmp/m2/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar \ +-pl sdk -Dtest=DurableMapOperationsTest,DurableParallelOperationsTest,DurableWaitForCallbackOperationsTest,DurableWaitForConditionOperationsTest,DurableWithRetryOperationsTest test +``` + +Expected: compilation fails because the split facades and supplier branch overloads do not exist. + +- [ ] **Step 3: Implement the five adapters** + +Each class is a stateless final utility class containing only its named family. Adapt callbacks as follows: + +```java +(item, index, ignored) -> { + try (var scope = MapItemContext.attach(index)) { + return function.apply(item); + } +} +``` + +```java +(callbackId, ignored) -> { + try (var scope = WaitForCallbackContext.attach(callbackId)) { + submitter.run(); + } +} +``` + +```java +(attempt, ignored) -> { + try (var scope = WithRetryContext.attach(attempt)) { + return operation.get(); + } +} +``` + +The condition adapter is `(state, ignored) -> checkFunction.apply(state)` because `WaitForConditionOperation` binds +the active `StepContext`. Add default supplier overloads to `ParallelDurableFuture` that delegate to its existing +`Function` core method. + +- [ ] **Step 4: Run focused tests and verify GREEN** + +Run the command from Step 2. Expected: all five facade tests pass. + +- [ ] **Step 5: Commit** + +```bash +git add sdk/src/main/java/software/amazon/lambda/durable \ + sdk/src/test/java/software/amazon/lambda/durable/Durable*OperationsTest.java +git commit -m "feat: split built-in extension operation facades" +``` + +### Task 5: Public Durable Future Completion Contract + +**Files:** +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/DurableFuture.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java` +- Modify: `sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java` + +**Interfaces:** +- Consumes: Existing `DurableFuture.anyOf` and SDK operation completion futures. +- Produces: `default CompletableFuture completionFuture()` for custom composed futures. + +- [ ] **Step 1: Write a failing custom-future combinator test** + +Create a test-only `DurableFuture` whose result future and completion signal are independent, override +`completionFuture()`, pass it to `DurableFuture.anyOf`, complete the signal, and assert `anyOf` completes without +requiring the future to extend `BaseDurableOperation`. + +- [ ] **Step 2: Run the focused test and verify RED** + +Run: + +```bash +JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ +/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ +-Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ +-Djacoco.skip=true \ +-pl sdk -Dtest=DurableFutureTest test +``` + +Expected: the test fails because `anyOf` still relies on an SDK-internal operation downcast or no public completion +method exists. + +- [ ] **Step 3: Implement the completion signal** + +Add the default method throwing `UnsupportedOperationException` to custom futures that do not opt in. Override it in +`BaseDurableOperation` by deriving `internalFuture.thenApply(ignored -> null)` so callers cannot complete or cancel +the underlying durable operation. Change `DurableFuture.anyOf` to collect `completionFuture()` values without an +internal type cast. + +- [ ] **Step 4: Run focused tests and verify GREEN** + +Run the command from Step 2. Expected: all durable future tests pass. + +- [ ] **Step 5: Commit** + +```bash +git add sdk/src/main/java/software/amazon/lambda/durable/DurableFuture.java \ + sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java \ + sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java +git commit -m "feat: expose durable future completion signals" +``` + +### Task 6: Separate-Module Proof Extension and Integration Semantics + +**Files:** +- Modify: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/extension/PairOperations.java` +- Modify: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java` +- Create: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java` +- Modify: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java` + +**Interfaces:** +- Consumes: Only public classes from the `sdk` artifact; the `sdk-integration-tests` Maven module is the external + compilation boundary. +- Produces: A proof extension that reserves two step identities in registration order and can launch them in either + order without changing replay identities. + +- [ ] **Step 1: Write failing external-module integration tests** + +Implement the test fixture call site before its production helper: + +```java +var result = PairOperations.pair( + "pair", + () -> DurableCoreOperations.step("left-value", String.class, () -> "left"), + () -> DurableCoreOperations.step("right-value", String.class, () -> "right"), + true); +assertEquals(new PairResult("left", "right"), result); +``` + +Add tests using `LocalDurableTestRunner` for first execution plus replay, a wait-based suspension/resume, reverse launch +order, same-scope nested extension calls, an explicitly reserved child context, all static facade families, and TLS +metadata. Update the plugin test to assert primitive names are observed and no synthetic extension lifecycle event is +emitted. + +- [ ] **Step 2: Run integration tests and verify RED** + +Run: + +```bash +JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ +/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ +-Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ +-Djacoco.skip=true \ +-DargLine=-javaagent:/home/czk/.codex-tmp/m2/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar \ +-pl sdk-integration-tests -am \ +-Dtest=ExtensionOperationIntegrationTest,StaticOperationsIntegrationTest,PluginIntegrationTest \ +-Dsurefire.failIfNoSpecifiedTests=false test +``` + +Expected: tests fail because the external fixture still targets the discarded universal runner and context-bearing +callbacks. + +- [ ] **Step 3: Implement the public-contract-only fixture** + +`PairOperations` may import only public SDK types. It obtains `ExtensionContext.getCurrentContext()`, reserves +`left` and `right`, launches the selected reservation order, and combines the results. It must not import any package +under `software.amazon.lambda.durable.context`, `.execution`, or `.operation`. + +- [ ] **Step 4: Run integration tests and verify GREEN** + +Run the command from Step 2. Expected: all extension, static operation, and plugin integration tests pass. + +- [ ] **Step 5: Commit** + +```bash +git add sdk-integration-tests/src/test/java/software/amazon/lambda/durable +git commit -m "test: verify custom extensions across module boundary" +``` + +### Task 7: Documentation, Formatting, and Reactor Verification + +**Files:** +- Modify: `README.md` +- Modify: `docs/advanced/extensions.md` +- Modify: `docs/adr/006-custom-extension-operations.md` + +**Interfaces:** +- Consumes: The final public API implemented by Tasks 1-6. +- Produces: User and extension-author documentation matching the code exactly. + +- [ ] **Step 1: Rewrite the extension guide** + +Document: + +- static import examples for `DurableCoreOperations` and every split extension facade +- TLS-only user functions and all three metadata contexts +- ordinary static extension methods without `DurableExtensions.run` +- same-scope direct composition versus explicit child contexts +- deterministic reservation order and variable launch order +- one-shot reservation behavior +- public-contract-only separate Maven modules +- primitive-only plugin lifecycle events +- `DurableFuture.completionFuture()` for custom composed futures + +Mark ADR-006 `Accepted` only after implementation and verification succeed. + +- [ ] **Step 2: Format all Java sources** + +Run: + +```bash +JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ +/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ +-Dmaven.repo.local=/home/czk/.codex-tmp/m2 spotless:apply +``` + +Expected: Spotless exits successfully and only intended Java files change. + +- [ ] **Step 3: Run focused SDK and integration verification** + +Run: + +```bash +JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ +/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ +-Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ +-Djacoco.skip=true \ +-DargLine=-javaagent:/home/czk/.codex-tmp/m2/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar \ +-pl sdk,sdk-integration-tests -am test +``` + +Expected: all tests in the dependency closure pass. + +- [ ] **Step 4: Run the full reactor** + +Run: + +```bash +JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ +/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ +-Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ +-Djacoco.skip=true \ +-DargLine=-javaagent:/home/czk/.codex-tmp/m2/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar \ +clean install +``` + +Expected: `BUILD SUCCESS`. Cloud example tests remain disabled. + +- [ ] **Step 5: Review public compatibility and commit** + +Verify: + +```bash +git diff 6962f5a -- sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java +rg -n "DurableExtensions|DurableOperations|runExtensionAsync" \ + sdk/src/main sdk/src/test sdk-integration-tests docs README.md +git status --short +``` + +Expected: `DurableContext` contains no new method signatures; discarded names have no live references; the status +contains only intended files. + +Then commit: + +```bash +git add README.md docs sdk sdk-integration-tests +git commit -m "docs: document custom extension operations" +``` From 6905f4b405d9d898d760a8b8637411286fb686ee Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 00:59:58 +0000 Subject: [PATCH 06/40] feat: add extension operation reservations --- .../amazon/lambda/durable/DurableContext.java | 16 +- .../lambda/durable/ExtensionContext.java | 42 +++++ .../lambda/durable/ExtensionOperation.java | 138 ++++++++++++++++ .../amazon/lambda/durable/MapItemContext.java | 31 ++++ .../durable/OperationContextStorage.java | 36 +++++ .../amazon/lambda/durable/StepContext.java | 11 +- .../durable/WaitForCallbackContext.java | 32 ++++ .../lambda/durable/WithRetryContext.java | 31 ++++ .../durable/context/BaseContextImpl.java | 19 +++ .../durable/context/DurableContextImpl.java | 88 ++++++++-- .../context/ExtensionOperationImpl.java | 67 ++++++++ .../durable/execution/DurableExecutor.java | 4 +- .../operation/ChildContextOperation.java | 4 +- .../durable/operation/StepOperation.java | 5 +- .../operation/WaitForConditionOperation.java | 4 +- .../lambda/durable/CurrentContextTest.java | 74 +++++++++ .../durable/OperationContextStorageTest.java | 63 ++++++++ .../context/ExtensionOperationImplTest.java | 152 ++++++++++++++++++ 18 files changed, 796 insertions(+), 21 deletions(-) create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/ExtensionContext.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/MapItemContext.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/OperationContextStorage.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/WaitForCallbackContext.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/WithRetryContext.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/CurrentContextTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/OperationContextStorageTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java index ce6eb7070..19fb7ebad 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java @@ -22,8 +22,22 @@ import software.amazon.lambda.durable.model.WaitForConditionResult; public interface DurableContext extends BaseContext { + /** + * Returns the durable context attached to the current SDK-managed context thread. + * + * @return the current durable context + * @throws IllegalStateException if called outside a durable context or from a step thread + */ static DurableContext getCurrentContext() { - return (DurableContext) BaseContext.getCurrentContext(); + var context = BaseContext.getCurrentContext(); + if (context instanceof DurableContext durableContext) { + return durableContext; + } + if (context == null) { + throw new IllegalStateException("No DurableContext is active on the current thread"); + } + throw new IllegalStateException( + "DurableContext is not available from a step thread; use StepContext.getCurrentContext() instead"); } /** Returns whether this context is currently replaying checkpointed durable operations. */ diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContext.java b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContext.java new file mode 100644 index 000000000..4d5379a19 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContext.java @@ -0,0 +1,42 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import software.amazon.lambda.durable.context.BaseContext; + +/** + * Public context available to custom extension operations. + * + *

This interface exposes replay state and opaque primitive reservations without exposing checkpoint internals or + * raw operation IDs. + */ +public interface ExtensionContext extends BaseContext { + /** + * Returns the extension context attached to the current SDK-managed context thread. + * + * @return the current extension context + * @throws IllegalStateException if called outside a durable context or from a step thread + */ + static ExtensionContext getCurrentContext() { + var context = BaseContext.getCurrentContext(); + if (context instanceof ExtensionContext extensionContext) { + return extensionContext; + } + if (context == null) { + throw new IllegalStateException("No ExtensionContext is active on the current thread"); + } + throw new IllegalStateException( + "ExtensionContext is not available from a step thread; use StepContext.getCurrentContext() instead"); + } + + /** Returns whether this extension scope is replaying checkpointed operations. */ + boolean isReplaying(); + + /** + * Reserves the next sequential primitive operation identity. + * + * @param name the primitive operation name + * @return an opaque one-shot reservation + */ + ExtensionOperation reserve(String name); +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java new file mode 100644 index 000000000..0a29376c9 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java @@ -0,0 +1,138 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import java.time.Duration; +import java.util.function.Supplier; +import software.amazon.lambda.durable.config.CallbackConfig; +import software.amazon.lambda.durable.config.InvokeConfig; +import software.amazon.lambda.durable.config.RunInChildContextConfig; +import software.amazon.lambda.durable.config.StepConfig; + +/** + * An opaque, one-shot reservation for a primitive operation. + * + *

The SDK allocates the operation ID when the reservation is created. Reserving operations in deterministic order + * allows an extension to launch them later in a different order without changing their IDs. + */ +public interface ExtensionOperation { + default T step(Class resultType, Supplier function) { + return step(TypeToken.get(resultType), function); + } + + default T step(TypeToken resultType, Supplier function) { + return step(resultType, function, StepConfig.builder().build()); + } + + default T step(Class resultType, Supplier function, StepConfig config) { + return step(TypeToken.get(resultType), function, config); + } + + default T step(TypeToken resultType, Supplier function, StepConfig config) { + return stepAsync(resultType, function, config).get(); + } + + default DurableFuture stepAsync(Class resultType, Supplier function) { + return stepAsync(TypeToken.get(resultType), function); + } + + default DurableFuture stepAsync(TypeToken resultType, Supplier function) { + return stepAsync(resultType, function, StepConfig.builder().build()); + } + + default DurableFuture stepAsync(Class resultType, Supplier function, StepConfig config) { + return stepAsync(TypeToken.get(resultType), function, config); + } + + DurableFuture stepAsync(TypeToken resultType, Supplier function, StepConfig config); + + default Void wait(Duration duration) { + return waitAsync(duration).get(); + } + + DurableFuture waitAsync(Duration duration); + + default T invoke(String functionName, U payload, Class resultType) { + return invoke(functionName, payload, TypeToken.get(resultType)); + } + + default T invoke(String functionName, U payload, TypeToken resultType) { + return invoke( + functionName, payload, resultType, InvokeConfig.builder().build()); + } + + default T invoke(String functionName, U payload, Class resultType, InvokeConfig config) { + return invoke(functionName, payload, TypeToken.get(resultType), config); + } + + default T invoke(String functionName, U payload, TypeToken resultType, InvokeConfig config) { + return invokeAsync(functionName, payload, resultType, config).get(); + } + + default DurableFuture invokeAsync(String functionName, U payload, Class resultType) { + return invokeAsync(functionName, payload, TypeToken.get(resultType)); + } + + default DurableFuture invokeAsync(String functionName, U payload, TypeToken resultType) { + return invokeAsync( + functionName, payload, resultType, InvokeConfig.builder().build()); + } + + default DurableFuture invokeAsync( + String functionName, U payload, Class resultType, InvokeConfig config) { + return invokeAsync(functionName, payload, TypeToken.get(resultType), config); + } + + DurableFuture invokeAsync( + String functionName, U payload, TypeToken resultType, InvokeConfig config); + + default DurableCallbackFuture createCallback(Class resultType) { + return createCallback(TypeToken.get(resultType)); + } + + default DurableCallbackFuture createCallback(TypeToken resultType) { + return createCallback(resultType, CallbackConfig.builder().build()); + } + + default DurableCallbackFuture createCallback(Class resultType, CallbackConfig config) { + return createCallback(TypeToken.get(resultType), config); + } + + DurableCallbackFuture createCallback(TypeToken resultType, CallbackConfig config); + + default T runInChildContext(Class resultType, Supplier function) { + return runInChildContext(TypeToken.get(resultType), function); + } + + default T runInChildContext(TypeToken resultType, Supplier function) { + return runInChildContext( + resultType, function, RunInChildContextConfig.builder().build()); + } + + default T runInChildContext( + Class resultType, Supplier function, RunInChildContextConfig config) { + return runInChildContext(TypeToken.get(resultType), function, config); + } + + default T runInChildContext( + TypeToken resultType, Supplier function, RunInChildContextConfig config) { + return runInChildContextAsync(resultType, function, config).get(); + } + + default DurableFuture runInChildContextAsync(Class resultType, Supplier function) { + return runInChildContextAsync(TypeToken.get(resultType), function); + } + + default DurableFuture runInChildContextAsync(TypeToken resultType, Supplier function) { + return runInChildContextAsync( + resultType, function, RunInChildContextConfig.builder().build()); + } + + default DurableFuture runInChildContextAsync( + Class resultType, Supplier function, RunInChildContextConfig config) { + return runInChildContextAsync(TypeToken.get(resultType), function, config); + } + + DurableFuture runInChildContextAsync( + TypeToken resultType, Supplier function, RunInChildContextConfig config); +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/MapItemContext.java b/sdk/src/main/java/software/amazon/lambda/durable/MapItemContext.java new file mode 100644 index 000000000..14b2689b7 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/MapItemContext.java @@ -0,0 +1,31 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import software.amazon.lambda.durable.model.SafeCloseable; + +/** Metadata for the map item function active on the current SDK-managed thread. */ +public final class MapItemContext { + private static final OperationContextStorage CURRENT = + new OperationContextStorage<>("MapItemContext"); + + private final int index; + + private MapItemContext(int index) { + this.index = index; + } + + /** Returns the map item context attached to the current SDK-managed thread. */ + public static MapItemContext getCurrentContext() { + return CURRENT.getCurrentContext(); + } + + /** Returns the zero-based index of the current map item. */ + public int getIndex() { + return index; + } + + static SafeCloseable attach(int index) { + return CURRENT.attach(new MapItemContext(index)); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/OperationContextStorage.java b/sdk/src/main/java/software/amazon/lambda/durable/OperationContextStorage.java new file mode 100644 index 000000000..a36bfc212 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/OperationContextStorage.java @@ -0,0 +1,36 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import java.util.Objects; +import software.amazon.lambda.durable.model.SafeCloseable; + +final class OperationContextStorage { + private final String contextName; + private final ThreadLocal current = new ThreadLocal<>(); + + OperationContextStorage(String contextName) { + this.contextName = contextName; + } + + T getCurrentContext() { + var context = current.get(); + if (context == null) { + throw new IllegalStateException(contextName + " is not active on the current thread"); + } + return context; + } + + SafeCloseable attach(T context) { + Objects.requireNonNull(context, "context cannot be null"); + var previous = current.get(); + current.set(context); + return () -> { + if (previous == null) { + current.remove(); + } else { + current.set(previous); + } + }; + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/StepContext.java b/sdk/src/main/java/software/amazon/lambda/durable/StepContext.java index 092897d10..f2469b8ca 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/StepContext.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/StepContext.java @@ -8,7 +8,16 @@ public interface StepContext extends BaseContext { /** Returns the current retry attempt number (0-based). */ int getAttempt(); + /** Returns the step context attached to the current SDK-managed thread. */ static StepContext getCurrentContext() { - return (StepContext) BaseContext.getCurrentContext(); + var context = BaseContext.getCurrentContext(); + if (context instanceof StepContext stepContext) { + return stepContext; + } + if (context == null) { + throw new IllegalStateException("No StepContext is active on the current thread"); + } + throw new IllegalStateException( + "StepContext is not available from a durable context thread; use DurableContext.getCurrentContext() instead"); } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/WaitForCallbackContext.java b/sdk/src/main/java/software/amazon/lambda/durable/WaitForCallbackContext.java new file mode 100644 index 000000000..68a528c84 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/WaitForCallbackContext.java @@ -0,0 +1,32 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import java.util.Objects; +import software.amazon.lambda.durable.model.SafeCloseable; + +/** Metadata for the callback submitter active on the current SDK-managed thread. */ +public final class WaitForCallbackContext { + private static final OperationContextStorage CURRENT = + new OperationContextStorage<>("WaitForCallbackContext"); + + private final String callbackId; + + private WaitForCallbackContext(String callbackId) { + this.callbackId = Objects.requireNonNull(callbackId, "callbackId cannot be null"); + } + + /** Returns the callback context attached to the current SDK-managed thread. */ + public static WaitForCallbackContext getCurrentContext() { + return CURRENT.getCurrentContext(); + } + + /** Returns the callback ID to send to the external system. */ + public String getCallbackId() { + return callbackId; + } + + static SafeCloseable attach(String callbackId) { + return CURRENT.attach(new WaitForCallbackContext(callbackId)); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/WithRetryContext.java b/sdk/src/main/java/software/amazon/lambda/durable/WithRetryContext.java new file mode 100644 index 000000000..8c4953bf8 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/WithRetryContext.java @@ -0,0 +1,31 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import software.amazon.lambda.durable.model.SafeCloseable; + +/** Metadata for the retry body active on the current SDK-managed thread. */ +public final class WithRetryContext { + private static final OperationContextStorage CURRENT = + new OperationContextStorage<>("WithRetryContext"); + + private final int attempt; + + private WithRetryContext(int attempt) { + this.attempt = attempt; + } + + /** Returns the retry context attached to the current SDK-managed thread. */ + public static WithRetryContext getCurrentContext() { + return CURRENT.getCurrentContext(); + } + + /** Returns the current one-based retry attempt. */ + public int getAttempt() { + return attempt; + } + + static SafeCloseable attach(int attempt) { + return CURRENT.attach(new WithRetryContext(attempt)); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/BaseContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/BaseContextImpl.java index 6379acdf1..1753f590e 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/BaseContextImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/BaseContextImpl.java @@ -8,6 +8,7 @@ import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.execution.ThreadType; import software.amazon.lambda.durable.logging.DurableLogger; +import software.amazon.lambda.durable.model.SafeCloseable; public abstract class BaseContextImpl implements BaseContext { private final ExecutionManager executionManager; @@ -109,4 +110,22 @@ public DurableLogger getLogger(Logger delegate) { public static void setCurrentContext(BaseContext context) { CONTEXT.set(context); } + + /** + * Sets the current SDK context until the returned scope is closed. + * + * @param context the context to attach + * @return a scope that restores the previous context + */ + public static SafeCloseable attachCurrentContext(BaseContext context) { + var previous = CONTEXT.get(); + CONTEXT.set(context); + return () -> { + if (previous == null) { + CONTEXT.remove(); + } else { + CONTEXT.set(previous); + } + }; + } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java index 0c79165ec..94e4fd3f7 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java @@ -14,6 +14,8 @@ import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.ExtensionContext; +import software.amazon.lambda.durable.ExtensionOperation; import software.amazon.lambda.durable.ParallelDurableFuture; import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.TypeToken; @@ -52,7 +54,7 @@ *

Provides methods for creating steps, waits, chained invokes, callbacks, and child contexts. Each method creates a * checkpoint-backed operation that survives Lambda interruptions. */ -public class DurableContextImpl extends BaseContextImpl implements DurableContext { +public class DurableContextImpl extends BaseContextImpl implements DurableContext, ExtensionContext { private static final String WAIT_FOR_CALLBACK_CALLBACK_SUFFIX = "-callback"; private static final String WAIT_FOR_CALLBACK_SUBMITTER_SUFFIX = "-submitter"; private static final int MAX_WAIT_FOR_CALLBACK_NAME_LENGTH = ParameterValidator.MAX_OPERATION_NAME_LENGTH @@ -136,12 +138,22 @@ public DurableFuture stepAsync( Objects.requireNonNull(config, "config cannot be null"); Objects.requireNonNull(resultType, "resultType cannot be null"); ParameterValidator.validateOperationName(name); + return stepAsyncWithId(nextOperationId(), name, resultType, func, config); + } + + DurableFuture stepAsyncWithId( + String operationId, + String name, + TypeToken resultType, + Function func, + StepConfig config) { + Objects.requireNonNull(config, "config cannot be null"); + Objects.requireNonNull(resultType, "resultType cannot be null"); + ParameterValidator.validateOperationName(name); if (config.serDes() == null) { config = config.toBuilder().serDes(getDurableConfig().getSerDes()).build(); } - var operationId = nextOperationId(); - // Create and start step operation with TypeToken var operation = new StepOperation<>( OperationIdentifier.of(operationId, name, OperationSubType.STEP), func, resultType, config, this); @@ -155,8 +167,12 @@ public DurableFuture stepAsync( public DurableFuture waitAsync(String name, Duration duration) { ParameterValidator.validateDuration(duration, "Wait duration"); ParameterValidator.validateOperationName(name); + return waitAsyncWithId(nextOperationId(), name, duration); + } - var operationId = nextOperationId(); + DurableFuture waitAsyncWithId(String operationId, String name, Duration duration) { + ParameterValidator.validateDuration(duration, "Wait duration"); + ParameterValidator.validateOperationName(name); // Create and start wait operation var operation = @@ -172,6 +188,19 @@ public DurableFuture invokeAsync( Objects.requireNonNull(config, "config cannot be null"); Objects.requireNonNull(resultType, "resultType cannot be null"); ParameterValidator.validateOperationName(name); + return invokeAsyncWithId(nextOperationId(), name, functionName, payload, resultType, config); + } + + DurableFuture invokeAsyncWithId( + String operationId, + String name, + String functionName, + U payload, + TypeToken resultType, + InvokeConfig config) { + Objects.requireNonNull(config, "config cannot be null"); + Objects.requireNonNull(resultType, "resultType cannot be null"); + ParameterValidator.validateOperationName(name); if (config.serDes() == null) { config = config.toBuilder().serDes(getDurableConfig().getSerDes()).build(); @@ -181,8 +210,6 @@ public DurableFuture invokeAsync( .payloadSerDes(getDurableConfig().getSerDes()) .build(); } - var operationId = nextOperationId(); - // Create and start invoke operation var operation = new InvokeOperation<>( OperationIdentifier.of(operationId, name, OperationSubType.CHAINED_INVOKE), @@ -198,12 +225,20 @@ public DurableFuture invokeAsync( @Override public DurableCallbackFuture createCallback(String name, TypeToken resultType, CallbackConfig config) { + Objects.requireNonNull(config, "config cannot be null"); + Objects.requireNonNull(resultType, "resultType cannot be null"); + ParameterValidator.validateOperationName(name); + return createCallbackWithId(nextOperationId(), name, resultType, config); + } + + DurableCallbackFuture createCallbackWithId( + String operationId, String name, TypeToken resultType, CallbackConfig config) { + Objects.requireNonNull(config, "config cannot be null"); + Objects.requireNonNull(resultType, "resultType cannot be null"); ParameterValidator.validateOperationName(name); if (config.serDes() == null) { config = config.toBuilder().serDes(getDurableConfig().getSerDes()).build(); } - var operationId = nextOperationId(); - var operation = new CallbackOperation<>( OperationIdentifier.of(operationId, name, OperationSubType.CALLBACK), resultType, config, this); operation.execute(); @@ -236,6 +271,31 @@ private DurableFuture runInChildContextAsync( RunInChildContextConfig config, OperationSubType subType) { Objects.requireNonNull(resultType, "resultType cannot be null"); + Objects.requireNonNull(func, "func cannot be null"); + Objects.requireNonNull(config, "RunInChildContextConfig cannot be null"); + ParameterValidator.validateOperationName(name); + return runInChildContextAsyncWithId(nextOperationId(), name, resultType, func, config, subType); + } + + DurableFuture runInChildContextAsyncWithId( + String operationId, + String name, + TypeToken resultType, + Function func, + RunInChildContextConfig config) { + return runInChildContextAsyncWithId( + operationId, name, resultType, func, config, OperationSubType.RUN_IN_CHILD_CONTEXT); + } + + private DurableFuture runInChildContextAsyncWithId( + String operationId, + String name, + TypeToken resultType, + Function func, + RunInChildContextConfig config, + OperationSubType subType) { + Objects.requireNonNull(resultType, "resultType cannot be null"); + Objects.requireNonNull(func, "func cannot be null"); Objects.requireNonNull(config, "RunInChildContextConfig cannot be null"); ParameterValidator.validateOperationName(name); @@ -243,8 +303,6 @@ private DurableFuture runInChildContextAsync( config = config.toBuilder().serDes(getDurableConfig().getSerDes()).build(); } - var operationId = nextOperationId(); - var operation = new ChildContextOperation<>( OperationIdentifier.of(operationId, name, subType), func, resultType, config, this); @@ -441,6 +499,16 @@ private String nextOperationId() { return operationIdGenerator.nextOperationId(); } + String reserveOperationId() { + return nextOperationId(); + } + + @Override + public ExtensionOperation reserve(String name) { + ParameterValidator.validateOperationName(name); + return new ExtensionOperationImpl(this, reserveOperationId(), name); + } + /** Returns whether this context is currently in replay mode. */ @Override public boolean isReplaying() { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java new file mode 100644 index 000000000..bf11713f3 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java @@ -0,0 +1,67 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.context; + +import java.time.Duration; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Supplier; +import software.amazon.lambda.durable.DurableCallbackFuture; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.ExtensionOperation; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.CallbackConfig; +import software.amazon.lambda.durable.config.InvokeConfig; +import software.amazon.lambda.durable.config.RunInChildContextConfig; +import software.amazon.lambda.durable.config.StepConfig; + +final class ExtensionOperationImpl implements ExtensionOperation { + private final DurableContextImpl context; + private final String operationId; + private final String name; + private final AtomicBoolean claimed = new AtomicBoolean(); + + ExtensionOperationImpl(DurableContextImpl context, String operationId, String name) { + this.context = context; + this.operationId = operationId; + this.name = name; + } + + @Override + public DurableFuture stepAsync(TypeToken resultType, Supplier function, StepConfig config) { + claim(); + return context.stepAsyncWithId(operationId, name, resultType, ignored -> function.get(), config); + } + + @Override + public DurableFuture waitAsync(Duration duration) { + claim(); + return context.waitAsyncWithId(operationId, name, duration); + } + + @Override + public DurableFuture invokeAsync( + String functionName, U payload, TypeToken resultType, InvokeConfig config) { + claim(); + return context.invokeAsyncWithId(operationId, name, functionName, payload, resultType, config); + } + + @Override + public DurableCallbackFuture createCallback(TypeToken resultType, CallbackConfig config) { + claim(); + return context.createCallbackWithId(operationId, name, resultType, config); + } + + @Override + public DurableFuture runInChildContextAsync( + TypeToken resultType, Supplier function, RunInChildContextConfig config) { + claim(); + return context.runInChildContextAsyncWithId( + operationId, name, resultType, ignored -> function.get(), config); + } + + private void claim() { + if (!claimed.compareAndSet(false, true)) { + throw new IllegalStateException("An extension operation reservation can only be used once"); + } + } +} 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 6ab6bdbc3..f90b3fb9f 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 @@ -76,9 +76,9 @@ public static DurableExecutionOutput execute( var userInput = extractUserInput( executionManager.getExecutionOperation(), config.getSerDes(), inputType); var context = DurableContextImpl.createRootContext(executionManager, config, lambdaContext); - DurableContextImpl.setCurrentContext(context); // use a try-with-resources to clear logger properties - try (var ignored = DurableLogger.attachContext()) { + try (var ignoredContext = DurableContextImpl.attachCurrentContext(context); + var ignoredLogger = DurableLogger.attachContext()) { return handler.apply(userInput, context); } }, diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java index 8c299cfa4..be2191726 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java @@ -132,8 +132,8 @@ private void executeChildContext() { // we notify the parent BEFORE closing the child context. This ensures the parent // can trigger the next queued branch while the current child context is still valid. var childContext = getContext().createChildContext(contextId, getName(), isVirtual); - DurableContextImpl.setCurrentContext(childContext); - try (var ignored = DurableLogger.attachContext()) { + try (var ignoredContext = DurableContextImpl.attachCurrentContext(childContext); + var ignoredLogger = DurableLogger.attachContext()) { try { // Run the user function inside the plugin hook boundary (attempt is null for contexts) // so a failure is reported through onUserFunctionEnd; checkpointing stays outside. diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java index 467a87b94..d2b3cc7f4 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java @@ -107,9 +107,8 @@ private void executeStepLogic(int attempt) { // - add thread id/type to thread local when the step starts // - clear logger properties when the step finishes StepContext stepContext = getContext().createStepContext(getOperationId(), getName(), attempt); - BaseContextImpl.setCurrentContext(stepContext); - - try (var ignored = DurableLogger.attachContext()) { + try (var ignoredContext = BaseContextImpl.attachCurrentContext(stepContext); + var ignoredLogger = DurableLogger.attachContext()) { try { checkpointStarted(); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java index 6a653c37b..c906bc5f6 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java @@ -115,8 +115,8 @@ private CompletableFuture pollReadyAndResumeCheckLoop(Operation existing) private void executeCheckLogic(T currentState, int attempt) { Runnable userHandler = () -> { var stepContext = getContext().createStepContext(getOperationId(), getName(), attempt); - BaseContextImpl.setCurrentContext(stepContext); - try (var ignored = DurableLogger.attachContext()) { + try (var ignoredContext = BaseContextImpl.attachCurrentContext(stepContext); + var ignoredLogger = DurableLogger.attachContext()) { try { // Checkpoint START if not already started var existing = getOperation(); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/CurrentContextTest.java b/sdk/src/test/java/software/amazon/lambda/durable/CurrentContextTest.java new file mode 100644 index 000000000..c0a54b9e3 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/CurrentContextTest.java @@ -0,0 +1,74 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.context.BaseContextImpl; + +class CurrentContextTest { + @AfterEach + void clearContext() { + BaseContextImpl.setCurrentContext(null); + } + + @Test + void durableContextFailsClearlyOutsideDurableThread() { + var exception = assertThrows(IllegalStateException.class, DurableContext::getCurrentContext); + + assertTrue(exception.getMessage().contains("No DurableContext")); + } + + @Test + void durableContextFailsClearlyOnStepThread() { + BaseContextImpl.setCurrentContext(mock(StepContext.class)); + + var exception = assertThrows(IllegalStateException.class, DurableContext::getCurrentContext); + + assertTrue(exception.getMessage().contains("step thread")); + } + + @Test + void stepContextFailsClearlyOutsideStepThread() { + var exception = assertThrows(IllegalStateException.class, StepContext::getCurrentContext); + + assertTrue(exception.getMessage().contains("No StepContext")); + } + + @Test + void extensionContextReturnsActiveExtensionContext() { + var context = mock(CurrentExtensionContext.class); + BaseContextImpl.setCurrentContext(context); + + assertSame(context, ExtensionContext.getCurrentContext()); + } + + @Test + void extensionContextFailsClearlyOnStepThread() { + BaseContextImpl.setCurrentContext(mock(StepContext.class)); + + var exception = assertThrows(IllegalStateException.class, ExtensionContext::getCurrentContext); + + assertTrue(exception.getMessage().contains("step thread")); + } + + @Test + void currentContextScopesRestorePreviousContext() { + var outer = mock(DurableContext.class); + var inner = mock(StepContext.class); + BaseContextImpl.setCurrentContext(outer); + + try (var ignored = BaseContextImpl.attachCurrentContext(inner)) { + assertSame(inner, StepContext.getCurrentContext()); + } + + assertSame(outer, DurableContext.getCurrentContext()); + } + + private interface CurrentExtensionContext extends DurableContext, ExtensionContext {} +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/OperationContextStorageTest.java b/sdk/src/test/java/software/amazon/lambda/durable/OperationContextStorageTest.java new file mode 100644 index 000000000..c47b6e072 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/OperationContextStorageTest.java @@ -0,0 +1,63 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import org.junit.jupiter.api.Test; + +class OperationContextStorageTest { + @Test + void mapItemContextRestoresNestedScope() { + assertThrows(IllegalStateException.class, MapItemContext::getCurrentContext); + + try (var outer = MapItemContext.attach(2)) { + assertEquals(2, MapItemContext.getCurrentContext().getIndex()); + try (var inner = MapItemContext.attach(7)) { + assertEquals(7, MapItemContext.getCurrentContext().getIndex()); + } + assertEquals(2, MapItemContext.getCurrentContext().getIndex()); + } + + assertThrows(IllegalStateException.class, MapItemContext::getCurrentContext); + } + + @Test + void waitForCallbackContextRestoresNestedScope() { + assertThrows(IllegalStateException.class, WaitForCallbackContext::getCurrentContext); + + try (var outer = WaitForCallbackContext.attach("outer")) { + assertEquals("outer", WaitForCallbackContext.getCurrentContext().getCallbackId()); + try (var inner = WaitForCallbackContext.attach("inner")) { + assertEquals("inner", WaitForCallbackContext.getCurrentContext().getCallbackId()); + } + assertEquals("outer", WaitForCallbackContext.getCurrentContext().getCallbackId()); + } + + assertThrows(IllegalStateException.class, WaitForCallbackContext::getCurrentContext); + } + + @Test + void withRetryContextRestoresNestedScope() { + assertThrows(IllegalStateException.class, WithRetryContext::getCurrentContext); + + try (var outer = WithRetryContext.attach(1)) { + assertEquals(1, WithRetryContext.getCurrentContext().getAttempt()); + try (var inner = WithRetryContext.attach(2)) { + assertEquals(2, WithRetryContext.getCurrentContext().getAttempt()); + } + assertEquals(1, WithRetryContext.getCurrentContext().getAttempt()); + } + + assertThrows(IllegalStateException.class, WithRetryContext::getCurrentContext); + } + + @Test + void operationContextFailureNamesRequestedContext() { + var exception = assertThrows(IllegalStateException.class, MapItemContext::getCurrentContext); + + assertTrue(exception.getMessage().contains("MapItemContext")); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java new file mode 100644 index 000000000..8984336ed --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java @@ -0,0 +1,152 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.context; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doCallRealMethod; +import static org.mockito.Mockito.inOrder; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.time.Duration; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Function; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import software.amazon.lambda.durable.DurableCallbackFuture; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.ExtensionOperation; +import software.amazon.lambda.durable.StepContext; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.CallbackConfig; +import software.amazon.lambda.durable.config.InvokeConfig; +import software.amazon.lambda.durable.config.RunInChildContextConfig; +import software.amazon.lambda.durable.config.StepConfig; + +class ExtensionOperationImplTest { + @Test + void reservationsKeepSequentialIdsWhenExecutedOutOfOrder() { + var context = mock(DurableContextImpl.class); + when(context.reserveOperationId()).thenReturn("sequential-1", "sequential-2"); + doCallRealMethod().when(context).reserve("first"); + doCallRealMethod().when(context).reserve("second"); + var duration = Duration.ofSeconds(1); + when(context.waitAsyncWithId("sequential-1", "first", duration)).thenReturn(mockFuture()); + when(context.waitAsyncWithId("sequential-2", "second", duration)).thenReturn(mockFuture()); + + var first = context.reserve("first"); + var second = context.reserve("second"); + second.waitAsync(duration); + first.waitAsync(duration); + + var ordered = inOrder(context); + ordered.verify(context).reserveOperationId(); + ordered.verify(context).reserveOperationId(); + ordered.verify(context).waitAsyncWithId("sequential-2", "second", duration); + ordered.verify(context).waitAsyncWithId("sequential-1", "first", duration); + } + + @Test + void reservedStepAdaptsSupplierToStepFunction() { + var context = mock(DurableContextImpl.class); + var future = mockStringFuture(); + var resultType = TypeToken.get(String.class); + var config = StepConfig.builder().build(); + var called = new AtomicBoolean(); + when(context.stepAsyncWithId(eq("1"), eq("step"), eq(resultType), any(), eq(config))) + .thenReturn(future); + var operation = new ExtensionOperationImpl(context, "1", "step"); + + assertEquals(future, operation.stepAsync(resultType, () -> { + called.set(true); + return "result"; + }, config)); + + @SuppressWarnings("unchecked") + var function = (ArgumentCaptor>) (ArgumentCaptor) ArgumentCaptor.forClass( + Function.class); + verify(context).stepAsyncWithId(eq("1"), eq("step"), eq(resultType), function.capture(), eq(config)); + assertEquals("result", function.getValue().apply(mock(StepContext.class))); + assertEquals(true, called.get()); + } + + @Test + void reservationDelegatesWaitInvokeAndCallback() { + var duration = Duration.ofSeconds(2); + var waitContext = mock(DurableContextImpl.class); + var waitFuture = mockFuture(); + when(waitContext.waitAsyncWithId("1", "wait", duration)).thenReturn(waitFuture); + assertEquals(waitFuture, new ExtensionOperationImpl(waitContext, "1", "wait").waitAsync(duration)); + + var invokeContext = mock(DurableContextImpl.class); + var invokeFuture = mockStringFuture(); + var invokeConfig = InvokeConfig.builder().build(); + var resultType = TypeToken.get(String.class); + when(invokeContext.invokeAsyncWithId("2", "invoke", "target", "payload", resultType, invokeConfig)) + .thenReturn(invokeFuture); + assertEquals( + invokeFuture, + new ExtensionOperationImpl(invokeContext, "2", "invoke") + .invokeAsync("target", "payload", resultType, invokeConfig)); + + var callbackContext = mock(DurableContextImpl.class); + @SuppressWarnings("unchecked") + var callbackFuture = (DurableCallbackFuture) mock(DurableCallbackFuture.class); + var callbackConfig = CallbackConfig.builder().build(); + when(callbackContext.createCallbackWithId("3", "callback", resultType, callbackConfig)) + .thenReturn(callbackFuture); + assertEquals( + callbackFuture, + new ExtensionOperationImpl(callbackContext, "3", "callback") + .createCallback(resultType, callbackConfig)); + } + + @Test + void reservedChildContextAdaptsSupplierToChildFunction() { + var context = mock(DurableContextImpl.class); + var future = mockStringFuture(); + var resultType = TypeToken.get(String.class); + var config = RunInChildContextConfig.builder().build(); + when(context.runInChildContextAsyncWithId(eq("1"), eq("child"), eq(resultType), any(), eq(config))) + .thenReturn(future); + var operation = new ExtensionOperationImpl(context, "1", "child"); + + assertEquals(future, operation.runInChildContextAsync(resultType, () -> "result", config)); + + @SuppressWarnings("unchecked") + var function = (ArgumentCaptor>) (ArgumentCaptor) ArgumentCaptor.forClass( + Function.class); + verify(context) + .runInChildContextAsyncWithId(eq("1"), eq("child"), eq(resultType), function.capture(), eq(config)); + assertEquals("result", function.getValue().apply(mock(DurableContext.class))); + } + + @Test + void reservationCanOnlyExecuteOnceAcrossPrimitiveSelectors() { + var context = mock(DurableContextImpl.class); + var duration = Duration.ofSeconds(1); + when(context.waitAsyncWithId("1", "only-once", duration)).thenReturn(mockFuture()); + ExtensionOperation operation = new ExtensionOperationImpl(context, "1", "only-once"); + + operation.waitAsync(duration); + + assertThrows( + IllegalStateException.class, + () -> operation.stepAsync(String.class, () -> "second")); + } + + @SuppressWarnings("unchecked") + private DurableFuture mockFuture() { + return mock(DurableFuture.class); + } + + @SuppressWarnings("unchecked") + private DurableFuture mockStringFuture() { + return mock(DurableFuture.class); + } +} From 59eb34d8e1d88234f3e2d3d95a4602b27b26411b Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 01:02:57 +0000 Subject: [PATCH 07/40] feat: add static durable operation facades --- .../lambda/durable/DurableCoreOperations.java | 159 ++++++++++++++++++ .../lambda/durable/DurableMapOperations.java | 83 +++++++++ .../durable/DurableParallelOperations.java | 18 ++ .../DurableWaitForCallbackOperations.java | 63 +++++++ .../DurableWaitForConditionOperations.java | 76 +++++++++ .../durable/DurableWithRetryOperations.java | 43 +++++ .../lambda/durable/ParallelDurableFuture.java | 22 +++ .../durable/DurableCoreOperationsTest.java | 113 +++++++++++++ .../durable/DurableMapOperationsTest.java | 40 +++++ .../DurableParallelOperationsTest.java | 62 +++++++ .../DurableWaitForCallbackOperationsTest.java | 39 +++++ ...DurableWaitForConditionOperationsTest.java | 44 +++++ .../DurableWithRetryOperationsTest.java | 38 +++++ 13 files changed, 800 insertions(+) create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/DurableCoreOperations.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/DurableMapOperations.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/DurableParallelOperations.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForCallbackOperations.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForConditionOperations.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/DurableWithRetryOperations.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/DurableCoreOperationsTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationsTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationsTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationsTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForConditionOperationsTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationsTest.java diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableCoreOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableCoreOperations.java new file mode 100644 index 000000000..a929b5e33 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableCoreOperations.java @@ -0,0 +1,159 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import java.time.Duration; +import java.util.function.Supplier; +import software.amazon.lambda.durable.config.CallbackConfig; +import software.amazon.lambda.durable.config.InvokeConfig; +import software.amazon.lambda.durable.config.RunInChildContextConfig; +import software.amazon.lambda.durable.config.StepConfig; + +/** + * Context-free static facades for SDK-owned primitive durable operations. + * + *

The equivalent instance methods on {@link DurableContext} remain supported for backward compatibility. + */ +public final class DurableCoreOperations { + private DurableCoreOperations() {} + + public static T step(String name, Class resultType, Supplier function) { + return currentContext().step(name, resultType, ignored -> function.get()); + } + + public static T step(String name, TypeToken resultType, Supplier function) { + return currentContext().step(name, resultType, ignored -> function.get()); + } + + public static T step(String name, Class resultType, Supplier function, StepConfig config) { + return currentContext().step(name, resultType, ignored -> function.get(), config); + } + + public static T step(String name, TypeToken resultType, Supplier function, StepConfig config) { + return currentContext().step(name, resultType, ignored -> function.get(), config); + } + + public static DurableFuture stepAsync(String name, Class resultType, Supplier function) { + return currentContext().stepAsync(name, resultType, ignored -> function.get()); + } + + public static DurableFuture stepAsync(String name, TypeToken resultType, Supplier function) { + return currentContext().stepAsync(name, resultType, ignored -> function.get()); + } + + public static DurableFuture stepAsync( + String name, Class resultType, Supplier function, StepConfig config) { + return currentContext().stepAsync(name, resultType, ignored -> function.get(), config); + } + + public static DurableFuture stepAsync( + String name, TypeToken resultType, Supplier function, StepConfig config) { + return currentContext().stepAsync(name, resultType, ignored -> function.get(), config); + } + + public static Void wait(String name, Duration duration) { + return currentContext().wait(name, duration); + } + + public static DurableFuture waitAsync(String name, Duration duration) { + return currentContext().waitAsync(name, duration); + } + + public static T invoke(String name, String functionName, U payload, Class resultType) { + return currentContext().invoke(name, functionName, payload, resultType); + } + + public static T invoke(String name, String functionName, U payload, TypeToken resultType) { + return currentContext().invoke(name, functionName, payload, resultType); + } + + public static T invoke( + String name, String functionName, U payload, Class resultType, InvokeConfig config) { + return currentContext().invoke(name, functionName, payload, resultType, config); + } + + public static T invoke( + String name, String functionName, U payload, TypeToken resultType, InvokeConfig config) { + return currentContext().invoke(name, functionName, payload, resultType, config); + } + + public static DurableFuture invokeAsync( + String name, String functionName, U payload, Class resultType) { + return currentContext().invokeAsync(name, functionName, payload, resultType); + } + + public static DurableFuture invokeAsync( + String name, String functionName, U payload, TypeToken resultType) { + return currentContext().invokeAsync(name, functionName, payload, resultType); + } + + public static DurableFuture invokeAsync( + String name, String functionName, U payload, Class resultType, InvokeConfig config) { + return currentContext().invokeAsync(name, functionName, payload, resultType, config); + } + + public static DurableFuture invokeAsync( + String name, String functionName, U payload, TypeToken resultType, InvokeConfig config) { + return currentContext().invokeAsync(name, functionName, payload, resultType, config); + } + + public static DurableCallbackFuture createCallback(String name, Class resultType) { + return currentContext().createCallback(name, resultType); + } + + public static DurableCallbackFuture createCallback(String name, TypeToken resultType) { + return currentContext().createCallback(name, resultType); + } + + public static DurableCallbackFuture createCallback( + String name, Class resultType, CallbackConfig config) { + return currentContext().createCallback(name, resultType, config); + } + + public static DurableCallbackFuture createCallback( + String name, TypeToken resultType, CallbackConfig config) { + return currentContext().createCallback(name, resultType, config); + } + + public static T runInChildContext(String name, Class resultType, Supplier function) { + return currentContext().runInChildContext(name, resultType, ignored -> function.get()); + } + + public static T runInChildContext(String name, TypeToken resultType, Supplier function) { + return currentContext().runInChildContext(name, resultType, ignored -> function.get()); + } + + public static T runInChildContext( + String name, Class resultType, Supplier function, RunInChildContextConfig config) { + return currentContext().runInChildContext(name, resultType, ignored -> function.get(), config); + } + + public static T runInChildContext( + String name, TypeToken resultType, Supplier function, RunInChildContextConfig config) { + return currentContext().runInChildContext(name, resultType, ignored -> function.get(), config); + } + + public static DurableFuture runInChildContextAsync( + String name, Class resultType, Supplier function) { + return currentContext().runInChildContextAsync(name, resultType, ignored -> function.get()); + } + + public static DurableFuture runInChildContextAsync( + String name, TypeToken resultType, Supplier function) { + return currentContext().runInChildContextAsync(name, resultType, ignored -> function.get()); + } + + public static DurableFuture runInChildContextAsync( + String name, Class resultType, Supplier function, RunInChildContextConfig config) { + return currentContext().runInChildContextAsync(name, resultType, ignored -> function.get(), config); + } + + public static DurableFuture runInChildContextAsync( + String name, TypeToken resultType, Supplier function, RunInChildContextConfig config) { + return currentContext().runInChildContextAsync(name, resultType, ignored -> function.get(), config); + } + + private static DurableContext currentContext() { + return DurableContext.getCurrentContext(); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableMapOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableMapOperations.java new file mode 100644 index 000000000..3ea76bbd1 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableMapOperations.java @@ -0,0 +1,83 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import java.util.Collection; +import java.util.Objects; +import java.util.function.Function; +import software.amazon.lambda.durable.config.MapConfig; +import software.amazon.lambda.durable.model.MapResult; + +/** Context-free static facades for durable map operations. */ +public final class DurableMapOperations { + private DurableMapOperations() {} + + public static MapResult map( + String name, Collection items, Class resultType, Function function) { + return currentContext().map(name, items, resultType, adapt(function)); + } + + public static MapResult map( + String name, Collection items, TypeToken resultType, Function function) { + return currentContext().map(name, items, resultType, adapt(function)); + } + + public static MapResult map( + String name, + Collection items, + Class resultType, + Function function, + MapConfig config) { + return currentContext().map(name, items, resultType, adapt(function), config); + } + + public static MapResult map( + String name, + Collection items, + TypeToken resultType, + Function function, + MapConfig config) { + return currentContext().map(name, items, resultType, adapt(function), config); + } + + public static DurableFuture> mapAsync( + String name, Collection items, Class resultType, Function function) { + return currentContext().mapAsync(name, items, resultType, adapt(function)); + } + + public static DurableFuture> mapAsync( + String name, Collection items, TypeToken resultType, Function function) { + return currentContext().mapAsync(name, items, resultType, adapt(function)); + } + + public static DurableFuture> mapAsync( + String name, + Collection items, + Class resultType, + Function function, + MapConfig config) { + return currentContext().mapAsync(name, items, resultType, adapt(function), config); + } + + public static DurableFuture> mapAsync( + String name, + Collection items, + TypeToken resultType, + Function function, + MapConfig config) { + return currentContext().mapAsync(name, items, resultType, adapt(function), config); + } + + private static DurableContext.MapFunction adapt(Function function) { + Objects.requireNonNull(function, "function cannot be null"); + return (item, index, ignored) -> { + try (var scope = MapItemContext.attach(index)) { + return function.apply(item); + } + }; + } + + private static DurableContext currentContext() { + return DurableContext.getCurrentContext(); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableParallelOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableParallelOperations.java new file mode 100644 index 000000000..18cf6d519 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableParallelOperations.java @@ -0,0 +1,18 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import software.amazon.lambda.durable.config.ParallelConfig; + +/** Context-free static facades for durable parallel operations. */ +public final class DurableParallelOperations { + private DurableParallelOperations() {} + + public static ParallelDurableFuture parallel(String name) { + return DurableContext.getCurrentContext().parallel(name); + } + + public static ParallelDurableFuture parallel(String name, ParallelConfig config) { + return DurableContext.getCurrentContext().parallel(name, config); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForCallbackOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForCallbackOperations.java new file mode 100644 index 000000000..fbefa764a --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForCallbackOperations.java @@ -0,0 +1,63 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import java.util.Objects; +import java.util.function.BiConsumer; +import software.amazon.lambda.durable.config.WaitForCallbackConfig; + +/** Context-free static facades for durable wait-for-callback operations. */ +public final class DurableWaitForCallbackOperations { + private DurableWaitForCallbackOperations() {} + + public static T waitForCallback(String name, Class resultType, Runnable submitter) { + return currentContext().waitForCallback(name, resultType, adapt(submitter)); + } + + public static T waitForCallback(String name, TypeToken resultType, Runnable submitter) { + return currentContext().waitForCallback(name, resultType, adapt(submitter)); + } + + public static T waitForCallback( + String name, Class resultType, Runnable submitter, WaitForCallbackConfig config) { + return currentContext().waitForCallback(name, resultType, adapt(submitter), config); + } + + public static T waitForCallback( + String name, TypeToken resultType, Runnable submitter, WaitForCallbackConfig config) { + return currentContext().waitForCallback(name, resultType, adapt(submitter), config); + } + + public static DurableFuture waitForCallbackAsync( + String name, Class resultType, Runnable submitter) { + return currentContext().waitForCallbackAsync(name, resultType, adapt(submitter)); + } + + public static DurableFuture waitForCallbackAsync( + String name, TypeToken resultType, Runnable submitter) { + return currentContext().waitForCallbackAsync(name, resultType, adapt(submitter)); + } + + public static DurableFuture waitForCallbackAsync( + String name, Class resultType, Runnable submitter, WaitForCallbackConfig config) { + return currentContext().waitForCallbackAsync(name, resultType, adapt(submitter), config); + } + + public static DurableFuture waitForCallbackAsync( + String name, TypeToken resultType, Runnable submitter, WaitForCallbackConfig config) { + return currentContext().waitForCallbackAsync(name, resultType, adapt(submitter), config); + } + + private static BiConsumer adapt(Runnable submitter) { + Objects.requireNonNull(submitter, "submitter cannot be null"); + return (callbackId, ignored) -> { + try (var scope = WaitForCallbackContext.attach(callbackId)) { + submitter.run(); + } + }; + } + + private static DurableContext currentContext() { + return DurableContext.getCurrentContext(); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForConditionOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForConditionOperations.java new file mode 100644 index 000000000..06784844f --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForConditionOperations.java @@ -0,0 +1,76 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import java.util.Objects; +import java.util.function.BiFunction; +import java.util.function.Function; +import software.amazon.lambda.durable.config.WaitForConditionConfig; +import software.amazon.lambda.durable.model.WaitForConditionResult; + +/** Context-free static facades for durable wait-for-condition operations. */ +public final class DurableWaitForConditionOperations { + private DurableWaitForConditionOperations() {} + + public static T waitForCondition( + String name, Class resultType, Function> checkFunction) { + return currentContext().waitForCondition(name, resultType, adapt(checkFunction)); + } + + public static T waitForCondition( + String name, TypeToken resultType, Function> checkFunction) { + return currentContext().waitForCondition(name, resultType, adapt(checkFunction)); + } + + public static T waitForCondition( + String name, + Class resultType, + Function> checkFunction, + WaitForConditionConfig config) { + return currentContext().waitForCondition(name, resultType, adapt(checkFunction), config); + } + + public static T waitForCondition( + String name, + TypeToken resultType, + Function> checkFunction, + WaitForConditionConfig config) { + return currentContext().waitForCondition(name, resultType, adapt(checkFunction), config); + } + + public static DurableFuture waitForConditionAsync( + String name, Class resultType, Function> checkFunction) { + return currentContext().waitForConditionAsync(name, resultType, adapt(checkFunction)); + } + + public static DurableFuture waitForConditionAsync( + String name, TypeToken resultType, Function> checkFunction) { + return currentContext().waitForConditionAsync(name, resultType, adapt(checkFunction)); + } + + public static DurableFuture waitForConditionAsync( + String name, + Class resultType, + Function> checkFunction, + WaitForConditionConfig config) { + return currentContext().waitForConditionAsync(name, resultType, adapt(checkFunction), config); + } + + public static DurableFuture waitForConditionAsync( + String name, + TypeToken resultType, + Function> checkFunction, + WaitForConditionConfig config) { + return currentContext().waitForConditionAsync(name, resultType, adapt(checkFunction), config); + } + + private static BiFunction> adapt( + Function> checkFunction) { + Objects.requireNonNull(checkFunction, "checkFunction cannot be null"); + return (state, ignored) -> checkFunction.apply(state); + } + + private static DurableContext currentContext() { + return DurableContext.getCurrentContext(); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableWithRetryOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableWithRetryOperations.java new file mode 100644 index 000000000..75151eb31 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableWithRetryOperations.java @@ -0,0 +1,43 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import java.util.Objects; +import java.util.function.BiFunction; +import java.util.function.Supplier; +import software.amazon.lambda.durable.config.WithRetryConfig; + +/** Context-free static facades for replay-safe retry operations. */ +public final class DurableWithRetryOperations { + private DurableWithRetryOperations() {} + + public static T withRetry(String name, Supplier operation) { + return currentContext().withRetry(name, adapt(operation)); + } + + public static T withRetry(String name, Supplier operation, WithRetryConfig config) { + return currentContext().withRetry(name, adapt(operation), config); + } + + public static DurableFuture withRetryAsync(String name, Supplier operation) { + return currentContext().withRetryAsync(name, adapt(operation)); + } + + public static DurableFuture withRetryAsync( + String name, Supplier operation, WithRetryConfig config) { + return currentContext().withRetryAsync(name, adapt(operation), config); + } + + private static BiFunction adapt(Supplier operation) { + Objects.requireNonNull(operation, "operation cannot be null"); + return (attempt, ignored) -> { + try (var scope = WithRetryContext.attach(attempt)) { + return operation.get(); + } + }; + } + + private static DurableContext currentContext() { + return DurableContext.getCurrentContext(); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java b/sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java index 5f3067eda..a1a1dae1c 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java @@ -3,12 +3,34 @@ package software.amazon.lambda.durable; import java.util.function.Function; +import java.util.function.Supplier; import software.amazon.lambda.durable.config.ParallelBranchConfig; import software.amazon.lambda.durable.model.ParallelResult; import software.amazon.lambda.durable.model.SafeCloseable; /** User-facing context for managing parallel branch execution within a durable function. */ public interface ParallelDurableFuture extends SafeCloseable, DurableFuture { + default DurableFuture branch(String name, Class resultType, Supplier function) { + return branch(name, TypeToken.get(resultType), function); + } + + default DurableFuture branch(String name, TypeToken resultType, Supplier function) { + return branch( + name, + resultType, + ignored -> function.get(), + ParallelBranchConfig.builder().build()); + } + + default DurableFuture branch( + String name, Class resultType, Supplier function, ParallelBranchConfig config) { + return branch(name, TypeToken.get(resultType), function, config); + } + + default DurableFuture branch( + String name, TypeToken resultType, Supplier function, ParallelBranchConfig config) { + return branch(name, resultType, ignored -> function.get(), config); + } /** * Registers and immediately starts a branch (respects maxConcurrency). diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableCoreOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableCoreOperationsTest.java new file mode 100644 index 000000000..b31e8afd4 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableCoreOperationsTest.java @@ -0,0 +1,113 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.time.Duration; +import java.util.function.Function; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import software.amazon.lambda.durable.config.RunInChildContextConfig; +import software.amazon.lambda.durable.config.StepConfig; +import software.amazon.lambda.durable.context.BaseContextImpl; + +class DurableCoreOperationsTest { + @AfterEach + void clearContext() { + BaseContextImpl.setCurrentContext(null); + } + + @Test + void stepAcceptsContextFreeSupplier() { + var context = mock(DurableContext.class); + BaseContextImpl.setCurrentContext(context); + + DurableCoreOperations.step("step", String.class, () -> "result"); + + @SuppressWarnings("unchecked") + var function = (ArgumentCaptor>) (ArgumentCaptor) ArgumentCaptor.forClass( + Function.class); + verify(context).step(eq("step"), eq(String.class), function.capture()); + assertEquals("result", function.getValue().apply(mock(StepContext.class))); + } + + @Test + void childContextAcceptsContextFreeSupplier() { + var context = mock(DurableContext.class); + BaseContextImpl.setCurrentContext(context); + + DurableCoreOperations.runInChildContext("child", String.class, () -> "result"); + + @SuppressWarnings("unchecked") + var function = (ArgumentCaptor>) (ArgumentCaptor) ArgumentCaptor.forClass( + Function.class); + verify(context).runInChildContext(eq("child"), eq(String.class), function.capture()); + assertEquals("result", function.getValue().apply(mock(DurableContext.class))); + } + + @Test + void coreValueOperationsDelegateToCurrentContext() { + var context = mock(DurableContext.class); + BaseContextImpl.setCurrentContext(context); + var duration = Duration.ofSeconds(1); + var waitFuture = mockFuture(); + var invokeFuture = mockStringFuture(); + @SuppressWarnings("unchecked") + var callbackFuture = (DurableCallbackFuture) mock(DurableCallbackFuture.class); + when(context.waitAsync("wait", duration)).thenReturn(waitFuture); + when(context.invokeAsync("invoke", "function", "payload", String.class)) + .thenReturn(invokeFuture); + when(context.createCallback("callback", String.class)) + .thenReturn(callbackFuture); + + assertEquals(waitFuture, DurableCoreOperations.waitAsync("wait", duration)); + assertEquals( + invokeFuture, + DurableCoreOperations.invokeAsync( + "invoke", "function", "payload", String.class)); + assertEquals( + callbackFuture, + DurableCoreOperations.createCallback("callback", String.class)); + } + + @Test + void configuredSupplierOverloadsDelegateToCurrentContext() { + var context = mock(DurableContext.class); + BaseContextImpl.setCurrentContext(context); + var stepConfig = StepConfig.builder().build(); + var childConfig = RunInChildContextConfig.builder().build(); + + DurableCoreOperations.stepAsync("step", new TypeToken() {}, () -> "step", stepConfig); + DurableCoreOperations.runInChildContextAsync( + "child", new TypeToken() {}, () -> "child", childConfig); + + verify(context).stepAsync(eq("step"), any(TypeToken.class), any(Function.class), eq(stepConfig)); + verify(context) + .runInChildContextAsync(eq("child"), any(TypeToken.class), any(Function.class), eq(childConfig)); + } + + @Test + void coreOperationsFailOutsideDurableContext() { + assertThrows( + IllegalStateException.class, + () -> DurableCoreOperations.step("step", String.class, () -> "result")); + } + + @SuppressWarnings("unchecked") + private DurableFuture mockFuture() { + return mock(DurableFuture.class); + } + + @SuppressWarnings("unchecked") + private DurableFuture mockStringFuture() { + return mock(DurableFuture.class); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationsTest.java new file mode 100644 index 000000000..7dc2d4c83 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationsTest.java @@ -0,0 +1,40 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; + +import java.util.List; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import software.amazon.lambda.durable.context.BaseContextImpl; + +class DurableMapOperationsTest { + @AfterEach + void clearContext() { + BaseContextImpl.setCurrentContext(null); + } + + @Test + void mapExposesItemIndexThroughScopedContext() { + var context = mock(DurableContext.class); + BaseContextImpl.setCurrentContext(context); + + DurableMapOperations.map("map", List.of("value"), String.class, item -> { + assertEquals(3, MapItemContext.getCurrentContext().getIndex()); + return item.toUpperCase(); + }); + + @SuppressWarnings("unchecked") + var function = (ArgumentCaptor>) (ArgumentCaptor) + ArgumentCaptor.forClass(DurableContext.MapFunction.class); + verify(context).map(eq("map"), eq(List.of("value")), eq(String.class), function.capture()); + assertEquals("VALUE", function.getValue().apply("value", 3, mock(DurableContext.class))); + assertThrows(IllegalStateException.class, MapItemContext::getCurrentContext); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationsTest.java new file mode 100644 index 000000000..b49b5eef0 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationsTest.java @@ -0,0 +1,62 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.mockito.Answers.CALLS_REAL_METHODS; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.function.Function; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import software.amazon.lambda.durable.config.ParallelBranchConfig; +import software.amazon.lambda.durable.context.BaseContextImpl; + +class DurableParallelOperationsTest { + @AfterEach + void clearContext() { + BaseContextImpl.setCurrentContext(null); + } + + @Test + void parallelBranchesAcceptContextFreeSuppliers() { + var context = mock(DurableContext.class); + var parallel = mock(ParallelDurableFuture.class, CALLS_REAL_METHODS); + var branchFuture = mockStringFuture(); + BaseContextImpl.setCurrentContext(context); + when(context.parallel("parallel")).thenReturn(parallel); + when(parallel.branch( + eq("branch"), + any(TypeToken.class), + any(Function.class), + any(ParallelBranchConfig.class))) + .thenReturn(branchFuture); + + var result = DurableParallelOperations.parallel("parallel"); + var resultFuture = result.branch("branch", String.class, () -> "result"); + + assertSame(parallel, result); + assertSame(branchFuture, resultFuture); + @SuppressWarnings("unchecked") + var function = (ArgumentCaptor>) (ArgumentCaptor) ArgumentCaptor.forClass( + Function.class); + verify(parallel) + .branch( + eq("branch"), + any(TypeToken.class), + function.capture(), + any(ParallelBranchConfig.class)); + assertEquals("result", function.getValue().apply(mock(DurableContext.class))); + } + + @SuppressWarnings("unchecked") + private DurableFuture mockStringFuture() { + return mock(DurableFuture.class); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationsTest.java new file mode 100644 index 000000000..65b825fee --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationsTest.java @@ -0,0 +1,39 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; + +import java.util.function.BiConsumer; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import software.amazon.lambda.durable.context.BaseContextImpl; + +class DurableWaitForCallbackOperationsTest { + @AfterEach + void clearContext() { + BaseContextImpl.setCurrentContext(null); + } + + @Test + void callbackSubmitterUsesRunnableAndScopedCallbackId() { + var context = mock(DurableContext.class); + BaseContextImpl.setCurrentContext(context); + + DurableWaitForCallbackOperations.waitForCallback("callback", String.class, () -> assertEquals( + "callback-id", + WaitForCallbackContext.getCurrentContext().getCallbackId())); + + @SuppressWarnings("unchecked") + var submitter = (ArgumentCaptor>) (ArgumentCaptor) + ArgumentCaptor.forClass(BiConsumer.class); + verify(context).waitForCallback(eq("callback"), eq(String.class), submitter.capture()); + submitter.getValue().accept("callback-id", mock(StepContext.class)); + assertThrows(IllegalStateException.class, WaitForCallbackContext::getCurrentContext); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForConditionOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForConditionOperationsTest.java new file mode 100644 index 000000000..5d2cc9a52 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForConditionOperationsTest.java @@ -0,0 +1,44 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; + +import java.util.function.BiFunction; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import software.amazon.lambda.durable.context.BaseContextImpl; +import software.amazon.lambda.durable.model.WaitForConditionResult; + +class DurableWaitForConditionOperationsTest { + @AfterEach + void clearContext() { + BaseContextImpl.setCurrentContext(null); + } + + @Test + void conditionFunctionReceivesOnlyStateAndUsesStepContextFromTls() { + var context = mock(DurableContext.class); + var stepContext = mock(StepContext.class); + BaseContextImpl.setCurrentContext(context); + + DurableWaitForConditionOperations.waitForCondition("condition", String.class, state -> { + assertEquals(stepContext, StepContext.getCurrentContext()); + return WaitForConditionResult.stopPolling(state.toUpperCase()); + }); + + @SuppressWarnings("unchecked") + var check = (ArgumentCaptor>>) + (ArgumentCaptor) ArgumentCaptor.forClass(BiFunction.class); + verify(context).waitForCondition(eq("condition"), eq(String.class), check.capture()); + try (var ignored = BaseContextImpl.attachCurrentContext(stepContext)) { + assertEquals( + WaitForConditionResult.stopPolling("VALUE"), + check.getValue().apply("value", stepContext)); + } + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationsTest.java new file mode 100644 index 000000000..e26571dc8 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationsTest.java @@ -0,0 +1,38 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; + +import java.util.function.BiFunction; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import software.amazon.lambda.durable.context.BaseContextImpl; + +class DurableWithRetryOperationsTest { + @AfterEach + void clearContext() { + BaseContextImpl.setCurrentContext(null); + } + + @Test + void retryBodyUsesSupplierAndScopedAttempt() { + var context = mock(DurableContext.class); + BaseContextImpl.setCurrentContext(context); + + DurableWithRetryOperations.withRetry("retry", () -> WithRetryContext.getCurrentContext() + .getAttempt()); + + @SuppressWarnings("unchecked") + var operation = (ArgumentCaptor>) (ArgumentCaptor) + ArgumentCaptor.forClass(BiFunction.class); + verify(context).withRetry(eq("retry"), operation.capture()); + assertEquals(2, operation.getValue().apply(2, mock(DurableContext.class))); + assertThrows(IllegalStateException.class, WithRetryContext::getCurrentContext); + } +} From 96eab15c29c4b0e1693d22ab9ceb1db2ba257df9 Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 01:03:22 +0000 Subject: [PATCH 08/40] feat: expose durable future completion signals --- .../amazon/lambda/durable/DurableFuture.java | 15 ++++++-- .../operation/BaseDurableOperation.java | 9 +++++ .../lambda/durable/DurableFutureTest.java | 35 +++++++++++++++++++ 3 files changed, 57 insertions(+), 2 deletions(-) diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableFuture.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableFuture.java index 51e7163ef..d297da318 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableFuture.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableFuture.java @@ -5,7 +5,6 @@ import java.util.Arrays; import java.util.List; import java.util.concurrent.CompletableFuture; -import software.amazon.lambda.durable.operation.BaseDurableOperation; /** * A future representing the result of an asynchronous durable operation. @@ -26,6 +25,18 @@ public interface DurableFuture { */ T get(); + /** + * Returns a completion signal for this durable future. + * + *

The returned future completes when the durable operation completes. Completing or cancelling the returned + * future does not affect the durable operation. + * + * @return a future that signals durable operation completion + */ + default CompletableFuture completionFuture() { + throw new UnsupportedOperationException("This DurableFuture does not expose a completion signal"); + } + /** * Waits for all provided futures to complete and returns their results in order. * @@ -63,7 +74,7 @@ static List allOf(List> futures) { */ static Object anyOf(DurableFuture... futures) { return CompletableFuture.anyOf(Arrays.stream(futures) - .map(f -> ((BaseDurableOperation) f).getCompletionFuture()) + .map(f -> f.completionFuture().thenApply(ignored -> f)) .toArray(CompletableFuture[]::new)) .thenApply(o -> (DurableFuture) o) .join() diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java index f0391567a..6c26f35fc 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java @@ -97,6 +97,15 @@ public CompletableFuture getCompletionFuture() { return completionFuture; } + /** + * Returns a non-mutating completion signal for public {@code DurableFuture} combinators. + * + * @return a future that completes with this operation + */ + public CompletableFuture completionFuture() { + return completionFuture.thenApply(ignored -> null); + } + /** Gets the operation sub-type (e.g. RUN_IN_CHILD_CONTEXT, WAIT_FOR_CALLBACK). */ public OperationSubType getSubType() { return operationIdentifier.subType(); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java index 91338f9b1..0f3bcab74 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java @@ -6,6 +6,7 @@ import static org.mockito.Mockito.*; import java.util.List; +import java.util.concurrent.CompletableFuture; import org.junit.jupiter.api.Test; import software.amazon.lambda.durable.operation.SerializableDurableOperation; @@ -69,10 +70,44 @@ void allOfPropagatesException() { assertThrows(RuntimeException.class, () -> DurableFuture.allOf(op1, op2)); } + @Test + void anyOfSupportsPublicDurableFutureImplementations() { + var pending = new TestFuture<>("pending"); + var completed = new TestFuture<>("completed"); + completed.complete(); + + var result = DurableFuture.anyOf(pending, completed); + + assertEquals("completed", result); + } + @SuppressWarnings("unchecked") private SerializableDurableOperation mockOperation(T result) { SerializableDurableOperation op = mock(SerializableDurableOperation.class); when(op.get()).thenReturn(result); return op; } + + private static final class TestFuture implements DurableFuture { + private final T result; + private final CompletableFuture completion = new CompletableFuture<>(); + + private TestFuture(T result) { + this.result = result; + } + + @Override + public T get() { + return result; + } + + @Override + public CompletableFuture completionFuture() { + return completion; + } + + private void complete() { + completion.complete(null); + } + } } From d3c77499ffc546a02e39f63da6f802b12052f4ce Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 01:07:23 +0000 Subject: [PATCH 09/40] test: verify extensions across module boundary --- .../ExtensionOperationIntegrationTest.java | 87 ++++++++++++ .../lambda/durable/PluginIntegrationTest.java | 24 ++++ .../StaticOperationsIntegrationTest.java | 129 ++++++++++++++++++ .../durable/extension/PairOperations.java | 49 +++++++ .../amazon/lambda/durable/StepContext.java | 2 +- .../durable/context/DurableContextImpl.java | 2 +- .../durable/context/StepContextImpl.java | 4 +- 7 files changed, 293 insertions(+), 4 deletions(-) create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/extension/PairOperations.java diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java new file mode 100644 index 000000000..41b96a4a5 --- /dev/null +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java @@ -0,0 +1,87 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.assertEquals; +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 software.amazon.lambda.durable.DurableCoreOperations.step; +import static software.amazon.lambda.durable.extension.PairOperations.pairAsync; + +import java.nio.charset.StandardCharsets; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.util.HexFormat; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class ExtensionOperationIntegrationTest { + @Test + void reservedOperationsReplayWhenLaunchOrderChanges() { + var extensionExecutions = new AtomicInteger(); + var runner = + LocalDurableTestRunner.create(String.class, (input, context) -> pairAsync("pair", extensionExecutions) + .get()); + + var result = runner.runUntilComplete("input"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("LR", result.getResult(String.class)); + assertTrue(extensionExecutions.get() >= 2); + + assertEquals(hash("1"), result.getOperation("pair-left").getId()); + assertEquals(hash("2"), result.getOperation("pair-right").getId()); + assertEquals(hash("3"), result.getOperation("pair-pause").getId()); + } + + @Test + void staticOperationsUseCurrentContextAndRejectStepThreads() { + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { + assertSame(context, DurableContext.getCurrentContext()); + assertSame(context, ExtensionContext.getCurrentContext()); + return step("current-context", String.class, () -> { + var stepContext = StepContext.getCurrentContext(); + assertSame(stepContext, StepContext.getCurrentContext()); + return assertThrows(IllegalStateException.class, DurableContext::getCurrentContext) + .getMessage(); + }); + }); + + var result = runner.runUntilComplete("input"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertTrue(result.getResult(String.class).contains("step thread")); + assertThrows(IllegalStateException.class, DurableContext::getCurrentContext); + } + + @Test + void extensionCanExplicitlyCreateChildContext() { + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { + var outer = ExtensionContext.getCurrentContext(); + return outer.reserve("child").runInChildContext(String.class, () -> { + var child = ExtensionContext.getCurrentContext(); + assertNotSame(outer, child); + return child.reserve("value").step(String.class, () -> "nested"); + }); + }); + + var result = runner.runUntilComplete("input"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("nested", result.getResult(String.class)); + assertEquals(hash("1"), result.getOperation("child").getId()); + } + + private static String hash(String value) { + try { + var digest = MessageDigest.getInstance("SHA-256"); + return HexFormat.of().formatHex(digest.digest(value.getBytes(StandardCharsets.UTF_8))); + } catch (NoSuchAlgorithmException e) { + throw new AssertionError(e); + } + } +} diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java index 5bd794b9c..7f5c40019 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java @@ -168,6 +168,30 @@ void plugin_receivesOperationStart_forMultipleSteps() { assertTrue(opNames.contains("step-c")); } + @Test + void plugin_receivesPrimitiveLifecycleForCustomExtension() { + var plugin = new RecordingPlugin(); + var config = DurableConfig.builder().withPlugins(plugin).build(); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> customExtension(), + config); + + var result = runner.runUntilComplete("input"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + var operationNames = + plugin.operationStarts.stream().map(OperationInfo::name).toList(); + assertTrue(operationNames.contains("inner-step")); + assertFalse(operationNames.contains("custom-extension")); + } + + private static String customExtension() { + return ExtensionContext.getCurrentContext() + .reserve("inner-step") + .step(String.class, () -> "done"); + } + @Test void plugin_operationEnd_notFiredOnReplay() { var plugin = new RecordingPlugin(); diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java new file mode 100644 index 000000000..0b351c6dc --- /dev/null +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java @@ -0,0 +1,129 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.config.WaitForConditionConfig; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.model.WaitForConditionResult; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class StaticOperationsIntegrationTest { + @Test + void coreOperationsExposeStepAndChildContextsThroughTls() { + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { + var root = ExtensionContext.getCurrentContext(); + var step = DurableCoreOperations.step( + "step", + String.class, + () -> "attempt-" + StepContext.getCurrentContext().getAttempt()); + var child = DurableCoreOperations.runInChildContext("child", String.class, () -> { + assertNotSame(root, ExtensionContext.getCurrentContext()); + return DurableCoreOperations.step("child-step", String.class, () -> "child"); + }); + return step + ":" + child; + }); + + var result = runner.runUntilComplete("input"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("attempt-1:child", result.getResult(String.class)); + } + + @Test + void mapAndParallelExposeContextFreeUserFunctions() { + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { + var mapResult = DurableMapOperations.map( + "map", + List.of("a", "b"), + String.class, + item -> { + var index = MapItemContext.getCurrentContext().getIndex(); + return DurableCoreOperations.step("map-step", String.class, () -> item + index); + }); + + var branchFutures = new ArrayList>(); + try (var parallel = DurableParallelOperations.parallel("parallel")) { + branchFutures.add(parallel.branch( + "left", + String.class, + () -> DurableCoreOperations.step("branch-step", String.class, () -> "L"))); + branchFutures.add(parallel.branch( + "right", + String.class, + () -> DurableCoreOperations.step("branch-step", String.class, () -> "R"))); + } + return mapResult.results() + ":" + DurableFuture.allOf(branchFutures); + }); + + var result = runner.runUntilComplete("input"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("[a0, b1]:[L, R]", result.getResult(String.class)); + } + + @Test + void conditionAndRetryExposeGeneratedMetadataThroughTls() { + var retryExecutions = new AtomicInteger(); + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { + var condition = DurableWaitForConditionOperations.waitForCondition( + "condition", + Integer.class, + state -> { + assertNotNull(StepContext.getCurrentContext()); + return WaitForConditionResult.stopPolling(state + 1); + }, + WaitForConditionConfig.builder() + .initialState(0) + .build()); + var retry = DurableWithRetryOperations.withRetry("retry", () -> { + var attempt = WithRetryContext.getCurrentContext().getAttempt(); + retryExecutions.incrementAndGet(); + if (attempt == 1) { + throw new IllegalStateException("retry"); + } + return attempt; + }); + return condition + ":" + retry; + }); + + var result = runner.runUntilComplete("input"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("1:2", result.getResult(String.class)); + assertTrue(retryExecutions.get() >= 2); + } + + @Test + void waitForCallbackExposesCallbackIdThroughTls() { + var submittedId = new AtomicReference(); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> DurableWaitForCallbackOperations.waitForCallback( + "approval", + String.class, + () -> submittedId.set( + WaitForCallbackContext.getCurrentContext().getCallbackId()))); + + var pending = runner.run("input"); + + assertEquals(ExecutionStatus.PENDING, pending.getStatus()); + var callbackId = runner.getCallbackId("approval-callback"); + assertEquals(callbackId, submittedId.get()); + runner.completeCallback(callbackId, "\"approved\""); + + var completed = runner.run("input"); + + assertEquals(ExecutionStatus.SUCCEEDED, completed.getStatus()); + assertEquals("approved", completed.getResult(String.class)); + } +} diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/extension/PairOperations.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/extension/PairOperations.java new file mode 100644 index 000000000..ac6444441 --- /dev/null +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/extension/PairOperations.java @@ -0,0 +1,49 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.extension; + +import java.time.Duration; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicInteger; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.ExtensionContext; + +/** Example extension library implemented only with public SDK contracts. */ +public final class PairOperations { + private PairOperations() {} + + public static DurableFuture pairAsync(String name, AtomicInteger extensionExecutions) { + var extension = ExtensionContext.getCurrentContext(); + var left = extension.reserve(name + "-left"); + var right = extension.reserve(name + "-right"); + var pause = extension.reserve(name + "-pause"); + + DurableFuture leftFuture; + DurableFuture rightFuture; + if (extensionExecutions.getAndIncrement() % 2 == 0) { + leftFuture = left.stepAsync(String.class, () -> "L"); + rightFuture = right.stepAsync(String.class, () -> "R"); + } else { + rightFuture = right.stepAsync(String.class, () -> "R"); + leftFuture = left.stepAsync(String.class, () -> "L"); + } + + return new PairFuture(leftFuture, rightFuture, pause.waitAsync(Duration.ofSeconds(1))); + } + + private record PairFuture( + DurableFuture left, DurableFuture right, DurableFuture pause) + implements DurableFuture { + @Override + public String get() { + pause.get(); + return left.get() + right.get(); + } + + @Override + public CompletableFuture completionFuture() { + return CompletableFuture.allOf( + left.completionFuture(), right.completionFuture(), pause.completionFuture()); + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/StepContext.java b/sdk/src/main/java/software/amazon/lambda/durable/StepContext.java index f2469b8ca..261b53a9a 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/StepContext.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/StepContext.java @@ -5,7 +5,7 @@ import software.amazon.lambda.durable.context.BaseContext; public interface StepContext extends BaseContext { - /** Returns the current retry attempt number (0-based). */ + /** Returns the current retry attempt number (1-based). */ int getAttempt(); /** Returns the step context attached to the current SDK-managed thread. */ diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java index 94e4fd3f7..303cbcdf6 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java @@ -119,7 +119,7 @@ public DurableContextImpl createChildContext(String childContextId, String child * * @param stepOperationId the ID of the step operation (used for thread registration) * @param stepOperationName the name of the step operation - * @param attempt the current retry attempt number (0-based) + * @param attempt the current retry attempt number (1-based) * @return a new StepContext instance */ public StepContextImpl createStepContext(String stepOperationId, String stepOperationName, int attempt) { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/StepContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/StepContextImpl.java index d2429dd8e..e406b7e56 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/StepContextImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/StepContextImpl.java @@ -25,7 +25,7 @@ public class StepContextImpl extends BaseContextImpl implements StepContext { * @param lambdaContext AWS Lambda runtime context * @param stepOperationId Unique identifier for this context instance that equals to step operation id * @param stepOperationName the name of the step operation - * @param attempt the current retry attempt number (0-based) + * @param attempt the current retry attempt number (1-based) */ protected StepContextImpl( ExecutionManager executionManager, @@ -38,7 +38,7 @@ protected StepContextImpl( this.attempt = attempt; } - /** Returns the current retry attempt number (0-based). */ + /** Returns the current retry attempt number (1-based). */ @Override public int getAttempt() { return attempt; From 781655b49f09091a992f3b457448b7412b5cf00d Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 01:22:32 +0000 Subject: [PATCH 10/40] docs: publish custom extension operations guide --- README.md | 1 + docs/adr/006-custom-extension-operations.md | 3 +- docs/advanced/extensions.md | 183 ++++++++++++++++++ .../lambda/durable/PluginIntegrationTest.java | 9 +- .../StaticOperationsIntegrationTest.java | 16 +- .../durable/extension/PairOperations.java | 6 +- .../lambda/durable/DurableCoreOperations.java | 6 +- .../lambda/durable/DurableMapOperations.java | 24 +-- .../DurableWaitForCallbackOperations.java | 6 +- .../durable/DurableWithRetryOperations.java | 3 +- .../lambda/durable/ExtensionContext.java | 4 +- .../lambda/durable/ExtensionOperation.java | 12 +- .../context/ExtensionOperationImpl.java | 3 +- .../durable/DurableCoreOperationsTest.java | 32 ++- .../durable/DurableMapOperationsTest.java | 4 +- .../DurableParallelOperationsTest.java | 16 +- .../DurableWaitForCallbackOperationsTest.java | 13 +- .../DurableWithRetryOperationsTest.java | 8 +- .../context/ExtensionOperationImplTest.java | 25 +-- 19 files changed, 254 insertions(+), 120 deletions(-) create mode 100644 docs/advanced/extensions.md diff --git a/README.md b/README.md index 9a3c25433..c84e2795e 100644 --- a/README.md +++ b/README.md @@ -111,6 +111,7 @@ See [Deploy Lambda durable functions with Infrastructure as Code](https://docs.a - [Configuration](docs/advanced/configuration.md) - Customize SDK behaviour - [Error Handling](docs/advanced/error-handling.md) - SDK exceptions for handling failures - [Logging](docs/advanced/logging.md) - How to use DurableLogger +- [Custom Extensions](docs/advanced/extensions.md) - Build reusable operations from supported primitives - [Migrating from 1.x to 2.x](docs/migration-1.x-to-2.x.md) - Upgrade guide for breaking changes since `v1.2.1` - [Release Process](RELEASE.md) - Prepare and publish Maven releases - [Testing](docs/advanced/testing.md) - Utilities for local development and cloud-based integration testing diff --git a/docs/adr/006-custom-extension-operations.md b/docs/adr/006-custom-extension-operations.md index 32c1186ac..d131ae26f 100644 --- a/docs/adr/006-custom-extension-operations.md +++ b/docs/adr/006-custom-extension-operations.md @@ -1,6 +1,7 @@ # ADR-006: Public API for Custom Extension Operations -**Status:** Proposed +**Status:** Accepted + **Date:** 2026-08-10 ## Context diff --git a/docs/advanced/extensions.md b/docs/advanced/extensions.md new file mode 100644 index 000000000..81a5dce42 --- /dev/null +++ b/docs/advanced/extensions.md @@ -0,0 +1,183 @@ +# Custom extension operations + +Extension operations are ordinary static Java methods that compose SDK-owned durable primitives. They can live in a +separate Maven module without defining backend operation types, sending checkpoint updates, or depending on SDK +implementation packages. + +Application code calls only the extension's API: + +```java +import static com.example.durable.PairOperations.pairAsync; + +var result = pairAsync("pair", left, right).get(); +``` + +The extension retrieves the active scope internally: + +```java +public final class PairOperations { + private PairOperations() {} + + public static DurableFuture pairAsync( + String name, Supplier leftFunction, Supplier rightFunction) { + var extension = ExtensionContext.getCurrentContext(); + var left = extension.reserve(name + "-left"); + var right = extension.reserve(name + "-right"); + + var leftFuture = left.stepAsync(String.class, leftFunction); + var rightFuture = right.stepAsync(String.class, rightFunction); + return new PairFuture(leftFuture, rightFuture); + } +} +``` + +There is no extension registration API and no automatic child-context boundary. The extension chooses whether to +compose primitives in the current scope or explicitly create a child context. + +## Static operation APIs + +New code can use context-free static facades: + +| Facade | Operations | +| --- | --- | +| `DurableCoreOperations` | `step`, `wait`, chained `invoke`, callbacks, child contexts | +| `DurableMapOperations` | `map`, `mapAsync` | +| `DurableParallelOperations` | `parallel` | +| `DurableWaitForCallbackOperations` | `waitForCallback`, `waitForCallbackAsync` | +| `DurableWaitForConditionOperations` | `waitForCondition`, `waitForConditionAsync` | +| `DurableWithRetryOperations` | `withRetry`, `withRetryAsync` | + +The existing `DurableContext` instance methods and callback signatures remain supported for backward compatibility. + +User functions in the static APIs do not receive SDK context objects: + +```java +var result = DurableCoreOperations.step("process", Result.class, () -> { + var step = StepContext.getCurrentContext(); + return process(step.getAttempt()); +}); +``` + +```java +var result = DurableMapOperations.map("process", items, Result.class, item -> { + var index = MapItemContext.getCurrentContext().getIndex(); + return process(item, index); +}); +``` + +```java +var result = DurableWaitForCallbackOperations.waitForCallback( + "approval", + Approval.class, + () -> submit(WaitForCallbackContext.getCurrentContext().getCallbackId())); +``` + +```java +var result = DurableWithRetryOperations.withRetry("transaction", () -> { + var attempt = WithRetryContext.getCurrentContext().getAttempt(); + return executeAttempt(attempt); +}); +``` + +Parallel branch functions are `Supplier`. Wait-for-condition functions receive only the durable state value and +obtain attempt metadata from `StepContext.getCurrentContext()`. + +## Current context scopes + +`DurableContext.getCurrentContext()` and `ExtensionContext.getCurrentContext()` are available on SDK-managed handler +and child-context threads. `StepContext.getCurrentContext()` is available inside step and wait-for-condition user +functions. + +`MapItemContext`, `WaitForCallbackContext`, and `WithRetryContext` are available only inside their corresponding user +function. Nested scopes restore the previous context when they close. + +Operation-specific TLS is not automatically propagated into a nested primitive's separate user-function thread. Read +the metadata in its owning function and capture any application value needed by the nested operation: + +```java +var result = DurableMapOperations.map("process", items, Result.class, item -> { + var index = MapItemContext.getCurrentContext().getIndex(); + return DurableCoreOperations.step("process-item", Result.class, () -> process(item, index)); +}); +``` + +Current context is not propagated to application-created threads. Durable primitives must be created from an +SDK-managed durable context thread. + +## Primitive reservations + +Extensions with deterministic call order can use `DurableCoreOperations` directly. Schedulers whose registration +order is deterministic but launch order may vary use `ExtensionContext.reserve(name)`. + +Each reservation immediately consumes the next sequential operation ID and returns an opaque, one-shot +`ExtensionOperation`: + +```java +var extension = ExtensionContext.getCurrentContext(); +var first = extension.reserve("first"); +var second = extension.reserve("second"); + +// Launch order can differ from reservation order. +var secondResult = second.stepAsync(String.class, () -> runSecond()); +var firstResult = first.stepAsync(String.class, () -> runFirst()); +``` + +A reservation can create exactly one primitive: step, wait, chained invoke, callback, or child context. Reuse throws +`IllegalStateException`. Raw operation IDs are never exposed. + +Create reservations in the same order on every replay. Reordering, inserting, or removing reservations is a workflow +compatibility change because it can associate existing checkpoints with different logical primitives. Launching +already reserved operations in a different order is supported. + +## Explicit child contexts + +An extension creates a child context only when its own semantics require isolation: + +```java +var result = ExtensionContext.getCurrentContext() + .reserve("isolated-work") + .runInChildContext(Result.class, () -> executeIsolatedWork()); +``` + +Inside the supplier, `DurableContext.getCurrentContext()` and `ExtensionContext.getCurrentContext()` return the child +context. + +## Custom durable futures + +An asynchronous extension may return an SDK primitive future or implement `DurableFuture`. Custom composed futures +that participate in `DurableFuture.anyOf` override `completionFuture()`: + +```java +private record PairFuture(DurableFuture left, DurableFuture right) + implements DurableFuture { + @Override + public String get() { + return left.get() + right.get(); + } + + @Override + public CompletableFuture completionFuture() { + return CompletableFuture.allOf(left.completionFuture(), right.completionFuture()); + } +} +``` + +Completing or cancelling the returned completion signal must not mutate the underlying durable operations. + +## Plugins and failures + +Extensions do not create an automatic plugin lifecycle event or checkpoint boundary. Plugins observe the primitive +operations created by the extension. If the extension explicitly creates a child context, plugins also observe that +child-context operation. + +Serialization, suspension, replay, cancellation, failures, and checkpointing retain the semantics of the underlying +primitive operations. + +## Module compatibility + +An extension Maven module should depend only on the public SDK artifact and import public types under +`software.amazon.lambda.durable`. Do not import SDK implementation packages such as `context`, `execution`, or +`operation`. + +The supported extension contracts are `ExtensionContext`, `ExtensionOperation`, the static operation facades, the +typed TLS contexts, and `DurableFuture.completionFuture()`. diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java index 7f5c40019..e89334278 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java @@ -172,10 +172,7 @@ void plugin_receivesOperationStart_forMultipleSteps() { void plugin_receivesPrimitiveLifecycleForCustomExtension() { var plugin = new RecordingPlugin(); var config = DurableConfig.builder().withPlugins(plugin).build(); - var runner = LocalDurableTestRunner.create( - String.class, - (input, context) -> customExtension(), - config); + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> customExtension(), config); var result = runner.runUntilComplete("input"); @@ -187,9 +184,7 @@ void plugin_receivesPrimitiveLifecycleForCustomExtension() { } private static String customExtension() { - return ExtensionContext.getCurrentContext() - .reserve("inner-step") - .step(String.class, () -> "done"); + return ExtensionContext.getCurrentContext().reserve("inner-step").step(String.class, () -> "done"); } @Test diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java index 0b351c6dc..0bd044632 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java @@ -42,14 +42,10 @@ void coreOperationsExposeStepAndChildContextsThroughTls() { @Test void mapAndParallelExposeContextFreeUserFunctions() { var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { - var mapResult = DurableMapOperations.map( - "map", - List.of("a", "b"), - String.class, - item -> { - var index = MapItemContext.getCurrentContext().getIndex(); - return DurableCoreOperations.step("map-step", String.class, () -> item + index); - }); + var mapResult = DurableMapOperations.map("map", List.of("a", "b"), String.class, item -> { + var index = MapItemContext.getCurrentContext().getIndex(); + return DurableCoreOperations.step("map-step", String.class, () -> item + index); + }); var branchFutures = new ArrayList>(); try (var parallel = DurableParallelOperations.parallel("parallel")) { @@ -82,9 +78,7 @@ void conditionAndRetryExposeGeneratedMetadataThroughTls() { assertNotNull(StepContext.getCurrentContext()); return WaitForConditionResult.stopPolling(state + 1); }, - WaitForConditionConfig.builder() - .initialState(0) - .build()); + WaitForConditionConfig.builder().initialState(0).build()); var retry = DurableWithRetryOperations.withRetry("retry", () -> { var attempt = WithRetryContext.getCurrentContext().getAttempt(); retryExecutions.incrementAndGet(); diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/extension/PairOperations.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/extension/PairOperations.java index ac6444441..d7a49ed5f 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/extension/PairOperations.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/extension/PairOperations.java @@ -31,8 +31,7 @@ public static DurableFuture pairAsync(String name, AtomicInteger extensi return new PairFuture(leftFuture, rightFuture, pause.waitAsync(Duration.ofSeconds(1))); } - private record PairFuture( - DurableFuture left, DurableFuture right, DurableFuture pause) + private record PairFuture(DurableFuture left, DurableFuture right, DurableFuture pause) implements DurableFuture { @Override public String get() { @@ -42,8 +41,7 @@ public String get() { @Override public CompletableFuture completionFuture() { - return CompletableFuture.allOf( - left.completionFuture(), right.completionFuture(), pause.completionFuture()); + return CompletableFuture.allOf(left.completionFuture(), right.completionFuture(), pause.completionFuture()); } } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableCoreOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableCoreOperations.java index a929b5e33..e4bfbad51 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableCoreOperations.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableCoreOperations.java @@ -105,8 +105,7 @@ public static DurableCallbackFuture createCallback(String name, TypeToken return currentContext().createCallback(name, resultType); } - public static DurableCallbackFuture createCallback( - String name, Class resultType, CallbackConfig config) { + public static DurableCallbackFuture createCallback(String name, Class resultType, CallbackConfig config) { return currentContext().createCallback(name, resultType, config); } @@ -133,8 +132,7 @@ public static T runInChildContext( return currentContext().runInChildContext(name, resultType, ignored -> function.get(), config); } - public static DurableFuture runInChildContextAsync( - String name, Class resultType, Supplier function) { + public static DurableFuture runInChildContextAsync(String name, Class resultType, Supplier function) { return currentContext().runInChildContextAsync(name, resultType, ignored -> function.get()); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableMapOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableMapOperations.java index 3ea76bbd1..e71dc1d9a 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableMapOperations.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableMapOperations.java @@ -23,20 +23,12 @@ public static MapResult map( } public static MapResult map( - String name, - Collection items, - Class resultType, - Function function, - MapConfig config) { + String name, Collection items, Class resultType, Function function, MapConfig config) { return currentContext().map(name, items, resultType, adapt(function), config); } public static MapResult map( - String name, - Collection items, - TypeToken resultType, - Function function, - MapConfig config) { + String name, Collection items, TypeToken resultType, Function function, MapConfig config) { return currentContext().map(name, items, resultType, adapt(function), config); } @@ -51,20 +43,12 @@ public static DurableFuture> mapAsync( } public static DurableFuture> mapAsync( - String name, - Collection items, - Class resultType, - Function function, - MapConfig config) { + String name, Collection items, Class resultType, Function function, MapConfig config) { return currentContext().mapAsync(name, items, resultType, adapt(function), config); } public static DurableFuture> mapAsync( - String name, - Collection items, - TypeToken resultType, - Function function, - MapConfig config) { + String name, Collection items, TypeToken resultType, Function function, MapConfig config) { return currentContext().mapAsync(name, items, resultType, adapt(function), config); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForCallbackOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForCallbackOperations.java index fbefa764a..f63c433e6 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForCallbackOperations.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForCallbackOperations.java @@ -28,13 +28,11 @@ public static T waitForCallback( return currentContext().waitForCallback(name, resultType, adapt(submitter), config); } - public static DurableFuture waitForCallbackAsync( - String name, Class resultType, Runnable submitter) { + public static DurableFuture waitForCallbackAsync(String name, Class resultType, Runnable submitter) { return currentContext().waitForCallbackAsync(name, resultType, adapt(submitter)); } - public static DurableFuture waitForCallbackAsync( - String name, TypeToken resultType, Runnable submitter) { + public static DurableFuture waitForCallbackAsync(String name, TypeToken resultType, Runnable submitter) { return currentContext().waitForCallbackAsync(name, resultType, adapt(submitter)); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableWithRetryOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableWithRetryOperations.java index 75151eb31..b65b2ee67 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableWithRetryOperations.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableWithRetryOperations.java @@ -23,8 +23,7 @@ public static DurableFuture withRetryAsync(String name, Supplier opera return currentContext().withRetryAsync(name, adapt(operation)); } - public static DurableFuture withRetryAsync( - String name, Supplier operation, WithRetryConfig config) { + public static DurableFuture withRetryAsync(String name, Supplier operation, WithRetryConfig config) { return currentContext().withRetryAsync(name, adapt(operation), config); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContext.java b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContext.java index 4d5379a19..ae34514ce 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContext.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContext.java @@ -7,8 +7,8 @@ /** * Public context available to custom extension operations. * - *

This interface exposes replay state and opaque primitive reservations without exposing checkpoint internals or - * raw operation IDs. + *

This interface exposes replay state and opaque primitive reservations without exposing checkpoint internals or raw + * operation IDs. */ public interface ExtensionContext extends BaseContext { /** diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java index 0a29376c9..efbd04ffa 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java @@ -57,8 +57,7 @@ default T invoke(String functionName, U payload, Class resultType) { } default T invoke(String functionName, U payload, TypeToken resultType) { - return invoke( - functionName, payload, resultType, InvokeConfig.builder().build()); + return invoke(functionName, payload, resultType, InvokeConfig.builder().build()); } default T invoke(String functionName, U payload, Class resultType, InvokeConfig config) { @@ -83,8 +82,7 @@ default DurableFuture invokeAsync( return invokeAsync(functionName, payload, TypeToken.get(resultType), config); } - DurableFuture invokeAsync( - String functionName, U payload, TypeToken resultType, InvokeConfig config); + DurableFuture invokeAsync(String functionName, U payload, TypeToken resultType, InvokeConfig config); default DurableCallbackFuture createCallback(Class resultType) { return createCallback(TypeToken.get(resultType)); @@ -109,13 +107,11 @@ default T runInChildContext(TypeToken resultType, Supplier function) { resultType, function, RunInChildContextConfig.builder().build()); } - default T runInChildContext( - Class resultType, Supplier function, RunInChildContextConfig config) { + default T runInChildContext(Class resultType, Supplier function, RunInChildContextConfig config) { return runInChildContext(TypeToken.get(resultType), function, config); } - default T runInChildContext( - TypeToken resultType, Supplier function, RunInChildContextConfig config) { + default T runInChildContext(TypeToken resultType, Supplier function, RunInChildContextConfig config) { return runInChildContextAsync(resultType, function, config).get(); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java index bf11713f3..ab34486a9 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java @@ -55,8 +55,7 @@ public DurableCallbackFuture createCallback(TypeToken resultType, Call public DurableFuture runInChildContextAsync( TypeToken resultType, Supplier function, RunInChildContextConfig config) { claim(); - return context.runInChildContextAsyncWithId( - operationId, name, resultType, ignored -> function.get(), config); + return context.runInChildContextAsyncWithId(operationId, name, resultType, ignored -> function.get(), config); } private void claim() { diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableCoreOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableCoreOperationsTest.java index b31e8afd4..423aa98b6 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableCoreOperationsTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableCoreOperationsTest.java @@ -33,8 +33,8 @@ void stepAcceptsContextFreeSupplier() { DurableCoreOperations.step("step", String.class, () -> "result"); @SuppressWarnings("unchecked") - var function = (ArgumentCaptor>) (ArgumentCaptor) ArgumentCaptor.forClass( - Function.class); + var function = (ArgumentCaptor>) + (ArgumentCaptor) ArgumentCaptor.forClass(Function.class); verify(context).step(eq("step"), eq(String.class), function.capture()); assertEquals("result", function.getValue().apply(mock(StepContext.class))); } @@ -47,8 +47,8 @@ void childContextAcceptsContextFreeSupplier() { DurableCoreOperations.runInChildContext("child", String.class, () -> "result"); @SuppressWarnings("unchecked") - var function = (ArgumentCaptor>) (ArgumentCaptor) ArgumentCaptor.forClass( - Function.class); + var function = (ArgumentCaptor>) + (ArgumentCaptor) ArgumentCaptor.forClass(Function.class); verify(context).runInChildContext(eq("child"), eq(String.class), function.capture()); assertEquals("result", function.getValue().apply(mock(DurableContext.class))); } @@ -63,19 +63,12 @@ void coreValueOperationsDelegateToCurrentContext() { @SuppressWarnings("unchecked") var callbackFuture = (DurableCallbackFuture) mock(DurableCallbackFuture.class); when(context.waitAsync("wait", duration)).thenReturn(waitFuture); - when(context.invokeAsync("invoke", "function", "payload", String.class)) - .thenReturn(invokeFuture); - when(context.createCallback("callback", String.class)) - .thenReturn(callbackFuture); + when(context.invokeAsync("invoke", "function", "payload", String.class)).thenReturn(invokeFuture); + when(context.createCallback("callback", String.class)).thenReturn(callbackFuture); assertEquals(waitFuture, DurableCoreOperations.waitAsync("wait", duration)); - assertEquals( - invokeFuture, - DurableCoreOperations.invokeAsync( - "invoke", "function", "payload", String.class)); - assertEquals( - callbackFuture, - DurableCoreOperations.createCallback("callback", String.class)); + assertEquals(invokeFuture, DurableCoreOperations.invokeAsync("invoke", "function", "payload", String.class)); + assertEquals(callbackFuture, DurableCoreOperations.createCallback("callback", String.class)); } @Test @@ -86,19 +79,16 @@ void configuredSupplierOverloadsDelegateToCurrentContext() { var childConfig = RunInChildContextConfig.builder().build(); DurableCoreOperations.stepAsync("step", new TypeToken() {}, () -> "step", stepConfig); - DurableCoreOperations.runInChildContextAsync( - "child", new TypeToken() {}, () -> "child", childConfig); + DurableCoreOperations.runInChildContextAsync("child", new TypeToken() {}, () -> "child", childConfig); verify(context).stepAsync(eq("step"), any(TypeToken.class), any(Function.class), eq(stepConfig)); - verify(context) - .runInChildContextAsync(eq("child"), any(TypeToken.class), any(Function.class), eq(childConfig)); + verify(context).runInChildContextAsync(eq("child"), any(TypeToken.class), any(Function.class), eq(childConfig)); } @Test void coreOperationsFailOutsideDurableContext() { assertThrows( - IllegalStateException.class, - () -> DurableCoreOperations.step("step", String.class, () -> "result")); + IllegalStateException.class, () -> DurableCoreOperations.step("step", String.class, () -> "result")); } @SuppressWarnings("unchecked") diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationsTest.java index 7dc2d4c83..334910c62 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationsTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationsTest.java @@ -31,8 +31,8 @@ void mapExposesItemIndexThroughScopedContext() { }); @SuppressWarnings("unchecked") - var function = (ArgumentCaptor>) (ArgumentCaptor) - ArgumentCaptor.forClass(DurableContext.MapFunction.class); + var function = (ArgumentCaptor>) + (ArgumentCaptor) ArgumentCaptor.forClass(DurableContext.MapFunction.class); verify(context).map(eq("map"), eq(List.of("value")), eq(String.class), function.capture()); assertEquals("VALUE", function.getValue().apply("value", 3, mock(DurableContext.class))); assertThrows(IllegalStateException.class, MapItemContext::getCurrentContext); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationsTest.java index b49b5eef0..413b7b5dd 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationsTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationsTest.java @@ -31,11 +31,7 @@ void parallelBranchesAcceptContextFreeSuppliers() { var branchFuture = mockStringFuture(); BaseContextImpl.setCurrentContext(context); when(context.parallel("parallel")).thenReturn(parallel); - when(parallel.branch( - eq("branch"), - any(TypeToken.class), - any(Function.class), - any(ParallelBranchConfig.class))) + when(parallel.branch(eq("branch"), any(TypeToken.class), any(Function.class), any(ParallelBranchConfig.class))) .thenReturn(branchFuture); var result = DurableParallelOperations.parallel("parallel"); @@ -44,14 +40,10 @@ void parallelBranchesAcceptContextFreeSuppliers() { assertSame(parallel, result); assertSame(branchFuture, resultFuture); @SuppressWarnings("unchecked") - var function = (ArgumentCaptor>) (ArgumentCaptor) ArgumentCaptor.forClass( - Function.class); + var function = (ArgumentCaptor>) + (ArgumentCaptor) ArgumentCaptor.forClass(Function.class); verify(parallel) - .branch( - eq("branch"), - any(TypeToken.class), - function.capture(), - any(ParallelBranchConfig.class)); + .branch(eq("branch"), any(TypeToken.class), function.capture(), any(ParallelBranchConfig.class)); assertEquals("result", function.getValue().apply(mock(DurableContext.class))); } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationsTest.java index 65b825fee..9da335025 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationsTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationsTest.java @@ -25,13 +25,16 @@ void callbackSubmitterUsesRunnableAndScopedCallbackId() { var context = mock(DurableContext.class); BaseContextImpl.setCurrentContext(context); - DurableWaitForCallbackOperations.waitForCallback("callback", String.class, () -> assertEquals( - "callback-id", - WaitForCallbackContext.getCurrentContext().getCallbackId())); + DurableWaitForCallbackOperations.waitForCallback( + "callback", + String.class, + () -> assertEquals( + "callback-id", + WaitForCallbackContext.getCurrentContext().getCallbackId())); @SuppressWarnings("unchecked") - var submitter = (ArgumentCaptor>) (ArgumentCaptor) - ArgumentCaptor.forClass(BiConsumer.class); + var submitter = (ArgumentCaptor>) + (ArgumentCaptor) ArgumentCaptor.forClass(BiConsumer.class); verify(context).waitForCallback(eq("callback"), eq(String.class), submitter.capture()); submitter.getValue().accept("callback-id", mock(StepContext.class)); assertThrows(IllegalStateException.class, WaitForCallbackContext::getCurrentContext); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationsTest.java index e26571dc8..e69c8aaa4 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationsTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationsTest.java @@ -25,12 +25,12 @@ void retryBodyUsesSupplierAndScopedAttempt() { var context = mock(DurableContext.class); BaseContextImpl.setCurrentContext(context); - DurableWithRetryOperations.withRetry("retry", () -> WithRetryContext.getCurrentContext() - .getAttempt()); + DurableWithRetryOperations.withRetry( + "retry", () -> WithRetryContext.getCurrentContext().getAttempt()); @SuppressWarnings("unchecked") - var operation = (ArgumentCaptor>) (ArgumentCaptor) - ArgumentCaptor.forClass(BiFunction.class); + var operation = (ArgumentCaptor>) + (ArgumentCaptor) ArgumentCaptor.forClass(BiFunction.class); verify(context).withRetry(eq("retry"), operation.capture()); assertEquals(2, operation.getValue().apply(2, mock(DurableContext.class))); assertThrows(IllegalStateException.class, WithRetryContext::getCurrentContext); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java index 8984336ed..ecc050d4a 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java @@ -62,14 +62,19 @@ void reservedStepAdaptsSupplierToStepFunction() { .thenReturn(future); var operation = new ExtensionOperationImpl(context, "1", "step"); - assertEquals(future, operation.stepAsync(resultType, () -> { - called.set(true); - return "result"; - }, config)); + assertEquals( + future, + operation.stepAsync( + resultType, + () -> { + called.set(true); + return "result"; + }, + config)); @SuppressWarnings("unchecked") - var function = (ArgumentCaptor>) (ArgumentCaptor) ArgumentCaptor.forClass( - Function.class); + var function = (ArgumentCaptor>) + (ArgumentCaptor) ArgumentCaptor.forClass(Function.class); verify(context).stepAsyncWithId(eq("1"), eq("step"), eq(resultType), function.capture(), eq(config)); assertEquals("result", function.getValue().apply(mock(StepContext.class))); assertEquals(true, called.get()); @@ -119,8 +124,8 @@ void reservedChildContextAdaptsSupplierToChildFunction() { assertEquals(future, operation.runInChildContextAsync(resultType, () -> "result", config)); @SuppressWarnings("unchecked") - var function = (ArgumentCaptor>) (ArgumentCaptor) ArgumentCaptor.forClass( - Function.class); + var function = (ArgumentCaptor>) + (ArgumentCaptor) ArgumentCaptor.forClass(Function.class); verify(context) .runInChildContextAsyncWithId(eq("1"), eq("child"), eq(resultType), function.capture(), eq(config)); assertEquals("result", function.getValue().apply(mock(DurableContext.class))); @@ -135,9 +140,7 @@ void reservationCanOnlyExecuteOnceAcrossPrimitiveSelectors() { operation.waitAsync(duration); - assertThrows( - IllegalStateException.class, - () -> operation.stepAsync(String.class, () -> "second")); + assertThrows(IllegalStateException.class, () -> operation.stepAsync(String.class, () -> "second")); } @SuppressWarnings("unchecked") From 44dd8acc9a390d267f2494e2eb82f4ea8d52e6a6 Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 03:01:49 +0000 Subject: [PATCH 11/40] docs: design built-in extension migration --- ...8-10-migrate-built-in-extensions-design.md | 570 ++++++++++++++++++ 1 file changed, 570 insertions(+) create mode 100644 docs/superpowers/specs/2026-08-10-migrate-built-in-extensions-design.md diff --git a/docs/superpowers/specs/2026-08-10-migrate-built-in-extensions-design.md b/docs/superpowers/specs/2026-08-10-migrate-built-in-extensions-design.md new file mode 100644 index 000000000..532b83895 --- /dev/null +++ b/docs/superpowers/specs/2026-08-10-migrate-built-in-extensions-design.md @@ -0,0 +1,570 @@ +# Migrate Built-In Operations to the Extension API + +## Objective + +Rewrite the SDK's existing non-primitive operations as built-in extensions using the public extension operation +model. Preserve every existing user-facing interface, configuration type, overload, result, exception, checkpoint +shape, replay behavior, plugin event, and concurrency behavior. + +The migrated operation families are: + +- map +- parallel +- wait for callback +- wait for condition +- with retry + +The existing `DurableContext` methods remain supported compatibility APIs. They delegate to the same implementations +used by the static built-in extension facades. + +## Compatibility Boundary + +The following existing public APIs remain unchanged: + +- `DurableContext` +- `ParallelDurableFuture` +- `MapConfig` +- `ParallelConfig` +- `WaitForCallbackConfig` +- `WaitForConditionConfig` +- `WithRetryConfig` +- `StepConfig` +- `RunInChildContextConfig` +- all existing result and exception types + +This migration may add public extension-specific interfaces, overloads, and configuration types. It must not add +methods to the existing operation interfaces or fields to their existing configuration types. + +The following observable behavior remains unchanged: + +- operation IDs and parent-child ID namespaces +- operation types, subtype strings, names, and tree shape +- checkpoint action sequences and replay validation +- serialized result and failure payloads +- map and parallel completion decisions +- concurrency limits and skipped item behavior +- nested and flat concurrency modes +- large-result replay-children behavior +- wait-for-condition state, attempts, and delays +- wait-for-callback failure and timeout translation +- retry backoff names and virtual-context behavior +- plugin operation and user-function event ordering + +## Primitive State Machines + +Extension authors may select operation subtype strings, but they may not define checkpoint state machines. + +Each primitive retains its SDK-owned lifecycle: + +| Primitive selector | Backend operation type | SDK-owned state machine | +| --- | --- | --- | +| step | `STEP` | start, retry, ready, succeed, fail | +| wait | `WAIT` | start, poll, succeed | +| invoke | `CHAINED_INVOKE` | start, poll, succeed, fail | +| callback | `CALLBACK` | start, poll, succeed, fail, timeout | +| child context | `CONTEXT` | start, execute/replay children, succeed, fail | + +The primitive selector determines the backend operation type. The supplied subtype is metadata used for checkpoint +validation, plugins, exception translation, and execution history. + +Subtype strings must be non-null and nonblank. The SDK does not restrict them to an allow-list because the backend +accepts arbitrary subtype strings. + +## Subtype-Aware Reservations + +`ExtensionOperation` gains subtype-aware overloads for every primitive. Existing overloads remain and use the current +standard subtype strings. + +Representative asynchronous signatures: + +```java + DurableFuture stepAsync( + String subType, + TypeToken resultType, + Supplier function, + StepConfig config); + +DurableFuture waitAsync( + String subType, + Duration duration); + + DurableFuture invokeAsync( + String subType, + String functionName, + U payload, + TypeToken resultType, + InvokeConfig config); + + DurableCallbackFuture createCallback( + String subType, + TypeToken resultType, + CallbackConfig config); + + DurableFuture runInChildContextAsync( + String subType, + TypeToken resultType, + Supplier function, + RunInChildContextConfig config); +``` + +The existing no-subtype methods delegate with these values: + +- `Step` +- `Wait` +- `ChainedInvoke` +- `Callback` +- `RunInChildContext` + +Synchronous, `Class`, and default-configuration overloads remain default methods. + +## Custom Local Operation IDs + +`ExtensionContext` retains sequential reservation and adds a custom-local-ID overload: + +```java +ExtensionOperation reserve(String name); + +ExtensionOperation reserve(String name, String localOperationId); +``` + +The custom value replaces the generated sequence number for that reservation. It is not the final backend operation +ID. + +Custom local IDs must be non-null and nonblank. They are otherwise treated as opaque UTF-8 strings. + +The SDK constructs the final ID using the current context namespace: + +```text +root context: sha256(localOperationId) +child context: sha256(parentContextId + "-" + localOperationId) +``` + +Every reservation occupies one position in the context's reservation sequence: + +1. A sequential reservation advances the counter until it finds an unused numeric local ID. +2. A custom reservation validates uniqueness, advances the counter once, and uses the supplied local ID. +3. A generated numeric local ID skips values already claimed by custom reservations. +4. Reusing a local ID in the same context fails immediately. + +Primitive operations created without a reservation use the same counter and local-ID registry. A custom reservation +therefore cannot reuse a numeric local ID already consumed by an ordinary core operation. + +Examples: + +```text +reserve("a", "node-a") -> hash("node-a") +reserve("b") -> hash("2") +reserve("c", "2") -> fails because "2" is already used +``` + +Inside a child context whose ID is `parentHash`: + +```text +reserve("a", "node-a") -> hash("parentHash-node-a") +``` + +Custom local IDs allow the custom-ID operations themselves to keep stable identities when definition order changes. +Sequential operations around them may still receive different IDs. Adding, removing, or changing an ID remains a +workflow compatibility change. + +## Stateful Step Extensions + +`waitForCondition` needs the existing STEP state machine with checkpointed state between retry attempts. The extension +API exposes this without exposing raw checkpoint actions. + +```java +@FunctionalInterface +public interface ExtensionStepFunction { + ExtensionStepResult apply(T state); +} +``` + +```java +public sealed interface ExtensionStepResult { + static ExtensionStepResult succeed(T value); + + static ExtensionStepResult retry(T state, Duration delay); +} +``` + +```java +public final class ExtensionStepConfig { + T initialState(); + + SerDes serDes(); +} +``` + +The subtype-aware stateful step selector is: + +```java + DurableFuture stepAsync( + String subType, + TypeToken resultType, + ExtensionStepFunction function, + ExtensionStepConfig config); +``` + +The SDK interprets results through the fixed STEP state machine: + +- `succeed(value)` serializes the value and checkpoints `SUCCEED`. +- `retry(state, delay)` serializes the state and checkpoints `RETRY`. +- a thrown exception checkpoints `FAIL`. +- internal suspension and unrecoverable control-flow exceptions propagate without conversion. + +`StepContext.getCurrentContext()` exposes the one-based attempt number. The extension function receives only its +application state. + +## Extension Context Results and Replay State + +A subtype-aware extension context may return a full application result and a smaller replay state: + +```java +public final class ExtensionContextResult { + static ExtensionContextResult completed(T result); + + static ExtensionContextResult replayChildren( + T result, + T replayState); + + static ExtensionContextResult replayChildrenAboveSize( + T result, + T replayState, + int thresholdBytes); +} +``` + +The application receives `result`. The checkpoint stores either the normal serialized result or `replayState`, +according to the selected factory. + +For `replayChildrenAboveSize`, the threshold is evaluated against the serialized full application result, before the +replay state is selected for checkpointing. + +When replay-children is enabled, the SDK reexecutes the extension context function and exposes the stored replay state +through a scoped extension context: + +```java +public final class ExtensionContextReplayContext { + static ExtensionContextReplayContext getCurrentContext(); + + boolean isReplayingChildren(); + + T getReplayState(); +} +``` + +The replay context is available only while the extension framework function is running. It is restored across nested +extension contexts and is not propagated to application-created threads. + +This preserves current map and parallel behavior: + +- A small map stores and replays its complete `MapResult`. +- A large map stores statuses and completion reason as replay state, then reconstructs values from iteration + checkpoints. +- Parallel stores its current `ParallelResult` as replay state and always replays branch children. + +## Extension Context Failure Translation + +The CONTEXT state machine always handles failures in the same way: + +1. serialize the thrown exception when possible +2. checkpoint `FAIL` +3. deserialize the original exception when the future is read or replayed + +Subtype-specific behavior is customizable only as a fallback when the original exception cannot be reconstructed. + +```java +@FunctionalInterface +public interface ExtensionContextErrorHandler { + Throwable translate(ExtensionContextFailure failure); +} +``` + +`ExtensionContextFailure` is a read-only view containing: + +- context name +- context subtype string +- deserialized original exception, when available +- serialized error metadata +- child operation summaries + +Each child summary contains: + +- operation type +- subtype string +- status +- serialized error metadata + +Resolution order is: + +1. rethrow a deserialized original exception +2. invoke the configured `ExtensionContextErrorHandler` +3. fall back to `ChildContextFailedException` + +The built-in extensions provide handlers that preserve current behavior: + +- wait for callback distinguishes callback failure, callback timeout, and submitter failure +- map iteration falls back to `MapIterationFailedException` +- parallel branch falls back to `ParallelBranchFailedException` +- with retry and ordinary child contexts fall back to `ChildContextFailedException` + +## Extension Context Configuration + +Existing `RunInChildContextConfig` remains unchanged and continues to provide SerDes and virtual-context settings. + +The subtype-aware extension context overload uses a new extension-specific wrapper: + +```java +public final class ExtensionContextConfig { + RunInChildContextConfig childContextConfig(); + + ExtensionContextErrorHandler errorHandler(); + + boolean emitUserFunctionEvents(); + + boolean suppressLateChildCheckpoints(); +} +``` + +The extension context function has a distinct type so it does not conflict by erasure with existing supplier +overloads: + +```java +@FunctionalInterface +public interface ExtensionContextFunction { + ExtensionContextResult apply(); +} +``` + +```java + DurableFuture runInChildContextAsync( + String subType, + TypeToken resultType, + ExtensionContextFunction function, + ExtensionContextConfig config); +``` + +The two additional booleans preserve existing family-specific behavior: + +- `emitUserFunctionEvents` controls whether the extension context function is reported as a user function. It defaults + to `true`, matching ordinary child contexts. +- `suppressLateChildCheckpoints` tracks extension-managed children and prevents them from writing checkpoints after + their parent extension context has completed. It defaults to `false`. + +Nested user step and child-context functions retain their own existing plugin hooks regardless of the parent setting. + +## Internal Operation Identity + +The existing public `OperationSubType` enum and enum-based identity factories remain unchanged. + +Internally, primitive operations use an identity containing: + +- operation ID +- name +- backend operation type +- subtype string + +Existing enum values convert to this representation. Custom subtype strings flow unchanged through: + +- operation updates +- replay validation +- plugin events +- logs +- failure views + +Replay validation compares both operation type and exact subtype string. + +## Built-In Extension Implementations + +Each family has one canonical implementation that accepts an `ExtensionContext`. The static facade obtains it from +TLS. The corresponding `DurableContextImpl` method passes `this`. + +```text +static facade --------------------+ + +-> built-in extension -> extension primitives +legacy DurableContext adapter ----+ +``` + +The canonical implementations use the existing public configurations and callback contracts after adapting them to +the extension-specific functions. + +### Wait for callback + +The extension: + +1. reserves a CONTEXT operation with subtype `WaitForCallback` +2. creates a CALLBACK child with subtype `Callback` +3. creates a STEP child with subtype `Step` +4. runs the submitter +5. waits for the callback result + +The parent extension context emits the same context user-function hooks as the current implementation. +The context failure handler preserves callback failure, timeout, and submitter exception translation. + +### With retry + +The extension: + +1. reserves a CONTEXT operation with subtype `WithRetry` +2. uses the current virtual or checkpointed behavior from `WithRetryConfig` +3. invokes the user operation with attempt metadata in `WithRetryContext` TLS +4. creates WAIT operations for backoff using the existing names and delays + +The retry context emits the same context user-function hooks as the current implementation. +Internal suspension and unrecoverable control-flow exceptions are never retried. + +### Wait for condition + +The extension reserves a stateful STEP operation with subtype `WaitForCondition`. + +The adapter: + +1. starts with `WaitForConditionConfig.initialState()` +2. invokes the existing check function +3. returns `succeed(value)` when polling completes +4. evaluates the existing wait strategy +5. returns `retry(value, delay)` when polling continues + +Attempt metadata remains available through `StepContext`. + +### Map + +The extension reserves a CONTEXT operation with subtype `Map`. + +Inside that context it: + +1. deterministically reserves all iteration contexts in input order +2. assigns subtype `MapIteration` +3. launches iterations through the shared concurrency coordinator +4. evaluates the existing `CompletionConfig` +5. constructs the existing `MapResult` +6. uses map replay state for large results + +Iteration reservations continue using sequential local IDs so existing operation IDs remain unchanged. + +The map parent does not emit context user-function hooks. Iteration contexts do emit them. The parent enables +late-child checkpoint suppression. + +Empty maps preserve `DurableConfig.shouldCheckpointEmptyMap()`: + +- when enabled, the map parent checkpoints `START` and `SUCCEED` +- when disabled, the reservation still consumes the same operation ID, the map emits the existing warning and plugin + lifecycle, and it completes with `MapResult.empty()` without a backend checkpoint + +### Parallel + +The extension reserves a CONTEXT operation with subtype `Parallel` and returns the existing +`ParallelDurableFuture`. + +Branch calls: + +1. reserve branch identities in registration order +2. assign subtype `ParallelBranch` +3. enqueue branch definitions in the parent extension context +4. launch through the shared concurrency coordinator + +`close()` and `get()` retain current join behavior. Branch registration after join still fails. + +Each branch call returns a deferred `DurableFuture` immediately. The coordinator binds it to the reserved child +context future when concurrency capacity permits. Its `get()` and `completionFuture()` retain the behavior of the +current branch future. + +The parallel parent does not emit context user-function hooks. Branch contexts do emit them. The parent always stores +replay state, replays children, and enables late-child checkpoint suppression. + +## Shared Concurrency Coordinator + +Map and parallel share a coordinator that is not itself a durable operation. + +It owns: + +- pending registration order +- max-concurrency enforcement +- running completion signals +- success and failure counts +- `CompletionConfig` evaluation +- skipped item tracking +- late-child checkpoint suppression +- the `allItemsRegistered` transition when map registration completes or parallel is joined + +It creates no operation type or checkpoint. All durable state belongs to the parent extension context and its reserved +child context primitives. + +The coordinator uses only `DurableFuture.completionFuture()` and reserved extension operations. It does not downcast +futures to SDK operation classes. + +## Removal of Specialized Engines + +After parity is proven, the following specialized engines are removed: + +- `MapOperation` +- `ParallelOperation` +- `ConcurrencyOperation` +- `WaitForConditionOperation` + +Their reusable primitive lifecycle behavior moves into the generalized STEP and CONTEXT primitive implementations. + +`ChildContextOperation`, `StepOperation`, and the other primitive operation classes remain as the SDK-owned state +machines. They are generalized to accept string subtypes and extension-specific result or failure policies. + +## Testing Strategy + +### Reservation tests + +Cover: + +- sequential reservations retain current IDs +- custom local IDs use the current context namespace +- custom reservations advance the sequence position +- generated numeric IDs skip reserved custom values +- duplicate local IDs fail +- nested contexts hash custom IDs with their parent ID +- custom IDs remain stable when reservation order changes + +### Primitive extension tests + +Cover: + +- arbitrary subtype strings for every primitive +- operation type remains determined by the primitive selector +- replay rejects type or subtype changes +- subtype strings reach plugin events unchanged +- stateful step success, retry, replay, state serialization, and failure +- context replay state and nested TLS restoration +- custom context failure translation and default fallback + +### Built-in parity tests + +For every operation family, compare legacy and static entry points for: + +- results and thrown exception types +- operation IDs, names, types, subtypes, and parent IDs +- checkpoint status and payload shape +- replay and suspension behavior +- plugin lifecycle ordering + +The existing map, parallel, callback, condition, retry, plugin, conformance, and example tests remain behavioral gates. +Tests formerly tied to specialized classes move to the generalized primitive and coordinator implementations without +weakening their assertions. + +### Completion gate + +Run: + +```bash +mvn spotless:apply +mvn clean install +``` + +Cloud example tests remain disabled unless their existing environment requirements are configured. + +## Documentation and ADR + +Update: + +- ADR-006 to include custom local IDs, arbitrary subtype strings, fixed primitive state machines, replay state, and + customizable context failure translation +- the custom extension guide with subtype-aware and custom-ID examples +- public Javadocs for every new extension-specific contract + +The documentation must state that operation IDs, subtypes, and replay state are workflow compatibility contracts. From b1ff1670d392f0e83dea435e79af35de48136408 Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 03:22:25 +0000 Subject: [PATCH 12/40] docs: amend custom extension operations ADR --- docs/adr/006-custom-extension-operations.md | 446 ++++++++++++-------- 1 file changed, 276 insertions(+), 170 deletions(-) diff --git a/docs/adr/006-custom-extension-operations.md b/docs/adr/006-custom-extension-operations.md index d131ae26f..772e87b32 100644 --- a/docs/adr/006-custom-extension-operations.md +++ b/docs/adr/006-custom-extension-operations.md @@ -6,73 +6,57 @@ ## Context -Issue [#571](https://github.com/aws/aws-durable-execution-sdk-java/issues/571) requests a supported way to -implement reusable durable operations in separate Maven modules without changing or rebuilding the core SDK. +Issue [#571](https://github.com/aws/aws-durable-execution-sdk-java/issues/571) requests a supported way to implement +reusable durable operations in separate Maven modules without changing or rebuilding the core SDK. -The SDK currently exposes all operations as instance methods on `DurableContext`. This creates several constraints: +The SDK historically exposed all operations as instance methods on `DurableContext`. This created several +constraints: -- Adding an optional operation requires changing the core `DurableContext` interface and `DurableContextImpl`. -- Extension libraries cannot create primitive operations with stable identities when registration order and execution - order differ. -- Extension code would need SDK implementation classes or public explicit-operation-ID methods to implement schedulers - such as DAG. -- Existing user-function APIs receive SDK-created context and metadata parameters, coupling new APIs to callback - signatures instead of the SDK-managed current context. -- A single facade containing every built-in extension would couple unrelated operation families and make independent - maintenance difficult. +- Adding an optional operation required changing `DurableContext` and `DurableContextImpl`. +- Extension libraries could not reserve stable operation identities before execution order was known. +- Extension code needed implementation classes to reproduce subtype-specific checkpoint and replay behavior. +- User-function APIs received SDK-created contexts and metadata as callback parameters. +- Built-in composed operations were split between static facades and dedicated operation engines instead of proving + the same extension model available to third parties. -Extension operations do not share one execution scope. Some extensions are direct primitive wrappers in the current -scope, while others deliberately create child contexts. The extension mechanism must not impose a child-context -boundary. +Extension operations do not share one execution scope. Some operate directly in the current context, while others +create child contexts. The extension mechanism must not impose a universal child-context boundary. -The SDK must continue to own primitive operation IDs, checkpointing, replay, suspension, serialization, failures, and -backend communication. This decision does not make backend operation types extensible. +The backend accepts arbitrary operation subtype strings, but each backend operation type has a fixed state machine. +Extensions need subtype control, replay state, and failure translation without receiving raw checkpoint actions or +defining new backend state machines. -The detailed API specification is in -[Custom Extension Operations Design](../superpowers/specs/2026-08-10-custom-extension-operations-design.md). +The detailed designs are: + +- [Custom Extension Operations Design](../superpowers/specs/2026-08-10-custom-extension-operations-design.md) +- [Migrate Built-In Operations to the Extension API](../superpowers/specs/2026-08-10-migrate-built-in-extensions-design.md) ## Decision -### Preserve DurableContext +### Preserve Existing Operation APIs -Keep the public `DurableContext` interface unchanged. Its existing instance methods and context-accepting callback -types remain supported for backward compatibility. +Keep every existing method signature, callback contract, configuration field, result type, exception, and behavior +unchanged. New capabilities are additive and limited to extension-specific overloads and types. -SDK-managed handler and child contexts additionally implement a new public `ExtensionContext` interface. Step -contexts do not implement it. +This includes: -```java -public interface ExtensionContext extends BaseContext { - static ExtensionContext getCurrentContext() { - var context = BaseContext.getCurrentContext(); - if (context instanceof ExtensionContext extensionContext) { - return extensionContext; - } - throw new IllegalStateException( - "ExtensionContext is only available from a durable handler or child-context thread"); - } - - boolean isReplaying(); - - ExtensionOperation reserve(String name); -} -``` +- `DurableContext` +- `ParallelDurableFuture` +- `MapConfig` +- `ParallelConfig` +- `WaitForCallbackConfig` +- `WaitForConditionConfig` +- `WithRetryConfig` +- `StepConfig` +- `RunInChildContextConfig` -Extension libraries expose ordinary static methods. There is no required registration mechanism and no universal -`DurableExtensions.run` method. +Existing methods on `ExtensionContext` and `ExtensionOperation` also remain unchanged; this decision adds overloads +rather than replacing them. -```java -public final class DagOperations { - public static DagResult dag(String name, Runnable definition) { - var extension = ExtensionContext.getCurrentContext(); - return executeDag(name, extension, definition); - } -} -``` +The existing `DurableContext` methods remain compatibility APIs. Their implementations delegate to the same built-in +extension implementations used by the static facades. -An extension decides whether to execute in the current scope or explicitly create a child context. - -### Separate Core and Extension Facades +### Expose Core and Built-In Extension Facades Expose primitive operations through `DurableCoreOperations`: @@ -92,190 +76,312 @@ Expose each built-in extension family through an independently maintained class: | `DurableWaitForConditionOperations` | `waitForCondition`, `waitForConditionAsync` | | `DurableWithRetryOperations` | `withRetry`, `withRetryAsync` | -These classes obtain the active durable context from SDK-managed thread-local storage and delegate to existing -operation implementations. The facade split does not require rewriting the established operation implementations. - -### Use TLS for SDK Context and Metadata +An extension is an ordinary static Java method. There is no registration API and no universal +`DurableExtensions.run` boundary. -User functions in the new static APIs receive only application-provided values or values from the application's -durable data flow. They do not receive `DurableContext`, `StepContext`, `ExtensionContext`, or SDK-generated metadata -as callback parameters. +### Use Scoped Current Context -Examples: +SDK-managed handler and child contexts implement `ExtensionContext`. Step contexts do not. -- Step and child-context functions use `Supplier`. -- Map functions receive the item; `MapItemContext.getCurrentContext()` provides the item index. -- Parallel branches use `Supplier`. -- Wait-for-callback submitters use `Runnable`; - `WaitForCallbackContext.getCurrentContext()` provides the callback ID. -- Wait-for-condition checks receive the durable state; - `StepContext.getCurrentContext()` provides attempt metadata. -- With-retry bodies use `Supplier`; - `WithRetryContext.getCurrentContext()` provides the attempt number. +User functions in the new APIs receive only application-provided values. SDK-created contexts and metadata are +retrieved from scoped thread-local contexts: -Operation-specific contexts use scoped SDK-managed thread-local storage. Nested scopes restore the previous value, and -the value is removed when no previous scope exists. Operation-specific metadata TLS is bound in addition to the base -durable or step context, allowing static core operations to resolve the correct active context. +- `DurableContext` +- `ExtensionContext` +- `StepContext` +- `MapItemContext` +- `WaitForCallbackContext` +- `WithRetryContext` +- extension replay contexts -Current context is available only on SDK-managed user-code threads. It is not propagated to application-created -threads. +Nested scopes restore the preceding value. Current context is available only on SDK-managed threads and is not +propagated to application-created threads. -### Reserve Primitive Identities +### Reserve Sequential or Custom Local Identities -`ExtensionContext.reserve(name)` immediately allocates the next sequential operation ID in the active durable scope -and returns an opaque, one-shot `ExtensionOperation`. +`ExtensionContext` supports sequential and custom-local-ID reservations: ```java -public interface ExtensionOperation { - DurableFuture stepAsync( - TypeToken resultType, - Supplier function, - StepConfig config); - - DurableFuture waitAsync(Duration duration); - - DurableFuture invokeAsync( - String functionName, - U payload, - TypeToken resultType, - InvokeConfig config); - - DurableCallbackFuture createCallback( - TypeToken resultType, - CallbackConfig config); - - DurableFuture runInChildContextAsync( - TypeToken resultType, - Supplier function, - RunInChildContextConfig config); -} +ExtensionOperation reserve(String name); + +ExtensionOperation reserve(String name, String localOperationId); ``` -The SDK binds the operation name and hidden ID to the reservation. The reservation can create exactly one primitive; -reusing it throws `IllegalStateException`. +Both forms return opaque, one-shot `ExtensionOperation` handles. -Extensions with deterministic invocation order can call `DurableCoreOperations` directly. Extensions such as DAG -reserve identities during deterministic definition, then execute the reservations in any dependency-valid order. +Sequential reservations use the next available numeric local ID. A custom reservation replaces the sequence number +for that position with a non-null, nonblank caller-provided local ID. -The implementation may add package-private explicit-ID primitive constructors. Extension code cannot access those -methods or raw IDs. +The SDK constructs the final backend ID: -### Support Composed Durable Futures +```text +root context: sha256(localOperationId) +child context: sha256(parentContextId + "-" + localOperationId) +``` + +Every operation allocation in a context shares one counter and local-ID registry. Custom reservations advance the +counter once, generated numeric IDs skip already claimed values, and duplicate local IDs fail immediately. + +Extension authors never provide or observe the final globally stored operation ID. + +### Allow Arbitrary Subtype Strings -Add a public, non-mutating completion signal to `DurableFuture`: +`ExtensionOperation` provides subtype-aware overloads for every primitive: ```java -default CompletableFuture completionFuture() { - throw new UnsupportedOperationException( - "This DurableFuture does not expose a completion signal"); -} +reservation.stepAsync("MyStep", ...); +reservation.waitAsync("MyWait", ...); +reservation.invokeAsync("MyInvoke", ...); +reservation.createCallback("MyCallback", ...); +reservation.runInChildContextAsync("MyContext", ...); ``` -SDK operations return a derived completion future that cannot mutate the underlying durable operation. -`DurableFuture.anyOf` uses this public contract instead of downcasting to `BaseDurableOperation`. Custom composed -futures override the method when they support `anyOf`. +The primitive selector determines the backend operation type. The string controls only the subtype recorded in +checkpoints, replay validation, plugins, logs, and error metadata. + +Subtype strings must be non-null and nonblank. They are not restricted to the existing `OperationSubType` enum. +Existing no-subtype overloads retain the standard subtype strings. + +### Keep Primitive State Machines Fixed + +Extension authors cannot send raw checkpoint updates or define arbitrary state machines. + +Each primitive retains its SDK-owned lifecycle: + +| Primitive | Backend operation type | Lifecycle | +| --- | --- | --- | +| step | `STEP` | start, retry, ready, succeed, fail | +| wait | `WAIT` | start, poll, succeed | +| invoke | `CHAINED_INVOKE` | start, poll, succeed, fail | +| callback | `CALLBACK` | start, poll, succeed, fail, timeout | +| child context | `CONTEXT` | start, execute or replay children, succeed, fail | + +A stateful extension STEP may return only: + +- `ExtensionStepResult.succeed(value)` +- `ExtensionStepResult.retry(state, delay)` + +The SDK maps those outcomes onto the fixed STEP lifecycle. Thrown exceptions follow the normal STEP failure path. +Attempt metadata remains available through `StepContext`. + +### Support Context Replay State + +A subtype-aware extension context returns `ExtensionContextResult`, which separates the application result from +optional replay state. + +Supported result policies are: + +- completed with the normal result +- always replay children and store a replay state +- replay children above a serialized-size threshold and store a replay state + +On replay, the framework function receives the stored replay state through scoped TLS. This supports large map +results and parallel branch reconstruction without exposing checkpoint APIs. + +`ExtensionContextConfig` composes the existing `RunInChildContextConfig` and adds extension-only behavior: + +- context failure translation +- whether the framework function emits user-function plugin events +- whether late child checkpoints are suppressed after parent completion + +Existing configuration classes are not changed. + +### Allow Context Failure Translation + +The CONTEXT state machine always serializes failures and checkpoints `FAIL`. Exception translation is customizable +when the original exception cannot be reconstructed. + +An `ExtensionContextErrorHandler` receives a read-only `ExtensionContextFailure` containing: + +- context name and subtype +- error metadata +- child operation type, subtype, status, and error summaries -### Preserve Primitive Plugin Semantics +Resolution order is: -Extensions do not create an automatic lifecycle or checkpoint boundary. Plugins observe the primitives created by an -extension. If an extension explicitly creates a child context, plugins also observe that context operation. +1. rethrow a deserialized original exception +2. invoke the configured error handler +3. fall back to `ChildContextFailedException` -No extension-specific backend operation type or raw checkpoint API is added. +This preserves the existing fallback behavior for callback failures and timeouts, map iterations, parallel branches, +with-retry contexts, and ordinary child contexts. + +### Support Composed Durable Futures + +`DurableFuture.completionFuture()` provides a public, non-mutating completion signal. Custom composed futures override +it when they support `DurableFuture.anyOf`. + +Map and parallel use a shared concurrency coordinator built from reserved extension context operations and public +completion signals. The coordinator is not a durable operation and creates no checkpoint of its own. + +### Implement Built-Ins Through Extensions + +Rewrite the existing composed operation families using the extension contract: + +- wait for callback uses a `WaitForCallback` context containing callback and step primitives +- with retry uses a virtual or checkpointed `WithRetry` context plus wait primitives +- wait for condition uses a stateful `WaitForCondition` step +- map uses a `Map` context and reserved `MapIteration` contexts +- parallel uses a `Parallel` context and dynamically registered `ParallelBranch` contexts + +The legacy `DurableContext` methods and static facades adapt into the same canonical family implementations. + +After behavior and checkpoint parity are proven, remove the specialized: + +- `MapOperation` +- `ParallelOperation` +- `ConcurrencyOperation` +- `WaitForConditionOperation` + +Primitive operation classes remain and are generalized for string subtypes, replay state, and extension failure +policies. + +### Preserve Checkpoint and Plugin Compatibility + +The migration preserves: + +- sequential operation IDs and parent-child namespaces +- existing type, subtype, name, and operation tree shape +- checkpoint actions, payloads, statuses, and replay validation +- map and parallel completion, skipped items, nesting, and large-result behavior +- wait-for-condition state, attempts, and delays +- wait-for-callback exception translation +- retry naming and virtual-context behavior +- plugin operation and user-function event ordering +- non-checkpointed empty-map behavior + +Extension framework callbacks emit user-function events only when configured. Nested application callbacks retain +their existing plugin events. ## Alternatives Considered -### Add extension methods to DurableContext +### Keep Dedicated Engines Behind Static Facades -Add `runExtensionAsync` or operation-specific methods to `DurableContext`. +Retain `MapOperation`, `ParallelOperation`, `ConcurrencyOperation`, and `WaitForConditionOperation`, while making only +the public facades look like extensions. **Rejected because:** -- It changes the interface that this decision must preserve. -- Optional extension families would continue to expand the core API. -- It makes extension execution appear to require a special runtime boundary. +- Built-ins would not validate that the public extension model is sufficient. +- Static and legacy APIs would continue to depend on a separate implementation architecture. +- Extension authors could not reproduce the capabilities exercised by built-ins. -### Require every extension to run in a child context +### Allow Arbitrary Checkpoint State Machines -Provide a universal `DurableExtensions.run` method that creates a child context. +Expose raw `START`, `RETRY`, `SUCCEED`, `FAIL`, polling, and operation-update APIs. **Rejected because:** -- Direct primitive wrappers do not need a child context. -- Child-context checkpoint and replay behavior would be imposed even when it is not part of the extension semantics. -- Extensions such as map or retry must remain responsible for selecting their own isolation strategy. +- Extensions could violate backend transition rules. +- Suspension, replay, and error handling would become extension-author responsibilities. +- The SDK would no longer own checkpoint correctness. -### Compose directly through DurableContext only +### Restrict Subtypes to OperationSubType -Let extension implementations retrieve `DurableContext` and invoke its existing methods. +Allow extension operations to use only the SDK's existing enum values. **Rejected because:** -- It exposes the entire legacy operation surface instead of a stable extension contract. -- Operations receive IDs when executed, so schedulers cannot register identities before varying launch order. -- Extension implementations remain coupled to context-accepting legacy callbacks. +- The backend accepts arbitrary subtype strings. +- Third-party extensions need distinct history and plugin identities. +- Adding an extension subtype would otherwise require a core SDK release. + +### Accept Exact Global Operation IDs + +Allow callers to provide the final backend operation ID. + +**Rejected because:** + +- Callers would need to understand context namespaces and hashing. +- Nested extensions could collide with unrelated operations. +- Backend identity details would become public workflow contracts. + +Custom IDs are therefore local values that the SDK namespaces and hashes. + +### Derive IDs from Operation Names + +Use operation names as local IDs automatically. + +**Rejected because:** + +- Names are not required to be unique. +- Changing a display name would silently change checkpoint identity. +- Explicit local IDs make the compatibility decision visible. + +### Require Every Extension to Run in a Child Context + +Provide a universal extension runner that always creates a child context. + +**Rejected because:** + +- Direct primitive wrappers do not need a child context. +- It imposes checkpoint and replay behavior unrelated to the extension's semantics. +- Each extension must select its own scope. -### Expose raw or name-derived operation IDs +### Add Extension Families to DurableContext -Allow extension libraries to supply operation IDs or derive them from operation names. +Continue adding new built-in or third-party operation methods to `DurableContext`. **Rejected because:** -- The SDK must retain ownership of global uniqueness and backend identity rules. -- Name-derived IDs introduce collision, normalization, and compatibility requirements. -- Public explicit-ID methods expose checkpoint protocol details. +- It expands the legacy interface for optional features. +- It prevents independently maintained extension modules. +- It retains context-bearing callback signatures. -### Pass contexts and generated metadata as callback arguments +### Pass SDK Contexts and Metadata as Callback Arguments -Mirror the existing `DurableContext` callback signatures in the new static APIs. +Mirror the existing `DurableContext` callback signatures in the new APIs. **Rejected because:** -- The new API uses SDK-managed current context consistently across core and extension operations. -- Generated values such as map index, retry attempt, and callback ID belong to typed operation contexts. -- Context-free callbacks make extension methods compose without threading SDK objects through application code. +- New APIs consistently use scoped current contexts. +- Generated values such as indexes, attempts, and callback IDs belong to typed metadata contexts. +- Context-free callbacks compose without threading SDK objects through application code. -### Use one built-in extension facade +### Use One Built-In Extension Facade -Place map, parallel, callback, condition, and retry methods in one `DurableExtensionOperations` class. +Place every built-in composed operation in one class. **Rejected because:** -- Unrelated overload sets, tests, and documentation would change together. -- Large operation families such as map and parallel need independent ownership. -- Separate classes align the public API with independently maintained extension implementations. +- Unrelated overloads, tests, and documentation would change together. +- Map and parallel require independently maintainable APIs. +- One class would become a second monolithic operation interface. ## Consequences **Positive:** -- Third-party Maven modules can publish static durable operations using supported public contracts. -- Application call sites do not pass or qualify `DurableContext`. -- `DurableContext` remains source- and binary-compatible. -- Extensions choose their own scope instead of inheriting a mandatory child-context boundary. -- Deterministic reservations support replay-safe schedulers whose launch order can vary. -- Primitive IDs, checkpointing, replay, and backend communication remain SDK-owned. -- Built-in extension families can evolve independently. -- New callback APIs consistently use TLS for SDK context and generated metadata. -- Custom composed futures work with public future combinators without internal downcasts. +- Third-party Maven modules can implement durable operations using supported public contracts. +- Custom subtype strings do not require SDK enum changes. +- Custom local IDs support replay-stable schedulers without exposing global IDs. +- `DurableContext` and existing operation configurations remain compatible. +- Built-in operations prove the same extension architecture available to third parties. +- The SDK continues to own all backend state machines and checkpoint transitions. +- Replay state and failure translation support advanced context extensions without raw checkpoint access. +- Static and legacy APIs share one implementation per operation family. **Negative:** -- The SDK must manage multiple scoped thread-local context types and restore them correctly across nested calls. -- Static APIs depend on execution from SDK-managed threads and fail from application-created threads. -- Reservations add package-private explicit-ID paths that must remain consistent with ordinary primitive creation. -- The static facade API duplicates overloads that remain on `DurableContext` for compatibility. -- Extension authors must understand that reservation order is part of workflow replay compatibility. +- Primitive operation implementations become more general and carry extension policies. +- Custom IDs require per-context collision tracking shared by reserved and direct operations. +- Extension authors must treat subtype strings, local IDs, and replay state as workflow compatibility contracts. +- Extension context failure handlers must remain deterministic and side-effect free. +- Map and parallel require a reusable concurrency coordinator and deferred futures. +- The SDK must preserve family-specific plugin-hook and late-child-checkpoint behavior through explicit policies. **Compatibility requirements:** -- Reservations must be created in the same deterministic order on every replay. -- Reordering, inserting, or removing reservations can rebind existing checkpoints and is a workflow compatibility - change. -- Launching already reserved operations in a different order is supported. -- Existing `DurableContext` methods and callback signatures remain unchanged. +- Existing operation interfaces, configs, overloads, results, exceptions, and behavior remain unchanged. +- Built-in migration must preserve exact operation IDs, topology, subtype strings, payloads, and replay behavior. +- Sequential reservations remain order-dependent. +- Custom-ID operations retain stable identities, but surrounding sequential IDs can change when definitions move. +- Reusing or changing a local ID is a workflow compatibility change. +- Changing a subtype string is a workflow compatibility change. +- Launching already reserved operations in a different order remains supported. **Deferred:** - A production DAG extension module. -- Reimplementing every built-in extension through the new public reservation contract. - Propagating current context to application-created threads. +- User-defined backend operation types or checkpoint state machines. From f543e8a5cc5ef24478beccaabb02837a9e5b9426 Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 03:25:53 +0000 Subject: [PATCH 13/40] docs: plan built-in extension migration --- .../2026-08-10-migrate-built-in-extensions.md | 995 ++++++++++++++++++ 1 file changed, 995 insertions(+) create mode 100644 docs/superpowers/plans/2026-08-10-migrate-built-in-extensions.md diff --git a/docs/superpowers/plans/2026-08-10-migrate-built-in-extensions.md b/docs/superpowers/plans/2026-08-10-migrate-built-in-extensions.md new file mode 100644 index 000000000..046a8f97e --- /dev/null +++ b/docs/superpowers/plans/2026-08-10-migrate-built-in-extensions.md @@ -0,0 +1,995 @@ +# Built-In Extension Migration Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Rewrite map, parallel, wait-for-callback, wait-for-condition, and with-retry through subtype-aware extension primitives while preserving all existing APIs, checkpoints, replay behavior, exceptions, and plugin events. + +**Architecture:** Extend reservations with custom local IDs and arbitrary subtype strings while keeping backend state machines SDK-owned. Generalize STEP and CONTEXT primitives with extension-only state, replay, and failure policies, then make legacy `DurableContext` methods and static facades delegate to one implementation per built-in family. Map and parallel share a non-operation concurrency coordinator that waits through suspension-aware public durable-future combinators. + +**Tech Stack:** Java 17, Maven reactor, JUnit 6, Mockito 5, `LocalDurableTestRunner`, Palantir Java Format through Spotless. + +## Global Constraints + +- Do not remove or change any existing method signature on `DurableContext`, `ParallelDurableFuture`, `ExtensionContext`, or `ExtensionOperation`. +- Do not add fields or methods to existing operation configuration classes. +- New subtype strings are non-null, nonblank, and are not restricted to `OperationSubType`. +- Primitive selectors determine backend operation types; extension code cannot emit raw checkpoint actions. +- Custom operation IDs are local values that replace a sequence number and are hashed with the current context prefix. +- Preserve exact built-in operation IDs, names, types, subtypes, parent IDs, payloads, statuses, replay behavior, exceptions, and plugin ordering. +- Do not add dependencies. +- Run `mvn spotless:apply` after Java changes. + +--- + +### Task 1: Custom Local Operation IDs + +**Files:** +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/execution/OperationIdGenerator.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionContext.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/execution/OperationIdGeneratorTest.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java` +- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java` + +**Interfaces:** +- Consumes: Existing `OperationIdGenerator.nextOperationId()` and `ExtensionContext.reserve(String)`. +- Produces: `OperationIdGenerator.nextOperationId(String localOperationId)` and `ExtensionContext.reserve(String name, String localOperationId)`. + +- [ ] **Step 1: Write failing generator tests** + +Add tests covering generated/custom interleaving: + +```java +@Test +void customLocalIdsUseContextPrefixAndAdvanceSequence() { + var root = new OperationIdGenerator(null); + + assertEquals(hashOperationId("node-a"), root.nextOperationId("node-a")); + assertEquals(hashOperationId("2"), root.nextOperationId()); +} + +@Test +void generatedIdsSkipCustomNumericIds() { + var generator = new OperationIdGenerator(null); + + assertEquals(hashOperationId("2"), generator.nextOperationId("2")); + assertEquals(hashOperationId("3"), generator.nextOperationId()); +} + +@Test +void duplicateLocalIdsFail() { + var generator = new OperationIdGenerator("parent"); + generator.nextOperationId("node"); + + assertThrows(IllegalArgumentException.class, () -> generator.nextOperationId("node")); +} +``` + +Also cover null, blank, a direct generated ID followed by the same custom numeric ID, and +`hashOperationId("parent-node")` for child contexts. + +- [ ] **Step 2: Run the generator tests and verify RED** + +Run: + +```bash +JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ +/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ +-Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ +-Djacoco.skip=true \ +-pl sdk -Dtest=OperationIdGeneratorTest test +``` + +Expected: compilation fails because the custom-local-ID overload does not exist. + +- [ ] **Step 3: Implement shared local-ID allocation** + +Use one atomic counter and concurrent local-ID set: + +```java +private final Set allocatedLocalIds = ConcurrentHashMap.newKeySet(); + +public String nextOperationId() { + String localId; + do { + localId = String.valueOf(operationCounter.incrementAndGet()); + } while (!allocatedLocalIds.add(localId)); + return hashOperationId(operationIdPrefix + localId); +} + +public String nextOperationId(String localOperationId) { + validateLocalOperationId(localOperationId); + if (!allocatedLocalIds.add(localOperationId)) { + throw new IllegalArgumentException("Local operation ID is already in use: " + localOperationId); + } + operationCounter.incrementAndGet(); + return hashOperationId(operationIdPrefix + localOperationId); +} +``` + +Validate before advancing the counter. + +- [ ] **Step 4: Add the reservation overload** + +Add an additive method to `ExtensionContext`: + +```java +ExtensionOperation reserve(String name, String localOperationId); +``` + +Implement it in `DurableContextImpl` by validating the name and allocating through the new generator overload. +Keep `reserve(String)` unchanged. + +- [ ] **Step 5: Add reservation and integration tests** + +Assert: + +- custom reservations are one-shot +- custom IDs remain stable when their registration order changes +- nested custom IDs use the child context ID prefix +- ordinary core operations and reservations share the same local-ID registry + +- [ ] **Step 6: Run focused tests and commit** + +Run: + +```bash +JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ +/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ +-Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ +-Djacoco.skip=true \ +-DargLine=-javaagent:/home/czk/.codex-tmp/m2/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar \ +-pl sdk,sdk-integration-tests -am \ +-Dtest=OperationIdGeneratorTest,ExtensionOperationImplTest,ExtensionOperationIntegrationTest \ +-Dsurefire.failIfNoSpecifiedTests=false test +``` + +Commit: + +```bash +git add sdk/src/main sdk/src/test sdk-integration-tests/src/test +git commit -m "feat: add custom extension operation ids" +``` + +--- + +### Task 2: Arbitrary Primitive Subtype Strings + +**Files:** +- Create: `sdk/src/main/java/software/amazon/lambda/durable/model/OperationDescriptor.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/WaitOperation.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/InvokeOperation.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/CallbackOperation.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginInfoConverter.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginInfoConverterTest.java` +- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java` + +**Interfaces:** +- Consumes: One-shot `ExtensionOperation` selectors and existing enum-based `OperationIdentifier`. +- Produces: Additive subtype overloads for all primitive selectors and an internal identity carrying `OperationType` plus exact subtype string. + +- [ ] **Step 1: Write failing subtype tests** + +Add one test per primitive: + +```java +var step = context.reserve("custom-step"); +step.stepAsync("AcmeStep", String.class, () -> "done"); + +verify(context).stepAsyncWithId( + eq("1"), + eq("custom-step"), + eq("AcmeStep"), + eq(TypeToken.get(String.class)), + any(), + any()); +``` + +Add integration assertions that checkpoints and plugin events contain `AcmeStep`, `AcmeWait`, `AcmeInvoke`, +`AcmeCallback`, and `AcmeContext`, while their operation types remain fixed by the selector. + +- [ ] **Step 2: Run focused tests and verify RED** + +Expected: compilation fails on missing subtype overloads. + +- [ ] **Step 3: Add internal string-based identity** + +Create: + +```java +public record OperationDescriptor( + String operationId, + String name, + OperationType operationType, + String subType) { + + public OperationDescriptor { + Objects.requireNonNull(operationId, "operationId cannot be null"); + Objects.requireNonNull(operationType, "operationType cannot be null"); + if (subType == null || subType.isBlank()) { + throw new IllegalArgumentException("subType cannot be null or blank"); + } + } + + public static OperationDescriptor from(OperationIdentifier identifier) { + return new OperationDescriptor( + identifier.operationId(), + identifier.name(), + identifier.operationType(), + identifier.subType().getValue()); + } +} +``` + +Keep `OperationIdentifier` unchanged. Add descriptor constructor overloads to primitive operation classes while +retaining enum-based constructors for current call sites and tests. + +- [ ] **Step 4: Generalize base replay and plugin paths** + +Store `OperationDescriptor` in `BaseDurableOperation`. Use `descriptor.subType()` for updates and replay validation. +Retain: + +```java +public OperationSubType getSubType() +``` + +for known enum-based operations, and add: + +```java +public String getSubTypeValue() +``` + +for arbitrary values. Add descriptor overloads to `PluginInfoConverter` without changing existing overloads. + +- [ ] **Step 5: Add subtype-aware selector overloads** + +For each primitive, add additive methods such as: + +```java + DurableFuture stepAsync( + String subType, + TypeToken resultType, + Supplier function, + StepConfig config); +``` + +Existing methods delegate using `OperationSubType.STEP.getValue()` and equivalent standard values. Validate subtype +before claiming the reservation so invalid input does not consume it. + +- [ ] **Step 6: Run focused tests and commit** + +Run the extension, primitive operation, replay-validation, and plugin converter unit tests plus +`ExtensionOperationIntegrationTest`. + +Commit: + +```bash +git add sdk/src/main sdk/src/test sdk-integration-tests/src/test +git commit -m "feat: support custom extension subtypes" +``` + +--- + +### Task 3: Stateful STEP Extension Primitive + +**Files:** +- Create: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionStepFunction.java` +- Create: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionStepResult.java` +- Create: `sdk/src/main/java/software/amazon/lambda/durable/config/ExtensionStepConfig.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/ExtensionStepResultTest.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionStepConfigTest.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/operation/StepOperationTest.java` +- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java` + +**Interfaces:** +- Consumes: Subtype-aware STEP reservations. +- Produces: A fixed STEP lifecycle whose user outcome is `succeed(value)` or `retry(state, delay)`. + +- [ ] **Step 1: Add failing API and validation tests** + +Test immutable factories and builder defaults: + +```java +var retry = ExtensionStepResult.retry("next", Duration.ofSeconds(2)); +assertEquals("next", retry.state()); +assertEquals(Duration.ofSeconds(2), retry.delay()); +``` + +Reject null delays, negative delays, null results, and missing `ExtensionStepConfig`. + +- [ ] **Step 2: Add failing state-machine tests** + +Exercise: + +- first attempt receives `initialState` +- retry checkpoints serialized state and delay +- READY replay resumes with checkpointed state and incremented attempt +- success returns the final state +- user exception checkpoints failure +- suspension and unrecoverable exceptions propagate + +- [ ] **Step 3: Implement extension types** + +Use a sealed result: + +```java +public sealed interface ExtensionStepResult + permits ExtensionStepResult.Succeeded, ExtensionStepResult.Retry { + record Succeeded(T value) implements ExtensionStepResult {} + record Retry(T state, Duration delay) implements ExtensionStepResult {} +} +``` + +Implement `ExtensionStepConfig` with builder fields `initialState` and `serDes`; null SerDes uses the durable +configuration default. + +- [ ] **Step 4: Generalize StepOperation** + +Introduce an internal attempt strategy inside `StepOperation`: + +```java +private interface AttemptBehavior { + AttemptOutcome execute(T state, StepContext context); +} +``` + +The existing constructor wraps the current function and retry strategy. The extension constructor maps +`ExtensionStepResult` onto the same START/RETRY/READY/SUCCEED/FAIL paths. Do not duplicate checkpoint sending or poll +logic. + +- [ ] **Step 5: Expose the reservation selector** + +Add: + +```java + DurableFuture stepAsync( + String subType, + TypeToken resultType, + ExtensionStepFunction function, + ExtensionStepConfig config); +``` + +The function receives state only. `StepContext` remains TLS-bound. + +- [ ] **Step 6: Run focused tests and commit** + +Run `ExtensionStepResultTest`, `ExtensionStepConfigTest`, `StepOperationTest`, +`ExtensionOperationImplTest`, and `ExtensionOperationIntegrationTest`. + +Commit: + +```bash +git add sdk/src/main sdk/src/test sdk-integration-tests/src/test +git commit -m "feat: add stateful extension steps" +``` + +--- + +### Task 4: Configurable CONTEXT Extension Primitive + +**Files:** +- Create: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextFunction.java` +- Create: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextResult.java` +- Create: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextReplayContext.java` +- Create: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextErrorHandler.java` +- Create: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextFailure.java` +- Create: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionChildOperationSummary.java` +- Create: `sdk/src/main/java/software/amazon/lambda/durable/config/ExtensionContextConfig.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/ExtensionContextResultTest.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/ExtensionContextReplayContextTest.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionContextConfigTest.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java` + +**Interfaces:** +- Consumes: Subtype-aware CONTEXT reservations. +- Produces: Replay-state result policies, scoped replay TLS, configurable fallback failure translation, plugin-hook policy, and late-child checkpoint suppression. + +- [ ] **Step 1: Write failing value and config tests** + +Test these factories: + +```java +ExtensionContextResult.completed(fullResult); +ExtensionContextResult.replayChildren(fullResult, replayState); +ExtensionContextResult.replayChildrenAboveSize(fullResult, replayState, 256 * 1024); +``` + +Test `ExtensionContextConfig.builder()` defaults: + +```java +assertTrue(config.emitUserFunctionEvents()); +assertFalse(config.suppressLateChildCheckpoints()); +assertNotNull(config.childContextConfig()); +``` + +- [ ] **Step 2: Write failing CONTEXT lifecycle tests** + +Add tests for: + +- normal result payload +- always replay-children with replay state +- threshold evaluated against serialized full result +- replay state available only inside `ExtensionContextReplayContext` +- nested replay scopes restore prior values +- framework hook emission enabled and disabled +- original exception reconstruction before fallback handler +- fallback handler receives child summaries +- default `ChildContextFailedException` +- a child finishing after a suppressing parent does not checkpoint + +- [ ] **Step 3: Implement extension context value types** + +Make all values immutable and defensively copy child summary lists. The replay TLS follows the existing +`OperationContextStorage` scoped-attachment pattern. + +- [ ] **Step 4: Generalize ChildContextOperation** + +Retain the existing `RunInChildContextConfig` constructor and adapt it to a standard context policy. Add an extension +constructor accepting `ExtensionContextFunction` and `ExtensionContextConfig`. + +On success: + +```java +var outcome = extensionFunction.apply(); +var full = serializeAndDeserializeResult(outcome.result()); +var checkpoint = selectCheckpointPayload(outcome, full.serialized()); +``` + +On replay with `replayChildren=true`, deserialize the stored replay state and attach it while rerunning the framework +function. Apply `emitUserFunctionEvents` only around the framework callback. Continue firing nested primitive hooks. + +- [ ] **Step 5: Generalize parent completion suppression** + +Replace the `ConcurrencyOperation`-specific parent constructor dependency with a general parent completion owner. +Store the owning extension context operation on child `DurableContextImpl` instances when +`suppressLateChildCheckpoints` is enabled. Nested extension child operations consult this owner before checkpointing. + +- [ ] **Step 6: Expose the advanced context selector** + +Add: + +```java + DurableFuture runInChildContextAsync( + String subType, + TypeToken resultType, + ExtensionContextFunction function, + ExtensionContextConfig config); +``` + +Keep the standard subtype plus `Supplier` overload and all existing methods unchanged. + +- [ ] **Step 7: Run focused tests and commit** + +Run the new extension context tests, `ChildContextOperationTest`, `ExtensionOperationImplTest`, plugin tests, and +extension integration tests. + +Commit: + +```bash +git add sdk/src/main sdk/src/test sdk-integration-tests/src/test +git commit -m "feat: add configurable extension contexts" +``` + +--- + +### Task 5: Migrate Wait for Callback + +**Files:** +- Create: `sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForCallbackExtension.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForCallbackOperations.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/extension/WaitForCallbackExtensionTest.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationsTest.java` +- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/CallbackIntegrationTest.java` +- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java` +- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java` + +**Interfaces:** +- Consumes: Advanced CONTEXT reservation, callback and step primitives, and configurable context failure translation. +- Produces: One canonical wait-for-callback implementation used by static and legacy APIs. + +- [ ] **Step 1: Write failing canonical-delegation tests** + +Assert both entry points call a handler with the same names and topology: + +```text +approval CONTEXT / WaitForCallback + approval-callback CALLBACK / Callback + approval-submitter STEP / Step +``` + +Cover callback failure, timeout, submitter failure, generic result types, custom SerDes, and plugin event ordering. + +- [ ] **Step 2: Implement WaitForCallbackExtension** + +Use: + +```java +public static DurableFuture execute( + ExtensionContext context, + String name, + TypeToken resultType, + BiConsumer submitter, + WaitForCallbackConfig config) +``` + +Reserve the parent first, then create callback and submitter reservations inside its framework function. Configure +parent user-function hooks as enabled. Supply an error handler that inspects child summaries and recreates +`CallbackFailedException`, `CallbackTimeoutException`, and `CallbackSubmitterException`. + +- [ ] **Step 3: Redirect both APIs** + +`DurableContextImpl.waitForCallbackAsync` delegates with `this`. The static facade resolves +`ExtensionContext.getCurrentContext()` and retains its existing TLS adapter around the `Runnable`. + +- [ ] **Step 4: Run callback suites and commit** + +Run callback unit, integration, retry-with-callback, static operations, and plugin tests. + +Commit: + +```bash +git add sdk/src/main sdk/src/test sdk-integration-tests/src/test +git commit -m "refactor: implement wait for callback as extension" +``` + +--- + +### Task 6: Migrate With Retry + +**Files:** +- Create: `sdk/src/main/java/software/amazon/lambda/durable/context/extension/WithRetryExtension.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/DurableWithRetryOperations.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/extension/WithRetryExtensionTest.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java` +- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java` +- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java` + +**Interfaces:** +- Consumes: Advanced `WithRetry` context reservation and ordinary wait primitives. +- Produces: One retry loop shared by static and legacy APIs. + +- [ ] **Step 1: Write failing parity tests** + +Cover first-attempt success, delayed retries, exhausted retries, null names, virtual versus checkpointed context, +control-flow exception propagation, and static TLS attempt metadata. + +- [ ] **Step 2: Implement WithRetryExtension** + +Move the retry loop out of `DurableContextImpl`: + +```java +public static DurableFuture execute( + ExtensionContext context, + String name, + BiFunction operation, + WithRetryConfig config) +``` + +Create a `WithRetry` extension context using the existing naming and virtual-context rules. Read the child +`DurableContext` through TLS, run the operation, and reserve ordinary waits with the existing backoff names. + +- [ ] **Step 3: Redirect APIs and remove old loop helpers** + +Delegate both legacy and static APIs to the canonical extension. Delete retry constants and loop methods from +`DurableContextImpl` after tests compile. + +- [ ] **Step 4: Run retry suites and commit** + +Commit: + +```bash +git add sdk/src/main sdk/src/test sdk-integration-tests/src/test +git commit -m "refactor: implement retry as extension" +``` + +--- + +### Task 7: Migrate Wait for Condition + +**Files:** +- Create: `sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForConditionExtension.java` +- Create: `sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForConditionFuture.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForConditionOperations.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java` +- Delete: `sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java` +- Move/replace test: `sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionOperationTest.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/extension/WaitForConditionExtensionTest.java` +- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionIntegrationTest.java` +- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java` + +**Interfaces:** +- Consumes: Stateful extension STEP. +- Produces: Existing wait-for-condition APIs and exception behavior through subtype `WaitForCondition`. + +- [ ] **Step 1: Write failing extension parity tests** + +Assert the existing single STEP checkpoint is retained across immediate success, multiple retries, READY replay, +initial state, custom strategy, custom SerDes, thrown checks, and plugin attempt numbers. + +- [ ] **Step 2: Implement WaitForConditionExtension** + +Map the existing result to fixed step outcomes: + +```java +var result = checkFunction.apply(state, StepContext.getCurrentContext()); +if (result.isDone()) { + return ExtensionStepResult.succeed(result.value()); +} +var delay = config.waitStrategy().evaluate( + result.value(), + StepContext.getCurrentContext().getAttempt()); +return ExtensionStepResult.retry(result.value(), delay); +``` + +Use subtype `WaitForCondition`, the existing initial state, and the existing SerDes defaulting. + +- [ ] **Step 3: Preserve fallback exception type** + +Wrap the stateful step future in `WaitForConditionFuture`. Delegate `completionFuture()`. In `get()`, let original +deserialized exceptions propagate; translate only fallback `StepFailedException` to +`WaitForConditionFailedException` using its operation. + +- [ ] **Step 4: Redirect APIs and remove the specialized operation** + +Delegate static and legacy methods to the canonical extension. Delete `WaitForConditionOperation` after all its +state-machine assertions have equivalent coverage in `StepOperationTest` and `WaitForConditionExtensionTest`. + +- [ ] **Step 5: Run condition suites and commit** + +Commit: + +```bash +git add -A sdk/src/main sdk/src/test sdk-integration-tests/src/test +git commit -m "refactor: implement wait for condition as extension" +``` + +--- + +### Task 8: Suspension-Aware Concurrency Coordination + +**Files:** +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/DurableFuture.java` +- Create: `sdk/src/main/java/software/amazon/lambda/durable/context/extension/DeferredDurableFuture.java` +- Create: `sdk/src/main/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinator.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/extension/DeferredDurableFutureTest.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinatorTest.java` +- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionConcurrencyIntegrationTest.java` + +**Interfaces:** +- Consumes: `DurableFuture.completionFuture()`, extension child reservations, and `CompletionConfig`. +- Produces: A non-operation coordinator usable by map and parallel without internal future downcasts. + +- [ ] **Step 1: Write failing suspension tests for anyOf** + +Use `LocalDurableTestRunner` with child futures waiting on callbacks. Assert the invocation reaches `PENDING` instead +of remaining active while `DurableFuture.anyOf` waits. + +- [ ] **Step 2: Make anyOf cooperate with SDK thread registration** + +When called on an SDK-managed context thread, delegate completion waiting to a context helper that: + +1. records the current thread context +2. registers a completion continuation +3. deregisters the active thread before blocking +4. re-registers when any completion signal fires + +Keep current behavior outside SDK threads. Do not change `completionFuture()` mutation isolation. + +- [ ] **Step 3: Implement DeferredDurableFuture** + +Provide a one-time `bind(DurableFuture)` method. `get()` waits for binding then delegates. `completionFuture()` +returns a stable future completed from the bound future. Reject a second binding. + +- [ ] **Step 4: Implement ExtensionConcurrencyCoordinator** + +The coordinator maintains: + +```java +record Item( + ExtensionOperation reservation, + Supplier> launcher, + DeferredDurableFuture exposedFuture) {} +``` + +It must: + +- register items in deterministic order +- launch no more than `maxConcurrency` +- wait through `DurableFuture.anyOf` +- count succeeded and failed items +- evaluate `CompletionConfig.completionDecisionFunction()` +- preserve `allItemsRegistered` +- mark pending/running incomplete items as skipped when completion occurs +- propagate suspension and unrecoverable control flow + +- [ ] **Step 5: Run focused and integration tests and commit** + +Commit: + +```bash +git add sdk/src/main sdk/src/test sdk-integration-tests/src/test +git commit -m "feat: add extension concurrency coordination" +``` + +--- + +### Task 9: Migrate Map + +**Files:** +- Create: `sdk/src/main/java/software/amazon/lambda/durable/context/extension/MapExtension.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/DurableMapOperations.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java` +- Delete: `sdk/src/main/java/software/amazon/lambda/durable/operation/MapOperation.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/extension/MapExtensionTest.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationsTest.java` +- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/MapIntegrationTest.java` +- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/MapInputValidationTest.java` +- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java` + +**Interfaces:** +- Consumes: Extension CONTEXT result policies and the shared concurrency coordinator. +- Produces: Existing map behavior with `Map` and `MapIteration` subtype-aware context primitives. + +- [ ] **Step 1: Add checkpoint-history parity tests** + +For legacy and static map calls, assert identical: + +- parent and iteration IDs +- `Map` and `MapIteration` subtypes +- nested/flat parent IDs +- small result payloads +- large result replay state and `replayChildren` +- early completion statuses and skipped iterations +- empty-map checkpoint flag behavior +- plugin events + +- [ ] **Step 2: Implement MapExtension** + +Use: + +```java +public static DurableFuture> execute( + ExtensionContext context, + String name, + Collection items, + TypeToken resultType, + DurableContext.MapFunction function, + MapConfig config) +``` + +Validate and copy items exactly as the current implementation. Reserve the `Map` parent first. Inside it, reserve +iterations in input order, attach `MapItemContext`, and launch through `ExtensionConcurrencyCoordinator`. + +- [ ] **Step 3: Preserve result and replay policies** + +Construct `MapResult` with existing success, failure, and skipped entries. For results at least 256 KB, use: + +```java +ExtensionContextResult.replayChildrenAboveSize( + fullResult, + stripMapResult(fullResult), + 256 * 1024); +``` + +On replay, use `ExtensionContextReplayContext` statuses to avoid launching previously skipped iterations and to +restore the prior completion decision. + +- [ ] **Step 4: Preserve empty-map and plugin behavior** + +Consume the parent reservation in all cases. Use a virtual `Map` extension context when empty-map checkpointing is +disabled, retain the warning, return `MapResult.empty()`, suppress parent framework user-function hooks, and keep +operation start/end plugin events balanced. + +- [ ] **Step 5: Redirect APIs and delete MapOperation** + +Delegate static and legacy map methods to `MapExtension`. Move reusable result assertions from operation tests into +`MapExtensionTest`. Delete `MapOperation`. + +- [ ] **Step 6: Run the complete map suite and commit** + +Run all map unit/integration tests and map-related plugin tests for both nesting modes. + +Commit: + +```bash +git add -A sdk/src/main sdk/src/test sdk-integration-tests/src/test +git commit -m "refactor: implement map as extension" +``` + +--- + +### Task 10: Migrate Parallel and Remove Specialized Concurrency + +**Files:** +- Create: `sdk/src/main/java/software/amazon/lambda/durable/context/extension/ParallelExtension.java` +- Create: `sdk/src/main/java/software/amazon/lambda/durable/context/extension/ParallelExtensionFuture.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/DurableParallelOperations.java` +- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java` +- Delete: `sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java` +- Delete: `sdk/src/main/java/software/amazon/lambda/durable/operation/ConcurrencyOperation.java` +- Delete/replace test: `sdk/src/test/java/software/amazon/lambda/durable/operation/ParallelOperationTest.java` +- Delete/replace test: `sdk/src/test/java/software/amazon/lambda/durable/operation/ConcurrencyOperationTest.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/extension/ParallelExtensionTest.java` +- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinatorTest.java` +- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ParallelIntegrationTest.java` +- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java` + +**Interfaces:** +- Consumes: Dynamic coordinator registration, deferred futures, and advanced CONTEXT replay policies. +- Produces: Existing `ParallelDurableFuture` backed entirely by extension primitives. + +- [ ] **Step 1: Add parallel checkpoint-history parity tests** + +Compare legacy and static APIs for empty, heterogeneous, max-concurrency, early-success, failure-tolerance, nested, +flat, replay, branches added after parent completion, and plugin scenarios. + +- [ ] **Step 2: Implement ParallelExtensionFuture** + +The future: + +- starts one `Parallel` extension context +- queues branch definitions in registration order +- returns a `DeferredDurableFuture` from each branch call +- rejects registration after `get()` or `close()` +- signals `allItemsRegistered` on join +- delegates its own `completionFuture()` and `get()` to the parent context future + +The parent framework function obtains its child `ExtensionContext`, drains registrations through the coordinator, +and returns `ExtensionContextResult.replayChildren(result, result)`. + +- [ ] **Step 3: Preserve replay and late-completion behavior** + +Use stored `ParallelResult.statuses()` to skip branches that did not exist or were previously skipped. Configure the +parent with framework user hooks disabled and late-child checkpoints suppressed. Configure branch contexts with +`ParallelBranch` fallback translation and virtual mode for flat nesting. + +- [ ] **Step 4: Redirect APIs and delete specialized classes** + +`DurableContextImpl.parallel` and `DurableParallelOperations.parallel` instantiate the same canonical extension +future. Delete `ParallelOperation` and `ConcurrencyOperation` after moving all shared assertions to coordinator and +extension tests. + +- [ ] **Step 5: Run parallel and broad integration suites and commit** + +Run parallel unit/integration tests, nested map/parallel tests, callbacks inside branches, condition operations inside +branches, and plugin tests. + +Commit: + +```bash +git add -A sdk/src/main sdk/src/test sdk-integration-tests/src/test +git commit -m "refactor: implement parallel as extension" +``` + +--- + +### Task 11: API Parity, Documentation, and Full Verification + +**Files:** +- Modify: `docs/advanced/extensions.md` +- Modify: `docs/adr/006-custom-extension-operations.md` +- Modify: `README.md` only if the extension guide link or description changes +- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java` +- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java` + +**Interfaces:** +- Consumes: All migrated built-in extensions. +- Produces: Final compatibility evidence and documented public extension contracts. + +- [ ] **Step 1: Add final legacy/static parity coverage** + +For each family, run the legacy and static forms with equivalent inputs and compare normalized operation history: + +```java +record OperationShape( + String name, + String type, + String subType, + String parentId, + String status) {} +``` + +Also verify custom subtype strings and custom local IDs from a separate Maven module fixture. + +- [ ] **Step 2: Verify public API compatibility** + +Confirm: + +```bash +git diff 1d3de02 -- sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java +git diff 1d3de02 -- sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java +git diff 1d3de02 -- sdk/src/main/java/software/amazon/lambda/durable/config +``` + +Expected: no removed or changed existing signatures or fields; only additive extension-specific files and overloads. + +- [ ] **Step 3: Update extension documentation** + +Document: + +- arbitrary subtype strings +- custom local ID hashing and collision rules +- stateful STEP outcomes +- context replay state +- context error handlers +- built-in operations as reference extensions +- subtype, local ID, and replay state compatibility warnings + +- [ ] **Step 4: Run Spotless** + +Run: + +```bash +JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ +/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ +-Dmaven.repo.local=/home/czk/.codex-tmp/m2 spotless:apply +``` + +Review and remove only unrelated formatter churn outside touched files. + +- [ ] **Step 5: Run focused dependency closure** + +Run: + +```bash +JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ +/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ +-Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ +-Djacoco.skip=true \ +-DargLine=-javaagent:/home/czk/.codex-tmp/m2/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar \ +-pl sdk,sdk-integration-tests -am test +``` + +Expected: all SDK, testing, and integration tests pass. + +- [ ] **Step 6: Run full reactor** + +Run: + +```bash +JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ +/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ +-Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ +-Djacoco.skip=true \ +-DargLine=-javaagent:/home/czk/.codex-tmp/m2/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar \ +clean install +``` + +Expected: all eight reactor modules succeed. Cloud tests remain disabled by their existing guard. + +- [ ] **Step 7: Final audit and commit** + +Run: + +```bash +git diff --check +git status --short +rg -n "MapOperation|ParallelOperation|ConcurrencyOperation|WaitForConditionOperation" sdk/src/main sdk/src/test +``` + +Expected: no obsolete specialized engine references and no unintended worktree changes. + +Commit: + +```bash +git add README.md docs sdk sdk-integration-tests +git commit -m "docs: document migrated built-in extensions" +``` From ef2ce97fd071c396e4b90838ba7d83ea430c4b23 Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 03:27:31 +0000 Subject: [PATCH 14/40] feat: add custom extension operation ids --- .../ExtensionOperationIntegrationTest.java | 29 ++++++++- .../lambda/durable/ExtensionContext.java | 12 ++++ .../durable/context/DurableContextImpl.java | 14 +++++ .../execution/OperationIdGenerator.java | 36 ++++++++++- .../context/ExtensionOperationImplTest.java | 17 +++++ .../execution/OperationIdGeneratorTest.java | 62 +++++++++++++++++++ 6 files changed, 167 insertions(+), 3 deletions(-) create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/execution/OperationIdGeneratorTest.java diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java index 41b96a4a5..cd9481970 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java @@ -13,6 +13,7 @@ import java.nio.charset.StandardCharsets; import java.security.MessageDigest; import java.security.NoSuchAlgorithmException; +import java.time.Duration; import java.util.HexFormat; import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.Test; @@ -38,6 +39,31 @@ void reservedOperationsReplayWhenLaunchOrderChanges() { assertEquals(hash("3"), result.getOperation("pair-pause").getId()); } + @Test + void customReservationsRemainStableWhenRegistrationOrderChanges() { + var invocations = new AtomicInteger(); + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { + var extension = ExtensionContext.getCurrentContext(); + var replay = invocations.incrementAndGet() > 1; + var first = replay + ? extension.reserve("right", "right") + : extension.reserve("left", "left"); + var second = replay + ? extension.reserve("left", "left") + : extension.reserve("right", "right"); + first.step(String.class, () -> first == second ? "invalid" : "first"); + second.step(String.class, () -> "second"); + context.wait("replay", Duration.ofSeconds(1)); + return "done"; + }); + + var result = runner.runUntilComplete("input"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals(hash("left"), result.getOperation("left").getId()); + assertEquals(hash("right"), result.getOperation("right").getId()); + } + @Test void staticOperationsUseCurrentContextAndRejectStepThreads() { var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { @@ -65,7 +91,7 @@ void extensionCanExplicitlyCreateChildContext() { return outer.reserve("child").runInChildContext(String.class, () -> { var child = ExtensionContext.getCurrentContext(); assertNotSame(outer, child); - return child.reserve("value").step(String.class, () -> "nested"); + return child.reserve("value", "node").step(String.class, () -> "nested"); }); }); @@ -74,6 +100,7 @@ void extensionCanExplicitlyCreateChildContext() { assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); assertEquals("nested", result.getResult(String.class)); assertEquals(hash("1"), result.getOperation("child").getId()); + assertEquals(hash(hash("1") + "-node"), result.getOperation("value").getId()); } private static String hash(String value) { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContext.java b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContext.java index ae34514ce..d776ee2c8 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContext.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContext.java @@ -39,4 +39,16 @@ static ExtensionContext getCurrentContext() { * @return an opaque one-shot reservation */ ExtensionOperation reserve(String name); + + /** + * Reserves a primitive operation identity using a caller-provided local ID. + * + *

The SDK namespaces and hashes the local ID with the current context. The local ID replaces one generated + * sequence value and must be unique within this context. + * + * @param name the primitive operation name + * @param localOperationId the stable local ID within the current context + * @return an opaque one-shot reservation + */ + ExtensionOperation reserve(String name, String localOperationId); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java index 303cbcdf6..4d09efe7f 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java @@ -499,16 +499,30 @@ private String nextOperationId() { return operationIdGenerator.nextOperationId(); } + private String nextOperationId(String localOperationId) { + return operationIdGenerator.nextOperationId(localOperationId); + } + String reserveOperationId() { return nextOperationId(); } + String reserveOperationId(String localOperationId) { + return nextOperationId(localOperationId); + } + @Override public ExtensionOperation reserve(String name) { ParameterValidator.validateOperationName(name); return new ExtensionOperationImpl(this, reserveOperationId(), name); } + @Override + public ExtensionOperation reserve(String name, String localOperationId) { + ParameterValidator.validateOperationName(name); + return new ExtensionOperationImpl(this, reserveOperationId(localOperationId), name); + } + /** Returns whether this context is currently in replay mode. */ @Override public boolean isReplaying() { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/execution/OperationIdGenerator.java b/sdk/src/main/java/software/amazon/lambda/durable/execution/OperationIdGenerator.java index 08ea883db..e24052396 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/execution/OperationIdGenerator.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/execution/OperationIdGenerator.java @@ -6,12 +6,16 @@ import java.security.MessageDigest; import java.security.NoSuchAlgorithmException; import java.util.HexFormat; +import java.util.Objects; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicInteger; /** Generates operation IDs for the durable operations. */ public class OperationIdGenerator { private final AtomicInteger operationCounter; private final String operationIdPrefix; + private final Set allocatedLocalIds = ConcurrentHashMap.newKeySet(); public OperationIdGenerator(String contextId) { this.operationCounter = new AtomicInteger(0); @@ -42,7 +46,35 @@ public static String hashOperationId(String rawId) { * {@code hash("-2")} inside a child context. */ public String nextOperationId() { - var counter = String.valueOf(operationCounter.incrementAndGet()); - return hashOperationId(operationIdPrefix + counter); + String localOperationId; + do { + localOperationId = String.valueOf(operationCounter.incrementAndGet()); + } while (!allocatedLocalIds.add(localOperationId)); + return hashOperationId(operationIdPrefix + localOperationId); + } + + /** + * Returns an operation ID derived from a caller-provided local ID. + * + *

The local ID replaces the generated sequence number for this allocation. It is namespaced by the current + * context prefix and advances the generated sequence once. + * + * @param localOperationId the caller-provided ID within the current context + * @return the hashed, context-scoped operation ID + */ + public String nextOperationId(String localOperationId) { + validateLocalOperationId(localOperationId); + if (!allocatedLocalIds.add(localOperationId)) { + throw new IllegalArgumentException("Local operation ID is already in use: " + localOperationId); + } + operationCounter.incrementAndGet(); + return hashOperationId(operationIdPrefix + localOperationId); + } + + private void validateLocalOperationId(String localOperationId) { + Objects.requireNonNull(localOperationId, "localOperationId cannot be null"); + if (localOperationId.isBlank()) { + throw new IllegalArgumentException("localOperationId cannot be blank"); + } } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java index ecc050d4a..653cb8e48 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java @@ -51,6 +51,23 @@ void reservationsKeepSequentialIdsWhenExecutedOutOfOrder() { ordered.verify(context).waitAsyncWithId("sequential-1", "first", duration); } + @Test + void customReservationUsesExplicitLocalOperationId() { + var context = mock(DurableContextImpl.class); + var duration = Duration.ofSeconds(1); + var expectedFuture = mockFuture(); + when(context.reserveOperationId("node-a")).thenReturn("custom-node-a"); + doCallRealMethod().when(context).reserve("custom", "node-a"); + when(context.waitAsyncWithId("custom-node-a", "custom", duration)).thenReturn(expectedFuture); + + var operation = context.reserve("custom", "node-a"); + var actualFuture = operation.waitAsync(duration); + + verify(context).reserveOperationId("node-a"); + verify(context).waitAsyncWithId("custom-node-a", "custom", duration); + assertEquals(expectedFuture, actualFuture); + } + @Test void reservedStepAdaptsSupplierToStepFunction() { var context = mock(DurableContextImpl.class); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/OperationIdGeneratorTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/OperationIdGeneratorTest.java new file mode 100644 index 000000000..0ad05cdbd --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/OperationIdGeneratorTest.java @@ -0,0 +1,62 @@ +// 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.assertThrows; +import static software.amazon.lambda.durable.execution.OperationIdGenerator.hashOperationId; + +import org.junit.jupiter.api.Test; + +class OperationIdGeneratorTest { + @Test + void customLocalIdUsesRootNamespaceAndAdvancesSequence() { + var generator = new OperationIdGenerator(null); + + assertEquals(hashOperationId("node-a"), generator.nextOperationId("node-a")); + assertEquals(hashOperationId("2"), generator.nextOperationId()); + } + + @Test + void customLocalIdUsesParentNamespace() { + var generator = new OperationIdGenerator("parent"); + + assertEquals(hashOperationId("parent-node-a"), generator.nextOperationId("node-a")); + } + + @Test + void generatedIdsSkipClaimedCustomNumericIds() { + var generator = new OperationIdGenerator(null); + + assertEquals(hashOperationId("2"), generator.nextOperationId("2")); + assertEquals(hashOperationId("3"), generator.nextOperationId()); + } + + @Test + void customIdsCannotClaimGeneratedNumericIds() { + var generator = new OperationIdGenerator(null); + generator.nextOperationId(); + + assertThrows(IllegalArgumentException.class, () -> generator.nextOperationId("1")); + assertEquals(hashOperationId("2"), generator.nextOperationId()); + } + + @Test + void duplicateCustomLocalIdsFailWithoutAdvancingSequence() { + var generator = new OperationIdGenerator("parent"); + generator.nextOperationId("node"); + + assertThrows(IllegalArgumentException.class, () -> generator.nextOperationId("node")); + assertEquals(hashOperationId("parent-2"), generator.nextOperationId()); + } + + @Test + void invalidCustomLocalIdsFailWithoutAdvancingSequence() { + var generator = new OperationIdGenerator(null); + + assertThrows(NullPointerException.class, () -> generator.nextOperationId(null)); + assertThrows(IllegalArgumentException.class, () -> generator.nextOperationId("")); + assertThrows(IllegalArgumentException.class, () -> generator.nextOperationId(" ")); + assertEquals(hashOperationId("1"), generator.nextOperationId()); + } +} From 5a2b5e74321add5b34a9a835e325db773870ca4e Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 03:31:31 +0000 Subject: [PATCH 15/40] feat: support custom extension subtypes --- .../ExtensionOperationIntegrationTest.java | 30 +++- .../lambda/durable/ExtensionOperation.java | 131 ++++++++++++++++++ .../durable/context/DurableContextImpl.java | 71 +++++++++- .../context/ExtensionOperationImpl.java | 47 +++++++ .../durable/model/OperationDescriptor.java | 46 ++++++ .../operation/BaseDurableOperation.java | 47 +++++-- .../durable/operation/CallbackOperation.java | 10 ++ .../operation/ChildContextOperation.java | 32 +++-- .../durable/operation/InvokeOperation.java | 15 ++ .../SerializableDurableOperation.java | 21 +++ .../durable/operation/StepOperation.java | 12 ++ .../durable/operation/WaitOperation.java | 7 + .../durable/plugin/PluginInfoConverter.java | 52 +++++-- .../context/ExtensionOperationImplTest.java | 72 ++++++++++ .../plugin/PluginInfoConverterTest.java | 12 ++ 15 files changed, 557 insertions(+), 48 deletions(-) create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/model/OperationDescriptor.java diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java index cd9481970..5e4f146d5 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java @@ -17,6 +17,7 @@ import java.util.HexFormat; import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.Test; +import software.amazon.awssdk.services.lambda.model.OperationType; import software.amazon.lambda.durable.model.ExecutionStatus; import software.amazon.lambda.durable.testing.LocalDurableTestRunner; @@ -45,12 +46,8 @@ void customReservationsRemainStableWhenRegistrationOrderChanges() { var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { var extension = ExtensionContext.getCurrentContext(); var replay = invocations.incrementAndGet() > 1; - var first = replay - ? extension.reserve("right", "right") - : extension.reserve("left", "left"); - var second = replay - ? extension.reserve("left", "left") - : extension.reserve("right", "right"); + var first = replay ? extension.reserve("right", "right") : extension.reserve("left", "left"); + var second = replay ? extension.reserve("left", "left") : extension.reserve("right", "right"); first.step(String.class, () -> first == second ? "invalid" : "first"); second.step(String.class, () -> "second"); context.wait("replay", Duration.ofSeconds(1)); @@ -103,6 +100,27 @@ void extensionCanExplicitlyCreateChildContext() { assertEquals(hash(hash("1") + "-node"), result.getOperation("value").getId()); } + @Test + void customPrimitiveSubtypesAreStoredWithoutChangingOperationTypes() { + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { + var extension = ExtensionContext.getCurrentContext(); + extension.reserve("custom-step").step("AcmeStep", String.class, () -> "step"); + extension.reserve("custom-wait").wait("AcmeWait", Duration.ofSeconds(1)); + return extension.reserve("custom-context").runInChildContext("AcmeContext", String.class, () -> "done"); + }); + + var result = runner.runUntilComplete("input"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals(OperationType.STEP, result.getOperation("custom-step").getType()); + assertEquals("AcmeStep", result.getOperation("custom-step").getSubtype()); + assertEquals(OperationType.WAIT, result.getOperation("custom-wait").getType()); + assertEquals("AcmeWait", result.getOperation("custom-wait").getSubtype()); + assertEquals( + OperationType.CONTEXT, result.getOperation("custom-context").getType()); + assertEquals("AcmeContext", result.getOperation("custom-context").getSubtype()); + } + private static String hash(String value) { try { var digest = MessageDigest.getInstance("SHA-256"); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java index efbd04ffa..65e25a6b8 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java @@ -46,12 +46,49 @@ default DurableFuture stepAsync(Class resultType, Supplier function DurableFuture stepAsync(TypeToken resultType, Supplier function, StepConfig config); + default T step(String subType, Class resultType, Supplier function) { + return step(subType, TypeToken.get(resultType), function); + } + + default T step(String subType, TypeToken resultType, Supplier function) { + return step(subType, resultType, function, StepConfig.builder().build()); + } + + default T step(String subType, Class resultType, Supplier function, StepConfig config) { + return step(subType, TypeToken.get(resultType), function, config); + } + + default T step(String subType, TypeToken resultType, Supplier function, StepConfig config) { + return stepAsync(subType, resultType, function, config).get(); + } + + default DurableFuture stepAsync(String subType, Class resultType, Supplier function) { + return stepAsync(subType, TypeToken.get(resultType), function); + } + + default DurableFuture stepAsync(String subType, TypeToken resultType, Supplier function) { + return stepAsync(subType, resultType, function, StepConfig.builder().build()); + } + + default DurableFuture stepAsync( + String subType, Class resultType, Supplier function, StepConfig config) { + return stepAsync(subType, TypeToken.get(resultType), function, config); + } + + DurableFuture stepAsync(String subType, TypeToken resultType, Supplier function, StepConfig config); + default Void wait(Duration duration) { return waitAsync(duration).get(); } DurableFuture waitAsync(Duration duration); + default Void wait(String subType, Duration duration) { + return waitAsync(subType, duration).get(); + } + + DurableFuture waitAsync(String subType, Duration duration); + default T invoke(String functionName, U payload, Class resultType) { return invoke(functionName, payload, TypeToken.get(resultType)); } @@ -84,6 +121,50 @@ default DurableFuture invokeAsync( DurableFuture invokeAsync(String functionName, U payload, TypeToken resultType, InvokeConfig config); + default T invoke(String subType, String functionName, U payload, Class resultType) { + return invoke(subType, functionName, payload, TypeToken.get(resultType)); + } + + default T invoke(String subType, String functionName, U payload, TypeToken resultType) { + return invoke( + subType, + functionName, + payload, + resultType, + InvokeConfig.builder().build()); + } + + default T invoke(String subType, String functionName, U payload, Class resultType, InvokeConfig config) { + return invoke(subType, functionName, payload, TypeToken.get(resultType), config); + } + + default T invoke( + String subType, String functionName, U payload, TypeToken resultType, InvokeConfig config) { + return invokeAsync(subType, functionName, payload, resultType, config).get(); + } + + default DurableFuture invokeAsync(String subType, String functionName, U payload, Class resultType) { + return invokeAsync(subType, functionName, payload, TypeToken.get(resultType)); + } + + default DurableFuture invokeAsync( + String subType, String functionName, U payload, TypeToken resultType) { + return invokeAsync( + subType, + functionName, + payload, + resultType, + InvokeConfig.builder().build()); + } + + default DurableFuture invokeAsync( + String subType, String functionName, U payload, Class resultType, InvokeConfig config) { + return invokeAsync(subType, functionName, payload, TypeToken.get(resultType), config); + } + + DurableFuture invokeAsync( + String subType, String functionName, U payload, TypeToken resultType, InvokeConfig config); + default DurableCallbackFuture createCallback(Class resultType) { return createCallback(TypeToken.get(resultType)); } @@ -98,6 +179,20 @@ default DurableCallbackFuture createCallback(Class resultType, Callbac DurableCallbackFuture createCallback(TypeToken resultType, CallbackConfig config); + default DurableCallbackFuture createCallback(String subType, Class resultType) { + return createCallback(subType, TypeToken.get(resultType)); + } + + default DurableCallbackFuture createCallback(String subType, TypeToken resultType) { + return createCallback(subType, resultType, CallbackConfig.builder().build()); + } + + default DurableCallbackFuture createCallback(String subType, Class resultType, CallbackConfig config) { + return createCallback(subType, TypeToken.get(resultType), config); + } + + DurableCallbackFuture createCallback(String subType, TypeToken resultType, CallbackConfig config); + default T runInChildContext(Class resultType, Supplier function) { return runInChildContext(TypeToken.get(resultType), function); } @@ -131,4 +226,40 @@ default DurableFuture runInChildContextAsync( DurableFuture runInChildContextAsync( TypeToken resultType, Supplier function, RunInChildContextConfig config); + + default T runInChildContext(String subType, Class resultType, Supplier function) { + return runInChildContext(subType, TypeToken.get(resultType), function); + } + + default T runInChildContext(String subType, TypeToken resultType, Supplier function) { + return runInChildContext( + subType, resultType, function, RunInChildContextConfig.builder().build()); + } + + default T runInChildContext( + String subType, Class resultType, Supplier function, RunInChildContextConfig config) { + return runInChildContext(subType, TypeToken.get(resultType), function, config); + } + + default T runInChildContext( + String subType, TypeToken resultType, Supplier function, RunInChildContextConfig config) { + return runInChildContextAsync(subType, resultType, function, config).get(); + } + + default DurableFuture runInChildContextAsync(String subType, Class resultType, Supplier function) { + return runInChildContextAsync(subType, TypeToken.get(resultType), function); + } + + default DurableFuture runInChildContextAsync(String subType, TypeToken resultType, Supplier function) { + return runInChildContextAsync( + subType, resultType, function, RunInChildContextConfig.builder().build()); + } + + default DurableFuture runInChildContextAsync( + String subType, Class resultType, Supplier function, RunInChildContextConfig config) { + return runInChildContextAsync(subType, TypeToken.get(resultType), function, config); + } + + DurableFuture runInChildContextAsync( + String subType, TypeToken resultType, Supplier function, RunInChildContextConfig config); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java index 4d09efe7f..b1a5a1e4d 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java @@ -10,6 +10,7 @@ import java.util.function.BiConsumer; import java.util.function.BiFunction; import java.util.function.Function; +import software.amazon.awssdk.services.lambda.model.OperationType; import software.amazon.lambda.durable.DurableCallbackFuture; import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.DurableContext; @@ -34,6 +35,7 @@ import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.execution.ThreadType; import software.amazon.lambda.durable.model.MapResult; +import software.amazon.lambda.durable.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.model.WaitForConditionResult; @@ -147,6 +149,16 @@ DurableFuture stepAsyncWithId( TypeToken resultType, Function func, StepConfig config) { + return stepAsyncWithId(operationId, name, OperationSubType.STEP.getValue(), resultType, func, config); + } + + DurableFuture stepAsyncWithId( + String operationId, + String name, + String subType, + TypeToken resultType, + Function func, + StepConfig config) { Objects.requireNonNull(config, "config cannot be null"); Objects.requireNonNull(resultType, "resultType cannot be null"); ParameterValidator.validateOperationName(name); @@ -156,7 +168,11 @@ DurableFuture stepAsyncWithId( } // Create and start step operation with TypeToken var operation = new StepOperation<>( - OperationIdentifier.of(operationId, name, OperationSubType.STEP), func, resultType, config, this); + new OperationDescriptor(operationId, name, OperationType.STEP, subType), + func, + resultType, + config, + this); operation.execute(); // Start the step (returns immediately) @@ -171,12 +187,16 @@ public DurableFuture waitAsync(String name, Duration duration) { } DurableFuture waitAsyncWithId(String operationId, String name, Duration duration) { + return waitAsyncWithId(operationId, name, OperationSubType.WAIT.getValue(), duration); + } + + DurableFuture waitAsyncWithId(String operationId, String name, String subType, Duration duration) { ParameterValidator.validateDuration(duration, "Wait duration"); ParameterValidator.validateOperationName(name); // Create and start wait operation - var operation = - new WaitOperation(OperationIdentifier.of(operationId, name, OperationSubType.WAIT), duration, this); + var operation = new WaitOperation( + new OperationDescriptor(operationId, name, OperationType.WAIT, subType), duration, this); operation.execute(); // Checkpoint the wait return operation; @@ -198,6 +218,24 @@ DurableFuture invokeAsyncWithId( U payload, TypeToken resultType, InvokeConfig config) { + return invokeAsyncWithId( + operationId, + name, + OperationSubType.CHAINED_INVOKE.getValue(), + functionName, + payload, + resultType, + config); + } + + DurableFuture invokeAsyncWithId( + String operationId, + String name, + String subType, + String functionName, + U payload, + TypeToken resultType, + InvokeConfig config) { Objects.requireNonNull(config, "config cannot be null"); Objects.requireNonNull(resultType, "resultType cannot be null"); ParameterValidator.validateOperationName(name); @@ -212,7 +250,7 @@ DurableFuture invokeAsyncWithId( } // Create and start invoke operation var operation = new InvokeOperation<>( - OperationIdentifier.of(operationId, name, OperationSubType.CHAINED_INVOKE), + new OperationDescriptor(operationId, name, OperationType.CHAINED_INVOKE, subType), functionName, payload, resultType, @@ -233,6 +271,11 @@ public DurableCallbackFuture createCallback(String name, TypeToken res DurableCallbackFuture createCallbackWithId( String operationId, String name, TypeToken resultType, CallbackConfig config) { + return createCallbackWithId(operationId, name, OperationSubType.CALLBACK.getValue(), resultType, config); + } + + DurableCallbackFuture createCallbackWithId( + String operationId, String name, String subType, TypeToken resultType, CallbackConfig config) { Objects.requireNonNull(config, "config cannot be null"); Objects.requireNonNull(resultType, "resultType cannot be null"); ParameterValidator.validateOperationName(name); @@ -240,7 +283,7 @@ DurableCallbackFuture createCallbackWithId( config = config.toBuilder().serDes(getDurableConfig().getSerDes()).build(); } var operation = new CallbackOperation<>( - OperationIdentifier.of(operationId, name, OperationSubType.CALLBACK), resultType, config, this); + new OperationDescriptor(operationId, name, OperationType.CALLBACK, subType), resultType, config, this); operation.execute(); return operation; @@ -284,7 +327,7 @@ DurableFuture runInChildContextAsyncWithId( Function func, RunInChildContextConfig config) { return runInChildContextAsyncWithId( - operationId, name, resultType, func, config, OperationSubType.RUN_IN_CHILD_CONTEXT); + operationId, name, OperationSubType.RUN_IN_CHILD_CONTEXT.getValue(), resultType, func, config); } private DurableFuture runInChildContextAsyncWithId( @@ -294,6 +337,16 @@ private DurableFuture runInChildContextAsyncWithId( Function func, RunInChildContextConfig config, OperationSubType subType) { + return runInChildContextAsyncWithId(operationId, name, subType.getValue(), resultType, func, config); + } + + DurableFuture runInChildContextAsyncWithId( + String operationId, + String name, + String subType, + TypeToken resultType, + Function func, + RunInChildContextConfig config) { Objects.requireNonNull(resultType, "resultType cannot be null"); Objects.requireNonNull(func, "func cannot be null"); Objects.requireNonNull(config, "RunInChildContextConfig cannot be null"); @@ -304,7 +357,11 @@ private DurableFuture runInChildContextAsyncWithId( } var operation = new ChildContextOperation<>( - OperationIdentifier.of(operationId, name, subType), func, resultType, config, this); + new OperationDescriptor(operationId, name, OperationType.CONTEXT, subType), + func, + resultType, + config, + this); operation.execute(); return operation; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java index ab34486a9..8c11850aa 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java @@ -3,6 +3,7 @@ package software.amazon.lambda.durable.context; import java.time.Duration; +import java.util.Objects; import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Supplier; import software.amazon.lambda.durable.DurableCallbackFuture; @@ -32,12 +33,27 @@ public DurableFuture stepAsync(TypeToken resultType, Supplier funct return context.stepAsyncWithId(operationId, name, resultType, ignored -> function.get(), config); } + @Override + public DurableFuture stepAsync( + String subType, TypeToken resultType, Supplier function, StepConfig config) { + validateSubType(subType); + claim(); + return context.stepAsyncWithId(operationId, name, subType, resultType, ignored -> function.get(), config); + } + @Override public DurableFuture waitAsync(Duration duration) { claim(); return context.waitAsyncWithId(operationId, name, duration); } + @Override + public DurableFuture waitAsync(String subType, Duration duration) { + validateSubType(subType); + claim(); + return context.waitAsyncWithId(operationId, name, subType, duration); + } + @Override public DurableFuture invokeAsync( String functionName, U payload, TypeToken resultType, InvokeConfig config) { @@ -45,12 +61,27 @@ public DurableFuture invokeAsync( return context.invokeAsyncWithId(operationId, name, functionName, payload, resultType, config); } + @Override + public DurableFuture invokeAsync( + String subType, String functionName, U payload, TypeToken resultType, InvokeConfig config) { + validateSubType(subType); + claim(); + return context.invokeAsyncWithId(operationId, name, subType, functionName, payload, resultType, config); + } + @Override public DurableCallbackFuture createCallback(TypeToken resultType, CallbackConfig config) { claim(); return context.createCallbackWithId(operationId, name, resultType, config); } + @Override + public DurableCallbackFuture createCallback(String subType, TypeToken resultType, CallbackConfig config) { + validateSubType(subType); + claim(); + return context.createCallbackWithId(operationId, name, subType, resultType, config); + } + @Override public DurableFuture runInChildContextAsync( TypeToken resultType, Supplier function, RunInChildContextConfig config) { @@ -58,6 +89,22 @@ public DurableFuture runInChildContextAsync( return context.runInChildContextAsyncWithId(operationId, name, resultType, ignored -> function.get(), config); } + @Override + public DurableFuture runInChildContextAsync( + String subType, TypeToken resultType, Supplier function, RunInChildContextConfig config) { + validateSubType(subType); + claim(); + return context.runInChildContextAsyncWithId( + operationId, name, subType, resultType, ignored -> function.get(), config); + } + + private void validateSubType(String subType) { + Objects.requireNonNull(subType, "subType cannot be null"); + if (subType.isBlank()) { + throw new IllegalArgumentException("subType cannot be blank"); + } + } + private void claim() { if (!claimed.compareAndSet(false, true)) { throw new IllegalStateException("An extension operation reservation can only be used once"); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/model/OperationDescriptor.java b/sdk/src/main/java/software/amazon/lambda/durable/model/OperationDescriptor.java new file mode 100644 index 000000000..ffc663d0d --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/model/OperationDescriptor.java @@ -0,0 +1,46 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.model; + +import java.util.Objects; +import software.amazon.awssdk.services.lambda.model.OperationType; + +/** + * Operation identity that permits extension-defined subtype strings. + * + * @param operationId globally unique operation ID + * @param name human-readable operation name + * @param operationType backend primitive operation type + * @param subType checkpoint subtype string + */ +public record OperationDescriptor(String operationId, String name, OperationType operationType, String subType) { + public OperationDescriptor { + Objects.requireNonNull(operationId, "operationId cannot be null"); + Objects.requireNonNull(operationType, "operationType cannot be null"); + Objects.requireNonNull(subType, "subType cannot be null"); + if (subType.isBlank()) { + throw new IllegalArgumentException("subType cannot be blank"); + } + } + + /** Converts an existing enum-based identity to a descriptor. */ + public static OperationDescriptor from(OperationIdentifier identifier) { + Objects.requireNonNull(identifier, "identifier cannot be null"); + return new OperationDescriptor( + identifier.operationId(), + identifier.name(), + identifier.operationType(), + identifier.subType().getValue()); + } + + /** Returns the matching SDK subtype, or {@code null} for an extension-defined value. */ + public OperationSubType standardSubType() { + for (var candidate : OperationSubType.values()) { + if (candidate.getOperationType() == operationType + && candidate.getValue().equals(subType)) { + return candidate; + } + } + return null; + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java index 6c26f35fc..9c2cd372d 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java @@ -25,6 +25,7 @@ import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.execution.ThreadContext; import software.amazon.lambda.durable.execution.ThreadType; +import software.amazon.lambda.durable.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.plugin.PluginInfoConverter; @@ -52,7 +53,7 @@ public abstract class BaseDurableOperation { private static final Logger logger = LoggerFactory.getLogger(BaseDurableOperation.class); - private final OperationIdentifier operationIdentifier; + private final OperationDescriptor operationDescriptor; protected final ExecutionManager executionManager; protected final CompletableFuture completionFuture; protected final BaseDurableOperation parentOperation; @@ -65,7 +66,14 @@ protected BaseDurableOperation( OperationIdentifier operationIdentifier, DurableContextImpl durableContext, BaseDurableOperation parentOperation) { - this(operationIdentifier, durableContext, parentOperation, false); + this(OperationDescriptor.from(operationIdentifier), durableContext, parentOperation, false); + } + + protected BaseDurableOperation( + OperationDescriptor operationDescriptor, + DurableContextImpl durableContext, + BaseDurableOperation parentOperation) { + this(operationDescriptor, durableContext, parentOperation, false); } /** @@ -81,7 +89,15 @@ protected BaseDurableOperation( DurableContextImpl durableContext, BaseDurableOperation parentOperation, boolean isVirtual) { - this.operationIdentifier = operationIdentifier; + this(OperationDescriptor.from(operationIdentifier), durableContext, parentOperation, isVirtual); + } + + protected BaseDurableOperation( + OperationDescriptor operationDescriptor, + DurableContextImpl durableContext, + BaseDurableOperation parentOperation, + boolean isVirtual) { + this.operationDescriptor = operationDescriptor; this.parentOperation = parentOperation; this.durableContext = durableContext; this.executionManager = durableContext.getExecutionManager(); @@ -108,17 +124,22 @@ public CompletableFuture completionFuture() { /** Gets the operation sub-type (e.g. RUN_IN_CHILD_CONTEXT, WAIT_FOR_CALLBACK). */ public OperationSubType getSubType() { - return operationIdentifier.subType(); + return operationDescriptor.standardSubType(); + } + + /** Gets the exact operation subtype string. */ + public String getSubTypeValue() { + return operationDescriptor.subType(); } /** Gets the unique identifier for this operation. */ public String getOperationId() { - return operationIdentifier.operationId(); + return operationDescriptor.operationId(); } /** Gets the operation name (may be null). */ public String getName() { - return operationIdentifier.name(); + return operationDescriptor.name(); } /** Gets the parent context. */ @@ -128,7 +149,7 @@ protected DurableContextImpl getContext() { /** Gets the operation type. */ public OperationType getType() { - return operationIdentifier.operationType(); + return operationDescriptor.operationType(); } /** @@ -359,7 +380,7 @@ protected void runUserHandler(Runnable runnable, ThreadType threadType) { protected T runUserFunction(Integer attempt, Supplier userFunction) { var pluginRunner = getPluginRunner(); var startInfo = PluginInfoConverter.toUserFunctionStartInfo( - operationIdentifier, durableContext.getParentId(), durableContext.isReplaying(), attempt); + operationDescriptor, durableContext.getParentId(), durableContext.isReplaying(), attempt); pluginRunner.onUserFunctionStart(startInfo); try { T result = userFunction.get(); @@ -492,7 +513,7 @@ protected CompletableFuture sendOperationUpdateAsync(OperationUpdate.Build var updateBuilder = builder.id(getOperationId()) .name(getName()) .type(getType()) - .subType(getSubType().getValue()) + .subType(getSubTypeValue()) .parentId(durableContext.getParentId()); var update = updateBuilder.build(); if (replayCompletedOperation.get()) { @@ -523,10 +544,10 @@ protected void validateReplay(Operation checkpointed) { getOperationId(), checkpointed.name(), getName()))); } - if (!Objects.equals(checkpointed.subType(), getSubType().getValue())) { + if (!Objects.equals(checkpointed.subType(), getSubTypeValue())) { throw terminateExecution(new NonDeterministicExecutionException(String.format( "Operation subType mismatch for \"%s\". Expected \"%s\", got \"%s\"", - getOperationId(), checkpointed.subType(), getSubType()))); + getOperationId(), checkpointed.subType(), getSubTypeValue()))); } } @@ -544,14 +565,14 @@ private PluginRunner getPluginRunner() { /** Fires onOperationStart plugin hook. */ private void fireOnOperationStart(Operation existing) { - var info = PluginInfoConverter.toOperationInfo(existing, operationIdentifier, durableContext.getParentId()); + var info = PluginInfoConverter.toOperationInfo(existing, operationDescriptor, durableContext.getParentId()); getPluginRunner().onOperationStart(info); } /** Fires onOperationEnd plugin hook when an operation reaches terminal status. */ protected void fireOnOperationEnd(Operation operation, Throwable error, boolean isReplay) { var info = PluginInfoConverter.toOperationEndInfo( - operation, operationIdentifier, durableContext.getParentId(), isReplay, error); + operation, operationDescriptor, durableContext.getParentId(), isReplay, error); getPluginRunner().onOperationEnd(info); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/CallbackOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/CallbackOperation.java index 9d9481fb9..b9dc73814 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/CallbackOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/CallbackOperation.java @@ -12,6 +12,7 @@ import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.exception.CallbackFailedException; import software.amazon.lambda.durable.exception.CallbackTimeoutException; +import software.amazon.lambda.durable.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; /** Durable operation for creating and waiting on external callbacks. */ @@ -30,6 +31,15 @@ public CallbackOperation( this.config = config; } + public CallbackOperation( + OperationDescriptor operationDescriptor, + TypeToken resultTypeToken, + CallbackConfig config, + DurableContextImpl durableContext) { + super(operationDescriptor, resultTypeToken, config.serDes(), durableContext); + this.config = config; + } + public String callbackId() { return callbackId; } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java index be2191726..9fde070c1 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java @@ -34,7 +34,9 @@ import software.amazon.lambda.durable.execution.ThreadType; import software.amazon.lambda.durable.logging.DurableLogger; import software.amazon.lambda.durable.model.DeserializedOperationResult; +import software.amazon.lambda.durable.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; +import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.util.ExceptionHelper; /** @@ -83,6 +85,16 @@ public ChildContextOperation( this.function = function; } + public ChildContextOperation( + OperationDescriptor operationDescriptor, + Function function, + TypeToken resultTypeToken, + RunInChildContextConfig config, + DurableContextImpl durableContext) { + super(operationDescriptor, resultTypeToken, config.serDes(), durableContext, null, config.isVirtual()); + this.function = function; + } + /** Starts the operation. */ @Override protected void start() { @@ -252,16 +264,16 @@ private Throwable translateException(Operation op, ErrorObject errorObject) { } // throw a general failed exception if a user exception is not reconstructed - return switch (getSubType()) { - case WAIT_FOR_CALLBACK -> handleWaitForCallbackFailure(); - case MAP_ITERATION -> new MapIterationFailedException(op); - case PARALLEL_BRANCH -> new ParallelBranchFailedException(op); - case RUN_IN_CHILD_CONTEXT, WITH_RETRY -> new ChildContextFailedException(op); - - // the following subtypes should not be able to reach here - case PARALLEL, MAP, WAIT_FOR_CONDITION, STEP, WAIT, CALLBACK, CHAINED_INVOKE -> - new IllegalStateException("Unexpected sub-type: " + getSubType()); - }; + if (OperationSubType.WAIT_FOR_CALLBACK.getValue().equals(getSubTypeValue())) { + return handleWaitForCallbackFailure(); + } + if (OperationSubType.MAP_ITERATION.getValue().equals(getSubTypeValue())) { + return new MapIterationFailedException(op); + } + if (OperationSubType.PARALLEL_BRANCH.getValue().equals(getSubTypeValue())) { + return new ParallelBranchFailedException(op); + } + return new ChildContextFailedException(op); } private Operation createVirtualOperation(ErrorObject errorObject) { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/InvokeOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/InvokeOperation.java index 9e2c54ace..27a301b6a 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/InvokeOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/InvokeOperation.java @@ -13,6 +13,7 @@ import software.amazon.lambda.durable.exception.InvokeFailedException; import software.amazon.lambda.durable.exception.InvokeStoppedException; import software.amazon.lambda.durable.exception.InvokeTimedOutException; +import software.amazon.lambda.durable.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.serde.SerDes; @@ -43,6 +44,20 @@ public InvokeOperation( this.payloadSerDes = config.payloadSerDes() != null ? config.payloadSerDes() : config.serDes(); } + public InvokeOperation( + OperationDescriptor operationDescriptor, + String functionName, + I payload, + TypeToken resultTypeToken, + InvokeConfig config, + DurableContextImpl durableContext) { + super(operationDescriptor, resultTypeToken, config.serDes(), durableContext); + this.functionName = functionName; + this.payload = payload; + this.invokeConfig = config; + this.payloadSerDes = config.payloadSerDes() != null ? config.payloadSerDes() : config.serDes(); + } + /** Starts the operation. */ @Override protected void start() { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/SerializableDurableOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/SerializableDurableOperation.java index 6457c996d..90bbcf64c 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/SerializableDurableOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/SerializableDurableOperation.java @@ -9,6 +9,7 @@ import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.exception.SerDesException; +import software.amazon.lambda.durable.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.serde.SerDes; import software.amazon.lambda.durable.util.ExceptionHelper; @@ -55,6 +56,14 @@ protected SerializableDurableOperation( this(operationIdentifier, resultTypeToken, resultSerDes, durableContext, null, false); } + protected SerializableDurableOperation( + OperationDescriptor operationDescriptor, + TypeToken resultTypeToken, + SerDes resultSerDes, + DurableContextImpl durableContext) { + this(operationDescriptor, resultTypeToken, resultSerDes, durableContext, null, false); + } + /** * Constructs a new durable operation. * @@ -77,6 +86,18 @@ protected SerializableDurableOperation( this.resultSerDes = resultSerDes; } + protected SerializableDurableOperation( + OperationDescriptor operationDescriptor, + TypeToken resultTypeToken, + SerDes resultSerDes, + DurableContextImpl durableContext, + BaseDurableOperation parentOperation, + boolean isVirtual) { + super(operationDescriptor, durableContext, parentOperation, isVirtual); + this.resultTypeToken = resultTypeToken; + this.resultSerDes = resultSerDes; + } + /** * Deserializes a result string into the operation's result type. * diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java index d2b3cc7f4..84a9edd58 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java @@ -24,6 +24,7 @@ import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.execution.ThreadType; import software.amazon.lambda.durable.logging.DurableLogger; +import software.amazon.lambda.durable.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.util.ExceptionHelper; @@ -53,6 +54,17 @@ public StepOperation( this.config = config; } + public StepOperation( + OperationDescriptor operationDescriptor, + Function function, + TypeToken resultTypeToken, + StepConfig config, + DurableContextImpl durableContext) { + super(operationDescriptor, resultTypeToken, config.serDes(), durableContext); + this.function = function; + this.config = config; + } + /** Starts the operation. */ @Override protected void start() { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitOperation.java index 3e53ff736..1cd0ac040 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitOperation.java @@ -13,6 +13,7 @@ import software.amazon.awssdk.services.lambda.model.WaitOptions; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; /** @@ -33,6 +34,12 @@ public WaitOperation( this.duration = duration; } + public WaitOperation( + OperationDescriptor operationDescriptor, Duration duration, DurableContextImpl durableContext) { + super(operationDescriptor, durableContext, null); + this.duration = duration; + } + /** Starts the operation. */ @Override protected void start() { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginInfoConverter.java b/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginInfoConverter.java index 58911f7be..5af7f3a2a 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginInfoConverter.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginInfoConverter.java @@ -6,6 +6,7 @@ import java.util.Collection; import java.util.stream.Collectors; import software.amazon.awssdk.services.lambda.model.Operation; +import software.amazon.lambda.durable.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.operation.BaseDurableOperation; @@ -28,11 +29,23 @@ private PluginInfoConverter() {} * @return an OperationInfo record */ public static OperationInfo toOperationInfo(Operation operation, OperationIdentifier identifier, String parentId) { + return toOperationInfo(operation, OperationDescriptor.from(identifier), parentId); + } + + /** + * Converts an SDK {@link Operation} to an {@link OperationInfo} using an {@link OperationDescriptor}. + * + * @param operation the SDK operation (may be null for first-start scenarios) + * @param descriptor the operation identity + * @param parentId the parent operation ID (may be null for root operations) + * @return an OperationInfo record + */ + public static OperationInfo toOperationInfo(Operation operation, OperationDescriptor descriptor, String parentId) { return new OperationInfo( - identifier.operationId(), - identifier.name(), - identifier.operationType() != null ? identifier.operationType().toString() : null, - identifier.subType() != null ? identifier.subType().getValue() : null, + descriptor.operationId(), + descriptor.name(), + descriptor.operationType().toString(), + descriptor.subType(), parentId, operation != null ? operation.startTimestamp() : Instant.now(), operation != null ? operation.endTimestamp() : null, @@ -54,11 +67,20 @@ public static OperationInfo toOperationInfo(Operation operation, OperationIdenti */ public static OperationEndInfo toOperationEndInfo( Operation operation, OperationIdentifier identifier, String parentId, boolean isReplay, Throwable error) { + return toOperationEndInfo(operation, OperationDescriptor.from(identifier), parentId, isReplay, error); + } + + /** + * Creates an {@link OperationEndInfo} from an SDK {@link Operation}, an {@link OperationDescriptor}, and an + * optional error. + */ + public static OperationEndInfo toOperationEndInfo( + Operation operation, OperationDescriptor descriptor, String parentId, boolean isReplay, Throwable error) { return new OperationEndInfo( - identifier.operationId(), - identifier.name(), - identifier.operationType() != null ? identifier.operationType().toString() : null, - identifier.subType() != null ? identifier.subType().getValue() : null, + descriptor.operationId(), + descriptor.name(), + descriptor.operationType().toString(), + descriptor.subType(), parentId, operation != null ? operation.startTimestamp() : null, operation != null ? operation.endTimestamp() : null, @@ -83,11 +105,17 @@ public static OperationEndInfo toOperationEndInfo( */ public static UserFunctionStartInfo toUserFunctionStartInfo( OperationIdentifier identifier, String parentId, boolean isReplayingChildren, Integer attempt) { + return toUserFunctionStartInfo(OperationDescriptor.from(identifier), parentId, isReplayingChildren, attempt); + } + + /** Creates a user-function start record from an operation descriptor. */ + public static UserFunctionStartInfo toUserFunctionStartInfo( + OperationDescriptor descriptor, String parentId, boolean isReplayingChildren, Integer attempt) { return new UserFunctionStartInfo( - identifier.operationId(), - identifier.name(), - identifier.operationType() != null ? identifier.operationType().toString() : null, - identifier.subType() != null ? identifier.subType().getValue() : null, + descriptor.operationId(), + descriptor.name(), + descriptor.operationType().toString(), + descriptor.subType(), parentId, Instant.now(), isReplayingChildren, diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java index 653cb8e48..c56baa5b4 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java @@ -97,6 +97,21 @@ void reservedStepAdaptsSupplierToStepFunction() { assertEquals(true, called.get()); } + @Test + void customSubtypeStepDelegatesExactSubtype() { + var context = mock(DurableContextImpl.class); + var future = mockStringFuture(); + var resultType = TypeToken.get(String.class); + var config = StepConfig.builder().build(); + when(context.stepAsyncWithId(eq("1"), eq("step"), eq("AcmeStep"), eq(resultType), any(), eq(config))) + .thenReturn(future); + + var actual = new ExtensionOperationImpl(context, "1", "step") + .stepAsync("AcmeStep", resultType, () -> "result", config); + + assertEquals(future, actual); + } + @Test void reservationDelegatesWaitInvokeAndCallback() { var duration = Duration.ofSeconds(2); @@ -128,6 +143,63 @@ void reservationDelegatesWaitInvokeAndCallback() { .createCallback(resultType, callbackConfig)); } + @Test + void customSubtypeSelectorsDelegateExactSubtype() { + var duration = Duration.ofSeconds(2); + var resultType = TypeToken.get(String.class); + + var waitContext = mock(DurableContextImpl.class); + var waitFuture = mockFuture(); + when(waitContext.waitAsyncWithId("1", "wait", "AcmeWait", duration)).thenReturn(waitFuture); + assertEquals(waitFuture, new ExtensionOperationImpl(waitContext, "1", "wait").waitAsync("AcmeWait", duration)); + + var invokeContext = mock(DurableContextImpl.class); + var invokeFuture = mockStringFuture(); + var invokeConfig = InvokeConfig.builder().build(); + when(invokeContext.invokeAsyncWithId( + "2", "invoke", "AcmeInvoke", "target", "payload", resultType, invokeConfig)) + .thenReturn(invokeFuture); + assertEquals( + invokeFuture, + new ExtensionOperationImpl(invokeContext, "2", "invoke") + .invokeAsync("AcmeInvoke", "target", "payload", resultType, invokeConfig)); + + var callbackContext = mock(DurableContextImpl.class); + @SuppressWarnings("unchecked") + var callbackFuture = (DurableCallbackFuture) mock(DurableCallbackFuture.class); + var callbackConfig = CallbackConfig.builder().build(); + when(callbackContext.createCallbackWithId("3", "callback", "AcmeCallback", resultType, callbackConfig)) + .thenReturn(callbackFuture); + assertEquals( + callbackFuture, + new ExtensionOperationImpl(callbackContext, "3", "callback") + .createCallback("AcmeCallback", resultType, callbackConfig)); + + var childContext = mock(DurableContextImpl.class); + var childFuture = mockStringFuture(); + var childConfig = RunInChildContextConfig.builder().build(); + when(childContext.runInChildContextAsyncWithId( + eq("4"), eq("child"), eq("AcmeContext"), eq(resultType), any(), eq(childConfig))) + .thenReturn(childFuture); + assertEquals( + childFuture, + new ExtensionOperationImpl(childContext, "4", "child") + .runInChildContextAsync("AcmeContext", resultType, () -> "result", childConfig)); + } + + @Test + void invalidSubtypeDoesNotClaimReservation() { + var context = mock(DurableContextImpl.class); + var duration = Duration.ofSeconds(1); + var future = mockFuture(); + when(context.waitAsyncWithId("1", "wait", duration)).thenReturn(future); + var operation = new ExtensionOperationImpl(context, "1", "wait"); + + assertThrows(NullPointerException.class, () -> operation.waitAsync(null, duration)); + assertThrows(IllegalArgumentException.class, () -> operation.waitAsync(" ", duration)); + assertEquals(future, operation.waitAsync(duration)); + } + @Test void reservedChildContextAdaptsSupplierToChildFunction() { var context = mock(DurableContextImpl.class); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginInfoConverterTest.java b/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginInfoConverterTest.java index dbba0870f..171c21203 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginInfoConverterTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginInfoConverterTest.java @@ -8,6 +8,8 @@ import org.junit.jupiter.api.Test; 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.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; @@ -50,6 +52,16 @@ void toOperationInfo_withIdentifier_mapsAllFields() { assertEquals("STARTED", info.status()); } + @Test + void toOperationInfo_withDescriptor_preservesCustomSubtype() { + var descriptor = new OperationDescriptor(OPERATION_ID, OPERATION_NAME, OperationType.STEP, "AcmeStep"); + + var info = PluginInfoConverter.toOperationInfo(null, descriptor, PARENT_ID); + + assertEquals("STEP", info.type()); + assertEquals("AcmeStep", info.subType()); + } + @Test void toOperationInfo_withIdentifier_nullOperation_usesCurrentTime() { var before = Instant.now(); From 783dcf06ba478a51941c3a27bd09c94189eaca71 Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 03:36:47 +0000 Subject: [PATCH 16/40] feat: add stateful extension steps --- .../ExtensionOperationIntegrationTest.java | 24 ++++ .../lambda/durable/ExtensionOperation.java | 19 +++ .../lambda/durable/ExtensionStepFunction.java | 19 +++ .../lambda/durable/ExtensionStepResult.java | 37 +++++ .../durable/config/ExtensionStepConfig.java | 60 +++++++++ .../durable/context/DurableContextImpl.java | 30 +++++ .../context/ExtensionOperationImpl.java | 12 ++ .../durable/operation/StepOperation.java | 127 +++++++++++++++++- .../durable/ExtensionStepResultTest.java | 32 +++++ .../config/ExtensionStepConfigTest.java | 31 +++++ .../context/ExtensionOperationImplTest.java | 19 +++ 11 files changed, 408 insertions(+), 2 deletions(-) create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/ExtensionStepFunction.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/ExtensionStepResult.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/config/ExtensionStepConfig.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/ExtensionStepResultTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionStepConfigTest.java diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java index 5e4f146d5..f56b4fa19 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java @@ -18,6 +18,7 @@ import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.Test; import software.amazon.awssdk.services.lambda.model.OperationType; +import software.amazon.lambda.durable.config.ExtensionStepConfig; import software.amazon.lambda.durable.model.ExecutionStatus; import software.amazon.lambda.durable.testing.LocalDurableTestRunner; @@ -121,6 +122,29 @@ void customPrimitiveSubtypesAreStoredWithoutChangingOperationTypes() { assertEquals("AcmeContext", result.getOperation("custom-context").getSubtype()); } + @Test + void statefulExtensionStepCheckpointsStateAcrossRetries() { + var runner = + LocalDurableTestRunner.create(Integer.class, (input, context) -> ExtensionContext.getCurrentContext() + .reserve("stateful") + .step( + "AcmeStateful", + Integer.class, + state -> state >= 2 + ? ExtensionStepResult.succeed(state) + : ExtensionStepResult.retry(state + 1, Duration.ofSeconds(1)), + ExtensionStepConfig.builder() + .initialState(0) + .build())); + + var result = runner.runUntilComplete(0); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals(2, result.getResult(Integer.class)); + assertEquals("AcmeStateful", result.getOperation("stateful").getSubtype()); + assertEquals(3, result.getOperation("stateful").getAttempt()); + } + private static String hash(String value) { try { var digest = MessageDigest.getInstance("SHA-256"); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java index 65e25a6b8..329acfc58 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java @@ -5,6 +5,7 @@ import java.time.Duration; import java.util.function.Supplier; import software.amazon.lambda.durable.config.CallbackConfig; +import software.amazon.lambda.durable.config.ExtensionStepConfig; import software.amazon.lambda.durable.config.InvokeConfig; import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.config.StepConfig; @@ -77,6 +78,24 @@ default DurableFuture stepAsync( DurableFuture stepAsync(String subType, TypeToken resultType, Supplier function, StepConfig config); + default T step( + String subType, Class resultType, ExtensionStepFunction function, ExtensionStepConfig config) { + return step(subType, TypeToken.get(resultType), function, config); + } + + default T step( + String subType, TypeToken resultType, ExtensionStepFunction function, ExtensionStepConfig config) { + return stepAsync(subType, resultType, function, config).get(); + } + + default DurableFuture stepAsync( + String subType, Class resultType, ExtensionStepFunction function, ExtensionStepConfig config) { + return stepAsync(subType, TypeToken.get(resultType), function, config); + } + + DurableFuture stepAsync( + String subType, TypeToken resultType, ExtensionStepFunction function, ExtensionStepConfig config); + default Void wait(Duration duration) { return waitAsync(duration).get(); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionStepFunction.java b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionStepFunction.java new file mode 100644 index 000000000..48a06fd08 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionStepFunction.java @@ -0,0 +1,19 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +/** + * Evaluates one attempt of a stateful extension step. + * + * @param the checkpointed state and final result type + */ +@FunctionalInterface +public interface ExtensionStepFunction { + /** + * Evaluates the current state. + * + * @param state state restored from the prior retry checkpoint, or the configured initial state + * @return a success or retry outcome + */ + ExtensionStepResult apply(T state); +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionStepResult.java b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionStepResult.java new file mode 100644 index 000000000..fb0404927 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionStepResult.java @@ -0,0 +1,37 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import java.time.Duration; +import java.util.Objects; + +/** + * Fixed outcomes supported by a stateful extension STEP primitive. + * + * @param the checkpointed state and final result type + */ +public sealed interface ExtensionStepResult permits ExtensionStepResult.Succeeded, ExtensionStepResult.Retry { + + /** Creates a terminal successful outcome. */ + static Succeeded succeed(T value) { + return new Succeeded<>(value); + } + + /** Creates a retry outcome with checkpointed state and delay. */ + static Retry retry(T state, Duration delay) { + return new Retry<>(state, delay); + } + + /** Terminal successful outcome. */ + record Succeeded(T value) implements ExtensionStepResult {} + + /** Retry outcome. */ + record Retry(T state, Duration delay) implements ExtensionStepResult { + public Retry { + Objects.requireNonNull(delay, "delay cannot be null"); + if (delay.isNegative()) { + throw new IllegalArgumentException("delay cannot be negative"); + } + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/ExtensionStepConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/ExtensionStepConfig.java new file mode 100644 index 000000000..fd6a06321 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/ExtensionStepConfig.java @@ -0,0 +1,60 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.config; + +import software.amazon.lambda.durable.serde.SerDes; + +/** + * Configuration for a stateful extension STEP primitive. + * + * @param the checkpointed state and final result type + */ +public final class ExtensionStepConfig { + private final T initialState; + private final SerDes serDes; + + private ExtensionStepConfig(Builder builder) { + initialState = builder.initialState; + serDes = builder.serDes; + } + + /** Returns the state supplied to the first attempt. */ + public T initialState() { + return initialState; + } + + /** Returns the custom serializer, or {@code null} to use the durable configuration default. */ + public SerDes serDes() { + return serDes; + } + + /** Returns a builder for a stateful extension step. */ + public static Builder builder() { + return new Builder<>(); + } + + /** Builder for {@link ExtensionStepConfig}. */ + public static final class Builder { + private T initialState; + private SerDes serDes; + + private Builder() {} + + /** Sets the state supplied to the first attempt. */ + public Builder initialState(T initialState) { + this.initialState = initialState; + return this; + } + + /** Sets the serializer for checkpointed state and the final result. */ + public Builder serDes(SerDes serDes) { + this.serDes = serDes; + return this; + } + + /** Builds the immutable configuration. */ + public ExtensionStepConfig build() { + return new ExtensionStepConfig<>(this); + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java index b1a5a1e4d..7c1138d14 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java @@ -17,10 +17,12 @@ import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.ExtensionContext; import software.amazon.lambda.durable.ExtensionOperation; +import software.amazon.lambda.durable.ExtensionStepFunction; import software.amazon.lambda.durable.ParallelDurableFuture; import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.CallbackConfig; +import software.amazon.lambda.durable.config.ExtensionStepConfig; import software.amazon.lambda.durable.config.InvokeConfig; import software.amazon.lambda.durable.config.MapConfig; import software.amazon.lambda.durable.config.ParallelConfig; @@ -179,6 +181,34 @@ DurableFuture stepAsyncWithId( return operation; } + DurableFuture extensionStepAsyncWithId( + String operationId, + String name, + String subType, + TypeToken resultType, + ExtensionStepFunction function, + ExtensionStepConfig config) { + Objects.requireNonNull(resultType, "resultType cannot be null"); + Objects.requireNonNull(function, "function cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + ParameterValidator.validateOperationName(name); + + if (config.serDes() == null) { + config = ExtensionStepConfig.builder() + .initialState(config.initialState()) + .serDes(getDurableConfig().getSerDes()) + .build(); + } + var operation = new StepOperation<>( + new OperationDescriptor(operationId, name, OperationType.STEP, subType), + function, + resultType, + config, + this); + operation.execute(); + return operation; + } + @Override public DurableFuture waitAsync(String name, Duration duration) { ParameterValidator.validateDuration(duration, "Wait duration"); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java index 8c11850aa..7fe416176 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java @@ -9,8 +9,10 @@ import software.amazon.lambda.durable.DurableCallbackFuture; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.ExtensionOperation; +import software.amazon.lambda.durable.ExtensionStepFunction; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.CallbackConfig; +import software.amazon.lambda.durable.config.ExtensionStepConfig; import software.amazon.lambda.durable.config.InvokeConfig; import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.config.StepConfig; @@ -41,6 +43,16 @@ public DurableFuture stepAsync( return context.stepAsyncWithId(operationId, name, subType, resultType, ignored -> function.get(), config); } + @Override + public DurableFuture stepAsync( + String subType, TypeToken resultType, ExtensionStepFunction function, ExtensionStepConfig config) { + validateSubType(subType); + Objects.requireNonNull(function, "function cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + claim(); + return context.extensionStepAsyncWithId(operationId, name, subType, resultType, function, config); + } + @Override public DurableFuture waitAsync(Duration duration) { claim(); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java index 84a9edd58..bf753b2f2 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java @@ -11,8 +11,11 @@ import software.amazon.awssdk.services.lambda.model.OperationStatus; import software.amazon.awssdk.services.lambda.model.OperationUpdate; import software.amazon.awssdk.services.lambda.model.StepOptions; +import software.amazon.lambda.durable.ExtensionStepFunction; +import software.amazon.lambda.durable.ExtensionStepResult; import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.ExtensionStepConfig; import software.amazon.lambda.durable.config.StepConfig; import software.amazon.lambda.durable.config.StepSemantics; import software.amazon.lambda.durable.context.BaseContextImpl; @@ -41,6 +44,8 @@ public class StepOperation extends SerializableDurableOperation { private final Function function; private final StepConfig config; + private final ExtensionStepFunction extensionFunction; + private final ExtensionStepConfig extensionConfig; public StepOperation( OperationIdentifier operationIdentifier, @@ -52,6 +57,8 @@ public StepOperation( this.function = function; this.config = config; + this.extensionFunction = null; + this.extensionConfig = null; } public StepOperation( @@ -63,17 +70,40 @@ public StepOperation( super(operationDescriptor, resultTypeToken, config.serDes(), durableContext); this.function = function; this.config = config; + this.extensionFunction = null; + this.extensionConfig = null; + } + + public StepOperation( + OperationDescriptor operationDescriptor, + ExtensionStepFunction function, + TypeToken resultTypeToken, + ExtensionStepConfig config, + DurableContextImpl durableContext) { + super(operationDescriptor, resultTypeToken, config.serDes(), durableContext); + this.function = null; + this.config = null; + this.extensionFunction = function; + this.extensionConfig = config; } /** Starts the operation. */ @Override protected void start() { - executeStepLogic(FIRST_ATTEMPT); + if (isExtensionStep()) { + executeExtensionStepLogic(extensionConfig.initialState(), FIRST_ATTEMPT); + } else { + executeStepLogic(FIRST_ATTEMPT); + } } /** Replays the operation. */ @Override protected void replay(Operation existing) { + if (isExtensionStep()) { + replayExtensionStep(existing); + return; + } var attempt = existing.stepDetails() != null && existing.stepDetails().attempt() != null ? existing.stepDetails().attempt() + 1 : FIRST_ATTEMPT; @@ -105,6 +135,34 @@ protected void replay(Operation existing) { } } + private void replayExtensionStep(Operation existing) { + switch (existing.status()) { + case SUCCEEDED, FAILED -> markAlreadyCompleted(); + case PENDING -> pollReadyAndResumeExtensionStep(); + case STARTED, READY -> resumeExtensionStep(existing); + default -> + throw terminateExecutionWithIllegalDurableOperationException( + "Unexpected extension step status: " + existing.status()); + } + } + + private void resumeExtensionStep(Operation existing) { + var details = existing.stepDetails(); + var attempt = details != null && details.attempt() != null ? details.attempt() + 1 : FIRST_ATTEMPT; + var state = details != null && details.result() != null + ? deserializeResult(details.result()) + : extensionConfig.initialState(); + executeExtensionStepLogic(state, attempt); + } + + private void pollReadyAndResumeExtensionStep() { + pollForOperationUpdates() + .thenCompose(op -> op.status() == OperationStatus.READY + ? CompletableFuture.completedFuture(op) + : pollForOperationUpdates()) + .thenAccept(this::resumeExtensionStep); + } + private void pollReadyAndExecuteStepLogic(Instant nextAttemptInstant, int attempt) { pollForOperationUpdates(nextAttemptInstant) .thenCompose(op -> op.status() == OperationStatus.READY @@ -139,6 +197,67 @@ private void executeStepLogic(int attempt) { runUserHandler(userHandler, ThreadType.STEP); } + private void executeExtensionStepLogic(T state, int attempt) { + Runnable userHandler = () -> { + var stepContext = getContext().createStepContext(getOperationId(), getName(), attempt); + try (var ignoredContext = BaseContextImpl.attachCurrentContext(stepContext); + var ignoredLogger = DurableLogger.attachContext()) { + try { + checkpointStarted(); + var result = runUserFunction(attempt, () -> extensionFunction.apply(state)); + handleExtensionStepResult(result, attempt); + } catch (Throwable e) { + handleExtensionStepFailure(e); + } + } + }; + runUserHandler(userHandler, ThreadType.STEP); + } + + private void handleExtensionStepResult(ExtensionStepResult result, int attempt) { + if (result == null) { + throw new NullPointerException("Extension step function result cannot be null"); + } + if (result instanceof ExtensionStepResult.Succeeded succeeded) { + handleStepSucceeded(succeeded.value()); + return; + } + var retry = (ExtensionStepResult.Retry) result; + var serializedState = serializeAndDeserializeResult(retry.state()); + var retryDelaySeconds = Math.toIntExact(retry.delay().toSeconds()); + var update = OperationUpdate.builder() + .action(OperationAction.RETRY) + .payload(serializedState.serialized()) + .stepOptions(StepOptions.builder() + .nextAttemptDelaySeconds(retryDelaySeconds) + .build()); + sendOperationUpdate(update); + pollReadyAndExecuteExtensionStep(serializedState.deserialized(), attempt + 1); + } + + private void pollReadyAndExecuteExtensionStep(T state, int attempt) { + pollForOperationUpdates() + .thenCompose(op -> op.status() == OperationStatus.READY + ? CompletableFuture.completedFuture(op) + : pollForOperationUpdates()) + .thenRun(() -> executeExtensionStepLogic(state, attempt)); + } + + private void handleExtensionStepFailure(Throwable exception) { + exception = ExceptionHelper.unwrapCompletableFuture(exception); + if (exception instanceof SuspendExecutionException suspendExecutionException) { + throw suspendExecutionException; + } + if (exception instanceof UnrecoverableDurableExecutionException unrecoverable) { + throw terminateExecution(unrecoverable); + } + var error = exception instanceof DurableOperationException durableOperationException + ? durableOperationException.getErrorObject() + : serializeException(exception); + sendOperationUpdate( + OperationUpdate.builder().action(OperationAction.FAIL).error(error)); + } + private void checkpointStarted() { // Check if we need to send START var existing = getOperation(); @@ -237,6 +356,10 @@ public T get() { } private boolean isAtMostOnce() { - return config.semanticsPerRetry() == StepSemantics.AT_MOST_ONCE_PER_RETRY; + return config != null && config.semanticsPerRetry() == StepSemantics.AT_MOST_ONCE_PER_RETRY; + } + + private boolean isExtensionStep() { + return extensionFunction != null; } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/ExtensionStepResultTest.java b/sdk/src/test/java/software/amazon/lambda/durable/ExtensionStepResultTest.java new file mode 100644 index 000000000..c18581628 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/ExtensionStepResultTest.java @@ -0,0 +1,32 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.time.Duration; +import org.junit.jupiter.api.Test; + +class ExtensionStepResultTest { + @Test + void succeedCarriesFinalValue() { + var result = ExtensionStepResult.succeed("done"); + + assertEquals("done", result.value()); + } + + @Test + void retryCarriesStateAndDelay() { + var result = ExtensionStepResult.retry("next", Duration.ofSeconds(2)); + + assertEquals("next", result.state()); + assertEquals(Duration.ofSeconds(2), result.delay()); + } + + @Test + void retryRejectsInvalidDelay() { + assertThrows(NullPointerException.class, () -> ExtensionStepResult.retry("next", null)); + assertThrows(IllegalArgumentException.class, () -> ExtensionStepResult.retry("next", Duration.ofSeconds(-1))); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionStepConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionStepConfigTest.java new file mode 100644 index 000000000..592858800 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionStepConfigTest.java @@ -0,0 +1,31 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.config; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.serde.JacksonSerDes; + +class ExtensionStepConfigTest { + @Test + void builderDefaultsToNullStateAndSerDes() { + var config = ExtensionStepConfig.builder().build(); + + assertNull(config.initialState()); + assertNull(config.serDes()); + } + + @Test + void builderRetainsStateAndSerDes() { + var serDes = new JacksonSerDes(); + var config = ExtensionStepConfig.builder() + .initialState(42) + .serDes(serDes) + .build(); + + assertEquals(42, config.initialState()); + assertEquals(serDes, config.serDes()); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java index c56baa5b4..ddb8867e2 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java @@ -21,9 +21,11 @@ import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.ExtensionOperation; +import software.amazon.lambda.durable.ExtensionStepResult; import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.CallbackConfig; +import software.amazon.lambda.durable.config.ExtensionStepConfig; import software.amazon.lambda.durable.config.InvokeConfig; import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.config.StepConfig; @@ -112,6 +114,23 @@ void customSubtypeStepDelegatesExactSubtype() { assertEquals(future, actual); } + @Test + void statefulStepDelegatesWithoutExposingStepContext() { + var context = mock(DurableContextImpl.class); + var future = mockStringFuture(); + var resultType = TypeToken.get(String.class); + var config = + ExtensionStepConfig.builder().initialState("initial").build(); + when(context.extensionStepAsyncWithId( + eq("1"), eq("step"), eq("AcmeStateful"), eq(resultType), any(), eq(config))) + .thenReturn(future); + + var actual = new ExtensionOperationImpl(context, "1", "step") + .stepAsync("AcmeStateful", resultType, state -> ExtensionStepResult.succeed(state + "-done"), config); + + assertEquals(future, actual); + } + @Test void reservationDelegatesWaitInvokeAndCallback() { var duration = Duration.ofSeconds(2); From 3f638b3053dd282b82f26ded5ac21537ef5fb535 Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 03:46:18 +0000 Subject: [PATCH 17/40] feat: add configurable extension contexts --- .../ExtensionOperationIntegrationTest.java | 34 ++++ .../ExtensionChildOperationSummary.java | 11 ++ .../durable/ExtensionContextErrorHandler.java | 10 ++ .../durable/ExtensionContextFailure.java | 48 +++++ .../durable/ExtensionContextFunction.java | 10 ++ .../ExtensionContextReplayContext.java | 40 +++++ .../durable/ExtensionContextResult.java | 58 ++++++ .../lambda/durable/ExtensionOperation.java | 25 +++ .../config/ExtensionContextConfig.java | 84 +++++++++ .../durable/context/DurableContextImpl.java | 55 +++++- .../context/ExtensionOperationImpl.java | 16 ++ .../operation/ChildContextOperation.java | 169 +++++++++++++++--- .../ExtensionContextReplayContextTest.java | 36 ++++ .../durable/ExtensionContextResultTest.java | 44 +++++ .../config/ExtensionContextConfigTest.java | 41 +++++ .../context/ExtensionOperationImplTest.java | 19 ++ .../operation/ChildContextOperationTest.java | 126 +++++++++++++ 17 files changed, 801 insertions(+), 25 deletions(-) create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/ExtensionChildOperationSummary.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextErrorHandler.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextFailure.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextFunction.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextReplayContext.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextResult.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/config/ExtensionContextConfig.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/ExtensionContextReplayContextTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/ExtensionContextResultTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionContextConfigTest.java diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java index f56b4fa19..0b87f0dcf 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java @@ -16,8 +16,10 @@ import java.time.Duration; import java.util.HexFormat; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.Test; import software.amazon.awssdk.services.lambda.model.OperationType; +import software.amazon.lambda.durable.config.ExtensionContextConfig; import software.amazon.lambda.durable.config.ExtensionStepConfig; import software.amazon.lambda.durable.model.ExecutionStatus; import software.amazon.lambda.durable.testing.LocalDurableTestRunner; @@ -145,6 +147,38 @@ void statefulExtensionStepCheckpointsStateAcrossRetries() { assertEquals(3, result.getOperation("stateful").getAttempt()); } + @Test + void extensionContextExposesStoredReplayStateWhileReplayingChildren() { + var replayState = new AtomicReference(); + var executions = new AtomicInteger(); + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { + var result = ExtensionContext.getCurrentContext() + .reserve("advanced") + .runInChildContext( + "AcmeContext", + String.class, + () -> { + executions.incrementAndGet(); + var replay = ExtensionContextReplayContext.getCurrentContext(); + if (replay.isReplayingChildren()) { + replayState.set(replay.getReplayState()); + } + return ExtensionContextResult.replayChildren("full", "stored"); + }, + ExtensionContextConfig.builder().build()); + context.wait("replay", Duration.ofSeconds(1)); + return result; + }); + + var result = runner.runUntilComplete("input"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("full", result.getResult(String.class)); + assertTrue(executions.get() >= 2); + assertEquals("stored", replayState.get()); + assertThrows(IllegalStateException.class, ExtensionContextReplayContext::getCurrentContext); + } + private static String hash(String value) { try { var digest = MessageDigest.getInstance("SHA-256"); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionChildOperationSummary.java b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionChildOperationSummary.java new file mode 100644 index 000000000..db0bb694f --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionChildOperationSummary.java @@ -0,0 +1,11 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import software.amazon.awssdk.services.lambda.model.ErrorObject; +import software.amazon.awssdk.services.lambda.model.OperationStatus; +import software.amazon.awssdk.services.lambda.model.OperationType; + +/** Read-only summary of a direct child operation involved in an extension CONTEXT failure. */ +public record ExtensionChildOperationSummary( + OperationType operationType, String subType, OperationStatus status, ErrorObject error) {} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextErrorHandler.java b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextErrorHandler.java new file mode 100644 index 000000000..a38fd500b --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextErrorHandler.java @@ -0,0 +1,10 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +/** Translates an extension CONTEXT failure when its original exception cannot be reconstructed. */ +@FunctionalInterface +public interface ExtensionContextErrorHandler { + /** Returns the exception exposed by the durable future. */ + Throwable translate(ExtensionContextFailure failure); +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextFailure.java b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextFailure.java new file mode 100644 index 000000000..b6fbb9126 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextFailure.java @@ -0,0 +1,48 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import java.util.List; +import software.amazon.awssdk.services.lambda.model.ErrorObject; + +/** Read-only failure information supplied to an extension CONTEXT error handler. */ +public final class ExtensionContextFailure { + private final String contextName; + private final String subType; + private final Throwable originalException; + private final ErrorObject error; + private final List childOperations; + + public ExtensionContextFailure( + String contextName, + String subType, + Throwable originalException, + ErrorObject error, + List childOperations) { + this.contextName = contextName; + this.subType = subType; + this.originalException = originalException; + this.error = error; + this.childOperations = List.copyOf(childOperations); + } + + public String contextName() { + return contextName; + } + + public String subType() { + return subType; + } + + public Throwable originalException() { + return originalException; + } + + public ErrorObject error() { + return error; + } + + public List childOperations() { + return childOperations; + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextFunction.java b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextFunction.java new file mode 100644 index 000000000..0ef626088 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextFunction.java @@ -0,0 +1,10 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +/** Framework callback for an advanced extension CONTEXT primitive. */ +@FunctionalInterface +public interface ExtensionContextFunction { + /** Executes the extension framework logic and returns its application result policy. */ + ExtensionContextResult apply(); +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextReplayContext.java b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextReplayContext.java new file mode 100644 index 000000000..f5fd1ae1c --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextReplayContext.java @@ -0,0 +1,40 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import software.amazon.lambda.durable.model.SafeCloseable; + +/** Replay metadata available while an advanced extension CONTEXT framework callback is running. */ +public final class ExtensionContextReplayContext { + private static final OperationContextStorage> CURRENT = + new OperationContextStorage<>("ExtensionContextReplayContext"); + + private final boolean replayingChildren; + private final T replayState; + + private ExtensionContextReplayContext(boolean replayingChildren, T replayState) { + this.replayingChildren = replayingChildren; + this.replayState = replayState; + } + + /** Returns the replay context attached to the current extension framework thread. */ + @SuppressWarnings("unchecked") + public static ExtensionContextReplayContext getCurrentContext() { + return (ExtensionContextReplayContext) CURRENT.getCurrentContext(); + } + + /** Returns whether a completed CONTEXT is replaying its child operations. */ + public boolean isReplayingChildren() { + return replayingChildren; + } + + /** Returns the checkpointed replay state, or {@code null} on initial execution. */ + public T getReplayState() { + return replayState; + } + + /** Attaches replay metadata for the duration of an SDK-managed framework callback. */ + public static SafeCloseable attach(boolean replayingChildren, T replayState) { + return CURRENT.attach(new ExtensionContextReplayContext<>(replayingChildren, replayState)); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextResult.java b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextResult.java new file mode 100644 index 000000000..bc86e0a5b --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextResult.java @@ -0,0 +1,58 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +/** Result and replay policy returned by an advanced extension CONTEXT primitive. */ +public final class ExtensionContextResult { + private enum ReplayPolicy { + NONE, + ALWAYS, + ABOVE_SIZE + } + + private final T result; + private final T replayState; + private final ReplayPolicy replayPolicy; + private final int thresholdBytes; + + private ExtensionContextResult(T result, T replayState, ReplayPolicy replayPolicy, int thresholdBytes) { + this.result = result; + this.replayState = replayState; + this.replayPolicy = replayPolicy; + this.thresholdBytes = thresholdBytes; + } + + /** Returns a normal context result that does not replay children. */ + public static ExtensionContextResult completed(T result) { + return new ExtensionContextResult<>(result, null, ReplayPolicy.NONE, 0); + } + + /** Returns a result that always replays children using the supplied replay state. */ + public static ExtensionContextResult replayChildren(T result, T replayState) { + return new ExtensionContextResult<>(result, replayState, ReplayPolicy.ALWAYS, 0); + } + + /** Returns a result that replays children when the serialized full result reaches the threshold. */ + public static ExtensionContextResult replayChildrenAboveSize(T result, T replayState, int thresholdBytes) { + if (thresholdBytes <= 0) { + throw new IllegalArgumentException("thresholdBytes must be greater than zero"); + } + return new ExtensionContextResult<>(result, replayState, ReplayPolicy.ABOVE_SIZE, thresholdBytes); + } + + /** Returns the application result. */ + public T result() { + return result; + } + + /** Returns the replay-only state. */ + public T replayState() { + return replayState; + } + + /** Returns whether children should replay for a serialized full result of the given size. */ + public boolean shouldReplayChildren(int serializedResultBytes) { + return replayPolicy == ReplayPolicy.ALWAYS + || replayPolicy == ReplayPolicy.ABOVE_SIZE && serializedResultBytes >= thresholdBytes; + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java index 329acfc58..b4901565e 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java @@ -5,6 +5,7 @@ import java.time.Duration; import java.util.function.Supplier; import software.amazon.lambda.durable.config.CallbackConfig; +import software.amazon.lambda.durable.config.ExtensionContextConfig; import software.amazon.lambda.durable.config.ExtensionStepConfig; import software.amazon.lambda.durable.config.InvokeConfig; import software.amazon.lambda.durable.config.RunInChildContextConfig; @@ -281,4 +282,28 @@ default DurableFuture runInChildContextAsync( DurableFuture runInChildContextAsync( String subType, TypeToken resultType, Supplier function, RunInChildContextConfig config); + + default T runInChildContext( + String subType, Class resultType, ExtensionContextFunction function, ExtensionContextConfig config) { + return runInChildContext(subType, TypeToken.get(resultType), function, config); + } + + default T runInChildContext( + String subType, + TypeToken resultType, + ExtensionContextFunction function, + ExtensionContextConfig config) { + return runInChildContextAsync(subType, resultType, function, config).get(); + } + + default DurableFuture runInChildContextAsync( + String subType, Class resultType, ExtensionContextFunction function, ExtensionContextConfig config) { + return runInChildContextAsync(subType, TypeToken.get(resultType), function, config); + } + + DurableFuture runInChildContextAsync( + String subType, + TypeToken resultType, + ExtensionContextFunction function, + ExtensionContextConfig config); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/ExtensionContextConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/ExtensionContextConfig.java new file mode 100644 index 000000000..177e5fdb9 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/ExtensionContextConfig.java @@ -0,0 +1,84 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.config; + +import java.util.Objects; +import software.amazon.lambda.durable.ExtensionContextErrorHandler; + +/** Extension-only policies for an advanced CONTEXT primitive. */ +public final class ExtensionContextConfig { + private final RunInChildContextConfig childContextConfig; + private final ExtensionContextErrorHandler errorHandler; + private final boolean emitUserFunctionEvents; + private final boolean suppressLateChildCheckpoints; + + private ExtensionContextConfig(Builder builder) { + childContextConfig = + Objects.requireNonNullElseGet(builder.childContextConfig, () -> RunInChildContextConfig.builder() + .build()); + errorHandler = builder.errorHandler; + emitUserFunctionEvents = builder.emitUserFunctionEvents; + suppressLateChildCheckpoints = builder.suppressLateChildCheckpoints; + } + + public RunInChildContextConfig childContextConfig() { + return childContextConfig; + } + + public ExtensionContextErrorHandler errorHandler() { + return errorHandler; + } + + public boolean emitUserFunctionEvents() { + return emitUserFunctionEvents; + } + + public boolean suppressLateChildCheckpoints() { + return suppressLateChildCheckpoints; + } + + public Builder toBuilder() { + return new Builder() + .childContextConfig(childContextConfig) + .errorHandler(errorHandler) + .emitUserFunctionEvents(emitUserFunctionEvents) + .suppressLateChildCheckpoints(suppressLateChildCheckpoints); + } + + public static Builder builder() { + return new Builder(); + } + + public static final class Builder { + private RunInChildContextConfig childContextConfig; + private ExtensionContextErrorHandler errorHandler; + private boolean emitUserFunctionEvents = true; + private boolean suppressLateChildCheckpoints; + + private Builder() {} + + public Builder childContextConfig(RunInChildContextConfig childContextConfig) { + this.childContextConfig = Objects.requireNonNull(childContextConfig, "childContextConfig cannot be null"); + return this; + } + + public Builder errorHandler(ExtensionContextErrorHandler errorHandler) { + this.errorHandler = errorHandler; + return this; + } + + public Builder emitUserFunctionEvents(boolean emitUserFunctionEvents) { + this.emitUserFunctionEvents = emitUserFunctionEvents; + return this; + } + + public Builder suppressLateChildCheckpoints(boolean suppressLateChildCheckpoints) { + this.suppressLateChildCheckpoints = suppressLateChildCheckpoints; + return this; + } + + public ExtensionContextConfig build() { + return new ExtensionContextConfig(this); + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java index 7c1138d14..157b547a9 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java @@ -16,12 +16,14 @@ import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.ExtensionContext; +import software.amazon.lambda.durable.ExtensionContextFunction; import software.amazon.lambda.durable.ExtensionOperation; import software.amazon.lambda.durable.ExtensionStepFunction; import software.amazon.lambda.durable.ParallelDurableFuture; import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.CallbackConfig; +import software.amazon.lambda.durable.config.ExtensionContextConfig; import software.amazon.lambda.durable.config.ExtensionStepConfig; import software.amazon.lambda.durable.config.InvokeConfig; import software.amazon.lambda.durable.config.MapConfig; @@ -41,6 +43,7 @@ import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.model.WaitForConditionResult; +import software.amazon.lambda.durable.operation.BaseDurableOperation; import software.amazon.lambda.durable.operation.CallbackOperation; import software.amazon.lambda.durable.operation.ChildContextOperation; import software.amazon.lambda.durable.operation.InvokeOperation; @@ -65,6 +68,7 @@ public class DurableContextImpl extends BaseContextImpl implements DurableContex - Math.max(WAIT_FOR_CALLBACK_CALLBACK_SUFFIX.length(), WAIT_FOR_CALLBACK_SUBMITTER_SUFFIX.length()); private final OperationIdGenerator operationIdGenerator; private final DurableContextImpl parentContext; + private final BaseDurableOperation lateCheckpointOwner; private final boolean isVirtual; private boolean isReplaying; @@ -76,10 +80,12 @@ private DurableContextImpl( String contextId, String contextName, boolean isVirtual, - DurableContextImpl parentContext) { + DurableContextImpl parentContext, + BaseDurableOperation lateCheckpointOwner) { super(executionManager, durableConfig, lambdaContext, contextId, contextName, ThreadType.CONTEXT); operationIdGenerator = new OperationIdGenerator(contextId); this.parentContext = parentContext; + this.lateCheckpointOwner = lateCheckpointOwner; this.isVirtual = isVirtual; this.isReplaying = executionManager.hasOperationsForContext(contextId); } @@ -96,7 +102,7 @@ private DurableContextImpl( */ public static DurableContextImpl createRootContext( ExecutionManager executionManager, DurableConfig durableConfig, Context lambdaContext) { - return new DurableContextImpl(executionManager, durableConfig, lambdaContext, null, null, false, null); + return new DurableContextImpl(executionManager, durableConfig, lambdaContext, null, null, false, null, null); } /** @@ -108,6 +114,14 @@ public static DurableContextImpl createRootContext( * @return a new DurableContext for the child context */ public DurableContextImpl createChildContext(String childContextId, String childContextName, boolean isVirtual) { + return createChildContext(childContextId, childContextName, isVirtual, null); + } + + public DurableContextImpl createChildContext( + String childContextId, + String childContextName, + boolean isVirtual, + BaseDurableOperation lateCheckpointOwner) { return new DurableContextImpl( getExecutionManager(), getDurableConfig(), @@ -115,7 +129,8 @@ public DurableContextImpl createChildContext(String childContextId, String child childContextId, childContextName, isVirtual, - this); + this, + lateCheckpointOwner); } /** @@ -391,8 +406,40 @@ DurableFuture runInChildContextAsyncWithId( func, resultType, config, - this); + this, + lateCheckpointOwner); + + operation.execute(); + return operation; + } + + DurableFuture extensionContextAsyncWithId( + String operationId, + String name, + String subType, + TypeToken resultType, + ExtensionContextFunction function, + ExtensionContextConfig config) { + Objects.requireNonNull(resultType, "resultType cannot be null"); + Objects.requireNonNull(function, "function cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + ParameterValidator.validateOperationName(name); + var childConfig = config.childContextConfig(); + if (childConfig.serDes() == null) { + childConfig = childConfig.toBuilder() + .serDes(getDurableConfig().getSerDes()) + .build(); + config = config.toBuilder().childContextConfig(childConfig).build(); + } + + var operation = new ChildContextOperation<>( + new OperationDescriptor(operationId, name, OperationType.CONTEXT, subType), + function, + resultType, + config, + this, + lateCheckpointOwner); operation.execute(); return operation; } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java index 7fe416176..c84611ed9 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java @@ -8,10 +8,12 @@ import java.util.function.Supplier; import software.amazon.lambda.durable.DurableCallbackFuture; import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.ExtensionContextFunction; import software.amazon.lambda.durable.ExtensionOperation; import software.amazon.lambda.durable.ExtensionStepFunction; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.CallbackConfig; +import software.amazon.lambda.durable.config.ExtensionContextConfig; import software.amazon.lambda.durable.config.ExtensionStepConfig; import software.amazon.lambda.durable.config.InvokeConfig; import software.amazon.lambda.durable.config.RunInChildContextConfig; @@ -110,6 +112,20 @@ public DurableFuture runInChildContextAsync( operationId, name, subType, resultType, ignored -> function.get(), config); } + @Override + public DurableFuture runInChildContextAsync( + String subType, + TypeToken resultType, + ExtensionContextFunction function, + ExtensionContextConfig config) { + validateSubType(subType); + Objects.requireNonNull(resultType, "resultType cannot be null"); + Objects.requireNonNull(function, "function cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + claim(); + return context.extensionContextAsyncWithId(operationId, name, subType, resultType, function, config); + } + private void validateSubType(String subType) { Objects.requireNonNull(subType, "subType cannot be null"); if (subType.isBlank()) { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java index 9fde070c1..38c984d3b 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java @@ -5,6 +5,8 @@ import static software.amazon.lambda.durable.execution.ExecutionManager.isTerminalStatus; import java.nio.charset.StandardCharsets; +import java.util.List; +import java.util.Objects; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; @@ -17,7 +19,13 @@ import software.amazon.awssdk.services.lambda.model.OperationType; import software.amazon.awssdk.services.lambda.model.OperationUpdate; import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.ExtensionChildOperationSummary; +import software.amazon.lambda.durable.ExtensionContextFailure; +import software.amazon.lambda.durable.ExtensionContextFunction; +import software.amazon.lambda.durable.ExtensionContextReplayContext; +import software.amazon.lambda.durable.ExtensionContextResult; import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.ExtensionContextConfig; import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.exception.CallbackFailedException; @@ -54,7 +62,10 @@ public class ChildContextOperation extends SerializableDurableOperation { private static final int LARGE_RESULT_THRESHOLD = 256 * 1024; private final Function function; + private final ExtensionContextFunction extensionFunction; + private final ExtensionContextConfig extensionConfig; private final AtomicBoolean replayChildren = new AtomicBoolean(false); + private final AtomicReference replayState = new AtomicReference<>(null); private final AtomicReference> cachedOperationResult = new AtomicReference<>(null); // child context for RunInChildContext @@ -83,6 +94,8 @@ public ChildContextOperation( parentOperation, config.isVirtual()); this.function = function; + this.extensionFunction = null; + this.extensionConfig = null; } public ChildContextOperation( @@ -91,8 +104,54 @@ public ChildContextOperation( TypeToken resultTypeToken, RunInChildContextConfig config, DurableContextImpl durableContext) { - super(operationDescriptor, resultTypeToken, config.serDes(), durableContext, null, config.isVirtual()); + this(operationDescriptor, function, resultTypeToken, config, durableContext, null); + } + + public ChildContextOperation( + OperationDescriptor operationDescriptor, + Function function, + TypeToken resultTypeToken, + RunInChildContextConfig config, + DurableContextImpl durableContext, + BaseDurableOperation parentOperation) { + super( + operationDescriptor, + resultTypeToken, + config.serDes(), + durableContext, + parentOperation, + config.isVirtual()); this.function = function; + this.extensionFunction = null; + this.extensionConfig = null; + } + + public ChildContextOperation( + OperationDescriptor operationDescriptor, + ExtensionContextFunction function, + TypeToken resultTypeToken, + ExtensionContextConfig config, + DurableContextImpl durableContext) { + this(operationDescriptor, function, resultTypeToken, config, durableContext, null); + } + + public ChildContextOperation( + OperationDescriptor operationDescriptor, + ExtensionContextFunction function, + TypeToken resultTypeToken, + ExtensionContextConfig config, + DurableContextImpl durableContext, + BaseDurableOperation parentOperation) { + super( + operationDescriptor, + resultTypeToken, + config.childContextConfig().serDes(), + durableContext, + parentOperation, + config.childContextConfig().isVirtual()); + this.function = null; + this.extensionFunction = function; + this.extensionConfig = config; } /** Starts the operation. */ @@ -112,8 +171,11 @@ protected void replay(Operation existing) { case SUCCEEDED -> { if (existing.contextDetails() != null && Boolean.TRUE.equals(existing.contextDetails().replayChildren())) { - // Large result: re-execute child context to reconstruct result replayChildren.set(true); + if (extensionFunction != null) { + replayState.set( + deserializeResult(existing.contextDetails().result())); + } executeChildContext(); } else { markAlreadyCompleted(); @@ -143,15 +205,11 @@ private void executeChildContext() { // When this child is part of a ConcurrencyOperation (parentOperation != null), // we notify the parent BEFORE closing the child context. This ensures the parent // can trigger the next queued branch while the current child context is still valid. - var childContext = getContext().createChildContext(contextId, getName(), isVirtual); + var childContext = createChildContext(contextId); try (var ignoredContext = DurableContextImpl.attachCurrentContext(childContext); var ignoredLogger = DurableLogger.attachContext()) { try { - // Run the user function inside the plugin hook boundary (attempt is null for contexts) - // so a failure is reported through onUserFunctionEnd; checkpointing stays outside. - T result = runUserFunction(null, () -> function.apply(childContext)); - - handleChildContextSuccess(result); + executeFunction(childContext); } catch (Throwable e) { handleChildContextFailure(e); } @@ -162,28 +220,75 @@ private void executeChildContext() { runUserHandler(userHandler, ThreadType.CONTEXT); } + private DurableContextImpl createChildContext(String contextId) { + if (extensionConfig != null && extensionConfig.suppressLateChildCheckpoints()) { + return getContext().createChildContext(contextId, getName(), isVirtual, this); + } + return getContext().createChildContext(contextId, getName(), isVirtual); + } + + private void executeFunction(DurableContextImpl childContext) { + if (extensionFunction == null) { + var result = runUserFunction(null, () -> function.apply(childContext)); + handleChildContextSuccess(result); + return; + } + + try (var ignoredReplayContext = ExtensionContextReplayContext.attach(replayChildren.get(), replayState.get())) { + var result = extensionConfig.emitUserFunctionEvents() + ? runUserFunction(null, extensionFunction::apply) + : extensionFunction.apply(); + handleExtensionContextSuccess( + Objects.requireNonNull(result, "Extension context function result cannot be null")); + } + } + private void handleChildContextSuccess(T result) { var serializedResult = serializeAndDeserializeResult(result); - if (replayChildren.get() || isVirtual || parentOperation != null && parentOperation.isOperationCompleted()) { - // Skip checkpointing if - // - parent ConcurrencyOperation has already completed, preventing race conditions where a child finishes - // after the parent has already completed. - // - replaying a SUCCEEDED child with replayChildren=true — skip checkpointing. - // - nestingType is FLAT - // Mark the completableFuture completed so get() doesn't block waiting for a checkpoint response. - cachedOperationResult.set(DeserializedOperationResult.succeeded(serializedResult.deserialized())); - if (isVirtual) { - fireOnOperationEnd(null, null, false); - } - markAlreadyCompleted(); + if (shouldSkipCheckpoint()) { + cacheSuccessAndComplete(serializedResult.deserialized()); } else { checkpointSuccess(serializedResult.deserialized(), serializedResult.serialized()); } } + private void handleExtensionContextSuccess(ExtensionContextResult result) { + var serializedResult = serializeAndDeserializeResult(result.result()); + if (shouldSkipCheckpoint()) { + cacheSuccessAndComplete(serializedResult.deserialized()); + return; + } + + var resultBytes = serializedSize(serializedResult.serialized()); + if (result.shouldReplayChildren(resultBytes)) { + var serializedReplayState = serializeAndDeserializeResult(result.replayState()); + cachedOperationResult.set(DeserializedOperationResult.succeeded(serializedResult.deserialized())); + sendOperationUpdate(OperationUpdate.builder() + .action(OperationAction.SUCCEED) + .payload(serializedReplayState.serialized()) + .contextOptions( + ContextOptions.builder().replayChildren(true).build())); + } else { + sendOperationUpdate( + OperationUpdate.builder().action(OperationAction.SUCCEED).payload(serializedResult.serialized())); + } + } + + private boolean shouldSkipCheckpoint() { + return replayChildren.get() || isVirtual || parentOperation != null && parentOperation.isOperationCompleted(); + } + + private void cacheSuccessAndComplete(T result) { + cachedOperationResult.set(DeserializedOperationResult.succeeded(result)); + if (isVirtual) { + fireOnOperationEnd(null, null, false); + } + markAlreadyCompleted(); + } + private void checkpointSuccess(T result, String serialized) { - if (serialized == null || serialized.getBytes(StandardCharsets.UTF_8).length < LARGE_RESULT_THRESHOLD) { + if (serializedSize(serialized) < LARGE_RESULT_THRESHOLD) { sendOperationUpdate( OperationUpdate.builder().action(OperationAction.SUCCEED).payload(serialized)); } else { @@ -198,6 +303,10 @@ private void checkpointSuccess(T result, String serialized) { } } + private int serializedSize(String serialized) { + return serialized == null ? 0 : serialized.getBytes(StandardCharsets.UTF_8).length; + } + private void handleChildContextFailure(Throwable exception) { exception = ExceptionHelper.unwrapCompletableFuture(exception); if (exception instanceof SuspendExecutionException suspendExecutionException) { @@ -263,6 +372,14 @@ private Throwable translateException(Operation op, ErrorObject errorObject) { return original; } + if (extensionConfig != null && extensionConfig.errorHandler() != null) { + var failure = new ExtensionContextFailure( + getName(), getSubTypeValue(), null, errorObject, getChildOperationSummaries()); + return Objects.requireNonNull( + extensionConfig.errorHandler().translate(failure), + "Extension context error handler result cannot be null"); + } + // throw a general failed exception if a user exception is not reconstructed if (OperationSubType.WAIT_FOR_CALLBACK.getValue().equals(getSubTypeValue())) { return handleWaitForCallbackFailure(); @@ -276,6 +393,16 @@ private Throwable translateException(Operation op, ErrorObject errorObject) { return new ChildContextFailedException(op); } + private List getChildOperationSummaries() { + return getChildOperations().stream() + .map(operation -> new ExtensionChildOperationSummary( + operation.type(), + operation.subType(), + operation.status(), + BaseDurableOperation.getErrorObject(operation))) + .toList(); + } + private Operation createVirtualOperation(ErrorObject errorObject) { return Operation.builder() .id(getOperationId()) diff --git a/sdk/src/test/java/software/amazon/lambda/durable/ExtensionContextReplayContextTest.java b/sdk/src/test/java/software/amazon/lambda/durable/ExtensionContextReplayContextTest.java new file mode 100644 index 000000000..ad2dbb7fd --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/ExtensionContextReplayContextTest.java @@ -0,0 +1,36 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import org.junit.jupiter.api.Test; + +class ExtensionContextReplayContextTest { + @Test + void lookupFailsOutsideExtensionContextFunction() { + assertThrows(IllegalStateException.class, ExtensionContextReplayContext::getCurrentContext); + } + + @Test + void nestedScopesRestorePreviousReplayState() { + try (var outer = ExtensionContextReplayContext.attach(false, "outer")) { + var outerContext = ExtensionContextReplayContext.getCurrentContext(); + assertFalse(outerContext.isReplayingChildren()); + assertEquals("outer", outerContext.getReplayState()); + + try (var inner = ExtensionContextReplayContext.attach(true, "inner")) { + var innerContext = ExtensionContextReplayContext.getCurrentContext(); + assertTrue(innerContext.isReplayingChildren()); + assertEquals("inner", innerContext.getReplayState()); + } + + assertEquals( + "outer", + ExtensionContextReplayContext.getCurrentContext().getReplayState()); + } + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/ExtensionContextResultTest.java b/sdk/src/test/java/software/amazon/lambda/durable/ExtensionContextResultTest.java new file mode 100644 index 000000000..54cd4da5a --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/ExtensionContextResultTest.java @@ -0,0 +1,44 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import org.junit.jupiter.api.Test; + +class ExtensionContextResultTest { + @Test + void completedStoresOnlyApplicationResult() { + var result = ExtensionContextResult.completed("full"); + + assertEquals("full", result.result()); + assertFalse(result.shouldReplayChildren(1024)); + } + + @Test + void replayChildrenStoresReplayState() { + var result = ExtensionContextResult.replayChildren("full", "state"); + + assertEquals("full", result.result()); + assertEquals("state", result.replayState()); + assertTrue(result.shouldReplayChildren(1)); + } + + @Test + void replayChildrenAboveSizeUsesSerializedFullResultSize() { + var result = ExtensionContextResult.replayChildrenAboveSize("full", "state", 5); + + assertFalse(result.shouldReplayChildren(4)); + assertTrue(result.shouldReplayChildren(5)); + } + + @Test + void replayChildrenAboveSizeRejectsInvalidThreshold() { + assertThrows( + IllegalArgumentException.class, + () -> ExtensionContextResult.replayChildrenAboveSize("full", "state", 0)); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionContextConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionContextConfigTest.java new file mode 100644 index 000000000..1f2a87dd6 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionContextConfigTest.java @@ -0,0 +1,41 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.config; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.ExtensionContextErrorHandler; + +class ExtensionContextConfigTest { + @Test + void builderUsesOrdinaryChildContextDefaults() { + var config = ExtensionContextConfig.builder().build(); + + assertNotNull(config.childContextConfig()); + assertNull(config.errorHandler()); + assertTrue(config.emitUserFunctionEvents()); + assertFalse(config.suppressLateChildCheckpoints()); + } + + @Test + void builderRetainsExtensionPolicies() { + var childConfig = RunInChildContextConfig.builder().isVirtual(true).build(); + ExtensionContextErrorHandler handler = failure -> new RuntimeException(failure.contextName()); + var config = ExtensionContextConfig.builder() + .childContextConfig(childConfig) + .errorHandler(handler) + .emitUserFunctionEvents(false) + .suppressLateChildCheckpoints(true) + .build(); + + assertEquals(childConfig, config.childContextConfig()); + assertEquals(handler, config.errorHandler()); + assertFalse(config.emitUserFunctionEvents()); + assertTrue(config.suppressLateChildCheckpoints()); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java index ddb8867e2..3fa38cf1b 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java @@ -20,11 +20,13 @@ import software.amazon.lambda.durable.DurableCallbackFuture; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.ExtensionContextResult; import software.amazon.lambda.durable.ExtensionOperation; import software.amazon.lambda.durable.ExtensionStepResult; import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.CallbackConfig; +import software.amazon.lambda.durable.config.ExtensionContextConfig; import software.amazon.lambda.durable.config.ExtensionStepConfig; import software.amazon.lambda.durable.config.InvokeConfig; import software.amazon.lambda.durable.config.RunInChildContextConfig; @@ -239,6 +241,23 @@ void reservedChildContextAdaptsSupplierToChildFunction() { assertEquals("result", function.getValue().apply(mock(DurableContext.class))); } + @Test + void advancedChildContextDelegatesFrameworkFunction() { + var context = mock(DurableContextImpl.class); + var future = mockStringFuture(); + var resultType = TypeToken.get(String.class); + var config = ExtensionContextConfig.builder().build(); + when(context.extensionContextAsyncWithId( + eq("1"), eq("child"), eq("AcmeContext"), eq(resultType), any(), eq(config))) + .thenReturn(future); + + var actual = new ExtensionOperationImpl(context, "1", "child") + .runInChildContextAsync( + "AcmeContext", resultType, () -> ExtensionContextResult.completed("result"), config); + + assertEquals(future, actual); + } + @Test void reservationCanOnlyExecuteOnceAcrossPrimitiveSelectors() { var context = mock(DurableContextImpl.class); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java index 99d994538..f91da9131 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java @@ -6,8 +6,10 @@ import static org.mockito.Mockito.*; import java.util.List; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executors; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -17,9 +19,13 @@ import software.amazon.awssdk.services.lambda.model.OperationAction; import software.amazon.awssdk.services.lambda.model.OperationStatus; import software.amazon.awssdk.services.lambda.model.OperationType; +import software.amazon.awssdk.services.lambda.model.StepDetails; import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.ExtensionContextFailure; +import software.amazon.lambda.durable.ExtensionContextResult; import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.ExtensionContextConfig; import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.exception.ChildContextFailedException; @@ -28,6 +34,7 @@ import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.execution.ThreadContext; import software.amazon.lambda.durable.execution.ThreadType; +import software.amazon.lambda.durable.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.serde.JacksonSerDes; @@ -126,6 +133,15 @@ private ChildContextOperation createOperationWithParent( parent); } + private ChildContextOperation createExtensionOperation(ExtensionContextConfig config) { + return new ChildContextOperation<>( + new OperationDescriptor("1", "test-context", OperationType.CONTEXT, "AcmeContext"), + () -> ExtensionContextResult.completed("unused"), + TypeToken.get(String.class), + config, + durableContext); + } + // ===== SUCCEEDED replay ===== /** SUCCEEDED replay returns cached result without re-executing the function. */ @@ -246,6 +262,116 @@ void replayFailedFallsBackToChildContextFailedException() { assertTrue(thrown.getMessage().contains("unknown error")); } + @Test + void replayFailedUsesExtensionErrorHandlerWithChildSummaries() { + var contextError = ErrorObject.builder() + .errorType("com.nonexistent.ContextException") + .errorMessage("context failed") + .build(); + var childError = ErrorObject.builder() + .errorType("com.nonexistent.ChildException") + .errorMessage("child failed") + .build(); + var failedContext = Operation.builder() + .id("1") + .name("test-context") + .type(OperationType.CONTEXT) + .subType("AcmeContext") + .status(OperationStatus.FAILED) + .contextDetails(ContextDetails.builder().error(contextError).build()) + .build(); + var failedChild = Operation.builder() + .id("1-1") + .name("child") + .type(OperationType.STEP) + .subType("AcmeChild") + .status(OperationStatus.FAILED) + .stepDetails(StepDetails.builder().error(childError).build()) + .build(); + when(executionManager.getOperationAndUpdateReplayState("1")).thenReturn(failedContext); + when(executionManager.getChildOperations("1")).thenReturn(List.of(failedChild)); + var capturedFailure = new AtomicReference(); + var translated = new IllegalStateException("translated"); + var config = ExtensionContextConfig.builder() + .childContextConfig( + RunInChildContextConfig.builder().serDes(SERDES).build()) + .errorHandler(failure -> { + capturedFailure.set(failure); + return translated; + }) + .build(); + + var operation = createExtensionOperation(config); + operation.execute(); + + assertSame(translated, assertThrows(IllegalStateException.class, operation::get)); + var failure = capturedFailure.get(); + assertEquals("test-context", failure.contextName()); + assertEquals("AcmeContext", failure.subType()); + assertEquals(contextError, failure.error()); + assertEquals(1, failure.childOperations().size()); + assertEquals(OperationType.STEP, failure.childOperations().get(0).operationType()); + assertEquals("AcmeChild", failure.childOperations().get(0).subType()); + assertEquals(childError, failure.childOperations().get(0).error()); + } + + @Test + void replayFailedPrefersReconstructedExceptionOverExtensionHandler() { + var originalException = new IllegalArgumentException("bad input"); + var failedContext = Operation.builder() + .id("1") + .name("test-context") + .type(OperationType.CONTEXT) + .subType("AcmeContext") + .status(OperationStatus.FAILED) + .contextDetails(ContextDetails.builder() + .error(ErrorObject.builder() + .errorType("java.lang.IllegalArgumentException") + .errorMessage("bad input") + .errorData(SERDES.serialize(originalException)) + .stackTrace(List.of("com.example.Test|method|Test.java|42")) + .build()) + .build()) + .build(); + when(executionManager.getOperationAndUpdateReplayState("1")).thenReturn(failedContext); + var handlerCalled = new AtomicBoolean(); + var config = ExtensionContextConfig.builder() + .childContextConfig( + RunInChildContextConfig.builder().serDes(SERDES).build()) + .errorHandler(failure -> { + handlerCalled.set(true); + return new IllegalStateException("translated"); + }) + .build(); + + var operation = createExtensionOperation(config); + operation.execute(); + + var thrown = assertThrows(IllegalArgumentException.class, operation::get); + assertEquals("bad input", thrown.getMessage()); + assertFalse(handlerCalled.get()); + } + + @Test + void suppressingExtensionContextPropagatesCompletionOwnerToChildContext() { + when(executionManager.getOperationAndUpdateReplayState("1")).thenReturn(null); + when(executionManager.sendOperationUpdate(any())).thenReturn(CompletableFuture.completedFuture(null)); + var childContext = mock(DurableContextImpl.class); + when(childContext.getDurableConfig()).thenReturn(createConfig()); + var config = ExtensionContextConfig.builder() + .childContextConfig( + RunInChildContextConfig.builder().serDes(SERDES).build()) + .suppressLateChildCheckpoints(true) + .build(); + var operation = createExtensionOperation(config); + when(durableContext.createChildContext("1", "test-context", false, operation)) + .thenReturn(childContext); + + operation.execute(); + + verify(durableContext, timeout(1000)).createChildContext("1", "test-context", false, operation); + } + // ===== Replay STARTED ===== /** STARTED replay re-executes the child context (interrupted mid-execution). */ From 2546635302365401d722cfa2bffde505310f83a2 Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 03:56:59 +0000 Subject: [PATCH 18/40] feat: organize extension APIs and migrate callbacks --- docs/adr/006-custom-extension-operations.md | 24 +++- docs/advanced/extensions.md | 15 ++- .../ExtensionOperationIntegrationTest.java | 8 +- .../lambda/durable/PluginIntegrationTest.java | 1 + .../StaticOperationsIntegrationTest.java | 1 + .../durable/extension/PairOperations.java | 1 - .../DurableWaitForCallbackOperations.java | 23 ++-- .../ExtensionChildOperationSummary.java | 11 -- .../durable/context/DurableContextImpl.java | 50 ++------ .../context/ExtensionOperationImpl.java | 10 +- .../extension/WaitForCallbackExtension.java | 117 +++++++++++++++++ .../ExtensionChildOperationSummary.java | 78 ++++++++++++ .../{ => extension}/ExtensionContext.java | 2 +- .../ExtensionContextConfig.java | 4 +- .../ExtensionContextErrorHandler.java | 2 +- .../ExtensionContextFailure.java | 2 +- .../ExtensionContextFunction.java | 2 +- .../ExtensionContextReplayContext.java | 21 ++- .../ExtensionContextResult.java | 2 +- .../{ => extension}/ExtensionOperation.java | 7 +- .../ExtensionStepConfig.java | 2 +- .../ExtensionStepFunction.java | 2 +- .../{ => extension}/ExtensionStepResult.java | 2 +- .../operation/ChildContextOperation.java | 18 +-- .../durable/operation/StepOperation.java | 6 +- .../lambda/durable/CurrentContextTest.java | 1 + .../DurableWaitForCallbackOperationsTest.java | 77 +++++++++-- .../ExtensionContextReplayContextTest.java | 1 + .../durable/ExtensionContextResultTest.java | 1 + .../durable/ExtensionStepResultTest.java | 1 + .../config/ExtensionContextConfigTest.java | 3 +- .../config/ExtensionStepConfigTest.java | 1 + .../context/ExtensionOperationImplTest.java | 10 +- .../WaitForCallbackExtensionTest.java | 120 ++++++++++++++++++ .../operation/ChildContextOperationTest.java | 6 +- 35 files changed, 499 insertions(+), 133 deletions(-) delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/ExtensionChildOperationSummary.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForCallbackExtension.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionChildOperationSummary.java rename sdk/src/main/java/software/amazon/lambda/durable/{ => extension}/ExtensionContext.java (97%) rename sdk/src/main/java/software/amazon/lambda/durable/{config => extension}/ExtensionContextConfig.java (96%) rename sdk/src/main/java/software/amazon/lambda/durable/{ => extension}/ExtensionContextErrorHandler.java (88%) rename sdk/src/main/java/software/amazon/lambda/durable/{ => extension}/ExtensionContextFailure.java (96%) rename sdk/src/main/java/software/amazon/lambda/durable/{ => extension}/ExtensionContextFunction.java (88%) rename sdk/src/main/java/software/amazon/lambda/durable/{ => extension}/ExtensionContextReplayContext.java (64%) rename sdk/src/main/java/software/amazon/lambda/durable/{ => extension}/ExtensionContextResult.java (97%) rename sdk/src/main/java/software/amazon/lambda/durable/{ => extension}/ExtensionOperation.java (98%) rename sdk/src/main/java/software/amazon/lambda/durable/{config => extension}/ExtensionStepConfig.java (97%) rename sdk/src/main/java/software/amazon/lambda/durable/{ => extension}/ExtensionStepFunction.java (91%) rename sdk/src/main/java/software/amazon/lambda/durable/{ => extension}/ExtensionStepResult.java (95%) create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/context/extension/WaitForCallbackExtensionTest.java diff --git a/docs/adr/006-custom-extension-operations.md b/docs/adr/006-custom-extension-operations.md index 772e87b32..9ec6cf734 100644 --- a/docs/adr/006-custom-extension-operations.md +++ b/docs/adr/006-custom-extension-operations.md @@ -26,11 +26,6 @@ The backend accepts arbitrary operation subtype strings, but each backend operat Extensions need subtype control, replay state, and failure translation without receiving raw checkpoint actions or defining new backend state machines. -The detailed designs are: - -- [Custom Extension Operations Design](../superpowers/specs/2026-08-10-custom-extension-operations-design.md) -- [Migrate Built-In Operations to the Extension API](../superpowers/specs/2026-08-10-migrate-built-in-extensions-design.md) - ## Decision ### Preserve Existing Operation APIs @@ -79,6 +74,25 @@ Expose each built-in extension family through an independently maintained class: An extension is an ordinary static Java method. There is no registration API and no universal `DurableExtensions.run` boundary. +### Separate Extension-Author Contracts + +Place contracts intended specifically for extension authors in: + +```text +software.amazon.lambda.durable.extension +``` + +This package contains `ExtensionContext`, `ExtensionOperation`, stateful-step contracts, and configurable +extension-context contracts. + +Keep the following customer-facing types in the root `software.amazon.lambda.durable` package: + +- static operation facades such as `DurableCoreOperations` and `DurableMapOperations` +- operation-specific TLS metadata such as `MapItemContext`, `WaitForCallbackContext`, and `WithRetryContext` +- established SDK types such as `DurableFuture`, `StepContext`, and `TypeToken` + +SDK implementations of built-in extensions remain internal and are not part of the extension-author API. + ### Use Scoped Current Context SDK-managed handler and child contexts implement `ExtensionContext`. Step contexts do not. diff --git a/docs/advanced/extensions.md b/docs/advanced/extensions.md index 81a5dce42..8f8c9d71e 100644 --- a/docs/advanced/extensions.md +++ b/docs/advanced/extensions.md @@ -4,6 +4,9 @@ Extension operations are ordinary static Java methods that compose SDK-owned dur separate Maven module without defining backend operation types, sending checkpoint updates, or depending on SDK implementation packages. +Extension-author contracts are in `software.amazon.lambda.durable.extension`. Static operation facades and +operation-specific TLS metadata contexts remain in `software.amazon.lambda.durable`. + Application code calls only the extension's API: ```java @@ -15,6 +18,9 @@ var result = pairAsync("pair", left, right).get(); The extension retrieves the active scope internally: ```java +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.extension.ExtensionContext; + public final class PairOperations { private PairOperations() {} @@ -176,8 +182,9 @@ primitive operations. ## Module compatibility An extension Maven module should depend only on the public SDK artifact and import public types under -`software.amazon.lambda.durable`. Do not import SDK implementation packages such as `context`, `execution`, or -`operation`. +`software.amazon.lambda.durable` and `software.amazon.lambda.durable.extension`. Do not import SDK implementation +packages such as `context`, `execution`, or `operation`. -The supported extension contracts are `ExtensionContext`, `ExtensionOperation`, the static operation facades, the -typed TLS contexts, and `DurableFuture.completionFuture()`. +The extension-author SPI includes `ExtensionContext`, `ExtensionOperation`, stateful-step contracts, and configurable +extension-context contracts under `software.amazon.lambda.durable.extension`. Static operation facades, typed TLS +contexts, and `DurableFuture.completionFuture()` remain under `software.amazon.lambda.durable`. diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java index 0b87f0dcf..5c69d2ee1 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java @@ -19,8 +19,12 @@ import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.Test; import software.amazon.awssdk.services.lambda.model.OperationType; -import software.amazon.lambda.durable.config.ExtensionContextConfig; -import software.amazon.lambda.durable.config.ExtensionStepConfig; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; +import software.amazon.lambda.durable.extension.ExtensionContextResult; +import software.amazon.lambda.durable.extension.ExtensionStepConfig; +import software.amazon.lambda.durable.extension.ExtensionStepResult; import software.amazon.lambda.durable.model.ExecutionStatus; import software.amazon.lambda.durable.testing.LocalDurableTestRunner; diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java index e89334278..45e99bfba 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java @@ -15,6 +15,7 @@ import software.amazon.lambda.durable.config.ParallelConfig; import software.amazon.lambda.durable.config.StepConfig; import software.amazon.lambda.durable.execution.SuspendExecutionException; +import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.model.ExecutionStatus; import software.amazon.lambda.durable.model.WaitForConditionResult; import software.amazon.lambda.durable.plugin.*; diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java index 0bd044632..aee903645 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java @@ -13,6 +13,7 @@ import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.Test; import software.amazon.lambda.durable.config.WaitForConditionConfig; +import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.model.ExecutionStatus; import software.amazon.lambda.durable.model.WaitForConditionResult; import software.amazon.lambda.durable.testing.LocalDurableTestRunner; diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/extension/PairOperations.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/extension/PairOperations.java index d7a49ed5f..770081e21 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/extension/PairOperations.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/extension/PairOperations.java @@ -6,7 +6,6 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.atomic.AtomicInteger; import software.amazon.lambda.durable.DurableFuture; -import software.amazon.lambda.durable.ExtensionContext; /** Example extension library implemented only with public SDK contracts. */ public final class PairOperations { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForCallbackOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForCallbackOperations.java index f63c433e6..f9cedb99e 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForCallbackOperations.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForCallbackOperations.java @@ -5,45 +5,48 @@ import java.util.Objects; import java.util.function.BiConsumer; import software.amazon.lambda.durable.config.WaitForCallbackConfig; +import software.amazon.lambda.durable.context.extension.WaitForCallbackExtension; +import software.amazon.lambda.durable.extension.ExtensionContext; /** Context-free static facades for durable wait-for-callback operations. */ public final class DurableWaitForCallbackOperations { private DurableWaitForCallbackOperations() {} public static T waitForCallback(String name, Class resultType, Runnable submitter) { - return currentContext().waitForCallback(name, resultType, adapt(submitter)); + return waitForCallbackAsync(name, resultType, submitter).get(); } public static T waitForCallback(String name, TypeToken resultType, Runnable submitter) { - return currentContext().waitForCallback(name, resultType, adapt(submitter)); + return waitForCallbackAsync(name, resultType, submitter).get(); } public static T waitForCallback( String name, Class resultType, Runnable submitter, WaitForCallbackConfig config) { - return currentContext().waitForCallback(name, resultType, adapt(submitter), config); + return waitForCallbackAsync(name, resultType, submitter, config).get(); } public static T waitForCallback( String name, TypeToken resultType, Runnable submitter, WaitForCallbackConfig config) { - return currentContext().waitForCallback(name, resultType, adapt(submitter), config); + return waitForCallbackAsync(name, resultType, submitter, config).get(); } public static DurableFuture waitForCallbackAsync(String name, Class resultType, Runnable submitter) { - return currentContext().waitForCallbackAsync(name, resultType, adapt(submitter)); + return waitForCallbackAsync(name, TypeToken.get(resultType), submitter); } public static DurableFuture waitForCallbackAsync(String name, TypeToken resultType, Runnable submitter) { - return currentContext().waitForCallbackAsync(name, resultType, adapt(submitter)); + return waitForCallbackAsync( + name, resultType, submitter, WaitForCallbackConfig.builder().build()); } public static DurableFuture waitForCallbackAsync( String name, Class resultType, Runnable submitter, WaitForCallbackConfig config) { - return currentContext().waitForCallbackAsync(name, resultType, adapt(submitter), config); + return waitForCallbackAsync(name, TypeToken.get(resultType), submitter, config); } public static DurableFuture waitForCallbackAsync( String name, TypeToken resultType, Runnable submitter, WaitForCallbackConfig config) { - return currentContext().waitForCallbackAsync(name, resultType, adapt(submitter), config); + return WaitForCallbackExtension.execute(currentContext(), name, resultType, adapt(submitter), config); } private static BiConsumer adapt(Runnable submitter) { @@ -55,7 +58,7 @@ private static BiConsumer adapt(Runnable submitter) { }; } - private static DurableContext currentContext() { - return DurableContext.getCurrentContext(); + private static ExtensionContext currentContext() { + return ExtensionContext.getCurrentContext(); } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionChildOperationSummary.java b/sdk/src/main/java/software/amazon/lambda/durable/ExtensionChildOperationSummary.java deleted file mode 100644 index db0bb694f..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionChildOperationSummary.java +++ /dev/null @@ -1,11 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; - -import software.amazon.awssdk.services.lambda.model.ErrorObject; -import software.amazon.awssdk.services.lambda.model.OperationStatus; -import software.amazon.awssdk.services.lambda.model.OperationType; - -/** Read-only summary of a direct child operation involved in an extension CONTEXT failure. */ -public record ExtensionChildOperationSummary( - OperationType operationType, String subType, OperationStatus status, ErrorObject error) {} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java index 157b547a9..8a917d061 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java @@ -15,16 +15,10 @@ import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; -import software.amazon.lambda.durable.ExtensionContext; -import software.amazon.lambda.durable.ExtensionContextFunction; -import software.amazon.lambda.durable.ExtensionOperation; -import software.amazon.lambda.durable.ExtensionStepFunction; import software.amazon.lambda.durable.ParallelDurableFuture; import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.CallbackConfig; -import software.amazon.lambda.durable.config.ExtensionContextConfig; -import software.amazon.lambda.durable.config.ExtensionStepConfig; import software.amazon.lambda.durable.config.InvokeConfig; import software.amazon.lambda.durable.config.MapConfig; import software.amazon.lambda.durable.config.ParallelConfig; @@ -33,11 +27,18 @@ import software.amazon.lambda.durable.config.WaitForCallbackConfig; import software.amazon.lambda.durable.config.WaitForConditionConfig; import software.amazon.lambda.durable.config.WithRetryConfig; +import software.amazon.lambda.durable.context.extension.WaitForCallbackExtension; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.execution.OperationIdGenerator; import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.execution.ThreadType; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextFunction; +import software.amazon.lambda.durable.extension.ExtensionOperation; +import software.amazon.lambda.durable.extension.ExtensionStepConfig; +import software.amazon.lambda.durable.extension.ExtensionStepFunction; import software.amazon.lambda.durable.model.MapResult; import software.amazon.lambda.durable.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; @@ -497,42 +498,7 @@ public DurableFuture waitForCallbackAsync( TypeToken resultType, BiConsumer func, WaitForCallbackConfig waitForCallbackConfig) { - Objects.requireNonNull(resultType, "resultType cannot be null"); - Objects.requireNonNull(waitForCallbackConfig, "waitForCallbackConfig cannot be null"); - // waitForCallback adds a suffix for the callback operation name and the submitter operation name so - // the length restriction of waitForCallback name is different from the other operations. - ParameterValidator.validateOperationName(name, MAX_WAIT_FOR_CALLBACK_NAME_LENGTH); - - var finalWaitForCallbackConfig = waitForCallbackConfig.stepConfig().serDes() == null - ? waitForCallbackConfig.toBuilder() - .stepConfig(waitForCallbackConfig.stepConfig().toBuilder() - .serDes(getDurableConfig().getSerDes()) - .build()) - .build() - : waitForCallbackConfig; - - return runInChildContextAsync( - name, - resultType, - childCtx -> { - var callback = childCtx.createCallback( - name + WAIT_FOR_CALLBACK_CALLBACK_SUFFIX, - resultType, - finalWaitForCallbackConfig.callbackConfig()); - childCtx.step( - name + WAIT_FOR_CALLBACK_SUBMITTER_SUFFIX, - Void.class, - stepCtx -> { - func.accept(callback.callbackId(), stepCtx); - return null; - }, - finalWaitForCallbackConfig.stepConfig()); - return callback.get(); - }, - RunInChildContextConfig.builder() - .serDes(finalWaitForCallbackConfig.stepConfig().serDes()) - .build(), - OperationSubType.WAIT_FOR_CALLBACK); + return WaitForCallbackExtension.execute(this, name, resultType, func, waitForCallbackConfig); } @Override diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java index c84611ed9..2928a841b 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java @@ -8,16 +8,16 @@ import java.util.function.Supplier; import software.amazon.lambda.durable.DurableCallbackFuture; import software.amazon.lambda.durable.DurableFuture; -import software.amazon.lambda.durable.ExtensionContextFunction; -import software.amazon.lambda.durable.ExtensionOperation; -import software.amazon.lambda.durable.ExtensionStepFunction; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.CallbackConfig; -import software.amazon.lambda.durable.config.ExtensionContextConfig; -import software.amazon.lambda.durable.config.ExtensionStepConfig; import software.amazon.lambda.durable.config.InvokeConfig; import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.config.StepConfig; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextFunction; +import software.amazon.lambda.durable.extension.ExtensionOperation; +import software.amazon.lambda.durable.extension.ExtensionStepConfig; +import software.amazon.lambda.durable.extension.ExtensionStepFunction; final class ExtensionOperationImpl implements ExtensionOperation { private final DurableContextImpl context; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForCallbackExtension.java b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForCallbackExtension.java new file mode 100644 index 000000000..c8253c46d --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForCallbackExtension.java @@ -0,0 +1,117 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.context.extension; + +import static software.amazon.lambda.durable.execution.ExecutionManager.isTerminalStatus; + +import java.util.Objects; +import java.util.function.BiConsumer; +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.DurableFuture; +import software.amazon.lambda.durable.StepContext; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.RunInChildContextConfig; +import software.amazon.lambda.durable.config.WaitForCallbackConfig; +import software.amazon.lambda.durable.exception.CallbackFailedException; +import software.amazon.lambda.durable.exception.CallbackSubmitterException; +import software.amazon.lambda.durable.exception.CallbackTimeoutException; +import software.amazon.lambda.durable.exception.StepFailedException; +import software.amazon.lambda.durable.exception.StepInterruptedException; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextFailure; +import software.amazon.lambda.durable.extension.ExtensionContextResult; +import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.util.ParameterValidator; + +/** Canonical implementation of the built-in wait-for-callback extension. */ +public final class WaitForCallbackExtension { + private static final String CALLBACK_SUFFIX = "-callback"; + private static final String SUBMITTER_SUFFIX = "-submitter"; + private static final int MAX_NAME_LENGTH = ParameterValidator.MAX_OPERATION_NAME_LENGTH + - Math.max(CALLBACK_SUFFIX.length(), SUBMITTER_SUFFIX.length()); + + private WaitForCallbackExtension() {} + + public static DurableFuture execute( + ExtensionContext context, + String name, + TypeToken resultType, + BiConsumer submitter, + WaitForCallbackConfig config) { + Objects.requireNonNull(context, "context cannot be null"); + Objects.requireNonNull(resultType, "resultType cannot be null"); + Objects.requireNonNull(submitter, "submitter cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + ParameterValidator.validateOperationName(name, MAX_NAME_LENGTH); + + var parent = context.reserve(name); + return parent.runInChildContextAsync( + OperationSubType.WAIT_FOR_CALLBACK.getValue(), + resultType, + () -> executeInChildContext(name, resultType, submitter, config), + extensionConfig(config)); + } + + private static ExtensionContextResult executeInChildContext( + String name, + TypeToken resultType, + BiConsumer submitter, + WaitForCallbackConfig config) { + var child = ExtensionContext.getCurrentContext(); + var callback = child.reserve(name + CALLBACK_SUFFIX).createCallback(resultType, config.callbackConfig()); + child.reserve(name + SUBMITTER_SUFFIX) + .step( + Void.class, + () -> { + submitter.accept(callback.callbackId(), StepContext.getCurrentContext()); + return null; + }, + config.stepConfig()); + return ExtensionContextResult.completed(callback.get()); + } + + private static ExtensionContextConfig extensionConfig(WaitForCallbackConfig config) { + return ExtensionContextConfig.builder() + .childContextConfig(RunInChildContextConfig.builder() + .serDes(config.stepConfig().serDes()) + .build()) + .errorHandler(WaitForCallbackExtension::translateFailure) + .build(); + } + + private static Throwable translateFailure(ExtensionContextFailure failure) { + var callback = findChild(failure, OperationType.CALLBACK); + var submitter = findChild(failure, OperationType.STEP); + if (callback != null && isTerminalStatus(callback.status())) { + if (callback.status() == OperationStatus.FAILED) { + return new CallbackFailedException(callback); + } + if (callback.status() == OperationStatus.TIMED_OUT) { + return new CallbackTimeoutException(callback); + } + } + if (callback != null + && submitter != null + && isTerminalStatus(submitter.status()) + && submitter.status() != OperationStatus.SUCCEEDED) { + var error = submitter.stepDetails().error(); + var cause = StepInterruptedException.isStepInterruptedException(error) + ? new StepInterruptedException(submitter) + : new StepFailedException(submitter); + return new CallbackSubmitterException(callback, cause); + } + return new IllegalStateException("Unknown waitForCallback status"); + } + + private static Operation findChild(ExtensionContextFailure failure, OperationType type) { + return failure.childOperations().stream() + .filter(summary -> summary.operationType() == type) + .map(summary -> summary.operation()) + .filter(Objects::nonNull) + .findFirst() + .orElse(null); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionChildOperationSummary.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionChildOperationSummary.java new file mode 100644 index 000000000..14f6a6299 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionChildOperationSummary.java @@ -0,0 +1,78 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.extension; + +import java.util.Objects; +import software.amazon.awssdk.services.lambda.model.ErrorObject; +import software.amazon.awssdk.services.lambda.model.Operation; +import software.amazon.awssdk.services.lambda.model.OperationStatus; +import software.amazon.awssdk.services.lambda.model.OperationType; + +/** Read-only summary of a direct child operation involved in an extension CONTEXT failure. */ +public final class ExtensionChildOperationSummary { + private final Operation operation; + private final OperationType operationType; + private final String subType; + private final OperationStatus status; + private final ErrorObject error; + + public ExtensionChildOperationSummary(Operation operation) { + this.operation = Objects.requireNonNull(operation, "operation cannot be null"); + operationType = operation.type(); + subType = operation.subType(); + status = operation.status(); + error = extractError(operation); + } + + public ExtensionChildOperationSummary( + OperationType operationType, String subType, OperationStatus status, ErrorObject error) { + operation = null; + this.operationType = operationType; + this.subType = subType; + this.status = status; + this.error = error; + } + + public Operation operation() { + return operation; + } + + public OperationType operationType() { + return operationType; + } + + public String subType() { + return subType; + } + + public OperationStatus status() { + return status; + } + + public ErrorObject error() { + return error; + } + + private static ErrorObject extractError(Operation operation) { + if (operation.type() == null) { + return null; + } + return switch (operation.type()) { + case STEP -> + operation.stepDetails() == null ? null : operation.stepDetails().error(); + case CHAINED_INVOKE -> + operation.chainedInvokeDetails() == null + ? null + : operation.chainedInvokeDetails().error(); + case CALLBACK -> + operation.callbackDetails() == null + ? null + : operation.callbackDetails().error(); + case CONTEXT -> + operation.contextDetails() == null + ? null + : operation.contextDetails().error(); + default -> null; + }; + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContext.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContext.java similarity index 97% rename from sdk/src/main/java/software/amazon/lambda/durable/ExtensionContext.java rename to sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContext.java index d776ee2c8..a136d945f 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContext.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContext.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; +package software.amazon.lambda.durable.extension; import software.amazon.lambda.durable.context.BaseContext; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/ExtensionContextConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextConfig.java similarity index 96% rename from sdk/src/main/java/software/amazon/lambda/durable/config/ExtensionContextConfig.java rename to sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextConfig.java index 177e5fdb9..3ce387d43 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/ExtensionContextConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextConfig.java @@ -1,9 +1,9 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.config; +package software.amazon.lambda.durable.extension; import java.util.Objects; -import software.amazon.lambda.durable.ExtensionContextErrorHandler; +import software.amazon.lambda.durable.config.RunInChildContextConfig; /** Extension-only policies for an advanced CONTEXT primitive. */ public final class ExtensionContextConfig { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextErrorHandler.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextErrorHandler.java similarity index 88% rename from sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextErrorHandler.java rename to sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextErrorHandler.java index a38fd500b..221762e7e 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextErrorHandler.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextErrorHandler.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; +package software.amazon.lambda.durable.extension; /** Translates an extension CONTEXT failure when its original exception cannot be reconstructed. */ @FunctionalInterface diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextFailure.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextFailure.java similarity index 96% rename from sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextFailure.java rename to sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextFailure.java index b6fbb9126..320e1c9aa 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextFailure.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextFailure.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; +package software.amazon.lambda.durable.extension; import java.util.List; import software.amazon.awssdk.services.lambda.model.ErrorObject; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextFunction.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextFunction.java similarity index 88% rename from sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextFunction.java rename to sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextFunction.java index 0ef626088..de59f6fec 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextFunction.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextFunction.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; +package software.amazon.lambda.durable.extension; /** Framework callback for an advanced extension CONTEXT primitive. */ @FunctionalInterface diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextReplayContext.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextReplayContext.java similarity index 64% rename from sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextReplayContext.java rename to sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextReplayContext.java index f5fd1ae1c..b19c147ae 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextReplayContext.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextReplayContext.java @@ -1,13 +1,12 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; +package software.amazon.lambda.durable.extension; import software.amazon.lambda.durable.model.SafeCloseable; /** Replay metadata available while an advanced extension CONTEXT framework callback is running. */ public final class ExtensionContextReplayContext { - private static final OperationContextStorage> CURRENT = - new OperationContextStorage<>("ExtensionContextReplayContext"); + private static final ThreadLocal> CURRENT = new ThreadLocal<>(); private final boolean replayingChildren; private final T replayState; @@ -20,7 +19,11 @@ private ExtensionContextReplayContext(boolean replayingChildren, T replayState) /** Returns the replay context attached to the current extension framework thread. */ @SuppressWarnings("unchecked") public static ExtensionContextReplayContext getCurrentContext() { - return (ExtensionContextReplayContext) CURRENT.getCurrentContext(); + var context = CURRENT.get(); + if (context == null) { + throw new IllegalStateException("ExtensionContextReplayContext is not active on the current thread"); + } + return (ExtensionContextReplayContext) context; } /** Returns whether a completed CONTEXT is replaying its child operations. */ @@ -35,6 +38,14 @@ public T getReplayState() { /** Attaches replay metadata for the duration of an SDK-managed framework callback. */ public static SafeCloseable attach(boolean replayingChildren, T replayState) { - return CURRENT.attach(new ExtensionContextReplayContext<>(replayingChildren, replayState)); + var previous = CURRENT.get(); + CURRENT.set(new ExtensionContextReplayContext<>(replayingChildren, replayState)); + return () -> { + if (previous == null) { + CURRENT.remove(); + } else { + CURRENT.set(previous); + } + }; } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextResult.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextResult.java similarity index 97% rename from sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextResult.java rename to sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextResult.java index bc86e0a5b..b444dd225 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextResult.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextResult.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; +package software.amazon.lambda.durable.extension; /** Result and replay policy returned by an advanced extension CONTEXT primitive. */ public final class ExtensionContextResult { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionOperation.java similarity index 98% rename from sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java rename to sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionOperation.java index b4901565e..6a205b710 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionOperation.java @@ -1,12 +1,13 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; +package software.amazon.lambda.durable.extension; import java.time.Duration; import java.util.function.Supplier; +import software.amazon.lambda.durable.DurableCallbackFuture; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.CallbackConfig; -import software.amazon.lambda.durable.config.ExtensionContextConfig; -import software.amazon.lambda.durable.config.ExtensionStepConfig; import software.amazon.lambda.durable.config.InvokeConfig; import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.config.StepConfig; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/ExtensionStepConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepConfig.java similarity index 97% rename from sdk/src/main/java/software/amazon/lambda/durable/config/ExtensionStepConfig.java rename to sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepConfig.java index fd6a06321..81c0c4284 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/ExtensionStepConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepConfig.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.config; +package software.amazon.lambda.durable.extension; import software.amazon.lambda.durable.serde.SerDes; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionStepFunction.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepFunction.java similarity index 91% rename from sdk/src/main/java/software/amazon/lambda/durable/ExtensionStepFunction.java rename to sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepFunction.java index 48a06fd08..866236f67 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionStepFunction.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepFunction.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; +package software.amazon.lambda.durable.extension; /** * Evaluates one attempt of a stateful extension step. diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionStepResult.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepResult.java similarity index 95% rename from sdk/src/main/java/software/amazon/lambda/durable/ExtensionStepResult.java rename to sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepResult.java index fb0404927..eb7067b26 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/ExtensionStepResult.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepResult.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; +package software.amazon.lambda.durable.extension; import java.time.Duration; import java.util.Objects; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java index 38c984d3b..3398e4ac9 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java @@ -19,13 +19,7 @@ import software.amazon.awssdk.services.lambda.model.OperationType; import software.amazon.awssdk.services.lambda.model.OperationUpdate; import software.amazon.lambda.durable.DurableContext; -import software.amazon.lambda.durable.ExtensionChildOperationSummary; -import software.amazon.lambda.durable.ExtensionContextFailure; -import software.amazon.lambda.durable.ExtensionContextFunction; -import software.amazon.lambda.durable.ExtensionContextReplayContext; -import software.amazon.lambda.durable.ExtensionContextResult; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.ExtensionContextConfig; import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.exception.CallbackFailedException; @@ -40,6 +34,12 @@ import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.execution.ThreadType; +import software.amazon.lambda.durable.extension.ExtensionChildOperationSummary; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextFailure; +import software.amazon.lambda.durable.extension.ExtensionContextFunction; +import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; +import software.amazon.lambda.durable.extension.ExtensionContextResult; import software.amazon.lambda.durable.logging.DurableLogger; import software.amazon.lambda.durable.model.DeserializedOperationResult; import software.amazon.lambda.durable.model.OperationDescriptor; @@ -395,11 +395,7 @@ private Throwable translateException(Operation op, ErrorObject errorObject) { private List getChildOperationSummaries() { return getChildOperations().stream() - .map(operation -> new ExtensionChildOperationSummary( - operation.type(), - operation.subType(), - operation.status(), - BaseDurableOperation.getErrorObject(operation))) + .map(ExtensionChildOperationSummary::new) .toList(); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java index bf753b2f2..67b8f0bf2 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java @@ -11,11 +11,8 @@ import software.amazon.awssdk.services.lambda.model.OperationStatus; import software.amazon.awssdk.services.lambda.model.OperationUpdate; import software.amazon.awssdk.services.lambda.model.StepOptions; -import software.amazon.lambda.durable.ExtensionStepFunction; -import software.amazon.lambda.durable.ExtensionStepResult; import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.ExtensionStepConfig; import software.amazon.lambda.durable.config.StepConfig; import software.amazon.lambda.durable.config.StepSemantics; import software.amazon.lambda.durable.context.BaseContextImpl; @@ -26,6 +23,9 @@ import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.execution.ThreadType; +import software.amazon.lambda.durable.extension.ExtensionStepConfig; +import software.amazon.lambda.durable.extension.ExtensionStepFunction; +import software.amazon.lambda.durable.extension.ExtensionStepResult; import software.amazon.lambda.durable.logging.DurableLogger; import software.amazon.lambda.durable.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; diff --git a/sdk/src/test/java/software/amazon/lambda/durable/CurrentContextTest.java b/sdk/src/test/java/software/amazon/lambda/durable/CurrentContextTest.java index c0a54b9e3..41c94c422 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/CurrentContextTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/CurrentContextTest.java @@ -10,6 +10,7 @@ import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import software.amazon.lambda.durable.context.BaseContextImpl; +import software.amazon.lambda.durable.extension.ExtensionContext; class CurrentContextTest { @AfterEach diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationsTest.java index 9da335025..677a5ef20 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationsTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationsTest.java @@ -4,15 +4,24 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; -import java.util.function.BiConsumer; +import java.util.function.Supplier; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; +import software.amazon.lambda.durable.config.CallbackConfig; +import software.amazon.lambda.durable.config.StepConfig; import software.amazon.lambda.durable.context.BaseContextImpl; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextFunction; +import software.amazon.lambda.durable.extension.ExtensionOperation; +import software.amazon.lambda.durable.model.OperationSubType; class DurableWaitForCallbackOperationsTest { @AfterEach @@ -22,21 +31,65 @@ void clearContext() { @Test void callbackSubmitterUsesRunnableAndScopedCallbackId() { - var context = mock(DurableContext.class); + var context = mock(ExtensionContext.class); + var parent = mock(ExtensionOperation.class); + var parentFuture = mockStringFuture(); + when(context.reserve("callback")).thenReturn(parent); + when(parentFuture.get()).thenReturn("approved"); + when(parent.runInChildContextAsync( + eq(OperationSubType.WAIT_FOR_CALLBACK.getValue()), + eq(TypeToken.get(String.class)), + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) + .thenReturn(parentFuture); BaseContextImpl.setCurrentContext(context); - DurableWaitForCallbackOperations.waitForCallback( - "callback", - String.class, - () -> assertEquals( - "callback-id", - WaitForCallbackContext.getCurrentContext().getCallbackId())); + assertEquals( + "approved", + DurableWaitForCallbackOperations.waitForCallback( + "callback", + String.class, + () -> assertEquals( + "callback-id", + WaitForCallbackContext.getCurrentContext().getCallbackId()))); @SuppressWarnings("unchecked") - var submitter = (ArgumentCaptor>) - (ArgumentCaptor) ArgumentCaptor.forClass(BiConsumer.class); - verify(context).waitForCallback(eq("callback"), eq(String.class), submitter.capture()); - submitter.getValue().accept("callback-id", mock(StepContext.class)); + var function = (ArgumentCaptor>) + (ArgumentCaptor) ArgumentCaptor.forClass(ExtensionContextFunction.class); + verify(parent) + .runInChildContextAsync( + eq(OperationSubType.WAIT_FOR_CALLBACK.getValue()), + eq(TypeToken.get(String.class)), + function.capture(), + any(ExtensionContextConfig.class)); + + var child = mock(ExtensionContext.class); + var callbackReservation = mock(ExtensionOperation.class); + var submitterReservation = mock(ExtensionOperation.class); + @SuppressWarnings("unchecked") + var callback = (DurableCallbackFuture) mock(DurableCallbackFuture.class); + when(child.reserve("callback-callback")).thenReturn(callbackReservation); + when(child.reserve("callback-submitter")).thenReturn(submitterReservation); + when(callbackReservation.createCallback(eq(TypeToken.get(String.class)), any(CallbackConfig.class))) + .thenReturn(callback); + when(callback.callbackId()).thenReturn("callback-id"); + when(callback.get()).thenReturn("approved"); + + try (var ignored = BaseContextImpl.attachCurrentContext(child)) { + assertEquals("approved", function.getValue().apply().result()); + } + + @SuppressWarnings("unchecked") + var submitter = (ArgumentCaptor>) (ArgumentCaptor) ArgumentCaptor.forClass(Supplier.class); + verify(submitterReservation).step(eq(Void.class), submitter.capture(), any(StepConfig.class)); + try (var ignored = BaseContextImpl.attachCurrentContext(mock(StepContext.class))) { + submitter.getValue().get(); + } assertThrows(IllegalStateException.class, WaitForCallbackContext::getCurrentContext); } + + @SuppressWarnings("unchecked") + private DurableFuture mockStringFuture() { + return mock(DurableFuture.class); + } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/ExtensionContextReplayContextTest.java b/sdk/src/test/java/software/amazon/lambda/durable/ExtensionContextReplayContextTest.java index ad2dbb7fd..28e58bd26 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/ExtensionContextReplayContextTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/ExtensionContextReplayContextTest.java @@ -8,6 +8,7 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; class ExtensionContextReplayContextTest { @Test diff --git a/sdk/src/test/java/software/amazon/lambda/durable/ExtensionContextResultTest.java b/sdk/src/test/java/software/amazon/lambda/durable/ExtensionContextResultTest.java index 54cd4da5a..f3aa4ae8c 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/ExtensionContextResultTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/ExtensionContextResultTest.java @@ -8,6 +8,7 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.extension.ExtensionContextResult; class ExtensionContextResultTest { @Test diff --git a/sdk/src/test/java/software/amazon/lambda/durable/ExtensionStepResultTest.java b/sdk/src/test/java/software/amazon/lambda/durable/ExtensionStepResultTest.java index c18581628..3963b2619 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/ExtensionStepResultTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/ExtensionStepResultTest.java @@ -7,6 +7,7 @@ import java.time.Duration; import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.extension.ExtensionStepResult; class ExtensionStepResultTest { @Test diff --git a/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionContextConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionContextConfigTest.java index 1f2a87dd6..69589039c 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionContextConfigTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionContextConfigTest.java @@ -9,7 +9,8 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import org.junit.jupiter.api.Test; -import software.amazon.lambda.durable.ExtensionContextErrorHandler; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextErrorHandler; class ExtensionContextConfigTest { @Test diff --git a/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionStepConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionStepConfigTest.java index 592858800..53370aee0 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionStepConfigTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionStepConfigTest.java @@ -6,6 +6,7 @@ import static org.junit.jupiter.api.Assertions.assertNull; import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.extension.ExtensionStepConfig; import software.amazon.lambda.durable.serde.JacksonSerDes; class ExtensionStepConfigTest { diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java index 3fa38cf1b..301d484da 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java @@ -20,17 +20,17 @@ import software.amazon.lambda.durable.DurableCallbackFuture; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; -import software.amazon.lambda.durable.ExtensionContextResult; -import software.amazon.lambda.durable.ExtensionOperation; -import software.amazon.lambda.durable.ExtensionStepResult; import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.CallbackConfig; -import software.amazon.lambda.durable.config.ExtensionContextConfig; -import software.amazon.lambda.durable.config.ExtensionStepConfig; import software.amazon.lambda.durable.config.InvokeConfig; import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.config.StepConfig; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextResult; +import software.amazon.lambda.durable.extension.ExtensionOperation; +import software.amazon.lambda.durable.extension.ExtensionStepConfig; +import software.amazon.lambda.durable.extension.ExtensionStepResult; class ExtensionOperationImplTest { @Test diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/WaitForCallbackExtensionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/extension/WaitForCallbackExtensionTest.java new file mode 100644 index 000000000..cb9203bed --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/context/extension/WaitForCallbackExtensionTest.java @@ -0,0 +1,120 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.context.extension; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.List; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import software.amazon.awssdk.services.lambda.model.CallbackDetails; +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.DurableFuture; +import software.amazon.lambda.durable.StepContext; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.StepConfig; +import software.amazon.lambda.durable.config.WaitForCallbackConfig; +import software.amazon.lambda.durable.exception.CallbackTimeoutException; +import software.amazon.lambda.durable.extension.ExtensionChildOperationSummary; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextFailure; +import software.amazon.lambda.durable.extension.ExtensionContextFunction; +import software.amazon.lambda.durable.extension.ExtensionOperation; +import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.serde.JacksonSerDes; + +class WaitForCallbackExtensionTest { + @Test + void executeCreatesExistingWaitForCallbackContextTopology() { + var context = mock(ExtensionContext.class); + var parent = mock(ExtensionOperation.class); + var future = mockStringFuture(); + var resultType = TypeToken.get(String.class); + var serDes = new JacksonSerDes(); + var config = WaitForCallbackConfig.builder() + .stepConfig(StepConfig.builder().serDes(serDes).build()) + .build(); + when(context.reserve("approval")).thenReturn(parent); + when(parent.runInChildContextAsync( + eq(OperationSubType.WAIT_FOR_CALLBACK.getValue()), + eq(resultType), + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) + .thenReturn(future); + + var actual = WaitForCallbackExtension.execute( + context, "approval", resultType, (callbackId, stepContext) -> {}, config); + + assertSame(future, actual); + var contextConfig = ArgumentCaptor.forClass(ExtensionContextConfig.class); + verify(parent) + .runInChildContextAsync( + eq(OperationSubType.WAIT_FOR_CALLBACK.getValue()), + eq(resultType), + any(ExtensionContextFunction.class), + contextConfig.capture()); + assertSame(serDes, contextConfig.getValue().childContextConfig().serDes()); + } + + @Test + void errorHandlerPreservesCallbackTimeoutException() { + var context = mock(ExtensionContext.class); + var parent = mock(ExtensionOperation.class); + var resultType = TypeToken.get(String.class); + when(context.reserve("approval")).thenReturn(parent); + when(parent.runInChildContextAsync( + eq(OperationSubType.WAIT_FOR_CALLBACK.getValue()), + eq(resultType), + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) + .thenReturn(mockStringFuture()); + WaitForCallbackExtension.execute( + context, + "approval", + resultType, + (String callbackId, StepContext stepContext) -> {}, + WaitForCallbackConfig.builder().build()); + var config = ArgumentCaptor.forClass(ExtensionContextConfig.class); + verify(parent) + .runInChildContextAsync( + eq(OperationSubType.WAIT_FOR_CALLBACK.getValue()), + eq(resultType), + any(ExtensionContextFunction.class), + config.capture()); + var callback = Operation.builder() + .id("callback-op") + .name("approval-callback") + .type(OperationType.CALLBACK) + .subType(OperationSubType.CALLBACK.getValue()) + .status(OperationStatus.TIMED_OUT) + .callbackDetails( + CallbackDetails.builder().callbackId("callback-id").build()) + .build(); + var failure = new ExtensionContextFailure( + "approval", + OperationSubType.WAIT_FOR_CALLBACK.getValue(), + null, + null, + List.of(new ExtensionChildOperationSummary(callback))); + + var translated = config.getValue().errorHandler().translate(failure); + + var timeout = assertInstanceOf(CallbackTimeoutException.class, translated); + assertEquals("callback-id", timeout.getCallbackId()); + } + + @SuppressWarnings("unchecked") + private DurableFuture mockStringFuture() { + return mock(DurableFuture.class); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java index f91da9131..b72af25aa 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java @@ -22,10 +22,7 @@ import software.amazon.awssdk.services.lambda.model.StepDetails; import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.DurableContext; -import software.amazon.lambda.durable.ExtensionContextFailure; -import software.amazon.lambda.durable.ExtensionContextResult; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.ExtensionContextConfig; import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.exception.ChildContextFailedException; @@ -34,6 +31,9 @@ import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.execution.ThreadContext; import software.amazon.lambda.durable.execution.ThreadType; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextFailure; +import software.amazon.lambda.durable.extension.ExtensionContextResult; import software.amazon.lambda.durable.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; From b5590fa15c4e9c72dacc1dbadab9b15835b98ec3 Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 04:06:40 +0000 Subject: [PATCH 19/40] refactor: implement retry and condition as extensions --- .../DurableWaitForConditionOperations.java | 26 +- .../durable/DurableWithRetryOperations.java | 14 +- .../durable/context/DurableContextImpl.java | 80 +--- .../extension/WaitForConditionExtension.java | 59 +++ .../extension/WaitForConditionFuture.java | 31 ++ .../context/extension/WithRetryExtension.java | 78 ++++ .../operation/WaitForConditionOperation.java | 189 --------- ...DurableWaitForConditionOperationsTest.java | 54 ++- .../DurableWithRetryOperationsTest.java | 51 ++- .../WaitForConditionExtensionTest.java | 147 +++++++ .../extension/WithRetryExtensionTest.java | 104 +++++ .../StatefulExtensionStepOperationTest.java | 262 ++++++++++++ .../WaitForConditionOperationTest.java | 392 ------------------ 13 files changed, 794 insertions(+), 693 deletions(-) create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForConditionExtension.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForConditionFuture.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/context/extension/WithRetryExtension.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/context/extension/WaitForConditionExtensionTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/context/extension/WithRetryExtensionTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/operation/StatefulExtensionStepOperationTest.java delete mode 100644 sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionOperationTest.java diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForConditionOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForConditionOperations.java index 06784844f..1d30320ef 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForConditionOperations.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForConditionOperations.java @@ -6,6 +6,8 @@ import java.util.function.BiFunction; import java.util.function.Function; import software.amazon.lambda.durable.config.WaitForConditionConfig; +import software.amazon.lambda.durable.context.extension.WaitForConditionExtension; +import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.model.WaitForConditionResult; /** Context-free static facades for durable wait-for-condition operations. */ @@ -14,12 +16,12 @@ private DurableWaitForConditionOperations() {} public static T waitForCondition( String name, Class resultType, Function> checkFunction) { - return currentContext().waitForCondition(name, resultType, adapt(checkFunction)); + return waitForConditionAsync(name, resultType, checkFunction).get(); } public static T waitForCondition( String name, TypeToken resultType, Function> checkFunction) { - return currentContext().waitForCondition(name, resultType, adapt(checkFunction)); + return waitForConditionAsync(name, resultType, checkFunction).get(); } public static T waitForCondition( @@ -27,7 +29,7 @@ public static T waitForCondition( Class resultType, Function> checkFunction, WaitForConditionConfig config) { - return currentContext().waitForCondition(name, resultType, adapt(checkFunction), config); + return waitForConditionAsync(name, resultType, checkFunction, config).get(); } public static T waitForCondition( @@ -35,17 +37,21 @@ public static T waitForCondition( TypeToken resultType, Function> checkFunction, WaitForConditionConfig config) { - return currentContext().waitForCondition(name, resultType, adapt(checkFunction), config); + return waitForConditionAsync(name, resultType, checkFunction, config).get(); } public static DurableFuture waitForConditionAsync( String name, Class resultType, Function> checkFunction) { - return currentContext().waitForConditionAsync(name, resultType, adapt(checkFunction)); + return waitForConditionAsync(name, TypeToken.get(resultType), checkFunction); } public static DurableFuture waitForConditionAsync( String name, TypeToken resultType, Function> checkFunction) { - return currentContext().waitForConditionAsync(name, resultType, adapt(checkFunction)); + return waitForConditionAsync( + name, + resultType, + checkFunction, + WaitForConditionConfig.builder().build()); } public static DurableFuture waitForConditionAsync( @@ -53,7 +59,7 @@ public static DurableFuture waitForConditionAsync( Class resultType, Function> checkFunction, WaitForConditionConfig config) { - return currentContext().waitForConditionAsync(name, resultType, adapt(checkFunction), config); + return waitForConditionAsync(name, TypeToken.get(resultType), checkFunction, config); } public static DurableFuture waitForConditionAsync( @@ -61,7 +67,7 @@ public static DurableFuture waitForConditionAsync( TypeToken resultType, Function> checkFunction, WaitForConditionConfig config) { - return currentContext().waitForConditionAsync(name, resultType, adapt(checkFunction), config); + return WaitForConditionExtension.execute(currentContext(), name, resultType, adapt(checkFunction), config); } private static BiFunction> adapt( @@ -70,7 +76,7 @@ private static BiFunction> adapt( return (state, ignored) -> checkFunction.apply(state); } - private static DurableContext currentContext() { - return DurableContext.getCurrentContext(); + private static ExtensionContext currentContext() { + return ExtensionContext.getCurrentContext(); } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableWithRetryOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableWithRetryOperations.java index b65b2ee67..c9fac5413 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableWithRetryOperations.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableWithRetryOperations.java @@ -6,25 +6,27 @@ import java.util.function.BiFunction; import java.util.function.Supplier; import software.amazon.lambda.durable.config.WithRetryConfig; +import software.amazon.lambda.durable.context.extension.WithRetryExtension; +import software.amazon.lambda.durable.extension.ExtensionContext; /** Context-free static facades for replay-safe retry operations. */ public final class DurableWithRetryOperations { private DurableWithRetryOperations() {} public static T withRetry(String name, Supplier operation) { - return currentContext().withRetry(name, adapt(operation)); + return withRetryAsync(name, operation).get(); } public static T withRetry(String name, Supplier operation, WithRetryConfig config) { - return currentContext().withRetry(name, adapt(operation), config); + return withRetryAsync(name, operation, config).get(); } public static DurableFuture withRetryAsync(String name, Supplier operation) { - return currentContext().withRetryAsync(name, adapt(operation)); + return withRetryAsync(name, operation, WithRetryConfig.builder().build()); } public static DurableFuture withRetryAsync(String name, Supplier operation, WithRetryConfig config) { - return currentContext().withRetryAsync(name, adapt(operation), config); + return WithRetryExtension.execute(currentContext(), name, adapt(operation), config); } private static BiFunction adapt(Supplier operation) { @@ -36,7 +38,7 @@ private static BiFunction adapt(Supplier oper }; } - private static DurableContext currentContext() { - return DurableContext.getCurrentContext(); + private static ExtensionContext currentContext() { + return ExtensionContext.getCurrentContext(); } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java index 8a917d061..d80466b60 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java @@ -28,10 +28,10 @@ import software.amazon.lambda.durable.config.WaitForConditionConfig; import software.amazon.lambda.durable.config.WithRetryConfig; import software.amazon.lambda.durable.context.extension.WaitForCallbackExtension; -import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.context.extension.WaitForConditionExtension; +import software.amazon.lambda.durable.context.extension.WithRetryExtension; import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.execution.OperationIdGenerator; -import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.execution.ThreadType; import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.extension.ExtensionContextConfig; @@ -51,9 +51,7 @@ import software.amazon.lambda.durable.operation.MapOperation; import software.amazon.lambda.durable.operation.ParallelOperation; import software.amazon.lambda.durable.operation.StepOperation; -import software.amazon.lambda.durable.operation.WaitForConditionOperation; import software.amazon.lambda.durable.operation.WaitOperation; -import software.amazon.lambda.durable.retry.RetryDecision; import software.amazon.lambda.durable.util.ParameterValidator; /** @@ -507,85 +505,15 @@ public DurableFuture waitForConditionAsync( TypeToken resultType, BiFunction> checkFunc, WaitForConditionConfig config) { - Objects.requireNonNull(config, "config cannot be null"); - Objects.requireNonNull(resultType, "resultType cannot be null"); - Objects.requireNonNull(checkFunc, "checkFunc cannot be null"); - ParameterValidator.validateOperationName(name); - - if (config.serDes() == null) { - config = config.toBuilder().serDes(getDurableConfig().getSerDes()).build(); - } - var operationId = nextOperationId(); - - var operation = new WaitForConditionOperation<>( - OperationIdentifier.of(operationId, name, OperationSubType.WAIT_FOR_CONDITION), - checkFunc, - resultType, - config, - this); - - operation.execute(); - - return operation; + return WaitForConditionExtension.execute(this, name, resultType, checkFunc, config); } // =============== withRetry ================ - private static final Duration DEFAULT_BACKOFF_DELAY = Duration.ofSeconds(1); - private static final String BACKOFF_SUFFIX = "-backoff-"; - private static final String ANONYMOUS_CHILD_CONTEXT_NAME = "retry"; - private static final String ANONYMOUS_BACKOFF_PREFIX = "retry-backoff-"; - @Override - @SuppressWarnings("unchecked") public DurableFuture withRetryAsync( String name, BiFunction operation, WithRetryConfig config) { - Objects.requireNonNull(operation, "operation cannot be null"); - Objects.requireNonNull(config, "config cannot be null"); - - var childContextName = name != null ? name : ANONYMOUS_CHILD_CONTEXT_NAME; - - return (DurableFuture) runInChildContextAsync( - childContextName, - new TypeToken() {}, - childCtx -> executeRetryLoop(childCtx, name, operation, config), - RunInChildContextConfig.builder() - .isVirtual(!config.wrapInChildContext()) - .build(), - OperationSubType.WITH_RETRY); - } - - /** - * Core retry loop. Replay-safe because every side-effect is a durable operation: the user's operation calls durable - * primitives, and backoff uses {@code context.wait()}. - * - *

{@link SuspendExecutionException} and {@link UnrecoverableDurableExecutionException} are never retried — they - * are internal SDK control flow signals that must propagate immediately. - */ - private static T executeRetryLoop( - DurableContext context, - String name, - BiFunction operation, - WithRetryConfig config) { - var attempt = 1; - while (true) { - try { - return operation.apply(attempt, context); - } catch (SuspendExecutionException | UnrecoverableDurableExecutionException e) { - // Internal SDK control flow — never retry, always propagate - throw e; - } catch (Exception e) { - RetryDecision decision = config.retryStrategy().makeRetryDecision(e, attempt); - if (!decision.shouldRetry()) { - throw e; - } - - var delay = decision.delay().isZero() ? DEFAULT_BACKOFF_DELAY : decision.delay(); - var waitName = name != null ? name + BACKOFF_SUFFIX + attempt : ANONYMOUS_BACKOFF_PREFIX + attempt; - context.wait(waitName, delay); - attempt++; - } - } + return WithRetryExtension.execute(this, name, operation, config); } // =============== accessors ================ diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForConditionExtension.java b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForConditionExtension.java new file mode 100644 index 000000000..5e8459794 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForConditionExtension.java @@ -0,0 +1,59 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.context.extension; + +import java.util.Objects; +import java.util.function.BiFunction; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.StepContext; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.WaitForConditionConfig; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionStepConfig; +import software.amazon.lambda.durable.extension.ExtensionStepResult; +import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.model.WaitForConditionResult; +import software.amazon.lambda.durable.util.ParameterValidator; + +/** Canonical implementation of the built-in wait-for-condition extension. */ +public final class WaitForConditionExtension { + private WaitForConditionExtension() {} + + public static DurableFuture execute( + ExtensionContext context, + String name, + TypeToken resultType, + BiFunction> checkFunction, + WaitForConditionConfig config) { + Objects.requireNonNull(context, "context cannot be null"); + Objects.requireNonNull(resultType, "resultType cannot be null"); + Objects.requireNonNull(checkFunction, "checkFunction cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + ParameterValidator.validateOperationName(name); + + var future = context.reserve(name) + .stepAsync( + OperationSubType.WAIT_FOR_CONDITION.getValue(), + resultType, + state -> evaluate(state, checkFunction, config), + ExtensionStepConfig.builder() + .initialState(config.initialState()) + .serDes(config.serDes()) + .build()); + return new WaitForConditionFuture<>(future); + } + + private static ExtensionStepResult evaluate( + T state, + BiFunction> checkFunction, + WaitForConditionConfig config) { + var stepContext = StepContext.getCurrentContext(); + var result = Objects.requireNonNull( + checkFunction.apply(state, stepContext), "waitForCondition check result cannot be null"); + if (result.isDone()) { + return ExtensionStepResult.succeed(result.value()); + } + var delay = config.waitStrategy().evaluate(result.value(), stepContext.getAttempt()); + return ExtensionStepResult.retry(result.value(), delay); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForConditionFuture.java b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForConditionFuture.java new file mode 100644 index 000000000..9cc0036b1 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForConditionFuture.java @@ -0,0 +1,31 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.context.extension; + +import java.util.Objects; +import java.util.concurrent.CompletableFuture; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.exception.StepFailedException; +import software.amazon.lambda.durable.exception.WaitForConditionFailedException; + +final class WaitForConditionFuture implements DurableFuture { + private final DurableFuture delegate; + + WaitForConditionFuture(DurableFuture delegate) { + this.delegate = Objects.requireNonNull(delegate, "delegate cannot be null"); + } + + @Override + public T get() { + try { + return delegate.get(); + } catch (StepFailedException e) { + throw new WaitForConditionFailedException(e.getOperation()); + } + } + + @Override + public CompletableFuture completionFuture() { + return delegate.completionFuture(); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WithRetryExtension.java b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WithRetryExtension.java new file mode 100644 index 000000000..2fba9ba1b --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WithRetryExtension.java @@ -0,0 +1,78 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.context.extension; + +import java.time.Duration; +import java.util.Objects; +import java.util.function.BiFunction; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.RunInChildContextConfig; +import software.amazon.lambda.durable.config.WithRetryConfig; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.execution.SuspendExecutionException; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextResult; +import software.amazon.lambda.durable.model.OperationSubType; + +/** Canonical implementation of the built-in with-retry extension. */ +public final class WithRetryExtension { + private static final Duration DEFAULT_BACKOFF_DELAY = Duration.ofSeconds(1); + private static final String BACKOFF_SUFFIX = "-backoff-"; + private static final String ANONYMOUS_CONTEXT_NAME = "retry"; + private static final String ANONYMOUS_BACKOFF_PREFIX = "retry-backoff-"; + + private WithRetryExtension() {} + + @SuppressWarnings("unchecked") + public static DurableFuture execute( + ExtensionContext context, + String name, + BiFunction operation, + WithRetryConfig config) { + Objects.requireNonNull(context, "context cannot be null"); + Objects.requireNonNull(operation, "operation cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + + var contextName = name != null ? name : ANONYMOUS_CONTEXT_NAME; + var future = context.reserve(contextName) + .runInChildContextAsync( + OperationSubType.WITH_RETRY.getValue(), + new TypeToken() {}, + () -> ExtensionContextResult.completed(executeRetryLoop(name, operation, config)), + ExtensionContextConfig.builder() + .childContextConfig(RunInChildContextConfig.builder() + .isVirtual(!config.wrapInChildContext()) + .build()) + .build()); + return (DurableFuture) future; + } + + private static T executeRetryLoop( + String name, BiFunction operation, WithRetryConfig config) { + var durableContext = DurableContext.getCurrentContext(); + var extensionContext = ExtensionContext.getCurrentContext(); + var attempt = 1; + while (true) { + try { + return operation.apply(attempt, durableContext); + } catch (SuspendExecutionException | UnrecoverableDurableExecutionException e) { + throw e; + } catch (Exception e) { + var decision = config.retryStrategy().makeRetryDecision(e, attempt); + if (!decision.shouldRetry()) { + throw e; + } + var delay = decision.delay().isZero() ? DEFAULT_BACKOFF_DELAY : decision.delay(); + extensionContext.reserve(backoffName(name, attempt)).wait(delay); + attempt++; + } + } + } + + private static String backoffName(String name, int attempt) { + return name != null ? name + BACKOFF_SUFFIX + attempt : ANONYMOUS_BACKOFF_PREFIX + attempt; + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java deleted file mode 100644 index c906bc5f6..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java +++ /dev/null @@ -1,189 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; - -import java.time.Duration; -import java.util.concurrent.CompletableFuture; -import java.util.function.BiFunction; -import software.amazon.awssdk.services.lambda.model.Operation; -import software.amazon.awssdk.services.lambda.model.OperationAction; -import software.amazon.awssdk.services.lambda.model.OperationStatus; -import software.amazon.awssdk.services.lambda.model.OperationUpdate; -import software.amazon.awssdk.services.lambda.model.StepOptions; -import software.amazon.lambda.durable.StepContext; -import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.WaitForConditionConfig; -import software.amazon.lambda.durable.context.BaseContextImpl; -import software.amazon.lambda.durable.context.DurableContextImpl; -import software.amazon.lambda.durable.exception.DurableOperationException; -import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; -import software.amazon.lambda.durable.exception.WaitForConditionFailedException; -import software.amazon.lambda.durable.execution.SuspendExecutionException; -import software.amazon.lambda.durable.execution.ThreadType; -import software.amazon.lambda.durable.logging.DurableLogger; -import software.amazon.lambda.durable.model.OperationIdentifier; -import software.amazon.lambda.durable.model.WaitForConditionResult; -import software.amazon.lambda.durable.util.ExceptionHelper; - -/** - * Durable operation that periodically checks a user-supplied condition function, using a configurable wait strategy to - * determine polling intervals and termination. - * - *

Uses {@link OperationType#STEP} with {@link OperationSubType#WAIT_FOR_CONDITION} subtype. Each polling iteration - * is checkpointed as a RETRY on the same STEP operation. - * - * @param the type of state being polled - */ -public class WaitForConditionOperation extends SerializableDurableOperation { - private static final Integer FIRST_ATTEMPT = 1; - - private final BiFunction> checkFunc; - private final WaitForConditionConfig config; - - public WaitForConditionOperation( - OperationIdentifier operationIdentifier, - BiFunction> checkFunc, - TypeToken resultTypeToken, - WaitForConditionConfig config, - DurableContextImpl durableContext) { - super(operationIdentifier, resultTypeToken, config.serDes(), durableContext); - - this.checkFunc = checkFunc; - this.config = config; - } - - @Override - protected void start() { - executeCheckLogic(config.initialState(), FIRST_ATTEMPT); - } - - @Override - protected void replay(Operation existing) { - switch (existing.status()) { - case SUCCEEDED, FAILED -> markAlreadyCompleted(); // Check if already completed / failed - case PENDING -> pollReadyAndResumeCheckLoop(existing); // Check if pending retry - case STARTED, READY -> resumeCheckLoop(existing); - default -> - throw terminateExecutionWithIllegalDurableOperationException( - "Unexpected waitForCondition status: " + existing.status()); - } - } - - @Override - public T get() { - var op = waitForOperationCompletion(); - - if (op.status() == OperationStatus.SUCCEEDED) { - var stepDetails = op.stepDetails(); - var result = (stepDetails != null) ? stepDetails.result() : null; - return deserializeResult(result); - } else { - var errorObject = op.stepDetails().error(); - - // Attempt to reconstruct and throw the original exception - Throwable original = deserializeException(errorObject); - if (original != null) { - ExceptionHelper.sneakyThrow(original); - } - // Fallback: wrap in WaitForConditionFailedException - throw new WaitForConditionFailedException(op); - } - } - - private void resumeCheckLoop(Operation existing) { - var stepDetails = existing.stepDetails(); - int attempt = - (stepDetails != null && stepDetails.attempt() != null) ? stepDetails.attempt() + 1 : FIRST_ATTEMPT; - var checkpointData = stepDetails != null ? stepDetails.result() : null; - T currentState; // Get current state - if (checkpointData != null) { - currentState = deserializeResult(checkpointData); - } else { - currentState = config.initialState(); - } - executeCheckLogic(currentState, attempt); - } - - private CompletableFuture pollReadyAndResumeCheckLoop(Operation existing) { - return pollForOperationUpdates() - .thenCompose(op -> op.status() == OperationStatus.READY - ? CompletableFuture.completedFuture(op) - : pollForOperationUpdates()) - .thenAccept(this::resumeCheckLoop); - } - - private void executeCheckLogic(T currentState, int attempt) { - Runnable userHandler = () -> { - var stepContext = getContext().createStepContext(getOperationId(), getName(), attempt); - try (var ignoredContext = BaseContextImpl.attachCurrentContext(stepContext); - var ignoredLogger = DurableLogger.attachContext()) { - try { - // Checkpoint START if not already started - var existing = getOperation(); - if (existing == null || existing.status() != OperationStatus.STARTED) { - var startUpdate = OperationUpdate.builder().action(OperationAction.START); - sendOperationUpdateAsync(startUpdate); - } - - // Execute check function inside the plugin hook boundary so a failure is reported - // through onUserFunctionEnd; checkpoint/poll handling stays outside the boundary. - WaitForConditionResult result = - runUserFunction(attempt, () -> checkFunc.apply(currentState, stepContext)); - - // Normalize the value through SerDes so first execution matches replay. - var serializedState = serializeAndDeserializeResult(result.value()); - T deserializedValue = serializedState.deserialized(); - - if (result.isDone()) { - // Condition met — checkpoint SUCCEED - var successUpdate = OperationUpdate.builder() - .action(OperationAction.SUCCEED) - .payload(serializedState.serialized()); - sendOperationUpdate(successUpdate); - } else { - // Compute delay from strategy - Duration delay = config.waitStrategy().evaluate(deserializedValue, attempt); - - // Checkpoint RETRY with delay - var retryUpdate = OperationUpdate.builder() - .action(OperationAction.RETRY) - .payload(serializedState.serialized()) - .stepOptions(StepOptions.builder() - .nextAttemptDelaySeconds(Math.toIntExact(delay.toSeconds())) - .build()); - sendOperationUpdate(retryUpdate); - - // Poll for READY, then continue the loop - pollForOperationUpdates() - .thenCompose(op -> op.status() == OperationStatus.READY - ? CompletableFuture.completedFuture(op) - : pollForOperationUpdates()) - .thenRun(() -> executeCheckLogic(deserializedValue, attempt + 1)); - } - } catch (Throwable e) { - handleCheckFailure(e); - } - } - }; - - runUserHandler(userHandler, ThreadType.STEP); - } - - private void handleCheckFailure(Throwable exception) { - exception = ExceptionHelper.unwrapCompletableFuture(exception); - if (exception instanceof SuspendExecutionException suspendExecutionException) { - throw suspendExecutionException; - } - if (exception instanceof UnrecoverableDurableExecutionException unrecoverable) { - throw terminateExecution(unrecoverable); - } - - final var errorObject = (exception instanceof DurableOperationException durableOpEx) - ? durableOpEx.getErrorObject() - : serializeException(exception); - - // Checkpoint FAIL - var failUpdate = OperationUpdate.builder().action(OperationAction.FAIL).error(errorObject); - sendOperationUpdate(failUpdate); - } -} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForConditionOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForConditionOperationsTest.java index 5d2cc9a52..40558124b 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForConditionOperationsTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForConditionOperationsTest.java @@ -3,15 +3,22 @@ package software.amazon.lambda.durable; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; -import java.util.function.BiFunction; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; import software.amazon.lambda.durable.context.BaseContextImpl; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionOperation; +import software.amazon.lambda.durable.extension.ExtensionStepConfig; +import software.amazon.lambda.durable.extension.ExtensionStepFunction; +import software.amazon.lambda.durable.extension.ExtensionStepResult; +import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.model.WaitForConditionResult; class DurableWaitForConditionOperationsTest { @@ -22,23 +29,48 @@ void clearContext() { @Test void conditionFunctionReceivesOnlyStateAndUsesStepContextFromTls() { - var context = mock(DurableContext.class); + var context = mock(CurrentExtensionContext.class); + var reservation = mock(ExtensionOperation.class); + var future = mockStringFuture(); var stepContext = mock(StepContext.class); + when(context.reserve("condition")).thenReturn(reservation); + when(future.get()).thenReturn("VALUE"); + when(reservation.stepAsync( + eq(OperationSubType.WAIT_FOR_CONDITION.getValue()), + eq(TypeToken.get(String.class)), + any(ExtensionStepFunction.class), + any(ExtensionStepConfig.class))) + .thenReturn(future); BaseContextImpl.setCurrentContext(context); - DurableWaitForConditionOperations.waitForCondition("condition", String.class, state -> { + assertEquals("VALUE", DurableWaitForConditionOperations.waitForCondition("condition", String.class, state -> { assertEquals(stepContext, StepContext.getCurrentContext()); return WaitForConditionResult.stopPolling(state.toUpperCase()); - }); + })); - @SuppressWarnings("unchecked") - var check = (ArgumentCaptor>>) - (ArgumentCaptor) ArgumentCaptor.forClass(BiFunction.class); - verify(context).waitForCondition(eq("condition"), eq(String.class), check.capture()); + var check = extensionFunction(); + verify(reservation) + .stepAsync( + eq(OperationSubType.WAIT_FOR_CONDITION.getValue()), + eq(TypeToken.get(String.class)), + check.capture(), + any(ExtensionStepConfig.class)); try (var ignored = BaseContextImpl.attachCurrentContext(stepContext)) { - assertEquals( - WaitForConditionResult.stopPolling("VALUE"), - check.getValue().apply("value", stepContext)); + var result = + (ExtensionStepResult.Succeeded) check.getValue().apply("value"); + assertEquals("VALUE", result.value()); } } + + @SuppressWarnings({"rawtypes", "unchecked"}) + private ArgumentCaptor> extensionFunction() { + return (ArgumentCaptor) ArgumentCaptor.forClass(ExtensionStepFunction.class); + } + + @SuppressWarnings("unchecked") + private DurableFuture mockStringFuture() { + return mock(DurableFuture.class); + } + + private interface CurrentExtensionContext extends DurableContext, ExtensionContext {} } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationsTest.java index e69c8aaa4..279a66994 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationsTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationsTest.java @@ -4,15 +4,21 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; -import java.util.function.BiFunction; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; import software.amazon.lambda.durable.context.BaseContextImpl; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextFunction; +import software.amazon.lambda.durable.extension.ExtensionOperation; +import software.amazon.lambda.durable.model.OperationSubType; class DurableWithRetryOperationsTest { @AfterEach @@ -22,17 +28,44 @@ void clearContext() { @Test void retryBodyUsesSupplierAndScopedAttempt() { - var context = mock(DurableContext.class); + var context = mock(CurrentExtensionContext.class); + var parent = mock(ExtensionOperation.class); + var future = mockIntegerFuture(); + when(context.reserve("retry")).thenReturn(parent); + when(future.get()).thenReturn(1); + when(parent.runInChildContextAsync( + eq(OperationSubType.WITH_RETRY.getValue()), + any(TypeToken.class), + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) + .thenReturn(future); BaseContextImpl.setCurrentContext(context); - DurableWithRetryOperations.withRetry( - "retry", () -> WithRetryContext.getCurrentContext().getAttempt()); + assertEquals(1, DurableWithRetryOperations.withRetry("retry", () -> WithRetryContext.getCurrentContext() + .getAttempt())); - @SuppressWarnings("unchecked") - var operation = (ArgumentCaptor>) - (ArgumentCaptor) ArgumentCaptor.forClass(BiFunction.class); - verify(context).withRetry(eq("retry"), operation.capture()); - assertEquals(2, operation.getValue().apply(2, mock(DurableContext.class))); + var function = extensionFunction(); + verify(parent) + .runInChildContextAsync( + eq(OperationSubType.WITH_RETRY.getValue()), + any(TypeToken.class), + function.capture(), + any(ExtensionContextConfig.class)); + try (var ignored = BaseContextImpl.attachCurrentContext(context)) { + assertEquals(1, function.getValue().apply().result()); + } assertThrows(IllegalStateException.class, WithRetryContext::getCurrentContext); } + + @SuppressWarnings({"rawtypes", "unchecked"}) + private ArgumentCaptor> extensionFunction() { + return (ArgumentCaptor) ArgumentCaptor.forClass(ExtensionContextFunction.class); + } + + @SuppressWarnings("unchecked") + private DurableFuture mockIntegerFuture() { + return mock(DurableFuture.class); + } + + private interface CurrentExtensionContext extends DurableContext, ExtensionContext {} } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/WaitForConditionExtensionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/extension/WaitForConditionExtensionTest.java new file mode 100644 index 000000000..f81a8794c --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/context/extension/WaitForConditionExtensionTest.java @@ -0,0 +1,147 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.context.extension; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.time.Duration; +import java.util.concurrent.CompletableFuture; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import software.amazon.awssdk.services.lambda.model.ErrorObject; +import software.amazon.awssdk.services.lambda.model.Operation; +import software.amazon.awssdk.services.lambda.model.OperationStatus; +import software.amazon.awssdk.services.lambda.model.OperationType; +import software.amazon.awssdk.services.lambda.model.StepDetails; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.StepContext; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.WaitForConditionConfig; +import software.amazon.lambda.durable.context.BaseContextImpl; +import software.amazon.lambda.durable.exception.StepFailedException; +import software.amazon.lambda.durable.exception.WaitForConditionFailedException; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionOperation; +import software.amazon.lambda.durable.extension.ExtensionStepConfig; +import software.amazon.lambda.durable.extension.ExtensionStepFunction; +import software.amazon.lambda.durable.extension.ExtensionStepResult; +import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.model.WaitForConditionResult; +import software.amazon.lambda.durable.serde.JacksonSerDes; + +class WaitForConditionExtensionTest { + @AfterEach + void clearContext() { + BaseContextImpl.setCurrentContext(null); + } + + @Test + void executeMapsPollingResultsToStatefulStepOutcomes() { + var context = mock(ExtensionContext.class); + var reservation = mock(ExtensionOperation.class); + var future = mockStringFuture(); + var resultType = TypeToken.get(String.class); + var serDes = new JacksonSerDes(); + var config = WaitForConditionConfig.builder() + .initialState("initial") + .serDes(serDes) + .waitStrategy((state, attempt) -> { + assertEquals("next", state); + assertEquals(2, attempt); + return Duration.ofSeconds(7); + }) + .build(); + when(context.reserve("ready")).thenReturn(reservation); + when(reservation.stepAsync( + eq(OperationSubType.WAIT_FOR_CONDITION.getValue()), + eq(resultType), + any(ExtensionStepFunction.class), + any(ExtensionStepConfig.class))) + .thenReturn(future); + + var actual = WaitForConditionExtension.execute( + context, "ready", resultType, (state, step) -> WaitForConditionResult.continuePolling("next"), config); + + assertEquals(future.get(), actual.get()); + var function = extensionFunction(); + var extensionConfig = ArgumentCaptor.forClass(ExtensionStepConfig.class); + verify(reservation) + .stepAsync( + eq(OperationSubType.WAIT_FOR_CONDITION.getValue()), + eq(resultType), + function.capture(), + extensionConfig.capture()); + assertEquals("initial", extensionConfig.getValue().initialState()); + assertSame(serDes, extensionConfig.getValue().serDes()); + + var stepContext = mock(StepContext.class); + when(stepContext.getAttempt()).thenReturn(2); + try (var ignored = BaseContextImpl.attachCurrentContext(stepContext)) { + var retry = assertInstanceOf( + ExtensionStepResult.Retry.class, function.getValue().apply("state")); + assertEquals("next", retry.state()); + assertEquals(Duration.ofSeconds(7), retry.delay()); + } + } + + @Test + void futureTranslatesOnlyFallbackStepFailure() { + var operation = Operation.builder() + .id("operation-id") + .name("ready") + .type(OperationType.STEP) + .subType(OperationSubType.WAIT_FOR_CONDITION.getValue()) + .status(OperationStatus.FAILED) + .stepDetails(StepDetails.builder() + .error(ErrorObject.builder().errorMessage("failed").build()) + .build()) + .build(); + DurableFuture delegate = () -> { + throw new StepFailedException(operation); + }; + + var future = new WaitForConditionFuture<>(delegate); + + var failure = assertThrows(WaitForConditionFailedException.class, future::get); + assertSame(operation, failure.getOperation()); + } + + @Test + void futureDelegatesCompletionSignal() { + var completion = new CompletableFuture(); + DurableFuture delegate = new DurableFuture<>() { + @Override + public String get() { + return "done"; + } + + @Override + public CompletableFuture completionFuture() { + return completion; + } + }; + + assertSame(completion, new WaitForConditionFuture<>(delegate).completionFuture()); + } + + @SuppressWarnings({"rawtypes", "unchecked"}) + private ArgumentCaptor> extensionFunction() { + return (ArgumentCaptor) ArgumentCaptor.forClass(ExtensionStepFunction.class); + } + + @SuppressWarnings("unchecked") + private DurableFuture mockStringFuture() { + var future = (DurableFuture) mock(DurableFuture.class); + when(future.get()).thenReturn("done"); + return future; + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/WithRetryExtensionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/extension/WithRetryExtensionTest.java new file mode 100644 index 000000000..1e05e1a82 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/context/extension/WithRetryExtensionTest.java @@ -0,0 +1,104 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.context.extension; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.time.Duration; +import java.util.ArrayList; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.WithRetryConfig; +import software.amazon.lambda.durable.context.BaseContextImpl; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextFunction; +import software.amazon.lambda.durable.extension.ExtensionOperation; +import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.retry.RetryDecision; + +class WithRetryExtensionTest { + @AfterEach + void clearContext() { + BaseContextImpl.setCurrentContext(null); + } + + @Test + void executePreservesContextTopologyAndDurableBackoff() { + var context = mock(CurrentExtensionContext.class); + var parent = mock(ExtensionOperation.class); + var future = mockObjectFuture(); + when(context.reserve("transaction")).thenReturn(parent); + when(parent.runInChildContextAsync( + eq(OperationSubType.WITH_RETRY.getValue()), + any(TypeToken.class), + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) + .thenReturn(future); + var attempts = new ArrayList(); + var config = WithRetryConfig.builder() + .retryStrategy((error, attempt) -> + attempt == 1 ? RetryDecision.retry(Duration.ofSeconds(5)) : RetryDecision.fail()) + .wrapInChildContext(true) + .build(); + + var actual = WithRetryExtension.execute( + context, + "transaction", + (attempt, child) -> { + attempts.add(attempt); + assertSame(child, DurableContext.getCurrentContext()); + if (attempt == 1) { + throw new RuntimeException("retry"); + } + return "done"; + }, + config); + + assertSame(future, actual); + var function = extensionFunction(); + var contextConfig = ArgumentCaptor.forClass(ExtensionContextConfig.class); + verify(parent) + .runInChildContextAsync( + eq(OperationSubType.WITH_RETRY.getValue()), + any(TypeToken.class), + function.capture(), + contextConfig.capture()); + assertFalse(contextConfig.getValue().childContextConfig().isVirtual()); + + var child = mock(CurrentExtensionContext.class); + var wait = mock(ExtensionOperation.class); + when(child.reserve("transaction-backoff-1")).thenReturn(wait); + BaseContextImpl.setCurrentContext(child); + + var result = function.getValue().apply(); + + assertEquals("done", result.result()); + assertEquals(1, attempts.get(0)); + assertEquals(2, attempts.get(1)); + verify(wait).wait(Duration.ofSeconds(5)); + } + + @SuppressWarnings({"rawtypes", "unchecked"}) + private ArgumentCaptor> extensionFunction() { + return (ArgumentCaptor) ArgumentCaptor.forClass(ExtensionContextFunction.class); + } + + @SuppressWarnings("unchecked") + private DurableFuture mockObjectFuture() { + return mock(DurableFuture.class); + } + + private interface CurrentExtensionContext extends DurableContext, ExtensionContext {} +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/StatefulExtensionStepOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/StatefulExtensionStepOperationTest.java new file mode 100644 index 000000000..4743e769e --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/StatefulExtensionStepOperationTest.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.operation; + +import static org.junit.jupiter.api.Assertions.assertEquals; +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 java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.services.lambda.model.ErrorObject; +import software.amazon.awssdk.services.lambda.model.Operation; +import software.amazon.awssdk.services.lambda.model.OperationStatus; +import software.amazon.awssdk.services.lambda.model.OperationType; +import software.amazon.awssdk.services.lambda.model.StepDetails; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.exception.IllegalDurableOperationException; +import software.amazon.lambda.durable.exception.NonDeterministicExecutionException; +import software.amazon.lambda.durable.exception.SerDesException; +import software.amazon.lambda.durable.exception.StepFailedException; +import software.amazon.lambda.durable.execution.ExecutionManager; +import software.amazon.lambda.durable.execution.ThreadContext; +import software.amazon.lambda.durable.execution.ThreadType; +import software.amazon.lambda.durable.extension.ExtensionStepConfig; +import software.amazon.lambda.durable.extension.ExtensionStepFunction; +import software.amazon.lambda.durable.extension.ExtensionStepResult; +import software.amazon.lambda.durable.model.OperationDescriptor; +import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.serde.JacksonSerDes; + +class StatefulExtensionStepOperationTest { + private static final String OPERATION_ID = "1"; + private static final String OPERATION_NAME = "test-wait-for-condition"; + private static final JacksonSerDes SERDES = new JacksonSerDes(); + + private ExecutionManager executionManager; + private DurableContextImpl durableContext; + + @BeforeEach + void setUp() { + executionManager = mock(ExecutionManager.class); + durableContext = mock(DurableContextImpl.class); + when(durableContext.getExecutionManager()).thenReturn(executionManager); + when(executionManager.getCurrentThreadContext()).thenReturn(new ThreadContext("handler", ThreadType.CONTEXT)); + when(durableContext.getDurableConfig()) + .thenReturn(DurableConfig.builder() + .withExecutorService(Executors.newCachedThreadPool()) + .build()); + } + + @Test + void replaySucceededReturnsCachedResultWithoutCallingFunction() { + var existing = operation( + OperationStatus.SUCCEEDED, StepDetails.builder().result("42").build()); + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(existing); + var called = new CountDownLatch(1); + var operation = createOperation(state -> { + called.countDown(); + return ExtensionStepResult.succeed(state); + }); + + operation.execute(); + + assertEquals(42, operation.get()); + assertEquals(1, called.getCount()); + } + + @Test + void replayFailedThrowsOriginalException() { + var original = new IllegalArgumentException("bad state"); + var error = ErrorObject.builder() + .errorType(IllegalArgumentException.class.getName()) + .errorMessage("bad state") + .errorData(SERDES.serialize(original)) + .stackTrace(List.of("com.example.Test|method|Test.java|42")) + .build(); + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)) + .thenReturn(operation( + OperationStatus.FAILED, + StepDetails.builder().error(error).build())); + + var operation = createOperation(ExtensionStepResult::succeed); + operation.execute(); + + assertEquals( + "bad state", + assertThrows(IllegalArgumentException.class, operation::get).getMessage()); + } + + @Test + void replayFailedFallsBackToStepFailedException() { + var error = ErrorObject.builder() + .errorType("com.example.MissingException") + .errorMessage("failed") + .stackTrace(List.of("com.example.Test|method|Test.java|1")) + .build(); + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)) + .thenReturn(operation( + OperationStatus.FAILED, + StepDetails.builder().error(error).build())); + + var operation = createOperation(ExtensionStepResult::succeed); + operation.execute(); + + assertThrows(StepFailedException.class, operation::get); + } + + @Test + void replayStartedAndReadyResumeWithCheckpointedState() throws Exception { + assertResumes(OperationStatus.STARTED, 10); + assertResumes(OperationStatus.READY, 5); + } + + @Test + void replayPendingPollsUntilReadyAndResumes() throws Exception { + var pending = operation( + OperationStatus.PENDING, + StepDetails.builder().attempt(1).result("5").build()); + var ready = operation( + OperationStatus.READY, + StepDetails.builder().attempt(1).result("5").build()); + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(pending); + when(executionManager.pollForOperationUpdates(OPERATION_ID)) + .thenReturn(CompletableFuture.completedFuture(ready)); + var called = new CountDownLatch(1); + var operation = createOperation(state -> { + called.countDown(); + return ExtensionStepResult.succeed(state); + }); + + operation.execute(); + + assertTrue(called.await(2, TimeUnit.SECONDS)); + } + + @Test + void replayWithoutCheckpointStateUsesInitialState() throws Exception { + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)) + .thenReturn(operation( + OperationStatus.STARTED, + StepDetails.builder().attempt(0).build())); + var called = new CountDownLatch(1); + var operation = createOperation( + state -> { + assertEquals(42, state); + called.countDown(); + return ExtensionStepResult.succeed(state); + }, + 42); + + operation.execute(); + + assertTrue(called.await(2, TimeUnit.SECONDS)); + } + + @Test + void corruptReplayStateFailsBeforeCallingFunction() { + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)) + .thenReturn(operation( + OperationStatus.STARTED, + StepDetails.builder() + .attempt(1) + .result("not-valid-json!!!") + .build())); + var called = new CountDownLatch(1); + var operation = createOperation(state -> { + called.countDown(); + return ExtensionStepResult.succeed(state); + }); + + assertThrows(SerDesException.class, operation::execute); + assertEquals(1, called.getCount()); + } + + @Test + void replayStillValidatesTypeNameAndStatus() { + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)) + .thenReturn(Operation.builder() + .id(OPERATION_ID) + .name(OPERATION_NAME) + .type(OperationType.WAIT) + .status(OperationStatus.SUCCEEDED) + .build()); + assertThrows(NonDeterministicExecutionException.class, () -> createOperation(ExtensionStepResult::succeed) + .execute()); + + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)) + .thenReturn(Operation.builder() + .id(OPERATION_ID) + .name("different") + .type(OperationType.STEP) + .status(OperationStatus.SUCCEEDED) + .build()); + assertThrows(NonDeterministicExecutionException.class, () -> createOperation(ExtensionStepResult::succeed) + .execute()); + + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)) + .thenReturn(operation(OperationStatus.UNKNOWN_TO_SDK_VERSION, null)); + assertThrows(IllegalDurableOperationException.class, () -> createOperation(ExtensionStepResult::succeed) + .execute()); + } + + private void assertResumes(OperationStatus status, int expectedState) throws Exception { + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)) + .thenReturn(operation( + status, + StepDetails.builder() + .attempt(1) + .result(String.valueOf(expectedState)) + .build())); + var called = new CountDownLatch(1); + var operation = createOperation(state -> { + assertEquals(expectedState, state); + called.countDown(); + return ExtensionStepResult.succeed(state); + }); + + operation.execute(); + + assertTrue(called.await(2, TimeUnit.SECONDS)); + } + + private StepOperation createOperation(ExtensionStepFunction function) { + return createOperation(function, null); + } + + private StepOperation createOperation(ExtensionStepFunction function, Integer initialState) { + return new StepOperation<>( + new OperationDescriptor( + OPERATION_ID, + OPERATION_NAME, + OperationType.STEP, + OperationSubType.WAIT_FOR_CONDITION.getValue()), + function, + TypeToken.get(Integer.class), + ExtensionStepConfig.builder() + .initialState(initialState) + .serDes(SERDES) + .build(), + durableContext); + } + + private Operation operation(OperationStatus status, StepDetails details) { + return Operation.builder() + .id(OPERATION_ID) + .name(OPERATION_NAME) + .type(OperationType.STEP) + .subType(OperationSubType.WAIT_FOR_CONDITION.getValue()) + .status(status) + .stepDetails(details) + .build(); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionOperationTest.java deleted file mode 100644 index 69502a3c3..000000000 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionOperationTest.java +++ /dev/null @@ -1,392 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; - -import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.Mockito.*; - -import java.util.List; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.Executors; -import java.util.concurrent.atomic.AtomicBoolean; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import software.amazon.awssdk.services.lambda.model.ErrorObject; -import software.amazon.awssdk.services.lambda.model.Operation; -import software.amazon.awssdk.services.lambda.model.OperationStatus; -import software.amazon.awssdk.services.lambda.model.OperationType; -import software.amazon.awssdk.services.lambda.model.StepDetails; -import software.amazon.lambda.durable.DurableConfig; -import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.WaitForConditionConfig; -import software.amazon.lambda.durable.context.DurableContextImpl; -import software.amazon.lambda.durable.exception.IllegalDurableOperationException; -import software.amazon.lambda.durable.exception.NonDeterministicExecutionException; -import software.amazon.lambda.durable.exception.SerDesException; -import software.amazon.lambda.durable.exception.WaitForConditionFailedException; -import software.amazon.lambda.durable.execution.ExecutionManager; -import software.amazon.lambda.durable.execution.ThreadContext; -import software.amazon.lambda.durable.execution.ThreadType; -import software.amazon.lambda.durable.model.OperationIdentifier; -import software.amazon.lambda.durable.model.OperationSubType; -import software.amazon.lambda.durable.model.WaitForConditionResult; -import software.amazon.lambda.durable.serde.JacksonSerDes; - -class WaitForConditionOperationTest { - - private static final String OPERATION_ID = "1"; - private static final String OPERATION_NAME = "test-wait-for-condition"; - private static final JacksonSerDes SERDES = new JacksonSerDes(); - - private ExecutionManager executionManager; - private DurableContextImpl durableContext; - - @BeforeEach - void setUp() { - executionManager = mock(ExecutionManager.class); - durableContext = mock(DurableContextImpl.class); - when(durableContext.getExecutionManager()).thenReturn(executionManager); - when(executionManager.getCurrentThreadContext()).thenReturn(new ThreadContext("handler", ThreadType.CONTEXT)); - when(durableContext.getDurableConfig()) - .thenReturn(DurableConfig.builder() - .withExecutorService(Executors.newCachedThreadPool()) - .build()); - } - - private WaitForConditionOperation createOperation( - java.util.function.BiFunction< - Integer, software.amazon.lambda.durable.StepContext, WaitForConditionResult> - checkFunc, - WaitForConditionConfig config) { - return new WaitForConditionOperation<>( - OperationIdentifier.of(OPERATION_ID, OPERATION_NAME, OperationSubType.WAIT_FOR_CONDITION), - checkFunc, - TypeToken.get(Integer.class), - config, - durableContext); - } - - // ===== Replay SUCCEEDED ===== - - @Test - void replaySucceededReturnsCachedResult() { - var op = Operation.builder() - .id(OPERATION_ID) - .name(OPERATION_NAME) - .type(OperationType.STEP) - .subType("WaitForCondition") - .status(OperationStatus.SUCCEEDED) - .stepDetails(StepDetails.builder().result("42").build()) - .build(); - when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); - - var functionCalled = new AtomicBoolean(false); - var config = WaitForConditionConfig.builder().serDes(SERDES).build(); - var operation = createOperation( - (state, ctx) -> { - functionCalled.set(true); - return WaitForConditionResult.stopPolling(state); - }, - config); - - operation.execute(); - - var result = operation.get(); - assertEquals(42, result); - assertFalse(functionCalled.get(), "Check function should not be called during SUCCEEDED replay"); - } - - // ===== Replay FAILED ===== - - @Test - void replayFailedThrowsOriginalException() { - var originalException = new IllegalArgumentException("bad state"); - var stackTrace = List.of("com.example.Test|method|Test.java|42"); - - var op = Operation.builder() - .id(OPERATION_ID) - .name(OPERATION_NAME) - .type(OperationType.STEP) - .subType("WaitForCondition") - .status(OperationStatus.FAILED) - .stepDetails(StepDetails.builder() - .error(ErrorObject.builder() - .errorType("java.lang.IllegalArgumentException") - .errorMessage("bad state") - .errorData(SERDES.serialize(originalException)) - .stackTrace(stackTrace) - .build()) - .build()) - .build(); - when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); - - var config = WaitForConditionConfig.builder().serDes(SERDES).build(); - var operation = createOperation((state, ctx) -> WaitForConditionResult.stopPolling(state), config); - - operation.execute(); - - var thrown = assertThrows(IllegalArgumentException.class, operation::get); - assertEquals("bad state", thrown.getMessage()); - } - - @Test - void replayFailedFallsBackToStepFailedException() { - var op = Operation.builder() - .id(OPERATION_ID) - .name(OPERATION_NAME) - .type(OperationType.STEP) - .subType("WaitForCondition") - .status(OperationStatus.FAILED) - .stepDetails(StepDetails.builder() - .error(ErrorObject.builder() - .errorType("com.nonexistent.SomeException") - .errorMessage("unknown error") - .stackTrace(List.of("com.example.Test|method|Test.java|1")) - .build()) - .build()) - .build(); - when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); - - var config = WaitForConditionConfig.builder().serDes(SERDES).build(); - var operation = createOperation((state, ctx) -> WaitForConditionResult.stopPolling(state), config); - - operation.execute(); - - assertThrows(WaitForConditionFailedException.class, operation::get); - } - - // ===== Replay STARTED ===== - - @Test - void replayStartedResumesCheckLoop() throws Exception { - var op = Operation.builder() - .id(OPERATION_ID) - .name(OPERATION_NAME) - .type(OperationType.STEP) - .subType("WaitForCondition") - .status(OperationStatus.STARTED) - .stepDetails(StepDetails.builder().attempt(2).result("10").build()) - .build(); - when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); - - var functionCalled = new AtomicBoolean(false); - var config = WaitForConditionConfig.builder().serDes(SERDES).build(); - var operation = createOperation( - (state, ctx) -> { - functionCalled.set(true); - return WaitForConditionResult.stopPolling(state + 1); - }, - config); - - operation.execute(); - - // Give the executor thread time to run - Thread.sleep(200); - assertTrue(functionCalled.get(), "Check function should be re-executed for STARTED replay"); - } - - // ===== Replay READY ===== - - @Test - void replayReadyResumesCheckLoop() throws Exception { - var op = Operation.builder() - .id(OPERATION_ID) - .name(OPERATION_NAME) - .type(OperationType.STEP) - .subType("WaitForCondition") - .status(OperationStatus.READY) - .stepDetails(StepDetails.builder().attempt(1).result("5").build()) - .build(); - when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); - - var functionCalled = new AtomicBoolean(false); - var config = WaitForConditionConfig.builder().serDes(SERDES).build(); - var operation = createOperation( - (state, ctx) -> { - functionCalled.set(true); - return WaitForConditionResult.stopPolling(state); - }, - config); - - operation.execute(); - - Thread.sleep(200); - assertTrue(functionCalled.get(), "Check function should be re-executed for READY replay"); - } - - // ===== Non-deterministic detection ===== - - @Test - void replayWithTypeMismatchTerminatesExecution() { - when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)) - .thenReturn(Operation.builder() - .id(OPERATION_ID) - .name(OPERATION_NAME) - .type(OperationType.WAIT) // Wrong type — should be STEP - .status(OperationStatus.SUCCEEDED) - .build()); - - var config = WaitForConditionConfig.builder().serDes(SERDES).build(); - var operation = createOperation((state, ctx) -> WaitForConditionResult.stopPolling(state), config); - - assertThrows(NonDeterministicExecutionException.class, operation::execute); - } - - @Test - void replayWithNameMismatchTerminatesExecution() { - when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)) - .thenReturn(Operation.builder() - .id(OPERATION_ID) - .name("different-name") - .type(OperationType.STEP) - .status(OperationStatus.SUCCEEDED) - .build()); - - var config = WaitForConditionConfig.builder().serDes(SERDES).build(); - var operation = createOperation((state, ctx) -> WaitForConditionResult.stopPolling(state), config); - - assertThrows(NonDeterministicExecutionException.class, operation::execute); - } - - // ===== get() with null error data ===== - - @Test - void getFailedWithNullErrorDataThrowsStepFailedException() { - var op = Operation.builder() - .id(OPERATION_ID) - .name(OPERATION_NAME) - .type(OperationType.STEP) - .subType("WaitForCondition") - .status(OperationStatus.FAILED) - .stepDetails(StepDetails.builder() - .error(ErrorObject.builder() - .errorType(RuntimeException.class.getName()) - .errorMessage("Something went wrong") - .stackTrace(List.of("com.example.Test|method|Test.java|42")) - .build()) - .build()) - .build(); - when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); - - var config = WaitForConditionConfig.builder().serDes(SERDES).build(); - var operation = createOperation((state, ctx) -> WaitForConditionResult.stopPolling(state), config); - - operation.execute(); - - assertThrows(WaitForConditionFailedException.class, operation::get); - } - - // ===== Replay PENDING ===== - - @Test - void replayPendingPollsAndResumesCheckLoop() throws Exception { - var pendingOp = Operation.builder() - .id(OPERATION_ID) - .name(OPERATION_NAME) - .type(OperationType.STEP) - .subType("WaitForCondition") - .status(OperationStatus.PENDING) - .stepDetails(StepDetails.builder().attempt(1).result("5").build()) - .build(); - when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(pendingOp); - - var readyOp = Operation.builder() - .id(OPERATION_ID) - .name(OPERATION_NAME) - .type(OperationType.STEP) - .subType("WaitForCondition") - .status(OperationStatus.READY) - .stepDetails(StepDetails.builder().attempt(1).result("5").build()) - .build(); - when(executionManager.pollForOperationUpdates(OPERATION_ID)) - .thenReturn(CompletableFuture.completedFuture(readyOp)); - - var functionCalled = new AtomicBoolean(false); - var config = WaitForConditionConfig.builder().serDes(SERDES).build(); - var operation = createOperation( - (state, ctx) -> { - functionCalled.set(true); - return WaitForConditionResult.stopPolling(state); - }, - config); - - operation.execute(); - - Thread.sleep(200); - assertTrue(functionCalled.get(), "Check function should be called after PENDING → READY transition"); - } - - // ===== Replay unexpected status ===== - - @Test - void replayWithUnexpectedStatusTerminatesExecution() { - when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)) - .thenReturn(Operation.builder() - .id(OPERATION_ID) - .name(OPERATION_NAME) - .type(OperationType.STEP) - .subType("WaitForCondition") - .status(OperationStatus.UNKNOWN_TO_SDK_VERSION) - .build()); - - var config = WaitForConditionConfig.builder().serDes(SERDES).build(); - var operation = createOperation((state, ctx) -> WaitForConditionResult.stopPolling(state), config); - - assertThrows(IllegalDurableOperationException.class, operation::execute); - } - - // ===== resumeCheckLoop with null checkpoint data ===== - - @Test - void replayStartedWithNullCheckpointDataUsesInitialState() throws Exception { - var op = Operation.builder() - .id(OPERATION_ID) - .name(OPERATION_NAME) - .type(OperationType.STEP) - .subType("WaitForCondition") - .status(OperationStatus.STARTED) - .stepDetails(StepDetails.builder().attempt(0).build()) // no result set - .build(); - when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); - - var receivedState = new java.util.concurrent.atomic.AtomicInteger(-1); - var config = WaitForConditionConfig.builder() - .serDes(SERDES) - .initialState(42) - .build(); - var operation = createOperation( - (state, ctx) -> { - receivedState.set(state); - return WaitForConditionResult.stopPolling(state); - }, - config); - - operation.execute(); - - Thread.sleep(200); - assertEquals(42, receivedState.get(), "Should use initialState when checkpoint data is null"); - } - - // ===== resumeCheckLoop checkpoint deserialize exception ===== - - @Test - void replayStartedWithCorruptCheckpointDataThrowsSerDesException() { - var op = Operation.builder() - .id(OPERATION_ID) - .name(OPERATION_NAME) - .type(OperationType.STEP) - .subType("WaitForCondition") - .status(OperationStatus.STARTED) - .stepDetails(StepDetails.builder() - .attempt(1) - .result("not-valid-json!!!") - .build()) - .build(); - when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); - - var config = WaitForConditionConfig.builder().serDes(SERDES).build(); - var operation = createOperation((state, ctx) -> WaitForConditionResult.stopPolling(state), config); - - assertThrows(SerDesException.class, operation::execute); - } -} From f126e4c8b6c349bfe91ed210d05158213f7348b3 Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 04:15:09 +0000 Subject: [PATCH 20/40] feat: add extension concurrency coordination --- .../ExtensionConcurrencyIntegrationTest.java | 30 +++ .../amazon/lambda/durable/DurableFuture.java | 13 +- .../extension/DeferredDurableFuture.java | 54 ++++ .../ExtensionConcurrencyCoordinator.java | 235 ++++++++++++++++++ .../durable/execution/ExecutionManager.java | 41 +++ .../lambda/durable/DurableFutureTest.java | 24 ++ .../extension/DeferredDurableFutureTest.java | 78 ++++++ .../ExtensionConcurrencyCoordinatorTest.java | 180 ++++++++++++++ .../execution/ExecutionManagerTest.java | 44 +++- 9 files changed, 694 insertions(+), 5 deletions(-) create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionConcurrencyIntegrationTest.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/context/extension/DeferredDurableFuture.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinator.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/context/extension/DeferredDurableFutureTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinatorTest.java diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionConcurrencyIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionConcurrencyIntegrationTest.java new file mode 100644 index 000000000..4730790bb --- /dev/null +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionConcurrencyIntegrationTest.java @@ -0,0 +1,30 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class ExtensionConcurrencyIntegrationTest { + @Test + void anyOfSuspendsWhileCallbacksArePending() { + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { + var first = context.createCallback("first", String.class); + var second = context.createCallback("second", String.class); + return (String) DurableFuture.anyOf(first, second); + }); + + var pending = runner.run("input"); + + assertEquals(ExecutionStatus.PENDING, pending.getStatus()); + + runner.completeCallback(runner.getCallbackId("second"), "\"second-result\""); + var completed = runner.run("input"); + + assertEquals(ExecutionStatus.SUCCEEDED, completed.getStatus()); + assertEquals("second-result", completed.getResult(String.class)); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableFuture.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableFuture.java index d297da318..1880e3439 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableFuture.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableFuture.java @@ -5,6 +5,8 @@ import java.util.Arrays; import java.util.List; import java.util.concurrent.CompletableFuture; +import software.amazon.lambda.durable.context.BaseContext; +import software.amazon.lambda.durable.context.BaseContextImpl; /** * A future representing the result of an asynchronous durable operation. @@ -73,11 +75,14 @@ static List allOf(List> futures) { * @return the result of the first future to complete */ static Object anyOf(DurableFuture... futures) { - return CompletableFuture.anyOf(Arrays.stream(futures) + var firstCompleted = CompletableFuture.anyOf(Arrays.stream(futures) .map(f -> f.completionFuture().thenApply(ignored -> f)) .toArray(CompletableFuture[]::new)) - .thenApply(o -> (DurableFuture) o) - .join() - .get(); + .thenApply(o -> (DurableFuture) o); + var context = BaseContext.getCurrentContext(); + var future = context instanceof BaseContextImpl contextImpl + ? contextImpl.getExecutionManager().awaitFuture(firstCompleted) + : firstCompleted.join(); + return future.get(); } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/DeferredDurableFuture.java b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/DeferredDurableFuture.java new file mode 100644 index 000000000..43bea59ab --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/DeferredDurableFuture.java @@ -0,0 +1,54 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.context.extension; + +import java.util.Objects; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicBoolean; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.context.BaseContext; +import software.amazon.lambda.durable.context.BaseContextImpl; + +final class DeferredDurableFuture implements DurableFuture { + private final AtomicBoolean bound = new AtomicBoolean(); + private final CompletableFuture> delegateFuture = new CompletableFuture<>(); + private final CompletableFuture completionSignal = new CompletableFuture<>(); + + void bind(DurableFuture delegate) { + Objects.requireNonNull(delegate, "delegate cannot be null"); + if (!bound.compareAndSet(false, true)) { + throw new IllegalStateException("A deferred durable future can only be bound once"); + } + + delegateFuture.complete(delegate); + delegate.completionFuture().whenComplete((ignored, throwable) -> { + if (throwable == null) { + completionSignal.complete(null); + } else { + completionSignal.completeExceptionally(throwable); + } + }); + } + + @Override + public T get() { + return awaitDelegate().get(); + } + + @Override + public CompletableFuture completionFuture() { + return completionSignal.thenApply(ignored -> null); + } + + boolean isDone() { + return completionSignal.isDone(); + } + + private DurableFuture awaitDelegate() { + var context = BaseContext.getCurrentContext(); + if (context instanceof BaseContextImpl contextImpl) { + return contextImpl.getExecutionManager().awaitFuture(delegateFuture); + } + return delegateFuture.join(); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinator.java b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinator.java new file mode 100644 index 000000000..f16775fd3 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinator.java @@ -0,0 +1,235 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.context.extension; + +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Objects; +import java.util.Queue; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.function.Function; +import java.util.function.Supplier; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.config.CompletionConfig; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.execution.SuspendExecutionException; + +final class ExtensionConcurrencyCoordinator { + enum ItemStatus { + PENDING, + RUNNING, + SUCCEEDED, + FAILED, + SKIPPED + } + + record ExpectedCompletionStatus(int completed, CompletionConfig.CompletionDecision completionDecision) { + ExpectedCompletionStatus { + if (completed < 0) { + throw new IllegalArgumentException("completed cannot be negative"); + } + Objects.requireNonNull(completionDecision, "completionDecision cannot be null"); + } + } + + record Completion(CompletionConfig.CompletionDecision completionDecision, List> items) { + Completion { + Objects.requireNonNull(completionDecision, "completionDecision cannot be null"); + items = List.copyOf(items); + } + } + + static final class Item { + private final Supplier> launcher; + private final DeferredDurableFuture future = new DeferredDurableFuture<>(); + private volatile ItemStatus status; + + private Item(Supplier> launcher, ItemStatus status) { + this.launcher = launcher; + this.status = status; + } + + DurableFuture future() { + return future; + } + + ItemStatus status() { + return status; + } + } + + private final Object lock = new Object(); + private final int maxConcurrency; + private final Function shouldComplete; + private final List> items = new ArrayList<>(); + private final Queue> pending = new ArrayDeque<>(); + private final Set> running = new LinkedHashSet<>(); + private CompletableFuture changed = new CompletableFuture<>(); + private boolean registrationClosed; + private int succeeded; + private int failed; + + ExtensionConcurrencyCoordinator(int maxConcurrency, CompletionConfig completionConfig) { + if (maxConcurrency < 1) { + throw new IllegalArgumentException("maxConcurrency must be at least 1"); + } + this.maxConcurrency = maxConcurrency; + this.shouldComplete = Objects.requireNonNull(completionConfig, "completionConfig cannot be null") + .completionDecisionFunction(); + } + + Item register(Supplier> launcher) { + return register(launcher, false); + } + + Item register(Supplier> launcher, boolean skipped) { + Objects.requireNonNull(launcher, "launcher cannot be null"); + synchronized (lock) { + if (registrationClosed) { + throw new IllegalStateException("Cannot register items after registration is closed"); + } + var item = new Item<>(launcher, skipped ? ItemStatus.SKIPPED : ItemStatus.PENDING); + items.add(item); + if (!skipped) { + pending.add(item); + } + notifyChanged(); + return item; + } + } + + void closeRegistration() { + synchronized (lock) { + registrationClosed = true; + notifyChanged(); + } + } + + Completion awaitCompletion() { + return awaitCompletion(null); + } + + Completion awaitCompletion(ExpectedCompletionStatus expectedCompletionStatus) { + while (true) { + DurableFuture[] waiters; + synchronized (lock) { + collectCompletedItems(); + var decision = completionDecision(expectedCompletionStatus); + if (decision != null) { + markIncompleteItemsSkipped(); + return new Completion(decision, items); + } + + launchPendingItems(); + collectCompletedItems(); + decision = completionDecision(expectedCompletionStatus); + if (decision != null) { + markIncompleteItemsSkipped(); + return new Completion(decision, items); + } + waiters = completionWaiters(); + } + DurableFuture.anyOf(waiters); + } + } + + private void launchPendingItems() { + while (running.size() < maxConcurrency && !pending.isEmpty()) { + var item = pending.remove(); + launch(item); + running.add(item); + } + } + + @SuppressWarnings("unchecked") + private void launch(Item untypedItem) { + var item = (Item) untypedItem; + var delegate = Objects.requireNonNull(item.launcher.get(), "launcher cannot return null"); + item.future.bind(delegate); + item.status = ItemStatus.RUNNING; + } + + private void collectCompletedItems() { + var completed = running.stream().filter(item -> item.future.isDone()).toList(); + for (var item : completed) { + running.remove(item); + complete(item); + } + } + + private void complete(Item item) { + try { + item.future.get(); + item.status = ItemStatus.SUCCEEDED; + succeeded++; + } catch (SuspendExecutionException | UnrecoverableDurableExecutionException exception) { + throw exception; + } catch (Throwable throwable) { + item.status = ItemStatus.FAILED; + failed++; + } + } + + private CompletionConfig.CompletionDecision completionDecision(ExpectedCompletionStatus expectedCompletionStatus) { + if (expectedCompletionStatus != null) { + return succeeded + failed >= expectedCompletionStatus.completed() + ? expectedCompletionStatus.completionDecision() + : null; + } + var status = new CompletionConfig.CompletionStatus( + succeeded, failed, succeeded + failed, items.size(), registrationClosed); + var decision = Objects.requireNonNull( + shouldComplete.apply(status), "shouldComplete must return a completion decision"); + return decision.shouldComplete() ? decision : null; + } + + private DurableFuture[] completionWaiters() { + if (changed.isDone()) { + changed = new CompletableFuture<>(); + } + var waiters = new ArrayList>(); + running.stream().map(item -> new CompletionOnlyFuture(item.future)).forEach(waiters::add); + waiters.add(new SignalFuture(changed)); + return waiters.toArray(DurableFuture[]::new); + } + + private void markIncompleteItemsSkipped() { + items.stream() + .filter(item -> item.status == ItemStatus.PENDING || item.status == ItemStatus.RUNNING) + .forEach(item -> item.status = ItemStatus.SKIPPED); + pending.clear(); + running.clear(); + } + + private void notifyChanged() { + changed.complete(null); + } + + private record CompletionOnlyFuture(DurableFuture delegate) implements DurableFuture { + @Override + public Void get() { + return null; + } + + @Override + public CompletableFuture completionFuture() { + return delegate.completionFuture(); + } + } + + private record SignalFuture(CompletableFuture signal) implements DurableFuture { + @Override + public Void get() { + signal.join(); + return null; + } + + @Override + public CompletableFuture completionFuture() { + return signal.thenApply(ignored -> null); + } + } +} 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 1c45cb0d6..788c7a9ae 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 @@ -29,6 +29,7 @@ import software.amazon.lambda.durable.model.SafeCloseable; import software.amazon.lambda.durable.operation.BaseDurableOperation; import software.amazon.lambda.durable.plugin.PluginInfoConverter; +import software.amazon.lambda.durable.util.ExceptionHelper; /** * Central manager for durable execution coordination. @@ -228,6 +229,46 @@ public ThreadContext getCurrentThreadContext() { return currentThreadContext.get(); } + /** + * Waits for a future from the current durable context thread. + * + *

The current thread is marked inactive while waiting so the execution can suspend. It is reactivated + * synchronously when the future completes. + * + * @param future the future to wait for + * @param the future result type + * @return the completed result + */ + public T awaitFuture(CompletableFuture future) { + var threadContext = getCurrentThreadContext(); + CompletableFuture awaitedFuture = future; + + if (threadContext != null && !future.isDone()) { + var coordinationLock = new Object(); + var deregistered = new boolean[1]; + synchronized (coordinationLock) { + awaitedFuture = future.whenComplete((ignored, throwable) -> { + synchronized (coordinationLock) { + if (deregistered[0] && !isExecutionCompletedExceptionally()) { + registerActiveThread(threadContext.threadId()); + } + } + }); + if (!future.isDone()) { + deregistered[0] = true; + deregisterActiveThread(threadContext.threadId()); + } + } + } + + try { + return awaitedFuture.join(); + } catch (Throwable throwable) { + ExceptionHelper.sneakyThrow(ExceptionHelper.unwrapCompletableFuture(throwable)); + return null; + } + } + /** * Registers a thread as active. * diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java index 0f3bcab74..2ccc1bbea 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java @@ -7,10 +7,17 @@ import java.util.List; import java.util.concurrent.CompletableFuture; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.context.BaseContextImpl; +import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.operation.SerializableDurableOperation; class DurableFutureTest { + @AfterEach + void clearContext() { + BaseContextImpl.setCurrentContext(null); + } @Test void allOfVarargsReturnsResultsInOrder() { @@ -81,6 +88,23 @@ void anyOfSupportsPublicDurableFutureImplementations() { assertEquals("completed", result); } + @Test + void anyOfUsesExecutionManagerWhenCalledFromDurableContext() { + var context = mock(BaseContextImpl.class); + var executionManager = mock(ExecutionManager.class); + when(context.getExecutionManager()).thenReturn(executionManager); + when(executionManager.awaitFuture(any())).thenAnswer(invocation -> { + CompletableFuture future = invocation.getArgument(0); + return future.join(); + }); + BaseContextImpl.setCurrentContext(context); + var completed = new TestFuture<>("completed"); + completed.complete(); + + assertEquals("completed", DurableFuture.anyOf(completed)); + verify(executionManager).awaitFuture(any()); + } + @SuppressWarnings("unchecked") private SerializableDurableOperation mockOperation(T result) { SerializableDurableOperation op = mock(SerializableDurableOperation.class); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/DeferredDurableFutureTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/extension/DeferredDurableFutureTest.java new file mode 100644 index 000000000..3fc7e51bc --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/context/extension/DeferredDurableFutureTest.java @@ -0,0 +1,78 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.context.extension; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.DurableFuture; + +class DeferredDurableFutureTest { + @Test + void getWaitsForBindingThenDelegates() throws Exception { + var deferred = new DeferredDurableFuture(); + var getStarted = new CountDownLatch(1); + var result = CompletableFuture.supplyAsync(() -> { + getStarted.countDown(); + return deferred.get(); + }); + getStarted.await(); + + assertFalse(result.isDone()); + + deferred.bind(new TestFuture<>("result")); + + assertEquals("result", result.join()); + } + + @Test + void completionSignalObtainedBeforeBindingTracksDelegate() { + var deferred = new DeferredDurableFuture(); + var completion = deferred.completionFuture(); + var delegate = new TestFuture<>("result"); + + deferred.bind(delegate); + assertFalse(completion.isDone()); + + delegate.complete(); + + completion.join(); + } + + @Test + void bindRejectsASecondDelegate() { + var deferred = new DeferredDurableFuture(); + deferred.bind(new TestFuture<>("first")); + + var exception = assertThrows(IllegalStateException.class, () -> deferred.bind(new TestFuture<>("second"))); + + assertEquals("A deferred durable future can only be bound once", exception.getMessage()); + } + + private static final class TestFuture implements DurableFuture { + private final T result; + private final CompletableFuture completion = new CompletableFuture<>(); + + private TestFuture(T result) { + this.result = result; + } + + @Override + public T get() { + return result; + } + + @Override + public CompletableFuture completionFuture() { + return completion; + } + + private void complete() { + completion.complete(null); + } + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinatorTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinatorTest.java new file mode 100644 index 000000000..60ae85499 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinatorTest.java @@ -0,0 +1,180 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.context.extension; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static software.amazon.lambda.durable.context.extension.ExtensionConcurrencyCoordinator.ItemStatus.FAILED; +import static software.amazon.lambda.durable.context.extension.ExtensionConcurrencyCoordinator.ItemStatus.SKIPPED; +import static software.amazon.lambda.durable.context.extension.ExtensionConcurrencyCoordinator.ItemStatus.SUCCEEDED; +import static software.amazon.lambda.durable.model.ConcurrencyCompletionStatus.ALL_COMPLETED; +import static software.amazon.lambda.durable.model.ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED; + +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.config.CompletionConfig; + +class ExtensionConcurrencyCoordinatorTest { + @Test + void launchesNoMoreThanMaxConcurrency() throws Exception { + var coordinator = new ExtensionConcurrencyCoordinator(2, CompletionConfig.allCompleted()); + var first = new TestFuture<>("first"); + var second = new TestFuture<>("second"); + var third = new TestFuture<>("third"); + var launched = new AtomicInteger(); + var firstTwoLaunched = new CountDownLatch(2); + var thirdLaunched = new CountDownLatch(1); + coordinator.register(() -> launch(first, launched, firstTwoLaunched)); + coordinator.register(() -> launch(second, launched, firstTwoLaunched)); + coordinator.register(() -> launch(third, launched, thirdLaunched)); + coordinator.closeRegistration(); + + var result = CompletableFuture.supplyAsync(coordinator::awaitCompletion); + firstTwoLaunched.await(); + + assertEquals(2, launched.get()); + + first.complete(); + thirdLaunched.await(); + second.complete(); + third.complete(); + + assertEquals(ALL_COMPLETED, result.join().completionDecision().completionStatus()); + } + + @Test + void earlyCompletionMarksUnlaunchedItemsSkipped() { + var coordinator = new ExtensionConcurrencyCoordinator(1, CompletionConfig.minSuccessful(1)); + var first = new TestFuture<>("first"); + var launched = new AtomicInteger(); + coordinator.register(() -> { + launched.incrementAndGet(); + return first; + }); + coordinator.register(() -> { + launched.incrementAndGet(); + return new TestFuture<>("second"); + }); + coordinator.register(() -> { + launched.incrementAndGet(); + return new TestFuture<>("third"); + }); + coordinator.closeRegistration(); + + var result = CompletableFuture.supplyAsync(coordinator::awaitCompletion); + first.complete(); + var completion = result.join(); + + assertEquals(MIN_SUCCESSFUL_REACHED, completion.completionDecision().completionStatus()); + assertEquals(1, launched.get()); + assertEquals( + List.of(SUCCEEDED, SKIPPED, SKIPPED), + completion.items().stream() + .map(ExtensionConcurrencyCoordinator.Item::status) + .toList()); + } + + @Test + void failedItemsContributeToCompletionStatus() { + var coordinator = new ExtensionConcurrencyCoordinator(2, CompletionConfig.allCompleted()); + var failed = new TestFuture(new IllegalStateException("failed")); + var succeeded = new TestFuture<>("succeeded"); + coordinator.register(() -> failed); + coordinator.register(() -> succeeded); + coordinator.closeRegistration(); + + var result = CompletableFuture.supplyAsync(coordinator::awaitCompletion); + failed.complete(); + succeeded.complete(); + var completion = result.join(); + + assertEquals(ALL_COMPLETED, completion.completionDecision().completionStatus()); + assertEquals( + List.of(FAILED, SUCCEEDED), + completion.items().stream() + .map(ExtensionConcurrencyCoordinator.Item::status) + .toList()); + } + + @Test + void dynamicRegistrationDoesNotCompleteUntilRegistrationCloses() throws Exception { + var completedBeforeClose = new CountDownLatch(1); + var completionConfig = CompletionConfig.shouldComplete(status -> { + if (status.completedCount() == 1 && !status.allItemsRegistered()) { + completedBeforeClose.countDown(); + } + return status.allCompleted() + ? CompletionConfig.CompletionDecision.complete(ALL_COMPLETED) + : CompletionConfig.CompletionDecision.continueExecution(); + }); + var coordinator = new ExtensionConcurrencyCoordinator(1, completionConfig); + var result = CompletableFuture.supplyAsync(coordinator::awaitCompletion); + var item = new TestFuture<>("result"); + + coordinator.register(() -> item); + item.complete(); + completedBeforeClose.await(); + + assertFalse(result.isDone()); + + coordinator.closeRegistration(); + + assertEquals(ALL_COMPLETED, result.join().completionDecision().completionStatus()); + } + + @Test + void registrationAfterCloseFails() { + var coordinator = new ExtensionConcurrencyCoordinator(1, CompletionConfig.allCompleted()); + coordinator.closeRegistration(); + + var exception = + assertThrows(IllegalStateException.class, () -> coordinator.register(() -> new TestFuture<>("late"))); + + assertEquals("Cannot register items after registration is closed", exception.getMessage()); + } + + private static DurableFuture launch( + TestFuture future, AtomicInteger launched, CountDownLatch launchLatch) { + launched.incrementAndGet(); + launchLatch.countDown(); + return future; + } + + private static final class TestFuture implements DurableFuture { + private final T result; + private final RuntimeException failure; + private final CompletableFuture completion = new CompletableFuture<>(); + + private TestFuture(T result) { + this.result = result; + this.failure = null; + } + + private TestFuture(RuntimeException failure) { + this.result = null; + this.failure = failure; + } + + @Override + public T get() { + if (failure != null) { + throw failure; + } + return result; + } + + @Override + public CompletableFuture completionFuture() { + return completion.thenApply(ignored -> null); + } + + private void complete() { + completion.complete(null); + } + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java index dea26d90b..72c825e5d 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java @@ -4,10 +4,16 @@ import static org.junit.jupiter.api.Assertions.*; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState; import software.amazon.awssdk.services.lambda.model.GetDurableExecutionStateResponse; @@ -26,15 +32,24 @@ class ExecutionManagerTest { + EXECUTION_NAME + "/" + EXECUTION_OP_ID; private DurableExecutionClient client; + private ExecutionManager executionManager; + + @AfterEach + void clearCurrentThreadContext() { + if (executionManager != null) { + executionManager.setCurrentThreadContext(null); + } + } private ExecutionManager createManager(List operations) { client = TestUtils.createMockClient(); var initialState = CheckpointUpdatedExecutionState.builder().operations(operations).build(); - return new ExecutionManager( + executionManager = new ExecutionManager( new DurableExecutionInput(EXECUTION_ARN, "test-token", initialState), DurableConfig.builder().withDurableExecutionClient(client).build(), null); + return executionManager; } private Operation executionOp() { @@ -199,4 +214,31 @@ void isOperationUpdatedSinceLastInvocation_handlesMultipleIds() { assertTrue(manager.isOperationUpdatedSinceLastInvocation("3")); assertFalse(manager.isOperationUpdatedSinceLastInvocation("4")); } + + @Test + void awaitFutureDeregistersAndReregistersCurrentContextThread() { + var manager = spy(createManager(List.of(executionOp()))); + var deregistered = new CountDownLatch(1); + doAnswer(invocation -> { + deregistered.countDown(); + return null; + }) + .when(manager) + .deregisterActiveThread("context"); + var future = new CompletableFuture(); + manager.setCurrentThreadContext(new ThreadContext("context", ThreadType.CONTEXT)); + CompletableFuture.runAsync(() -> { + try { + deregistered.await(); + future.complete("done"); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + future.completeExceptionally(e); + } + }); + + assertEquals("done", manager.awaitFuture(future)); + verify(manager).deregisterActiveThread("context"); + verify(manager).registerActiveThread("context"); + } } From ca88c938d6c261e22eb0cc03fec4e2db39c9ab78 Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 04:26:39 +0000 Subject: [PATCH 21/40] refactor: implement map as an extension --- .../lambda/durable/PluginIntegrationTest.java | 28 ++ .../StaticOperationsIntegrationTest.java | 49 +++ .../lambda/durable/DurableMapOperations.java | 22 +- .../durable/context/DurableContextImpl.java | 30 +- .../ExtensionConcurrencyCoordinator.java | 3 + .../context/extension/MapExtension.java | 225 ++++++++++++ .../operation/ConcurrencyOperation.java | 2 +- .../durable/operation/MapOperation.java | 324 ------------------ .../durable/DurableMapOperationsTest.java | 85 ++++- .../ExtensionConcurrencyCoordinatorTest.java | 24 ++ .../context/extension/MapExtensionTest.java | 130 +++++++ 11 files changed, 552 insertions(+), 370 deletions(-) create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/context/extension/MapExtension.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/operation/MapOperation.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/context/extension/MapExtensionTest.java diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java index 45e99bfba..d610d4eed 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java @@ -408,6 +408,34 @@ void plugin_operationStartAndEnd_balanced_forEmptyMap() { assertEquals(1, endCount, "empty map should fire onOperationEnd exactly once"); } + @Test + void plugin_preservesMapAndIterationHookBoundaries() { + var plugin = new RecordingPlugin(); + var config = DurableConfig.builder().withPlugins(plugin).build(); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> DurableMapOperations.map( + "items", List.of("a", "b"), String.class, String::toUpperCase) + .results() + .toString(), + config); + + var result = runner.runUntilComplete("input"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + var operationNames = + plugin.operationStarts.stream().map(OperationInfo::name).toList(); + assertTrue(operationNames.contains("items")); + assertTrue(operationNames.contains("items-iteration-0")); + assertTrue(operationNames.contains("items-iteration-1")); + var userFunctionNames = plugin.userFunctionStarts.stream() + .map(UserFunctionStartInfo::name) + .toList(); + assertFalse(userFunctionNames.contains("items")); + assertTrue(userFunctionNames.contains("items-iteration-0")); + assertTrue(userFunctionNames.contains("items-iteration-1")); + } + // ─── Operation change hook ─────────────────────────────────────────── @Test diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java index aee903645..06ee1c9dc 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java @@ -12,11 +12,14 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.config.MapConfig; +import software.amazon.lambda.durable.config.NestingType; import software.amazon.lambda.durable.config.WaitForConditionConfig; import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.model.ExecutionStatus; import software.amazon.lambda.durable.model.WaitForConditionResult; import software.amazon.lambda.durable.testing.LocalDurableTestRunner; +import software.amazon.lambda.durable.testing.TestResult; class StaticOperationsIntegrationTest { @Test @@ -68,6 +71,40 @@ void mapAndParallelExposeContextFreeUserFunctions() { assertEquals("[a0, b1]:[L, R]", result.getResult(String.class)); } + @Test + void staticMapMatchesLegacyCheckpointHistory() { + var mapConfig = MapConfig.builder() + .maxConcurrency(1) + .nestingType(NestingType.NESTED) + .build(); + var legacyRunner = LocalDurableTestRunner.create(String.class, (input, context) -> context.map( + "map", + List.of("a", "b"), + String.class, + (item, index, child) -> child.step("work", String.class, step -> item + index), + mapConfig) + .results() + .toString()); + var staticRunner = LocalDurableTestRunner.create(String.class, (input, context) -> DurableMapOperations.map( + "map", + List.of("a", "b"), + String.class, + item -> { + var index = MapItemContext.getCurrentContext().getIndex(); + return DurableCoreOperations.step("work", String.class, () -> item + index); + }, + mapConfig) + .results() + .toString()); + + var legacyResult = legacyRunner.runUntilComplete("input"); + var staticResult = staticRunner.runUntilComplete("input"); + + assertEquals("[a0, b1]", legacyResult.getResult(String.class)); + assertEquals(legacyResult.getResult(String.class), staticResult.getResult(String.class)); + assertEquals(operationHistory(legacyResult), operationHistory(staticResult)); + } + @Test void conditionAndRetryExposeGeneratedMetadataThroughTls() { var retryExecutions = new AtomicInteger(); @@ -121,4 +158,16 @@ void waitForCallbackExposesCallbackIdThroughTls() { assertEquals(ExecutionStatus.SUCCEEDED, completed.getStatus()); assertEquals("approved", completed.getResult(String.class)); } + + private static List operationHistory(TestResult result) { + return result.getOperations().stream() + .map(operation -> String.join( + ":", + operation.getId(), + operation.getName(), + operation.getType().toString(), + operation.getSubtype(), + operation.getStatus().toString())) + .toList(); + } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableMapOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableMapOperations.java index e71dc1d9a..de1761e8f 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableMapOperations.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableMapOperations.java @@ -6,6 +6,8 @@ import java.util.Objects; import java.util.function.Function; import software.amazon.lambda.durable.config.MapConfig; +import software.amazon.lambda.durable.context.extension.MapExtension; +import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.model.MapResult; /** Context-free static facades for durable map operations. */ @@ -14,42 +16,42 @@ private DurableMapOperations() {} public static MapResult map( String name, Collection items, Class resultType, Function function) { - return currentContext().map(name, items, resultType, adapt(function)); + return mapAsync(name, items, resultType, function).get(); } public static MapResult map( String name, Collection items, TypeToken resultType, Function function) { - return currentContext().map(name, items, resultType, adapt(function)); + return mapAsync(name, items, resultType, function).get(); } public static MapResult map( String name, Collection items, Class resultType, Function function, MapConfig config) { - return currentContext().map(name, items, resultType, adapt(function), config); + return mapAsync(name, items, resultType, function, config).get(); } public static MapResult map( String name, Collection items, TypeToken resultType, Function function, MapConfig config) { - return currentContext().map(name, items, resultType, adapt(function), config); + return mapAsync(name, items, resultType, function, config).get(); } public static DurableFuture> mapAsync( String name, Collection items, Class resultType, Function function) { - return currentContext().mapAsync(name, items, resultType, adapt(function)); + return mapAsync(name, items, TypeToken.get(resultType), function); } public static DurableFuture> mapAsync( String name, Collection items, TypeToken resultType, Function function) { - return currentContext().mapAsync(name, items, resultType, adapt(function)); + return mapAsync(name, items, resultType, function, MapConfig.builder().build()); } public static DurableFuture> mapAsync( String name, Collection items, Class resultType, Function function, MapConfig config) { - return currentContext().mapAsync(name, items, resultType, adapt(function), config); + return mapAsync(name, items, TypeToken.get(resultType), function, config); } public static DurableFuture> mapAsync( String name, Collection items, TypeToken resultType, Function function, MapConfig config) { - return currentContext().mapAsync(name, items, resultType, adapt(function), config); + return MapExtension.execute(currentContext(), name, items, resultType, adapt(function), config); } private static DurableContext.MapFunction adapt(Function function) { @@ -61,7 +63,7 @@ private static DurableContext.MapFunction adapt(Function func }; } - private static DurableContext currentContext() { - return DurableContext.getCurrentContext(); + private static ExtensionContext currentContext() { + return ExtensionContext.getCurrentContext(); } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java index d80466b60..35726edc1 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java @@ -5,7 +5,6 @@ import com.amazonaws.services.lambda.runtime.Context; import java.time.Duration; import java.util.Collection; -import java.util.List; import java.util.Objects; import java.util.function.BiConsumer; import java.util.function.BiFunction; @@ -27,6 +26,7 @@ import software.amazon.lambda.durable.config.WaitForCallbackConfig; import software.amazon.lambda.durable.config.WaitForConditionConfig; import software.amazon.lambda.durable.config.WithRetryConfig; +import software.amazon.lambda.durable.context.extension.MapExtension; import software.amazon.lambda.durable.context.extension.WaitForCallbackExtension; import software.amazon.lambda.durable.context.extension.WaitForConditionExtension; import software.amazon.lambda.durable.context.extension.WithRetryExtension; @@ -48,7 +48,6 @@ import software.amazon.lambda.durable.operation.CallbackOperation; import software.amazon.lambda.durable.operation.ChildContextOperation; import software.amazon.lambda.durable.operation.InvokeOperation; -import software.amazon.lambda.durable.operation.MapOperation; import software.amazon.lambda.durable.operation.ParallelOperation; import software.amazon.lambda.durable.operation.StepOperation; import software.amazon.lambda.durable.operation.WaitOperation; @@ -446,32 +445,7 @@ DurableFuture extensionContextAsyncWithId( @Override public DurableFuture> mapAsync( String name, Collection items, TypeToken resultType, MapFunction function, MapConfig config) { - Objects.requireNonNull(items, "items cannot be null"); - Objects.requireNonNull(function, "function cannot be null"); - Objects.requireNonNull(resultType, "resultType cannot be null"); - Objects.requireNonNull(config, "config cannot be null"); - ParameterValidator.validateOperationName(name); - ParameterValidator.validateOrderedCollection(items); - - if (config.serDes() == null) { - config = config.toBuilder().serDes(getDurableConfig().getSerDes()).build(); - } - - // Convert to List for deterministic index-based access - var itemList = List.copyOf(items); - var iterationNames = MapOperation.resolveIterationNames(name, itemList, config); - var operationId = nextOperationId(); - - var operation = new MapOperation<>( - OperationIdentifier.of(operationId, name, OperationSubType.MAP), - itemList, - function, - resultType, - config, - iterationNames, - this); - operation.execute(); - return operation; + return MapExtension.execute(this, name, items, resultType, function, config); } @Override diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinator.java b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinator.java index f16775fd3..8b4f82399 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinator.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinator.java @@ -130,6 +130,9 @@ Completion awaitCompletion(ExpectedCompletionStatus expectedCompletionStatus) { markIncompleteItemsSkipped(); return new Completion(decision, items); } + if (running.size() < maxConcurrency && !pending.isEmpty()) { + continue; + } waiters = completionWaiters(); } DurableFuture.anyOf(waiters); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/MapExtension.java b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/MapExtension.java new file mode 100644 index 000000000..a05be6659 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/MapExtension.java @@ -0,0 +1,225 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.context.extension; + +import static software.amazon.lambda.durable.config.NestingType.FLAT; +import static software.amazon.lambda.durable.context.extension.ExtensionConcurrencyCoordinator.ItemStatus.FAILED; +import static software.amazon.lambda.durable.context.extension.ExtensionConcurrencyCoordinator.ItemStatus.SKIPPED; +import static software.amazon.lambda.durable.model.OperationSubType.MAP; +import static software.amazon.lambda.durable.model.OperationSubType.MAP_ITERATION; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.CompletionConfig; +import software.amazon.lambda.durable.config.MapConfig; +import software.amazon.lambda.durable.config.RunInChildContextConfig; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.execution.SuspendExecutionException; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; +import software.amazon.lambda.durable.extension.ExtensionContextResult; +import software.amazon.lambda.durable.model.MapResult; +import software.amazon.lambda.durable.util.ExceptionHelper; +import software.amazon.lambda.durable.util.ParameterValidator; + +/** Canonical implementation of the built-in map extension. */ +public final class MapExtension { + private static final int LARGE_RESULT_THRESHOLD = 256 * 1024; + + private MapExtension() {} + + public static DurableFuture> execute( + ExtensionContext context, + String name, + Collection items, + TypeToken resultType, + DurableContext.MapFunction function, + MapConfig config) { + Objects.requireNonNull(context, "context cannot be null"); + Objects.requireNonNull(items, "items cannot be null"); + Objects.requireNonNull(function, "function cannot be null"); + Objects.requireNonNull(resultType, "resultType cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + ParameterValidator.validateOperationName(name); + ParameterValidator.validateOrderedCollection(items); + + if (config.serDes() == null) { + config = config.toBuilder() + .serDes(context.getDurableConfig().getSerDes()) + .build(); + } + var itemList = List.copyOf(items); + var iterationNames = resolveIterationNames(name, itemList, config); + var parent = context.reserve(name); + validateMinSuccessful(itemList, config); + + var mapConfig = config; + var virtualEmptyMap = itemList.isEmpty() && !context.getDurableConfig().shouldCheckpointEmptyMap(); + return parent.runInChildContextAsync( + MAP.getValue(), + mapResultType(), + () -> executeInChildContext( + name, itemList, iterationNames, resultType, function, mapConfig, virtualEmptyMap), + parentConfig(mapConfig, virtualEmptyMap)); + } + + private static ExtensionContextResult> executeInChildContext( + String name, + List items, + List iterationNames, + TypeToken resultType, + DurableContext.MapFunction function, + MapConfig config, + boolean virtualEmptyMap) { + if (virtualEmptyMap) { + ExtensionContext.getCurrentContext() + .getLogger() + .warn( + "Empty map operation '{}' is not checkpointed by default. This behavior is unintended and" + + " may affect replay and plugin instrumentation. Enable" + + " DurableConfig.withCheckpointEmptyMap(true) to checkpoint empty maps.", + name); + return ExtensionContextResult.completed(MapResult.empty()); + } + + var replay = ExtensionContextReplayContext.>getCurrentContext(); + var replayState = replay.isReplayingChildren() ? replay.getReplayState() : null; + if (replay.isReplayingChildren() && replayState == null) { + throw new IllegalStateException("Missing result in completed Map operation"); + } + + var coordinator = new ExtensionConcurrencyCoordinator(config.maxConcurrency(), config.completionConfig()); + var registeredItems = + registerItems(coordinator, items, iterationNames, resultType, function, config, replayState); + coordinator.closeRegistration(); + var completion = replayState == null + ? coordinator.awaitCompletion() + : coordinator.awaitCompletion(expectedCompletion(replayState)); + var result = constructResult(registeredItems, completion.completionDecision()); + var strippedResult = stripMapResult(result); + return config.itemNamer() == null + ? ExtensionContextResult.replayChildrenAboveSize(result, strippedResult, LARGE_RESULT_THRESHOLD) + : ExtensionContextResult.replayChildren(result, strippedResult); + } + + private static List> registerItems( + ExtensionConcurrencyCoordinator coordinator, + List items, + List iterationNames, + TypeToken resultType, + DurableContext.MapFunction function, + MapConfig config, + MapResult replayState) { + var context = ExtensionContext.getCurrentContext(); + var registeredItems = new ArrayList>(items.size()); + var iterationConfig = RunInChildContextConfig.builder() + .serDes(config.serDes()) + .isVirtual(config.nestingType() == FLAT) + .build(); + + for (int index = 0; index < items.size(); index++) { + var item = items.get(index); + var itemIndex = index; + var reservation = context.reserve(iterationNames.get(index)); + var skipped = replayState != null + && replayState.getItem(index).status() == MapResult.MapResultItem.Status.SKIPPED; + registeredItems.add(coordinator.register( + () -> reservation.runInChildContextAsync( + MAP_ITERATION.getValue(), + resultType, + () -> function.apply(item, itemIndex, DurableContext.getCurrentContext()), + iterationConfig), + skipped)); + } + return registeredItems; + } + + private static List resolveIterationNames(String mapName, List items, MapConfig config) { + var namer = config.itemNamer(); + var prefix = mapName == null ? "map-iteration-" : mapName + "-iteration-"; + var names = new ArrayList(items.size()); + for (int index = 0; index < items.size(); index++) { + var iterationName = namer == null ? prefix + index : namer.apply(items.get(index), index); + ParameterValidator.validateOperationName(iterationName); + names.add(iterationName); + } + return names; + } + + private static ExtensionConcurrencyCoordinator.ExpectedCompletionStatus expectedCompletion( + MapResult replayState) { + return new ExtensionConcurrencyCoordinator.ExpectedCompletionStatus( + replayState.succeeded().size() + replayState.failed().size(), + CompletionConfig.CompletionDecision.complete(replayState.completionReason())); + } + + private static MapResult constructResult( + List> items, + CompletionConfig.CompletionDecision completionDecision) { + var results = new ArrayList>(Collections.nCopies(items.size(), null)); + for (int index = 0; index < items.size(); index++) { + var item = items.get(index); + if (item.status() == SKIPPED) { + results.set(index, MapResult.MapResultItem.skipped()); + } else if (item.status() == FAILED) { + results.set(index, failedResult(item)); + } else { + results.set( + index, MapResult.MapResultItem.succeeded(item.future().get())); + } + } + return new MapResult<>(results, completionDecision.completionStatus()); + } + + private static MapResult.MapResultItem failedResult(ExtensionConcurrencyCoordinator.Item item) { + try { + item.future().get(); + throw new IllegalStateException("Failed map item completed successfully"); + } catch (SuspendExecutionException | UnrecoverableDurableExecutionException exception) { + throw exception; + } catch (Throwable throwable) { + return MapResult.MapResultItem.failed( + MapResult.MapError.of(ExceptionHelper.unwrapCompletableFuture(throwable))); + } + } + + private static MapResult stripMapResult(MapResult result) { + return new MapResult<>( + result.items().stream() + .map(item -> new MapResult.MapResultItem(item.status(), null, null)) + .toList(), + result.completionReason()); + } + + private static ExtensionContextConfig parentConfig(MapConfig config, boolean virtualEmptyMap) { + return ExtensionContextConfig.builder() + .childContextConfig(RunInChildContextConfig.builder() + .serDes(config.serDes()) + .isVirtual(virtualEmptyMap) + .build()) + .emitUserFunctionEvents(false) + .suppressLateChildCheckpoints(true) + .build(); + } + + private static void validateMinSuccessful(List items, MapConfig config) { + var completionConfig = config.completionConfig(); + if (!completionConfig.hasCustomShouldComplete() + && completionConfig.minSuccessful() != null + && completionConfig.minSuccessful() > items.size()) { + throw new IllegalArgumentException("minSuccessful cannot be greater than total items: " + + completionConfig.minSuccessful() + " > " + items.size()); + } + } + + private static TypeToken> mapResultType() { + return new TypeToken<>() {}; + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ConcurrencyOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ConcurrencyOperation.java index 321894bd3..913103a4c 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ConcurrencyOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/ConcurrencyOperation.java @@ -36,7 +36,7 @@ * Abstract base class for concurrent execution of multiple child context operations. * *

Encapsulates shared concurrency logic: queue-based concurrency control, success/failure counting, and completion - * checking. Both {@code ParallelOperation} and {@code MapOperation} extend this base. + * checking. This remains the legacy execution engine for {@code ParallelOperation}. * *

Key design points: * diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/MapOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/MapOperation.java deleted file mode 100644 index 2f665547a..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/MapOperation.java +++ /dev/null @@ -1,324 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; - -import java.nio.charset.StandardCharsets; -import java.util.ArrayList; -import java.util.Collections; -import java.util.HashMap; -import java.util.List; -import java.util.Objects; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; -import software.amazon.awssdk.services.lambda.model.ContextOptions; -import software.amazon.awssdk.services.lambda.model.Operation; -import software.amazon.awssdk.services.lambda.model.OperationAction; -import software.amazon.awssdk.services.lambda.model.OperationUpdate; -import software.amazon.lambda.durable.DurableContext; -import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.CompletionConfig; -import software.amazon.lambda.durable.config.MapConfig; -import software.amazon.lambda.durable.context.DurableContextImpl; -import software.amazon.lambda.durable.exception.NonDeterministicExecutionException; -import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; -import software.amazon.lambda.durable.execution.SuspendExecutionException; -import software.amazon.lambda.durable.model.ConcurrencyCompletionStatus; -import software.amazon.lambda.durable.model.MapResult; -import software.amazon.lambda.durable.model.OperationIdentifier; -import software.amazon.lambda.durable.model.OperationSubType; -import software.amazon.lambda.durable.serde.SerDes; -import software.amazon.lambda.durable.util.ExceptionHelper; -import software.amazon.lambda.durable.util.ParameterValidator; - -/** - * Executes a map operation: applies a function to each item in a collection concurrently, with each item running in its - * own child context. - * - *

Extends {@link ConcurrencyOperation} following the same pattern as {@link ParallelOperation}. All branches are - * created upfront in {@code start()}/{@code replay()}, and results are aggregated into a {@link MapResult} in - * {@code get()}. - * - * @param the input item type - * @param the output result type per item - */ -public class MapOperation extends ConcurrencyOperation> { - - private static final Logger logger = LoggerFactory.getLogger(MapOperation.class); - private static final int LARGE_RESULT_THRESHOLD = 256 * 1024; - - private final List items; - private final DurableContext.MapFunction function; - private final TypeToken itemResultType; - private final SerDes serDes; - private final List iterationNames; - private volatile MapResult cachedResult; - - public MapOperation( - OperationIdentifier operationIdentifier, - List items, - DurableContext.MapFunction function, - TypeToken itemResultType, - MapConfig config, - DurableContextImpl durableContext) { - this( - operationIdentifier, - items, - function, - itemResultType, - config, - resolveIterationNames(operationIdentifier.name(), items, config), - durableContext); - } - - public MapOperation( - OperationIdentifier operationIdentifier, - List items, - DurableContext.MapFunction function, - TypeToken itemResultType, - MapConfig config, - List iterationNames, - DurableContextImpl durableContext) { - super( - operationIdentifier, - new TypeToken<>() {}, - config.serDes(), - durableContext, - config.maxConcurrency(), - config.completionConfig().completionDecisionFunction(), - config.nestingType()); - if (!config.completionConfig().hasCustomShouldComplete() - && config.completionConfig().minSuccessful() != null - && config.completionConfig().minSuccessful() > items.size()) { - throw new IllegalArgumentException("minSuccessful cannot be greater than total items: " - + config.completionConfig().minSuccessful() + " > " + items.size()); - } - this.items = List.copyOf(items); - this.function = function; - this.itemResultType = itemResultType; - this.serDes = config.serDes(); - this.iterationNames = Collections.unmodifiableList(new ArrayList<>(iterationNames)); - if (this.iterationNames.size() != this.items.size()) { - throw new IllegalArgumentException("iterationNames must have one entry per item"); - } - } - - /** - * Resolves the operation name for every iteration of a map, applying the config's item namer when present and the - * default {@code "-iteration-N"} naming otherwise. A namer that returns null yields an unnamed iteration; - * any non-null name is validated here. - * - *

SDK-internal. This is the single source of iteration naming for both construction paths: the caller resolves - * names before an operation ID is allocated, and the legacy constructor resolves them on behalf of callers that do - * not. - * - * @param mapName the map operation's name, or null - * @param items the map's items, in iteration order - * @param config the map configuration supplying the optional item namer - * @return one name per item, in iteration order - */ - public static List resolveIterationNames(String mapName, List items, MapConfig config) { - var namer = config.itemNamer(); - var branchPrefix = mapName == null ? "map-iteration-" : mapName + "-iteration-"; - var names = new ArrayList(items.size()); - for (int i = 0; i < items.size(); i++) { - if (namer == null) { - names.add(branchPrefix + i); - } else { - var iterationName = namer.apply(items.get(i), i); - ParameterValidator.validateOperationName(iterationName); - names.add(iterationName); - } - } - return names; - } - - private void addAllItems() { - addUnskippedItems(Collections.nCopies(items.size(), null)); - } - - private void addUnskippedItems(List resultItems) { - // Enqueue all items first. - // If the map is completed when replaying, mapResult != null and the items that have been skipped - // will be skipped during replay. - for (int i = 0; i < items.size(); i++) { - var index = i; - var item = items.get(i); - var status = resultItems.get(i); - // the item will be skipped by ConcurrencyOperation if skip=true - var skip = status == MapResult.MapResultItem.Status.SKIPPED; - - enqueueItem( - iterationNames.get(i), - childCtx -> function.apply(item, index, childCtx), - itemResultType, - serDes, - OperationSubType.MAP_ITERATION, - skip); - } - } - - private void validateIterationNamesAgainstCheckpoint() { - var checkpointedById = new HashMap(); - for (var child : getChildOperations()) { - checkpointedById.put(child.id(), child); - } - for (var branch : getBranches()) { - var checkpointed = checkpointedById.get(branch.getOperationId()); - if (checkpointed != null && !Objects.equals(checkpointed.name(), branch.getName())) { - throw terminateExecution(new NonDeterministicExecutionException(String.format( - "Map iteration name mismatch for \"%s\". Expected \"%s\", got \"%s\"", - branch.getOperationId(), checkpointed.name(), branch.getName()))); - } - } - } - - @Override - protected void start() { - if (items.isEmpty()) { - // TODO: Remove the checkpointEmptyMap flag and the non-checkpointing branch in a future major version, - // making checkpointing the default. - if (getContext().getDurableConfig().shouldCheckpointEmptyMap()) { - // Checkpoint START + SUCCEED to produce a complete map operation with an empty result. - sendOperationUpdate(OperationUpdate.builder() - .action(OperationAction.START) - .subType(getSubType().getValue())); - handleCompletion( - CompletionConfig.CompletionDecision.complete(ConcurrencyCompletionStatus.ALL_COMPLETED)); - } else { - // Default: complete without checkpointing. Fire onOperationEnd so plugin hooks stay balanced. - logger.warn( - "Empty map operation '{}' is not checkpointed by default. This behavior is unintended and may" - + " affect replay and plugin instrumentation. Enable" - + " DurableConfig.withCheckpointEmptyMap(true) to checkpoint empty maps.", - getName()); - cachedResult = MapResult.empty(); - fireOnOperationEnd(null, null, false); - markAlreadyCompleted(); - } - return; - } - sendOperationUpdateAsync(OperationUpdate.builder() - .action(OperationAction.START) - .subType(getSubType().getValue())); - - addAllItems(); - executeItems(); - } - - @Override - protected void replay(Operation existing) { - switch (existing.status()) { - case SUCCEEDED -> { - var result = existing.contextDetails() != null - ? existing.contextDetails().result() - : null; - var deserializedResult = result != null ? deserializeResult(result) : null; - if (deserializedResult != null) { - addUnskippedItems(deserializedResult.items().stream() - .map(MapResult.MapResultItem::status) - .toList()); - } else { - throw terminateExecutionWithIllegalDurableOperationException( - "Missing result in completed Map operation"); - } - validateIterationNamesAgainstCheckpoint(); - if (Boolean.TRUE.equals(existing.contextDetails().replayChildren())) { - // Large result: re-execute children to reconstruct MapResult - var expected = new ExpectedCompletionStatus( - deserializedResult.succeeded().size() - + deserializedResult.failed().size(), - CompletionConfig.CompletionDecision.complete(deserializedResult.completionReason())); - executeItems(expected); - } else { - // Small result: MapResult is in the payload, skip child replay - cachedResult = deserializedResult; - markAlreadyCompleted(); - } - } - case STARTED -> { - // Map was in progress when interrupted — re-create children without sending - // another START (the backend rejects duplicate START for existing operations) - addAllItems(); - validateIterationNamesAgainstCheckpoint(); - executeItems(); - } - default -> - throw terminateExecutionWithIllegalDurableOperationException( - "Unexpected map operation status: " + existing.status()); - } - } - - @Override - protected void handleCompletion(CompletionConfig.CompletionDecision completionDecision) { - this.cachedResult = constructMapResult(completionDecision); - var serializedResult = serializeAndDeserializeResult(cachedResult); - this.cachedResult = serializedResult.deserialized(); - var serializedBytes = serializedResult.serialized().getBytes(StandardCharsets.UTF_8); - - if (serializedBytes.length < LARGE_RESULT_THRESHOLD) { - sendOperationUpdate(OperationUpdate.builder() - .action(OperationAction.SUCCEED) - .subType(getSubType().getValue()) - .payload(serializedResult.serialized())); - } else { - // Large result: checkpoint with stripped payload + replayChildren flag - var strippedResult = serializeAndDeserializeResult(stripMapResult(cachedResult)); - sendOperationUpdate(OperationUpdate.builder() - .action(OperationAction.SUCCEED) - .subType(getSubType().getValue()) - .payload(strippedResult.serialized()) - .contextOptions( - ContextOptions.builder().replayChildren(true).build())); - } - } - - private MapResult stripMapResult(MapResult result) { - return new MapResult<>( - result.items().stream() - .map(item -> new MapResult.MapResultItem(item.status(), null, null)) - .toList(), - result.completionReason()); - } - - @SuppressWarnings("unchecked") - private MapResult constructMapResult(CompletionConfig.CompletionDecision completionDecision) { - var children = getBranches(); - var resultItems = new ArrayList>(Collections.nCopies(items.size(), null)); - - for (int i = 0; i < children.size(); i++) { - var branch = (ChildContextOperation) children.get(i); - if (!branch.isOperationCompleted()) { - resultItems.set(i, MapResult.MapResultItem.skipped()); - } else { - try { - resultItems.set(i, MapResult.MapResultItem.succeeded(branch.get())); - } catch (Throwable exception) { - Throwable throwable = ExceptionHelper.unwrapCompletableFuture(exception); - if (throwable instanceof SuspendExecutionException suspendExecutionException) { - // Rethrow Error immediately — do not checkpoint - throw suspendExecutionException; - } - if (throwable - instanceof UnrecoverableDurableExecutionException unrecoverableDurableExecutionException) { - // terminate the execution and throw the exception if it's not recoverable - throw terminateExecution(unrecoverableDurableExecutionException); - } - resultItems.set(i, MapResult.MapResultItem.failed(MapResult.MapError.of(throwable))); - } - } - } - return new MapResult<>(resultItems, completionDecision.completionStatus()); - } - - @Override - public MapResult get() { - // Non-checkpointed empty map: no stored operation, so skip join() and return the result set in start(). - // Tied to the temporary checkpointEmptyMap flag; remove with it in a future major version. - if (items.isEmpty() && !getContext().getDurableConfig().shouldCheckpointEmptyMap()) { - return cachedResult; - } - join(); - // cachedResult is always set upon successful completion - return cachedResult; - } -} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationsTest.java index 334910c62..23e279abb 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationsTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationsTest.java @@ -3,16 +3,29 @@ package software.amazon.lambda.durable; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.function.Supplier; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; +import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.context.BaseContextImpl; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextFunction; +import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; +import software.amazon.lambda.durable.extension.ExtensionOperation; +import software.amazon.lambda.durable.model.MapResult; +import software.amazon.lambda.durable.model.OperationSubType; class DurableMapOperationsTest { @AfterEach @@ -22,19 +35,77 @@ void clearContext() { @Test void mapExposesItemIndexThroughScopedContext() { - var context = mock(DurableContext.class); + var context = mock(CurrentContext.class); + var parent = mock(ExtensionOperation.class); + var iteration = mock(ExtensionOperation.class); + var expected = new MapResult<>( + List.of(MapResult.MapResultItem.succeeded("VALUE")), + software.amazon.lambda.durable.model.ConcurrencyCompletionStatus.ALL_COMPLETED); + var parentFuture = new CompletedFuture<>(expected); + when(context.getDurableConfig()).thenReturn(DurableConfig.builder().build()); + when(context.reserve("map")).thenReturn(parent); + when(parent.runInChildContextAsync( + eq(OperationSubType.MAP.getValue()), + any(TypeToken.class), + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) + .thenReturn(parentFuture); BaseContextImpl.setCurrentContext(context); - DurableMapOperations.map("map", List.of("value"), String.class, item -> { - assertEquals(3, MapItemContext.getCurrentContext().getIndex()); + var actual = DurableMapOperations.map("map", List.of("value"), String.class, item -> { + assertEquals(0, MapItemContext.getCurrentContext().getIndex()); return item.toUpperCase(); }); + assertSame(expected, actual); @SuppressWarnings("unchecked") - var function = (ArgumentCaptor>) - (ArgumentCaptor) ArgumentCaptor.forClass(DurableContext.MapFunction.class); - verify(context).map(eq("map"), eq(List.of("value")), eq(String.class), function.capture()); - assertEquals("VALUE", function.getValue().apply("value", 3, mock(DurableContext.class))); + var parentFunction = (ArgumentCaptor>>) + (ArgumentCaptor) ArgumentCaptor.forClass(ExtensionContextFunction.class); + verify(parent) + .runInChildContextAsync( + eq(OperationSubType.MAP.getValue()), + any(TypeToken.class), + parentFunction.capture(), + any(ExtensionContextConfig.class)); + + when(context.reserve("map-iteration-0")).thenReturn(iteration); + when(iteration.runInChildContextAsync( + eq(OperationSubType.MAP_ITERATION.getValue()), + eq(TypeToken.get(String.class)), + any(Supplier.class), + any(RunInChildContextConfig.class))) + .thenReturn(new CompletedFuture<>("VALUE")); + try (var ignoredContext = BaseContextImpl.attachCurrentContext(context); + var ignoredReplay = ExtensionContextReplayContext.attach(false, null)) { + parentFunction.getValue().apply(); + } + + @SuppressWarnings("unchecked") + var itemFunction = + (ArgumentCaptor>) (ArgumentCaptor) ArgumentCaptor.forClass(Supplier.class); + verify(iteration) + .runInChildContextAsync( + eq(OperationSubType.MAP_ITERATION.getValue()), + eq(TypeToken.get(String.class)), + itemFunction.capture(), + any(RunInChildContextConfig.class)); + try (var ignored = BaseContextImpl.attachCurrentContext(context)) { + assertEquals("VALUE", itemFunction.getValue().get()); + } assertThrows(IllegalStateException.class, MapItemContext::getCurrentContext); } + + private interface CurrentContext extends DurableContext, ExtensionContext {} + + private record CompletedFuture(T result) implements DurableFuture { + @Override + public T get() { + return result; + } + + @Override + public CompletableFuture completionFuture() { + return CompletableFuture.completedFuture(null); + } + } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinatorTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinatorTest.java index 60ae85499..2b70c62f1 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinatorTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinatorTest.java @@ -5,6 +5,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; import static software.amazon.lambda.durable.context.extension.ExtensionConcurrencyCoordinator.ItemStatus.FAILED; import static software.amazon.lambda.durable.context.extension.ExtensionConcurrencyCoordinator.ItemStatus.SKIPPED; import static software.amazon.lambda.durable.context.extension.ExtensionConcurrencyCoordinator.ItemStatus.SUCCEEDED; @@ -14,6 +15,7 @@ import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.Test; import software.amazon.lambda.durable.DurableFuture; @@ -47,6 +49,28 @@ void launchesNoMoreThanMaxConcurrency() throws Exception { assertEquals(ALL_COMPLETED, result.join().completionDecision().completionStatus()); } + @Test + void launchesNextItemAfterSynchronousReplayCompletion() throws Exception { + var coordinator = new ExtensionConcurrencyCoordinator(1, CompletionConfig.allCompleted()); + var replayed = new TestFuture<>("replayed"); + replayed.complete(); + var next = new TestFuture<>("next"); + var nextLaunched = new CountDownLatch(1); + coordinator.register(() -> replayed); + coordinator.register(() -> { + nextLaunched.countDown(); + return next; + }); + + var result = CompletableFuture.supplyAsync(coordinator::awaitCompletion); + var launchedWithoutExternalSignal = nextLaunched.await(200, TimeUnit.MILLISECONDS); + coordinator.closeRegistration(); + next.complete(); + result.join(); + + assertTrue(launchedWithoutExternalSignal); + } + @Test void earlyCompletionMarksUnlaunchedItemsSkipped() { var coordinator = new ExtensionConcurrencyCoordinator(1, CompletionConfig.minSuccessful(1)); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/MapExtensionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/extension/MapExtensionTest.java new file mode 100644 index 000000000..750660f49 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/context/extension/MapExtensionTest.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.context.extension; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static software.amazon.lambda.durable.model.OperationSubType.MAP; +import static software.amazon.lambda.durable.model.OperationSubType.MAP_ITERATION; + +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.function.Supplier; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.MapConfig; +import software.amazon.lambda.durable.config.NestingType; +import software.amazon.lambda.durable.config.RunInChildContextConfig; +import software.amazon.lambda.durable.context.BaseContextImpl; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextFunction; +import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; +import software.amazon.lambda.durable.extension.ExtensionOperation; +import software.amazon.lambda.durable.model.MapResult; +import software.amazon.lambda.durable.serde.JacksonSerDes; + +class MapExtensionTest { + @Test + void executeBuildsMapAndIterationContextsFromReservations() { + var context = mock(ExtensionContext.class); + var parent = mock(ExtensionOperation.class); + var parentFuture = mockMapFuture(); + var serDes = new JacksonSerDes(); + var config = + MapConfig.builder().serDes(serDes).nestingType(NestingType.FLAT).build(); + when(context.reserve("map")).thenReturn(parent); + when(parent.runInChildContextAsync( + eq(MAP.getValue()), + any(TypeToken.class), + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) + .thenReturn(parentFuture); + + var actual = MapExtension.execute( + context, + "map", + List.of("a", "b"), + TypeToken.get(String.class), + (item, index, child) -> item + index, + config); + + assertSame(parentFuture, actual); + var function = extensionFunction(); + var parentConfig = ArgumentCaptor.forClass(ExtensionContextConfig.class); + verify(parent) + .runInChildContextAsync( + eq(MAP.getValue()), any(TypeToken.class), function.capture(), parentConfig.capture()); + assertSame(serDes, parentConfig.getValue().childContextConfig().serDes()); + assertFalse(parentConfig.getValue().emitUserFunctionEvents()); + assertTrue(parentConfig.getValue().suppressLateChildCheckpoints()); + + var child = mock(CurrentContext.class); + var first = mock(ExtensionOperation.class); + var second = mock(ExtensionOperation.class); + when(child.reserve("map-iteration-0")).thenReturn(first); + when(child.reserve("map-iteration-1")).thenReturn(second); + when(first.runInChildContextAsync( + eq(MAP_ITERATION.getValue()), + eq(TypeToken.get(String.class)), + any(Supplier.class), + any(RunInChildContextConfig.class))) + .thenReturn(new CompletedFuture<>("a0")); + when(second.runInChildContextAsync( + eq(MAP_ITERATION.getValue()), + eq(TypeToken.get(String.class)), + any(Supplier.class), + any(RunInChildContextConfig.class))) + .thenReturn(new CompletedFuture<>("b1")); + + try (var ignoredContext = BaseContextImpl.attachCurrentContext(child); + var ignoredReplay = ExtensionContextReplayContext.attach(false, null)) { + var result = function.getValue().apply().result(); + assertEquals(List.of("a0", "b1"), result.results()); + } + + var iterationConfig = ArgumentCaptor.forClass(RunInChildContextConfig.class); + verify(first) + .runInChildContextAsync( + eq(MAP_ITERATION.getValue()), + eq(TypeToken.get(String.class)), + any(Supplier.class), + iterationConfig.capture()); + assertTrue(iterationConfig.getValue().isVirtual()); + assertSame(serDes, iterationConfig.getValue().serDes()); + } + + @SuppressWarnings({"rawtypes", "unchecked"}) + private ArgumentCaptor>> extensionFunction() { + return (ArgumentCaptor) ArgumentCaptor.forClass(ExtensionContextFunction.class); + } + + @SuppressWarnings("unchecked") + private DurableFuture> mockMapFuture() { + return mock(DurableFuture.class); + } + + private interface CurrentContext extends DurableContext, ExtensionContext {} + + private record CompletedFuture(T result) implements DurableFuture { + @Override + public T get() { + return result; + } + + @Override + public CompletableFuture completionFuture() { + return CompletableFuture.completedFuture(null); + } + } +} From 109c0caa6cf8572e33f8d868d7537e3c26f39d7c Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 04:37:17 +0000 Subject: [PATCH 22/40] refactor: implement parallel as an extension --- docs/design.md | 27 +- .../StaticOperationsIntegrationTest.java | 32 + .../durable/DurableParallelOperations.java | 6 +- .../durable/context/DurableContextImpl.java | 16 +- .../context/extension/ParallelExtension.java | 21 + .../extension/ParallelExtensionFuture.java | 245 ++++++++ .../operation/BaseDurableOperation.java | 5 +- .../operation/ChildContextOperation.java | 15 +- .../operation/ConcurrencyOperation.java | 363 ----------- .../durable/operation/ParallelOperation.java | 192 ------ .../SerializableDurableOperation.java | 2 +- .../DurableParallelOperationsTest.java | 47 +- .../extension/ParallelExtensionTest.java | 263 ++++++++ .../operation/ChildContextOperationTest.java | 30 +- .../operation/ConcurrencyOperationTest.java | 359 ----------- .../operation/ParallelOperationTest.java | 579 ------------------ 16 files changed, 640 insertions(+), 1562 deletions(-) create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/context/extension/ParallelExtension.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/context/extension/ParallelExtensionFuture.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/operation/ConcurrencyOperation.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/context/extension/ParallelExtensionTest.java delete mode 100644 sdk/src/test/java/software/amazon/lambda/durable/operation/ConcurrencyOperationTest.java delete mode 100644 sdk/src/test/java/software/amazon/lambda/durable/operation/ParallelOperationTest.java diff --git a/docs/design.md b/docs/design.md index eaba54af0..50a2c5325 100644 --- a/docs/design.md +++ b/docs/design.md @@ -253,10 +253,6 @@ context.step("name", Type.class, stepCtx -> doWork(), │ - WaitOperation │ │ - Batches API calls (750KB) │ │ - InvokeOperation │ │ │ │ - CallbackOperation │ │ - Notifies via callback │ -│ - WaitForConditionOperation │ └──────────────────────────────┘ -│ - ConcurrencyOperation │ -│ - MapOperation │ -│ - ParallelOperation │ │ - ChildContextOperation │ │ - execute() / get() │ └──────────────────────────────┘ @@ -295,7 +291,22 @@ software.amazon.lambda.durable │ └── CompletionConfig # Completion criteria for map/parallel │ ├── context/ -│ └── BaseContext # Base interface for DurableContext +│ ├── BaseContext # Base interface for DurableContext +│ └── extension/ # Built-in composed operation implementations +│ ├── MapExtension +│ ├── ParallelExtension +│ ├── WaitForCallbackExtension +│ ├── WaitForConditionExtension +│ ├── WithRetryExtension +│ └── ExtensionConcurrencyCoordinator +│ +├── extension/ # Public SPI for extension authors +│ ├── ExtensionContext +│ ├── ExtensionOperation +│ ├── ExtensionStepConfig +│ ├── ExtensionStepResult +│ ├── ExtensionContextConfig +│ └── ExtensionContextResult │ ├── execution/ │ ├── ExecutionManager # Central coordinator @@ -311,11 +322,7 @@ software.amazon.lambda.durable │ ├── InvokeOperation # Invoke logic │ ├── CallbackOperation # Callback logic │ ├── WaitOperation # Wait logic -│ ├── WaitForConditionOperation # Polling condition logic -│ ├── ConcurrencyOperation # Shared base for map/parallel -│ ├── MapOperation # Map operation logic -│ ├── ParallelOperation # Parallel operation logic -│ └── ChildContextOperation # Per-item child context execution +│ └── ChildContextOperation # Child context primitive │ ├── logging/ │ ├── DurableLogger # Context-aware logger wrapper (MDC-based) diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java index 06ee1c9dc..b51cddece 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java @@ -14,6 +14,7 @@ import org.junit.jupiter.api.Test; import software.amazon.lambda.durable.config.MapConfig; import software.amazon.lambda.durable.config.NestingType; +import software.amazon.lambda.durable.config.ParallelConfig; import software.amazon.lambda.durable.config.WaitForConditionConfig; import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.model.ExecutionStatus; @@ -105,6 +106,37 @@ void staticMapMatchesLegacyCheckpointHistory() { assertEquals(operationHistory(legacyResult), operationHistory(staticResult)); } + @Test + void staticParallelMatchesLegacyCheckpointHistory() { + var parallelConfig = ParallelConfig.builder() + .maxConcurrency(1) + .nestingType(NestingType.NESTED) + .build(); + var legacyRunner = LocalDurableTestRunner.create(String.class, (input, context) -> { + try (var parallel = context.parallel("parallel", parallelConfig)) { + parallel.branch("left", String.class, child -> child.step("work", String.class, step -> "L")); + parallel.branch("right", String.class, child -> child.step("work", String.class, step -> "R")); + return parallel.get().statuses().toString(); + } + }); + var staticRunner = LocalDurableTestRunner.create(String.class, (input, context) -> { + try (var parallel = DurableParallelOperations.parallel("parallel", parallelConfig)) { + parallel.branch( + "left", String.class, () -> DurableCoreOperations.step("work", String.class, () -> "L")); + parallel.branch( + "right", String.class, () -> DurableCoreOperations.step("work", String.class, () -> "R")); + return parallel.get().statuses().toString(); + } + }); + + var legacyResult = legacyRunner.runUntilComplete("input"); + var staticResult = staticRunner.runUntilComplete("input"); + + assertEquals("[SUCCEEDED, SUCCEEDED]", legacyResult.getResult(String.class)); + assertEquals(legacyResult.getResult(String.class), staticResult.getResult(String.class)); + assertEquals(operationHistory(legacyResult), operationHistory(staticResult)); + } + @Test void conditionAndRetryExposeGeneratedMetadataThroughTls() { var retryExecutions = new AtomicInteger(); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableParallelOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableParallelOperations.java index 18cf6d519..25c7751b6 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableParallelOperations.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableParallelOperations.java @@ -3,16 +3,18 @@ package software.amazon.lambda.durable; import software.amazon.lambda.durable.config.ParallelConfig; +import software.amazon.lambda.durable.context.extension.ParallelExtension; +import software.amazon.lambda.durable.extension.ExtensionContext; /** Context-free static facades for durable parallel operations. */ public final class DurableParallelOperations { private DurableParallelOperations() {} public static ParallelDurableFuture parallel(String name) { - return DurableContext.getCurrentContext().parallel(name); + return parallel(name, ParallelConfig.builder().build()); } public static ParallelDurableFuture parallel(String name, ParallelConfig config) { - return DurableContext.getCurrentContext().parallel(name, config); + return ParallelExtension.execute(ExtensionContext.getCurrentContext(), name, config); } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java index 35726edc1..ac5de46c3 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java @@ -27,6 +27,7 @@ import software.amazon.lambda.durable.config.WaitForConditionConfig; import software.amazon.lambda.durable.config.WithRetryConfig; import software.amazon.lambda.durable.context.extension.MapExtension; +import software.amazon.lambda.durable.context.extension.ParallelExtension; import software.amazon.lambda.durable.context.extension.WaitForCallbackExtension; import software.amazon.lambda.durable.context.extension.WaitForConditionExtension; import software.amazon.lambda.durable.context.extension.WithRetryExtension; @@ -41,14 +42,12 @@ import software.amazon.lambda.durable.extension.ExtensionStepFunction; import software.amazon.lambda.durable.model.MapResult; import software.amazon.lambda.durable.model.OperationDescriptor; -import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.model.WaitForConditionResult; import software.amazon.lambda.durable.operation.BaseDurableOperation; import software.amazon.lambda.durable.operation.CallbackOperation; import software.amazon.lambda.durable.operation.ChildContextOperation; import software.amazon.lambda.durable.operation.InvokeOperation; -import software.amazon.lambda.durable.operation.ParallelOperation; import software.amazon.lambda.durable.operation.StepOperation; import software.amazon.lambda.durable.operation.WaitOperation; import software.amazon.lambda.durable.util.ParameterValidator; @@ -450,18 +449,7 @@ public DurableFuture> mapAsync( @Override public ParallelDurableFuture parallel(String name, ParallelConfig config) { - Objects.requireNonNull(config, "config cannot be null"); - var operationId = nextOperationId(); - - var parallelOp = new ParallelOperation( - OperationIdentifier.of(operationId, name, OperationSubType.PARALLEL), - getDurableConfig().getSerDes(), - this, - config); - - parallelOp.execute(); - - return parallelOp; + return ParallelExtension.execute(this, name, config); } @Override diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/ParallelExtension.java b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/ParallelExtension.java new file mode 100644 index 000000000..bbf62779f --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/ParallelExtension.java @@ -0,0 +1,21 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.context.extension; + +import java.util.Objects; +import software.amazon.lambda.durable.ParallelDurableFuture; +import software.amazon.lambda.durable.config.ParallelConfig; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.util.ParameterValidator; + +/** Canonical implementation of the built-in parallel extension. */ +public final class ParallelExtension { + private ParallelExtension() {} + + public static ParallelDurableFuture execute(ExtensionContext context, String name, ParallelConfig config) { + Objects.requireNonNull(context, "context cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + ParameterValidator.validateOperationName(name); + return new ParallelExtensionFuture(context, name, config); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/ParallelExtensionFuture.java b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/ParallelExtensionFuture.java new file mode 100644 index 000000000..778fa11d2 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/ParallelExtensionFuture.java @@ -0,0 +1,245 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.context.extension; + +import static software.amazon.lambda.durable.config.NestingType.FLAT; +import static software.amazon.lambda.durable.context.extension.ExtensionConcurrencyCoordinator.ItemStatus.FAILED; +import static software.amazon.lambda.durable.context.extension.ExtensionConcurrencyCoordinator.ItemStatus.SKIPPED; +import static software.amazon.lambda.durable.model.OperationSubType.PARALLEL; +import static software.amazon.lambda.durable.model.OperationSubType.PARALLEL_BRANCH; + +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.CompletableFuture; +import java.util.function.Function; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.ParallelDurableFuture; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.CompletionConfig; +import software.amazon.lambda.durable.config.ParallelBranchConfig; +import software.amazon.lambda.durable.config.ParallelConfig; +import software.amazon.lambda.durable.config.RunInChildContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; +import software.amazon.lambda.durable.extension.ExtensionContextResult; +import software.amazon.lambda.durable.model.ParallelResult; +import software.amazon.lambda.durable.serde.SerDes; +import software.amazon.lambda.durable.util.ParameterValidator; + +final class ParallelExtensionFuture implements ParallelDurableFuture { + private final Object lock = new Object(); + private final ParallelConfig config; + private final SerDes defaultSerDes; + private final List> branches = new ArrayList<>(); + private final DurableFuture parentFuture; + private ExtensionContext childContext; + private ExtensionConcurrencyCoordinator coordinator; + private ParallelResult replayState; + private boolean registrationClosed; + + ParallelExtensionFuture(ExtensionContext context, String name, ParallelConfig config) { + this.config = config; + this.defaultSerDes = context.getDurableConfig().getSerDes(); + var parent = context.reserve(name); + this.parentFuture = parent.runInChildContextAsync( + PARALLEL.getValue(), parallelResultType(), this::executeInChildContext, parentConfig(defaultSerDes)); + } + + @Override + public DurableFuture branch( + String name, TypeToken resultType, Function function, ParallelBranchConfig config) { + Objects.requireNonNull(resultType, "resultType cannot be null"); + Objects.requireNonNull(function, "function cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + ParameterValidator.validateOperationName(name); + + synchronized (lock) { + ensureRegistrationOpen(); + var definition = new BranchDefinition<>(name, resultType, function, config); + branches.add(definition); + if (coordinator != null) { + registerBranch(definition, branches.size() - 1); + } + return definition.future; + } + } + + @Override + public ParallelResult get() { + closeRegistration(); + return rebuildResult(parentFuture.get()); + } + + @Override + public CompletableFuture completionFuture() { + return parentFuture.completionFuture(); + } + + @Override + public void close() { + if (closeRegistration()) { + parentFuture.get(); + } + } + + private ExtensionContextResult executeInChildContext() { + var replay = ExtensionContextReplayContext.getCurrentContext(); + initializeCoordinator(ExtensionContext.getCurrentContext(), replay); + var completion = replayState == null + ? coordinator.awaitCompletion() + : coordinator.awaitCompletion(expectedCompletion(replayState)); + var result = constructResult(completion); + return ExtensionContextResult.replayChildren(result, result); + } + + private void initializeCoordinator( + ExtensionContext context, ExtensionContextReplayContext replayContext) { + synchronized (lock) { + childContext = context; + replayState = replayContext.isReplayingChildren() ? replayContext.getReplayState() : null; + if (replayContext.isReplayingChildren() && replayState == null) { + throw new IllegalStateException("Missing result in completed Parallel operation"); + } + coordinator = new ExtensionConcurrencyCoordinator(config.maxConcurrency(), config.completionConfig()); + for (int index = 0; index < branches.size(); index++) { + registerBranch(branches.get(index), index); + } + if (registrationClosed) { + coordinator.closeRegistration(); + } + } + } + + private void registerBranch(BranchDefinition definition, int index) { + var reservation = childContext.reserve(definition.name); + var skipped = shouldSkip(index); + var item = coordinator.register( + () -> reservation.runInChildContextAsync( + PARALLEL_BRANCH.getValue(), + definition.resultType, + () -> definition.function.apply(DurableContext.getCurrentContext()), + branchConfig(definition.config)), + skipped); + definition.future.bind(item.future()); + } + + private boolean shouldSkip(int index) { + return replayState != null + && (replayState.statuses().size() <= index + || replayState.statuses().get(index) == ParallelResult.Status.SKIPPED); + } + + private boolean closeRegistration() { + synchronized (lock) { + if (registrationClosed) { + return false; + } + registrationClosed = true; + if (coordinator != null) { + coordinator.closeRegistration(); + } + return true; + } + } + + private void ensureRegistrationOpen() { + if (registrationClosed) { + throw new IllegalStateException("Cannot add branches after join() has been called"); + } + } + + private ParallelResult rebuildResult(ParallelResult result) { + synchronized (lock) { + if (result == null) { + return null; + } + var statuses = new ArrayList<>(result.statuses()); + while (statuses.size() < branches.size()) { + statuses.add(ParallelResult.Status.SKIPPED); + } + var succeeded = Math.toIntExact(statuses.stream() + .filter(status -> status == ParallelResult.Status.SUCCEEDED) + .count()); + var failed = Math.toIntExact(statuses.stream() + .filter(status -> status == ParallelResult.Status.FAILED) + .count()); + return new ParallelResult( + statuses.size(), + succeeded, + failed, + statuses.size() - succeeded - failed, + result.completionStatus(), + List.copyOf(statuses)); + } + } + + private static ParallelResult constructResult(ExtensionConcurrencyCoordinator.Completion completion) { + var statuses = completion.items().stream() + .map(item -> item.status() == FAILED + ? ParallelResult.Status.FAILED + : item.status() == SKIPPED ? ParallelResult.Status.SKIPPED : ParallelResult.Status.SUCCEEDED) + .toList(); + var succeeded = Math.toIntExact(statuses.stream() + .filter(status -> status == ParallelResult.Status.SUCCEEDED) + .count()); + var failed = Math.toIntExact(statuses.stream() + .filter(status -> status == ParallelResult.Status.FAILED) + .count()); + return new ParallelResult( + statuses.size(), + succeeded, + failed, + statuses.size() - succeeded - failed, + completion.completionDecision().completionStatus(), + statuses); + } + + private RunInChildContextConfig branchConfig(ParallelBranchConfig branchConfig) { + return RunInChildContextConfig.builder() + .serDes(branchConfig.serDes() == null ? defaultSerDes : branchConfig.serDes()) + .isVirtual(config.nestingType() == FLAT) + .build(); + } + + private static ExtensionConcurrencyCoordinator.ExpectedCompletionStatus expectedCompletion( + ParallelResult replayState) { + return new ExtensionConcurrencyCoordinator.ExpectedCompletionStatus( + replayState.succeeded() + replayState.failed(), + CompletionConfig.CompletionDecision.complete(replayState.completionStatus())); + } + + private static ExtensionContextConfig parentConfig(SerDes serDes) { + return ExtensionContextConfig.builder() + .childContextConfig( + RunInChildContextConfig.builder().serDes(serDes).build()) + .emitUserFunctionEvents(false) + .suppressLateChildCheckpoints(true) + .build(); + } + + private static TypeToken parallelResultType() { + return TypeToken.get(ParallelResult.class); + } + + private static final class BranchDefinition { + private final String name; + private final TypeToken resultType; + private final Function function; + private final ParallelBranchConfig config; + private final DeferredDurableFuture future = new DeferredDurableFuture<>(); + + private BranchDefinition( + String name, + TypeToken resultType, + Function function, + ParallelBranchConfig config) { + this.name = name; + this.resultType = resultType; + this.function = function; + this.config = config; + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java index 9c2cd372d..fee1953f9 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java @@ -81,7 +81,7 @@ protected BaseDurableOperation( * * @param operationIdentifier the unique identifier for this operation * @param durableContext the parent context this operation belongs to - * @param parentOperation the parent operation if this is a branch/iteration of a ConcurrencyOperation + * @param parentOperation the operation that owns late-checkpoint suppression, if any * @param isVirtual whether this is a virtual operation that should not be persisted */ protected BaseDurableOperation( @@ -261,8 +261,7 @@ protected Operation waitForOperationCompletion() { // It's important that we synchronize access to the future. Otherwise, a race condition could happen if the // completionFuture is completed by a user thread (a step or child context thread) when the execution here // is between `isOperationCompleted` and `thenRun`. - // If this operation is a branch/iteration of a ConcurrencyOperation (map or parallel), the branches/iterations - // must be completed sequentially to avoid race conditions. + // Operations sharing a late-checkpoint owner must complete sequentially to avoid races with parent completion. synchronized (parentOperation == null ? completionFuture : parentOperation.completionFuture) { if (!isOperationCompleted()) { // Add a completion stage to completionFuture so that when the completionFuture is completed, diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java index 3398e4ac9..0f51de20f 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java @@ -53,9 +53,7 @@ *

A child context runs a user function in a separate thread with its own operation counter and checkpoint log. * Operations within the child context use the child's context ID as their parentId. * - *

When created as part of a {@link ConcurrencyOperation} (e.g., parallel execution), the child notifies its parent - * on completion via {@code onItemComplete()} BEFORE closing its own child context. It also skips checkpointing if the - * parent operation has already succeeded. + *

When created with a parent operation, the child skips checkpointing if that parent has already completed. */ public class ChildContextOperation extends SerializableDurableOperation { @@ -78,14 +76,14 @@ public ChildContextOperation( this(operationIdentifier, function, resultTypeToken, config, durableContext, null); } - // child context for a ConcurrencyOperation branch + // child context with a late-checkpoint owner public ChildContextOperation( OperationIdentifier operationIdentifier, Function function, TypeToken resultTypeToken, RunInChildContextConfig config, DurableContextImpl durableContext, - ConcurrencyOperation parentOperation) { + BaseDurableOperation parentOperation) { super( operationIdentifier, resultTypeToken, @@ -202,9 +200,7 @@ private void executeChildContext() { // - add thread id/type to thread local when the step starts // - clear logger properties when the step finishes // - // When this child is part of a ConcurrencyOperation (parentOperation != null), - // we notify the parent BEFORE closing the child context. This ensures the parent - // can trigger the next queued branch while the current child context is still valid. + // A parent operation may own late-checkpoint suppression for this child. var childContext = createChildContext(contextId); try (var ignoredContext = DurableContextImpl.attachCurrentContext(childContext); var ignoredLogger = DurableLogger.attachContext()) { @@ -329,8 +325,7 @@ private void handleChildContextFailure(Throwable exception) { cachedOperationResult.set(DeserializedOperationResult.failed(translateException(op, errorObject))); // Skip checkpointing if - // - parent ConcurrencyOperation has already completed, preventing race conditions where a child finishes after - // the parent has already succeeded. + // - the owning parent operation has already completed, preventing a late child checkpoint. // - this child is not a direct child of a parent context (i.e. nestingType == FLAT), such as a parallel branch. if ((parentOperation != null && parentOperation.isOperationCompleted()) || isVirtual) { if (isVirtual) { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ConcurrencyOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ConcurrencyOperation.java deleted file mode 100644 index 913103a4c..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ConcurrencyOperation.java +++ /dev/null @@ -1,363 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; - -import java.util.ArrayList; -import java.util.Collections; -import java.util.HashSet; -import java.util.List; -import java.util.Objects; -import java.util.Queue; -import java.util.Set; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ConcurrentLinkedDeque; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Function; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; -import software.amazon.lambda.durable.DurableContext; -import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.CompletionConfig; -import software.amazon.lambda.durable.config.NestingType; -import software.amazon.lambda.durable.config.RunInChildContextConfig; -import software.amazon.lambda.durable.context.DurableContextImpl; -import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; -import software.amazon.lambda.durable.execution.OperationIdGenerator; -import software.amazon.lambda.durable.execution.SuspendExecutionException; -import software.amazon.lambda.durable.execution.ThreadType; -import software.amazon.lambda.durable.model.OperationIdentifier; -import software.amazon.lambda.durable.model.OperationSubType; -import software.amazon.lambda.durable.serde.SerDes; -import software.amazon.lambda.durable.util.ExceptionHelper; - -/** - * Abstract base class for concurrent execution of multiple child context operations. - * - *

Encapsulates shared concurrency logic: queue-based concurrency control, success/failure counting, and completion - * checking. This remains the legacy execution engine for {@code ParallelOperation}. - * - *

Key design points: - * - *

    - *
  • Does NOT register its own thread — child context threads handle all suspension - *
  • Uses a pending queue + running counter for concurrency control - *
  • Completion is determined by {@link CompletionConfig#completionDecisionFunction()} - *
  • When a child suspends, the running count is NOT decremented - *
- * - * @param the result type of this operation - */ -public abstract class ConcurrencyOperation extends SerializableDurableOperation { - - protected record ExpectedCompletionStatus(int completed, CompletionConfig.CompletionDecision completionDecision) {} - - private static final Logger logger = LoggerFactory.getLogger(ConcurrencyOperation.class); - - private final int maxConcurrency; - private final Function shouldComplete; - private final OperationIdGenerator operationIdGenerator; - private final DurableContextImpl rootContext; - private final NestingType nestingType; - - // access by context thread only - private final List> branches = Collections.synchronizedList(new ArrayList<>()); - - // put only by context thread and consume only by consumer thread - private final Queue> pendingQueue = new ConcurrentLinkedDeque<>(); - - // set by context thread and used by consumer thread - protected final AtomicBoolean isJoined = new AtomicBoolean(false); - - // used to wake up consumer thread for either new items or checking completion condition (isJoined changed) - private final AtomicReference> consumerThreadListener; - - protected ConcurrencyOperation( - OperationIdentifier operationIdentifier, - TypeToken resultTypeToken, - SerDes resultSerDes, - DurableContextImpl durableContext, - int maxConcurrency, - Function shouldComplete, - NestingType nestingType) { - super(operationIdentifier, resultTypeToken, resultSerDes, durableContext); - this.maxConcurrency = maxConcurrency; - this.shouldComplete = Objects.requireNonNull(shouldComplete, "shouldComplete cannot be null"); - this.operationIdGenerator = new OperationIdGenerator(getOperationId()); - // root context of the concurrency operation is always non-virtual - this.rootContext = durableContext.createChildContext(getOperationId(), getName(), false); - this.consumerThreadListener = new AtomicReference<>(new CompletableFuture<>()); - this.nestingType = nestingType; - } - - // ========== Template methods for subclasses ========== - - /** - * Creates a child context operation for a single item (branch or iteration). - * - * @param operationId the unique operation ID for this item - * @param name the name of this item - * @param function the user function to execute - * @param resultType the result type token - * @param branchSubType the sub-type of the branch operation - * @param the result type of the child operation - * @return a new ChildContextOperation - */ - protected ChildContextOperation createItem( - String operationId, - String name, - Function function, - TypeToken resultType, - SerDes serDes, - OperationSubType branchSubType) { - return new ChildContextOperation<>( - OperationIdentifier.of(operationId, name, branchSubType), - function, - resultType, - RunInChildContextConfig.builder() - .serDes(serDes) - .isVirtual(nestingType == NestingType.FLAT) - .build(), - rootContext, - this); - } - - /** Called when the concurrency operation completes. Subclasses define checkpointing behavior. */ - protected abstract void handleCompletion(CompletionConfig.CompletionDecision completionDecision); - - // ========== Concurrency control ========== - - /** - * Creates and enqueues an item without starting execution. Use {@link #executeItems(ExpectedCompletionStatus)} to - * begin execution after all items have been enqueued. This prevents early termination from blocking item creation - * when all items are known upfront (e.g., map operations). - */ - protected ChildContextOperation enqueueItem( - String name, - Function function, - TypeToken resultType, - SerDes serDes, - OperationSubType branchSubType, - boolean skipped) { - var operationId = this.operationIdGenerator.nextOperationId(); - var childOp = createItem(operationId, name, function, resultType, serDes, branchSubType); - branches.add(childOp); - if (!skipped) { - logger.debug("Item enqueued {}", name); - pendingQueue.add(childOp); - } - // notify the consumer thread a new item is available - notifyConsumerThread(); - return childOp; - } - - private void notifyConsumerThread() { - synchronized (completionFuture) { - consumerThreadListener.get().complete(null); - } - } - - /** Starts execution of all enqueued items. */ - protected void executeItems() { - executeItems(null); - } - - /** Starts execution of all enqueued items until the expectedCompletionStatus is met. */ - protected void executeItems(ExpectedCompletionStatus expectedCompletionStatus) { - // variables accessed only by the consumer thread. Put them here to avoid accidentally used by other threads - Set runningChildren = new HashSet<>(); - AtomicInteger succeededCount = new AtomicInteger(0); - AtomicInteger failedCount = new AtomicInteger(0); - - Runnable consumer = () -> { - try { - while (true) { - // Set a new future if it's completed so that it will be able to receive a notification of - // new items when the thread is checking completion condition and processing - // the queued items below. - synchronized (completionFuture) { - if (consumerThreadListener.get() != null - && consumerThreadListener.get().isDone()) { - consumerThreadListener.set(new CompletableFuture<>()); - } - } - - // Process completion condition. Quit the loop if the condition is met. - if (isOperationCompleted()) { - return; - } - var completionDecision = canComplete(succeededCount, failedCount, expectedCompletionStatus); - if (completionDecision != null) { - handleCompletion(completionDecision); - return; - } - - // process new items in the queue - while (runningChildren.size() < maxConcurrency && !pendingQueue.isEmpty()) { - var next = pendingQueue.poll(); - runningChildren.add(next); - logger.debug("Executing operation {}", next.getName()); - next.execute(); - } - - // If consumerThreadListener has been completed when processing above, waitForChildCompletion will - // immediately return null and repeat the above again - var child = waitForChildCompletion( - succeededCount, failedCount, runningChildren, expectedCompletionStatus); - - // child may be null if the consumer thread is woken up due to new items added or completion - // condition - // changed - if (child != null) { - if (runningChildren.contains(child)) { - runningChildren.remove(child); - onItemComplete(succeededCount, failedCount, (ChildContextOperation) child); - } else { - throw new IllegalStateException("Unexpected completion: " + child); - } - } - } - } catch (Throwable ex) { - handleException(ex); - } - }; - // run consumer in the user thread pool, although it's not a real user thread - runUserHandler(consumer, ThreadType.CONTEXT); - } - - private void handleException(Throwable ex) { - Throwable throwable = ExceptionHelper.unwrapCompletableFuture(ex); - if (throwable instanceof SuspendExecutionException suspendExecutionException) { - // Rethrow Error immediately — do not checkpoint - throw suspendExecutionException; - } - if (throwable instanceof UnrecoverableDurableExecutionException unrecoverableDurableExecutionException) { - throw terminateExecution(unrecoverableDurableExecutionException); - } - - throw terminateExecutionWithIllegalDurableOperationException( - String.format("Unexpected exception in concurrency operation: %s", throwable)); - } - - private BaseDurableOperation waitForChildCompletion( - AtomicInteger succeededCount, - AtomicInteger failedCount, - Set runningChildren, - ExpectedCompletionStatus expectedCompletionStatus) { - var threadContext = getCurrentThreadContext(); - CompletableFuture future; - - synchronized (completionFuture) { - // check again in synchronized block to prevent race conditions - if (isOperationCompleted()) { - return null; - } - var completionDecision = canComplete(succeededCount, failedCount, expectedCompletionStatus); - if (completionDecision != null) { - return null; - } - ArrayList> futures; - futures = new ArrayList<>(runningChildren.stream() - .map(BaseDurableOperation::getCompletionFuture) - .toList()); - if (futures.size() < maxConcurrency) { - // add a future to listen to the new items if there is a vacancy - consumerThreadListener.compareAndSet(null, new CompletableFuture<>()); - futures.add(consumerThreadListener.get()); - } - - // future will be completed immediately if any future of the list is already completed - future = CompletableFuture.anyOf(futures.toArray(CompletableFuture[]::new)); - // skip deregistering the current thread if there is more completed future to process - if (!future.isDone()) { - future = future.thenApply(o -> { - registerActiveThread(threadContext.threadId()); - return o; - }); - // Deregister the current thread to allow suspension - deregisterActiveThread(threadContext.threadId()); - } - } - try { - return future.thenApply(o -> (BaseDurableOperation) o).join(); - } catch (Throwable throwable) { - ExceptionHelper.sneakyThrow(ExceptionHelper.unwrapCompletableFuture(throwable)); - throw throwable; - } - } - - /** - * Called by a ChildContextOperation BEFORE it closes its child context. Updates counters, checks completion - * criteria, and either triggers the next queued item or completes the operation. - * - * @param child the child operation that completed - */ - private void onItemComplete( - AtomicInteger succeededCount, AtomicInteger failedCount, ChildContextOperation child) { - // Evaluate child result outside the lock — child.get() may block waiting for a checkpoint response. - logger.debug("OnItemComplete called by {}, Id: {}", child.getName(), child.getOperationId()); - try { - child.get(); - logger.debug("Result succeeded - {}", child.getName()); - succeededCount.incrementAndGet(); - } catch (Throwable e) { - logger.debug("Child operation {} failed: {}", child.getOperationId(), e.getMessage()); - failedCount.incrementAndGet(); - } - } - - // ========== Completion logic ========== - /** - * Checks whether the concurrency operation can be considered complete. - * - * @return the completion status if the operation is complete, or null if it should continue - */ - private CompletionConfig.CompletionDecision canComplete( - AtomicInteger succeededCount, - AtomicInteger failedCount, - ExpectedCompletionStatus expectedCompletionStatus) { - int succeeded = succeededCount.get(); - int failed = failedCount.get(); - - if (expectedCompletionStatus != null) { - if (succeeded + failed >= expectedCompletionStatus.completed) { - return expectedCompletionStatus.completionDecision; - } - - // if expected completion status is not null, we always complete all the children previously completed - return null; - } - - var decision = Objects.requireNonNull( - shouldComplete.apply(completionStatus(succeeded, failed)), - "shouldComplete must return a completion decision"); - return decision.shouldComplete() ? decision : null; - } - - private CompletionConfig.CompletionStatus completionStatus(int succeeded, int failed) { - return new CompletionConfig.CompletionStatus( - succeeded, failed, succeeded + failed, branches.size(), allItemsRegistered()); - } - - private boolean allItemsRegistered() { - return isJoined.get(); - } - - /** - * Blocks the calling thread until the concurrency operation reaches a terminal state. Validates item count, handles - * zero-branch case, then delegates to {@code waitForOperationCompletion()} from BaseDurableOperation. - */ - protected void join() { - isJoined.set(true); - - // Notify the consumer thread this concurrency operation is joined. Consumer thread need to check the - // completion condition again. - notifyConsumerThread(); - waitForOperationCompletion(); - } - - protected List> getBranches() { - return branches; - } -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java deleted file mode 100644 index 41a687e6b..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java +++ /dev/null @@ -1,192 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; - -import java.util.ArrayList; -import java.util.List; -import java.util.function.Function; -import software.amazon.awssdk.services.lambda.model.ContextOptions; -import software.amazon.awssdk.services.lambda.model.Operation; -import software.amazon.awssdk.services.lambda.model.OperationAction; -import software.amazon.awssdk.services.lambda.model.OperationUpdate; -import software.amazon.lambda.durable.DurableContext; -import software.amazon.lambda.durable.DurableFuture; -import software.amazon.lambda.durable.ParallelDurableFuture; -import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.CompletionConfig; -import software.amazon.lambda.durable.config.ParallelBranchConfig; -import software.amazon.lambda.durable.config.ParallelConfig; -import software.amazon.lambda.durable.context.DurableContextImpl; -import software.amazon.lambda.durable.execution.ExecutionManager; -import software.amazon.lambda.durable.model.OperationIdentifier; -import software.amazon.lambda.durable.model.OperationSubType; -import software.amazon.lambda.durable.model.ParallelResult; -import software.amazon.lambda.durable.serde.SerDes; - -/** - * Manages parallel execution of multiple branches as child context operations. - * - *

Extends {@link ConcurrencyOperation} to provide parallel-specific behavior: - * - *

    - *
  • Creates branches as {@link ChildContextOperation} with {@link OperationSubType#PARALLEL_BRANCH} - *
  • Checkpoints SUCCESS on the parallel context when completion criteria are met - *
  • Returns a {@link ParallelResult} summarising branch outcomes - *
- * - *

Context hierarchy: - * - * 

- * DurableContext (root)
- *   └── ParallelOperation context (ChildContextOperation with PARALLEL subtype)
- *         ├── Branch 1 context (ChildContextOperation with PARALLEL_BRANCH)
- *         ├── Branch 2 context (ChildContextOperation with PARALLEL_BRANCH)
- *         └── Branch N context (ChildContextOperation with PARALLEL_BRANCH)
- *
- */ -public class ParallelOperation extends ConcurrencyOperation implements ParallelDurableFuture { - - // this field could be written and read in different threads - private volatile ParallelResult cachedResult; - private volatile ParallelResult partialResult; - - public ParallelOperation( - OperationIdentifier operationIdentifier, - SerDes resultSerDes, - DurableContextImpl durableContext, - ParallelConfig config) { - super( - operationIdentifier, - TypeToken.get(ParallelResult.class), - resultSerDes, - durableContext, - config.maxConcurrency(), - config.completionConfig().completionDecisionFunction(), - config.nestingType()); - } - - @Override - protected void handleCompletion(CompletionConfig.CompletionDecision completionDecision) { - - var items = List.copyOf(getBranches()); - var statuses = items.stream().map(this::getParallelItemStatus).toList(); - int succeededCount = Math.toIntExact(statuses.stream() - .filter(s -> s == ParallelResult.Status.SUCCEEDED) - .count()); - int failedCount = Math.toIntExact( - statuses.stream().filter(s -> s == ParallelResult.Status.FAILED).count()); - int skippedCount = items.size() - succeededCount - failedCount; - cachedResult = new ParallelResult( - items.size(), - succeededCount, - failedCount, - skippedCount, - completionDecision.completionStatus(), - statuses); - var serializedResult = serializeAndDeserializeResult(cachedResult); - cachedResult = serializedResult.deserialized(); - - // Branches added after checkpoint will not exist in the checkpointed result, but they'll be in the returned - // value from get() method. - sendOperationUpdate(OperationUpdate.builder() - .action(OperationAction.SUCCEED) - .subType(getSubType().getValue()) - .payload(serializedResult.serialized()) - .contextOptions(ContextOptions.builder().replayChildren(true).build())); - } - - private ParallelResult.Status getParallelItemStatus(ChildContextOperation childContextOperation) { - if (!childContextOperation.isOperationCompleted()) { - return ParallelResult.Status.SKIPPED; - } - try { - childContextOperation.get(); - return ParallelResult.Status.SUCCEEDED; - } catch (Throwable t) { - return ParallelResult.Status.FAILED; - } - } - - private ParallelResult rebuildParallelResult() { - var statuses = new ArrayList<>(cachedResult.statuses()); - while (statuses.size() < getBranches().size()) { - statuses.add(ParallelResult.Status.SKIPPED); - } - - int succeededCount = Math.toIntExact(statuses.stream() - .filter(status -> status == ParallelResult.Status.SUCCEEDED) - .count()); - int failedCount = Math.toIntExact(statuses.stream() - .filter(status -> status == ParallelResult.Status.FAILED) - .count()); - return new ParallelResult( - statuses.size(), - succeededCount, - failedCount, - statuses.size() - succeededCount - failedCount, - cachedResult.completionStatus(), - List.copyOf(statuses)); - } - - @Override - protected void start() { - sendOperationUpdateAsync(OperationUpdate.builder() - .action(OperationAction.START) - .subType(getSubType().getValue())); - - executeItems(); - } - - @Override - protected void replay(Operation existing) { - // No-op: child branches handle their own replay via ChildContextOperation.replay(). - // Set replaying=true so handleSuccess() skips re-checkpointing the already-completed parallel context. - if (ExecutionManager.isTerminalStatus(existing.status())) { - // the operation is already completed, extract the branch completion status from the partialResult - partialResult = existing.contextDetails() != null - ? deserializeResult(existing.contextDetails().result()) - : null; - if (partialResult != null) { - var expected = new ExpectedCompletionStatus( - partialResult.succeeded() + partialResult.failed(), - CompletionConfig.CompletionDecision.complete(partialResult.completionStatus())); - executeItems(expected); - return; - } - } - executeItems(); - } - - @Override - public ParallelResult get() { - join(); - return rebuildParallelResult(); - } - - /** Calls {@link #get()} if not already called. Guarantees that the context is closed. */ - @Override - public void close() { - if (isJoined.get()) { - return; - } - join(); - } - - public DurableFuture branch( - String name, TypeToken resultType, Function func, ParallelBranchConfig config) { - if (isJoined.get()) { - throw new IllegalStateException("Cannot add branches after join() has been called"); - } - - var nextBranchIndex = getBranches().size(); - - // ConcurrencyOperation will skip this branch if skip=true: - // 1. if the parallel operation is already completed (partialResult is not null) - // 2. if the branch is already skipped in the partialResult or nonexistent in the partialResult - var skip = partialResult != null - && (partialResult.statuses().size() <= nextBranchIndex - || partialResult.statuses().get(nextBranchIndex) == ParallelResult.Status.SKIPPED); - var serDes = config.serDes() == null ? getContext().getDurableConfig().getSerDes() : config.serDes(); - return enqueueItem(name, func, resultType, serDes, OperationSubType.PARALLEL_BRANCH, skip); - } -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/SerializableDurableOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/SerializableDurableOperation.java index 90bbcf64c..1ef424d39 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/SerializableDurableOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/SerializableDurableOperation.java @@ -72,7 +72,7 @@ protected SerializableDurableOperation( * @param resultSerDes the serializer/deserializer for the result * @param durableContext the parent context this operation belongs to * @param isVirtual whether this is a virtual operation that should not be persisted - * @param parentOperation the parent operation if this is a branch/iteration of a ConcurrencyOperation + * @param parentOperation the operation that owns late-checkpoint suppression, if any */ protected SerializableDurableOperation( OperationIdentifier operationIdentifier, diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationsTest.java index 413b7b5dd..cd1278c84 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationsTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationsTest.java @@ -2,21 +2,21 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertSame; -import static org.mockito.Answers.CALLS_REAL_METHODS; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import java.util.function.Function; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; -import org.mockito.ArgumentCaptor; -import software.amazon.lambda.durable.config.ParallelBranchConfig; +import software.amazon.lambda.durable.config.ParallelConfig; import software.amazon.lambda.durable.context.BaseContextImpl; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextFunction; +import software.amazon.lambda.durable.extension.ExtensionOperation; class DurableParallelOperationsTest { @AfterEach @@ -26,29 +26,30 @@ void clearContext() { @Test void parallelBranchesAcceptContextFreeSuppliers() { - var context = mock(DurableContext.class); - var parallel = mock(ParallelDurableFuture.class, CALLS_REAL_METHODS); - var branchFuture = mockStringFuture(); + var context = mock(CurrentContext.class); + var parent = mock(ExtensionOperation.class); + var parentFuture = mockParallelResultFuture(); BaseContextImpl.setCurrentContext(context); - when(context.parallel("parallel")).thenReturn(parallel); - when(parallel.branch(eq("branch"), any(TypeToken.class), any(Function.class), any(ParallelBranchConfig.class))) - .thenReturn(branchFuture); + when(context.getDurableConfig()).thenReturn(DurableConfig.builder().build()); + when(context.reserve("parallel")).thenReturn(parent); + when(parent.runInChildContextAsync( + any(String.class), + any(TypeToken.class), + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) + .thenReturn(parentFuture); var result = DurableParallelOperations.parallel("parallel"); - var resultFuture = result.branch("branch", String.class, () -> "result"); - - assertSame(parallel, result); - assertSame(branchFuture, resultFuture); - @SuppressWarnings("unchecked") - var function = (ArgumentCaptor>) - (ArgumentCaptor) ArgumentCaptor.forClass(Function.class); - verify(parallel) - .branch(eq("branch"), any(TypeToken.class), function.capture(), any(ParallelBranchConfig.class)); - assertEquals("result", function.getValue().apply(mock(DurableContext.class))); + result.branch("branch", String.class, () -> "result"); + + verify(context).reserve("parallel"); + verify(context, never()).parallel(eq("parallel"), any(ParallelConfig.class)); } @SuppressWarnings("unchecked") - private DurableFuture mockStringFuture() { + private DurableFuture mockParallelResultFuture() { return mock(DurableFuture.class); } + + private interface CurrentContext extends DurableContext, ExtensionContext {} } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/ParallelExtensionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/extension/ParallelExtensionTest.java new file mode 100644 index 000000000..b602c2494 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/context/extension/ParallelExtensionTest.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.context.extension; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static software.amazon.lambda.durable.model.OperationSubType.PARALLEL; +import static software.amazon.lambda.durable.model.OperationSubType.PARALLEL_BRANCH; + +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.function.Supplier; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.NestingType; +import software.amazon.lambda.durable.config.ParallelConfig; +import software.amazon.lambda.durable.config.RunInChildContextConfig; +import software.amazon.lambda.durable.context.BaseContextImpl; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextFunction; +import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; +import software.amazon.lambda.durable.extension.ExtensionOperation; +import software.amazon.lambda.durable.model.ConcurrencyCompletionStatus; +import software.amazon.lambda.durable.model.ParallelResult; +import software.amazon.lambda.durable.serde.JacksonSerDes; + +class ParallelExtensionTest { + @Test + void executeBuildsParallelAndBranchContextsFromReservations() { + var context = mock(ExtensionContext.class); + var parent = mock(ExtensionOperation.class); + var parentFuture = mockParallelResultFuture(); + var parentCompletion = new CompletableFuture(); + var serDes = new JacksonSerDes(); + when(context.getDurableConfig()) + .thenReturn(DurableConfig.builder().withSerDes(serDes).build()); + when(context.reserve("parallel")).thenReturn(parent); + when(parent.runInChildContextAsync( + eq(PARALLEL.getValue()), + any(TypeToken.class), + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) + .thenReturn(parentFuture); + when(parentFuture.completionFuture()).thenReturn(parentCompletion); + + var parallel = ParallelExtension.execute( + context, + "parallel", + ParallelConfig.builder().nestingType(NestingType.FLAT).build()); + var firstFuture = parallel.branch("first", String.class, child -> "first"); + var secondFuture = parallel.branch("second", String.class, child -> "second"); + assertSame(parentCompletion, parallel.completionFuture()); + parallel.close(); + verify(parentFuture).get(); + + var parentFunction = extensionFunction(); + var parentConfig = ArgumentCaptor.forClass(ExtensionContextConfig.class); + verify(parent) + .runInChildContextAsync( + eq(PARALLEL.getValue()), + any(TypeToken.class), + parentFunction.capture(), + parentConfig.capture()); + assertFalse(parentConfig.getValue().emitUserFunctionEvents()); + assertTrue(parentConfig.getValue().suppressLateChildCheckpoints()); + assertSame(serDes, parentConfig.getValue().childContextConfig().serDes()); + + var child = mock(CurrentContext.class); + var first = mock(ExtensionOperation.class); + var second = mock(ExtensionOperation.class); + when(child.reserve("first")).thenReturn(first); + when(child.reserve("second")).thenReturn(second); + when(first.runInChildContextAsync( + eq(PARALLEL_BRANCH.getValue()), + eq(TypeToken.get(String.class)), + any(Supplier.class), + any(RunInChildContextConfig.class))) + .thenReturn(new CompletedFuture<>("first")); + when(second.runInChildContextAsync( + eq(PARALLEL_BRANCH.getValue()), + eq(TypeToken.get(String.class)), + any(Supplier.class), + any(RunInChildContextConfig.class))) + .thenReturn(new CompletedFuture<>("second")); + + ParallelResult result; + try (var ignoredContext = BaseContextImpl.attachCurrentContext(child); + var ignoredReplay = ExtensionContextReplayContext.attach(false, null)) { + result = parentFunction.getValue().apply().result(); + } + + assertEquals(2, result.size()); + assertEquals(2, result.succeeded()); + assertEquals(0, result.failed()); + assertEquals(0, result.skipped()); + assertEquals(ConcurrencyCompletionStatus.ALL_COMPLETED, result.completionStatus()); + assertEquals(List.of(ParallelResult.Status.SUCCEEDED, ParallelResult.Status.SUCCEEDED), result.statuses()); + assertEquals("first", firstFuture.get()); + assertEquals("second", secondFuture.get()); + + var branchConfig = ArgumentCaptor.forClass(RunInChildContextConfig.class); + verify(first) + .runInChildContextAsync( + eq(PARALLEL_BRANCH.getValue()), + eq(TypeToken.get(String.class)), + any(Supplier.class), + branchConfig.capture()); + assertTrue(branchConfig.getValue().isVirtual()); + assertSame(serDes, branchConfig.getValue().serDes()); + } + + @Test + void replaySkipsBranchesMissingFromCompletedResult() { + var context = mock(ExtensionContext.class); + var parent = mock(ExtensionOperation.class); + when(context.getDurableConfig()).thenReturn(DurableConfig.builder().build()); + when(context.reserve("parallel")).thenReturn(parent); + when(parent.runInChildContextAsync( + eq(PARALLEL.getValue()), + any(TypeToken.class), + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) + .thenReturn(mockParallelResultFuture()); + + var parallel = ParallelExtension.execute( + context, "parallel", ParallelConfig.builder().build()); + parallel.branch("skipped", String.class, child -> "skipped"); + parallel.branch("completed", String.class, child -> "completed"); + parallel.close(); + + var parentFunction = extensionFunction(); + verify(parent) + .runInChildContextAsync(eq(PARALLEL.getValue()), any(TypeToken.class), parentFunction.capture(), any()); + var child = mock(CurrentContext.class); + var skipped = mock(ExtensionOperation.class); + var completed = mock(ExtensionOperation.class); + when(child.reserve("skipped")).thenReturn(skipped); + when(child.reserve("completed")).thenReturn(completed); + when(completed.runInChildContextAsync( + eq(PARALLEL_BRANCH.getValue()), + eq(TypeToken.get(String.class)), + any(Supplier.class), + any(RunInChildContextConfig.class))) + .thenReturn(new CompletedFuture<>("completed")); + var replayState = new ParallelResult( + 2, + 1, + 0, + 1, + ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED, + List.of(ParallelResult.Status.SKIPPED, ParallelResult.Status.SUCCEEDED)); + + ParallelResult result; + try (var ignoredContext = BaseContextImpl.attachCurrentContext(child); + var ignoredReplay = ExtensionContextReplayContext.attach(true, replayState)) { + result = parentFunction.getValue().apply().result(); + } + + assertEquals(replayState, result); + verify(skipped, never()) + .runInChildContextAsync( + any(String.class), + any(TypeToken.class), + any(Supplier.class), + any(RunInChildContextConfig.class)); + } + + @Test + void getIncludesLateBranchesAsSkipped() { + var context = mock(ExtensionContext.class); + var parent = mock(ExtensionOperation.class); + var parentFuture = mockParallelResultFuture(); + when(context.getDurableConfig()).thenReturn(DurableConfig.builder().build()); + when(context.reserve("parallel")).thenReturn(parent); + when(parent.runInChildContextAsync( + eq(PARALLEL.getValue()), + any(TypeToken.class), + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) + .thenReturn(parentFuture); + when(parentFuture.get()) + .thenReturn(new ParallelResult( + 1, + 1, + 0, + 0, + ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED, + List.of(ParallelResult.Status.SUCCEEDED))); + + var parallel = ParallelExtension.execute( + context, "parallel", ParallelConfig.builder().build()); + parallel.branch("completed", String.class, child -> "completed"); + parallel.branch("late", String.class, child -> "late"); + + var result = parallel.get(); + + assertEquals(2, result.size()); + assertEquals(1, result.succeeded()); + assertEquals(0, result.failed()); + assertEquals(1, result.skipped()); + assertEquals(List.of(ParallelResult.Status.SUCCEEDED, ParallelResult.Status.SKIPPED), result.statuses()); + } + + @Test + void branchAfterCloseFails() { + var context = mock(ExtensionContext.class); + var parent = mock(ExtensionOperation.class); + when(context.getDurableConfig()).thenReturn(DurableConfig.builder().build()); + when(context.reserve("parallel")).thenReturn(parent); + when(parent.runInChildContextAsync( + eq(PARALLEL.getValue()), + any(TypeToken.class), + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) + .thenReturn(mockParallelResultFuture()); + var parallel = ParallelExtension.execute( + context, "parallel", ParallelConfig.builder().build()); + + parallel.close(); + + var exception = + assertThrows(IllegalStateException.class, () -> parallel.branch("late", String.class, child -> "late")); + assertEquals("Cannot add branches after join() has been called", exception.getMessage()); + } + + @SuppressWarnings({"rawtypes", "unchecked"}) + private ArgumentCaptor> extensionFunction() { + return (ArgumentCaptor) ArgumentCaptor.forClass(ExtensionContextFunction.class); + } + + @SuppressWarnings("unchecked") + private DurableFuture mockParallelResultFuture() { + return mock(DurableFuture.class); + } + + private interface CurrentContext extends DurableContext, ExtensionContext {} + + private record CompletedFuture(T result) implements DurableFuture { + @Override + public T get() { + return result; + } + + @Override + public CompletableFuture completionFuture() { + return CompletableFuture.completedFuture(null); + } + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java index b72af25aa..5f5d336e9 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java @@ -123,7 +123,7 @@ private ChildContextOperation createVirtualOperation(Function createOperationWithParent( - Function func, ConcurrencyOperation parent) { + Function func, BaseDurableOperation parent) { return new ChildContextOperation<>( OPERATION_IDENTIFIER, func, @@ -467,15 +467,14 @@ void replayWithNameMismatchTerminatesExecution() { assertThrows(NonDeterministicExecutionException.class, operation::execute); } - // ===== Parent ConcurrencyOperation support ===== + // ===== Parent operation support ===== /** Child skips success checkpoint when parent operation has already completed. */ @Test void childSkipsSuccessCheckpointWhenParentAlreadyCompleted() throws Exception { when(executionManager.getOperationAndUpdateReplayState("1")).thenReturn(null); - var parent = mock(ConcurrencyOperation.class); - when(parent.isOperationCompleted()).thenReturn(true); + var parent = new CompletedParentOperation(durableContext); var operation = createOperationWithParent(ctx -> "result", parent); operation.execute(); @@ -525,8 +524,7 @@ void virtualChildSucceedsWhenResultValidationDisabled() throws Exception { void childSkipsFailureCheckpointWhenParentAlreadyCompleted() throws Exception { when(executionManager.getOperationAndUpdateReplayState("1")).thenReturn(null); - var parent = mock(ConcurrencyOperation.class); - when(parent.isOperationCompleted()).thenReturn(true); + var parent = new CompletedParentOperation(durableContext); var operation = createOperationWithParent( ctx -> { @@ -540,4 +538,24 @@ void childSkipsFailureCheckpointWhenParentAlreadyCompleted() throws Exception { verify(executionManager, never()) .sendOperationUpdate(argThat(update -> update.action() == OperationAction.FAIL)); } + + private static final class CompletedParentOperation extends BaseDurableOperation { + private CompletedParentOperation(DurableContextImpl durableContext) { + super( + OperationIdentifier.of("parent", "parent", OperationSubType.RUN_IN_CHILD_CONTEXT), + durableContext, + null); + } + + @Override + protected void start() {} + + @Override + protected void replay(Operation existing) {} + + @Override + protected boolean isOperationCompleted() { + return true; + } + } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/ConcurrencyOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/ConcurrencyOperationTest.java deleted file mode 100644 index e51bd13da..000000000 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/ConcurrencyOperationTest.java +++ /dev/null @@ -1,359 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; - -import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.Mockito.*; - -import java.lang.reflect.Field; -import java.lang.reflect.InvocationTargetException; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.Executors; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicInteger; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import software.amazon.awssdk.services.lambda.model.ContextDetails; -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.config.CompletionConfig; -import software.amazon.lambda.durable.config.NestingType; -import software.amazon.lambda.durable.context.DurableContextImpl; -import software.amazon.lambda.durable.execution.ExecutionManager; -import software.amazon.lambda.durable.execution.OperationIdGenerator; -import software.amazon.lambda.durable.execution.ThreadContext; -import software.amazon.lambda.durable.execution.ThreadType; -import software.amazon.lambda.durable.model.OperationIdentifier; -import software.amazon.lambda.durable.model.OperationSubType; -import software.amazon.lambda.durable.serde.JacksonSerDes; -import software.amazon.lambda.durable.serde.SerDes; - -class ConcurrencyOperationTest { - - private static final SerDes SER_DES = new JacksonSerDes(); - private static final String OPERATION_ID = "op-1"; - private static final String CHILD_OP_1 = TestUtils.hashOperationId(OPERATION_ID + "-1"); - private static final String CHILD_OP_2 = TestUtils.hashOperationId(OPERATION_ID + "-2"); - private static final TypeToken RESULT_TYPE = TypeToken.get(Void.class); - - private DurableContextImpl durableContext; - private DurableContextImpl childContext; - private ExecutionManager executionManager; - - @BeforeEach - void setUp() { - durableContext = mock(DurableContextImpl.class); - executionManager = mock(ExecutionManager.class); - - var childContext = mock(DurableContextImpl.class); - this.childContext = childContext; - when(childContext.getExecutionManager()).thenReturn(executionManager); - when(childContext.getDurableConfig()) - .thenReturn(DurableConfig.builder() - .withExecutorService(Executors.newCachedThreadPool()) - .build()); - - when(durableContext.getExecutionManager()).thenReturn(executionManager); - when(durableContext.getDurableConfig()) - .thenReturn(DurableConfig.builder() - .withExecutorService(Executors.newCachedThreadPool()) - .build()); - when(durableContext.createChildContext(anyString(), anyString(), anyBoolean())) - .thenReturn(childContext); - when(executionManager.getCurrentThreadContext()).thenReturn(new ThreadContext("Root", ThreadType.CONTEXT)); - // All child operations are NOT in replay - when(executionManager.getOperationAndUpdateReplayState(anyString())).thenReturn(null); - // Simulate the real backend: the parent concurrency operation is available in storage after completion - // so that waitForOperationCompletion() can find it. TestConcurrencyOperation.handleSuccess/Failure are no-ops - // (no checkpoint sent), so we stub this unconditionally for OPERATION_ID. - when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)) - .thenReturn(Operation.builder() - .id(OPERATION_ID) - .name("test-concurrency") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL.getValue()) - .status(OperationStatus.SUCCEEDED) - .build()); - when(executionManager.sendOperationUpdate(any())).thenReturn(CompletableFuture.completedFuture(null)); - } - - private TestConcurrencyOperation createOperation(CompletionConfig completionConfig) throws Exception { - return new TestConcurrencyOperation( - OperationIdentifier.of(OPERATION_ID, "test-concurrency", OperationSubType.PARALLEL), - RESULT_TYPE, - SER_DES, - durableContext, - Integer.MAX_VALUE, - completionConfig); - } - - private void setOperationIdGenerator(ConcurrencyOperation op, OperationIdGenerator mockGenerator) - throws Exception { - Field field = ConcurrencyOperation.class.getDeclaredField("operationIdGenerator"); - field.setAccessible(true); - field.set(op, mockGenerator); - } - - private CompletionConfig.CompletionDecision canComplete( - ConcurrencyOperation op, int succeededCount, int failedCount) throws Exception { - var method = ConcurrencyOperation.class.getDeclaredMethod( - "canComplete", - AtomicInteger.class, - AtomicInteger.class, - ConcurrencyOperation.ExpectedCompletionStatus.class); - method.setAccessible(true); - try { - return (CompletionConfig.CompletionDecision) - method.invoke(op, new AtomicInteger(succeededCount), new AtomicInteger(failedCount), null); - } catch (InvocationTargetException e) { - var cause = e.getCause(); - if (cause instanceof RuntimeException runtimeException) { - throw runtimeException; - } - if (cause instanceof Error error) { - throw error; - } - throw e; - } - } - - // ===== Callback cycle tests ===== - - @Test - void allChildrenAlreadySucceed_callsHandleSuccess() throws Exception { - when(executionManager.getOperationAndUpdateReplayState(CHILD_OP_1)) - .thenReturn(Operation.builder() - .id(CHILD_OP_1) - .name("branch-1") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL_BRANCH.getValue()) - .status(OperationStatus.SUCCEEDED) - .contextDetails( - ContextDetails.builder().result("\"result-1\"").build()) - .build()); - when(executionManager.getOperationAndUpdateReplayState(CHILD_OP_2)) - .thenReturn(Operation.builder() - .id(CHILD_OP_2) - .name("branch-2") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL_BRANCH.getValue()) - .status(OperationStatus.SUCCEEDED) - .contextDetails( - ContextDetails.builder().result("\"result-2\"").build()) - .build()); - - var functionCalled = new AtomicBoolean(false); - var op = createOperation(CompletionConfig.allSuccessful()); - op.execute(); - op.enqueueItem( - "branch-1", - ctx1 -> { - functionCalled.set(true); - return "result-1"; - }, - TypeToken.get(String.class), - SER_DES, - OperationSubType.PARALLEL_BRANCH, - false); - op.enqueueItem( - "branch-2", - ctx -> { - functionCalled.set(true); - return "result-2"; - }, - TypeToken.get(String.class), - SER_DES, - OperationSubType.PARALLEL_BRANCH, - false); - - op.exposedJoin(); - - assertTrue(op.isSuccessHandled()); - assertFalse(op.isFailureHandled()); - var items = op.getBranches(); - assertEquals(2, items.size()); - assertTrue(items.stream().allMatch(b -> b.getOperation().status().equals(OperationStatus.SUCCEEDED))); - assertFalse(functionCalled.get(), "Functions should not be called during SUCCEEDED replay"); - } - - @Test - void someChildrenSkipped_skippedChildrenNotExecuted() throws Exception { - when(executionManager.getOperationAndUpdateReplayState(CHILD_OP_1)) - .thenReturn(Operation.builder() - .id(CHILD_OP_1) - .name("branch-1") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL_BRANCH.getValue()) - .status(OperationStatus.STARTED) - .build()); - when(executionManager.getOperationAndUpdateReplayState(CHILD_OP_2)) - .thenReturn(Operation.builder() - .id(CHILD_OP_2) - .name("branch-2") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL_BRANCH.getValue()) - .status(OperationStatus.SUCCEEDED) - .contextDetails( - ContextDetails.builder().result("\"result-2\"").build()) - .build()); - - var functionCalled = new AtomicBoolean(false); - var op = createOperation(CompletionConfig.minSuccessful(1)); - op.execute(); - op.enqueueItem( - "branch-1", - ctx1 -> { - functionCalled.set(true); - return "result-1"; - }, - TypeToken.get(String.class), - SER_DES, - OperationSubType.PARALLEL_BRANCH, - true); - op.enqueueItem( - "branch-2", - ctx -> { - functionCalled.set(true); - return "result-2"; - }, - TypeToken.get(String.class), - SER_DES, - OperationSubType.PARALLEL_BRANCH, - false); - - op.exposedJoin(); - - assertTrue(op.isSuccessHandled()); - assertFalse(op.isFailureHandled()); - var items = op.getBranches(); - assertEquals(2, items.size()); - assertFalse(items.get(0).isOperationCompleted()); - assertEquals(OperationStatus.STARTED, items.get(0).getOperation().status()); - assertEquals(OperationStatus.SUCCEEDED, items.get(1).getOperation().status()); - assertFalse(functionCalled.get(), "Functions should not be called during SUCCEEDED replay"); - } - - @Test - void singleChildAlreadySucceeds_fullCycle() throws Exception { - when(executionManager.getOperationAndUpdateReplayState(CHILD_OP_1)) - .thenReturn(Operation.builder() - .id(CHILD_OP_1) - .name("only-branch") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL_BRANCH.getValue()) - .status(OperationStatus.SUCCEEDED) - .contextDetails( - ContextDetails.builder().result("\"done\"").build()) - .build()); - - var functionCalled = new AtomicBoolean(false); - var op = createOperation(CompletionConfig.minSuccessful(1)); - op.enqueueItem( - "only-branch", - ctx -> { - functionCalled.set(true); - return "done"; - }, - TypeToken.get(String.class), - SER_DES, - OperationSubType.PARALLEL_BRANCH, - false); - - op.execute(); - op.exposedJoin(); - - assertTrue(op.isSuccessHandled()); - var items = op.getBranches(); - assertEquals(1, items.size()); - assertEquals(OperationStatus.SUCCEEDED, items.get(0).getOperation().status()); - assertFalse(functionCalled.get(), "Function should not be called during SUCCEEDED replay"); - } - - @Test - void canComplete_whenShouldCompleteReturnsNull_shouldThrow() throws Exception { - var op = createOperation(CompletionConfig.shouldComplete(status -> null)); - - var exception = assertThrows(NullPointerException.class, () -> canComplete(op, 0, 0)); - - assertEquals("shouldComplete must return a completion decision", exception.getMessage()); - } - - // ===== Test subclass ===== - - static class TestConcurrencyOperation extends ConcurrencyOperation { - - private boolean successHandled = false; - private boolean failureHandled = false; - private final AtomicInteger executingCount = new AtomicInteger(0); - private DurableContextImpl lastParentContext; - - TestConcurrencyOperation( - OperationIdentifier operationIdentifier, - TypeToken resultTypeToken, - SerDes resultSerDes, - DurableContextImpl durableContext, - int maxConcurrency, - CompletionConfig completionConfig) { - super( - operationIdentifier, - resultTypeToken, - resultSerDes, - durableContext, - maxConcurrency, - completionConfig.completionDecisionFunction(), - NestingType.NESTED); - } - - @Override - protected void handleCompletion(CompletionConfig.CompletionDecision completionDecision) { - successHandled = true; - // Simulate the checkpoint ACK that a real subclass would receive after sendOperationUpdate. - // This drives completionFuture to completion so waitForOperationCompletion() unblocks. - onCheckpointComplete(Operation.builder() - .id(getOperationId()) - .status(OperationStatus.SUCCEEDED) - .build()); - } - - @Override - protected void start() { - executeItems(); - } - - @Override - protected void replay(Operation existing) { - executeItems(); - } - - @Override - public Void get() { - return null; - } - - void exposedJoin() { - join(); - } - - int getExecutingCount() { - return executingCount.get(); - } - - boolean isSuccessHandled() { - return successHandled; - } - - boolean isFailureHandled() { - return failureHandled; - } - - DurableContextImpl getLastParentContext() { - return lastParentContext; - } - } -} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/ParallelOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/ParallelOperationTest.java deleted file mode 100644 index e02287cdc..000000000 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/ParallelOperationTest.java +++ /dev/null @@ -1,579 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; - -import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.ArgumentMatchers.argThat; -import static org.mockito.Mockito.*; - -import java.util.List; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.Executors; -import java.util.concurrent.TimeUnit; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import software.amazon.awssdk.services.lambda.model.ContextDetails; -import software.amazon.awssdk.services.lambda.model.Operation; -import software.amazon.awssdk.services.lambda.model.OperationAction; -import software.amazon.awssdk.services.lambda.model.OperationStatus; -import software.amazon.awssdk.services.lambda.model.OperationType; -import software.amazon.awssdk.services.lambda.model.OperationUpdate; -import software.amazon.lambda.durable.DurableConfig; -import software.amazon.lambda.durable.TestUtils; -import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.CompletionConfig; -import software.amazon.lambda.durable.config.ParallelBranchConfig; -import software.amazon.lambda.durable.config.ParallelConfig; -import software.amazon.lambda.durable.context.DurableContextImpl; -import software.amazon.lambda.durable.execution.ExecutionManager; -import software.amazon.lambda.durable.execution.ThreadContext; -import software.amazon.lambda.durable.execution.ThreadType; -import software.amazon.lambda.durable.model.ConcurrencyCompletionStatus; -import software.amazon.lambda.durable.model.OperationIdentifier; -import software.amazon.lambda.durable.model.OperationSubType; -import software.amazon.lambda.durable.model.ParallelResult; -import software.amazon.lambda.durable.serde.JacksonSerDes; -import software.amazon.lambda.durable.serde.SerDes; - -class ParallelOperationTest { - - private static final SerDes SER_DES = new JacksonSerDes(); - private static final String OPERATION_ID = "parallel-op-1"; - private static final String CHILD_OP_1 = TestUtils.hashOperationId(OPERATION_ID + "-1"); - private static final String CHILD_OP_2 = TestUtils.hashOperationId(OPERATION_ID + "-2"); - - private DurableContextImpl durableContext; - private ExecutionManager executionManager; - // Thread-safe backing store for getOperationAndUpdateReplayState. - // Tests pre-populate this; doAnswer writes here before firing onCheckpointComplete, - // guaranteeing visibility to any thread that reads after the future unblocks. - private ConcurrentHashMap operationStore; - private CountDownLatch parallelCheckpointLatch; - - @BeforeEach - void setUp() { - durableContext = mock(DurableContextImpl.class); - executionManager = mock(ExecutionManager.class); - operationStore = new ConcurrentHashMap<>(); - parallelCheckpointLatch = new CountDownLatch(1); - - when(executionManager.getCurrentThreadContext()).thenReturn(new ThreadContext(null, ThreadType.CONTEXT)); - // Delegate to operationStore so all reads see the latest write, regardless of thread. - when(executionManager.getOperationAndUpdateReplayState(anyString())) - .thenAnswer(inv -> operationStore.get(inv.getArgument(0))); - - var childContext = mock(DurableContextImpl.class); - when(childContext.getExecutionManager()).thenReturn(executionManager); - when(childContext.getDurableConfig()) - .thenReturn(DurableConfig.builder() - .withExecutorService(Executors.newCachedThreadPool()) - .build()); - - when(durableContext.getExecutionManager()).thenReturn(executionManager); - when(durableContext.getDurableConfig()) - .thenReturn(DurableConfig.builder() - .withExecutorService(Executors.newCachedThreadPool()) - .build()); - when(durableContext.createChildContext(anyString(), anyString(), eq(false))) - .thenReturn(childContext); - - // Capture registered operations so we can drive onCheckpointComplete callbacks. - var registeredOps = new ConcurrentHashMap(); - doAnswer(inv -> { - BaseDurableOperation op = inv.getArgument(0); - registeredOps.put(op.getOperationId(), op); - return null; - }) - .when(executionManager) - .registerOperation(any()); - - // Simulate the real backend for all sendOperationUpdate calls. - // For SUCCEED on the parallel op: write to operationStore first (establishes happens-before - // via ConcurrentHashMap's volatile semantics), then fire onCheckpointComplete to unblock join(). - // This ordering guarantees getOperationAndUpdateReplayState() never returns null after unblocking. - var succeededParallelOp = Operation.builder() - .id(OPERATION_ID) - .name("test-parallel") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL.getValue()) - .status(OperationStatus.SUCCEEDED) - .build(); - doAnswer(inv -> { - var update = (OperationUpdate) inv.getArgument(0); - - if (OPERATION_ID.equals(update.id()) && update.action() == OperationAction.SUCCEED) { - // Write before completing the future — ConcurrentHashMap guarantees visibility. - operationStore.put(OPERATION_ID, succeededParallelOp); - var op = registeredOps.get(OPERATION_ID); - if (op != null) { - op.onCheckpointComplete(succeededParallelOp); - } - parallelCheckpointLatch.countDown(); - } - return CompletableFuture.completedFuture(null); - }) - .when(executionManager) - .sendOperationUpdate(any()); - } - - private ParallelOperation createOperation(CompletionConfig completionConfig) { - var op = new ParallelOperation( - OperationIdentifier.of(OPERATION_ID, "test-parallel", OperationSubType.PARALLEL), - SER_DES, - durableContext, - ParallelConfig.builder().completionConfig(completionConfig).build()); - - op.execute(); - return op; - } - - // ===== Branch creation delegates to ConcurrencyOperation ===== - - @Test - void branchCreation_createsBranchWithParallelBranchSubType() { - var op = createOperation(CompletionConfig.allSuccessful()); - - var childOp = op.enqueueItem( - "branch-1", - ctx -> "result", - TypeToken.get(String.class), - SER_DES, - OperationSubType.PARALLEL_BRANCH, - false); - - assertNotNull(childOp); - assertEquals(OperationSubType.PARALLEL_BRANCH, childOp.getSubType()); - } - - @Test - void branchCreation_multipleBranchesAllCreated() { - var op = createOperation(CompletionConfig.allSuccessful()); - - op.enqueueItem( - "branch-1", - ctx2 -> "r1", - TypeToken.get(String.class), - SER_DES, - OperationSubType.PARALLEL_BRANCH, - false); - op.enqueueItem( - "branch-2", - ctx1 -> "r2", - TypeToken.get(String.class), - SER_DES, - OperationSubType.PARALLEL_BRANCH, - false); - op.enqueueItem( - "branch-3", ctx -> "r3", TypeToken.get(String.class), SER_DES, OperationSubType.PARALLEL_BRANCH, false); - - assertEquals(3, op.getBranches().size()); - } - - @Test - void branchCreation_childOperationHasParentReference() throws Exception { - var op = createOperation(CompletionConfig.allSuccessful()); - - // The child operation should be a ChildContextOperation with this op as parent - var childOp = op.enqueueItem( - "branch-1", - ctx -> "result", - TypeToken.get(String.class), - SER_DES, - OperationSubType.PARALLEL_BRANCH, - false); - - assertNotNull(childOp); - // Verify it's a ChildContextOperation (the concrete type returned by createItem) - assertInstanceOf(ChildContextOperation.class, childOp); - } - - // ===== All branches succeed ===== - - @Test - void allBranchesSucceed_sendsSucceedCheckpointAndReturnsCorrectResult() throws Exception { - when(executionManager.getOperationAndUpdateReplayState(CHILD_OP_1)) - .thenReturn(Operation.builder() - .id(CHILD_OP_1) - .name("branch-1") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL_BRANCH.getValue()) - .status(OperationStatus.SUCCEEDED) - .contextDetails( - ContextDetails.builder().result("\"r1\"").build()) - .build()); - when(executionManager.getOperationAndUpdateReplayState(CHILD_OP_2)) - .thenReturn(Operation.builder() - .id(CHILD_OP_2) - .name("branch-2") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL_BRANCH.getValue()) - .status(OperationStatus.SUCCEEDED) - .contextDetails( - ContextDetails.builder().result("\"r2\"").build()) - .build()); - - var op = createOperation(CompletionConfig.allSuccessful()); - op.enqueueItem( - "branch-1", - ctx1 -> "r1", - TypeToken.get(String.class), - SER_DES, - OperationSubType.PARALLEL_BRANCH, - false); - op.enqueueItem( - "branch-2", ctx -> "r2", TypeToken.get(String.class), SER_DES, OperationSubType.PARALLEL_BRANCH, false); - - var result = op.get(); - - verify(executionManager).sendOperationUpdate(argThat(update -> update.action() == OperationAction.SUCCEED)); - assertEquals(2, result.size()); - assertEquals(2, result.succeeded()); - assertEquals(0, result.failed()); - assertEquals(ConcurrencyCompletionStatus.ALL_COMPLETED, result.completionStatus()); - assertTrue(result.completionStatus().isSucceeded()); - } - - // ===== MinSuccessful satisfaction ===== - - @Test - void minSuccessful_completesWhenThresholdMetAndReturnsResult() throws Exception { - when(executionManager.getOperationAndUpdateReplayState(CHILD_OP_1)) - .thenReturn(Operation.builder() - .id(CHILD_OP_1) - .name("branch-1") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL_BRANCH.getValue()) - .status(OperationStatus.SUCCEEDED) - .contextDetails( - ContextDetails.builder().result("\"r1\"").build()) - .build()); - - var op = createOperation(CompletionConfig.minSuccessful(1)); - op.enqueueItem( - "branch-1", ctx -> "r1", TypeToken.get(String.class), SER_DES, OperationSubType.PARALLEL_BRANCH, false); - - var result = op.get(); - - verify(executionManager).sendOperationUpdate(argThat(update -> update.action() == OperationAction.SUCCEED)); - assertEquals(1, result.size()); - assertEquals(1, result.succeeded()); - assertEquals(0, result.failed()); - assertEquals(ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED, result.completionStatus()); - assertTrue(result.completionStatus().isSucceeded()); - } - - @Test - void minSuccessful_branchRegisteredAfterCheckpointIsIncludedAsSkipped() throws Exception { - when(executionManager.getOperationAndUpdateReplayState(CHILD_OP_1)) - .thenReturn(Operation.builder() - .id(CHILD_OP_1) - .name("branch-1") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL_BRANCH.getValue()) - .status(OperationStatus.SUCCEEDED) - .contextDetails( - ContextDetails.builder().result("\"r1\"").build()) - .build()); - - var op = createOperation(CompletionConfig.minSuccessful(1)); - op.enqueueItem( - "branch-1", ctx -> "r1", TypeToken.get(String.class), SER_DES, OperationSubType.PARALLEL_BRANCH, false); - - assertTrue(parallelCheckpointLatch.await(5, TimeUnit.SECONDS)); - - op.enqueueItem( - "branch-2", ctx -> "r2", TypeToken.get(String.class), SER_DES, OperationSubType.PARALLEL_BRANCH, false); - - var result = op.get(); - - assertEquals(2, result.size()); - assertEquals(1, result.succeeded()); - assertEquals(0, result.failed()); - assertEquals(1, result.skipped()); - assertEquals(ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED, result.completionStatus()); - assertEquals(List.of(ParallelResult.Status.SUCCEEDED, ParallelResult.Status.SKIPPED), result.statuses()); - assertEquals(result.size(), result.statuses().size()); - assertEquals(result.size(), result.succeeded() + result.failed() + result.skipped()); - } - - @Test - void minSuccessful_notExecuteSkippedBranchWhenReplay() { - when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)) - .thenReturn(Operation.builder() - .id(OPERATION_ID) - .name("test-parallel") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL.getValue()) - .status(OperationStatus.SUCCEEDED) - .contextDetails(ContextDetails.builder() - .result( - "{\"size\": 2, \"skipped\": 1, \"succeeded\": 1, \"completionStatus\": \"MIN_SUCCESSFUL_REACHED\", \"statuses\":[\"SKIPPED\", \"SUCCEEDED\"]}") - .build()) - .build()); - when(executionManager.getOperationAndUpdateReplayState(CHILD_OP_2)) - .thenReturn(Operation.builder() - .id(CHILD_OP_2) - .name("branch-2") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL_BRANCH.getValue()) - .status(OperationStatus.SUCCEEDED) - .contextDetails( - ContextDetails.builder().result("\"r2\"").build()) - .build()); - - var op = createOperation(CompletionConfig.minSuccessful(1)); - op.branch( - "branch-1", - TypeToken.get(String.class), - ctx -> "r1", - ParallelBranchConfig.builder().serDes(SER_DES).build()); - op.branch( - "branch-2", - TypeToken.get(String.class), - ctx -> "r2", - ParallelBranchConfig.builder().serDes(SER_DES).build()); - - var result = op.get(); - - verify(executionManager, never()).sendOperationUpdate(any()); - assertEquals(2, result.size()); - assertEquals(1, result.succeeded()); - assertEquals(0, result.failed()); - assertEquals(1, result.skipped()); - assertEquals(ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED, result.completionStatus()); - assertTrue(result.completionStatus().isSucceeded()); - } - - // ===== Context hierarchy ===== - - @Test - void contextHierarchy_branchesUseParallelContextAsParent() throws Exception { - // Verify that branches are created with the parallel operation's context (durableContext) - // as their parent — not some other context - var op = createOperation(CompletionConfig.allSuccessful()); - - var childOp = op.enqueueItem( - "branch-1", - ctx -> "result", - TypeToken.get(String.class), - SER_DES, - OperationSubType.PARALLEL_BRANCH, - false); - - // The child operation should be registered in the execution manager - // (BaseDurableOperation constructor calls executionManager.registerOperation) - verify(executionManager, atLeastOnce()).registerOperation(any()); - assertNotNull(childOp); - } - - // ===== Replay ===== - - @Test - void replay_fromStartedState_sendsSucceedCheckpointAndReturnsResult() throws Exception { - when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)) - .thenReturn(Operation.builder() - .id(OPERATION_ID) - .name("test-parallel") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL.getValue()) - .status(OperationStatus.STARTED) - .build()); - when(executionManager.getOperationAndUpdateReplayState(CHILD_OP_1)) - .thenReturn(Operation.builder() - .id(CHILD_OP_1) - .name("branch-1") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL_BRANCH.getValue()) - .status(OperationStatus.SUCCEEDED) - .contextDetails( - ContextDetails.builder().result("\"r1\"").build()) - .build()); - when(executionManager.getOperationAndUpdateReplayState(CHILD_OP_2)) - .thenReturn(Operation.builder() - .id(CHILD_OP_2) - .name("branch-2") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL_BRANCH.getValue()) - .status(OperationStatus.SUCCEEDED) - .contextDetails( - ContextDetails.builder().result("\"r2\"").build()) - .build()); - - var op = createOperation(CompletionConfig.allSuccessful()); - op.enqueueItem( - "branch-1", - ctx1 -> "r1", - TypeToken.get(String.class), - SER_DES, - OperationSubType.PARALLEL_BRANCH, - false); - op.enqueueItem( - "branch-2", ctx -> "r2", TypeToken.get(String.class), SER_DES, OperationSubType.PARALLEL_BRANCH, false); - - var result = op.get(); - - verify(executionManager, never()) - .sendOperationUpdate(argThat(update -> update.action() == OperationAction.START)); - verify(executionManager, times(1)) - .sendOperationUpdate(argThat(update -> update.action() == OperationAction.SUCCEED)); - assertEquals(2, result.size()); - assertEquals(2, result.succeeded()); - assertEquals(0, result.failed()); - assertEquals(ConcurrencyCompletionStatus.ALL_COMPLETED, result.completionStatus()); - } - - @Test - void replay_fromSucceededState_skipsCheckpointAndReturnsResult() throws Exception { - when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)) - .thenReturn(Operation.builder() - .id(OPERATION_ID) - .name("test-parallel") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL.getValue()) - .status(OperationStatus.SUCCEEDED) - .build()); - when(executionManager.getOperationAndUpdateReplayState(CHILD_OP_1)) - .thenReturn(Operation.builder() - .id(CHILD_OP_1) - .name("branch-1") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL_BRANCH.getValue()) - .status(OperationStatus.SUCCEEDED) - .contextDetails( - ContextDetails.builder().result("\"r1\"").build()) - .build()); - when(executionManager.getOperationAndUpdateReplayState(CHILD_OP_2)) - .thenReturn(Operation.builder() - .id(CHILD_OP_2) - .name("branch-2") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL_BRANCH.getValue()) - .status(OperationStatus.SUCCEEDED) - .contextDetails( - ContextDetails.builder().result("\"r2\"").build()) - .build()); - - var op = createOperation(CompletionConfig.allSuccessful()); - op.enqueueItem( - "branch-1", - ctx1 -> "r1", - TypeToken.get(String.class), - SER_DES, - OperationSubType.PARALLEL_BRANCH, - false); - op.enqueueItem( - "branch-2", ctx -> "r2", TypeToken.get(String.class), SER_DES, OperationSubType.PARALLEL_BRANCH, false); - - var result = op.get(); - - verify(executionManager, never()) - .sendOperationUpdate(argThat(update -> update.action() == OperationAction.START)); - verify(executionManager, never()) - .sendOperationUpdate(argThat(update -> update.action() == OperationAction.SUCCEED)); - assertEquals(2, result.size()); - assertEquals(2, result.succeeded()); - assertEquals(0, result.failed()); - assertEquals(ConcurrencyCompletionStatus.ALL_COMPLETED, result.completionStatus()); - } - - // ===== Branch failure sends SUCCEED checkpoint and returns result ===== - - @Test - void branchFailure_sendsSucceedCheckpointAndReturnsFailureCounts() { - when(executionManager.getOperationAndUpdateReplayState(CHILD_OP_1)) - .thenReturn(Operation.builder() - .id(CHILD_OP_1) - .name("branch-1") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL_BRANCH.getValue()) - .status(OperationStatus.FAILED) - .build()); - - var op = createOperation(CompletionConfig.allSuccessful()); - op.enqueueItem( - "branch-1", - ctx -> { - throw new RuntimeException("branch failed"); - }, - TypeToken.get(String.class), - SER_DES, - OperationSubType.PARALLEL_BRANCH, - false); - - var result = assertDoesNotThrow(op::get); - - verify(executionManager).sendOperationUpdate(argThat(update -> update.action() == OperationAction.SUCCEED)); - verify(executionManager, never()) - .sendOperationUpdate(argThat(update -> update.action() == OperationAction.FAIL)); - assertEquals(1, result.size()); - assertEquals(0, result.succeeded()); - assertEquals(1, result.failed()); - assertFalse(result.completionStatus().isSucceeded()); - } - - @Test - void get_someBranchesFail_returnsCorrectCountsAndFailureStatus() throws Exception { - when(executionManager.getOperationAndUpdateReplayState(CHILD_OP_1)) - .thenReturn(Operation.builder() - .id(CHILD_OP_1) - .name("branch-1") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL_BRANCH.getValue()) - .status(OperationStatus.SUCCEEDED) - .contextDetails( - ContextDetails.builder().result("\"r1\"").build()) - .build()); - when(executionManager.getOperationAndUpdateReplayState(CHILD_OP_2)) - .thenReturn(Operation.builder() - .id(CHILD_OP_2) - .name("branch-2") - .type(OperationType.CONTEXT) - .subType(OperationSubType.PARALLEL_BRANCH.getValue()) - .status(OperationStatus.FAILED) - .build()); - - // toleratedFailureCount=1 so the operation completes after both branches finish - var op = createOperation(CompletionConfig.toleratedFailureCount(1)); - op.enqueueItem( - "branch-1", - ctx1 -> "r1", - TypeToken.get(String.class), - SER_DES, - OperationSubType.PARALLEL_BRANCH, - false); - op.enqueueItem( - "branch-2", - ctx -> { - throw new RuntimeException("branch failed"); - }, - TypeToken.get(String.class), - SER_DES, - OperationSubType.PARALLEL_BRANCH, - false); - - var result = op.get(); - - verify(executionManager).sendOperationUpdate(argThat(update -> update.action() == OperationAction.SUCCEED)); - assertEquals(2, result.size()); - assertEquals(1, result.succeeded()); - assertEquals(1, result.failed()); - assertTrue(result.completionStatus().isSucceeded()); - } - - @Test - void get_zeroBranches_returnsAllZerosAndAllCompletedStatus() throws Exception { - var op = createOperation(CompletionConfig.allSuccessful()); - - var result = op.get(); - - assertEquals(0, result.size()); - assertEquals(0, result.succeeded()); - assertEquals(0, result.failed()); - assertEquals(ConcurrencyCompletionStatus.ALL_COMPLETED, result.completionStatus()); - verify(executionManager).sendOperationUpdate(argThat(update -> update.action() == OperationAction.SUCCEED)); - } -} From 000e2908888ba944afa16f0e6adfd7d0f5ce7f27 Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 07:41:15 +0000 Subject: [PATCH 23/40] refactor: unify durable operation extension architecture --- docs/adr/006-custom-extension-operations.md | 102 +- docs/advanced/extensions.md | 242 ++++- docs/design.md | 114 +- .../2026-08-10-custom-extension-operations.md | 608 ----------- .../2026-08-10-migrate-built-in-extensions.md | 995 ------------------ ...8-10-custom-extension-operations-design.md | 409 ------- ...8-10-migrate-built-in-extensions-design.md | 570 ---------- .../DeserializationFailedParallelExample.java | 5 +- .../ExtensionOperationIntegrationTest.java | 97 +- .../lambda/durable/PluginIntegrationTest.java | 17 +- .../StaticOperationsIntegrationTest.java | 165 ++- .../durable/extension/PairOperations.java | 66 +- .../amazon/lambda/durable/DurableContext.java | 36 +- .../lambda/durable/DurableCoreOperations.java | 157 --- .../amazon/lambda/durable/DurableHandler.java | 28 +- .../lambda/durable/DurableMapOperations.java | 69 -- .../durable/DurableParallelOperations.java | 20 - .../DurableWaitForCallbackOperations.java | 64 -- .../DurableWaitForConditionOperations.java | 82 -- .../durable/DurableWithRetryOperations.java | 44 - .../amazon/lambda/durable/MapItemContext.java | 3 +- .../lambda/durable/ParallelDurableFuture.java | 2 +- .../durable/WaitForCallbackContext.java | 3 +- .../lambda/durable/WithRetryContext.java | 3 +- .../lambda/durable/config/CallbackConfig.java | 10 + .../lambda/durable/config/InvokeConfig.java | 10 + .../lambda/durable/config/MapConfig.java | 12 + .../durable/config/ParallelBranchConfig.java | 8 + .../lambda/durable/config/ParallelConfig.java | 10 + .../config/RunInChildContextConfig.java | 9 + .../lambda/durable/config/StepConfig.java | 10 + .../durable/config/WaitForCallbackConfig.java | 10 + .../config/WaitForConditionConfig.java | 10 + .../durable/config/WithRetryConfig.java | 9 + .../durable/context/DurableContextImpl.java | 351 +----- .../context/ExtensionOperationImpl.java | 141 --- .../context/extension/ParallelExtension.java | 21 - .../extension/WaitForCallbackExtension.java | 117 -- .../extension/WaitForConditionExtension.java | 59 -- .../durable/execution/ExecutionManager.java | 8 +- .../execution/OperationIdGenerator.java | 8 +- .../extension/ExtensionCallbackConfig.java | 70 ++ .../extension/ExtensionInvokeConfig.java | 66 ++ .../durable/extension/ExtensionOperation.java | 283 +---- .../extension/ExtensionOperationImpl.java | 160 +++ .../extension/ExtensionStepConfig.java | 39 + .../durable/model/OperationDescriptor.java | 46 - .../durable/model/OperationIdentifier.java | 33 +- .../DeferredDurableFuture.java | 2 +- .../operation/DurableCallbackOperation.java | 119 +++ .../operation/DurableContextOperation.java | 137 +++ .../operation/DurableInvokeOperation.java | 142 +++ .../DurableMapOperation.java} | 201 +++- .../operation/DurableParallelOperation.java | 140 +++ .../operation/DurableStepOperation.java | 149 +++ .../DurableWaitForCallbackOperation.java | 221 ++++ .../DurableWaitForConditionOperation.java | 186 ++++ .../operation/DurableWaitOperation.java | 31 + .../DurableWithRetryOperation.java} | 94 +- .../OperationConcurrencyCoordinator.java} | 6 +- .../operation/OperationConfigAdapters.java | 16 + .../ParallelOperationFuture.java} | 41 +- .../WaitForConditionFuture.java | 2 +- .../durable/plugin/PluginInfoConverter.java | 56 +- .../BasePrimitive.java} | 60 +- .../CallbackPrimitive.java} | 22 +- .../ChildContextPrimitive.java} | 53 +- .../InvokePrimitive.java} | 27 +- .../SerializablePrimitive.java} | 33 +- .../StepPrimitive.java} | 206 +--- .../WaitPrimitive.java} | 15 +- .../durable/DurableCoreOperationsTest.java | 103 -- .../lambda/durable/DurableFutureTest.java | 8 +- .../lambda/durable/DurableHandlerTest.java | 16 + ...Test.java => DurableMapOperationTest.java} | 19 +- .../durable/DurableOperationFacadeTest.java | 384 +++++++ ...java => DurableParallelOperationTest.java} | 5 +- ... DurableWaitForCallbackOperationTest.java} | 40 +- ...DurableWaitForConditionOperationTest.java} | 5 +- ...ava => DurableWithRetryOperationTest.java} | 5 +- .../DurationValidationIntegrationTest.java | 19 + .../context/ExtensionOperationImplTest.java | 342 +++--- .../execution/OperationIdGeneratorTest.java | 4 +- .../extension/ExtensionOperationTest.java | 37 + .../model/OperationIdentifierTest.java | 38 + .../DeferredDurableFutureTest.java | 2 +- ...urableMapOperationImplementationTest.java} | 26 +- .../operation/DurableOperationConfigTest.java | 180 ++++ ...eParallelOperationImplementationTest.java} | 40 +- ...rCallbackOperationImplementationTest.java} | 12 +- ...ConditionOperationImplementationTest.java} | 12 +- ...WithRetryOperationImplementationTest.java} | 18 +- .../OperationConcurrencyCoordinatorTest.java} | 26 +- .../plugin/PluginInfoConverterTest.java | 7 +- .../BasePrimitivePluginTest.java} | 14 +- .../CallbackPrimitiveTest.java} | 46 +- .../ChildContextPrimitiveTest.java} | 33 +- .../InvokePrimitiveTest.java} | 26 +- .../SerializablePrimitiveTest.java} | 84 +- .../StatefulExtensionStepPrimitiveTest.java} | 65 +- .../StepPrimitiveTest.java} | 71 +- .../WaitPrimitiveTest.java} | 10 +- 102 files changed, 3854 insertions(+), 5203 deletions(-) delete mode 100644 docs/superpowers/plans/2026-08-10-custom-extension-operations.md delete mode 100644 docs/superpowers/plans/2026-08-10-migrate-built-in-extensions.md delete mode 100644 docs/superpowers/specs/2026-08-10-custom-extension-operations-design.md delete mode 100644 docs/superpowers/specs/2026-08-10-migrate-built-in-extensions-design.md delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/DurableCoreOperations.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/DurableMapOperations.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/DurableParallelOperations.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForCallbackOperations.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForConditionOperations.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/DurableWithRetryOperations.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/context/extension/ParallelExtension.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForCallbackExtension.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForConditionExtension.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionCallbackConfig.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionInvokeConfig.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionOperationImpl.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/model/OperationDescriptor.java rename sdk/src/main/java/software/amazon/lambda/durable/{context/extension => operation}/DeferredDurableFuture.java (97%) create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/operation/DurableCallbackOperation.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/operation/DurableContextOperation.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/operation/DurableInvokeOperation.java rename sdk/src/main/java/software/amazon/lambda/durable/{context/extension/MapExtension.java => operation/DurableMapOperation.java} (55%) create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/operation/DurableParallelOperation.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/operation/DurableStepOperation.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperation.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperation.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitOperation.java rename sdk/src/main/java/software/amazon/lambda/durable/{context/extension/WithRetryExtension.java => operation/DurableWithRetryOperation.java} (50%) rename sdk/src/main/java/software/amazon/lambda/durable/{context/extension/ExtensionConcurrencyCoordinator.java => operation/OperationConcurrencyCoordinator.java} (98%) create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/operation/OperationConfigAdapters.java rename sdk/src/main/java/software/amazon/lambda/durable/{context/extension/ParallelExtensionFuture.java => operation/ParallelOperationFuture.java} (85%) rename sdk/src/main/java/software/amazon/lambda/durable/{context/extension => operation}/WaitForConditionFuture.java (94%) rename sdk/src/main/java/software/amazon/lambda/durable/{operation/BaseDurableOperation.java => primitive/BasePrimitive.java} (93%) rename sdk/src/main/java/software/amazon/lambda/durable/{operation/CallbackOperation.java => primitive/CallbackPrimitive.java} (82%) rename sdk/src/main/java/software/amazon/lambda/durable/{operation/ChildContextOperation.java => primitive/ChildContextPrimitive.java} (91%) rename sdk/src/main/java/software/amazon/lambda/durable/{operation/InvokeOperation.java => primitive/InvokePrimitive.java} (79%) rename sdk/src/main/java/software/amazon/lambda/durable/{operation/SerializableDurableOperation.java => primitive/SerializablePrimitive.java} (84%) rename sdk/src/main/java/software/amazon/lambda/durable/{operation/StepOperation.java => primitive/StepPrimitive.java} (53%) rename sdk/src/main/java/software/amazon/lambda/durable/{operation/WaitOperation.java => primitive/WaitPrimitive.java} (84%) delete mode 100644 sdk/src/test/java/software/amazon/lambda/durable/DurableCoreOperationsTest.java rename sdk/src/test/java/software/amazon/lambda/durable/{DurableMapOperationsTest.java => DurableMapOperationTest.java} (87%) create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/DurableOperationFacadeTest.java rename sdk/src/test/java/software/amazon/lambda/durable/{DurableParallelOperationsTest.java => DurableParallelOperationTest.java} (92%) rename sdk/src/test/java/software/amazon/lambda/durable/{DurableWaitForCallbackOperationsTest.java => DurableWaitForCallbackOperationTest.java} (69%) rename sdk/src/test/java/software/amazon/lambda/durable/{DurableWaitForConditionOperationsTest.java => DurableWaitForConditionOperationTest.java} (92%) rename sdk/src/test/java/software/amazon/lambda/durable/{DurableWithRetryOperationsTest.java => DurableWithRetryOperationTest.java} (92%) create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/extension/ExtensionOperationTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/model/OperationIdentifierTest.java rename sdk/src/test/java/software/amazon/lambda/durable/{context/extension => operation}/DeferredDurableFutureTest.java (97%) rename sdk/src/test/java/software/amazon/lambda/durable/{context/extension/MapExtensionTest.java => operation/DurableMapOperationImplementationTest.java} (87%) create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/operation/DurableOperationConfigTest.java rename sdk/src/test/java/software/amazon/lambda/durable/{context/extension/ParallelExtensionTest.java => operation/DurableParallelOperationImplementationTest.java} (90%) rename sdk/src/test/java/software/amazon/lambda/durable/{context/extension/WaitForCallbackExtensionTest.java => operation/DurableWaitForCallbackOperationImplementationTest.java} (93%) rename sdk/src/test/java/software/amazon/lambda/durable/{context/extension/WaitForConditionExtensionTest.java => operation/DurableWaitForConditionOperationImplementationTest.java} (93%) rename sdk/src/test/java/software/amazon/lambda/durable/{context/extension/WithRetryExtensionTest.java => operation/DurableWithRetryOperationImplementationTest.java} (87%) rename sdk/src/test/java/software/amazon/lambda/durable/{context/extension/ExtensionConcurrencyCoordinatorTest.java => operation/OperationConcurrencyCoordinatorTest.java} (88%) rename sdk/src/test/java/software/amazon/lambda/durable/{operation/BaseDurableOperationPluginTest.java => primitive/BasePrimitivePluginTest.java} (94%) rename sdk/src/test/java/software/amazon/lambda/durable/{operation/CallbackOperationTest.java => primitive/CallbackPrimitiveTest.java} (90%) rename sdk/src/test/java/software/amazon/lambda/durable/{operation/ChildContextOperationTest.java => primitive/ChildContextPrimitiveTest.java} (96%) rename sdk/src/test/java/software/amazon/lambda/durable/{operation/InvokeOperationTest.java => primitive/InvokePrimitiveTest.java} (90%) rename sdk/src/test/java/software/amazon/lambda/durable/{operation/SerializableDurableOperationTest.java => primitive/SerializablePrimitiveTest.java} (85%) rename sdk/src/test/java/software/amazon/lambda/durable/{operation/StatefulExtensionStepOperationTest.java => primitive/StatefulExtensionStepPrimitiveTest.java} (82%) rename sdk/src/test/java/software/amazon/lambda/durable/{operation/StepOperationTest.java => primitive/StepPrimitiveTest.java} (81%) rename sdk/src/test/java/software/amazon/lambda/durable/{operation/WaitOperationTest.java => primitive/WaitPrimitiveTest.java} (92%) diff --git a/docs/adr/006-custom-extension-operations.md b/docs/adr/006-custom-extension-operations.md index 9ec6cf734..98f46cdd3 100644 --- a/docs/adr/006-custom-extension-operations.md +++ b/docs/adr/006-custom-extension-operations.md @@ -28,7 +28,7 @@ defining new backend state machines. ## Decision -### Preserve Existing Operation APIs +### Preserve Existing Customer Operation APIs Keep every existing method signature, callback contract, configuration field, result type, exception, and behavior unchanged. New capabilities are additive and limited to extension-specific overloads and types. @@ -45,31 +45,53 @@ This includes: - `StepConfig` - `RunInChildContextConfig` -Existing methods on `ExtensionContext` and `ExtensionOperation` also remain unchanged; this decision adds overloads -rather than replacing them. +`ExtensionContext` and `ExtensionOperation` are new extension-author SPI contracts. `ExtensionOperation` deliberately +exposes only one fully specified asynchronous method per primitive instead of mirroring the customer-facing overload +families. Extension libraries can add their own conveniences, and deterministic operations can use the matching +built-in operation directly. The existing `DurableContext` methods remain compatibility APIs. Their implementations delegate to the same built-in -extension implementations used by the static facades. +operation classes used by the static APIs. -### Expose Core and Built-In Extension Facades +Each built-in operation owns its public nested configuration type, such as +`DurableStepOperation.StepConfig` or `DurableMapOperation.MapConfig`. The compatibility types under +`software.amazon.lambda.durable.config` convert through `toOperationConfig()` at the `DurableContext` boundary. -Expose primitive operations through `DurableCoreOperations`: +### Expose Primitive and Built-In Extension Facades -- `step` -- `wait` -- chained `invoke` -- `createCallback` -- `runInChildContext` +Expose each primitive through an independently maintained class: + +| Facade | Primitive | +| --- | --- | +| `DurableStepOperation` | STEP | +| `DurableWaitOperation` | WAIT | +| `DurableInvokeOperation` | CHAINED_INVOKE | +| `DurableCallbackOperation` | CALLBACK | +| `DurableContextOperation` | CONTEXT | + +The customer-facing primitive APIs use the extension SPI internally. For example: + +```text +DurableContext.step + -> DurableStepOperation.step + -> ExtensionOperation.stepAsync + -> primitive.StepPrimitive +``` + +The other primitives follow the same dependency direction through their matching merged operation class. +`DurableContextImpl` provides the durable scope and reservation mechanism. Only +`extension.ExtensionOperationImpl` constructs +the concrete primitive operation engines, so customer APIs and third-party extensions share one backend boundary. Expose each built-in extension family through an independently maintained class: | Facade | Operation family | | --- | --- | -| `DurableMapOperations` | `map`, `mapAsync` | -| `DurableParallelOperations` | `parallel` and branch construction | -| `DurableWaitForCallbackOperations` | `waitForCallback`, `waitForCallbackAsync` | -| `DurableWaitForConditionOperations` | `waitForCondition`, `waitForConditionAsync` | -| `DurableWithRetryOperations` | `withRetry`, `withRetryAsync` | +| `DurableMapOperation` | `map`, `mapAsync` | +| `DurableParallelOperation` | `parallel` and branch construction | +| `DurableWaitForCallbackOperation` | `waitForCallback`, `waitForCallbackAsync` | +| `DurableWaitForConditionOperation` | `waitForCondition`, `waitForConditionAsync` | +| `DurableWithRetryOperation` | `withRetry`, `withRetryAsync` | An extension is an ordinary static Java method. There is no registration API and no universal `DurableExtensions.run` boundary. @@ -85,13 +107,21 @@ software.amazon.lambda.durable.extension This package contains `ExtensionContext`, `ExtensionOperation`, stateful-step contracts, and configurable extension-context contracts. +Place merged built-in operation APIs in: + +```text +software.amazon.lambda.durable.operation +``` + +Each `Durable*Operation` class owns its context-free overloads and its canonical `ExtensionContext` implementation. +There are no separate built-in `*Extension` classes. + Keep the following customer-facing types in the root `software.amazon.lambda.durable` package: -- static operation facades such as `DurableCoreOperations` and `DurableMapOperations` - operation-specific TLS metadata such as `MapItemContext`, `WaitForCallbackContext`, and `WithRetryContext` - established SDK types such as `DurableFuture`, `StepContext`, and `TypeToken` -SDK implementations of built-in extensions remain internal and are not part of the extension-author API. +Backend primitive engines remain internal under `software.amazon.lambda.durable.primitive`. ### Use Scoped Current Context @@ -140,21 +170,41 @@ Extension authors never provide or observe the final globally stored operation I ### Allow Arbitrary Subtype Strings -`ExtensionOperation` provides subtype-aware overloads for every primitive: +`ExtensionOperation` provides one subtype-aware method for every primitive: ```java -reservation.stepAsync("MyStep", ...); -reservation.waitAsync("MyWait", ...); -reservation.invokeAsync("MyInvoke", ...); -reservation.createCallback("MyCallback", ...); -reservation.runInChildContextAsync("MyContext", ...); + DurableFuture stepAsync( + String subType, + TypeToken resultType, + ExtensionStepFunction function, + ExtensionStepConfig config); + +DurableFuture waitAsync(String subType, Duration duration); + + DurableFuture invokeAsync( + String subType, + String functionName, + U payload, + TypeToken resultType, + ExtensionInvokeConfig config); + + DurableCallbackFuture createCallback( + String subType, + TypeToken resultType, + ExtensionCallbackConfig config); + + DurableFuture runInChildContextAsync( + String subType, + TypeToken resultType, + ExtensionContextFunction function, + ExtensionContextConfig config); ``` The primitive selector determines the backend operation type. The string controls only the subtype recorded in checkpoints, replay validation, plugins, logs, and error metadata. Subtype strings must be non-null and nonblank. They are not restricted to the existing `OperationSubType` enum. -Existing no-subtype overloads retain the standard subtype strings. +Extensions use the corresponding `OperationSubType` value when they want a standard subtype. ### Keep Primitive State Machines Fixed @@ -238,7 +288,7 @@ Rewrite the existing composed operation families using the extension contract: - map uses a `Map` context and reserved `MapIteration` contexts - parallel uses a `Parallel` context and dynamically registered `ParallelBranch` contexts -The legacy `DurableContext` methods and static facades adapt into the same canonical family implementations. +The legacy `DurableContext` methods and static APIs adapt into the same canonical family implementations. After behavior and checkpoint parity are proven, remove the specialized: diff --git a/docs/advanced/extensions.md b/docs/advanced/extensions.md index 8f8c9d71e..5cc89362f 100644 --- a/docs/advanced/extensions.md +++ b/docs/advanced/extensions.md @@ -4,8 +4,9 @@ Extension operations are ordinary static Java methods that compose SDK-owned dur separate Maven module without defining backend operation types, sending checkpoint updates, or depending on SDK implementation packages. -Extension-author contracts are in `software.amazon.lambda.durable.extension`. Static operation facades and -operation-specific TLS metadata contexts remain in `software.amazon.lambda.durable`. +Extension-author contracts are in `software.amazon.lambda.durable.extension`. Built-in operation APIs are in +`software.amazon.lambda.durable.operation`, while operation-specific TLS metadata contexts remain in +`software.amazon.lambda.durable`. Application code calls only the extension's API: @@ -29,9 +30,18 @@ public final class PairOperations { var extension = ExtensionContext.getCurrentContext(); var left = extension.reserve(name + "-left"); var right = extension.reserve(name + "-right"); - - var leftFuture = left.stepAsync(String.class, leftFunction); - var rightFuture = right.stepAsync(String.class, rightFunction); + var config = ExtensionStepConfig.builder().build(); + + var leftFuture = left.stepAsync( + "PairStep", + TypeToken.get(String.class), + state -> ExtensionStepResult.succeed(leftFunction.get()), + config); + var rightFuture = right.stepAsync( + "PairStep", + TypeToken.get(String.class), + state -> ExtensionStepResult.succeed(rightFunction.get()), + config); return new PairFuture(leftFuture, rightFuture); } } @@ -42,44 +52,95 @@ compose primitives in the current scope or explicitly create a child context. ## Static operation APIs -New code can use context-free static facades: +New code can use context-free static operations from `software.amazon.lambda.durable.operation`: | Facade | Operations | | --- | --- | -| `DurableCoreOperations` | `step`, `wait`, chained `invoke`, callbacks, child contexts | -| `DurableMapOperations` | `map`, `mapAsync` | -| `DurableParallelOperations` | `parallel` | -| `DurableWaitForCallbackOperations` | `waitForCallback`, `waitForCallbackAsync` | -| `DurableWaitForConditionOperations` | `waitForCondition`, `waitForConditionAsync` | -| `DurableWithRetryOperations` | `withRetry`, `withRetryAsync` | +| `DurableStepOperation` | `step`, `stepAsync` | +| `DurableWaitOperation` | `wait`, `waitAsync` | +| `DurableInvokeOperation` | `invoke`, `invokeAsync` | +| `DurableCallbackOperation` | `createCallback` | +| `DurableContextOperation` | `runInChildContext`, `runInChildContextAsync` | +| `DurableMapOperation` | `map`, `mapAsync` | +| `DurableParallelOperation` | `parallel` | +| `DurableWaitForCallbackOperation` | `waitForCallback`, `waitForCallbackAsync` | +| `DurableWaitForConditionOperation` | `waitForCondition`, `waitForConditionAsync` | +| `DurableWithRetryOperation` | `withRetry`, `withRetryAsync` | The existing `DurableContext` instance methods and callback signatures remain supported for backward compatibility. +Each static operation owns its configuration type. For example: + +```java +var config = DurableStepOperation.StepConfig.builder() + .retryStrategy(RetryStrategies.Presets.DEFAULT) + .build(); +var result = DurableStepOperation.step("process", Result.class, () -> process(), config); +``` + +The same pattern applies to `DurableInvokeOperation.InvokeConfig`, +`DurableCallbackOperation.CallbackConfig`, `DurableContextOperation.RunInChildContextConfig`, +`DurableMapOperation.MapConfig`, `DurableParallelOperation.ParallelConfig`, +`DurableParallelOperation.ParallelBranchConfig`, `DurableWaitForCallbackOperation.WaitForCallbackConfig`, +`DurableWaitForConditionOperation.WaitForConditionConfig`, and +`DurableWithRetryOperation.WithRetryConfig`. + +The compatibility types in `software.amazon.lambda.durable.config` remain accepted by `DurableContext`. They can be +passed to a static operation through `toOperationConfig()`: + +```java +var legacyConfig = software.amazon.lambda.durable.config.StepConfig.builder().build(); +DurableStepOperation.step("process", Result.class, () -> process(), legacyConfig.toOperationConfig()); +``` + +## Primitive implementation path + +The instance APIs and static operations share the same canonical implementation: + +```text +DurableContext.step + -> DurableStepOperation.step + -> ExtensionOperation.stepAsync + -> primitive.StepPrimitive +``` + +WAIT, CHAINED_INVOKE, CALLBACK, and CONTEXT follow the same path through `DurableWaitOperation`, +`DurableInvokeOperation`, `DurableCallbackOperation`, and `DurableContextOperation`. Each class owns both its +context-free overloads and its `ExtensionContext` implementation; there are no separate built-in extension classes. + +`DurableContextImpl` owns the current durable scope and reservations, but it does not construct primitive operation +engines. `extension.ExtensionOperationImpl` is the single internal boundary that creates `StepPrimitive`, `WaitPrimitive`, +`InvokePrimitive`, `CallbackPrimitive`, and `ChildContextPrimitive`. + +The extension SPI uses extension-specific configuration types. `ExtensionInvokeConfig` and +`ExtensionCallbackConfig` isolate extension authors and primitive engines from operation-owned configuration; the +corresponding `Durable*Operation` class performs that conversion. + User functions in the static APIs do not receive SDK context objects: ```java -var result = DurableCoreOperations.step("process", Result.class, () -> { +var result = DurableStepOperation.step("process", Result.class, () -> { var step = StepContext.getCurrentContext(); return process(step.getAttempt()); }); ``` ```java -var result = DurableMapOperations.map("process", items, Result.class, item -> { +var result = DurableMapOperation.map("process", items, Result.class, item -> { var index = MapItemContext.getCurrentContext().getIndex(); return process(item, index); }); ``` ```java -var result = DurableWaitForCallbackOperations.waitForCallback( +var result = DurableWaitForCallbackOperation.waitForCallback( "approval", Approval.class, () -> submit(WaitForCallbackContext.getCurrentContext().getCallbackId())); ``` ```java -var result = DurableWithRetryOperations.withRetry("transaction", () -> { +var result = DurableWithRetryOperation.withRetry("transaction", () -> { var attempt = WithRetryContext.getCurrentContext().getAttempt(); return executeAttempt(attempt); }); @@ -101,9 +162,9 @@ Operation-specific TLS is not automatically propagated into a nested primitive's the metadata in its owning function and capture any application value needed by the nested operation: ```java -var result = DurableMapOperations.map("process", items, Result.class, item -> { +var result = DurableMapOperation.map("process", items, Result.class, item -> { var index = MapItemContext.getCurrentContext().getIndex(); - return DurableCoreOperations.step("process-item", Result.class, () -> process(item, index)); + return DurableStepOperation.step("process-item", Result.class, () -> process(item, index)); }); ``` @@ -112,8 +173,8 @@ SDK-managed durable context thread. ## Primitive reservations -Extensions with deterministic call order can use `DurableCoreOperations` directly. Schedulers whose registration -order is deterministic but launch order may vary use `ExtensionContext.reserve(name)`. +Extensions with deterministic call order can use the matching built-in operation directly. Schedulers whose +registration order is deterministic but launch order may vary use `ExtensionContext.reserve(name)`. Each reservation immediately consumes the next sequential operation ID and returns an opaque, one-shot `ExtensionOperation`: @@ -122,19 +183,133 @@ Each reservation immediately consumes the next sequential operation ID and retur var extension = ExtensionContext.getCurrentContext(); var first = extension.reserve("first"); var second = extension.reserve("second"); +var config = ExtensionStepConfig.builder().build(); // Launch order can differ from reservation order. -var secondResult = second.stepAsync(String.class, () -> runSecond()); -var firstResult = first.stepAsync(String.class, () -> runFirst()); +var secondResult = second.stepAsync( + "ScheduledStep", + TypeToken.get(String.class), + state -> ExtensionStepResult.succeed(runSecond()), + config); +var firstResult = first.stepAsync( + "ScheduledStep", + TypeToken.get(String.class), + state -> ExtensionStepResult.succeed(runFirst()), + config); ``` A reservation can create exactly one primitive: step, wait, chained invoke, callback, or child context. Reuse throws `IllegalStateException`. Raw operation IDs are never exposed. +`ExtensionOperation` exposes one fully specified asynchronous method for each primitive. Callers always provide the +subtype, use `TypeToken` for typed results, and supply the complete primitive configuration. Extensions can call +`get()` when they need blocking behavior or use the matching built-in operation when reservation-time ID allocation is +not needed. + Create reservations in the same order on every replay. Reordering, inserting, or removing reservations is a workflow compatibility change because it can associate existing checkpoints with different logical primitives. Launching already reserved operations in a different order is supported. +### Custom local IDs + +Schedulers whose registration order can change may reserve an explicit local ID: + +```java +var node = ExtensionContext.getCurrentContext().reserve("process-node", "node-a"); +var result = node.stepAsync( + "ProcessNode", + TypeToken.get(NodeResult.class), + state -> ExtensionStepResult.succeed(processNode("node-a")), + ExtensionStepConfig.builder().build()); +``` + +The local ID must be non-null, nonblank, and unique within the current context. It is never used directly as the +backend operation ID. The SDK namespaces and hashes it as follows: + +```text +root context: sha256(localOperationId) +child context: sha256(parentContextId + "-" + localOperationId) +``` + +Custom reservations and generated sequential operations share one local-ID registry. A custom reservation advances +the sequence once, generated numeric IDs skip values that were already claimed, and collisions fail immediately. +Changing or reusing a local ID is a workflow compatibility change. + +### Custom primitive subtypes + +Subtype-aware reservation overloads record an extension-specific identity while retaining the selected primitive's +SDK-owned state machine: + +```java +var result = ExtensionContext.getCurrentContext() + .reserve("process-node", "node-a") + .stepAsync( + "AcmeNode", + TypeToken.get(NodeResult.class), + state -> ExtensionStepResult.succeed(processNode("node-a")), + ExtensionStepConfig.builder().build()); +``` + +The selector still determines the backend operation type: `stepAsync` creates a `STEP`, `waitAsync` creates a `WAIT`, +and so on. The subtype must be non-null and nonblank. It appears in checkpoints, replay validation, plugins, logs, +and failure metadata, so changing it is a workflow compatibility change. + +## Stateful extension steps + +A stateful extension STEP can checkpoint application state between attempts without exposing raw checkpoint actions: + +```java +var result = ExtensionContext.getCurrentContext() + .reserve("poll") + .stepAsync( + "AcmePoll", + TypeToken.get(PollState.class), + state -> state.complete() + ? ExtensionStepResult.succeed(state) + : ExtensionStepResult.retry(refresh(state), Duration.ofSeconds(5)), + ExtensionStepConfig.builder() + .initialState(initialState) + .build()); +``` + +The function may return only `ExtensionStepResult.succeed(value)` or +`ExtensionStepResult.retry(state, delay)`. Retry state uses the configured `SerDes`; attempt metadata remains +available through `StepContext.getCurrentContext()`. Thrown exceptions follow the normal STEP failure path. + +## Configurable extension contexts + +An advanced CONTEXT primitive separates the application result from optional replay state: + +```java +var result = ExtensionContext.getCurrentContext() + .reserve("batch") + .runInChildContextAsync( + "AcmeBatch", + TypeToken.get(BatchResult.class), + () -> { + var replay = ExtensionContextReplayContext.getCurrentContext(); + var previous = replay.isReplayingChildren() ? replay.getReplayState() : null; + var current = rebuildBatch(previous); + return ExtensionContextResult.replayChildren(current, compact(current)); + }, + ExtensionContextConfig.builder() + .emitUserFunctionEvents(false) + .suppressLateChildCheckpoints(true) + .errorHandler(failure -> new IllegalStateException( + "Batch context failed: " + failure.contextName())) + .build()); +``` + +Use `ExtensionContextResult.completed(result)` when children never need to replay, +`replayChildren(result, replayState)` to always replay them, or +`replayChildrenAboveSize(result, replayState, thresholdBytes)` to replay only when the serialized full result reaches +the threshold. Replay metadata is scoped to the framework callback through `ExtensionContextReplayContext`. + +`ExtensionContextConfig` also composes `RunInChildContextConfig`, controls framework user-function plugin events, and +can suppress child checkpoints that finish after the parent. If a context fails, the SDK first rethrows a +deserialized original exception, then calls the configured error handler, and finally falls back to +`ChildContextFailedException`. The handler receives read-only context metadata and child-operation summaries. + ## Explicit child contexts An extension creates a child context only when its own semantics require isolation: @@ -142,10 +317,15 @@ An extension creates a child context only when its own semantics require isolati ```java var result = ExtensionContext.getCurrentContext() .reserve("isolated-work") - .runInChildContext(Result.class, () -> executeIsolatedWork()); + .runInChildContextAsync( + "IsolatedWork", + TypeToken.get(Result.class), + () -> ExtensionContextResult.completed(executeIsolatedWork()), + ExtensionContextConfig.builder().build()) + .get(); ``` -Inside the supplier, `DurableContext.getCurrentContext()` and `ExtensionContext.getCurrentContext()` return the child +Inside the function, `DurableContext.getCurrentContext()` and `ExtensionContext.getCurrentContext()` return the child context. ## Custom durable futures @@ -179,12 +359,18 @@ child-context operation. Serialization, suspension, replay, cancellation, failures, and checkpointing retain the semantics of the underlying primitive operations. +The built-in map, parallel, wait-for-callback, wait-for-condition, and with-retry families are implemented through +the same extension primitives. Their legacy `DurableContext` methods and context-free static APIs share one +canonical implementation while preserving their established checkpoint topology and plugin behavior. + ## Module compatibility An extension Maven module should depend only on the public SDK artifact and import public types under -`software.amazon.lambda.durable` and `software.amazon.lambda.durable.extension`. Do not import SDK implementation -packages such as `context`, `execution`, or `operation`. +`software.amazon.lambda.durable`, `software.amazon.lambda.durable.operation`, and +`software.amazon.lambda.durable.extension`. Do not import SDK implementation packages such as `context`, `execution`, +or `primitive`. The extension-author SPI includes `ExtensionContext`, `ExtensionOperation`, stateful-step contracts, and configurable -extension-context contracts under `software.amazon.lambda.durable.extension`. Static operation facades, typed TLS -contexts, and `DurableFuture.completionFuture()` remain under `software.amazon.lambda.durable`. +extension-context contracts under `software.amazon.lambda.durable.extension`. Static operation APIs are under +`software.amazon.lambda.durable.operation`; typed TLS contexts and `DurableFuture.completionFuture()` remain under +`software.amazon.lambda.durable`. diff --git a/docs/design.md b/docs/design.md index 50a2c5325..9fbe89042 100644 --- a/docs/design.md +++ b/docs/design.md @@ -249,11 +249,11 @@ context.step("name", Type.class, stepCtx -> doWork(), ▼ ▼ ┌──────────────────────────────┐ ┌──────────────────────────────┐ │ Operations │ │ CheckpointBatcher │ -│ - StepOperation │ │ - Queues requests │ -│ - WaitOperation │ │ - Batches API calls (750KB) │ -│ - InvokeOperation │ │ │ -│ - CallbackOperation │ │ - Notifies via callback │ -│ - ChildContextOperation │ +│ - StepPrimitive │ │ - Queues requests │ +│ - WaitPrimitive │ │ - Batches API calls (750KB) │ +│ - InvokePrimitive │ │ │ +│ - CallbackPrimitive │ │ - Notifies via callback │ +│ - ChildContextPrimitive │ │ - execute() / get() │ └──────────────────────────────┘ │ @@ -282,23 +282,35 @@ software.amazon.lambda.durable │ ├── StepConfig # Step configuration (retry, semantics, serDes) │ ├── InvokeConfig # Invoke configuration (payload/result serDes, tenantId) │ ├── CallbackConfig # Callback configuration (timeout, heartbeat, serDes) -│ ├── WaitForCallbackConfig # Composite callback + step config -│ ├── MapConfig # Map configuration (concurrency, completion, serDes) -│ ├── ParallelConfig # Parallel configuration (concurrency, completion) -│ ├── ParallelBranchConfig # Per-branch configuration -│ ├── RunInChildContextConfig # Child context configuration -│ ├── WaitForConditionConfig # Polling configuration (wait strategy, serDes, initialState) +│ ├── WaitForCallbackConfig # DurableContext compatibility config +│ ├── MapConfig # DurableContext compatibility config +│ ├── ParallelConfig # DurableContext compatibility config +│ ├── ParallelBranchConfig # ParallelDurableFuture compatibility config +│ ├── RunInChildContextConfig # DurableContext compatibility config +│ ├── WaitForConditionConfig # DurableContext compatibility config │ └── CompletionConfig # Completion criteria for map/parallel │ ├── context/ -│ ├── BaseContext # Base interface for DurableContext -│ └── extension/ # Built-in composed operation implementations -│ ├── MapExtension -│ ├── ParallelExtension -│ ├── WaitForCallbackExtension -│ ├── WaitForConditionExtension -│ ├── WithRetryExtension -│ └── ExtensionConcurrencyCoordinator +│ └── BaseContext # Base interface for DurableContext +│ +├── operation/ # Public built-in operation APIs + implementations +│ ├── DurableStepOperation # Owns nested StepConfig +│ ├── DurableWaitOperation +│ ├── DurableInvokeOperation +│ ├── DurableCallbackOperation +│ ├── DurableContextOperation +│ ├── DurableMapOperation +│ ├── DurableParallelOperation # Owns nested ParallelConfig and ParallelBranchConfig +│ ├── DurableWaitForCallbackOperation +│ ├── DurableWaitForConditionOperation +│ └── DurableWithRetryOperation +│ +├── primitive/ # Internal checkpoint-backed operation engines +│ ├── StepPrimitive +│ ├── WaitPrimitive +│ ├── InvokePrimitive +│ ├── CallbackPrimitive +│ └── ChildContextPrimitive │ ├── extension/ # Public SPI for extension authors │ ├── ExtensionContext @@ -317,12 +329,12 @@ software.amazon.lambda.durable │ └── ThreadType # CONTEXT, STEP │ ├── operation/ -│ ├── BaseDurableOperation # Common operation logic -│ ├── StepOperation # Step logic -│ ├── InvokeOperation # Invoke logic -│ ├── CallbackOperation # Callback logic -│ ├── WaitOperation # Wait logic -│ └── ChildContextOperation # Child context primitive +│ ├── BasePrimitive # Common operation logic +│ ├── StepPrimitive # Step logic +│ ├── InvokePrimitive # Invoke logic +│ ├── CallbackPrimitive # Callback logic +│ ├── WaitPrimitive # Wait logic +│ └── ChildContextPrimitive # Child context primitive │ ├── logging/ │ ├── DurableLogger # Context-aware logger wrapper (MDC-based) @@ -392,12 +404,12 @@ software.amazon.lambda.durable sequenceDiagram participant UC as User Code participant DC as DurableContext - participant SO as StepOperation + participant SO as StepPrimitive participant EM as ExecutionManager participant Backend UC->>DC: step("name", Type.class, stepCtx -> doWork()) - DC->>SO: new StepOperation(...) + DC->>SO: new StepPrimitive(...) DC->>SO: execute() SO->>EM: sendOperationUpdate(START) EM->>Backend: checkpoint(START) @@ -420,7 +432,7 @@ sequenceDiagram participant DE as DurableExecutor participant UC as User Code participant DC as DurableContext - participant SO as StepOperation + participant SO as StepPrimitive participant EM as ExecutionManager Note over LR: Re-invocation with existing state @@ -444,7 +456,7 @@ sequenceDiagram sequenceDiagram participant UC as User Code participant DC as DurableContext - participant WO as WaitOperation + participant WO as WaitPrimitive participant EM as ExecutionManager participant Backend @@ -543,7 +555,7 @@ This is a one-way transition (REPLAY → EXECUTION, never back). `DurableLogger` **Context Flow:** 1. `DurableLogger` constructor sets execution-level MDC (ARN, requestId) on the handler thread -2. `StepOperation.executeStepLogic()` calls `durableLogger.setOperationContext()` before user code runs +2. `StepPrimitive.executeStepLogic()` calls `durableLogger.setOperationContext()` before user code runs 3. User code logs via `context.getLogger()` - MDC values automatically included 4. `clearOperationContext()` called in finally block after step completes @@ -576,11 +588,11 @@ Multiple concurrent operations may checkpoint simultaneously. `CheckpointBatcher The `checkpointDelay` configuration option (default: 0) controls how long the batcher waits before flushing, allowing more operations to accumulate in a single batch. For functions with many concurrent operations, setting a small delay (e.g., 10ms) can significantly reduce the number of API calls. ``` -StepOperation 1 ──┐ +StepPrimitive 1 ──┐ │ -StepOperation 2 ──┼──► CheckpointBatcher ──► Backend +StepPrimitive 2 ──┼──► CheckpointBatcher ──► Backend │ -WaitOperation ────┘ +WaitPrimitive ────┘ ``` Callback mechanism avoids cyclic dependency between `ExecutionManager` and `CheckpointBatcher`: @@ -687,8 +699,8 @@ The SDK uses a threaded execution model where the handler runs on a user-configu | ThreadType | Identifier (threadId) | Created By | Purpose | |------------|--------------------------------------------------------------------------------|------------|---------| -| `CONTEXT` | `null` for root context; the operation ID for child contexts (e.g. `"hash(1)"`) | `DurableExecutor` (root), `ChildContextOperation` (child) | Runs the handler function body or a child context function body. Orchestrates operations. | -| `STEP` | The step's operation ID (e.g. `"hash(2)"`) | `StepOperation` | Runs user-provided step code (`Function`). | +| `CONTEXT` | `null` for root context; the operation ID for child contexts (e.g. `"hash(1)"`) | `DurableExecutor` (root), `ChildContextPrimitive` (child) | Runs the handler function body or a child context function body. Orchestrates operations. | +| `STEP` | The step's operation ID (e.g. `"hash(2)"`) | `StepPrimitive` | Runs user-provided step code (`Function`). | Each thread has a `ThreadContext` record (threadId + threadType) stored in a `ThreadLocal` so operations can identify which context they belong to. @@ -741,10 +753,10 @@ A thread deregisters when it cannot make forward progress — typically when it ### The `waitForOperationCompletion()` Pattern -This method in `BaseDurableOperation` is the core coordination primitive. It is called by every operation's `get()` method (step, wait, invoke, callback, child context): +This method in `BasePrimitive` is the core coordination primitive. It is called by every operation's `get()` method (step, wait, invoke, callback, child context): ```java -// BaseDurableOperation.waitForOperationCompletion() +// BasePrimitive.waitForOperationCompletion() protected Operation waitForOperationCompletion() { var threadContext = getCurrentThreadContext(); @@ -772,7 +784,7 @@ The re-registration callback (`thenRun`) runs synchronously on the thread that c When `CheckpointManager` receives a checkpoint response, it calls `ExecutionManager.onCheckpointComplete()`, which notifies each registered operation: ```java -// BaseDurableOperation.onCheckpointComplete() +// BasePrimitive.onCheckpointComplete() public void onCheckpointComplete(Operation operation) { if (ExecutionManager.isTerminalStatus(operation.status())) { synchronized (completionFuture) { @@ -786,12 +798,12 @@ Completing the future triggers the `thenRun` callback (re-registers the waiting ### Operation-Specific Threading -#### StepOperation +#### StepPrimitive Steps run user code on a separate thread via the user executor: ```java -// StepOperation.executeStepLogic() +// StepPrimitive.executeStepLogic() registerActiveThread(getOperationId()); // register BEFORE submitting to executor CompletableFuture.runAsync(() -> { @@ -809,12 +821,12 @@ Key details: - The step thread is implicitly deregistered when it finishes — it never calls `deregisterActiveThread` directly. Instead, the step thread's work is done after checkpointing, and the checkpoint response completes the `completionFuture`, which re-registers the waiting context thread. - For retries, the step sends a RETRY checkpoint and then polls for the READY status before re-executing. If no other threads are active during the retry delay, the execution suspends. -#### WaitOperation +#### WaitPrimitive Waits checkpoint a WAIT action with a duration, then poll for completion: ```java -// WaitOperation.start() +// WaitPrimitive.start() sendOperationUpdate(OperationUpdate.builder() .action(OperationAction.START) .waitOptions(WaitOptions.builder().waitSeconds((int) duration.toSeconds()).build())); @@ -823,20 +835,20 @@ pollForOperationUpdates(remainingWaitTime); The wait itself doesn't deregister any thread. Suspension happens when the context thread calls `wait()` (synchronous) which calls `get()`, which calls `waitForOperationCompletion()`, which deregisters the context thread. If no other threads are active, the execution suspends and the Lambda returns PENDING. On re-invocation, the wait replays: if the wait period has elapsed, `markAlreadyCompleted()` is called; otherwise, polling resumes with the remaining duration. -#### InvokeOperation +#### InvokePrimitive -Invokes checkpoint a START action with the target function name and payload, then poll for the result. The threading model is identical to WaitOperation — the invoke itself doesn't create a new thread. The context thread deregisters when it calls `get()` on the invoke future. +Invokes checkpoint a START action with the target function name and payload, then poll for the result. The threading model is identical to WaitPrimitive — the invoke itself doesn't create a new thread. The context thread deregisters when it calls `get()` on the invoke future. -#### CallbackOperation +#### CallbackPrimitive Callbacks checkpoint a START action to obtain a `callbackId`, then poll for an external system to complete the callback. Like waits and invokes, the context thread deregisters when it calls `get()`. The callback can complete via an external API call (success, failure, or heartbeat timeout). -#### ChildContextOperation +#### ChildContextPrimitive Child contexts run a user function in a separate thread with its own `DurableContext` and operation counter: ```java -// ChildContextOperation.executeChildContext() +// ChildContextPrimitive.executeChildContext() var contextId = getOperationId(); // Register on PARENT thread — prevents race with parent deregistration @@ -867,7 +879,7 @@ When a context thread calls `ctx.step(...)`, the following coordination occurs: | Seq | Context Thread | Step Thread | System Thread (CheckpointManager) | |-----|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| 1 | Create `StepOperation` + `completionFuture`. Call `execute()`. `execute()` calls `start()` which registers step thread and submits to user executor. Checkpoint START (sync or async depending on semantics). | — | (idle) | +| 1 | Create `StepPrimitive` + `completionFuture`. Call `execute()`. `execute()` calls `start()` which registers step thread and submits to user executor. Checkpoint START (sync or async depending on semantics). | — | (idle) | | 2 | `step()` calls `get()` → `waitForOperationCompletion()`. Attach `thenRun(re-register)` to `completionFuture`. Deregister context thread. Block on `join()`. | User code begins executing. Execute `function.apply(stepContext)`. | (idle) | | 3 | (blocked) | User code completes. Call `handleStepSucceeded(result)` → `sendOperationUpdate(SUCCEED)` (synchronous — blocks until checkpoint response). | Process checkpoint API call. On terminal response, call `onCheckpointComplete()` → `completionFuture.complete(null)`. `thenRun` fires: re-register context thread. | | 4 | `join()` returns. Retrieve result from operation. | Call `deregisterActiveThread` to deregister Step thread. Step thread ends. | (idle) | @@ -878,7 +890,7 @@ When a context thread calls `ctx.step(...)`, the following coordination occurs: | Seq | Context Thread | System Thread | |-----|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------|------------------------| -| 1 | Create `WaitOperation` + `completionFuture`. Call `execute()`. `execute()` calls `start()` → checkpoint WAIT with duration → `pollForOperationUpdates(remainingWaitTime)`. | Begin polling backend. | +| 1 | Create `WaitPrimitive` + `completionFuture`. Call `execute()`. `execute()` calls `start()` → checkpoint WAIT with duration → `pollForOperationUpdates(remainingWaitTime)`. | Begin polling backend. | | 2 | `wait()` calls `get()` → `waitForOperationCompletion()`. Attach `thenRun(re-register)`. Deregister context thread. | (polling) | | 3 | `activeThreads` is empty → `suspendExecution()` → `executionExceptionFuture.completeExceptionally(SuspendExecutionException)`. | — | | 4 | `runUntilCompleteOrSuspend` resolves with `SuspendExecutionException` → return `PENDING`. | — | @@ -895,8 +907,8 @@ var result = stepFuture.get(); | Seq | Context Thread | Step Thread | System Thread | |-----|--------------------------------------------------------------------|--------------------------------|---------------------------------------------------------------------------------------------------------| -| 1 | Create `StepOperation`, register step thread, submit to executor. | — | — | -| 2 | Create `WaitOperation`, checkpoint WAIT, start polling. | User code begins. | Begin polling for wait. | +| 1 | Create `StepPrimitive`, register step thread, submit to executor. | — | — | +| 2 | Create `WaitPrimitive`, checkpoint WAIT, start polling. | User code begins. | Begin polling for wait. | | 3 | `wait()` calls `get()` → deregister context thread. | (running) | (polling) | | 4 | (blocked — but step thread is still active, so no suspension) | Complete → checkpoint SUCCEED. | Process step checkpoint. | | 5 | (blocked) | — | Wait poll returns SUCCEEDED → `completionFuture.complete(null)` for wait. Context thread re-registered. | diff --git a/docs/superpowers/plans/2026-08-10-custom-extension-operations.md b/docs/superpowers/plans/2026-08-10-custom-extension-operations.md deleted file mode 100644 index 12dfc111f..000000000 --- a/docs/superpowers/plans/2026-08-10-custom-extension-operations.md +++ /dev/null @@ -1,608 +0,0 @@ -# Custom Extension Operations Implementation Plan - -> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. - -**Goal:** Add a supported public API for context-free static durable operations and replay-safe custom extensions while leaving the existing `DurableContext` interface surface unchanged. - -**Architecture:** SDK-managed handler and child contexts implement a minimal `ExtensionContext` that can reserve opaque, one-shot primitive identities. `DurableCoreOperations` and one facade per built-in extension family adapt context-free user callbacks to the existing `DurableContext` implementation, while typed scoped TLS contexts expose SDK-generated metadata. Existing primitive operation classes continue to own checkpointing, replay, suspension, serialization, and plugin events. - -**Tech Stack:** Java 17, Maven reactor, JUnit 6, Mockito 5, `LocalDurableTestRunner`, Palantir Java Format through Spotless. - -## Global Constraints - -- Keep all existing `DurableContext` methods and callback signatures unchanged for backward compatibility. -- Do not add `runExtensionAsync` or any other extension entry point to `DurableContext`. -- New user callbacks receive only application-provided values; SDK contexts and generated metadata are retrieved through TLS. -- Keep primitive operation IDs opaque and SDK-owned. -- A reservation is one-shot and allocates its sequential ID when `reserve` is called, not when the primitive is launched. -- Extensions do not automatically create child contexts. -- Do not add dependencies. -- Run `mvn spotless:apply` after Java changes. - ---- - -### Task 1: Scoped Current Contexts - -**Files:** -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/BaseContextImpl.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/execution/DurableExecutor.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java` -- Create: `sdk/src/main/java/software/amazon/lambda/durable/OperationContextStorage.java` -- Create: `sdk/src/main/java/software/amazon/lambda/durable/MapItemContext.java` -- Create: `sdk/src/main/java/software/amazon/lambda/durable/WaitForCallbackContext.java` -- Create: `sdk/src/main/java/software/amazon/lambda/durable/WithRetryContext.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/CurrentContextTest.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/OperationContextStorageTest.java` - -**Interfaces:** -- Consumes: Existing `BaseContext.getCurrentContext()`, `DurableContext.getCurrentContext()`, and `StepContext.getCurrentContext()`. -- Produces: `MapItemContext.getCurrentContext().getIndex()`, `WaitForCallbackContext.getCurrentContext().getCallbackId()`, and `WithRetryContext.getCurrentContext().getAttempt()`. - -- [ ] **Step 1: Write failing lookup and restoration tests** - -Add tests that establish these behaviors: - -```java -@Test -void mapItemContextRestoresNestedScope() { - assertThrows(IllegalStateException.class, MapItemContext::getCurrentContext); - try (var outer = MapItemContext.attach(2)) { - assertEquals(2, MapItemContext.getCurrentContext().getIndex()); - try (var inner = MapItemContext.attach(7)) { - assertEquals(7, MapItemContext.getCurrentContext().getIndex()); - } - assertEquals(2, MapItemContext.getCurrentContext().getIndex()); - } - assertThrows(IllegalStateException.class, MapItemContext::getCurrentContext); -} -``` - -Add equivalent outside-scope and nested-restoration assertions for callback IDs and retry attempts. In -`CurrentContextTest`, assert that handler/child contexts resolve as `DurableContext`, step scopes reject -`DurableContext` with guidance to use `StepContext`, and all scopes restore the preceding base context. - -- [ ] **Step 2: Run the focused tests and verify RED** - -Run: - -```bash -JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ -/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ --Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ --Djacoco.skip=true \ --DargLine=-javaagent:/home/czk/.codex-tmp/m2/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar \ --pl sdk -Dtest=CurrentContextTest,OperationContextStorageTest test -``` - -Expected: compilation fails because the three operation-specific context classes and their scoped attachment methods -do not exist. - -- [ ] **Step 3: Implement scoped storage and SDK binding** - -Implement a package-private generic storage helper: - -```java -final class OperationContextStorage { - private final String contextName; - private final ThreadLocal current = new ThreadLocal<>(); - - T getCurrentContext() { - var context = current.get(); - if (context == null) { - throw new IllegalStateException(contextName + " is not active on the current thread"); - } - return context; - } - - SafeCloseable attach(T context) { - var previous = current.get(); - current.set(context); - return () -> { - if (previous == null) { - current.remove(); - } else { - current.set(previous); - } - }; - } -} -``` - -Each public final metadata context owns a private static storage, a private immutable value, a public static lookup, -a public getter, and a package-private `attach` used by same-package facades. Preserve the existing -`DurableContext` signatures while clarifying its current-context failure behavior. Bind base contexts with -try-with-resources around handler, child, step, and condition user functions so nested calls restore rather than -blindly clear TLS. - -- [ ] **Step 4: Run focused tests and verify GREEN** - -Run the command from Step 2. Expected: all current-context and operation-context tests pass. - -- [ ] **Step 5: Commit** - -```bash -git add sdk/src/main/java/software/amazon/lambda/durable \ - sdk/src/test/java/software/amazon/lambda/durable/CurrentContextTest.java \ - sdk/src/test/java/software/amazon/lambda/durable/OperationContextStorageTest.java -git commit -m "feat: add scoped durable operation contexts" -``` - -### Task 2: Opaque Primitive Reservations - -**Files:** -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionContext.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java` -- Delete: `sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionContextImpl.java` -- Delete: `sdk/src/main/java/software/amazon/lambda/durable/DurableExtensions.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java` -- Delete: `sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionContextImplTest.java` - -**Interfaces:** -- Consumes: `DurableContextImpl` primitive construction and `OperationIdGenerator`. -- Produces: - -```java -public interface ExtensionContext extends BaseContext { - static ExtensionContext getCurrentContext(); - boolean isReplaying(); - ExtensionOperation reserve(String name); -} -``` - -```java -public interface ExtensionOperation { - DurableFuture stepAsync( - TypeToken resultType, Supplier function, StepConfig config); - DurableFuture waitAsync(Duration duration); - DurableFuture invokeAsync( - String functionName, U payload, TypeToken resultType, InvokeConfig config); - DurableCallbackFuture createCallback( - TypeToken resultType, CallbackConfig config); - DurableFuture runInChildContextAsync( - TypeToken resultType, Supplier function, RunInChildContextConfig config); -} -``` - -- [ ] **Step 1: Write failing reservation tests** - -Create tests that use a real `DurableContextImpl` with mocked dependencies to prove: - -```java -var first = context.reserve("first"); -var second = context.reserve("second"); - -second.stepAsync(String.class, () -> "second"); -first.stepAsync(String.class, () -> "first"); - -verify(operationFactory).createStepOperation("2", "second", ...); -verify(operationFactory).createStepOperation("1", "first", ...); -``` - -Also test every reserved primitive path, `ExtensionContext.getCurrentContext()` on handler/child versus step/no scope, -and a second use of the same reservation throwing: - -```java -assertThrows(IllegalStateException.class, () -> reservation.waitAsync(Duration.ZERO)); -``` - -- [ ] **Step 2: Run focused tests and verify RED** - -Run: - -```bash -JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ -/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ --Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ --Djacoco.skip=true \ --DargLine=-javaagent:/home/czk/.codex-tmp/m2/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar \ --pl sdk -Dtest=ExtensionOperationImplTest,CurrentContextTest test -``` - -Expected: compilation fails because reservations still accept context-bearing functions and -`DurableContextImpl` does not directly implement the final `ExtensionContext` contract. - -- [ ] **Step 3: Implement the minimal reservation path** - -Make `DurableContextImpl` implement `ExtensionContext`, with: - -```java -@Override -public ExtensionOperation reserve(String name) { - return new ExtensionOperationImpl(this, operationIdGenerator.next(), name); -} -``` - -Retain package-private explicit-ID helpers for step, wait, invoke, callback, and child context. Adapt `Supplier` to -the existing primitive callbacks inside `ExtensionOperationImpl`; current TLS is already attached when the user -supplier executes. Guard all primitive selectors with a single `AtomicBoolean.compareAndSet(false, true)`. - -Remove `runExtensionAsync` from `DurableContext`, remove the universal `DurableExtensions` facade, remove the wrapper -`ExtensionContextImpl`, and keep `ExtensionContext` limited to current lookup, replay state, and reservation. - -- [ ] **Step 4: Run focused tests and verify GREEN** - -Run the command from Step 2. Expected: all reservation and current-context tests pass. - -- [ ] **Step 5: Commit** - -```bash -git add sdk/src/main/java/software/amazon/lambda/durable \ - sdk/src/test/java/software/amazon/lambda/durable/context -git commit -m "feat: add opaque extension operation reservations" -``` - -### Task 3: Context-Free Core Static Facade - -**Files:** -- Create: `sdk/src/main/java/software/amazon/lambda/durable/DurableCoreOperations.java` -- Delete: `sdk/src/main/java/software/amazon/lambda/durable/DurableOperations.java` -- Create: `sdk/src/test/java/software/amazon/lambda/durable/DurableCoreOperationsTest.java` - -**Interfaces:** -- Consumes: `DurableContext.getCurrentContext()` and all existing primitive instance methods. -- Produces: sync/async, `Class`/`TypeToken`, default/custom config overloads for `step`, `wait`, chained - `invoke`, `createCallback`, and `runInChildContext`. - -- [ ] **Step 1: Write failing facade tests** - -Write tests against a TLS-bound mocked `DurableContext` proving the core overloads delegate and strip SDK context -parameters: - -```java -var result = DurableCoreOperations.step("step", String.class, () -> { - assertSame(stepContext, StepContext.getCurrentContext()); - return "done"; -}); -assertEquals("done", result); -``` - -For child contexts, verify a zero-argument supplier can obtain both `DurableContext` and `ExtensionContext` from TLS. -Also assert every facade family throws `IllegalStateException` outside a durable context. - -- [ ] **Step 2: Run the focused test and verify RED** - -Run: - -```bash -JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ -/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ --Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ --Djacoco.skip=true \ --DargLine=-javaagent:/home/czk/.codex-tmp/m2/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar \ --pl sdk -Dtest=DurableCoreOperationsTest test -``` - -Expected: compilation fails because `DurableCoreOperations` does not exist. - -- [ ] **Step 3: Implement only primitive facade overloads** - -Create a stateless final utility class whose step methods accept `Supplier` and delegate with -`ignored -> function.get()`. Child-context methods also accept `Supplier` and delegate with -`ignored -> function.get()`. Invoke, wait, and callback methods delegate values unchanged. Do not include map, -parallel, callback composition, condition polling, or retry methods. - -- [ ] **Step 4: Run focused tests and verify GREEN** - -Run the command from Step 2. Expected: all core facade tests pass. - -- [ ] **Step 5: Commit** - -```bash -git add sdk/src/main/java/software/amazon/lambda/durable/DurableCoreOperations.java \ - sdk/src/main/java/software/amazon/lambda/durable/DurableOperations.java \ - sdk/src/test/java/software/amazon/lambda/durable/DurableCoreOperationsTest.java -git commit -m "feat: add context-free core operation facade" -``` - -### Task 4: Independently Maintained Extension Facades - -**Files:** -- Create: `sdk/src/main/java/software/amazon/lambda/durable/DurableMapOperations.java` -- Create: `sdk/src/main/java/software/amazon/lambda/durable/DurableParallelOperations.java` -- Create: `sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForCallbackOperations.java` -- Create: `sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForConditionOperations.java` -- Create: `sdk/src/main/java/software/amazon/lambda/durable/DurableWithRetryOperations.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java` -- Create: `sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationsTest.java` -- Create: `sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationsTest.java` -- Create: `sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationsTest.java` -- Create: `sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForConditionOperationsTest.java` -- Create: `sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationsTest.java` - -**Interfaces:** -- Consumes: Existing `DurableContext` map, parallel, wait-for-callback, wait-for-condition, and with-retry methods. -- Produces: - - map callbacks as `Function` with index in `MapItemContext` - - parallel branch callbacks as `Supplier` - - callback submitters as `Runnable` with ID in `WaitForCallbackContext` - - condition checks as `Function>` - - retry bodies as `Supplier` with attempt in `WithRetryContext` - -- [ ] **Step 1: Write one failing metadata test per facade** - -Tests must invoke the adapted legacy callback and assert the public callback sees TLS metadata: - -```java -DurableMapOperations.map("map", List.of("a"), String.class, item -> { - assertEquals(3, MapItemContext.getCurrentContext().getIndex()); - return item.toUpperCase(); -}); -``` - -Use the mocked legacy callback to supply index `3`; repeat for callback ID `"cb-1"` and retry attempt `2`. For -wait-for-condition, bind a `StepContext` and prove the check function receives only the state value. For parallel, -compile and execute `parallel.branch("branch", String.class, () -> "done")`. - -- [ ] **Step 2: Run the five focused tests and verify RED** - -Run: - -```bash -JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ -/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ --Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ --Djacoco.skip=true \ --DargLine=-javaagent:/home/czk/.codex-tmp/m2/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar \ --pl sdk -Dtest=DurableMapOperationsTest,DurableParallelOperationsTest,DurableWaitForCallbackOperationsTest,DurableWaitForConditionOperationsTest,DurableWithRetryOperationsTest test -``` - -Expected: compilation fails because the split facades and supplier branch overloads do not exist. - -- [ ] **Step 3: Implement the five adapters** - -Each class is a stateless final utility class containing only its named family. Adapt callbacks as follows: - -```java -(item, index, ignored) -> { - try (var scope = MapItemContext.attach(index)) { - return function.apply(item); - } -} -``` - -```java -(callbackId, ignored) -> { - try (var scope = WaitForCallbackContext.attach(callbackId)) { - submitter.run(); - } -} -``` - -```java -(attempt, ignored) -> { - try (var scope = WithRetryContext.attach(attempt)) { - return operation.get(); - } -} -``` - -The condition adapter is `(state, ignored) -> checkFunction.apply(state)` because `WaitForConditionOperation` binds -the active `StepContext`. Add default supplier overloads to `ParallelDurableFuture` that delegate to its existing -`Function` core method. - -- [ ] **Step 4: Run focused tests and verify GREEN** - -Run the command from Step 2. Expected: all five facade tests pass. - -- [ ] **Step 5: Commit** - -```bash -git add sdk/src/main/java/software/amazon/lambda/durable \ - sdk/src/test/java/software/amazon/lambda/durable/Durable*OperationsTest.java -git commit -m "feat: split built-in extension operation facades" -``` - -### Task 5: Public Durable Future Completion Contract - -**Files:** -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/DurableFuture.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java` -- Modify: `sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java` - -**Interfaces:** -- Consumes: Existing `DurableFuture.anyOf` and SDK operation completion futures. -- Produces: `default CompletableFuture completionFuture()` for custom composed futures. - -- [ ] **Step 1: Write a failing custom-future combinator test** - -Create a test-only `DurableFuture` whose result future and completion signal are independent, override -`completionFuture()`, pass it to `DurableFuture.anyOf`, complete the signal, and assert `anyOf` completes without -requiring the future to extend `BaseDurableOperation`. - -- [ ] **Step 2: Run the focused test and verify RED** - -Run: - -```bash -JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ -/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ --Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ --Djacoco.skip=true \ --pl sdk -Dtest=DurableFutureTest test -``` - -Expected: the test fails because `anyOf` still relies on an SDK-internal operation downcast or no public completion -method exists. - -- [ ] **Step 3: Implement the completion signal** - -Add the default method throwing `UnsupportedOperationException` to custom futures that do not opt in. Override it in -`BaseDurableOperation` by deriving `internalFuture.thenApply(ignored -> null)` so callers cannot complete or cancel -the underlying durable operation. Change `DurableFuture.anyOf` to collect `completionFuture()` values without an -internal type cast. - -- [ ] **Step 4: Run focused tests and verify GREEN** - -Run the command from Step 2. Expected: all durable future tests pass. - -- [ ] **Step 5: Commit** - -```bash -git add sdk/src/main/java/software/amazon/lambda/durable/DurableFuture.java \ - sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java \ - sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java -git commit -m "feat: expose durable future completion signals" -``` - -### Task 6: Separate-Module Proof Extension and Integration Semantics - -**Files:** -- Modify: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/extension/PairOperations.java` -- Modify: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java` -- Create: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java` -- Modify: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java` - -**Interfaces:** -- Consumes: Only public classes from the `sdk` artifact; the `sdk-integration-tests` Maven module is the external - compilation boundary. -- Produces: A proof extension that reserves two step identities in registration order and can launch them in either - order without changing replay identities. - -- [ ] **Step 1: Write failing external-module integration tests** - -Implement the test fixture call site before its production helper: - -```java -var result = PairOperations.pair( - "pair", - () -> DurableCoreOperations.step("left-value", String.class, () -> "left"), - () -> DurableCoreOperations.step("right-value", String.class, () -> "right"), - true); -assertEquals(new PairResult("left", "right"), result); -``` - -Add tests using `LocalDurableTestRunner` for first execution plus replay, a wait-based suspension/resume, reverse launch -order, same-scope nested extension calls, an explicitly reserved child context, all static facade families, and TLS -metadata. Update the plugin test to assert primitive names are observed and no synthetic extension lifecycle event is -emitted. - -- [ ] **Step 2: Run integration tests and verify RED** - -Run: - -```bash -JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ -/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ --Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ --Djacoco.skip=true \ --DargLine=-javaagent:/home/czk/.codex-tmp/m2/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar \ --pl sdk-integration-tests -am \ --Dtest=ExtensionOperationIntegrationTest,StaticOperationsIntegrationTest,PluginIntegrationTest \ --Dsurefire.failIfNoSpecifiedTests=false test -``` - -Expected: tests fail because the external fixture still targets the discarded universal runner and context-bearing -callbacks. - -- [ ] **Step 3: Implement the public-contract-only fixture** - -`PairOperations` may import only public SDK types. It obtains `ExtensionContext.getCurrentContext()`, reserves -`left` and `right`, launches the selected reservation order, and combines the results. It must not import any package -under `software.amazon.lambda.durable.context`, `.execution`, or `.operation`. - -- [ ] **Step 4: Run integration tests and verify GREEN** - -Run the command from Step 2. Expected: all extension, static operation, and plugin integration tests pass. - -- [ ] **Step 5: Commit** - -```bash -git add sdk-integration-tests/src/test/java/software/amazon/lambda/durable -git commit -m "test: verify custom extensions across module boundary" -``` - -### Task 7: Documentation, Formatting, and Reactor Verification - -**Files:** -- Modify: `README.md` -- Modify: `docs/advanced/extensions.md` -- Modify: `docs/adr/006-custom-extension-operations.md` - -**Interfaces:** -- Consumes: The final public API implemented by Tasks 1-6. -- Produces: User and extension-author documentation matching the code exactly. - -- [ ] **Step 1: Rewrite the extension guide** - -Document: - -- static import examples for `DurableCoreOperations` and every split extension facade -- TLS-only user functions and all three metadata contexts -- ordinary static extension methods without `DurableExtensions.run` -- same-scope direct composition versus explicit child contexts -- deterministic reservation order and variable launch order -- one-shot reservation behavior -- public-contract-only separate Maven modules -- primitive-only plugin lifecycle events -- `DurableFuture.completionFuture()` for custom composed futures - -Mark ADR-006 `Accepted` only after implementation and verification succeed. - -- [ ] **Step 2: Format all Java sources** - -Run: - -```bash -JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ -/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ --Dmaven.repo.local=/home/czk/.codex-tmp/m2 spotless:apply -``` - -Expected: Spotless exits successfully and only intended Java files change. - -- [ ] **Step 3: Run focused SDK and integration verification** - -Run: - -```bash -JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ -/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ --Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ --Djacoco.skip=true \ --DargLine=-javaagent:/home/czk/.codex-tmp/m2/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar \ --pl sdk,sdk-integration-tests -am test -``` - -Expected: all tests in the dependency closure pass. - -- [ ] **Step 4: Run the full reactor** - -Run: - -```bash -JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ -/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ --Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ --Djacoco.skip=true \ --DargLine=-javaagent:/home/czk/.codex-tmp/m2/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar \ -clean install -``` - -Expected: `BUILD SUCCESS`. Cloud example tests remain disabled. - -- [ ] **Step 5: Review public compatibility and commit** - -Verify: - -```bash -git diff 6962f5a -- sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java -rg -n "DurableExtensions|DurableOperations|runExtensionAsync" \ - sdk/src/main sdk/src/test sdk-integration-tests docs README.md -git status --short -``` - -Expected: `DurableContext` contains no new method signatures; discarded names have no live references; the status -contains only intended files. - -Then commit: - -```bash -git add README.md docs sdk sdk-integration-tests -git commit -m "docs: document custom extension operations" -``` diff --git a/docs/superpowers/plans/2026-08-10-migrate-built-in-extensions.md b/docs/superpowers/plans/2026-08-10-migrate-built-in-extensions.md deleted file mode 100644 index 046a8f97e..000000000 --- a/docs/superpowers/plans/2026-08-10-migrate-built-in-extensions.md +++ /dev/null @@ -1,995 +0,0 @@ -# Built-In Extension Migration Implementation Plan - -> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. - -**Goal:** Rewrite map, parallel, wait-for-callback, wait-for-condition, and with-retry through subtype-aware extension primitives while preserving all existing APIs, checkpoints, replay behavior, exceptions, and plugin events. - -**Architecture:** Extend reservations with custom local IDs and arbitrary subtype strings while keeping backend state machines SDK-owned. Generalize STEP and CONTEXT primitives with extension-only state, replay, and failure policies, then make legacy `DurableContext` methods and static facades delegate to one implementation per built-in family. Map and parallel share a non-operation concurrency coordinator that waits through suspension-aware public durable-future combinators. - -**Tech Stack:** Java 17, Maven reactor, JUnit 6, Mockito 5, `LocalDurableTestRunner`, Palantir Java Format through Spotless. - -## Global Constraints - -- Do not remove or change any existing method signature on `DurableContext`, `ParallelDurableFuture`, `ExtensionContext`, or `ExtensionOperation`. -- Do not add fields or methods to existing operation configuration classes. -- New subtype strings are non-null, nonblank, and are not restricted to `OperationSubType`. -- Primitive selectors determine backend operation types; extension code cannot emit raw checkpoint actions. -- Custom operation IDs are local values that replace a sequence number and are hashed with the current context prefix. -- Preserve exact built-in operation IDs, names, types, subtypes, parent IDs, payloads, statuses, replay behavior, exceptions, and plugin ordering. -- Do not add dependencies. -- Run `mvn spotless:apply` after Java changes. - ---- - -### Task 1: Custom Local Operation IDs - -**Files:** -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/execution/OperationIdGenerator.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionContext.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/execution/OperationIdGeneratorTest.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java` -- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java` - -**Interfaces:** -- Consumes: Existing `OperationIdGenerator.nextOperationId()` and `ExtensionContext.reserve(String)`. -- Produces: `OperationIdGenerator.nextOperationId(String localOperationId)` and `ExtensionContext.reserve(String name, String localOperationId)`. - -- [ ] **Step 1: Write failing generator tests** - -Add tests covering generated/custom interleaving: - -```java -@Test -void customLocalIdsUseContextPrefixAndAdvanceSequence() { - var root = new OperationIdGenerator(null); - - assertEquals(hashOperationId("node-a"), root.nextOperationId("node-a")); - assertEquals(hashOperationId("2"), root.nextOperationId()); -} - -@Test -void generatedIdsSkipCustomNumericIds() { - var generator = new OperationIdGenerator(null); - - assertEquals(hashOperationId("2"), generator.nextOperationId("2")); - assertEquals(hashOperationId("3"), generator.nextOperationId()); -} - -@Test -void duplicateLocalIdsFail() { - var generator = new OperationIdGenerator("parent"); - generator.nextOperationId("node"); - - assertThrows(IllegalArgumentException.class, () -> generator.nextOperationId("node")); -} -``` - -Also cover null, blank, a direct generated ID followed by the same custom numeric ID, and -`hashOperationId("parent-node")` for child contexts. - -- [ ] **Step 2: Run the generator tests and verify RED** - -Run: - -```bash -JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ -/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ --Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ --Djacoco.skip=true \ --pl sdk -Dtest=OperationIdGeneratorTest test -``` - -Expected: compilation fails because the custom-local-ID overload does not exist. - -- [ ] **Step 3: Implement shared local-ID allocation** - -Use one atomic counter and concurrent local-ID set: - -```java -private final Set allocatedLocalIds = ConcurrentHashMap.newKeySet(); - -public String nextOperationId() { - String localId; - do { - localId = String.valueOf(operationCounter.incrementAndGet()); - } while (!allocatedLocalIds.add(localId)); - return hashOperationId(operationIdPrefix + localId); -} - -public String nextOperationId(String localOperationId) { - validateLocalOperationId(localOperationId); - if (!allocatedLocalIds.add(localOperationId)) { - throw new IllegalArgumentException("Local operation ID is already in use: " + localOperationId); - } - operationCounter.incrementAndGet(); - return hashOperationId(operationIdPrefix + localOperationId); -} -``` - -Validate before advancing the counter. - -- [ ] **Step 4: Add the reservation overload** - -Add an additive method to `ExtensionContext`: - -```java -ExtensionOperation reserve(String name, String localOperationId); -``` - -Implement it in `DurableContextImpl` by validating the name and allocating through the new generator overload. -Keep `reserve(String)` unchanged. - -- [ ] **Step 5: Add reservation and integration tests** - -Assert: - -- custom reservations are one-shot -- custom IDs remain stable when their registration order changes -- nested custom IDs use the child context ID prefix -- ordinary core operations and reservations share the same local-ID registry - -- [ ] **Step 6: Run focused tests and commit** - -Run: - -```bash -JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ -/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ --Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ --Djacoco.skip=true \ --DargLine=-javaagent:/home/czk/.codex-tmp/m2/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar \ --pl sdk,sdk-integration-tests -am \ --Dtest=OperationIdGeneratorTest,ExtensionOperationImplTest,ExtensionOperationIntegrationTest \ --Dsurefire.failIfNoSpecifiedTests=false test -``` - -Commit: - -```bash -git add sdk/src/main sdk/src/test sdk-integration-tests/src/test -git commit -m "feat: add custom extension operation ids" -``` - ---- - -### Task 2: Arbitrary Primitive Subtype Strings - -**Files:** -- Create: `sdk/src/main/java/software/amazon/lambda/durable/model/OperationDescriptor.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/WaitOperation.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/InvokeOperation.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/CallbackOperation.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginInfoConverter.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginInfoConverterTest.java` -- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java` - -**Interfaces:** -- Consumes: One-shot `ExtensionOperation` selectors and existing enum-based `OperationIdentifier`. -- Produces: Additive subtype overloads for all primitive selectors and an internal identity carrying `OperationType` plus exact subtype string. - -- [ ] **Step 1: Write failing subtype tests** - -Add one test per primitive: - -```java -var step = context.reserve("custom-step"); -step.stepAsync("AcmeStep", String.class, () -> "done"); - -verify(context).stepAsyncWithId( - eq("1"), - eq("custom-step"), - eq("AcmeStep"), - eq(TypeToken.get(String.class)), - any(), - any()); -``` - -Add integration assertions that checkpoints and plugin events contain `AcmeStep`, `AcmeWait`, `AcmeInvoke`, -`AcmeCallback`, and `AcmeContext`, while their operation types remain fixed by the selector. - -- [ ] **Step 2: Run focused tests and verify RED** - -Expected: compilation fails on missing subtype overloads. - -- [ ] **Step 3: Add internal string-based identity** - -Create: - -```java -public record OperationDescriptor( - String operationId, - String name, - OperationType operationType, - String subType) { - - public OperationDescriptor { - Objects.requireNonNull(operationId, "operationId cannot be null"); - Objects.requireNonNull(operationType, "operationType cannot be null"); - if (subType == null || subType.isBlank()) { - throw new IllegalArgumentException("subType cannot be null or blank"); - } - } - - public static OperationDescriptor from(OperationIdentifier identifier) { - return new OperationDescriptor( - identifier.operationId(), - identifier.name(), - identifier.operationType(), - identifier.subType().getValue()); - } -} -``` - -Keep `OperationIdentifier` unchanged. Add descriptor constructor overloads to primitive operation classes while -retaining enum-based constructors for current call sites and tests. - -- [ ] **Step 4: Generalize base replay and plugin paths** - -Store `OperationDescriptor` in `BaseDurableOperation`. Use `descriptor.subType()` for updates and replay validation. -Retain: - -```java -public OperationSubType getSubType() -``` - -for known enum-based operations, and add: - -```java -public String getSubTypeValue() -``` - -for arbitrary values. Add descriptor overloads to `PluginInfoConverter` without changing existing overloads. - -- [ ] **Step 5: Add subtype-aware selector overloads** - -For each primitive, add additive methods such as: - -```java - DurableFuture stepAsync( - String subType, - TypeToken resultType, - Supplier function, - StepConfig config); -``` - -Existing methods delegate using `OperationSubType.STEP.getValue()` and equivalent standard values. Validate subtype -before claiming the reservation so invalid input does not consume it. - -- [ ] **Step 6: Run focused tests and commit** - -Run the extension, primitive operation, replay-validation, and plugin converter unit tests plus -`ExtensionOperationIntegrationTest`. - -Commit: - -```bash -git add sdk/src/main sdk/src/test sdk-integration-tests/src/test -git commit -m "feat: support custom extension subtypes" -``` - ---- - -### Task 3: Stateful STEP Extension Primitive - -**Files:** -- Create: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionStepFunction.java` -- Create: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionStepResult.java` -- Create: `sdk/src/main/java/software/amazon/lambda/durable/config/ExtensionStepConfig.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/ExtensionStepResultTest.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionStepConfigTest.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/operation/StepOperationTest.java` -- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java` - -**Interfaces:** -- Consumes: Subtype-aware STEP reservations. -- Produces: A fixed STEP lifecycle whose user outcome is `succeed(value)` or `retry(state, delay)`. - -- [ ] **Step 1: Add failing API and validation tests** - -Test immutable factories and builder defaults: - -```java -var retry = ExtensionStepResult.retry("next", Duration.ofSeconds(2)); -assertEquals("next", retry.state()); -assertEquals(Duration.ofSeconds(2), retry.delay()); -``` - -Reject null delays, negative delays, null results, and missing `ExtensionStepConfig`. - -- [ ] **Step 2: Add failing state-machine tests** - -Exercise: - -- first attempt receives `initialState` -- retry checkpoints serialized state and delay -- READY replay resumes with checkpointed state and incremented attempt -- success returns the final state -- user exception checkpoints failure -- suspension and unrecoverable exceptions propagate - -- [ ] **Step 3: Implement extension types** - -Use a sealed result: - -```java -public sealed interface ExtensionStepResult - permits ExtensionStepResult.Succeeded, ExtensionStepResult.Retry { - record Succeeded(T value) implements ExtensionStepResult {} - record Retry(T state, Duration delay) implements ExtensionStepResult {} -} -``` - -Implement `ExtensionStepConfig` with builder fields `initialState` and `serDes`; null SerDes uses the durable -configuration default. - -- [ ] **Step 4: Generalize StepOperation** - -Introduce an internal attempt strategy inside `StepOperation`: - -```java -private interface AttemptBehavior { - AttemptOutcome execute(T state, StepContext context); -} -``` - -The existing constructor wraps the current function and retry strategy. The extension constructor maps -`ExtensionStepResult` onto the same START/RETRY/READY/SUCCEED/FAIL paths. Do not duplicate checkpoint sending or poll -logic. - -- [ ] **Step 5: Expose the reservation selector** - -Add: - -```java - DurableFuture stepAsync( - String subType, - TypeToken resultType, - ExtensionStepFunction function, - ExtensionStepConfig config); -``` - -The function receives state only. `StepContext` remains TLS-bound. - -- [ ] **Step 6: Run focused tests and commit** - -Run `ExtensionStepResultTest`, `ExtensionStepConfigTest`, `StepOperationTest`, -`ExtensionOperationImplTest`, and `ExtensionOperationIntegrationTest`. - -Commit: - -```bash -git add sdk/src/main sdk/src/test sdk-integration-tests/src/test -git commit -m "feat: add stateful extension steps" -``` - ---- - -### Task 4: Configurable CONTEXT Extension Primitive - -**Files:** -- Create: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextFunction.java` -- Create: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextResult.java` -- Create: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextReplayContext.java` -- Create: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextErrorHandler.java` -- Create: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionContextFailure.java` -- Create: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionChildOperationSummary.java` -- Create: `sdk/src/main/java/software/amazon/lambda/durable/config/ExtensionContextConfig.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/ExtensionOperation.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/ExtensionContextResultTest.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/ExtensionContextReplayContextTest.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionContextConfigTest.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java` - -**Interfaces:** -- Consumes: Subtype-aware CONTEXT reservations. -- Produces: Replay-state result policies, scoped replay TLS, configurable fallback failure translation, plugin-hook policy, and late-child checkpoint suppression. - -- [ ] **Step 1: Write failing value and config tests** - -Test these factories: - -```java -ExtensionContextResult.completed(fullResult); -ExtensionContextResult.replayChildren(fullResult, replayState); -ExtensionContextResult.replayChildrenAboveSize(fullResult, replayState, 256 * 1024); -``` - -Test `ExtensionContextConfig.builder()` defaults: - -```java -assertTrue(config.emitUserFunctionEvents()); -assertFalse(config.suppressLateChildCheckpoints()); -assertNotNull(config.childContextConfig()); -``` - -- [ ] **Step 2: Write failing CONTEXT lifecycle tests** - -Add tests for: - -- normal result payload -- always replay-children with replay state -- threshold evaluated against serialized full result -- replay state available only inside `ExtensionContextReplayContext` -- nested replay scopes restore prior values -- framework hook emission enabled and disabled -- original exception reconstruction before fallback handler -- fallback handler receives child summaries -- default `ChildContextFailedException` -- a child finishing after a suppressing parent does not checkpoint - -- [ ] **Step 3: Implement extension context value types** - -Make all values immutable and defensively copy child summary lists. The replay TLS follows the existing -`OperationContextStorage` scoped-attachment pattern. - -- [ ] **Step 4: Generalize ChildContextOperation** - -Retain the existing `RunInChildContextConfig` constructor and adapt it to a standard context policy. Add an extension -constructor accepting `ExtensionContextFunction` and `ExtensionContextConfig`. - -On success: - -```java -var outcome = extensionFunction.apply(); -var full = serializeAndDeserializeResult(outcome.result()); -var checkpoint = selectCheckpointPayload(outcome, full.serialized()); -``` - -On replay with `replayChildren=true`, deserialize the stored replay state and attach it while rerunning the framework -function. Apply `emitUserFunctionEvents` only around the framework callback. Continue firing nested primitive hooks. - -- [ ] **Step 5: Generalize parent completion suppression** - -Replace the `ConcurrencyOperation`-specific parent constructor dependency with a general parent completion owner. -Store the owning extension context operation on child `DurableContextImpl` instances when -`suppressLateChildCheckpoints` is enabled. Nested extension child operations consult this owner before checkpointing. - -- [ ] **Step 6: Expose the advanced context selector** - -Add: - -```java - DurableFuture runInChildContextAsync( - String subType, - TypeToken resultType, - ExtensionContextFunction function, - ExtensionContextConfig config); -``` - -Keep the standard subtype plus `Supplier` overload and all existing methods unchanged. - -- [ ] **Step 7: Run focused tests and commit** - -Run the new extension context tests, `ChildContextOperationTest`, `ExtensionOperationImplTest`, plugin tests, and -extension integration tests. - -Commit: - -```bash -git add sdk/src/main sdk/src/test sdk-integration-tests/src/test -git commit -m "feat: add configurable extension contexts" -``` - ---- - -### Task 5: Migrate Wait for Callback - -**Files:** -- Create: `sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForCallbackExtension.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForCallbackOperations.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/extension/WaitForCallbackExtensionTest.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationsTest.java` -- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/CallbackIntegrationTest.java` -- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java` -- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java` - -**Interfaces:** -- Consumes: Advanced CONTEXT reservation, callback and step primitives, and configurable context failure translation. -- Produces: One canonical wait-for-callback implementation used by static and legacy APIs. - -- [ ] **Step 1: Write failing canonical-delegation tests** - -Assert both entry points call a handler with the same names and topology: - -```text -approval CONTEXT / WaitForCallback - approval-callback CALLBACK / Callback - approval-submitter STEP / Step -``` - -Cover callback failure, timeout, submitter failure, generic result types, custom SerDes, and plugin event ordering. - -- [ ] **Step 2: Implement WaitForCallbackExtension** - -Use: - -```java -public static DurableFuture execute( - ExtensionContext context, - String name, - TypeToken resultType, - BiConsumer submitter, - WaitForCallbackConfig config) -``` - -Reserve the parent first, then create callback and submitter reservations inside its framework function. Configure -parent user-function hooks as enabled. Supply an error handler that inspects child summaries and recreates -`CallbackFailedException`, `CallbackTimeoutException`, and `CallbackSubmitterException`. - -- [ ] **Step 3: Redirect both APIs** - -`DurableContextImpl.waitForCallbackAsync` delegates with `this`. The static facade resolves -`ExtensionContext.getCurrentContext()` and retains its existing TLS adapter around the `Runnable`. - -- [ ] **Step 4: Run callback suites and commit** - -Run callback unit, integration, retry-with-callback, static operations, and plugin tests. - -Commit: - -```bash -git add sdk/src/main sdk/src/test sdk-integration-tests/src/test -git commit -m "refactor: implement wait for callback as extension" -``` - ---- - -### Task 6: Migrate With Retry - -**Files:** -- Create: `sdk/src/main/java/software/amazon/lambda/durable/context/extension/WithRetryExtension.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/DurableWithRetryOperations.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/extension/WithRetryExtensionTest.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java` -- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java` -- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java` - -**Interfaces:** -- Consumes: Advanced `WithRetry` context reservation and ordinary wait primitives. -- Produces: One retry loop shared by static and legacy APIs. - -- [ ] **Step 1: Write failing parity tests** - -Cover first-attempt success, delayed retries, exhausted retries, null names, virtual versus checkpointed context, -control-flow exception propagation, and static TLS attempt metadata. - -- [ ] **Step 2: Implement WithRetryExtension** - -Move the retry loop out of `DurableContextImpl`: - -```java -public static DurableFuture execute( - ExtensionContext context, - String name, - BiFunction operation, - WithRetryConfig config) -``` - -Create a `WithRetry` extension context using the existing naming and virtual-context rules. Read the child -`DurableContext` through TLS, run the operation, and reserve ordinary waits with the existing backoff names. - -- [ ] **Step 3: Redirect APIs and remove old loop helpers** - -Delegate both legacy and static APIs to the canonical extension. Delete retry constants and loop methods from -`DurableContextImpl` after tests compile. - -- [ ] **Step 4: Run retry suites and commit** - -Commit: - -```bash -git add sdk/src/main sdk/src/test sdk-integration-tests/src/test -git commit -m "refactor: implement retry as extension" -``` - ---- - -### Task 7: Migrate Wait for Condition - -**Files:** -- Create: `sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForConditionExtension.java` -- Create: `sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForConditionFuture.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForConditionOperations.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java` -- Delete: `sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java` -- Move/replace test: `sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionOperationTest.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/extension/WaitForConditionExtensionTest.java` -- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionIntegrationTest.java` -- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java` - -**Interfaces:** -- Consumes: Stateful extension STEP. -- Produces: Existing wait-for-condition APIs and exception behavior through subtype `WaitForCondition`. - -- [ ] **Step 1: Write failing extension parity tests** - -Assert the existing single STEP checkpoint is retained across immediate success, multiple retries, READY replay, -initial state, custom strategy, custom SerDes, thrown checks, and plugin attempt numbers. - -- [ ] **Step 2: Implement WaitForConditionExtension** - -Map the existing result to fixed step outcomes: - -```java -var result = checkFunction.apply(state, StepContext.getCurrentContext()); -if (result.isDone()) { - return ExtensionStepResult.succeed(result.value()); -} -var delay = config.waitStrategy().evaluate( - result.value(), - StepContext.getCurrentContext().getAttempt()); -return ExtensionStepResult.retry(result.value(), delay); -``` - -Use subtype `WaitForCondition`, the existing initial state, and the existing SerDes defaulting. - -- [ ] **Step 3: Preserve fallback exception type** - -Wrap the stateful step future in `WaitForConditionFuture`. Delegate `completionFuture()`. In `get()`, let original -deserialized exceptions propagate; translate only fallback `StepFailedException` to -`WaitForConditionFailedException` using its operation. - -- [ ] **Step 4: Redirect APIs and remove the specialized operation** - -Delegate static and legacy methods to the canonical extension. Delete `WaitForConditionOperation` after all its -state-machine assertions have equivalent coverage in `StepOperationTest` and `WaitForConditionExtensionTest`. - -- [ ] **Step 5: Run condition suites and commit** - -Commit: - -```bash -git add -A sdk/src/main sdk/src/test sdk-integration-tests/src/test -git commit -m "refactor: implement wait for condition as extension" -``` - ---- - -### Task 8: Suspension-Aware Concurrency Coordination - -**Files:** -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/DurableFuture.java` -- Create: `sdk/src/main/java/software/amazon/lambda/durable/context/extension/DeferredDurableFuture.java` -- Create: `sdk/src/main/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinator.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/extension/DeferredDurableFutureTest.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinatorTest.java` -- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionConcurrencyIntegrationTest.java` - -**Interfaces:** -- Consumes: `DurableFuture.completionFuture()`, extension child reservations, and `CompletionConfig`. -- Produces: A non-operation coordinator usable by map and parallel without internal future downcasts. - -- [ ] **Step 1: Write failing suspension tests for anyOf** - -Use `LocalDurableTestRunner` with child futures waiting on callbacks. Assert the invocation reaches `PENDING` instead -of remaining active while `DurableFuture.anyOf` waits. - -- [ ] **Step 2: Make anyOf cooperate with SDK thread registration** - -When called on an SDK-managed context thread, delegate completion waiting to a context helper that: - -1. records the current thread context -2. registers a completion continuation -3. deregisters the active thread before blocking -4. re-registers when any completion signal fires - -Keep current behavior outside SDK threads. Do not change `completionFuture()` mutation isolation. - -- [ ] **Step 3: Implement DeferredDurableFuture** - -Provide a one-time `bind(DurableFuture)` method. `get()` waits for binding then delegates. `completionFuture()` -returns a stable future completed from the bound future. Reject a second binding. - -- [ ] **Step 4: Implement ExtensionConcurrencyCoordinator** - -The coordinator maintains: - -```java -record Item( - ExtensionOperation reservation, - Supplier> launcher, - DeferredDurableFuture exposedFuture) {} -``` - -It must: - -- register items in deterministic order -- launch no more than `maxConcurrency` -- wait through `DurableFuture.anyOf` -- count succeeded and failed items -- evaluate `CompletionConfig.completionDecisionFunction()` -- preserve `allItemsRegistered` -- mark pending/running incomplete items as skipped when completion occurs -- propagate suspension and unrecoverable control flow - -- [ ] **Step 5: Run focused and integration tests and commit** - -Commit: - -```bash -git add sdk/src/main sdk/src/test sdk-integration-tests/src/test -git commit -m "feat: add extension concurrency coordination" -``` - ---- - -### Task 9: Migrate Map - -**Files:** -- Create: `sdk/src/main/java/software/amazon/lambda/durable/context/extension/MapExtension.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/DurableMapOperations.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java` -- Delete: `sdk/src/main/java/software/amazon/lambda/durable/operation/MapOperation.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/extension/MapExtensionTest.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationsTest.java` -- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/MapIntegrationTest.java` -- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/MapInputValidationTest.java` -- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java` - -**Interfaces:** -- Consumes: Extension CONTEXT result policies and the shared concurrency coordinator. -- Produces: Existing map behavior with `Map` and `MapIteration` subtype-aware context primitives. - -- [ ] **Step 1: Add checkpoint-history parity tests** - -For legacy and static map calls, assert identical: - -- parent and iteration IDs -- `Map` and `MapIteration` subtypes -- nested/flat parent IDs -- small result payloads -- large result replay state and `replayChildren` -- early completion statuses and skipped iterations -- empty-map checkpoint flag behavior -- plugin events - -- [ ] **Step 2: Implement MapExtension** - -Use: - -```java -public static DurableFuture> execute( - ExtensionContext context, - String name, - Collection items, - TypeToken resultType, - DurableContext.MapFunction function, - MapConfig config) -``` - -Validate and copy items exactly as the current implementation. Reserve the `Map` parent first. Inside it, reserve -iterations in input order, attach `MapItemContext`, and launch through `ExtensionConcurrencyCoordinator`. - -- [ ] **Step 3: Preserve result and replay policies** - -Construct `MapResult` with existing success, failure, and skipped entries. For results at least 256 KB, use: - -```java -ExtensionContextResult.replayChildrenAboveSize( - fullResult, - stripMapResult(fullResult), - 256 * 1024); -``` - -On replay, use `ExtensionContextReplayContext` statuses to avoid launching previously skipped iterations and to -restore the prior completion decision. - -- [ ] **Step 4: Preserve empty-map and plugin behavior** - -Consume the parent reservation in all cases. Use a virtual `Map` extension context when empty-map checkpointing is -disabled, retain the warning, return `MapResult.empty()`, suppress parent framework user-function hooks, and keep -operation start/end plugin events balanced. - -- [ ] **Step 5: Redirect APIs and delete MapOperation** - -Delegate static and legacy map methods to `MapExtension`. Move reusable result assertions from operation tests into -`MapExtensionTest`. Delete `MapOperation`. - -- [ ] **Step 6: Run the complete map suite and commit** - -Run all map unit/integration tests and map-related plugin tests for both nesting modes. - -Commit: - -```bash -git add -A sdk/src/main sdk/src/test sdk-integration-tests/src/test -git commit -m "refactor: implement map as extension" -``` - ---- - -### Task 10: Migrate Parallel and Remove Specialized Concurrency - -**Files:** -- Create: `sdk/src/main/java/software/amazon/lambda/durable/context/extension/ParallelExtension.java` -- Create: `sdk/src/main/java/software/amazon/lambda/durable/context/extension/ParallelExtensionFuture.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/DurableParallelOperations.java` -- Modify: `sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java` -- Delete: `sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java` -- Delete: `sdk/src/main/java/software/amazon/lambda/durable/operation/ConcurrencyOperation.java` -- Delete/replace test: `sdk/src/test/java/software/amazon/lambda/durable/operation/ParallelOperationTest.java` -- Delete/replace test: `sdk/src/test/java/software/amazon/lambda/durable/operation/ConcurrencyOperationTest.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/extension/ParallelExtensionTest.java` -- Test: `sdk/src/test/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinatorTest.java` -- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ParallelIntegrationTest.java` -- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java` - -**Interfaces:** -- Consumes: Dynamic coordinator registration, deferred futures, and advanced CONTEXT replay policies. -- Produces: Existing `ParallelDurableFuture` backed entirely by extension primitives. - -- [ ] **Step 1: Add parallel checkpoint-history parity tests** - -Compare legacy and static APIs for empty, heterogeneous, max-concurrency, early-success, failure-tolerance, nested, -flat, replay, branches added after parent completion, and plugin scenarios. - -- [ ] **Step 2: Implement ParallelExtensionFuture** - -The future: - -- starts one `Parallel` extension context -- queues branch definitions in registration order -- returns a `DeferredDurableFuture` from each branch call -- rejects registration after `get()` or `close()` -- signals `allItemsRegistered` on join -- delegates its own `completionFuture()` and `get()` to the parent context future - -The parent framework function obtains its child `ExtensionContext`, drains registrations through the coordinator, -and returns `ExtensionContextResult.replayChildren(result, result)`. - -- [ ] **Step 3: Preserve replay and late-completion behavior** - -Use stored `ParallelResult.statuses()` to skip branches that did not exist or were previously skipped. Configure the -parent with framework user hooks disabled and late-child checkpoints suppressed. Configure branch contexts with -`ParallelBranch` fallback translation and virtual mode for flat nesting. - -- [ ] **Step 4: Redirect APIs and delete specialized classes** - -`DurableContextImpl.parallel` and `DurableParallelOperations.parallel` instantiate the same canonical extension -future. Delete `ParallelOperation` and `ConcurrencyOperation` after moving all shared assertions to coordinator and -extension tests. - -- [ ] **Step 5: Run parallel and broad integration suites and commit** - -Run parallel unit/integration tests, nested map/parallel tests, callbacks inside branches, condition operations inside -branches, and plugin tests. - -Commit: - -```bash -git add -A sdk/src/main sdk/src/test sdk-integration-tests/src/test -git commit -m "refactor: implement parallel as extension" -``` - ---- - -### Task 11: API Parity, Documentation, and Full Verification - -**Files:** -- Modify: `docs/advanced/extensions.md` -- Modify: `docs/adr/006-custom-extension-operations.md` -- Modify: `README.md` only if the extension guide link or description changes -- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java` -- Test: `sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java` - -**Interfaces:** -- Consumes: All migrated built-in extensions. -- Produces: Final compatibility evidence and documented public extension contracts. - -- [ ] **Step 1: Add final legacy/static parity coverage** - -For each family, run the legacy and static forms with equivalent inputs and compare normalized operation history: - -```java -record OperationShape( - String name, - String type, - String subType, - String parentId, - String status) {} -``` - -Also verify custom subtype strings and custom local IDs from a separate Maven module fixture. - -- [ ] **Step 2: Verify public API compatibility** - -Confirm: - -```bash -git diff 1d3de02 -- sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java -git diff 1d3de02 -- sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java -git diff 1d3de02 -- sdk/src/main/java/software/amazon/lambda/durable/config -``` - -Expected: no removed or changed existing signatures or fields; only additive extension-specific files and overloads. - -- [ ] **Step 3: Update extension documentation** - -Document: - -- arbitrary subtype strings -- custom local ID hashing and collision rules -- stateful STEP outcomes -- context replay state -- context error handlers -- built-in operations as reference extensions -- subtype, local ID, and replay state compatibility warnings - -- [ ] **Step 4: Run Spotless** - -Run: - -```bash -JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ -/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ --Dmaven.repo.local=/home/czk/.codex-tmp/m2 spotless:apply -``` - -Review and remove only unrelated formatter churn outside touched files. - -- [ ] **Step 5: Run focused dependency closure** - -Run: - -```bash -JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ -/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ --Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ --Djacoco.skip=true \ --DargLine=-javaagent:/home/czk/.codex-tmp/m2/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar \ --pl sdk,sdk-integration-tests -am test -``` - -Expected: all SDK, testing, and integration tests pass. - -- [ ] **Step 6: Run full reactor** - -Run: - -```bash -JAVA_HOME=/home/czk/.codex-tmp/jdk17 \ -/home/czk/.codex-tmp/maven-3.9.11/bin/mvn \ --Dmaven.repo.local=/home/czk/.codex-tmp/m2 \ --Djacoco.skip=true \ --DargLine=-javaagent:/home/czk/.codex-tmp/m2/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar \ -clean install -``` - -Expected: all eight reactor modules succeed. Cloud tests remain disabled by their existing guard. - -- [ ] **Step 7: Final audit and commit** - -Run: - -```bash -git diff --check -git status --short -rg -n "MapOperation|ParallelOperation|ConcurrencyOperation|WaitForConditionOperation" sdk/src/main sdk/src/test -``` - -Expected: no obsolete specialized engine references and no unintended worktree changes. - -Commit: - -```bash -git add README.md docs sdk sdk-integration-tests -git commit -m "docs: document migrated built-in extensions" -``` diff --git a/docs/superpowers/specs/2026-08-10-custom-extension-operations-design.md b/docs/superpowers/specs/2026-08-10-custom-extension-operations-design.md deleted file mode 100644 index 7d0291a0b..000000000 --- a/docs/superpowers/specs/2026-08-10-custom-extension-operations-design.md +++ /dev/null @@ -1,409 +0,0 @@ -# Custom Extension Operations Design - -## Objective - -Provide a supported public API for implementing reusable durable extension operations in a separate Maven module. -Extensions compose SDK-owned primitive operations without adding backend operation types, sending raw checkpoint -updates, or depending on SDK implementation classes. - -Application code calls ordinary static extension methods: - -```java -import static software.amazon.lambda.durable.dag.DagOperations.dag; - -var result = dag("etl", () -> { - var dag = DagContext.getCurrentContext(); - // Define DAG nodes through the current DAG context. -}); -``` - -The public `DurableContext` interface remains unchanged. Its existing instance methods continue to work for backward -compatibility. - -## Operation Categories - -### Core operations - -Core operations correspond to SDK-owned primitive behavior: - -- `step` -- `wait` -- chained `invoke` -- `createCallback` -- `runInChildContext` - -`DurableCoreOperations` exposes context-free static facades for these operations. The facades obtain the active -`DurableContext` from SDK-managed current-context storage and delegate to the existing instance methods. User -functions in the new APIs do not receive SDK context objects. For example, static step methods accept `Supplier`; -step code obtains `StepContext` through `StepContext.getCurrentContext()`. - -### Extension operations - -Extension operations compose core operations: - -- `waitForCallback` -- `waitForCondition` -- `withRetry` -- `map` -- `parallel` -- third-party operations such as DAG - -Each built-in extension has an independently maintained static facade: - -- `DurableMapOperations` -- `DurableParallelOperations` -- `DurableWaitForCallbackOperations` -- `DurableWaitForConditionOperations` -- `DurableWithRetryOperations` - -Each class owns only its operation's overloads, tests, and documentation. Existing `DurableContext` instance methods -and their behavior remain unchanged. - -An extension does not automatically create a child context. Each extension chooses its scope: - -- Replay-safe value helpers can create a step directly in the current context. -- Recursive invocation helpers can delegate directly to `invoke`. -- `map`, `parallel`, or `withRetry` can explicitly create child contexts when isolation is part of their semantics. -- DAG can reserve primitive identities in the current scope or explicitly create a child context if the DAG contract - requires one. - -## Extension Authoring Contract - -There is no universal `DurableExtensions.run` or `runAsync` method. An extension is an ordinary public static method -that uses the static operation facades and, when it needs stable deferred identities, the active `ExtensionContext`. - -```java -public interface ExtensionContext extends BaseContext { - static ExtensionContext getCurrentContext() { - var context = BaseContext.getCurrentContext(); - if (context instanceof ExtensionContext extensionContext) { - return extensionContext; - } - throw new IllegalStateException( - "ExtensionContext is only available from a durable handler or child-context thread"); - } - - boolean isReplaying(); - - ExtensionOperation reserve(String name); -} -``` - -SDK-managed handler and child contexts implement both `DurableContext` and `ExtensionContext`. Step contexts do not. -`ExtensionContext.getCurrentContext()` therefore succeeds only on supported handler and child-context threads. - -`ExtensionContext` exposes metadata through `BaseContext`, replay state, and deterministic primitive reservations. It -does not expose execution managers, checkpoint models, backend operation types, operation updates, or raw operation -IDs. - -## Primitive Reservations - -`ExtensionContext.reserve(name)` immediately consumes the next sequential operation ID in the active durable scope and -returns an opaque, one-shot `ExtensionOperation`. The ID remains hidden from extension code. - -```java -public interface ExtensionOperation { - DurableFuture stepAsync( - TypeToken resultType, - Supplier function, - StepConfig config); - - DurableFuture waitAsync(Duration duration); - - DurableFuture invokeAsync( - String functionName, - U payload, - TypeToken resultType, - InvokeConfig config); - - DurableCallbackFuture createCallback( - TypeToken resultType, - CallbackConfig config); - - DurableFuture runInChildContextAsync( - TypeToken resultType, - Supplier function, - RunInChildContextConfig config); - - // Class, synchronous, and default-configuration overloads are default methods. -} -``` - -The operation name is bound by `reserve` and is not repeated when selecting the primitive. A reservation can execute -exactly one primitive operation. Reuse fails with `IllegalStateException`. - -The SDK implements reservations by allocating an ID through the current context's normal `OperationIdGenerator`. -Package-private explicit-ID variants of primitive creation methods consume the reserved ID. These internal methods are -not part of the extension API. - -## DAG Usage - -A DAG module uses only public SDK contracts: - -```java -public static DagResult dag(String name, Runnable register) { - var extension = ExtensionContext.getCurrentContext(); - var dag = new DagContext(name, extension); - try (var ignored = DagContext.attach(dag)) { - register.run(); - } - return dag.execute(); -} -``` - -During its deterministic definition phase, the DAG reserves primitive positions: - -```java -var extract = extension.reserve("extract"); -var transform = extension.reserve("transform"); -var load = extension.reserve("load"); -``` - -The scheduler can later execute those reservations in any dependency-valid order: - -```java -var transformFuture = transform.stepAsync(String.class, () -> transformData()); -var extractFuture = extract.runInChildContextAsync( - ExtractResult.class, - () -> executeExtraction()); -``` - -Registration order determines IDs; launch order does not. This supports graph scheduling without name-derived IDs or -public explicit-ID APIs. - -The production DAG module is outside this issue. A small extension fixture in a separate repository Maven module -proves that an external module can compile and execute using only the supported contracts. - -## Current Context - -The SDK binds and restores current context around: - -- the durable handler -- child-context functions -- step functions -- wait-for-condition check functions -- map item functions -- wait-for-callback submitters -- with-retry bodies - -`DurableContext.getCurrentContext()` continues to use its existing signature. Its failure behavior is clarified: - -- Handler or child-context thread: returns the active durable context. -- Step thread: throws `IllegalStateException` directing callers to `StepContext`. -- Unsupported or application-created thread: throws `IllegalStateException` explaining that no durable context is - active. - -`ExtensionContext.getCurrentContext()` returns the active extension-capable handler or child context. It throws a -clear `IllegalStateException` from step threads and unsupported threads. - -Current context is not propagated to application-created threads. Extensions must create durable primitives on -SDK-managed durable context threads. - -## User Function Signatures - -New static APIs and extension reservations never pass SDK-created context or metadata values as user-function -arguments. User functions receive only values supplied by the application or values from the application's durable -data flow. - -Examples: - -```java -var result = DurableCoreOperations.step( - "process", - Result.class, - () -> { - var step = StepContext.getCurrentContext(); - return process(step.getAttempt()); - }); -``` - -```java -var result = DurableMapOperations.map( - "process", - items, - Result.class, - item -> { - var mapItem = MapItemContext.getCurrentContext(); - return process(item, mapItem.getIndex()); - }); -``` - -```java -var result = DurableWaitForCallbackOperations.waitForCallback( - "approval", - Approval.class, - () -> { - var callback = WaitForCallbackContext.getCurrentContext(); - submitApproval(callback.getCallbackId()); - }); -``` - -```java -var result = DurableWithRetryOperations.withRetry( - "transaction", - () -> { - var retry = WithRetryContext.getCurrentContext(); - return executeAttempt(retry.getAttempt()); - }); -``` - -The new callback shapes are: - -- step and child-context functions: `Supplier` -- map item functions: `Function`; `MapItemContext` exposes the item index -- parallel branch functions: `Supplier` -- wait-for-callback submitters: `Runnable`; `WaitForCallbackContext` exposes the callback ID -- wait-for-condition checks: receive only the durable state value; attempt metadata is available from - `StepContext.getCurrentContext()` -- with-retry bodies: `Supplier`; `WithRetryContext` exposes the attempt number -- extension child-context reservations: `Supplier`; `ExtensionContext` is obtained through TLS - -`MapItemContext`, `WaitForCallbackContext`, `WithRetryContext`, and any equivalent operation-specific context provide -`getCurrentContext()` static accessors. Each accessor fails clearly outside its matching user-function scope. These -operation-specific contexts are bound in addition to the base durable or step context so static core operations -continue to resolve the active `DurableContext` or `StepContext`. - -The initial operation-specific metadata contracts are: - -```java -public interface MapItemContext { - static MapItemContext getCurrentContext() { - return OperationContextStorage.get(MapItemContext.class); - } - - int getIndex(); -} - -public interface WaitForCallbackContext { - static WaitForCallbackContext getCurrentContext() { - return OperationContextStorage.get(WaitForCallbackContext.class); - } - - String getCallbackId(); -} - -public interface WithRetryContext { - static WithRetryContext getCurrentContext() { - return OperationContextStorage.get(WithRetryContext.class); - } - - int getAttempt(); -} -``` - -Each context uses a scoped SDK-managed `ThreadLocal`. Entering a nested operation stores the previous value, and -closing the scope restores it. The thread-local value is removed when no previous value exists. -`OperationContextStorage` is a package-private SDK implementation detail. - -Existing context-accepting functions on `DurableContext`, including `Function`, -`Function`, and existing map/retry callback types, remain unchanged for backward compatibility. - -## Static Operation Facades - -`DurableCoreOperations` contains only core operations. Each method obtains the current durable context internally. -Its step and child-context static methods accept context-free suppliers. Code inside those callbacks uses typed -current-context accessors when it needs SDK metadata. - -Each built-in extension facade contains only one operation family: - -| Facade | Methods | -| --- | --- | -| `DurableMapOperations` | `map`, `mapAsync` | -| `DurableParallelOperations` | `parallel` and its branch-building API | -| `DurableWaitForCallbackOperations` | `waitForCallback`, `waitForCallbackAsync` | -| `DurableWaitForConditionOperations` | `waitForCondition`, `waitForConditionAsync` | -| `DurableWithRetryOperations` | `withRetry`, `withRetryAsync` | - -The initial implementations delegate to the existing `DurableContext` methods to preserve behavior. The -classification and static API do not require rewriting each established operation implementation in this change. -Each facade is a stateless utility class. Calling a facade outside a supported durable context produces the same clear -failure as `DurableContext.getCurrentContext()`. - -## Durable Futures - -Asynchronous extension methods may return SDK operation futures or custom composed `DurableFuture` implementations. -`DurableFuture` therefore exposes a public non-mutating completion signal: - -```java -default CompletableFuture completionFuture() { - throw new UnsupportedOperationException( - "This DurableFuture does not expose a completion signal"); -} -``` - -SDK operation implementations return a derived completion future whose completion or cancellation cannot mutate the -durable operation. `DurableFuture.anyOf` uses this public contract instead of downcasting to -`BaseDurableOperation`. Custom futures that support `anyOf` override `completionFuture()`. - -## Replay and Compatibility - -Reservations must be created in the same deterministic order on every replay. Reordering, inserting, or removing -reservations can associate existing checkpoints with different logical primitives and is a workflow compatibility -change. After registration, executing reservations in a different order is supported. - -Direct static core calls allocate IDs when invoked, matching existing `DurableContext` semantics. They are appropriate -when call order is deterministic. - -Nested extension calls execute in the active scope unless an extension explicitly creates a child context. No -extension-specific recursion limit is introduced. - -Public compatibility guarantees apply to: - -- `DurableCoreOperations` -- `DurableMapOperations` -- `DurableParallelOperations` -- `DurableWaitForCallbackOperations` -- `DurableWaitForConditionOperations` -- `DurableWithRetryOperations` -- `ExtensionContext` -- `ExtensionOperation` -- `MapItemContext` -- `WaitForCallbackContext` -- `WithRetryContext` -- `DurableFuture.completionFuture()` - -Compatible SDK releases may add new default overloads or new primitive capabilities. Existing reservation ordering, -one-shot behavior, and primitive semantics change only in a breaking release. - -## Plugins and Failures - -There is no automatic extension lifecycle boundary because an extension is an ordinary composition method. Plugins -observe every primitive created by the extension. If the extension explicitly creates a child context, plugins also -observe that child-context operation. - -Primitive serialization, exception, suspension, cancellation, retry, and checkpoint behavior remain owned by the -existing primitive implementation. Extension code cannot send checkpoint updates or define backend operation -subtypes. - -Invalid names and null arguments use the existing SDK validators. Reusing a reservation or requesting current context -from an unsupported thread fails before creating a primitive. - -## Verification - -Unit tests cover: - -- current durable, step, and extension context lookup -- operation-specific context lookup and scope validation -- restoration of nested current-context bindings -- deterministic reservation allocation -- out-of-order reservation execution -- one-shot reservation enforcement -- each reserved primitive delegation path -- custom `DurableFuture` participation in `anyOf` -- static facade failure outside a durable context - -Integration tests in `sdk-integration-tests` cover: - -- an extension fixture compiled in a separate Maven module -- initial execution and replay -- suspension and resume -- reservations launched in different orders across replays -- nested extensions in the same scope -- extensions that explicitly create child contexts -- static core and built-in extension facades -- context-free user functions with TLS-based metadata access -- primitive plugin lifecycle events - -Formatting runs through `mvn spotless:apply`. Verification starts with focused SDK and integration tests, then expands -to the full reactor because the change affects public APIs, execution context propagation, replay identity, and -durable futures. diff --git a/docs/superpowers/specs/2026-08-10-migrate-built-in-extensions-design.md b/docs/superpowers/specs/2026-08-10-migrate-built-in-extensions-design.md deleted file mode 100644 index 532b83895..000000000 --- a/docs/superpowers/specs/2026-08-10-migrate-built-in-extensions-design.md +++ /dev/null @@ -1,570 +0,0 @@ -# Migrate Built-In Operations to the Extension API - -## Objective - -Rewrite the SDK's existing non-primitive operations as built-in extensions using the public extension operation -model. Preserve every existing user-facing interface, configuration type, overload, result, exception, checkpoint -shape, replay behavior, plugin event, and concurrency behavior. - -The migrated operation families are: - -- map -- parallel -- wait for callback -- wait for condition -- with retry - -The existing `DurableContext` methods remain supported compatibility APIs. They delegate to the same implementations -used by the static built-in extension facades. - -## Compatibility Boundary - -The following existing public APIs remain unchanged: - -- `DurableContext` -- `ParallelDurableFuture` -- `MapConfig` -- `ParallelConfig` -- `WaitForCallbackConfig` -- `WaitForConditionConfig` -- `WithRetryConfig` -- `StepConfig` -- `RunInChildContextConfig` -- all existing result and exception types - -This migration may add public extension-specific interfaces, overloads, and configuration types. It must not add -methods to the existing operation interfaces or fields to their existing configuration types. - -The following observable behavior remains unchanged: - -- operation IDs and parent-child ID namespaces -- operation types, subtype strings, names, and tree shape -- checkpoint action sequences and replay validation -- serialized result and failure payloads -- map and parallel completion decisions -- concurrency limits and skipped item behavior -- nested and flat concurrency modes -- large-result replay-children behavior -- wait-for-condition state, attempts, and delays -- wait-for-callback failure and timeout translation -- retry backoff names and virtual-context behavior -- plugin operation and user-function event ordering - -## Primitive State Machines - -Extension authors may select operation subtype strings, but they may not define checkpoint state machines. - -Each primitive retains its SDK-owned lifecycle: - -| Primitive selector | Backend operation type | SDK-owned state machine | -| --- | --- | --- | -| step | `STEP` | start, retry, ready, succeed, fail | -| wait | `WAIT` | start, poll, succeed | -| invoke | `CHAINED_INVOKE` | start, poll, succeed, fail | -| callback | `CALLBACK` | start, poll, succeed, fail, timeout | -| child context | `CONTEXT` | start, execute/replay children, succeed, fail | - -The primitive selector determines the backend operation type. The supplied subtype is metadata used for checkpoint -validation, plugins, exception translation, and execution history. - -Subtype strings must be non-null and nonblank. The SDK does not restrict them to an allow-list because the backend -accepts arbitrary subtype strings. - -## Subtype-Aware Reservations - -`ExtensionOperation` gains subtype-aware overloads for every primitive. Existing overloads remain and use the current -standard subtype strings. - -Representative asynchronous signatures: - -```java - DurableFuture stepAsync( - String subType, - TypeToken resultType, - Supplier function, - StepConfig config); - -DurableFuture waitAsync( - String subType, - Duration duration); - - DurableFuture invokeAsync( - String subType, - String functionName, - U payload, - TypeToken resultType, - InvokeConfig config); - - DurableCallbackFuture createCallback( - String subType, - TypeToken resultType, - CallbackConfig config); - - DurableFuture runInChildContextAsync( - String subType, - TypeToken resultType, - Supplier function, - RunInChildContextConfig config); -``` - -The existing no-subtype methods delegate with these values: - -- `Step` -- `Wait` -- `ChainedInvoke` -- `Callback` -- `RunInChildContext` - -Synchronous, `Class`, and default-configuration overloads remain default methods. - -## Custom Local Operation IDs - -`ExtensionContext` retains sequential reservation and adds a custom-local-ID overload: - -```java -ExtensionOperation reserve(String name); - -ExtensionOperation reserve(String name, String localOperationId); -``` - -The custom value replaces the generated sequence number for that reservation. It is not the final backend operation -ID. - -Custom local IDs must be non-null and nonblank. They are otherwise treated as opaque UTF-8 strings. - -The SDK constructs the final ID using the current context namespace: - -```text -root context: sha256(localOperationId) -child context: sha256(parentContextId + "-" + localOperationId) -``` - -Every reservation occupies one position in the context's reservation sequence: - -1. A sequential reservation advances the counter until it finds an unused numeric local ID. -2. A custom reservation validates uniqueness, advances the counter once, and uses the supplied local ID. -3. A generated numeric local ID skips values already claimed by custom reservations. -4. Reusing a local ID in the same context fails immediately. - -Primitive operations created without a reservation use the same counter and local-ID registry. A custom reservation -therefore cannot reuse a numeric local ID already consumed by an ordinary core operation. - -Examples: - -```text -reserve("a", "node-a") -> hash("node-a") -reserve("b") -> hash("2") -reserve("c", "2") -> fails because "2" is already used -``` - -Inside a child context whose ID is `parentHash`: - -```text -reserve("a", "node-a") -> hash("parentHash-node-a") -``` - -Custom local IDs allow the custom-ID operations themselves to keep stable identities when definition order changes. -Sequential operations around them may still receive different IDs. Adding, removing, or changing an ID remains a -workflow compatibility change. - -## Stateful Step Extensions - -`waitForCondition` needs the existing STEP state machine with checkpointed state between retry attempts. The extension -API exposes this without exposing raw checkpoint actions. - -```java -@FunctionalInterface -public interface ExtensionStepFunction { - ExtensionStepResult apply(T state); -} -``` - -```java -public sealed interface ExtensionStepResult { - static ExtensionStepResult succeed(T value); - - static ExtensionStepResult retry(T state, Duration delay); -} -``` - -```java -public final class ExtensionStepConfig { - T initialState(); - - SerDes serDes(); -} -``` - -The subtype-aware stateful step selector is: - -```java - DurableFuture stepAsync( - String subType, - TypeToken resultType, - ExtensionStepFunction function, - ExtensionStepConfig config); -``` - -The SDK interprets results through the fixed STEP state machine: - -- `succeed(value)` serializes the value and checkpoints `SUCCEED`. -- `retry(state, delay)` serializes the state and checkpoints `RETRY`. -- a thrown exception checkpoints `FAIL`. -- internal suspension and unrecoverable control-flow exceptions propagate without conversion. - -`StepContext.getCurrentContext()` exposes the one-based attempt number. The extension function receives only its -application state. - -## Extension Context Results and Replay State - -A subtype-aware extension context may return a full application result and a smaller replay state: - -```java -public final class ExtensionContextResult { - static ExtensionContextResult completed(T result); - - static ExtensionContextResult replayChildren( - T result, - T replayState); - - static ExtensionContextResult replayChildrenAboveSize( - T result, - T replayState, - int thresholdBytes); -} -``` - -The application receives `result`. The checkpoint stores either the normal serialized result or `replayState`, -according to the selected factory. - -For `replayChildrenAboveSize`, the threshold is evaluated against the serialized full application result, before the -replay state is selected for checkpointing. - -When replay-children is enabled, the SDK reexecutes the extension context function and exposes the stored replay state -through a scoped extension context: - -```java -public final class ExtensionContextReplayContext { - static ExtensionContextReplayContext getCurrentContext(); - - boolean isReplayingChildren(); - - T getReplayState(); -} -``` - -The replay context is available only while the extension framework function is running. It is restored across nested -extension contexts and is not propagated to application-created threads. - -This preserves current map and parallel behavior: - -- A small map stores and replays its complete `MapResult`. -- A large map stores statuses and completion reason as replay state, then reconstructs values from iteration - checkpoints. -- Parallel stores its current `ParallelResult` as replay state and always replays branch children. - -## Extension Context Failure Translation - -The CONTEXT state machine always handles failures in the same way: - -1. serialize the thrown exception when possible -2. checkpoint `FAIL` -3. deserialize the original exception when the future is read or replayed - -Subtype-specific behavior is customizable only as a fallback when the original exception cannot be reconstructed. - -```java -@FunctionalInterface -public interface ExtensionContextErrorHandler { - Throwable translate(ExtensionContextFailure failure); -} -``` - -`ExtensionContextFailure` is a read-only view containing: - -- context name -- context subtype string -- deserialized original exception, when available -- serialized error metadata -- child operation summaries - -Each child summary contains: - -- operation type -- subtype string -- status -- serialized error metadata - -Resolution order is: - -1. rethrow a deserialized original exception -2. invoke the configured `ExtensionContextErrorHandler` -3. fall back to `ChildContextFailedException` - -The built-in extensions provide handlers that preserve current behavior: - -- wait for callback distinguishes callback failure, callback timeout, and submitter failure -- map iteration falls back to `MapIterationFailedException` -- parallel branch falls back to `ParallelBranchFailedException` -- with retry and ordinary child contexts fall back to `ChildContextFailedException` - -## Extension Context Configuration - -Existing `RunInChildContextConfig` remains unchanged and continues to provide SerDes and virtual-context settings. - -The subtype-aware extension context overload uses a new extension-specific wrapper: - -```java -public final class ExtensionContextConfig { - RunInChildContextConfig childContextConfig(); - - ExtensionContextErrorHandler errorHandler(); - - boolean emitUserFunctionEvents(); - - boolean suppressLateChildCheckpoints(); -} -``` - -The extension context function has a distinct type so it does not conflict by erasure with existing supplier -overloads: - -```java -@FunctionalInterface -public interface ExtensionContextFunction { - ExtensionContextResult apply(); -} -``` - -```java - DurableFuture runInChildContextAsync( - String subType, - TypeToken resultType, - ExtensionContextFunction function, - ExtensionContextConfig config); -``` - -The two additional booleans preserve existing family-specific behavior: - -- `emitUserFunctionEvents` controls whether the extension context function is reported as a user function. It defaults - to `true`, matching ordinary child contexts. -- `suppressLateChildCheckpoints` tracks extension-managed children and prevents them from writing checkpoints after - their parent extension context has completed. It defaults to `false`. - -Nested user step and child-context functions retain their own existing plugin hooks regardless of the parent setting. - -## Internal Operation Identity - -The existing public `OperationSubType` enum and enum-based identity factories remain unchanged. - -Internally, primitive operations use an identity containing: - -- operation ID -- name -- backend operation type -- subtype string - -Existing enum values convert to this representation. Custom subtype strings flow unchanged through: - -- operation updates -- replay validation -- plugin events -- logs -- failure views - -Replay validation compares both operation type and exact subtype string. - -## Built-In Extension Implementations - -Each family has one canonical implementation that accepts an `ExtensionContext`. The static facade obtains it from -TLS. The corresponding `DurableContextImpl` method passes `this`. - -```text -static facade --------------------+ - +-> built-in extension -> extension primitives -legacy DurableContext adapter ----+ -``` - -The canonical implementations use the existing public configurations and callback contracts after adapting them to -the extension-specific functions. - -### Wait for callback - -The extension: - -1. reserves a CONTEXT operation with subtype `WaitForCallback` -2. creates a CALLBACK child with subtype `Callback` -3. creates a STEP child with subtype `Step` -4. runs the submitter -5. waits for the callback result - -The parent extension context emits the same context user-function hooks as the current implementation. -The context failure handler preserves callback failure, timeout, and submitter exception translation. - -### With retry - -The extension: - -1. reserves a CONTEXT operation with subtype `WithRetry` -2. uses the current virtual or checkpointed behavior from `WithRetryConfig` -3. invokes the user operation with attempt metadata in `WithRetryContext` TLS -4. creates WAIT operations for backoff using the existing names and delays - -The retry context emits the same context user-function hooks as the current implementation. -Internal suspension and unrecoverable control-flow exceptions are never retried. - -### Wait for condition - -The extension reserves a stateful STEP operation with subtype `WaitForCondition`. - -The adapter: - -1. starts with `WaitForConditionConfig.initialState()` -2. invokes the existing check function -3. returns `succeed(value)` when polling completes -4. evaluates the existing wait strategy -5. returns `retry(value, delay)` when polling continues - -Attempt metadata remains available through `StepContext`. - -### Map - -The extension reserves a CONTEXT operation with subtype `Map`. - -Inside that context it: - -1. deterministically reserves all iteration contexts in input order -2. assigns subtype `MapIteration` -3. launches iterations through the shared concurrency coordinator -4. evaluates the existing `CompletionConfig` -5. constructs the existing `MapResult` -6. uses map replay state for large results - -Iteration reservations continue using sequential local IDs so existing operation IDs remain unchanged. - -The map parent does not emit context user-function hooks. Iteration contexts do emit them. The parent enables -late-child checkpoint suppression. - -Empty maps preserve `DurableConfig.shouldCheckpointEmptyMap()`: - -- when enabled, the map parent checkpoints `START` and `SUCCEED` -- when disabled, the reservation still consumes the same operation ID, the map emits the existing warning and plugin - lifecycle, and it completes with `MapResult.empty()` without a backend checkpoint - -### Parallel - -The extension reserves a CONTEXT operation with subtype `Parallel` and returns the existing -`ParallelDurableFuture`. - -Branch calls: - -1. reserve branch identities in registration order -2. assign subtype `ParallelBranch` -3. enqueue branch definitions in the parent extension context -4. launch through the shared concurrency coordinator - -`close()` and `get()` retain current join behavior. Branch registration after join still fails. - -Each branch call returns a deferred `DurableFuture` immediately. The coordinator binds it to the reserved child -context future when concurrency capacity permits. Its `get()` and `completionFuture()` retain the behavior of the -current branch future. - -The parallel parent does not emit context user-function hooks. Branch contexts do emit them. The parent always stores -replay state, replays children, and enables late-child checkpoint suppression. - -## Shared Concurrency Coordinator - -Map and parallel share a coordinator that is not itself a durable operation. - -It owns: - -- pending registration order -- max-concurrency enforcement -- running completion signals -- success and failure counts -- `CompletionConfig` evaluation -- skipped item tracking -- late-child checkpoint suppression -- the `allItemsRegistered` transition when map registration completes or parallel is joined - -It creates no operation type or checkpoint. All durable state belongs to the parent extension context and its reserved -child context primitives. - -The coordinator uses only `DurableFuture.completionFuture()` and reserved extension operations. It does not downcast -futures to SDK operation classes. - -## Removal of Specialized Engines - -After parity is proven, the following specialized engines are removed: - -- `MapOperation` -- `ParallelOperation` -- `ConcurrencyOperation` -- `WaitForConditionOperation` - -Their reusable primitive lifecycle behavior moves into the generalized STEP and CONTEXT primitive implementations. - -`ChildContextOperation`, `StepOperation`, and the other primitive operation classes remain as the SDK-owned state -machines. They are generalized to accept string subtypes and extension-specific result or failure policies. - -## Testing Strategy - -### Reservation tests - -Cover: - -- sequential reservations retain current IDs -- custom local IDs use the current context namespace -- custom reservations advance the sequence position -- generated numeric IDs skip reserved custom values -- duplicate local IDs fail -- nested contexts hash custom IDs with their parent ID -- custom IDs remain stable when reservation order changes - -### Primitive extension tests - -Cover: - -- arbitrary subtype strings for every primitive -- operation type remains determined by the primitive selector -- replay rejects type or subtype changes -- subtype strings reach plugin events unchanged -- stateful step success, retry, replay, state serialization, and failure -- context replay state and nested TLS restoration -- custom context failure translation and default fallback - -### Built-in parity tests - -For every operation family, compare legacy and static entry points for: - -- results and thrown exception types -- operation IDs, names, types, subtypes, and parent IDs -- checkpoint status and payload shape -- replay and suspension behavior -- plugin lifecycle ordering - -The existing map, parallel, callback, condition, retry, plugin, conformance, and example tests remain behavioral gates. -Tests formerly tied to specialized classes move to the generalized primitive and coordinator implementations without -weakening their assertions. - -### Completion gate - -Run: - -```bash -mvn spotless:apply -mvn clean install -``` - -Cloud example tests remain disabled unless their existing environment requirements are configured. - -## Documentation and ADR - -Update: - -- ADR-006 to include custom local IDs, arbitrary subtype strings, fixed primitive state machines, replay state, and - customizable context failure translation -- the custom extension guide with subtype-aware and custom-ID examples -- public Javadocs for every new extension-specific contract - -The documentation must state that operation IDs, subtypes, and replay state are workflow compatibility contracts. diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/parallel/DeserializationFailedParallelExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/parallel/DeserializationFailedParallelExample.java index a54aedfab..0c938f9a0 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/parallel/DeserializationFailedParallelExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/parallel/DeserializationFailedParallelExample.java @@ -51,7 +51,10 @@ public String handleRequest(Input input, DurableContext context) { throw new RuntimeException("Intentional failure for transform"); }); }, - ParallelBranchConfig.builder().serDes(new FailedSerDes()).build()); + ParallelBranchConfig.builder() + .serDes(new FailedSerDes()) + .build() + .toOperationConfig()); parallel.get(); try { diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java index 5c69d2ee1..e3108a556 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java @@ -7,8 +7,9 @@ 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 software.amazon.lambda.durable.DurableCoreOperations.step; +import static software.amazon.lambda.durable.extension.PairOperations.customOperationsAsync; import static software.amazon.lambda.durable.extension.PairOperations.pairAsync; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; import java.nio.charset.StandardCharsets; import java.security.MessageDigest; @@ -26,6 +27,7 @@ import software.amazon.lambda.durable.extension.ExtensionStepConfig; import software.amazon.lambda.durable.extension.ExtensionStepResult; import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.testing.LocalDurableTestRunner; class ExtensionOperationIntegrationTest { @@ -48,24 +50,28 @@ void reservedOperationsReplayWhenLaunchOrderChanges() { } @Test - void customReservationsRemainStableWhenRegistrationOrderChanges() { - var invocations = new AtomicInteger(); - var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { - var extension = ExtensionContext.getCurrentContext(); - var replay = invocations.incrementAndGet() > 1; - var first = replay ? extension.reserve("right", "right") : extension.reserve("left", "left"); - var second = replay ? extension.reserve("left", "left") : extension.reserve("right", "right"); - first.step(String.class, () -> first == second ? "invalid" : "first"); - second.step(String.class, () -> "second"); - context.wait("replay", Duration.ofSeconds(1)); - return "done"; - }); + void customExtensionFixtureSupportsLocalIdsAndSubtypesAcrossReplay() { + var extensionExecutions = new AtomicInteger(); + var runner = LocalDurableTestRunner.create( + String.class, (input, context) -> customOperationsAsync("custom", extensionExecutions) + .get()); var result = runner.runUntilComplete("input"); assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); - assertEquals(hash("left"), result.getOperation("left").getId()); - assertEquals(hash("right"), result.getOperation("right").getId()); + assertEquals("step:context", result.getResult(String.class)); + assertTrue(extensionExecutions.get() >= 2); + assertEquals(hash("custom-step-id"), result.getOperation("custom-step").getId()); + assertEquals(OperationType.STEP, result.getOperation("custom-step").getType()); + assertEquals("AcmeStep", result.getOperation("custom-step").getSubtype()); + assertEquals(hash("custom-wait-id"), result.getOperation("custom-wait").getId()); + assertEquals(OperationType.WAIT, result.getOperation("custom-wait").getType()); + assertEquals("AcmeWait", result.getOperation("custom-wait").getSubtype()); + assertEquals( + hash("custom-context-id"), result.getOperation("custom-context").getId()); + assertEquals( + OperationType.CONTEXT, result.getOperation("custom-context").getType()); + assertEquals("AcmeContext", result.getOperation("custom-context").getSubtype()); } @Test @@ -92,11 +98,25 @@ void staticOperationsUseCurrentContextAndRejectStepThreads() { void extensionCanExplicitlyCreateChildContext() { var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { var outer = ExtensionContext.getCurrentContext(); - return outer.reserve("child").runInChildContext(String.class, () -> { - var child = ExtensionContext.getCurrentContext(); - assertNotSame(outer, child); - return child.reserve("value", "node").step(String.class, () -> "nested"); - }); + return outer.reserve("child") + .runInChildContextAsync( + OperationSubType.RUN_IN_CHILD_CONTEXT.getValue(), + TypeToken.get(String.class), + () -> { + var child = ExtensionContext.getCurrentContext(); + assertNotSame(outer, child); + var nested = child.reserve("value", "node") + .stepAsync( + OperationSubType.STEP.getValue(), + TypeToken.get(String.class), + state -> ExtensionStepResult.succeed("nested"), + ExtensionStepConfig.builder() + .build()) + .get(); + return ExtensionContextResult.completed(nested); + }, + ExtensionContextConfig.builder().build()) + .get(); }); var result = runner.runUntilComplete("input"); @@ -111,9 +131,26 @@ void extensionCanExplicitlyCreateChildContext() { void customPrimitiveSubtypesAreStoredWithoutChangingOperationTypes() { var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { var extension = ExtensionContext.getCurrentContext(); - extension.reserve("custom-step").step("AcmeStep", String.class, () -> "step"); - extension.reserve("custom-wait").wait("AcmeWait", Duration.ofSeconds(1)); - return extension.reserve("custom-context").runInChildContext("AcmeContext", String.class, () -> "done"); + extension + .reserve("custom-step") + .stepAsync( + "AcmeStep", + TypeToken.get(String.class), + state -> ExtensionStepResult.succeed("step"), + ExtensionStepConfig.builder().build()) + .get(); + extension + .reserve("custom-wait") + .waitAsync("AcmeWait", Duration.ofSeconds(1)) + .get(); + return extension + .reserve("custom-context") + .runInChildContextAsync( + "AcmeContext", + TypeToken.get(String.class), + () -> ExtensionContextResult.completed("done"), + ExtensionContextConfig.builder().build()) + .get(); }); var result = runner.runUntilComplete("input"); @@ -133,15 +170,16 @@ void statefulExtensionStepCheckpointsStateAcrossRetries() { var runner = LocalDurableTestRunner.create(Integer.class, (input, context) -> ExtensionContext.getCurrentContext() .reserve("stateful") - .step( + .stepAsync( "AcmeStateful", - Integer.class, + TypeToken.get(Integer.class), state -> state >= 2 ? ExtensionStepResult.succeed(state) : ExtensionStepResult.retry(state + 1, Duration.ofSeconds(1)), ExtensionStepConfig.builder() .initialState(0) - .build())); + .build()) + .get()); var result = runner.runUntilComplete(0); @@ -158,9 +196,9 @@ void extensionContextExposesStoredReplayStateWhileReplayingChildren() { var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { var result = ExtensionContext.getCurrentContext() .reserve("advanced") - .runInChildContext( + .runInChildContextAsync( "AcmeContext", - String.class, + TypeToken.get(String.class), () -> { executions.incrementAndGet(); var replay = ExtensionContextReplayContext.getCurrentContext(); @@ -169,7 +207,8 @@ void extensionContextExposesStoredReplayStateWhileReplayingChildren() { } return ExtensionContextResult.replayChildren("full", "stored"); }, - ExtensionContextConfig.builder().build()); + ExtensionContextConfig.builder().build()) + .get(); context.wait("replay", Duration.ofSeconds(1)); return result; }); diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java index d610d4eed..d396dea4c 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java @@ -16,8 +16,12 @@ import software.amazon.lambda.durable.config.StepConfig; import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionStepConfig; +import software.amazon.lambda.durable.extension.ExtensionStepResult; import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.model.WaitForConditionResult; +import software.amazon.lambda.durable.operation.DurableMapOperation; import software.amazon.lambda.durable.plugin.*; import software.amazon.lambda.durable.retry.RetryStrategies; import software.amazon.lambda.durable.testing.LocalDurableTestRunner; @@ -185,7 +189,14 @@ void plugin_receivesPrimitiveLifecycleForCustomExtension() { } private static String customExtension() { - return ExtensionContext.getCurrentContext().reserve("inner-step").step(String.class, () -> "done"); + return ExtensionContext.getCurrentContext() + .reserve("inner-step") + .stepAsync( + OperationSubType.STEP.getValue(), + TypeToken.get(String.class), + state -> ExtensionStepResult.succeed("done"), + ExtensionStepConfig.builder().build()) + .get(); } @Test @@ -414,7 +425,7 @@ void plugin_preservesMapAndIterationHookBoundaries() { var config = DurableConfig.builder().withPlugins(plugin).build(); var runner = LocalDurableTestRunner.create( String.class, - (input, context) -> DurableMapOperations.map( + (input, context) -> DurableMapOperation.map( "items", List.of("a", "b"), String.class, String::toUpperCase) .results() .toString(), @@ -714,7 +725,7 @@ void plugin_reportsFailedThenSucceededAttempts_forRetriedStep() { .toList(); assertEquals(2, flakyEnds.size(), "Expected two attempts: one failed, one succeeded"); - // First attempt failed — its exception is handled internally by StepOperation, but the + // First attempt failed — its exception is handled internally by StepPrimitive, but the // onUserFunctionEnd hook must still report it as failed. Look up by attempt number since the // retry may run nested within the failed attempt, making end-hook ordering non-deterministic. var firstAttempt = flakyEnds.stream() diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java index b51cddece..9e0c0e140 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java @@ -7,18 +7,32 @@ import static org.junit.jupiter.api.Assertions.assertNotSame; import static org.junit.jupiter.api.Assertions.assertTrue; +import java.time.Duration; import java.util.ArrayList; import java.util.List; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Supplier; import org.junit.jupiter.api.Test; +import software.amazon.awssdk.services.lambda.model.OperationStatus; +import software.amazon.awssdk.services.lambda.model.OperationType; import software.amazon.lambda.durable.config.MapConfig; import software.amazon.lambda.durable.config.NestingType; import software.amazon.lambda.durable.config.ParallelConfig; import software.amazon.lambda.durable.config.WaitForConditionConfig; +import software.amazon.lambda.durable.config.WithRetryConfig; import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.model.ExecutionStatus; import software.amazon.lambda.durable.model.WaitForConditionResult; +import software.amazon.lambda.durable.operation.DurableContextOperation; +import software.amazon.lambda.durable.operation.DurableMapOperation; +import software.amazon.lambda.durable.operation.DurableParallelOperation; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation; +import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation; +import software.amazon.lambda.durable.operation.DurableWithRetryOperation; +import software.amazon.lambda.durable.retry.RetryStrategies; +import software.amazon.lambda.durable.retry.WaitStrategies; import software.amazon.lambda.durable.testing.LocalDurableTestRunner; import software.amazon.lambda.durable.testing.TestResult; @@ -27,13 +41,13 @@ class StaticOperationsIntegrationTest { void coreOperationsExposeStepAndChildContextsThroughTls() { var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { var root = ExtensionContext.getCurrentContext(); - var step = DurableCoreOperations.step( + var step = DurableStepOperation.step( "step", String.class, () -> "attempt-" + StepContext.getCurrentContext().getAttempt()); - var child = DurableCoreOperations.runInChildContext("child", String.class, () -> { + var child = DurableContextOperation.runInChildContext("child", String.class, () -> { assertNotSame(root, ExtensionContext.getCurrentContext()); - return DurableCoreOperations.step("child-step", String.class, () -> "child"); + return DurableStepOperation.step("child-step", String.class, () -> "child"); }); return step + ":" + child; }); @@ -47,21 +61,19 @@ void coreOperationsExposeStepAndChildContextsThroughTls() { @Test void mapAndParallelExposeContextFreeUserFunctions() { var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { - var mapResult = DurableMapOperations.map("map", List.of("a", "b"), String.class, item -> { + var mapResult = DurableMapOperation.map("map", List.of("a", "b"), String.class, item -> { var index = MapItemContext.getCurrentContext().getIndex(); - return DurableCoreOperations.step("map-step", String.class, () -> item + index); + return DurableStepOperation.step("map-step", String.class, () -> item + index); }); var branchFutures = new ArrayList>(); - try (var parallel = DurableParallelOperations.parallel("parallel")) { + try (var parallel = DurableParallelOperation.parallel("parallel")) { branchFutures.add(parallel.branch( - "left", - String.class, - () -> DurableCoreOperations.step("branch-step", String.class, () -> "L"))); + "left", String.class, () -> DurableStepOperation.step("branch-step", String.class, () -> "L"))); branchFutures.add(parallel.branch( "right", String.class, - () -> DurableCoreOperations.step("branch-step", String.class, () -> "R"))); + () -> DurableStepOperation.step("branch-step", String.class, () -> "R"))); } return mapResult.results() + ":" + DurableFuture.allOf(branchFutures); }); @@ -86,15 +98,15 @@ void staticMapMatchesLegacyCheckpointHistory() { mapConfig) .results() .toString()); - var staticRunner = LocalDurableTestRunner.create(String.class, (input, context) -> DurableMapOperations.map( + var staticRunner = LocalDurableTestRunner.create(String.class, (input, context) -> DurableMapOperation.map( "map", List.of("a", "b"), String.class, item -> { var index = MapItemContext.getCurrentContext().getIndex(); - return DurableCoreOperations.step("work", String.class, () -> item + index); + return DurableStepOperation.step("work", String.class, () -> item + index); }, - mapConfig) + mapConfig.toOperationConfig()) .results() .toString()); @@ -120,11 +132,10 @@ void staticParallelMatchesLegacyCheckpointHistory() { } }); var staticRunner = LocalDurableTestRunner.create(String.class, (input, context) -> { - try (var parallel = DurableParallelOperations.parallel("parallel", parallelConfig)) { - parallel.branch( - "left", String.class, () -> DurableCoreOperations.step("work", String.class, () -> "L")); + try (var parallel = DurableParallelOperation.parallel("parallel", parallelConfig.toOperationConfig())) { + parallel.branch("left", String.class, () -> DurableStepOperation.step("work", String.class, () -> "L")); parallel.branch( - "right", String.class, () -> DurableCoreOperations.step("work", String.class, () -> "R")); + "right", String.class, () -> DurableStepOperation.step("work", String.class, () -> "R")); return parallel.get().statuses().toString(); } }); @@ -137,19 +148,106 @@ void staticParallelMatchesLegacyCheckpointHistory() { assertEquals(operationHistory(legacyResult), operationHistory(staticResult)); } + @Test + void staticWaitForCallbackMatchesLegacyCheckpointHistory() { + var legacyRunner = LocalDurableTestRunner.create( + String.class, + (input, context) -> context.waitForCallback("approval", String.class, (callbackId, stepContext) -> {})); + var staticRunner = LocalDurableTestRunner.create( + String.class, + (input, context) -> + DurableWaitForCallbackOperation.waitForCallback("approval", String.class, () -> {})); + + var legacyPending = legacyRunner.run("input"); + var staticPending = staticRunner.run("input"); + assertEquals(ExecutionStatus.PENDING, legacyPending.getStatus()); + assertEquals(ExecutionStatus.PENDING, staticPending.getStatus()); + legacyRunner.completeCallback(legacyRunner.getCallbackId("approval-callback"), "\"approved\""); + staticRunner.completeCallback(staticRunner.getCallbackId("approval-callback"), "\"approved\""); + + var legacyResult = legacyRunner.runUntilComplete("input"); + var staticResult = staticRunner.runUntilComplete("input"); + + assertEquals("approved", legacyResult.getResult(String.class)); + assertEquals(legacyResult.getResult(String.class), staticResult.getResult(String.class)); + assertEquals(operationHistory(legacyResult), operationHistory(staticResult)); + } + + @Test + void staticWaitForConditionMatchesLegacyCheckpointHistory() { + var config = WaitForConditionConfig.builder() + .initialState(0) + .waitStrategy(WaitStrategies.fixedDelay(3, Duration.ofSeconds(1))) + .build(); + var legacyRunner = LocalDurableTestRunner.create( + String.class, + (input, context) -> context.waitForCondition( + "condition", Integer.class, (state, stepContext) -> nextConditionState(state), config)); + var staticRunner = LocalDurableTestRunner.create( + String.class, + (input, context) -> DurableWaitForConditionOperation.waitForCondition( + "condition", + Integer.class, + StaticOperationsIntegrationTest::nextConditionState, + config.toOperationConfig())); + + var legacyResult = legacyRunner.runUntilComplete("input"); + var staticResult = staticRunner.runUntilComplete("input"); + + assertEquals(2, legacyResult.getResult(Integer.class)); + assertEquals(legacyResult.getResult(Integer.class), staticResult.getResult(Integer.class)); + assertEquals(operationHistory(legacyResult), operationHistory(staticResult)); + } + + @Test + void staticWithRetryMatchesLegacyCheckpointHistory() { + var config = WithRetryConfig.builder() + .retryStrategy(RetryStrategies.fixedDelay(2, Duration.ofSeconds(1))) + .wrapInChildContext(true) + .build(); + var legacyRunner = LocalDurableTestRunner.create( + String.class, + (input, context) -> context.withRetry( + "retry", + (attempt, child) -> retryAttempt( + attempt, () -> child.step("work", String.class, stepContext -> "attempt-" + attempt)), + config)); + var staticRunner = LocalDurableTestRunner.create( + String.class, + (input, context) -> DurableWithRetryOperation.withRetry( + "retry", + () -> { + var attempt = WithRetryContext.getCurrentContext().getAttempt(); + return retryAttempt( + attempt, + () -> DurableStepOperation.step("work", String.class, () -> "attempt-" + attempt)); + }, + config.toOperationConfig())); + + var legacyResult = legacyRunner.runUntilComplete("input"); + var staticResult = staticRunner.runUntilComplete("input"); + + assertEquals("attempt-2", legacyResult.getResult(String.class)); + assertEquals(legacyResult.getResult(String.class), staticResult.getResult(String.class)); + assertEquals(operationHistory(legacyResult), operationHistory(staticResult)); + } + @Test void conditionAndRetryExposeGeneratedMetadataThroughTls() { var retryExecutions = new AtomicInteger(); var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { - var condition = DurableWaitForConditionOperations.waitForCondition( + var condition = DurableWaitForConditionOperation.waitForCondition( "condition", Integer.class, state -> { assertNotNull(StepContext.getCurrentContext()); return WaitForConditionResult.stopPolling(state + 1); }, - WaitForConditionConfig.builder().initialState(0).build()); - var retry = DurableWithRetryOperations.withRetry("retry", () -> { + WaitForConditionConfig.builder() + .initialState(0) + .build() + .toOperationConfig()); + var retry = DurableWithRetryOperation.withRetry("retry", () -> { var attempt = WithRetryContext.getCurrentContext().getAttempt(); retryExecutions.incrementAndGet(); if (attempt == 1) { @@ -172,7 +270,7 @@ void waitForCallbackExposesCallbackIdThroughTls() { var submittedId = new AtomicReference(); var runner = LocalDurableTestRunner.create( String.class, - (input, context) -> DurableWaitForCallbackOperations.waitForCallback( + (input, context) -> DurableWaitForCallbackOperation.waitForCallback( "approval", String.class, () -> submittedId.set( @@ -191,15 +289,30 @@ void waitForCallbackExposesCallbackIdThroughTls() { assertEquals("approved", completed.getResult(String.class)); } - private static List operationHistory(TestResult result) { + private static WaitForConditionResult nextConditionState(int state) { + var next = state + 1; + return next >= 2 ? WaitForConditionResult.stopPolling(next) : WaitForConditionResult.continuePolling(next); + } + + private static String retryAttempt(int attempt, Supplier operation) { + if (attempt == 1) { + throw new IllegalStateException("retry"); + } + return operation.get(); + } + + private static List operationHistory(TestResult result) { return result.getOperations().stream() - .map(operation -> String.join( - ":", + .map(operation -> new OperationShape( operation.getId(), operation.getName(), - operation.getType().toString(), + operation.getType(), operation.getSubtype(), - operation.getStatus().toString())) + operation.getEvents().get(0).parentId(), + operation.getStatus())) .toList(); } + + private record OperationShape( + String id, String name, OperationType type, String subtype, String parentId, OperationStatus status) {} } diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/extension/PairOperations.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/extension/PairOperations.java index 770081e21..555da4e7e 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/extension/PairOperations.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/extension/PairOperations.java @@ -6,6 +6,8 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.atomic.AtomicInteger; import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.model.OperationSubType; /** Example extension library implemented only with public SDK contracts. */ public final class PairOperations { @@ -20,14 +22,52 @@ public static DurableFuture pairAsync(String name, AtomicInteger extensi DurableFuture leftFuture; DurableFuture rightFuture; if (extensionExecutions.getAndIncrement() % 2 == 0) { - leftFuture = left.stepAsync(String.class, () -> "L"); - rightFuture = right.stepAsync(String.class, () -> "R"); + leftFuture = stepAsync(left, "L"); + rightFuture = stepAsync(right, "R"); } else { - rightFuture = right.stepAsync(String.class, () -> "R"); - leftFuture = left.stepAsync(String.class, () -> "L"); + rightFuture = stepAsync(right, "R"); + leftFuture = stepAsync(left, "L"); } - return new PairFuture(leftFuture, rightFuture, pause.waitAsync(Duration.ofSeconds(1))); + return new PairFuture( + leftFuture, rightFuture, pause.waitAsync(OperationSubType.WAIT.getValue(), Duration.ofSeconds(1))); + } + + public static DurableFuture customOperationsAsync(String name, AtomicInteger extensionExecutions) { + var extension = ExtensionContext.getCurrentContext(); + ExtensionOperation step; + ExtensionOperation wait; + ExtensionOperation context; + if (extensionExecutions.getAndIncrement() % 2 == 0) { + step = extension.reserve(name + "-step", name + "-step-id"); + wait = extension.reserve(name + "-wait", name + "-wait-id"); + context = extension.reserve(name + "-context", name + "-context-id"); + } else { + context = extension.reserve(name + "-context", name + "-context-id"); + wait = extension.reserve(name + "-wait", name + "-wait-id"); + step = extension.reserve(name + "-step", name + "-step-id"); + } + + var stepFuture = step.stepAsync( + "AcmeStep", + TypeToken.get(String.class), + state -> ExtensionStepResult.succeed("step"), + ExtensionStepConfig.builder().build()); + var waitFuture = wait.waitAsync("AcmeWait", Duration.ofSeconds(1)); + var contextFuture = context.runInChildContextAsync( + "AcmeContext", + TypeToken.get(String.class), + () -> ExtensionContextResult.completed("context"), + ExtensionContextConfig.builder().build()); + return new CustomOperationsFuture(stepFuture, waitFuture, contextFuture); + } + + private static DurableFuture stepAsync(ExtensionOperation operation, String result) { + return operation.stepAsync( + OperationSubType.STEP.getValue(), + TypeToken.get(String.class), + state -> ExtensionStepResult.succeed(result), + ExtensionStepConfig.builder().build()); } private record PairFuture(DurableFuture left, DurableFuture right, DurableFuture pause) @@ -43,4 +83,20 @@ public CompletableFuture completionFuture() { return CompletableFuture.allOf(left.completionFuture(), right.completionFuture(), pause.completionFuture()); } } + + private record CustomOperationsFuture( + DurableFuture step, DurableFuture pause, DurableFuture context) + implements DurableFuture { + @Override + public String get() { + pause.get(); + return step.get() + ":" + context.get(); + } + + @Override + public CompletableFuture completionFuture() { + return CompletableFuture.allOf( + step.completionFuture(), pause.completionFuture(), context.completionFuture()); + } + } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java index 19fb7ebad..a295d8b37 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java @@ -18,8 +18,14 @@ import software.amazon.lambda.durable.config.WaitForConditionConfig; import software.amazon.lambda.durable.config.WithRetryConfig; import software.amazon.lambda.durable.context.BaseContext; +import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.model.MapResult; import software.amazon.lambda.durable.model.WaitForConditionResult; +import software.amazon.lambda.durable.operation.DurableCallbackOperation; +import software.amazon.lambda.durable.operation.DurableContextOperation; +import software.amazon.lambda.durable.operation.DurableInvokeOperation; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableWaitOperation; public interface DurableContext extends BaseContext { /** @@ -157,8 +163,11 @@ default DurableFuture stepAsync(String name, TypeToken resultType, Fun * @param config the step configuration (retry strategy, semantics, custom SerDes) * @return a future representing the step result */ - DurableFuture stepAsync( - String name, TypeToken resultType, Function func, StepConfig config); + default DurableFuture stepAsync( + String name, TypeToken resultType, Function func, StepConfig config) { + return DurableStepOperation.stepAsync( + (ExtensionContext) this, name, resultType, func, config.toOperationConfig()); + } /** @deprecated use the variants accepting StepContext instead */ @Deprecated @@ -236,7 +245,9 @@ default Void wait(String name, Duration duration) { * @param duration the duration to wait * @return a future that completes when the wait duration has elapsed */ - DurableFuture waitAsync(String name, Duration duration); + default DurableFuture waitAsync(String name, Duration duration) { + return DurableWaitOperation.waitAsync((ExtensionContext) this, name, duration); + } /** * Invokes another Lambda function by name and blocks until the result is available. @@ -320,8 +331,11 @@ default DurableFuture invokeAsync(String name, String functionName, U * @param config the invoke configuration (custom SerDes for result and payload) * @return a future representing the invocation result */ - DurableFuture invokeAsync( - String name, String functionName, U payload, TypeToken resultType, InvokeConfig config); + default DurableFuture invokeAsync( + String name, String functionName, U payload, TypeToken resultType, InvokeConfig config) { + return DurableInvokeOperation.invokeAsync( + (ExtensionContext) this, name, functionName, payload, resultType, config.toOperationConfig()); + } /** Creates a callback with custom configuration. */ default DurableCallbackFuture createCallback(String name, Class resultType, CallbackConfig config) { @@ -351,7 +365,10 @@ default DurableCallbackFuture createCallback(String name, Class result * @param config the callback configuration (custom SerDes) * @return a future containing the callback ID and eventual result */ - DurableCallbackFuture createCallback(String name, TypeToken resultType, CallbackConfig config); + default DurableCallbackFuture createCallback(String name, TypeToken resultType, CallbackConfig config) { + return DurableCallbackOperation.createCallback( + (ExtensionContext) this, name, resultType, config.toOperationConfig()); + } /** * Runs a function in a child context, blocking until it completes. @@ -487,8 +504,11 @@ default DurableFuture runInChildContextAsync( * @param config the configuration for the child context * @return the DurableFuture wrapping the child context result */ - DurableFuture runInChildContextAsync( - String name, TypeToken resultType, Function func, RunInChildContextConfig config); + default DurableFuture runInChildContextAsync( + String name, TypeToken resultType, Function func, RunInChildContextConfig config) { + return DurableContextOperation.runInChildContextAsync( + (ExtensionContext) this, name, resultType, func, config.toOperationConfig()); + } default MapResult map(String name, Collection items, Class resultType, MapFunction function) { return mapAsync( diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableCoreOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableCoreOperations.java deleted file mode 100644 index e4bfbad51..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableCoreOperations.java +++ /dev/null @@ -1,157 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; - -import java.time.Duration; -import java.util.function.Supplier; -import software.amazon.lambda.durable.config.CallbackConfig; -import software.amazon.lambda.durable.config.InvokeConfig; -import software.amazon.lambda.durable.config.RunInChildContextConfig; -import software.amazon.lambda.durable.config.StepConfig; - -/** - * Context-free static facades for SDK-owned primitive durable operations. - * - *

The equivalent instance methods on {@link DurableContext} remain supported for backward compatibility. - */ -public final class DurableCoreOperations { - private DurableCoreOperations() {} - - public static T step(String name, Class resultType, Supplier function) { - return currentContext().step(name, resultType, ignored -> function.get()); - } - - public static T step(String name, TypeToken resultType, Supplier function) { - return currentContext().step(name, resultType, ignored -> function.get()); - } - - public static T step(String name, Class resultType, Supplier function, StepConfig config) { - return currentContext().step(name, resultType, ignored -> function.get(), config); - } - - public static T step(String name, TypeToken resultType, Supplier function, StepConfig config) { - return currentContext().step(name, resultType, ignored -> function.get(), config); - } - - public static DurableFuture stepAsync(String name, Class resultType, Supplier function) { - return currentContext().stepAsync(name, resultType, ignored -> function.get()); - } - - public static DurableFuture stepAsync(String name, TypeToken resultType, Supplier function) { - return currentContext().stepAsync(name, resultType, ignored -> function.get()); - } - - public static DurableFuture stepAsync( - String name, Class resultType, Supplier function, StepConfig config) { - return currentContext().stepAsync(name, resultType, ignored -> function.get(), config); - } - - public static DurableFuture stepAsync( - String name, TypeToken resultType, Supplier function, StepConfig config) { - return currentContext().stepAsync(name, resultType, ignored -> function.get(), config); - } - - public static Void wait(String name, Duration duration) { - return currentContext().wait(name, duration); - } - - public static DurableFuture waitAsync(String name, Duration duration) { - return currentContext().waitAsync(name, duration); - } - - public static T invoke(String name, String functionName, U payload, Class resultType) { - return currentContext().invoke(name, functionName, payload, resultType); - } - - public static T invoke(String name, String functionName, U payload, TypeToken resultType) { - return currentContext().invoke(name, functionName, payload, resultType); - } - - public static T invoke( - String name, String functionName, U payload, Class resultType, InvokeConfig config) { - return currentContext().invoke(name, functionName, payload, resultType, config); - } - - public static T invoke( - String name, String functionName, U payload, TypeToken resultType, InvokeConfig config) { - return currentContext().invoke(name, functionName, payload, resultType, config); - } - - public static DurableFuture invokeAsync( - String name, String functionName, U payload, Class resultType) { - return currentContext().invokeAsync(name, functionName, payload, resultType); - } - - public static DurableFuture invokeAsync( - String name, String functionName, U payload, TypeToken resultType) { - return currentContext().invokeAsync(name, functionName, payload, resultType); - } - - public static DurableFuture invokeAsync( - String name, String functionName, U payload, Class resultType, InvokeConfig config) { - return currentContext().invokeAsync(name, functionName, payload, resultType, config); - } - - public static DurableFuture invokeAsync( - String name, String functionName, U payload, TypeToken resultType, InvokeConfig config) { - return currentContext().invokeAsync(name, functionName, payload, resultType, config); - } - - public static DurableCallbackFuture createCallback(String name, Class resultType) { - return currentContext().createCallback(name, resultType); - } - - public static DurableCallbackFuture createCallback(String name, TypeToken resultType) { - return currentContext().createCallback(name, resultType); - } - - public static DurableCallbackFuture createCallback(String name, Class resultType, CallbackConfig config) { - return currentContext().createCallback(name, resultType, config); - } - - public static DurableCallbackFuture createCallback( - String name, TypeToken resultType, CallbackConfig config) { - return currentContext().createCallback(name, resultType, config); - } - - public static T runInChildContext(String name, Class resultType, Supplier function) { - return currentContext().runInChildContext(name, resultType, ignored -> function.get()); - } - - public static T runInChildContext(String name, TypeToken resultType, Supplier function) { - return currentContext().runInChildContext(name, resultType, ignored -> function.get()); - } - - public static T runInChildContext( - String name, Class resultType, Supplier function, RunInChildContextConfig config) { - return currentContext().runInChildContext(name, resultType, ignored -> function.get(), config); - } - - public static T runInChildContext( - String name, TypeToken resultType, Supplier function, RunInChildContextConfig config) { - return currentContext().runInChildContext(name, resultType, ignored -> function.get(), config); - } - - public static DurableFuture runInChildContextAsync(String name, Class resultType, Supplier function) { - return currentContext().runInChildContextAsync(name, resultType, ignored -> function.get()); - } - - public static DurableFuture runInChildContextAsync( - String name, TypeToken resultType, Supplier function) { - return currentContext().runInChildContextAsync(name, resultType, ignored -> function.get()); - } - - public static DurableFuture runInChildContextAsync( - String name, Class resultType, Supplier function, RunInChildContextConfig config) { - return currentContext().runInChildContextAsync(name, resultType, ignored -> function.get(), config); - } - - public static DurableFuture runInChildContextAsync( - String name, TypeToken resultType, Supplier function, RunInChildContextConfig config) { - return currentContext().runInChildContextAsync(name, resultType, ignored -> function.get(), config); - } - - private static DurableContext currentContext() { - return DurableContext.getCurrentContext(); - } -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableHandler.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableHandler.java index 2342c70ac..fd065545b 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableHandler.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableHandler.java @@ -16,9 +16,9 @@ /** * Abstract base class for Lambda handlers that use durable execution. * - *

Extend this class and implement {@link #handleRequest(Object, DurableContext)} to build resilient, multi-step - * workflows. The handler automatically manages checkpoint-and-replay, input deserialization, and communication with the - * Lambda Durable Functions backend. + *

Extend this class and implement either {@link #handleRequest(Object)} or {@link #handleRequest(Object, + * DurableContext)} to build resilient, multi-step workflows. The handler automatically manages checkpoint-and-replay, + * input deserialization, and communication with the Lambda Durable Functions backend. * * @param the input type * @param the output type @@ -152,11 +152,29 @@ public final void handleRequest(InputStream inputStream, OutputStream outputStre } /** - * Handle the durable execution. + * Handles the durable execution without receiving the durable context directly. + * + *

Override either this method or {@link #handleRequest(Object, DurableContext)}. + * + * @param input User input + * @return Result + */ + public O handleRequest(I input) { + throw new UnsupportedOperationException( + "DurableHandler must override handleRequest(input) or handleRequest(input, context)"); + } + + /** + * Handles the durable execution with access to the durable context. + * + *

Override this method when the handler needs direct access to the context. By default, it delegates to + * {@link #handleRequest(Object)}. * * @param input User input * @param context Durable context for operations * @return Result */ - public abstract O handleRequest(I input, DurableContext context); + public O handleRequest(I input, DurableContext context) { + return handleRequest(input); + } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableMapOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableMapOperations.java deleted file mode 100644 index de1761e8f..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableMapOperations.java +++ /dev/null @@ -1,69 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; - -import java.util.Collection; -import java.util.Objects; -import java.util.function.Function; -import software.amazon.lambda.durable.config.MapConfig; -import software.amazon.lambda.durable.context.extension.MapExtension; -import software.amazon.lambda.durable.extension.ExtensionContext; -import software.amazon.lambda.durable.model.MapResult; - -/** Context-free static facades for durable map operations. */ -public final class DurableMapOperations { - private DurableMapOperations() {} - - public static MapResult map( - String name, Collection items, Class resultType, Function function) { - return mapAsync(name, items, resultType, function).get(); - } - - public static MapResult map( - String name, Collection items, TypeToken resultType, Function function) { - return mapAsync(name, items, resultType, function).get(); - } - - public static MapResult map( - String name, Collection items, Class resultType, Function function, MapConfig config) { - return mapAsync(name, items, resultType, function, config).get(); - } - - public static MapResult map( - String name, Collection items, TypeToken resultType, Function function, MapConfig config) { - return mapAsync(name, items, resultType, function, config).get(); - } - - public static DurableFuture> mapAsync( - String name, Collection items, Class resultType, Function function) { - return mapAsync(name, items, TypeToken.get(resultType), function); - } - - public static DurableFuture> mapAsync( - String name, Collection items, TypeToken resultType, Function function) { - return mapAsync(name, items, resultType, function, MapConfig.builder().build()); - } - - public static DurableFuture> mapAsync( - String name, Collection items, Class resultType, Function function, MapConfig config) { - return mapAsync(name, items, TypeToken.get(resultType), function, config); - } - - public static DurableFuture> mapAsync( - String name, Collection items, TypeToken resultType, Function function, MapConfig config) { - return MapExtension.execute(currentContext(), name, items, resultType, adapt(function), config); - } - - private static DurableContext.MapFunction adapt(Function function) { - Objects.requireNonNull(function, "function cannot be null"); - return (item, index, ignored) -> { - try (var scope = MapItemContext.attach(index)) { - return function.apply(item); - } - }; - } - - private static ExtensionContext currentContext() { - return ExtensionContext.getCurrentContext(); - } -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableParallelOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableParallelOperations.java deleted file mode 100644 index 25c7751b6..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableParallelOperations.java +++ /dev/null @@ -1,20 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; - -import software.amazon.lambda.durable.config.ParallelConfig; -import software.amazon.lambda.durable.context.extension.ParallelExtension; -import software.amazon.lambda.durable.extension.ExtensionContext; - -/** Context-free static facades for durable parallel operations. */ -public final class DurableParallelOperations { - private DurableParallelOperations() {} - - public static ParallelDurableFuture parallel(String name) { - return parallel(name, ParallelConfig.builder().build()); - } - - public static ParallelDurableFuture parallel(String name, ParallelConfig config) { - return ParallelExtension.execute(ExtensionContext.getCurrentContext(), name, config); - } -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForCallbackOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForCallbackOperations.java deleted file mode 100644 index f9cedb99e..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForCallbackOperations.java +++ /dev/null @@ -1,64 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; - -import java.util.Objects; -import java.util.function.BiConsumer; -import software.amazon.lambda.durable.config.WaitForCallbackConfig; -import software.amazon.lambda.durable.context.extension.WaitForCallbackExtension; -import software.amazon.lambda.durable.extension.ExtensionContext; - -/** Context-free static facades for durable wait-for-callback operations. */ -public final class DurableWaitForCallbackOperations { - private DurableWaitForCallbackOperations() {} - - public static T waitForCallback(String name, Class resultType, Runnable submitter) { - return waitForCallbackAsync(name, resultType, submitter).get(); - } - - public static T waitForCallback(String name, TypeToken resultType, Runnable submitter) { - return waitForCallbackAsync(name, resultType, submitter).get(); - } - - public static T waitForCallback( - String name, Class resultType, Runnable submitter, WaitForCallbackConfig config) { - return waitForCallbackAsync(name, resultType, submitter, config).get(); - } - - public static T waitForCallback( - String name, TypeToken resultType, Runnable submitter, WaitForCallbackConfig config) { - return waitForCallbackAsync(name, resultType, submitter, config).get(); - } - - public static DurableFuture waitForCallbackAsync(String name, Class resultType, Runnable submitter) { - return waitForCallbackAsync(name, TypeToken.get(resultType), submitter); - } - - public static DurableFuture waitForCallbackAsync(String name, TypeToken resultType, Runnable submitter) { - return waitForCallbackAsync( - name, resultType, submitter, WaitForCallbackConfig.builder().build()); - } - - public static DurableFuture waitForCallbackAsync( - String name, Class resultType, Runnable submitter, WaitForCallbackConfig config) { - return waitForCallbackAsync(name, TypeToken.get(resultType), submitter, config); - } - - public static DurableFuture waitForCallbackAsync( - String name, TypeToken resultType, Runnable submitter, WaitForCallbackConfig config) { - return WaitForCallbackExtension.execute(currentContext(), name, resultType, adapt(submitter), config); - } - - private static BiConsumer adapt(Runnable submitter) { - Objects.requireNonNull(submitter, "submitter cannot be null"); - return (callbackId, ignored) -> { - try (var scope = WaitForCallbackContext.attach(callbackId)) { - submitter.run(); - } - }; - } - - private static ExtensionContext currentContext() { - return ExtensionContext.getCurrentContext(); - } -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForConditionOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForConditionOperations.java deleted file mode 100644 index 1d30320ef..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableWaitForConditionOperations.java +++ /dev/null @@ -1,82 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; - -import java.util.Objects; -import java.util.function.BiFunction; -import java.util.function.Function; -import software.amazon.lambda.durable.config.WaitForConditionConfig; -import software.amazon.lambda.durable.context.extension.WaitForConditionExtension; -import software.amazon.lambda.durable.extension.ExtensionContext; -import software.amazon.lambda.durable.model.WaitForConditionResult; - -/** Context-free static facades for durable wait-for-condition operations. */ -public final class DurableWaitForConditionOperations { - private DurableWaitForConditionOperations() {} - - public static T waitForCondition( - String name, Class resultType, Function> checkFunction) { - return waitForConditionAsync(name, resultType, checkFunction).get(); - } - - public static T waitForCondition( - String name, TypeToken resultType, Function> checkFunction) { - return waitForConditionAsync(name, resultType, checkFunction).get(); - } - - public static T waitForCondition( - String name, - Class resultType, - Function> checkFunction, - WaitForConditionConfig config) { - return waitForConditionAsync(name, resultType, checkFunction, config).get(); - } - - public static T waitForCondition( - String name, - TypeToken resultType, - Function> checkFunction, - WaitForConditionConfig config) { - return waitForConditionAsync(name, resultType, checkFunction, config).get(); - } - - public static DurableFuture waitForConditionAsync( - String name, Class resultType, Function> checkFunction) { - return waitForConditionAsync(name, TypeToken.get(resultType), checkFunction); - } - - public static DurableFuture waitForConditionAsync( - String name, TypeToken resultType, Function> checkFunction) { - return waitForConditionAsync( - name, - resultType, - checkFunction, - WaitForConditionConfig.builder().build()); - } - - public static DurableFuture waitForConditionAsync( - String name, - Class resultType, - Function> checkFunction, - WaitForConditionConfig config) { - return waitForConditionAsync(name, TypeToken.get(resultType), checkFunction, config); - } - - public static DurableFuture waitForConditionAsync( - String name, - TypeToken resultType, - Function> checkFunction, - WaitForConditionConfig config) { - return WaitForConditionExtension.execute(currentContext(), name, resultType, adapt(checkFunction), config); - } - - private static BiFunction> adapt( - Function> checkFunction) { - Objects.requireNonNull(checkFunction, "checkFunction cannot be null"); - return (state, ignored) -> checkFunction.apply(state); - } - - private static ExtensionContext currentContext() { - return ExtensionContext.getCurrentContext(); - } -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableWithRetryOperations.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableWithRetryOperations.java deleted file mode 100644 index c9fac5413..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableWithRetryOperations.java +++ /dev/null @@ -1,44 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; - -import java.util.Objects; -import java.util.function.BiFunction; -import java.util.function.Supplier; -import software.amazon.lambda.durable.config.WithRetryConfig; -import software.amazon.lambda.durable.context.extension.WithRetryExtension; -import software.amazon.lambda.durable.extension.ExtensionContext; - -/** Context-free static facades for replay-safe retry operations. */ -public final class DurableWithRetryOperations { - private DurableWithRetryOperations() {} - - public static T withRetry(String name, Supplier operation) { - return withRetryAsync(name, operation).get(); - } - - public static T withRetry(String name, Supplier operation, WithRetryConfig config) { - return withRetryAsync(name, operation, config).get(); - } - - public static DurableFuture withRetryAsync(String name, Supplier operation) { - return withRetryAsync(name, operation, WithRetryConfig.builder().build()); - } - - public static DurableFuture withRetryAsync(String name, Supplier operation, WithRetryConfig config) { - return WithRetryExtension.execute(currentContext(), name, adapt(operation), config); - } - - private static BiFunction adapt(Supplier operation) { - Objects.requireNonNull(operation, "operation cannot be null"); - return (attempt, ignored) -> { - try (var scope = WithRetryContext.attach(attempt)) { - return operation.get(); - } - }; - } - - private static ExtensionContext currentContext() { - return ExtensionContext.getCurrentContext(); - } -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/MapItemContext.java b/sdk/src/main/java/software/amazon/lambda/durable/MapItemContext.java index 14b2689b7..99174a388 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/MapItemContext.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/MapItemContext.java @@ -25,7 +25,8 @@ public int getIndex() { return index; } - static SafeCloseable attach(int index) { + /** Attaches map item metadata for the duration of the returned scope. */ + public static SafeCloseable attach(int index) { return CURRENT.attach(new MapItemContext(index)); } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java b/sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java index a1a1dae1c..d18b2234a 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java @@ -4,9 +4,9 @@ import java.util.function.Function; import java.util.function.Supplier; -import software.amazon.lambda.durable.config.ParallelBranchConfig; import software.amazon.lambda.durable.model.ParallelResult; import software.amazon.lambda.durable.model.SafeCloseable; +import software.amazon.lambda.durable.operation.DurableParallelOperation.ParallelBranchConfig; /** User-facing context for managing parallel branch execution within a durable function. */ public interface ParallelDurableFuture extends SafeCloseable, DurableFuture { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/WaitForCallbackContext.java b/sdk/src/main/java/software/amazon/lambda/durable/WaitForCallbackContext.java index 68a528c84..53a2a6721 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/WaitForCallbackContext.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/WaitForCallbackContext.java @@ -26,7 +26,8 @@ public String getCallbackId() { return callbackId; } - static SafeCloseable attach(String callbackId) { + /** Attaches callback metadata for the duration of the returned scope. */ + public static SafeCloseable attach(String callbackId) { return CURRENT.attach(new WaitForCallbackContext(callbackId)); } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/WithRetryContext.java b/sdk/src/main/java/software/amazon/lambda/durable/WithRetryContext.java index 8c4953bf8..d84370b21 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/WithRetryContext.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/WithRetryContext.java @@ -25,7 +25,8 @@ public int getAttempt() { return attempt; } - static SafeCloseable attach(int attempt) { + /** Attaches retry metadata for the duration of the returned scope. */ + public static SafeCloseable attach(int attempt) { return CURRENT.attach(new WithRetryContext(attempt)); } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/CallbackConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/CallbackConfig.java index a71962aff..ea7977808 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/CallbackConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/CallbackConfig.java @@ -3,6 +3,7 @@ package software.amazon.lambda.durable.config; import java.time.Duration; +import software.amazon.lambda.durable.operation.DurableCallbackOperation; import software.amazon.lambda.durable.serde.SerDes; import software.amazon.lambda.durable.util.ParameterValidator; @@ -51,6 +52,15 @@ public Builder toBuilder() { return new Builder(timeout, heartbeatTimeout, serDes); } + /** Converts this compatibility config to the operation-owned config. */ + public DurableCallbackOperation.CallbackConfig toOperationConfig() { + return DurableCallbackOperation.CallbackConfig.builder() + .timeout(timeout()) + .heartbeatTimeout(heartbeatTimeout()) + .serDes(serDes()) + .build(); + } + /** Builder for {@link CallbackConfig}. */ public static class Builder { private Duration timeout; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/InvokeConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/InvokeConfig.java index e9dc7af24..9d94073b9 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/InvokeConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/InvokeConfig.java @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.config; +import software.amazon.lambda.durable.operation.DurableInvokeOperation; import software.amazon.lambda.durable.serde.SerDes; /** @@ -40,6 +41,15 @@ public Builder toBuilder() { return new Builder(payloadSerDes, resultSerDes, tenantId); } + /** Converts this compatibility config to the operation-owned config. */ + public DurableInvokeOperation.InvokeConfig toOperationConfig() { + return DurableInvokeOperation.InvokeConfig.builder() + .payloadSerDes(payloadSerDes()) + .serDes(serDes()) + .tenantId(tenantId()) + .build(); + } + /** Builder for creating InvokeConfig instances. */ public static class Builder { private SerDes payloadSerDes; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/MapConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/MapConfig.java index 78bdcc293..75a40b2bf 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/MapConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/MapConfig.java @@ -4,6 +4,7 @@ import java.util.Objects; import java.util.function.BiFunction; +import software.amazon.lambda.durable.operation.DurableMapOperation; import software.amazon.lambda.durable.serde.SerDes; /** @@ -74,6 +75,17 @@ public Builder toBuilder() { .itemNamer(itemNamer); } + /** Converts this compatibility config to the operation-owned config. */ + public DurableMapOperation.MapConfig toOperationConfig() { + return DurableMapOperation.MapConfig.builder() + .maxConcurrency(maxConcurrency()) + .completionConfig(completionConfig()) + .serDes(serDes()) + .nestingType(nestingType()) + .itemNamer(itemNamer()) + .build(); + } + /** Builder for creating MapConfig instances. */ public static class Builder { public NestingType nestingType; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/ParallelBranchConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/ParallelBranchConfig.java index 689f9aa54..846289e5b 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/ParallelBranchConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/ParallelBranchConfig.java @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.config; +import software.amazon.lambda.durable.operation.DurableParallelOperation; import software.amazon.lambda.durable.serde.SerDes; /** @@ -25,6 +26,13 @@ public Builder toBuilder() { return new Builder(serDes); } + /** Converts this compatibility config to the operation-owned config. */ + public DurableParallelOperation.ParallelBranchConfig toOperationConfig() { + return DurableParallelOperation.ParallelBranchConfig.builder() + .serDes(serDes()) + .build(); + } + /** * Creates a new builder for ParallelBranchConfig. * diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/ParallelConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/ParallelConfig.java index 863d36972..2cadd8100 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/ParallelConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/ParallelConfig.java @@ -3,6 +3,7 @@ package software.amazon.lambda.durable.config; import java.util.Objects; +import software.amazon.lambda.durable.operation.DurableParallelOperation; /** * Configuration options for parallel operations in durable executions. @@ -52,6 +53,15 @@ public Builder toBuilder() { .nestingType(nestingType); } + /** Converts this compatibility config to the operation-owned config. */ + public DurableParallelOperation.ParallelConfig toOperationConfig() { + return DurableParallelOperation.ParallelConfig.builder() + .maxConcurrency(maxConcurrency()) + .completionConfig(completionConfig()) + .nestingType(nestingType()) + .build(); + } + /** Builder for creating ParallelConfig instances. */ public static class Builder { private Integer maxConcurrency; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/RunInChildContextConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/RunInChildContextConfig.java index 93fde31e8..d1eff46c8 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/RunInChildContextConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/RunInChildContextConfig.java @@ -3,6 +3,7 @@ package software.amazon.lambda.durable.config; import java.util.Objects; +import software.amazon.lambda.durable.operation.DurableContextOperation; import software.amazon.lambda.durable.serde.SerDes; /** @@ -36,6 +37,14 @@ public Builder toBuilder() { return new Builder().serDes(serDes).isVirtual(isVirtual); } + /** Converts this compatibility config to the operation-owned config. */ + public DurableContextOperation.RunInChildContextConfig toOperationConfig() { + return DurableContextOperation.RunInChildContextConfig.builder() + .serDes(serDes()) + .isVirtual(isVirtual()) + .build(); + } + /** * Creates a new builder for RunInChildContextConfig. * diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/StepConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/StepConfig.java index 92a90a643..28610c8ea 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/StepConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/StepConfig.java @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.config; +import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.retry.RetryStrategies; import software.amazon.lambda.durable.retry.RetryStrategy; import software.amazon.lambda.durable.serde.SerDes; @@ -42,6 +43,15 @@ public Builder toBuilder() { return new Builder(retryStrategy, semanticsPerRetry, serDes); } + /** Converts this compatibility config to the operation-owned config. */ + public DurableStepOperation.StepConfig toOperationConfig() { + return DurableStepOperation.StepConfig.builder() + .retryStrategy(retryStrategy()) + .semanticsPerRetry(semanticsPerRetry()) + .serDes(serDes()) + .build(); + } + /** * Creates a new builder for StepConfig. * diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/WaitForCallbackConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/WaitForCallbackConfig.java index e3bd57f44..0cd072112 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/WaitForCallbackConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/WaitForCallbackConfig.java @@ -2,6 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.config; +import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation; + /** * Configuration for the {@code waitForCallback} composite operation. * @@ -38,6 +40,14 @@ public Builder toBuilder() { return new Builder().stepConfig(this.stepConfig).callbackConfig(this.callbackConfig); } + /** Converts this compatibility config to the operation-owned config. */ + public DurableWaitForCallbackOperation.WaitForCallbackConfig toOperationConfig() { + return DurableWaitForCallbackOperation.WaitForCallbackConfig.builder() + .stepConfig(stepConfig().toOperationConfig()) + .callbackConfig(callbackConfig().toOperationConfig()) + .build(); + } + /** Builder for {@link WaitForCallbackConfig}. */ public static class Builder { private StepConfig stepConfig; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/WaitForConditionConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/WaitForConditionConfig.java index 1561199f9..9da044629 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/WaitForConditionConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/WaitForConditionConfig.java @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.config; +import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation; import software.amazon.lambda.durable.retry.WaitForConditionWaitStrategy; import software.amazon.lambda.durable.retry.WaitStrategies; import software.amazon.lambda.durable.serde.SerDes; @@ -56,6 +57,15 @@ public Builder toBuilder() { return b; } + /** Converts this compatibility config to the operation-owned config. */ + public DurableWaitForConditionOperation.WaitForConditionConfig toOperationConfig() { + return DurableWaitForConditionOperation.WaitForConditionConfig.builder() + .waitStrategy(waitStrategy()) + .serDes(serDes()) + .initialState(initialState()) + .build(); + } + /** * Creates a new builder for {@code WaitForConditionConfig}. All fields are optional. * diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/WithRetryConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/WithRetryConfig.java index 23ff43830..f3dcca334 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/WithRetryConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/WithRetryConfig.java @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.config; +import software.amazon.lambda.durable.operation.DurableWithRetryOperation; import software.amazon.lambda.durable.retry.RetryStrategies; import software.amazon.lambda.durable.retry.RetryStrategy; @@ -54,6 +55,14 @@ public static Builder builder() { return new Builder(); } + /** Converts this compatibility config to the operation-owned config. */ + public DurableWithRetryOperation.WithRetryConfig toOperationConfig() { + return DurableWithRetryOperation.WithRetryConfig.builder() + .retryStrategy(retryStrategy()) + .wrapInChildContext(wrapInChildContext()) + .build(); + } + /** Builder for creating {@link WithRetryConfig} instances. */ public static class Builder { private RetryStrategy retryStrategy; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java index ac5de46c3..28a92ae3b 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java @@ -3,53 +3,34 @@ package software.amazon.lambda.durable.context; import com.amazonaws.services.lambda.runtime.Context; -import java.time.Duration; import java.util.Collection; -import java.util.Objects; import java.util.function.BiConsumer; import java.util.function.BiFunction; -import java.util.function.Function; -import software.amazon.awssdk.services.lambda.model.OperationType; -import software.amazon.lambda.durable.DurableCallbackFuture; import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.ParallelDurableFuture; import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.CallbackConfig; -import software.amazon.lambda.durable.config.InvokeConfig; import software.amazon.lambda.durable.config.MapConfig; import software.amazon.lambda.durable.config.ParallelConfig; -import software.amazon.lambda.durable.config.RunInChildContextConfig; -import software.amazon.lambda.durable.config.StepConfig; import software.amazon.lambda.durable.config.WaitForCallbackConfig; import software.amazon.lambda.durable.config.WaitForConditionConfig; import software.amazon.lambda.durable.config.WithRetryConfig; -import software.amazon.lambda.durable.context.extension.MapExtension; -import software.amazon.lambda.durable.context.extension.ParallelExtension; -import software.amazon.lambda.durable.context.extension.WaitForCallbackExtension; -import software.amazon.lambda.durable.context.extension.WaitForConditionExtension; -import software.amazon.lambda.durable.context.extension.WithRetryExtension; import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.execution.OperationIdGenerator; import software.amazon.lambda.durable.execution.ThreadType; import software.amazon.lambda.durable.extension.ExtensionContext; -import software.amazon.lambda.durable.extension.ExtensionContextConfig; -import software.amazon.lambda.durable.extension.ExtensionContextFunction; import software.amazon.lambda.durable.extension.ExtensionOperation; -import software.amazon.lambda.durable.extension.ExtensionStepConfig; -import software.amazon.lambda.durable.extension.ExtensionStepFunction; +import software.amazon.lambda.durable.extension.ExtensionOperationImpl; import software.amazon.lambda.durable.model.MapResult; -import software.amazon.lambda.durable.model.OperationDescriptor; -import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.model.WaitForConditionResult; -import software.amazon.lambda.durable.operation.BaseDurableOperation; -import software.amazon.lambda.durable.operation.CallbackOperation; -import software.amazon.lambda.durable.operation.ChildContextOperation; -import software.amazon.lambda.durable.operation.InvokeOperation; -import software.amazon.lambda.durable.operation.StepOperation; -import software.amazon.lambda.durable.operation.WaitOperation; +import software.amazon.lambda.durable.operation.DurableMapOperation; +import software.amazon.lambda.durable.operation.DurableParallelOperation; +import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation; +import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation; +import software.amazon.lambda.durable.operation.DurableWithRetryOperation; +import software.amazon.lambda.durable.primitive.BasePrimitive; import software.amazon.lambda.durable.util.ParameterValidator; /** @@ -65,7 +46,7 @@ public class DurableContextImpl extends BaseContextImpl implements DurableContex - Math.max(WAIT_FOR_CALLBACK_CALLBACK_SUFFIX.length(), WAIT_FOR_CALLBACK_SUBMITTER_SUFFIX.length()); private final OperationIdGenerator operationIdGenerator; private final DurableContextImpl parentContext; - private final BaseDurableOperation lateCheckpointOwner; + private final BasePrimitive lateCheckpointOwner; private final boolean isVirtual; private boolean isReplaying; @@ -78,7 +59,7 @@ private DurableContextImpl( String contextName, boolean isVirtual, DurableContextImpl parentContext, - BaseDurableOperation lateCheckpointOwner) { + BasePrimitive lateCheckpointOwner) { super(executionManager, durableConfig, lambdaContext, contextId, contextName, ThreadType.CONTEXT); operationIdGenerator = new OperationIdGenerator(contextId); this.parentContext = parentContext; @@ -115,10 +96,7 @@ public DurableContextImpl createChildContext(String childContextId, String child } public DurableContextImpl createChildContext( - String childContextId, - String childContextName, - boolean isVirtual, - BaseDurableOperation lateCheckpointOwner) { + String childContextId, String childContextName, boolean isVirtual, BasePrimitive lateCheckpointOwner) { return new DurableContextImpl( getExecutionManager(), getDurableConfig(), @@ -148,308 +126,19 @@ public StepContextImpl createStepContext(String stepOperationId, String stepOper attempt); } - @Override - public DurableFuture stepAsync( - String name, TypeToken resultType, Function func, StepConfig config) { - Objects.requireNonNull(config, "config cannot be null"); - Objects.requireNonNull(resultType, "resultType cannot be null"); - ParameterValidator.validateOperationName(name); - return stepAsyncWithId(nextOperationId(), name, resultType, func, config); - } - - DurableFuture stepAsyncWithId( - String operationId, - String name, - TypeToken resultType, - Function func, - StepConfig config) { - return stepAsyncWithId(operationId, name, OperationSubType.STEP.getValue(), resultType, func, config); - } - - DurableFuture stepAsyncWithId( - String operationId, - String name, - String subType, - TypeToken resultType, - Function func, - StepConfig config) { - Objects.requireNonNull(config, "config cannot be null"); - Objects.requireNonNull(resultType, "resultType cannot be null"); - ParameterValidator.validateOperationName(name); - - if (config.serDes() == null) { - config = config.toBuilder().serDes(getDurableConfig().getSerDes()).build(); - } - // Create and start step operation with TypeToken - var operation = new StepOperation<>( - new OperationDescriptor(operationId, name, OperationType.STEP, subType), - func, - resultType, - config, - this); - - operation.execute(); // Start the step (returns immediately) - - return operation; - } - - DurableFuture extensionStepAsyncWithId( - String operationId, - String name, - String subType, - TypeToken resultType, - ExtensionStepFunction function, - ExtensionStepConfig config) { - Objects.requireNonNull(resultType, "resultType cannot be null"); - Objects.requireNonNull(function, "function cannot be null"); - Objects.requireNonNull(config, "config cannot be null"); - ParameterValidator.validateOperationName(name); - - if (config.serDes() == null) { - config = ExtensionStepConfig.builder() - .initialState(config.initialState()) - .serDes(getDurableConfig().getSerDes()) - .build(); - } - var operation = new StepOperation<>( - new OperationDescriptor(operationId, name, OperationType.STEP, subType), - function, - resultType, - config, - this); - operation.execute(); - return operation; - } - - @Override - public DurableFuture waitAsync(String name, Duration duration) { - ParameterValidator.validateDuration(duration, "Wait duration"); - ParameterValidator.validateOperationName(name); - return waitAsyncWithId(nextOperationId(), name, duration); - } - - DurableFuture waitAsyncWithId(String operationId, String name, Duration duration) { - return waitAsyncWithId(operationId, name, OperationSubType.WAIT.getValue(), duration); - } - - DurableFuture waitAsyncWithId(String operationId, String name, String subType, Duration duration) { - ParameterValidator.validateDuration(duration, "Wait duration"); - ParameterValidator.validateOperationName(name); - - // Create and start wait operation - var operation = new WaitOperation( - new OperationDescriptor(operationId, name, OperationType.WAIT, subType), duration, this); - - operation.execute(); // Checkpoint the wait - return operation; - } - - @Override - public DurableFuture invokeAsync( - String name, String functionName, U payload, TypeToken resultType, InvokeConfig config) { - Objects.requireNonNull(config, "config cannot be null"); - Objects.requireNonNull(resultType, "resultType cannot be null"); - ParameterValidator.validateOperationName(name); - return invokeAsyncWithId(nextOperationId(), name, functionName, payload, resultType, config); - } - - DurableFuture invokeAsyncWithId( - String operationId, - String name, - String functionName, - U payload, - TypeToken resultType, - InvokeConfig config) { - return invokeAsyncWithId( - operationId, - name, - OperationSubType.CHAINED_INVOKE.getValue(), - functionName, - payload, - resultType, - config); - } - - DurableFuture invokeAsyncWithId( - String operationId, - String name, - String subType, - String functionName, - U payload, - TypeToken resultType, - InvokeConfig config) { - Objects.requireNonNull(config, "config cannot be null"); - Objects.requireNonNull(resultType, "resultType cannot be null"); - ParameterValidator.validateOperationName(name); - - if (config.serDes() == null) { - config = config.toBuilder().serDes(getDurableConfig().getSerDes()).build(); - } - if (config.payloadSerDes() == null) { - config = config.toBuilder() - .payloadSerDes(getDurableConfig().getSerDes()) - .build(); - } - // Create and start invoke operation - var operation = new InvokeOperation<>( - new OperationDescriptor(operationId, name, OperationType.CHAINED_INVOKE, subType), - functionName, - payload, - resultType, - config, - this); - - operation.execute(); // checkpoint the invoke operation - return operation; // Block (will throw SuspendExecutionException if needed) - } - - @Override - public DurableCallbackFuture createCallback(String name, TypeToken resultType, CallbackConfig config) { - Objects.requireNonNull(config, "config cannot be null"); - Objects.requireNonNull(resultType, "resultType cannot be null"); - ParameterValidator.validateOperationName(name); - return createCallbackWithId(nextOperationId(), name, resultType, config); - } - - DurableCallbackFuture createCallbackWithId( - String operationId, String name, TypeToken resultType, CallbackConfig config) { - return createCallbackWithId(operationId, name, OperationSubType.CALLBACK.getValue(), resultType, config); - } - - DurableCallbackFuture createCallbackWithId( - String operationId, String name, String subType, TypeToken resultType, CallbackConfig config) { - Objects.requireNonNull(config, "config cannot be null"); - Objects.requireNonNull(resultType, "resultType cannot be null"); - ParameterValidator.validateOperationName(name); - if (config.serDes() == null) { - config = config.toBuilder().serDes(getDurableConfig().getSerDes()).build(); - } - var operation = new CallbackOperation<>( - new OperationDescriptor(operationId, name, OperationType.CALLBACK, subType), resultType, config, this); - operation.execute(); - - return operation; - } - - /** - * Runs a function in a child context, blocking until it completes. - * - *

Child contexts provide isolated operation ID namespaces, allowing nested workflows to be composed without ID - * collisions. On replay, the child context's operations are replayed independently. - * - * @param name the operation name within this context - * @param resultType the result class for deserialization - * @param func the function to execute, receiving a child {@link DurableContext} - * @param config the configuration for the child context - * @return the DurableFuture wrapping the child context result - */ - @Override - public DurableFuture runInChildContextAsync( - String name, TypeToken resultType, Function func, RunInChildContextConfig config) { - return runInChildContextAsync(name, resultType, func, config, OperationSubType.RUN_IN_CHILD_CONTEXT); - } - - private DurableFuture runInChildContextAsync( - String name, - TypeToken resultType, - Function func, - RunInChildContextConfig config, - OperationSubType subType) { - Objects.requireNonNull(resultType, "resultType cannot be null"); - Objects.requireNonNull(func, "func cannot be null"); - Objects.requireNonNull(config, "RunInChildContextConfig cannot be null"); - ParameterValidator.validateOperationName(name); - return runInChildContextAsyncWithId(nextOperationId(), name, resultType, func, config, subType); - } - - DurableFuture runInChildContextAsyncWithId( - String operationId, - String name, - TypeToken resultType, - Function func, - RunInChildContextConfig config) { - return runInChildContextAsyncWithId( - operationId, name, OperationSubType.RUN_IN_CHILD_CONTEXT.getValue(), resultType, func, config); - } - - private DurableFuture runInChildContextAsyncWithId( - String operationId, - String name, - TypeToken resultType, - Function func, - RunInChildContextConfig config, - OperationSubType subType) { - return runInChildContextAsyncWithId(operationId, name, subType.getValue(), resultType, func, config); - } - - DurableFuture runInChildContextAsyncWithId( - String operationId, - String name, - String subType, - TypeToken resultType, - Function func, - RunInChildContextConfig config) { - Objects.requireNonNull(resultType, "resultType cannot be null"); - Objects.requireNonNull(func, "func cannot be null"); - Objects.requireNonNull(config, "RunInChildContextConfig cannot be null"); - ParameterValidator.validateOperationName(name); - - if (config.serDes() == null) { - config = config.toBuilder().serDes(getDurableConfig().getSerDes()).build(); - } - - var operation = new ChildContextOperation<>( - new OperationDescriptor(operationId, name, OperationType.CONTEXT, subType), - func, - resultType, - config, - this, - lateCheckpointOwner); - - operation.execute(); - return operation; - } - - DurableFuture extensionContextAsyncWithId( - String operationId, - String name, - String subType, - TypeToken resultType, - ExtensionContextFunction function, - ExtensionContextConfig config) { - Objects.requireNonNull(resultType, "resultType cannot be null"); - Objects.requireNonNull(function, "function cannot be null"); - Objects.requireNonNull(config, "config cannot be null"); - ParameterValidator.validateOperationName(name); - - var childConfig = config.childContextConfig(); - if (childConfig.serDes() == null) { - childConfig = childConfig.toBuilder() - .serDes(getDurableConfig().getSerDes()) - .build(); - config = config.toBuilder().childContextConfig(childConfig).build(); - } - - var operation = new ChildContextOperation<>( - new OperationDescriptor(operationId, name, OperationType.CONTEXT, subType), - function, - resultType, - config, - this, - lateCheckpointOwner); - operation.execute(); - return operation; + BasePrimitive getLateCheckpointOwner() { + return lateCheckpointOwner; } @Override public DurableFuture> mapAsync( String name, Collection items, TypeToken resultType, MapFunction function, MapConfig config) { - return MapExtension.execute(this, name, items, resultType, function, config); + return DurableMapOperation.mapAsync(this, name, items, resultType, function, config.toOperationConfig()); } @Override public ParallelDurableFuture parallel(String name, ParallelConfig config) { - return ParallelExtension.execute(this, name, config); + return DurableParallelOperation.parallel(this, name, config.toOperationConfig()); } @Override @@ -458,7 +147,8 @@ public DurableFuture waitForCallbackAsync( TypeToken resultType, BiConsumer func, WaitForCallbackConfig waitForCallbackConfig) { - return WaitForCallbackExtension.execute(this, name, resultType, func, waitForCallbackConfig); + return DurableWaitForCallbackOperation.waitForCallbackAsync( + this, name, resultType, func, waitForCallbackConfig.toOperationConfig()); } @Override @@ -467,7 +157,8 @@ public DurableFuture waitForConditionAsync( TypeToken resultType, BiFunction> checkFunc, WaitForConditionConfig config) { - return WaitForConditionExtension.execute(this, name, resultType, checkFunc, config); + return DurableWaitForConditionOperation.waitForConditionAsync( + this, name, resultType, checkFunc, config.toOperationConfig()); } // =============== withRetry ================ @@ -475,7 +166,7 @@ public DurableFuture waitForConditionAsync( @Override public DurableFuture withRetryAsync( String name, BiFunction operation, WithRetryConfig config) { - return WithRetryExtension.execute(this, name, operation, config); + return DurableWithRetryOperation.withRetryAsync(this, name, operation, config.toOperationConfig()); } // =============== accessors ================ @@ -504,13 +195,13 @@ String reserveOperationId(String localOperationId) { @Override public ExtensionOperation reserve(String name) { ParameterValidator.validateOperationName(name); - return new ExtensionOperationImpl(this, reserveOperationId(), name); + return new ExtensionOperationImpl(this, reserveOperationId(), name, lateCheckpointOwner); } @Override public ExtensionOperation reserve(String name, String localOperationId) { ParameterValidator.validateOperationName(name); - return new ExtensionOperationImpl(this, reserveOperationId(localOperationId), name); + return new ExtensionOperationImpl(this, reserveOperationId(localOperationId), name, lateCheckpointOwner); } /** Returns whether this context is currently in replay mode. */ diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java deleted file mode 100644 index 2928a841b..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/ExtensionOperationImpl.java +++ /dev/null @@ -1,141 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.context; - -import java.time.Duration; -import java.util.Objects; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.function.Supplier; -import software.amazon.lambda.durable.DurableCallbackFuture; -import software.amazon.lambda.durable.DurableFuture; -import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.CallbackConfig; -import software.amazon.lambda.durable.config.InvokeConfig; -import software.amazon.lambda.durable.config.RunInChildContextConfig; -import software.amazon.lambda.durable.config.StepConfig; -import software.amazon.lambda.durable.extension.ExtensionContextConfig; -import software.amazon.lambda.durable.extension.ExtensionContextFunction; -import software.amazon.lambda.durable.extension.ExtensionOperation; -import software.amazon.lambda.durable.extension.ExtensionStepConfig; -import software.amazon.lambda.durable.extension.ExtensionStepFunction; - -final class ExtensionOperationImpl implements ExtensionOperation { - private final DurableContextImpl context; - private final String operationId; - private final String name; - private final AtomicBoolean claimed = new AtomicBoolean(); - - ExtensionOperationImpl(DurableContextImpl context, String operationId, String name) { - this.context = context; - this.operationId = operationId; - this.name = name; - } - - @Override - public DurableFuture stepAsync(TypeToken resultType, Supplier function, StepConfig config) { - claim(); - return context.stepAsyncWithId(operationId, name, resultType, ignored -> function.get(), config); - } - - @Override - public DurableFuture stepAsync( - String subType, TypeToken resultType, Supplier function, StepConfig config) { - validateSubType(subType); - claim(); - return context.stepAsyncWithId(operationId, name, subType, resultType, ignored -> function.get(), config); - } - - @Override - public DurableFuture stepAsync( - String subType, TypeToken resultType, ExtensionStepFunction function, ExtensionStepConfig config) { - validateSubType(subType); - Objects.requireNonNull(function, "function cannot be null"); - Objects.requireNonNull(config, "config cannot be null"); - claim(); - return context.extensionStepAsyncWithId(operationId, name, subType, resultType, function, config); - } - - @Override - public DurableFuture waitAsync(Duration duration) { - claim(); - return context.waitAsyncWithId(operationId, name, duration); - } - - @Override - public DurableFuture waitAsync(String subType, Duration duration) { - validateSubType(subType); - claim(); - return context.waitAsyncWithId(operationId, name, subType, duration); - } - - @Override - public DurableFuture invokeAsync( - String functionName, U payload, TypeToken resultType, InvokeConfig config) { - claim(); - return context.invokeAsyncWithId(operationId, name, functionName, payload, resultType, config); - } - - @Override - public DurableFuture invokeAsync( - String subType, String functionName, U payload, TypeToken resultType, InvokeConfig config) { - validateSubType(subType); - claim(); - return context.invokeAsyncWithId(operationId, name, subType, functionName, payload, resultType, config); - } - - @Override - public DurableCallbackFuture createCallback(TypeToken resultType, CallbackConfig config) { - claim(); - return context.createCallbackWithId(operationId, name, resultType, config); - } - - @Override - public DurableCallbackFuture createCallback(String subType, TypeToken resultType, CallbackConfig config) { - validateSubType(subType); - claim(); - return context.createCallbackWithId(operationId, name, subType, resultType, config); - } - - @Override - public DurableFuture runInChildContextAsync( - TypeToken resultType, Supplier function, RunInChildContextConfig config) { - claim(); - return context.runInChildContextAsyncWithId(operationId, name, resultType, ignored -> function.get(), config); - } - - @Override - public DurableFuture runInChildContextAsync( - String subType, TypeToken resultType, Supplier function, RunInChildContextConfig config) { - validateSubType(subType); - claim(); - return context.runInChildContextAsyncWithId( - operationId, name, subType, resultType, ignored -> function.get(), config); - } - - @Override - public DurableFuture runInChildContextAsync( - String subType, - TypeToken resultType, - ExtensionContextFunction function, - ExtensionContextConfig config) { - validateSubType(subType); - Objects.requireNonNull(resultType, "resultType cannot be null"); - Objects.requireNonNull(function, "function cannot be null"); - Objects.requireNonNull(config, "config cannot be null"); - claim(); - return context.extensionContextAsyncWithId(operationId, name, subType, resultType, function, config); - } - - private void validateSubType(String subType) { - Objects.requireNonNull(subType, "subType cannot be null"); - if (subType.isBlank()) { - throw new IllegalArgumentException("subType cannot be blank"); - } - } - - private void claim() { - if (!claimed.compareAndSet(false, true)) { - throw new IllegalStateException("An extension operation reservation can only be used once"); - } - } -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/ParallelExtension.java b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/ParallelExtension.java deleted file mode 100644 index bbf62779f..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/ParallelExtension.java +++ /dev/null @@ -1,21 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.context.extension; - -import java.util.Objects; -import software.amazon.lambda.durable.ParallelDurableFuture; -import software.amazon.lambda.durable.config.ParallelConfig; -import software.amazon.lambda.durable.extension.ExtensionContext; -import software.amazon.lambda.durable.util.ParameterValidator; - -/** Canonical implementation of the built-in parallel extension. */ -public final class ParallelExtension { - private ParallelExtension() {} - - public static ParallelDurableFuture execute(ExtensionContext context, String name, ParallelConfig config) { - Objects.requireNonNull(context, "context cannot be null"); - Objects.requireNonNull(config, "config cannot be null"); - ParameterValidator.validateOperationName(name); - return new ParallelExtensionFuture(context, name, config); - } -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForCallbackExtension.java b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForCallbackExtension.java deleted file mode 100644 index c8253c46d..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForCallbackExtension.java +++ /dev/null @@ -1,117 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.context.extension; - -import static software.amazon.lambda.durable.execution.ExecutionManager.isTerminalStatus; - -import java.util.Objects; -import java.util.function.BiConsumer; -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.DurableFuture; -import software.amazon.lambda.durable.StepContext; -import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.RunInChildContextConfig; -import software.amazon.lambda.durable.config.WaitForCallbackConfig; -import software.amazon.lambda.durable.exception.CallbackFailedException; -import software.amazon.lambda.durable.exception.CallbackSubmitterException; -import software.amazon.lambda.durable.exception.CallbackTimeoutException; -import software.amazon.lambda.durable.exception.StepFailedException; -import software.amazon.lambda.durable.exception.StepInterruptedException; -import software.amazon.lambda.durable.extension.ExtensionContext; -import software.amazon.lambda.durable.extension.ExtensionContextConfig; -import software.amazon.lambda.durable.extension.ExtensionContextFailure; -import software.amazon.lambda.durable.extension.ExtensionContextResult; -import software.amazon.lambda.durable.model.OperationSubType; -import software.amazon.lambda.durable.util.ParameterValidator; - -/** Canonical implementation of the built-in wait-for-callback extension. */ -public final class WaitForCallbackExtension { - private static final String CALLBACK_SUFFIX = "-callback"; - private static final String SUBMITTER_SUFFIX = "-submitter"; - private static final int MAX_NAME_LENGTH = ParameterValidator.MAX_OPERATION_NAME_LENGTH - - Math.max(CALLBACK_SUFFIX.length(), SUBMITTER_SUFFIX.length()); - - private WaitForCallbackExtension() {} - - public static DurableFuture execute( - ExtensionContext context, - String name, - TypeToken resultType, - BiConsumer submitter, - WaitForCallbackConfig config) { - Objects.requireNonNull(context, "context cannot be null"); - Objects.requireNonNull(resultType, "resultType cannot be null"); - Objects.requireNonNull(submitter, "submitter cannot be null"); - Objects.requireNonNull(config, "config cannot be null"); - ParameterValidator.validateOperationName(name, MAX_NAME_LENGTH); - - var parent = context.reserve(name); - return parent.runInChildContextAsync( - OperationSubType.WAIT_FOR_CALLBACK.getValue(), - resultType, - () -> executeInChildContext(name, resultType, submitter, config), - extensionConfig(config)); - } - - private static ExtensionContextResult executeInChildContext( - String name, - TypeToken resultType, - BiConsumer submitter, - WaitForCallbackConfig config) { - var child = ExtensionContext.getCurrentContext(); - var callback = child.reserve(name + CALLBACK_SUFFIX).createCallback(resultType, config.callbackConfig()); - child.reserve(name + SUBMITTER_SUFFIX) - .step( - Void.class, - () -> { - submitter.accept(callback.callbackId(), StepContext.getCurrentContext()); - return null; - }, - config.stepConfig()); - return ExtensionContextResult.completed(callback.get()); - } - - private static ExtensionContextConfig extensionConfig(WaitForCallbackConfig config) { - return ExtensionContextConfig.builder() - .childContextConfig(RunInChildContextConfig.builder() - .serDes(config.stepConfig().serDes()) - .build()) - .errorHandler(WaitForCallbackExtension::translateFailure) - .build(); - } - - private static Throwable translateFailure(ExtensionContextFailure failure) { - var callback = findChild(failure, OperationType.CALLBACK); - var submitter = findChild(failure, OperationType.STEP); - if (callback != null && isTerminalStatus(callback.status())) { - if (callback.status() == OperationStatus.FAILED) { - return new CallbackFailedException(callback); - } - if (callback.status() == OperationStatus.TIMED_OUT) { - return new CallbackTimeoutException(callback); - } - } - if (callback != null - && submitter != null - && isTerminalStatus(submitter.status()) - && submitter.status() != OperationStatus.SUCCEEDED) { - var error = submitter.stepDetails().error(); - var cause = StepInterruptedException.isStepInterruptedException(error) - ? new StepInterruptedException(submitter) - : new StepFailedException(submitter); - return new CallbackSubmitterException(callback, cause); - } - return new IllegalStateException("Unknown waitForCallback status"); - } - - private static Operation findChild(ExtensionContextFailure failure, OperationType type) { - return failure.childOperations().stream() - .filter(summary -> summary.operationType() == type) - .map(summary -> summary.operation()) - .filter(Objects::nonNull) - .findFirst() - .orElse(null); - } -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForConditionExtension.java b/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForConditionExtension.java deleted file mode 100644 index 5e8459794..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForConditionExtension.java +++ /dev/null @@ -1,59 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.context.extension; - -import java.util.Objects; -import java.util.function.BiFunction; -import software.amazon.lambda.durable.DurableFuture; -import software.amazon.lambda.durable.StepContext; -import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.WaitForConditionConfig; -import software.amazon.lambda.durable.extension.ExtensionContext; -import software.amazon.lambda.durable.extension.ExtensionStepConfig; -import software.amazon.lambda.durable.extension.ExtensionStepResult; -import software.amazon.lambda.durable.model.OperationSubType; -import software.amazon.lambda.durable.model.WaitForConditionResult; -import software.amazon.lambda.durable.util.ParameterValidator; - -/** Canonical implementation of the built-in wait-for-condition extension. */ -public final class WaitForConditionExtension { - private WaitForConditionExtension() {} - - public static DurableFuture execute( - ExtensionContext context, - String name, - TypeToken resultType, - BiFunction> checkFunction, - WaitForConditionConfig config) { - Objects.requireNonNull(context, "context cannot be null"); - Objects.requireNonNull(resultType, "resultType cannot be null"); - Objects.requireNonNull(checkFunction, "checkFunction cannot be null"); - Objects.requireNonNull(config, "config cannot be null"); - ParameterValidator.validateOperationName(name); - - var future = context.reserve(name) - .stepAsync( - OperationSubType.WAIT_FOR_CONDITION.getValue(), - resultType, - state -> evaluate(state, checkFunction, config), - ExtensionStepConfig.builder() - .initialState(config.initialState()) - .serDes(config.serDes()) - .build()); - return new WaitForConditionFuture<>(future); - } - - private static ExtensionStepResult evaluate( - T state, - BiFunction> checkFunction, - WaitForConditionConfig config) { - var stepContext = StepContext.getCurrentContext(); - var result = Objects.requireNonNull( - checkFunction.apply(state, stepContext), "waitForCondition check result cannot be null"); - if (result.isDone()) { - return ExtensionStepResult.succeed(result.value()); - } - var delay = config.waitStrategy().evaluate(result.value(), stepContext.getAttempt()); - return ExtensionStepResult.retry(result.value(), delay); - } -} 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 788c7a9ae..49caa5489 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 @@ -27,8 +27,8 @@ import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.model.DurableExecutionInput; import software.amazon.lambda.durable.model.SafeCloseable; -import software.amazon.lambda.durable.operation.BaseDurableOperation; import software.amazon.lambda.durable.plugin.PluginInfoConverter; +import software.amazon.lambda.durable.primitive.BasePrimitive; import software.amazon.lambda.durable.util.ExceptionHelper; /** @@ -66,7 +66,7 @@ public class ExecutionManager implements SafeCloseable { private final Set updatedOperationIdsSinceLastInvocation; // ===== Thread Coordination ===== - private final Map registeredOperations = new ConcurrentHashMap<>(); + private final Map registeredOperations = new ConcurrentHashMap<>(); private final Set activeThreads = Collections.synchronizedSet(new HashSet<>()); private static final ThreadLocal currentThreadContext = new ThreadLocal<>(); private final CompletableFuture executionExceptionFuture = new CompletableFuture<>(); @@ -134,7 +134,7 @@ public boolean isOperationUpdatedSinceLastInvocation(String operationId) { } /** Registers an operation so it can receive checkpoint completion notifications. */ - public void registerOperation(BaseDurableOperation operation) { + public void registerOperation(BasePrimitive operation) { registeredOperations.put(operation.getOperationId(), operation); } @@ -367,7 +367,7 @@ public void close() { private void validateRunningThreads() { // This will detect stuck user thread and thread leaks in the thread pool - for (BaseDurableOperation op : registeredOperations.values()) { + for (BasePrimitive op : registeredOperations.values()) { var userHandlerFuture = op.getRunningUserHandler(); if (userHandlerFuture != null && !userHandlerFuture.isDone()) { // Some user threads can still be running because diff --git a/sdk/src/main/java/software/amazon/lambda/durable/execution/OperationIdGenerator.java b/sdk/src/main/java/software/amazon/lambda/durable/execution/OperationIdGenerator.java index e24052396..d21004b8d 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/execution/OperationIdGenerator.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/execution/OperationIdGenerator.java @@ -46,10 +46,10 @@ public static String hashOperationId(String rawId) { * {@code hash("-2")} inside a child context. */ public String nextOperationId() { - String localOperationId; - do { - localOperationId = String.valueOf(operationCounter.incrementAndGet()); - } while (!allocatedLocalIds.add(localOperationId)); + var localOperationId = String.valueOf(operationCounter.incrementAndGet()); + if (!allocatedLocalIds.add(localOperationId)) { + throw new IllegalArgumentException("Local operation ID is already in use: " + localOperationId); + } return hashOperationId(operationIdPrefix + localOperationId); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionCallbackConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionCallbackConfig.java new file mode 100644 index 000000000..543648496 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionCallbackConfig.java @@ -0,0 +1,70 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.extension; + +import java.time.Duration; +import software.amazon.lambda.durable.serde.SerDes; +import software.amazon.lambda.durable.util.ParameterValidator; + +/** Configuration for an extension CALLBACK primitive. */ +public final class ExtensionCallbackConfig { + private final Duration timeout; + private final Duration heartbeatTimeout; + private final SerDes serDes; + + private ExtensionCallbackConfig(Builder builder) { + timeout = builder.timeout; + heartbeatTimeout = builder.heartbeatTimeout; + serDes = builder.serDes; + } + + public Duration timeout() { + return timeout; + } + + public Duration heartbeatTimeout() { + return heartbeatTimeout; + } + + public SerDes serDes() { + return serDes; + } + + public Builder toBuilder() { + return new Builder().timeout(timeout).heartbeatTimeout(heartbeatTimeout).serDes(serDes); + } + + public static Builder builder() { + return new Builder(); + } + + /** Builder for {@link ExtensionCallbackConfig}. */ + public static final class Builder { + private Duration timeout; + private Duration heartbeatTimeout; + private SerDes serDes; + + private Builder() {} + + public Builder timeout(Duration timeout) { + ParameterValidator.validateOptionalDuration(timeout, "Callback timeout"); + this.timeout = timeout; + return this; + } + + public Builder heartbeatTimeout(Duration heartbeatTimeout) { + ParameterValidator.validateOptionalDuration(heartbeatTimeout, "Heartbeat timeout"); + this.heartbeatTimeout = heartbeatTimeout; + return this; + } + + public Builder serDes(SerDes serDes) { + this.serDes = serDes; + return this; + } + + public ExtensionCallbackConfig build() { + return new ExtensionCallbackConfig(this); + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionInvokeConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionInvokeConfig.java new file mode 100644 index 000000000..1628f14d3 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionInvokeConfig.java @@ -0,0 +1,66 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.extension; + +import software.amazon.lambda.durable.serde.SerDes; + +/** Configuration for an extension CHAINED_INVOKE primitive. */ +public final class ExtensionInvokeConfig { + private final SerDes payloadSerDes; + private final SerDes resultSerDes; + private final String tenantId; + + private ExtensionInvokeConfig(Builder builder) { + payloadSerDes = builder.payloadSerDes; + resultSerDes = builder.resultSerDes; + tenantId = builder.tenantId; + } + + public SerDes payloadSerDes() { + return payloadSerDes; + } + + public SerDes serDes() { + return resultSerDes; + } + + public String tenantId() { + return tenantId; + } + + public Builder toBuilder() { + return new Builder().payloadSerDes(payloadSerDes).serDes(resultSerDes).tenantId(tenantId); + } + + public static Builder builder() { + return new Builder(); + } + + /** Builder for {@link ExtensionInvokeConfig}. */ + public static final class Builder { + private SerDes payloadSerDes; + private SerDes resultSerDes; + private String tenantId; + + private Builder() {} + + public Builder payloadSerDes(SerDes payloadSerDes) { + this.payloadSerDes = payloadSerDes; + return this; + } + + public Builder serDes(SerDes resultSerDes) { + this.resultSerDes = resultSerDes; + return this; + } + + public Builder tenantId(String tenantId) { + this.tenantId = tenantId; + return this; + } + + public ExtensionInvokeConfig build() { + return new ExtensionInvokeConfig(this); + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionOperation.java index 6a205b710..4ac0fc833 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionOperation.java @@ -3,14 +3,9 @@ package software.amazon.lambda.durable.extension; import java.time.Duration; -import java.util.function.Supplier; import software.amazon.lambda.durable.DurableCallbackFuture; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.CallbackConfig; -import software.amazon.lambda.durable.config.InvokeConfig; -import software.amazon.lambda.durable.config.RunInChildContextConfig; -import software.amazon.lambda.durable.config.StepConfig; /** * An opaque, one-shot reservation for a primitive operation. @@ -19,288 +14,16 @@ * allows an extension to launch them later in a different order without changing their IDs. */ public interface ExtensionOperation { - default T step(Class resultType, Supplier function) { - return step(TypeToken.get(resultType), function); - } - - default T step(TypeToken resultType, Supplier function) { - return step(resultType, function, StepConfig.builder().build()); - } - - default T step(Class resultType, Supplier function, StepConfig config) { - return step(TypeToken.get(resultType), function, config); - } - - default T step(TypeToken resultType, Supplier function, StepConfig config) { - return stepAsync(resultType, function, config).get(); - } - - default DurableFuture stepAsync(Class resultType, Supplier function) { - return stepAsync(TypeToken.get(resultType), function); - } - - default DurableFuture stepAsync(TypeToken resultType, Supplier function) { - return stepAsync(resultType, function, StepConfig.builder().build()); - } - - default DurableFuture stepAsync(Class resultType, Supplier function, StepConfig config) { - return stepAsync(TypeToken.get(resultType), function, config); - } - - DurableFuture stepAsync(TypeToken resultType, Supplier function, StepConfig config); - - default T step(String subType, Class resultType, Supplier function) { - return step(subType, TypeToken.get(resultType), function); - } - - default T step(String subType, TypeToken resultType, Supplier function) { - return step(subType, resultType, function, StepConfig.builder().build()); - } - - default T step(String subType, Class resultType, Supplier function, StepConfig config) { - return step(subType, TypeToken.get(resultType), function, config); - } - - default T step(String subType, TypeToken resultType, Supplier function, StepConfig config) { - return stepAsync(subType, resultType, function, config).get(); - } - - default DurableFuture stepAsync(String subType, Class resultType, Supplier function) { - return stepAsync(subType, TypeToken.get(resultType), function); - } - - default DurableFuture stepAsync(String subType, TypeToken resultType, Supplier function) { - return stepAsync(subType, resultType, function, StepConfig.builder().build()); - } - - default DurableFuture stepAsync( - String subType, Class resultType, Supplier function, StepConfig config) { - return stepAsync(subType, TypeToken.get(resultType), function, config); - } - - DurableFuture stepAsync(String subType, TypeToken resultType, Supplier function, StepConfig config); - - default T step( - String subType, Class resultType, ExtensionStepFunction function, ExtensionStepConfig config) { - return step(subType, TypeToken.get(resultType), function, config); - } - - default T step( - String subType, TypeToken resultType, ExtensionStepFunction function, ExtensionStepConfig config) { - return stepAsync(subType, resultType, function, config).get(); - } - - default DurableFuture stepAsync( - String subType, Class resultType, ExtensionStepFunction function, ExtensionStepConfig config) { - return stepAsync(subType, TypeToken.get(resultType), function, config); - } - DurableFuture stepAsync( String subType, TypeToken resultType, ExtensionStepFunction function, ExtensionStepConfig config); - default Void wait(Duration duration) { - return waitAsync(duration).get(); - } - - DurableFuture waitAsync(Duration duration); - - default Void wait(String subType, Duration duration) { - return waitAsync(subType, duration).get(); - } - DurableFuture waitAsync(String subType, Duration duration); - default T invoke(String functionName, U payload, Class resultType) { - return invoke(functionName, payload, TypeToken.get(resultType)); - } - - default T invoke(String functionName, U payload, TypeToken resultType) { - return invoke(functionName, payload, resultType, InvokeConfig.builder().build()); - } - - default T invoke(String functionName, U payload, Class resultType, InvokeConfig config) { - return invoke(functionName, payload, TypeToken.get(resultType), config); - } - - default T invoke(String functionName, U payload, TypeToken resultType, InvokeConfig config) { - return invokeAsync(functionName, payload, resultType, config).get(); - } - - default DurableFuture invokeAsync(String functionName, U payload, Class resultType) { - return invokeAsync(functionName, payload, TypeToken.get(resultType)); - } - - default DurableFuture invokeAsync(String functionName, U payload, TypeToken resultType) { - return invokeAsync( - functionName, payload, resultType, InvokeConfig.builder().build()); - } - - default DurableFuture invokeAsync( - String functionName, U payload, Class resultType, InvokeConfig config) { - return invokeAsync(functionName, payload, TypeToken.get(resultType), config); - } - - DurableFuture invokeAsync(String functionName, U payload, TypeToken resultType, InvokeConfig config); - - default T invoke(String subType, String functionName, U payload, Class resultType) { - return invoke(subType, functionName, payload, TypeToken.get(resultType)); - } - - default T invoke(String subType, String functionName, U payload, TypeToken resultType) { - return invoke( - subType, - functionName, - payload, - resultType, - InvokeConfig.builder().build()); - } - - default T invoke(String subType, String functionName, U payload, Class resultType, InvokeConfig config) { - return invoke(subType, functionName, payload, TypeToken.get(resultType), config); - } - - default T invoke( - String subType, String functionName, U payload, TypeToken resultType, InvokeConfig config) { - return invokeAsync(subType, functionName, payload, resultType, config).get(); - } - - default DurableFuture invokeAsync(String subType, String functionName, U payload, Class resultType) { - return invokeAsync(subType, functionName, payload, TypeToken.get(resultType)); - } - - default DurableFuture invokeAsync( - String subType, String functionName, U payload, TypeToken resultType) { - return invokeAsync( - subType, - functionName, - payload, - resultType, - InvokeConfig.builder().build()); - } - - default DurableFuture invokeAsync( - String subType, String functionName, U payload, Class resultType, InvokeConfig config) { - return invokeAsync(subType, functionName, payload, TypeToken.get(resultType), config); - } - DurableFuture invokeAsync( - String subType, String functionName, U payload, TypeToken resultType, InvokeConfig config); - - default DurableCallbackFuture createCallback(Class resultType) { - return createCallback(TypeToken.get(resultType)); - } - - default DurableCallbackFuture createCallback(TypeToken resultType) { - return createCallback(resultType, CallbackConfig.builder().build()); - } - - default DurableCallbackFuture createCallback(Class resultType, CallbackConfig config) { - return createCallback(TypeToken.get(resultType), config); - } - - DurableCallbackFuture createCallback(TypeToken resultType, CallbackConfig config); - - default DurableCallbackFuture createCallback(String subType, Class resultType) { - return createCallback(subType, TypeToken.get(resultType)); - } - - default DurableCallbackFuture createCallback(String subType, TypeToken resultType) { - return createCallback(subType, resultType, CallbackConfig.builder().build()); - } - - default DurableCallbackFuture createCallback(String subType, Class resultType, CallbackConfig config) { - return createCallback(subType, TypeToken.get(resultType), config); - } - - DurableCallbackFuture createCallback(String subType, TypeToken resultType, CallbackConfig config); - - default T runInChildContext(Class resultType, Supplier function) { - return runInChildContext(TypeToken.get(resultType), function); - } - - default T runInChildContext(TypeToken resultType, Supplier function) { - return runInChildContext( - resultType, function, RunInChildContextConfig.builder().build()); - } - - default T runInChildContext(Class resultType, Supplier function, RunInChildContextConfig config) { - return runInChildContext(TypeToken.get(resultType), function, config); - } - - default T runInChildContext(TypeToken resultType, Supplier function, RunInChildContextConfig config) { - return runInChildContextAsync(resultType, function, config).get(); - } - - default DurableFuture runInChildContextAsync(Class resultType, Supplier function) { - return runInChildContextAsync(TypeToken.get(resultType), function); - } - - default DurableFuture runInChildContextAsync(TypeToken resultType, Supplier function) { - return runInChildContextAsync( - resultType, function, RunInChildContextConfig.builder().build()); - } - - default DurableFuture runInChildContextAsync( - Class resultType, Supplier function, RunInChildContextConfig config) { - return runInChildContextAsync(TypeToken.get(resultType), function, config); - } - - DurableFuture runInChildContextAsync( - TypeToken resultType, Supplier function, RunInChildContextConfig config); - - default T runInChildContext(String subType, Class resultType, Supplier function) { - return runInChildContext(subType, TypeToken.get(resultType), function); - } - - default T runInChildContext(String subType, TypeToken resultType, Supplier function) { - return runInChildContext( - subType, resultType, function, RunInChildContextConfig.builder().build()); - } - - default T runInChildContext( - String subType, Class resultType, Supplier function, RunInChildContextConfig config) { - return runInChildContext(subType, TypeToken.get(resultType), function, config); - } - - default T runInChildContext( - String subType, TypeToken resultType, Supplier function, RunInChildContextConfig config) { - return runInChildContextAsync(subType, resultType, function, config).get(); - } - - default DurableFuture runInChildContextAsync(String subType, Class resultType, Supplier function) { - return runInChildContextAsync(subType, TypeToken.get(resultType), function); - } - - default DurableFuture runInChildContextAsync(String subType, TypeToken resultType, Supplier function) { - return runInChildContextAsync( - subType, resultType, function, RunInChildContextConfig.builder().build()); - } - - default DurableFuture runInChildContextAsync( - String subType, Class resultType, Supplier function, RunInChildContextConfig config) { - return runInChildContextAsync(subType, TypeToken.get(resultType), function, config); - } - - DurableFuture runInChildContextAsync( - String subType, TypeToken resultType, Supplier function, RunInChildContextConfig config); - - default T runInChildContext( - String subType, Class resultType, ExtensionContextFunction function, ExtensionContextConfig config) { - return runInChildContext(subType, TypeToken.get(resultType), function, config); - } - - default T runInChildContext( - String subType, - TypeToken resultType, - ExtensionContextFunction function, - ExtensionContextConfig config) { - return runInChildContextAsync(subType, resultType, function, config).get(); - } + String subType, String functionName, U payload, TypeToken resultType, ExtensionInvokeConfig config); - default DurableFuture runInChildContextAsync( - String subType, Class resultType, ExtensionContextFunction function, ExtensionContextConfig config) { - return runInChildContextAsync(subType, TypeToken.get(resultType), function, config); - } + DurableCallbackFuture createCallback( + String subType, TypeToken resultType, ExtensionCallbackConfig config); DurableFuture runInChildContextAsync( String subType, diff --git a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionOperationImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionOperationImpl.java new file mode 100644 index 000000000..299d53192 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionOperationImpl.java @@ -0,0 +1,160 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.extension; + +import java.time.Duration; +import java.util.Objects; +import java.util.concurrent.atomic.AtomicBoolean; +import software.amazon.awssdk.services.lambda.model.OperationType; +import software.amazon.lambda.durable.DurableCallbackFuture; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.model.OperationIdentifier; +import software.amazon.lambda.durable.primitive.BasePrimitive; +import software.amazon.lambda.durable.primitive.CallbackPrimitive; +import software.amazon.lambda.durable.primitive.ChildContextPrimitive; +import software.amazon.lambda.durable.primitive.InvokePrimitive; +import software.amazon.lambda.durable.primitive.StepPrimitive; +import software.amazon.lambda.durable.primitive.WaitPrimitive; +import software.amazon.lambda.durable.util.ParameterValidator; + +public final class ExtensionOperationImpl implements ExtensionOperation { + private final DurableContextImpl context; + private final String operationId; + private final String name; + private final BasePrimitive lateCheckpointOwner; + private final AtomicBoolean claimed = new AtomicBoolean(); + + public ExtensionOperationImpl( + DurableContextImpl context, String operationId, String name, BasePrimitive lateCheckpointOwner) { + this.context = context; + this.operationId = operationId; + this.name = name; + this.lateCheckpointOwner = lateCheckpointOwner; + } + + @Override + public DurableFuture stepAsync( + String subType, TypeToken resultType, ExtensionStepFunction function, ExtensionStepConfig config) { + validateSubType(subType); + Objects.requireNonNull(function, "function cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + claim(); + if (config.serDes() == null) { + config = config.toBuilder() + .serDes(context.getDurableConfig().getSerDes()) + .build(); + } + var operation = new StepPrimitive<>( + new OperationIdentifier(operationId, name, OperationType.STEP, subType), + function, + resultType, + config, + context); + operation.execute(); + return operation; + } + + @Override + public DurableFuture waitAsync(String subType, Duration duration) { + validateSubType(subType); + claim(); + ParameterValidator.validateDuration(duration, "Wait duration"); + var operation = new WaitPrimitive( + new OperationIdentifier(operationId, name, OperationType.WAIT, subType), duration, context); + operation.execute(); + return operation; + } + + @Override + public DurableFuture invokeAsync( + String subType, String functionName, U payload, TypeToken resultType, ExtensionInvokeConfig config) { + validateSubType(subType); + Objects.requireNonNull(resultType, "resultType cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + claim(); + if (config.serDes() == null) { + config = config.toBuilder() + .serDes(context.getDurableConfig().getSerDes()) + .build(); + } + if (config.payloadSerDes() == null) { + config = config.toBuilder() + .payloadSerDes(context.getDurableConfig().getSerDes()) + .build(); + } + var operation = new InvokePrimitive<>( + new OperationIdentifier(operationId, name, OperationType.CHAINED_INVOKE, subType), + functionName, + payload, + resultType, + config, + context); + operation.execute(); + return operation; + } + + @Override + public DurableCallbackFuture createCallback( + String subType, TypeToken resultType, ExtensionCallbackConfig config) { + validateSubType(subType); + Objects.requireNonNull(resultType, "resultType cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + claim(); + if (config.serDes() == null) { + config = config.toBuilder() + .serDes(context.getDurableConfig().getSerDes()) + .build(); + } + var operation = new CallbackPrimitive<>( + new OperationIdentifier(operationId, name, OperationType.CALLBACK, subType), + resultType, + config, + context); + operation.execute(); + return operation; + } + + @Override + public DurableFuture runInChildContextAsync( + String subType, + TypeToken resultType, + ExtensionContextFunction function, + ExtensionContextConfig config) { + validateSubType(subType); + Objects.requireNonNull(resultType, "resultType cannot be null"); + Objects.requireNonNull(function, "function cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + claim(); + var childConfig = config.childContextConfig(); + if (childConfig.serDes() == null) { + childConfig = childConfig.toBuilder() + .serDes(context.getDurableConfig().getSerDes()) + .build(); + config = config.toBuilder().childContextConfig(childConfig).build(); + } + var operation = new ChildContextPrimitive<>( + new OperationIdentifier(operationId, name, OperationType.CONTEXT, subType), + function, + resultType, + config, + context, + lateCheckpointOwner); + operation.execute(); + return operation; + } + + private void validateSubType(String subType) { + Objects.requireNonNull(subType, "subType cannot be null"); + if (subType.isBlank()) { + throw new IllegalArgumentException("subType cannot be blank"); + } + } + + private void claim() { + if (!claimed.compareAndSet(false, true)) { + throw new IllegalStateException("An extension operation reservation can only be used once"); + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepConfig.java index 81c0c4284..2b0e9527c 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepConfig.java @@ -2,6 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.extension; +import software.amazon.lambda.durable.config.StepSemantics; +import software.amazon.lambda.durable.retry.RetryStrategy; import software.amazon.lambda.durable.serde.SerDes; /** @@ -12,10 +14,14 @@ public final class ExtensionStepConfig { private final T initialState; private final SerDes serDes; + private final RetryStrategy retryStrategy; + private final StepSemantics semanticsPerRetry; private ExtensionStepConfig(Builder builder) { initialState = builder.initialState; serDes = builder.serDes; + retryStrategy = builder.retryStrategy; + semanticsPerRetry = builder.semanticsPerRetry; } /** Returns the state supplied to the first attempt. */ @@ -28,6 +34,25 @@ public SerDes serDes() { return serDes; } + /** Returns the exception retry strategy, or {@code null} when thrown exceptions are terminal. */ + public RetryStrategy retryStrategy() { + return retryStrategy; + } + + /** Returns the delivery semantics used for each attempt. */ + public StepSemantics semanticsPerRetry() { + return semanticsPerRetry != null ? semanticsPerRetry : StepSemantics.AT_LEAST_ONCE_PER_RETRY; + } + + /** Returns a builder initialized from this configuration. */ + public Builder toBuilder() { + return new Builder() + .initialState(initialState) + .serDes(serDes) + .retryStrategy(retryStrategy) + .semanticsPerRetry(semanticsPerRetry); + } + /** Returns a builder for a stateful extension step. */ public static Builder builder() { return new Builder<>(); @@ -37,6 +62,8 @@ public static Builder builder() { public static final class Builder { private T initialState; private SerDes serDes; + private RetryStrategy retryStrategy; + private StepSemantics semanticsPerRetry; private Builder() {} @@ -52,6 +79,18 @@ public Builder serDes(SerDes serDes) { return this; } + /** Sets the retry strategy used when the extension function throws. */ + public Builder retryStrategy(RetryStrategy retryStrategy) { + this.retryStrategy = retryStrategy; + return this; + } + + /** Sets the delivery semantics used for each attempt. */ + public Builder semanticsPerRetry(StepSemantics semanticsPerRetry) { + this.semanticsPerRetry = semanticsPerRetry; + return this; + } + /** Builds the immutable configuration. */ public ExtensionStepConfig build() { return new ExtensionStepConfig<>(this); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/model/OperationDescriptor.java b/sdk/src/main/java/software/amazon/lambda/durable/model/OperationDescriptor.java deleted file mode 100644 index ffc663d0d..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/model/OperationDescriptor.java +++ /dev/null @@ -1,46 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.model; - -import java.util.Objects; -import software.amazon.awssdk.services.lambda.model.OperationType; - -/** - * Operation identity that permits extension-defined subtype strings. - * - * @param operationId globally unique operation ID - * @param name human-readable operation name - * @param operationType backend primitive operation type - * @param subType checkpoint subtype string - */ -public record OperationDescriptor(String operationId, String name, OperationType operationType, String subType) { - public OperationDescriptor { - Objects.requireNonNull(operationId, "operationId cannot be null"); - Objects.requireNonNull(operationType, "operationType cannot be null"); - Objects.requireNonNull(subType, "subType cannot be null"); - if (subType.isBlank()) { - throw new IllegalArgumentException("subType cannot be blank"); - } - } - - /** Converts an existing enum-based identity to a descriptor. */ - public static OperationDescriptor from(OperationIdentifier identifier) { - Objects.requireNonNull(identifier, "identifier cannot be null"); - return new OperationDescriptor( - identifier.operationId(), - identifier.name(), - identifier.operationType(), - identifier.subType().getValue()); - } - - /** Returns the matching SDK subtype, or {@code null} for an extension-defined value. */ - public OperationSubType standardSubType() { - for (var candidate : OperationSubType.values()) { - if (candidate.getOperationType() == operationType - && candidate.getValue().equals(subType)) { - return candidate; - } - } - return null; - } -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/model/OperationIdentifier.java b/sdk/src/main/java/software/amazon/lambda/durable/model/OperationIdentifier.java index 9f986f4ce..e7b84a25e 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/model/OperationIdentifier.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/model/OperationIdentifier.java @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.model; +import java.util.Objects; import software.amazon.awssdk.services.lambda.model.OperationType; /** @@ -9,17 +10,33 @@ * * @param operationId unique sequential identifier for the operation within an execution * @param name human-readable name for the operation - * @param subType the operation sub-type which also determines the operation type + * @param operationType backend primitive operation type + * @param subType checkpoint subtype string */ -public record OperationIdentifier(String operationId, String name, OperationSubType subType) { - - /** Returns the operation type derived from the sub-type. */ - public OperationType operationType() { - return subType.getOperationType(); +public record OperationIdentifier(String operationId, String name, OperationType operationType, String subType) { + public OperationIdentifier { + Objects.requireNonNull(operationId, "operationId cannot be null"); + Objects.requireNonNull(operationType, "operationType cannot be null"); + Objects.requireNonNull(subType, "subType cannot be null"); + if (subType.isBlank()) { + throw new IllegalArgumentException("subType cannot be blank"); + } } - /** Creates an identifier with the given sub-type. */ + /** Creates an identifier for a standard SDK operation sub-type. */ public static OperationIdentifier of(String operationId, String name, OperationSubType subType) { - return new OperationIdentifier(operationId, name, subType); + Objects.requireNonNull(subType, "subType cannot be null"); + return new OperationIdentifier(operationId, name, subType.getOperationType(), subType.getValue()); + } + + /** Returns the matching SDK subtype, or {@code null} for an extension-defined value. */ + public OperationSubType standardSubType() { + for (var candidate : OperationSubType.values()) { + if (candidate.getOperationType() == operationType + && candidate.getValue().equals(subType)) { + return candidate; + } + } + return null; } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/DeferredDurableFuture.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DeferredDurableFuture.java similarity index 97% rename from sdk/src/main/java/software/amazon/lambda/durable/context/extension/DeferredDurableFuture.java rename to sdk/src/main/java/software/amazon/lambda/durable/operation/DeferredDurableFuture.java index 43bea59ab..7bf04fc7c 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/DeferredDurableFuture.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DeferredDurableFuture.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.context.extension; +package software.amazon.lambda.durable.operation; import java.util.Objects; import java.util.concurrent.CompletableFuture; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableCallbackOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableCallbackOperation.java new file mode 100644 index 000000000..7d4880102 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableCallbackOperation.java @@ -0,0 +1,119 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.operation; + +import static software.amazon.lambda.durable.model.OperationSubType.CALLBACK; + +import java.time.Duration; +import java.util.Objects; +import software.amazon.lambda.durable.DurableCallbackFuture; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.extension.ExtensionCallbackConfig; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.serde.SerDes; +import software.amazon.lambda.durable.util.ParameterValidator; + +/** Context-free static facade and canonical implementation of durable CALLBACK operations. */ +public final class DurableCallbackOperation { + private DurableCallbackOperation() {} + + public static DurableCallbackFuture createCallback(String name, Class resultType) { + return createCallback(name, TypeToken.get(resultType)); + } + + public static DurableCallbackFuture createCallback(String name, TypeToken resultType) { + return createCallback(name, resultType, CallbackConfig.builder().build()); + } + + public static DurableCallbackFuture createCallback(String name, Class resultType, CallbackConfig config) { + return createCallback(name, TypeToken.get(resultType), config); + } + + public static DurableCallbackFuture createCallback( + String name, TypeToken resultType, CallbackConfig config) { + return createCallback(ExtensionContext.getCurrentContext(), name, resultType, config); + } + + public static DurableCallbackFuture createCallback( + ExtensionContext context, String name, TypeToken resultType, CallbackConfig config) { + Objects.requireNonNull(context, "context cannot be null"); + Objects.requireNonNull(resultType, "resultType cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + ParameterValidator.validateOperationName(name); + return context.reserve(name).createCallback(CALLBACK.getValue(), resultType, extensionConfig(config)); + } + + static ExtensionCallbackConfig extensionConfig(CallbackConfig config) { + return ExtensionCallbackConfig.builder() + .timeout(config.timeout()) + .heartbeatTimeout(config.heartbeatTimeout()) + .serDes(config.serDes()) + .build(); + } + + /** Configuration for durable CALLBACK operations. */ + public static final class CallbackConfig { + private final Duration timeout; + private final Duration heartbeatTimeout; + private final SerDes serDes; + + private CallbackConfig(Builder builder) { + timeout = builder.timeout; + heartbeatTimeout = builder.heartbeatTimeout; + serDes = builder.serDes; + } + + public Duration timeout() { + return timeout; + } + + public Duration heartbeatTimeout() { + return heartbeatTimeout; + } + + public SerDes serDes() { + return serDes; + } + + public static Builder builder() { + return new Builder(); + } + + public Builder toBuilder() { + return new Builder() + .timeout(timeout) + .heartbeatTimeout(heartbeatTimeout) + .serDes(serDes); + } + + /** Builder for {@link CallbackConfig}. */ + public static final class Builder { + private Duration timeout; + private Duration heartbeatTimeout; + private SerDes serDes; + + private Builder() {} + + public Builder timeout(Duration timeout) { + ParameterValidator.validateOptionalDuration(timeout, "Callback timeout"); + this.timeout = timeout; + return this; + } + + public Builder heartbeatTimeout(Duration heartbeatTimeout) { + ParameterValidator.validateOptionalDuration(heartbeatTimeout, "Heartbeat timeout"); + this.heartbeatTimeout = heartbeatTimeout; + return this; + } + + public Builder serDes(SerDes serDes) { + this.serDes = serDes; + return this; + } + + public CallbackConfig build() { + return new CallbackConfig(this); + } + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableContextOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableContextOperation.java new file mode 100644 index 000000000..c055e6722 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableContextOperation.java @@ -0,0 +1,137 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.operation; + +import static software.amazon.lambda.durable.model.OperationSubType.RUN_IN_CHILD_CONTEXT; + +import java.util.Objects; +import java.util.function.Function; +import java.util.function.Supplier; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextResult; +import software.amazon.lambda.durable.serde.SerDes; +import software.amazon.lambda.durable.util.ParameterValidator; + +/** Context-free static facade and canonical implementation of durable CONTEXT operations. */ +public final class DurableContextOperation { + private static final int LARGE_RESULT_THRESHOLD = 256 * 1024; + + private DurableContextOperation() {} + + public static T runInChildContext(String name, Class resultType, Supplier function) { + return runInChildContext(name, TypeToken.get(resultType), function); + } + + public static T runInChildContext(String name, TypeToken resultType, Supplier function) { + return runInChildContext( + name, resultType, function, RunInChildContextConfig.builder().build()); + } + + public static T runInChildContext( + String name, Class resultType, Supplier function, RunInChildContextConfig config) { + return runInChildContext(name, TypeToken.get(resultType), function, config); + } + + public static T runInChildContext( + String name, TypeToken resultType, Supplier function, RunInChildContextConfig config) { + return runInChildContextAsync(name, resultType, function, config).get(); + } + + public static DurableFuture runInChildContextAsync(String name, Class resultType, Supplier function) { + return runInChildContextAsync(name, TypeToken.get(resultType), function); + } + + public static DurableFuture runInChildContextAsync( + String name, TypeToken resultType, Supplier function) { + return runInChildContextAsync( + name, resultType, function, RunInChildContextConfig.builder().build()); + } + + public static DurableFuture runInChildContextAsync( + String name, Class resultType, Supplier function, RunInChildContextConfig config) { + return runInChildContextAsync(name, TypeToken.get(resultType), function, config); + } + + public static DurableFuture runInChildContextAsync( + String name, TypeToken resultType, Supplier function, RunInChildContextConfig config) { + Objects.requireNonNull(function, "function cannot be null"); + return runInChildContextAsync( + ExtensionContext.getCurrentContext(), name, resultType, ignored -> function.get(), config); + } + + public static DurableFuture runInChildContextAsync( + ExtensionContext context, + String name, + TypeToken resultType, + Function function, + RunInChildContextConfig config) { + Objects.requireNonNull(context, "context cannot be null"); + Objects.requireNonNull(resultType, "resultType cannot be null"); + Objects.requireNonNull(function, "function cannot be null"); + Objects.requireNonNull(config, "RunInChildContextConfig cannot be null"); + ParameterValidator.validateOperationName(name); + + return context.reserve(name) + .runInChildContextAsync( + RUN_IN_CHILD_CONTEXT.getValue(), + resultType, + () -> ExtensionContextResult.replayChildrenAboveSize( + function.apply(DurableContext.getCurrentContext()), null, LARGE_RESULT_THRESHOLD), + ExtensionContextConfig.builder() + .childContextConfig(OperationConfigAdapters.toLegacy(config)) + .build()); + } + + /** Configuration for durable CONTEXT operations. */ + public static final class RunInChildContextConfig { + private final SerDes serDes; + private final boolean virtual; + + private RunInChildContextConfig(Builder builder) { + serDes = builder.serDes; + virtual = Objects.requireNonNullElse(builder.virtual, false); + } + + public SerDes serDes() { + return serDes; + } + + public Boolean isVirtual() { + return virtual; + } + + public static Builder builder() { + return new Builder(); + } + + public Builder toBuilder() { + return new Builder().serDes(serDes).isVirtual(virtual); + } + + /** Builder for {@link RunInChildContextConfig}. */ + public static final class Builder { + private SerDes serDes; + private Boolean virtual; + + private Builder() {} + + public Builder serDes(SerDes serDes) { + this.serDes = serDes; + return this; + } + + public Builder isVirtual(Boolean virtual) { + this.virtual = virtual; + return this; + } + + public RunInChildContextConfig build() { + return new RunInChildContextConfig(this); + } + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableInvokeOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableInvokeOperation.java new file mode 100644 index 000000000..28e0713ad --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableInvokeOperation.java @@ -0,0 +1,142 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.operation; + +import static software.amazon.lambda.durable.model.OperationSubType.CHAINED_INVOKE; + +import java.util.Objects; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionInvokeConfig; +import software.amazon.lambda.durable.serde.SerDes; +import software.amazon.lambda.durable.util.ParameterValidator; + +/** Context-free static facade and canonical implementation of durable CHAINED_INVOKE operations. */ +public final class DurableInvokeOperation { + private DurableInvokeOperation() {} + + public static T invoke(String name, String functionName, U payload, Class resultType) { + return invoke(name, functionName, payload, TypeToken.get(resultType)); + } + + public static T invoke(String name, String functionName, U payload, TypeToken resultType) { + return invoke( + name, functionName, payload, resultType, InvokeConfig.builder().build()); + } + + public static T invoke( + String name, String functionName, U payload, Class resultType, InvokeConfig config) { + return invoke(name, functionName, payload, TypeToken.get(resultType), config); + } + + public static T invoke( + String name, String functionName, U payload, TypeToken resultType, InvokeConfig config) { + return invokeAsync(name, functionName, payload, resultType, config).get(); + } + + public static DurableFuture invokeAsync( + String name, String functionName, U payload, Class resultType) { + return invokeAsync(name, functionName, payload, TypeToken.get(resultType)); + } + + public static DurableFuture invokeAsync( + String name, String functionName, U payload, TypeToken resultType) { + return invokeAsync( + name, functionName, payload, resultType, InvokeConfig.builder().build()); + } + + public static DurableFuture invokeAsync( + String name, String functionName, U payload, Class resultType, InvokeConfig config) { + return invokeAsync(name, functionName, payload, TypeToken.get(resultType), config); + } + + public static DurableFuture invokeAsync( + String name, String functionName, U payload, TypeToken resultType, InvokeConfig config) { + return invokeAsync(ExtensionContext.getCurrentContext(), name, functionName, payload, resultType, config); + } + + public static DurableFuture invokeAsync( + ExtensionContext context, + String name, + String functionName, + U payload, + TypeToken resultType, + InvokeConfig config) { + Objects.requireNonNull(context, "context cannot be null"); + Objects.requireNonNull(resultType, "resultType cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + ParameterValidator.validateOperationName(name); + return context.reserve(name) + .invokeAsync(CHAINED_INVOKE.getValue(), functionName, payload, resultType, extensionConfig(config)); + } + + private static ExtensionInvokeConfig extensionConfig(InvokeConfig config) { + return ExtensionInvokeConfig.builder() + .payloadSerDes(config.payloadSerDes()) + .serDes(config.serDes()) + .tenantId(config.tenantId()) + .build(); + } + + /** Configuration for durable CHAINED_INVOKE operations. */ + public static final class InvokeConfig { + private final SerDes payloadSerDes; + private final SerDes serDes; + private final String tenantId; + + private InvokeConfig(Builder builder) { + payloadSerDes = builder.payloadSerDes; + serDes = builder.serDes; + tenantId = builder.tenantId; + } + + public SerDes payloadSerDes() { + return payloadSerDes; + } + + public SerDes serDes() { + return serDes; + } + + public String tenantId() { + return tenantId; + } + + public static Builder builder() { + return new Builder(); + } + + public Builder toBuilder() { + return new Builder().payloadSerDes(payloadSerDes).serDes(serDes).tenantId(tenantId); + } + + /** Builder for {@link InvokeConfig}. */ + public static final class Builder { + private SerDes payloadSerDes; + private SerDes serDes; + private String tenantId; + + private Builder() {} + + public Builder payloadSerDes(SerDes payloadSerDes) { + this.payloadSerDes = payloadSerDes; + return this; + } + + public Builder serDes(SerDes serDes) { + this.serDes = serDes; + return this; + } + + public Builder tenantId(String tenantId) { + this.tenantId = tenantId; + return this; + } + + public InvokeConfig build() { + return new InvokeConfig(this); + } + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/MapExtension.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java similarity index 55% rename from sdk/src/main/java/software/amazon/lambda/durable/context/extension/MapExtension.java rename to sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java index a05be6659..4b348ab61 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/MapExtension.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java @@ -1,23 +1,26 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.context.extension; +package software.amazon.lambda.durable.operation; import static software.amazon.lambda.durable.config.NestingType.FLAT; -import static software.amazon.lambda.durable.context.extension.ExtensionConcurrencyCoordinator.ItemStatus.FAILED; -import static software.amazon.lambda.durable.context.extension.ExtensionConcurrencyCoordinator.ItemStatus.SKIPPED; import static software.amazon.lambda.durable.model.OperationSubType.MAP; import static software.amazon.lambda.durable.model.OperationSubType.MAP_ITERATION; +import static software.amazon.lambda.durable.operation.OperationConcurrencyCoordinator.ItemStatus.FAILED; +import static software.amazon.lambda.durable.operation.OperationConcurrencyCoordinator.ItemStatus.SKIPPED; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; import java.util.List; import java.util.Objects; +import java.util.function.BiFunction; +import java.util.function.Function; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.MapItemContext; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.CompletionConfig; -import software.amazon.lambda.durable.config.MapConfig; +import software.amazon.lambda.durable.config.NestingType; import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.SuspendExecutionException; @@ -26,16 +29,57 @@ import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; import software.amazon.lambda.durable.extension.ExtensionContextResult; import software.amazon.lambda.durable.model.MapResult; +import software.amazon.lambda.durable.serde.SerDes; import software.amazon.lambda.durable.util.ExceptionHelper; import software.amazon.lambda.durable.util.ParameterValidator; -/** Canonical implementation of the built-in map extension. */ -public final class MapExtension { +/** Context-free static facade and canonical implementation of durable MAP operations. */ +public final class DurableMapOperation { private static final int LARGE_RESULT_THRESHOLD = 256 * 1024; - private MapExtension() {} + private DurableMapOperation() {} - public static DurableFuture> execute( + public static MapResult map( + String name, Collection items, Class resultType, Function function) { + return mapAsync(name, items, resultType, function).get(); + } + + public static MapResult map( + String name, Collection items, TypeToken resultType, Function function) { + return mapAsync(name, items, resultType, function).get(); + } + + public static MapResult map( + String name, Collection items, Class resultType, Function function, MapConfig config) { + return mapAsync(name, items, resultType, function, config).get(); + } + + public static MapResult map( + String name, Collection items, TypeToken resultType, Function function, MapConfig config) { + return mapAsync(name, items, resultType, function, config).get(); + } + + public static DurableFuture> mapAsync( + String name, Collection items, Class resultType, Function function) { + return mapAsync(name, items, TypeToken.get(resultType), function); + } + + public static DurableFuture> mapAsync( + String name, Collection items, TypeToken resultType, Function function) { + return mapAsync(name, items, resultType, function, MapConfig.builder().build()); + } + + public static DurableFuture> mapAsync( + String name, Collection items, Class resultType, Function function, MapConfig config) { + return mapAsync(name, items, TypeToken.get(resultType), function, config); + } + + public static DurableFuture> mapAsync( + String name, Collection items, TypeToken resultType, Function function, MapConfig config) { + return mapAsync(ExtensionContext.getCurrentContext(), name, items, resultType, adapt(function), config); + } + + public static DurableFuture> mapAsync( ExtensionContext context, String name, Collection items, @@ -70,6 +114,15 @@ public static DurableFuture> execute( parentConfig(mapConfig, virtualEmptyMap)); } + private static DurableContext.MapFunction adapt(Function function) { + Objects.requireNonNull(function, "function cannot be null"); + return (item, index, ignored) -> { + try (var scope = MapItemContext.attach(index)) { + return function.apply(item); + } + }; + } + private static ExtensionContextResult> executeInChildContext( String name, List items, @@ -95,7 +148,7 @@ private static ExtensionContextResult> executeInChildContext throw new IllegalStateException("Missing result in completed Map operation"); } - var coordinator = new ExtensionConcurrencyCoordinator(config.maxConcurrency(), config.completionConfig()); + var coordinator = new OperationConcurrencyCoordinator(config.maxConcurrency(), config.completionConfig()); var registeredItems = registerItems(coordinator, items, iterationNames, resultType, function, config, replayState); coordinator.closeRegistration(); @@ -109,8 +162,8 @@ private static ExtensionContextResult> executeInChildContext : ExtensionContextResult.replayChildren(result, strippedResult); } - private static List> registerItems( - ExtensionConcurrencyCoordinator coordinator, + private static List> registerItems( + OperationConcurrencyCoordinator coordinator, List items, List iterationNames, TypeToken resultType, @@ -118,10 +171,12 @@ private static List> registerItem MapConfig config, MapResult replayState) { var context = ExtensionContext.getCurrentContext(); - var registeredItems = new ArrayList>(items.size()); - var iterationConfig = RunInChildContextConfig.builder() - .serDes(config.serDes()) - .isVirtual(config.nestingType() == FLAT) + var registeredItems = new ArrayList>(items.size()); + var iterationConfig = ExtensionContextConfig.builder() + .childContextConfig(RunInChildContextConfig.builder() + .serDes(config.serDes()) + .isVirtual(config.nestingType() == FLAT) + .build()) .build(); for (int index = 0; index < items.size(); index++) { @@ -134,7 +189,10 @@ private static List> registerItem () -> reservation.runInChildContextAsync( MAP_ITERATION.getValue(), resultType, - () -> function.apply(item, itemIndex, DurableContext.getCurrentContext()), + () -> ExtensionContextResult.replayChildrenAboveSize( + function.apply(item, itemIndex, DurableContext.getCurrentContext()), + null, + LARGE_RESULT_THRESHOLD), iterationConfig), skipped)); } @@ -153,15 +211,15 @@ private static List resolveIterationNames(String mapName, List items, return names; } - private static ExtensionConcurrencyCoordinator.ExpectedCompletionStatus expectedCompletion( + private static OperationConcurrencyCoordinator.ExpectedCompletionStatus expectedCompletion( MapResult replayState) { - return new ExtensionConcurrencyCoordinator.ExpectedCompletionStatus( + return new OperationConcurrencyCoordinator.ExpectedCompletionStatus( replayState.succeeded().size() + replayState.failed().size(), CompletionConfig.CompletionDecision.complete(replayState.completionReason())); } private static MapResult constructResult( - List> items, + List> items, CompletionConfig.CompletionDecision completionDecision) { var results = new ArrayList>(Collections.nCopies(items.size(), null)); for (int index = 0; index < items.size(); index++) { @@ -178,7 +236,7 @@ private static MapResult constructResult( return new MapResult<>(results, completionDecision.completionStatus()); } - private static MapResult.MapResultItem failedResult(ExtensionConcurrencyCoordinator.Item item) { + private static MapResult.MapResultItem failedResult(OperationConcurrencyCoordinator.Item item) { try { item.future().get(); throw new IllegalStateException("Failed map item completed successfully"); @@ -222,4 +280,107 @@ private static void validateMinSuccessful(List items, MapConfig config) { private static TypeToken> mapResultType() { return new TypeToken<>() {}; } + + /** Configuration for durable MAP operations. */ + public static final class MapConfig { + private final int maxConcurrency; + private final CompletionConfig completionConfig; + private final SerDes serDes; + private final NestingType nestingType; + private final BiFunction itemNamer; + + private MapConfig(Builder builder) { + maxConcurrency = Objects.requireNonNullElse(builder.maxConcurrency, Integer.MAX_VALUE); + completionConfig = Objects.requireNonNullElseGet(builder.completionConfig, CompletionConfig::allCompleted); + serDes = builder.serDes; + nestingType = Objects.requireNonNullElse(builder.nestingType, NestingType.NESTED); + itemNamer = builder.itemNamer; + if (itemNamer != null && nestingType == FLAT) { + throw new IllegalArgumentException("itemNamer is not supported with FLAT map nesting"); + } + } + + public Integer maxConcurrency() { + return maxConcurrency; + } + + public CompletionConfig completionConfig() { + return completionConfig; + } + + public SerDes serDes() { + return serDes; + } + + public NestingType nestingType() { + return nestingType; + } + + public BiFunction itemNamer() { + return itemNamer; + } + + public static Builder builder() { + return new Builder(); + } + + public Builder toBuilder() { + return new Builder() + .maxConcurrency(maxConcurrency) + .completionConfig(completionConfig) + .serDes(serDes) + .nestingType(nestingType) + .itemNamer(itemNamer); + } + + /** Builder for {@link MapConfig}. */ + public static final class Builder { + private Integer maxConcurrency; + private CompletionConfig completionConfig; + private SerDes serDes; + private NestingType nestingType; + private BiFunction itemNamer; + + private Builder() {} + + public Builder maxConcurrency(Integer maxConcurrency) { + if (maxConcurrency != null && maxConcurrency < 1) { + throw new IllegalArgumentException("maxConcurrency must be at least 1, got: " + maxConcurrency); + } + this.maxConcurrency = maxConcurrency; + return this; + } + + public Builder completionConfig(CompletionConfig completionConfig) { + this.completionConfig = completionConfig; + return this; + } + + public Builder serDes(SerDes serDes) { + this.serDes = serDes; + return this; + } + + public Builder nestingType(NestingType nestingType) { + this.nestingType = nestingType; + return this; + } + + public Builder itemNamer(BiFunction itemNamer) { + this.itemNamer = itemNamer; + return this; + } + + public Builder itemNamer(Class itemType, BiFunction itemNamer) { + Objects.requireNonNull(itemType, "itemType cannot be null"); + this.itemNamer = + itemNamer == null ? null : (item, index) -> itemNamer.apply(itemType.cast(item), index); + return this; + } + + public MapConfig build() { + return new MapConfig(this); + } + } + } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableParallelOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableParallelOperation.java new file mode 100644 index 000000000..0e16f8aa8 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableParallelOperation.java @@ -0,0 +1,140 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.operation; + +import java.util.Objects; +import software.amazon.lambda.durable.ParallelDurableFuture; +import software.amazon.lambda.durable.config.CompletionConfig; +import software.amazon.lambda.durable.config.NestingType; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.serde.SerDes; +import software.amazon.lambda.durable.util.ParameterValidator; + +/** Context-free static facade and canonical implementation of durable PARALLEL operations. */ +public final class DurableParallelOperation { + private DurableParallelOperation() {} + + public static ParallelDurableFuture parallel(String name) { + return parallel(name, ParallelConfig.builder().build()); + } + + public static ParallelDurableFuture parallel(String name, ParallelConfig config) { + return parallel(ExtensionContext.getCurrentContext(), name, config); + } + + public static ParallelDurableFuture parallel(ExtensionContext context, String name, ParallelConfig config) { + Objects.requireNonNull(context, "context cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + ParameterValidator.validateOperationName(name); + return new ParallelOperationFuture(context, name, config); + } + + /** Configuration for durable PARALLEL operations. */ + public static final class ParallelConfig { + private final int maxConcurrency; + private final CompletionConfig completionConfig; + private final NestingType nestingType; + + private ParallelConfig(Builder builder) { + maxConcurrency = Objects.requireNonNullElse(builder.maxConcurrency, Integer.MAX_VALUE); + completionConfig = Objects.requireNonNullElseGet(builder.completionConfig, CompletionConfig::allCompleted); + nestingType = Objects.requireNonNullElse(builder.nestingType, NestingType.NESTED); + } + + public int maxConcurrency() { + return maxConcurrency; + } + + public CompletionConfig completionConfig() { + return completionConfig; + } + + public NestingType nestingType() { + return nestingType; + } + + public static Builder builder() { + return new Builder(); + } + + public Builder toBuilder() { + return new Builder() + .maxConcurrency(maxConcurrency) + .completionConfig(completionConfig) + .nestingType(nestingType); + } + + /** Builder for {@link ParallelConfig}. */ + public static final class Builder { + private Integer maxConcurrency; + private CompletionConfig completionConfig; + private NestingType nestingType; + + private Builder() {} + + public Builder maxConcurrency(Integer maxConcurrency) { + if (maxConcurrency != null && maxConcurrency < 1) { + throw new IllegalArgumentException("maxConcurrency must be at least 1, got: " + maxConcurrency); + } + this.maxConcurrency = maxConcurrency; + return this; + } + + public Builder completionConfig(CompletionConfig completionConfig) { + if (completionConfig != null + && !completionConfig.hasCustomShouldComplete() + && completionConfig.toleratedFailurePercentage() != null) { + throw new IllegalArgumentException("ParallelConfig does not support toleratedFailurePercentage"); + } + this.completionConfig = completionConfig; + return this; + } + + public Builder nestingType(NestingType nestingType) { + this.nestingType = nestingType; + return this; + } + + public ParallelConfig build() { + return new ParallelConfig(this); + } + } + } + + /** Configuration for a durable PARALLEL branch. */ + public static final class ParallelBranchConfig { + private final SerDes serDes; + + private ParallelBranchConfig(Builder builder) { + serDes = builder.serDes; + } + + public SerDes serDes() { + return serDes; + } + + public static Builder builder() { + return new Builder(); + } + + public Builder toBuilder() { + return new Builder().serDes(serDes); + } + + /** Builder for {@link ParallelBranchConfig}. */ + public static final class Builder { + private SerDes serDes; + + private Builder() {} + + public Builder serDes(SerDes serDes) { + this.serDes = serDes; + return this; + } + + public ParallelBranchConfig build() { + return new ParallelBranchConfig(this); + } + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableStepOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableStepOperation.java new file mode 100644 index 000000000..82cdb3269 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableStepOperation.java @@ -0,0 +1,149 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.operation; + +import static software.amazon.lambda.durable.model.OperationSubType.STEP; + +import java.util.Objects; +import java.util.function.Function; +import java.util.function.Supplier; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.StepContext; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.StepSemantics; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionStepConfig; +import software.amazon.lambda.durable.extension.ExtensionStepResult; +import software.amazon.lambda.durable.retry.RetryStrategies; +import software.amazon.lambda.durable.retry.RetryStrategy; +import software.amazon.lambda.durable.serde.SerDes; +import software.amazon.lambda.durable.util.ParameterValidator; + +/** Context-free static facade and canonical implementation of durable STEP operations. */ +public final class DurableStepOperation { + private DurableStepOperation() {} + + public static T step(String name, Class resultType, Supplier function) { + return step(name, TypeToken.get(resultType), function); + } + + public static T step(String name, TypeToken resultType, Supplier function) { + return step(name, resultType, function, StepConfig.builder().build()); + } + + public static T step(String name, Class resultType, Supplier function, StepConfig config) { + return step(name, TypeToken.get(resultType), function, config); + } + + public static T step(String name, TypeToken resultType, Supplier function, StepConfig config) { + return stepAsync(name, resultType, function, config).get(); + } + + public static DurableFuture stepAsync(String name, Class resultType, Supplier function) { + return stepAsync(name, TypeToken.get(resultType), function); + } + + public static DurableFuture stepAsync(String name, TypeToken resultType, Supplier function) { + return stepAsync(name, resultType, function, StepConfig.builder().build()); + } + + public static DurableFuture stepAsync( + String name, Class resultType, Supplier function, StepConfig config) { + return stepAsync(name, TypeToken.get(resultType), function, config); + } + + public static DurableFuture stepAsync( + String name, TypeToken resultType, Supplier function, StepConfig config) { + Objects.requireNonNull(function, "function cannot be null"); + return stepAsync(ExtensionContext.getCurrentContext(), name, resultType, ignored -> function.get(), config); + } + + public static DurableFuture stepAsync( + ExtensionContext context, + String name, + TypeToken resultType, + Function function, + StepConfig config) { + Objects.requireNonNull(context, "context cannot be null"); + Objects.requireNonNull(resultType, "resultType cannot be null"); + Objects.requireNonNull(function, "function cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + ParameterValidator.validateOperationName(name); + + return context.reserve(name) + .stepAsync( + STEP.getValue(), + resultType, + ignored -> ExtensionStepResult.succeed(function.apply(StepContext.getCurrentContext())), + ExtensionStepConfig.builder() + .serDes(config.serDes()) + .retryStrategy(config.retryStrategy()) + .semanticsPerRetry(config.semanticsPerRetry()) + .build()); + } + + /** Configuration for durable STEP operations. */ + public static final class StepConfig { + private final RetryStrategy retryStrategy; + private final StepSemantics semanticsPerRetry; + private final SerDes serDes; + + private StepConfig(Builder builder) { + retryStrategy = Objects.requireNonNullElse(builder.retryStrategy, RetryStrategies.Presets.DEFAULT); + semanticsPerRetry = + Objects.requireNonNullElse(builder.semanticsPerRetry, StepSemantics.AT_LEAST_ONCE_PER_RETRY); + serDes = builder.serDes; + } + + public RetryStrategy retryStrategy() { + return retryStrategy; + } + + public StepSemantics semanticsPerRetry() { + return semanticsPerRetry; + } + + public SerDes serDes() { + return serDes; + } + + public static Builder builder() { + return new Builder(); + } + + public Builder toBuilder() { + return new Builder() + .retryStrategy(retryStrategy) + .semanticsPerRetry(semanticsPerRetry) + .serDes(serDes); + } + + /** Builder for {@link StepConfig}. */ + public static final class Builder { + private RetryStrategy retryStrategy; + private StepSemantics semanticsPerRetry; + private SerDes serDes; + + private Builder() {} + + public Builder retryStrategy(RetryStrategy retryStrategy) { + this.retryStrategy = retryStrategy; + return this; + } + + public Builder semanticsPerRetry(StepSemantics semanticsPerRetry) { + this.semanticsPerRetry = semanticsPerRetry; + return this; + } + + public Builder serDes(SerDes serDes) { + this.serDes = serDes; + return this; + } + + public StepConfig build() { + return new StepConfig(this); + } + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperation.java new file mode 100644 index 000000000..2235a6b4a --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperation.java @@ -0,0 +1,221 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.operation; + +import static software.amazon.lambda.durable.execution.ExecutionManager.isTerminalStatus; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + +import java.util.Objects; +import java.util.function.BiConsumer; +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.DurableFuture; +import software.amazon.lambda.durable.StepContext; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.WaitForCallbackContext; +import software.amazon.lambda.durable.config.RunInChildContextConfig; +import software.amazon.lambda.durable.exception.CallbackFailedException; +import software.amazon.lambda.durable.exception.CallbackSubmitterException; +import software.amazon.lambda.durable.exception.CallbackTimeoutException; +import software.amazon.lambda.durable.exception.StepFailedException; +import software.amazon.lambda.durable.exception.StepInterruptedException; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextFailure; +import software.amazon.lambda.durable.extension.ExtensionContextResult; +import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.util.ParameterValidator; + +/** Context-free static facade and canonical implementation of durable wait-for-callback operations. */ +public final class DurableWaitForCallbackOperation { + private static final String CALLBACK_SUFFIX = "-callback"; + private static final String SUBMITTER_SUFFIX = "-submitter"; + private static final int MAX_NAME_LENGTH = ParameterValidator.MAX_OPERATION_NAME_LENGTH + - Math.max(CALLBACK_SUFFIX.length(), SUBMITTER_SUFFIX.length()); + + private DurableWaitForCallbackOperation() {} + + public static T waitForCallback(String name, Class resultType, Runnable submitter) { + return waitForCallbackAsync(name, resultType, submitter).get(); + } + + public static T waitForCallback(String name, TypeToken resultType, Runnable submitter) { + return waitForCallbackAsync(name, resultType, submitter).get(); + } + + public static T waitForCallback( + String name, Class resultType, Runnable submitter, WaitForCallbackConfig config) { + return waitForCallbackAsync(name, resultType, submitter, config).get(); + } + + public static T waitForCallback( + String name, TypeToken resultType, Runnable submitter, WaitForCallbackConfig config) { + return waitForCallbackAsync(name, resultType, submitter, config).get(); + } + + public static DurableFuture waitForCallbackAsync(String name, Class resultType, Runnable submitter) { + return waitForCallbackAsync(name, TypeToken.get(resultType), submitter); + } + + public static DurableFuture waitForCallbackAsync(String name, TypeToken resultType, Runnable submitter) { + return waitForCallbackAsync( + name, resultType, submitter, WaitForCallbackConfig.builder().build()); + } + + public static DurableFuture waitForCallbackAsync( + String name, Class resultType, Runnable submitter, WaitForCallbackConfig config) { + return waitForCallbackAsync(name, TypeToken.get(resultType), submitter, config); + } + + public static DurableFuture waitForCallbackAsync( + String name, TypeToken resultType, Runnable submitter, WaitForCallbackConfig config) { + return waitForCallbackAsync(ExtensionContext.getCurrentContext(), name, resultType, adapt(submitter), config); + } + + public static DurableFuture waitForCallbackAsync( + ExtensionContext context, + String name, + TypeToken resultType, + BiConsumer submitter, + WaitForCallbackConfig config) { + Objects.requireNonNull(context, "context cannot be null"); + Objects.requireNonNull(resultType, "resultType cannot be null"); + Objects.requireNonNull(submitter, "submitter cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + ParameterValidator.validateOperationName(name, MAX_NAME_LENGTH); + + var parent = context.reserve(name); + return parent.runInChildContextAsync( + OperationSubType.WAIT_FOR_CALLBACK.getValue(), + resultType, + () -> executeInChildContext(name, resultType, submitter, config), + extensionConfig(config)); + } + + private static BiConsumer adapt(Runnable submitter) { + Objects.requireNonNull(submitter, "submitter cannot be null"); + return (callbackId, ignored) -> { + try (var scope = WaitForCallbackContext.attach(callbackId)) { + submitter.run(); + } + }; + } + + private static ExtensionContextResult executeInChildContext( + String name, + TypeToken resultType, + BiConsumer submitter, + WaitForCallbackConfig config) { + var child = ExtensionContext.getCurrentContext(); + var callback = child.reserve(name + CALLBACK_SUFFIX) + .createCallback( + OperationSubType.CALLBACK.getValue(), + resultType, + DurableCallbackOperation.extensionConfig(config.callbackConfig())); + step( + name + SUBMITTER_SUFFIX, + Void.class, + () -> { + submitter.accept(callback.callbackId(), StepContext.getCurrentContext()); + return null; + }, + config.stepConfig()); + return ExtensionContextResult.completed(callback.get()); + } + + private static ExtensionContextConfig extensionConfig(WaitForCallbackConfig config) { + return ExtensionContextConfig.builder() + .childContextConfig(RunInChildContextConfig.builder() + .serDes(config.stepConfig().serDes()) + .build()) + .errorHandler(DurableWaitForCallbackOperation::translateFailure) + .build(); + } + + private static Throwable translateFailure(ExtensionContextFailure failure) { + var callback = findChild(failure, OperationType.CALLBACK); + var submitter = findChild(failure, OperationType.STEP); + if (callback != null && isTerminalStatus(callback.status())) { + if (callback.status() == OperationStatus.FAILED) { + return new CallbackFailedException(callback); + } + if (callback.status() == OperationStatus.TIMED_OUT) { + return new CallbackTimeoutException(callback); + } + } + if (callback != null + && submitter != null + && isTerminalStatus(submitter.status()) + && submitter.status() != OperationStatus.SUCCEEDED) { + var error = submitter.stepDetails().error(); + var cause = StepInterruptedException.isStepInterruptedException(error) + ? new StepInterruptedException(submitter) + : new StepFailedException(submitter); + return new CallbackSubmitterException(callback, cause); + } + return new IllegalStateException("Unknown waitForCallback status"); + } + + private static Operation findChild(ExtensionContextFailure failure, OperationType type) { + return failure.childOperations().stream() + .filter(summary -> summary.operationType() == type) + .map(summary -> summary.operation()) + .filter(Objects::nonNull) + .findFirst() + .orElse(null); + } + + /** Configuration for durable wait-for-callback operations. */ + public static final class WaitForCallbackConfig { + private final DurableStepOperation.StepConfig stepConfig; + private final DurableCallbackOperation.CallbackConfig callbackConfig; + + private WaitForCallbackConfig(Builder builder) { + stepConfig = + Objects.requireNonNullElseGet(builder.stepConfig, () -> DurableStepOperation.StepConfig.builder() + .build()); + callbackConfig = Objects.requireNonNullElseGet( + builder.callbackConfig, + () -> DurableCallbackOperation.CallbackConfig.builder().build()); + } + + public DurableStepOperation.StepConfig stepConfig() { + return stepConfig; + } + + public DurableCallbackOperation.CallbackConfig callbackConfig() { + return callbackConfig; + } + + public static Builder builder() { + return new Builder(); + } + + public Builder toBuilder() { + return new Builder().stepConfig(stepConfig).callbackConfig(callbackConfig); + } + + /** Builder for {@link WaitForCallbackConfig}. */ + public static final class Builder { + private DurableStepOperation.StepConfig stepConfig; + private DurableCallbackOperation.CallbackConfig callbackConfig; + + private Builder() {} + + public Builder stepConfig(DurableStepOperation.StepConfig stepConfig) { + this.stepConfig = stepConfig; + return this; + } + + public Builder callbackConfig(DurableCallbackOperation.CallbackConfig callbackConfig) { + this.callbackConfig = callbackConfig; + return this; + } + + public WaitForCallbackConfig build() { + return new WaitForCallbackConfig(this); + } + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperation.java new file mode 100644 index 000000000..c662141f8 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperation.java @@ -0,0 +1,186 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.operation; + +import java.util.Objects; +import java.util.function.BiFunction; +import java.util.function.Function; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.StepContext; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionStepConfig; +import software.amazon.lambda.durable.extension.ExtensionStepResult; +import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.model.WaitForConditionResult; +import software.amazon.lambda.durable.retry.WaitForConditionWaitStrategy; +import software.amazon.lambda.durable.retry.WaitStrategies; +import software.amazon.lambda.durable.serde.SerDes; +import software.amazon.lambda.durable.util.ParameterValidator; + +/** Context-free static facade and canonical implementation of durable wait-for-condition operations. */ +public final class DurableWaitForConditionOperation { + private DurableWaitForConditionOperation() {} + + public static T waitForCondition( + String name, Class resultType, Function> checkFunction) { + return waitForConditionAsync(name, resultType, checkFunction).get(); + } + + public static T waitForCondition( + String name, TypeToken resultType, Function> checkFunction) { + return waitForConditionAsync(name, resultType, checkFunction).get(); + } + + public static T waitForCondition( + String name, + Class resultType, + Function> checkFunction, + WaitForConditionConfig config) { + return waitForConditionAsync(name, resultType, checkFunction, config).get(); + } + + public static T waitForCondition( + String name, + TypeToken resultType, + Function> checkFunction, + WaitForConditionConfig config) { + return waitForConditionAsync(name, resultType, checkFunction, config).get(); + } + + public static DurableFuture waitForConditionAsync( + String name, Class resultType, Function> checkFunction) { + return waitForConditionAsync(name, TypeToken.get(resultType), checkFunction); + } + + public static DurableFuture waitForConditionAsync( + String name, TypeToken resultType, Function> checkFunction) { + return waitForConditionAsync( + name, + resultType, + checkFunction, + WaitForConditionConfig.builder().build()); + } + + public static DurableFuture waitForConditionAsync( + String name, + Class resultType, + Function> checkFunction, + WaitForConditionConfig config) { + return waitForConditionAsync(name, TypeToken.get(resultType), checkFunction, config); + } + + public static DurableFuture waitForConditionAsync( + String name, + TypeToken resultType, + Function> checkFunction, + WaitForConditionConfig config) { + return waitForConditionAsync( + ExtensionContext.getCurrentContext(), name, resultType, adapt(checkFunction), config); + } + + public static DurableFuture waitForConditionAsync( + ExtensionContext context, + String name, + TypeToken resultType, + BiFunction> checkFunction, + WaitForConditionConfig config) { + Objects.requireNonNull(context, "context cannot be null"); + Objects.requireNonNull(resultType, "resultType cannot be null"); + Objects.requireNonNull(checkFunction, "checkFunction cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + ParameterValidator.validateOperationName(name); + + var future = context.reserve(name) + .stepAsync( + OperationSubType.WAIT_FOR_CONDITION.getValue(), + resultType, + state -> evaluate(state, checkFunction, config), + ExtensionStepConfig.builder() + .initialState(config.initialState()) + .serDes(config.serDes()) + .build()); + return new WaitForConditionFuture<>(future); + } + + private static BiFunction> adapt( + Function> checkFunction) { + Objects.requireNonNull(checkFunction, "checkFunction cannot be null"); + return (state, ignored) -> checkFunction.apply(state); + } + + private static ExtensionStepResult evaluate( + T state, + BiFunction> checkFunction, + WaitForConditionConfig config) { + var stepContext = StepContext.getCurrentContext(); + var result = Objects.requireNonNull( + checkFunction.apply(state, stepContext), "waitForCondition check result cannot be null"); + if (result.isDone()) { + return ExtensionStepResult.succeed(result.value()); + } + var delay = config.waitStrategy().evaluate(result.value(), stepContext.getAttempt()); + return ExtensionStepResult.retry(result.value(), delay); + } + + /** Configuration for durable wait-for-condition operations. */ + public static final class WaitForConditionConfig { + private final WaitForConditionWaitStrategy waitStrategy; + private final SerDes serDes; + private final T initialState; + + private WaitForConditionConfig(Builder builder) { + waitStrategy = Objects.requireNonNullElseGet(builder.waitStrategy, WaitStrategies::defaultStrategy); + serDes = builder.serDes; + initialState = builder.initialState; + } + + public WaitForConditionWaitStrategy waitStrategy() { + return waitStrategy; + } + + public SerDes serDes() { + return serDes; + } + + public T initialState() { + return initialState; + } + + public static Builder builder() { + return new Builder<>(); + } + + public Builder toBuilder() { + return new Builder().waitStrategy(waitStrategy).serDes(serDes).initialState(initialState); + } + + /** Builder for {@link WaitForConditionConfig}. */ + public static final class Builder { + private WaitForConditionWaitStrategy waitStrategy; + private SerDes serDes; + private T initialState; + + private Builder() {} + + public Builder waitStrategy(WaitForConditionWaitStrategy waitStrategy) { + this.waitStrategy = waitStrategy; + return this; + } + + public Builder serDes(SerDes serDes) { + this.serDes = serDes; + return this; + } + + public Builder initialState(T initialState) { + this.initialState = initialState; + return this; + } + + public WaitForConditionConfig build() { + return new WaitForConditionConfig<>(this); + } + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitOperation.java new file mode 100644 index 000000000..8e9d204ba --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitOperation.java @@ -0,0 +1,31 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.operation; + +import static software.amazon.lambda.durable.model.OperationSubType.WAIT; + +import java.time.Duration; +import java.util.Objects; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.util.ParameterValidator; + +/** Context-free static facade and canonical implementation of durable WAIT operations. */ +public final class DurableWaitOperation { + private DurableWaitOperation() {} + + public static Void wait(String name, Duration duration) { + return waitAsync(name, duration).get(); + } + + public static DurableFuture waitAsync(String name, Duration duration) { + return waitAsync(ExtensionContext.getCurrentContext(), name, duration); + } + + public static DurableFuture waitAsync(ExtensionContext context, String name, Duration duration) { + Objects.requireNonNull(context, "context cannot be null"); + ParameterValidator.validateOperationName(name); + ParameterValidator.validateDuration(duration, "Wait duration"); + return context.reserve(name).waitAsync(WAIT.getValue(), duration); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WithRetryExtension.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWithRetryOperation.java similarity index 50% rename from sdk/src/main/java/software/amazon/lambda/durable/context/extension/WithRetryExtension.java rename to sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWithRetryOperation.java index 2fba9ba1b..b4f68d790 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WithRetryExtension.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWithRetryOperation.java @@ -1,33 +1,52 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.context.extension; +package software.amazon.lambda.durable.operation; import java.time.Duration; import java.util.Objects; import java.util.function.BiFunction; +import java.util.function.Supplier; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.WithRetryContext; import software.amazon.lambda.durable.config.RunInChildContextConfig; -import software.amazon.lambda.durable.config.WithRetryConfig; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.extension.ExtensionContextConfig; import software.amazon.lambda.durable.extension.ExtensionContextResult; import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.retry.RetryStrategies; +import software.amazon.lambda.durable.retry.RetryStrategy; -/** Canonical implementation of the built-in with-retry extension. */ -public final class WithRetryExtension { +/** Context-free static facade and canonical implementation of replay-safe retry operations. */ +public final class DurableWithRetryOperation { private static final Duration DEFAULT_BACKOFF_DELAY = Duration.ofSeconds(1); private static final String BACKOFF_SUFFIX = "-backoff-"; private static final String ANONYMOUS_CONTEXT_NAME = "retry"; private static final String ANONYMOUS_BACKOFF_PREFIX = "retry-backoff-"; - private WithRetryExtension() {} + private DurableWithRetryOperation() {} + + public static T withRetry(String name, Supplier operation) { + return withRetryAsync(name, operation).get(); + } + + public static T withRetry(String name, Supplier operation, WithRetryConfig config) { + return withRetryAsync(name, operation, config).get(); + } + + public static DurableFuture withRetryAsync(String name, Supplier operation) { + return withRetryAsync(name, operation, WithRetryConfig.builder().build()); + } + + public static DurableFuture withRetryAsync(String name, Supplier operation, WithRetryConfig config) { + return withRetryAsync(ExtensionContext.getCurrentContext(), name, adapt(operation), config); + } @SuppressWarnings("unchecked") - public static DurableFuture execute( + public static DurableFuture withRetryAsync( ExtensionContext context, String name, BiFunction operation, @@ -50,6 +69,15 @@ public static DurableFuture execute( return (DurableFuture) future; } + private static BiFunction adapt(Supplier operation) { + Objects.requireNonNull(operation, "operation cannot be null"); + return (attempt, ignored) -> { + try (var scope = WithRetryContext.attach(attempt)) { + return operation.get(); + } + }; + } + private static T executeRetryLoop( String name, BiFunction operation, WithRetryConfig config) { var durableContext = DurableContext.getCurrentContext(); @@ -66,7 +94,10 @@ private static T executeRetryLoop( throw e; } var delay = decision.delay().isZero() ? DEFAULT_BACKOFF_DELAY : decision.delay(); - extensionContext.reserve(backoffName(name, attempt)).wait(delay); + extensionContext + .reserve(backoffName(name, attempt)) + .waitAsync(OperationSubType.WAIT.getValue(), delay) + .get(); attempt++; } } @@ -75,4 +106,53 @@ private static T executeRetryLoop( private static String backoffName(String name, int attempt) { return name != null ? name + BACKOFF_SUFFIX + attempt : ANONYMOUS_BACKOFF_PREFIX + attempt; } + + /** Configuration for replay-safe retry operations. */ + public static final class WithRetryConfig { + private final RetryStrategy retryStrategy; + private final boolean wrapInChildContext; + + private WithRetryConfig(Builder builder) { + retryStrategy = Objects.requireNonNullElse(builder.retryStrategy, RetryStrategies.Presets.DEFAULT); + wrapInChildContext = builder.wrapInChildContext; + } + + public RetryStrategy retryStrategy() { + return retryStrategy; + } + + public boolean wrapInChildContext() { + return wrapInChildContext; + } + + public static Builder builder() { + return new Builder(); + } + + public Builder toBuilder() { + return new Builder().retryStrategy(retryStrategy).wrapInChildContext(wrapInChildContext); + } + + /** Builder for {@link WithRetryConfig}. */ + public static final class Builder { + private RetryStrategy retryStrategy; + private boolean wrapInChildContext; + + private Builder() {} + + public Builder retryStrategy(RetryStrategy retryStrategy) { + this.retryStrategy = retryStrategy; + return this; + } + + public Builder wrapInChildContext(boolean wrapInChildContext) { + this.wrapInChildContext = wrapInChildContext; + return this; + } + + public WithRetryConfig build() { + return new WithRetryConfig(this); + } + } + } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinator.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/OperationConcurrencyCoordinator.java similarity index 98% rename from sdk/src/main/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinator.java rename to sdk/src/main/java/software/amazon/lambda/durable/operation/OperationConcurrencyCoordinator.java index 8b4f82399..264fb8af3 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinator.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/OperationConcurrencyCoordinator.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.context.extension; +package software.amazon.lambda.durable.operation; import java.util.ArrayDeque; import java.util.ArrayList; @@ -17,7 +17,7 @@ import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.SuspendExecutionException; -final class ExtensionConcurrencyCoordinator { +final class OperationConcurrencyCoordinator { enum ItemStatus { PENDING, RUNNING, @@ -72,7 +72,7 @@ ItemStatus status() { private int succeeded; private int failed; - ExtensionConcurrencyCoordinator(int maxConcurrency, CompletionConfig completionConfig) { + OperationConcurrencyCoordinator(int maxConcurrency, CompletionConfig completionConfig) { if (maxConcurrency < 1) { throw new IllegalArgumentException("maxConcurrency must be at least 1"); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/OperationConfigAdapters.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/OperationConfigAdapters.java new file mode 100644 index 000000000..9128f1995 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/OperationConfigAdapters.java @@ -0,0 +1,16 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.operation; + +import software.amazon.lambda.durable.config.RunInChildContextConfig; + +final class OperationConfigAdapters { + private OperationConfigAdapters() {} + + static RunInChildContextConfig toLegacy(DurableContextOperation.RunInChildContextConfig config) { + return RunInChildContextConfig.builder() + .serDes(config.serDes()) + .isVirtual(config.isVirtual()) + .build(); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/ParallelExtensionFuture.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperationFuture.java similarity index 85% rename from sdk/src/main/java/software/amazon/lambda/durable/context/extension/ParallelExtensionFuture.java rename to sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperationFuture.java index 778fa11d2..d42e70311 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/ParallelExtensionFuture.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperationFuture.java @@ -1,12 +1,12 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.context.extension; +package software.amazon.lambda.durable.operation; import static software.amazon.lambda.durable.config.NestingType.FLAT; -import static software.amazon.lambda.durable.context.extension.ExtensionConcurrencyCoordinator.ItemStatus.FAILED; -import static software.amazon.lambda.durable.context.extension.ExtensionConcurrencyCoordinator.ItemStatus.SKIPPED; import static software.amazon.lambda.durable.model.OperationSubType.PARALLEL; import static software.amazon.lambda.durable.model.OperationSubType.PARALLEL_BRANCH; +import static software.amazon.lambda.durable.operation.OperationConcurrencyCoordinator.ItemStatus.FAILED; +import static software.amazon.lambda.durable.operation.OperationConcurrencyCoordinator.ItemStatus.SKIPPED; import java.util.ArrayList; import java.util.List; @@ -18,29 +18,31 @@ import software.amazon.lambda.durable.ParallelDurableFuture; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.CompletionConfig; -import software.amazon.lambda.durable.config.ParallelBranchConfig; -import software.amazon.lambda.durable.config.ParallelConfig; import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.extension.ExtensionContextConfig; import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; import software.amazon.lambda.durable.extension.ExtensionContextResult; import software.amazon.lambda.durable.model.ParallelResult; +import software.amazon.lambda.durable.operation.DurableParallelOperation.ParallelBranchConfig; +import software.amazon.lambda.durable.operation.DurableParallelOperation.ParallelConfig; import software.amazon.lambda.durable.serde.SerDes; import software.amazon.lambda.durable.util.ParameterValidator; -final class ParallelExtensionFuture implements ParallelDurableFuture { +final class ParallelOperationFuture implements ParallelDurableFuture { + private static final int LARGE_RESULT_THRESHOLD = 256 * 1024; + private final Object lock = new Object(); private final ParallelConfig config; private final SerDes defaultSerDes; private final List> branches = new ArrayList<>(); private final DurableFuture parentFuture; private ExtensionContext childContext; - private ExtensionConcurrencyCoordinator coordinator; + private OperationConcurrencyCoordinator coordinator; private ParallelResult replayState; private boolean registrationClosed; - ParallelExtensionFuture(ExtensionContext context, String name, ParallelConfig config) { + ParallelOperationFuture(ExtensionContext context, String name, ParallelConfig config) { this.config = config; this.defaultSerDes = context.getDurableConfig().getSerDes(); var parent = context.reserve(name); @@ -103,7 +105,7 @@ private void initializeCoordinator( if (replayContext.isReplayingChildren() && replayState == null) { throw new IllegalStateException("Missing result in completed Parallel operation"); } - coordinator = new ExtensionConcurrencyCoordinator(config.maxConcurrency(), config.completionConfig()); + coordinator = new OperationConcurrencyCoordinator(config.maxConcurrency(), config.completionConfig()); for (int index = 0; index < branches.size(); index++) { registerBranch(branches.get(index), index); } @@ -120,7 +122,10 @@ private void registerBranch(BranchDefinition definition, int index) { () -> reservation.runInChildContextAsync( PARALLEL_BRANCH.getValue(), definition.resultType, - () -> definition.function.apply(DurableContext.getCurrentContext()), + () -> ExtensionContextResult.replayChildrenAboveSize( + definition.function.apply(DurableContext.getCurrentContext()), + null, + LARGE_RESULT_THRESHOLD), branchConfig(definition.config)), skipped); definition.future.bind(item.future()); @@ -176,7 +181,7 @@ private ParallelResult rebuildResult(ParallelResult result) { } } - private static ParallelResult constructResult(ExtensionConcurrencyCoordinator.Completion completion) { + private static ParallelResult constructResult(OperationConcurrencyCoordinator.Completion completion) { var statuses = completion.items().stream() .map(item -> item.status() == FAILED ? ParallelResult.Status.FAILED @@ -197,16 +202,18 @@ private static ParallelResult constructResult(ExtensionConcurrencyCoordinator.Co statuses); } - private RunInChildContextConfig branchConfig(ParallelBranchConfig branchConfig) { - return RunInChildContextConfig.builder() - .serDes(branchConfig.serDes() == null ? defaultSerDes : branchConfig.serDes()) - .isVirtual(config.nestingType() == FLAT) + private ExtensionContextConfig branchConfig(ParallelBranchConfig branchConfig) { + return ExtensionContextConfig.builder() + .childContextConfig(RunInChildContextConfig.builder() + .serDes(branchConfig.serDes() == null ? defaultSerDes : branchConfig.serDes()) + .isVirtual(config.nestingType() == FLAT) + .build()) .build(); } - private static ExtensionConcurrencyCoordinator.ExpectedCompletionStatus expectedCompletion( + private static OperationConcurrencyCoordinator.ExpectedCompletionStatus expectedCompletion( ParallelResult replayState) { - return new ExtensionConcurrencyCoordinator.ExpectedCompletionStatus( + return new OperationConcurrencyCoordinator.ExpectedCompletionStatus( replayState.succeeded() + replayState.failed(), CompletionConfig.CompletionDecision.complete(replayState.completionStatus())); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForConditionFuture.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionFuture.java similarity index 94% rename from sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForConditionFuture.java rename to sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionFuture.java index 9cc0036b1..c3848e74f 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/extension/WaitForConditionFuture.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionFuture.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.context.extension; +package software.amazon.lambda.durable.operation; import java.util.Objects; import java.util.concurrent.CompletableFuture; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginInfoConverter.java b/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginInfoConverter.java index 5af7f3a2a..a3d9e44a5 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginInfoConverter.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/plugin/PluginInfoConverter.java @@ -6,9 +6,8 @@ import java.util.Collection; import java.util.stream.Collectors; import software.amazon.awssdk.services.lambda.model.Operation; -import software.amazon.lambda.durable.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; -import software.amazon.lambda.durable.operation.BaseDurableOperation; +import software.amazon.lambda.durable.primitive.BasePrimitive; /** * Utility methods for converting SDK internal types to plugin info records. @@ -29,23 +28,11 @@ private PluginInfoConverter() {} * @return an OperationInfo record */ public static OperationInfo toOperationInfo(Operation operation, OperationIdentifier identifier, String parentId) { - return toOperationInfo(operation, OperationDescriptor.from(identifier), parentId); - } - - /** - * Converts an SDK {@link Operation} to an {@link OperationInfo} using an {@link OperationDescriptor}. - * - * @param operation the SDK operation (may be null for first-start scenarios) - * @param descriptor the operation identity - * @param parentId the parent operation ID (may be null for root operations) - * @return an OperationInfo record - */ - public static OperationInfo toOperationInfo(Operation operation, OperationDescriptor descriptor, String parentId) { return new OperationInfo( - descriptor.operationId(), - descriptor.name(), - descriptor.operationType().toString(), - descriptor.subType(), + identifier.operationId(), + identifier.name(), + identifier.operationType().toString(), + identifier.subType(), parentId, operation != null ? operation.startTimestamp() : Instant.now(), operation != null ? operation.endTimestamp() : null, @@ -67,20 +54,11 @@ public static OperationInfo toOperationInfo(Operation operation, OperationDescri */ public static OperationEndInfo toOperationEndInfo( Operation operation, OperationIdentifier identifier, String parentId, boolean isReplay, Throwable error) { - return toOperationEndInfo(operation, OperationDescriptor.from(identifier), parentId, isReplay, error); - } - - /** - * Creates an {@link OperationEndInfo} from an SDK {@link Operation}, an {@link OperationDescriptor}, and an - * optional error. - */ - public static OperationEndInfo toOperationEndInfo( - Operation operation, OperationDescriptor descriptor, String parentId, boolean isReplay, Throwable error) { return new OperationEndInfo( - descriptor.operationId(), - descriptor.name(), - descriptor.operationType().toString(), - descriptor.subType(), + identifier.operationId(), + identifier.name(), + identifier.operationType().toString(), + identifier.subType(), parentId, operation != null ? operation.startTimestamp() : null, operation != null ? operation.endTimestamp() : null, @@ -105,17 +83,11 @@ public static OperationEndInfo toOperationEndInfo( */ public static UserFunctionStartInfo toUserFunctionStartInfo( OperationIdentifier identifier, String parentId, boolean isReplayingChildren, Integer attempt) { - return toUserFunctionStartInfo(OperationDescriptor.from(identifier), parentId, isReplayingChildren, attempt); - } - - /** Creates a user-function start record from an operation descriptor. */ - public static UserFunctionStartInfo toUserFunctionStartInfo( - OperationDescriptor descriptor, String parentId, boolean isReplayingChildren, Integer attempt) { return new UserFunctionStartInfo( - descriptor.operationId(), - descriptor.name(), - descriptor.operationType().toString(), - descriptor.subType(), + identifier.operationId(), + identifier.name(), + identifier.operationType().toString(), + identifier.subType(), parentId, Instant.now(), isReplayingChildren, @@ -181,7 +153,7 @@ private static OperationChangeItemInfo toOperationChangeItemInfo(Operation opera operation.parentId(), operation.startTimestamp(), operation.endTimestamp(), - BaseDurableOperation.extractErrorFromOperation(operation), + BasePrimitive.extractErrorFromOperation(operation), operation.status()); } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/primitive/BasePrimitive.java similarity index 93% rename from sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java rename to sdk/src/main/java/software/amazon/lambda/durable/primitive/BasePrimitive.java index fee1953f9..abed04ccf 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/primitive/BasePrimitive.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; +package software.amazon.lambda.durable.primitive; import java.time.Instant; import java.util.List; @@ -25,7 +25,6 @@ import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.execution.ThreadContext; import software.amazon.lambda.durable.execution.ThreadType; -import software.amazon.lambda.durable.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.plugin.PluginInfoConverter; @@ -50,30 +49,21 @@ *

  • Proper thread coordination via future * */ -public abstract class BaseDurableOperation { - private static final Logger logger = LoggerFactory.getLogger(BaseDurableOperation.class); +public abstract class BasePrimitive { + private static final Logger logger = LoggerFactory.getLogger(BasePrimitive.class); - private final OperationDescriptor operationDescriptor; + private final OperationIdentifier operationIdentifier; protected final ExecutionManager executionManager; - protected final CompletableFuture completionFuture; - protected final BaseDurableOperation parentOperation; + protected final CompletableFuture completionFuture; + protected final BasePrimitive parentOperation; protected final boolean isVirtual; protected final AtomicBoolean replayCompletedOperation = new AtomicBoolean(false); private final DurableContextImpl durableContext; private final AtomicReference> runningUserHandler = new AtomicReference<>(null); - protected BaseDurableOperation( - OperationIdentifier operationIdentifier, - DurableContextImpl durableContext, - BaseDurableOperation parentOperation) { - this(OperationDescriptor.from(operationIdentifier), durableContext, parentOperation, false); - } - - protected BaseDurableOperation( - OperationDescriptor operationDescriptor, - DurableContextImpl durableContext, - BaseDurableOperation parentOperation) { - this(operationDescriptor, durableContext, parentOperation, false); + protected BasePrimitive( + OperationIdentifier operationIdentifier, DurableContextImpl durableContext, BasePrimitive parentOperation) { + this(operationIdentifier, durableContext, parentOperation, false); } /** @@ -84,20 +74,12 @@ protected BaseDurableOperation( * @param parentOperation the operation that owns late-checkpoint suppression, if any * @param isVirtual whether this is a virtual operation that should not be persisted */ - protected BaseDurableOperation( + protected BasePrimitive( OperationIdentifier operationIdentifier, DurableContextImpl durableContext, - BaseDurableOperation parentOperation, - boolean isVirtual) { - this(OperationDescriptor.from(operationIdentifier), durableContext, parentOperation, isVirtual); - } - - protected BaseDurableOperation( - OperationDescriptor operationDescriptor, - DurableContextImpl durableContext, - BaseDurableOperation parentOperation, + BasePrimitive parentOperation, boolean isVirtual) { - this.operationDescriptor = operationDescriptor; + this.operationIdentifier = operationIdentifier; this.parentOperation = parentOperation; this.durableContext = durableContext; this.executionManager = durableContext.getExecutionManager(); @@ -109,7 +91,7 @@ protected BaseDurableOperation( executionManager.registerOperation(this); } - public CompletableFuture getCompletionFuture() { + public CompletableFuture getCompletionFuture() { return completionFuture; } @@ -124,22 +106,22 @@ public CompletableFuture completionFuture() { /** Gets the operation sub-type (e.g. RUN_IN_CHILD_CONTEXT, WAIT_FOR_CALLBACK). */ public OperationSubType getSubType() { - return operationDescriptor.standardSubType(); + return operationIdentifier.standardSubType(); } /** Gets the exact operation subtype string. */ public String getSubTypeValue() { - return operationDescriptor.subType(); + return operationIdentifier.subType(); } /** Gets the unique identifier for this operation. */ public String getOperationId() { - return operationDescriptor.operationId(); + return operationIdentifier.operationId(); } /** Gets the operation name (may be null). */ public String getName() { - return operationDescriptor.name(); + return operationIdentifier.name(); } /** Gets the parent context. */ @@ -149,7 +131,7 @@ protected DurableContextImpl getContext() { /** Gets the operation type. */ public OperationType getType() { - return operationDescriptor.operationType(); + return operationIdentifier.operationType(); } /** @@ -379,7 +361,7 @@ protected void runUserHandler(Runnable runnable, ThreadType threadType) { protected T runUserFunction(Integer attempt, Supplier userFunction) { var pluginRunner = getPluginRunner(); var startInfo = PluginInfoConverter.toUserFunctionStartInfo( - operationDescriptor, durableContext.getParentId(), durableContext.isReplaying(), attempt); + operationIdentifier, durableContext.getParentId(), durableContext.isReplaying(), attempt); pluginRunner.onUserFunctionStart(startInfo); try { T result = userFunction.get(); @@ -564,14 +546,14 @@ private PluginRunner getPluginRunner() { /** Fires onOperationStart plugin hook. */ private void fireOnOperationStart(Operation existing) { - var info = PluginInfoConverter.toOperationInfo(existing, operationDescriptor, durableContext.getParentId()); + var info = PluginInfoConverter.toOperationInfo(existing, operationIdentifier, durableContext.getParentId()); getPluginRunner().onOperationStart(info); } /** Fires onOperationEnd plugin hook when an operation reaches terminal status. */ protected void fireOnOperationEnd(Operation operation, Throwable error, boolean isReplay) { var info = PluginInfoConverter.toOperationEndInfo( - operation, operationDescriptor, durableContext.getParentId(), isReplay, error); + operation, operationIdentifier, durableContext.getParentId(), isReplay, error); getPluginRunner().onOperationEnd(info); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/CallbackOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/primitive/CallbackPrimitive.java similarity index 82% rename from sdk/src/main/java/software/amazon/lambda/durable/operation/CallbackOperation.java rename to sdk/src/main/java/software/amazon/lambda/durable/primitive/CallbackPrimitive.java index b9dc73814..c4f359e69 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/CallbackOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/primitive/CallbackPrimitive.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; +package software.amazon.lambda.durable.primitive; import software.amazon.awssdk.services.lambda.model.CallbackOptions; import software.amazon.awssdk.services.lambda.model.Operation; @@ -8,38 +8,28 @@ import software.amazon.awssdk.services.lambda.model.OperationUpdate; import software.amazon.lambda.durable.DurableCallbackFuture; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.CallbackConfig; import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.exception.CallbackFailedException; import software.amazon.lambda.durable.exception.CallbackTimeoutException; -import software.amazon.lambda.durable.model.OperationDescriptor; +import software.amazon.lambda.durable.extension.ExtensionCallbackConfig; import software.amazon.lambda.durable.model.OperationIdentifier; /** Durable operation for creating and waiting on external callbacks. */ -public class CallbackOperation extends SerializableDurableOperation implements DurableCallbackFuture { +public class CallbackPrimitive extends SerializablePrimitive implements DurableCallbackFuture { - private final CallbackConfig config; + private final ExtensionCallbackConfig config; private String callbackId; - public CallbackOperation( + public CallbackPrimitive( OperationIdentifier operationIdentifier, TypeToken resultTypeToken, - CallbackConfig config, + ExtensionCallbackConfig config, DurableContextImpl durableContext) { super(operationIdentifier, resultTypeToken, config.serDes(), durableContext); this.config = config; } - public CallbackOperation( - OperationDescriptor operationDescriptor, - TypeToken resultTypeToken, - CallbackConfig config, - DurableContextImpl durableContext) { - super(operationDescriptor, resultTypeToken, config.serDes(), durableContext); - this.config = config; - } - public String callbackId() { return callbackId; } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/primitive/ChildContextPrimitive.java similarity index 91% rename from sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java rename to sdk/src/main/java/software/amazon/lambda/durable/primitive/ChildContextPrimitive.java index 0f51de20f..a60ea4d02 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/primitive/ChildContextPrimitive.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; +package software.amazon.lambda.durable.primitive; import static software.amazon.lambda.durable.execution.ExecutionManager.isTerminalStatus; @@ -42,7 +42,6 @@ import software.amazon.lambda.durable.extension.ExtensionContextResult; import software.amazon.lambda.durable.logging.DurableLogger; import software.amazon.lambda.durable.model.DeserializedOperationResult; -import software.amazon.lambda.durable.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.util.ExceptionHelper; @@ -55,7 +54,7 @@ * *

    When created with a parent operation, the child skips checkpointing if that parent has already completed. */ -public class ChildContextOperation extends SerializableDurableOperation { +public class ChildContextPrimitive extends SerializablePrimitive { private static final int LARGE_RESULT_THRESHOLD = 256 * 1024; @@ -67,7 +66,7 @@ public class ChildContextOperation extends SerializableDurableOperation { private final AtomicReference> cachedOperationResult = new AtomicReference<>(null); // child context for RunInChildContext - public ChildContextOperation( + public ChildContextPrimitive( OperationIdentifier operationIdentifier, Function function, TypeToken resultTypeToken, @@ -77,13 +76,13 @@ public ChildContextOperation( } // child context with a late-checkpoint owner - public ChildContextOperation( + public ChildContextPrimitive( OperationIdentifier operationIdentifier, Function function, TypeToken resultTypeToken, RunInChildContextConfig config, DurableContextImpl durableContext, - BaseDurableOperation parentOperation) { + BasePrimitive parentOperation) { super( operationIdentifier, resultTypeToken, @@ -96,52 +95,24 @@ public ChildContextOperation( this.extensionConfig = null; } - public ChildContextOperation( - OperationDescriptor operationDescriptor, - Function function, - TypeToken resultTypeToken, - RunInChildContextConfig config, - DurableContextImpl durableContext) { - this(operationDescriptor, function, resultTypeToken, config, durableContext, null); - } - - public ChildContextOperation( - OperationDescriptor operationDescriptor, - Function function, - TypeToken resultTypeToken, - RunInChildContextConfig config, - DurableContextImpl durableContext, - BaseDurableOperation parentOperation) { - super( - operationDescriptor, - resultTypeToken, - config.serDes(), - durableContext, - parentOperation, - config.isVirtual()); - this.function = function; - this.extensionFunction = null; - this.extensionConfig = null; - } - - public ChildContextOperation( - OperationDescriptor operationDescriptor, + public ChildContextPrimitive( + OperationIdentifier operationIdentifier, ExtensionContextFunction function, TypeToken resultTypeToken, ExtensionContextConfig config, DurableContextImpl durableContext) { - this(operationDescriptor, function, resultTypeToken, config, durableContext, null); + this(operationIdentifier, function, resultTypeToken, config, durableContext, null); } - public ChildContextOperation( - OperationDescriptor operationDescriptor, + public ChildContextPrimitive( + OperationIdentifier operationIdentifier, ExtensionContextFunction function, TypeToken resultTypeToken, ExtensionContextConfig config, DurableContextImpl durableContext, - BaseDurableOperation parentOperation) { + BasePrimitive parentOperation) { super( - operationDescriptor, + operationIdentifier, resultTypeToken, config.childContextConfig().serDes(), durableContext, diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/InvokeOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/primitive/InvokePrimitive.java similarity index 79% rename from sdk/src/main/java/software/amazon/lambda/durable/operation/InvokeOperation.java rename to sdk/src/main/java/software/amazon/lambda/durable/primitive/InvokePrimitive.java index 27a301b6a..0a5d09976 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/InvokeOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/primitive/InvokePrimitive.java @@ -1,19 +1,18 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; +package software.amazon.lambda.durable.primitive; import software.amazon.awssdk.services.lambda.model.ChainedInvokeOptions; import software.amazon.awssdk.services.lambda.model.Operation; import software.amazon.awssdk.services.lambda.model.OperationAction; import software.amazon.awssdk.services.lambda.model.OperationUpdate; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.InvokeConfig; import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.exception.InvokeException; import software.amazon.lambda.durable.exception.InvokeFailedException; import software.amazon.lambda.durable.exception.InvokeStoppedException; import software.amazon.lambda.durable.exception.InvokeTimedOutException; -import software.amazon.lambda.durable.model.OperationDescriptor; +import software.amazon.lambda.durable.extension.ExtensionInvokeConfig; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.serde.SerDes; @@ -23,18 +22,18 @@ * @param the result type from the invoked function * @param the payload type sent to the invoked function */ -public class InvokeOperation extends SerializableDurableOperation { +public class InvokePrimitive extends SerializablePrimitive { private final String functionName; private final I payload; - private final InvokeConfig invokeConfig; + private final ExtensionInvokeConfig invokeConfig; private final SerDes payloadSerDes; - public InvokeOperation( + public InvokePrimitive( OperationIdentifier operationIdentifier, String functionName, I payload, TypeToken resultTypeToken, - InvokeConfig config, + ExtensionInvokeConfig config, DurableContextImpl durableContext) { super(operationIdentifier, resultTypeToken, config.serDes(), durableContext); @@ -44,20 +43,6 @@ public InvokeOperation( this.payloadSerDes = config.payloadSerDes() != null ? config.payloadSerDes() : config.serDes(); } - public InvokeOperation( - OperationDescriptor operationDescriptor, - String functionName, - I payload, - TypeToken resultTypeToken, - InvokeConfig config, - DurableContextImpl durableContext) { - super(operationDescriptor, resultTypeToken, config.serDes(), durableContext); - this.functionName = functionName; - this.payload = payload; - this.invokeConfig = config; - this.payloadSerDes = config.payloadSerDes() != null ? config.payloadSerDes() : config.serDes(); - } - /** Starts the operation. */ @Override protected void start() { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/SerializableDurableOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/primitive/SerializablePrimitive.java similarity index 84% rename from sdk/src/main/java/software/amazon/lambda/durable/operation/SerializableDurableOperation.java rename to sdk/src/main/java/software/amazon/lambda/durable/primitive/SerializablePrimitive.java index 1ef424d39..5a318bbb1 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/SerializableDurableOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/primitive/SerializablePrimitive.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; +package software.amazon.lambda.durable.primitive; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -9,7 +9,6 @@ import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.exception.SerDesException; -import software.amazon.lambda.durable.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.serde.SerDes; import software.amazon.lambda.durable.util.ExceptionHelper; @@ -32,8 +31,8 @@ *

  • Proper thread coordination via future * */ -public abstract class SerializableDurableOperation extends BaseDurableOperation implements DurableFuture { - private static final Logger logger = LoggerFactory.getLogger(SerializableDurableOperation.class); +public abstract class SerializablePrimitive extends BasePrimitive implements DurableFuture { + private static final Logger logger = LoggerFactory.getLogger(SerializablePrimitive.class); protected record SerializedResult(String serialized, T deserialized) {} @@ -48,7 +47,7 @@ protected record SerializedResult(String serialized, T deserialized) {} * @param resultSerDes the serializer/deserializer for the result * @param durableContext the parent context this operation belongs to */ - protected SerializableDurableOperation( + protected SerializablePrimitive( OperationIdentifier operationIdentifier, TypeToken resultTypeToken, SerDes resultSerDes, @@ -56,14 +55,6 @@ protected SerializableDurableOperation( this(operationIdentifier, resultTypeToken, resultSerDes, durableContext, null, false); } - protected SerializableDurableOperation( - OperationDescriptor operationDescriptor, - TypeToken resultTypeToken, - SerDes resultSerDes, - DurableContextImpl durableContext) { - this(operationDescriptor, resultTypeToken, resultSerDes, durableContext, null, false); - } - /** * Constructs a new durable operation. * @@ -74,30 +65,18 @@ protected SerializableDurableOperation( * @param isVirtual whether this is a virtual operation that should not be persisted * @param parentOperation the operation that owns late-checkpoint suppression, if any */ - protected SerializableDurableOperation( + protected SerializablePrimitive( OperationIdentifier operationIdentifier, TypeToken resultTypeToken, SerDes resultSerDes, DurableContextImpl durableContext, - BaseDurableOperation parentOperation, + BasePrimitive parentOperation, boolean isVirtual) { super(operationIdentifier, durableContext, parentOperation, isVirtual); this.resultTypeToken = resultTypeToken; this.resultSerDes = resultSerDes; } - protected SerializableDurableOperation( - OperationDescriptor operationDescriptor, - TypeToken resultTypeToken, - SerDes resultSerDes, - DurableContextImpl durableContext, - BaseDurableOperation parentOperation, - boolean isVirtual) { - super(operationDescriptor, durableContext, parentOperation, isVirtual); - this.resultTypeToken = resultTypeToken; - this.resultSerDes = resultSerDes; - } - /** * Deserializes a result string into the operation's result type. * diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/primitive/StepPrimitive.java similarity index 53% rename from sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java rename to sdk/src/main/java/software/amazon/lambda/durable/primitive/StepPrimitive.java index 67b8f0bf2..165b009e0 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/primitive/StepPrimitive.java @@ -1,19 +1,14 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; +package software.amazon.lambda.durable.primitive; -import java.time.Instant; import java.util.concurrent.CompletableFuture; -import java.util.function.Function; -import software.amazon.awssdk.services.lambda.model.ErrorObject; import software.amazon.awssdk.services.lambda.model.Operation; import software.amazon.awssdk.services.lambda.model.OperationAction; import software.amazon.awssdk.services.lambda.model.OperationStatus; import software.amazon.awssdk.services.lambda.model.OperationUpdate; import software.amazon.awssdk.services.lambda.model.StepOptions; -import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.StepConfig; import software.amazon.lambda.durable.config.StepSemantics; import software.amazon.lambda.durable.context.BaseContextImpl; import software.amazon.lambda.durable.context.DurableContextImpl; @@ -27,7 +22,6 @@ import software.amazon.lambda.durable.extension.ExtensionStepFunction; import software.amazon.lambda.durable.extension.ExtensionStepResult; import software.amazon.lambda.durable.logging.DurableLogger; -import software.amazon.lambda.durable.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.util.ExceptionHelper; @@ -39,50 +33,19 @@ * * @param the result type of the step function */ -public class StepOperation extends SerializableDurableOperation { +public class StepPrimitive extends SerializablePrimitive { private static final Integer FIRST_ATTEMPT = 1; - private final Function function; - private final StepConfig config; private final ExtensionStepFunction extensionFunction; private final ExtensionStepConfig extensionConfig; - public StepOperation( + public StepPrimitive( OperationIdentifier operationIdentifier, - Function function, - TypeToken resultTypeToken, - StepConfig config, - DurableContextImpl durableContext) { - super(operationIdentifier, resultTypeToken, config.serDes(), durableContext); - - this.function = function; - this.config = config; - this.extensionFunction = null; - this.extensionConfig = null; - } - - public StepOperation( - OperationDescriptor operationDescriptor, - Function function, - TypeToken resultTypeToken, - StepConfig config, - DurableContextImpl durableContext) { - super(operationDescriptor, resultTypeToken, config.serDes(), durableContext); - this.function = function; - this.config = config; - this.extensionFunction = null; - this.extensionConfig = null; - } - - public StepOperation( - OperationDescriptor operationDescriptor, ExtensionStepFunction function, TypeToken resultTypeToken, ExtensionStepConfig config, DurableContextImpl durableContext) { - super(operationDescriptor, resultTypeToken, config.serDes(), durableContext); - this.function = null; - this.config = null; + super(operationIdentifier, resultTypeToken, config.serDes(), durableContext); this.extensionFunction = function; this.extensionConfig = config; } @@ -90,56 +53,24 @@ public StepOperation( /** Starts the operation. */ @Override protected void start() { - if (isExtensionStep()) { - executeExtensionStepLogic(extensionConfig.initialState(), FIRST_ATTEMPT); - } else { - executeStepLogic(FIRST_ATTEMPT); - } + executeExtensionStepLogic(extensionConfig.initialState(), FIRST_ATTEMPT); } /** Replays the operation. */ @Override protected void replay(Operation existing) { - if (isExtensionStep()) { - replayExtensionStep(existing); - return; - } - var attempt = existing.stepDetails() != null && existing.stepDetails().attempt() != null - ? existing.stepDetails().attempt() + 1 - : FIRST_ATTEMPT; switch (existing.status()) { case SUCCEEDED, FAILED -> markAlreadyCompleted(); + case PENDING -> pollReadyAndResumeExtensionStep(); case STARTED -> { if (isAtMostOnce()) { - // AT_MOST_ONCE: treat as interrupted, go through retry logic - handleStepFailure(new StepInterruptedException(existing), attempt); - } else { - // AT_LEAST_ONCE: re-execute the step - executeStepLogic(attempt); - } - } - // Step is pending retry - Start polling for PENDING -> READY transition - case PENDING -> { - if (existing.stepDetails() != null && existing.stepDetails().nextAttemptTimestamp() != null) { - pollReadyAndExecuteStepLogic(existing.stepDetails().nextAttemptTimestamp(), attempt); + handleExtensionStepFailure( + new StepInterruptedException(existing), extensionState(existing), nextAttempt(existing)); } else { - throw terminateExecutionWithIllegalDurableOperationException( - "Unexpected PENDING step without nextAttemptTimestamp: " + getOperationId()); + resumeExtensionStep(existing); } } - // Execute with current attempt - case READY -> executeStepLogic(attempt); - default -> - throw terminateExecutionWithIllegalDurableOperationException( - "Unexpected step status: " + existing.status()); - } - } - - private void replayExtensionStep(Operation existing) { - switch (existing.status()) { - case SUCCEEDED, FAILED -> markAlreadyCompleted(); - case PENDING -> pollReadyAndResumeExtensionStep(); - case STARTED, READY -> resumeExtensionStep(existing); + case READY -> resumeExtensionStep(existing); default -> throw terminateExecutionWithIllegalDurableOperationException( "Unexpected extension step status: " + existing.status()); @@ -147,12 +78,19 @@ private void replayExtensionStep(Operation existing) { } private void resumeExtensionStep(Operation existing) { + executeExtensionStepLogic(extensionState(existing), nextAttempt(existing)); + } + + private int nextAttempt(Operation existing) { + var details = existing.stepDetails(); + return details != null && details.attempt() != null ? details.attempt() + 1 : FIRST_ATTEMPT; + } + + private T extensionState(Operation existing) { var details = existing.stepDetails(); - var attempt = details != null && details.attempt() != null ? details.attempt() + 1 : FIRST_ATTEMPT; - var state = details != null && details.result() != null + return details != null && details.result() != null ? deserializeResult(details.result()) : extensionConfig.initialState(); - executeExtensionStepLogic(state, attempt); } private void pollReadyAndResumeExtensionStep() { @@ -163,40 +101,6 @@ private void pollReadyAndResumeExtensionStep() { .thenAccept(this::resumeExtensionStep); } - private void pollReadyAndExecuteStepLogic(Instant nextAttemptInstant, int attempt) { - pollForOperationUpdates(nextAttemptInstant) - .thenCompose(op -> op.status() == OperationStatus.READY - ? CompletableFuture.completedFuture(op) - : pollForOperationUpdates(nextAttemptInstant)) - .thenRun(() -> executeStepLogic(attempt)); - } - - private void executeStepLogic(int attempt) { - Runnable userHandler = () -> { - // use a try-with-resources to - // - add thread id/type to thread local when the step starts - // - clear logger properties when the step finishes - StepContext stepContext = getContext().createStepContext(getOperationId(), getName(), attempt); - try (var ignoredContext = BaseContextImpl.attachCurrentContext(stepContext); - var ignoredLogger = DurableLogger.attachContext()) { - try { - checkpointStarted(); - - // Execute the user function inside the plugin hook boundary so a failure is reported - // through onUserFunctionEnd; retry/checkpoint handling stays outside the boundary. - T result = runUserFunction(attempt, () -> function.apply(stepContext)); - - handleStepSucceeded(result); - } catch (Throwable e) { - handleStepFailure(e, attempt); - } - } - }; - - // Execute user provided step code in user-configured executor - runUserHandler(userHandler, ThreadType.STEP); - } - private void executeExtensionStepLogic(T state, int attempt) { Runnable userHandler = () -> { var stepContext = getContext().createStepContext(getOperationId(), getName(), attempt); @@ -207,7 +111,7 @@ private void executeExtensionStepLogic(T state, int attempt) { var result = runUserFunction(attempt, () -> extensionFunction.apply(state)); handleExtensionStepResult(result, attempt); } catch (Throwable e) { - handleExtensionStepFailure(e); + handleExtensionStepFailure(e, state, attempt); } } }; @@ -243,7 +147,7 @@ private void pollReadyAndExecuteExtensionStep(T state, int attempt) { .thenRun(() -> executeExtensionStepLogic(state, attempt)); } - private void handleExtensionStepFailure(Throwable exception) { + private void handleExtensionStepFailure(Throwable exception, T state, int attempt) { exception = ExceptionHelper.unwrapCompletableFuture(exception); if (exception instanceof SuspendExecutionException suspendExecutionException) { throw suspendExecutionException; @@ -254,6 +158,24 @@ private void handleExtensionStepFailure(Throwable exception) { var error = exception instanceof DurableOperationException durableOperationException ? durableOperationException.getErrorObject() : serializeException(exception); + + var retryStrategy = extensionConfig.retryStrategy(); + if (retryStrategy != null) { + var decision = retryStrategy.makeRetryDecision(exception, attempt); + if (decision.shouldRetry()) { + var serializedState = serializeAndDeserializeResult(state); + var retryDelaySeconds = Math.toIntExact(decision.delay().toSeconds()); + sendOperationUpdate(OperationUpdate.builder() + .action(OperationAction.RETRY) + .payload(serializedState.serialized()) + .error(error) + .stepOptions(StepOptions.builder() + .nextAttemptDelaySeconds(retryDelaySeconds) + .build())); + pollReadyAndExecuteExtensionStep(serializedState.deserialized(), attempt + 1); + return; + } + } sendOperationUpdate( OperationUpdate.builder().action(OperationAction.FAIL).error(error)); } @@ -286,48 +208,6 @@ private void handleStepSucceeded(T result) { sendOperationUpdate(successUpdate); } - private void handleStepFailure(Throwable exception, int attempt) { - exception = ExceptionHelper.unwrapCompletableFuture(exception); - if (exception instanceof SuspendExecutionException suspendExecutionException) { - throw suspendExecutionException; - } - if (exception instanceof UnrecoverableDurableExecutionException unrecoverableDurableExecutionException) { - // terminate the execution and throw the exception if it's not recoverable - throw terminateExecution(unrecoverableDurableExecutionException); - } - - final ErrorObject errorObject; - if (exception instanceof DurableOperationException durableOperationException) { - errorObject = durableOperationException.getErrorObject(); - } else { - errorObject = serializeException(exception); - } - - var retryDecision = config.retryStrategy().makeRetryDecision(exception, attempt); - - if (retryDecision.shouldRetry()) { - // Send RETRY - var retryDelayInSeconds = Math.toIntExact(retryDecision.delay().toSeconds()); - var retryUpdate = OperationUpdate.builder() - .action(OperationAction.RETRY) - .error(errorObject) - .stepOptions(StepOptions.builder() - // RetryDecisions always produce integer number of seconds greater or equals to - // 1 (no sub-second numbers) - .nextAttemptDelaySeconds(retryDelayInSeconds) - .build()); - sendOperationUpdate(retryUpdate); - - // Poll for READY status and then execute the step again - pollReadyAndExecuteStepLogic(Instant.now().plusSeconds(retryDelayInSeconds), attempt + 1); - } else { - // Send FAIL - retries exhausted - var failUpdate = - OperationUpdate.builder().action(OperationAction.FAIL).error(errorObject); - sendOperationUpdate(failUpdate); - } - } - @Override public T get() { var op = waitForOperationCompletion(); @@ -356,10 +236,6 @@ public T get() { } private boolean isAtMostOnce() { - return config != null && config.semanticsPerRetry() == StepSemantics.AT_MOST_ONCE_PER_RETRY; - } - - private boolean isExtensionStep() { - return extensionFunction != null; + return extensionConfig.semanticsPerRetry() == StepSemantics.AT_MOST_ONCE_PER_RETRY; } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/primitive/WaitPrimitive.java similarity index 84% rename from sdk/src/main/java/software/amazon/lambda/durable/operation/WaitOperation.java rename to sdk/src/main/java/software/amazon/lambda/durable/primitive/WaitPrimitive.java index 1cd0ac040..930968123 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/primitive/WaitPrimitive.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; +package software.amazon.lambda.durable.primitive; import java.time.Duration; import java.time.Instant; @@ -13,7 +13,6 @@ import software.amazon.awssdk.services.lambda.model.WaitOptions; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.context.DurableContextImpl; -import software.amazon.lambda.durable.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; /** @@ -22,24 +21,18 @@ *

    The wait is checkpointed and the Lambda is suspended. On re-invocation after the wait period, execution resumes * from where it left off. */ -public class WaitOperation extends BaseDurableOperation implements DurableFuture { +public class WaitPrimitive extends BasePrimitive implements DurableFuture { - private static final Logger logger = LoggerFactory.getLogger(WaitOperation.class); + private static final Logger logger = LoggerFactory.getLogger(WaitPrimitive.class); private final Duration duration; - public WaitOperation( + public WaitPrimitive( OperationIdentifier operationIdentifier, Duration duration, DurableContextImpl durableContext) { super(operationIdentifier, durableContext, null); this.duration = duration; } - public WaitOperation( - OperationDescriptor operationDescriptor, Duration duration, DurableContextImpl durableContext) { - super(operationDescriptor, durableContext, null); - this.duration = duration; - } - /** Starts the operation. */ @Override protected void start() { diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableCoreOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableCoreOperationsTest.java deleted file mode 100644 index 423aa98b6..000000000 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableCoreOperationsTest.java +++ /dev/null @@ -1,103 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -import java.time.Duration; -import java.util.function.Function; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.Test; -import org.mockito.ArgumentCaptor; -import software.amazon.lambda.durable.config.RunInChildContextConfig; -import software.amazon.lambda.durable.config.StepConfig; -import software.amazon.lambda.durable.context.BaseContextImpl; - -class DurableCoreOperationsTest { - @AfterEach - void clearContext() { - BaseContextImpl.setCurrentContext(null); - } - - @Test - void stepAcceptsContextFreeSupplier() { - var context = mock(DurableContext.class); - BaseContextImpl.setCurrentContext(context); - - DurableCoreOperations.step("step", String.class, () -> "result"); - - @SuppressWarnings("unchecked") - var function = (ArgumentCaptor>) - (ArgumentCaptor) ArgumentCaptor.forClass(Function.class); - verify(context).step(eq("step"), eq(String.class), function.capture()); - assertEquals("result", function.getValue().apply(mock(StepContext.class))); - } - - @Test - void childContextAcceptsContextFreeSupplier() { - var context = mock(DurableContext.class); - BaseContextImpl.setCurrentContext(context); - - DurableCoreOperations.runInChildContext("child", String.class, () -> "result"); - - @SuppressWarnings("unchecked") - var function = (ArgumentCaptor>) - (ArgumentCaptor) ArgumentCaptor.forClass(Function.class); - verify(context).runInChildContext(eq("child"), eq(String.class), function.capture()); - assertEquals("result", function.getValue().apply(mock(DurableContext.class))); - } - - @Test - void coreValueOperationsDelegateToCurrentContext() { - var context = mock(DurableContext.class); - BaseContextImpl.setCurrentContext(context); - var duration = Duration.ofSeconds(1); - var waitFuture = mockFuture(); - var invokeFuture = mockStringFuture(); - @SuppressWarnings("unchecked") - var callbackFuture = (DurableCallbackFuture) mock(DurableCallbackFuture.class); - when(context.waitAsync("wait", duration)).thenReturn(waitFuture); - when(context.invokeAsync("invoke", "function", "payload", String.class)).thenReturn(invokeFuture); - when(context.createCallback("callback", String.class)).thenReturn(callbackFuture); - - assertEquals(waitFuture, DurableCoreOperations.waitAsync("wait", duration)); - assertEquals(invokeFuture, DurableCoreOperations.invokeAsync("invoke", "function", "payload", String.class)); - assertEquals(callbackFuture, DurableCoreOperations.createCallback("callback", String.class)); - } - - @Test - void configuredSupplierOverloadsDelegateToCurrentContext() { - var context = mock(DurableContext.class); - BaseContextImpl.setCurrentContext(context); - var stepConfig = StepConfig.builder().build(); - var childConfig = RunInChildContextConfig.builder().build(); - - DurableCoreOperations.stepAsync("step", new TypeToken() {}, () -> "step", stepConfig); - DurableCoreOperations.runInChildContextAsync("child", new TypeToken() {}, () -> "child", childConfig); - - verify(context).stepAsync(eq("step"), any(TypeToken.class), any(Function.class), eq(stepConfig)); - verify(context).runInChildContextAsync(eq("child"), any(TypeToken.class), any(Function.class), eq(childConfig)); - } - - @Test - void coreOperationsFailOutsideDurableContext() { - assertThrows( - IllegalStateException.class, () -> DurableCoreOperations.step("step", String.class, () -> "result")); - } - - @SuppressWarnings("unchecked") - private DurableFuture mockFuture() { - return mock(DurableFuture.class); - } - - @SuppressWarnings("unchecked") - private DurableFuture mockStringFuture() { - return mock(DurableFuture.class); - } -} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java index 2ccc1bbea..db57b0471 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java @@ -11,7 +11,7 @@ import org.junit.jupiter.api.Test; import software.amazon.lambda.durable.context.BaseContextImpl; import software.amazon.lambda.durable.execution.ExecutionManager; -import software.amazon.lambda.durable.operation.SerializableDurableOperation; +import software.amazon.lambda.durable.primitive.SerializablePrimitive; class DurableFutureTest { @AfterEach @@ -71,7 +71,7 @@ void allOfSingleFutureReturnsSingleResult() { void allOfPropagatesException() { var op1 = mockOperation("first"); @SuppressWarnings("unchecked") - SerializableDurableOperation op2 = mock(SerializableDurableOperation.class); + SerializablePrimitive op2 = mock(SerializablePrimitive.class); when(op2.get()).thenThrow(new RuntimeException("Step failed")); assertThrows(RuntimeException.class, () -> DurableFuture.allOf(op1, op2)); @@ -106,8 +106,8 @@ void anyOfUsesExecutionManagerWhenCalledFromDurableContext() { } @SuppressWarnings("unchecked") - private SerializableDurableOperation mockOperation(T result) { - SerializableDurableOperation op = mock(SerializableDurableOperation.class); + private SerializablePrimitive mockOperation(T result) { + SerializablePrimitive op = mock(SerializablePrimitive.class); when(op.get()).thenReturn(result); return op; } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableHandlerTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableHandlerTest.java index 0502d1334..c8e5b48c6 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableHandlerTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableHandlerTest.java @@ -77,6 +77,15 @@ void testIndirectDurableHandlerInheritance() { assertEquals("indirect: test-input", result); } + @Test + void testHandlerCanOverrideContextFreeMethod() { + var handler = new ContextFreeHandler(); + + var result = handler.handleRequest("test-input", null); + + assertEquals("context-free: test-input", result); + } + // Test handler implementation private static class TestDurableHandler extends DurableHandler { @Override @@ -85,6 +94,13 @@ public String handleRequest(String input, DurableContext context) { } } + private static class ContextFreeHandler extends DurableHandler { + @Override + public String handleRequest(String input) { + return "context-free: " + input; + } + } + // Intermediate handler that forwards an explicit input type private abstract static class AbstractIndirectHandler extends DurableHandler { protected AbstractIndirectHandler() { diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationTest.java similarity index 87% rename from sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationsTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationTest.java index 23e279abb..47ce4c22c 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationsTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationTest.java @@ -13,11 +13,9 @@ import java.util.List; import java.util.concurrent.CompletableFuture; -import java.util.function.Supplier; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; -import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.context.BaseContextImpl; import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.extension.ExtensionContextConfig; @@ -26,8 +24,9 @@ import software.amazon.lambda.durable.extension.ExtensionOperation; import software.amazon.lambda.durable.model.MapResult; import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.operation.DurableMapOperation; -class DurableMapOperationsTest { +class DurableMapOperationTest { @AfterEach void clearContext() { BaseContextImpl.setCurrentContext(null); @@ -52,7 +51,7 @@ void mapExposesItemIndexThroughScopedContext() { .thenReturn(parentFuture); BaseContextImpl.setCurrentContext(context); - var actual = DurableMapOperations.map("map", List.of("value"), String.class, item -> { + var actual = DurableMapOperation.map("map", List.of("value"), String.class, item -> { assertEquals(0, MapItemContext.getCurrentContext().getIndex()); return item.toUpperCase(); }); @@ -72,8 +71,8 @@ void mapExposesItemIndexThroughScopedContext() { when(iteration.runInChildContextAsync( eq(OperationSubType.MAP_ITERATION.getValue()), eq(TypeToken.get(String.class)), - any(Supplier.class), - any(RunInChildContextConfig.class))) + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) .thenReturn(new CompletedFuture<>("VALUE")); try (var ignoredContext = BaseContextImpl.attachCurrentContext(context); var ignoredReplay = ExtensionContextReplayContext.attach(false, null)) { @@ -81,16 +80,16 @@ void mapExposesItemIndexThroughScopedContext() { } @SuppressWarnings("unchecked") - var itemFunction = - (ArgumentCaptor>) (ArgumentCaptor) ArgumentCaptor.forClass(Supplier.class); + var itemFunction = (ArgumentCaptor>) + (ArgumentCaptor) ArgumentCaptor.forClass(ExtensionContextFunction.class); verify(iteration) .runInChildContextAsync( eq(OperationSubType.MAP_ITERATION.getValue()), eq(TypeToken.get(String.class)), itemFunction.capture(), - any(RunInChildContextConfig.class)); + any(ExtensionContextConfig.class)); try (var ignored = BaseContextImpl.attachCurrentContext(context)) { - assertEquals("VALUE", itemFunction.getValue().get()); + assertEquals("VALUE", itemFunction.getValue().apply().result()); } assertThrows(IllegalStateException.class, MapItemContext::getCurrentContext); } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableOperationFacadeTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableOperationFacadeTest.java new file mode 100644 index 000000000..d1937cbee --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableOperationFacadeTest.java @@ -0,0 +1,384 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Answers.CALLS_REAL_METHODS; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.mockito.Mockito.withSettings; +import static software.amazon.lambda.durable.model.OperationSubType.CALLBACK; +import static software.amazon.lambda.durable.model.OperationSubType.CHAINED_INVOKE; +import static software.amazon.lambda.durable.model.OperationSubType.RUN_IN_CHILD_CONTEXT; +import static software.amazon.lambda.durable.model.OperationSubType.STEP; +import static software.amazon.lambda.durable.model.OperationSubType.WAIT; + +import java.time.Duration; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import software.amazon.lambda.durable.config.CallbackConfig; +import software.amazon.lambda.durable.config.InvokeConfig; +import software.amazon.lambda.durable.config.RunInChildContextConfig; +import software.amazon.lambda.durable.config.StepConfig; +import software.amazon.lambda.durable.context.BaseContextImpl; +import software.amazon.lambda.durable.extension.ExtensionCallbackConfig; +import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextFunction; +import software.amazon.lambda.durable.extension.ExtensionInvokeConfig; +import software.amazon.lambda.durable.extension.ExtensionOperation; +import software.amazon.lambda.durable.extension.ExtensionStepConfig; +import software.amazon.lambda.durable.extension.ExtensionStepFunction; +import software.amazon.lambda.durable.extension.ExtensionStepResult; +import software.amazon.lambda.durable.operation.DurableCallbackOperation; +import software.amazon.lambda.durable.operation.DurableContextOperation; +import software.amazon.lambda.durable.operation.DurableInvokeOperation; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableWaitOperation; +import software.amazon.lambda.durable.serde.SerDes; + +class DurableOperationFacadeTest { + private static final String PACKAGE_NAME = "software.amazon.lambda.durable."; + + @AfterEach + void clearContext() { + BaseContextImpl.setCurrentContext(null); + } + + @Test + void operationFacadesUseSingularClassNames() { + assertFacadeRenamed("DurableStepOperation", "DurableStepOperations"); + assertFacadeRenamed("DurableWaitOperation", "DurableWaitOperations"); + assertFacadeRenamed("DurableInvokeOperation", "DurableInvokeOperations"); + assertFacadeRenamed("DurableCallbackOperation", "DurableCallbackOperations"); + assertFacadeRenamed("DurableContextOperation", "DurableContextOperations"); + assertFacadeRenamed("DurableMapOperation", "DurableMapOperations"); + assertFacadeRenamed("DurableParallelOperation", "DurableParallelOperations"); + assertFacadeRenamed("DurableWaitForCallbackOperation", "DurableWaitForCallbackOperations"); + assertFacadeRenamed("DurableWaitForConditionOperation", "DurableWaitForConditionOperations"); + assertFacadeRenamed("DurableWithRetryOperation", "DurableWithRetryOperations"); + } + + @Test + void backendPrimitivesUsePrimitiveClassNames() { + assertClassMoved("primitive.BasePrimitive", "primitive.BaseDurableOperation"); + assertClassMoved("primitive.SerializablePrimitive", "primitive.SerializableDurableOperation"); + assertClassMoved("primitive.StepPrimitive", "primitive.StepOperation"); + assertClassMoved("primitive.WaitPrimitive", "primitive.WaitOperation"); + assertClassMoved("primitive.InvokePrimitive", "primitive.InvokeOperation"); + assertClassMoved("primitive.CallbackPrimitive", "primitive.CallbackOperation"); + assertClassMoved("primitive.ChildContextPrimitive", "primitive.ChildContextOperation"); + } + + @Test + void extensionOperationImplementationLivesWithExtensionSpi() { + assertClassMoved("extension.ExtensionOperationImpl", "context.ExtensionOperationImpl"); + } + + @Test + void operationFacadesAbsorbTheirExtensions() { + assertMergedOperation("DurableStepOperation", "StepExtension"); + assertMergedOperation("DurableWaitOperation", "WaitExtension"); + assertMergedOperation("DurableInvokeOperation", "InvokeExtension"); + assertMergedOperation("DurableCallbackOperation", "CallbackExtension"); + assertMergedOperation("DurableContextOperation", "ContextExtension"); + assertMergedOperation("DurableMapOperation", "MapExtension"); + assertMergedOperation("DurableParallelOperation", "ParallelExtension"); + assertMergedOperation("DurableWaitForCallbackOperation", "WaitForCallbackExtension"); + assertMergedOperation("DurableWaitForConditionOperation", "WaitForConditionExtension"); + assertMergedOperation("DurableWithRetryOperation", "WithRetryExtension"); + } + + @Test + void stepAcceptsContextFreeSupplier() { + var context = mockDurableContext(); + var reservation = mock(ExtensionOperation.class); + var future = mockStringFuture(); + BaseContextImpl.setCurrentContext(context); + when(((ExtensionContext) context).reserve("step")).thenReturn(reservation); + when(reservation.stepAsync( + eq(STEP.getValue()), + any(TypeToken.class), + any(ExtensionStepFunction.class), + any(ExtensionStepConfig.class))) + .thenReturn(future); + + assertSame(future, DurableStepOperation.stepAsync("step", String.class, () -> "result")); + + @SuppressWarnings("unchecked") + var function = (ArgumentCaptor>) + (ArgumentCaptor) ArgumentCaptor.forClass(ExtensionStepFunction.class); + verify(reservation) + .stepAsync( + eq(STEP.getValue()), any(TypeToken.class), function.capture(), any(ExtensionStepConfig.class)); + try (var ignored = BaseContextImpl.attachCurrentContext(mock(StepContext.class))) { + var result = assertInstanceOf( + ExtensionStepResult.Succeeded.class, function.getValue().apply(null)); + assertEquals("result", result.value()); + } + } + + @Test + void durableContextStepUsesPrimitiveExtension() { + var future = mockStringFuture(); + var reservation = mock(ExtensionOperation.class); + var context = mock( + DurableContext.class, + withSettings().extraInterfaces(ExtensionContext.class).defaultAnswer(CALLS_REAL_METHODS)); + when(((ExtensionContext) context).reserve("step")).thenReturn(reservation); + when(reservation.stepAsync( + eq(STEP.getValue()), + any(TypeToken.class), + any(ExtensionStepFunction.class), + any(ExtensionStepConfig.class))) + .thenReturn(future); + + var result = context.stepAsync( + "step", + TypeToken.get(String.class), + ignored -> "result", + StepConfig.builder().build()); + + assertSame(future, result); + } + + @Test + void durableContextWaitUsesPrimitiveExtension() { + var future = mockFuture(); + var reservation = mock(ExtensionOperation.class); + var context = mockDurableContext(); + var duration = Duration.ofSeconds(1); + when(((ExtensionContext) context).reserve("wait")).thenReturn(reservation); + when(reservation.waitAsync(WAIT.getValue(), duration)).thenReturn(future); + + var result = context.waitAsync("wait", duration); + + assertSame(future, result); + } + + @Test + void durableContextInvokeUsesPrimitiveExtension() { + var future = mockStringFuture(); + var reservation = mock(ExtensionOperation.class); + var context = mockDurableContext(); + var payloadSerDes = mock(SerDes.class); + var resultSerDes = mock(SerDes.class); + var config = InvokeConfig.builder() + .payloadSerDes(payloadSerDes) + .serDes(resultSerDes) + .tenantId("tenant") + .build(); + when(((ExtensionContext) context).reserve("invoke")).thenReturn(reservation); + when(reservation.invokeAsync( + eq(CHAINED_INVOKE.getValue()), eq("function"), eq("payload"), any(TypeToken.class), any())) + .thenReturn(future); + + var result = context.invokeAsync("invoke", "function", "payload", TypeToken.get(String.class), config); + + assertSame(future, result); + var extensionConfig = ArgumentCaptor.forClass(ExtensionInvokeConfig.class); + verify(reservation) + .invokeAsync( + eq(CHAINED_INVOKE.getValue()), + eq("function"), + eq("payload"), + eq(TypeToken.get(String.class)), + extensionConfig.capture()); + assertSame(payloadSerDes, extensionConfig.getValue().payloadSerDes()); + assertSame(resultSerDes, extensionConfig.getValue().serDes()); + assertEquals("tenant", extensionConfig.getValue().tenantId()); + } + + @Test + void durableContextCallbackUsesPrimitiveExtension() { + @SuppressWarnings("unchecked") + var future = (DurableCallbackFuture) mock(DurableCallbackFuture.class); + var reservation = mock(ExtensionOperation.class); + var context = mockDurableContext(); + var serDes = mock(SerDes.class); + var config = CallbackConfig.builder() + .timeout(Duration.ofMinutes(5)) + .heartbeatTimeout(Duration.ofMinutes(1)) + .serDes(serDes) + .build(); + when(((ExtensionContext) context).reserve("callback")).thenReturn(reservation); + when(reservation.createCallback(eq(CALLBACK.getValue()), any(TypeToken.class), any())) + .thenReturn(future); + + var result = context.createCallback("callback", TypeToken.get(String.class), config); + + assertSame(future, result); + var extensionConfig = ArgumentCaptor.forClass(ExtensionCallbackConfig.class); + verify(reservation) + .createCallback(eq(CALLBACK.getValue()), eq(TypeToken.get(String.class)), extensionConfig.capture()); + assertEquals(Duration.ofMinutes(5), extensionConfig.getValue().timeout()); + assertEquals(Duration.ofMinutes(1), extensionConfig.getValue().heartbeatTimeout()); + assertSame(serDes, extensionConfig.getValue().serDes()); + } + + @Test + void durableContextChildContextUsesPrimitiveExtension() { + var future = mockStringFuture(); + var reservation = mock(ExtensionOperation.class); + var context = mockDurableContext(); + when(((ExtensionContext) context).reserve("child")).thenReturn(reservation); + when(reservation.runInChildContextAsync( + eq(RUN_IN_CHILD_CONTEXT.getValue()), + any(TypeToken.class), + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) + .thenReturn(future); + + var result = context.runInChildContextAsync( + "child", + TypeToken.get(String.class), + ignored -> "result", + RunInChildContextConfig.builder().build()); + + assertSame(future, result); + } + + @Test + void childContextAcceptsContextFreeSupplier() { + var context = mockDurableContext(); + var reservation = mock(ExtensionOperation.class); + var future = mockStringFuture(); + BaseContextImpl.setCurrentContext(context); + when(((ExtensionContext) context).reserve("child")).thenReturn(reservation); + when(reservation.runInChildContextAsync( + eq(RUN_IN_CHILD_CONTEXT.getValue()), + any(TypeToken.class), + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) + .thenReturn(future); + + assertSame(future, DurableContextOperation.runInChildContextAsync("child", String.class, () -> "result")); + + @SuppressWarnings("unchecked") + var function = (ArgumentCaptor>) + (ArgumentCaptor) ArgumentCaptor.forClass(ExtensionContextFunction.class); + verify(reservation) + .runInChildContextAsync( + eq(RUN_IN_CHILD_CONTEXT.getValue()), + any(TypeToken.class), + function.capture(), + any(ExtensionContextConfig.class)); + try (var ignored = BaseContextImpl.attachCurrentContext(context)) { + assertEquals("result", function.getValue().apply().result()); + } + } + + @Test + void valueOperationsDelegateToCurrentContext() { + var context = mockDurableContext(); + BaseContextImpl.setCurrentContext(context); + var duration = Duration.ofSeconds(1); + var waitFuture = mockFuture(); + var invokeFuture = mockStringFuture(); + @SuppressWarnings("unchecked") + var callbackFuture = (DurableCallbackFuture) mock(DurableCallbackFuture.class); + var waitReservation = mock(ExtensionOperation.class); + var invokeReservation = mock(ExtensionOperation.class); + var callbackReservation = mock(ExtensionOperation.class); + when(((ExtensionContext) context).reserve("wait")).thenReturn(waitReservation); + when(((ExtensionContext) context).reserve("invoke")).thenReturn(invokeReservation); + when(((ExtensionContext) context).reserve("callback")).thenReturn(callbackReservation); + when(waitReservation.waitAsync(WAIT.getValue(), duration)).thenReturn(waitFuture); + when(invokeReservation.invokeAsync( + eq(CHAINED_INVOKE.getValue()), eq("function"), eq("payload"), any(TypeToken.class), any())) + .thenReturn(invokeFuture); + when(callbackReservation.createCallback(eq(CALLBACK.getValue()), any(TypeToken.class), any())) + .thenReturn(callbackFuture); + + assertEquals(waitFuture, DurableWaitOperation.waitAsync("wait", duration)); + assertEquals(invokeFuture, DurableInvokeOperation.invokeAsync("invoke", "function", "payload", String.class)); + assertEquals(callbackFuture, DurableCallbackOperation.createCallback("callback", String.class)); + } + + @Test + void configuredSupplierOverloadsDelegateToCurrentContext() { + var context = mockDurableContext(); + var stepReservation = mock(ExtensionOperation.class); + var childReservation = mock(ExtensionOperation.class); + BaseContextImpl.setCurrentContext(context); + var stepConfig = StepConfig.builder().build(); + var childConfig = RunInChildContextConfig.builder().build(); + when(((ExtensionContext) context).reserve("step")).thenReturn(stepReservation); + when(((ExtensionContext) context).reserve("child")).thenReturn(childReservation); + + DurableStepOperation.stepAsync( + "step", new TypeToken() {}, () -> "step", stepConfig.toOperationConfig()); + DurableContextOperation.runInChildContextAsync( + "child", new TypeToken() {}, () -> "child", childConfig.toOperationConfig()); + + verify(stepReservation) + .stepAsync( + eq(STEP.getValue()), + any(TypeToken.class), + any(ExtensionStepFunction.class), + any(ExtensionStepConfig.class)); + verify(childReservation) + .runInChildContextAsync( + eq(RUN_IN_CHILD_CONTEXT.getValue()), + any(TypeToken.class), + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class)); + } + + @Test + void primitiveOperationsFailOutsideDurableContext() { + assertThrows( + IllegalStateException.class, () -> DurableStepOperation.step("step", String.class, () -> "result")); + } + + @SuppressWarnings("unchecked") + private DurableFuture mockFuture() { + return mock(DurableFuture.class); + } + + @SuppressWarnings("unchecked") + private DurableFuture mockStringFuture() { + return mock(DurableFuture.class); + } + + private DurableContext mockDurableContext() { + return mock( + DurableContext.class, + withSettings().extraInterfaces(ExtensionContext.class).defaultAnswer(CALLS_REAL_METHODS)); + } + + private void assertFacadeRenamed(String singularName, String pluralName) { + assertTrue(classExists("operation." + singularName), singularName + " must be part of the public API"); + assertFalse(classExists(singularName), singularName + " must be removed from the root package"); + assertFalse(classExists(pluralName), pluralName + " must be removed from the public API"); + assertFalse(classExists("operation." + pluralName), pluralName + " must be removed from the operation package"); + } + + private void assertClassMoved(String newName, String oldName) { + assertTrue(classExists(newName), newName + " must exist"); + assertFalse(classExists(oldName), oldName + " must be removed"); + } + + private void assertMergedOperation(String operationName, String extensionName) { + assertClassMoved("operation." + operationName, operationName); + assertFalse(classExists("operation." + extensionName), extensionName + " must be merged into " + operationName); + } + + private boolean classExists(String simpleName) { + try { + Class.forName(PACKAGE_NAME + simpleName); + return true; + } catch (ClassNotFoundException exception) { + return false; + } + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationTest.java similarity index 92% rename from sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationsTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationTest.java index cd1278c84..3a665b8f9 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationsTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationTest.java @@ -17,8 +17,9 @@ import software.amazon.lambda.durable.extension.ExtensionContextConfig; import software.amazon.lambda.durable.extension.ExtensionContextFunction; import software.amazon.lambda.durable.extension.ExtensionOperation; +import software.amazon.lambda.durable.operation.DurableParallelOperation; -class DurableParallelOperationsTest { +class DurableParallelOperationTest { @AfterEach void clearContext() { BaseContextImpl.setCurrentContext(null); @@ -39,7 +40,7 @@ void parallelBranchesAcceptContextFreeSuppliers() { any(ExtensionContextConfig.class))) .thenReturn(parentFuture); - var result = DurableParallelOperations.parallel("parallel"); + var result = DurableParallelOperation.parallel("parallel"); result.branch("branch", String.class, () -> "result"); verify(context).reserve("parallel"); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationTest.java similarity index 69% rename from sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationsTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationTest.java index 677a5ef20..12395b6d5 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationsTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationTest.java @@ -10,20 +10,21 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import java.util.function.Supplier; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; -import software.amazon.lambda.durable.config.CallbackConfig; -import software.amazon.lambda.durable.config.StepConfig; import software.amazon.lambda.durable.context.BaseContextImpl; +import software.amazon.lambda.durable.extension.ExtensionCallbackConfig; import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.extension.ExtensionContextConfig; import software.amazon.lambda.durable.extension.ExtensionContextFunction; import software.amazon.lambda.durable.extension.ExtensionOperation; +import software.amazon.lambda.durable.extension.ExtensionStepConfig; +import software.amazon.lambda.durable.extension.ExtensionStepFunction; import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation; -class DurableWaitForCallbackOperationsTest { +class DurableWaitForCallbackOperationTest { @AfterEach void clearContext() { BaseContextImpl.setCurrentContext(null); @@ -46,7 +47,7 @@ void callbackSubmitterUsesRunnableAndScopedCallbackId() { assertEquals( "approved", - DurableWaitForCallbackOperations.waitForCallback( + DurableWaitForCallbackOperation.waitForCallback( "callback", String.class, () -> assertEquals( @@ -63,15 +64,26 @@ void callbackSubmitterUsesRunnableAndScopedCallbackId() { function.capture(), any(ExtensionContextConfig.class)); - var child = mock(ExtensionContext.class); + var child = mock(CurrentContext.class); var callbackReservation = mock(ExtensionOperation.class); var submitterReservation = mock(ExtensionOperation.class); @SuppressWarnings("unchecked") var callback = (DurableCallbackFuture) mock(DurableCallbackFuture.class); + @SuppressWarnings("unchecked") + var submitterFuture = (DurableFuture) mock(DurableFuture.class); when(child.reserve("callback-callback")).thenReturn(callbackReservation); when(child.reserve("callback-submitter")).thenReturn(submitterReservation); - when(callbackReservation.createCallback(eq(TypeToken.get(String.class)), any(CallbackConfig.class))) + when(callbackReservation.createCallback( + eq(OperationSubType.CALLBACK.getValue()), + eq(TypeToken.get(String.class)), + any(ExtensionCallbackConfig.class))) .thenReturn(callback); + when(submitterReservation.stepAsync( + eq(OperationSubType.STEP.getValue()), + eq(TypeToken.get(Void.class)), + any(ExtensionStepFunction.class), + any(ExtensionStepConfig.class))) + .thenReturn(submitterFuture); when(callback.callbackId()).thenReturn("callback-id"); when(callback.get()).thenReturn("approved"); @@ -80,10 +92,16 @@ void callbackSubmitterUsesRunnableAndScopedCallbackId() { } @SuppressWarnings("unchecked") - var submitter = (ArgumentCaptor>) (ArgumentCaptor) ArgumentCaptor.forClass(Supplier.class); - verify(submitterReservation).step(eq(Void.class), submitter.capture(), any(StepConfig.class)); + var submitter = (ArgumentCaptor>) + (ArgumentCaptor) ArgumentCaptor.forClass(ExtensionStepFunction.class); + verify(submitterReservation) + .stepAsync( + eq(OperationSubType.STEP.getValue()), + eq(TypeToken.get(Void.class)), + submitter.capture(), + any(ExtensionStepConfig.class)); try (var ignored = BaseContextImpl.attachCurrentContext(mock(StepContext.class))) { - submitter.getValue().get(); + submitter.getValue().apply(null); } assertThrows(IllegalStateException.class, WaitForCallbackContext::getCurrentContext); } @@ -92,4 +110,6 @@ void callbackSubmitterUsesRunnableAndScopedCallbackId() { private DurableFuture mockStringFuture() { return mock(DurableFuture.class); } + + private interface CurrentContext extends DurableContext, ExtensionContext {} } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForConditionOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForConditionOperationTest.java similarity index 92% rename from sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForConditionOperationsTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForConditionOperationTest.java index 40558124b..a3cd25c83 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForConditionOperationsTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForConditionOperationTest.java @@ -20,8 +20,9 @@ import software.amazon.lambda.durable.extension.ExtensionStepResult; import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.model.WaitForConditionResult; +import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation; -class DurableWaitForConditionOperationsTest { +class DurableWaitForConditionOperationTest { @AfterEach void clearContext() { BaseContextImpl.setCurrentContext(null); @@ -43,7 +44,7 @@ void conditionFunctionReceivesOnlyStateAndUsesStepContextFromTls() { .thenReturn(future); BaseContextImpl.setCurrentContext(context); - assertEquals("VALUE", DurableWaitForConditionOperations.waitForCondition("condition", String.class, state -> { + assertEquals("VALUE", DurableWaitForConditionOperation.waitForCondition("condition", String.class, state -> { assertEquals(stepContext, StepContext.getCurrentContext()); return WaitForConditionResult.stopPolling(state.toUpperCase()); })); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationsTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationTest.java similarity index 92% rename from sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationsTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationTest.java index 279a66994..ab13ced6d 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationsTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationTest.java @@ -19,8 +19,9 @@ import software.amazon.lambda.durable.extension.ExtensionContextFunction; import software.amazon.lambda.durable.extension.ExtensionOperation; import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.operation.DurableWithRetryOperation; -class DurableWithRetryOperationsTest { +class DurableWithRetryOperationTest { @AfterEach void clearContext() { BaseContextImpl.setCurrentContext(null); @@ -41,7 +42,7 @@ void retryBodyUsesSupplierAndScopedAttempt() { .thenReturn(future); BaseContextImpl.setCurrentContext(context); - assertEquals(1, DurableWithRetryOperations.withRetry("retry", () -> WithRetryContext.getCurrentContext() + assertEquals(1, DurableWithRetryOperation.withRetry("retry", () -> WithRetryContext.getCurrentContext() .getAttempt())); var function = extensionFunction(); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurationValidationIntegrationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurationValidationIntegrationTest.java index 2deda7107..225ae7f0d 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurationValidationIntegrationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurationValidationIntegrationTest.java @@ -7,6 +7,7 @@ import java.time.Duration; import org.junit.jupiter.api.Test; import software.amazon.lambda.durable.config.CallbackConfig; +import software.amazon.lambda.durable.extension.ExtensionCallbackConfig; class DurationValidationIntegrationTest { @@ -43,4 +44,22 @@ void callbackConfig_withNullTimeouts_shouldPass() { assertDoesNotThrow(() -> CallbackConfig.builder().timeout(null).heartbeatTimeout(null).build()); } + + @Test + void extensionCallbackConfig_withInvalidTimeout_shouldThrow() { + var exception = assertThrows(IllegalArgumentException.class, () -> ExtensionCallbackConfig.builder() + .timeout(Duration.ofMillis(500)) + .build()); + + assertTrue(exception.getMessage().contains("Callback timeout")); + assertTrue(exception.getMessage().contains("at least 1 second")); + } + + @Test + void extensionCallbackConfig_withValidTimeouts_shouldPass() { + assertDoesNotThrow(() -> ExtensionCallbackConfig.builder() + .timeout(Duration.ofSeconds(30)) + .heartbeatTimeout(Duration.ofSeconds(10)) + .build()); + } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java index 301d484da..f0e0e0489 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/context/ExtensionOperationImplTest.java @@ -3,280 +3,196 @@ package software.amazon.lambda.durable.context; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doCallRealMethod; -import static org.mockito.Mockito.inOrder; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import java.time.Duration; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.function.Function; +import java.util.concurrent.Executors; import org.junit.jupiter.api.Test; -import org.mockito.ArgumentCaptor; -import software.amazon.lambda.durable.DurableCallbackFuture; -import software.amazon.lambda.durable.DurableContext; -import software.amazon.lambda.durable.DurableFuture; -import software.amazon.lambda.durable.StepContext; +import software.amazon.awssdk.services.lambda.model.CallbackDetails; +import software.amazon.awssdk.services.lambda.model.ContextDetails; +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.awssdk.services.lambda.model.StepDetails; +import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.CallbackConfig; -import software.amazon.lambda.durable.config.InvokeConfig; -import software.amazon.lambda.durable.config.RunInChildContextConfig; -import software.amazon.lambda.durable.config.StepConfig; +import software.amazon.lambda.durable.execution.ExecutionManager; +import software.amazon.lambda.durable.extension.ExtensionCallbackConfig; import software.amazon.lambda.durable.extension.ExtensionContextConfig; import software.amazon.lambda.durable.extension.ExtensionContextResult; +import software.amazon.lambda.durable.extension.ExtensionInvokeConfig; import software.amazon.lambda.durable.extension.ExtensionOperation; +import software.amazon.lambda.durable.extension.ExtensionOperationImpl; import software.amazon.lambda.durable.extension.ExtensionStepConfig; import software.amazon.lambda.durable.extension.ExtensionStepResult; +import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.primitive.BasePrimitive; +import software.amazon.lambda.durable.primitive.CallbackPrimitive; +import software.amazon.lambda.durable.primitive.ChildContextPrimitive; +import software.amazon.lambda.durable.primitive.InvokePrimitive; +import software.amazon.lambda.durable.primitive.StepPrimitive; +import software.amazon.lambda.durable.primitive.WaitPrimitive; class ExtensionOperationImplTest { @Test void reservationsKeepSequentialIdsWhenExecutedOutOfOrder() { - var context = mock(DurableContextImpl.class); + var context = context(); when(context.reserveOperationId()).thenReturn("sequential-1", "sequential-2"); doCallRealMethod().when(context).reserve("first"); doCallRealMethod().when(context).reserve("second"); - var duration = Duration.ofSeconds(1); - when(context.waitAsyncWithId("sequential-1", "first", duration)).thenReturn(mockFuture()); - when(context.waitAsyncWithId("sequential-2", "second", duration)).thenReturn(mockFuture()); + replay(context, "sequential-1", "first", OperationType.WAIT, OperationSubType.WAIT.getValue()); + replay(context, "sequential-2", "second", OperationType.WAIT, OperationSubType.WAIT.getValue()); var first = context.reserve("first"); var second = context.reserve("second"); - second.waitAsync(duration); - first.waitAsync(duration); + var secondFuture = second.waitAsync(OperationSubType.WAIT.getValue(), Duration.ofSeconds(1)); + var firstFuture = first.waitAsync(OperationSubType.WAIT.getValue(), Duration.ofSeconds(1)); - var ordered = inOrder(context); - ordered.verify(context).reserveOperationId(); - ordered.verify(context).reserveOperationId(); - ordered.verify(context).waitAsyncWithId("sequential-2", "second", duration); - ordered.verify(context).waitAsyncWithId("sequential-1", "first", duration); + assertEquals( + "sequential-2", + assertInstanceOf(BasePrimitive.class, secondFuture).getOperationId()); + assertEquals( + "sequential-1", + assertInstanceOf(BasePrimitive.class, firstFuture).getOperationId()); } @Test void customReservationUsesExplicitLocalOperationId() { - var context = mock(DurableContextImpl.class); - var duration = Duration.ofSeconds(1); - var expectedFuture = mockFuture(); + var context = context(); when(context.reserveOperationId("node-a")).thenReturn("custom-node-a"); doCallRealMethod().when(context).reserve("custom", "node-a"); - when(context.waitAsyncWithId("custom-node-a", "custom", duration)).thenReturn(expectedFuture); - - var operation = context.reserve("custom", "node-a"); - var actualFuture = operation.waitAsync(duration); - - verify(context).reserveOperationId("node-a"); - verify(context).waitAsyncWithId("custom-node-a", "custom", duration); - assertEquals(expectedFuture, actualFuture); - } + replay(context, "custom-node-a", "custom", OperationType.WAIT, OperationSubType.WAIT.getValue()); - @Test - void reservedStepAdaptsSupplierToStepFunction() { - var context = mock(DurableContextImpl.class); - var future = mockStringFuture(); - var resultType = TypeToken.get(String.class); - var config = StepConfig.builder().build(); - var called = new AtomicBoolean(); - when(context.stepAsyncWithId(eq("1"), eq("step"), eq(resultType), any(), eq(config))) - .thenReturn(future); - var operation = new ExtensionOperationImpl(context, "1", "step"); + var future = + context.reserve("custom", "node-a").waitAsync(OperationSubType.WAIT.getValue(), Duration.ofSeconds(1)); assertEquals( - future, - operation.stepAsync( - resultType, - () -> { - called.set(true); - return "result"; - }, - config)); - - @SuppressWarnings("unchecked") - var function = (ArgumentCaptor>) - (ArgumentCaptor) ArgumentCaptor.forClass(Function.class); - verify(context).stepAsyncWithId(eq("1"), eq("step"), eq(resultType), function.capture(), eq(config)); - assertEquals("result", function.getValue().apply(mock(StepContext.class))); - assertEquals(true, called.get()); - } - - @Test - void customSubtypeStepDelegatesExactSubtype() { - var context = mock(DurableContextImpl.class); - var future = mockStringFuture(); - var resultType = TypeToken.get(String.class); - var config = StepConfig.builder().build(); - when(context.stepAsyncWithId(eq("1"), eq("step"), eq("AcmeStep"), eq(resultType), any(), eq(config))) - .thenReturn(future); - - var actual = new ExtensionOperationImpl(context, "1", "step") - .stepAsync("AcmeStep", resultType, () -> "result", config); - - assertEquals(future, actual); + "custom-node-a", assertInstanceOf(BasePrimitive.class, future).getOperationId()); } @Test - void statefulStepDelegatesWithoutExposingStepContext() { - var context = mock(DurableContextImpl.class); - var future = mockStringFuture(); + void statefulStepCreatesStepOperationWithExactSubtype() { + var context = context(); + replay(context, "1", "step", OperationType.STEP, "AcmeStateful"); var resultType = TypeToken.get(String.class); var config = ExtensionStepConfig.builder().initialState("initial").build(); - when(context.extensionStepAsyncWithId( - eq("1"), eq("step"), eq("AcmeStateful"), eq(resultType), any(), eq(config))) - .thenReturn(future); - var actual = new ExtensionOperationImpl(context, "1", "step") + var future = new ExtensionOperationImpl(context, "1", "step", null) .stepAsync("AcmeStateful", resultType, state -> ExtensionStepResult.succeed(state + "-done"), config); - assertEquals(future, actual); - } - - @Test - void reservationDelegatesWaitInvokeAndCallback() { - var duration = Duration.ofSeconds(2); - var waitContext = mock(DurableContextImpl.class); - var waitFuture = mockFuture(); - when(waitContext.waitAsyncWithId("1", "wait", duration)).thenReturn(waitFuture); - assertEquals(waitFuture, new ExtensionOperationImpl(waitContext, "1", "wait").waitAsync(duration)); - - var invokeContext = mock(DurableContextImpl.class); - var invokeFuture = mockStringFuture(); - var invokeConfig = InvokeConfig.builder().build(); - var resultType = TypeToken.get(String.class); - when(invokeContext.invokeAsyncWithId("2", "invoke", "target", "payload", resultType, invokeConfig)) - .thenReturn(invokeFuture); - assertEquals( - invokeFuture, - new ExtensionOperationImpl(invokeContext, "2", "invoke") - .invokeAsync("target", "payload", resultType, invokeConfig)); - - var callbackContext = mock(DurableContextImpl.class); - @SuppressWarnings("unchecked") - var callbackFuture = (DurableCallbackFuture) mock(DurableCallbackFuture.class); - var callbackConfig = CallbackConfig.builder().build(); - when(callbackContext.createCallbackWithId("3", "callback", resultType, callbackConfig)) - .thenReturn(callbackFuture); - assertEquals( - callbackFuture, - new ExtensionOperationImpl(callbackContext, "3", "callback") - .createCallback(resultType, callbackConfig)); + assertOperation(future, StepPrimitive.class, "1", "step", "AcmeStateful"); } @Test - void customSubtypeSelectorsDelegateExactSubtype() { - var duration = Duration.ofSeconds(2); + void primitiveSelectorsCreateMatchingOperationTypesAndExactSubtypes() { + var context = context(); var resultType = TypeToken.get(String.class); + replay(context, "1", "wait", OperationType.WAIT, "AcmeWait"); + replay(context, "2", "invoke", OperationType.CHAINED_INVOKE, "AcmeInvoke"); + replay(context, "3", "callback", OperationType.CALLBACK, "AcmeCallback"); + replay(context, "4", "child", OperationType.CONTEXT, "AcmeContext"); + + var wait = new ExtensionOperationImpl(context, "1", "wait", null).waitAsync("AcmeWait", Duration.ofSeconds(2)); + var invoke = new ExtensionOperationImpl(context, "2", "invoke", null) + .invokeAsync( + "AcmeInvoke", + "target", + "payload", + resultType, + ExtensionInvokeConfig.builder().build()); + var callback = new ExtensionOperationImpl(context, "3", "callback", null) + .createCallback( + "AcmeCallback", + resultType, + ExtensionCallbackConfig.builder().build()); + var child = new ExtensionOperationImpl(context, "4", "child", null) + .runInChildContextAsync( + "AcmeContext", + resultType, + () -> ExtensionContextResult.completed("result"), + ExtensionContextConfig.builder().build()); - var waitContext = mock(DurableContextImpl.class); - var waitFuture = mockFuture(); - when(waitContext.waitAsyncWithId("1", "wait", "AcmeWait", duration)).thenReturn(waitFuture); - assertEquals(waitFuture, new ExtensionOperationImpl(waitContext, "1", "wait").waitAsync("AcmeWait", duration)); - - var invokeContext = mock(DurableContextImpl.class); - var invokeFuture = mockStringFuture(); - var invokeConfig = InvokeConfig.builder().build(); - when(invokeContext.invokeAsyncWithId( - "2", "invoke", "AcmeInvoke", "target", "payload", resultType, invokeConfig)) - .thenReturn(invokeFuture); - assertEquals( - invokeFuture, - new ExtensionOperationImpl(invokeContext, "2", "invoke") - .invokeAsync("AcmeInvoke", "target", "payload", resultType, invokeConfig)); - - var callbackContext = mock(DurableContextImpl.class); - @SuppressWarnings("unchecked") - var callbackFuture = (DurableCallbackFuture) mock(DurableCallbackFuture.class); - var callbackConfig = CallbackConfig.builder().build(); - when(callbackContext.createCallbackWithId("3", "callback", "AcmeCallback", resultType, callbackConfig)) - .thenReturn(callbackFuture); - assertEquals( - callbackFuture, - new ExtensionOperationImpl(callbackContext, "3", "callback") - .createCallback("AcmeCallback", resultType, callbackConfig)); - - var childContext = mock(DurableContextImpl.class); - var childFuture = mockStringFuture(); - var childConfig = RunInChildContextConfig.builder().build(); - when(childContext.runInChildContextAsyncWithId( - eq("4"), eq("child"), eq("AcmeContext"), eq(resultType), any(), eq(childConfig))) - .thenReturn(childFuture); - assertEquals( - childFuture, - new ExtensionOperationImpl(childContext, "4", "child") - .runInChildContextAsync("AcmeContext", resultType, () -> "result", childConfig)); + assertOperation(wait, WaitPrimitive.class, "1", "wait", "AcmeWait"); + assertOperation(invoke, InvokePrimitive.class, "2", "invoke", "AcmeInvoke"); + assertOperation(callback, CallbackPrimitive.class, "3", "callback", "AcmeCallback"); + assertOperation(child, ChildContextPrimitive.class, "4", "child", "AcmeContext"); } @Test void invalidSubtypeDoesNotClaimReservation() { - var context = mock(DurableContextImpl.class); - var duration = Duration.ofSeconds(1); - var future = mockFuture(); - when(context.waitAsyncWithId("1", "wait", duration)).thenReturn(future); - var operation = new ExtensionOperationImpl(context, "1", "wait"); - - assertThrows(NullPointerException.class, () -> operation.waitAsync(null, duration)); - assertThrows(IllegalArgumentException.class, () -> operation.waitAsync(" ", duration)); - assertEquals(future, operation.waitAsync(duration)); - } - - @Test - void reservedChildContextAdaptsSupplierToChildFunction() { - var context = mock(DurableContextImpl.class); - var future = mockStringFuture(); - var resultType = TypeToken.get(String.class); - var config = RunInChildContextConfig.builder().build(); - when(context.runInChildContextAsyncWithId(eq("1"), eq("child"), eq(resultType), any(), eq(config))) - .thenReturn(future); - var operation = new ExtensionOperationImpl(context, "1", "child"); - - assertEquals(future, operation.runInChildContextAsync(resultType, () -> "result", config)); + var context = context(); + replay(context, "1", "wait", OperationType.WAIT, "Wait"); + var operation = new ExtensionOperationImpl(context, "1", "wait", null); - @SuppressWarnings("unchecked") - var function = (ArgumentCaptor>) - (ArgumentCaptor) ArgumentCaptor.forClass(Function.class); - verify(context) - .runInChildContextAsyncWithId(eq("1"), eq("child"), eq(resultType), function.capture(), eq(config)); - assertEquals("result", function.getValue().apply(mock(DurableContext.class))); + assertThrows(NullPointerException.class, () -> operation.waitAsync(null, Duration.ofSeconds(1))); + assertThrows(IllegalArgumentException.class, () -> operation.waitAsync(" ", Duration.ofSeconds(1))); + assertOperation(operation.waitAsync("Wait", Duration.ofSeconds(1)), WaitPrimitive.class, "1", "wait", "Wait"); } @Test - void advancedChildContextDelegatesFrameworkFunction() { - var context = mock(DurableContextImpl.class); - var future = mockStringFuture(); - var resultType = TypeToken.get(String.class); - var config = ExtensionContextConfig.builder().build(); - when(context.extensionContextAsyncWithId( - eq("1"), eq("child"), eq("AcmeContext"), eq(resultType), any(), eq(config))) - .thenReturn(future); - - var actual = new ExtensionOperationImpl(context, "1", "child") - .runInChildContextAsync( - "AcmeContext", resultType, () -> ExtensionContextResult.completed("result"), config); - - assertEquals(future, actual); + void reservationCanOnlyExecuteOnceAcrossPrimitiveSelectors() { + var context = context(); + replay(context, "1", "only-once", OperationType.WAIT, "Wait"); + ExtensionOperation operation = new ExtensionOperationImpl(context, "1", "only-once", null); + + operation.waitAsync("Wait", Duration.ofSeconds(1)); + + assertThrows( + IllegalStateException.class, + () -> operation.stepAsync( + "Step", + TypeToken.get(String.class), + state -> ExtensionStepResult.succeed("second"), + ExtensionStepConfig.builder().build())); } - @Test - void reservationCanOnlyExecuteOnceAcrossPrimitiveSelectors() { + private DurableContextImpl context() { var context = mock(DurableContextImpl.class); - var duration = Duration.ofSeconds(1); - when(context.waitAsyncWithId("1", "only-once", duration)).thenReturn(mockFuture()); - ExtensionOperation operation = new ExtensionOperationImpl(context, "1", "only-once"); - - operation.waitAsync(duration); - - assertThrows(IllegalStateException.class, () -> operation.stepAsync(String.class, () -> "second")); + var executionManager = mock(ExecutionManager.class); + when(context.getExecutionManager()).thenReturn(executionManager); + when(context.getDurableConfig()) + .thenReturn(DurableConfig.builder() + .withExecutorService(Executors.newCachedThreadPool()) + .build()); + return context; } - @SuppressWarnings("unchecked") - private DurableFuture mockFuture() { - return mock(DurableFuture.class); + private void replay(DurableContextImpl context, String id, String name, OperationType type, String subType) { + var details = type == OperationType.STEP + ? StepDetails.builder().result("\"result\"").build() + : null; + var contextDetails = type == OperationType.CONTEXT + ? ContextDetails.builder().result("\"result\"").build() + : null; + var callbackDetails = type == OperationType.CALLBACK + ? CallbackDetails.builder() + .callbackId("callback-id") + .result("\"result\"") + .build() + : null; + when(context.getExecutionManager().getOperationAndUpdateReplayState(id)) + .thenReturn(Operation.builder() + .id(id) + .name(name) + .type(type) + .subType(subType) + .status(OperationStatus.SUCCEEDED) + .stepDetails(details) + .contextDetails(contextDetails) + .callbackDetails(callbackDetails) + .build()); } - @SuppressWarnings("unchecked") - private DurableFuture mockStringFuture() { - return mock(DurableFuture.class); + private void assertOperation( + Object future, Class type, String id, String name, String subType) { + var operation = assertInstanceOf(type, future); + assertEquals(id, operation.getOperationId()); + assertEquals(name, operation.getName()); + assertEquals(subType, operation.getSubTypeValue()); } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/OperationIdGeneratorTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/OperationIdGeneratorTest.java index 0ad05cdbd..09a187a6d 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/OperationIdGeneratorTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/OperationIdGeneratorTest.java @@ -25,10 +25,12 @@ void customLocalIdUsesParentNamespace() { } @Test - void generatedIdsSkipClaimedCustomNumericIds() { + void generatedIdsFailWhenNextNumericIdWasClaimed() { var generator = new OperationIdGenerator(null); assertEquals(hashOperationId("2"), generator.nextOperationId("2")); + var exception = assertThrows(IllegalArgumentException.class, generator::nextOperationId); + assertEquals("Local operation ID is already in use: 2", exception.getMessage()); assertEquals(hashOperationId("3"), generator.nextOperationId()); } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/extension/ExtensionOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/extension/ExtensionOperationTest.java new file mode 100644 index 000000000..ff15e57b6 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/extension/ExtensionOperationTest.java @@ -0,0 +1,37 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.extension; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import java.lang.reflect.Method; +import java.util.Arrays; +import java.util.Set; +import java.util.stream.Collectors; +import org.junit.jupiter.api.Test; + +class ExtensionOperationTest { + @Test + void exposesOnlyTheFullySpecifiedMethodForEachPrimitive() { + var expected = Set.of( + "createCallback(String,TypeToken,ExtensionCallbackConfig)", + "invokeAsync(String,String,Object,TypeToken,ExtensionInvokeConfig)", + "runInChildContextAsync(String,TypeToken,ExtensionContextFunction,ExtensionContextConfig)", + "stepAsync(String,TypeToken,ExtensionStepFunction,ExtensionStepConfig)", + "waitAsync(String,Duration)"); + + var actual = Arrays.stream(ExtensionOperation.class.getDeclaredMethods()) + .filter(method -> !method.isSynthetic()) + .map(ExtensionOperationTest::signature) + .collect(Collectors.toSet()); + + assertEquals(expected, actual); + } + + private static String signature(Method method) { + var parameters = Arrays.stream(method.getParameterTypes()) + .map(Class::getSimpleName) + .collect(Collectors.joining(",")); + return method.getName() + "(" + parameters + ")"; + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/model/OperationIdentifierTest.java b/sdk/src/test/java/software/amazon/lambda/durable/model/OperationIdentifierTest.java new file mode 100644 index 000000000..24c5a7b1c --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/model/OperationIdentifierTest.java @@ -0,0 +1,38 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.model; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static software.amazon.lambda.durable.model.OperationSubType.WAIT_FOR_CONDITION; + +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.services.lambda.model.OperationType; + +class OperationIdentifierTest { + + @Test + void supportsExtensionDefinedSubtypeStrings() { + var identifier = new OperationIdentifier("operation-1", "custom", OperationType.STEP, "CustomStep"); + + assertEquals(OperationType.STEP, identifier.operationType()); + assertEquals("CustomStep", identifier.subType()); + } + + @Test + void standardSubtypeFactoryStoresWireValue() { + var identifier = OperationIdentifier.of("operation-1", "condition", WAIT_FOR_CONDITION); + + assertEquals(OperationType.STEP, identifier.operationType()); + assertEquals("WaitForCondition", identifier.subType()); + } + + @Test + void rejectsBlankSubtype() { + var exception = assertThrows( + IllegalArgumentException.class, + () -> new OperationIdentifier("operation-1", "custom", OperationType.STEP, " ")); + + assertEquals("subType cannot be blank", exception.getMessage()); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/DeferredDurableFutureTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DeferredDurableFutureTest.java similarity index 97% rename from sdk/src/test/java/software/amazon/lambda/durable/context/extension/DeferredDurableFutureTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/operation/DeferredDurableFutureTest.java index 3fc7e51bc..fdaff4bb6 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/DeferredDurableFutureTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DeferredDurableFutureTest.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.context.extension; +package software.amazon.lambda.durable.operation; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/MapExtensionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java similarity index 87% rename from sdk/src/test/java/software/amazon/lambda/durable/context/extension/MapExtensionTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java index 750660f49..976064faa 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/MapExtensionTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.context.extension; +package software.amazon.lambda.durable.operation; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; @@ -16,7 +16,6 @@ import java.util.List; import java.util.concurrent.CompletableFuture; -import java.util.function.Supplier; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; import software.amazon.lambda.durable.DurableContext; @@ -24,7 +23,6 @@ import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.MapConfig; import software.amazon.lambda.durable.config.NestingType; -import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.context.BaseContextImpl; import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.extension.ExtensionContextConfig; @@ -34,7 +32,7 @@ import software.amazon.lambda.durable.model.MapResult; import software.amazon.lambda.durable.serde.JacksonSerDes; -class MapExtensionTest { +class DurableMapOperationImplementationTest { @Test void executeBuildsMapAndIterationContextsFromReservations() { var context = mock(ExtensionContext.class); @@ -51,13 +49,13 @@ void executeBuildsMapAndIterationContextsFromReservations() { any(ExtensionContextConfig.class))) .thenReturn(parentFuture); - var actual = MapExtension.execute( + var actual = DurableMapOperation.mapAsync( context, "map", List.of("a", "b"), TypeToken.get(String.class), (item, index, child) -> item + index, - config); + config.toOperationConfig()); assertSame(parentFuture, actual); var function = extensionFunction(); @@ -77,14 +75,14 @@ void executeBuildsMapAndIterationContextsFromReservations() { when(first.runInChildContextAsync( eq(MAP_ITERATION.getValue()), eq(TypeToken.get(String.class)), - any(Supplier.class), - any(RunInChildContextConfig.class))) + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) .thenReturn(new CompletedFuture<>("a0")); when(second.runInChildContextAsync( eq(MAP_ITERATION.getValue()), eq(TypeToken.get(String.class)), - any(Supplier.class), - any(RunInChildContextConfig.class))) + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) .thenReturn(new CompletedFuture<>("b1")); try (var ignoredContext = BaseContextImpl.attachCurrentContext(child); @@ -93,15 +91,15 @@ void executeBuildsMapAndIterationContextsFromReservations() { assertEquals(List.of("a0", "b1"), result.results()); } - var iterationConfig = ArgumentCaptor.forClass(RunInChildContextConfig.class); + var iterationConfig = ArgumentCaptor.forClass(ExtensionContextConfig.class); verify(first) .runInChildContextAsync( eq(MAP_ITERATION.getValue()), eq(TypeToken.get(String.class)), - any(Supplier.class), + any(ExtensionContextFunction.class), iterationConfig.capture()); - assertTrue(iterationConfig.getValue().isVirtual()); - assertSame(serDes, iterationConfig.getValue().serDes()); + assertTrue(iterationConfig.getValue().childContextConfig().isVirtual()); + assertSame(serDes, iterationConfig.getValue().childContextConfig().serDes()); } @SuppressWarnings({"rawtypes", "unchecked"}) diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableOperationConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableOperationConfigTest.java new file mode 100644 index 000000000..e66e38b92 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableOperationConfigTest.java @@ -0,0 +1,180 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.operation; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; + +import java.lang.reflect.Modifier; +import java.time.Duration; +import java.util.Arrays; +import java.util.function.BiFunction; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.ParallelDurableFuture; +import software.amazon.lambda.durable.config.CallbackConfig; +import software.amazon.lambda.durable.config.CompletionConfig; +import software.amazon.lambda.durable.config.InvokeConfig; +import software.amazon.lambda.durable.config.MapConfig; +import software.amazon.lambda.durable.config.NestingType; +import software.amazon.lambda.durable.config.ParallelBranchConfig; +import software.amazon.lambda.durable.config.ParallelConfig; +import software.amazon.lambda.durable.config.RunInChildContextConfig; +import software.amazon.lambda.durable.config.StepConfig; +import software.amazon.lambda.durable.config.StepSemantics; +import software.amazon.lambda.durable.config.WaitForCallbackConfig; +import software.amazon.lambda.durable.config.WaitForConditionConfig; +import software.amazon.lambda.durable.config.WithRetryConfig; +import software.amazon.lambda.durable.retry.RetryStrategy; +import software.amazon.lambda.durable.retry.WaitForConditionWaitStrategy; +import software.amazon.lambda.durable.serde.SerDes; + +class DurableOperationConfigTest { + @Test + void operationApisOwnTheirConfigTypes() throws Exception { + assertOperationConfig(DurableStepOperation.class, "StepConfig", StepConfig.class); + assertOperationConfig(DurableInvokeOperation.class, "InvokeConfig", InvokeConfig.class); + assertOperationConfig(DurableCallbackOperation.class, "CallbackConfig", CallbackConfig.class); + assertOperationConfig(DurableContextOperation.class, "RunInChildContextConfig", RunInChildContextConfig.class); + assertOperationConfig(DurableMapOperation.class, "MapConfig", MapConfig.class); + assertOperationConfig(DurableParallelOperation.class, "ParallelConfig", ParallelConfig.class); + assertOperationConfig( + ParallelDurableFuture.class, + DurableParallelOperation.class, + "ParallelBranchConfig", + ParallelBranchConfig.class); + assertOperationConfig( + DurableWaitForCallbackOperation.class, "WaitForCallbackConfig", WaitForCallbackConfig.class); + assertOperationConfig( + DurableWaitForConditionOperation.class, "WaitForConditionConfig", WaitForConditionConfig.class); + assertOperationConfig(DurableWithRetryOperation.class, "WithRetryConfig", WithRetryConfig.class); + } + + @Test + void legacyPrimitiveConfigsConvertWithoutLosingValues() throws Exception { + var retryStrategy = mock(RetryStrategy.class); + var serDes = mock(SerDes.class); + var step = convert(StepConfig.builder() + .retryStrategy(retryStrategy) + .semanticsPerRetry(StepSemantics.AT_MOST_ONCE_PER_RETRY) + .serDes(serDes) + .build()); + assertSame(retryStrategy, value(step, "retryStrategy")); + assertEquals(StepSemantics.AT_MOST_ONCE_PER_RETRY, value(step, "semanticsPerRetry")); + assertSame(serDes, value(step, "serDes")); + + var payloadSerDes = mock(SerDes.class); + var invoke = convert(InvokeConfig.builder() + .payloadSerDes(payloadSerDes) + .serDes(serDes) + .tenantId("tenant") + .build()); + assertSame(payloadSerDes, value(invoke, "payloadSerDes")); + assertSame(serDes, value(invoke, "serDes")); + assertEquals("tenant", value(invoke, "tenantId")); + + var callback = convert(CallbackConfig.builder() + .timeout(Duration.ofMinutes(5)) + .heartbeatTimeout(Duration.ofMinutes(1)) + .serDes(serDes) + .build()); + assertEquals(Duration.ofMinutes(5), value(callback, "timeout")); + assertEquals(Duration.ofMinutes(1), value(callback, "heartbeatTimeout")); + assertSame(serDes, value(callback, "serDes")); + + var child = convert( + RunInChildContextConfig.builder().serDes(serDes).isVirtual(true).build()); + assertSame(serDes, value(child, "serDes")); + assertEquals(true, value(child, "isVirtual")); + } + + @Test + void legacyCompositeConfigsConvertWithoutLosingValues() throws Exception { + var serDes = mock(SerDes.class); + var completionConfig = CompletionConfig.firstSuccessful(); + BiFunction itemNamer = (item, index) -> item + "-" + index; + var map = convert(MapConfig.builder() + .maxConcurrency(3) + .completionConfig(completionConfig) + .serDes(serDes) + .nestingType(NestingType.NESTED) + .itemNamer(itemNamer) + .build()); + assertEquals(3, value(map, "maxConcurrency")); + assertSame(completionConfig, value(map, "completionConfig")); + assertSame(serDes, value(map, "serDes")); + assertEquals(NestingType.NESTED, value(map, "nestingType")); + assertSame(itemNamer, value(map, "itemNamer")); + + var parallel = convert(ParallelConfig.builder() + .maxConcurrency(2) + .completionConfig(completionConfig) + .nestingType(NestingType.FLAT) + .build()); + assertEquals(2, value(parallel, "maxConcurrency")); + assertSame(completionConfig, value(parallel, "completionConfig")); + assertEquals(NestingType.FLAT, value(parallel, "nestingType")); + + var branch = convert(ParallelBranchConfig.builder().serDes(serDes).build()); + assertSame(serDes, value(branch, "serDes")); + + var retryStrategy = mock(RetryStrategy.class); + var retry = convert(WithRetryConfig.builder() + .retryStrategy(retryStrategy) + .wrapInChildContext(true) + .build()); + assertSame(retryStrategy, value(retry, "retryStrategy")); + assertEquals(true, value(retry, "wrapInChildContext")); + } + + @Test + void legacyStatefulConfigsConvertWithoutLosingValues() throws Exception { + var serDes = mock(SerDes.class); + var stepConfig = StepConfig.builder().serDes(serDes).build(); + var callbackConfig = + CallbackConfig.builder().timeout(Duration.ofMinutes(2)).build(); + var callback = convert(WaitForCallbackConfig.builder() + .stepConfig(stepConfig) + .callbackConfig(callbackConfig) + .build()); + assertSame(serDes, value(value(callback, "stepConfig"), "serDes")); + assertEquals(Duration.ofMinutes(2), value(value(callback, "callbackConfig"), "timeout")); + + @SuppressWarnings("unchecked") + var waitStrategy = (WaitForConditionWaitStrategy) mock(WaitForConditionWaitStrategy.class); + var condition = convert(WaitForConditionConfig.builder() + .waitStrategy(waitStrategy) + .serDes(serDes) + .initialState("initial") + .build()); + assertSame(waitStrategy, value(condition, "waitStrategy")); + assertSame(serDes, value(condition, "serDes")); + assertEquals("initial", value(condition, "initialState")); + } + + private static void assertOperationConfig(Class operationClass, String nestedName, Class legacyClass) + throws Exception { + assertOperationConfig(operationClass, operationClass, nestedName, legacyClass); + } + + private static void assertOperationConfig( + Class apiClass, Class operationClass, String nestedName, Class legacyClass) throws Exception { + var nestedClass = Class.forName(operationClass.getName() + "$" + nestedName); + assertTrue(Modifier.isPublic(nestedClass.getModifiers())); + assertTrue(Modifier.isStatic(nestedClass.getModifiers())); + assertEquals(nestedClass, legacyClass.getMethod("toOperationConfig").getReturnType()); + assertFalse(Arrays.stream(apiClass.getMethods()) + .flatMap(method -> Arrays.stream(method.getParameterTypes())) + .anyMatch(legacyClass::equals)); + } + + private static Object convert(Object legacyConfig) throws Exception { + return legacyConfig.getClass().getMethod("toOperationConfig").invoke(legacyConfig); + } + + private static Object value(Object config, String methodName) throws Exception { + return config.getClass().getMethod(methodName).invoke(config); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/ParallelExtensionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableParallelOperationImplementationTest.java similarity index 90% rename from sdk/src/test/java/software/amazon/lambda/durable/context/extension/ParallelExtensionTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/operation/DurableParallelOperationImplementationTest.java index b602c2494..3765fb030 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/ParallelExtensionTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableParallelOperationImplementationTest.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.context.extension; +package software.amazon.lambda.durable.operation; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; @@ -18,7 +18,6 @@ import java.util.List; import java.util.concurrent.CompletableFuture; -import java.util.function.Supplier; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; import software.amazon.lambda.durable.DurableConfig; @@ -26,8 +25,6 @@ import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.NestingType; -import software.amazon.lambda.durable.config.ParallelConfig; -import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.context.BaseContextImpl; import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.extension.ExtensionContextConfig; @@ -36,9 +33,10 @@ import software.amazon.lambda.durable.extension.ExtensionOperation; import software.amazon.lambda.durable.model.ConcurrencyCompletionStatus; import software.amazon.lambda.durable.model.ParallelResult; +import software.amazon.lambda.durable.operation.DurableParallelOperation.ParallelConfig; import software.amazon.lambda.durable.serde.JacksonSerDes; -class ParallelExtensionTest { +class DurableParallelOperationImplementationTest { @Test void executeBuildsParallelAndBranchContextsFromReservations() { var context = mock(ExtensionContext.class); @@ -57,7 +55,7 @@ void executeBuildsParallelAndBranchContextsFromReservations() { .thenReturn(parentFuture); when(parentFuture.completionFuture()).thenReturn(parentCompletion); - var parallel = ParallelExtension.execute( + var parallel = DurableParallelOperation.parallel( context, "parallel", ParallelConfig.builder().nestingType(NestingType.FLAT).build()); @@ -87,14 +85,14 @@ void executeBuildsParallelAndBranchContextsFromReservations() { when(first.runInChildContextAsync( eq(PARALLEL_BRANCH.getValue()), eq(TypeToken.get(String.class)), - any(Supplier.class), - any(RunInChildContextConfig.class))) + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) .thenReturn(new CompletedFuture<>("first")); when(second.runInChildContextAsync( eq(PARALLEL_BRANCH.getValue()), eq(TypeToken.get(String.class)), - any(Supplier.class), - any(RunInChildContextConfig.class))) + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) .thenReturn(new CompletedFuture<>("second")); ParallelResult result; @@ -112,15 +110,15 @@ void executeBuildsParallelAndBranchContextsFromReservations() { assertEquals("first", firstFuture.get()); assertEquals("second", secondFuture.get()); - var branchConfig = ArgumentCaptor.forClass(RunInChildContextConfig.class); + var branchConfig = ArgumentCaptor.forClass(ExtensionContextConfig.class); verify(first) .runInChildContextAsync( eq(PARALLEL_BRANCH.getValue()), eq(TypeToken.get(String.class)), - any(Supplier.class), + any(ExtensionContextFunction.class), branchConfig.capture()); - assertTrue(branchConfig.getValue().isVirtual()); - assertSame(serDes, branchConfig.getValue().serDes()); + assertTrue(branchConfig.getValue().childContextConfig().isVirtual()); + assertSame(serDes, branchConfig.getValue().childContextConfig().serDes()); } @Test @@ -136,7 +134,7 @@ void replaySkipsBranchesMissingFromCompletedResult() { any(ExtensionContextConfig.class))) .thenReturn(mockParallelResultFuture()); - var parallel = ParallelExtension.execute( + var parallel = DurableParallelOperation.parallel( context, "parallel", ParallelConfig.builder().build()); parallel.branch("skipped", String.class, child -> "skipped"); parallel.branch("completed", String.class, child -> "completed"); @@ -153,8 +151,8 @@ void replaySkipsBranchesMissingFromCompletedResult() { when(completed.runInChildContextAsync( eq(PARALLEL_BRANCH.getValue()), eq(TypeToken.get(String.class)), - any(Supplier.class), - any(RunInChildContextConfig.class))) + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) .thenReturn(new CompletedFuture<>("completed")); var replayState = new ParallelResult( 2, @@ -175,8 +173,8 @@ void replaySkipsBranchesMissingFromCompletedResult() { .runInChildContextAsync( any(String.class), any(TypeToken.class), - any(Supplier.class), - any(RunInChildContextConfig.class)); + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class)); } @Test @@ -201,7 +199,7 @@ void getIncludesLateBranchesAsSkipped() { ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED, List.of(ParallelResult.Status.SUCCEEDED))); - var parallel = ParallelExtension.execute( + var parallel = DurableParallelOperation.parallel( context, "parallel", ParallelConfig.builder().build()); parallel.branch("completed", String.class, child -> "completed"); parallel.branch("late", String.class, child -> "late"); @@ -227,7 +225,7 @@ void branchAfterCloseFails() { any(ExtensionContextFunction.class), any(ExtensionContextConfig.class))) .thenReturn(mockParallelResultFuture()); - var parallel = ParallelExtension.execute( + var parallel = DurableParallelOperation.parallel( context, "parallel", ParallelConfig.builder().build()); parallel.close(); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/WaitForCallbackExtensionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperationImplementationTest.java similarity index 93% rename from sdk/src/test/java/software/amazon/lambda/durable/context/extension/WaitForCallbackExtensionTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperationImplementationTest.java index cb9203bed..fe0f72b74 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/WaitForCallbackExtensionTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperationImplementationTest.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.context.extension; +package software.amazon.lambda.durable.operation; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertInstanceOf; @@ -33,7 +33,7 @@ import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.serde.JacksonSerDes; -class WaitForCallbackExtensionTest { +class DurableWaitForCallbackOperationImplementationTest { @Test void executeCreatesExistingWaitForCallbackContextTopology() { var context = mock(ExtensionContext.class); @@ -52,8 +52,8 @@ void executeCreatesExistingWaitForCallbackContextTopology() { any(ExtensionContextConfig.class))) .thenReturn(future); - var actual = WaitForCallbackExtension.execute( - context, "approval", resultType, (callbackId, stepContext) -> {}, config); + var actual = DurableWaitForCallbackOperation.waitForCallbackAsync( + context, "approval", resultType, (callbackId, stepContext) -> {}, config.toOperationConfig()); assertSame(future, actual); var contextConfig = ArgumentCaptor.forClass(ExtensionContextConfig.class); @@ -78,12 +78,12 @@ void errorHandlerPreservesCallbackTimeoutException() { any(ExtensionContextFunction.class), any(ExtensionContextConfig.class))) .thenReturn(mockStringFuture()); - WaitForCallbackExtension.execute( + DurableWaitForCallbackOperation.waitForCallbackAsync( context, "approval", resultType, (String callbackId, StepContext stepContext) -> {}, - WaitForCallbackConfig.builder().build()); + WaitForCallbackConfig.builder().build().toOperationConfig()); var config = ArgumentCaptor.forClass(ExtensionContextConfig.class); verify(parent) .runInChildContextAsync( diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/WaitForConditionExtensionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperationImplementationTest.java similarity index 93% rename from sdk/src/test/java/software/amazon/lambda/durable/context/extension/WaitForConditionExtensionTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperationImplementationTest.java index f81a8794c..e72d20d45 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/WaitForConditionExtensionTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperationImplementationTest.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.context.extension; +package software.amazon.lambda.durable.operation; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertInstanceOf; @@ -38,7 +38,7 @@ import software.amazon.lambda.durable.model.WaitForConditionResult; import software.amazon.lambda.durable.serde.JacksonSerDes; -class WaitForConditionExtensionTest { +class DurableWaitForConditionOperationImplementationTest { @AfterEach void clearContext() { BaseContextImpl.setCurrentContext(null); @@ -68,8 +68,12 @@ void executeMapsPollingResultsToStatefulStepOutcomes() { any(ExtensionStepConfig.class))) .thenReturn(future); - var actual = WaitForConditionExtension.execute( - context, "ready", resultType, (state, step) -> WaitForConditionResult.continuePolling("next"), config); + var actual = DurableWaitForConditionOperation.waitForConditionAsync( + context, + "ready", + resultType, + (state, step) -> WaitForConditionResult.continuePolling("next"), + config.toOperationConfig()); assertEquals(future.get(), actual.get()); var function = extensionFunction(); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/WithRetryExtensionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWithRetryOperationImplementationTest.java similarity index 87% rename from sdk/src/test/java/software/amazon/lambda/durable/context/extension/WithRetryExtensionTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWithRetryOperationImplementationTest.java index 1e05e1a82..f163fa8e0 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/WithRetryExtensionTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWithRetryOperationImplementationTest.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.context.extension; +package software.amazon.lambda.durable.operation; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; @@ -28,7 +28,7 @@ import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.retry.RetryDecision; -class WithRetryExtensionTest { +class DurableWithRetryOperationImplementationTest { @AfterEach void clearContext() { BaseContextImpl.setCurrentContext(null); @@ -53,7 +53,7 @@ void executePreservesContextTopologyAndDurableBackoff() { .wrapInChildContext(true) .build(); - var actual = WithRetryExtension.execute( + var actual = DurableWithRetryOperation.withRetryAsync( context, "transaction", (attempt, child) -> { @@ -64,7 +64,7 @@ void executePreservesContextTopologyAndDurableBackoff() { } return "done"; }, - config); + config.toOperationConfig()); assertSame(future, actual); var function = extensionFunction(); @@ -79,7 +79,10 @@ void executePreservesContextTopologyAndDurableBackoff() { var child = mock(CurrentExtensionContext.class); var wait = mock(ExtensionOperation.class); + var waitFuture = mockVoidFuture(); when(child.reserve("transaction-backoff-1")).thenReturn(wait); + when(wait.waitAsync(OperationSubType.WAIT.getValue(), Duration.ofSeconds(5))) + .thenReturn(waitFuture); BaseContextImpl.setCurrentContext(child); var result = function.getValue().apply(); @@ -87,7 +90,7 @@ void executePreservesContextTopologyAndDurableBackoff() { assertEquals("done", result.result()); assertEquals(1, attempts.get(0)); assertEquals(2, attempts.get(1)); - verify(wait).wait(Duration.ofSeconds(5)); + verify(wait).waitAsync(OperationSubType.WAIT.getValue(), Duration.ofSeconds(5)); } @SuppressWarnings({"rawtypes", "unchecked"}) @@ -100,5 +103,10 @@ private DurableFuture mockObjectFuture() { return mock(DurableFuture.class); } + @SuppressWarnings("unchecked") + private DurableFuture mockVoidFuture() { + return mock(DurableFuture.class); + } + private interface CurrentExtensionContext extends DurableContext, ExtensionContext {} } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinatorTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/OperationConcurrencyCoordinatorTest.java similarity index 88% rename from sdk/src/test/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinatorTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/operation/OperationConcurrencyCoordinatorTest.java index 2b70c62f1..fcaab2c51 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/context/extension/ExtensionConcurrencyCoordinatorTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/OperationConcurrencyCoordinatorTest.java @@ -1,16 +1,16 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.context.extension; +package software.amazon.lambda.durable.operation; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; -import static software.amazon.lambda.durable.context.extension.ExtensionConcurrencyCoordinator.ItemStatus.FAILED; -import static software.amazon.lambda.durable.context.extension.ExtensionConcurrencyCoordinator.ItemStatus.SKIPPED; -import static software.amazon.lambda.durable.context.extension.ExtensionConcurrencyCoordinator.ItemStatus.SUCCEEDED; import static software.amazon.lambda.durable.model.ConcurrencyCompletionStatus.ALL_COMPLETED; import static software.amazon.lambda.durable.model.ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED; +import static software.amazon.lambda.durable.operation.OperationConcurrencyCoordinator.ItemStatus.FAILED; +import static software.amazon.lambda.durable.operation.OperationConcurrencyCoordinator.ItemStatus.SKIPPED; +import static software.amazon.lambda.durable.operation.OperationConcurrencyCoordinator.ItemStatus.SUCCEEDED; import java.util.List; import java.util.concurrent.CompletableFuture; @@ -21,10 +21,10 @@ import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.config.CompletionConfig; -class ExtensionConcurrencyCoordinatorTest { +class OperationConcurrencyCoordinatorTest { @Test void launchesNoMoreThanMaxConcurrency() throws Exception { - var coordinator = new ExtensionConcurrencyCoordinator(2, CompletionConfig.allCompleted()); + var coordinator = new OperationConcurrencyCoordinator(2, CompletionConfig.allCompleted()); var first = new TestFuture<>("first"); var second = new TestFuture<>("second"); var third = new TestFuture<>("third"); @@ -51,7 +51,7 @@ void launchesNoMoreThanMaxConcurrency() throws Exception { @Test void launchesNextItemAfterSynchronousReplayCompletion() throws Exception { - var coordinator = new ExtensionConcurrencyCoordinator(1, CompletionConfig.allCompleted()); + var coordinator = new OperationConcurrencyCoordinator(1, CompletionConfig.allCompleted()); var replayed = new TestFuture<>("replayed"); replayed.complete(); var next = new TestFuture<>("next"); @@ -73,7 +73,7 @@ void launchesNextItemAfterSynchronousReplayCompletion() throws Exception { @Test void earlyCompletionMarksUnlaunchedItemsSkipped() { - var coordinator = new ExtensionConcurrencyCoordinator(1, CompletionConfig.minSuccessful(1)); + var coordinator = new OperationConcurrencyCoordinator(1, CompletionConfig.minSuccessful(1)); var first = new TestFuture<>("first"); var launched = new AtomicInteger(); coordinator.register(() -> { @@ -99,13 +99,13 @@ void earlyCompletionMarksUnlaunchedItemsSkipped() { assertEquals( List.of(SUCCEEDED, SKIPPED, SKIPPED), completion.items().stream() - .map(ExtensionConcurrencyCoordinator.Item::status) + .map(OperationConcurrencyCoordinator.Item::status) .toList()); } @Test void failedItemsContributeToCompletionStatus() { - var coordinator = new ExtensionConcurrencyCoordinator(2, CompletionConfig.allCompleted()); + var coordinator = new OperationConcurrencyCoordinator(2, CompletionConfig.allCompleted()); var failed = new TestFuture(new IllegalStateException("failed")); var succeeded = new TestFuture<>("succeeded"); coordinator.register(() -> failed); @@ -121,7 +121,7 @@ void failedItemsContributeToCompletionStatus() { assertEquals( List.of(FAILED, SUCCEEDED), completion.items().stream() - .map(ExtensionConcurrencyCoordinator.Item::status) + .map(OperationConcurrencyCoordinator.Item::status) .toList()); } @@ -136,7 +136,7 @@ void dynamicRegistrationDoesNotCompleteUntilRegistrationCloses() throws Exceptio ? CompletionConfig.CompletionDecision.complete(ALL_COMPLETED) : CompletionConfig.CompletionDecision.continueExecution(); }); - var coordinator = new ExtensionConcurrencyCoordinator(1, completionConfig); + var coordinator = new OperationConcurrencyCoordinator(1, completionConfig); var result = CompletableFuture.supplyAsync(coordinator::awaitCompletion); var item = new TestFuture<>("result"); @@ -153,7 +153,7 @@ void dynamicRegistrationDoesNotCompleteUntilRegistrationCloses() throws Exceptio @Test void registrationAfterCloseFails() { - var coordinator = new ExtensionConcurrencyCoordinator(1, CompletionConfig.allCompleted()); + var coordinator = new OperationConcurrencyCoordinator(1, CompletionConfig.allCompleted()); coordinator.closeRegistration(); var exception = diff --git a/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginInfoConverterTest.java b/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginInfoConverterTest.java index 171c21203..1070bbd28 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginInfoConverterTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginInfoConverterTest.java @@ -9,7 +9,6 @@ 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.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; @@ -53,10 +52,10 @@ void toOperationInfo_withIdentifier_mapsAllFields() { } @Test - void toOperationInfo_withDescriptor_preservesCustomSubtype() { - var descriptor = new OperationDescriptor(OPERATION_ID, OPERATION_NAME, OperationType.STEP, "AcmeStep"); + void toOperationInfo_withCustomIdentifier_preservesCustomSubtype() { + var identifier = new OperationIdentifier(OPERATION_ID, OPERATION_NAME, OperationType.STEP, "AcmeStep"); - var info = PluginInfoConverter.toOperationInfo(null, descriptor, PARENT_ID); + var info = PluginInfoConverter.toOperationInfo(null, identifier, PARENT_ID); assertEquals("STEP", info.type()); assertEquals("AcmeStep", info.subType()); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/BaseDurableOperationPluginTest.java b/sdk/src/test/java/software/amazon/lambda/durable/primitive/BasePrimitivePluginTest.java similarity index 94% rename from sdk/src/test/java/software/amazon/lambda/durable/operation/BaseDurableOperationPluginTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/primitive/BasePrimitivePluginTest.java index 463a824ea..d0140e5f0 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/BaseDurableOperationPluginTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/primitive/BasePrimitivePluginTest.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; +package software.amazon.lambda.durable.primitive; import static org.junit.jupiter.api.Assertions.*; import static org.mockito.Mockito.mock; @@ -30,12 +30,12 @@ import software.amazon.lambda.durable.plugin.OperationInfo; /** - * Unit tests verifying that BaseDurableOperation.execute() fires onOperationStart with isReplay=true for all - * non-terminal operations during replay, regardless of operation type. + * Unit tests verifying that BasePrimitive.execute() fires onOperationStart with isReplay=true for all non-terminal + * operations during replay, regardless of operation type. * *

    This mirrors the Python SDK's TestPluginExecutorOnOperationReplay tests. */ -class BaseDurableOperationPluginTest { +class BasePrimitivePluginTest { private static final String EXECUTION_OP_ID = "exec-123"; private static final String EXECUTION_ARN = @@ -57,7 +57,7 @@ void execute_firesOnOperationStart_withIsReplayTrue_forNonTerminalWait() { var executionManager = createExecutionManager(List.of(waitOp), plugin); var durableContext = mockDurableContext(executionManager, plugin); - var operation = new WaitOperation( + var operation = new WaitPrimitive( OperationIdentifier.of(OPERATION_ID, OPERATION_NAME, OperationSubType.WAIT), Duration.ofMinutes(5), durableContext); @@ -89,7 +89,7 @@ void execute_doesNotFireOnOperationStart_forTerminalOperation(OperationStatus te var executionManager = createExecutionManager(List.of(waitOp), plugin); var durableContext = mockDurableContext(executionManager, plugin); - var operation = new WaitOperation( + var operation = new WaitPrimitive( OperationIdentifier.of(OPERATION_ID, OPERATION_NAME, OperationSubType.WAIT), Duration.ofMinutes(5), durableContext); @@ -109,7 +109,7 @@ void execute_firesOnOperationStart_withIsReplayFalse_forFirstExecution() { var executionManager = createExecutionManager(List.of(), plugin); var durableContext = mockDurableContext(executionManager, plugin); - var operation = new WaitOperation( + var operation = new WaitPrimitive( OperationIdentifier.of(OPERATION_ID, OPERATION_NAME, OperationSubType.WAIT), Duration.ofMinutes(5), durableContext); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/CallbackOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/primitive/CallbackPrimitiveTest.java similarity index 90% rename from sdk/src/test/java/software/amazon/lambda/durable/operation/CallbackOperationTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/primitive/CallbackPrimitiveTest.java index c8cea7934..af9a9b8f1 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/CallbackOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/primitive/CallbackPrimitiveTest.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; +package software.amazon.lambda.durable.primitive; import static org.junit.jupiter.api.Assertions.*; import static org.mockito.Mockito.mock; @@ -16,7 +16,6 @@ import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.TestUtils; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.CallbackConfig; import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.exception.CallbackFailedException; import software.amazon.lambda.durable.exception.CallbackTimeoutException; @@ -24,13 +23,14 @@ import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.execution.ThreadContext; import software.amazon.lambda.durable.execution.ThreadType; +import software.amazon.lambda.durable.extension.ExtensionCallbackConfig; import software.amazon.lambda.durable.model.DurableExecutionInput; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.serde.JacksonSerDes; import software.amazon.lambda.durable.serde.SerDes; -class CallbackOperationTest { +class CallbackPrimitiveTest { private static final String OPERATION_ID = TestUtils.hashOperationId("1"); private static final String OPERATION_NAME = "approval"; @@ -107,10 +107,10 @@ void executeCreatesCheckpointAndGetsCallbackId() { when(durableContext.getExecutionManager()).thenReturn(executionManager); var serDes = new JacksonSerDes(); - var operation = new CallbackOperation<>( + var operation = new CallbackPrimitive<>( OPERATION_IDENTIFIER, TypeToken.get(String.class), - CallbackConfig.builder().serDes(serDes).build(), + ExtensionCallbackConfig.builder().serDes(serDes).build(), durableContext); operation.execute(); @@ -122,14 +122,14 @@ void executeWithConfigSetsOptions() { var executionManager = createExecutionManager(List.of()); when(durableContext.getExecutionManager()).thenReturn(executionManager); var serDes = new JacksonSerDes(); - var config = CallbackConfig.builder() + var config = ExtensionCallbackConfig.builder() .timeout(Duration.ofMinutes(5)) .heartbeatTimeout(Duration.ofSeconds(30)) .serDes(serDes) .build(); var operation = - new CallbackOperation<>(OPERATION_IDENTIFIER, TypeToken.get(String.class), config, durableContext); + new CallbackPrimitive<>(OPERATION_IDENTIFIER, TypeToken.get(String.class), config, durableContext); operation.execute(); assertNotNull(operation.callbackId()); @@ -152,10 +152,10 @@ void replayReturnsExistingCallbackIdWhenSucceeded() { when(durableContext.getExecutionManager()).thenReturn(executionManager); var serDes = new JacksonSerDes(); - var operation = new CallbackOperation<>( + var operation = new CallbackPrimitive<>( OPERATION_IDENTIFIER, TypeToken.get(String.class), - CallbackConfig.builder().serDes(serDes).build(), + ExtensionCallbackConfig.builder().serDes(serDes).build(), durableContext); operation.execute(); @@ -180,10 +180,10 @@ void getReturnsDeserializedResultWhenSucceeded() { var serDes = new JacksonSerDes(); - var operation = new CallbackOperation<>( + var operation = new CallbackPrimitive<>( OPERATION_IDENTIFIER, TypeToken.get(String.class), - CallbackConfig.builder().serDes(serDes).build(), + ExtensionCallbackConfig.builder().serDes(serDes).build(), durableContext); operation.execute(); var result = operation.get(); @@ -212,10 +212,10 @@ void getThrowsCallbackExceptionWhenFailed() { var serDes = new JacksonSerDes(); - var operation = new CallbackOperation<>( + var operation = new CallbackPrimitive<>( OPERATION_IDENTIFIER, TypeToken.get(String.class), - CallbackConfig.builder().serDes(serDes).build(), + ExtensionCallbackConfig.builder().serDes(serDes).build(), durableContext); operation.execute(); @@ -239,10 +239,10 @@ void getThrowsCallbackTimeoutExceptionWhenTimedOut() { var serDes = new JacksonSerDes(); - var operation = new CallbackOperation<>( + var operation = new CallbackPrimitive<>( OPERATION_IDENTIFIER, TypeToken.get(String.class), - CallbackConfig.builder().serDes(serDes).build(), + ExtensionCallbackConfig.builder().serDes(serDes).build(), durableContext); operation.execute(); @@ -268,9 +268,9 @@ void operationUsesCustomSerDesWhenConfigContainsOne() { var executionManager = createExecutionManager(List.of(existingCallback)); when(durableContext.getExecutionManager()).thenReturn(executionManager); - var config = CallbackConfig.builder().serDes(customSerDes).build(); + var config = ExtensionCallbackConfig.builder().serDes(customSerDes).build(); var operation = - new CallbackOperation<>(OPERATION_IDENTIFIER, TypeToken.get(String.class), config, durableContext); + new CallbackPrimitive<>(OPERATION_IDENTIFIER, TypeToken.get(String.class), config, durableContext); operation.execute(); var result = operation.get(); @@ -297,10 +297,10 @@ void operationUsesDefaultSerDesWhenConfigIsNull() { var executionManager = createExecutionManager(List.of(existingCallback)); when(durableContext.getExecutionManager()).thenReturn(executionManager); - var operation = new CallbackOperation<>( + var operation = new CallbackPrimitive<>( OPERATION_IDENTIFIER, TypeToken.get(String.class), - CallbackConfig.builder().serDes(customSerDes).build(), + ExtensionCallbackConfig.builder().serDes(customSerDes).build(), durableContext); operation.execute(); var result = operation.get(); @@ -328,9 +328,9 @@ void operationUsesDefaultSerDesWhenConfigSerDesIsNull() { var executionManager = createExecutionManager(List.of(existingCallback)); when(durableContext.getExecutionManager()).thenReturn(executionManager); - var config = CallbackConfig.builder().serDes(customSerDes).build(); + var config = ExtensionCallbackConfig.builder().serDes(customSerDes).build(); var operation = - new CallbackOperation<>(OPERATION_IDENTIFIER, TypeToken.get(String.class), config, durableContext); + new CallbackPrimitive<>(OPERATION_IDENTIFIER, TypeToken.get(String.class), config, durableContext); operation.execute(); var result = operation.get(); @@ -357,10 +357,10 @@ void getThrowsSerDesExceptionWithHelpfulMessageWhenDeserializationFails() { var executionManager = createExecutionManager(List.of(existingCallback)); when(durableContext.getExecutionManager()).thenReturn(executionManager); - var operation = new CallbackOperation<>( + var operation = new CallbackPrimitive<>( OPERATION_IDENTIFIER, TypeToken.get(String.class), - CallbackConfig.builder().serDes(failingSerDes).build(), + ExtensionCallbackConfig.builder().serDes(failingSerDes).build(), durableContext); operation.execute(); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/primitive/ChildContextPrimitiveTest.java similarity index 96% rename from sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/primitive/ChildContextPrimitiveTest.java index 5f5d336e9..55220a93a 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/primitive/ChildContextPrimitiveTest.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; +package software.amazon.lambda.durable.primitive; import static org.junit.jupiter.api.Assertions.*; import static org.mockito.Mockito.*; @@ -34,14 +34,13 @@ import software.amazon.lambda.durable.extension.ExtensionContextConfig; import software.amazon.lambda.durable.extension.ExtensionContextFailure; import software.amazon.lambda.durable.extension.ExtensionContextResult; -import software.amazon.lambda.durable.model.OperationDescriptor; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.serde.JacksonSerDes; import software.amazon.lambda.durable.serde.SerDes; -/** Unit tests for ChildContextOperation. */ -class ChildContextOperationTest { +/** Unit tests for ChildContextPrimitive. */ +class ChildContextPrimitiveTest { private static final class SerializationOnlySerDes implements SerDes { @Override @@ -96,12 +95,12 @@ private DurableConfig createConfig(boolean deserializeAfterSerialization) { private static final OperationIdentifier OPERATION_IDENTIFIER = OperationIdentifier.of("1", "test-context", OperationSubType.RUN_IN_CHILD_CONTEXT); - private ChildContextOperation createOperation(Function func) { + private ChildContextPrimitive createOperation(Function func) { return createOperation(func, SERDES); } - private ChildContextOperation createOperation(Function func, SerDes serDes) { - return new ChildContextOperation<>( + private ChildContextPrimitive createOperation(Function func, SerDes serDes) { + return new ChildContextPrimitive<>( OPERATION_IDENTIFIER, func, TypeToken.get(String.class), @@ -109,12 +108,12 @@ private ChildContextOperation createOperation(Function createVirtualOperation(Function func) { + private ChildContextPrimitive createVirtualOperation(Function func) { return createVirtualOperation(func, SERDES); } - private ChildContextOperation createVirtualOperation(Function func, SerDes serDes) { - return new ChildContextOperation<>( + private ChildContextPrimitive createVirtualOperation(Function func, SerDes serDes) { + return new ChildContextPrimitive<>( OPERATION_IDENTIFIER, func, TypeToken.get(String.class), @@ -122,9 +121,9 @@ private ChildContextOperation createVirtualOperation(Function createOperationWithParent( - Function func, BaseDurableOperation parent) { - return new ChildContextOperation<>( + private ChildContextPrimitive createOperationWithParent( + Function func, BasePrimitive parent) { + return new ChildContextPrimitive<>( OPERATION_IDENTIFIER, func, TypeToken.get(String.class), @@ -133,9 +132,9 @@ private ChildContextOperation createOperationWithParent( parent); } - private ChildContextOperation createExtensionOperation(ExtensionContextConfig config) { - return new ChildContextOperation<>( - new OperationDescriptor("1", "test-context", OperationType.CONTEXT, "AcmeContext"), + private ChildContextPrimitive createExtensionOperation(ExtensionContextConfig config) { + return new ChildContextPrimitive<>( + new OperationIdentifier("1", "test-context", OperationType.CONTEXT, "AcmeContext"), () -> ExtensionContextResult.completed("unused"), TypeToken.get(String.class), config, @@ -539,7 +538,7 @@ void childSkipsFailureCheckpointWhenParentAlreadyCompleted() throws Exception { .sendOperationUpdate(argThat(update -> update.action() == OperationAction.FAIL)); } - private static final class CompletedParentOperation extends BaseDurableOperation { + private static final class CompletedParentOperation extends BasePrimitive { private CompletedParentOperation(DurableContextImpl durableContext) { super( OperationIdentifier.of("parent", "parent", OperationSubType.RUN_IN_CHILD_CONTEXT), diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/InvokeOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/primitive/InvokePrimitiveTest.java similarity index 90% rename from sdk/src/test/java/software/amazon/lambda/durable/operation/InvokeOperationTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/primitive/InvokePrimitiveTest.java index 2c1d76c74..76158b0e9 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/InvokeOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/primitive/InvokePrimitiveTest.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; +package software.amazon.lambda.durable.primitive; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -14,7 +14,6 @@ import software.amazon.awssdk.services.lambda.model.Operation; import software.amazon.awssdk.services.lambda.model.OperationStatus; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.InvokeConfig; import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.exception.InvokeException; import software.amazon.lambda.durable.exception.InvokeFailedException; @@ -23,11 +22,12 @@ import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.execution.ThreadContext; import software.amazon.lambda.durable.execution.ThreadType; +import software.amazon.lambda.durable.extension.ExtensionInvokeConfig; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.serde.JacksonSerDes; -class InvokeOperationTest { +class InvokePrimitiveTest { private static final String OPERATION_ID = "2"; private static final String OPERATION_NAME = "test-invoke"; private static final OperationIdentifier OPERATION_IDENTIFIER = @@ -56,12 +56,12 @@ void getDoesNotThrowWhenCalledFromHandlerContext() { .build(); when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); - var operation = new InvokeOperation<>( + var operation = new InvokePrimitive<>( OPERATION_IDENTIFIER, "test-function", "{}", TypeToken.get(String.class), - InvokeConfig.builder().serDes(new JacksonSerDes()).build(), + ExtensionInvokeConfig.builder().serDes(new JacksonSerDes()).build(), durableContext); operation.onCheckpointComplete(op); @@ -85,12 +85,12 @@ void getInvokeFailedExceptionWhenInvocationFailed() { .build(); when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); - var operation = new InvokeOperation<>( + var operation = new InvokePrimitive<>( OPERATION_IDENTIFIER, "test-function", "{}", TypeToken.get(String.class), - InvokeConfig.builder().serDes(new JacksonSerDes()).build(), + ExtensionInvokeConfig.builder().serDes(new JacksonSerDes()).build(), durableContext); operation.onCheckpointComplete(op); @@ -116,12 +116,12 @@ void getInvokeTimedOutExceptionWhenInvocationTimedOut() { .build(); when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); - var operation = new InvokeOperation<>( + var operation = new InvokePrimitive<>( OPERATION_IDENTIFIER, "test-function", "{}", TypeToken.get(String.class), - InvokeConfig.builder().serDes(new JacksonSerDes()).build(), + ExtensionInvokeConfig.builder().serDes(new JacksonSerDes()).build(), durableContext); operation.onCheckpointComplete(op); @@ -147,12 +147,12 @@ void getInvokeStoppedExceptionWhenInvocationTimedOut() { .build(); when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); - var operation = new InvokeOperation<>( + var operation = new InvokePrimitive<>( OPERATION_IDENTIFIER, "test-function", "{}", TypeToken.get(String.class), - InvokeConfig.builder().serDes(new JacksonSerDes()).build(), + ExtensionInvokeConfig.builder().serDes(new JacksonSerDes()).build(), durableContext); operation.onCheckpointComplete(op); @@ -178,12 +178,12 @@ void getInvokeFailedExceptionWhenInvocationEndedUnexpectedly() { .build(); when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); - var operation = new InvokeOperation<>( + var operation = new InvokePrimitive<>( OPERATION_IDENTIFIER, "test-function", "{}", TypeToken.get(String.class), - InvokeConfig.builder().serDes(new JacksonSerDes()).build(), + ExtensionInvokeConfig.builder().serDes(new JacksonSerDes()).build(), durableContext); operation.onCheckpointComplete(op); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/SerializableDurableOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/primitive/SerializablePrimitiveTest.java similarity index 85% rename from sdk/src/test/java/software/amazon/lambda/durable/operation/SerializableDurableOperationTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/primitive/SerializablePrimitiveTest.java index bc9e940b8..139df7475 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/SerializableDurableOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/primitive/SerializablePrimitiveTest.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; +package software.amazon.lambda.durable.primitive; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertInstanceOf; @@ -43,7 +43,7 @@ import software.amazon.lambda.durable.serde.JacksonSerDes; import software.amazon.lambda.durable.serde.SerDes; -class SerializableDurableOperationTest { +class SerializablePrimitiveTest { private static final class TrackingSerDes extends JacksonSerDes { private final AtomicInteger deserializeCount = new AtomicInteger(0); @@ -110,8 +110,8 @@ void setUp() { @Test void getOperation() { - SerializableDurableOperation op = - new SerializableDurableOperation<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { + SerializablePrimitive op = + new SerializablePrimitive<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { @Override protected void start() {} @@ -134,8 +134,8 @@ public String get() { @Test void waitForOperationCompletionThrowsIfOperationMissing() { when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(null); - SerializableDurableOperation op = - new SerializableDurableOperation<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { + SerializablePrimitive op = + new SerializablePrimitive<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { @Override protected void start() { markAlreadyCompleted(); @@ -159,8 +159,8 @@ public String get() { void waitForOperationCompletionThrowsIllegalStateExceptionWhenCalledFromStepThread() { when(executionManager.getCurrentThreadContext()).thenReturn(new ThreadContext(CONTEXT_ID, ThreadType.STEP)); - SerializableDurableOperation op = - new SerializableDurableOperation<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { + SerializablePrimitive op = + new SerializablePrimitive<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { @Override protected void start() { markAlreadyCompleted(); @@ -183,8 +183,8 @@ public String get() { @Test void waitForOperationCompletionWhenRunningAndReadyToComplete() throws InterruptedException, ExecutionException, TimeoutException { - SerializableDurableOperation op = - new SerializableDurableOperation<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { + SerializablePrimitive op = + new SerializablePrimitive<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { @Override protected void start() {} @@ -215,8 +215,8 @@ public String get() { @Test void waitForOperationCompletionWhenAlreadyCompleted() { - SerializableDurableOperation op = - new SerializableDurableOperation<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { + SerializablePrimitive op = + new SerializablePrimitive<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { @Override protected void start() { markAlreadyCompleted(); @@ -239,8 +239,8 @@ public String get() { @Test void markAlreadyCompleted() { - SerializableDurableOperation op = - new SerializableDurableOperation<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { + SerializablePrimitive op = + new SerializablePrimitive<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { @Override protected void start() { markAlreadyCompleted(); @@ -266,8 +266,8 @@ void validateReplayThrowsWhenTypeMismatch() { .thenReturn( Operation.builder().type(OperationType.CHAINED_INVOKE).build()); - SerializableDurableOperation op = - new SerializableDurableOperation<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { + SerializablePrimitive op = + new SerializablePrimitive<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { @Override protected void start() { validateReplay(getOperation()); @@ -293,8 +293,8 @@ void validateReplayThrowsWhenNameMismatch() { .type(OPERATION_TYPE) .build()); - SerializableDurableOperation op = - new SerializableDurableOperation<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { + SerializablePrimitive op = + new SerializablePrimitive<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { @Override protected void start() { validateReplay(getOperation()); @@ -316,8 +316,8 @@ public String get() { void validateReplayDoesNotThrowWhenNoOperation() { when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(null); - SerializableDurableOperation op = - new SerializableDurableOperation<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { + SerializablePrimitive op = + new SerializablePrimitive<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { @Override protected void start() { validateReplay(getOperation()); @@ -343,8 +343,8 @@ void validateReplayDoesNotThrowWhenNameAndTypeMatch() { .subType(OperationSubType.STEP.getValue()) .build()); - SerializableDurableOperation op = - new SerializableDurableOperation<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { + SerializablePrimitive op = + new SerializablePrimitive<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { @Override protected void start() { validateReplay(getOperation()); @@ -363,8 +363,8 @@ public String get() { @Test void deserializeResult() { - SerializableDurableOperation op = - new SerializableDurableOperation<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { + SerializablePrimitive op = + new SerializablePrimitive<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { @Override protected void start() {} @@ -385,8 +385,8 @@ public String get() { @Test void serializeAndDeserializeResultDeserializesResult() { var serDes = new TrackingSerDes(); - SerializableDurableOperation op = - new SerializableDurableOperation<>(OPERATION_IDENTIFIER, RESULT_TYPE, serDes, durableContext) { + SerializablePrimitive op = + new SerializablePrimitive<>(OPERATION_IDENTIFIER, RESULT_TYPE, serDes, durableContext) { @Override protected void start() {} @@ -409,8 +409,8 @@ public String get() { @Test void serializeAndDeserializeResultThrowsWhenDeserializeFails() { var serDes = new SerializationOnlySerDes(); - SerializableDurableOperation op = - new SerializableDurableOperation<>(OPERATION_IDENTIFIER, RESULT_TYPE, serDes, durableContext) { + SerializablePrimitive op = + new SerializablePrimitive<>(OPERATION_IDENTIFIER, RESULT_TYPE, serDes, durableContext) { @Override protected void start() {} @@ -433,8 +433,8 @@ void serializeAndDeserializeResultReturnsRawResultWhenDeserializationDisabled() when(durableContext.getDurableConfig()).thenReturn(configWithDeserializeAfterSerialization(false)); var serDes = new NormalizingSerDes(); - SerializableDurableOperation op = - new SerializableDurableOperation<>(OPERATION_IDENTIFIER, RESULT_TYPE, serDes, durableContext) { + SerializablePrimitive op = + new SerializablePrimitive<>(OPERATION_IDENTIFIER, RESULT_TYPE, serDes, durableContext) { @Override protected void start() {} @@ -455,8 +455,8 @@ public String get() { @Test void serializeAndDeserializeResultReturnsDeserializedValue() { - SerializableDurableOperation op = - new SerializableDurableOperation<>( + SerializablePrimitive op = + new SerializablePrimitive<>( OPERATION_IDENTIFIER, RESULT_TYPE, new NormalizingSerDes(), durableContext) { @Override protected void start() {} @@ -478,8 +478,8 @@ public String get() { @Test void deserializeException() { - SerializableDurableOperation op = - new SerializableDurableOperation<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { + SerializablePrimitive op = + new SerializablePrimitive<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { @Override protected void start() {} @@ -505,8 +505,8 @@ public String get() { @Test void serializeExceptionValidatesRoundTrip() { var serDes = new TrackingSerDes(); - SerializableDurableOperation op = - new SerializableDurableOperation<>(OPERATION_IDENTIFIER, RESULT_TYPE, serDes, durableContext) { + SerializablePrimitive op = + new SerializablePrimitive<>(OPERATION_IDENTIFIER, RESULT_TYPE, serDes, durableContext) { @Override protected void start() {} @@ -530,8 +530,8 @@ void serializeExceptionSkipsRoundTripValidationWhenDisabled() { when(durableContext.getDurableConfig()).thenReturn(configWithDeserializeAfterSerialization(false)); var serDes = new SerializationOnlySerDes(); - SerializableDurableOperation op = - new SerializableDurableOperation<>(OPERATION_IDENTIFIER, RESULT_TYPE, serDes, durableContext) { + SerializablePrimitive op = + new SerializablePrimitive<>(OPERATION_IDENTIFIER, RESULT_TYPE, serDes, durableContext) { @Override protected void start() {} @@ -551,8 +551,8 @@ public String get() { @Test void polling() { - SerializableDurableOperation op = - new SerializableDurableOperation<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { + SerializablePrimitive op = + new SerializablePrimitive<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { @Override protected void start() {} @@ -575,8 +575,8 @@ public String get() { void sendOperationUpdate() { var update = OperationUpdate.builder(); - SerializableDurableOperation op = - new SerializableDurableOperation<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { + SerializablePrimitive op = + new SerializablePrimitive<>(OPERATION_IDENTIFIER, RESULT_TYPE, SER_DES, durableContext) { @Override protected void start() {} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/StatefulExtensionStepOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java similarity index 82% rename from sdk/src/test/java/software/amazon/lambda/durable/operation/StatefulExtensionStepOperationTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java index 4743e769e..57c216db9 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/StatefulExtensionStepOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java @@ -1,24 +1,30 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; +package software.amazon.lambda.durable.primitive; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import java.time.Duration; import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import software.amazon.awssdk.services.lambda.model.ErrorObject; import software.amazon.awssdk.services.lambda.model.Operation; +import software.amazon.awssdk.services.lambda.model.OperationAction; import software.amazon.awssdk.services.lambda.model.OperationStatus; import software.amazon.awssdk.services.lambda.model.OperationType; +import software.amazon.awssdk.services.lambda.model.OperationUpdate; import software.amazon.awssdk.services.lambda.model.StepDetails; import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.TypeToken; @@ -33,11 +39,12 @@ import software.amazon.lambda.durable.extension.ExtensionStepConfig; import software.amazon.lambda.durable.extension.ExtensionStepFunction; import software.amazon.lambda.durable.extension.ExtensionStepResult; -import software.amazon.lambda.durable.model.OperationDescriptor; +import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.retry.RetryDecision; import software.amazon.lambda.durable.serde.JacksonSerDes; -class StatefulExtensionStepOperationTest { +class StatefulExtensionStepPrimitiveTest { private static final String OPERATION_ID = "1"; private static final String OPERATION_NAME = "test-wait-for-condition"; private static final JacksonSerDes SERDES = new JacksonSerDes(); @@ -162,6 +169,35 @@ void replayWithoutCheckpointStateUsesInitialState() throws Exception { assertTrue(called.await(2, TimeUnit.SECONDS)); } + @Test + void exceptionRetryWithoutStateDoesNotCheckpointPayload() throws Exception { + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(null); + when(executionManager.pollForOperationUpdates(OPERATION_ID)).thenReturn(new CompletableFuture<>()); + var retryUpdate = new AtomicReference(); + var retrySent = new CountDownLatch(1); + when(executionManager.sendOperationUpdate(any())).thenAnswer(invocation -> { + var update = invocation.getArgument(0); + if (update.action() == OperationAction.RETRY) { + retryUpdate.set(update); + retrySent.countDown(); + } + return CompletableFuture.completedFuture(null); + }); + var operation = createOperationWithConfig( + state -> { + throw new IllegalStateException("retry"); + }, + ExtensionStepConfig.builder() + .serDes(SERDES) + .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(1))) + .build()); + + operation.execute(); + + assertTrue(retrySent.await(2, TimeUnit.SECONDS)); + assertNull(retryUpdate.get().payload()); + } + @Test void corruptReplayStateFailsBeforeCallingFunction() { when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)) @@ -229,23 +265,30 @@ private void assertResumes(OperationStatus status, int expectedState) throws Exc assertTrue(called.await(2, TimeUnit.SECONDS)); } - private StepOperation createOperation(ExtensionStepFunction function) { + private StepPrimitive createOperation(ExtensionStepFunction function) { return createOperation(function, null); } - private StepOperation createOperation(ExtensionStepFunction function, Integer initialState) { - return new StepOperation<>( - new OperationDescriptor( + private StepPrimitive createOperation(ExtensionStepFunction function, Integer initialState) { + return createOperationWithConfig( + function, + ExtensionStepConfig.builder() + .initialState(initialState) + .serDes(SERDES) + .build()); + } + + private StepPrimitive createOperationWithConfig( + ExtensionStepFunction function, ExtensionStepConfig config) { + return new StepPrimitive<>( + new OperationIdentifier( OPERATION_ID, OPERATION_NAME, OperationType.STEP, OperationSubType.WAIT_FOR_CONDITION.getValue()), function, TypeToken.get(Integer.class), - ExtensionStepConfig.builder() - .initialState(initialState) - .serDes(SERDES) - .build(), + config, durableContext); } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/StepOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/primitive/StepPrimitiveTest.java similarity index 81% rename from sdk/src/test/java/software/amazon/lambda/durable/operation/StepOperationTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/primitive/StepPrimitiveTest.java index be4962d71..9190c4e40 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/StepOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/primitive/StepPrimitiveTest.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; +package software.amazon.lambda.durable.primitive; import static org.junit.jupiter.api.Assertions.*; import static org.mockito.Mockito.*; @@ -15,18 +15,20 @@ import software.amazon.awssdk.services.lambda.model.StepDetails; import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.StepConfig; import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.exception.StepFailedException; import software.amazon.lambda.durable.exception.StepInterruptedException; import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.execution.ThreadContext; import software.amazon.lambda.durable.execution.ThreadType; +import software.amazon.lambda.durable.extension.ExtensionStepConfig; +import software.amazon.lambda.durable.extension.ExtensionStepResult; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.serde.JacksonSerDes; +import software.amazon.lambda.durable.serde.SerDes; -class StepOperationTest { +class StepPrimitiveTest { private static final String OPERATION_ID = "1"; private static final String OPERATION_NAME = "test-step"; @@ -36,6 +38,11 @@ class StepOperationTest { private ExecutionManager executionManager; private DurableContextImpl durableContext; + @Test + void exposesOnlyTheExtensionConstructor() { + assertEquals(1, StepPrimitive.class.getDeclaredConstructors().length); + } + @BeforeEach void setUp() { executionManager = mock(ExecutionManager.class); @@ -71,6 +78,15 @@ private void mockFailedOperation( when(executionManager.getOperationAndUpdateReplayState("1")).thenReturn(operation); } + private StepPrimitive createOperation(SerDes serDes) { + return new StepPrimitive<>( + OPERATION_IDENTIFIER, + ignored -> ExtensionStepResult.succeed(RESULT), + TypeToken.get(String.class), + ExtensionStepConfig.builder().serDes(serDes).build(), + durableContext); + } + @Test void getDoesNotThrowWhenCalledFromHandlerContext() { var op = Operation.builder() @@ -82,12 +98,7 @@ void getDoesNotThrowWhenCalledFromHandlerContext() { when(executionManager.getCurrentThreadContext()).thenReturn(new ThreadContext("handler", ThreadType.CONTEXT)); when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); - var operation = new StepOperation<>( - OPERATION_IDENTIFIER, - (ctx) -> RESULT, - TypeToken.get(String.class), - StepConfig.builder().serDes(new JacksonSerDes()).build(), - durableContext); + var operation = createOperation(new JacksonSerDes()); operation.onCheckpointComplete(op); var result = operation.get(); @@ -107,12 +118,7 @@ void getThrowsOriginalExceptionWhenClassIsAvailable() { serDes.serialize(originalException), stackTrace); - var operation = new StepOperation<>( - OPERATION_IDENTIFIER, - (ctx) -> RESULT, - TypeToken.get(String.class), - StepConfig.builder().serDes(serDes).build(), - durableContext); + var operation = createOperation(serDes); operation.execute(); @@ -136,12 +142,7 @@ void getThrowsOriginalCustomExceptionWhenClassIsAvailable() { serDes.serialize(originalException), stackTrace); - var operation = new StepOperation<>( - OPERATION_IDENTIFIER, - (ctx) -> RESULT, - TypeToken.get(String.class), - StepConfig.builder().serDes(serDes).build(), - durableContext); + var operation = createOperation(serDes); operation.execute(); @@ -156,12 +157,7 @@ void getFallsBackToStepFailedExceptionWhenClassNotFound() { mockFailedOperation(executionManager, "NonExistentException", "This class doesn't exist", "{}", stackTrace); - var operation = new StepOperation<>( - OPERATION_IDENTIFIER, - (ctx) -> RESULT, - TypeToken.get(String.class), - StepConfig.builder().serDes(new JacksonSerDes()).build(), - durableContext); + var operation = createOperation(new JacksonSerDes()); operation.execute(); @@ -182,12 +178,7 @@ void getFallsBackToStepFailedExceptionWhenDeserializationFails() { "invalid-json-{{{", stackTrace); - var operation = new StepOperation<>( - OPERATION_IDENTIFIER, - (ctx) -> RESULT, - TypeToken.get(String.class), - StepConfig.builder().serDes(new JacksonSerDes()).build(), - durableContext); + var operation = createOperation(new JacksonSerDes()); operation.execute(); @@ -203,12 +194,7 @@ void getFallsBackToStepFailedExceptionWhenErrorDataIsNull() { mockFailedOperation( executionManager, RuntimeException.class.getName(), "Something went wrong", null, stackTrace); - var operation = new StepOperation<>( - OPERATION_IDENTIFIER, - (ctx) -> RESULT, - TypeToken.get(String.class), - StepConfig.builder().serDes(new JacksonSerDes()).build(), - durableContext); + var operation = createOperation(new JacksonSerDes()); operation.execute(); @@ -224,12 +210,7 @@ void getThrowsStepInterruptedExceptionDirectly() { mockFailedOperation( executionManager, StepInterruptedException.class.getName(), "Step was interrupted", null, stackTrace); - var operation = new StepOperation<>( - OPERATION_IDENTIFIER, - (ctx) -> RESULT, - TypeToken.get(String.class), - StepConfig.builder().serDes(new JacksonSerDes()).build(), - durableContext); + var operation = createOperation(new JacksonSerDes()); operation.execute(); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/WaitOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/primitive/WaitPrimitiveTest.java similarity index 92% rename from sdk/src/test/java/software/amazon/lambda/durable/operation/WaitOperationTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/primitive/WaitPrimitiveTest.java index d620595ec..1f50fcde7 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/WaitOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/primitive/WaitPrimitiveTest.java @@ -1,6 +1,6 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; +package software.amazon.lambda.durable.primitive; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNull; @@ -20,7 +20,7 @@ import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; -class WaitOperationTest { +class WaitPrimitiveTest { private static final String OPERATION_ID = "2"; private static final String CONTEXT_ID = "handler"; private static final String OPERATION_NAME = "test-wait"; @@ -38,7 +38,7 @@ void setUp() { @Test void constructor_withValidDuration_shouldPass() { - var operation = new WaitOperation(OPERATION_IDENTIFIER, Duration.ofSeconds(10), durableContext); + var operation = new WaitPrimitive(OPERATION_IDENTIFIER, Duration.ofSeconds(10), durableContext); assertEquals(OPERATION_ID, operation.getOperationId()); } @@ -54,7 +54,7 @@ void getDoesNotThrowWhenCalledFromHandlerContext() { when(executionManager.getCurrentThreadContext()).thenReturn(new ThreadContext(CONTEXT_ID, ThreadType.CONTEXT)); when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); - var operation = new WaitOperation(OPERATION_IDENTIFIER, Duration.ofSeconds(10), durableContext); + var operation = new WaitPrimitive(OPERATION_IDENTIFIER, Duration.ofSeconds(10), durableContext); operation.onCheckpointComplete(op); var result = operation.get(); @@ -71,7 +71,7 @@ void getSucceededWhenStarted() { when(executionManager.getCurrentThreadContext()).thenReturn(new ThreadContext(CONTEXT_ID, ThreadType.CONTEXT)); when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); - var operation = new WaitOperation(OPERATION_IDENTIFIER, Duration.ofSeconds(10), durableContext); + var operation = new WaitPrimitive(OPERATION_IDENTIFIER, Duration.ofSeconds(10), durableContext); operation.onCheckpointComplete(op); // we currently don't check the operation status at all, so it's not blocked or failed From c362182e0610e5864c23575620d39962d6a01a06 Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 08:05:12 +0000 Subject: [PATCH 24/40] fix: preserve parallel branch config compatibility --- .../DeserializationFailedParallelExample.java | 5 +---- .../lambda/durable/ParallelDurableFuture.java | 2 +- .../operation/ParallelOperationFuture.java | 10 +++++----- .../operation/DurableOperationConfigTest.java | 19 ++++++++++++++----- 4 files changed, 21 insertions(+), 15 deletions(-) diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/parallel/DeserializationFailedParallelExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/parallel/DeserializationFailedParallelExample.java index 0c938f9a0..a54aedfab 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/parallel/DeserializationFailedParallelExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/parallel/DeserializationFailedParallelExample.java @@ -51,10 +51,7 @@ public String handleRequest(Input input, DurableContext context) { throw new RuntimeException("Intentional failure for transform"); }); }, - ParallelBranchConfig.builder() - .serDes(new FailedSerDes()) - .build() - .toOperationConfig()); + ParallelBranchConfig.builder().serDes(new FailedSerDes()).build()); parallel.get(); try { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java b/sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java index d18b2234a..a1a1dae1c 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java @@ -4,9 +4,9 @@ import java.util.function.Function; import java.util.function.Supplier; +import software.amazon.lambda.durable.config.ParallelBranchConfig; import software.amazon.lambda.durable.model.ParallelResult; import software.amazon.lambda.durable.model.SafeCloseable; -import software.amazon.lambda.durable.operation.DurableParallelOperation.ParallelBranchConfig; /** User-facing context for managing parallel branch execution within a durable function. */ public interface ParallelDurableFuture extends SafeCloseable, DurableFuture { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperationFuture.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperationFuture.java index d42e70311..16b235ca8 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperationFuture.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperationFuture.java @@ -18,13 +18,13 @@ import software.amazon.lambda.durable.ParallelDurableFuture; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.CompletionConfig; +import software.amazon.lambda.durable.config.ParallelBranchConfig; import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.extension.ExtensionContextConfig; import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; import software.amazon.lambda.durable.extension.ExtensionContextResult; import software.amazon.lambda.durable.model.ParallelResult; -import software.amazon.lambda.durable.operation.DurableParallelOperation.ParallelBranchConfig; import software.amazon.lambda.durable.operation.DurableParallelOperation.ParallelConfig; import software.amazon.lambda.durable.serde.SerDes; import software.amazon.lambda.durable.util.ParameterValidator; @@ -60,7 +60,7 @@ public DurableFuture branch( synchronized (lock) { ensureRegistrationOpen(); - var definition = new BranchDefinition<>(name, resultType, function, config); + var definition = new BranchDefinition<>(name, resultType, function, config.toOperationConfig()); branches.add(definition); if (coordinator != null) { registerBranch(definition, branches.size() - 1); @@ -202,7 +202,7 @@ private static ParallelResult constructResult(OperationConcurrencyCoordinator.Co statuses); } - private ExtensionContextConfig branchConfig(ParallelBranchConfig branchConfig) { + private ExtensionContextConfig branchConfig(DurableParallelOperation.ParallelBranchConfig branchConfig) { return ExtensionContextConfig.builder() .childContextConfig(RunInChildContextConfig.builder() .serDes(branchConfig.serDes() == null ? defaultSerDes : branchConfig.serDes()) @@ -235,14 +235,14 @@ private static final class BranchDefinition { private final String name; private final TypeToken resultType; private final Function function; - private final ParallelBranchConfig config; + private final DurableParallelOperation.ParallelBranchConfig config; private final DeferredDurableFuture future = new DeferredDurableFuture<>(); private BranchDefinition( String name, TypeToken resultType, Function function, - ParallelBranchConfig config) { + DurableParallelOperation.ParallelBranchConfig config) { this.name = name; this.resultType = resultType; this.function = function; diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableOperationConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableOperationConfigTest.java index e66e38b92..964ed5cdd 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableOperationConfigTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableOperationConfigTest.java @@ -40,11 +40,8 @@ void operationApisOwnTheirConfigTypes() throws Exception { assertOperationConfig(DurableContextOperation.class, "RunInChildContextConfig", RunInChildContextConfig.class); assertOperationConfig(DurableMapOperation.class, "MapConfig", MapConfig.class); assertOperationConfig(DurableParallelOperation.class, "ParallelConfig", ParallelConfig.class); - assertOperationConfig( - ParallelDurableFuture.class, - DurableParallelOperation.class, - "ParallelBranchConfig", - ParallelBranchConfig.class); + assertOperationConfig(DurableParallelOperation.class, "ParallelBranchConfig", ParallelBranchConfig.class); + assertParallelFutureUsesCompatibilityBranchConfig(); assertOperationConfig( DurableWaitForCallbackOperation.class, "WaitForCallbackConfig", WaitForCallbackConfig.class); assertOperationConfig( @@ -159,6 +156,18 @@ private static void assertOperationConfig(Class operationClass, String nested assertOperationConfig(operationClass, operationClass, nestedName, legacyClass); } + private static void assertParallelFutureUsesCompatibilityBranchConfig() { + var nestedConfig = DurableParallelOperation.ParallelBranchConfig.class; + assertTrue(Arrays.stream(ParallelDurableFuture.class.getMethods()) + .filter(method -> method.getName().equals("branch")) + .flatMap(method -> Arrays.stream(method.getParameterTypes())) + .anyMatch(ParallelBranchConfig.class::equals)); + assertFalse(Arrays.stream(ParallelDurableFuture.class.getMethods()) + .filter(method -> method.getName().equals("branch")) + .flatMap(method -> Arrays.stream(method.getParameterTypes())) + .anyMatch(nestedConfig::equals)); + } + private static void assertOperationConfig( Class apiClass, Class operationClass, String nestedName, Class legacyClass) throws Exception { var nestedClass = Class.forName(operationClass.getName() + "$" + nestedName); From 14057145c22410fb1ca16de2b126d1d499a19ef4 Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Mon, 10 Aug 2026 16:46:00 +0000 Subject: [PATCH 25/40] test: isolate stateful step replay cases --- .../primitive/StatefulExtensionStepPrimitiveTest.java | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java b/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java index 57c216db9..2d0eb6108 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java @@ -19,6 +19,8 @@ import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; import software.amazon.awssdk.services.lambda.model.ErrorObject; import software.amazon.awssdk.services.lambda.model.Operation; import software.amazon.awssdk.services.lambda.model.OperationAction; @@ -121,10 +123,10 @@ void replayFailedFallsBackToStepFailedException() { assertThrows(StepFailedException.class, operation::get); } - @Test - void replayStartedAndReadyResumeWithCheckpointedState() throws Exception { - assertResumes(OperationStatus.STARTED, 10); - assertResumes(OperationStatus.READY, 5); + @ParameterizedTest(name = "{0}") + @CsvSource({"STARTED, 10", "READY, 5"}) + void replayStartedOrReadyResumesWithCheckpointedState(OperationStatus status, int expectedState) throws Exception { + assertResumes(status, expectedState); } @Test From 085bd6c9f1e2d7ef5a1b4a4e5bdc05f77adbc5fc Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Mon, 10 Aug 2026 11:16:45 -0700 Subject: [PATCH 26/40] refactor: nest operation context types --- docs/adr/006-custom-extension-operations.md | 14 ++++--- docs/advanced/extensions.md | 15 ++++--- .../StaticOperationsIntegrationTest.java | 3 ++ .../amazon/lambda/durable/MapItemContext.java | 32 -------------- .../durable/OperationContextStorage.java | 36 ---------------- .../durable/WaitForCallbackContext.java | 33 --------------- .../lambda/durable/WithRetryContext.java | 32 -------------- .../operation/DurableMapOperation.java | 42 ++++++++++++++++++- .../DurableWaitForCallbackOperation.java | 42 ++++++++++++++++++- .../operation/DurableWithRetryOperation.java | 42 ++++++++++++++++++- .../durable/DurableMapOperationTest.java | 1 + .../DurableWaitForCallbackOperationTest.java | 1 + .../DurableWithRetryOperationTest.java | 1 + ...ageTest.java => OperationContextTest.java} | 5 ++- 14 files changed, 150 insertions(+), 149 deletions(-) delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/MapItemContext.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/OperationContextStorage.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/WaitForCallbackContext.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/WithRetryContext.java rename sdk/src/test/java/software/amazon/lambda/durable/{OperationContextStorageTest.java => OperationContextTest.java} (89%) diff --git a/docs/adr/006-custom-extension-operations.md b/docs/adr/006-custom-extension-operations.md index 98f46cdd3..5c1d478d1 100644 --- a/docs/adr/006-custom-extension-operations.md +++ b/docs/adr/006-custom-extension-operations.md @@ -116,10 +116,12 @@ software.amazon.lambda.durable.operation Each `Durable*Operation` class owns its context-free overloads and its canonical `ExtensionContext` implementation. There are no separate built-in `*Extension` classes. -Keep the following customer-facing types in the root `software.amazon.lambda.durable` package: +Keep established SDK types such as `DurableFuture`, `StepContext`, and `TypeToken` in the root +`software.amazon.lambda.durable` package. Keep operation-specific TLS metadata nested under its owning operation: -- operation-specific TLS metadata such as `MapItemContext`, `WaitForCallbackContext`, and `WithRetryContext` -- established SDK types such as `DurableFuture`, `StepContext`, and `TypeToken` +- `DurableMapOperation.MapItemContext` +- `DurableWaitForCallbackOperation.WaitForCallbackContext` +- `DurableWithRetryOperation.WithRetryContext` Backend primitive engines remain internal under `software.amazon.lambda.durable.primitive`. @@ -133,9 +135,9 @@ retrieved from scoped thread-local contexts: - `DurableContext` - `ExtensionContext` - `StepContext` -- `MapItemContext` -- `WaitForCallbackContext` -- `WithRetryContext` +- `DurableMapOperation.MapItemContext` +- `DurableWaitForCallbackOperation.WaitForCallbackContext` +- `DurableWithRetryOperation.WithRetryContext` - extension replay contexts Nested scopes restore the preceding value. Current context is available only on SDK-managed threads and is not diff --git a/docs/advanced/extensions.md b/docs/advanced/extensions.md index 5cc89362f..96a1195d7 100644 --- a/docs/advanced/extensions.md +++ b/docs/advanced/extensions.md @@ -127,7 +127,7 @@ var result = DurableStepOperation.step("process", Result.class, () -> { ```java var result = DurableMapOperation.map("process", items, Result.class, item -> { - var index = MapItemContext.getCurrentContext().getIndex(); + var index = DurableMapOperation.MapItemContext.getCurrentContext().getIndex(); return process(item, index); }); ``` @@ -136,12 +136,14 @@ var result = DurableMapOperation.map("process", items, Result.class, item -> { var result = DurableWaitForCallbackOperation.waitForCallback( "approval", Approval.class, - () -> submit(WaitForCallbackContext.getCurrentContext().getCallbackId())); + () -> submit(DurableWaitForCallbackOperation.WaitForCallbackContext + .getCurrentContext() + .getCallbackId())); ``` ```java var result = DurableWithRetryOperation.withRetry("transaction", () -> { - var attempt = WithRetryContext.getCurrentContext().getAttempt(); + var attempt = DurableWithRetryOperation.WithRetryContext.getCurrentContext().getAttempt(); return executeAttempt(attempt); }); ``` @@ -155,15 +157,16 @@ obtain attempt metadata from `StepContext.getCurrentContext()`. and child-context threads. `StepContext.getCurrentContext()` is available inside step and wait-for-condition user functions. -`MapItemContext`, `WaitForCallbackContext`, and `WithRetryContext` are available only inside their corresponding user -function. Nested scopes restore the previous context when they close. +`DurableMapOperation.MapItemContext`, `DurableWaitForCallbackOperation.WaitForCallbackContext`, and +`DurableWithRetryOperation.WithRetryContext` are available only inside their corresponding user function. Nested +scopes restore the previous context when they close. Operation-specific TLS is not automatically propagated into a nested primitive's separate user-function thread. Read the metadata in its owning function and capture any application value needed by the nested operation: ```java var result = DurableMapOperation.map("process", items, Result.class, item -> { - var index = MapItemContext.getCurrentContext().getIndex(); + var index = DurableMapOperation.MapItemContext.getCurrentContext().getIndex(); return DurableStepOperation.step("process-item", Result.class, () -> process(item, index)); }); ``` diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java index 9e0c0e140..23755482d 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java @@ -26,11 +26,14 @@ import software.amazon.lambda.durable.model.WaitForConditionResult; import software.amazon.lambda.durable.operation.DurableContextOperation; import software.amazon.lambda.durable.operation.DurableMapOperation; +import software.amazon.lambda.durable.operation.DurableMapOperation.MapItemContext; import software.amazon.lambda.durable.operation.DurableParallelOperation; import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation; +import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation.WaitForCallbackContext; import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation; import software.amazon.lambda.durable.operation.DurableWithRetryOperation; +import software.amazon.lambda.durable.operation.DurableWithRetryOperation.WithRetryContext; import software.amazon.lambda.durable.retry.RetryStrategies; import software.amazon.lambda.durable.retry.WaitStrategies; import software.amazon.lambda.durable.testing.LocalDurableTestRunner; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/MapItemContext.java b/sdk/src/main/java/software/amazon/lambda/durable/MapItemContext.java deleted file mode 100644 index 99174a388..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/MapItemContext.java +++ /dev/null @@ -1,32 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; - -import software.amazon.lambda.durable.model.SafeCloseable; - -/** Metadata for the map item function active on the current SDK-managed thread. */ -public final class MapItemContext { - private static final OperationContextStorage CURRENT = - new OperationContextStorage<>("MapItemContext"); - - private final int index; - - private MapItemContext(int index) { - this.index = index; - } - - /** Returns the map item context attached to the current SDK-managed thread. */ - public static MapItemContext getCurrentContext() { - return CURRENT.getCurrentContext(); - } - - /** Returns the zero-based index of the current map item. */ - public int getIndex() { - return index; - } - - /** Attaches map item metadata for the duration of the returned scope. */ - public static SafeCloseable attach(int index) { - return CURRENT.attach(new MapItemContext(index)); - } -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/OperationContextStorage.java b/sdk/src/main/java/software/amazon/lambda/durable/OperationContextStorage.java deleted file mode 100644 index a36bfc212..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/OperationContextStorage.java +++ /dev/null @@ -1,36 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; - -import java.util.Objects; -import software.amazon.lambda.durable.model.SafeCloseable; - -final class OperationContextStorage { - private final String contextName; - private final ThreadLocal current = new ThreadLocal<>(); - - OperationContextStorage(String contextName) { - this.contextName = contextName; - } - - T getCurrentContext() { - var context = current.get(); - if (context == null) { - throw new IllegalStateException(contextName + " is not active on the current thread"); - } - return context; - } - - SafeCloseable attach(T context) { - Objects.requireNonNull(context, "context cannot be null"); - var previous = current.get(); - current.set(context); - return () -> { - if (previous == null) { - current.remove(); - } else { - current.set(previous); - } - }; - } -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/WaitForCallbackContext.java b/sdk/src/main/java/software/amazon/lambda/durable/WaitForCallbackContext.java deleted file mode 100644 index 53a2a6721..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/WaitForCallbackContext.java +++ /dev/null @@ -1,33 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; - -import java.util.Objects; -import software.amazon.lambda.durable.model.SafeCloseable; - -/** Metadata for the callback submitter active on the current SDK-managed thread. */ -public final class WaitForCallbackContext { - private static final OperationContextStorage CURRENT = - new OperationContextStorage<>("WaitForCallbackContext"); - - private final String callbackId; - - private WaitForCallbackContext(String callbackId) { - this.callbackId = Objects.requireNonNull(callbackId, "callbackId cannot be null"); - } - - /** Returns the callback context attached to the current SDK-managed thread. */ - public static WaitForCallbackContext getCurrentContext() { - return CURRENT.getCurrentContext(); - } - - /** Returns the callback ID to send to the external system. */ - public String getCallbackId() { - return callbackId; - } - - /** Attaches callback metadata for the duration of the returned scope. */ - public static SafeCloseable attach(String callbackId) { - return CURRENT.attach(new WaitForCallbackContext(callbackId)); - } -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/WithRetryContext.java b/sdk/src/main/java/software/amazon/lambda/durable/WithRetryContext.java deleted file mode 100644 index d84370b21..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/WithRetryContext.java +++ /dev/null @@ -1,32 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; - -import software.amazon.lambda.durable.model.SafeCloseable; - -/** Metadata for the retry body active on the current SDK-managed thread. */ -public final class WithRetryContext { - private static final OperationContextStorage CURRENT = - new OperationContextStorage<>("WithRetryContext"); - - private final int attempt; - - private WithRetryContext(int attempt) { - this.attempt = attempt; - } - - /** Returns the retry context attached to the current SDK-managed thread. */ - public static WithRetryContext getCurrentContext() { - return CURRENT.getCurrentContext(); - } - - /** Returns the current one-based retry attempt. */ - public int getAttempt() { - return attempt; - } - - /** Attaches retry metadata for the duration of the returned scope. */ - public static SafeCloseable attach(int attempt) { - return CURRENT.attach(new WithRetryContext(attempt)); - } -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java index 4b348ab61..eecaf9c26 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java @@ -17,7 +17,6 @@ import java.util.function.Function; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; -import software.amazon.lambda.durable.MapItemContext; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.CompletionConfig; import software.amazon.lambda.durable.config.NestingType; @@ -29,6 +28,7 @@ import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; import software.amazon.lambda.durable.extension.ExtensionContextResult; import software.amazon.lambda.durable.model.MapResult; +import software.amazon.lambda.durable.model.SafeCloseable; import software.amazon.lambda.durable.serde.SerDes; import software.amazon.lambda.durable.util.ExceptionHelper; import software.amazon.lambda.durable.util.ParameterValidator; @@ -281,6 +281,46 @@ private static TypeToken> mapResultType() { return new TypeToken<>() {}; } + /** Metadata for the map item function active on the current SDK-managed thread. */ + public static final class MapItemContext { + private static final ThreadLocal CURRENT = new ThreadLocal<>(); + + private final int index; + + private MapItemContext(int index) { + this.index = index; + } + + /** Returns the map item context attached to the current SDK-managed thread. */ + public static MapItemContext getCurrentContext() { + var context = CURRENT.get(); + if (context == null) { + throw new IllegalStateException("MapItemContext is not active on the current thread"); + } + return context; + } + + /** Returns the zero-based index of the current map item. */ + public int getIndex() { + return index; + } + + /** Attaches map item metadata for the duration of the returned scope. */ + public static SafeCloseable attach(int index) { + var previous = CURRENT.get(); + CURRENT.set(new MapItemContext(index)); + return () -> restore(previous); + } + + private static void restore(MapItemContext previous) { + if (previous == null) { + CURRENT.remove(); + } else { + CURRENT.set(previous); + } + } + } + /** Configuration for durable MAP operations. */ public static final class MapConfig { private final int maxConcurrency; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperation.java index 2235a6b4a..d0908b318 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperation.java @@ -13,7 +13,6 @@ import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.WaitForCallbackContext; import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.exception.CallbackFailedException; import software.amazon.lambda.durable.exception.CallbackSubmitterException; @@ -25,6 +24,7 @@ import software.amazon.lambda.durable.extension.ExtensionContextFailure; import software.amazon.lambda.durable.extension.ExtensionContextResult; import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.model.SafeCloseable; import software.amazon.lambda.durable.util.ParameterValidator; /** Context-free static facade and canonical implementation of durable wait-for-callback operations. */ @@ -166,6 +166,46 @@ private static Operation findChild(ExtensionContextFailure failure, OperationTyp .orElse(null); } + /** Metadata for the callback submitter active on the current SDK-managed thread. */ + public static final class WaitForCallbackContext { + private static final ThreadLocal CURRENT = new ThreadLocal<>(); + + private final String callbackId; + + private WaitForCallbackContext(String callbackId) { + this.callbackId = Objects.requireNonNull(callbackId, "callbackId cannot be null"); + } + + /** Returns the callback context attached to the current SDK-managed thread. */ + public static WaitForCallbackContext getCurrentContext() { + var context = CURRENT.get(); + if (context == null) { + throw new IllegalStateException("WaitForCallbackContext is not active on the current thread"); + } + return context; + } + + /** Returns the callback ID to send to the external system. */ + public String getCallbackId() { + return callbackId; + } + + /** Attaches callback metadata for the duration of the returned scope. */ + public static SafeCloseable attach(String callbackId) { + var previous = CURRENT.get(); + CURRENT.set(new WaitForCallbackContext(callbackId)); + return () -> restore(previous); + } + + private static void restore(WaitForCallbackContext previous) { + if (previous == null) { + CURRENT.remove(); + } else { + CURRENT.set(previous); + } + } + } + /** Configuration for durable wait-for-callback operations. */ public static final class WaitForCallbackConfig { private final DurableStepOperation.StepConfig stepConfig; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWithRetryOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWithRetryOperation.java index b4f68d790..a0d959e4c 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWithRetryOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWithRetryOperation.java @@ -9,7 +9,6 @@ import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.WithRetryContext; import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.SuspendExecutionException; @@ -17,6 +16,7 @@ import software.amazon.lambda.durable.extension.ExtensionContextConfig; import software.amazon.lambda.durable.extension.ExtensionContextResult; import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.model.SafeCloseable; import software.amazon.lambda.durable.retry.RetryStrategies; import software.amazon.lambda.durable.retry.RetryStrategy; @@ -107,6 +107,46 @@ private static String backoffName(String name, int attempt) { return name != null ? name + BACKOFF_SUFFIX + attempt : ANONYMOUS_BACKOFF_PREFIX + attempt; } + /** Metadata for the retry body active on the current SDK-managed thread. */ + public static final class WithRetryContext { + private static final ThreadLocal CURRENT = new ThreadLocal<>(); + + private final int attempt; + + private WithRetryContext(int attempt) { + this.attempt = attempt; + } + + /** Returns the retry context attached to the current SDK-managed thread. */ + public static WithRetryContext getCurrentContext() { + var context = CURRENT.get(); + if (context == null) { + throw new IllegalStateException("WithRetryContext is not active on the current thread"); + } + return context; + } + + /** Returns the current one-based retry attempt. */ + public int getAttempt() { + return attempt; + } + + /** Attaches retry metadata for the duration of the returned scope. */ + public static SafeCloseable attach(int attempt) { + var previous = CURRENT.get(); + CURRENT.set(new WithRetryContext(attempt)); + return () -> restore(previous); + } + + private static void restore(WithRetryContext previous) { + if (previous == null) { + CURRENT.remove(); + } else { + CURRENT.set(previous); + } + } + } + /** Configuration for replay-safe retry operations. */ public static final class WithRetryConfig { private final RetryStrategy retryStrategy; diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationTest.java index 47ce4c22c..2cccdd4df 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableMapOperationTest.java @@ -25,6 +25,7 @@ import software.amazon.lambda.durable.model.MapResult; import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.operation.DurableMapOperation; +import software.amazon.lambda.durable.operation.DurableMapOperation.MapItemContext; class DurableMapOperationTest { @AfterEach diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationTest.java index 12395b6d5..f343b5f5a 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForCallbackOperationTest.java @@ -23,6 +23,7 @@ import software.amazon.lambda.durable.extension.ExtensionStepFunction; import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation; +import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation.WaitForCallbackContext; class DurableWaitForCallbackOperationTest { @AfterEach diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationTest.java index ab13ced6d..bc2a9986c 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableWithRetryOperationTest.java @@ -20,6 +20,7 @@ import software.amazon.lambda.durable.extension.ExtensionOperation; import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.operation.DurableWithRetryOperation; +import software.amazon.lambda.durable.operation.DurableWithRetryOperation.WithRetryContext; class DurableWithRetryOperationTest { @AfterEach diff --git a/sdk/src/test/java/software/amazon/lambda/durable/OperationContextStorageTest.java b/sdk/src/test/java/software/amazon/lambda/durable/OperationContextTest.java similarity index 89% rename from sdk/src/test/java/software/amazon/lambda/durable/OperationContextStorageTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/OperationContextTest.java index c47b6e072..e994e03a3 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/OperationContextStorageTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/OperationContextTest.java @@ -7,8 +7,11 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.operation.DurableMapOperation.MapItemContext; +import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation.WaitForCallbackContext; +import software.amazon.lambda.durable.operation.DurableWithRetryOperation.WithRetryContext; -class OperationContextStorageTest { +class OperationContextTest { @Test void mapItemContextRestoresNestedScope() { assertThrows(IllegalStateException.class, MapItemContext::getCurrentContext); From 480121076bc369a6e662fe5684ec33c87c957c3b Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Mon, 10 Aug 2026 11:41:10 -0700 Subject: [PATCH 27/40] refactor: simplify extension operation internals --- docs/adr/006-custom-extension-operations.md | 3 +- docs/advanced/extensions.md | 9 +- .../extension/ExtensionContextConfig.java | 34 ++- .../extension/ExtensionOperationImpl.java | 6 +- .../operation/DurableContextOperation.java | 3 +- .../operation/DurableMapOperation.java | 13 +- .../operation/DurableParallelOperation.java | 244 +++++++++++++++++ .../DurableWaitForCallbackOperation.java | 5 +- .../DurableWaitForConditionOperation.java | 25 ++ .../operation/DurableWithRetryOperation.java | 5 +- .../operation/OperationConfigAdapters.java | 16 -- .../operation/ParallelOperationFuture.java | 252 ------------------ .../operation/WaitForConditionFuture.java | 31 --- .../primitive/ChildContextPrimitive.java | 4 +- .../durable/DurableOperationFacadeTest.java | 12 +- .../config/ExtensionContextConfigTest.java | 14 +- ...DurableMapOperationImplementationTest.java | 6 +- ...leParallelOperationImplementationTest.java | 6 +- ...orCallbackOperationImplementationTest.java | 2 +- ...rConditionOperationImplementationTest.java | 23 +- ...eWithRetryOperationImplementationTest.java | 2 +- .../primitive/ChildContextPrimitiveTest.java | 9 +- 22 files changed, 362 insertions(+), 362 deletions(-) delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/operation/OperationConfigAdapters.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperationFuture.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionFuture.java diff --git a/docs/adr/006-custom-extension-operations.md b/docs/adr/006-custom-extension-operations.md index 5c1d478d1..c408d4283 100644 --- a/docs/adr/006-custom-extension-operations.md +++ b/docs/adr/006-custom-extension-operations.md @@ -244,7 +244,8 @@ Supported result policies are: On replay, the framework function receives the stored replay state through scoped TLS. This supports large map results and parallel branch reconstruction without exposing checkpoint APIs. -`ExtensionContextConfig` composes the existing `RunInChildContextConfig` and adds extension-only behavior: +`ExtensionContextConfig` directly owns the child context serializer and virtual-context flag, plus extension-only +behavior: - context failure translation - whether the framework function emits user-function plugin events diff --git a/docs/advanced/extensions.md b/docs/advanced/extensions.md index 96a1195d7..bc860ca03 100644 --- a/docs/advanced/extensions.md +++ b/docs/advanced/extensions.md @@ -308,10 +308,11 @@ Use `ExtensionContextResult.completed(result)` when children never need to repla `replayChildrenAboveSize(result, replayState, thresholdBytes)` to replay only when the serialized full result reaches the threshold. Replay metadata is scoped to the framework callback through `ExtensionContextReplayContext`. -`ExtensionContextConfig` also composes `RunInChildContextConfig`, controls framework user-function plugin events, and -can suppress child checkpoints that finish after the parent. If a context fails, the SDK first rethrows a -deserialized original exception, then calls the configured error handler, and finally falls back to -`ChildContextFailedException`. The handler receives read-only context metadata and child-operation summaries. +`ExtensionContextConfig` directly configures the context serializer and whether the context is virtual. It also +controls framework user-function plugin events and can suppress child checkpoints that finish after the parent. If a +context fails, the SDK first rethrows a deserialized original exception, then calls the configured error handler, and +finally falls back to `ChildContextFailedException`. The handler receives read-only context metadata and +child-operation summaries. ## Explicit child contexts diff --git a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextConfig.java index 3ce387d43..86faeb2fa 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextConfig.java @@ -2,27 +2,30 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.extension; -import java.util.Objects; -import software.amazon.lambda.durable.config.RunInChildContextConfig; +import software.amazon.lambda.durable.serde.SerDes; /** Extension-only policies for an advanced CONTEXT primitive. */ public final class ExtensionContextConfig { - private final RunInChildContextConfig childContextConfig; + private final SerDes serDes; + private final boolean virtual; private final ExtensionContextErrorHandler errorHandler; private final boolean emitUserFunctionEvents; private final boolean suppressLateChildCheckpoints; private ExtensionContextConfig(Builder builder) { - childContextConfig = - Objects.requireNonNullElseGet(builder.childContextConfig, () -> RunInChildContextConfig.builder() - .build()); + serDes = builder.serDes; + virtual = builder.virtual; errorHandler = builder.errorHandler; emitUserFunctionEvents = builder.emitUserFunctionEvents; suppressLateChildCheckpoints = builder.suppressLateChildCheckpoints; } - public RunInChildContextConfig childContextConfig() { - return childContextConfig; + public SerDes serDes() { + return serDes; + } + + public boolean isVirtual() { + return virtual; } public ExtensionContextErrorHandler errorHandler() { @@ -39,7 +42,8 @@ public boolean suppressLateChildCheckpoints() { public Builder toBuilder() { return new Builder() - .childContextConfig(childContextConfig) + .serDes(serDes) + .isVirtual(virtual) .errorHandler(errorHandler) .emitUserFunctionEvents(emitUserFunctionEvents) .suppressLateChildCheckpoints(suppressLateChildCheckpoints); @@ -50,15 +54,21 @@ public static Builder builder() { } public static final class Builder { - private RunInChildContextConfig childContextConfig; + private SerDes serDes; + private boolean virtual; private ExtensionContextErrorHandler errorHandler; private boolean emitUserFunctionEvents = true; private boolean suppressLateChildCheckpoints; private Builder() {} - public Builder childContextConfig(RunInChildContextConfig childContextConfig) { - this.childContextConfig = Objects.requireNonNull(childContextConfig, "childContextConfig cannot be null"); + public Builder serDes(SerDes serDes) { + this.serDes = serDes; + return this; + } + + public Builder isVirtual(boolean virtual) { + this.virtual = virtual; return this; } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionOperationImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionOperationImpl.java index 299d53192..546c2d331 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionOperationImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionOperationImpl.java @@ -127,12 +127,10 @@ public DurableFuture runInChildContextAsync( Objects.requireNonNull(function, "function cannot be null"); Objects.requireNonNull(config, "config cannot be null"); claim(); - var childConfig = config.childContextConfig(); - if (childConfig.serDes() == null) { - childConfig = childConfig.toBuilder() + if (config.serDes() == null) { + config = config.toBuilder() .serDes(context.getDurableConfig().getSerDes()) .build(); - config = config.toBuilder().childContextConfig(childConfig).build(); } var operation = new ChildContextPrimitive<>( new OperationIdentifier(operationId, name, OperationType.CONTEXT, subType), diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableContextOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableContextOperation.java index c055e6722..7d0454f36 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableContextOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableContextOperation.java @@ -82,7 +82,8 @@ public static DurableFuture runInChildContextAsync( () -> ExtensionContextResult.replayChildrenAboveSize( function.apply(DurableContext.getCurrentContext()), null, LARGE_RESULT_THRESHOLD), ExtensionContextConfig.builder() - .childContextConfig(OperationConfigAdapters.toLegacy(config)) + .serDes(config.serDes()) + .isVirtual(config.isVirtual()) .build()); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java index eecaf9c26..9e9b7dc5e 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java @@ -20,7 +20,6 @@ import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.CompletionConfig; import software.amazon.lambda.durable.config.NestingType; -import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.extension.ExtensionContext; @@ -173,10 +172,8 @@ private static List> registerItem var context = ExtensionContext.getCurrentContext(); var registeredItems = new ArrayList>(items.size()); var iterationConfig = ExtensionContextConfig.builder() - .childContextConfig(RunInChildContextConfig.builder() - .serDes(config.serDes()) - .isVirtual(config.nestingType() == FLAT) - .build()) + .serDes(config.serDes()) + .isVirtual(config.nestingType() == FLAT) .build(); for (int index = 0; index < items.size(); index++) { @@ -258,10 +255,8 @@ private static MapResult stripMapResult(MapResult result) { private static ExtensionContextConfig parentConfig(MapConfig config, boolean virtualEmptyMap) { return ExtensionContextConfig.builder() - .childContextConfig(RunInChildContextConfig.builder() - .serDes(config.serDes()) - .isVirtual(virtualEmptyMap) - .build()) + .serDes(config.serDes()) + .isVirtual(virtualEmptyMap) .emitUserFunctionEvents(false) .suppressLateChildCheckpoints(true) .build(); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableParallelOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableParallelOperation.java index 0e16f8aa8..95e356a69 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableParallelOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableParallelOperation.java @@ -2,11 +2,28 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.operation; +import static software.amazon.lambda.durable.config.NestingType.FLAT; +import static software.amazon.lambda.durable.model.OperationSubType.PARALLEL; +import static software.amazon.lambda.durable.model.OperationSubType.PARALLEL_BRANCH; +import static software.amazon.lambda.durable.operation.OperationConcurrencyCoordinator.ItemStatus.FAILED; +import static software.amazon.lambda.durable.operation.OperationConcurrencyCoordinator.ItemStatus.SKIPPED; + +import java.util.ArrayList; +import java.util.List; import java.util.Objects; +import java.util.concurrent.CompletableFuture; +import java.util.function.Function; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.ParallelDurableFuture; +import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.CompletionConfig; import software.amazon.lambda.durable.config.NestingType; import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; +import software.amazon.lambda.durable.extension.ExtensionContextResult; +import software.amazon.lambda.durable.model.ParallelResult; import software.amazon.lambda.durable.serde.SerDes; import software.amazon.lambda.durable.util.ParameterValidator; @@ -29,6 +46,233 @@ public static ParallelDurableFuture parallel(ExtensionContext context, String na return new ParallelOperationFuture(context, name, config); } + private static final class ParallelOperationFuture implements ParallelDurableFuture { + private static final int LARGE_RESULT_THRESHOLD = 256 * 1024; + + private final Object lock = new Object(); + private final ParallelConfig config; + private final SerDes defaultSerDes; + private final List> branches = new ArrayList<>(); + private final DurableFuture parentFuture; + private ExtensionContext childContext; + private OperationConcurrencyCoordinator coordinator; + private ParallelResult replayState; + private boolean registrationClosed; + + ParallelOperationFuture(ExtensionContext context, String name, ParallelConfig config) { + this.config = config; + defaultSerDes = context.getDurableConfig().getSerDes(); + var parent = context.reserve(name); + parentFuture = parent.runInChildContextAsync( + PARALLEL.getValue(), + parallelResultType(), + this::executeInChildContext, + parentConfig(defaultSerDes)); + } + + @Override + public DurableFuture branch( + String name, + TypeToken resultType, + Function function, + software.amazon.lambda.durable.config.ParallelBranchConfig config) { + Objects.requireNonNull(resultType, "resultType cannot be null"); + Objects.requireNonNull(function, "function cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + ParameterValidator.validateOperationName(name); + + synchronized (lock) { + ensureRegistrationOpen(); + var definition = new BranchDefinition<>(name, resultType, function, config.toOperationConfig()); + branches.add(definition); + if (coordinator != null) { + registerBranch(definition, branches.size() - 1); + } + return definition.future; + } + } + + @Override + public ParallelResult get() { + closeRegistration(); + return rebuildResult(parentFuture.get()); + } + + @Override + public CompletableFuture completionFuture() { + return parentFuture.completionFuture(); + } + + @Override + public void close() { + if (closeRegistration()) { + parentFuture.get(); + } + } + + private ExtensionContextResult executeInChildContext() { + var replay = ExtensionContextReplayContext.getCurrentContext(); + initializeCoordinator(ExtensionContext.getCurrentContext(), replay); + var completion = replayState == null + ? coordinator.awaitCompletion() + : coordinator.awaitCompletion(expectedCompletion(replayState)); + var result = constructResult(completion); + return ExtensionContextResult.replayChildren(result, result); + } + + private void initializeCoordinator( + ExtensionContext context, ExtensionContextReplayContext replayContext) { + synchronized (lock) { + childContext = context; + replayState = replayContext.isReplayingChildren() ? replayContext.getReplayState() : null; + if (replayContext.isReplayingChildren() && replayState == null) { + throw new IllegalStateException("Missing result in completed Parallel operation"); + } + coordinator = new OperationConcurrencyCoordinator(config.maxConcurrency(), config.completionConfig()); + for (int index = 0; index < branches.size(); index++) { + registerBranch(branches.get(index), index); + } + if (registrationClosed) { + coordinator.closeRegistration(); + } + } + } + + private void registerBranch(BranchDefinition definition, int index) { + var reservation = childContext.reserve(definition.name); + var skipped = shouldSkip(index); + var item = coordinator.register( + () -> reservation.runInChildContextAsync( + PARALLEL_BRANCH.getValue(), + definition.resultType, + () -> ExtensionContextResult.replayChildrenAboveSize( + definition.function.apply(DurableContext.getCurrentContext()), + null, + LARGE_RESULT_THRESHOLD), + branchConfig(definition.config)), + skipped); + definition.future.bind(item.future()); + } + + private boolean shouldSkip(int index) { + return replayState != null + && (replayState.statuses().size() <= index + || replayState.statuses().get(index) == ParallelResult.Status.SKIPPED); + } + + private boolean closeRegistration() { + synchronized (lock) { + if (registrationClosed) { + return false; + } + registrationClosed = true; + if (coordinator != null) { + coordinator.closeRegistration(); + } + return true; + } + } + + private void ensureRegistrationOpen() { + if (registrationClosed) { + throw new IllegalStateException("Cannot add branches after join() has been called"); + } + } + + private ParallelResult rebuildResult(ParallelResult result) { + synchronized (lock) { + if (result == null) { + return null; + } + var statuses = new ArrayList<>(result.statuses()); + while (statuses.size() < branches.size()) { + statuses.add(ParallelResult.Status.SKIPPED); + } + var succeeded = Math.toIntExact(statuses.stream() + .filter(status -> status == ParallelResult.Status.SUCCEEDED) + .count()); + var failed = Math.toIntExact(statuses.stream() + .filter(status -> status == ParallelResult.Status.FAILED) + .count()); + return new ParallelResult( + statuses.size(), + succeeded, + failed, + statuses.size() - succeeded - failed, + result.completionStatus(), + List.copyOf(statuses)); + } + } + + private static ParallelResult constructResult(OperationConcurrencyCoordinator.Completion completion) { + var statuses = completion.items().stream() + .map(item -> item.status() == FAILED + ? ParallelResult.Status.FAILED + : item.status() == SKIPPED + ? ParallelResult.Status.SKIPPED + : ParallelResult.Status.SUCCEEDED) + .toList(); + var succeeded = Math.toIntExact(statuses.stream() + .filter(status -> status == ParallelResult.Status.SUCCEEDED) + .count()); + var failed = Math.toIntExact(statuses.stream() + .filter(status -> status == ParallelResult.Status.FAILED) + .count()); + return new ParallelResult( + statuses.size(), + succeeded, + failed, + statuses.size() - succeeded - failed, + completion.completionDecision().completionStatus(), + statuses); + } + + private ExtensionContextConfig branchConfig(ParallelBranchConfig branchConfig) { + return ExtensionContextConfig.builder() + .serDes(branchConfig.serDes() == null ? defaultSerDes : branchConfig.serDes()) + .isVirtual(config.nestingType() == FLAT) + .build(); + } + + private static OperationConcurrencyCoordinator.ExpectedCompletionStatus expectedCompletion( + ParallelResult replayState) { + return new OperationConcurrencyCoordinator.ExpectedCompletionStatus( + replayState.succeeded() + replayState.failed(), + CompletionConfig.CompletionDecision.complete(replayState.completionStatus())); + } + + private static ExtensionContextConfig parentConfig(SerDes serDes) { + return ExtensionContextConfig.builder() + .serDes(serDes) + .emitUserFunctionEvents(false) + .suppressLateChildCheckpoints(true) + .build(); + } + + private static TypeToken parallelResultType() { + return TypeToken.get(ParallelResult.class); + } + + private static final class BranchDefinition { + private final String name; + private final TypeToken resultType; + private final Function function; + private final ParallelBranchConfig config; + private final DeferredDurableFuture future = new DeferredDurableFuture<>(); + + private BranchDefinition( + String name, + TypeToken resultType, + Function function, + ParallelBranchConfig config) { + this.name = name; + this.resultType = resultType; + this.function = function; + this.config = config; + } + } + } + /** Configuration for durable PARALLEL operations. */ public static final class ParallelConfig { private final int maxConcurrency; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperation.java index d0908b318..31299d9cf 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperation.java @@ -13,7 +13,6 @@ import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.exception.CallbackFailedException; import software.amazon.lambda.durable.exception.CallbackSubmitterException; import software.amazon.lambda.durable.exception.CallbackTimeoutException; @@ -126,9 +125,7 @@ private static ExtensionContextResult executeInChildContext( private static ExtensionContextConfig extensionConfig(WaitForCallbackConfig config) { return ExtensionContextConfig.builder() - .childContextConfig(RunInChildContextConfig.builder() - .serDes(config.stepConfig().serDes()) - .build()) + .serDes(config.stepConfig().serDes()) .errorHandler(DurableWaitForCallbackOperation::translateFailure) .build(); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperation.java index c662141f8..0a04c6e93 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperation.java @@ -3,11 +3,14 @@ package software.amazon.lambda.durable.operation; import java.util.Objects; +import java.util.concurrent.CompletableFuture; import java.util.function.BiFunction; import java.util.function.Function; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.exception.StepFailedException; +import software.amazon.lambda.durable.exception.WaitForConditionFailedException; import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.extension.ExtensionStepConfig; import software.amazon.lambda.durable.extension.ExtensionStepResult; @@ -123,6 +126,28 @@ private static ExtensionStepResult evaluate( return ExtensionStepResult.retry(result.value(), delay); } + private static final class WaitForConditionFuture implements DurableFuture { + private final DurableFuture delegate; + + WaitForConditionFuture(DurableFuture delegate) { + this.delegate = Objects.requireNonNull(delegate, "delegate cannot be null"); + } + + @Override + public T get() { + try { + return delegate.get(); + } catch (StepFailedException e) { + throw new WaitForConditionFailedException(e.getOperation()); + } + } + + @Override + public CompletableFuture completionFuture() { + return delegate.completionFuture(); + } + } + /** Configuration for durable wait-for-condition operations. */ public static final class WaitForConditionConfig { private final WaitForConditionWaitStrategy waitStrategy; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWithRetryOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWithRetryOperation.java index a0d959e4c..af021a94e 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWithRetryOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWithRetryOperation.java @@ -9,7 +9,6 @@ import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.extension.ExtensionContext; @@ -62,9 +61,7 @@ public static DurableFuture withRetryAsync( new TypeToken() {}, () -> ExtensionContextResult.completed(executeRetryLoop(name, operation, config)), ExtensionContextConfig.builder() - .childContextConfig(RunInChildContextConfig.builder() - .isVirtual(!config.wrapInChildContext()) - .build()) + .isVirtual(!config.wrapInChildContext()) .build()); return (DurableFuture) future; } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/OperationConfigAdapters.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/OperationConfigAdapters.java deleted file mode 100644 index 9128f1995..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/OperationConfigAdapters.java +++ /dev/null @@ -1,16 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; - -import software.amazon.lambda.durable.config.RunInChildContextConfig; - -final class OperationConfigAdapters { - private OperationConfigAdapters() {} - - static RunInChildContextConfig toLegacy(DurableContextOperation.RunInChildContextConfig config) { - return RunInChildContextConfig.builder() - .serDes(config.serDes()) - .isVirtual(config.isVirtual()) - .build(); - } -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperationFuture.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperationFuture.java deleted file mode 100644 index 16b235ca8..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperationFuture.java +++ /dev/null @@ -1,252 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; - -import static software.amazon.lambda.durable.config.NestingType.FLAT; -import static software.amazon.lambda.durable.model.OperationSubType.PARALLEL; -import static software.amazon.lambda.durable.model.OperationSubType.PARALLEL_BRANCH; -import static software.amazon.lambda.durable.operation.OperationConcurrencyCoordinator.ItemStatus.FAILED; -import static software.amazon.lambda.durable.operation.OperationConcurrencyCoordinator.ItemStatus.SKIPPED; - -import java.util.ArrayList; -import java.util.List; -import java.util.Objects; -import java.util.concurrent.CompletableFuture; -import java.util.function.Function; -import software.amazon.lambda.durable.DurableContext; -import software.amazon.lambda.durable.DurableFuture; -import software.amazon.lambda.durable.ParallelDurableFuture; -import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.CompletionConfig; -import software.amazon.lambda.durable.config.ParallelBranchConfig; -import software.amazon.lambda.durable.config.RunInChildContextConfig; -import software.amazon.lambda.durable.extension.ExtensionContext; -import software.amazon.lambda.durable.extension.ExtensionContextConfig; -import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; -import software.amazon.lambda.durable.extension.ExtensionContextResult; -import software.amazon.lambda.durable.model.ParallelResult; -import software.amazon.lambda.durable.operation.DurableParallelOperation.ParallelConfig; -import software.amazon.lambda.durable.serde.SerDes; -import software.amazon.lambda.durable.util.ParameterValidator; - -final class ParallelOperationFuture implements ParallelDurableFuture { - private static final int LARGE_RESULT_THRESHOLD = 256 * 1024; - - private final Object lock = new Object(); - private final ParallelConfig config; - private final SerDes defaultSerDes; - private final List> branches = new ArrayList<>(); - private final DurableFuture parentFuture; - private ExtensionContext childContext; - private OperationConcurrencyCoordinator coordinator; - private ParallelResult replayState; - private boolean registrationClosed; - - ParallelOperationFuture(ExtensionContext context, String name, ParallelConfig config) { - this.config = config; - this.defaultSerDes = context.getDurableConfig().getSerDes(); - var parent = context.reserve(name); - this.parentFuture = parent.runInChildContextAsync( - PARALLEL.getValue(), parallelResultType(), this::executeInChildContext, parentConfig(defaultSerDes)); - } - - @Override - public DurableFuture branch( - String name, TypeToken resultType, Function function, ParallelBranchConfig config) { - Objects.requireNonNull(resultType, "resultType cannot be null"); - Objects.requireNonNull(function, "function cannot be null"); - Objects.requireNonNull(config, "config cannot be null"); - ParameterValidator.validateOperationName(name); - - synchronized (lock) { - ensureRegistrationOpen(); - var definition = new BranchDefinition<>(name, resultType, function, config.toOperationConfig()); - branches.add(definition); - if (coordinator != null) { - registerBranch(definition, branches.size() - 1); - } - return definition.future; - } - } - - @Override - public ParallelResult get() { - closeRegistration(); - return rebuildResult(parentFuture.get()); - } - - @Override - public CompletableFuture completionFuture() { - return parentFuture.completionFuture(); - } - - @Override - public void close() { - if (closeRegistration()) { - parentFuture.get(); - } - } - - private ExtensionContextResult executeInChildContext() { - var replay = ExtensionContextReplayContext.getCurrentContext(); - initializeCoordinator(ExtensionContext.getCurrentContext(), replay); - var completion = replayState == null - ? coordinator.awaitCompletion() - : coordinator.awaitCompletion(expectedCompletion(replayState)); - var result = constructResult(completion); - return ExtensionContextResult.replayChildren(result, result); - } - - private void initializeCoordinator( - ExtensionContext context, ExtensionContextReplayContext replayContext) { - synchronized (lock) { - childContext = context; - replayState = replayContext.isReplayingChildren() ? replayContext.getReplayState() : null; - if (replayContext.isReplayingChildren() && replayState == null) { - throw new IllegalStateException("Missing result in completed Parallel operation"); - } - coordinator = new OperationConcurrencyCoordinator(config.maxConcurrency(), config.completionConfig()); - for (int index = 0; index < branches.size(); index++) { - registerBranch(branches.get(index), index); - } - if (registrationClosed) { - coordinator.closeRegistration(); - } - } - } - - private void registerBranch(BranchDefinition definition, int index) { - var reservation = childContext.reserve(definition.name); - var skipped = shouldSkip(index); - var item = coordinator.register( - () -> reservation.runInChildContextAsync( - PARALLEL_BRANCH.getValue(), - definition.resultType, - () -> ExtensionContextResult.replayChildrenAboveSize( - definition.function.apply(DurableContext.getCurrentContext()), - null, - LARGE_RESULT_THRESHOLD), - branchConfig(definition.config)), - skipped); - definition.future.bind(item.future()); - } - - private boolean shouldSkip(int index) { - return replayState != null - && (replayState.statuses().size() <= index - || replayState.statuses().get(index) == ParallelResult.Status.SKIPPED); - } - - private boolean closeRegistration() { - synchronized (lock) { - if (registrationClosed) { - return false; - } - registrationClosed = true; - if (coordinator != null) { - coordinator.closeRegistration(); - } - return true; - } - } - - private void ensureRegistrationOpen() { - if (registrationClosed) { - throw new IllegalStateException("Cannot add branches after join() has been called"); - } - } - - private ParallelResult rebuildResult(ParallelResult result) { - synchronized (lock) { - if (result == null) { - return null; - } - var statuses = new ArrayList<>(result.statuses()); - while (statuses.size() < branches.size()) { - statuses.add(ParallelResult.Status.SKIPPED); - } - var succeeded = Math.toIntExact(statuses.stream() - .filter(status -> status == ParallelResult.Status.SUCCEEDED) - .count()); - var failed = Math.toIntExact(statuses.stream() - .filter(status -> status == ParallelResult.Status.FAILED) - .count()); - return new ParallelResult( - statuses.size(), - succeeded, - failed, - statuses.size() - succeeded - failed, - result.completionStatus(), - List.copyOf(statuses)); - } - } - - private static ParallelResult constructResult(OperationConcurrencyCoordinator.Completion completion) { - var statuses = completion.items().stream() - .map(item -> item.status() == FAILED - ? ParallelResult.Status.FAILED - : item.status() == SKIPPED ? ParallelResult.Status.SKIPPED : ParallelResult.Status.SUCCEEDED) - .toList(); - var succeeded = Math.toIntExact(statuses.stream() - .filter(status -> status == ParallelResult.Status.SUCCEEDED) - .count()); - var failed = Math.toIntExact(statuses.stream() - .filter(status -> status == ParallelResult.Status.FAILED) - .count()); - return new ParallelResult( - statuses.size(), - succeeded, - failed, - statuses.size() - succeeded - failed, - completion.completionDecision().completionStatus(), - statuses); - } - - private ExtensionContextConfig branchConfig(DurableParallelOperation.ParallelBranchConfig branchConfig) { - return ExtensionContextConfig.builder() - .childContextConfig(RunInChildContextConfig.builder() - .serDes(branchConfig.serDes() == null ? defaultSerDes : branchConfig.serDes()) - .isVirtual(config.nestingType() == FLAT) - .build()) - .build(); - } - - private static OperationConcurrencyCoordinator.ExpectedCompletionStatus expectedCompletion( - ParallelResult replayState) { - return new OperationConcurrencyCoordinator.ExpectedCompletionStatus( - replayState.succeeded() + replayState.failed(), - CompletionConfig.CompletionDecision.complete(replayState.completionStatus())); - } - - private static ExtensionContextConfig parentConfig(SerDes serDes) { - return ExtensionContextConfig.builder() - .childContextConfig( - RunInChildContextConfig.builder().serDes(serDes).build()) - .emitUserFunctionEvents(false) - .suppressLateChildCheckpoints(true) - .build(); - } - - private static TypeToken parallelResultType() { - return TypeToken.get(ParallelResult.class); - } - - private static final class BranchDefinition { - private final String name; - private final TypeToken resultType; - private final Function function; - private final DurableParallelOperation.ParallelBranchConfig config; - private final DeferredDurableFuture future = new DeferredDurableFuture<>(); - - private BranchDefinition( - String name, - TypeToken resultType, - Function function, - DurableParallelOperation.ParallelBranchConfig config) { - this.name = name; - this.resultType = resultType; - this.function = function; - this.config = config; - } - } -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionFuture.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionFuture.java deleted file mode 100644 index c3848e74f..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionFuture.java +++ /dev/null @@ -1,31 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; - -import java.util.Objects; -import java.util.concurrent.CompletableFuture; -import software.amazon.lambda.durable.DurableFuture; -import software.amazon.lambda.durable.exception.StepFailedException; -import software.amazon.lambda.durable.exception.WaitForConditionFailedException; - -final class WaitForConditionFuture implements DurableFuture { - private final DurableFuture delegate; - - WaitForConditionFuture(DurableFuture delegate) { - this.delegate = Objects.requireNonNull(delegate, "delegate cannot be null"); - } - - @Override - public T get() { - try { - return delegate.get(); - } catch (StepFailedException e) { - throw new WaitForConditionFailedException(e.getOperation()); - } - } - - @Override - public CompletableFuture completionFuture() { - return delegate.completionFuture(); - } -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/primitive/ChildContextPrimitive.java b/sdk/src/main/java/software/amazon/lambda/durable/primitive/ChildContextPrimitive.java index a60ea4d02..30532d7d2 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/primitive/ChildContextPrimitive.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/primitive/ChildContextPrimitive.java @@ -114,10 +114,10 @@ public ChildContextPrimitive( super( operationIdentifier, resultTypeToken, - config.childContextConfig().serDes(), + config.serDes(), durableContext, parentOperation, - config.childContextConfig().isVirtual()); + config.isVirtual()); this.function = null; this.extensionFunction = function; this.extensionConfig = config; diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableOperationFacadeTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableOperationFacadeTest.java index d1937cbee..05973859d 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableOperationFacadeTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableOperationFacadeTest.java @@ -230,6 +230,7 @@ void durableContextChildContextUsesPrimitiveExtension() { var future = mockStringFuture(); var reservation = mock(ExtensionOperation.class); var context = mockDurableContext(); + var serDes = mock(SerDes.class); when(((ExtensionContext) context).reserve("child")).thenReturn(reservation); when(reservation.runInChildContextAsync( eq(RUN_IN_CHILD_CONTEXT.getValue()), @@ -242,9 +243,18 @@ void durableContextChildContextUsesPrimitiveExtension() { "child", TypeToken.get(String.class), ignored -> "result", - RunInChildContextConfig.builder().build()); + RunInChildContextConfig.builder().serDes(serDes).isVirtual(true).build()); assertSame(future, result); + var extensionConfig = ArgumentCaptor.forClass(ExtensionContextConfig.class); + verify(reservation) + .runInChildContextAsync( + eq(RUN_IN_CHILD_CONTEXT.getValue()), + eq(TypeToken.get(String.class)), + any(ExtensionContextFunction.class), + extensionConfig.capture()); + assertSame(serDes, extensionConfig.getValue().serDes()); + assertTrue(extensionConfig.getValue().isVirtual()); } @Test diff --git a/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionContextConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionContextConfigTest.java index 69589039c..83cec9aed 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionContextConfigTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionContextConfigTest.java @@ -4,20 +4,22 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; import org.junit.jupiter.api.Test; import software.amazon.lambda.durable.extension.ExtensionContextConfig; import software.amazon.lambda.durable.extension.ExtensionContextErrorHandler; +import software.amazon.lambda.durable.serde.JacksonSerDes; class ExtensionContextConfigTest { @Test void builderUsesOrdinaryChildContextDefaults() { var config = ExtensionContextConfig.builder().build(); - assertNotNull(config.childContextConfig()); + assertNull(config.serDes()); + assertFalse(config.isVirtual()); assertNull(config.errorHandler()); assertTrue(config.emitUserFunctionEvents()); assertFalse(config.suppressLateChildCheckpoints()); @@ -25,16 +27,18 @@ void builderUsesOrdinaryChildContextDefaults() { @Test void builderRetainsExtensionPolicies() { - var childConfig = RunInChildContextConfig.builder().isVirtual(true).build(); + var serDes = new JacksonSerDes(); ExtensionContextErrorHandler handler = failure -> new RuntimeException(failure.contextName()); var config = ExtensionContextConfig.builder() - .childContextConfig(childConfig) + .serDes(serDes) + .isVirtual(true) .errorHandler(handler) .emitUserFunctionEvents(false) .suppressLateChildCheckpoints(true) .build(); - assertEquals(childConfig, config.childContextConfig()); + assertSame(serDes, config.serDes()); + assertTrue(config.isVirtual()); assertEquals(handler, config.errorHandler()); assertFalse(config.emitUserFunctionEvents()); assertTrue(config.suppressLateChildCheckpoints()); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java index 976064faa..d475ade2b 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java @@ -63,7 +63,7 @@ void executeBuildsMapAndIterationContextsFromReservations() { verify(parent) .runInChildContextAsync( eq(MAP.getValue()), any(TypeToken.class), function.capture(), parentConfig.capture()); - assertSame(serDes, parentConfig.getValue().childContextConfig().serDes()); + assertSame(serDes, parentConfig.getValue().serDes()); assertFalse(parentConfig.getValue().emitUserFunctionEvents()); assertTrue(parentConfig.getValue().suppressLateChildCheckpoints()); @@ -98,8 +98,8 @@ void executeBuildsMapAndIterationContextsFromReservations() { eq(TypeToken.get(String.class)), any(ExtensionContextFunction.class), iterationConfig.capture()); - assertTrue(iterationConfig.getValue().childContextConfig().isVirtual()); - assertSame(serDes, iterationConfig.getValue().childContextConfig().serDes()); + assertTrue(iterationConfig.getValue().isVirtual()); + assertSame(serDes, iterationConfig.getValue().serDes()); } @SuppressWarnings({"rawtypes", "unchecked"}) diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableParallelOperationImplementationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableParallelOperationImplementationTest.java index 3765fb030..27f47324f 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableParallelOperationImplementationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableParallelOperationImplementationTest.java @@ -75,7 +75,7 @@ void executeBuildsParallelAndBranchContextsFromReservations() { parentConfig.capture()); assertFalse(parentConfig.getValue().emitUserFunctionEvents()); assertTrue(parentConfig.getValue().suppressLateChildCheckpoints()); - assertSame(serDes, parentConfig.getValue().childContextConfig().serDes()); + assertSame(serDes, parentConfig.getValue().serDes()); var child = mock(CurrentContext.class); var first = mock(ExtensionOperation.class); @@ -117,8 +117,8 @@ void executeBuildsParallelAndBranchContextsFromReservations() { eq(TypeToken.get(String.class)), any(ExtensionContextFunction.class), branchConfig.capture()); - assertTrue(branchConfig.getValue().childContextConfig().isVirtual()); - assertSame(serDes, branchConfig.getValue().childContextConfig().serDes()); + assertTrue(branchConfig.getValue().isVirtual()); + assertSame(serDes, branchConfig.getValue().serDes()); } @Test diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperationImplementationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperationImplementationTest.java index fe0f72b74..b9eb342de 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperationImplementationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperationImplementationTest.java @@ -63,7 +63,7 @@ void executeCreatesExistingWaitForCallbackContextTopology() { eq(resultType), any(ExtensionContextFunction.class), contextConfig.capture()); - assertSame(serDes, contextConfig.getValue().childContextConfig().serDes()); + assertSame(serDes, contextConfig.getValue().serDes()); } @Test diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperationImplementationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperationImplementationTest.java index e72d20d45..8ffd6e45a 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperationImplementationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperationImplementationTest.java @@ -113,7 +113,7 @@ void futureTranslatesOnlyFallbackStepFailure() { throw new StepFailedException(operation); }; - var future = new WaitForConditionFuture<>(delegate); + var future = createFuture(delegate); var failure = assertThrows(WaitForConditionFailedException.class, future::get); assertSame(operation, failure.getOperation()); @@ -134,7 +134,26 @@ public CompletableFuture completionFuture() { } }; - assertSame(completion, new WaitForConditionFuture<>(delegate).completionFuture()); + assertSame(completion, createFuture(delegate).completionFuture()); + } + + private DurableFuture createFuture(DurableFuture delegate) { + var context = mock(ExtensionContext.class); + var reservation = mock(ExtensionOperation.class); + var resultType = TypeToken.get(String.class); + when(context.reserve("ready")).thenReturn(reservation); + when(reservation.stepAsync( + eq(OperationSubType.WAIT_FOR_CONDITION.getValue()), + eq(resultType), + any(ExtensionStepFunction.class), + any(ExtensionStepConfig.class))) + .thenReturn(delegate); + return DurableWaitForConditionOperation.waitForConditionAsync( + context, + "ready", + resultType, + (state, step) -> WaitForConditionResult.stopPolling(state), + WaitForConditionConfig.builder().build().toOperationConfig()); } @SuppressWarnings({"rawtypes", "unchecked"}) diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWithRetryOperationImplementationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWithRetryOperationImplementationTest.java index f163fa8e0..566cba23c 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWithRetryOperationImplementationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWithRetryOperationImplementationTest.java @@ -75,7 +75,7 @@ void executePreservesContextTopologyAndDurableBackoff() { any(TypeToken.class), function.capture(), contextConfig.capture()); - assertFalse(contextConfig.getValue().childContextConfig().isVirtual()); + assertFalse(contextConfig.getValue().isVirtual()); var child = mock(CurrentExtensionContext.class); var wait = mock(ExtensionOperation.class); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/primitive/ChildContextPrimitiveTest.java b/sdk/src/test/java/software/amazon/lambda/durable/primitive/ChildContextPrimitiveTest.java index 55220a93a..3a1ba3d6f 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/primitive/ChildContextPrimitiveTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/primitive/ChildContextPrimitiveTest.java @@ -292,8 +292,7 @@ void replayFailedUsesExtensionErrorHandlerWithChildSummaries() { var capturedFailure = new AtomicReference(); var translated = new IllegalStateException("translated"); var config = ExtensionContextConfig.builder() - .childContextConfig( - RunInChildContextConfig.builder().serDes(SERDES).build()) + .serDes(SERDES) .errorHandler(failure -> { capturedFailure.set(failure); return translated; @@ -335,8 +334,7 @@ void replayFailedPrefersReconstructedExceptionOverExtensionHandler() { when(executionManager.getOperationAndUpdateReplayState("1")).thenReturn(failedContext); var handlerCalled = new AtomicBoolean(); var config = ExtensionContextConfig.builder() - .childContextConfig( - RunInChildContextConfig.builder().serDes(SERDES).build()) + .serDes(SERDES) .errorHandler(failure -> { handlerCalled.set(true); return new IllegalStateException("translated"); @@ -358,8 +356,7 @@ void suppressingExtensionContextPropagatesCompletionOwnerToChildContext() { var childContext = mock(DurableContextImpl.class); when(childContext.getDurableConfig()).thenReturn(createConfig()); var config = ExtensionContextConfig.builder() - .childContextConfig( - RunInChildContextConfig.builder().serDes(SERDES).build()) + .serDes(SERDES) .suppressLateChildCheckpoints(true) .build(); var operation = createExtensionOperation(config); From 4583bd508bea573163c94857641b8a926a335c38 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Mon, 10 Aug 2026 12:11:25 -0700 Subject: [PATCH 28/40] refactor: consolidate concurrent operation internals --- docs/adr/006-custom-extension-operations.md | 3 + docs/advanced/extensions.md | 3 + docs/design.md | 5 +- .../durable/config/CompletionConfig.java | 22 + .../lambda/durable/config/MapConfig.java | 4 +- .../lambda/durable/config/NestingType.java | 9 +- .../lambda/durable/config/ParallelConfig.java | 4 +- .../operation/DeferredDurableFuture.java | 54 -- .../DurableConcurrencyOperation.java | 482 ++++++++++++++++++ .../operation/DurableMapOperation.java | 30 +- .../operation/DurableParallelOperation.java | 27 +- .../OperationConcurrencyCoordinator.java | 238 --------- .../operation/DeferredDurableFutureTest.java | 6 +- ...DurableMapOperationImplementationTest.java | 10 +- .../operation/DurableOperationConfigTest.java | 42 +- ...leParallelOperationImplementationTest.java | 5 +- .../OperationConcurrencyCoordinatorTest.java | 9 +- 17 files changed, 591 insertions(+), 362 deletions(-) delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/operation/DeferredDurableFuture.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/operation/DurableConcurrencyOperation.java delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/operation/OperationConcurrencyCoordinator.java diff --git a/docs/adr/006-custom-extension-operations.md b/docs/adr/006-custom-extension-operations.md index c408d4283..e8f116396 100644 --- a/docs/adr/006-custom-extension-operations.md +++ b/docs/adr/006-custom-extension-operations.md @@ -56,6 +56,8 @@ operation classes used by the static APIs. Each built-in operation owns its public nested configuration type, such as `DurableStepOperation.StepConfig` or `DurableMapOperation.MapConfig`. The compatibility types under `software.amazon.lambda.durable.config` convert through `toOperationConfig()` at the `DurableContext` boundary. +Map and parallel share `DurableConcurrencyOperation.CompletionConfig` and +`DurableConcurrencyOperation.NestingType`; their legacy config counterparts convert to those operation-owned types. ### Expose Primitive and Built-In Extension Facades @@ -87,6 +89,7 @@ Expose each built-in extension family through an independently maintained class: | Facade | Operation family | | --- | --- | +| `DurableConcurrencyOperation` | Shared map/parallel completion, nesting, and coordination support | | `DurableMapOperation` | `map`, `mapAsync` | | `DurableParallelOperation` | `parallel` and branch construction | | `DurableWaitForCallbackOperation` | `waitForCallback`, `waitForCallbackAsync` | diff --git a/docs/advanced/extensions.md b/docs/advanced/extensions.md index bc860ca03..af411f361 100644 --- a/docs/advanced/extensions.md +++ b/docs/advanced/extensions.md @@ -85,6 +85,9 @@ The same pattern applies to `DurableInvokeOperation.InvokeConfig`, `DurableWaitForConditionOperation.WaitForConditionConfig`, and `DurableWithRetryOperation.WithRetryConfig`. +Map and parallel extend `DurableConcurrencyOperation` and use its shared +`DurableConcurrencyOperation.CompletionConfig` and `DurableConcurrencyOperation.NestingType` configuration types. + The compatibility types in `software.amazon.lambda.durable.config` remain accepted by `DurableContext`. They can be passed to a static operation through `toOperationConfig()`: diff --git a/docs/design.md b/docs/design.md index 9fbe89042..f3501c872 100644 --- a/docs/design.md +++ b/docs/design.md @@ -294,13 +294,14 @@ software.amazon.lambda.durable │ └── BaseContext # Base interface for DurableContext │ ├── operation/ # Public built-in operation APIs + implementations +│ ├── DurableConcurrencyOperation # Shared map/parallel config, futures, and coordination │ ├── DurableStepOperation # Owns nested StepConfig │ ├── DurableWaitOperation │ ├── DurableInvokeOperation │ ├── DurableCallbackOperation │ ├── DurableContextOperation -│ ├── DurableMapOperation -│ ├── DurableParallelOperation # Owns nested ParallelConfig and ParallelBranchConfig +│ ├── DurableMapOperation # Extends DurableConcurrencyOperation; owns MapConfig +│ ├── DurableParallelOperation # Extends DurableConcurrencyOperation; owns parallel configs │ ├── DurableWaitForCallbackOperation │ ├── DurableWaitForConditionOperation │ └── DurableWithRetryOperation diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/CompletionConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/CompletionConfig.java index 869f39870..ca6ccde93 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/CompletionConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/CompletionConfig.java @@ -5,6 +5,7 @@ import java.util.Objects; import java.util.function.Function; import software.amazon.lambda.durable.model.ConcurrencyCompletionStatus; +import software.amazon.lambda.durable.operation.DurableConcurrencyOperation; /** * Controls when a concurrent operation (map or parallel) completes. @@ -155,6 +156,27 @@ public boolean hasCustomShouldComplete() { return shouldComplete != null; } + /** Converts this compatibility config to the operation-owned config. */ + public DurableConcurrencyOperation.CompletionConfig toOperationConfig() { + if (shouldComplete == null) { + return new DurableConcurrencyOperation.CompletionConfig( + minSuccessful, toleratedFailureCount, toleratedFailurePercentage); + } + return DurableConcurrencyOperation.CompletionConfig.shouldComplete(status -> { + var decision = shouldComplete.apply(new CompletionStatus( + status.successCount(), + status.failureCount(), + status.completedCount(), + status.totalCount(), + status.allItemsRegistered())); + if (decision == null) { + return null; + } + return new DurableConcurrencyOperation.CompletionConfig.CompletionDecision( + decision.shouldComplete(), decision.completionStatus()); + }); + } + private Function thresholdBasedShouldComplete() { return status -> { if (minSuccessful != null) { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/MapConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/MapConfig.java index 75a40b2bf..7e8bc54b5 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/MapConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/MapConfig.java @@ -79,9 +79,9 @@ public Builder toBuilder() { public DurableMapOperation.MapConfig toOperationConfig() { return DurableMapOperation.MapConfig.builder() .maxConcurrency(maxConcurrency()) - .completionConfig(completionConfig()) + .completionConfig(completionConfig().toOperationConfig()) .serDes(serDes()) - .nestingType(nestingType()) + .nestingType(nestingType().toOperationType()) .itemNamer(itemNamer()) .build(); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/NestingType.java b/sdk/src/main/java/software/amazon/lambda/durable/config/NestingType.java index 4749aaa3d..07f903352 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/NestingType.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/NestingType.java @@ -2,6 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.config; +import software.amazon.lambda.durable.operation.DurableConcurrencyOperation; + public enum NestingType { /** * Create CONTEXT operations for each branch/iteration with full checkpointing. Operations within each @@ -17,5 +19,10 @@ public enum NestingType { * - **Cost**: ~30% lower - reduces operation consumption by skipping CONTEXT overhead - **Scale**: Higher maximum * iterations possible within operation limits */ - FLAT, + FLAT; + + /** Converts this compatibility type to the operation-owned type. */ + public DurableConcurrencyOperation.NestingType toOperationType() { + return DurableConcurrencyOperation.NestingType.valueOf(name()); + } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/ParallelConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/ParallelConfig.java index 2cadd8100..41acc8336 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/ParallelConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/ParallelConfig.java @@ -57,8 +57,8 @@ public Builder toBuilder() { public DurableParallelOperation.ParallelConfig toOperationConfig() { return DurableParallelOperation.ParallelConfig.builder() .maxConcurrency(maxConcurrency()) - .completionConfig(completionConfig()) - .nestingType(nestingType()) + .completionConfig(completionConfig().toOperationConfig()) + .nestingType(nestingType().toOperationType()) .build(); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DeferredDurableFuture.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DeferredDurableFuture.java deleted file mode 100644 index 7bf04fc7c..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DeferredDurableFuture.java +++ /dev/null @@ -1,54 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; - -import java.util.Objects; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.atomic.AtomicBoolean; -import software.amazon.lambda.durable.DurableFuture; -import software.amazon.lambda.durable.context.BaseContext; -import software.amazon.lambda.durable.context.BaseContextImpl; - -final class DeferredDurableFuture implements DurableFuture { - private final AtomicBoolean bound = new AtomicBoolean(); - private final CompletableFuture> delegateFuture = new CompletableFuture<>(); - private final CompletableFuture completionSignal = new CompletableFuture<>(); - - void bind(DurableFuture delegate) { - Objects.requireNonNull(delegate, "delegate cannot be null"); - if (!bound.compareAndSet(false, true)) { - throw new IllegalStateException("A deferred durable future can only be bound once"); - } - - delegateFuture.complete(delegate); - delegate.completionFuture().whenComplete((ignored, throwable) -> { - if (throwable == null) { - completionSignal.complete(null); - } else { - completionSignal.completeExceptionally(throwable); - } - }); - } - - @Override - public T get() { - return awaitDelegate().get(); - } - - @Override - public CompletableFuture completionFuture() { - return completionSignal.thenApply(ignored -> null); - } - - boolean isDone() { - return completionSignal.isDone(); - } - - private DurableFuture awaitDelegate() { - var context = BaseContext.getCurrentContext(); - if (context instanceof BaseContextImpl contextImpl) { - return contextImpl.getExecutionManager().awaitFuture(delegateFuture); - } - return delegateFuture.join(); - } -} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableConcurrencyOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableConcurrencyOperation.java new file mode 100644 index 000000000..50a30cff8 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableConcurrencyOperation.java @@ -0,0 +1,482 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.operation; + +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Objects; +import java.util.Queue; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Function; +import java.util.function.Supplier; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.context.BaseContext; +import software.amazon.lambda.durable.context.BaseContextImpl; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.execution.SuspendExecutionException; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.model.ConcurrencyCompletionStatus; +import software.amazon.lambda.durable.serde.SerDes; + +/** Shared implementation and configuration types for durable concurrent operations. */ +public abstract class DurableConcurrencyOperation { + protected static final int LARGE_RESULT_THRESHOLD = 256 * 1024; + + DurableConcurrencyOperation() {} + + protected static ExtensionContextConfig childContextConfig(SerDes serDes, NestingType nestingType) { + return ExtensionContextConfig.builder() + .serDes(serDes) + .isVirtual(nestingType == NestingType.FLAT) + .build(); + } + + protected static ExtensionContextConfig parentContextConfig(SerDes serDes) { + return parentContextConfig(serDes, false); + } + + protected static ExtensionContextConfig parentContextConfig(SerDes serDes, boolean isVirtual) { + return ExtensionContextConfig.builder() + .serDes(serDes) + .isVirtual(isVirtual) + .emitUserFunctionEvents(false) + .suppressLateChildCheckpoints(true) + .build(); + } + + /** Controls when a concurrent operation completes. */ + public record CompletionConfig( + Integer minSuccessful, + Integer toleratedFailureCount, + Double toleratedFailurePercentage, + Function shouldComplete) { + + public CompletionConfig( + Integer minSuccessful, Integer toleratedFailureCount, Double toleratedFailurePercentage) { + this(minSuccessful, toleratedFailureCount, toleratedFailurePercentage, null); + } + + public CompletionConfig { + if (shouldComplete != null + && (minSuccessful != null || toleratedFailureCount != null || toleratedFailurePercentage != null)) { + throw new IllegalArgumentException( + "shouldComplete is mutually exclusive with minSuccessful, toleratedFailureCount, and toleratedFailurePercentage"); + } + } + + /** Live completion progress for a concurrent operation. */ + public record CompletionStatus( + int successCount, int failureCount, int completedCount, int totalCount, boolean allItemsRegistered) { + public CompletionStatus(int successCount, int failureCount, int completedCount, int totalCount) { + this(successCount, failureCount, completedCount, totalCount, completedCount == totalCount); + } + + public CompletionStatus { + if (successCount < 0 || failureCount < 0 || completedCount < 0 || totalCount < 0) { + throw new IllegalArgumentException("completion counts must be non-negative"); + } + if (completedCount != successCount + failureCount) { + throw new IllegalArgumentException("completedCount must equal successCount + failureCount"); + } + if (completedCount > totalCount) { + throw new IllegalArgumentException("completedCount cannot exceed totalCount"); + } + } + + public boolean allCompleted() { + return allItemsRegistered && completedCount == totalCount; + } + } + + /** The completion decision returned by {@link #completionDecisionFunction()}. */ + public record CompletionDecision(boolean shouldComplete, ConcurrencyCompletionStatus completionStatus) { + public CompletionDecision { + if (shouldComplete && completionStatus == null) { + throw new IllegalArgumentException("completionStatus is required when shouldComplete is true"); + } + if (!shouldComplete && completionStatus != null) { + throw new IllegalArgumentException("completionStatus must be null when shouldComplete is false"); + } + } + + public static CompletionDecision complete(ConcurrencyCompletionStatus completionStatus) { + return new CompletionDecision(true, completionStatus); + } + + public static CompletionDecision continueExecution() { + return new CompletionDecision(false, null); + } + + public boolean isSucceeded() { + return shouldComplete && completionStatus.isSucceeded(); + } + } + + /** All items must succeed. Zero failures are tolerated. */ + public static CompletionConfig allSuccessful() { + return new CompletionConfig(null, 0, null); + } + + /** All items run regardless of failures. */ + public static CompletionConfig allCompleted() { + return new CompletionConfig(null, null, null); + } + + /** Complete as soon as the first item succeeds. */ + public static CompletionConfig firstSuccessful() { + return new CompletionConfig(1, null, null); + } + + /** Complete when the specified number of items have succeeded. */ + public static CompletionConfig minSuccessful(int count) { + if (count < 1) { + throw new IllegalArgumentException("minSuccessful must be at least 1, got: " + count); + } + return new CompletionConfig(count, null, null); + } + + /** Complete when more than the specified number of failures have occurred. */ + public static CompletionConfig toleratedFailureCount(int count) { + if (count < 0) { + throw new IllegalArgumentException("toleratedFailureCount must be non-negative, got: " + count); + } + return new CompletionConfig(null, count, null); + } + + /** Complete when the failure percentage exceeds the specified threshold. */ + public static CompletionConfig toleratedFailurePercentage(double percentage) { + if (percentage < 0.0 || percentage > 1.0) { + throw new IllegalArgumentException( + "toleratedFailurePercentage must be between 0.0 and 1.0, got: " + percentage); + } + return new CompletionConfig(null, null, percentage); + } + + /** Complete when the function returns a completing decision. */ + public static CompletionConfig shouldComplete(Function shouldComplete) { + Objects.requireNonNull(shouldComplete, "shouldComplete cannot be null"); + return new CompletionConfig(null, null, null, shouldComplete); + } + + /** Returns the configured completion decision function. */ + public Function completionDecisionFunction() { + return shouldComplete != null ? shouldComplete : thresholdBasedShouldComplete(); + } + + public boolean hasCustomShouldComplete() { + return shouldComplete != null; + } + + private Function thresholdBasedShouldComplete() { + return status -> { + if (minSuccessful != null) { + if (status.successCount() >= minSuccessful) { + return CompletionDecision.complete(ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED); + } + if (status.allItemsRegistered() && status.totalCount() < minSuccessful) { + throw new IllegalStateException("minSuccessful (" + minSuccessful + + ") exceeds the number of registered items (" + status.totalCount() + ")"); + } + } + + var toleratedFailures = toleratedFailureLimit(status.totalCount()); + if (toleratedFailures != null && status.failureCount() > toleratedFailures) { + return CompletionDecision.complete(ConcurrencyCompletionStatus.FAILURE_TOLERANCE_EXCEEDED); + } + + if (status.allCompleted()) { + return CompletionDecision.complete(ConcurrencyCompletionStatus.ALL_COMPLETED); + } + + return CompletionDecision.continueExecution(); + }; + } + + private Integer toleratedFailureLimit(int totalCount) { + if (toleratedFailureCount == null && toleratedFailurePercentage == null) { + return null; + } + var count = toleratedFailureCount != null ? toleratedFailureCount : Integer.MAX_VALUE; + var percentageCount = toleratedFailurePercentage != null + ? (int) Math.floor(totalCount * toleratedFailurePercentage) + : Integer.MAX_VALUE; + return Math.min(count, percentageCount); + } + } + + /** Controls whether each item is represented by a checkpointed child context. */ + public enum NestingType { + NESTED, + FLAT + } + + protected static final class OperationConcurrencyCoordinator { + enum ItemStatus { + PENDING, + RUNNING, + SUCCEEDED, + FAILED, + SKIPPED + } + + record ExpectedCompletionStatus(int completed, CompletionConfig.CompletionDecision completionDecision) { + ExpectedCompletionStatus { + if (completed < 0) { + throw new IllegalArgumentException("completed cannot be negative"); + } + Objects.requireNonNull(completionDecision, "completionDecision cannot be null"); + } + } + + record Completion(CompletionConfig.CompletionDecision completionDecision, List> items) { + Completion { + Objects.requireNonNull(completionDecision, "completionDecision cannot be null"); + items = List.copyOf(items); + } + } + + static final class Item { + private final Supplier> launcher; + private final DeferredDurableFuture future = new DeferredDurableFuture<>(); + private volatile ItemStatus status; + + private Item(Supplier> launcher, ItemStatus status) { + this.launcher = launcher; + this.status = status; + } + + DurableFuture future() { + return future; + } + + ItemStatus status() { + return status; + } + } + + private final Object lock = new Object(); + private final int maxConcurrency; + private final Function shouldComplete; + private final List> items = new ArrayList<>(); + private final Queue> pending = new ArrayDeque<>(); + private final Set> running = new LinkedHashSet<>(); + private CompletableFuture changed = new CompletableFuture<>(); + private boolean registrationClosed; + private int succeeded; + private int failed; + + OperationConcurrencyCoordinator(int maxConcurrency, CompletionConfig completionConfig) { + if (maxConcurrency < 1) { + throw new IllegalArgumentException("maxConcurrency must be at least 1"); + } + this.maxConcurrency = maxConcurrency; + this.shouldComplete = Objects.requireNonNull(completionConfig, "completionConfig cannot be null") + .completionDecisionFunction(); + } + + Item register(Supplier> launcher) { + return register(launcher, false); + } + + Item register(Supplier> launcher, boolean skipped) { + Objects.requireNonNull(launcher, "launcher cannot be null"); + synchronized (lock) { + if (registrationClosed) { + throw new IllegalStateException("Cannot register items after registration is closed"); + } + var item = new Item<>(launcher, skipped ? ItemStatus.SKIPPED : ItemStatus.PENDING); + items.add(item); + if (!skipped) { + pending.add(item); + } + notifyChanged(); + return item; + } + } + + void closeRegistration() { + synchronized (lock) { + registrationClosed = true; + notifyChanged(); + } + } + + Completion awaitCompletion() { + return awaitCompletion(null); + } + + Completion awaitCompletion(ExpectedCompletionStatus expectedCompletionStatus) { + while (true) { + DurableFuture[] waiters; + synchronized (lock) { + collectCompletedItems(); + var decision = completionDecision(expectedCompletionStatus); + if (decision != null) { + markIncompleteItemsSkipped(); + return new Completion(decision, items); + } + + launchPendingItems(); + collectCompletedItems(); + decision = completionDecision(expectedCompletionStatus); + if (decision != null) { + markIncompleteItemsSkipped(); + return new Completion(decision, items); + } + if (running.size() < maxConcurrency && !pending.isEmpty()) { + continue; + } + waiters = completionWaiters(); + } + DurableFuture.anyOf(waiters); + } + } + + private void launchPendingItems() { + while (running.size() < maxConcurrency && !pending.isEmpty()) { + var item = pending.remove(); + launch(item); + running.add(item); + } + } + + @SuppressWarnings("unchecked") + private void launch(Item untypedItem) { + var item = (Item) untypedItem; + var delegate = Objects.requireNonNull(item.launcher.get(), "launcher cannot return null"); + item.future.bind(delegate); + item.status = ItemStatus.RUNNING; + } + + private void collectCompletedItems() { + var completed = + running.stream().filter(item -> item.future.isDone()).toList(); + for (var item : completed) { + running.remove(item); + complete(item); + } + } + + private void complete(Item item) { + try { + item.future.get(); + item.status = ItemStatus.SUCCEEDED; + succeeded++; + } catch (SuspendExecutionException | UnrecoverableDurableExecutionException exception) { + throw exception; + } catch (Throwable throwable) { + item.status = ItemStatus.FAILED; + failed++; + } + } + + private CompletionConfig.CompletionDecision completionDecision( + ExpectedCompletionStatus expectedCompletionStatus) { + if (expectedCompletionStatus != null) { + return succeeded + failed >= expectedCompletionStatus.completed() + ? expectedCompletionStatus.completionDecision() + : null; + } + var status = new CompletionConfig.CompletionStatus( + succeeded, failed, succeeded + failed, items.size(), registrationClosed); + var decision = Objects.requireNonNull( + shouldComplete.apply(status), "shouldComplete must return a completion decision"); + return decision.shouldComplete() ? decision : null; + } + + private DurableFuture[] completionWaiters() { + if (changed.isDone()) { + changed = new CompletableFuture<>(); + } + var waiters = new ArrayList>(); + running.stream().map(item -> new CompletionOnlyFuture(item.future)).forEach(waiters::add); + waiters.add(new SignalFuture(changed)); + return waiters.toArray(DurableFuture[]::new); + } + + private void markIncompleteItemsSkipped() { + items.stream() + .filter(item -> item.status == ItemStatus.PENDING || item.status == ItemStatus.RUNNING) + .forEach(item -> item.status = ItemStatus.SKIPPED); + pending.clear(); + running.clear(); + } + + private void notifyChanged() { + changed.complete(null); + } + + private record CompletionOnlyFuture(DurableFuture delegate) implements DurableFuture { + @Override + public Void get() { + return null; + } + + @Override + public CompletableFuture completionFuture() { + return delegate.completionFuture(); + } + } + + private record SignalFuture(CompletableFuture signal) implements DurableFuture { + @Override + public Void get() { + signal.join(); + return null; + } + + @Override + public CompletableFuture completionFuture() { + return signal.thenApply(ignored -> null); + } + } + } + + protected static final class DeferredDurableFuture implements DurableFuture { + private final AtomicBoolean bound = new AtomicBoolean(); + private final CompletableFuture> delegateFuture = new CompletableFuture<>(); + private final CompletableFuture completionSignal = new CompletableFuture<>(); + + void bind(DurableFuture delegate) { + Objects.requireNonNull(delegate, "delegate cannot be null"); + if (!bound.compareAndSet(false, true)) { + throw new IllegalStateException("A deferred durable future can only be bound once"); + } + + delegateFuture.complete(delegate); + delegate.completionFuture().whenComplete((ignored, throwable) -> { + if (throwable == null) { + completionSignal.complete(null); + } else { + completionSignal.completeExceptionally(throwable); + } + }); + } + + @Override + public T get() { + return awaitDelegate().get(); + } + + @Override + public CompletableFuture completionFuture() { + return completionSignal.thenApply(ignored -> null); + } + + boolean isDone() { + return completionSignal.isDone(); + } + + private DurableFuture awaitDelegate() { + var context = BaseContext.getCurrentContext(); + if (context instanceof BaseContextImpl contextImpl) { + return contextImpl.getExecutionManager().awaitFuture(delegateFuture); + } + return delegateFuture.join(); + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java index 9e9b7dc5e..955e93610 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java @@ -2,11 +2,10 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.operation; -import static software.amazon.lambda.durable.config.NestingType.FLAT; import static software.amazon.lambda.durable.model.OperationSubType.MAP; import static software.amazon.lambda.durable.model.OperationSubType.MAP_ITERATION; -import static software.amazon.lambda.durable.operation.OperationConcurrencyCoordinator.ItemStatus.FAILED; -import static software.amazon.lambda.durable.operation.OperationConcurrencyCoordinator.ItemStatus.SKIPPED; +import static software.amazon.lambda.durable.operation.DurableConcurrencyOperation.OperationConcurrencyCoordinator.ItemStatus.FAILED; +import static software.amazon.lambda.durable.operation.DurableConcurrencyOperation.OperationConcurrencyCoordinator.ItemStatus.SKIPPED; import java.util.ArrayList; import java.util.Collection; @@ -18,12 +17,9 @@ import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.CompletionConfig; -import software.amazon.lambda.durable.config.NestingType; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.extension.ExtensionContext; -import software.amazon.lambda.durable.extension.ExtensionContextConfig; import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; import software.amazon.lambda.durable.extension.ExtensionContextResult; import software.amazon.lambda.durable.model.MapResult; @@ -33,9 +29,7 @@ import software.amazon.lambda.durable.util.ParameterValidator; /** Context-free static facade and canonical implementation of durable MAP operations. */ -public final class DurableMapOperation { - private static final int LARGE_RESULT_THRESHOLD = 256 * 1024; - +public final class DurableMapOperation extends DurableConcurrencyOperation { private DurableMapOperation() {} public static MapResult map( @@ -110,7 +104,7 @@ public static DurableFuture> mapAsync( mapResultType(), () -> executeInChildContext( name, itemList, iterationNames, resultType, function, mapConfig, virtualEmptyMap), - parentConfig(mapConfig, virtualEmptyMap)); + parentContextConfig(mapConfig.serDes(), virtualEmptyMap)); } private static DurableContext.MapFunction adapt(Function function) { @@ -171,10 +165,7 @@ private static List> registerItem MapResult replayState) { var context = ExtensionContext.getCurrentContext(); var registeredItems = new ArrayList>(items.size()); - var iterationConfig = ExtensionContextConfig.builder() - .serDes(config.serDes()) - .isVirtual(config.nestingType() == FLAT) - .build(); + var iterationConfig = childContextConfig(config.serDes(), config.nestingType()); for (int index = 0; index < items.size(); index++) { var item = items.get(index); @@ -253,15 +244,6 @@ private static MapResult stripMapResult(MapResult result) { result.completionReason()); } - private static ExtensionContextConfig parentConfig(MapConfig config, boolean virtualEmptyMap) { - return ExtensionContextConfig.builder() - .serDes(config.serDes()) - .isVirtual(virtualEmptyMap) - .emitUserFunctionEvents(false) - .suppressLateChildCheckpoints(true) - .build(); - } - private static void validateMinSuccessful(List items, MapConfig config) { var completionConfig = config.completionConfig(); if (!completionConfig.hasCustomShouldComplete() @@ -330,7 +312,7 @@ private MapConfig(Builder builder) { serDes = builder.serDes; nestingType = Objects.requireNonNullElse(builder.nestingType, NestingType.NESTED); itemNamer = builder.itemNamer; - if (itemNamer != null && nestingType == FLAT) { + if (itemNamer != null && nestingType == NestingType.FLAT) { throw new IllegalArgumentException("itemNamer is not supported with FLAT map nesting"); } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableParallelOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableParallelOperation.java index 95e356a69..4b25dd7d7 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableParallelOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableParallelOperation.java @@ -2,11 +2,10 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.operation; -import static software.amazon.lambda.durable.config.NestingType.FLAT; import static software.amazon.lambda.durable.model.OperationSubType.PARALLEL; import static software.amazon.lambda.durable.model.OperationSubType.PARALLEL_BRANCH; -import static software.amazon.lambda.durable.operation.OperationConcurrencyCoordinator.ItemStatus.FAILED; -import static software.amazon.lambda.durable.operation.OperationConcurrencyCoordinator.ItemStatus.SKIPPED; +import static software.amazon.lambda.durable.operation.DurableConcurrencyOperation.OperationConcurrencyCoordinator.ItemStatus.FAILED; +import static software.amazon.lambda.durable.operation.DurableConcurrencyOperation.OperationConcurrencyCoordinator.ItemStatus.SKIPPED; import java.util.ArrayList; import java.util.List; @@ -17,8 +16,6 @@ import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.ParallelDurableFuture; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.CompletionConfig; -import software.amazon.lambda.durable.config.NestingType; import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.extension.ExtensionContextConfig; import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; @@ -28,7 +25,7 @@ import software.amazon.lambda.durable.util.ParameterValidator; /** Context-free static facade and canonical implementation of durable PARALLEL operations. */ -public final class DurableParallelOperation { +public final class DurableParallelOperation extends DurableConcurrencyOperation { private DurableParallelOperation() {} public static ParallelDurableFuture parallel(String name) { @@ -47,8 +44,6 @@ public static ParallelDurableFuture parallel(ExtensionContext context, String na } private static final class ParallelOperationFuture implements ParallelDurableFuture { - private static final int LARGE_RESULT_THRESHOLD = 256 * 1024; - private final Object lock = new Object(); private final ParallelConfig config; private final SerDes defaultSerDes; @@ -67,7 +62,7 @@ private static final class ParallelOperationFuture implements ParallelDurableFut PARALLEL.getValue(), parallelResultType(), this::executeInChildContext, - parentConfig(defaultSerDes)); + parentContextConfig(defaultSerDes)); } @Override @@ -228,10 +223,8 @@ private static ParallelResult constructResult(OperationConcurrencyCoordinator.Co } private ExtensionContextConfig branchConfig(ParallelBranchConfig branchConfig) { - return ExtensionContextConfig.builder() - .serDes(branchConfig.serDes() == null ? defaultSerDes : branchConfig.serDes()) - .isVirtual(config.nestingType() == FLAT) - .build(); + var serDes = branchConfig.serDes() == null ? defaultSerDes : branchConfig.serDes(); + return childContextConfig(serDes, config.nestingType()); } private static OperationConcurrencyCoordinator.ExpectedCompletionStatus expectedCompletion( @@ -241,14 +234,6 @@ private static OperationConcurrencyCoordinator.ExpectedCompletionStatus expected CompletionConfig.CompletionDecision.complete(replayState.completionStatus())); } - private static ExtensionContextConfig parentConfig(SerDes serDes) { - return ExtensionContextConfig.builder() - .serDes(serDes) - .emitUserFunctionEvents(false) - .suppressLateChildCheckpoints(true) - .build(); - } - private static TypeToken parallelResultType() { return TypeToken.get(ParallelResult.class); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/OperationConcurrencyCoordinator.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/OperationConcurrencyCoordinator.java deleted file mode 100644 index 264fb8af3..000000000 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/OperationConcurrencyCoordinator.java +++ /dev/null @@ -1,238 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.operation; - -import java.util.ArrayDeque; -import java.util.ArrayList; -import java.util.LinkedHashSet; -import java.util.List; -import java.util.Objects; -import java.util.Queue; -import java.util.Set; -import java.util.concurrent.CompletableFuture; -import java.util.function.Function; -import java.util.function.Supplier; -import software.amazon.lambda.durable.DurableFuture; -import software.amazon.lambda.durable.config.CompletionConfig; -import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; -import software.amazon.lambda.durable.execution.SuspendExecutionException; - -final class OperationConcurrencyCoordinator { - enum ItemStatus { - PENDING, - RUNNING, - SUCCEEDED, - FAILED, - SKIPPED - } - - record ExpectedCompletionStatus(int completed, CompletionConfig.CompletionDecision completionDecision) { - ExpectedCompletionStatus { - if (completed < 0) { - throw new IllegalArgumentException("completed cannot be negative"); - } - Objects.requireNonNull(completionDecision, "completionDecision cannot be null"); - } - } - - record Completion(CompletionConfig.CompletionDecision completionDecision, List> items) { - Completion { - Objects.requireNonNull(completionDecision, "completionDecision cannot be null"); - items = List.copyOf(items); - } - } - - static final class Item { - private final Supplier> launcher; - private final DeferredDurableFuture future = new DeferredDurableFuture<>(); - private volatile ItemStatus status; - - private Item(Supplier> launcher, ItemStatus status) { - this.launcher = launcher; - this.status = status; - } - - DurableFuture future() { - return future; - } - - ItemStatus status() { - return status; - } - } - - private final Object lock = new Object(); - private final int maxConcurrency; - private final Function shouldComplete; - private final List> items = new ArrayList<>(); - private final Queue> pending = new ArrayDeque<>(); - private final Set> running = new LinkedHashSet<>(); - private CompletableFuture changed = new CompletableFuture<>(); - private boolean registrationClosed; - private int succeeded; - private int failed; - - OperationConcurrencyCoordinator(int maxConcurrency, CompletionConfig completionConfig) { - if (maxConcurrency < 1) { - throw new IllegalArgumentException("maxConcurrency must be at least 1"); - } - this.maxConcurrency = maxConcurrency; - this.shouldComplete = Objects.requireNonNull(completionConfig, "completionConfig cannot be null") - .completionDecisionFunction(); - } - - Item register(Supplier> launcher) { - return register(launcher, false); - } - - Item register(Supplier> launcher, boolean skipped) { - Objects.requireNonNull(launcher, "launcher cannot be null"); - synchronized (lock) { - if (registrationClosed) { - throw new IllegalStateException("Cannot register items after registration is closed"); - } - var item = new Item<>(launcher, skipped ? ItemStatus.SKIPPED : ItemStatus.PENDING); - items.add(item); - if (!skipped) { - pending.add(item); - } - notifyChanged(); - return item; - } - } - - void closeRegistration() { - synchronized (lock) { - registrationClosed = true; - notifyChanged(); - } - } - - Completion awaitCompletion() { - return awaitCompletion(null); - } - - Completion awaitCompletion(ExpectedCompletionStatus expectedCompletionStatus) { - while (true) { - DurableFuture[] waiters; - synchronized (lock) { - collectCompletedItems(); - var decision = completionDecision(expectedCompletionStatus); - if (decision != null) { - markIncompleteItemsSkipped(); - return new Completion(decision, items); - } - - launchPendingItems(); - collectCompletedItems(); - decision = completionDecision(expectedCompletionStatus); - if (decision != null) { - markIncompleteItemsSkipped(); - return new Completion(decision, items); - } - if (running.size() < maxConcurrency && !pending.isEmpty()) { - continue; - } - waiters = completionWaiters(); - } - DurableFuture.anyOf(waiters); - } - } - - private void launchPendingItems() { - while (running.size() < maxConcurrency && !pending.isEmpty()) { - var item = pending.remove(); - launch(item); - running.add(item); - } - } - - @SuppressWarnings("unchecked") - private void launch(Item untypedItem) { - var item = (Item) untypedItem; - var delegate = Objects.requireNonNull(item.launcher.get(), "launcher cannot return null"); - item.future.bind(delegate); - item.status = ItemStatus.RUNNING; - } - - private void collectCompletedItems() { - var completed = running.stream().filter(item -> item.future.isDone()).toList(); - for (var item : completed) { - running.remove(item); - complete(item); - } - } - - private void complete(Item item) { - try { - item.future.get(); - item.status = ItemStatus.SUCCEEDED; - succeeded++; - } catch (SuspendExecutionException | UnrecoverableDurableExecutionException exception) { - throw exception; - } catch (Throwable throwable) { - item.status = ItemStatus.FAILED; - failed++; - } - } - - private CompletionConfig.CompletionDecision completionDecision(ExpectedCompletionStatus expectedCompletionStatus) { - if (expectedCompletionStatus != null) { - return succeeded + failed >= expectedCompletionStatus.completed() - ? expectedCompletionStatus.completionDecision() - : null; - } - var status = new CompletionConfig.CompletionStatus( - succeeded, failed, succeeded + failed, items.size(), registrationClosed); - var decision = Objects.requireNonNull( - shouldComplete.apply(status), "shouldComplete must return a completion decision"); - return decision.shouldComplete() ? decision : null; - } - - private DurableFuture[] completionWaiters() { - if (changed.isDone()) { - changed = new CompletableFuture<>(); - } - var waiters = new ArrayList>(); - running.stream().map(item -> new CompletionOnlyFuture(item.future)).forEach(waiters::add); - waiters.add(new SignalFuture(changed)); - return waiters.toArray(DurableFuture[]::new); - } - - private void markIncompleteItemsSkipped() { - items.stream() - .filter(item -> item.status == ItemStatus.PENDING || item.status == ItemStatus.RUNNING) - .forEach(item -> item.status = ItemStatus.SKIPPED); - pending.clear(); - running.clear(); - } - - private void notifyChanged() { - changed.complete(null); - } - - private record CompletionOnlyFuture(DurableFuture delegate) implements DurableFuture { - @Override - public Void get() { - return null; - } - - @Override - public CompletableFuture completionFuture() { - return delegate.completionFuture(); - } - } - - private record SignalFuture(CompletableFuture signal) implements DurableFuture { - @Override - public Void get() { - signal.join(); - return null; - } - - @Override - public CompletableFuture completionFuture() { - return signal.thenApply(ignored -> null); - } - } -} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DeferredDurableFutureTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DeferredDurableFutureTest.java index fdaff4bb6..a8cf0cf3d 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/DeferredDurableFutureTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DeferredDurableFutureTest.java @@ -14,7 +14,7 @@ class DeferredDurableFutureTest { @Test void getWaitsForBindingThenDelegates() throws Exception { - var deferred = new DeferredDurableFuture(); + var deferred = new DurableConcurrencyOperation.DeferredDurableFuture(); var getStarted = new CountDownLatch(1); var result = CompletableFuture.supplyAsync(() -> { getStarted.countDown(); @@ -31,7 +31,7 @@ void getWaitsForBindingThenDelegates() throws Exception { @Test void completionSignalObtainedBeforeBindingTracksDelegate() { - var deferred = new DeferredDurableFuture(); + var deferred = new DurableConcurrencyOperation.DeferredDurableFuture(); var completion = deferred.completionFuture(); var delegate = new TestFuture<>("result"); @@ -45,7 +45,7 @@ void completionSignalObtainedBeforeBindingTracksDelegate() { @Test void bindRejectsASecondDelegate() { - var deferred = new DeferredDurableFuture(); + var deferred = new DurableConcurrencyOperation.DeferredDurableFuture(); deferred.bind(new TestFuture<>("first")); var exception = assertThrows(IllegalStateException.class, () -> deferred.bind(new TestFuture<>("second"))); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java index d475ade2b..31ffd1686 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java @@ -21,8 +21,6 @@ import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.MapConfig; -import software.amazon.lambda.durable.config.NestingType; import software.amazon.lambda.durable.context.BaseContextImpl; import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.extension.ExtensionContextConfig; @@ -39,8 +37,10 @@ void executeBuildsMapAndIterationContextsFromReservations() { var parent = mock(ExtensionOperation.class); var parentFuture = mockMapFuture(); var serDes = new JacksonSerDes(); - var config = - MapConfig.builder().serDes(serDes).nestingType(NestingType.FLAT).build(); + var config = DurableMapOperation.MapConfig.builder() + .serDes(serDes) + .nestingType(DurableConcurrencyOperation.NestingType.FLAT) + .build(); when(context.reserve("map")).thenReturn(parent); when(parent.runInChildContextAsync( eq(MAP.getValue()), @@ -55,7 +55,7 @@ void executeBuildsMapAndIterationContextsFromReservations() { List.of("a", "b"), TypeToken.get(String.class), (item, index, child) -> item + index, - config.toOperationConfig()); + config); assertSame(parentFuture, actual); var function = extensionFunction(); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableOperationConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableOperationConfigTest.java index 964ed5cdd..6069e4e2b 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableOperationConfigTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableOperationConfigTest.java @@ -7,6 +7,7 @@ import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.mock; +import static software.amazon.lambda.durable.model.ConcurrencyCompletionStatus.CUSTOM_COMPLETION_SUCCEEDED; import java.lang.reflect.Modifier; import java.time.Duration; @@ -34,6 +35,12 @@ class DurableOperationConfigTest { @Test void operationApisOwnTheirConfigTypes() throws Exception { + assertEquals(DurableConcurrencyOperation.class, DurableMapOperation.class.getSuperclass()); + assertEquals(DurableConcurrencyOperation.class, DurableParallelOperation.class.getSuperclass()); + assertPublicStaticNestedType(DurableConcurrencyOperation.class, "CompletionConfig"); + assertPublicStaticNestedType(DurableConcurrencyOperation.class, "NestingType"); + assertProtectedStaticNestedType(DurableConcurrencyOperation.class, "OperationConcurrencyCoordinator"); + assertProtectedStaticNestedType(DurableConcurrencyOperation.class, "DeferredDurableFuture"); assertOperationConfig(DurableStepOperation.class, "StepConfig", StepConfig.class); assertOperationConfig(DurableInvokeOperation.class, "InvokeConfig", InvokeConfig.class); assertOperationConfig(DurableCallbackOperation.class, "CallbackConfig", CallbackConfig.class); @@ -100,9 +107,9 @@ void legacyCompositeConfigsConvertWithoutLosingValues() throws Exception { .itemNamer(itemNamer) .build()); assertEquals(3, value(map, "maxConcurrency")); - assertSame(completionConfig, value(map, "completionConfig")); + assertEquals(completionConfig.toOperationConfig(), value(map, "completionConfig")); assertSame(serDes, value(map, "serDes")); - assertEquals(NestingType.NESTED, value(map, "nestingType")); + assertEquals(DurableConcurrencyOperation.NestingType.NESTED, value(map, "nestingType")); assertSame(itemNamer, value(map, "itemNamer")); var parallel = convert(ParallelConfig.builder() @@ -111,8 +118,8 @@ void legacyCompositeConfigsConvertWithoutLosingValues() throws Exception { .nestingType(NestingType.FLAT) .build()); assertEquals(2, value(parallel, "maxConcurrency")); - assertSame(completionConfig, value(parallel, "completionConfig")); - assertEquals(NestingType.FLAT, value(parallel, "nestingType")); + assertEquals(completionConfig.toOperationConfig(), value(parallel, "completionConfig")); + assertEquals(DurableConcurrencyOperation.NestingType.FLAT, value(parallel, "nestingType")); var branch = convert(ParallelBranchConfig.builder().serDes(serDes).build()); assertSame(serDes, value(branch, "serDes")); @@ -126,6 +133,21 @@ void legacyCompositeConfigsConvertWithoutLosingValues() throws Exception { assertEquals(true, value(retry, "wrapInChildContext")); } + @Test + void legacyCustomCompletionConfigConvertsStatusAndDecision() { + var legacy = CompletionConfig.shouldComplete(status -> status.successCount() == 2 && status.allItemsRegistered() + ? CompletionConfig.CompletionDecision.complete(CUSTOM_COMPLETION_SUCCEEDED) + : CompletionConfig.CompletionDecision.continueExecution()); + var operationConfig = legacy.toOperationConfig(); + + var decision = operationConfig + .completionDecisionFunction() + .apply(new DurableConcurrencyOperation.CompletionConfig.CompletionStatus(2, 1, 3, 3, true)); + + assertTrue(decision.shouldComplete()); + assertEquals(CUSTOM_COMPLETION_SUCCEEDED, decision.completionStatus()); + } + @Test void legacyStatefulConfigsConvertWithoutLosingValues() throws Exception { var serDes = mock(SerDes.class); @@ -156,6 +178,18 @@ private static void assertOperationConfig(Class operationClass, String nested assertOperationConfig(operationClass, operationClass, nestedName, legacyClass); } + private static void assertPublicStaticNestedType(Class owner, String nestedName) throws Exception { + var nestedClass = Class.forName(owner.getName() + "$" + nestedName); + assertTrue(Modifier.isPublic(nestedClass.getModifiers())); + assertTrue(Modifier.isStatic(nestedClass.getModifiers())); + } + + private static void assertProtectedStaticNestedType(Class owner, String nestedName) throws Exception { + var nestedClass = Class.forName(owner.getName() + "$" + nestedName); + assertTrue(Modifier.isProtected(nestedClass.getModifiers())); + assertTrue(Modifier.isStatic(nestedClass.getModifiers())); + } + private static void assertParallelFutureUsesCompatibilityBranchConfig() { var nestedConfig = DurableParallelOperation.ParallelBranchConfig.class; assertTrue(Arrays.stream(ParallelDurableFuture.class.getMethods()) diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableParallelOperationImplementationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableParallelOperationImplementationTest.java index 27f47324f..aea53dde0 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableParallelOperationImplementationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableParallelOperationImplementationTest.java @@ -24,7 +24,6 @@ import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.NestingType; import software.amazon.lambda.durable.context.BaseContextImpl; import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.extension.ExtensionContextConfig; @@ -58,7 +57,9 @@ void executeBuildsParallelAndBranchContextsFromReservations() { var parallel = DurableParallelOperation.parallel( context, "parallel", - ParallelConfig.builder().nestingType(NestingType.FLAT).build()); + ParallelConfig.builder() + .nestingType(DurableConcurrencyOperation.NestingType.FLAT) + .build()); var firstFuture = parallel.branch("first", String.class, child -> "first"); var secondFuture = parallel.branch("second", String.class, child -> "second"); assertSame(parentCompletion, parallel.completionFuture()); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/OperationConcurrencyCoordinatorTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/OperationConcurrencyCoordinatorTest.java index fcaab2c51..7f8310127 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/OperationConcurrencyCoordinatorTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/OperationConcurrencyCoordinatorTest.java @@ -8,9 +8,9 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import static software.amazon.lambda.durable.model.ConcurrencyCompletionStatus.ALL_COMPLETED; import static software.amazon.lambda.durable.model.ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED; -import static software.amazon.lambda.durable.operation.OperationConcurrencyCoordinator.ItemStatus.FAILED; -import static software.amazon.lambda.durable.operation.OperationConcurrencyCoordinator.ItemStatus.SKIPPED; -import static software.amazon.lambda.durable.operation.OperationConcurrencyCoordinator.ItemStatus.SUCCEEDED; +import static software.amazon.lambda.durable.operation.DurableConcurrencyOperation.OperationConcurrencyCoordinator.ItemStatus.FAILED; +import static software.amazon.lambda.durable.operation.DurableConcurrencyOperation.OperationConcurrencyCoordinator.ItemStatus.SKIPPED; +import static software.amazon.lambda.durable.operation.DurableConcurrencyOperation.OperationConcurrencyCoordinator.ItemStatus.SUCCEEDED; import java.util.List; import java.util.concurrent.CompletableFuture; @@ -19,7 +19,8 @@ import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.Test; import software.amazon.lambda.durable.DurableFuture; -import software.amazon.lambda.durable.config.CompletionConfig; +import software.amazon.lambda.durable.operation.DurableConcurrencyOperation.CompletionConfig; +import software.amazon.lambda.durable.operation.DurableConcurrencyOperation.OperationConcurrencyCoordinator; class OperationConcurrencyCoordinatorTest { @Test From b986357d8ed57f33d2952fb960fd867c6b91b902 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Mon, 10 Aug 2026 12:28:10 -0700 Subject: [PATCH 29/40] refactor: decouple extension step configuration --- docs/adr/006-custom-extension-operations.md | 4 ++ docs/advanced/extensions.md | 11 +++ docs/design.md | 2 +- .../ExtensionOperationIntegrationTest.java | 30 ++++++++ .../extension/ExtensionStepConfig.java | 55 ++++++++++++++- .../operation/DurableStepOperation.java | 20 +++++- .../durable/primitive/StepPrimitive.java | 3 +- .../durable/DurableOperationFacadeTest.java | 40 +++++++++++ .../config/ExtensionStepConfigTest.java | 32 --------- .../extension/ExtensionStepConfigTest.java | 69 +++++++++++++++++++ .../StatefulExtensionStepPrimitiveTest.java | 2 +- 11 files changed, 228 insertions(+), 40 deletions(-) delete mode 100644 sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionStepConfigTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/extension/ExtensionStepConfigTest.java diff --git a/docs/adr/006-custom-extension-operations.md b/docs/adr/006-custom-extension-operations.md index e8f116396..cec9dbb44 100644 --- a/docs/adr/006-custom-extension-operations.md +++ b/docs/adr/006-custom-extension-operations.md @@ -205,6 +205,10 @@ DurableFuture waitAsync(String subType, Duration duration); ExtensionContextConfig config); ``` +`ExtensionStepConfig` owns its nested `StepSemantics`, `RetryStrategy`, and `RetryDecision` contracts. Built-in step +operations adapt the customer-facing config and retry types at the operation boundary, keeping the extension SPI +independent of those packages. + The primitive selector determines the backend operation type. The string controls only the subtype recorded in checkpoints, replay validation, plugins, logs, and error metadata. diff --git a/docs/advanced/extensions.md b/docs/advanced/extensions.md index af411f361..ced4adcd6 100644 --- a/docs/advanced/extensions.md +++ b/docs/advanced/extensions.md @@ -281,6 +281,17 @@ var result = ExtensionContext.getCurrentContext() The function may return only `ExtensionStepResult.succeed(value)` or `ExtensionStepResult.retry(state, delay)`. Retry state uses the configured `SerDes`; attempt metadata remains available through `StepContext.getCurrentContext()`. Thrown exceptions follow the normal STEP failure path. +`ExtensionStepConfig` owns its retry contracts, so extension libraries can configure exception retries and delivery +semantics without depending on the customer-facing config or retry packages: + +```java +ExtensionStepConfig.builder() + .retryStrategy((error, attempt) -> attempt < 3 + ? ExtensionStepConfig.RetryDecision.retry(Duration.ofSeconds(1)) + : ExtensionStepConfig.RetryDecision.fail()) + .semanticsPerRetry(ExtensionStepConfig.StepSemantics.AT_MOST_ONCE_PER_RETRY) + .build(); +``` ## Configurable extension contexts diff --git a/docs/design.md b/docs/design.md index f3501c872..ac9760df9 100644 --- a/docs/design.md +++ b/docs/design.md @@ -316,7 +316,7 @@ software.amazon.lambda.durable ├── extension/ # Public SPI for extension authors │ ├── ExtensionContext │ ├── ExtensionOperation -│ ├── ExtensionStepConfig +│ ├── ExtensionStepConfig # Owns extension StepSemantics and retry contracts │ ├── ExtensionStepResult │ ├── ExtensionContextConfig │ └── ExtensionContextResult diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java index e3108a556..730ed9bcb 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java @@ -189,6 +189,36 @@ void statefulExtensionStepCheckpointsStateAcrossRetries() { assertEquals(3, result.getOperation("stateful").getAttempt()); } + @Test + void extensionStepRetriesExceptionsWithExtensionOwnedStrategy() { + var attempts = new AtomicInteger(); + var runner = + LocalDurableTestRunner.create(String.class, (input, context) -> ExtensionContext.getCurrentContext() + .reserve("retry") + .stepAsync( + "AcmeRetry", + TypeToken.get(String.class), + state -> { + if (attempts.incrementAndGet() == 1) { + throw new IllegalStateException("retry"); + } + return ExtensionStepResult.succeed("done"); + }, + ExtensionStepConfig.builder() + .retryStrategy((error, attempt) -> attempt < 2 + ? ExtensionStepConfig.RetryDecision.retry(Duration.ofSeconds(1)) + : ExtensionStepConfig.RetryDecision.fail()) + .build()) + .get()); + + var result = runner.runUntilComplete("input"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("done", result.getResult(String.class)); + assertEquals(2, attempts.get()); + assertEquals(2, result.getOperation("retry").getAttempt()); + } + @Test void extensionContextExposesStoredReplayStateWhileReplayingChildren() { var replayState = new AtomicReference(); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepConfig.java index 2b0e9527c..341c5daa7 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepConfig.java @@ -2,8 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.extension; -import software.amazon.lambda.durable.config.StepSemantics; -import software.amazon.lambda.durable.retry.RetryStrategy; +import java.time.Duration; import software.amazon.lambda.durable.serde.SerDes; /** @@ -58,6 +57,58 @@ public static Builder builder() { return new Builder<>(); } + /** Delivery semantics for each extension step attempt. */ + public enum StepSemantics { + /** The step may be re-executed if an attempt is interrupted. */ + AT_LEAST_ONCE_PER_RETRY, + + /** The START checkpoint is awaited so an interrupted attempt is not re-executed. */ + AT_MOST_ONCE_PER_RETRY + } + + /** Determines whether a thrown exception should retry the extension step. */ + @FunctionalInterface + public interface RetryStrategy { + /** + * Returns the retry decision for a failed attempt. + * + * @param error the thrown exception + * @param attempt the current one-based attempt number + */ + RetryDecision makeRetryDecision(Throwable error, int attempt); + } + + /** A retry decision and the delay before the next attempt. */ + public static final class RetryDecision { + private final boolean shouldRetry; + private final Duration delay; + + private RetryDecision(boolean shouldRetry, Duration delay) { + this.shouldRetry = shouldRetry; + this.delay = delay != null ? delay : Duration.ZERO; + } + + /** Returns a decision to retry after the supplied delay. */ + public static RetryDecision retry(Duration delay) { + return new RetryDecision(true, delay); + } + + /** Returns a decision to fail without retrying. */ + public static RetryDecision fail() { + return new RetryDecision(false, Duration.ZERO); + } + + /** Returns whether another attempt should run. */ + public boolean shouldRetry() { + return shouldRetry; + } + + /** Returns the delay before the next attempt. */ + public Duration delay() { + return delay; + } + } + /** Builder for {@link ExtensionStepConfig}. */ public static final class Builder { private T initialState; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableStepOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableStepOperation.java index 82cdb3269..8aa52d181 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableStepOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableStepOperation.java @@ -77,11 +77,27 @@ public static DurableFuture stepAsync( ignored -> ExtensionStepResult.succeed(function.apply(StepContext.getCurrentContext())), ExtensionStepConfig.builder() .serDes(config.serDes()) - .retryStrategy(config.retryStrategy()) - .semanticsPerRetry(config.semanticsPerRetry()) + .retryStrategy(adapt(config.retryStrategy())) + .semanticsPerRetry(adapt(config.semanticsPerRetry())) .build()); } + private static ExtensionStepConfig.RetryStrategy adapt(RetryStrategy retryStrategy) { + return (error, attempt) -> { + var decision = retryStrategy.makeRetryDecision(error, attempt); + return decision.shouldRetry() + ? ExtensionStepConfig.RetryDecision.retry(decision.delay()) + : ExtensionStepConfig.RetryDecision.fail(); + }; + } + + private static ExtensionStepConfig.StepSemantics adapt(StepSemantics semantics) { + return switch (semantics) { + case AT_LEAST_ONCE_PER_RETRY -> ExtensionStepConfig.StepSemantics.AT_LEAST_ONCE_PER_RETRY; + case AT_MOST_ONCE_PER_RETRY -> ExtensionStepConfig.StepSemantics.AT_MOST_ONCE_PER_RETRY; + }; + } + /** Configuration for durable STEP operations. */ public static final class StepConfig { private final RetryStrategy retryStrategy; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/primitive/StepPrimitive.java b/sdk/src/main/java/software/amazon/lambda/durable/primitive/StepPrimitive.java index 165b009e0..d141c0864 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/primitive/StepPrimitive.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/primitive/StepPrimitive.java @@ -9,7 +9,6 @@ import software.amazon.awssdk.services.lambda.model.OperationUpdate; import software.amazon.awssdk.services.lambda.model.StepOptions; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.StepSemantics; import software.amazon.lambda.durable.context.BaseContextImpl; import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.exception.DurableOperationException; @@ -236,6 +235,6 @@ public T get() { } private boolean isAtMostOnce() { - return extensionConfig.semanticsPerRetry() == StepSemantics.AT_MOST_ONCE_PER_RETRY; + return extensionConfig.semanticsPerRetry() == ExtensionStepConfig.StepSemantics.AT_MOST_ONCE_PER_RETRY; } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableOperationFacadeTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableOperationFacadeTest.java index 05973859d..79b4c3695 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableOperationFacadeTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableOperationFacadeTest.java @@ -29,6 +29,7 @@ import software.amazon.lambda.durable.config.InvokeConfig; import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.config.StepConfig; +import software.amazon.lambda.durable.config.StepSemantics; import software.amazon.lambda.durable.context.BaseContextImpl; import software.amazon.lambda.durable.extension.ExtensionCallbackConfig; import software.amazon.lambda.durable.extension.ExtensionContext; @@ -44,6 +45,7 @@ import software.amazon.lambda.durable.operation.DurableInvokeOperation; import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.operation.DurableWaitOperation; +import software.amazon.lambda.durable.retry.RetryDecision; import software.amazon.lambda.durable.serde.SerDes; class DurableOperationFacadeTest { @@ -127,6 +129,44 @@ void stepAcceptsContextFreeSupplier() { } } + @Test + void stepAdaptsUserFacingRetryConfigToExtensionSpi() { + var context = mock(ExtensionContext.class); + var reservation = mock(ExtensionOperation.class); + var future = mockStringFuture(); + var resultType = TypeToken.get(String.class); + when(context.reserve("step")).thenReturn(reservation); + when(reservation.stepAsync( + eq(STEP.getValue()), + eq(resultType), + any(ExtensionStepFunction.class), + any(ExtensionStepConfig.class))) + .thenReturn(future); + var config = DurableStepOperation.StepConfig.builder() + .retryStrategy((error, attempt) -> + attempt == 1 ? RetryDecision.retry(Duration.ofSeconds(3)) : RetryDecision.fail()) + .semanticsPerRetry(StepSemantics.AT_MOST_ONCE_PER_RETRY) + .build(); + + assertSame(future, DurableStepOperation.stepAsync(context, "step", resultType, ignored -> "result", config)); + + var extensionConfig = ArgumentCaptor.forClass(ExtensionStepConfig.class); + verify(reservation) + .stepAsync( + eq(STEP.getValue()), + eq(resultType), + any(ExtensionStepFunction.class), + extensionConfig.capture()); + assertEquals( + ExtensionStepConfig.StepSemantics.AT_MOST_ONCE_PER_RETRY, + extensionConfig.getValue().semanticsPerRetry()); + var retry = extensionConfig.getValue().retryStrategy().makeRetryDecision(new IllegalStateException("retry"), 1); + var fail = extensionConfig.getValue().retryStrategy().makeRetryDecision(new IllegalStateException("fail"), 2); + assertTrue(retry.shouldRetry()); + assertEquals(Duration.ofSeconds(3), retry.delay()); + assertFalse(fail.shouldRetry()); + } + @Test void durableContextStepUsesPrimitiveExtension() { var future = mockStringFuture(); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionStepConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionStepConfigTest.java deleted file mode 100644 index 53370aee0..000000000 --- a/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionStepConfigTest.java +++ /dev/null @@ -1,32 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable.config; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNull; - -import org.junit.jupiter.api.Test; -import software.amazon.lambda.durable.extension.ExtensionStepConfig; -import software.amazon.lambda.durable.serde.JacksonSerDes; - -class ExtensionStepConfigTest { - @Test - void builderDefaultsToNullStateAndSerDes() { - var config = ExtensionStepConfig.builder().build(); - - assertNull(config.initialState()); - assertNull(config.serDes()); - } - - @Test - void builderRetainsStateAndSerDes() { - var serDes = new JacksonSerDes(); - var config = ExtensionStepConfig.builder() - .initialState(42) - .serDes(serDes) - .build(); - - assertEquals(42, config.initialState()); - assertEquals(serDes, config.serDes()); - } -} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/extension/ExtensionStepConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/extension/ExtensionStepConfigTest.java new file mode 100644 index 000000000..9f7a13115 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/extension/ExtensionStepConfigTest.java @@ -0,0 +1,69 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.extension; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.time.Duration; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.serde.JacksonSerDes; + +class ExtensionStepConfigTest { + @Test + void builderDefaultsToNullStateAndSerDes() { + var config = ExtensionStepConfig.builder().build(); + + assertNull(config.initialState()); + assertNull(config.serDes()); + assertNull(config.retryStrategy()); + assertEquals(ExtensionStepConfig.StepSemantics.AT_LEAST_ONCE_PER_RETRY, config.semanticsPerRetry()); + } + + @Test + void builderRetainsConfiguredValues() { + var serDes = new JacksonSerDes(); + ExtensionStepConfig.RetryStrategy retryStrategy = + (error, attempt) -> ExtensionStepConfig.RetryDecision.retry(Duration.ofSeconds(attempt)); + var config = ExtensionStepConfig.builder() + .initialState(42) + .serDes(serDes) + .retryStrategy(retryStrategy) + .semanticsPerRetry(ExtensionStepConfig.StepSemantics.AT_MOST_ONCE_PER_RETRY) + .build(); + + assertEquals(42, config.initialState()); + assertEquals(serDes, config.serDes()); + assertSame(retryStrategy, config.retryStrategy()); + assertEquals(ExtensionStepConfig.StepSemantics.AT_MOST_ONCE_PER_RETRY, config.semanticsPerRetry()); + } + + @Test + void retryDecisionFactoriesExposeExtensionOwnedDecision() { + var retry = ExtensionStepConfig.RetryDecision.retry(Duration.ofSeconds(3)); + var fail = ExtensionStepConfig.RetryDecision.fail(); + + assertTrue(retry.shouldRetry()); + assertEquals(Duration.ofSeconds(3), retry.delay()); + assertFalse(fail.shouldRetry()); + assertEquals(Duration.ZERO, fail.delay()); + } + + @Test + void retryAndSemanticsContractsAreOwnedByExtensionStepConfig() throws Exception { + assertEquals( + ExtensionStepConfig.RetryStrategy.class, + ExtensionStepConfig.class.getMethod("retryStrategy").getReturnType()); + assertEquals( + ExtensionStepConfig.StepSemantics.class, + ExtensionStepConfig.class.getMethod("semanticsPerRetry").getReturnType()); + assertEquals( + ExtensionStepConfig.RetryDecision.class, + ExtensionStepConfig.RetryStrategy.class + .getMethod("makeRetryDecision", Throwable.class, int.class) + .getReturnType()); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java b/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java index 2d0eb6108..a161b6609 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java @@ -39,11 +39,11 @@ import software.amazon.lambda.durable.execution.ThreadContext; import software.amazon.lambda.durable.execution.ThreadType; import software.amazon.lambda.durable.extension.ExtensionStepConfig; +import software.amazon.lambda.durable.extension.ExtensionStepConfig.RetryDecision; import software.amazon.lambda.durable.extension.ExtensionStepFunction; import software.amazon.lambda.durable.extension.ExtensionStepResult; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; -import software.amazon.lambda.durable.retry.RetryDecision; import software.amazon.lambda.durable.serde.JacksonSerDes; class StatefulExtensionStepPrimitiveTest { From d265b80ad9082c261c2f330f1142c6fd994654fe Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Mon, 10 Aug 2026 12:53:53 -0700 Subject: [PATCH 30/40] refactor: unify extension step retry outcomes --- docs/adr/006-custom-extension-operations.md | 7 +-- docs/advanced/extensions.md | 11 +++-- .../ExtensionOperationIntegrationTest.java | 15 ++++-- .../extension/ExtensionStepConfig.java | 46 ++++--------------- .../extension/ExtensionStepResult.java | 13 +++++- .../operation/DurableStepOperation.java | 8 ++-- .../durable/primitive/StepPrimitive.java | 23 +++++----- .../durable/DurableOperationFacadeTest.java | 15 ++++-- .../durable/ExtensionStepResultTest.java | 3 ++ .../extension/ExtensionStepConfigTest.java | 21 ++------- .../StatefulExtensionStepPrimitiveTest.java | 6 ++- 11 files changed, 79 insertions(+), 89 deletions(-) diff --git a/docs/adr/006-custom-extension-operations.md b/docs/adr/006-custom-extension-operations.md index cec9dbb44..ecbe53fcc 100644 --- a/docs/adr/006-custom-extension-operations.md +++ b/docs/adr/006-custom-extension-operations.md @@ -205,9 +205,10 @@ DurableFuture waitAsync(String subType, Duration duration); ExtensionContextConfig config); ``` -`ExtensionStepConfig` owns its nested `StepSemantics`, `RetryStrategy`, and `RetryDecision` contracts. Built-in step -operations adapt the customer-facing config and retry types at the operation boundary, keeping the extension SPI -independent of those packages. +`ExtensionStepConfig` owns its nested `StepSemantics` and `RetryStrategy` contracts. Retry strategies reuse +`ExtensionStepResult.retry(state, delay)` as their retry decision, so stateful continuations and failed attempts share +one retry representation. Built-in step operations adapt the customer-facing config and retry types at the operation +boundary, keeping the extension SPI independent of those packages. The primitive selector determines the backend operation type. The string controls only the subtype recorded in checkpoints, replay validation, plugins, logs, and error metadata. diff --git a/docs/advanced/extensions.md b/docs/advanced/extensions.md index ced4adcd6..affddd037 100644 --- a/docs/advanced/extensions.md +++ b/docs/advanced/extensions.md @@ -281,14 +281,15 @@ var result = ExtensionContext.getCurrentContext() The function may return only `ExtensionStepResult.succeed(value)` or `ExtensionStepResult.retry(state, delay)`. Retry state uses the configured `SerDes`; attempt metadata remains available through `StepContext.getCurrentContext()`. Thrown exceptions follow the normal STEP failure path. -`ExtensionStepConfig` owns its retry contracts, so extension libraries can configure exception retries and delivery -semantics without depending on the customer-facing config or retry packages: +`ExtensionStepConfig` owns its retry strategy, which returns the same retry outcome used by stateful continuations. +Extension libraries can therefore configure exception retries and delivery semantics without depending on the +customer-facing config or retry packages: ```java ExtensionStepConfig.builder() - .retryStrategy((error, attempt) -> attempt < 3 - ? ExtensionStepConfig.RetryDecision.retry(Duration.ofSeconds(1)) - : ExtensionStepConfig.RetryDecision.fail()) + .retryStrategy((error, state, attempt) -> attempt < 3 + ? ExtensionStepResult.retry(state, Duration.ofSeconds(1)) + : ExtensionStepResult.doNotRetry()) .semanticsPerRetry(ExtensionStepConfig.StepSemantics.AT_MOST_ONCE_PER_RETRY) .build(); ``` diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java index 730ed9bcb..d8e98bf56 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java @@ -192,6 +192,8 @@ void statefulExtensionStepCheckpointsStateAcrossRetries() { @Test void extensionStepRetriesExceptionsWithExtensionOwnedStrategy() { var attempts = new AtomicInteger(); + var failedState = new AtomicReference(); + var resumedState = new AtomicReference(); var runner = LocalDurableTestRunner.create(String.class, (input, context) -> ExtensionContext.getCurrentContext() .reserve("retry") @@ -202,12 +204,17 @@ void extensionStepRetriesExceptionsWithExtensionOwnedStrategy() { if (attempts.incrementAndGet() == 1) { throw new IllegalStateException("retry"); } + resumedState.set(state); return ExtensionStepResult.succeed("done"); }, ExtensionStepConfig.builder() - .retryStrategy((error, attempt) -> attempt < 2 - ? ExtensionStepConfig.RetryDecision.retry(Duration.ofSeconds(1)) - : ExtensionStepConfig.RetryDecision.fail()) + .initialState("initial") + .retryStrategy((error, state, attempt) -> { + failedState.set(state); + return attempt < 2 + ? ExtensionStepResult.retry("retried", Duration.ofSeconds(1)) + : ExtensionStepResult.doNotRetry(); + }) .build()) .get()); @@ -217,6 +224,8 @@ void extensionStepRetriesExceptionsWithExtensionOwnedStrategy() { assertEquals("done", result.getResult(String.class)); assertEquals(2, attempts.get()); assertEquals(2, result.getOperation("retry").getAttempt()); + assertEquals("initial", failedState.get()); + assertEquals("retried", resumedState.get()); } @Test diff --git a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepConfig.java index 341c5daa7..02f022abf 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepConfig.java @@ -2,7 +2,6 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.extension; -import java.time.Duration; import software.amazon.lambda.durable.serde.SerDes; /** @@ -13,7 +12,7 @@ public final class ExtensionStepConfig { private final T initialState; private final SerDes serDes; - private final RetryStrategy retryStrategy; + private final RetryStrategy retryStrategy; private final StepSemantics semanticsPerRetry; private ExtensionStepConfig(Builder builder) { @@ -34,7 +33,7 @@ public SerDes serDes() { } /** Returns the exception retry strategy, or {@code null} when thrown exceptions are terminal. */ - public RetryStrategy retryStrategy() { + public RetryStrategy retryStrategy() { return retryStrategy; } @@ -68,52 +67,23 @@ public enum StepSemantics { /** Determines whether a thrown exception should retry the extension step. */ @FunctionalInterface - public interface RetryStrategy { + public interface RetryStrategy { /** * Returns the retry decision for a failed attempt. * * @param error the thrown exception + * @param state state supplied to the failed attempt * @param attempt the current one-based attempt number + * @return a retry outcome with the next state and delay, or a do-not-retry decision */ - RetryDecision makeRetryDecision(Throwable error, int attempt); - } - - /** A retry decision and the delay before the next attempt. */ - public static final class RetryDecision { - private final boolean shouldRetry; - private final Duration delay; - - private RetryDecision(boolean shouldRetry, Duration delay) { - this.shouldRetry = shouldRetry; - this.delay = delay != null ? delay : Duration.ZERO; - } - - /** Returns a decision to retry after the supplied delay. */ - public static RetryDecision retry(Duration delay) { - return new RetryDecision(true, delay); - } - - /** Returns a decision to fail without retrying. */ - public static RetryDecision fail() { - return new RetryDecision(false, Duration.ZERO); - } - - /** Returns whether another attempt should run. */ - public boolean shouldRetry() { - return shouldRetry; - } - - /** Returns the delay before the next attempt. */ - public Duration delay() { - return delay; - } + ExtensionStepResult.RetryDecision makeRetryDecision(Throwable error, T state, int attempt); } /** Builder for {@link ExtensionStepConfig}. */ public static final class Builder { private T initialState; private SerDes serDes; - private RetryStrategy retryStrategy; + private RetryStrategy retryStrategy; private StepSemantics semanticsPerRetry; private Builder() {} @@ -131,7 +101,7 @@ public Builder serDes(SerDes serDes) { } /** Sets the retry strategy used when the extension function throws. */ - public Builder retryStrategy(RetryStrategy retryStrategy) { + public Builder retryStrategy(RetryStrategy retryStrategy) { this.retryStrategy = retryStrategy; return this; } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepResult.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepResult.java index eb7067b26..e144295ad 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepResult.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepResult.java @@ -22,11 +22,19 @@ static Retry retry(T state, Duration delay) { return new Retry<>(state, delay); } + /** Creates a decision that a failed attempt should not be retried. */ + static DoNotRetry doNotRetry() { + return new DoNotRetry<>(); + } + + /** Outcomes supported when deciding whether to retry a failed attempt. */ + sealed interface RetryDecision permits Retry, DoNotRetry {} + /** Terminal successful outcome. */ record Succeeded(T value) implements ExtensionStepResult {} /** Retry outcome. */ - record Retry(T state, Duration delay) implements ExtensionStepResult { + record Retry(T state, Duration delay) implements ExtensionStepResult, RetryDecision { public Retry { Objects.requireNonNull(delay, "delay cannot be null"); if (delay.isNegative()) { @@ -34,4 +42,7 @@ record Retry(T state, Duration delay) implements ExtensionStepResult { } } } + + /** Decision that a failed attempt should not be retried. */ + record DoNotRetry() implements RetryDecision {} } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableStepOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableStepOperation.java index 8aa52d181..4733ea52b 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableStepOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableStepOperation.java @@ -82,12 +82,12 @@ public static DurableFuture stepAsync( .build()); } - private static ExtensionStepConfig.RetryStrategy adapt(RetryStrategy retryStrategy) { - return (error, attempt) -> { + private static ExtensionStepConfig.RetryStrategy adapt(RetryStrategy retryStrategy) { + return (error, state, attempt) -> { var decision = retryStrategy.makeRetryDecision(error, attempt); return decision.shouldRetry() - ? ExtensionStepConfig.RetryDecision.retry(decision.delay()) - : ExtensionStepConfig.RetryDecision.fail(); + ? ExtensionStepResult.retry(state, decision.delay()) + : ExtensionStepResult.doNotRetry(); }; } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/primitive/StepPrimitive.java b/sdk/src/main/java/software/amazon/lambda/durable/primitive/StepPrimitive.java index d141c0864..553bc351b 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/primitive/StepPrimitive.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/primitive/StepPrimitive.java @@ -3,6 +3,7 @@ package software.amazon.lambda.durable.primitive; import java.util.concurrent.CompletableFuture; +import software.amazon.awssdk.services.lambda.model.ErrorObject; import software.amazon.awssdk.services.lambda.model.Operation; import software.amazon.awssdk.services.lambda.model.OperationAction; import software.amazon.awssdk.services.lambda.model.OperationStatus; @@ -126,6 +127,10 @@ private void handleExtensionStepResult(ExtensionStepResult result, int attemp return; } var retry = (ExtensionStepResult.Retry) result; + handleExtensionStepRetry(retry, null, attempt); + } + + private void handleExtensionStepRetry(ExtensionStepResult.Retry retry, ErrorObject error, int attempt) { var serializedState = serializeAndDeserializeResult(retry.state()); var retryDelaySeconds = Math.toIntExact(retry.delay().toSeconds()); var update = OperationUpdate.builder() @@ -134,6 +139,9 @@ private void handleExtensionStepResult(ExtensionStepResult result, int attemp .stepOptions(StepOptions.builder() .nextAttemptDelaySeconds(retryDelaySeconds) .build()); + if (error != null) { + update.error(error); + } sendOperationUpdate(update); pollReadyAndExecuteExtensionStep(serializedState.deserialized(), attempt + 1); } @@ -160,18 +168,9 @@ private void handleExtensionStepFailure(Throwable exception, T state, int attemp var retryStrategy = extensionConfig.retryStrategy(); if (retryStrategy != null) { - var decision = retryStrategy.makeRetryDecision(exception, attempt); - if (decision.shouldRetry()) { - var serializedState = serializeAndDeserializeResult(state); - var retryDelaySeconds = Math.toIntExact(decision.delay().toSeconds()); - sendOperationUpdate(OperationUpdate.builder() - .action(OperationAction.RETRY) - .payload(serializedState.serialized()) - .error(error) - .stepOptions(StepOptions.builder() - .nextAttemptDelaySeconds(retryDelaySeconds) - .build())); - pollReadyAndExecuteExtensionStep(serializedState.deserialized(), attempt + 1); + var decision = retryStrategy.makeRetryDecision(exception, state, attempt); + if (decision instanceof ExtensionStepResult.Retry retry) { + handleExtensionStepRetry(retry, error, attempt); return; } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableOperationFacadeTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableOperationFacadeTest.java index 79b4c3695..d8f58119a 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableOperationFacadeTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableOperationFacadeTest.java @@ -160,11 +160,18 @@ void stepAdaptsUserFacingRetryConfigToExtensionSpi() { assertEquals( ExtensionStepConfig.StepSemantics.AT_MOST_ONCE_PER_RETRY, extensionConfig.getValue().semanticsPerRetry()); - var retry = extensionConfig.getValue().retryStrategy().makeRetryDecision(new IllegalStateException("retry"), 1); - var fail = extensionConfig.getValue().retryStrategy().makeRetryDecision(new IllegalStateException("fail"), 2); - assertTrue(retry.shouldRetry()); + var retry = assertInstanceOf( + ExtensionStepResult.Retry.class, + extensionConfig + .getValue() + .retryStrategy() + .makeRetryDecision(new IllegalStateException("retry"), null, 1)); + var doNotRetry = extensionConfig + .getValue() + .retryStrategy() + .makeRetryDecision(new IllegalStateException("fail"), null, 2); assertEquals(Duration.ofSeconds(3), retry.delay()); - assertFalse(fail.shouldRetry()); + assertInstanceOf(ExtensionStepResult.DoNotRetry.class, doNotRetry); } @Test diff --git a/sdk/src/test/java/software/amazon/lambda/durable/ExtensionStepResultTest.java b/sdk/src/test/java/software/amazon/lambda/durable/ExtensionStepResultTest.java index 3963b2619..199e05363 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/ExtensionStepResultTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/ExtensionStepResultTest.java @@ -3,6 +3,7 @@ package software.amazon.lambda.durable; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertThrows; import java.time.Duration; @@ -23,6 +24,8 @@ void retryCarriesStateAndDelay() { assertEquals("next", result.state()); assertEquals(Duration.ofSeconds(2), result.delay()); + assertInstanceOf(ExtensionStepResult.RetryDecision.class, result); + assertInstanceOf(ExtensionStepResult.RetryDecision.class, ExtensionStepResult.doNotRetry()); } @Test diff --git a/sdk/src/test/java/software/amazon/lambda/durable/extension/ExtensionStepConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/extension/ExtensionStepConfigTest.java index 9f7a13115..20d98a897 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/extension/ExtensionStepConfigTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/extension/ExtensionStepConfigTest.java @@ -3,10 +3,8 @@ package software.amazon.lambda.durable.extension; import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertSame; -import static org.junit.jupiter.api.Assertions.assertTrue; import java.time.Duration; import org.junit.jupiter.api.Test; @@ -26,8 +24,8 @@ void builderDefaultsToNullStateAndSerDes() { @Test void builderRetainsConfiguredValues() { var serDes = new JacksonSerDes(); - ExtensionStepConfig.RetryStrategy retryStrategy = - (error, attempt) -> ExtensionStepConfig.RetryDecision.retry(Duration.ofSeconds(attempt)); + ExtensionStepConfig.RetryStrategy retryStrategy = + (error, state, attempt) -> ExtensionStepResult.retry(state, Duration.ofSeconds(attempt)); var config = ExtensionStepConfig.builder() .initialState(42) .serDes(serDes) @@ -41,17 +39,6 @@ void builderRetainsConfiguredValues() { assertEquals(ExtensionStepConfig.StepSemantics.AT_MOST_ONCE_PER_RETRY, config.semanticsPerRetry()); } - @Test - void retryDecisionFactoriesExposeExtensionOwnedDecision() { - var retry = ExtensionStepConfig.RetryDecision.retry(Duration.ofSeconds(3)); - var fail = ExtensionStepConfig.RetryDecision.fail(); - - assertTrue(retry.shouldRetry()); - assertEquals(Duration.ofSeconds(3), retry.delay()); - assertFalse(fail.shouldRetry()); - assertEquals(Duration.ZERO, fail.delay()); - } - @Test void retryAndSemanticsContractsAreOwnedByExtensionStepConfig() throws Exception { assertEquals( @@ -61,9 +48,9 @@ void retryAndSemanticsContractsAreOwnedByExtensionStepConfig() throws Exception ExtensionStepConfig.StepSemantics.class, ExtensionStepConfig.class.getMethod("semanticsPerRetry").getReturnType()); assertEquals( - ExtensionStepConfig.RetryDecision.class, + ExtensionStepResult.RetryDecision.class, ExtensionStepConfig.RetryStrategy.class - .getMethod("makeRetryDecision", Throwable.class, int.class) + .getMethod("makeRetryDecision", Throwable.class, Object.class, int.class) .getReturnType()); } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java b/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java index a161b6609..e112bf29e 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java @@ -39,7 +39,6 @@ import software.amazon.lambda.durable.execution.ThreadContext; import software.amazon.lambda.durable.execution.ThreadType; import software.amazon.lambda.durable.extension.ExtensionStepConfig; -import software.amazon.lambda.durable.extension.ExtensionStepConfig.RetryDecision; import software.amazon.lambda.durable.extension.ExtensionStepFunction; import software.amazon.lambda.durable.extension.ExtensionStepResult; import software.amazon.lambda.durable.model.OperationIdentifier; @@ -191,13 +190,16 @@ void exceptionRetryWithoutStateDoesNotCheckpointPayload() throws Exception { }, ExtensionStepConfig.builder() .serDes(SERDES) - .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(1))) + .retryStrategy( + (error, state, attempt) -> ExtensionStepResult.retry(state, Duration.ofSeconds(1))) .build()); operation.execute(); assertTrue(retrySent.await(2, TimeUnit.SECONDS)); assertNull(retryUpdate.get().payload()); + assertEquals( + IllegalStateException.class.getName(), retryUpdate.get().error().errorType()); } @Test From 954ed3138ee3a868552603917582fe3b31aa2c73 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Mon, 10 Aug 2026 13:18:44 -0700 Subject: [PATCH 31/40] refactor: delegate child context failure translation --- docs/adr/006-custom-extension-operations.md | 3 +- docs/advanced/extensions.md | 4 +- .../extension/ExtensionContextFailure.java | 43 +++++++++---- .../DurableConcurrencyOperation.java | 5 +- .../operation/DurableMapOperation.java | 4 +- .../operation/DurableParallelOperation.java | 4 +- .../primitive/ChildContextPrimitive.java | 63 +------------------ ...DurableMapOperationImplementationTest.java | 19 ++++++ ...leParallelOperationImplementationTest.java | 19 ++++++ ...orCallbackOperationImplementationTest.java | 10 ++- .../primitive/ChildContextPrimitiveTest.java | 32 +++++++++- 11 files changed, 123 insertions(+), 83 deletions(-) diff --git a/docs/adr/006-custom-extension-operations.md b/docs/adr/006-custom-extension-operations.md index ecbe53fcc..29e380841 100644 --- a/docs/adr/006-custom-extension-operations.md +++ b/docs/adr/006-custom-extension-operations.md @@ -268,8 +268,7 @@ when the original exception cannot be reconstructed. An `ExtensionContextErrorHandler` receives a read-only `ExtensionContextFailure` containing: -- context name and subtype -- error metadata +- the failed context operation, including its ID, name, subtype, status, and error metadata - child operation type, subtype, status, and error summaries Resolution order is: diff --git a/docs/advanced/extensions.md b/docs/advanced/extensions.md index affddd037..0d55210c5 100644 --- a/docs/advanced/extensions.md +++ b/docs/advanced/extensions.md @@ -326,8 +326,8 @@ the threshold. Replay metadata is scoped to the framework callback through `Exte `ExtensionContextConfig` directly configures the context serializer and whether the context is virtual. It also controls framework user-function plugin events and can suppress child checkpoints that finish after the parent. If a context fails, the SDK first rethrows a deserialized original exception, then calls the configured error handler, and -finally falls back to `ChildContextFailedException`. The handler receives read-only context metadata and -child-operation summaries. +finally falls back to `ChildContextFailedException`. The handler receives the complete failed context operation +through `failure.operation()`, plus read-only child-operation summaries. ## Explicit child contexts diff --git a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextFailure.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextFailure.java index 320e1c9aa..ce177be07 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextFailure.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextFailure.java @@ -3,35 +3,54 @@ package software.amazon.lambda.durable.extension; import java.util.List; +import java.util.Objects; +import software.amazon.awssdk.services.lambda.model.ContextDetails; import software.amazon.awssdk.services.lambda.model.ErrorObject; +import software.amazon.awssdk.services.lambda.model.Operation; +import software.amazon.awssdk.services.lambda.model.OperationStatus; +import software.amazon.awssdk.services.lambda.model.OperationType; /** Read-only failure information supplied to an extension CONTEXT error handler. */ public final class ExtensionContextFailure { - private final String contextName; - private final String subType; + private final Operation operation; private final Throwable originalException; - private final ErrorObject error; private final List childOperations; + public ExtensionContextFailure( + Operation operation, Throwable originalException, List childOperations) { + this.operation = Objects.requireNonNull(operation, "operation cannot be null"); + this.originalException = originalException; + this.childOperations = List.copyOf(childOperations); + } + public ExtensionContextFailure( String contextName, String subType, Throwable originalException, ErrorObject error, List childOperations) { - this.contextName = contextName; - this.subType = subType; - this.originalException = originalException; - this.error = error; - this.childOperations = List.copyOf(childOperations); + this( + Operation.builder() + .name(contextName) + .type(OperationType.CONTEXT) + .subType(subType) + .status(OperationStatus.FAILED) + .contextDetails(ContextDetails.builder().error(error).build()) + .build(), + originalException, + childOperations); + } + + public Operation operation() { + return operation; } public String contextName() { - return contextName; + return operation.name(); } public String subType() { - return subType; + return operation.subType(); } public Throwable originalException() { @@ -39,7 +58,9 @@ public Throwable originalException() { } public ErrorObject error() { - return error; + return operation.contextDetails() == null + ? null + : operation.contextDetails().error(); } public List childOperations() { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableConcurrencyOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableConcurrencyOperation.java index 50a30cff8..99b68b8bb 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableConcurrencyOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableConcurrencyOperation.java @@ -19,6 +19,7 @@ import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextErrorHandler; import software.amazon.lambda.durable.model.ConcurrencyCompletionStatus; import software.amazon.lambda.durable.serde.SerDes; @@ -28,10 +29,12 @@ public abstract class DurableConcurrencyOperation { DurableConcurrencyOperation() {} - protected static ExtensionContextConfig childContextConfig(SerDes serDes, NestingType nestingType) { + protected static ExtensionContextConfig childContextConfig( + SerDes serDes, NestingType nestingType, ExtensionContextErrorHandler errorHandler) { return ExtensionContextConfig.builder() .serDes(serDes) .isVirtual(nestingType == NestingType.FLAT) + .errorHandler(errorHandler) .build(); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java index 955e93610..3b8ac4f33 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java @@ -17,6 +17,7 @@ import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.exception.MapIterationFailedException; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.extension.ExtensionContext; @@ -165,7 +166,8 @@ private static List> registerItem MapResult replayState) { var context = ExtensionContext.getCurrentContext(); var registeredItems = new ArrayList>(items.size()); - var iterationConfig = childContextConfig(config.serDes(), config.nestingType()); + var iterationConfig = childContextConfig( + config.serDes(), config.nestingType(), failure -> new MapIterationFailedException(failure.operation())); for (int index = 0; index < items.size(); index++) { var item = items.get(index); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableParallelOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableParallelOperation.java index 4b25dd7d7..7871ef84d 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableParallelOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableParallelOperation.java @@ -16,6 +16,7 @@ import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.ParallelDurableFuture; import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.exception.ParallelBranchFailedException; import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.extension.ExtensionContextConfig; import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; @@ -224,7 +225,8 @@ private static ParallelResult constructResult(OperationConcurrencyCoordinator.Co private ExtensionContextConfig branchConfig(ParallelBranchConfig branchConfig) { var serDes = branchConfig.serDes() == null ? defaultSerDes : branchConfig.serDes(); - return childContextConfig(serDes, config.nestingType()); + return childContextConfig( + serDes, config.nestingType(), failure -> new ParallelBranchFailedException(failure.operation())); } private static OperationConcurrencyCoordinator.ExpectedCompletionStatus expectedCompletion( diff --git a/sdk/src/main/java/software/amazon/lambda/durable/primitive/ChildContextPrimitive.java b/sdk/src/main/java/software/amazon/lambda/durable/primitive/ChildContextPrimitive.java index 30532d7d2..84076b1b5 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/primitive/ChildContextPrimitive.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/primitive/ChildContextPrimitive.java @@ -2,8 +2,6 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.primitive; -import static software.amazon.lambda.durable.execution.ExecutionManager.isTerminalStatus; - import java.nio.charset.StandardCharsets; import java.util.List; import java.util.Objects; @@ -22,15 +20,8 @@ import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.context.DurableContextImpl; -import software.amazon.lambda.durable.exception.CallbackFailedException; -import software.amazon.lambda.durable.exception.CallbackSubmitterException; -import software.amazon.lambda.durable.exception.CallbackTimeoutException; import software.amazon.lambda.durable.exception.ChildContextFailedException; import software.amazon.lambda.durable.exception.DurableOperationException; -import software.amazon.lambda.durable.exception.MapIterationFailedException; -import software.amazon.lambda.durable.exception.ParallelBranchFailedException; -import software.amazon.lambda.durable.exception.StepFailedException; -import software.amazon.lambda.durable.exception.StepInterruptedException; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.execution.ThreadType; @@ -43,7 +34,6 @@ import software.amazon.lambda.durable.logging.DurableLogger; import software.amazon.lambda.durable.model.DeserializedOperationResult; import software.amazon.lambda.durable.model.OperationIdentifier; -import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.util.ExceptionHelper; /** @@ -339,23 +329,12 @@ private Throwable translateException(Operation op, ErrorObject errorObject) { } if (extensionConfig != null && extensionConfig.errorHandler() != null) { - var failure = new ExtensionContextFailure( - getName(), getSubTypeValue(), null, errorObject, getChildOperationSummaries()); + var failure = new ExtensionContextFailure(op, null, getChildOperationSummaries()); return Objects.requireNonNull( extensionConfig.errorHandler().translate(failure), "Extension context error handler result cannot be null"); } - // throw a general failed exception if a user exception is not reconstructed - if (OperationSubType.WAIT_FOR_CALLBACK.getValue().equals(getSubTypeValue())) { - return handleWaitForCallbackFailure(); - } - if (OperationSubType.MAP_ITERATION.getValue().equals(getSubTypeValue())) { - return new MapIterationFailedException(op); - } - if (OperationSubType.PARALLEL_BRANCH.getValue().equals(getSubTypeValue())) { - return new ParallelBranchFailedException(op); - } return new ChildContextFailedException(op); } @@ -370,47 +349,9 @@ private Operation createVirtualOperation(ErrorObject errorObject) { .id(getOperationId()) .name(getName()) .type(OperationType.CONTEXT) + .subType(getSubTypeValue()) .status(OperationStatus.FAILED) .contextDetails(ContextDetails.builder().error(errorObject).build()) .build(); } - - private Throwable handleWaitForCallbackFailure() { - var childrenOps = getChildOperations(); - var callbackOp = childrenOps.stream() - .filter(o -> o.type() == OperationType.CALLBACK) - .findFirst() - .orElse(null); - var submitterOp = childrenOps.stream() - .filter(o -> o.type() == OperationType.STEP) - .findFirst() - .orElse(null); - if (callbackOp != null) { - // if callback failed - if (isTerminalStatus(callbackOp.status())) { - switch (callbackOp.status()) { - case FAILED -> { - return new CallbackFailedException(callbackOp); - } - case TIMED_OUT -> { - return new CallbackTimeoutException(callbackOp); - } - } - } - - // if submitter failed - if (submitterOp != null - && isTerminalStatus(submitterOp.status()) - && submitterOp.status() != OperationStatus.SUCCEEDED) { - var stepError = submitterOp.stepDetails().error(); - if (StepInterruptedException.isStepInterruptedException(stepError)) { - return new CallbackSubmitterException(callbackOp, new StepInterruptedException(submitterOp)); - } else { - return new CallbackSubmitterException(callbackOp, new StepFailedException(submitterOp)); - } - } - } - - return new IllegalStateException("Unknown waitForCallback status"); - } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java index 31ffd1686..9c299541d 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java @@ -4,6 +4,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; @@ -18,12 +19,17 @@ import java.util.concurrent.CompletableFuture; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; +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.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.context.BaseContextImpl; +import software.amazon.lambda.durable.exception.MapIterationFailedException; import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextFailure; import software.amazon.lambda.durable.extension.ExtensionContextFunction; import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; import software.amazon.lambda.durable.extension.ExtensionOperation; @@ -100,6 +106,19 @@ void executeBuildsMapAndIterationContextsFromReservations() { iterationConfig.capture()); assertTrue(iterationConfig.getValue().isVirtual()); assertSame(serDes, iterationConfig.getValue().serDes()); + var failedIteration = Operation.builder() + .id("iteration-id") + .name("map-iteration-0") + .type(OperationType.CONTEXT) + .subType(MAP_ITERATION.getValue()) + .status(OperationStatus.FAILED) + .build(); + var translated = iterationConfig + .getValue() + .errorHandler() + .translate(new ExtensionContextFailure(failedIteration, null, List.of())); + var failure = assertInstanceOf(MapIterationFailedException.class, translated); + assertSame(failedIteration, failure.getOperation()); } @SuppressWarnings({"rawtypes", "unchecked"}) diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableParallelOperationImplementationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableParallelOperationImplementationTest.java index aea53dde0..59f6a2b1b 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableParallelOperationImplementationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableParallelOperationImplementationTest.java @@ -4,6 +4,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -20,13 +21,18 @@ import java.util.concurrent.CompletableFuture; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; +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.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.context.BaseContextImpl; +import software.amazon.lambda.durable.exception.ParallelBranchFailedException; import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.extension.ExtensionContextConfig; +import software.amazon.lambda.durable.extension.ExtensionContextFailure; import software.amazon.lambda.durable.extension.ExtensionContextFunction; import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; import software.amazon.lambda.durable.extension.ExtensionOperation; @@ -120,6 +126,19 @@ void executeBuildsParallelAndBranchContextsFromReservations() { branchConfig.capture()); assertTrue(branchConfig.getValue().isVirtual()); assertSame(serDes, branchConfig.getValue().serDes()); + var failedBranch = Operation.builder() + .id("branch-id") + .name("first") + .type(OperationType.CONTEXT) + .subType(PARALLEL_BRANCH.getValue()) + .status(OperationStatus.FAILED) + .build(); + var translated = branchConfig + .getValue() + .errorHandler() + .translate(new ExtensionContextFailure(failedBranch, null, List.of())); + var failure = assertInstanceOf(ParallelBranchFailedException.class, translated); + assertSame(failedBranch, failure.getOperation()); } @Test diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperationImplementationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperationImplementationTest.java index b9eb342de..04b460de8 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperationImplementationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperationImplementationTest.java @@ -101,9 +101,13 @@ void errorHandlerPreservesCallbackTimeoutException() { CallbackDetails.builder().callbackId("callback-id").build()) .build(); var failure = new ExtensionContextFailure( - "approval", - OperationSubType.WAIT_FOR_CALLBACK.getValue(), - null, + Operation.builder() + .id("approval") + .name("approval") + .type(OperationType.CONTEXT) + .subType(OperationSubType.WAIT_FOR_CALLBACK.getValue()) + .status(OperationStatus.FAILED) + .build(), null, List.of(new ExtensionChildOperationSummary(callback))); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/primitive/ChildContextPrimitiveTest.java b/sdk/src/test/java/software/amazon/lambda/durable/primitive/ChildContextPrimitiveTest.java index 3a1ba3d6f..bd60e65ea 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/primitive/ChildContextPrimitiveTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/primitive/ChildContextPrimitiveTest.java @@ -133,8 +133,12 @@ private ChildContextPrimitive createOperationWithParent( } private ChildContextPrimitive createExtensionOperation(ExtensionContextConfig config) { + return createExtensionOperation("AcmeContext", config); + } + + private ChildContextPrimitive createExtensionOperation(String subType, ExtensionContextConfig config) { return new ChildContextPrimitive<>( - new OperationIdentifier("1", "test-context", OperationType.CONTEXT, "AcmeContext"), + new OperationIdentifier("1", "test-context", OperationType.CONTEXT, subType), () -> ExtensionContextResult.completed("unused"), TypeToken.get(String.class), config, @@ -261,6 +265,31 @@ void replayFailedFallsBackToChildContextFailedException() { assertTrue(thrown.getMessage().contains("unknown error")); } + @Test + void replayKnownSubtypeWithoutErrorHandlerUsesGenericFailure() { + var failedContext = Operation.builder() + .id("1") + .name("test-context") + .type(OperationType.CONTEXT) + .subType(OperationSubType.MAP_ITERATION.getValue()) + .status(OperationStatus.FAILED) + .contextDetails(ContextDetails.builder() + .error(ErrorObject.builder() + .errorType("com.nonexistent.SomeException") + .errorMessage("unknown error") + .build()) + .build()) + .build(); + when(executionManager.getOperationAndUpdateReplayState("1")).thenReturn(failedContext); + var config = ExtensionContextConfig.builder().serDes(SERDES).build(); + var operation = createExtensionOperation(OperationSubType.MAP_ITERATION.getValue(), config); + + operation.execute(); + + var thrown = assertThrows(ChildContextFailedException.class, operation::get); + assertSame(failedContext, thrown.getOperation()); + } + @Test void replayFailedUsesExtensionErrorHandlerWithChildSummaries() { var contextError = ErrorObject.builder() @@ -304,6 +333,7 @@ void replayFailedUsesExtensionErrorHandlerWithChildSummaries() { assertSame(translated, assertThrows(IllegalStateException.class, operation::get)); var failure = capturedFailure.get(); + assertSame(failedContext, failure.operation()); assertEquals("test-context", failure.contextName()); assertEquals("AcmeContext", failure.subType()); assertEquals(contextError, failure.error()); From 8093c6a593b17ec49cbd3148d73924a1a51d2999 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Mon, 10 Aug 2026 13:47:10 -0700 Subject: [PATCH 32/40] refactor: adapt durable step facade in context --- .../amazon/lambda/durable/DurableContext.java | 9 +++-- .../operation/DurableStepOperation.java | 16 ++------- .../durable/DurableOperationFacadeTest.java | 34 ++++++++++++++++--- 3 files changed, 39 insertions(+), 20 deletions(-) diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java index a295d8b37..6e2a5c52e 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java @@ -4,6 +4,7 @@ import java.time.Duration; import java.util.Collection; +import java.util.Objects; import java.util.function.BiConsumer; import java.util.function.BiFunction; import java.util.function.Function; @@ -18,6 +19,7 @@ import software.amazon.lambda.durable.config.WaitForConditionConfig; import software.amazon.lambda.durable.config.WithRetryConfig; import software.amazon.lambda.durable.context.BaseContext; +import software.amazon.lambda.durable.context.BaseContextImpl; import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.model.MapResult; import software.amazon.lambda.durable.model.WaitForConditionResult; @@ -165,8 +167,11 @@ default DurableFuture stepAsync(String name, TypeToken resultType, Fun */ default DurableFuture stepAsync( String name, TypeToken resultType, Function func, StepConfig config) { - return DurableStepOperation.stepAsync( - (ExtensionContext) this, name, resultType, func, config.toOperationConfig()); + Objects.requireNonNull(func, "func cannot be null"); + try (var ignored = BaseContextImpl.attachCurrentContext(this)) { + return DurableStepOperation.stepAsync( + name, resultType, () -> func.apply(StepContext.getCurrentContext()), config.toOperationConfig()); + } } /** @deprecated use the variants accepting StepContext instead */ diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableStepOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableStepOperation.java index 4733ea52b..954a08346 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableStepOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableStepOperation.java @@ -5,10 +5,8 @@ import static software.amazon.lambda.durable.model.OperationSubType.STEP; import java.util.Objects; -import java.util.function.Function; import java.util.function.Supplier; import software.amazon.lambda.durable.DurableFuture; -import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.StepSemantics; import software.amazon.lambda.durable.extension.ExtensionContext; @@ -55,26 +53,16 @@ public static DurableFuture stepAsync( public static DurableFuture stepAsync( String name, TypeToken resultType, Supplier function, StepConfig config) { Objects.requireNonNull(function, "function cannot be null"); - return stepAsync(ExtensionContext.getCurrentContext(), name, resultType, ignored -> function.get(), config); - } - - public static DurableFuture stepAsync( - ExtensionContext context, - String name, - TypeToken resultType, - Function function, - StepConfig config) { - Objects.requireNonNull(context, "context cannot be null"); Objects.requireNonNull(resultType, "resultType cannot be null"); - Objects.requireNonNull(function, "function cannot be null"); Objects.requireNonNull(config, "config cannot be null"); ParameterValidator.validateOperationName(name); + var context = ExtensionContext.getCurrentContext(); return context.reserve(name) .stepAsync( STEP.getValue(), resultType, - ignored -> ExtensionStepResult.succeed(function.apply(StepContext.getCurrentContext())), + ignored -> ExtensionStepResult.succeed(function.get()), ExtensionStepConfig.builder() .serDes(config.serDes()) .retryStrategy(adapt(config.retryStrategy())) diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableOperationFacadeTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableOperationFacadeTest.java index d8f58119a..0ca83db54 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableOperationFacadeTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableOperationFacadeTest.java @@ -22,6 +22,7 @@ import static software.amazon.lambda.durable.model.OperationSubType.WAIT; import java.time.Duration; +import java.util.function.Function; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; @@ -100,6 +101,19 @@ void operationFacadesAbsorbTheirExtensions() { assertMergedOperation("DurableWithRetryOperation", "WithRetryExtension"); } + @Test + void stepFacadeDoesNotAcceptExplicitContext() { + assertThrows( + NoSuchMethodException.class, + () -> DurableStepOperation.class.getMethod( + "stepAsync", + ExtensionContext.class, + String.class, + TypeToken.class, + Function.class, + DurableStepOperation.StepConfig.class)); + } + @Test void stepAcceptsContextFreeSupplier() { var context = mockDurableContext(); @@ -130,12 +144,13 @@ void stepAcceptsContextFreeSupplier() { } @Test - void stepAdaptsUserFacingRetryConfigToExtensionSpi() { - var context = mock(ExtensionContext.class); + void stepAdaptsOperationConfigToExtensionSpi() { + var context = mockDurableContext(); var reservation = mock(ExtensionOperation.class); var future = mockStringFuture(); var resultType = TypeToken.get(String.class); - when(context.reserve("step")).thenReturn(reservation); + BaseContextImpl.setCurrentContext(context); + when(((ExtensionContext) context).reserve("step")).thenReturn(reservation); when(reservation.stepAsync( eq(STEP.getValue()), eq(resultType), @@ -148,7 +163,7 @@ void stepAdaptsUserFacingRetryConfigToExtensionSpi() { .semanticsPerRetry(StepSemantics.AT_MOST_ONCE_PER_RETRY) .build(); - assertSame(future, DurableStepOperation.stepAsync(context, "step", resultType, ignored -> "result", config)); + assertSame(future, DurableStepOperation.stepAsync("step", resultType, () -> "result", config)); var extensionConfig = ArgumentCaptor.forClass(ExtensionStepConfig.class); verify(reservation) @@ -196,6 +211,17 @@ void durableContextStepUsesPrimitiveExtension() { StepConfig.builder().build()); assertSame(future, result); + @SuppressWarnings("unchecked") + var function = (ArgumentCaptor>) + (ArgumentCaptor) ArgumentCaptor.forClass(ExtensionStepFunction.class); + verify(reservation) + .stepAsync( + eq(STEP.getValue()), any(TypeToken.class), function.capture(), any(ExtensionStepConfig.class)); + try (var ignored = BaseContextImpl.attachCurrentContext(mock(StepContext.class))) { + var stepResult = assertInstanceOf( + ExtensionStepResult.Succeeded.class, function.getValue().apply(null)); + assertEquals("result", stepResult.value()); + } } @Test From 5a8491c7fb19c5961475b1d1b469f572da31c850 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Mon, 10 Aug 2026 14:10:03 -0700 Subject: [PATCH 33/40] refactor: remove parallel supplier overloads --- .../StaticOperationsIntegrationTest.java | 11 ++-- .../lambda/durable/ParallelDurableFuture.java | 23 -------- .../durable/DurableParallelOperationTest.java | 56 ------------------- 3 files changed, 7 insertions(+), 83 deletions(-) delete mode 100644 sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationTest.java diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java index 23755482d..89838b1b0 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java @@ -72,11 +72,13 @@ void mapAndParallelExposeContextFreeUserFunctions() { var branchFutures = new ArrayList>(); try (var parallel = DurableParallelOperation.parallel("parallel")) { branchFutures.add(parallel.branch( - "left", String.class, () -> DurableStepOperation.step("branch-step", String.class, () -> "L"))); + "left", + String.class, + ignored -> DurableStepOperation.step("branch-step", String.class, () -> "L"))); branchFutures.add(parallel.branch( "right", String.class, - () -> DurableStepOperation.step("branch-step", String.class, () -> "R"))); + ignored -> DurableStepOperation.step("branch-step", String.class, () -> "R"))); } return mapResult.results() + ":" + DurableFuture.allOf(branchFutures); }); @@ -136,9 +138,10 @@ void staticParallelMatchesLegacyCheckpointHistory() { }); var staticRunner = LocalDurableTestRunner.create(String.class, (input, context) -> { try (var parallel = DurableParallelOperation.parallel("parallel", parallelConfig.toOperationConfig())) { - parallel.branch("left", String.class, () -> DurableStepOperation.step("work", String.class, () -> "L")); parallel.branch( - "right", String.class, () -> DurableStepOperation.step("work", String.class, () -> "R")); + "left", String.class, ignored -> DurableStepOperation.step("work", String.class, () -> "L")); + parallel.branch( + "right", String.class, ignored -> DurableStepOperation.step("work", String.class, () -> "R")); return parallel.get().statuses().toString(); } }); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java b/sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java index a1a1dae1c..09b5b6f78 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java @@ -3,35 +3,12 @@ package software.amazon.lambda.durable; import java.util.function.Function; -import java.util.function.Supplier; import software.amazon.lambda.durable.config.ParallelBranchConfig; import software.amazon.lambda.durable.model.ParallelResult; import software.amazon.lambda.durable.model.SafeCloseable; /** User-facing context for managing parallel branch execution within a durable function. */ public interface ParallelDurableFuture extends SafeCloseable, DurableFuture { - default DurableFuture branch(String name, Class resultType, Supplier function) { - return branch(name, TypeToken.get(resultType), function); - } - - default DurableFuture branch(String name, TypeToken resultType, Supplier function) { - return branch( - name, - resultType, - ignored -> function.get(), - ParallelBranchConfig.builder().build()); - } - - default DurableFuture branch( - String name, Class resultType, Supplier function, ParallelBranchConfig config) { - return branch(name, TypeToken.get(resultType), function, config); - } - - default DurableFuture branch( - String name, TypeToken resultType, Supplier function, ParallelBranchConfig config) { - return branch(name, resultType, ignored -> function.get(), config); - } - /** * Registers and immediately starts a branch (respects maxConcurrency). * diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationTest.java deleted file mode 100644 index 3a665b8f9..000000000 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableParallelOperationTest.java +++ /dev/null @@ -1,56 +0,0 @@ -// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 -package software.amazon.lambda.durable; - -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.Test; -import software.amazon.lambda.durable.config.ParallelConfig; -import software.amazon.lambda.durable.context.BaseContextImpl; -import software.amazon.lambda.durable.extension.ExtensionContext; -import software.amazon.lambda.durable.extension.ExtensionContextConfig; -import software.amazon.lambda.durable.extension.ExtensionContextFunction; -import software.amazon.lambda.durable.extension.ExtensionOperation; -import software.amazon.lambda.durable.operation.DurableParallelOperation; - -class DurableParallelOperationTest { - @AfterEach - void clearContext() { - BaseContextImpl.setCurrentContext(null); - } - - @Test - void parallelBranchesAcceptContextFreeSuppliers() { - var context = mock(CurrentContext.class); - var parent = mock(ExtensionOperation.class); - var parentFuture = mockParallelResultFuture(); - BaseContextImpl.setCurrentContext(context); - when(context.getDurableConfig()).thenReturn(DurableConfig.builder().build()); - when(context.reserve("parallel")).thenReturn(parent); - when(parent.runInChildContextAsync( - any(String.class), - any(TypeToken.class), - any(ExtensionContextFunction.class), - any(ExtensionContextConfig.class))) - .thenReturn(parentFuture); - - var result = DurableParallelOperation.parallel("parallel"); - result.branch("branch", String.class, () -> "result"); - - verify(context).reserve("parallel"); - verify(context, never()).parallel(eq("parallel"), any(ParallelConfig.class)); - } - - @SuppressWarnings("unchecked") - private DurableFuture mockParallelResultFuture() { - return mock(DurableFuture.class); - } - - private interface CurrentContext extends DurableContext, ExtensionContext {} -} From 5d8fbe0c0e7d8696e0ecb7d8932a9c11c63f4bce Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Mon, 10 Aug 2026 14:58:55 -0700 Subject: [PATCH 34/40] feat: add operation-based example suite --- .github/workflows/build.yml | 28 ++++ .github/workflows/e2e-tests.yml | 75 +++++++++ examples/README.md | 18 ++- examples/generate-template.py | 21 ++- .../operation/callback/CallbackExample.java | 82 ++++++++++ .../callback/RetryWaitForCallbackExample.java | 62 ++++++++ .../WaitForCallbackFailedExample.java | 58 +++++++ .../operation/child/ChildContextExample.java | 81 ++++++++++ .../child/ManyAsyncChildContextExample.java | 79 ++++++++++ .../child/VirtualChildContextExample.java | 97 ++++++++++++ .../general/CustomConfigExample.java | 138 ++++++++++++++++ .../general/CustomPollingExample.java | 67 ++++++++ .../general/ErrorHandlingExample.java | 100 ++++++++++++ .../general/GenericInputOutputExample.java | 52 ++++++ .../general/GenericTypesExample.java | 93 +++++++++++ .../operation/general/LoggingExample.java | 43 +++++ .../operation/general/NoopExample.java | 19 +++ .../operation/general/OtelExample.java | 65 ++++++++ .../operation/general/PluginExample.java | 102 ++++++++++++ .../operation/invoke/RetryInvokeExample.java | 45 ++++++ .../operation/invoke/SimpleInvokeExample.java | 37 +++++ .../operation/map/ComplexFlatMapExample.java | 81 ++++++++++ .../operation/map/ComplexMapExample.java | 72 +++++++++ .../map/CustomShouldCompleteMapExample.java | 88 +++++++++++ .../map/DeserializationFailedMapExample.java | 72 +++++++++ .../operation/map/SimpleMapExample.java | 44 ++++++ .../operation/otel/OtelXRayExamples.java | 109 +++++++++++++ .../otel/OtelXRayExecutionStepExample.java | 40 +++++ .../otel/OtelXRayExecutionWaitExample.java | 44 ++++++ .../operation/otel/OtelXRayStepExample.java | 54 +++++++ .../operation/otel/OtelXRayWaitExample.java | 70 +++++++++ .../DeserializationFailedParallelExample.java | 77 +++++++++ .../operation/parallel/ParallelExample.java | 68 ++++++++ .../ParallelFailureToleranceExample.java | 80 ++++++++++ .../parallel/ParallelWithWaitExample.java | 76 +++++++++ .../step/DeserializationFailureExample.java | 45 ++++++ .../operation/step/ManyAsyncStepsExample.java | 74 +++++++++ .../examples/operation/step/RetryExample.java | 88 +++++++++++ .../operation/step/RetryInProcessExample.java | 90 +++++++++++ .../operation/step/SimpleStepExample.java | 35 +++++ ...anyAsyncStepsVirtualThreadPoolExample.java | 78 +++++++++ .../ConcurrentWaitForConditionExample.java | 60 +++++++ .../operation/wait/WaitAsyncExample.java | 47 ++++++ .../operation/wait/WaitAtLeastExample.java | 69 ++++++++ .../wait/WaitAtLeastInProcessExample.java | 69 ++++++++ .../examples/operation/wait/WaitExample.java | 62 ++++++++ .../wait/WaitForConditionExample.java | 38 +++++ .../callback/CallbackExampleTest.java | 100 ++++++++++++ .../RetryWaitForCallbackExampleTest.java | 148 ++++++++++++++++++ .../WaitForCallbackFailedExampleTest.java | 31 ++++ .../child/ChildContextExampleTest.java | 56 +++++++ .../ManyAsyncChildContextExampleTest.java | 65 ++++++++ .../child/VirtualChildContextExampleTest.java | 56 +++++++ .../general/CustomConfigExampleTest.java | 43 +++++ .../general/CustomPollingExampleTest.java | 32 ++++ .../general/ErrorHandlingExampleTest.java | 54 +++++++ .../GenericInputOutputExampleTest.java | 57 +++++++ .../general/GenericTypesExampleTest.java | 68 ++++++++ .../operation/general/LoggingExampleTest.java | 24 +++ .../operation/general/OtelExampleTest.java | 28 ++++ .../operation/general/PluginExampleTest.java | 42 +++++ .../operation/invoke/InvokeExampleTest.java | 95 +++++++++++ .../invoke/RetryInvokeExampleTest.java | 130 +++++++++++++++ .../map/ComplexFlatMapExampleTest.java | 54 +++++++ .../operation/map/ComplexMapExampleTest.java | 53 +++++++ .../CustomShouldCompleteMapExampleTest.java | 61 ++++++++ .../DeserializationFailedMapExampleTest.java | 26 +++ .../operation/map/SimpleMapExampleTest.java | 49 ++++++ .../otel/OtelXRayExampleTestSupport.java | 30 ++++ .../operation/otel/OtelXRayExamplesTest.java | 51 ++++++ .../OtelXRayExecutionStepExampleTest.java | 48 ++++++ .../OtelXRayExecutionWaitExampleTest.java | 48 ++++++ .../otel/OtelXRayStepExampleTest.java | 59 +++++++ .../otel/OtelXRayWaitExampleTest.java | 61 ++++++++ ...erializationFailedParallelExampleTest.java | 29 ++++ .../parallel/ParallelExampleTest.java | 60 +++++++ .../ParallelFailureToleranceExampleTest.java | 59 +++++++ .../parallel/ParallelWithWaitExampleTest.java | 34 ++++ .../DeserializationFailureExampleTest.java | 31 ++++ .../step/ManyAsyncStepsExampleTest.java | 60 +++++++ .../operation/step/RetryExampleTest.java | 67 ++++++++ .../operation/step/SimpleStepExampleTest.java | 77 +++++++++ ...syncStepsVirtualThreadPoolExampleTest.java | 64 ++++++++ ...ConcurrentWaitForConditionExampleTest.java | 29 ++++ .../operation/wait/WaitAsyncExampleTest.java | 41 +++++ .../operation/wait/WaitExampleTest.java | 30 ++++ .../wait/WaitForConditionExampleTest.java | 23 +++ .../StaticOperationsIntegrationTest.java | 12 +- .../durable/context/DurableContextImpl.java | 12 +- .../DurableWaitForConditionOperation.java | 20 ++- .../DurableWaitForConditionOperationTest.java | 15 +- .../operation/DurableOperationConfigTest.java | 12 ++ ...rConditionOperationImplementationTest.java | 12 +- 93 files changed, 5329 insertions(+), 19 deletions(-) create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/CallbackExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/RetryWaitForCallbackExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/WaitForCallbackFailedExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ChildContextExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ManyAsyncChildContextExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/VirtualChildContextExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/CustomConfigExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/CustomPollingExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/ErrorHandlingExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/GenericInputOutputExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/GenericTypesExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/LoggingExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/NoopExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/OtelExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/PluginExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/invoke/RetryInvokeExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/invoke/SimpleInvokeExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexFlatMapExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexMapExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/CustomShouldCompleteMapExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/DeserializationFailedMapExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/SimpleMapExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExamples.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionStepExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionWaitExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayStepExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayWaitExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/DeserializationFailedParallelExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelFailureToleranceExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelWithWaitExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/DeserializationFailureExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/ManyAsyncStepsExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/RetryExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/RetryInProcessExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/SimpleStepExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/vt/ManyAsyncStepsVirtualThreadPoolExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/ConcurrentWaitForConditionExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAsyncExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAtLeastExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAtLeastInProcessExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitForConditionExample.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/callback/CallbackExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/callback/RetryWaitForCallbackExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/callback/WaitForCallbackFailedExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/child/ChildContextExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/child/ManyAsyncChildContextExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/child/VirtualChildContextExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/CustomConfigExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/CustomPollingExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/ErrorHandlingExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/GenericInputOutputExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/GenericTypesExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/LoggingExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/OtelExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/PluginExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/invoke/InvokeExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/invoke/RetryInvokeExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/map/ComplexFlatMapExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/map/ComplexMapExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/map/CustomShouldCompleteMapExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/map/DeserializationFailedMapExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/map/SimpleMapExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExampleTestSupport.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExamplesTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionStepExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionWaitExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayStepExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayWaitExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/parallel/DeserializationFailedParallelExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelFailureToleranceExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelWithWaitExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/step/DeserializationFailureExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/step/ManyAsyncStepsExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/step/RetryExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/step/SimpleStepExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/vt/ManyAsyncStepsVirtualThreadPoolExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/wait/ConcurrentWaitForConditionExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/wait/WaitAsyncExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/wait/WaitExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/operation/wait/WaitForConditionExampleTest.java diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 26e154242..b36580103 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -126,6 +126,34 @@ jobs: with: path: github-pages + operation-examples: + name: Operation examples (Java ${{ matrix.java }}) + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + java: + - 17 + - 21 + - 25 + steps: + - name: Checkout repository + uses: actions/checkout@v7 + + - name: Setup Java ${{ matrix.java }} + uses: actions/setup-java@v5 + with: + distribution: corretto + java-version: ${{ matrix.java }} + cache: maven + + - name: Test operation examples + run: >- + mvn -B -pl examples -am + -Dtest='software.amazon.lambda.durable.examples.operation.**.*Test' + -Dsurefire.failIfNoSpecifiedTests=false + test + deploy-pages: if: ${{ github.event_name == 'milestone' || github.ref == 'refs/heads/main' }} needs: build diff --git a/.github/workflows/e2e-tests.yml b/.github/workflows/e2e-tests.yml index b8bbe991d..b7ff944ac 100644 --- a/.github/workflows/e2e-tests.yml +++ b/.github/workflows/e2e-tests.yml @@ -108,3 +108,78 @@ jobs: JAVA_VERSION: ${{ matrix.java }} run: python3 ../.github/scripts/publish_e2e_test_summary.py working-directory: ./examples + + operation-e2e-tests: + name: Operation examples E2E (Java ${{ matrix.java }}) + env: + AWS_REGION: us-west-2 + E2E_TEST_PARALLELISM: 4 + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + java: + - 17 + - 21 + - 25 + steps: + - name: Checkout repository + uses: actions/checkout@v7 + - name: Setup AWS SAM CLI + uses: aws-actions/setup-sam@v3 + with: + use-installer: true + - name: Configure AWS credentials + uses: aws-actions/configure-aws-credentials@e6de054238d6b7531b4efff3b6587d9aade6a06c # v6.2.3 + with: + role-to-assume: "${{ secrets.TEST_ROLE_ARN }}" + role-session-name: java-language-sdk-operation-test + aws-region: ${{ env.AWS_REGION }} + allowed-account-ids: ${{ secrets.TEST_ACCOUNT_ID }} + - name: Setup Java ${{ matrix.java }} + uses: actions/setup-java@v5 + with: + distribution: corretto + java-version: ${{ matrix.java }} + cache: maven + - name: Build locally + run: mvn -B -q -Dmaven.test.skip=true install --file pom.xml + - name: Generate operation SAM template + run: python3 generate-template.py --suite operation + working-directory: ./examples + - name: sam build + env: + MAVEN_OPTS: -DskipTests=true -Dmaven.test.skip=true + run: | + sam build --debug --parameter-overrides \ + 'ParameterKey=Architecture,ParameterValue=x86_64 ParameterKey=JavaVersion,ParameterValue=java${{ matrix.java }} ParameterKey=FunctionNamePrefix,ParameterValue=Java${{ matrix.java }}-Operation- ParameterKey=RoleArn,ParameterValue=${{ secrets.TEST_LAMBDA_EXECUTION_ROLE_ARN }}' + working-directory: ./examples + - name: Clean up unmanaged Lambda log groups + run: | + .github/scripts/cleanup_e2e_unmanaged_log_groups.sh --execute \ + --stack-name Java${{ matrix.java }}-JavaSDKOperationCloudBasedIntegrationTestStack + - name: sam deploy + run: | + sam deploy --stack-name Java${{ matrix.java }}-JavaSDKOperationCloudBasedIntegrationTestStack \ + --resolve-image-repos --resolve-s3 --parameter-overrides \ + 'ParameterKey=Architecture,ParameterValue=x86_64 ParameterKey=JavaVersion,ParameterValue=java${{ matrix.java }} ParameterKey=FunctionNamePrefix,ParameterValue=Java${{ matrix.java }}-Operation- ParameterKey=RoleArn,ParameterValue=${{ secrets.TEST_LAMBDA_EXECUTION_ROLE_ARN }}' + working-directory: ./examples + - name: Cloud Based Integration Tests + run: | + mvn clean test -B \ + -Dtest.cloud.enabled=true \ + -Dtest.aws.account='${{ secrets.TEST_ACCOUNT_ID }}' \ + -Dtest="CloudBasedIntegrationTest,CloudBasedOtelIntegrationTest" \ + -Dtest.function.name.prefix='Java${{ matrix.java }}-Operation-' \ + -Djunit.jupiter.execution.parallel.enabled=true \ + -Djunit.jupiter.execution.parallel.mode.default=concurrent \ + -Djunit.jupiter.execution.parallel.mode.classes.default=concurrent \ + -Djunit.jupiter.execution.parallel.config.strategy=fixed \ + -Djunit.jupiter.execution.parallel.config.fixed.parallelism=${{ env.E2E_TEST_PARALLELISM }} + working-directory: ./examples + - name: Publish test case summary + if: always() + env: + JAVA_VERSION: ${{ matrix.java }} Operation Examples + run: python3 ../.github/scripts/publish_e2e_test_summary.py + working-directory: ./examples diff --git a/examples/README.md b/examples/README.md index 3c84477e6..397d24fd1 100644 --- a/examples/README.md +++ b/examples/README.md @@ -22,12 +22,20 @@ cd examples # Run all tests mvn test -# Run specific test +# Run the operation-based example suite +mvn test -Dtest='software.amazon.lambda.durable.examples.operation.**.*Test' + +# Run a specific test mvn test -Dtest=SimpleStepExampleTest ``` The local runner executes in-memory and skips wait durations—ideal for fast iteration and CI/CD. +The operation-based suite under +[`examples/operation`](src/main/java/software/amazon/lambda/durable/examples/operation) mirrors the context-based +examples one-to-one. Its handlers use the context-free `Durable*Operation` facades and the one-argument +`DurableHandler.handleRequest(input)` hook. + ## Deploy to AWS ```bash @@ -38,6 +46,14 @@ sam build sam deploy --guided ``` +Generate and deploy the operation-based suite instead: + +```bash +python3 generate-template.py --suite operation +sam build +sam deploy --guided +``` + On first deploy, SAM will prompt for stack name and region. Subsequent deploys use saved config: ```bash diff --git a/examples/generate-template.py b/examples/generate-template.py index d223d37d4..5907525c7 100755 --- a/examples/generate-template.py +++ b/examples/generate-template.py @@ -12,9 +12,11 @@ EXAMPLES_DIR = Path(__file__).resolve().parent SOURCE_ROOT = EXAMPLES_DIR / "src/main/java" EXAMPLE_PACKAGE_ROOT = SOURCE_ROOT / "software/amazon/lambda/durable/examples" +OPERATION_PACKAGE_ROOT = EXAMPLE_PACKAGE_ROOT / "operation" DEFAULT_OUTPUT = EXAMPLES_DIR / "template.yaml" TEMPLATE_ANNOTATION = "ExampleTemplate" POM_NAMESPACE = {"m": "http://maven.apache.org/POM/4.0.0"} +SUITES = ("context", "operation") def read_otel_plugin_jar_path() -> str: @@ -93,9 +95,12 @@ def read_template_annotation(source: str, class_name: str) -> tuple[str | None, return condition, tracing, java_agent -def discover_examples() -> list[ExampleFunction]: +def discover_examples(suite: str) -> list[ExampleFunction]: + package_root = OPERATION_PACKAGE_ROOT if suite == "operation" else EXAMPLE_PACKAGE_ROOT examples = [] - for path in sorted(EXAMPLE_PACKAGE_ROOT.rglob("*.java")): + for path in sorted(package_root.rglob("*.java")): + if suite == "context" and OPERATION_PACKAGE_ROOT in path.parents: + continue source = path.read_text(encoding="utf-8") class_name = path.stem if not is_top_level_durable_handler(source, class_name): @@ -259,14 +264,20 @@ def render_template(examples: list[ExampleFunction]) -> str: def main() -> None: parser = argparse.ArgumentParser(description="Generate the examples SAM template from Java example handlers.") parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT, help="Path to write the generated template.") + parser.add_argument( + "--suite", + choices=SUITES, + default="context", + help="Example suite to deploy: context-based or operation-based.", + ) args = parser.parse_args() - examples = discover_examples() + examples = discover_examples(args.suite) if not examples: - raise RuntimeError("No DurableHandler examples found") + raise RuntimeError(f"No DurableHandler examples found for suite {args.suite}") args.output.write_text(render_template(examples), encoding="utf-8") - print(f"Generated {args.output} with {len(examples)} Lambda functions.") + print(f"Generated {args.output} with {len(examples)} {args.suite} Lambda functions.") if __name__ == "__main__": diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/CallbackExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/CallbackExample.java new file mode 100644 index 000000000..82c4fe861 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/CallbackExample.java @@ -0,0 +1,82 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.callback; + +import java.time.Duration; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.StepContext; +import software.amazon.lambda.durable.examples.types.ApprovalRequest; +import software.amazon.lambda.durable.operation.DurableCallbackOperation; +import software.amazon.lambda.durable.operation.DurableCallbackOperation.CallbackConfig; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation; +import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation.WaitForCallbackContext; + +/** + * Example demonstrating callback operations for external system integration. + * + *

    This handler demonstrates a human approval workflow: + * + *

      + *
    1. Prepare the request for approval + *
    2. Create a callback and send the callback ID to an external approval system + *
    3. Suspend execution until the external system responds + *
    4. Process the approval result + *
    + * + *

    External systems respond using AWS Lambda APIs: + * + *

      + *
    • {@code SendDurableExecutionCallbackSuccess} - approve with result + *
    • {@code SendDurableExecutionCallbackFailure} - reject with error + *
    • {@code SendDurableExecutionCallbackHeartbeat} - keep callback alive + *
    + */ +public class CallbackExample extends DurableHandler { + + @Override + public String handleRequest(ApprovalRequest input) { + // Step 1: Prepare the approval request + var prepared = DurableStepOperation.step( + "prepare", + String.class, + () -> "Approval request for: " + input.description() + " ($" + input.amount() + ")"); + + // Step 2: Create callback for external approval + // Use timeout from input if provided, otherwise default to 5 minutes + var timeout = + input.timeoutSeconds() != null ? Duration.ofSeconds(input.timeoutSeconds()) : Duration.ofMinutes(5); + + var config = CallbackConfig.builder().timeout(timeout).build(); + + var preapprovalCallback = + DurableWaitForCallbackOperation.waitForCallbackAsync("preapproval", String.class, () -> { + var callbackId = WaitForCallbackContext.getCurrentContext().getCallbackId(); + StepContext.getCurrentContext() + .getLogger() + .info("Sending callback {} to preapproval system", callbackId); + }); + + var callback = DurableCallbackOperation.createCallback("approval", String.class, config); + + // Step 2.5: Log AWS CLI command to complete the callback + DurableStepOperation.step("log-callback-command", Void.class, () -> { + var callbackId = callback.callbackId(); + // The result must be base64-encoded JSON + var command = String.format( + "aws lambda send-durable-execution-callback-success --callback-id %s --result $(echo -n '\"approved\"' | base64)", + callbackId); + StepContext.getCurrentContext().getLogger().info("To complete this callback, run: {}", command); + return null; + }); + + var preapprovalResult = preapprovalCallback.get(); + + // Step 3: Wait for external approval (suspends execution) + var approvalResult = callback.get(); + + // Step 4: Process the approval + return DurableStepOperation.step( + "process-approval", String.class, () -> prepared + " - " + preapprovalResult + " - " + approvalResult); + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/RetryWaitForCallbackExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/RetryWaitForCallbackExample.java new file mode 100644 index 000000000..f590bc21b --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/RetryWaitForCallbackExample.java @@ -0,0 +1,62 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.callback; + +import java.time.Duration; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.StepContext; +import software.amazon.lambda.durable.examples.types.ApprovalRequest; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation; +import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation.WaitForCallbackContext; +import software.amazon.lambda.durable.operation.DurableWithRetryOperation; +import software.amazon.lambda.durable.operation.DurableWithRetryOperation.WithRetryConfig; +import software.amazon.lambda.durable.operation.DurableWithRetryOperation.WithRetryContext; +import software.amazon.lambda.durable.retry.RetryDecision; + +/** + * Example demonstrating {@link DurableWithRetryOperation} with {@link DurableWaitForCallbackOperation}. + * + *

    Submits an approval request to an external system via a callback. If the callback fails (e.g., the external system + * rejects the request), the helper retries the entire waitForCallback cycle — creating a fresh callback with a new ID + * each time. + * + *

    Each attempt uses a unique callback name ({@code "approval-1"}, {@code "approval-2"}, etc.) so the execution + * history stays clean and replay-safe. A {@code null} name is used, so attempts are grouped under a default-named child + * context. + */ +public class RetryWaitForCallbackExample extends DurableHandler { + + private static final int MAX_ATTEMPTS = 3; + + @Override + public String handleRequest(ApprovalRequest input) { + // Step 1: Prepare the approval request + var prepared = DurableStepOperation.step( + "prepare", String.class, () -> "Approval for: " + input.description() + " ($" + input.amount() + ")"); + + // Step 2: waitForCallback with retry — if the external system fails, try again with a fresh callback + var approvalResult = DurableWithRetryOperation.withRetry( + null, + () -> { + var attempt = WithRetryContext.getCurrentContext().getAttempt(); + return DurableWaitForCallbackOperation.waitForCallback( + "approval-" + attempt, String.class, () -> StepContext.getCurrentContext() + .getLogger() + .info( + "Attempt {}: sending callback {} to approval system", + attempt, + WaitForCallbackContext.getCurrentContext() + .getCallbackId())); + }, + WithRetryConfig.builder() + .retryStrategy((error, attempt) -> attempt < MAX_ATTEMPTS + ? RetryDecision.retry(Duration.ofSeconds(2)) + : RetryDecision.fail()) + .build()); + + // Step 3: Process the result + return DurableStepOperation.step( + "process-result", String.class, () -> prepared + " - Result: " + approvalResult); + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/WaitForCallbackFailedExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/WaitForCallbackFailedExample.java new file mode 100644 index 000000000..8843b8ae0 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/WaitForCallbackFailedExample.java @@ -0,0 +1,58 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.callback; + +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.StepContext; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.examples.types.ApprovalRequest; +import software.amazon.lambda.durable.exception.SerDesException; +import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; +import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation; +import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation.WaitForCallbackConfig; +import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation.WaitForCallbackContext; +import software.amazon.lambda.durable.serde.JacksonSerDes; + +public class WaitForCallbackFailedExample extends DurableHandler { + + @Override + public String handleRequest(ApprovalRequest input) { + + String approvalResult; + + try { + approvalResult = DurableWaitForCallbackOperation.waitForCallback( + "preapproval", + String.class, + () -> { + StepContext.getCurrentContext() + .getLogger() + .info( + "Sending callback {} to preapproval system", + WaitForCallbackContext.getCurrentContext() + .getCallbackId()); + throw new RuntimeException("Submitter failed with an exception"); + }, + WaitForCallbackConfig.builder() + .stepConfig(StepConfig.builder() + .serDes(new FailedSerDes()) + .build()) + .build()); + } catch (Exception ex) { + return ex.getClass().getSimpleName() + ":" + ex.getMessage(); + } + + return approvalResult; + } + + private static class FailedSerDes extends JacksonSerDes { + @Override + public T deserialize(String json, TypeToken typeToken) { + T result = super.deserialize(json, typeToken); + if (result instanceof RuntimeException ex) { + throw new SerDesException("Deserialization failed", ex); + } + return result; + } + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ChildContextExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ChildContextExample.java new file mode 100644 index 000000000..d0f82acfb --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ChildContextExample.java @@ -0,0 +1,81 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.child; + +import java.time.Duration; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.operation.DurableContextOperation; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableWaitOperation; + +/** + * Example demonstrating child context workflows with the Durable Execution SDK. + * + *

    This handler runs three concurrent child contexts using {@code runInChildContextAsync}: + * + *

      + *
    1. Order validation — performs a step then suspends via {@code wait()} before completing + *
    2. Inventory check — performs a step then suspends via {@code wait()} before completing + *
    3. Shipping estimate — nests another child context inside it to demonstrate hierarchical contexts + *
    + * + *

    All three child contexts run concurrently. Results are collected with {@link DurableFuture#allOf} and combined + * into a summary string. + */ +public class ChildContextExample extends DurableHandler { + + @Override + public String handleRequest(GreetingRequest input) { + var context = DurableContext.getCurrentContext(); + var name = input.getName(); + context.getLogger().info("Starting child context workflow for {}", name); + + // Child context 1: Order validation — step + wait + step + var orderFuture = DurableContextOperation.runInChildContextAsync("order-validation", String.class, () -> { + var prepared = DurableStepOperation.step("prepare-order", String.class, () -> "Order for " + name); + DurableContext.getCurrentContext().getLogger().info("Order prepared, waiting for validation"); + + DurableWaitOperation.wait("validation-delay", Duration.ofSeconds(5)); + + return DurableStepOperation.step("validate-order", String.class, () -> prepared + " [validated]"); + }); + + // Child context 2: Inventory check — step + wait + step + var inventoryFuture = DurableContextOperation.runInChildContextAsync("inventory-check", String.class, () -> { + var stock = DurableStepOperation.step("check-stock", String.class, () -> "Stock available for " + name); + DurableContext.getCurrentContext().getLogger().info("Stock checked, waiting for confirmation"); + + DurableWaitOperation.wait("confirmation-delay", Duration.ofSeconds(3)); + + return DurableStepOperation.step("confirm-inventory", String.class, () -> stock + " [confirmed]"); + }); + + // Child context 3: Shipping estimate — nests a child context inside it + var shippingFuture = DurableContextOperation.runInChildContextAsync("shipping-estimate", String.class, () -> { + var baseRate = + DurableStepOperation.step("calculate-base-rate", String.class, () -> "Base rate for " + name); + + // Nested child context: calculate regional adjustment + var adjustment = DurableContextOperation.runInChildContext( + "regional-adjustment", + String.class, + () -> DurableStepOperation.step( + "lookup-region", String.class, () -> baseRate + " + regional adjustment")); + + return DurableStepOperation.step("finalize-shipping", String.class, () -> adjustment + " [shipping ready]"); + }); + + // Collect all results using allOf + context.getLogger().info("Waiting for all child contexts to complete"); + var results = DurableFuture.allOf(orderFuture, inventoryFuture, shippingFuture); + + // Combine into summary + var summary = String.join(" | ", results); + context.getLogger().info("All child contexts complete: {}", summary); + + return summary; + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ManyAsyncChildContextExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ManyAsyncChildContextExample.java new file mode 100644 index 000000000..89ceed290 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ManyAsyncChildContextExample.java @@ -0,0 +1,79 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.child; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.concurrent.TimeUnit; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.types.ManyAsyncStepsInput; +import software.amazon.lambda.durable.examples.types.ManyAsyncStepsOutput; +import software.amazon.lambda.durable.operation.DurableContextOperation; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableWaitOperation; + +/** + * Performance test example demonstrating concurrent async child contexts. + * + *

    This example tests the SDK's ability to handle many concurrent operations: + * + *

      + *
    • Creates async child context in a loop + *
    • Each child context performs a simple computation in a step + *
    • All results are collected using {@link DurableFuture#allOf} + *
    + */ +public class ManyAsyncChildContextExample extends DurableHandler { + + @Override + public ManyAsyncStepsOutput handleRequest(ManyAsyncStepsInput input) { + var startTime = System.nanoTime(); + var multiplier = input.multiplier(); + var steps = input.steps(); + var logger = DurableContext.getCurrentContext().getLogger(); + + logger.info("Starting {} async child context with multiplier {}", steps, multiplier); + + // Create async steps + var futures = new ArrayList>(steps); + for (var i = 0; i < steps; i++) { + var index = i; + var future = DurableContextOperation.runInChildContextAsync("child-" + i, Integer.class, () -> { + // create a step inside the child context, which doubles the number of threads + return DurableStepOperation.step("compute-" + index, Integer.class, () -> index * multiplier); + }); + futures.add(future); + } + + logger.info("All {} async child context created, collecting results", steps); + + // Collect all results using allOf + var results = DurableFuture.allOf(futures); + var totalSum = results.stream().mapToInt(Integer::intValue).sum(); + + // checkpoint the executionTime so that we can have the same value when replay + var executionTimeMs = DurableStepOperation.step( + "execution-time", Long.class, () -> TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime)); + logger.info( + "Completed {} child context, total sum: {}, execution time: {}ms", steps, totalSum, executionTimeMs); + + // Wait 2 seconds to test replay + DurableWaitOperation.wait("post-compute-wait", Duration.ofSeconds(2)); + + var replayTimeMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime); + + return new ManyAsyncStepsOutput(totalSum, executionTimeMs, replayTimeMs); + } + + @Override + protected DurableConfig createConfiguration() { + // Add a small checkpoint delay to help batch the checkpoint requests and reduce the overall latencies + // when the function has many concurrent operations + return DurableConfig.builder() + .withCheckpointDelay(Duration.ofMillis(10)) + .build(); + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/VirtualChildContextExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/VirtualChildContextExample.java new file mode 100644 index 000000000..32031d490 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/VirtualChildContextExample.java @@ -0,0 +1,97 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.child; + +import java.time.Duration; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.operation.DurableContextOperation; +import software.amazon.lambda.durable.operation.DurableContextOperation.RunInChildContextConfig; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableWaitOperation; + +/** + * Example demonstrating virtual child context workflows with the Durable Execution SDK. + * + *

    This handler runs three concurrent child contexts using {@code runInChildContextAsync}: + * + *

      + *
    1. Order validation — performs a step then suspends via {@code wait()} before completing + *
    2. Inventory check — performs a step then suspends via {@code wait()} before completing + *
    3. Shipping estimate — nests another child context inside it to demonstrate hierarchical contexts + *
    + * + *

    All three child contexts run concurrently. Results are collected with {@link DurableFuture#allOf} and combined + * into a summary string. + */ +public class VirtualChildContextExample extends DurableHandler { + + @Override + public String handleRequest(GreetingRequest input) { + var context = DurableContext.getCurrentContext(); + var name = input.getName(); + context.getLogger().info("Starting child context workflow for {}", name); + + // Child context 1: Order validation — step + wait + step + var orderFuture = DurableContextOperation.runInChildContextAsync( + "order-validation", + String.class, + () -> { + var prepared = DurableStepOperation.step("prepare-order", String.class, () -> "Order for " + name); + DurableContext.getCurrentContext().getLogger().info("Order prepared, waiting for validation"); + + DurableWaitOperation.wait("validation-delay", Duration.ofSeconds(5)); + + return DurableStepOperation.step("validate-order", String.class, () -> prepared + " [validated]"); + }, + RunInChildContextConfig.builder().isVirtual(true).build()); + + // Child context 2: Inventory check — step + wait + step + var inventoryFuture = DurableContextOperation.runInChildContextAsync( + "inventory-check", + String.class, + () -> { + var stock = + DurableStepOperation.step("check-stock", String.class, () -> "Stock available for " + name); + DurableContext.getCurrentContext().getLogger().info("Stock checked, waiting for confirmation"); + + DurableWaitOperation.wait("confirmation-delay", Duration.ofSeconds(3)); + + return DurableStepOperation.step("confirm-inventory", String.class, () -> stock + " [confirmed]"); + }, + RunInChildContextConfig.builder().isVirtual(true).build()); + + // Child context 3: Shipping estimate — nests a child context inside it + var shippingFuture = DurableContextOperation.runInChildContextAsync( + "shipping-estimate", + String.class, + () -> { + var baseRate = DurableStepOperation.step( + "calculate-base-rate", String.class, () -> "Base rate for " + name); + + // Nested child context: calculate regional adjustment + var adjustment = DurableContextOperation.runInChildContext( + "regional-adjustment", + String.class, + () -> DurableStepOperation.step( + "lookup-region", String.class, () -> baseRate + " + regional adjustment"), + RunInChildContextConfig.builder().isVirtual(true).build()); + + return DurableStepOperation.step( + "finalize-shipping", String.class, () -> adjustment + " [shipping ready]"); + }, + RunInChildContextConfig.builder().isVirtual(true).build()); + + // Collect all results using allOf + context.getLogger().info("Waiting for all child contexts to complete"); + var results = DurableFuture.allOf(orderFuture, inventoryFuture, shippingFuture); + + // Combine into summary + var summary = String.join(" | ", results); + context.getLogger().info("All child contexts complete: {}", summary); + + return summary; + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/CustomConfigExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/CustomConfigExample.java new file mode 100644 index 000000000..ecdfad5f6 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/CustomConfigExample.java @@ -0,0 +1,138 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.general; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.PropertyNamingStrategies; +import java.io.IOException; +import java.time.Duration; +import software.amazon.awssdk.auth.credentials.EnvironmentVariableCredentialsProvider; +import software.amazon.awssdk.core.SdkSystemSetting; +import software.amazon.awssdk.http.apache.ApacheHttpClient; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.lambda.LambdaClient; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.client.LambdaDurableFunctionsClient; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.serde.SerDes; + +/** + * Example demonstrating custom configuration with both custom HTTP client and custom SerDes. Shows how to configure a + * custom Apache HTTP client for the Lambda client while maintaining automatic credentials detection and region + * fallback, plus a custom SerDes with snake_case naming. + * + *

    This example demonstrates: + * + *

      + *
    • Custom Apache HTTP client configuration for improved performance + *
    • Automatic region detection with fallback to us-east-1 for testing environments + *
    • Environment variable credentials provider + *
    • Custom SerDes with snake_case property naming + *
    • Optional post-serialization deserialization toggle for performance-sensitive workloads + *
    + */ +public class CustomConfigExample extends DurableHandler { + + @Override + protected DurableConfig createConfiguration() { + // Configure custom Apache HTTP client for better performance + var httpClient = ApacheHttpClient.builder() + .maxConnections(50) + .connectionTimeout(Duration.ofSeconds(30)) + .socketTimeout(Duration.ofSeconds(60)) + .build(); + + // Get region with fallback to us-east-1 if AWS_REGION not set + // This prevents initialization failures in testing environments + var region = System.getenv(SdkSystemSetting.AWS_REGION.environmentVariable()); + if (region == null || region.isEmpty()) { + region = "us-east-1"; + } + + // Create Lambda client with custom HTTP client + // Uses automatic credentials detection and region fallback + var lambdaClient = LambdaClient.builder() + .httpClient(httpClient) + .credentialsProvider(EnvironmentVariableCredentialsProvider.create()) + .region(Region.of(region)) + .build(); + + // Wrap the Lambda client with LambdaDurableFunctionsClient + var durableClient = new LambdaDurableFunctionsClient(lambdaClient); + + // Create custom SerDes with snake_case naming + var customSerDes = new SnakeCaseSerDes(); + + return DurableConfig.builder() + .withDurableExecutionClient(durableClient) + .withSerDes(customSerDes) + // Disable the extra deserialize pass if your workload is sensitive to the added cost. + .withDeserializeAfterSerialization(false) + .build(); + } + + @Override + public String handleRequest(String input) { + // Step 1: Create a custom object with camelCase fields to demonstrate snake_case serialization + var customObject = DurableStepOperation.step( + "create-custom-object", + CustomData.class, + () -> new CustomData("user123", "John Doe", 25, "john.doe@example.com")); + + return "Created custom object: " + customObject.userId + ", " + customObject.fullName + ", " + + customObject.userAge + ", " + customObject.emailAddress; + } + + /** + * Custom data class with camelCase field names to demonstrate snake_case serialization. The SerDes will convert + * these field names to snake_case in the JSON output. + */ + public static class CustomData { + public String userId; + public String fullName; + public int userAge; + public String emailAddress; + + public CustomData() {} + + public CustomData(String userId, String fullName, int userAge, String emailAddress) { + this.userId = userId; + this.fullName = fullName; + this.userAge = userAge; + this.emailAddress = emailAddress; + } + } + + /** + * Custom SerDes implementation using snake_case property naming. Demonstrates how to provide custom serialization + * behavior. + */ + private static class SnakeCaseSerDes implements SerDes { + private final ObjectMapper objectMapper; + + public SnakeCaseSerDes() { + this.objectMapper = new ObjectMapper().setPropertyNamingStrategy(PropertyNamingStrategies.SNAKE_CASE); + } + + @Override + public String serialize(Object obj) { + try { + return objectMapper.writeValueAsString(obj); + } catch (JsonProcessingException e) { + throw new RuntimeException("Serialization failed", e); + } + } + + @Override + public T deserialize(String json, TypeToken typeToken) { + try { + return objectMapper.readValue(json, objectMapper.constructType(typeToken.getType())); + } catch (IOException e) { + throw new RuntimeException("Deserialization failed", e); + } + } + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/CustomPollingExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/CustomPollingExample.java new file mode 100644 index 000000000..2165093de --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/CustomPollingExample.java @@ -0,0 +1,67 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.general; + +import java.time.Duration; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.operation.DurableInvokeOperation; +import software.amazon.lambda.durable.operation.DurableInvokeOperation.InvokeConfig; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.retry.JitterStrategy; +import software.amazon.lambda.durable.retry.PollingStrategies; + +/** + * Example demonstrating custom polling strategy configuration. + * + *

    The polling strategy controls how the SDK polls for async operation results. By default, the SDK uses exponential + * backoff (1s base, 2x rate, full jitter). This example shows how to customize the polling behavior. + * + *

    This example configures: + * + *

      + *
    • Exponential backoff with 500ms base interval + *
    • 1.5x backoff rate for gentler growth + *
    • Half jitter to balance between consistency and thundering herd avoidance + *
    + */ +public class CustomPollingExample extends DurableHandler { + + @Override + protected DurableConfig createConfiguration() { + return DurableConfig.builder() + .withPollingStrategy(PollingStrategies.exponentialBackoff( + Duration.ofMillis(500), 1.5, JitterStrategy.HALF, Duration.ofSeconds(5))) + .build(); + } + + @Override + public String handleRequest(GreetingRequest input) { + var context = DurableContext.getCurrentContext(); + context.getLogger().info("Starting workflow with input: {}", input); + + // Step 1: low case the input + var lowered = DurableStepOperation.stepAsync("validate", String.class, () -> { + try { + // prevent the execution from suspension + Thread.sleep(5000); + } catch (InterruptedException e) { + e.printStackTrace(); + } + return input.getName().toLowerCase(); + }); + + // Step 2: Invoke async + var future = DurableInvokeOperation.invokeAsync( + "call-greeting", + "simple-step-example" + input.getName() + ":$LATEST", + input, + String.class, + InvokeConfig.builder().build()); + // because we are sleeping 5 seconds in the first async step, the function will not be suspened. The invoke + // function will have to poll for completion. + return future.get() + lowered.get(); + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/ErrorHandlingExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/ErrorHandlingExample.java new file mode 100644 index 000000000..3d0ea70eb --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/ErrorHandlingExample.java @@ -0,0 +1,100 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.general; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.config.StepSemantics; +import software.amazon.lambda.durable.exception.StepFailedException; +import software.amazon.lambda.durable.exception.StepInterruptedException; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; +import software.amazon.lambda.durable.retry.RetryStrategies; + +/** + * Example demonstrating error handling patterns with the Durable Execution SDK. + * + *

    This example shows how to handle: + * + *

      + *
    • {@link StepFailedException} - when a step exhausts all retry attempts + *
    • {@link StepInterruptedException} - when an AT_MOST_ONCE step is interrupted + *
    • Custom exceptions - original exception types are preserved and can be caught directly + *
    + * + *

    Note: {@code NonDeterministicExecutionException} is thrown by the SDK when code changes between executions (e.g., + * step order/names changed). It should be fixed in code, not caught. + */ +public class ErrorHandlingExample extends DurableHandler { + + private static final Logger logger = LoggerFactory.getLogger(ErrorHandlingExample.class); + + /** Custom exception to demonstrate that original exception types are preserved across checkpoints. */ + public static class ServiceUnavailableException extends RuntimeException { + private String serviceName; + + /** Default constructor required for Jackson deserialization. */ + public ServiceUnavailableException() { + super(); + } + + public ServiceUnavailableException(String serviceName, String message) { + super(message); + this.serviceName = serviceName; + } + + public String getServiceName() { + return serviceName; + } + } + + @Override + public String handleRequest(Object input) { + // Example 1: Catching a custom exception type with fallback logic + // The SDK preserves the original exception type, so you can catch specific exceptions directly. + // NOTE: Exception type needs to be serializable by your SerDes implementation. + String primaryResult; + try { + primaryResult = DurableStepOperation.step( + "call-primary-service", + String.class, + () -> { + throw new ServiceUnavailableException("primary-api", "Primary service unavailable"); + }, + StepConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .build()); + } catch (ServiceUnavailableException e) { + // Catch the specific custom exception type - the SDK reconstructs the original exception + logger.warn("Service '{}' unavailable, using fallback: {}", e.getServiceName(), e.getMessage()); + primaryResult = DurableStepOperation.step("call-fallback-service", String.class, () -> "fallback-result"); + } + + // Example 2: Handling StepInterruptedException for AT_MOST_ONCE operations + // StepInterruptedException is thrown when an AT_MOST_ONCE step was started + // but the function was interrupted before the step completed on every attempt. + // In normal execution, this step succeeds. The catch block handles the + // interruption scenario that occurs during replay after an unexpected termination. + String paymentResult; + try { + paymentResult = DurableStepOperation.step( + "charge-payment", + String.class, + () -> "payment-" + input, + StepConfig.builder() + .semanticsPerRetry(StepSemantics.AT_MOST_ONCE_PER_RETRY) + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .build()); + } catch (StepInterruptedException e) { + logger.warn( + "Payment step interrupted, checking external status: {}", + e.getOperation().id()); + // In real code: check payment provider for transaction status + // If payment went through, return success; otherwise, handle appropriately + paymentResult = DurableStepOperation.step("verify-payment-status", String.class, () -> "verified-payment"); + } + + return "Completed: " + primaryResult + ", " + paymentResult; + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/GenericInputOutputExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/GenericInputOutputExample.java new file mode 100644 index 000000000..f30d8e203 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/GenericInputOutputExample.java @@ -0,0 +1,52 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.general; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; +import software.amazon.lambda.durable.retry.RetryStrategies; + +/** + * Example demonstrating a durable Lambda function that uses generic types in input and output. + * + *

    This example shows how to use TypeToken to work with generic types like {@code List}, {@code Map>}, and nested generics that cannot be represented by simple Class objects. + */ +public class GenericInputOutputExample + extends DurableHandler, Map>>> { + + private static final Logger logger = LoggerFactory.getLogger(GenericInputOutputExample.class); + + @Override + public Map>> handleRequest(Map input) { + logger.info("Starting generic types example for user: {}", input.get("userId")); + + // Fetch nested generic type with retry (Map>) + Map> categories = DurableStepOperation.step( + "fetch-categories", + new TypeToken>>() {}, + () -> { + logger.info("Fetching category details"); + var result = new HashMap>(); + result.put("electronics", List.of("laptop", "phone")); + result.put("books", List.of("fiction")); + result.put("clothing", List.of("shirt")); + return result; + }, + StepConfig.builder() + .retryStrategy(RetryStrategies.Presets.DEFAULT) + .build()); + logger.info("Fetched {} category details", categories.size()); + logger.info("Generic types example completed successfully"); + + // return a result of Map>> + return new HashMap<>(Map.of("categories", categories)); + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/GenericTypesExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/GenericTypesExample.java new file mode 100644 index 000000000..8e8db2faa --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/GenericTypesExample.java @@ -0,0 +1,93 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.general; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; +import software.amazon.lambda.durable.retry.RetryStrategies; + +/** + * Example demonstrating TypeToken support for complex generic types. + * + *

    This example shows how to use TypeToken to work with generic types like {@code List}, {@code Map}, and nested generics that cannot be represented by simple Class objects. + */ +public class GenericTypesExample extends DurableHandler { + + private static final Logger logger = LoggerFactory.getLogger(GenericTypesExample.class); + + public static class Input { + public String userId; + + public Input() {} + + public Input(String userId) { + this.userId = userId; + } + } + + public static class Output { + public List items; + public Map counts; + public Map> categories; + + public Output() {} + + public Output(List items, Map counts, Map> categories) { + this.items = items; + this.counts = counts; + this.categories = categories; + } + } + + @Override + public Output handleRequest(Input input) { + logger.info("Starting generic types example for user: {}", input.userId); + + // Step 1: Fetch a list of items (List) + List items = DurableStepOperation.step("fetch-items", new TypeToken>() {}, () -> { + logger.info("Fetching items for user: {}", input.userId); + return List.of("item1", "item2", "item3", "item4"); + }); + logger.info("Fetched {} items", items.size()); + + // Step 2: Count items by category (Map) + Map counts = + DurableStepOperation.step("count-by-category", new TypeToken>() {}, () -> { + logger.info("Counting items by category"); + var result = new HashMap(); + result.put("electronics", 2); + result.put("books", 1); + result.put("clothing", 1); + return result; + }); + logger.info("Counted {} categories", counts.size()); + + // Step 3: Fetch nested generic type with retry (Map>) + Map> categories = DurableStepOperation.step( + "fetch-categories", + new TypeToken>>() {}, + () -> { + logger.info("Fetching category details"); + var result = new HashMap>(); + result.put("electronics", List.of("laptop", "phone")); + result.put("books", List.of("fiction")); + result.put("clothing", List.of("shirt")); + return result; + }, + StepConfig.builder() + .retryStrategy(RetryStrategies.Presets.DEFAULT) + .build()); + logger.info("Fetched {} category details", categories.size()); + + logger.info("Generic types example completed successfully"); + return new Output(items, counts, categories); + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/LoggingExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/LoggingExample.java new file mode 100644 index 000000000..440d1e301 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/LoggingExample.java @@ -0,0 +1,43 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.general; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.StepContext; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.operation.DurableStepOperation; + +/** + * Example demonstrating DurableLogger usage for structured logging with execution context. + * + *

    The logger automatically includes execution metadata (durableExecutionArn, requestId, operationId, operationName) + * in log entries via MDC. By default, logs are suppressed during replay to avoid duplicates. + */ +public class LoggingExample extends DurableHandler { + Logger logger = LoggerFactory.getLogger(LoggingExample.class); + + @Override + public String handleRequest(GreetingRequest input) { + var context = DurableContext.getCurrentContext(); + // Log at execution level (outside any step) + context.getLogger(logger).info("Processing greeting for: {}", input.getName()); + + // Step 1: Create greeting - logs inside step include operation context + var greeting = DurableStepOperation.step("create-greeting", String.class, () -> { + StepContext.getCurrentContext().getLogger(logger).info("Creating greeting message"); + return "Hello, " + input.getName(); + }); + + // Step 2: Transform + var result = DurableStepOperation.step("transform", String.class, () -> { + StepContext.getCurrentContext().getLogger().info("Transforming greeting to uppercase"); + return greeting.toUpperCase() + "!"; + }); + + context.getLogger().info("Completed processing, result: {}", result); + return result; + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/NoopExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/NoopExample.java new file mode 100644 index 000000000..b35e75bd1 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/NoopExample.java @@ -0,0 +1,19 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.general; + +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.types.GreetingRequest; + +/** + * Simple example demonstrating a durable function doesn't have any durable operation + * + *

    This handler processes a greeting request and returns a greeting message + */ +public class NoopExample extends DurableHandler { + + @Override + public String handleRequest(GreetingRequest input) { + return "HELLO, " + input.getName() + "!"; + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/OtelExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/OtelExample.java new file mode 100644 index 000000000..12911957f --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/OtelExample.java @@ -0,0 +1,65 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.general; + +import io.opentelemetry.exporter.logging.LoggingSpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.otel.InvocationOtelPlugin; + +/** + * Example demonstrating OpenTelemetry instrumentation with the Durable Execution SDK. + * + *

    This handler configures the OTel plugin with: + * + *

      + *
    • Deterministic trace/span IDs (all invocations of the same execution share one trace) + *
    • MDC log enrichment (traceId, spanId, traceSampled in every log line) + *
    • Logging exporter (spans printed to stdout → CloudWatch Logs) + *
    + * + *

    In production, replace {@code LoggingSpanExporter} with {@code OtlpGrpcSpanExporter} to send spans to an OTLP + * collector (X-Ray, Datadog, etc.). + * + *

    Expected trace structure: + * + * 

    + * durable.invocation
+ * ├── durable.step:create-greeting [attempt 1]
+ * ├── durable.step:create-greeting (operation, backfilled)
+ * ├── durable.step:transform [attempt 1]
+ * └── durable.step:transform (operation, backfilled)
+ *
    + */ +public class OtelExample extends DurableHandler { + + @Override + protected DurableConfig createConfiguration() { + var otelPlugin = new InvocationOtelPlugin( + SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(LoggingSpanExporter.create()))); + + return DurableConfig.builder().withPlugins(otelPlugin).build(); + } + + @Override + public String handleRequest(GreetingRequest input) { + var context = DurableContext.getCurrentContext(); + // Log with MDC — traceId and spanId will be in the JSON output + context.getLogger().info("Starting OTel example for {}", input.getName()); + + var greeting = DurableStepOperation.step("create-greeting", String.class, () -> { + context.getLogger().info("Inside step — this log has trace context in MDC"); + return "Hello, " + input.getName(); + }); + + var result = DurableStepOperation.step("transform", String.class, () -> greeting.toUpperCase() + "!"); + + context.getLogger().info("OTel example complete: {}", result); + return result; + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/PluginExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/PluginExample.java new file mode 100644 index 000000000..6f9cee147 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/PluginExample.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.examples.operation.general; + +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.plugin.*; + +/** + * Example demonstrating plugin instrumentation with the Durable Execution SDK. + * + *

    This handler registers a simple logging plugin that prints lifecycle events to stdout (which appears in CloudWatch + * Logs when deployed). Deploy this and check CloudWatch to verify all hooks fire at the right times. + * + *

    Expected output for a successful run: + * + * 

    + * [PLUGIN] onInvocationStart: requestId=..., durableExecutionArn=..., firstInvocation=true
+ * [PLUGIN] onOperationStart: name=create-greeting, type=STEP
+ * [PLUGIN] onUserFunctionStart: name=create-greeting, attempt=1
+ * [PLUGIN] onUserFunctionEnd: name=create-greeting, succeeded=true
+ * [PLUGIN] onOperationEnd: name=create-greeting
+ * [PLUGIN] onOperationStart: name=transform, type=STEP
+ * [PLUGIN] onUserFunctionStart: name=transform, attempt=1
+ * [PLUGIN] onUserFunctionEnd: name=transform, succeeded=true
+ * [PLUGIN] onOperationEnd: name=transform
+ * [PLUGIN] onInvocationEnd: status=SUCCEEDED
+ *
    + */ +public class PluginExample extends DurableHandler { + + @Override + protected DurableConfig createConfiguration() { + return DurableConfig.builder().withPlugins(new LoggingPlugin()).build(); + } + + @Override + public String handleRequest(GreetingRequest input) { + var context = DurableContext.getCurrentContext(); + context.getLogger().info("Starting plugin example for {}", input.getName()); + + var greeting = DurableStepOperation.step("create-greeting", String.class, () -> "Hello, " + input.getName()); + + var result = DurableStepOperation.step("transform", String.class, () -> greeting.toUpperCase() + "!"); + + context.getLogger().info("Plugin example complete: {}", result); + return result; + } + + /** A simple plugin that logs all lifecycle events to stdout. In Lambda, stdout goes to CloudWatch Logs. */ + static class LoggingPlugin implements DurableExecutionPlugin { + + @Override + public void onInvocationStart(InvocationInfo info) { + System.out.printf( + "[PLUGIN] onInvocationStart: requestId=%s, durableExecutionArn=%s, firstInvocation=%s%n", + info.requestId(), info.durableExecutionArn(), info.isFirstInvocation()); + } + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + System.out.printf( + "[PLUGIN] onInvocationEnd: status=%s, error=%s%n", + info.invocationStatus(), + info.executionError() != null ? info.executionError().getMessage() : null); + } + + @Override + public void onOperationStart(OperationInfo info) { + System.out.printf( + "[PLUGIN] onOperationStart: name=%s, type=%s, id=%s%n", info.name(), info.type(), info.id()); + } + + @Override + public void onOperationEnd(OperationEndInfo info) { + System.out.printf( + "[PLUGIN] onOperationEnd: name=%s, type=%s, error=%s%n", + info.name(), + info.type(), + info.error() != null ? info.error().getMessage() : null); + } + + @Override + public void onUserFunctionStart(UserFunctionStartInfo info) { + System.out.printf( + "[PLUGIN] onUserFunctionStart: name=%s, type=%s, attempt=%s, isReplayingChildren=%s%n", + info.name(), info.type(), info.attempt(), info.isReplayingChildren()); + } + + @Override + public void onUserFunctionEnd(UserFunctionEndInfo info) { + System.out.printf( + "[PLUGIN] onUserFunctionEnd: name=%s, succeeded=%s, error=%s%n", + info.name(), + info.succeeded(), + info.error() != null ? info.error().getMessage() : null); + } + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/invoke/RetryInvokeExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/invoke/RetryInvokeExample.java new file mode 100644 index 000000000..a8dbd2fa2 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/invoke/RetryInvokeExample.java @@ -0,0 +1,45 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.invoke; + +import java.time.Duration; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.operation.DurableInvokeOperation; +import software.amazon.lambda.durable.operation.DurableWithRetryOperation; +import software.amazon.lambda.durable.operation.DurableWithRetryOperation.WithRetryConfig; +import software.amazon.lambda.durable.operation.DurableWithRetryOperation.WithRetryContext; +import software.amazon.lambda.durable.retry.RetryDecision; + +/** + * Example demonstrating {@link DurableWithRetryOperation} with {@link DurableInvokeOperation}. + * + *

    Retries a chained Lambda invocation up to 3 times with a fixed 2-second backoff between attempts. Each attempt + * uses a unique operation name ({@code "call-greeting-1"}, {@code "call-greeting-2"}, etc.) so the execution history + * stays clean and replay-safe. + * + *

    A {@code null} name is used, so attempts are grouped under a default-named child context. + */ +public class RetryInvokeExample extends DurableHandler { + + private static final int MAX_ATTEMPTS = 3; + + @Override + public String handleRequest(GreetingRequest input) { + var targetFunctionName = + System.getenv().getOrDefault("FUNCTION_NAME_PREFIX", "") + "simple-step-example:$LATEST"; + + return DurableWithRetryOperation.withRetry( + null, + () -> { + var attempt = WithRetryContext.getCurrentContext().getAttempt(); + return DurableInvokeOperation.invoke( + "call-greeting-" + attempt, targetFunctionName, input, String.class); + }, + WithRetryConfig.builder() + .retryStrategy((error, attempt) -> attempt < MAX_ATTEMPTS + ? RetryDecision.retry(Duration.ofSeconds(2)) + : RetryDecision.fail()) + .build()); + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/invoke/SimpleInvokeExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/invoke/SimpleInvokeExample.java new file mode 100644 index 000000000..d4ca1d68d --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/invoke/SimpleInvokeExample.java @@ -0,0 +1,37 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.invoke; + +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.operation.DurableInvokeOperation; +import software.amazon.lambda.durable.operation.DurableInvokeOperation.InvokeConfig; + +/** + * Simple example demonstrating basic invoke execution with the Durable Execution SDK. + * + *

    This handler invokes another Lambda function, such as simple-step-example. + */ +public class SimpleInvokeExample extends DurableHandler { + + @Override + public String handleRequest(GreetingRequest input) { + var targetFunctionName = + System.getenv().getOrDefault("FUNCTION_NAME_PREFIX", "") + "simple-step-example:$LATEST"; + + // Invoke the `simple-step-example` function. + var future = DurableInvokeOperation.invokeAsync( + "call-greeting1", + targetFunctionName, + input, + String.class, + InvokeConfig.builder().build()); + var result2 = DurableInvokeOperation.invoke( + "call-greeting2", + targetFunctionName, + input, + String.class, + InvokeConfig.builder().build()); + return future.get() + result2; + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexFlatMapExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexFlatMapExample.java new file mode 100644 index 000000000..03a4335a9 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexFlatMapExample.java @@ -0,0 +1,81 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.map; + +import java.time.Duration; +import java.util.List; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.operation.DurableConcurrencyOperation.CompletionConfig; +import software.amazon.lambda.durable.operation.DurableConcurrencyOperation.NestingType; +import software.amazon.lambda.durable.operation.DurableMapOperation; +import software.amazon.lambda.durable.operation.DurableMapOperation.MapConfig; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableWaitOperation; + +/** + * Example demonstrating advanced map features: wait operations inside branches, error handling, and early termination. + * + *

      + *
    1. Concurrent map with step + wait + step inside each branch — simulates multi-stage order processing with a + * cooldown between stages + *
    2. Early termination with {@code minSuccessful(2)} — finds 2 healthy servers then stops + *
    + */ +public class ComplexFlatMapExample extends DurableHandler { + + @Override + public String handleRequest(Integer input) { + var context = DurableContext.getCurrentContext(); + context.getLogger().info("Starting complex map example with {} items", input); + + // Part 1: Concurrent map with step + wait inside each branch + var orderIds = IntStream.range(1, input + 1).mapToObj(x -> "order-" + x).collect(Collectors.toList()); + + var orderResult = DurableMapOperation.map( + "process-orders", + orderIds, + String.class, + orderId -> { + var index = DurableMapOperation.MapItemContext.getCurrentContext() + .getIndex(); + // Step 1: validate the order + var validated = + DurableStepOperation.step("validate-" + index, String.class, () -> "validated:" + orderId); + + // Wait between stages (simulates a cooldown or external dependency) + DurableWaitOperation.wait("cooldown-" + index, Duration.ofSeconds(1)); + + // Step 2: finalize the order + return DurableStepOperation.step("finalize-" + index, String.class, () -> "done:" + validated); + }, + MapConfig.builder().nestingType(NestingType.FLAT).build()); + + var orderSummary = String.join(", ", orderResult.results()); + + // Part 2: Early termination — find 2 healthy servers then stop + var servers = List.of("server-1", "server-2", "server-3", "server-4", "server-5"); + var earlyTermConfig = MapConfig.builder() + .completionConfig(CompletionConfig.minSuccessful(2)) + .nestingType(NestingType.FLAT) + .build(); + + var serverResult = DurableMapOperation.map( + "find-healthy-servers", + servers, + String.class, + server -> { + var index = DurableMapOperation.MapItemContext.getCurrentContext() + .getIndex(); + return DurableStepOperation.step("health-check-" + index, String.class, () -> server + ":healthy"); + }, + earlyTermConfig); + + var healthyServers = serverResult.succeeded().stream().collect(Collectors.joining(", ")); + + return String.format( + "orders=[%s] | servers=[%s] reason=%s", orderSummary, healthyServers, serverResult.completionReason()); + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexMapExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexMapExample.java new file mode 100644 index 000000000..c466faffb --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexMapExample.java @@ -0,0 +1,72 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.map; + +import java.time.Duration; +import java.util.List; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.operation.DurableConcurrencyOperation.CompletionConfig; +import software.amazon.lambda.durable.operation.DurableMapOperation; +import software.amazon.lambda.durable.operation.DurableMapOperation.MapConfig; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableWaitOperation; + +/** + * Example demonstrating advanced map features: wait operations inside branches, error handling, and early termination. + * + *
      + *
    1. Concurrent map with step + wait + step inside each branch — simulates multi-stage order processing with a + * cooldown between stages + *
    2. Early termination with {@code minSuccessful(2)} — finds 2 healthy servers then stops + *
    + */ +public class ComplexMapExample extends DurableHandler { + + @Override + public String handleRequest(Integer input) { + var context = DurableContext.getCurrentContext(); + context.getLogger().info("Starting complex map example with {} items", input); + + // Part 1: Concurrent map with step + wait inside each branch + var orderIds = IntStream.range(1, input + 1).mapToObj(x -> "order-" + x).collect(Collectors.toList()); + + var orderResult = DurableMapOperation.map("process-orders", orderIds, String.class, orderId -> { + var index = DurableMapOperation.MapItemContext.getCurrentContext().getIndex(); + // Step 1: validate the order + var validated = DurableStepOperation.step("validate-" + index, String.class, () -> "validated:" + orderId); + + // Wait between stages (simulates a cooldown or external dependency) + DurableWaitOperation.wait("cooldown-" + index, Duration.ofSeconds(1)); + + // Step 2: finalize the order + return DurableStepOperation.step("finalize-" + index, String.class, () -> "done:" + validated); + }); + + var orderSummary = String.join(", ", orderResult.results()); + + // Part 2: Early termination — find 2 healthy servers then stop + var servers = List.of("server-1", "server-2", "server-3", "server-4", "server-5"); + var earlyTermConfig = MapConfig.builder() + .completionConfig(CompletionConfig.minSuccessful(2)) + .build(); + + var serverResult = DurableMapOperation.map( + "find-healthy-servers", + servers, + String.class, + server -> { + var index = DurableMapOperation.MapItemContext.getCurrentContext() + .getIndex(); + return DurableStepOperation.step("health-check-" + index, String.class, () -> server + ":healthy"); + }, + earlyTermConfig); + + var healthyServers = serverResult.succeeded().stream().collect(Collectors.joining(", ")); + + return String.format( + "orders=[%s] | servers=[%s] reason=%s", orderSummary, healthyServers, serverResult.completionReason()); + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/CustomShouldCompleteMapExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/CustomShouldCompleteMapExample.java new file mode 100644 index 000000000..fa42bef79 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/CustomShouldCompleteMapExample.java @@ -0,0 +1,88 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.map; + +import static software.amazon.lambda.durable.model.ConcurrencyCompletionStatus.CUSTOM_COMPLETION_FAILED; +import static software.amazon.lambda.durable.model.ConcurrencyCompletionStatus.CUSTOM_COMPLETION_SUCCEEDED; +import static software.amazon.lambda.durable.model.MapResult.MapResultItem.Status.SKIPPED; +import static software.amazon.lambda.durable.operation.DurableConcurrencyOperation.CompletionConfig.CompletionDecision.complete; +import static software.amazon.lambda.durable.operation.DurableConcurrencyOperation.CompletionConfig.CompletionDecision.continueExecution; + +import java.util.List; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.operation.DurableConcurrencyOperation.CompletionConfig; +import software.amazon.lambda.durable.operation.DurableMapOperation; +import software.amazon.lambda.durable.operation.DurableMapOperation.MapConfig; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; +import software.amazon.lambda.durable.retry.RetryStrategies; + +/** + * Example demonstrating a custom {@code shouldComplete} condition for a map operation. + * + *

    The operation completes successfully once enough providers respond, or completes unsuccessfully once too many + * providers fail. The custom decision chooses both when the map completes and which completion status the caller sees. + */ +public class CustomShouldCompleteMapExample + extends DurableHandler { + + public record Input(List providers, int requiredSuccesses, int failureLimit) { + public Input { + providers = providers == null + ? List.of("primary", "secondary", "bad-cache", "tertiary") + : List.copyOf(providers); + requiredSuccesses = requiredSuccesses > 0 ? requiredSuccesses : 2; + failureLimit = failureLimit > 0 ? failureLimit : 2; + } + } + + public record Output( + String completionStatus, boolean completionSucceeded, List responses, int failed, int skipped) {} + + @Override + public Output handleRequest(Input input) { + var config = MapConfig.builder() + .maxConcurrency(1) + .completionConfig(CompletionConfig.shouldComplete(status -> { + if (status.successCount() >= input.requiredSuccesses()) { + return complete(CUSTOM_COMPLETION_SUCCEEDED); + } + if (status.failureCount() >= input.failureLimit()) { + return complete(CUSTOM_COMPLETION_FAILED); + } + return continueExecution(); + })) + .build(); + + var result = DurableMapOperation.map( + "query-providers", + input.providers(), + String.class, + provider -> { + var index = DurableMapOperation.MapItemContext.getCurrentContext() + .getIndex(); + return DurableStepOperation.step( + "query-" + index, + String.class, + () -> { + if (provider.startsWith("bad-")) { + throw new RuntimeException("Provider unavailable: " + provider); + } + return "response:" + provider; + }, + StepConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .build()); + }, + config); + + var skipped = (int) + result.items().stream().filter(item -> item.status() == SKIPPED).count(); + return new Output( + result.completionReason().name(), + result.completionReason().isSucceeded(), + result.succeeded(), + result.failed().size(), + skipped); + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/DeserializationFailedMapExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/DeserializationFailedMapExample.java new file mode 100644 index 000000000..58f883d3a --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/DeserializationFailedMapExample.java @@ -0,0 +1,72 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.map; + +import java.time.Duration; +import java.util.List; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.exception.SerDesException; +import software.amazon.lambda.durable.operation.DurableMapOperation; +import software.amazon.lambda.durable.operation.DurableMapOperation.MapConfig; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableWaitOperation; +import software.amazon.lambda.durable.serde.JacksonSerDes; + +/** + * Example demonstrating the map operation with the Durable Execution SDK. + * + *

    This handler processes a list of names concurrently using {@code map()}, where each item runs in its own child + * context with full checkpoint-and-replay support. + * + *

      + *
    1. Create a list of names from the input + *
    2. Map over each name concurrently, applying a greeting transformation via a durable step + *
    3. Collect and join the results + *
    + */ +public class DeserializationFailedMapExample extends DurableHandler { + + @Override + public String handleRequest(GreetingRequest input) { + var context = DurableContext.getCurrentContext(); + var name = input.getName(); + context.getLogger().info("Starting map example for {}", name); + + var names = List.of(name, name.toUpperCase(), name.toLowerCase()); + + // Map over each name concurrently — each iteration runs in its own child context + var result = DurableMapOperation.map( + "greet-all", + names, + String.class, + item -> { + var index = DurableMapOperation.MapItemContext.getCurrentContext() + .getIndex(); + return DurableStepOperation.step("greet-" + index, String.class, () -> { + throw new RuntimeException("Failure from " + item + "!"); + }); + }, + MapConfig.builder().serDes(new FailedSerDes()).build()); + + context.getLogger().info("Map completed: allSucceeded={}, size={}", result.allSucceeded(), result.size()); + + DurableWaitOperation.wait("suspend and replay", Duration.ofSeconds(1)); + + return result.getError(0).errorMessage(); + } + + private static class FailedSerDes extends JacksonSerDes { + + @Override + public T deserialize(String json, TypeToken typeToken) { + T result = super.deserialize(json, typeToken); + if (result instanceof RuntimeException ex) { + throw new SerDesException("Deserialization failed", ex); + } + return result; + } + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/SimpleMapExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/SimpleMapExample.java new file mode 100644 index 000000000..b9c1f440e --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/SimpleMapExample.java @@ -0,0 +1,44 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.map; + +import java.util.List; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.operation.DurableMapOperation; +import software.amazon.lambda.durable.operation.DurableStepOperation; + +/** + * Example demonstrating the map operation with the Durable Execution SDK. + * + *

    This handler processes a list of names concurrently using {@code map()}, where each item runs in its own child + * context with full checkpoint-and-replay support. + * + *

      + *
    1. Create a list of names from the input + *
    2. Map over each name concurrently, applying a greeting transformation via a durable step + *
    3. Collect and join the results + *
    + */ +public class SimpleMapExample extends DurableHandler { + + @Override + public String handleRequest(GreetingRequest input) { + var context = DurableContext.getCurrentContext(); + var name = input.getName(); + context.getLogger().info("Starting map example for {}", name); + + var names = List.of(name, name.toUpperCase(), name.toLowerCase()); + + // Map over each name concurrently — each iteration runs in its own child context + var result = DurableMapOperation.map("greet-all", names, String.class, item -> { + var index = DurableMapOperation.MapItemContext.getCurrentContext().getIndex(); + return DurableStepOperation.step("greet-" + index, String.class, () -> "Hello, " + item + "!"); + }); + + context.getLogger().info("Map completed: allSucceeded={}, size={}", result.allSucceeded(), result.size()); + + return String.join(" | ", result.results()); + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExamples.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExamples.java new file mode 100644 index 000000000..6b238e13d --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExamples.java @@ -0,0 +1,109 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.otel; + +import io.opentelemetry.exporter.logging.LoggingSpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import java.util.List; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.operation.DurableContextOperation; +import software.amazon.lambda.durable.operation.DurableMapOperation; +import software.amazon.lambda.durable.operation.DurableParallelOperation; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.otel.InvocationOtelPlugin; + +/** + * OTel examples for map, parallel, and nested context operations. These are local-only examples (not deployed to + * Lambda) used to verify the OTel plugin doesn't break execution for these patterns. + */ +public final class OtelXRayExamples { + + private OtelXRayExamples() {} + + private static DurableConfig localOtelConfig() { + var otelPlugin = new InvocationOtelPlugin( + SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(LoggingSpanExporter.create()))); + return DurableConfig.builder().withPlugins(otelPlugin).build(); + } + + /** Map operation that processes items concurrently. */ + public static class MapExample extends DurableHandler { + + @Override + protected DurableConfig createConfiguration() { + return localOtelConfig(); + } + + @Override + public String handleRequest(GreetingRequest input) { + var context = DurableContext.getCurrentContext(); + context.getLogger().info("Starting OTel X-Ray map example for {}", input.getName()); + + var items = List.of("alpha", "beta", "gamma"); + var result = DurableMapOperation.map( + "process-items", + items, + String.class, + item -> DurableStepOperation.step("transform-" + item, String.class, () -> item.toUpperCase())); + + return "Mapped " + result.succeeded().size() + " items"; + } + } + + /** Parallel operation with multiple branches. */ + public static class ParallelExample extends DurableHandler { + + @Override + protected DurableConfig createConfiguration() { + return localOtelConfig(); + } + + @Override + public String handleRequest(GreetingRequest input) { + var context = DurableContext.getCurrentContext(); + context.getLogger().info("Starting OTel X-Ray parallel example for {}", input.getName()); + + var parallel = DurableParallelOperation.parallel("fan-out"); + try (parallel) { + parallel.branch( + "branch-a", + String.class, + childCtx -> DurableStepOperation.step("step-a", String.class, () -> "A: " + input.getName())); + parallel.branch( + "branch-b", + String.class, + childCtx -> DurableStepOperation.step("step-b", String.class, () -> "B: " + input.getName())); + } + var result = parallel.get(); + + return "Parallel completed: " + result.succeeded() + " branches"; + } + } + + /** Nested child contexts with inner steps. */ + public static class NestedContextExample extends DurableHandler { + + @Override + protected DurableConfig createConfiguration() { + return localOtelConfig(); + } + + @Override + public String handleRequest(GreetingRequest input) { + var context = DurableContext.getCurrentContext(); + context.getLogger().info("Starting OTel X-Ray nested context example for {}", input.getName()); + + return DurableContextOperation.runInChildContext("outer", String.class, () -> { + var intermediate = + DurableStepOperation.step("outer-step", String.class, () -> "Hello, " + input.getName()); + return DurableContextOperation.runInChildContext("inner", String.class, () -> { + return DurableStepOperation.step("deep-step", String.class, () -> intermediate.toUpperCase() + "!"); + }); + }); + } + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionStepExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionStepExample.java new file mode 100644 index 000000000..2efffb5c5 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionStepExample.java @@ -0,0 +1,40 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.otel; + +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.ExampleTemplate; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.otel.ExecutionOtelPlugin; + +/** + * OTel + X-Ray example using the ExecutionOtelPlugin with the no-arg constructor. + * + *

    {@link ExecutionOtelPlugin#ExecutionOtelPlugin()} uses the global provider initialized by the ADOT Java agent. The + * ExecutionOtelPlugin renders the Workflow span as the trace root with operations as siblings of the invocation span. + */ +@ExampleTemplate(tracing = true, javaAgent = true) +public class OtelXRayExecutionStepExample extends DurableHandler { + + @Override + protected DurableConfig createConfiguration() { + return DurableConfig.builder().withPlugins(new ExecutionOtelPlugin()).build(); + } + + @Override + public String handleRequest(GreetingRequest input) { + var context = DurableContext.getCurrentContext(); + context.getLogger().info("Starting OTel X-Ray execution view example for {}", input.getName()); + + var greeting = + DurableStepOperation.step("exec-create-greeting", String.class, () -> "Hello, " + input.getName()); + + var result = DurableStepOperation.step("exec-transform", String.class, () -> greeting.toUpperCase() + "!"); + + context.getLogger().info("OTel X-Ray execution view example complete: {}", result); + return result; + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionWaitExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionWaitExample.java new file mode 100644 index 000000000..c1f2965b9 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionWaitExample.java @@ -0,0 +1,44 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.otel; + +import java.time.Duration; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.ExampleTemplate; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableWaitOperation; +import software.amazon.lambda.durable.otel.ExecutionOtelPlugin; + +/** + * OTel + X-Ray example using ExecutionOtelPlugin with a step → wait → step pattern. + * + *

    Exercises the multi-invocation tracing scenario with the workflow-rooted trace structure. The Workflow span is + * only exported on the terminal invocation, producing a clean single-execution trace. + */ +@ExampleTemplate(tracing = true, javaAgent = true) +public class OtelXRayExecutionWaitExample extends DurableHandler { + + @Override + protected DurableConfig createConfiguration() { + return DurableConfig.builder().withPlugins(new ExecutionOtelPlugin()).build(); + } + + @Override + public String handleRequest(GreetingRequest input) { + var context = DurableContext.getCurrentContext(); + context.getLogger().info("Starting OTel X-Ray execution view wait example for {}", input.getName()); + + var before = DurableStepOperation.step("exec-before-wait", String.class, () -> "Prepared: " + input.getName()); + + DurableWaitOperation.wait("exec-pause", Duration.ofSeconds(5)); + + var after = + DurableStepOperation.step("exec-after-wait", String.class, () -> before + " | Resumed and completed"); + + context.getLogger().info("OTel X-Ray execution view wait example complete: {}", after); + return after; + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayStepExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayStepExample.java new file mode 100644 index 000000000..e81ca360e --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayStepExample.java @@ -0,0 +1,54 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.otel; + +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.ExampleTemplate; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.otel.InvocationOtelPlugin; + +/** + * OTel + X-Ray example: simple steps in a single invocation. + * + *

    Exports spans through the ADOT Java agent global OpenTelemetry provider. Requires: + * + *

      + *
    • {@code Tracing: Active} on the Lambda function + *
    • ADOT Lambda Layer added to the function + *
    • {@code OTEL_JAVAAGENT_EXTENSIONS} pointing at the OTel plugin jar + *
    + * + *

    Expected trace structure in X-Ray: + * + * 

    + * invocation
+ * ├── create-greeting
+ * │   └── create-greeting attempt 1
+ * └── transform
+ *     └── transform attempt 1
+ *
    + */ +@ExampleTemplate(tracing = true, javaAgent = true) +public class OtelXRayStepExample extends DurableHandler { + + @Override + protected DurableConfig createConfiguration() { + return DurableConfig.builder().withPlugins(new InvocationOtelPlugin()).build(); + } + + @Override + public String handleRequest(GreetingRequest input) { + var context = DurableContext.getCurrentContext(); + context.getLogger().info("Starting OTel X-Ray step example for {}", input.getName()); + + var greeting = DurableStepOperation.step("create-greeting", String.class, () -> "Hello, " + input.getName()); + + var result = DurableStepOperation.step("transform", String.class, () -> greeting.toUpperCase() + "!"); + + context.getLogger().info("OTel X-Ray step example complete: {}", result); + return result; + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayWaitExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayWaitExample.java new file mode 100644 index 000000000..6cd1eb8ac --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayWaitExample.java @@ -0,0 +1,70 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.otel; + +import java.time.Duration; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.ExampleTemplate; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableWaitOperation; +import software.amazon.lambda.durable.otel.InvocationOtelPlugin; + +/** + * OTel + X-Ray example: step → wait → step pattern that forces multiple Lambda invocations. + * + *

    This handler exercises the critical multi-invocation tracing scenario: + * + *

      + *
    1. Invocation 1: "before-wait" step completes → wait suspends execution + *
    2. Invocation 2: replays "before-wait" (no-op) → wait completes → "after-wait" step runs + *
    + * + *

    Exports spans through the ADOT Java agent global OpenTelemetry provider. Requires: + * + *

      + *
    • {@code Tracing: Active} on the Lambda function + *
    • ADOT Lambda Layer added to the function + *
    • {@code OTEL_JAVAAGENT_EXTENSIONS} pointing at the OTel plugin jar + *
    + * + *

    Expected trace structure in X-Ray (all under one trace ID — backend propagates same Root): + * + * 

    + * Trace (single trace ID across both invocations)
+ * ├── invocation (invocation 1)
+ * │   ├── before-wait
+ * │   │   └── before-wait attempt 1
+ * │   └── pause (ended as PENDING)
+ * └── invocation (invocation 2)
+ *     ├── pause (completed)
+ *     └── after-wait
+ *         └── after-wait attempt 1
+ *
    + */ +@ExampleTemplate(tracing = true, javaAgent = true) +public class OtelXRayWaitExample extends DurableHandler { + + @Override + protected DurableConfig createConfiguration() { + return DurableConfig.builder().withPlugins(new InvocationOtelPlugin()).build(); + } + + @Override + public String handleRequest(GreetingRequest input) { + var context = DurableContext.getCurrentContext(); + context.getLogger().info("Starting OTel X-Ray wait example for {}", input.getName()); + + var before = DurableStepOperation.step("before-wait", String.class, () -> "Prepared: " + input.getName()); + + // This wait forces Lambda to suspend and re-invoke after the duration + DurableWaitOperation.wait("pause", Duration.ofSeconds(5)); + + var after = DurableStepOperation.step("after-wait", String.class, () -> before + " | Resumed and completed"); + + context.getLogger().info("OTel X-Ray wait example complete: {}", after); + return after; + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/DeserializationFailedParallelExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/DeserializationFailedParallelExample.java new file mode 100644 index 000000000..ba63c55c5 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/DeserializationFailedParallelExample.java @@ -0,0 +1,77 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.parallel; + +import java.util.List; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.ParallelDurableFuture; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.ParallelBranchConfig; +import software.amazon.lambda.durable.exception.SerDesException; +import software.amazon.lambda.durable.operation.DurableParallelOperation; +import software.amazon.lambda.durable.operation.DurableParallelOperation.ParallelConfig; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.serde.JacksonSerDes; + +/** + * Example demonstrating parallel branch execution with the Durable Execution SDK. + * + *

    This handler processes a list of items concurrently using {@code DurableParallelOperation.parallel()}: + * + *

      + *
    1. Each item is processed in its own branch (child context) + *
    2. All branches run concurrently and their results are collected + *
    3. A final step combines the results into a summary + *
    + * + *

    The {@link ParallelDurableFuture} implements {@link AutoCloseable}, so try-with-resources guarantees + * {@code join()} is called even if an exception occurs. + */ +public class DeserializationFailedParallelExample + extends DurableHandler { + + public record Input(List items) {} + + @Override + public String handleRequest(Input input) { + var logger = DurableContext.getCurrentContext().getLogger(); + var items = input.items(); + logger.info("Starting parallel processing of {} items", items.size()); + + var config = ParallelConfig.builder().build(); + + var parallel = DurableParallelOperation.parallel("process-items", config); + + try (parallel) { + var future = parallel.branch( + "process", + String.class, + branchCtx -> { + return DurableStepOperation.step("transform", String.class, () -> { + throw new RuntimeException("Intentional failure for transform"); + }); + }, + ParallelBranchConfig.builder().serDes(new FailedSerDes()).build()); + + parallel.get(); + try { + return future.get(); + } catch (Exception e) { + return e.getMessage(); + } + } + } + + private static class FailedSerDes extends JacksonSerDes { + + @Override + public T deserialize(String json, TypeToken typeToken) { + T result = super.deserialize(json, typeToken); + if (result instanceof RuntimeException ex) { + throw new SerDesException("Deserialization failed", ex); + } + return result; + } + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelExample.java new file mode 100644 index 000000000..6e35a7249 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelExample.java @@ -0,0 +1,68 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.parallel; + +import java.util.ArrayList; +import java.util.List; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.ParallelDurableFuture; +import software.amazon.lambda.durable.model.ParallelResult; +import software.amazon.lambda.durable.operation.DurableParallelOperation; +import software.amazon.lambda.durable.operation.DurableParallelOperation.ParallelConfig; +import software.amazon.lambda.durable.operation.DurableStepOperation; + +/** + * Example demonstrating parallel branch execution with the Durable Execution SDK. + * + *

    This handler processes a list of items concurrently using {@code DurableParallelOperation.parallel()}: + * + *

      + *
    1. Each item is processed in its own branch (child context) + *
    2. All branches run concurrently and their results are collected + *
    3. A final step combines the results into a summary + *
    + * + *

    The {@link ParallelDurableFuture} implements {@link AutoCloseable}, so try-with-resources guarantees + * {@code join()} is called even if an exception occurs. + */ +public class ParallelExample extends DurableHandler { + + public record Input(List items) {} + + public record Output(List results, int totalProcessed) {} + + @Override + public Output handleRequest(Input input) { + var logger = DurableContext.getCurrentContext().getLogger(); + var items = input.items(); + logger.info("Starting parallel processing of {} items", items.size()); + + var config = ParallelConfig.builder().build(); + + var futures = new ArrayList>(items.size()); + var parallel = DurableParallelOperation.parallel("process-items", config); + + try (parallel) { + for (var item : items) { + var future = parallel.branch("process-" + item, String.class, branchCtx -> { + branchCtx.getLogger().info("Processing item: {}", item); + return DurableStepOperation.step("transform-" + item, String.class, () -> item.toUpperCase()); + }); + futures.add(future); + } + } // join() called here via AutoCloseable + + ParallelResult parallelResult = parallel.get(); + logger.info( + "Parallel complete: total={}, succeeded={}, failed={}", + parallelResult.size(), + parallelResult.succeeded(), + parallelResult.failed()); + + var results = futures.stream().map(DurableFuture::get).toList(); + + return new Output(results, results.size()); + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelFailureToleranceExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelFailureToleranceExample.java new file mode 100644 index 000000000..844eb8516 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelFailureToleranceExample.java @@ -0,0 +1,80 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.parallel; + +import java.util.ArrayList; +import java.util.List; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.model.ParallelResult; +import software.amazon.lambda.durable.operation.DurableConcurrencyOperation.CompletionConfig; +import software.amazon.lambda.durable.operation.DurableParallelOperation; +import software.amazon.lambda.durable.operation.DurableParallelOperation.ParallelConfig; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; +import software.amazon.lambda.durable.retry.RetryStrategies; + +/** + * Example demonstrating parallel execution with failure tolerance. + * + *

    When {@code toleratedFailureCount} is set, the parallel operation completes successfully even if some branches + * fail — as long as the number of failures does not exceed the threshold. Failed branches produce {@code null} results + * that callers can filter out. + * + *

    Use this pattern when partial success is acceptable, for example: sending notifications to multiple channels where + * some channels may be unavailable. + */ +public class ParallelFailureToleranceExample + extends DurableHandler { + + public record Input(List services, Integer toleratedFailures, Integer minSuccessful) {} + + public record Output(int succeeded, int failed) {} + + @Override + public Output handleRequest(Input input) { + var logger = DurableContext.getCurrentContext().getLogger(); + logger.info("Starting parallel execution with toleratedFailureCount={}", input.toleratedFailures()); + + var config = ParallelConfig.builder() + .completionConfig(new CompletionConfig(input.minSuccessful, input.toleratedFailures, null)) + .build(); + + var futures = new ArrayList>(input.services().size()); + var parallel = DurableParallelOperation.parallel("call-services", config); + + try (parallel) { + for (var service : input.services()) { + var future = parallel.branch("call-" + service, String.class, branchCtx -> { + return DurableStepOperation.step( + "invoke-" + service, + String.class, + () -> { + if (service.startsWith("bad-")) { + throw new RuntimeException("Service unavailable: " + service); + } + return "ok:" + service; + }, + StepConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .build()); + }); + futures.add(future); + } + } + + ParallelResult parallelResult = parallel.get(); + logger.info( + "Parallel complete: succeeded={}, failed={}, status={}", + parallelResult.succeeded(), + parallelResult.failed(), + parallelResult.completionStatus().isSucceeded() ? "succeeded" : "failed"); + + var succeeded = parallelResult.succeeded(); + var failed = parallelResult.failed(); + + logger.info("Completed: {} succeeded, {} failed", succeeded, failed); + return new Output(succeeded, failed); + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelWithWaitExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelWithWaitExample.java new file mode 100644 index 000000000..1565a7a8a --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelWithWaitExample.java @@ -0,0 +1,76 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.parallel; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.model.ParallelResult; +import software.amazon.lambda.durable.operation.DurableParallelOperation; +import software.amazon.lambda.durable.operation.DurableParallelOperation.ParallelConfig; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableWaitOperation; + +/** + * Example demonstrating parallel branches where some branches include wait operations. + * + *

    This models a notification fan-out pattern where different channels have different delivery delays: + * + *

      + *
    • Email — sent immediately + *
    • SMS — waits for a rate-limit window before sending + *
    • Push notification — waits for a quiet-hours window before sending + *
    + * + *

    All three branches run concurrently. Branches with waits suspend without consuming compute resources and resume + * automatically once the wait elapses. The parallel operation completes once all branches finish. + */ +public class ParallelWithWaitExample + extends DurableHandler { + + public record Input(String userId, String message) {} + + public record Output(List deliveries, int success, int faiure) {} + + @Override + public Output handleRequest(Input input) { + var logger = DurableContext.getCurrentContext().getLogger(); + logger.info("Sending notifications to user {}", input.userId()); + + var config = ParallelConfig.builder().build(); + var futures = new ArrayList>(3); + var parallel = DurableParallelOperation.parallel("notify", config); + + try (parallel) { + + // Branch 1: email — no wait, deliver immediately + futures.add(parallel.branch("email", String.class, ctx -> { + DurableWaitOperation.wait("email-rate-limit-delay", Duration.ofSeconds(10)); + return DurableStepOperation.step("send-email", String.class, () -> "email:" + input.message()); + })); + + // Branch 2: SMS — wait for rate-limit window, then send + futures.add(parallel.branch("sms", String.class, ctx -> { + DurableWaitOperation.wait("sms-rate-limit-delay", Duration.ofSeconds(10)); + return DurableStepOperation.step("send-sms", String.class, () -> "sms:" + input.message()); + })); + + // Branch 3: push notification — wait for quiet-hours window, then send + futures.add(parallel.branch("push", String.class, ctx -> { + DurableWaitOperation.wait("push-quiet-delay", Duration.ofSeconds(10)); + return DurableStepOperation.step("send-push", String.class, () -> "push:" + input.message()); + })); + } + + ParallelResult result = parallel.get(); + + var deliveries = futures.stream().map(DurableFuture::get).toList(); + logger.info("All {} notifications delivered", deliveries.size()); + // Test replay + DurableWaitOperation.wait("wait for finalization", Duration.ofSeconds(5)); + return new Output(deliveries, result.succeeded(), result.failed()); + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/DeserializationFailureExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/DeserializationFailureExample.java new file mode 100644 index 000000000..6e1a1636f --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/DeserializationFailureExample.java @@ -0,0 +1,45 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.step; + +import java.time.Duration; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.exception.SerDesException; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; +import software.amazon.lambda.durable.operation.DurableWaitOperation; +import software.amazon.lambda.durable.serde.JacksonSerDes; + +public class DeserializationFailureExample extends DurableHandler { + + @Override + public String handleRequest(String input) { + try { + DurableStepOperation.step( + "fail-step", + String.class, + () -> { + throw new RuntimeException("this is a test"); + }, + StepConfig.builder().serDes(new FailedSerDes()).build()); + } catch (Exception e) { + DurableWaitOperation.wait("suspend and replay", Duration.ofSeconds(1)); + return e.getClass().getSimpleName() + ":" + e.getMessage(); + } + + throw new IllegalStateException("should not reach here"); + } + + private static class FailedSerDes extends JacksonSerDes { + + @Override + public T deserialize(String json, TypeToken typeToken) { + T result = super.deserialize(json, typeToken); + if (result instanceof RuntimeException ex) { + throw new SerDesException("Deserialization failed", ex); + } + return result; + } + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/ManyAsyncStepsExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/ManyAsyncStepsExample.java new file mode 100644 index 000000000..e8bd80297 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/ManyAsyncStepsExample.java @@ -0,0 +1,74 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.step; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.concurrent.TimeUnit; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.types.ManyAsyncStepsInput; +import software.amazon.lambda.durable.examples.types.ManyAsyncStepsOutput; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableWaitOperation; + +/** + * Performance test example demonstrating concurrent async steps. + * + *

    This example tests the SDK's ability to handle many concurrent operations: + * + *

      + *
    • Creates async steps in a loop + *
    • Each step performs a simple computation + *
    • All results are collected using {@link DurableFuture#allOf} + *
    + */ +public class ManyAsyncStepsExample extends DurableHandler { + + @Override + public ManyAsyncStepsOutput handleRequest(ManyAsyncStepsInput input) { + var startTime = System.nanoTime(); + var multiplier = input.multiplier(); + var steps = input.steps(); + var logger = DurableContext.getCurrentContext().getLogger(); + + logger.info("Starting {} async steps with multiplier {}", steps, multiplier); + + // Create async steps + var futures = new ArrayList>(steps); + for (var i = 0; i < steps; i++) { + var index = i; + var future = DurableStepOperation.stepAsync("compute-" + i, Integer.class, () -> index * multiplier); + futures.add(future); + } + + logger.info("All {} async steps created, collecting results", steps); + + // Collect all results using allOf + var results = DurableFuture.allOf(futures); + var totalSum = results.stream().mapToInt(Integer::intValue).sum(); + + // checkpoint the executionTime so that we can have the same value when replay + var executionTimeMs = DurableStepOperation.step( + "execution-time", Long.class, () -> TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime)); + logger.info("Completed {} steps, total sum: {}, execution time: {}ms", steps, totalSum, executionTimeMs); + + // Wait 2 seconds to test replay + DurableWaitOperation.wait("post-compute-wait", Duration.ofSeconds(2)); + + var replayTimeMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime); + + return new ManyAsyncStepsOutput(totalSum, executionTimeMs, replayTimeMs); + } + + @Override + protected DurableConfig createConfiguration() { + // Add a small checkpoint delay to help batch the checkpoint requests and reduce the overall latencies + // when the function has many concurrent operations + return DurableConfig.builder() + .withCheckpointDelay(Duration.ofMillis(10)) + .build(); + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/RetryExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/RetryExample.java new file mode 100644 index 000000000..c5e59bcff --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/RetryExample.java @@ -0,0 +1,88 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.step; + +import java.time.Duration; +import java.time.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; +import software.amazon.lambda.durable.operation.DurableWaitOperation; +import software.amazon.lambda.durable.retry.RetryStrategies; + +/** + * Simple example demonstrating retry strategies with a flaky API. + * + *

    This example shows: + * + *

      + *
    • A step that never retries (fails immediately) + *
    • A step that retries with default exponential backoff + *
    • Time-based failure simulation for realistic retry behavior + *
    + */ +public class RetryExample extends DurableHandler { + + private static final Logger logger = LoggerFactory.getLogger(RetryExample.class); + + private Instant startTime; + + @Override + public String handleRequest(Object input) { + // Step 1: Record start time + startTime = DurableStepOperation.step("record-start-time", Instant.class, () -> Instant.now()); + logger.info("Recorded start time: {}", startTime); + + // Step 2: Call that never retries (fails immediately) + try { + DurableStepOperation.step( + "no-retry-call", + Void.class, + () -> { + throw new RuntimeException("This operation never retries"); + }, + StepConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .build()); + } catch (Exception e) { + logger.info("No-retry step failed as expected: {}", e.getMessage()); + } + + // Step 3: Flaky API call that succeeds after retries + var result = DurableStepOperation.step( + "flaky-api-call", + String.class, + () -> { + // Fail for first 8 seconds, then succeed + var failForMillis = 8000; + var elapsed = Duration.between(startTime, Instant.now()); + + if (elapsed.toMillis() < failForMillis) { + var message = String.format( + "Flaky API failing - elapsed time (%.1fs) < %.1fs", + elapsed.toMillis() / 1000.0, failForMillis / 1000.0); + logger.warn(message); + throw new RuntimeException(message); + } else { + var message = String.format( + "Flaky API succeeded - elapsed time (%.1fs) >= %.1fs", + elapsed.toMillis() / 1000.0, failForMillis / 1000.0); + logger.info(message); + return message; + } + }, + StepConfig.builder() + .retryStrategy(RetryStrategies.Presets.DEFAULT) + .build()); + + logger.info("Flaky API result: {}", result); + + // Step 4: Wait a bit before finishing + DurableWaitOperation.wait(null, Duration.ofSeconds(2)); + + logger.info("Retry example completed successfully"); + return "Retry example completed: " + result; + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/RetryInProcessExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/RetryInProcessExample.java new file mode 100644 index 000000000..9fa2c1a27 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/RetryInProcessExample.java @@ -0,0 +1,90 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.step; + +import java.time.Duration; +import java.util.concurrent.atomic.AtomicInteger; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; +import software.amazon.lambda.durable.retry.JitterStrategy; +import software.amazon.lambda.durable.retry.RetryStrategies; + +/** + * Example demonstrating in-process retry behavior with concurrent operations. + * + *

    This example shows: + * + *

      + *
    • An async step that fails and retries while other work continues + *
    • A long-running synchronous step that keeps the process busy + *
    • Retry happens in-process without suspension because main thread is active + *
    + */ +public class RetryInProcessExample extends DurableHandler { + + private static final Logger logger = LoggerFactory.getLogger(RetryInProcessExample.class); + + private final AtomicInteger attemptCount = new AtomicInteger(0); + + @Override + public String handleRequest(Object input) { + logger.info("Starting retry in-process example"); + + // Start async step that will fail and retry + DurableFuture asyncStep = DurableStepOperation.stepAsync( + "flaky-async-operation", + String.class, + () -> { + int attempt = attemptCount.incrementAndGet(); + logger.info( + "Async operation attempt #{} in thread: {}", + attempt, + Thread.currentThread().getName()); + + // Fail first 2 attempts, succeed on 3rd + if (attempt < 3) { + var message = "Async operation failing on attempt " + attempt; + logger.warn(message); + throw new RuntimeException(message); + } else { + var message = "Async operation succeeded on attempt " + attempt; + logger.info(message); + return message; + } + }, + StepConfig.builder() + .retryStrategy(RetryStrategies.exponentialBackoff( + 5, Duration.ofSeconds(1), Duration.ofSeconds(10), 2.0, JitterStrategy.NONE)) + .build()); + + // Long-running synchronous step that keeps process busy + // This prevents suspension during async step retries + String syncResult = DurableStepOperation.step("long-running-operation", String.class, () -> { + logger.info( + "Starting long-running operation (10 seconds) in thread: {}", + Thread.currentThread().getName()); + try { + Thread.sleep(10000); // 10 seconds + logger.info("Long-running operation completed"); + return "Long operation completed"; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Long operation interrupted", e); + } + }); + + // Get async step result (should be ready by now due to retries during sync + // step) + logger.info("Getting async step result"); + String asyncResult = asyncStep.get(); + + logger.info("Sync result: {}", syncResult); + logger.info("Async result: {}", asyncResult); + + return "Retry in-process completed - Sync: " + syncResult + ", Async: " + asyncResult; + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/SimpleStepExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/SimpleStepExample.java new file mode 100644 index 000000000..b4fc969f0 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/SimpleStepExample.java @@ -0,0 +1,35 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.step; + +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.operation.DurableStepOperation; + +/** + * Simple example demonstrating basic step execution with the Durable Execution SDK. + * + *

    This handler processes a greeting request through three sequential steps: + * + *

      + *
    1. Create greeting message + *
    2. Transform to uppercase + *
    3. Add punctuation + *
    + */ +public class SimpleStepExample extends DurableHandler { + + @Override + public String handleRequest(GreetingRequest input) { + // Step 1: Create greeting + var greeting = DurableStepOperation.step("create-greeting", String.class, () -> "Hello, " + input.getName()); + + // Step 2: Transform to uppercase + var uppercase = DurableStepOperation.step("to-uppercase", String.class, () -> greeting.toUpperCase()); + + // Step 3: Add punctuation + var result = DurableStepOperation.step("add-punctuation", String.class, () -> uppercase + "!"); + + return result; + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/vt/ManyAsyncStepsVirtualThreadPoolExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/vt/ManyAsyncStepsVirtualThreadPoolExample.java new file mode 100644 index 000000000..5eee36829 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/vt/ManyAsyncStepsVirtualThreadPoolExample.java @@ -0,0 +1,78 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.vt; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.ExampleTemplate; +import software.amazon.lambda.durable.examples.types.ManyAsyncStepsInput; +import software.amazon.lambda.durable.examples.types.ManyAsyncStepsOutput; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableWaitOperation; + +/** + * Performance test example demonstrating concurrent async steps. + * + *

    This example tests the SDK's ability to handle many concurrent operations: + * + *

      + *
    • Creates async steps in a loop + *
    • Each step performs a simple computation + *
    • All results are collected using {@link DurableFuture#allOf} + *
    + */ +@ExampleTemplate(condition = "IsJava21OrLater") +public class ManyAsyncStepsVirtualThreadPoolExample extends DurableHandler { + + @Override + public ManyAsyncStepsOutput handleRequest(ManyAsyncStepsInput input) { + var startTime = System.nanoTime(); + var multiplier = input.multiplier(); + var steps = input.steps(); + var logger = DurableContext.getCurrentContext().getLogger(); + + logger.info("Starting {} async steps with multiplier {}", steps, multiplier); + + // Create async steps + var futures = new ArrayList>(steps); + for (var i = 0; i < steps; i++) { + var index = i; + var future = DurableStepOperation.stepAsync("compute-" + i, Integer.class, () -> index * multiplier); + futures.add(future); + } + + logger.info("All {} async steps created, collecting results", steps); + + // Collect all results using allOf + var results = DurableFuture.allOf(futures); + var totalSum = results.stream().mapToInt(Integer::intValue).sum(); + + // checkpoint the executionTime so that we can have the same value when replay + var executionTimeMs = DurableStepOperation.step( + "execution-time", Long.class, () -> TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime)); + logger.info("Completed {} steps, total sum: {}, execution time: {}ms", steps, totalSum, executionTimeMs); + + // Wait 2 seconds to test replay + DurableWaitOperation.wait("post-compute-wait", Duration.ofSeconds(2)); + + var replayTimeMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime); + + return new ManyAsyncStepsOutput(totalSum, executionTimeMs, replayTimeMs); + } + + @Override + protected DurableConfig createConfiguration() { + // Add a small checkpoint delay to help batch the checkpoint requests and reduce the overall latencies + // when the function has many concurrent operations + return DurableConfig.builder() + .withCheckpointDelay(Duration.ofMillis(10)) + .withExecutorService(Executors.newVirtualThreadPerTaskExecutor()) + .build(); + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/ConcurrentWaitForConditionExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/ConcurrentWaitForConditionExample.java new file mode 100644 index 000000000..6cab805ad --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/ConcurrentWaitForConditionExample.java @@ -0,0 +1,60 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.wait; + +import java.util.stream.IntStream; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.operation.DurableMapOperation; +import software.amazon.lambda.durable.operation.DurableMapOperation.MapConfig; +import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation; +import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation.WaitForConditionConfig; +import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation.WaitForConditionResult; + +/** + * Example demonstrating concurrent waitForCondition operations using map. + * + *

    Runs many (totalOperations) waitForCondition operations concurrently (maxConcurrency). Each operation: + * + *

      + *
    1. Uses attempt count as state (replay-safe). + *
    2. Fails and retries until the attempt count reaches the given threshold, and then succeeds + *
    + */ +public class ConcurrentWaitForConditionExample extends DurableHandler { + + public record Input(int threshold, int totalOperations, int maxConcurrency) {} + + @Override + public String handleRequest(Input input) { + var items = IntStream.range(0, input.totalOperations()).boxed().toList(); + + var config = MapConfig.builder().maxConcurrency(input.maxConcurrency()).build(); + + var result = DurableMapOperation.map( + "concurrent-wait-for-conditions", + items, + String.class, + item -> { + var index = DurableMapOperation.MapItemContext.getCurrentContext() + .getIndex(); + var conditionConfig = WaitForConditionConfig.builder() + .initialState(1) + .build(); + // Poll until the counter reaches the input threshold + var count = DurableWaitForConditionOperation.waitForCondition( + "condition-" + index, + Integer.class, + callCount -> { + if (callCount >= input.threshold()) { + return WaitForConditionResult.stopPolling(callCount); + } + return WaitForConditionResult.continuePolling(callCount + 1); + }, + conditionConfig); + return String.valueOf(count); + }, + config); + + return String.join(" | ", result.results()); + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAsyncExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAsyncExample.java new file mode 100644 index 000000000..2ed1d8daf --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAsyncExample.java @@ -0,0 +1,47 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.wait; + +import java.time.Duration; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableWaitOperation; + +/** + * Example demonstrating non-blocking wait with waitAsync(). + * + *

    This handler starts a wait and a step concurrently, then collects both results. The wait acts as a minimum + * duration guarantee — the step runs in parallel, and the handler only proceeds once both the step completes and the + * wait elapses. + * + *

      + *
    1. Start a 5-second async wait (non-blocking) + *
    2. Start an async step concurrently + *
    3. Collect both results — ensures at least 5 seconds have passed + *
    + */ +public class WaitAsyncExample extends DurableHandler { + + @Override + public String handleRequest(GreetingRequest input) { + var context = DurableContext.getCurrentContext(); + context.getLogger().info("Starting waitAsync example for {}", input.getName()); + + // Start a non-blocking wait — returns immediately + DurableFuture waitFuture = DurableWaitOperation.waitAsync("min-delay", Duration.ofSeconds(5)); + + // Run a step concurrently while the wait timer is ticking + DurableFuture stepFuture = + DurableStepOperation.stepAsync("process", String.class, () -> "Processed: " + input.getName()); + + // Block until both complete — guarantees at least 5 seconds elapsed + waitFuture.get(); + var result = stepFuture.get(); + + context.getLogger().info("Both wait and step complete: {}", result); + return result; + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAtLeastExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAtLeastExample.java new file mode 100644 index 000000000..ab723885a --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAtLeastExample.java @@ -0,0 +1,69 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.wait; + +import java.time.Duration; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; +import software.amazon.lambda.durable.operation.DurableWaitOperation; +import software.amazon.lambda.durable.retry.RetryStrategies; + +/** + * Example demonstrating concurrent stepAsync() with wait() operations. + * + *

    This example shows suspension behavior with pending async steps: + * + *

      + *
    • stepAsync() starts a background operation (takes 2 seconds) + *
    • wait() is called immediately (3 second duration) + *
    • The step completes successfully before suspension + *
    • Execution suspends for the wait time + *
    + */ +public class WaitAtLeastExample extends DurableHandler { + + private static final Logger logger = LoggerFactory.getLogger(WaitAtLeastExample.class); + + @Override + public String handleRequest(GreetingRequest input) { + logger.info("Starting concurrent step + wait example for: {}", input.getName()); + + // Start an async step that takes 2 seconds + DurableFuture asyncStep = DurableStepOperation.stepAsync( + "async-operation", + String.class, + () -> { + logger.info( + "Async operation starting in thread: {}", + Thread.currentThread().getName()); + try { + Thread.sleep(2000); // 2 seconds + logger.info("Async operation completed successfully"); + return "Processed: " + input.getName(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Operation interrupted", e); + } + }, + StepConfig.builder() + .retryStrategy(RetryStrategies.Presets.DEFAULT) + .build()); + + // Immediately wait for 3 seconds + // The async step will complete during this wait + logger.info("Waiting 3 seconds (async step will complete in 2s)"); + DurableWaitOperation.wait("wait-3-seconds", Duration.ofSeconds(3)); + + // After wait, get the async step result + logger.info("Resumed after wait"); + String result = asyncStep.get(); + logger.info("Final result: {}", result); + + return result; + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAtLeastInProcessExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAtLeastInProcessExample.java new file mode 100644 index 000000000..e6174f328 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAtLeastInProcessExample.java @@ -0,0 +1,69 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.wait; + +import java.time.Duration; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; +import software.amazon.lambda.durable.operation.DurableWaitOperation; +import software.amazon.lambda.durable.retry.RetryStrategies; + +/** + * Example demonstrating concurrent stepAsync() with wait() operations where no suspension occurs. + * + *

    This example shows in-process wait behavior: + * + *

      + *
    • stepAsync() starts a background operation (takes 10 seconds) + *
    • wait() is called immediately (3 second duration) + *
    • The async step takes longer than the wait duration + *
    • No suspension occurs because we've already waited long enough + *
    + */ +public class WaitAtLeastInProcessExample extends DurableHandler { + + private static final Logger logger = LoggerFactory.getLogger(WaitAtLeastInProcessExample.class); + + @Override + public String handleRequest(GreetingRequest input) { + logger.info("Starting concurrent step + wait example for: {}", input.getName()); + + // Start an async step that takes 10 seconds + DurableFuture asyncStep = DurableStepOperation.stepAsync( + "async-operation", + String.class, + () -> { + logger.info( + "Async operation starting in thread: {}", + Thread.currentThread().getName()); + try { + Thread.sleep(10000); // 10 seconds + logger.info("Async operation completed successfully"); + return "Processed: " + input.getName(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Operation interrupted", e); + } + }, + StepConfig.builder() + .retryStrategy(RetryStrategies.Presets.DEFAULT) + .build()); + + // Immediately wait for 3 seconds + // The async step will still be running and will complete after the wait + logger.info("Waiting 3 seconds (async step will complete in 10s - no suspension expected)"); + DurableWaitOperation.wait("wait-3-seconds", Duration.ofSeconds(3)); + + // After wait, get the async step result + logger.info("Wait completed, getting async result"); + String result = asyncStep.get(); + logger.info("Final result: {}", result); + + return result; + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitExample.java new file mode 100644 index 000000000..0e7796d90 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitExample.java @@ -0,0 +1,62 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.wait; + +import java.time.Duration; +import software.amazon.lambda.durable.DurableFuture; +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.operation.DurableContextOperation; +import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableWaitOperation; + +/** + * Example demonstrating step execution with wait operations. + * + *

    This handler processes a request through steps with delays: + * + *

      + *
    1. Start processing + *
    2. Wait 10 seconds + *
    3. Continue processing + *
    4. Wait 5 seconds + *
    5. Complete + *
    + */ +public class WaitExample extends DurableHandler { + + @Override + public String handleRequest(GreetingRequest input) { + // Step 1: Start processing + var started = DurableStepOperation.step( + "start-processing", String.class, () -> "Started processing for " + input.getName()); + + // Wait 10 seconds + DurableWaitOperation.wait(null, Duration.ofSeconds(10)); + + // Step 2: Continue processing + var continued = DurableStepOperation.stepAsync("continue-processing", String.class, () -> { + try { + Thread.sleep(10000); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + return started + " - continued after 10s"; + }); + + // Wait at most seconds + var wait5seconds = DurableContextOperation.runInChildContextAsync("wait-5-seconds", String.class, () -> { + DurableWaitOperation.wait("wait-5-seconds", Duration.ofSeconds(5)); + + return started + " - waited 5 seconds"; + }); + + var step2 = DurableFuture.anyOf(continued, wait5seconds); + + // Step 3: Complete + var result = DurableStepOperation.step( + "complete-processing", String.class, () -> step2 + " - completed after 5s more"); + + return result; + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitForConditionExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitForConditionExample.java new file mode 100644 index 000000000..127c548bd --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitForConditionExample.java @@ -0,0 +1,38 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.wait; + +import software.amazon.lambda.durable.DurableHandler; +import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation; +import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation.WaitForConditionConfig; +import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation.WaitForConditionResult; + +/** + * Example demonstrating the waitForCondition operation. + * + *

    This handler polls a condition function until it signals completion: + * + *

      + *
    1. The attempt count is used as a state (replay safe) + *
    2. Fails and retries until the attempt count reaches the given threshold, and then succeeds + *
    + */ +public class WaitForConditionExample extends DurableHandler { + + @Override + public Integer handleRequest(Integer threshold) { + // Poll until the counter reaches the input threshold + return DurableWaitForConditionOperation.waitForCondition( + "wait-for-condition", + Integer.class, + callCount -> { + if (callCount >= threshold) { + // Condition met, stop polling + return WaitForConditionResult.stopPolling(callCount); + } + // Condition not met, keep polling + return WaitForConditionResult.continuePolling(callCount + 1); + }, + WaitForConditionConfig.builder().initialState(1).build()); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/callback/CallbackExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/callback/CallbackExampleTest.java new file mode 100644 index 000000000..6519d4a07 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/callback/CallbackExampleTest.java @@ -0,0 +1,100 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.callback; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.services.lambda.model.ErrorObject; +import software.amazon.awssdk.services.lambda.model.OperationStatus; +import software.amazon.awssdk.services.lambda.model.OperationType; +import software.amazon.lambda.durable.examples.types.ApprovalRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class CallbackExampleTest { + + @Test + void testCallbackExampleSuspendsForApproval() { + var handler = new CallbackExample(); + var runner = LocalDurableTestRunner.create(ApprovalRequest.class, handler); + + var input = new ApprovalRequest("New laptop", 1500.00); + + // First run - prepares request and creates callback, then suspends + var result = runner.run(input); + + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Verify the callback was created + var callbackOp = runner.getOperation("approval"); + assertNotNull(callbackOp); + assertEquals(OperationType.CALLBACK, callbackOp.getType()); + assertEquals(OperationStatus.STARTED, callbackOp.getStatus()); + } + + @Test + void testCallbackExampleCompletesAfterApproval() { + var handler = new CallbackExample(); + var runner = LocalDurableTestRunner.create(ApprovalRequest.class, handler); + + var input = new ApprovalRequest("New laptop", 1500.00); + + // First run - suspends waiting for callback + var result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Simulate external system approving the request + var callbackId = runner.getCallbackId("approval"); + runner.completeCallback(callbackId, "\"Approved by manager\""); + + result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // second run - pending preapproval + var preapprovalCallbackId = runner.getCallbackId("preapproval-callback"); + runner.completeCallback(preapprovalCallbackId, "\"Sent to preapprover\""); + + // third run - callback complete, finishes processing + result = runner.run(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals( + "Approval request for: New laptop ($1500.0) - Sent to preapprover - Approved by manager", + result.getResult(String.class)); + } + + @Test + void testCallbackExampleFail() { + var handler = new CallbackExample(); + var runner = LocalDurableTestRunner.create(ApprovalRequest.class, handler); + + var input = new ApprovalRequest("New laptop", 1500.00); + + // First run - suspends waiting for callback + var result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Simulate external system approving the request + var callbackId = runner.getCallbackId("approval"); + runner.completeCallback(callbackId, "\"Approved by manager\""); + + result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // second run - pending preapproval + var preapprovalCallbackId = runner.getCallbackId("preapproval-callback"); + runner.failCallback( + preapprovalCallbackId, + ErrorObject.builder() + .errorType("error type") + .errorMessage("error message") + .build()); + + // third run - callback complete, finishes processing + result = runner.run(input); + + assertEquals(ExecutionStatus.FAILED, result.getStatus()); + assertEquals("error message", result.getError().get().errorMessage()); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/callback/RetryWaitForCallbackExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/callback/RetryWaitForCallbackExampleTest.java new file mode 100644 index 000000000..34bac5d29 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/callback/RetryWaitForCallbackExampleTest.java @@ -0,0 +1,148 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.callback; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.services.lambda.model.ErrorObject; +import software.amazon.lambda.durable.examples.types.ApprovalRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class RetryWaitForCallbackExampleTest { + + @Test + void succeedsOnFirstAttempt() { + var handler = new RetryWaitForCallbackExample(); + var runner = LocalDurableTestRunner.create(ApprovalRequest.class, handler); + var input = new ApprovalRequest("New laptop", 1500.00); + + // First run — prepares request, starts waitForCallback, suspends + var result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Complete the callback (waitForCallback names it "approval-1-callback" internally) + var callbackId = runner.getCallbackId("approval-1-callback"); + assertNotNull(callbackId, "Callback 'approval-1-callback' should have been created"); + runner.completeCallback(callbackId, "\"Approved by manager\""); + + // Run to completion + result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals( + "Approval for: New laptop ($1500.0) - Result: Approved by manager", result.getResult(String.class)); + } + + @Test + void retriesAfterFirstCallbackFails() { + var handler = new RetryWaitForCallbackExample(); + var runner = LocalDurableTestRunner.create(ApprovalRequest.class, handler); + var input = new ApprovalRequest("Server upgrade", 5000.00); + + // First run — prepares, starts waitForCallback attempt 1, suspends + var result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Fail the first callback + var callbackId1 = runner.getCallbackId("approval-1-callback"); + assertNotNull(callbackId1); + runner.failCallback( + callbackId1, + ErrorObject.builder() + .errorType("RejectedError") + .errorMessage("Rejected by first reviewer") + .build()); + + // Run — processes failure, hits backoff wait, suspends + result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Advance past the backoff wait + runner.advanceTime(); + + // Run — starts waitForCallback attempt 2, suspends + result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Complete the second callback + var callbackId2 = runner.getCallbackId("approval-2-callback"); + assertNotNull(callbackId2, "Callback 'approval-2-callback' should have been created after retry"); + runner.completeCallback(callbackId2, "\"Approved on second try\""); + + // Run to completion + result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals( + "Approval for: Server upgrade ($5000.0) - Result: Approved on second try", + result.getResult(String.class)); + } + + @Test + void failsAfterAllRetriesExhausted() { + var handler = new RetryWaitForCallbackExample(); + var runner = LocalDurableTestRunner.create(ApprovalRequest.class, handler); + var input = new ApprovalRequest("Expensive item", 10000.00); + + // First run — starts waitForCallback attempt 1 + var result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Fail callback attempt 1 + var callbackId1 = runner.getCallbackId("approval-1-callback"); + runner.failCallback( + callbackId1, + ErrorObject.builder() + .errorType("Rejected") + .errorMessage("fail 1") + .build()); + result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Advance past backoff 1, run to start attempt 2 + runner.advanceTime(); + result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Fail callback attempt 2 + var callbackId2 = runner.getCallbackId("approval-2-callback"); + runner.failCallback( + callbackId2, + ErrorObject.builder() + .errorType("Rejected") + .errorMessage("fail 2") + .build()); + result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Advance past backoff 2, run to start attempt 3 + runner.advanceTime(); + result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Fail callback attempt 3 — last attempt, retryStrategy returns fail() + var callbackId3 = runner.getCallbackId("approval-3-callback"); + runner.failCallback( + callbackId3, + ErrorObject.builder() + .errorType("Rejected") + .errorMessage("fail 3") + .build()); + result = runner.run(input); + + assertEquals(ExecutionStatus.FAILED, result.getStatus()); + } + + @Test + void suspendsOnFirstRun() { + var handler = new RetryWaitForCallbackExample(); + var runner = LocalDurableTestRunner.create(ApprovalRequest.class, handler); + var input = new ApprovalRequest("Test item", 100.00); + + var result = runner.run(input); + + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/callback/WaitForCallbackFailedExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/callback/WaitForCallbackFailedExampleTest.java new file mode 100644 index 000000000..7efadc07b --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/callback/WaitForCallbackFailedExampleTest.java @@ -0,0 +1,31 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.callback; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.examples.types.ApprovalRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class WaitForCallbackFailedExampleTest { + + @Test + void testWaitForCallbackFailedExample() { + var handler = new WaitForCallbackFailedExample(); + var runner = LocalDurableTestRunner.create(ApprovalRequest.class, handler); + + var input = new ApprovalRequest("New laptop", 1500.00); + + // First run - prepares request and creates callback, then suspends + var result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + // Verify the callback was created + assertEquals( + "CallbackSubmitterException:Step failed with error of type java.lang.RuntimeException. Message: Submitter failed with an exception", + result.getResult(String.class)); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/child/ChildContextExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/child/ChildContextExampleTest.java new file mode 100644 index 000000000..b4b4d95a1 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/child/ChildContextExampleTest.java @@ -0,0 +1,56 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.child; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class ChildContextExampleTest { + + @Test + void testChildContextExampleRunsToCompletion() { + var handler = new ChildContextExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var input = new GreetingRequest("Alice"); + var result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals( + "Order for Alice [validated] | Stock available for Alice [confirmed] | Base rate for Alice + regional adjustment [shipping ready]", + result.getResult(String.class)); + } + + @Test + void testChildContextExampleSuspendsOnFirstRun() { + var handler = new ChildContextExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var input = new GreetingRequest("Bob"); + + // First run should suspend due to wait operations inside child contexts + var result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + } + + @Test + void testChildContextExampleReplay() { + var handler = new ChildContextExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var input = new GreetingRequest("Alice"); + + // First full execution + var result1 = runner.runUntilComplete(input); + assertEquals(ExecutionStatus.SUCCEEDED, result1.getStatus()); + + // Replay — should return cached results + var result2 = runner.run(input); + assertEquals(ExecutionStatus.SUCCEEDED, result2.getStatus()); + assertEquals(result1.getResult(String.class), result2.getResult(String.class)); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/child/ManyAsyncChildContextExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/child/ManyAsyncChildContextExampleTest.java new file mode 100644 index 000000000..d0d0ee9f3 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/child/ManyAsyncChildContextExampleTest.java @@ -0,0 +1,65 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.child; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.examples.types.ManyAsyncStepsInput; +import software.amazon.lambda.durable.examples.types.ManyAsyncStepsOutput; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class ManyAsyncChildContextExampleTest { + + @Test + void testManyAsyncSteps() { + var handler = new ManyAsyncChildContextExample(); + var runner = LocalDurableTestRunner.create(ManyAsyncStepsInput.class, handler); + + var input = new ManyAsyncStepsInput(2, 500); + var result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + var output = result.getResult(ManyAsyncStepsOutput.class); + assertNotNull(output); + + // Sum of 0..499 * 2 = 499 * 500 / 2 * 2 = 249500 + assertEquals(249500, output.result()); + } + + @Test + void testManyAsyncStepsWithDefaultMultiplier() { + var handler = new ManyAsyncChildContextExample(); + var runner = LocalDurableTestRunner.create(ManyAsyncStepsInput.class, handler); + + var input = new ManyAsyncStepsInput(1, 500); + var result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + // Sum of 0..499 = 499 * 500 / 2 = 124750 + assertEquals(124750, result.getResult(ManyAsyncStepsOutput.class).result()); + } + + @Test + void testOperationsAreTracked() { + var handler = new ManyAsyncChildContextExample(); + var runner = LocalDurableTestRunner.create(ManyAsyncStepsInput.class, handler); + + var result = runner.runUntilComplete(new ManyAsyncStepsInput(1, 500)); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + // Verify some operations are tracked + assertNotNull(result.getOperation("compute-0")); + assertNotNull(result.getOperation("compute-499")); + assertNotNull(result.getOperation("compute-250")); + + assertNotNull(result.getOperation("child-0")); + assertNotNull(result.getOperation("child-499")); + assertNotNull(result.getOperation("child-250")); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/child/VirtualChildContextExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/child/VirtualChildContextExampleTest.java new file mode 100644 index 000000000..cc47994f1 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/child/VirtualChildContextExampleTest.java @@ -0,0 +1,56 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.child; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class VirtualChildContextExampleTest { + + @Test + void testVirtualChildContextExampleRunsToCompletion() { + var handler = new VirtualChildContextExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var input = new GreetingRequest("Alice"); + var result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals( + "Order for Alice [validated] | Stock available for Alice [confirmed] | Base rate for Alice + regional adjustment [shipping ready]", + result.getResult(String.class)); + } + + @Test + void testVirtualChildContextExampleSuspendsOnFirstRun() { + var handler = new VirtualChildContextExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var input = new GreetingRequest("Bob"); + + // First run should suspend due to wait operations inside child contexts + var result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + } + + @Test + void testVirtualChildContextExampleReplay() { + var handler = new VirtualChildContextExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var input = new GreetingRequest("Alice"); + + // First full execution + var result1 = runner.runUntilComplete(input); + assertEquals(ExecutionStatus.SUCCEEDED, result1.getStatus()); + + // Replay — should return cached results + var result2 = runner.run(input); + assertEquals(ExecutionStatus.SUCCEEDED, result2.getStatus()); + assertEquals(result1.getResult(String.class), result2.getResult(String.class)); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/CustomConfigExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/CustomConfigExampleTest.java new file mode 100644 index 000000000..28a99f6e0 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/CustomConfigExampleTest.java @@ -0,0 +1,43 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.general; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class CustomConfigExampleTest { + + @Test + void testCustomConfigExample() { + var handler = new CustomConfigExample(); + + // Create test runner from handler (automatically extracts config) + var runner = LocalDurableTestRunner.create(String.class, handler); + + // Run with input + var result = runner.run("test-input"); + + // Verify result + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + // Get the operation result. This is the serialized result stored in the DAR backend + var operation = result.getOperation("create-custom-object"); + var operationResult = operation.getStepDetails().result(); + + // Assert that the stepDetails result uses snake_case (based on the custom serializer) + assertTrue(operationResult.contains("user_id"), "Should contain snake_case 'user_id' key"); + assertTrue(operationResult.contains("full_name"), "Should contain snake_case 'full_name' key"); + assertTrue(operationResult.contains("user_age"), "Should contain snake_case 'user_age' key"); + assertTrue(operationResult.contains("email_address"), "Should contain snake_case 'email_address' key"); + + // Verify that we got the expected output + var output = result.getResult(String.class); + assertNotNull(output); + assertEquals("Created custom object: user123, John Doe, 25, john.doe@example.com", output); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/CustomPollingExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/CustomPollingExampleTest.java new file mode 100644 index 000000000..cdbb188a8 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/CustomPollingExampleTest.java @@ -0,0 +1,32 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.general; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class CustomPollingExampleTest { + + @Test + void testCustomPollingExample() { + var handler = new CustomPollingExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + // First run: executes validate step, then pending at wait + var input = new GreetingRequest("world"); + var output1 = runner.run(input); + + assertEquals(ExecutionStatus.PENDING, output1.getStatus()); + + // Second run: + runner.completeChainedInvoke("call-greeting", "\"hello\""); + var output2 = runner.run(input); + + assertEquals(ExecutionStatus.SUCCEEDED, output2.getStatus()); + assertEquals("helloworld", output2.getResult(String.class)); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/ErrorHandlingExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/ErrorHandlingExampleTest.java new file mode 100644 index 000000000..23ffae34e --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/ErrorHandlingExampleTest.java @@ -0,0 +1,54 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.general; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class ErrorHandlingExampleTest { + + @Test + void testErrorHandlingWithFallback() { + var handler = new ErrorHandlingExample(); + var runner = LocalDurableTestRunner.create(Object.class, handler); + + var result = runner.run("test-input"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertTrue(result.getResult(String.class).contains("fallback-result")); + } + + @Test + void testPaymentStepCompletes() { + var handler = new ErrorHandlingExample(); + var runner = LocalDurableTestRunner.create(Object.class, handler); + + var result = runner.run("order-123"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + // Normal execution: payment step succeeds with "payment-order-123" + assertTrue(result.getResult(String.class).contains("payment-order-123")); + } + + @Test + void testPaymentStepInterruptedRecovery() { + var handler = new ErrorHandlingExample(); + var runner = LocalDurableTestRunner.create(Object.class, handler); + + // First run: both steps complete normally + var result1 = runner.run("order-456"); + assertEquals(ExecutionStatus.SUCCEEDED, result1.getStatus()); + + // Simulate interruption: reset payment step to STARTED state + runner.resetCheckpointToStarted("charge-payment"); + + // Second run: StepInterruptedException is caught, recovery step executes + var result2 = runner.run("order-456"); + + assertEquals(ExecutionStatus.SUCCEEDED, result2.getStatus()); + assertTrue(result2.getResult(String.class).contains("verified-payment")); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/GenericInputOutputExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/GenericInputOutputExampleTest.java new file mode 100644 index 000000000..3a883186b --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/GenericInputOutputExampleTest.java @@ -0,0 +1,57 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.general; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class GenericInputOutputExampleTest { + + private static final TypeToken>>> resultType = new TypeToken<>() {}; + private static final TypeToken> inputType = new TypeToken<>() {}; + + @Test + void testGenericTypesExample() { + var handler = new GenericInputOutputExample(); + var runner = LocalDurableTestRunner.create(inputType, handler); + + var input = new HashMap<>(Map.of("userId", "user123")); + var result = runner.run(input); + + assertNotNull(result); + var output = result.getResult(resultType); + assertNotNull(output); + + // Verify categories nested map + var categories = output.get("categories"); + assertNotNull(categories); + assertEquals(3, categories.size()); + assertEquals(2, categories.get("electronics").size()); + assertTrue(categories.get("electronics").contains("laptop")); + assertTrue(categories.get("electronics").contains("phone")); + assertEquals(1, categories.get("books").size()); + assertTrue(categories.get("books").contains("fiction")); + } + + @Test + void testOperationTracking() { + var handler = new GenericInputOutputExample(); + var runner = LocalDurableTestRunner.create(inputType, handler); + + var input = new HashMap<>(Map.of("userId", "user123")); + var result = runner.run(input); + + // Verify all operations were executed + var fetchCategories = result.getOperation("fetch-categories"); + assertNotNull(fetchCategories); + assertEquals("fetch-categories", fetchCategories.getName()); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/GenericTypesExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/GenericTypesExampleTest.java new file mode 100644 index 000000000..65a8b8420 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/GenericTypesExampleTest.java @@ -0,0 +1,68 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.general; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class GenericTypesExampleTest { + + @Test + void testGenericTypesExample() { + var handler = new GenericTypesExample(); + var runner = LocalDurableTestRunner.create(GenericTypesExample.Input.class, handler); + + var input = new GenericTypesExample.Input("user123"); + var result = runner.run(input); + + assertNotNull(result); + GenericTypesExample.Output output = result.getResult(GenericTypesExample.Output.class); + assertNotNull(output); + + // Verify items list + assertNotNull(output.items); + assertEquals(4, output.items.size()); + assertTrue(output.items.contains("item1")); + assertTrue(output.items.contains("item4")); + + // Verify counts map + assertNotNull(output.counts); + assertEquals(3, output.counts.size()); + assertEquals(2, output.counts.get("electronics")); + assertEquals(1, output.counts.get("books")); + assertEquals(1, output.counts.get("clothing")); + + // Verify categories nested map + assertNotNull(output.categories); + assertEquals(3, output.categories.size()); + assertEquals(2, output.categories.get("electronics").size()); + assertTrue(output.categories.get("electronics").contains("laptop")); + assertTrue(output.categories.get("electronics").contains("phone")); + assertEquals(1, output.categories.get("books").size()); + assertTrue(output.categories.get("books").contains("fiction")); + } + + @Test + void testOperationTracking() { + var handler = new GenericTypesExample(); + var runner = LocalDurableTestRunner.create(GenericTypesExample.Input.class, handler); + + var input = new GenericTypesExample.Input("user456"); + var result = runner.run(input); + + // Verify all operations were executed + var fetchItems = result.getOperation("fetch-items"); + assertNotNull(fetchItems); + assertEquals("fetch-items", fetchItems.getName()); + + var countByCategory = result.getOperation("count-by-category"); + assertNotNull(countByCategory); + assertEquals("count-by-category", countByCategory.getName()); + + var fetchCategories = result.getOperation("fetch-categories"); + assertNotNull(fetchCategories); + assertEquals("fetch-categories", fetchCategories.getName()); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/LoggingExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/LoggingExampleTest.java new file mode 100644 index 000000000..75b5be39f --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/LoggingExampleTest.java @@ -0,0 +1,24 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.general; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class LoggingExampleTest { + + @Test + void testLoggingExample() { + var handler = new LoggingExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var result = runner.run(new GreetingRequest("Alice")); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("HELLO, ALICE!", result.getResult(String.class)); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/OtelExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/OtelExampleTest.java new file mode 100644 index 000000000..6f43d48df --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/OtelExampleTest.java @@ -0,0 +1,28 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.general; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class OtelExampleTest { + + @Test + void testOtelExample_executesSuccessfully() { + var handler = new OtelExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var result = runner.run(new GreetingRequest("World")); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("HELLO, WORLD!", result.getResult(String.class)); + + assertNotNull(result.getOperation("create-greeting")); + assertNotNull(result.getOperation("transform")); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/PluginExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/PluginExampleTest.java new file mode 100644 index 000000000..8ece49c10 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/general/PluginExampleTest.java @@ -0,0 +1,42 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.general; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class PluginExampleTest { + + @Test + void testPluginExample_executesSuccessfully() { + var handler = new PluginExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var result = runner.run(new GreetingRequest("World")); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("HELLO, WORLD!", result.getResult(String.class)); + + // Verify operations were tracked + assertNotNull(result.getOperation("create-greeting")); + assertNotNull(result.getOperation("transform")); + } + + @Test + void testPluginExample_pluginHooksFire() { + // This test verifies that the plugin hooks fire without error. + // Check stdout/CloudWatch for [PLUGIN] log lines when deployed. + var handler = new PluginExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var result = runner.runUntilComplete(new GreetingRequest("Test")); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("HELLO, TEST!", result.getResult(String.class)); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/invoke/InvokeExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/invoke/InvokeExampleTest.java new file mode 100644 index 000000000..922690962 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/invoke/InvokeExampleTest.java @@ -0,0 +1,95 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.invoke; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.services.lambda.model.ErrorObject; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class InvokeExampleTest { + + @Test + void testSimpleInvokeExample_completeSequentially() { + var handler = new SimpleInvokeExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + // First run + var input = new GreetingRequest("world"); + var output1 = runner.run(input); + + assertEquals(ExecutionStatus.PENDING, output1.getStatus()); + + // Second run + runner.completeChainedInvoke("call-greeting1", "\"hello\""); + var output2 = runner.run(input); + assertEquals(ExecutionStatus.PENDING, output2.getStatus()); + + // Third run + runner.completeChainedInvoke("call-greeting2", "\"world\""); + var output3 = runner.run(input); + assertEquals(ExecutionStatus.SUCCEEDED, output3.getStatus()); + assertEquals("helloworld", output3.getResult(String.class)); + } + + @Test + void testSimpleInvokeExample_completeConcurrently() { + var handler = new SimpleInvokeExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + // First run + var input = new GreetingRequest("world"); + var output1 = runner.run(input); + + assertEquals(ExecutionStatus.PENDING, output1.getStatus()); + + // Second run + runner.completeChainedInvoke("call-greeting1", "\"hello\""); + runner.completeChainedInvoke("call-greeting2", "\"world\""); + var output2 = runner.run(input); + assertEquals(ExecutionStatus.SUCCEEDED, output2.getStatus()); + assertEquals("helloworld", output2.getResult(String.class)); + } + + @Test + void testSimpleInvokeExample_failFirst() { + var handler = new SimpleInvokeExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + // First run + var input = new GreetingRequest("world"); + var output1 = runner.run(input); + + assertEquals(ExecutionStatus.PENDING, output1.getStatus()); + + // Second run, fail the async invoke + runner.failChainedInvoke("call-greeting1", ErrorObject.builder().build()); + var output2 = runner.run(input); + assertEquals(ExecutionStatus.PENDING, output2.getStatus()); + + // Third run + runner.completeChainedInvoke("call-greeting2", "\"world\""); + var output3 = runner.run(input); + assertEquals(ExecutionStatus.FAILED, output3.getStatus()); + } + + @Test + void testSimpleInvokeExample_failSecond() { + var handler = new SimpleInvokeExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + // First run + var input = new GreetingRequest("world"); + var output1 = runner.run(input); + + assertEquals(ExecutionStatus.PENDING, output1.getStatus()); + + // Second run, fail the async invoke + runner.failChainedInvoke("call-greeting2", ErrorObject.builder().build()); + var output2 = runner.run(input); + assertEquals(ExecutionStatus.FAILED, output2.getStatus()); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/invoke/RetryInvokeExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/invoke/RetryInvokeExampleTest.java new file mode 100644 index 000000000..080293434 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/invoke/RetryInvokeExampleTest.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.examples.operation.invoke; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.services.lambda.model.ErrorObject; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class RetryInvokeExampleTest { + + @Test + void succeedsOnFirstAttempt() { + var handler = new RetryInvokeExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + var input = new GreetingRequest("world"); + + // First run — starts the invoke, suspends waiting for result + var result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Complete the first invoke attempt + runner.completeChainedInvoke("call-greeting-1", "\"hello world\""); + result = runner.run(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("hello world", result.getResult(String.class)); + } + + @Test + void retriesAfterFirstAttemptFails() { + var handler = new RetryInvokeExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + var input = new GreetingRequest("world"); + + // First run — starts invoke attempt 1 + var result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Fail the first invoke attempt + runner.failChainedInvoke( + "call-greeting-1", + ErrorObject.builder() + .errorType("TransientError") + .errorMessage("Service unavailable") + .build()); + + // Second run — processes the failure, does backoff wait, starts invoke attempt 2 + result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Advance past the backoff wait + runner.advanceTime(); + result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Complete the second invoke attempt + runner.completeChainedInvoke("call-greeting-2", "\"hello on retry\""); + result = runner.run(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("hello on retry", result.getResult(String.class)); + } + + @Test + void failsAfterAllRetriesExhausted() { + var handler = new RetryInvokeExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + var input = new GreetingRequest("world"); + + // First run — starts invoke attempt 1 + var result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Fail attempt 1 + runner.failChainedInvoke( + "call-greeting-1", + ErrorObject.builder() + .errorType("TransientError") + .errorMessage("fail 1") + .build()); + result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Advance past backoff wait 1 + runner.advanceTime(); + result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Fail attempt 2 + runner.failChainedInvoke( + "call-greeting-2", + ErrorObject.builder() + .errorType("TransientError") + .errorMessage("fail 2") + .build()); + result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Advance past backoff wait 2 + runner.advanceTime(); + result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Fail attempt 3 — this is the last attempt, retryStrategy returns fail() + runner.failChainedInvoke( + "call-greeting-3", + ErrorObject.builder() + .errorType("TransientError") + .errorMessage("fail 3") + .build()); + result = runner.run(input); + + assertEquals(ExecutionStatus.FAILED, result.getStatus()); + } + + @Test + void suspendsOnFirstRun() { + var handler = new RetryInvokeExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + var input = new GreetingRequest("test"); + + var result = runner.run(input); + + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/map/ComplexFlatMapExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/map/ComplexFlatMapExampleTest.java new file mode 100644 index 000000000..e86f472ae --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/map/ComplexFlatMapExampleTest.java @@ -0,0 +1,54 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.map; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class ComplexFlatMapExampleTest { + + @Test + void testComplexMapExample() { + var handler = new ComplexFlatMapExample(); + var runner = LocalDurableTestRunner.create(Integer.class, handler); + + var result = runner.runUntilComplete(50); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + var output = result.getResult(String.class); + + // Part 1: all 3 orders processed with step + wait + step + assertTrue(output.contains("done:validated:order-1")); + assertTrue(output.contains("done:validated:order-2")); + assertTrue(output.contains("done:validated:order-50")); + + // Part 2: early termination after 2 healthy servers + assertTrue(output.contains("reason=MIN_SUCCESSFUL_REACHED")); + assertTrue(output.contains("healthy")); + } + + @Test + void testReplay() { + var handler = new ComplexFlatMapExample(); + var runner = LocalDurableTestRunner.create(Integer.class, handler); + + var result1 = runner.runUntilComplete(50); + assertEquals(ExecutionStatus.SUCCEEDED, result1.getStatus()); + + // Replay — should use cached results. + // Structural assertion because the first map has wait() inside branches with unlimited + // concurrency, which can cause non-deterministic thread scheduling across invocations. + var result2 = runner.runUntilComplete(50); + assertEquals(ExecutionStatus.SUCCEEDED, result2.getStatus()); + var output = result2.getResult(String.class); + assertTrue(output.contains("done:validated:order-1")); + assertTrue(output.contains("done:validated:order-2")); + assertTrue(output.contains("done:validated:order-50")); + assertTrue(output.contains("reason=MIN_SUCCESSFUL_REACHED")); + assertTrue(output.contains("healthy")); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/map/ComplexMapExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/map/ComplexMapExampleTest.java new file mode 100644 index 000000000..50eb6ed3c --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/map/ComplexMapExampleTest.java @@ -0,0 +1,53 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.map; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class ComplexMapExampleTest { + + @Test + void testComplexMapExample() { + var handler = new ComplexMapExample(); + var runner = LocalDurableTestRunner.create(Integer.class, handler); + + var result = runner.runUntilComplete(50); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + var output = result.getResult(String.class); + + // Part 1: all 3 orders processed with step + wait + step + assertTrue(output.contains("done:validated:order-1")); + assertTrue(output.contains("done:validated:order-2")); + assertTrue(output.contains("done:validated:order-50")); + + // Part 2: early termination after 2 healthy servers + assertTrue(output.contains("reason=MIN_SUCCESSFUL_REACHED")); + assertTrue(output.contains("healthy")); + } + + @Test + void testReplay() { + var handler = new ComplexMapExample(); + var runner = LocalDurableTestRunner.create(Integer.class, handler); + + var result1 = runner.runUntilComplete(50); + assertEquals(ExecutionStatus.SUCCEEDED, result1.getStatus()); + + // Replay — should use cached results. + // Structural assertion because the first map has wait() inside branches with unlimited + // concurrency, which can cause non-deterministic thread scheduling across invocations. + var result2 = runner.runUntilComplete(50); + assertEquals(ExecutionStatus.SUCCEEDED, result2.getStatus()); + var output = result2.getResult(String.class); + assertTrue(output.contains("done:validated:order-1")); + assertTrue(output.contains("done:validated:order-2")); + assertTrue(output.contains("done:validated:order-50")); + assertTrue(output.contains("reason=MIN_SUCCESSFUL_REACHED")); + assertTrue(output.contains("healthy")); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/map/CustomShouldCompleteMapExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/map/CustomShouldCompleteMapExampleTest.java new file mode 100644 index 000000000..c22f33fb0 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/map/CustomShouldCompleteMapExampleTest.java @@ -0,0 +1,61 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.map; + +import static org.junit.jupiter.api.Assertions.*; + +import java.util.List; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class CustomShouldCompleteMapExampleTest { + + @Test + void testCustomShouldCompleteSucceedsAfterRequiredResponses() { + var handler = new CustomShouldCompleteMapExample(); + var runner = LocalDurableTestRunner.create(CustomShouldCompleteMapExample.Input.class, handler); + + var input = new CustomShouldCompleteMapExample.Input(List.of("primary", "secondary", "bad-cache"), 2, 2); + var result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + var output = result.getResult(CustomShouldCompleteMapExample.Output.class); + assertEquals("CUSTOM_COMPLETION_SUCCEEDED", output.completionStatus()); + assertTrue(output.completionSucceeded()); + assertEquals(List.of("response:primary", "response:secondary"), output.responses()); + assertEquals(0, output.failed()); + assertEquals(1, output.skipped()); + } + + @Test + void testCustomShouldCompleteCanCompleteAsFailed() { + var handler = new CustomShouldCompleteMapExample(); + var runner = LocalDurableTestRunner.create(CustomShouldCompleteMapExample.Input.class, handler); + + var input = new CustomShouldCompleteMapExample.Input(List.of("bad-primary", "bad-secondary", "primary"), 2, 2); + var result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + var output = result.getResult(CustomShouldCompleteMapExample.Output.class); + assertEquals("CUSTOM_COMPLETION_FAILED", output.completionStatus()); + assertFalse(output.completionSucceeded()); + assertTrue(output.responses().isEmpty()); + assertEquals(2, output.failed()); + assertEquals(1, output.skipped()); + } + + @Test + void testReplay() { + var handler = new CustomShouldCompleteMapExample(); + var runner = LocalDurableTestRunner.create(CustomShouldCompleteMapExample.Input.class, handler); + + var input = new CustomShouldCompleteMapExample.Input(List.of("primary", "secondary", "bad-cache"), 2, 2); + var result1 = runner.runUntilComplete(input); + var result2 = runner.runUntilComplete(input); + + assertEquals( + result1.getResult(CustomShouldCompleteMapExample.Output.class), + result2.getResult(CustomShouldCompleteMapExample.Output.class)); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/map/DeserializationFailedMapExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/map/DeserializationFailedMapExampleTest.java new file mode 100644 index 000000000..f959d4285 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/map/DeserializationFailedMapExampleTest.java @@ -0,0 +1,26 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.map; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class DeserializationFailedMapExampleTest { + + @Test + void testDeserializationFailedExample() { + var handler = new DeserializationFailedMapExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var result = runner.runUntilComplete(new GreetingRequest("Alice")); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals( + "Map iteration failed with error of type java.lang.RuntimeException. Message: Failure from Alice!", + result.getResult(String.class)); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/map/SimpleMapExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/map/SimpleMapExampleTest.java new file mode 100644 index 000000000..02e62e3db --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/map/SimpleMapExampleTest.java @@ -0,0 +1,49 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.map; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class SimpleMapExampleTest { + + @Test + void testSimpleMapExample() { + var handler = new SimpleMapExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var result = runner.runUntilComplete(new GreetingRequest("Alice")); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("Hello, Alice! | Hello, ALICE! | Hello, alice!", result.getResult(String.class)); + } + + @Test + void testWithDefaultName() { + var handler = new SimpleMapExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var result = runner.runUntilComplete(new GreetingRequest()); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("Hello, World! | Hello, WORLD! | Hello, world!", result.getResult(String.class)); + } + + @Test + void testReplay() { + var handler = new SimpleMapExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var input = new GreetingRequest("Bob"); + var result1 = runner.runUntilComplete(input); + assertEquals("Hello, Bob! | Hello, BOB! | Hello, bob!", result1.getResult(String.class)); + + // Replay — should use cached results + var result2 = runner.runUntilComplete(input); + assertEquals(result1.getResult(String.class), result2.getResult(String.class)); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExampleTestSupport.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExampleTestSupport.java new file mode 100644 index 000000000..c168150c6 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExampleTestSupport.java @@ -0,0 +1,30 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.otel; + +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.exporter.logging.LoggingSpanExporter; +import io.opentelemetry.sdk.OpenTelemetrySdk; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; + +final class OtelXRayExampleTestSupport { + + private static final String SPI_INSTALLED_PROPERTY = + "software.amazon.lambda.durable.otel.autoConfigurationCustomizerProviderInstalled"; + + private OtelXRayExampleTestSupport() {} + + static void installGlobalOpenTelemetry() { + System.setProperty(SPI_INSTALLED_PROPERTY, Boolean.TRUE.toString()); + var tracerProvider = SdkTracerProvider.builder() + .addSpanProcessor(SimpleSpanProcessor.create(LoggingSpanExporter.create())) + .build(); + OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).buildAndRegisterGlobal(); + } + + static void resetGlobalOpenTelemetry() { + GlobalOpenTelemetry.resetForTest(); + System.clearProperty(SPI_INSTALLED_PROPERTY); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExamplesTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExamplesTest.java new file mode 100644 index 000000000..513ae6b58 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExamplesTest.java @@ -0,0 +1,51 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.otel; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +/** + * Local tests for OTel examples using LocalDurableTestRunner. + * + *

    These verify that the OTel plugin doesn't break execution for map, parallel, and nested context scenarios. + */ +class OtelXRayExamplesTest { + + @Test + void mapExample_executesSuccessfully() { + var handler = new OtelXRayExamples.MapExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var result = runner.runUntilComplete(new GreetingRequest("test")); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("Mapped 3 items", result.getResult(String.class)); + } + + @Test + void parallelExample_executesSuccessfully() { + var handler = new OtelXRayExamples.ParallelExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var result = runner.runUntilComplete(new GreetingRequest("test")); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertTrue(result.getResult(String.class).contains("Parallel completed: 2 branches")); + } + + @Test + void nestedContextExample_executesSuccessfully() { + var handler = new OtelXRayExamples.NestedContextExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var result = runner.runUntilComplete(new GreetingRequest("World")); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("HELLO, WORLD!", result.getResult(String.class)); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionStepExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionStepExampleTest.java new file mode 100644 index 000000000..f2b4cfb03 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionStepExampleTest.java @@ -0,0 +1,48 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.otel; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class OtelXRayExecutionStepExampleTest { + + @BeforeEach + void setUp() { + OtelXRayExampleTestSupport.installGlobalOpenTelemetry(); + } + + @AfterEach + void tearDown() { + OtelXRayExampleTestSupport.resetGlobalOpenTelemetry(); + } + + @Test + void testSimpleSteps_succeeds() { + var handler = new OtelXRayExecutionStepExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var result = runner.runUntilComplete(new GreetingRequest("Alice")); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("HELLO, ALICE!", result.getResult(String.class)); + } + + @Test + void testReplay_returnsSameResult() { + var handler = new OtelXRayExecutionStepExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var input = new GreetingRequest("Bob"); + var result1 = runner.runUntilComplete(input); + var result2 = runner.runUntilComplete(input); + + assertEquals(result1.getResult(String.class), result2.getResult(String.class)); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionWaitExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionWaitExampleTest.java new file mode 100644 index 000000000..fd7ed672d --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionWaitExampleTest.java @@ -0,0 +1,48 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.otel; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class OtelXRayExecutionWaitExampleTest { + + @BeforeEach + void setUp() { + OtelXRayExampleTestSupport.installGlobalOpenTelemetry(); + } + + @AfterEach + void tearDown() { + OtelXRayExampleTestSupport.resetGlobalOpenTelemetry(); + } + + @Test + void testFirstInvocation_suspendsOnWait() { + var handler = new OtelXRayExecutionWaitExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var result = runner.run(new GreetingRequest("Alice")); + + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + } + + @Test + void testFullExecution_completesAfterWait() { + var handler = new OtelXRayExecutionWaitExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var result = runner.runUntilComplete(new GreetingRequest("Alice")); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertTrue( + result.getResult(String.class).contains("Resumed and completed"), + "Expected result to contain 'Resumed and completed', got: " + result.getResult(String.class)); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayStepExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayStepExampleTest.java new file mode 100644 index 000000000..322968d87 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayStepExampleTest.java @@ -0,0 +1,59 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.otel; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class OtelXRayStepExampleTest { + + @BeforeEach + void setUp() { + OtelXRayExampleTestSupport.installGlobalOpenTelemetry(); + } + + @AfterEach + void tearDown() { + OtelXRayExampleTestSupport.resetGlobalOpenTelemetry(); + } + + @Test + void testSimpleSteps_succeeds() { + var handler = new OtelXRayStepExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var result = runner.run(new GreetingRequest("Alice")); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("HELLO, ALICE!", result.getResult(String.class)); + } + + @Test + void testReplay_returnsSameResult() { + var handler = new OtelXRayStepExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var input = new GreetingRequest("Bob"); + var result1 = runner.run(input); + var result2 = runner.run(input); + + assertEquals(result1.getResult(String.class), result2.getResult(String.class)); + } + + @Test + void testDefaultName() { + var handler = new OtelXRayStepExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var result = runner.run(new GreetingRequest()); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("HELLO, WORLD!", result.getResult(String.class)); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayWaitExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayWaitExampleTest.java new file mode 100644 index 000000000..5a98fd778 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayWaitExampleTest.java @@ -0,0 +1,61 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.otel; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class OtelXRayWaitExampleTest { + + @BeforeEach + void setUp() { + OtelXRayExampleTestSupport.installGlobalOpenTelemetry(); + } + + @AfterEach + void tearDown() { + OtelXRayExampleTestSupport.resetGlobalOpenTelemetry(); + } + + @Test + void testFirstInvocation_suspendsOnWait() { + var handler = new OtelXRayWaitExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var result = runner.run(new GreetingRequest("Alice")); + + // First invocation hits the wait and suspends + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + } + + @Test + void testFullExecution_completesAfterWait() { + var handler = new OtelXRayWaitExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var result = runner.runUntilComplete(new GreetingRequest("Alice")); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertTrue( + result.getResult(String.class).contains("Resumed and completed"), + "Expected result to contain 'Resumed and completed', got: " + result.getResult(String.class)); + } + + @Test + void testReplay_returnsSameResult() { + var handler = new OtelXRayWaitExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var input = new GreetingRequest("Bob"); + var result1 = runner.runUntilComplete(input); + var result2 = runner.runUntilComplete(input); + + assertEquals(result1.getResult(String.class), result2.getResult(String.class)); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/parallel/DeserializationFailedParallelExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/parallel/DeserializationFailedParallelExampleTest.java new file mode 100644 index 000000000..13b1ac876 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/parallel/DeserializationFailedParallelExampleTest.java @@ -0,0 +1,29 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.parallel; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import java.util.List; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class DeserializationFailedParallelExampleTest { + + @Test + void testDeserializationFailedParallelExample() { + var handler = new DeserializationFailedParallelExample(); + var runner = LocalDurableTestRunner.create(DeserializationFailedParallelExample.Input.class, handler); + + var input = new DeserializationFailedParallelExample.Input(List.of("apple", "banana", "cherry")); + var result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + var output = result.getResult(String.class); + assertEquals( + "Parallel branch failed with error of type java.lang.RuntimeException. Message: Intentional failure for transform", + output); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelExampleTest.java new file mode 100644 index 000000000..02b132ba7 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelExampleTest.java @@ -0,0 +1,60 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.parallel; + +import static org.junit.jupiter.api.Assertions.*; + +import java.util.List; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class ParallelExampleTest { + + @Test + void testParallelExampleRunsSuccessfully() { + var handler = new ParallelExample(); + var runner = LocalDurableTestRunner.create(ParallelExample.Input.class, handler); + + var input = new ParallelExample.Input(List.of("apple", "banana", "cherry")); + var result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + var output = result.getResult(ParallelExample.Output.class); + assertEquals(3, output.totalProcessed()); + assertTrue(output.results().contains("APPLE")); + assertTrue(output.results().contains("BANANA")); + assertTrue(output.results().contains("CHERRY")); + } + + @Test + void testParallelExampleWithSingleItem() { + var handler = new ParallelExample(); + var runner = LocalDurableTestRunner.create(ParallelExample.Input.class, handler); + + var input = new ParallelExample.Input(List.of("hello")); + var result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + var output = result.getResult(ParallelExample.Output.class); + assertEquals(1, output.totalProcessed()); + assertEquals(List.of("HELLO"), output.results()); + } + + @Test + void testParallelExampleWithEmptyInput() { + var handler = new ParallelExample(); + var runner = LocalDurableTestRunner.create(ParallelExample.Input.class, handler); + + var input = new ParallelExample.Input(List.of()); + var result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + var output = result.getResult(ParallelExample.Output.class); + assertEquals(0, output.totalProcessed()); + assertTrue(output.results().isEmpty()); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelFailureToleranceExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelFailureToleranceExampleTest.java new file mode 100644 index 000000000..ee45f4222 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelFailureToleranceExampleTest.java @@ -0,0 +1,59 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.parallel; + +import static org.junit.jupiter.api.Assertions.*; + +import java.util.List; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class ParallelFailureToleranceExampleTest { + + @Test + void succeedsWhenFailuresAreWithinTolerance() { + var handler = new ParallelFailureToleranceExample(); + var runner = LocalDurableTestRunner.create(ParallelFailureToleranceExample.Input.class, handler); + + // 2 good services, 1 bad — toleratedFailureCount=1 so the parallel op still succeeds + var input = new ParallelFailureToleranceExample.Input(List.of("svc-a", "bad-svc-b", "svc-c"), 1, null); + var result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + var output = result.getResult(ParallelFailureToleranceExample.Output.class); + assertEquals(2, output.succeeded()); + assertEquals(1, output.failed()); + } + + @Test + void succeedsWhenAllBranchesSucceed() { + var handler = new ParallelFailureToleranceExample(); + var runner = LocalDurableTestRunner.create(ParallelFailureToleranceExample.Input.class, handler); + + var input = new ParallelFailureToleranceExample.Input(List.of("svc-a", "svc-b", "svc-c"), 2, null); + var result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + var output = result.getResult(ParallelFailureToleranceExample.Output.class); + assertEquals(3, output.succeeded()); + } + + @Test + void failsWhenFailuresExceedTolerance() { + var handler = new ParallelFailureToleranceExample(); + var runner = LocalDurableTestRunner.create(ParallelFailureToleranceExample.Input.class, handler); + + // 2 bad services, toleratedFailureCount=1 — second failure exceeds tolerance + var input = new ParallelFailureToleranceExample.Input(List.of("svc-a", "bad-svc-b", "bad-svc-c"), 1, 2); + var result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + var output = result.getResult(ParallelFailureToleranceExample.Output.class); + assertEquals(2, output.failed()); + assertEquals(1, output.succeeded()); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelWithWaitExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelWithWaitExampleTest.java new file mode 100644 index 000000000..49523b9e6 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelWithWaitExampleTest.java @@ -0,0 +1,34 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.parallel; + +import static org.junit.jupiter.api.Assertions.*; + +import java.util.List; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class ParallelWithWaitExampleTest { + @Test + void completesAfterManuallyAdvancingWaits() { + var handler = new ParallelWithWaitExample(); + var runner = LocalDurableTestRunner.create(ParallelWithWaitExample.Input.class, handler); + + var input = new ParallelWithWaitExample.Input("user-456", "world"); + + // First run suspends on wait branches + var first = runner.run(input); + assertEquals(ExecutionStatus.PENDING, first.getStatus()); + + // Advance waits and re-run to completion + runner.advanceTime(); + var result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + var output = result.getResult(ParallelWithWaitExample.Output.class); + assertEquals(List.of("email:world", "sms:world", "push:world"), output.deliveries()); + assertEquals(3, output.success()); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/step/DeserializationFailureExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/step/DeserializationFailureExampleTest.java new file mode 100644 index 000000000..1e09296ce --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/step/DeserializationFailureExampleTest.java @@ -0,0 +1,31 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.step; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class DeserializationFailureExampleTest { + + @Test + void testDeserializationExample() { + var handler = new DeserializationFailureExample(); + + // Create test runner from handler (automatically extracts config) + var runner = LocalDurableTestRunner.create(String.class, handler); + + // Run with input + var result = runner.runUntilComplete("test-input"); + + // Verify result + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + // assert StepFailedException is thrown when SerDes fails to deserialize the exception + assertEquals( + "StepFailedException:Step failed with error of type java.lang.RuntimeException. Message: this is a test", + result.getResult(String.class)); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/step/ManyAsyncStepsExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/step/ManyAsyncStepsExampleTest.java new file mode 100644 index 000000000..203f885ac --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/step/ManyAsyncStepsExampleTest.java @@ -0,0 +1,60 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.step; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.examples.types.ManyAsyncStepsInput; +import software.amazon.lambda.durable.examples.types.ManyAsyncStepsOutput; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class ManyAsyncStepsExampleTest { + + @Test + void testManyAsyncSteps() { + var handler = new ManyAsyncStepsExample(); + var runner = LocalDurableTestRunner.create(ManyAsyncStepsInput.class, handler); + + var input = new ManyAsyncStepsInput(2, 500); + var result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + var output = result.getResult(ManyAsyncStepsOutput.class); + assertNotNull(output); + + // Sum of 0..499 * 2 = 499 * 500 / 2 * 2 = 249500 + assertEquals(249500, result.getResult(ManyAsyncStepsOutput.class).result()); + } + + @Test + void testManyAsyncStepsWithDefaultMultiplier() { + var handler = new ManyAsyncStepsExample(); + var runner = LocalDurableTestRunner.create(ManyAsyncStepsInput.class, handler); + + var input = new ManyAsyncStepsInput(1, 500); + var result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + // Sum of 0..499 = 499 * 500 / 2 = 124750 + assertEquals(124750, result.getResult(ManyAsyncStepsOutput.class).result()); + } + + @Test + void testOperationsAreTracked() { + var handler = new ManyAsyncStepsExample(); + var runner = LocalDurableTestRunner.create(ManyAsyncStepsInput.class, handler); + + var result = runner.runUntilComplete(new ManyAsyncStepsInput(1, 500)); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + // Verify some operations are tracked + assertNotNull(result.getOperation("compute-0")); + assertNotNull(result.getOperation("compute-499")); + assertNotNull(result.getOperation("compute-250")); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/step/RetryExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/step/RetryExampleTest.java new file mode 100644 index 000000000..58fba3f84 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/step/RetryExampleTest.java @@ -0,0 +1,67 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.step; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class RetryExampleTest { + + @Test + void testRetryExampleWithTimeBasedFailure() { + var handler = new RetryExample(); + + // Test the retry example with time-based failure simulation + var runner = LocalDurableTestRunner.create(Object.class, handler); + + var result = runner.run("test-input"); + + // The test will likely result in PENDING due to the time-based failure + // This demonstrates the retry mechanism in action + System.out.println("Test result status: " + result.getStatus()); + + // In a real scenario with actual time delays, this would eventually succeed + // The LocalDurableTestRunner simulates the retry behavior + assertNotNull(result); + } + + @Test + void testRetryExampleDemonstration() { + // This test demonstrates the retry behavior without strict assertions + // It's useful for observing the retry mechanism in action + var handler = new RetryExample(); + + // Test the retry example with time-based failure simulation + var runner = LocalDurableTestRunner.create(Object.class, handler); + + var result = runner.run("demo-input"); + + System.out.println("Demo execution status: " + result.getStatus()); + + // This test always passes - it's just for demonstration + assertNotNull(result); + assertTrue(result.getStatus().toString().matches("PENDING|FAILED|SUCCEEDED")); + } + + @Test + void testRetryExampleShowsRetryBehavior() { + // Test that shows the different retry behaviors + var handler = new RetryExample(); + + // Test the retry example with time-based failure simulation + var runner = LocalDurableTestRunner.create(Object.class, handler); + + var result = runner.run("retry-behavior-test"); + + System.out.println("Retry behavior test status: " + result.getStatus()); + + // The example demonstrates: + // 1. No-retry step that fails immediately + // 2. Retry step that uses default exponential backoff + // 3. Time-based failure that would eventually succeed with retries + + assertNotNull(result); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/step/SimpleStepExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/step/SimpleStepExampleTest.java new file mode 100644 index 000000000..380f4e834 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/step/SimpleStepExampleTest.java @@ -0,0 +1,77 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.step; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class SimpleStepExampleTest { + + @Test + void testSimpleStepExample() { + // Create handler + var handler = new SimpleStepExample(); + + // Create test runner + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + // Run with input + var input = new GreetingRequest("Alice"); + var result = runner.run(input); + + // Verify result + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("HELLO, ALICE!", result.getResult(String.class)); + } + + @Test + void testWithLargePayload() { + // Create handler + var handler = new SimpleStepExample(); + + // Create test runner + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + // 6MB large input + var largeInput = "A".repeat(1024).repeat(1024).repeat(6); + + // Run with input + var input = new GreetingRequest(largeInput); + var result = runner.run(input); + + // Verify result + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("HELLO, " + largeInput + "!", result.getResult(String.class)); + } + + @Test + void testWithDefaultName() { + var handler = new SimpleStepExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var input = new GreetingRequest(); + var result = runner.run(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("HELLO, WORLD!", result.getResult(String.class)); + } + + @Test + void testReplay() { + var handler = new SimpleStepExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + // First execution + var input = new GreetingRequest("Bob"); + var result1 = runner.run(input); + assertEquals("HELLO, BOB!", result1.getResult(String.class)); + + // Second execution (replay) - should use cached results + var result2 = runner.run(input); + assertEquals("HELLO, BOB!", result2.getResult(String.class)); + assertEquals(result1.getResult(String.class), result2.getResult(String.class)); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/vt/ManyAsyncStepsVirtualThreadPoolExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/vt/ManyAsyncStepsVirtualThreadPoolExampleTest.java new file mode 100644 index 000000000..2d66744e5 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/vt/ManyAsyncStepsVirtualThreadPoolExampleTest.java @@ -0,0 +1,64 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.vt; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledForJreRange; +import org.junit.jupiter.api.condition.JRE; +import software.amazon.lambda.durable.examples.types.ManyAsyncStepsInput; +import software.amazon.lambda.durable.examples.types.ManyAsyncStepsOutput; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +@EnabledForJreRange(min = JRE.JAVA_21) +class ManyAsyncStepsVirtualThreadPoolExampleTest { + + @Test + void testManyAsyncSteps() { + var handler = new ManyAsyncStepsVirtualThreadPoolExample(); + var runner = LocalDurableTestRunner.create(ManyAsyncStepsInput.class, handler); + + var input = new ManyAsyncStepsInput(2, 500); + var result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + var output = result.getResult(ManyAsyncStepsOutput.class); + assertNotNull(output); + + // Sum of 0..499 * 2 = 499 * 500 / 2 * 2 = 249500 + assertEquals(249500, result.getResult(ManyAsyncStepsOutput.class).result()); + } + + @Test + void testManyAsyncStepsWithDefaultMultiplier() { + var handler = new ManyAsyncStepsVirtualThreadPoolExample(); + var runner = LocalDurableTestRunner.create(ManyAsyncStepsInput.class, handler); + + var input = new ManyAsyncStepsInput(1, 500); + var result = runner.runUntilComplete(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + // Sum of 0..499 = 499 * 500 / 2 = 124750 + assertEquals(124750, result.getResult(ManyAsyncStepsOutput.class).result()); + } + + @Test + void testOperationsAreTracked() { + var handler = new ManyAsyncStepsVirtualThreadPoolExample(); + var runner = LocalDurableTestRunner.create(ManyAsyncStepsInput.class, handler); + + var result = runner.runUntilComplete(new ManyAsyncStepsInput(1, 500)); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + // Verify some operations are tracked + assertNotNull(result.getOperation("compute-0")); + assertNotNull(result.getOperation("compute-499")); + assertNotNull(result.getOperation("compute-250")); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/wait/ConcurrentWaitForConditionExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/wait/ConcurrentWaitForConditionExampleTest.java new file mode 100644 index 000000000..8c140a501 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/wait/ConcurrentWaitForConditionExampleTest.java @@ -0,0 +1,29 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.wait; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class ConcurrentWaitForConditionExampleTest { + + @Test + void testConcurrentWaitForConditionExample() { + var handler = new ConcurrentWaitForConditionExample(); + var runner = LocalDurableTestRunner.create(ConcurrentWaitForConditionExample.Input.class, handler); + + var result = runner.runUntilComplete(new ConcurrentWaitForConditionExample.Input(3, 100, 50)); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + var allOperationsOutput = result.getResult(String.class); + var operationResults = allOperationsOutput.split(" \\| "); + assertEquals(100, operationResults.length); + for (var operationResult : operationResults) { + assertEquals("3", operationResult); + } + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/wait/WaitAsyncExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/wait/WaitAsyncExampleTest.java new file mode 100644 index 000000000..f679c6a37 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/wait/WaitAsyncExampleTest.java @@ -0,0 +1,41 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.wait; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class WaitAsyncExampleTest { + + @Test + void testWaitAsyncExampleCompletesSuccessfully() { + var handler = new WaitAsyncExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var result = runner.runUntilComplete(new GreetingRequest("Alice")); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("Processed: Alice", result.getResult(String.class)); + } + + @Test + void testWaitAsyncExampleSuspendsOnFirstRun() { + var handler = new WaitAsyncExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + // First run suspends because the wait hasn't elapsed yet + var result = runner.run(new GreetingRequest("Bob")); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Advance time so the wait completes, then re-run to finish + runner.advanceTime(); + + var result2 = runner.runUntilComplete(new GreetingRequest("Bob")); + assertEquals(ExecutionStatus.SUCCEEDED, result2.getStatus()); + assertEquals("Processed: Bob", result2.getResult(String.class)); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/wait/WaitExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/wait/WaitExampleTest.java new file mode 100644 index 000000000..d1ac0ea5a --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/wait/WaitExampleTest.java @@ -0,0 +1,30 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.wait; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class WaitExampleTest { + + @Test + void testWaitExampleStartsAndWaits() { + var handler = new WaitExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + + var input = new GreetingRequest("Bob"); + + // First run - executes first step and hits first wait + var result = runner.run(input); + + // Should be PENDING because of wait operation + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Note: In real Lambda, the function would be re-invoked after the wait period + // The LocalDurableTestRunner demonstrates the wait behavior but doesn't simulate time + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/operation/wait/WaitForConditionExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/wait/WaitForConditionExampleTest.java new file mode 100644 index 000000000..30a531961 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/operation/wait/WaitForConditionExampleTest.java @@ -0,0 +1,23 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.operation.wait; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class WaitForConditionExampleTest { + + @Test + void testWaitForConditionExample() { + var handler = new WaitForConditionExample(); + var runner = LocalDurableTestRunner.create(Integer.class, handler); + + var result = runner.runUntilComplete(3); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals(3, result.getResult(Integer.class)); + } +} diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java index 89838b1b0..f4408c112 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StaticOperationsIntegrationTest.java @@ -194,7 +194,7 @@ void staticWaitForConditionMatchesLegacyCheckpointHistory() { (input, context) -> DurableWaitForConditionOperation.waitForCondition( "condition", Integer.class, - StaticOperationsIntegrationTest::nextConditionState, + StaticOperationsIntegrationTest::nextOperationConditionState, config.toOperationConfig())); var legacyResult = legacyRunner.runUntilComplete("input"); @@ -247,7 +247,7 @@ void conditionAndRetryExposeGeneratedMetadataThroughTls() { Integer.class, state -> { assertNotNull(StepContext.getCurrentContext()); - return WaitForConditionResult.stopPolling(state + 1); + return DurableWaitForConditionOperation.WaitForConditionResult.stopPolling(state + 1); }, WaitForConditionConfig.builder() .initialState(0) @@ -300,6 +300,14 @@ private static WaitForConditionResult nextConditionState(int state) { return next >= 2 ? WaitForConditionResult.stopPolling(next) : WaitForConditionResult.continuePolling(next); } + private static DurableWaitForConditionOperation.WaitForConditionResult nextOperationConditionState( + int state) { + var next = state + 1; + return next >= 2 + ? DurableWaitForConditionOperation.WaitForConditionResult.stopPolling(next) + : DurableWaitForConditionOperation.WaitForConditionResult.continuePolling(next); + } + private static String retryAttempt(int attempt, Supplier operation) { if (attempt == 1) { throw new IllegalStateException("retry"); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java index 28a92ae3b..7422a28df 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java @@ -158,7 +158,17 @@ public DurableFuture waitForConditionAsync( BiFunction> checkFunc, WaitForConditionConfig config) { return DurableWaitForConditionOperation.waitForConditionAsync( - this, name, resultType, checkFunc, config.toOperationConfig()); + this, + name, + resultType, + (state, stepContext) -> { + var result = checkFunc.apply(state, stepContext); + return result == null + ? null + : new DurableWaitForConditionOperation.WaitForConditionResult<>( + result.value(), result.isDone()); + }, + config.toOperationConfig()); } // =============== withRetry ================ diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperation.java index 0a04c6e93..116e8890f 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperation.java @@ -15,7 +15,6 @@ import software.amazon.lambda.durable.extension.ExtensionStepConfig; import software.amazon.lambda.durable.extension.ExtensionStepResult; import software.amazon.lambda.durable.model.OperationSubType; -import software.amazon.lambda.durable.model.WaitForConditionResult; import software.amazon.lambda.durable.retry.WaitForConditionWaitStrategy; import software.amazon.lambda.durable.retry.WaitStrategies; import software.amazon.lambda.durable.serde.SerDes; @@ -148,6 +147,25 @@ public CompletableFuture completionFuture() { } } + /** + * Result returned by a wait-for-condition check function. + * + * @param value the current state after evaluation + * @param isDone true to stop polling, false to continue polling + * @param the state type + */ + public record WaitForConditionResult(T value, boolean isDone) { + /** Returns a result that stops polling with the supplied final value. */ + public static WaitForConditionResult stopPolling(T value) { + return new WaitForConditionResult<>(value, true); + } + + /** Returns a result that continues polling with the supplied state. */ + public static WaitForConditionResult continuePolling(T value) { + return new WaitForConditionResult<>(value, false); + } + } + /** Configuration for durable wait-for-condition operations. */ public static final class WaitForConditionConfig { private final WaitForConditionWaitStrategy waitStrategy; diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForConditionOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForConditionOperationTest.java index a3cd25c83..8cea63c7f 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForConditionOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableWaitForConditionOperationTest.java @@ -3,6 +3,8 @@ package software.amazon.lambda.durable; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; @@ -19,8 +21,8 @@ import software.amazon.lambda.durable.extension.ExtensionStepFunction; import software.amazon.lambda.durable.extension.ExtensionStepResult; import software.amazon.lambda.durable.model.OperationSubType; -import software.amazon.lambda.durable.model.WaitForConditionResult; import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation; +import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation.WaitForConditionResult; class DurableWaitForConditionOperationTest { @AfterEach @@ -63,6 +65,17 @@ void conditionFunctionReceivesOnlyStateAndUsesStepContextFromTls() { } } + @Test + void resultFactoriesRepresentPollingDecision() { + var completed = WaitForConditionResult.stopPolling("done"); + var pending = WaitForConditionResult.continuePolling("next"); + + assertEquals("done", completed.value()); + assertTrue(completed.isDone()); + assertEquals("next", pending.value()); + assertFalse(pending.isDone()); + } + @SuppressWarnings({"rawtypes", "unchecked"}) private ArgumentCaptor> extensionFunction() { return (ArgumentCaptor) ArgumentCaptor.forClass(ExtensionStepFunction.class); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableOperationConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableOperationConfigTest.java index 6069e4e2b..0c8836440 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableOperationConfigTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableOperationConfigTest.java @@ -9,6 +9,7 @@ import static org.mockito.Mockito.mock; import static software.amazon.lambda.durable.model.ConcurrencyCompletionStatus.CUSTOM_COMPLETION_SUCCEEDED; +import java.lang.reflect.Method; import java.lang.reflect.Modifier; import java.time.Duration; import java.util.Arrays; @@ -28,6 +29,7 @@ import software.amazon.lambda.durable.config.WaitForCallbackConfig; import software.amazon.lambda.durable.config.WaitForConditionConfig; import software.amazon.lambda.durable.config.WithRetryConfig; +import software.amazon.lambda.durable.model.WaitForConditionResult; import software.amazon.lambda.durable.retry.RetryStrategy; import software.amazon.lambda.durable.retry.WaitForConditionWaitStrategy; import software.amazon.lambda.durable.serde.SerDes; @@ -53,6 +55,8 @@ void operationApisOwnTheirConfigTypes() throws Exception { DurableWaitForCallbackOperation.class, "WaitForCallbackConfig", WaitForCallbackConfig.class); assertOperationConfig( DurableWaitForConditionOperation.class, "WaitForConditionConfig", WaitForConditionConfig.class); + assertOperationOwnedType( + DurableWaitForConditionOperation.class, "WaitForConditionResult", WaitForConditionResult.class); assertOperationConfig(DurableWithRetryOperation.class, "WithRetryConfig", WithRetryConfig.class); } @@ -190,6 +194,14 @@ private static void assertProtectedStaticNestedType(Class owner, String neste assertTrue(Modifier.isStatic(nestedClass.getModifiers())); } + private static void assertOperationOwnedType(Class operationClass, String nestedName, Class legacyClass) + throws Exception { + assertPublicStaticNestedType(operationClass, nestedName); + assertFalse(Arrays.stream(operationClass.getMethods()) + .map(Method::toGenericString) + .anyMatch(signature -> signature.contains(legacyClass.getName()))); + } + private static void assertParallelFutureUsesCompatibilityBranchConfig() { var nestedConfig = DurableParallelOperation.ParallelBranchConfig.class; assertTrue(Arrays.stream(ParallelDurableFuture.class.getMethods()) diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperationImplementationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperationImplementationTest.java index 8ffd6e45a..6d3d8d8cb 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperationImplementationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperationImplementationTest.java @@ -25,7 +25,6 @@ import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.WaitForConditionConfig; import software.amazon.lambda.durable.context.BaseContextImpl; import software.amazon.lambda.durable.exception.StepFailedException; import software.amazon.lambda.durable.exception.WaitForConditionFailedException; @@ -35,7 +34,8 @@ import software.amazon.lambda.durable.extension.ExtensionStepFunction; import software.amazon.lambda.durable.extension.ExtensionStepResult; import software.amazon.lambda.durable.model.OperationSubType; -import software.amazon.lambda.durable.model.WaitForConditionResult; +import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation.WaitForConditionConfig; +import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation.WaitForConditionResult; import software.amazon.lambda.durable.serde.JacksonSerDes; class DurableWaitForConditionOperationImplementationTest { @@ -69,11 +69,7 @@ void executeMapsPollingResultsToStatefulStepOutcomes() { .thenReturn(future); var actual = DurableWaitForConditionOperation.waitForConditionAsync( - context, - "ready", - resultType, - (state, step) -> WaitForConditionResult.continuePolling("next"), - config.toOperationConfig()); + context, "ready", resultType, (state, step) -> WaitForConditionResult.continuePolling("next"), config); assertEquals(future.get(), actual.get()); var function = extensionFunction(); @@ -153,7 +149,7 @@ private DurableFuture createFuture(DurableFuture delegate) { "ready", resultType, (state, step) -> WaitForConditionResult.stopPolling(state), - WaitForConditionConfig.builder().build().toOperationConfig()); + WaitForConditionConfig.builder().build()); } @SuppressWarnings({"rawtypes", "unchecked"}) From 461f6467356a60b1b0c3debcf8d9150174659812 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Mon, 10 Aug 2026 16:00:56 -0700 Subject: [PATCH 35/40] fix: preserve extension operation compatibility --- docs/adr/006-custom-extension-operations.md | 32 ++- docs/advanced/extensions.md | 55 +++-- docs/design.md | 211 ++++++++++-------- .../operation/callback/CallbackExample.java | 26 +-- .../callback/RetryWaitForCallbackExample.java | 30 ++- .../WaitForCallbackFailedExample.java | 5 +- .../operation/child/ChildContextExample.java | 36 +-- .../child/ManyAsyncChildContextExample.java | 13 +- .../child/VirtualChildContextExample.java | 38 ++-- .../general/CustomConfigExample.java | 5 +- .../general/CustomPollingExample.java | 9 +- .../general/ErrorHandlingExample.java | 11 +- .../general/GenericInputOutputExample.java | 5 +- .../general/GenericTypesExample.java | 24 +- .../operation/general/LoggingExample.java | 7 +- .../operation/general/OtelExample.java | 7 +- .../operation/general/PluginExample.java | 7 +- .../operation/invoke/RetryInvokeExample.java | 12 +- .../operation/invoke/SimpleInvokeExample.java | 8 +- .../operation/map/ComplexFlatMapExample.java | 23 +- .../operation/map/ComplexMapExample.java | 21 +- .../map/CustomShouldCompleteMapExample.java | 12 +- .../map/DeserializationFailedMapExample.java | 13 +- .../operation/map/SimpleMapExample.java | 12 +- .../operation/otel/OtelXRayExamples.java | 28 +-- .../otel/OtelXRayExecutionStepExample.java | 8 +- .../otel/OtelXRayExecutionWaitExample.java | 8 +- .../operation/otel/OtelXRayStepExample.java | 7 +- .../operation/otel/OtelXRayWaitExample.java | 7 +- .../DeserializationFailedParallelExample.java | 11 +- .../operation/parallel/ParallelExample.java | 11 +- .../ParallelFailureToleranceExample.java | 9 +- .../parallel/ParallelWithWaitExample.java | 13 +- .../step/DeserializationFailureExample.java | 5 +- .../operation/step/ManyAsyncStepsExample.java | 10 +- .../examples/operation/step/RetryExample.java | 9 +- .../operation/step/RetryInProcessExample.java | 8 +- .../operation/step/SimpleStepExample.java | 9 +- ...anyAsyncStepsVirtualThreadPoolExample.java | 10 +- .../ConcurrentWaitForConditionExample.java | 13 +- .../operation/wait/WaitAsyncExample.java | 10 +- .../operation/wait/WaitAtLeastExample.java | 9 +- .../wait/WaitAtLeastInProcessExample.java | 9 +- .../examples/operation/wait/WaitExample.java | 16 +- .../wait/WaitForConditionExample.java | 5 +- .../ExtensionOperationIntegrationTest.java | 3 +- .../WaitForConditionIntegrationTest.java | 46 ++++ .../amazon/lambda/durable/DurableContext.java | 25 ++- .../amazon/lambda/durable/StepContext.java | 23 +- .../extension/ExtensionStepResult.java | 33 ++- .../operation/DurableContextOperation.java | 2 +- .../operation/DurableMapOperation.java | 2 +- .../operation/DurableParallelOperation.java | 2 +- .../DurableWaitForCallbackOperation.java | 5 +- .../DurableWaitForConditionOperation.java | 7 +- .../operation/DurableWithRetryOperation.java | 6 +- .../primitive/ChildContextPrimitive.java | 15 +- .../durable/primitive/StepPrimitive.java | 20 +- .../lambda/durable/CurrentContextTest.java | 21 +- .../durable/ExtensionStepResultTest.java | 20 ++ ...orCallbackOperationImplementationTest.java | 86 +++++++ ...rConditionOperationImplementationTest.java | 14 +- ...eWithRetryOperationImplementationTest.java | 3 + .../primitive/ChildContextPrimitiveTest.java | 69 +++++- .../StatefulExtensionStepPrimitiveTest.java | 53 +++++ 65 files changed, 880 insertions(+), 412 deletions(-) diff --git a/docs/adr/006-custom-extension-operations.md b/docs/adr/006-custom-extension-operations.md index 29e380841..b6aa1827b 100644 --- a/docs/adr/006-custom-extension-operations.md +++ b/docs/adr/006-custom-extension-operations.md @@ -146,6 +146,12 @@ retrieved from scoped thread-local contexts: Nested scopes restore the preceding value. Current context is available only on SDK-managed threads and is not propagated to application-created threads. +`DurableContext.getCurrentContext()` and `StepContext.getCurrentContext()` preserve their nullable-probe behavior: +they return `null` when no SDK context is active. Their `requireCurrentContext()` variants throw when the expected +context is absent and are used when an operation requires an active SDK scope. Calling either getter from the wrong +SDK context type still throws. `ExtensionContext.getCurrentContext()` always requires an active durable context and +throws outside one. + ### Reserve Sequential or Custom Local Identities `ExtensionContext` supports sequential and custom-local-ID reservations: @@ -205,10 +211,11 @@ DurableFuture waitAsync(String subType, Duration duration); ExtensionContextConfig config); ``` -`ExtensionStepConfig` owns its nested `StepSemantics` and `RetryStrategy` contracts. Retry strategies reuse -`ExtensionStepResult.retry(state, delay)` as their retry decision, so stateful continuations and failed attempts share -one retry representation. Built-in step operations adapt the customer-facing config and retry types at the operation -boundary, keeping the extension SPI independent of those packages. +`ExtensionStepConfig` owns its nested `StepSemantics` and `RetryStrategy` contracts. Exception retry strategies reuse +the fixed-delay `ExtensionStepResult.retry(state, delay)` outcome as their retry decision. Stateful continuations may +also use `ExtensionStepResult.retryAfterNormalization(state, delayStrategy)` when the delay depends on the state after +its configured SerDes round trip. Built-in step operations adapt the customer-facing config and retry types at the +operation boundary, keeping the extension SPI independent of those packages. The primitive selector determines the backend operation type. The string controls only the subtype recorded in checkpoints, replay validation, plugins, logs, and error metadata. @@ -234,9 +241,13 @@ A stateful extension STEP may return only: - `ExtensionStepResult.succeed(value)` - `ExtensionStepResult.retry(state, delay)` +- `ExtensionStepResult.retryAfterNormalization(state, delayStrategy)` -The SDK maps those outcomes onto the fixed STEP lifecycle. Thrown exceptions follow the normal STEP failure path. -Attempt metadata remains available through `StepContext`. +The fixed-delay retry checkpoints the supplied state and delay. The normalization-aware retry first serializes and +deserializes the state with the configured SerDes, evaluates the delay strategy against that normalized state, then +checkpoints the normalized state and selected delay. This keeps first execution and replay behavior consistent for +normalizing serializers. The SDK maps both outcomes onto the fixed STEP lifecycle. Thrown exceptions follow the +normal STEP failure path. Attempt metadata remains available through `StepContext`. ### Support Context Replay State @@ -252,6 +263,11 @@ Supported result policies are: On replay, the framework function receives the stored replay state through scoped TLS. This supports large map results and parallel branch reconstruction without exposing checkpoint APIs. +When child replay is requested with a `null` replay state, the SDK checkpoints an empty payload. On replay, an empty +payload is interpreted as a `null` replay state rather than passed to the SerDes. This preserves compatibility with +child-context checkpoints produced by earlier SDK versions, which stored large results as an empty payload with +`replayChildren=true`. + `ExtensionContextConfig` directly owns the child context serializer and virtual-context flag, plus extension-only behavior: @@ -318,7 +334,9 @@ The migration preserves: - existing type, subtype, name, and operation tree shape - checkpoint actions, payloads, statuses, and replay validation - map and parallel completion, skipped items, nesting, and large-result behavior -- wait-for-condition state, attempts, and delays +- child contexts, map iterations, parallel branches, wait-for-callback, and checkpointed with-retry contexts replay + children when their serialized result is at least 256 KiB +- wait-for-condition state, attempts, and delays, including delay evaluation against SerDes-normalized state - wait-for-callback exception translation - retry naming and virtual-context behavior - plugin operation and user-function event ordering diff --git a/docs/advanced/extensions.md b/docs/advanced/extensions.md index 0d55210c5..ca7113171 100644 --- a/docs/advanced/extensions.md +++ b/docs/advanced/extensions.md @@ -5,8 +5,8 @@ separate Maven module without defining backend operation types, sending checkpoi implementation packages. Extension-author contracts are in `software.amazon.lambda.durable.extension`. Built-in operation APIs are in -`software.amazon.lambda.durable.operation`, while operation-specific TLS metadata contexts remain in -`software.amazon.lambda.durable`. +`software.amazon.lambda.durable.operation`, with operation-specific TLS metadata contexts nested under their owning +operation classes. Application code calls only the extension's API: @@ -85,6 +85,10 @@ The same pattern applies to `DurableInvokeOperation.InvokeConfig`, `DurableWaitForConditionOperation.WaitForConditionConfig`, and `DurableWithRetryOperation.WithRetryConfig`. +`DurableWaitForConditionOperation` also owns the `WaitForConditionResult` used by its static APIs. The existing +`DurableContext` methods continue to use the compatibility result type under +`software.amazon.lambda.durable.model`. + Map and parallel extend `DurableConcurrencyOperation` and use its shared `DurableConcurrencyOperation.CompletionConfig` and `DurableConcurrencyOperation.NestingType` configuration types. @@ -123,7 +127,7 @@ User functions in the static APIs do not receive SDK context objects: ```java var result = DurableStepOperation.step("process", Result.class, () -> { - var step = StepContext.getCurrentContext(); + var step = StepContext.requireCurrentContext(); return process(step.getAttempt()); }); ``` @@ -152,13 +156,18 @@ var result = DurableWithRetryOperation.withRetry("transaction", () -> { ``` Parallel branch functions are `Supplier`. Wait-for-condition functions receive only the durable state value and -obtain attempt metadata from `StepContext.getCurrentContext()`. +obtain attempt metadata from `StepContext.requireCurrentContext()`. ## Current context scopes -`DurableContext.getCurrentContext()` and `ExtensionContext.getCurrentContext()` are available on SDK-managed handler -and child-context threads. `StepContext.getCurrentContext()` is available inside step and wait-for-condition user -functions. +`DurableContext.getCurrentContext()` is available on SDK-managed handler and child-context threads, while +`StepContext.getCurrentContext()` is available inside step and wait-for-condition user functions. Both methods return +`null` when no SDK context is active, preserving their use as availability probes. Use +`DurableContext.requireCurrentContext()` or `StepContext.requireCurrentContext()` when absence should be an error. +Calling either getter from the wrong SDK context type throws. + +`ExtensionContext.getCurrentContext()` is available on SDK-managed handler and child-context threads and throws when +no durable extension scope is active or when called from a step thread. `DurableMapOperation.MapItemContext`, `DurableWaitForCallbackOperation.WaitForCallbackContext`, and `DurableWithRetryOperation.WithRetryContext` are available only inside their corresponding user function. Nested @@ -278,11 +287,16 @@ var result = ExtensionContext.getCurrentContext() .build()); ``` -The function may return only `ExtensionStepResult.succeed(value)` or -`ExtensionStepResult.retry(state, delay)`. Retry state uses the configured `SerDes`; attempt metadata remains -available through `StepContext.getCurrentContext()`. Thrown exceptions follow the normal STEP failure path. -`ExtensionStepConfig` owns its retry strategy, which returns the same retry outcome used by stateful continuations. -Extension libraries can therefore configure exception retries and delivery semantics without depending on the +The function may return `ExtensionStepResult.succeed(value)`, +`ExtensionStepResult.retry(state, delay)`, or +`ExtensionStepResult.retryAfterNormalization(state, delayStrategy)`. The fixed-delay form uses the supplied delay. +The normalization-aware form first round-trips the state through the configured `SerDes`, then evaluates the delay +strategy against that checkpoint-normalized state. Both retry forms checkpoint the normalized state; attempt metadata +remains available through `StepContext.requireCurrentContext()`. Thrown exceptions follow the normal STEP failure +path. + +`ExtensionStepConfig` owns its exception retry strategy. It returns the fixed-delay retry outcome or a do-not-retry +decision, allowing extension libraries to configure exception retries and delivery semantics without depending on the customer-facing config or retry packages: ```java @@ -323,6 +337,10 @@ Use `ExtensionContextResult.completed(result)` when children never need to repla `replayChildrenAboveSize(result, replayState, thresholdBytes)` to replay only when the serialized full result reaches the threshold. Replay metadata is scoped to the framework callback through `ExtensionContextReplayContext`. +If replay is selected with a `null` replay state, the SDK checkpoints an empty payload. During replay, an empty payload +is exposed as a `null` replay state instead of being deserialized. This also allows extension contexts to replay +large-result checkpoints written by earlier SDK versions. + `ExtensionContextConfig` directly configures the context serializer and whether the context is virtual. It also controls framework user-function plugin events and can suppress child checkpoints that finish after the parent. If a context fails, the SDK first rethrows a deserialized original exception, then calls the configured error handler, and @@ -344,8 +362,8 @@ var result = ExtensionContext.getCurrentContext() .get(); ``` -Inside the function, `DurableContext.getCurrentContext()` and `ExtensionContext.getCurrentContext()` return the child -context. +Inside the function, `DurableContext.requireCurrentContext()` and `ExtensionContext.getCurrentContext()` return the +child context. ## Custom durable futures @@ -382,6 +400,11 @@ The built-in map, parallel, wait-for-callback, wait-for-condition, and with-retr the same extension primitives. Their legacy `DurableContext` methods and context-free static APIs share one canonical implementation while preserving their established checkpoint topology and plugin behavior. +Built-in child contexts, map iterations, parallel branches, wait-for-callback contexts, and checkpointed with-retry +contexts replay children when the serialized result is at least 256 KiB. They checkpoint an empty replay-state payload +when no compact state is required. Map and parallel parent contexts may also always replay children when their +completion or item-naming behavior requires reconstruction. + ## Module compatibility An extension Maven module should depend only on the public SDK artifact and import public types under @@ -391,5 +414,5 @@ or `primitive`. The extension-author SPI includes `ExtensionContext`, `ExtensionOperation`, stateful-step contracts, and configurable extension-context contracts under `software.amazon.lambda.durable.extension`. Static operation APIs are under -`software.amazon.lambda.durable.operation`; typed TLS contexts and `DurableFuture.completionFuture()` remain under -`software.amazon.lambda.durable`. +`software.amazon.lambda.durable.operation`; operation-specific TLS contexts are nested under their owning operation +classes, while `DurableFuture.completionFuture()` remains on the root SDK interface. diff --git a/docs/design.md b/docs/design.md index ac9760df9..044992755 100644 --- a/docs/design.md +++ b/docs/design.md @@ -248,13 +248,14 @@ context.step("name", Type.class, stepCtx -> doWork(), │ │ ▼ ▼ ┌──────────────────────────────┐ ┌──────────────────────────────┐ -│ Operations │ │ CheckpointBatcher │ -│ - StepPrimitive │ │ - Queues requests │ -│ - WaitPrimitive │ │ - Batches API calls (750KB) │ -│ - InvokePrimitive │ │ │ -│ - CallbackPrimitive │ │ - Notifies via callback │ -│ - ChildContextPrimitive │ -│ - execute() / get() │ +│ Operation facades + SPI │ │ CheckpointManager │ +│ - Durable*Operation │ │ - Queues requests │ +│ - ExtensionOperationImpl │ │ - Batches API calls (750KB) │ +│ │ │ │ - Polls operation updates │ +│ ▼ │ └──────────────────────────────┘ +│ Primitive engines │ +│ - Step/Wait/Invoke │ +│ - Callback/ChildContext │ └──────────────────────────────┘ │ ▼ @@ -270,7 +271,6 @@ context.step("name", Type.class, stepCtx -> doWork(), ``` software.amazon.lambda.durable ├── DurableHandler # Entry point -├── DurableExecutor # Lifecycle orchestration ├── DurableContext # User API (interface) ├── DurableFuture # Async handle ├── DurableCallbackFuture # Callback future with callbackId @@ -288,13 +288,17 @@ software.amazon.lambda.durable │ ├── ParallelBranchConfig # ParallelDurableFuture compatibility config │ ├── RunInChildContextConfig # DurableContext compatibility config │ ├── WaitForConditionConfig # DurableContext compatibility config -│ └── CompletionConfig # Completion criteria for map/parallel +│ ├── WithRetryConfig # DurableContext compatibility config +│ ├── CompletionConfig # Completion criteria for map/parallel +│ ├── NestingType # DurableContext compatibility nesting mode +│ └── StepSemantics # DurableContext compatibility step semantics │ ├── context/ -│ └── BaseContext # Base interface for DurableContext +│ ├── BaseContext # Shared DurableContext/StepContext interface +│ └── BaseContextImpl # Scoped current-context attachment │ ├── operation/ # Public built-in operation APIs + implementations -│ ├── DurableConcurrencyOperation # Shared map/parallel config, futures, and coordination +│ ├── DurableConcurrencyOperation # Shared map/parallel config, futures, and coordinator │ ├── DurableStepOperation # Owns nested StepConfig │ ├── DurableWaitOperation │ ├── DurableInvokeOperation @@ -303,40 +307,44 @@ software.amazon.lambda.durable │ ├── DurableMapOperation # Extends DurableConcurrencyOperation; owns MapConfig │ ├── DurableParallelOperation # Extends DurableConcurrencyOperation; owns parallel configs │ ├── DurableWaitForCallbackOperation -│ ├── DurableWaitForConditionOperation +│ ├── DurableWaitForConditionOperation # Owns config, result, and future adapter │ └── DurableWithRetryOperation │ ├── primitive/ # Internal checkpoint-backed operation engines -│ ├── StepPrimitive +│ ├── BasePrimitive +│ ├── SerializablePrimitive +│ ├── StepPrimitive │ ├── WaitPrimitive -│ ├── InvokePrimitive -│ ├── CallbackPrimitive -│ └── ChildContextPrimitive +│ ├── InvokePrimitive +│ ├── CallbackPrimitive +│ └── ChildContextPrimitive │ -├── extension/ # Public SPI for extension authors +├── extension/ # Public SPI for extension authors plus its internal bridge │ ├── ExtensionContext │ ├── ExtensionOperation -│ ├── ExtensionStepConfig # Owns extension StepSemantics and retry contracts -│ ├── ExtensionStepResult +│ ├── ExtensionOperationImpl # Internal bridge to primitive engines +│ ├── ExtensionStepFunction +│ ├── ExtensionStepConfig # Owns extension StepSemantics and retry contracts +│ ├── ExtensionStepResult +│ ├── ExtensionInvokeConfig +│ ├── ExtensionCallbackConfig +│ ├── ExtensionContextFunction │ ├── ExtensionContextConfig -│ └── ExtensionContextResult +│ ├── ExtensionContextResult +│ ├── ExtensionContextReplayContext +│ ├── ExtensionContextErrorHandler +│ ├── ExtensionContextFailure +│ └── ExtensionChildOperationSummary │ ├── execution/ +│ ├── DurableExecutor # Lifecycle orchestration │ ├── ExecutionManager # Central coordinator │ ├── ExecutionMode # REPLAY or EXECUTION state -│ ├── CheckpointBatcher # Batching (package-private) -│ ├── CheckpointCallback # Callback interface +│ ├── CheckpointManager # Checkpoint batching and polling +│ ├── ApiRequestDelayedBatcher # Shared delayed request batching │ ├── SuspendExecutionException │ └── ThreadType # CONTEXT, STEP │ -├── operation/ -│ ├── BasePrimitive # Common operation logic -│ ├── StepPrimitive # Step logic -│ ├── InvokePrimitive # Invoke logic -│ ├── CallbackPrimitive # Callback logic -│ ├── WaitPrimitive # Wait logic -│ └── ChildContextPrimitive # Child context primitive -│ ├── logging/ │ ├── DurableLogger # Context-aware logger wrapper (MDC-based) │ └── LoggerConfig # Replay suppression config @@ -405,23 +413,26 @@ software.amazon.lambda.durable sequenceDiagram participant UC as User Code participant DC as DurableContext - participant SO as StepPrimitive + participant DSO as DurableStepOperation + participant EO as ExtensionOperationImpl + participant SP as StepPrimitive participant EM as ExecutionManager participant Backend UC->>DC: step("name", Type.class, stepCtx -> doWork()) - DC->>SO: new StepPrimitive(...) - DC->>SO: execute() - SO->>EM: sendOperationUpdate(START) + DC->>DSO: stepAsync(...) + DSO->>EO: reserve(...).stepAsync(...) + EO->>SP: new StepPrimitive(...) + execute() + SP->>EM: sendOperationUpdate(START) EM->>Backend: checkpoint(START) - SO->>SO: func.apply(stepContext) [execute user code] + SP->>SP: func.apply(stepContext) [execute user code] - SO->>EM: sendOperationUpdate(SUCCEED) + SP->>EM: sendOperationUpdate(SUCCEED) EM->>Backend: checkpoint(SUCCEED) - DC->>SO: get() - SO-->>DC: result + DC->>SP: get() + SP-->>DC: result DC-->>UC: result ``` @@ -433,7 +444,9 @@ sequenceDiagram participant DE as DurableExecutor participant UC as User Code participant DC as DurableContext - participant SO as StepPrimitive + participant DSO as DurableStepOperation + participant EO as ExtensionOperationImpl + participant SP as StepPrimitive participant EM as ExecutionManager Note over LR: Re-invocation with existing state @@ -442,12 +455,14 @@ sequenceDiagram DE->>EM: new ExecutionManager(existingOps) UC->>DC: step("step1", ...) - DC->>SO: execute() - SO->>EM: getOperation("1") - EM-->>SO: existing op (SUCCEEDED) - Note over SO: Skip execution - DC->>SO: get() - SO-->>DC: cached result + DC->>DSO: stepAsync(...) + DSO->>EO: reserve(...).stepAsync(...) + EO->>SP: new StepPrimitive(...) + execute() + SP->>EM: getOperation("1") + EM-->>SP: existing op (SUCCEEDED) + Note over SP: Skip user function + DC->>SP: get() + SP-->>DC: cached result DC-->>UC: result ``` @@ -457,21 +472,25 @@ sequenceDiagram sequenceDiagram participant UC as User Code participant DC as DurableContext - participant WO as WaitPrimitive + participant DWO as DurableWaitOperation + participant EO as ExtensionOperationImpl + participant WP as WaitPrimitive participant EM as ExecutionManager participant Backend UC->>DC: wait(null, Duration.ofMinutes(5)) - DC->>WO: execute() - WO->>EM: sendOperationUpdate(WAIT, duration) + DC->>DWO: waitAsync(...) + DWO->>EO: reserve(...).waitAsync(...) + EO->>WP: new WaitPrimitive(...) + execute() + WP->>EM: sendOperationUpdate(START, waitOptions) EM->>Backend: checkpoint - DC->>WO: get() - WO->>EM: deregisterActiveThread("Root") + DC->>WP: get() + WP->>EM: deregisterActiveThread("Root") Note over EM: No active threads! EM->>EM: executionExceptionFuture.completeExceptionally(SuspendExecutionException) - EM-->>WO: throw SuspendExecutionException + EM-->>WP: throw SuspendExecutionException Note over UC: Execution suspended, returns PENDING ``` @@ -548,17 +567,17 @@ This is a one-way transition (REPLAY → EXECUTION, never back). `DurableLogger` **MDC Keys:** | Key | Set When | Description | |-----|----------|-------------| -| `durableExecutionArn` | Logger construction | Execution ARN | -| `requestId` | Logger construction | Lambda request ID | -| `operationId` | Step start | Current operation ID | -| `operationName` | Step start | Step name | -| `attempt` | Step start | Retry attempt number | +| `executionArn` | Logger scope attachment | Execution ARN (`durableExecutionArn` with legacy key names) | +| `requestId` | Logger scope attachment | Lambda request ID | +| `operationId` | Logger scope attachment | Current operation or child-context ID (`contextId` for legacy child-context keys) | +| `operationName` | Logger scope attachment | Current operation or child-context name (`contextName` for legacy child-context keys) | +| `attempt` | Step logger scope attachment | Retry attempt number | **Context Flow:** -1. `DurableLogger` constructor sets execution-level MDC (ARN, requestId) on the handler thread -2. `StepPrimitive.executeStepLogic()` calls `durableLogger.setOperationContext()` before user code runs -3. User code logs via `context.getLogger()` - MDC values automatically included -4. `clearOperationContext()` called in finally block after step completes +1. `DurableExecutor` or a primitive attaches the current `BaseContext` on its SDK-managed thread +2. `DurableLogger.attachContext()` derives execution, context, operation, and attempt MDC values from that scope +3. User code logs via `context.getLogger()` with the MDC values already attached +4. Closing the logger scope clears MDC when the handler, step, or child-context function finishes **Log Pattern Example (Log4j2):** ```xml @@ -584,25 +603,21 @@ If result > 6MB Lambda limit: ### Checkpoint Batching -Multiple concurrent operations may checkpoint simultaneously. `CheckpointBatcher` batches these into single API calls to reduce latency and stay within the 750KB request limit. +Multiple concurrent operations may checkpoint simultaneously. `CheckpointManager` uses +`ApiRequestDelayedBatcher` to combine them into API calls that stay within the 750KB request limit. The `checkpointDelay` configuration option (default: 0) controls how long the batcher waits before flushing, allowing more operations to accumulate in a single batch. For functions with many concurrent operations, setting a small delay (e.g., 10ms) can significantly reduce the number of API calls. ``` StepPrimitive 1 ──┐ │ -StepPrimitive 2 ──┼──► CheckpointBatcher ──► Backend +StepPrimitive 2 ──┼──► CheckpointManager ──► Backend │ WaitPrimitive ────┘ ``` -Callback mechanism avoids cyclic dependency between `ExecutionManager` and `CheckpointBatcher`: - -```java -interface CheckpointCallback { - void onComplete(String newToken, List operations); -} -``` +`CheckpointManager` sends completed operation updates back to `ExecutionManager`, which refreshes operation state +and notifies the registered primitive futures. --- @@ -804,22 +819,24 @@ Completing the future triggers the `thenRun` callback (re-registers the waiting Steps run user code on a separate thread via the user executor: ```java -// StepPrimitive.executeStepLogic() -registerActiveThread(getOperationId()); // register BEFORE submitting to executor - -CompletableFuture.runAsync(() -> { - try (StepContext stepContext = getContext().createStepContext(...)) { - T result = function.apply(stepContext); - handleStepSucceeded(result); // checkpoint SUCCEED synchronously - } catch (Throwable e) { - handleStepFailure(e, attempt); // checkpoint RETRY or FAIL +// StepPrimitive.executeExtensionStepLogic() +Runnable userHandler = () -> { + var stepContext = getContext().createStepContext(getOperationId(), getName(), attempt); + try (var ignoredContext = BaseContextImpl.attachCurrentContext(stepContext); + var ignoredLogger = DurableLogger.attachContext()) { + checkpointStarted(); + var result = runUserFunction(attempt, () -> extensionFunction.apply(state)); + handleExtensionStepResult(result, attempt); } -}, userExecutor); +}; +runUserHandler(userHandler, ThreadType.STEP); ``` Key details: -- `registerActiveThread` is called on the *parent* thread before `runAsync`, preventing a race where the parent deregisters (triggering suspension) before the step thread starts. -- The step thread is implicitly deregistered when it finishes — it never calls `deregisterActiveThread` directly. Instead, the step thread's work is done after checkpointing, and the checkpoint response completes the `completionFuture`, which re-registers the waiting context thread. +- `runUserHandler` registers the step on the *parent* thread before submitting it, preventing a race where the parent + deregisters before the step thread starts. +- The wrapper deregisters the step thread in `finally`; terminal checkpoint completion re-registers and wakes a + context thread waiting on the primitive future. - For retries, the step sends a RETRY checkpoint and then polls for the READY status before re-executing. If no other threads are active during the retry delay, the execution suspends. #### WaitPrimitive @@ -852,23 +869,25 @@ Child contexts run a user function in a separate thread with its own `DurableCon // ChildContextPrimitive.executeChildContext() var contextId = getOperationId(); -// Register on PARENT thread — prevents race with parent deregistration -registerActiveThread(contextId); - -CompletableFuture.runAsync(() -> { - try (var childContext = getContext().createChildContext(contextId, getName())) { - T result = function.apply(childContext); - handleChildContextSuccess(result); - } catch (Throwable e) { - handleChildContextFailure(e); +Runnable userHandler = () -> { + var childContext = createChildContext(contextId); + try (var ignoredContext = DurableContextImpl.attachCurrentContext(childContext); + var ignoredLogger = DurableLogger.attachContext()) { + executeFunction(childContext); } -}, userExecutor); +}; +runUserHandler(userHandler, ThreadType.CONTEXT); ``` Key details: +- `runUserHandler` registers the child context on the parent thread before submitting it and deregisters it when the + child thread finishes. - The child context thread runs as `ThreadType.CONTEXT` (not STEP), so it can itself create steps, waits, invokes, callbacks, and nested child contexts. - Operations within the child context use the child's `contextId` as their `parentId`, and operation IDs are prefixed with the context path (e.g. `"hash(1)"` for first-level, `"hash(hash(1)-2)"` for second-level). -- On replay, if the child context completed with a large result (> 256KB), the SDK re-executes the child context to reconstruct the result in memory rather than storing it in the checkpoint payload. +- A serialized result smaller than 256 KiB is checkpointed directly. At 256 KiB or larger, the SDK checkpoints an + empty payload with `replayChildren=true` and re-executes the child context on replay to reconstruct the result. +- Extension contexts use the same empty payload for a `null` replay state. Legacy empty replay payloads are therefore + interpreted as `null` instead of being deserialized. ### In-Process Completion @@ -880,7 +899,7 @@ When a context thread calls `ctx.step(...)`, the following coordination occurs: | Seq | Context Thread | Step Thread | System Thread (CheckpointManager) | |-----|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| 1 | Create `StepPrimitive` + `completionFuture`. Call `execute()`. `execute()` calls `start()` which registers step thread and submits to user executor. Checkpoint START (sync or async depending on semantics). | — | (idle) | +| 1 | `DurableStepOperation` reserves through `ExtensionOperationImpl`, creates `StepPrimitive` + `completionFuture`, and calls `execute()`. `start()` registers the step thread and submits to the user executor. Checkpoint START is sync or async depending on semantics. | — | (idle) | | 2 | `step()` calls `get()` → `waitForOperationCompletion()`. Attach `thenRun(re-register)` to `completionFuture`. Deregister context thread. Block on `join()`. | User code begins executing. Execute `function.apply(stepContext)`. | (idle) | | 3 | (blocked) | User code completes. Call `handleStepSucceeded(result)` → `sendOperationUpdate(SUCCEED)` (synchronous — blocks until checkpoint response). | Process checkpoint API call. On terminal response, call `onCheckpointComplete()` → `completionFuture.complete(null)`. `thenRun` fires: re-register context thread. | | 4 | `join()` returns. Retrieve result from operation. | Call `deregisterActiveThread` to deregister Step thread. Step thread ends. | (idle) | @@ -891,7 +910,7 @@ When a context thread calls `ctx.step(...)`, the following coordination occurs: | Seq | Context Thread | System Thread | |-----|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------|------------------------| -| 1 | Create `WaitPrimitive` + `completionFuture`. Call `execute()`. `execute()` calls `start()` → checkpoint WAIT with duration → `pollForOperationUpdates(remainingWaitTime)`. | Begin polling backend. | +| 1 | `DurableWaitOperation` reserves through `ExtensionOperationImpl`, creates `WaitPrimitive` + `completionFuture`, and calls `execute()` → checkpoint START with wait options → `pollForOperationUpdates(remainingWaitTime)`. | Begin polling backend. | | 2 | `wait()` calls `get()` → `waitForOperationCompletion()`. Attach `thenRun(re-register)`. Deregister context thread. | (polling) | | 3 | `activeThreads` is empty → `suspendExecution()` → `executionExceptionFuture.completeExceptionally(SuspendExecutionException)`. | — | | 4 | `runUntilCompleteOrSuspend` resolves with `SuspendExecutionException` → return `PENDING`. | — | @@ -908,8 +927,8 @@ var result = stepFuture.get(); | Seq | Context Thread | Step Thread | System Thread | |-----|--------------------------------------------------------------------|--------------------------------|---------------------------------------------------------------------------------------------------------| -| 1 | Create `StepPrimitive`, register step thread, submit to executor. | — | — | -| 2 | Create `WaitPrimitive`, checkpoint WAIT, start polling. | User code begins. | Begin polling for wait. | +| 1 | Reserve and create `StepPrimitive`, register step thread, submit to executor. | — | — | +| 2 | Reserve and create `WaitPrimitive`, checkpoint START with wait options, start polling. | User code begins. | Begin polling for wait. | | 3 | `wait()` calls `get()` → deregister context thread. | (running) | (polling) | | 4 | (blocked — but step thread is still active, so no suspension) | Complete → checkpoint SUCCEED. | Process step checkpoint. | | 5 | (blocked) | — | Wait poll returns SUCCEEDED → `completionFuture.complete(null)` for wait. Context thread re-registered. | diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/CallbackExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/CallbackExample.java index 82c4fe861..80eb3b177 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/CallbackExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/CallbackExample.java @@ -2,14 +2,15 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.callback; +import static software.amazon.lambda.durable.operation.DurableCallbackOperation.createCallback; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; +import static software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation.waitForCallbackAsync; + import java.time.Duration; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.examples.types.ApprovalRequest; -import software.amazon.lambda.durable.operation.DurableCallbackOperation; import software.amazon.lambda.durable.operation.DurableCallbackOperation.CallbackConfig; -import software.amazon.lambda.durable.operation.DurableStepOperation; -import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation; import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation.WaitForCallbackContext; /** @@ -37,7 +38,7 @@ public class CallbackExample extends DurableHandler { @Override public String handleRequest(ApprovalRequest input) { // Step 1: Prepare the approval request - var prepared = DurableStepOperation.step( + var prepared = step( "prepare", String.class, () -> "Approval request for: " + input.description() + " ($" + input.amount() + ")"); @@ -49,18 +50,15 @@ public String handleRequest(ApprovalRequest input) { var config = CallbackConfig.builder().timeout(timeout).build(); - var preapprovalCallback = - DurableWaitForCallbackOperation.waitForCallbackAsync("preapproval", String.class, () -> { - var callbackId = WaitForCallbackContext.getCurrentContext().getCallbackId(); - StepContext.getCurrentContext() - .getLogger() - .info("Sending callback {} to preapproval system", callbackId); - }); + var preapprovalCallback = waitForCallbackAsync("preapproval", String.class, () -> { + var callbackId = WaitForCallbackContext.getCurrentContext().getCallbackId(); + StepContext.getCurrentContext().getLogger().info("Sending callback {} to preapproval system", callbackId); + }); - var callback = DurableCallbackOperation.createCallback("approval", String.class, config); + var callback = createCallback("approval", String.class, config); // Step 2.5: Log AWS CLI command to complete the callback - DurableStepOperation.step("log-callback-command", Void.class, () -> { + step("log-callback-command", Void.class, () -> { var callbackId = callback.callbackId(); // The result must be base64-encoded JSON var command = String.format( @@ -76,7 +74,7 @@ public String handleRequest(ApprovalRequest input) { var approvalResult = callback.get(); // Step 4: Process the approval - return DurableStepOperation.step( + return step( "process-approval", String.class, () -> prepared + " - " + preapprovalResult + " - " + approvalResult); } } diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/RetryWaitForCallbackExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/RetryWaitForCallbackExample.java index f590bc21b..55775581f 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/RetryWaitForCallbackExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/RetryWaitForCallbackExample.java @@ -2,20 +2,21 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.callback; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; +import static software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation.waitForCallback; +import static software.amazon.lambda.durable.operation.DurableWithRetryOperation.withRetry; + import java.time.Duration; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.examples.types.ApprovalRequest; -import software.amazon.lambda.durable.operation.DurableStepOperation; -import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation; import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation.WaitForCallbackContext; -import software.amazon.lambda.durable.operation.DurableWithRetryOperation; import software.amazon.lambda.durable.operation.DurableWithRetryOperation.WithRetryConfig; import software.amazon.lambda.durable.operation.DurableWithRetryOperation.WithRetryContext; import software.amazon.lambda.durable.retry.RetryDecision; /** - * Example demonstrating {@link DurableWithRetryOperation} with {@link DurableWaitForCallbackOperation}. + * Example demonstrating {@code withRetry} with {@code waitForCallback}. * *

    Submits an approval request to an external system via a callback. If the callback fails (e.g., the external system * rejects the request), the helper retries the entire waitForCallback cycle — creating a fresh callback with a new ID @@ -32,22 +33,20 @@ public class RetryWaitForCallbackExample extends DurableHandler "Approval for: " + input.description() + " ($" + input.amount() + ")"); // Step 2: waitForCallback with retry — if the external system fails, try again with a fresh callback - var approvalResult = DurableWithRetryOperation.withRetry( + var approvalResult = withRetry( null, () -> { var attempt = WithRetryContext.getCurrentContext().getAttempt(); - return DurableWaitForCallbackOperation.waitForCallback( - "approval-" + attempt, String.class, () -> StepContext.getCurrentContext() - .getLogger() - .info( - "Attempt {}: sending callback {} to approval system", - attempt, - WaitForCallbackContext.getCurrentContext() - .getCallbackId())); + return waitForCallback("approval-" + attempt, String.class, () -> StepContext.getCurrentContext() + .getLogger() + .info( + "Attempt {}: sending callback {} to approval system", + attempt, + WaitForCallbackContext.getCurrentContext().getCallbackId())); }, WithRetryConfig.builder() .retryStrategy((error, attempt) -> attempt < MAX_ATTEMPTS @@ -56,7 +55,6 @@ public String handleRequest(ApprovalRequest input) { .build()); // Step 3: Process the result - return DurableStepOperation.step( - "process-result", String.class, () -> prepared + " - Result: " + approvalResult); + return step("process-result", String.class, () -> prepared + " - Result: " + approvalResult); } } diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/WaitForCallbackFailedExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/WaitForCallbackFailedExample.java index 8843b8ae0..3bf57e916 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/WaitForCallbackFailedExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/WaitForCallbackFailedExample.java @@ -2,13 +2,14 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.callback; +import static software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation.waitForCallback; + import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.examples.types.ApprovalRequest; import software.amazon.lambda.durable.exception.SerDesException; import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; -import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation; import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation.WaitForCallbackConfig; import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation.WaitForCallbackContext; import software.amazon.lambda.durable.serde.JacksonSerDes; @@ -21,7 +22,7 @@ public String handleRequest(ApprovalRequest input) { String approvalResult; try { - approvalResult = DurableWaitForCallbackOperation.waitForCallback( + approvalResult = waitForCallback( "preapproval", String.class, () -> { diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ChildContextExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ChildContextExample.java index d0f82acfb..58bc9f3ef 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ChildContextExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ChildContextExample.java @@ -2,13 +2,15 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.child; +import static software.amazon.lambda.durable.operation.DurableContextOperation.runInChildContext; +import static software.amazon.lambda.durable.operation.DurableContextOperation.runInChildContextAsync; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import java.time.Duration; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; -import software.amazon.lambda.durable.operation.DurableContextOperation; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.operation.DurableWaitOperation; /** @@ -17,8 +19,10 @@ *

    This handler runs three concurrent child contexts using {@code runInChildContextAsync}: * *

      - *
    1. Order validation — performs a step then suspends via {@code wait()} before completing - *
    2. Inventory check — performs a step then suspends via {@code wait()} before completing + *
    3. Order validation — performs a step then suspends via {@code DurableWaitOperation.wait()} before + * completing + *
    4. Inventory check — performs a step then suspends via {@code DurableWaitOperation.wait()} before + * completing *
    5. Shipping estimate — nests another child context inside it to demonstrate hierarchical contexts *
    * @@ -34,38 +38,36 @@ public String handleRequest(GreetingRequest input) { context.getLogger().info("Starting child context workflow for {}", name); // Child context 1: Order validation — step + wait + step - var orderFuture = DurableContextOperation.runInChildContextAsync("order-validation", String.class, () -> { - var prepared = DurableStepOperation.step("prepare-order", String.class, () -> "Order for " + name); + var orderFuture = runInChildContextAsync("order-validation", String.class, () -> { + var prepared = step("prepare-order", String.class, () -> "Order for " + name); DurableContext.getCurrentContext().getLogger().info("Order prepared, waiting for validation"); DurableWaitOperation.wait("validation-delay", Duration.ofSeconds(5)); - return DurableStepOperation.step("validate-order", String.class, () -> prepared + " [validated]"); + return step("validate-order", String.class, () -> prepared + " [validated]"); }); // Child context 2: Inventory check — step + wait + step - var inventoryFuture = DurableContextOperation.runInChildContextAsync("inventory-check", String.class, () -> { - var stock = DurableStepOperation.step("check-stock", String.class, () -> "Stock available for " + name); + var inventoryFuture = runInChildContextAsync("inventory-check", String.class, () -> { + var stock = step("check-stock", String.class, () -> "Stock available for " + name); DurableContext.getCurrentContext().getLogger().info("Stock checked, waiting for confirmation"); DurableWaitOperation.wait("confirmation-delay", Duration.ofSeconds(3)); - return DurableStepOperation.step("confirm-inventory", String.class, () -> stock + " [confirmed]"); + return step("confirm-inventory", String.class, () -> stock + " [confirmed]"); }); // Child context 3: Shipping estimate — nests a child context inside it - var shippingFuture = DurableContextOperation.runInChildContextAsync("shipping-estimate", String.class, () -> { - var baseRate = - DurableStepOperation.step("calculate-base-rate", String.class, () -> "Base rate for " + name); + var shippingFuture = runInChildContextAsync("shipping-estimate", String.class, () -> { + var baseRate = step("calculate-base-rate", String.class, () -> "Base rate for " + name); // Nested child context: calculate regional adjustment - var adjustment = DurableContextOperation.runInChildContext( + var adjustment = runInChildContext( "regional-adjustment", String.class, - () -> DurableStepOperation.step( - "lookup-region", String.class, () -> baseRate + " + regional adjustment")); + () -> step("lookup-region", String.class, () -> baseRate + " + regional adjustment")); - return DurableStepOperation.step("finalize-shipping", String.class, () -> adjustment + " [shipping ready]"); + return step("finalize-shipping", String.class, () -> adjustment + " [shipping ready]"); }); // Collect all results using allOf diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ManyAsyncChildContextExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ManyAsyncChildContextExample.java index 89ceed290..b4a6b6914 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ManyAsyncChildContextExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ManyAsyncChildContextExample.java @@ -2,6 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.child; +import static software.amazon.lambda.durable.operation.DurableContextOperation.runInChildContextAsync; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import java.time.Duration; import java.util.ArrayList; import java.util.concurrent.TimeUnit; @@ -11,8 +14,6 @@ import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.ManyAsyncStepsInput; import software.amazon.lambda.durable.examples.types.ManyAsyncStepsOutput; -import software.amazon.lambda.durable.operation.DurableContextOperation; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.operation.DurableWaitOperation; /** @@ -41,9 +42,9 @@ public ManyAsyncStepsOutput handleRequest(ManyAsyncStepsInput input) { var futures = new ArrayList>(steps); for (var i = 0; i < steps; i++) { var index = i; - var future = DurableContextOperation.runInChildContextAsync("child-" + i, Integer.class, () -> { + var future = runInChildContextAsync("child-" + i, Integer.class, () -> { // create a step inside the child context, which doubles the number of threads - return DurableStepOperation.step("compute-" + index, Integer.class, () -> index * multiplier); + return step("compute-" + index, Integer.class, () -> index * multiplier); }); futures.add(future); } @@ -55,8 +56,8 @@ public ManyAsyncStepsOutput handleRequest(ManyAsyncStepsInput input) { var totalSum = results.stream().mapToInt(Integer::intValue).sum(); // checkpoint the executionTime so that we can have the same value when replay - var executionTimeMs = DurableStepOperation.step( - "execution-time", Long.class, () -> TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime)); + var executionTimeMs = + step("execution-time", Long.class, () -> TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime)); logger.info( "Completed {} child context, total sum: {}, execution time: {}ms", steps, totalSum, executionTimeMs); diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/VirtualChildContextExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/VirtualChildContextExample.java index 32031d490..bbc088f5c 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/VirtualChildContextExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/VirtualChildContextExample.java @@ -2,14 +2,16 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.child; +import static software.amazon.lambda.durable.operation.DurableContextOperation.runInChildContext; +import static software.amazon.lambda.durable.operation.DurableContextOperation.runInChildContextAsync; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import java.time.Duration; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; -import software.amazon.lambda.durable.operation.DurableContextOperation; import software.amazon.lambda.durable.operation.DurableContextOperation.RunInChildContextConfig; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.operation.DurableWaitOperation; /** @@ -18,8 +20,10 @@ *

    This handler runs three concurrent child contexts using {@code runInChildContextAsync}: * *

      - *
    1. Order validation — performs a step then suspends via {@code wait()} before completing - *
    2. Inventory check — performs a step then suspends via {@code wait()} before completing + *
    3. Order validation — performs a step then suspends via {@code DurableWaitOperation.wait()} before + * completing + *
    4. Inventory check — performs a step then suspends via {@code DurableWaitOperation.wait()} before + * completing *
    5. Shipping estimate — nests another child context inside it to demonstrate hierarchical contexts *
    * @@ -35,52 +39,48 @@ public String handleRequest(GreetingRequest input) { context.getLogger().info("Starting child context workflow for {}", name); // Child context 1: Order validation — step + wait + step - var orderFuture = DurableContextOperation.runInChildContextAsync( + var orderFuture = runInChildContextAsync( "order-validation", String.class, () -> { - var prepared = DurableStepOperation.step("prepare-order", String.class, () -> "Order for " + name); + var prepared = step("prepare-order", String.class, () -> "Order for " + name); DurableContext.getCurrentContext().getLogger().info("Order prepared, waiting for validation"); DurableWaitOperation.wait("validation-delay", Duration.ofSeconds(5)); - return DurableStepOperation.step("validate-order", String.class, () -> prepared + " [validated]"); + return step("validate-order", String.class, () -> prepared + " [validated]"); }, RunInChildContextConfig.builder().isVirtual(true).build()); // Child context 2: Inventory check — step + wait + step - var inventoryFuture = DurableContextOperation.runInChildContextAsync( + var inventoryFuture = runInChildContextAsync( "inventory-check", String.class, () -> { - var stock = - DurableStepOperation.step("check-stock", String.class, () -> "Stock available for " + name); + var stock = step("check-stock", String.class, () -> "Stock available for " + name); DurableContext.getCurrentContext().getLogger().info("Stock checked, waiting for confirmation"); DurableWaitOperation.wait("confirmation-delay", Duration.ofSeconds(3)); - return DurableStepOperation.step("confirm-inventory", String.class, () -> stock + " [confirmed]"); + return step("confirm-inventory", String.class, () -> stock + " [confirmed]"); }, RunInChildContextConfig.builder().isVirtual(true).build()); // Child context 3: Shipping estimate — nests a child context inside it - var shippingFuture = DurableContextOperation.runInChildContextAsync( + var shippingFuture = runInChildContextAsync( "shipping-estimate", String.class, () -> { - var baseRate = DurableStepOperation.step( - "calculate-base-rate", String.class, () -> "Base rate for " + name); + var baseRate = step("calculate-base-rate", String.class, () -> "Base rate for " + name); // Nested child context: calculate regional adjustment - var adjustment = DurableContextOperation.runInChildContext( + var adjustment = runInChildContext( "regional-adjustment", String.class, - () -> DurableStepOperation.step( - "lookup-region", String.class, () -> baseRate + " + regional adjustment"), + () -> step("lookup-region", String.class, () -> baseRate + " + regional adjustment"), RunInChildContextConfig.builder().isVirtual(true).build()); - return DurableStepOperation.step( - "finalize-shipping", String.class, () -> adjustment + " [shipping ready]"); + return step("finalize-shipping", String.class, () -> adjustment + " [shipping ready]"); }, RunInChildContextConfig.builder().isVirtual(true).build()); diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/CustomConfigExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/CustomConfigExample.java index ecdfad5f6..7171a4747 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/CustomConfigExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/CustomConfigExample.java @@ -2,6 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.general; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.PropertyNamingStrategies; @@ -16,7 +18,6 @@ import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.client.LambdaDurableFunctionsClient; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.serde.SerDes; /** @@ -77,7 +78,7 @@ protected DurableConfig createConfiguration() { @Override public String handleRequest(String input) { // Step 1: Create a custom object with camelCase fields to demonstrate snake_case serialization - var customObject = DurableStepOperation.step( + var customObject = step( "create-custom-object", CustomData.class, () -> new CustomData("user123", "John Doe", 25, "john.doe@example.com")); diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/CustomPollingExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/CustomPollingExample.java index 2165093de..8933d4e3a 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/CustomPollingExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/CustomPollingExample.java @@ -2,14 +2,15 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.general; +import static software.amazon.lambda.durable.operation.DurableInvokeOperation.invokeAsync; +import static software.amazon.lambda.durable.operation.DurableStepOperation.stepAsync; + import java.time.Duration; import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; -import software.amazon.lambda.durable.operation.DurableInvokeOperation; import software.amazon.lambda.durable.operation.DurableInvokeOperation.InvokeConfig; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.retry.JitterStrategy; import software.amazon.lambda.durable.retry.PollingStrategies; @@ -43,7 +44,7 @@ public String handleRequest(GreetingRequest input) { context.getLogger().info("Starting workflow with input: {}", input); // Step 1: low case the input - var lowered = DurableStepOperation.stepAsync("validate", String.class, () -> { + var lowered = stepAsync("validate", String.class, () -> { try { // prevent the execution from suspension Thread.sleep(5000); @@ -54,7 +55,7 @@ public String handleRequest(GreetingRequest input) { }); // Step 2: Invoke async - var future = DurableInvokeOperation.invokeAsync( + var future = invokeAsync( "call-greeting", "simple-step-example" + input.getName() + ":$LATEST", input, diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/ErrorHandlingExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/ErrorHandlingExample.java index 3d0ea70eb..5a131e047 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/ErrorHandlingExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/ErrorHandlingExample.java @@ -2,13 +2,14 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.general; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.config.StepSemantics; import software.amazon.lambda.durable.exception.StepFailedException; import software.amazon.lambda.durable.exception.StepInterruptedException; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; import software.amazon.lambda.durable.retry.RetryStrategies; @@ -56,7 +57,7 @@ public String handleRequest(Object input) { // NOTE: Exception type needs to be serializable by your SerDes implementation. String primaryResult; try { - primaryResult = DurableStepOperation.step( + primaryResult = step( "call-primary-service", String.class, () -> { @@ -68,7 +69,7 @@ public String handleRequest(Object input) { } catch (ServiceUnavailableException e) { // Catch the specific custom exception type - the SDK reconstructs the original exception logger.warn("Service '{}' unavailable, using fallback: {}", e.getServiceName(), e.getMessage()); - primaryResult = DurableStepOperation.step("call-fallback-service", String.class, () -> "fallback-result"); + primaryResult = step("call-fallback-service", String.class, () -> "fallback-result"); } // Example 2: Handling StepInterruptedException for AT_MOST_ONCE operations @@ -78,7 +79,7 @@ public String handleRequest(Object input) { // interruption scenario that occurs during replay after an unexpected termination. String paymentResult; try { - paymentResult = DurableStepOperation.step( + paymentResult = step( "charge-payment", String.class, () -> "payment-" + input, @@ -92,7 +93,7 @@ public String handleRequest(Object input) { e.getOperation().id()); // In real code: check payment provider for transaction status // If payment went through, return success; otherwise, handle appropriately - paymentResult = DurableStepOperation.step("verify-payment-status", String.class, () -> "verified-payment"); + paymentResult = step("verify-payment-status", String.class, () -> "verified-payment"); } return "Completed: " + primaryResult + ", " + paymentResult; diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/GenericInputOutputExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/GenericInputOutputExample.java index f30d8e203..e0c538ddd 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/GenericInputOutputExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/GenericInputOutputExample.java @@ -2,6 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.general; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import java.util.HashMap; import java.util.List; import java.util.Map; @@ -9,7 +11,6 @@ import org.slf4j.LoggerFactory; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; import software.amazon.lambda.durable.retry.RetryStrategies; @@ -29,7 +30,7 @@ public Map>> handleRequest(Map logger.info("Starting generic types example for user: {}", input.get("userId")); // Fetch nested generic type with retry (Map>) - Map> categories = DurableStepOperation.step( + Map> categories = step( "fetch-categories", new TypeToken>>() {}, () -> { diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/GenericTypesExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/GenericTypesExample.java index 8e8db2faa..60ee051b0 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/GenericTypesExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/GenericTypesExample.java @@ -2,6 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.general; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import java.util.HashMap; import java.util.List; import java.util.Map; @@ -9,7 +11,6 @@ import org.slf4j.LoggerFactory; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; import software.amazon.lambda.durable.retry.RetryStrategies; @@ -52,26 +53,25 @@ public Output handleRequest(Input input) { logger.info("Starting generic types example for user: {}", input.userId); // Step 1: Fetch a list of items (List) - List items = DurableStepOperation.step("fetch-items", new TypeToken>() {}, () -> { + List items = step("fetch-items", new TypeToken>() {}, () -> { logger.info("Fetching items for user: {}", input.userId); return List.of("item1", "item2", "item3", "item4"); }); logger.info("Fetched {} items", items.size()); // Step 2: Count items by category (Map) - Map counts = - DurableStepOperation.step("count-by-category", new TypeToken>() {}, () -> { - logger.info("Counting items by category"); - var result = new HashMap(); - result.put("electronics", 2); - result.put("books", 1); - result.put("clothing", 1); - return result; - }); + Map counts = step("count-by-category", new TypeToken>() {}, () -> { + logger.info("Counting items by category"); + var result = new HashMap(); + result.put("electronics", 2); + result.put("books", 1); + result.put("clothing", 1); + return result; + }); logger.info("Counted {} categories", counts.size()); // Step 3: Fetch nested generic type with retry (Map>) - Map> categories = DurableStepOperation.step( + Map> categories = step( "fetch-categories", new TypeToken>>() {}, () -> { diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/LoggingExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/LoggingExample.java index 440d1e301..118b8e8dc 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/LoggingExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/LoggingExample.java @@ -2,13 +2,14 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.general; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.examples.types.GreetingRequest; -import software.amazon.lambda.durable.operation.DurableStepOperation; /** * Example demonstrating DurableLogger usage for structured logging with execution context. @@ -26,13 +27,13 @@ public String handleRequest(GreetingRequest input) { context.getLogger(logger).info("Processing greeting for: {}", input.getName()); // Step 1: Create greeting - logs inside step include operation context - var greeting = DurableStepOperation.step("create-greeting", String.class, () -> { + var greeting = step("create-greeting", String.class, () -> { StepContext.getCurrentContext().getLogger(logger).info("Creating greeting message"); return "Hello, " + input.getName(); }); // Step 2: Transform - var result = DurableStepOperation.step("transform", String.class, () -> { + var result = step("transform", String.class, () -> { StepContext.getCurrentContext().getLogger().info("Transforming greeting to uppercase"); return greeting.toUpperCase() + "!"; }); diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/OtelExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/OtelExample.java index 12911957f..2f47a65c3 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/OtelExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/OtelExample.java @@ -2,6 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.general; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import io.opentelemetry.exporter.logging.LoggingSpanExporter; import io.opentelemetry.sdk.trace.SdkTracerProvider; import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; @@ -9,7 +11,6 @@ import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.otel.InvocationOtelPlugin; /** @@ -52,12 +53,12 @@ public String handleRequest(GreetingRequest input) { // Log with MDC — traceId and spanId will be in the JSON output context.getLogger().info("Starting OTel example for {}", input.getName()); - var greeting = DurableStepOperation.step("create-greeting", String.class, () -> { + var greeting = step("create-greeting", String.class, () -> { context.getLogger().info("Inside step — this log has trace context in MDC"); return "Hello, " + input.getName(); }); - var result = DurableStepOperation.step("transform", String.class, () -> greeting.toUpperCase() + "!"); + var result = step("transform", String.class, () -> greeting.toUpperCase() + "!"); context.getLogger().info("OTel example complete: {}", result); return result; diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/PluginExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/PluginExample.java index 6f9cee147..174e3b73d 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/PluginExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/PluginExample.java @@ -2,11 +2,12 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.general; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.plugin.*; /** @@ -42,9 +43,9 @@ public String handleRequest(GreetingRequest input) { var context = DurableContext.getCurrentContext(); context.getLogger().info("Starting plugin example for {}", input.getName()); - var greeting = DurableStepOperation.step("create-greeting", String.class, () -> "Hello, " + input.getName()); + var greeting = step("create-greeting", String.class, () -> "Hello, " + input.getName()); - var result = DurableStepOperation.step("transform", String.class, () -> greeting.toUpperCase() + "!"); + var result = step("transform", String.class, () -> greeting.toUpperCase() + "!"); context.getLogger().info("Plugin example complete: {}", result); return result; diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/invoke/RetryInvokeExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/invoke/RetryInvokeExample.java index a8dbd2fa2..da54bc233 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/invoke/RetryInvokeExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/invoke/RetryInvokeExample.java @@ -2,17 +2,18 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.invoke; +import static software.amazon.lambda.durable.operation.DurableInvokeOperation.invoke; +import static software.amazon.lambda.durable.operation.DurableWithRetryOperation.withRetry; + import java.time.Duration; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; -import software.amazon.lambda.durable.operation.DurableInvokeOperation; -import software.amazon.lambda.durable.operation.DurableWithRetryOperation; import software.amazon.lambda.durable.operation.DurableWithRetryOperation.WithRetryConfig; import software.amazon.lambda.durable.operation.DurableWithRetryOperation.WithRetryContext; import software.amazon.lambda.durable.retry.RetryDecision; /** - * Example demonstrating {@link DurableWithRetryOperation} with {@link DurableInvokeOperation}. + * Example demonstrating {@code withRetry} with {@code invoke}. * *

    Retries a chained Lambda invocation up to 3 times with a fixed 2-second backoff between attempts. Each attempt * uses a unique operation name ({@code "call-greeting-1"}, {@code "call-greeting-2"}, etc.) so the execution history @@ -29,12 +30,11 @@ public String handleRequest(GreetingRequest input) { var targetFunctionName = System.getenv().getOrDefault("FUNCTION_NAME_PREFIX", "") + "simple-step-example:$LATEST"; - return DurableWithRetryOperation.withRetry( + return withRetry( null, () -> { var attempt = WithRetryContext.getCurrentContext().getAttempt(); - return DurableInvokeOperation.invoke( - "call-greeting-" + attempt, targetFunctionName, input, String.class); + return invoke("call-greeting-" + attempt, targetFunctionName, input, String.class); }, WithRetryConfig.builder() .retryStrategy((error, attempt) -> attempt < MAX_ATTEMPTS diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/invoke/SimpleInvokeExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/invoke/SimpleInvokeExample.java index d4ca1d68d..233951319 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/invoke/SimpleInvokeExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/invoke/SimpleInvokeExample.java @@ -2,9 +2,11 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.invoke; +import static software.amazon.lambda.durable.operation.DurableInvokeOperation.invoke; +import static software.amazon.lambda.durable.operation.DurableInvokeOperation.invokeAsync; + import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; -import software.amazon.lambda.durable.operation.DurableInvokeOperation; import software.amazon.lambda.durable.operation.DurableInvokeOperation.InvokeConfig; /** @@ -20,13 +22,13 @@ public String handleRequest(GreetingRequest input) { System.getenv().getOrDefault("FUNCTION_NAME_PREFIX", "") + "simple-step-example:$LATEST"; // Invoke the `simple-step-example` function. - var future = DurableInvokeOperation.invokeAsync( + var future = invokeAsync( "call-greeting1", targetFunctionName, input, String.class, InvokeConfig.builder().build()); - var result2 = DurableInvokeOperation.invoke( + var result2 = invoke( "call-greeting2", targetFunctionName, input, diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexFlatMapExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexFlatMapExample.java index 03a4335a9..f0314048a 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexFlatMapExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexFlatMapExample.java @@ -2,6 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.map; +import static software.amazon.lambda.durable.operation.DurableMapOperation.map; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import java.time.Duration; import java.util.List; import java.util.stream.Collectors; @@ -10,9 +13,8 @@ import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.operation.DurableConcurrencyOperation.CompletionConfig; import software.amazon.lambda.durable.operation.DurableConcurrencyOperation.NestingType; -import software.amazon.lambda.durable.operation.DurableMapOperation; import software.amazon.lambda.durable.operation.DurableMapOperation.MapConfig; -import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableMapOperation.MapItemContext; import software.amazon.lambda.durable.operation.DurableWaitOperation; /** @@ -34,22 +36,20 @@ public String handleRequest(Integer input) { // Part 1: Concurrent map with step + wait inside each branch var orderIds = IntStream.range(1, input + 1).mapToObj(x -> "order-" + x).collect(Collectors.toList()); - var orderResult = DurableMapOperation.map( + var orderResult = map( "process-orders", orderIds, String.class, orderId -> { - var index = DurableMapOperation.MapItemContext.getCurrentContext() - .getIndex(); + var index = MapItemContext.getCurrentContext().getIndex(); // Step 1: validate the order - var validated = - DurableStepOperation.step("validate-" + index, String.class, () -> "validated:" + orderId); + var validated = step("validate-" + index, String.class, () -> "validated:" + orderId); // Wait between stages (simulates a cooldown or external dependency) DurableWaitOperation.wait("cooldown-" + index, Duration.ofSeconds(1)); // Step 2: finalize the order - return DurableStepOperation.step("finalize-" + index, String.class, () -> "done:" + validated); + return step("finalize-" + index, String.class, () -> "done:" + validated); }, MapConfig.builder().nestingType(NestingType.FLAT).build()); @@ -62,14 +62,13 @@ public String handleRequest(Integer input) { .nestingType(NestingType.FLAT) .build(); - var serverResult = DurableMapOperation.map( + var serverResult = map( "find-healthy-servers", servers, String.class, server -> { - var index = DurableMapOperation.MapItemContext.getCurrentContext() - .getIndex(); - return DurableStepOperation.step("health-check-" + index, String.class, () -> server + ":healthy"); + var index = MapItemContext.getCurrentContext().getIndex(); + return step("health-check-" + index, String.class, () -> server + ":healthy"); }, earlyTermConfig); diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexMapExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexMapExample.java index c466faffb..bfdcde1df 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexMapExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexMapExample.java @@ -2,6 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.map; +import static software.amazon.lambda.durable.operation.DurableMapOperation.map; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import java.time.Duration; import java.util.List; import java.util.stream.Collectors; @@ -9,9 +12,8 @@ import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.operation.DurableConcurrencyOperation.CompletionConfig; -import software.amazon.lambda.durable.operation.DurableMapOperation; import software.amazon.lambda.durable.operation.DurableMapOperation.MapConfig; -import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableMapOperation.MapItemContext; import software.amazon.lambda.durable.operation.DurableWaitOperation; /** @@ -33,16 +35,16 @@ public String handleRequest(Integer input) { // Part 1: Concurrent map with step + wait inside each branch var orderIds = IntStream.range(1, input + 1).mapToObj(x -> "order-" + x).collect(Collectors.toList()); - var orderResult = DurableMapOperation.map("process-orders", orderIds, String.class, orderId -> { - var index = DurableMapOperation.MapItemContext.getCurrentContext().getIndex(); + var orderResult = map("process-orders", orderIds, String.class, orderId -> { + var index = MapItemContext.getCurrentContext().getIndex(); // Step 1: validate the order - var validated = DurableStepOperation.step("validate-" + index, String.class, () -> "validated:" + orderId); + var validated = step("validate-" + index, String.class, () -> "validated:" + orderId); // Wait between stages (simulates a cooldown or external dependency) DurableWaitOperation.wait("cooldown-" + index, Duration.ofSeconds(1)); // Step 2: finalize the order - return DurableStepOperation.step("finalize-" + index, String.class, () -> "done:" + validated); + return step("finalize-" + index, String.class, () -> "done:" + validated); }); var orderSummary = String.join(", ", orderResult.results()); @@ -53,14 +55,13 @@ public String handleRequest(Integer input) { .completionConfig(CompletionConfig.minSuccessful(2)) .build(); - var serverResult = DurableMapOperation.map( + var serverResult = map( "find-healthy-servers", servers, String.class, server -> { - var index = DurableMapOperation.MapItemContext.getCurrentContext() - .getIndex(); - return DurableStepOperation.step("health-check-" + index, String.class, () -> server + ":healthy"); + var index = MapItemContext.getCurrentContext().getIndex(); + return step("health-check-" + index, String.class, () -> server + ":healthy"); }, earlyTermConfig); diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/CustomShouldCompleteMapExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/CustomShouldCompleteMapExample.java index fa42bef79..14394d4f7 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/CustomShouldCompleteMapExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/CustomShouldCompleteMapExample.java @@ -7,13 +7,14 @@ import static software.amazon.lambda.durable.model.MapResult.MapResultItem.Status.SKIPPED; import static software.amazon.lambda.durable.operation.DurableConcurrencyOperation.CompletionConfig.CompletionDecision.complete; import static software.amazon.lambda.durable.operation.DurableConcurrencyOperation.CompletionConfig.CompletionDecision.continueExecution; +import static software.amazon.lambda.durable.operation.DurableMapOperation.map; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; import java.util.List; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.operation.DurableConcurrencyOperation.CompletionConfig; -import software.amazon.lambda.durable.operation.DurableMapOperation; import software.amazon.lambda.durable.operation.DurableMapOperation.MapConfig; -import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableMapOperation.MapItemContext; import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; import software.amazon.lambda.durable.retry.RetryStrategies; @@ -54,14 +55,13 @@ public Output handleRequest(Input input) { })) .build(); - var result = DurableMapOperation.map( + var result = map( "query-providers", input.providers(), String.class, provider -> { - var index = DurableMapOperation.MapItemContext.getCurrentContext() - .getIndex(); - return DurableStepOperation.step( + var index = MapItemContext.getCurrentContext().getIndex(); + return step( "query-" + index, String.class, () -> { diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/DeserializationFailedMapExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/DeserializationFailedMapExample.java index 58f883d3a..a743fa5c0 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/DeserializationFailedMapExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/DeserializationFailedMapExample.java @@ -2,6 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.map; +import static software.amazon.lambda.durable.operation.DurableMapOperation.map; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import java.time.Duration; import java.util.List; import software.amazon.lambda.durable.DurableContext; @@ -9,9 +12,8 @@ import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.examples.types.GreetingRequest; import software.amazon.lambda.durable.exception.SerDesException; -import software.amazon.lambda.durable.operation.DurableMapOperation; import software.amazon.lambda.durable.operation.DurableMapOperation.MapConfig; -import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableMapOperation.MapItemContext; import software.amazon.lambda.durable.operation.DurableWaitOperation; import software.amazon.lambda.durable.serde.JacksonSerDes; @@ -38,14 +40,13 @@ public String handleRequest(GreetingRequest input) { var names = List.of(name, name.toUpperCase(), name.toLowerCase()); // Map over each name concurrently — each iteration runs in its own child context - var result = DurableMapOperation.map( + var result = map( "greet-all", names, String.class, item -> { - var index = DurableMapOperation.MapItemContext.getCurrentContext() - .getIndex(); - return DurableStepOperation.step("greet-" + index, String.class, () -> { + var index = MapItemContext.getCurrentContext().getIndex(); + return step("greet-" + index, String.class, () -> { throw new RuntimeException("Failure from " + item + "!"); }); }, diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/SimpleMapExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/SimpleMapExample.java index b9c1f440e..fb04f2ae6 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/SimpleMapExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/SimpleMapExample.java @@ -2,12 +2,14 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.map; +import static software.amazon.lambda.durable.operation.DurableMapOperation.map; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import java.util.List; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; -import software.amazon.lambda.durable.operation.DurableMapOperation; -import software.amazon.lambda.durable.operation.DurableStepOperation; +import software.amazon.lambda.durable.operation.DurableMapOperation.MapItemContext; /** * Example demonstrating the map operation with the Durable Execution SDK. @@ -32,9 +34,9 @@ public String handleRequest(GreetingRequest input) { var names = List.of(name, name.toUpperCase(), name.toLowerCase()); // Map over each name concurrently — each iteration runs in its own child context - var result = DurableMapOperation.map("greet-all", names, String.class, item -> { - var index = DurableMapOperation.MapItemContext.getCurrentContext().getIndex(); - return DurableStepOperation.step("greet-" + index, String.class, () -> "Hello, " + item + "!"); + var result = map("greet-all", names, String.class, item -> { + var index = MapItemContext.getCurrentContext().getIndex(); + return step("greet-" + index, String.class, () -> "Hello, " + item + "!"); }); context.getLogger().info("Map completed: allSucceeded={}, size={}", result.allSucceeded(), result.size()); diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExamples.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExamples.java index 6b238e13d..f64e73b93 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExamples.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExamples.java @@ -2,6 +2,11 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.otel; +import static software.amazon.lambda.durable.operation.DurableContextOperation.runInChildContext; +import static software.amazon.lambda.durable.operation.DurableMapOperation.map; +import static software.amazon.lambda.durable.operation.DurableParallelOperation.parallel; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import io.opentelemetry.exporter.logging.LoggingSpanExporter; import io.opentelemetry.sdk.trace.SdkTracerProvider; import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; @@ -10,10 +15,6 @@ import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; -import software.amazon.lambda.durable.operation.DurableContextOperation; -import software.amazon.lambda.durable.operation.DurableMapOperation; -import software.amazon.lambda.durable.operation.DurableParallelOperation; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.otel.InvocationOtelPlugin; /** @@ -44,11 +45,11 @@ public String handleRequest(GreetingRequest input) { context.getLogger().info("Starting OTel X-Ray map example for {}", input.getName()); var items = List.of("alpha", "beta", "gamma"); - var result = DurableMapOperation.map( + var result = map( "process-items", items, String.class, - item -> DurableStepOperation.step("transform-" + item, String.class, () -> item.toUpperCase())); + item -> step("transform-" + item, String.class, () -> item.toUpperCase())); return "Mapped " + result.succeeded().size() + " items"; } @@ -67,16 +68,16 @@ public String handleRequest(GreetingRequest input) { var context = DurableContext.getCurrentContext(); context.getLogger().info("Starting OTel X-Ray parallel example for {}", input.getName()); - var parallel = DurableParallelOperation.parallel("fan-out"); + var parallel = parallel("fan-out"); try (parallel) { parallel.branch( "branch-a", String.class, - childCtx -> DurableStepOperation.step("step-a", String.class, () -> "A: " + input.getName())); + childCtx -> step("step-a", String.class, () -> "A: " + input.getName())); parallel.branch( "branch-b", String.class, - childCtx -> DurableStepOperation.step("step-b", String.class, () -> "B: " + input.getName())); + childCtx -> step("step-b", String.class, () -> "B: " + input.getName())); } var result = parallel.get(); @@ -97,11 +98,10 @@ public String handleRequest(GreetingRequest input) { var context = DurableContext.getCurrentContext(); context.getLogger().info("Starting OTel X-Ray nested context example for {}", input.getName()); - return DurableContextOperation.runInChildContext("outer", String.class, () -> { - var intermediate = - DurableStepOperation.step("outer-step", String.class, () -> "Hello, " + input.getName()); - return DurableContextOperation.runInChildContext("inner", String.class, () -> { - return DurableStepOperation.step("deep-step", String.class, () -> intermediate.toUpperCase() + "!"); + return runInChildContext("outer", String.class, () -> { + var intermediate = step("outer-step", String.class, () -> "Hello, " + input.getName()); + return runInChildContext("inner", String.class, () -> { + return step("deep-step", String.class, () -> intermediate.toUpperCase() + "!"); }); }); } diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionStepExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionStepExample.java index 2efffb5c5..64ce30d7f 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionStepExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionStepExample.java @@ -2,12 +2,13 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.otel; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.ExampleTemplate; import software.amazon.lambda.durable.examples.types.GreetingRequest; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.otel.ExecutionOtelPlugin; /** @@ -29,10 +30,9 @@ public String handleRequest(GreetingRequest input) { var context = DurableContext.getCurrentContext(); context.getLogger().info("Starting OTel X-Ray execution view example for {}", input.getName()); - var greeting = - DurableStepOperation.step("exec-create-greeting", String.class, () -> "Hello, " + input.getName()); + var greeting = step("exec-create-greeting", String.class, () -> "Hello, " + input.getName()); - var result = DurableStepOperation.step("exec-transform", String.class, () -> greeting.toUpperCase() + "!"); + var result = step("exec-transform", String.class, () -> greeting.toUpperCase() + "!"); context.getLogger().info("OTel X-Ray execution view example complete: {}", result); return result; diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionWaitExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionWaitExample.java index c1f2965b9..af92c9018 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionWaitExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionWaitExample.java @@ -2,13 +2,14 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.otel; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import java.time.Duration; import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.ExampleTemplate; import software.amazon.lambda.durable.examples.types.GreetingRequest; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.operation.DurableWaitOperation; import software.amazon.lambda.durable.otel.ExecutionOtelPlugin; @@ -31,12 +32,11 @@ public String handleRequest(GreetingRequest input) { var context = DurableContext.getCurrentContext(); context.getLogger().info("Starting OTel X-Ray execution view wait example for {}", input.getName()); - var before = DurableStepOperation.step("exec-before-wait", String.class, () -> "Prepared: " + input.getName()); + var before = step("exec-before-wait", String.class, () -> "Prepared: " + input.getName()); DurableWaitOperation.wait("exec-pause", Duration.ofSeconds(5)); - var after = - DurableStepOperation.step("exec-after-wait", String.class, () -> before + " | Resumed and completed"); + var after = step("exec-after-wait", String.class, () -> before + " | Resumed and completed"); context.getLogger().info("OTel X-Ray execution view wait example complete: {}", after); return after; diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayStepExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayStepExample.java index e81ca360e..1e5a0e1b4 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayStepExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayStepExample.java @@ -2,12 +2,13 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.otel; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.ExampleTemplate; import software.amazon.lambda.durable.examples.types.GreetingRequest; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.otel.InvocationOtelPlugin; /** @@ -44,9 +45,9 @@ public String handleRequest(GreetingRequest input) { var context = DurableContext.getCurrentContext(); context.getLogger().info("Starting OTel X-Ray step example for {}", input.getName()); - var greeting = DurableStepOperation.step("create-greeting", String.class, () -> "Hello, " + input.getName()); + var greeting = step("create-greeting", String.class, () -> "Hello, " + input.getName()); - var result = DurableStepOperation.step("transform", String.class, () -> greeting.toUpperCase() + "!"); + var result = step("transform", String.class, () -> greeting.toUpperCase() + "!"); context.getLogger().info("OTel X-Ray step example complete: {}", result); return result; diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayWaitExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayWaitExample.java index 6cd1eb8ac..16d72ec6c 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayWaitExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayWaitExample.java @@ -2,13 +2,14 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.otel; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import java.time.Duration; import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.ExampleTemplate; import software.amazon.lambda.durable.examples.types.GreetingRequest; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.operation.DurableWaitOperation; import software.amazon.lambda.durable.otel.InvocationOtelPlugin; @@ -57,12 +58,12 @@ public String handleRequest(GreetingRequest input) { var context = DurableContext.getCurrentContext(); context.getLogger().info("Starting OTel X-Ray wait example for {}", input.getName()); - var before = DurableStepOperation.step("before-wait", String.class, () -> "Prepared: " + input.getName()); + var before = step("before-wait", String.class, () -> "Prepared: " + input.getName()); // This wait forces Lambda to suspend and re-invoke after the duration DurableWaitOperation.wait("pause", Duration.ofSeconds(5)); - var after = DurableStepOperation.step("after-wait", String.class, () -> before + " | Resumed and completed"); + var after = step("after-wait", String.class, () -> before + " | Resumed and completed"); context.getLogger().info("OTel X-Ray wait example complete: {}", after); return after; diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/DeserializationFailedParallelExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/DeserializationFailedParallelExample.java index ba63c55c5..4b9b63c10 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/DeserializationFailedParallelExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/DeserializationFailedParallelExample.java @@ -2,6 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.parallel; +import static software.amazon.lambda.durable.operation.DurableParallelOperation.parallel; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import java.util.List; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; @@ -9,15 +12,13 @@ import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.ParallelBranchConfig; import software.amazon.lambda.durable.exception.SerDesException; -import software.amazon.lambda.durable.operation.DurableParallelOperation; import software.amazon.lambda.durable.operation.DurableParallelOperation.ParallelConfig; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.serde.JacksonSerDes; /** * Example demonstrating parallel branch execution with the Durable Execution SDK. * - *

    This handler processes a list of items concurrently using {@code DurableParallelOperation.parallel()}: + *

    This handler processes a list of items concurrently using {@code parallel()}: * *

      *
    1. Each item is processed in its own branch (child context) @@ -41,14 +42,14 @@ public String handleRequest(Input input) { var config = ParallelConfig.builder().build(); - var parallel = DurableParallelOperation.parallel("process-items", config); + var parallel = parallel("process-items", config); try (parallel) { var future = parallel.branch( "process", String.class, branchCtx -> { - return DurableStepOperation.step("transform", String.class, () -> { + return step("transform", String.class, () -> { throw new RuntimeException("Intentional failure for transform"); }); }, diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelExample.java index 6e35a7249..75bb793ae 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelExample.java @@ -2,6 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.parallel; +import static software.amazon.lambda.durable.operation.DurableParallelOperation.parallel; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import java.util.ArrayList; import java.util.List; import software.amazon.lambda.durable.DurableContext; @@ -9,14 +12,12 @@ import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.ParallelDurableFuture; import software.amazon.lambda.durable.model.ParallelResult; -import software.amazon.lambda.durable.operation.DurableParallelOperation; import software.amazon.lambda.durable.operation.DurableParallelOperation.ParallelConfig; -import software.amazon.lambda.durable.operation.DurableStepOperation; /** * Example demonstrating parallel branch execution with the Durable Execution SDK. * - *

      This handler processes a list of items concurrently using {@code DurableParallelOperation.parallel()}: + *

      This handler processes a list of items concurrently using {@code parallel()}: * *

        *
      1. Each item is processed in its own branch (child context) @@ -42,13 +43,13 @@ public Output handleRequest(Input input) { var config = ParallelConfig.builder().build(); var futures = new ArrayList>(items.size()); - var parallel = DurableParallelOperation.parallel("process-items", config); + var parallel = parallel("process-items", config); try (parallel) { for (var item : items) { var future = parallel.branch("process-" + item, String.class, branchCtx -> { branchCtx.getLogger().info("Processing item: {}", item); - return DurableStepOperation.step("transform-" + item, String.class, () -> item.toUpperCase()); + return step("transform-" + item, String.class, () -> item.toUpperCase()); }); futures.add(future); } diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelFailureToleranceExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelFailureToleranceExample.java index 844eb8516..c060a5225 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelFailureToleranceExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelFailureToleranceExample.java @@ -2,6 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.parallel; +import static software.amazon.lambda.durable.operation.DurableParallelOperation.parallel; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import java.util.ArrayList; import java.util.List; import software.amazon.lambda.durable.DurableContext; @@ -9,9 +12,7 @@ import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.model.ParallelResult; import software.amazon.lambda.durable.operation.DurableConcurrencyOperation.CompletionConfig; -import software.amazon.lambda.durable.operation.DurableParallelOperation; import software.amazon.lambda.durable.operation.DurableParallelOperation.ParallelConfig; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; import software.amazon.lambda.durable.retry.RetryStrategies; @@ -42,12 +43,12 @@ public Output handleRequest(Input input) { .build(); var futures = new ArrayList>(input.services().size()); - var parallel = DurableParallelOperation.parallel("call-services", config); + var parallel = parallel("call-services", config); try (parallel) { for (var service : input.services()) { var future = parallel.branch("call-" + service, String.class, branchCtx -> { - return DurableStepOperation.step( + return step( "invoke-" + service, String.class, () -> { diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelWithWaitExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelWithWaitExample.java index 1565a7a8a..06f6510d4 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelWithWaitExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelWithWaitExample.java @@ -2,6 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.parallel; +import static software.amazon.lambda.durable.operation.DurableParallelOperation.parallel; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import java.time.Duration; import java.util.ArrayList; import java.util.List; @@ -9,9 +12,7 @@ import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.model.ParallelResult; -import software.amazon.lambda.durable.operation.DurableParallelOperation; import software.amazon.lambda.durable.operation.DurableParallelOperation.ParallelConfig; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.operation.DurableWaitOperation; /** @@ -42,26 +43,26 @@ public Output handleRequest(Input input) { var config = ParallelConfig.builder().build(); var futures = new ArrayList>(3); - var parallel = DurableParallelOperation.parallel("notify", config); + var parallel = parallel("notify", config); try (parallel) { // Branch 1: email — no wait, deliver immediately futures.add(parallel.branch("email", String.class, ctx -> { DurableWaitOperation.wait("email-rate-limit-delay", Duration.ofSeconds(10)); - return DurableStepOperation.step("send-email", String.class, () -> "email:" + input.message()); + return step("send-email", String.class, () -> "email:" + input.message()); })); // Branch 2: SMS — wait for rate-limit window, then send futures.add(parallel.branch("sms", String.class, ctx -> { DurableWaitOperation.wait("sms-rate-limit-delay", Duration.ofSeconds(10)); - return DurableStepOperation.step("send-sms", String.class, () -> "sms:" + input.message()); + return step("send-sms", String.class, () -> "sms:" + input.message()); })); // Branch 3: push notification — wait for quiet-hours window, then send futures.add(parallel.branch("push", String.class, ctx -> { DurableWaitOperation.wait("push-quiet-delay", Duration.ofSeconds(10)); - return DurableStepOperation.step("send-push", String.class, () -> "push:" + input.message()); + return step("send-push", String.class, () -> "push:" + input.message()); })); } diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/DeserializationFailureExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/DeserializationFailureExample.java index 6e1a1636f..c71a3e3f2 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/DeserializationFailureExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/DeserializationFailureExample.java @@ -2,11 +2,12 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.step; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import java.time.Duration; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.exception.SerDesException; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; import software.amazon.lambda.durable.operation.DurableWaitOperation; import software.amazon.lambda.durable.serde.JacksonSerDes; @@ -16,7 +17,7 @@ public class DeserializationFailureExample extends DurableHandler { diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/ManyAsyncStepsExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/ManyAsyncStepsExample.java index e8bd80297..a3274094d 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/ManyAsyncStepsExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/ManyAsyncStepsExample.java @@ -2,6 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.step; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; +import static software.amazon.lambda.durable.operation.DurableStepOperation.stepAsync; + import java.time.Duration; import java.util.ArrayList; import java.util.concurrent.TimeUnit; @@ -11,7 +14,6 @@ import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.ManyAsyncStepsInput; import software.amazon.lambda.durable.examples.types.ManyAsyncStepsOutput; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.operation.DurableWaitOperation; /** @@ -40,7 +42,7 @@ public ManyAsyncStepsOutput handleRequest(ManyAsyncStepsInput input) { var futures = new ArrayList>(steps); for (var i = 0; i < steps; i++) { var index = i; - var future = DurableStepOperation.stepAsync("compute-" + i, Integer.class, () -> index * multiplier); + var future = stepAsync("compute-" + i, Integer.class, () -> index * multiplier); futures.add(future); } @@ -51,8 +53,8 @@ public ManyAsyncStepsOutput handleRequest(ManyAsyncStepsInput input) { var totalSum = results.stream().mapToInt(Integer::intValue).sum(); // checkpoint the executionTime so that we can have the same value when replay - var executionTimeMs = DurableStepOperation.step( - "execution-time", Long.class, () -> TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime)); + var executionTimeMs = + step("execution-time", Long.class, () -> TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime)); logger.info("Completed {} steps, total sum: {}, execution time: {}ms", steps, totalSum, executionTimeMs); // Wait 2 seconds to test replay diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/RetryExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/RetryExample.java index c5e59bcff..6982c870f 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/RetryExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/RetryExample.java @@ -2,12 +2,13 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.step; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import java.time.Duration; import java.time.Instant; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.lambda.durable.DurableHandler; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; import software.amazon.lambda.durable.operation.DurableWaitOperation; import software.amazon.lambda.durable.retry.RetryStrategies; @@ -32,12 +33,12 @@ public class RetryExample extends DurableHandler { @Override public String handleRequest(Object input) { // Step 1: Record start time - startTime = DurableStepOperation.step("record-start-time", Instant.class, () -> Instant.now()); + startTime = step("record-start-time", Instant.class, () -> Instant.now()); logger.info("Recorded start time: {}", startTime); // Step 2: Call that never retries (fails immediately) try { - DurableStepOperation.step( + step( "no-retry-call", Void.class, () -> { @@ -51,7 +52,7 @@ public String handleRequest(Object input) { } // Step 3: Flaky API call that succeeds after retries - var result = DurableStepOperation.step( + var result = step( "flaky-api-call", String.class, () -> { diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/RetryInProcessExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/RetryInProcessExample.java index 9fa2c1a27..76aee29f2 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/RetryInProcessExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/RetryInProcessExample.java @@ -2,13 +2,15 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.step; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; +import static software.amazon.lambda.durable.operation.DurableStepOperation.stepAsync; + import java.time.Duration; import java.util.concurrent.atomic.AtomicInteger; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.DurableHandler; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; import software.amazon.lambda.durable.retry.JitterStrategy; import software.amazon.lambda.durable.retry.RetryStrategies; @@ -35,7 +37,7 @@ public String handleRequest(Object input) { logger.info("Starting retry in-process example"); // Start async step that will fail and retry - DurableFuture asyncStep = DurableStepOperation.stepAsync( + DurableFuture asyncStep = stepAsync( "flaky-async-operation", String.class, () -> { @@ -63,7 +65,7 @@ public String handleRequest(Object input) { // Long-running synchronous step that keeps process busy // This prevents suspension during async step retries - String syncResult = DurableStepOperation.step("long-running-operation", String.class, () -> { + String syncResult = step("long-running-operation", String.class, () -> { logger.info( "Starting long-running operation (10 seconds) in thread: {}", Thread.currentThread().getName()); diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/SimpleStepExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/SimpleStepExample.java index b4fc969f0..df51b0b08 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/SimpleStepExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/SimpleStepExample.java @@ -2,9 +2,10 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.step; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; + import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; -import software.amazon.lambda.durable.operation.DurableStepOperation; /** * Simple example demonstrating basic step execution with the Durable Execution SDK. @@ -22,13 +23,13 @@ public class SimpleStepExample extends DurableHandler { @Override public String handleRequest(GreetingRequest input) { // Step 1: Create greeting - var greeting = DurableStepOperation.step("create-greeting", String.class, () -> "Hello, " + input.getName()); + var greeting = step("create-greeting", String.class, () -> "Hello, " + input.getName()); // Step 2: Transform to uppercase - var uppercase = DurableStepOperation.step("to-uppercase", String.class, () -> greeting.toUpperCase()); + var uppercase = step("to-uppercase", String.class, () -> greeting.toUpperCase()); // Step 3: Add punctuation - var result = DurableStepOperation.step("add-punctuation", String.class, () -> uppercase + "!"); + var result = step("add-punctuation", String.class, () -> uppercase + "!"); return result; } diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/vt/ManyAsyncStepsVirtualThreadPoolExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/vt/ManyAsyncStepsVirtualThreadPoolExample.java index 5eee36829..a7604f6fd 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/vt/ManyAsyncStepsVirtualThreadPoolExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/vt/ManyAsyncStepsVirtualThreadPoolExample.java @@ -2,6 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.vt; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; +import static software.amazon.lambda.durable.operation.DurableStepOperation.stepAsync; + import java.time.Duration; import java.util.ArrayList; import java.util.concurrent.Executors; @@ -13,7 +16,6 @@ import software.amazon.lambda.durable.examples.ExampleTemplate; import software.amazon.lambda.durable.examples.types.ManyAsyncStepsInput; import software.amazon.lambda.durable.examples.types.ManyAsyncStepsOutput; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.operation.DurableWaitOperation; /** @@ -43,7 +45,7 @@ public ManyAsyncStepsOutput handleRequest(ManyAsyncStepsInput input) { var futures = new ArrayList>(steps); for (var i = 0; i < steps; i++) { var index = i; - var future = DurableStepOperation.stepAsync("compute-" + i, Integer.class, () -> index * multiplier); + var future = stepAsync("compute-" + i, Integer.class, () -> index * multiplier); futures.add(future); } @@ -54,8 +56,8 @@ public ManyAsyncStepsOutput handleRequest(ManyAsyncStepsInput input) { var totalSum = results.stream().mapToInt(Integer::intValue).sum(); // checkpoint the executionTime so that we can have the same value when replay - var executionTimeMs = DurableStepOperation.step( - "execution-time", Long.class, () -> TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime)); + var executionTimeMs = + step("execution-time", Long.class, () -> TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime)); logger.info("Completed {} steps, total sum: {}, execution time: {}ms", steps, totalSum, executionTimeMs); // Wait 2 seconds to test replay diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/ConcurrentWaitForConditionExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/ConcurrentWaitForConditionExample.java index 6cab805ad..e1e779b56 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/ConcurrentWaitForConditionExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/ConcurrentWaitForConditionExample.java @@ -2,11 +2,13 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.wait; +import static software.amazon.lambda.durable.operation.DurableMapOperation.map; +import static software.amazon.lambda.durable.operation.DurableWaitForConditionOperation.waitForCondition; + import java.util.stream.IntStream; import software.amazon.lambda.durable.DurableHandler; -import software.amazon.lambda.durable.operation.DurableMapOperation; import software.amazon.lambda.durable.operation.DurableMapOperation.MapConfig; -import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation; +import software.amazon.lambda.durable.operation.DurableMapOperation.MapItemContext; import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation.WaitForConditionConfig; import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation.WaitForConditionResult; @@ -30,18 +32,17 @@ public String handleRequest(Input input) { var config = MapConfig.builder().maxConcurrency(input.maxConcurrency()).build(); - var result = DurableMapOperation.map( + var result = map( "concurrent-wait-for-conditions", items, String.class, item -> { - var index = DurableMapOperation.MapItemContext.getCurrentContext() - .getIndex(); + var index = MapItemContext.getCurrentContext().getIndex(); var conditionConfig = WaitForConditionConfig.builder() .initialState(1) .build(); // Poll until the counter reaches the input threshold - var count = DurableWaitForConditionOperation.waitForCondition( + var count = waitForCondition( "condition-" + index, Integer.class, callCount -> { diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAsyncExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAsyncExample.java index 2ed1d8daf..db131e0e7 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAsyncExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAsyncExample.java @@ -2,13 +2,14 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.wait; +import static software.amazon.lambda.durable.operation.DurableStepOperation.stepAsync; +import static software.amazon.lambda.durable.operation.DurableWaitOperation.waitAsync; + import java.time.Duration; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; -import software.amazon.lambda.durable.operation.DurableStepOperation; -import software.amazon.lambda.durable.operation.DurableWaitOperation; /** * Example demonstrating non-blocking wait with waitAsync(). @@ -31,11 +32,10 @@ public String handleRequest(GreetingRequest input) { context.getLogger().info("Starting waitAsync example for {}", input.getName()); // Start a non-blocking wait — returns immediately - DurableFuture waitFuture = DurableWaitOperation.waitAsync("min-delay", Duration.ofSeconds(5)); + DurableFuture waitFuture = waitAsync("min-delay", Duration.ofSeconds(5)); // Run a step concurrently while the wait timer is ticking - DurableFuture stepFuture = - DurableStepOperation.stepAsync("process", String.class, () -> "Processed: " + input.getName()); + DurableFuture stepFuture = stepAsync("process", String.class, () -> "Processed: " + input.getName()); // Block until both complete — guarantees at least 5 seconds elapsed waitFuture.get(); diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAtLeastExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAtLeastExample.java index ab723885a..509ef1add 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAtLeastExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAtLeastExample.java @@ -2,25 +2,26 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.wait; +import static software.amazon.lambda.durable.operation.DurableStepOperation.stepAsync; + import java.time.Duration; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; import software.amazon.lambda.durable.operation.DurableWaitOperation; import software.amazon.lambda.durable.retry.RetryStrategies; /** - * Example demonstrating concurrent stepAsync() with wait() operations. + * Example demonstrating concurrent stepAsync() with DurableWaitOperation.wait() operations. * *

        This example shows suspension behavior with pending async steps: * *

          *
        • stepAsync() starts a background operation (takes 2 seconds) - *
        • wait() is called immediately (3 second duration) + *
        • DurableWaitOperation.wait() is called immediately (3 second duration) *
        • The step completes successfully before suspension *
        • Execution suspends for the wait time *
        @@ -34,7 +35,7 @@ public String handleRequest(GreetingRequest input) { logger.info("Starting concurrent step + wait example for: {}", input.getName()); // Start an async step that takes 2 seconds - DurableFuture asyncStep = DurableStepOperation.stepAsync( + DurableFuture asyncStep = stepAsync( "async-operation", String.class, () -> { diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAtLeastInProcessExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAtLeastInProcessExample.java index e6174f328..67da484d2 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAtLeastInProcessExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAtLeastInProcessExample.java @@ -2,25 +2,26 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.wait; +import static software.amazon.lambda.durable.operation.DurableStepOperation.stepAsync; + import java.time.Duration; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.operation.DurableStepOperation.StepConfig; import software.amazon.lambda.durable.operation.DurableWaitOperation; import software.amazon.lambda.durable.retry.RetryStrategies; /** - * Example demonstrating concurrent stepAsync() with wait() operations where no suspension occurs. + * Example demonstrating concurrent stepAsync() with DurableWaitOperation.wait() operations where no suspension occurs. * *

        This example shows in-process wait behavior: * *

          *
        • stepAsync() starts a background operation (takes 10 seconds) - *
        • wait() is called immediately (3 second duration) + *
        • DurableWaitOperation.wait() is called immediately (3 second duration) *
        • The async step takes longer than the wait duration *
        • No suspension occurs because we've already waited long enough *
        @@ -34,7 +35,7 @@ public String handleRequest(GreetingRequest input) { logger.info("Starting concurrent step + wait example for: {}", input.getName()); // Start an async step that takes 10 seconds - DurableFuture asyncStep = DurableStepOperation.stepAsync( + DurableFuture asyncStep = stepAsync( "async-operation", String.class, () -> { diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitExample.java index 0e7796d90..a91cea29b 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitExample.java @@ -2,12 +2,14 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.wait; +import static software.amazon.lambda.durable.operation.DurableContextOperation.runInChildContextAsync; +import static software.amazon.lambda.durable.operation.DurableStepOperation.step; +import static software.amazon.lambda.durable.operation.DurableStepOperation.stepAsync; + import java.time.Duration; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; -import software.amazon.lambda.durable.operation.DurableContextOperation; -import software.amazon.lambda.durable.operation.DurableStepOperation; import software.amazon.lambda.durable.operation.DurableWaitOperation; /** @@ -28,14 +30,13 @@ public class WaitExample extends DurableHandler { @Override public String handleRequest(GreetingRequest input) { // Step 1: Start processing - var started = DurableStepOperation.step( - "start-processing", String.class, () -> "Started processing for " + input.getName()); + var started = step("start-processing", String.class, () -> "Started processing for " + input.getName()); // Wait 10 seconds DurableWaitOperation.wait(null, Duration.ofSeconds(10)); // Step 2: Continue processing - var continued = DurableStepOperation.stepAsync("continue-processing", String.class, () -> { + var continued = stepAsync("continue-processing", String.class, () -> { try { Thread.sleep(10000); } catch (InterruptedException e) { @@ -45,7 +46,7 @@ public String handleRequest(GreetingRequest input) { }); // Wait at most seconds - var wait5seconds = DurableContextOperation.runInChildContextAsync("wait-5-seconds", String.class, () -> { + var wait5seconds = runInChildContextAsync("wait-5-seconds", String.class, () -> { DurableWaitOperation.wait("wait-5-seconds", Duration.ofSeconds(5)); return started + " - waited 5 seconds"; @@ -54,8 +55,7 @@ public String handleRequest(GreetingRequest input) { var step2 = DurableFuture.anyOf(continued, wait5seconds); // Step 3: Complete - var result = DurableStepOperation.step( - "complete-processing", String.class, () -> step2 + " - completed after 5s more"); + var result = step("complete-processing", String.class, () -> step2 + " - completed after 5s more"); return result; } diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitForConditionExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitForConditionExample.java index 127c548bd..16cbc15ff 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitForConditionExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitForConditionExample.java @@ -2,8 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.wait; +import static software.amazon.lambda.durable.operation.DurableWaitForConditionOperation.waitForCondition; + import software.amazon.lambda.durable.DurableHandler; -import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation; import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation.WaitForConditionConfig; import software.amazon.lambda.durable.operation.DurableWaitForConditionOperation.WaitForConditionResult; @@ -22,7 +23,7 @@ public class WaitForConditionExample extends DurableHandler { @Override public Integer handleRequest(Integer threshold) { // Poll until the counter reaches the input threshold - return DurableWaitForConditionOperation.waitForCondition( + return waitForCondition( "wait-for-condition", Integer.class, callCount -> { diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java index d8e98bf56..26e8c5584 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ExtensionOperationIntegrationTest.java @@ -4,6 +4,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -91,7 +92,7 @@ void staticOperationsUseCurrentContextAndRejectStepThreads() { assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); assertTrue(result.getResult(String.class).contains("step thread")); - assertThrows(IllegalStateException.class, DurableContext::getCurrentContext); + assertNull(DurableContext.getCurrentContext()); } @Test diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionIntegrationTest.java index ae4d9bfdd..8a5a84baa 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/WaitForConditionIntegrationTest.java @@ -8,6 +8,7 @@ import java.util.ArrayList; import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.RepeatedTest; import org.junit.jupiter.api.Test; import software.amazon.lambda.durable.config.WaitForConditionConfig; @@ -16,9 +17,26 @@ import software.amazon.lambda.durable.retry.JitterStrategy; import software.amazon.lambda.durable.retry.WaitForConditionWaitStrategy; import software.amazon.lambda.durable.retry.WaitStrategies; +import software.amazon.lambda.durable.serde.JacksonSerDes; +import software.amazon.lambda.durable.serde.SerDes; import software.amazon.lambda.durable.testing.LocalDurableTestRunner; class WaitForConditionIntegrationTest { + private static final class NormalizingStringSerDes implements SerDes { + private final SerDes delegate = new JacksonSerDes(); + + @Override + public String serialize(Object value) { + return delegate.serialize(value); + } + + @Override + @SuppressWarnings("unchecked") + public T deserialize(String data, TypeToken typeToken) { + var value = delegate.deserialize(data, typeToken); + return "raw".equals(value) ? (T) "normalized" : value; + } + } // ---- Basic integration tests ---- @@ -89,6 +107,34 @@ void testCustomWaitStrategy() { assertEquals("done", result.getResult(String.class)); } + @Test + void testWaitStrategyReceivesSerDesNormalizedState() { + var strategyState = new AtomicReference(); + var runner = LocalDurableTestRunner.create(String.class, (input, ctx) -> { + var config = WaitForConditionConfig.builder() + .serDes(new NormalizingStringSerDes()) + .waitStrategy((state, attempt) -> { + strategyState.set(state); + return Duration.ofSeconds(1); + }) + .build(); + + return ctx.waitForCondition( + "normalized-strategy-state", + String.class, + (state, stepCtx) -> state == null + ? WaitForConditionResult.continuePolling("raw") + : WaitForConditionResult.stopPolling("done"), + config); + }); + + var result = runner.runUntilComplete("test"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("done", result.getResult(String.class)); + assertEquals("normalized", strategyState.get()); + } + @Test void testMaxAttemptsExceeded() { var runner = LocalDurableTestRunner.create(String.class, (input, ctx) -> { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java index 6e2a5c52e..f54958e9f 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java @@ -33,8 +33,8 @@ public interface DurableContext extends BaseContext { /** * Returns the durable context attached to the current SDK-managed context thread. * - * @return the current durable context - * @throws IllegalStateException if called outside a durable context or from a step thread + * @return the current durable context, or {@code null} when no SDK context is active + * @throws IllegalStateException if called from a step thread */ static DurableContext getCurrentContext() { var context = BaseContext.getCurrentContext(); @@ -42,12 +42,26 @@ static DurableContext getCurrentContext() { return durableContext; } if (context == null) { - throw new IllegalStateException("No DurableContext is active on the current thread"); + return null; } throw new IllegalStateException( "DurableContext is not available from a step thread; use StepContext.getCurrentContext() instead"); } + /** + * Requires the durable context attached to the current SDK-managed context thread. + * + * @return the current durable context + * @throws IllegalStateException if called outside a durable context or from a step thread + */ + static DurableContext requireCurrentContext() { + var context = getCurrentContext(); + if (context == null) { + throw new IllegalStateException("No DurableContext is active on the current thread"); + } + return context; + } + /** Returns whether this context is currently replaying checkpointed durable operations. */ boolean isReplaying(); @@ -170,7 +184,10 @@ default DurableFuture stepAsync( Objects.requireNonNull(func, "func cannot be null"); try (var ignored = BaseContextImpl.attachCurrentContext(this)) { return DurableStepOperation.stepAsync( - name, resultType, () -> func.apply(StepContext.getCurrentContext()), config.toOperationConfig()); + name, + resultType, + () -> func.apply(StepContext.requireCurrentContext()), + config.toOperationConfig()); } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/StepContext.java b/sdk/src/main/java/software/amazon/lambda/durable/StepContext.java index 261b53a9a..c51fb7cfa 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/StepContext.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/StepContext.java @@ -8,16 +8,35 @@ public interface StepContext extends BaseContext { /** Returns the current retry attempt number (1-based). */ int getAttempt(); - /** Returns the step context attached to the current SDK-managed thread. */ + /** + * Returns the step context attached to the current SDK-managed thread. + * + * @return the current step context, or {@code null} when no SDK context is active + * @throws IllegalStateException if called from a durable context thread + */ static StepContext getCurrentContext() { var context = BaseContext.getCurrentContext(); if (context instanceof StepContext stepContext) { return stepContext; } if (context == null) { - throw new IllegalStateException("No StepContext is active on the current thread"); + return null; } throw new IllegalStateException( "StepContext is not available from a durable context thread; use DurableContext.getCurrentContext() instead"); } + + /** + * Requires the step context attached to the current SDK-managed thread. + * + * @return the current step context + * @throws IllegalStateException if called outside a step thread or from a durable context thread + */ + static StepContext requireCurrentContext() { + var context = getCurrentContext(); + if (context == null) { + throw new IllegalStateException("No StepContext is active on the current thread"); + } + return context; + } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepResult.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepResult.java index e144295ad..dac437518 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepResult.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionStepResult.java @@ -4,13 +4,15 @@ import java.time.Duration; import java.util.Objects; +import java.util.function.Function; /** * Fixed outcomes supported by a stateful extension STEP primitive. * * @param the checkpointed state and final result type */ -public sealed interface ExtensionStepResult permits ExtensionStepResult.Succeeded, ExtensionStepResult.Retry { +public sealed interface ExtensionStepResult + permits ExtensionStepResult.Succeeded, ExtensionStepResult.Retry, ExtensionStepResult.RetryAfterNormalization { /** Creates a terminal successful outcome. */ static Succeeded succeed(T value) { @@ -22,6 +24,11 @@ static Retry retry(T state, Duration delay) { return new Retry<>(state, delay); } + /** Creates a retry whose delay is evaluated from the checkpoint-normalized state. */ + static RetryAfterNormalization retryAfterNormalization(T state, Function delayStrategy) { + return new RetryAfterNormalization<>(state, delayStrategy); + } + /** Creates a decision that a failed attempt should not be retried. */ static DoNotRetry doNotRetry() { return new DoNotRetry<>(); @@ -36,13 +43,29 @@ record Succeeded(T value) implements ExtensionStepResult {} /** Retry outcome. */ record Retry(T state, Duration delay) implements ExtensionStepResult, RetryDecision { public Retry { - Objects.requireNonNull(delay, "delay cannot be null"); - if (delay.isNegative()) { - throw new IllegalArgumentException("delay cannot be negative"); - } + validateDelay(delay); + } + } + + /** Retry outcome whose delay is evaluated after the state completes its SerDes round trip. */ + record RetryAfterNormalization(T state, Function delayStrategy) implements ExtensionStepResult { + public RetryAfterNormalization { + Objects.requireNonNull(delayStrategy, "delayStrategy cannot be null"); + } + + public Duration delay(T normalizedState) { + return validateDelay(delayStrategy.apply(normalizedState)); } } /** Decision that a failed attempt should not be retried. */ record DoNotRetry() implements RetryDecision {} + + private static Duration validateDelay(Duration delay) { + Objects.requireNonNull(delay, "delay cannot be null"); + if (delay.isNegative()) { + throw new IllegalArgumentException("delay cannot be negative"); + } + return delay; + } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableContextOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableContextOperation.java index 7d0454f36..6c698f61a 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableContextOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableContextOperation.java @@ -80,7 +80,7 @@ public static DurableFuture runInChildContextAsync( RUN_IN_CHILD_CONTEXT.getValue(), resultType, () -> ExtensionContextResult.replayChildrenAboveSize( - function.apply(DurableContext.getCurrentContext()), null, LARGE_RESULT_THRESHOLD), + function.apply(DurableContext.requireCurrentContext()), null, LARGE_RESULT_THRESHOLD), ExtensionContextConfig.builder() .serDes(config.serDes()) .isVirtual(config.isVirtual()) diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java index 3b8ac4f33..21d62abf8 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java @@ -180,7 +180,7 @@ private static List> registerItem MAP_ITERATION.getValue(), resultType, () -> ExtensionContextResult.replayChildrenAboveSize( - function.apply(item, itemIndex, DurableContext.getCurrentContext()), + function.apply(item, itemIndex, DurableContext.requireCurrentContext()), null, LARGE_RESULT_THRESHOLD), iterationConfig), diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableParallelOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableParallelOperation.java index 7871ef84d..865344091 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableParallelOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableParallelOperation.java @@ -142,7 +142,7 @@ private void registerBranch(BranchDefinition definition, int index) { PARALLEL_BRANCH.getValue(), definition.resultType, () -> ExtensionContextResult.replayChildrenAboveSize( - definition.function.apply(DurableContext.getCurrentContext()), + definition.function.apply(DurableContext.requireCurrentContext()), null, LARGE_RESULT_THRESHOLD), branchConfig(definition.config)), diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperation.java index 31299d9cf..68b4a1bf6 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperation.java @@ -30,6 +30,7 @@ public final class DurableWaitForCallbackOperation { private static final String CALLBACK_SUFFIX = "-callback"; private static final String SUBMITTER_SUFFIX = "-submitter"; + private static final int LARGE_RESULT_THRESHOLD = 256 * 1024; private static final int MAX_NAME_LENGTH = ParameterValidator.MAX_OPERATION_NAME_LENGTH - Math.max(CALLBACK_SUFFIX.length(), SUBMITTER_SUFFIX.length()); @@ -116,11 +117,11 @@ private static ExtensionContextResult executeInChildContext( name + SUBMITTER_SUFFIX, Void.class, () -> { - submitter.accept(callback.callbackId(), StepContext.getCurrentContext()); + submitter.accept(callback.callbackId(), StepContext.requireCurrentContext()); return null; }, config.stepConfig()); - return ExtensionContextResult.completed(callback.get()); + return ExtensionContextResult.replayChildrenAboveSize(callback.get(), null, LARGE_RESULT_THRESHOLD); } private static ExtensionContextConfig extensionConfig(WaitForCallbackConfig config) { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperation.java index 116e8890f..7ed460ec9 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperation.java @@ -115,14 +115,15 @@ private static ExtensionStepResult evaluate( T state, BiFunction> checkFunction, WaitForConditionConfig config) { - var stepContext = StepContext.getCurrentContext(); + var stepContext = StepContext.requireCurrentContext(); var result = Objects.requireNonNull( checkFunction.apply(state, stepContext), "waitForCondition check result cannot be null"); if (result.isDone()) { return ExtensionStepResult.succeed(result.value()); } - var delay = config.waitStrategy().evaluate(result.value(), stepContext.getAttempt()); - return ExtensionStepResult.retry(result.value(), delay); + var attempt = stepContext.getAttempt(); + return ExtensionStepResult.retryAfterNormalization( + result.value(), normalizedState -> config.waitStrategy().evaluate(normalizedState, attempt)); } private static final class WaitForConditionFuture implements DurableFuture { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWithRetryOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWithRetryOperation.java index af021a94e..61e47d765 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWithRetryOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableWithRetryOperation.java @@ -25,6 +25,7 @@ public final class DurableWithRetryOperation { private static final String BACKOFF_SUFFIX = "-backoff-"; private static final String ANONYMOUS_CONTEXT_NAME = "retry"; private static final String ANONYMOUS_BACKOFF_PREFIX = "retry-backoff-"; + private static final int LARGE_RESULT_THRESHOLD = 256 * 1024; private DurableWithRetryOperation() {} @@ -59,7 +60,8 @@ public static DurableFuture withRetryAsync( .runInChildContextAsync( OperationSubType.WITH_RETRY.getValue(), new TypeToken() {}, - () -> ExtensionContextResult.completed(executeRetryLoop(name, operation, config)), + () -> ExtensionContextResult.replayChildrenAboveSize( + executeRetryLoop(name, operation, config), null, LARGE_RESULT_THRESHOLD), ExtensionContextConfig.builder() .isVirtual(!config.wrapInChildContext()) .build()); @@ -77,7 +79,7 @@ private static BiFunction adapt(Supplier oper private static T executeRetryLoop( String name, BiFunction operation, WithRetryConfig config) { - var durableContext = DurableContext.getCurrentContext(); + var durableContext = DurableContext.requireCurrentContext(); var extensionContext = ExtensionContext.getCurrentContext(); var attempt = 1; while (true) { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/primitive/ChildContextPrimitive.java b/sdk/src/main/java/software/amazon/lambda/durable/primitive/ChildContextPrimitive.java index 84076b1b5..7db0542c8 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/primitive/ChildContextPrimitive.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/primitive/ChildContextPrimitive.java @@ -131,9 +131,9 @@ protected void replay(Operation existing) { if (existing.contextDetails() != null && Boolean.TRUE.equals(existing.contextDetails().replayChildren())) { replayChildren.set(true); - if (extensionFunction != null) { - replayState.set( - deserializeResult(existing.contextDetails().result())); + var result = existing.contextDetails().result(); + if (extensionFunction != null && result != null && !result.isEmpty()) { + replayState.set(deserializeResult(result)); } executeChildContext(); } else { @@ -219,11 +219,10 @@ private void handleExtensionContextSuccess(ExtensionContextResult result) { var resultBytes = serializedSize(serializedResult.serialized()); if (result.shouldReplayChildren(resultBytes)) { - var serializedReplayState = serializeAndDeserializeResult(result.replayState()); cachedOperationResult.set(DeserializedOperationResult.succeeded(serializedResult.deserialized())); sendOperationUpdate(OperationUpdate.builder() .action(OperationAction.SUCCEED) - .payload(serializedReplayState.serialized()) + .payload(serializeReplayState(result.replayState())) .contextOptions( ContextOptions.builder().replayChildren(true).build())); } else { @@ -232,6 +231,12 @@ private void handleExtensionContextSuccess(ExtensionContextResult result) { } } + private String serializeReplayState(T replayState) { + return replayState == null + ? "" + : serializeAndDeserializeResult(replayState).serialized(); + } + private boolean shouldSkipCheckpoint() { return replayChildren.get() || isVirtual || parentOperation != null && parentOperation.isOperationCompleted(); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/primitive/StepPrimitive.java b/sdk/src/main/java/software/amazon/lambda/durable/primitive/StepPrimitive.java index 553bc351b..14533205d 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/primitive/StepPrimitive.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/primitive/StepPrimitive.java @@ -2,7 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.primitive; +import java.time.Duration; import java.util.concurrent.CompletableFuture; +import java.util.function.Function; import software.amazon.awssdk.services.lambda.model.ErrorObject; import software.amazon.awssdk.services.lambda.model.Operation; import software.amazon.awssdk.services.lambda.model.OperationAction; @@ -126,13 +128,23 @@ private void handleExtensionStepResult(ExtensionStepResult result, int attemp handleStepSucceeded(succeeded.value()); return; } - var retry = (ExtensionStepResult.Retry) result; - handleExtensionStepRetry(retry, null, attempt); + if (result instanceof ExtensionStepResult.Retry retry) { + handleExtensionStepRetry(retry.state(), ignored -> retry.delay(), null, attempt); + return; + } + var retry = (ExtensionStepResult.RetryAfterNormalization) result; + handleExtensionStepRetry(retry.state(), retry::delay, null, attempt); } private void handleExtensionStepRetry(ExtensionStepResult.Retry retry, ErrorObject error, int attempt) { - var serializedState = serializeAndDeserializeResult(retry.state()); - var retryDelaySeconds = Math.toIntExact(retry.delay().toSeconds()); + handleExtensionStepRetry(retry.state(), ignored -> retry.delay(), error, attempt); + } + + private void handleExtensionStepRetry( + T state, Function delayStrategy, ErrorObject error, int attempt) { + var serializedState = serializeAndDeserializeResult(state); + var delay = delayStrategy.apply(serializedState.deserialized()); + var retryDelaySeconds = Math.toIntExact(delay.toSeconds()); var update = OperationUpdate.builder() .action(OperationAction.RETRY) .payload(serializedState.serialized()) diff --git a/sdk/src/test/java/software/amazon/lambda/durable/CurrentContextTest.java b/sdk/src/test/java/software/amazon/lambda/durable/CurrentContextTest.java index 41c94c422..bb4979d9f 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/CurrentContextTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/CurrentContextTest.java @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -19,8 +20,13 @@ void clearContext() { } @Test - void durableContextFailsClearlyOutsideDurableThread() { - var exception = assertThrows(IllegalStateException.class, DurableContext::getCurrentContext); + void durableContextReturnsNullOutsideDurableThread() { + assertNull(DurableContext.getCurrentContext()); + } + + @Test + void durableContextCanBeRequiredOutsideDurableThread() { + var exception = assertThrows(IllegalStateException.class, DurableContext::requireCurrentContext); assertTrue(exception.getMessage().contains("No DurableContext")); } @@ -35,8 +41,13 @@ void durableContextFailsClearlyOnStepThread() { } @Test - void stepContextFailsClearlyOutsideStepThread() { - var exception = assertThrows(IllegalStateException.class, StepContext::getCurrentContext); + void stepContextReturnsNullOutsideStepThread() { + assertNull(StepContext.getCurrentContext()); + } + + @Test + void stepContextCanBeRequiredOutsideStepThread() { + var exception = assertThrows(IllegalStateException.class, StepContext::requireCurrentContext); assertTrue(exception.getMessage().contains("No StepContext")); } @@ -66,9 +77,11 @@ void currentContextScopesRestorePreviousContext() { try (var ignored = BaseContextImpl.attachCurrentContext(inner)) { assertSame(inner, StepContext.getCurrentContext()); + assertSame(inner, StepContext.requireCurrentContext()); } assertSame(outer, DurableContext.getCurrentContext()); + assertSame(outer, DurableContext.requireCurrentContext()); } private interface CurrentExtensionContext extends DurableContext, ExtensionContext {} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/ExtensionStepResultTest.java b/sdk/src/test/java/software/amazon/lambda/durable/ExtensionStepResultTest.java index 199e05363..8adc35ba6 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/ExtensionStepResultTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/ExtensionStepResultTest.java @@ -33,4 +33,24 @@ void retryRejectsInvalidDelay() { assertThrows(NullPointerException.class, () -> ExtensionStepResult.retry("next", null)); assertThrows(IllegalArgumentException.class, () -> ExtensionStepResult.retry("next", Duration.ofSeconds(-1))); } + + @Test + void retryAfterNormalizationEvaluatesDelayFromNormalizedState() { + var result = ExtensionStepResult.retryAfterNormalization( + "raw", state -> "normalized".equals(state) ? Duration.ofSeconds(3) : Duration.ofSeconds(1)); + + assertEquals("raw", result.state()); + assertEquals(Duration.ofSeconds(3), result.delay("normalized")); + } + + @Test + void retryAfterNormalizationRejectsInvalidStrategyOrDelay() { + assertThrows(NullPointerException.class, () -> ExtensionStepResult.retryAfterNormalization("next", null)); + assertThrows( + NullPointerException.class, () -> ExtensionStepResult.retryAfterNormalization("next", state -> null) + .delay("normalized")); + assertThrows(IllegalArgumentException.class, () -> ExtensionStepResult.retryAfterNormalization( + "next", state -> Duration.ofSeconds(-1)) + .delay("normalized")); + } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperationImplementationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperationImplementationTest.java index 04b460de8..87ed7c98a 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperationImplementationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForCallbackOperationImplementationTest.java @@ -3,8 +3,10 @@ package software.amazon.lambda.durable.operation; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; @@ -12,28 +14,40 @@ import static org.mockito.Mockito.when; import java.util.List; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; import software.amazon.awssdk.services.lambda.model.CallbackDetails; 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.DurableCallbackFuture; +import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.StepConfig; import software.amazon.lambda.durable.config.WaitForCallbackConfig; +import software.amazon.lambda.durable.context.BaseContextImpl; import software.amazon.lambda.durable.exception.CallbackTimeoutException; +import software.amazon.lambda.durable.extension.ExtensionCallbackConfig; import software.amazon.lambda.durable.extension.ExtensionChildOperationSummary; import software.amazon.lambda.durable.extension.ExtensionContext; import software.amazon.lambda.durable.extension.ExtensionContextConfig; import software.amazon.lambda.durable.extension.ExtensionContextFailure; import software.amazon.lambda.durable.extension.ExtensionContextFunction; import software.amazon.lambda.durable.extension.ExtensionOperation; +import software.amazon.lambda.durable.extension.ExtensionStepConfig; +import software.amazon.lambda.durable.extension.ExtensionStepFunction; import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.serde.JacksonSerDes; class DurableWaitForCallbackOperationImplementationTest { + @AfterEach + void clearContext() { + BaseContextImpl.setCurrentContext(null); + } + @Test void executeCreatesExistingWaitForCallbackContextTopology() { var context = mock(ExtensionContext.class); @@ -66,6 +80,61 @@ void executeCreatesExistingWaitForCallbackContextTopology() { assertSame(serDes, contextConfig.getValue().serDes()); } + @Test + void largeResultReplaysCallbackChildren() { + var context = mock(ExtensionContext.class); + var parent = mock(ExtensionOperation.class); + var resultType = TypeToken.get(String.class); + when(context.reserve("approval")).thenReturn(parent); + when(parent.runInChildContextAsync( + eq(OperationSubType.WAIT_FOR_CALLBACK.getValue()), + eq(resultType), + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) + .thenReturn(mockStringFuture()); + + DurableWaitForCallbackOperation.waitForCallbackAsync( + context, + "approval", + resultType, + (callbackId, stepContext) -> {}, + WaitForCallbackConfig.builder().build().toOperationConfig()); + + var function = extensionFunction(); + verify(parent) + .runInChildContextAsync( + eq(OperationSubType.WAIT_FOR_CALLBACK.getValue()), + eq(resultType), + function.capture(), + any(ExtensionContextConfig.class)); + var child = mock(CurrentExtensionContext.class); + var callbackReservation = mock(ExtensionOperation.class); + var submitterReservation = mock(ExtensionOperation.class); + var callback = mockStringCallback(); + var submitter = mockVoidFuture(); + var largeResult = "x".repeat(256 * 1024); + when(child.reserve("approval-callback")).thenReturn(callbackReservation); + when(child.reserve("approval-submitter")).thenReturn(submitterReservation); + when(callbackReservation.createCallback( + eq(OperationSubType.CALLBACK.getValue()), eq(resultType), any(ExtensionCallbackConfig.class))) + .thenReturn(callback); + when(callback.callbackId()).thenReturn("callback-id"); + when(callback.get()).thenReturn(largeResult); + when(submitterReservation.stepAsync( + eq(OperationSubType.STEP.getValue()), + eq(TypeToken.get(Void.class)), + any(ExtensionStepFunction.class), + any(ExtensionStepConfig.class))) + .thenReturn(submitter); + BaseContextImpl.setCurrentContext(child); + + var result = function.getValue().apply(); + + assertSame(largeResult, result.result()); + assertFalse(result.shouldReplayChildren(256 * 1024 - 1)); + assertTrue(result.shouldReplayChildren(256 * 1024)); + } + @Test void errorHandlerPreservesCallbackTimeoutException() { var context = mock(ExtensionContext.class); @@ -121,4 +190,21 @@ void errorHandlerPreservesCallbackTimeoutException() { private DurableFuture mockStringFuture() { return mock(DurableFuture.class); } + + @SuppressWarnings("unchecked") + private DurableCallbackFuture mockStringCallback() { + return mock(DurableCallbackFuture.class); + } + + @SuppressWarnings("unchecked") + private DurableFuture mockVoidFuture() { + return mock(DurableFuture.class); + } + + @SuppressWarnings({"rawtypes", "unchecked"}) + private ArgumentCaptor> extensionFunction() { + return (ArgumentCaptor) ArgumentCaptor.forClass(ExtensionContextFunction.class); + } + + private interface CurrentExtensionContext extends DurableContext, ExtensionContext {} } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperationImplementationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperationImplementationTest.java index 6d3d8d8cb..9bf4588a0 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperationImplementationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWaitForConditionOperationImplementationTest.java @@ -3,9 +3,11 @@ package software.amazon.lambda.durable.operation; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; @@ -14,6 +16,7 @@ import java.time.Duration; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicBoolean; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; @@ -51,11 +54,13 @@ void executeMapsPollingResultsToStatefulStepOutcomes() { var future = mockStringFuture(); var resultType = TypeToken.get(String.class); var serDes = new JacksonSerDes(); + var strategyCalled = new AtomicBoolean(); var config = WaitForConditionConfig.builder() .initialState("initial") .serDes(serDes) .waitStrategy((state, attempt) -> { - assertEquals("next", state); + strategyCalled.set(true); + assertEquals("normalized", state); assertEquals(2, attempt); return Duration.ofSeconds(7); }) @@ -87,9 +92,12 @@ void executeMapsPollingResultsToStatefulStepOutcomes() { when(stepContext.getAttempt()).thenReturn(2); try (var ignored = BaseContextImpl.attachCurrentContext(stepContext)) { var retry = assertInstanceOf( - ExtensionStepResult.Retry.class, function.getValue().apply("state")); + ExtensionStepResult.RetryAfterNormalization.class, + function.getValue().apply("state")); assertEquals("next", retry.state()); - assertEquals(Duration.ofSeconds(7), retry.delay()); + assertFalse(strategyCalled.get()); + assertEquals(Duration.ofSeconds(7), retry.delay("normalized")); + assertTrue(strategyCalled.get()); } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWithRetryOperationImplementationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWithRetryOperationImplementationTest.java index 566cba23c..f1c35819b 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWithRetryOperationImplementationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableWithRetryOperationImplementationTest.java @@ -5,6 +5,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; @@ -88,6 +89,8 @@ void executePreservesContextTopologyAndDurableBackoff() { var result = function.getValue().apply(); assertEquals("done", result.result()); + assertFalse(result.shouldReplayChildren(256 * 1024 - 1)); + assertTrue(result.shouldReplayChildren(256 * 1024)); assertEquals(1, attempts.get(0)); assertEquals(2, attempts.get(1)); verify(wait).waitAsync(OperationSubType.WAIT.getValue(), Duration.ofSeconds(5)); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/primitive/ChildContextPrimitiveTest.java b/sdk/src/test/java/software/amazon/lambda/durable/primitive/ChildContextPrimitiveTest.java index bd60e65ea..dae7f68da 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/primitive/ChildContextPrimitiveTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/primitive/ChildContextPrimitiveTest.java @@ -7,7 +7,9 @@ import java.util.List; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; @@ -19,6 +21,7 @@ import software.amazon.awssdk.services.lambda.model.OperationAction; import software.amazon.awssdk.services.lambda.model.OperationStatus; import software.amazon.awssdk.services.lambda.model.OperationType; +import software.amazon.awssdk.services.lambda.model.OperationUpdate; import software.amazon.awssdk.services.lambda.model.StepDetails; import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.DurableContext; @@ -33,6 +36,8 @@ import software.amazon.lambda.durable.execution.ThreadType; import software.amazon.lambda.durable.extension.ExtensionContextConfig; import software.amazon.lambda.durable.extension.ExtensionContextFailure; +import software.amazon.lambda.durable.extension.ExtensionContextFunction; +import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; import software.amazon.lambda.durable.extension.ExtensionContextResult; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; @@ -137,9 +142,14 @@ private ChildContextPrimitive createExtensionOperation(ExtensionContextC } private ChildContextPrimitive createExtensionOperation(String subType, ExtensionContextConfig config) { + return createExtensionOperation(subType, () -> ExtensionContextResult.completed("unused"), config); + } + + private ChildContextPrimitive createExtensionOperation( + String subType, ExtensionContextFunction function, ExtensionContextConfig config) { return new ChildContextPrimitive<>( new OperationIdentifier("1", "test-context", OperationType.CONTEXT, subType), - () -> ExtensionContextResult.completed("unused"), + function, TypeToken.get(String.class), config, durableContext); @@ -457,6 +467,63 @@ void replayChildrenReExecutesToReconstructResult() throws Exception { assertTrue(functionCalled.get(), "Function should be re-executed for replayChildren path"); } + @Test + void extensionReplayChildrenAcceptsLegacyEmptyResultPayload() { + when(executionManager.getOperationAndUpdateReplayState("1")) + .thenReturn(Operation.builder() + .id("1") + .name("test-context") + .type(OperationType.CONTEXT) + .subType("AcmeContext") + .status(OperationStatus.SUCCEEDED) + .contextDetails(ContextDetails.builder() + .result("") + .replayChildren(true) + .build()) + .build()); + when(executionManager.hasOperationsForContext("1")).thenReturn(false); + var replayState = new AtomicReference(); + var config = ExtensionContextConfig.builder().serDes(SERDES).build(); + var operation = createExtensionOperation( + "AcmeContext", + () -> { + var replayContext = ExtensionContextReplayContext.getCurrentContext(); + assertTrue(replayContext.isReplayingChildren()); + replayState.set(replayContext.getReplayState()); + return ExtensionContextResult.completed("reconstructed"); + }, + config); + + operation.execute(); + + assertEquals("reconstructed", operation.get()); + assertNull(replayState.get()); + } + + @Test + void extensionReplayChildrenWithoutStateCheckpointsLegacyEmptyPayload() throws Exception { + when(executionManager.getOperationAndUpdateReplayState("1")).thenReturn(null); + var successUpdate = new AtomicReference(); + var successSent = new CountDownLatch(1); + when(executionManager.sendOperationUpdate(any())).thenAnswer(invocation -> { + var update = invocation.getArgument(0); + if (update.action() == OperationAction.SUCCEED) { + successUpdate.set(update); + successSent.countDown(); + } + return CompletableFuture.completedFuture(null); + }); + var config = ExtensionContextConfig.builder().serDes(SERDES).build(); + var operation = createExtensionOperation( + "AcmeContext", () -> ExtensionContextResult.replayChildrenAboveSize("large", null, 1), config); + + operation.execute(); + + assertTrue(successSent.await(2, TimeUnit.SECONDS)); + assertEquals("", successUpdate.get().payload()); + assertTrue(successUpdate.get().contextOptions().replayChildren()); + } + // ===== Non-deterministic detection ===== /** Type mismatch during replay terminates execution. */ diff --git a/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java b/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java index e112bf29e..2799c9083 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java @@ -44,6 +44,7 @@ import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.serde.JacksonSerDes; +import software.amazon.lambda.durable.serde.SerDes; class StatefulExtensionStepPrimitiveTest { private static final String OPERATION_ID = "1"; @@ -53,6 +54,19 @@ class StatefulExtensionStepPrimitiveTest { private ExecutionManager executionManager; private DurableContextImpl durableContext; + private static final class NormalizingSerDes implements SerDes { + @Override + public String serialize(Object value) { + return "\"raw\""; + } + + @Override + @SuppressWarnings("unchecked") + public T deserialize(String data, TypeToken typeToken) { + return (T) "normalized"; + } + } + @BeforeEach void setUp() { executionManager = mock(ExecutionManager.class); @@ -202,6 +216,45 @@ void exceptionRetryWithoutStateDoesNotCheckpointPayload() throws Exception { IllegalStateException.class.getName(), retryUpdate.get().error().errorType()); } + @Test + void retryDelayUsesCheckpointNormalizedState() throws Exception { + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(null); + when(executionManager.pollForOperationUpdates(OPERATION_ID)).thenReturn(new CompletableFuture<>()); + var retryUpdate = new AtomicReference(); + var retrySent = new CountDownLatch(1); + when(executionManager.sendOperationUpdate(any())).thenAnswer(invocation -> { + var update = invocation.getArgument(0); + if (update.action() == OperationAction.RETRY) { + retryUpdate.set(update); + retrySent.countDown(); + } + return CompletableFuture.completedFuture(null); + }); + var strategyState = new AtomicReference(); + var operation = new StepPrimitive<>( + new OperationIdentifier( + OPERATION_ID, + OPERATION_NAME, + OperationType.STEP, + OperationSubType.WAIT_FOR_CONDITION.getValue()), + state -> ExtensionStepResult.retryAfterNormalization("raw", normalized -> { + strategyState.set(normalized); + return Duration.ofSeconds(7); + }), + TypeToken.get(String.class), + ExtensionStepConfig.builder() + .serDes(new NormalizingSerDes()) + .build(), + durableContext); + + operation.execute(); + + assertTrue(retrySent.await(2, TimeUnit.SECONDS)); + assertEquals("normalized", strategyState.get()); + assertEquals("\"raw\"", retryUpdate.get().payload()); + assertEquals(7, retryUpdate.get().stepOptions().nextAttemptDelaySeconds()); + } + @Test void corruptReplayStateFailsBeforeCallingFunction() { when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)) From e8b929f331ac70a9adeecfed7f740bb473c330ca Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Mon, 10 Aug 2026 16:35:13 -0700 Subject: [PATCH 36/40] Add static durable logger accessors --- docs/advanced/logging.md | 14 +++++++++----- docs/design.md | 2 +- .../operation/callback/CallbackExample.java | 6 +++--- .../callback/RetryWaitForCallbackExample.java | 5 ++--- .../callback/WaitForCallbackFailedExample.java | 5 ++--- .../operation/child/ChildContextExample.java | 13 ++++++------- .../child/ManyAsyncChildContextExample.java | 15 +++++++++------ .../child/VirtualChildContextExample.java | 13 ++++++------- .../general/CustomPollingExample.java | 5 ++--- .../operation/general/LoggingExample.java | 12 +++++------- .../operation/general/OtelExample.java | 9 ++++----- .../operation/general/PluginExample.java | 7 +++---- .../operation/map/ComplexFlatMapExample.java | 5 ++--- .../operation/map/ComplexMapExample.java | 5 ++--- .../map/DeserializationFailedMapExample.java | 7 +++---- .../operation/map/SimpleMapExample.java | 7 +++---- .../operation/otel/OtelXRayExamples.java | 11 ++++------- .../otel/OtelXRayExecutionStepExample.java | 7 +++---- .../otel/OtelXRayExecutionWaitExample.java | 7 +++---- .../operation/otel/OtelXRayStepExample.java | 7 +++---- .../operation/otel/OtelXRayWaitExample.java | 7 +++---- .../DeserializationFailedParallelExample.java | 5 ++--- .../operation/parallel/ParallelExample.java | 18 +++++++++--------- .../ParallelFailureToleranceExample.java | 18 +++++++++--------- .../parallel/ParallelWithWaitExample.java | 7 +++---- .../operation/step/ManyAsyncStepsExample.java | 9 ++++----- ...ManyAsyncStepsVirtualThreadPoolExample.java | 9 ++++----- .../operation/wait/WaitAsyncExample.java | 7 +++---- .../durable/context/BaseContextImpl.java | 4 ++-- .../lambda/durable/logging/DurableLogger.java | 14 ++++++++++++++ .../durable/logging/DurableLoggerTest.java | 15 +++++++++++++++ 31 files changed, 143 insertions(+), 132 deletions(-) diff --git a/docs/advanced/logging.md b/docs/advanced/logging.md index fe4762235..58c12a09e 100644 --- a/docs/advanced/logging.md +++ b/docs/advanced/logging.md @@ -1,24 +1,28 @@ ## Logging -The SDK provides a `DurableLogger` via `ctx.getLogger()` that automatically includes execution metadata in log entries and suppresses duplicate logs during replay. +The SDK provides a shared `DurableLogger` via `getLogger()` that automatically includes execution metadata in log entries and suppresses duplicate logs during replay. ### Basic Usage ```java +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; + @Override protected OrderResult handleRequest(Order order, DurableContext ctx) { - ctx.getLogger().info("Processing order: {}", order.getId()); + getLogger().info("Processing order: {}", order.getId()); var result = ctx.step("validate", String.class, stepCtx -> { - stepCtx.getLogger().debug("Validating order details"); + getLogger().debug("Validating order details"); return validate(order); }); - ctx.getLogger().info("Order processed successfully"); + getLogger().info("Order processed successfully"); return new OrderResult(result); } ``` +Use `getLogger(existingLogger)` to wrap an existing SLF4J logger while retaining durable execution context. + ### Log Output Logs include execution context via MDC (works with any SLF4J-compatible logging framework): @@ -59,4 +63,4 @@ protected DurableConfig createConfiguration() { .withLoggerConfig(LoggerConfig.withReplayLogging()) .build(); } -``` \ No newline at end of file +``` diff --git a/docs/design.md b/docs/design.md index 044992755..04980d02e 100644 --- a/docs/design.md +++ b/docs/design.md @@ -576,7 +576,7 @@ This is a one-way transition (REPLAY → EXECUTION, never back). `DurableLogger` **Context Flow:** 1. `DurableExecutor` or a primitive attaches the current `BaseContext` on its SDK-managed thread 2. `DurableLogger.attachContext()` derives execution, context, operation, and attempt MDC values from that scope -3. User code logs via `context.getLogger()` with the MDC values already attached +3. User code logs via `DurableLogger.getLogger()` with the MDC values already attached 4. Closing the logger scope clears MDC when the handler, step, or child-context function finishes **Log Pattern Example (Log4j2):** diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/CallbackExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/CallbackExample.java index 80eb3b177..541bf3050 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/CallbackExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/CallbackExample.java @@ -2,13 +2,13 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.callback; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableCallbackOperation.createCallback; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; import static software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation.waitForCallbackAsync; import java.time.Duration; import software.amazon.lambda.durable.DurableHandler; -import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.examples.types.ApprovalRequest; import software.amazon.lambda.durable.operation.DurableCallbackOperation.CallbackConfig; import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation.WaitForCallbackContext; @@ -52,7 +52,7 @@ public String handleRequest(ApprovalRequest input) { var preapprovalCallback = waitForCallbackAsync("preapproval", String.class, () -> { var callbackId = WaitForCallbackContext.getCurrentContext().getCallbackId(); - StepContext.getCurrentContext().getLogger().info("Sending callback {} to preapproval system", callbackId); + getLogger().info("Sending callback {} to preapproval system", callbackId); }); var callback = createCallback("approval", String.class, config); @@ -64,7 +64,7 @@ public String handleRequest(ApprovalRequest input) { var command = String.format( "aws lambda send-durable-execution-callback-success --callback-id %s --result $(echo -n '\"approved\"' | base64)", callbackId); - StepContext.getCurrentContext().getLogger().info("To complete this callback, run: {}", command); + getLogger().info("To complete this callback, run: {}", command); return null; }); diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/RetryWaitForCallbackExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/RetryWaitForCallbackExample.java index 55775581f..f5115e33e 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/RetryWaitForCallbackExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/RetryWaitForCallbackExample.java @@ -2,13 +2,13 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.callback; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; import static software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation.waitForCallback; import static software.amazon.lambda.durable.operation.DurableWithRetryOperation.withRetry; import java.time.Duration; import software.amazon.lambda.durable.DurableHandler; -import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.examples.types.ApprovalRequest; import software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation.WaitForCallbackContext; import software.amazon.lambda.durable.operation.DurableWithRetryOperation.WithRetryConfig; @@ -41,8 +41,7 @@ public String handleRequest(ApprovalRequest input) { null, () -> { var attempt = WithRetryContext.getCurrentContext().getAttempt(); - return waitForCallback("approval-" + attempt, String.class, () -> StepContext.getCurrentContext() - .getLogger() + return waitForCallback("approval-" + attempt, String.class, () -> getLogger() .info( "Attempt {}: sending callback {} to approval system", attempt, diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/WaitForCallbackFailedExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/WaitForCallbackFailedExample.java index 3bf57e916..44f57d37c 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/WaitForCallbackFailedExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/callback/WaitForCallbackFailedExample.java @@ -2,10 +2,10 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.callback; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableWaitForCallbackOperation.waitForCallback; import software.amazon.lambda.durable.DurableHandler; -import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.examples.types.ApprovalRequest; import software.amazon.lambda.durable.exception.SerDesException; @@ -26,8 +26,7 @@ public String handleRequest(ApprovalRequest input) { "preapproval", String.class, () -> { - StepContext.getCurrentContext() - .getLogger() + getLogger() .info( "Sending callback {} to preapproval system", WaitForCallbackContext.getCurrentContext() diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ChildContextExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ChildContextExample.java index 58bc9f3ef..5f412a448 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ChildContextExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ChildContextExample.java @@ -2,12 +2,12 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.child; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableContextOperation.runInChildContext; import static software.amazon.lambda.durable.operation.DurableContextOperation.runInChildContextAsync; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; import java.time.Duration; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; @@ -33,14 +33,13 @@ public class ChildContextExample extends DurableHandler @Override public String handleRequest(GreetingRequest input) { - var context = DurableContext.getCurrentContext(); var name = input.getName(); - context.getLogger().info("Starting child context workflow for {}", name); + getLogger().info("Starting child context workflow for {}", name); // Child context 1: Order validation — step + wait + step var orderFuture = runInChildContextAsync("order-validation", String.class, () -> { var prepared = step("prepare-order", String.class, () -> "Order for " + name); - DurableContext.getCurrentContext().getLogger().info("Order prepared, waiting for validation"); + getLogger().info("Order prepared, waiting for validation"); DurableWaitOperation.wait("validation-delay", Duration.ofSeconds(5)); @@ -50,7 +49,7 @@ public String handleRequest(GreetingRequest input) { // Child context 2: Inventory check — step + wait + step var inventoryFuture = runInChildContextAsync("inventory-check", String.class, () -> { var stock = step("check-stock", String.class, () -> "Stock available for " + name); - DurableContext.getCurrentContext().getLogger().info("Stock checked, waiting for confirmation"); + getLogger().info("Stock checked, waiting for confirmation"); DurableWaitOperation.wait("confirmation-delay", Duration.ofSeconds(3)); @@ -71,12 +70,12 @@ public String handleRequest(GreetingRequest input) { }); // Collect all results using allOf - context.getLogger().info("Waiting for all child contexts to complete"); + getLogger().info("Waiting for all child contexts to complete"); var results = DurableFuture.allOf(orderFuture, inventoryFuture, shippingFuture); // Combine into summary var summary = String.join(" | ", results); - context.getLogger().info("All child contexts complete: {}", summary); + getLogger().info("All child contexts complete: {}", summary); return summary; } diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ManyAsyncChildContextExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ManyAsyncChildContextExample.java index b4a6b6914..ac39ac36b 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ManyAsyncChildContextExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/ManyAsyncChildContextExample.java @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.child; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableContextOperation.runInChildContextAsync; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; @@ -9,7 +10,6 @@ import java.util.ArrayList; import java.util.concurrent.TimeUnit; import software.amazon.lambda.durable.DurableConfig; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.ManyAsyncStepsInput; @@ -34,9 +34,8 @@ public ManyAsyncStepsOutput handleRequest(ManyAsyncStepsInput input) { var startTime = System.nanoTime(); var multiplier = input.multiplier(); var steps = input.steps(); - var logger = DurableContext.getCurrentContext().getLogger(); - logger.info("Starting {} async child context with multiplier {}", steps, multiplier); + getLogger().info("Starting {} async child context with multiplier {}", steps, multiplier); // Create async steps var futures = new ArrayList>(steps); @@ -49,7 +48,7 @@ public ManyAsyncStepsOutput handleRequest(ManyAsyncStepsInput input) { futures.add(future); } - logger.info("All {} async child context created, collecting results", steps); + getLogger().info("All {} async child context created, collecting results", steps); // Collect all results using allOf var results = DurableFuture.allOf(futures); @@ -58,8 +57,12 @@ public ManyAsyncStepsOutput handleRequest(ManyAsyncStepsInput input) { // checkpoint the executionTime so that we can have the same value when replay var executionTimeMs = step("execution-time", Long.class, () -> TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime)); - logger.info( - "Completed {} child context, total sum: {}, execution time: {}ms", steps, totalSum, executionTimeMs); + getLogger() + .info( + "Completed {} child context, total sum: {}, execution time: {}ms", + steps, + totalSum, + executionTimeMs); // Wait 2 seconds to test replay DurableWaitOperation.wait("post-compute-wait", Duration.ofSeconds(2)); diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/VirtualChildContextExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/VirtualChildContextExample.java index bbc088f5c..43b1276e6 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/VirtualChildContextExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/child/VirtualChildContextExample.java @@ -2,12 +2,12 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.child; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableContextOperation.runInChildContext; import static software.amazon.lambda.durable.operation.DurableContextOperation.runInChildContextAsync; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; import java.time.Duration; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; @@ -34,9 +34,8 @@ public class VirtualChildContextExample extends DurableHandler { var prepared = step("prepare-order", String.class, () -> "Order for " + name); - DurableContext.getCurrentContext().getLogger().info("Order prepared, waiting for validation"); + getLogger().info("Order prepared, waiting for validation"); DurableWaitOperation.wait("validation-delay", Duration.ofSeconds(5)); @@ -58,7 +57,7 @@ public String handleRequest(GreetingRequest input) { String.class, () -> { var stock = step("check-stock", String.class, () -> "Stock available for " + name); - DurableContext.getCurrentContext().getLogger().info("Stock checked, waiting for confirmation"); + getLogger().info("Stock checked, waiting for confirmation"); DurableWaitOperation.wait("confirmation-delay", Duration.ofSeconds(3)); @@ -85,12 +84,12 @@ public String handleRequest(GreetingRequest input) { RunInChildContextConfig.builder().isVirtual(true).build()); // Collect all results using allOf - context.getLogger().info("Waiting for all child contexts to complete"); + getLogger().info("Waiting for all child contexts to complete"); var results = DurableFuture.allOf(orderFuture, inventoryFuture, shippingFuture); // Combine into summary var summary = String.join(" | ", results); - context.getLogger().info("All child contexts complete: {}", summary); + getLogger().info("All child contexts complete: {}", summary); return summary; } diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/CustomPollingExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/CustomPollingExample.java index 8933d4e3a..d00a0f0da 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/CustomPollingExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/CustomPollingExample.java @@ -2,12 +2,12 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.general; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableInvokeOperation.invokeAsync; import static software.amazon.lambda.durable.operation.DurableStepOperation.stepAsync; import java.time.Duration; import software.amazon.lambda.durable.DurableConfig; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; import software.amazon.lambda.durable.operation.DurableInvokeOperation.InvokeConfig; @@ -40,8 +40,7 @@ protected DurableConfig createConfiguration() { @Override public String handleRequest(GreetingRequest input) { - var context = DurableContext.getCurrentContext(); - context.getLogger().info("Starting workflow with input: {}", input); + getLogger().info("Starting workflow with input: {}", input); // Step 1: low case the input var lowered = stepAsync("validate", String.class, () -> { diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/LoggingExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/LoggingExample.java index 118b8e8dc..6a3bab4bd 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/LoggingExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/LoggingExample.java @@ -2,13 +2,12 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.general; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; -import software.amazon.lambda.durable.StepContext; import software.amazon.lambda.durable.examples.types.GreetingRequest; /** @@ -22,23 +21,22 @@ public class LoggingExample extends DurableHandler { @Override public String handleRequest(GreetingRequest input) { - var context = DurableContext.getCurrentContext(); // Log at execution level (outside any step) - context.getLogger(logger).info("Processing greeting for: {}", input.getName()); + getLogger(logger).info("Processing greeting for: {}", input.getName()); // Step 1: Create greeting - logs inside step include operation context var greeting = step("create-greeting", String.class, () -> { - StepContext.getCurrentContext().getLogger(logger).info("Creating greeting message"); + getLogger(logger).info("Creating greeting message"); return "Hello, " + input.getName(); }); // Step 2: Transform var result = step("transform", String.class, () -> { - StepContext.getCurrentContext().getLogger().info("Transforming greeting to uppercase"); + getLogger().info("Transforming greeting to uppercase"); return greeting.toUpperCase() + "!"; }); - context.getLogger().info("Completed processing, result: {}", result); + getLogger().info("Completed processing, result: {}", result); return result; } } diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/OtelExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/OtelExample.java index 2f47a65c3..093626a7e 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/OtelExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/OtelExample.java @@ -2,13 +2,13 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.general; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; import io.opentelemetry.exporter.logging.LoggingSpanExporter; import io.opentelemetry.sdk.trace.SdkTracerProvider; import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; import software.amazon.lambda.durable.DurableConfig; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; import software.amazon.lambda.durable.otel.InvocationOtelPlugin; @@ -49,18 +49,17 @@ protected DurableConfig createConfiguration() { @Override public String handleRequest(GreetingRequest input) { - var context = DurableContext.getCurrentContext(); // Log with MDC — traceId and spanId will be in the JSON output - context.getLogger().info("Starting OTel example for {}", input.getName()); + getLogger().info("Starting OTel example for {}", input.getName()); var greeting = step("create-greeting", String.class, () -> { - context.getLogger().info("Inside step — this log has trace context in MDC"); + getLogger().info("Inside step — this log has trace context in MDC"); return "Hello, " + input.getName(); }); var result = step("transform", String.class, () -> greeting.toUpperCase() + "!"); - context.getLogger().info("OTel example complete: {}", result); + getLogger().info("OTel example complete: {}", result); return result; } } diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/PluginExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/PluginExample.java index 174e3b73d..95268da20 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/PluginExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/general/PluginExample.java @@ -2,10 +2,10 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.general; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; import software.amazon.lambda.durable.DurableConfig; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; import software.amazon.lambda.durable.plugin.*; @@ -40,14 +40,13 @@ protected DurableConfig createConfiguration() { @Override public String handleRequest(GreetingRequest input) { - var context = DurableContext.getCurrentContext(); - context.getLogger().info("Starting plugin example for {}", input.getName()); + getLogger().info("Starting plugin example for {}", input.getName()); var greeting = step("create-greeting", String.class, () -> "Hello, " + input.getName()); var result = step("transform", String.class, () -> greeting.toUpperCase() + "!"); - context.getLogger().info("Plugin example complete: {}", result); + getLogger().info("Plugin example complete: {}", result); return result; } diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexFlatMapExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexFlatMapExample.java index f0314048a..188bd1cac 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexFlatMapExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexFlatMapExample.java @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.map; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableMapOperation.map; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; @@ -9,7 +10,6 @@ import java.util.List; import java.util.stream.Collectors; import java.util.stream.IntStream; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.operation.DurableConcurrencyOperation.CompletionConfig; import software.amazon.lambda.durable.operation.DurableConcurrencyOperation.NestingType; @@ -30,8 +30,7 @@ public class ComplexFlatMapExample extends DurableHandler { @Override public String handleRequest(Integer input) { - var context = DurableContext.getCurrentContext(); - context.getLogger().info("Starting complex map example with {} items", input); + getLogger().info("Starting complex map example with {} items", input); // Part 1: Concurrent map with step + wait inside each branch var orderIds = IntStream.range(1, input + 1).mapToObj(x -> "order-" + x).collect(Collectors.toList()); diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexMapExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexMapExample.java index bfdcde1df..87426910c 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexMapExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/ComplexMapExample.java @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.map; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableMapOperation.map; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; @@ -9,7 +10,6 @@ import java.util.List; import java.util.stream.Collectors; import java.util.stream.IntStream; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.operation.DurableConcurrencyOperation.CompletionConfig; import software.amazon.lambda.durable.operation.DurableMapOperation.MapConfig; @@ -29,8 +29,7 @@ public class ComplexMapExample extends DurableHandler { @Override public String handleRequest(Integer input) { - var context = DurableContext.getCurrentContext(); - context.getLogger().info("Starting complex map example with {} items", input); + getLogger().info("Starting complex map example with {} items", input); // Part 1: Concurrent map with step + wait inside each branch var orderIds = IntStream.range(1, input + 1).mapToObj(x -> "order-" + x).collect(Collectors.toList()); diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/DeserializationFailedMapExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/DeserializationFailedMapExample.java index a743fa5c0..76e192535 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/DeserializationFailedMapExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/map/DeserializationFailedMapExample.java @@ -2,12 +2,12 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.map; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableMapOperation.map; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; import java.time.Duration; import java.util.List; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.examples.types.GreetingRequest; @@ -33,9 +33,8 @@ public class DeserializationFailedMapExample extends DurableHandler { @Override public String handleRequest(GreetingRequest input) { - var context = DurableContext.getCurrentContext(); var name = input.getName(); - context.getLogger().info("Starting map example for {}", name); + getLogger().info("Starting map example for {}", name); var names = List.of(name, name.toUpperCase(), name.toLowerCase()); @@ -39,7 +38,7 @@ public String handleRequest(GreetingRequest input) { return step("greet-" + index, String.class, () -> "Hello, " + item + "!"); }); - context.getLogger().info("Map completed: allSucceeded={}, size={}", result.allSucceeded(), result.size()); + getLogger().info("Map completed: allSucceeded={}, size={}", result.allSucceeded(), result.size()); return String.join(" | ", result.results()); } diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExamples.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExamples.java index f64e73b93..2bb76b314 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExamples.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExamples.java @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.otel; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableContextOperation.runInChildContext; import static software.amazon.lambda.durable.operation.DurableMapOperation.map; import static software.amazon.lambda.durable.operation.DurableParallelOperation.parallel; @@ -12,7 +13,6 @@ import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; import java.util.List; import software.amazon.lambda.durable.DurableConfig; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; import software.amazon.lambda.durable.otel.InvocationOtelPlugin; @@ -41,8 +41,7 @@ protected DurableConfig createConfiguration() { @Override public String handleRequest(GreetingRequest input) { - var context = DurableContext.getCurrentContext(); - context.getLogger().info("Starting OTel X-Ray map example for {}", input.getName()); + getLogger().info("Starting OTel X-Ray map example for {}", input.getName()); var items = List.of("alpha", "beta", "gamma"); var result = map( @@ -65,8 +64,7 @@ protected DurableConfig createConfiguration() { @Override public String handleRequest(GreetingRequest input) { - var context = DurableContext.getCurrentContext(); - context.getLogger().info("Starting OTel X-Ray parallel example for {}", input.getName()); + getLogger().info("Starting OTel X-Ray parallel example for {}", input.getName()); var parallel = parallel("fan-out"); try (parallel) { @@ -95,8 +93,7 @@ protected DurableConfig createConfiguration() { @Override public String handleRequest(GreetingRequest input) { - var context = DurableContext.getCurrentContext(); - context.getLogger().info("Starting OTel X-Ray nested context example for {}", input.getName()); + getLogger().info("Starting OTel X-Ray nested context example for {}", input.getName()); return runInChildContext("outer", String.class, () -> { var intermediate = step("outer-step", String.class, () -> "Hello, " + input.getName()); diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionStepExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionStepExample.java index 64ce30d7f..f595b9ff5 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionStepExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionStepExample.java @@ -2,10 +2,10 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.otel; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; import software.amazon.lambda.durable.DurableConfig; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.ExampleTemplate; import software.amazon.lambda.durable.examples.types.GreetingRequest; @@ -27,14 +27,13 @@ protected DurableConfig createConfiguration() { @Override public String handleRequest(GreetingRequest input) { - var context = DurableContext.getCurrentContext(); - context.getLogger().info("Starting OTel X-Ray execution view example for {}", input.getName()); + getLogger().info("Starting OTel X-Ray execution view example for {}", input.getName()); var greeting = step("exec-create-greeting", String.class, () -> "Hello, " + input.getName()); var result = step("exec-transform", String.class, () -> greeting.toUpperCase() + "!"); - context.getLogger().info("OTel X-Ray execution view example complete: {}", result); + getLogger().info("OTel X-Ray execution view example complete: {}", result); return result; } } diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionWaitExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionWaitExample.java index af92c9018..f0d0c84dd 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionWaitExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayExecutionWaitExample.java @@ -2,11 +2,11 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.otel; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; import java.time.Duration; import software.amazon.lambda.durable.DurableConfig; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.ExampleTemplate; import software.amazon.lambda.durable.examples.types.GreetingRequest; @@ -29,8 +29,7 @@ protected DurableConfig createConfiguration() { @Override public String handleRequest(GreetingRequest input) { - var context = DurableContext.getCurrentContext(); - context.getLogger().info("Starting OTel X-Ray execution view wait example for {}", input.getName()); + getLogger().info("Starting OTel X-Ray execution view wait example for {}", input.getName()); var before = step("exec-before-wait", String.class, () -> "Prepared: " + input.getName()); @@ -38,7 +37,7 @@ public String handleRequest(GreetingRequest input) { var after = step("exec-after-wait", String.class, () -> before + " | Resumed and completed"); - context.getLogger().info("OTel X-Ray execution view wait example complete: {}", after); + getLogger().info("OTel X-Ray execution view wait example complete: {}", after); return after; } } diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayStepExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayStepExample.java index 1e5a0e1b4..1662e59e0 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayStepExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayStepExample.java @@ -2,10 +2,10 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.otel; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; import software.amazon.lambda.durable.DurableConfig; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.ExampleTemplate; import software.amazon.lambda.durable.examples.types.GreetingRequest; @@ -42,14 +42,13 @@ protected DurableConfig createConfiguration() { @Override public String handleRequest(GreetingRequest input) { - var context = DurableContext.getCurrentContext(); - context.getLogger().info("Starting OTel X-Ray step example for {}", input.getName()); + getLogger().info("Starting OTel X-Ray step example for {}", input.getName()); var greeting = step("create-greeting", String.class, () -> "Hello, " + input.getName()); var result = step("transform", String.class, () -> greeting.toUpperCase() + "!"); - context.getLogger().info("OTel X-Ray step example complete: {}", result); + getLogger().info("OTel X-Ray step example complete: {}", result); return result; } } diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayWaitExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayWaitExample.java index 16d72ec6c..ccabe5a2c 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayWaitExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/otel/OtelXRayWaitExample.java @@ -2,11 +2,11 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.otel; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; import java.time.Duration; import software.amazon.lambda.durable.DurableConfig; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.ExampleTemplate; import software.amazon.lambda.durable.examples.types.GreetingRequest; @@ -55,8 +55,7 @@ protected DurableConfig createConfiguration() { @Override public String handleRequest(GreetingRequest input) { - var context = DurableContext.getCurrentContext(); - context.getLogger().info("Starting OTel X-Ray wait example for {}", input.getName()); + getLogger().info("Starting OTel X-Ray wait example for {}", input.getName()); var before = step("before-wait", String.class, () -> "Prepared: " + input.getName()); @@ -65,7 +64,7 @@ public String handleRequest(GreetingRequest input) { var after = step("after-wait", String.class, () -> before + " | Resumed and completed"); - context.getLogger().info("OTel X-Ray wait example complete: {}", after); + getLogger().info("OTel X-Ray wait example complete: {}", after); return after; } } diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/DeserializationFailedParallelExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/DeserializationFailedParallelExample.java index 4b9b63c10..d19d74f30 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/DeserializationFailedParallelExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/DeserializationFailedParallelExample.java @@ -2,11 +2,11 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.parallel; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableParallelOperation.parallel; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; import java.util.List; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.ParallelDurableFuture; import software.amazon.lambda.durable.TypeToken; @@ -36,9 +36,8 @@ public record Input(List items) {} @Override public String handleRequest(Input input) { - var logger = DurableContext.getCurrentContext().getLogger(); var items = input.items(); - logger.info("Starting parallel processing of {} items", items.size()); + getLogger().info("Starting parallel processing of {} items", items.size()); var config = ParallelConfig.builder().build(); diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelExample.java index 75bb793ae..e29ab2a66 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelExample.java @@ -2,12 +2,12 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.parallel; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableParallelOperation.parallel; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; import java.util.ArrayList; import java.util.List; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.ParallelDurableFuture; @@ -36,9 +36,8 @@ public record Output(List results, int totalProcessed) {} @Override public Output handleRequest(Input input) { - var logger = DurableContext.getCurrentContext().getLogger(); var items = input.items(); - logger.info("Starting parallel processing of {} items", items.size()); + getLogger().info("Starting parallel processing of {} items", items.size()); var config = ParallelConfig.builder().build(); @@ -48,7 +47,7 @@ public Output handleRequest(Input input) { try (parallel) { for (var item : items) { var future = parallel.branch("process-" + item, String.class, branchCtx -> { - branchCtx.getLogger().info("Processing item: {}", item); + getLogger().info("Processing item: {}", item); return step("transform-" + item, String.class, () -> item.toUpperCase()); }); futures.add(future); @@ -56,11 +55,12 @@ public Output handleRequest(Input input) { } // join() called here via AutoCloseable ParallelResult parallelResult = parallel.get(); - logger.info( - "Parallel complete: total={}, succeeded={}, failed={}", - parallelResult.size(), - parallelResult.succeeded(), - parallelResult.failed()); + getLogger() + .info( + "Parallel complete: total={}, succeeded={}, failed={}", + parallelResult.size(), + parallelResult.succeeded(), + parallelResult.failed()); var results = futures.stream().map(DurableFuture::get).toList(); diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelFailureToleranceExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelFailureToleranceExample.java index c060a5225..1f293693d 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelFailureToleranceExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelFailureToleranceExample.java @@ -2,12 +2,12 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.parallel; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableParallelOperation.parallel; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; import java.util.ArrayList; import java.util.List; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.model.ParallelResult; @@ -35,8 +35,7 @@ public record Output(int succeeded, int failed) {} @Override public Output handleRequest(Input input) { - var logger = DurableContext.getCurrentContext().getLogger(); - logger.info("Starting parallel execution with toleratedFailureCount={}", input.toleratedFailures()); + getLogger().info("Starting parallel execution with toleratedFailureCount={}", input.toleratedFailures()); var config = ParallelConfig.builder() .completionConfig(new CompletionConfig(input.minSuccessful, input.toleratedFailures, null)) @@ -66,16 +65,17 @@ public Output handleRequest(Input input) { } ParallelResult parallelResult = parallel.get(); - logger.info( - "Parallel complete: succeeded={}, failed={}, status={}", - parallelResult.succeeded(), - parallelResult.failed(), - parallelResult.completionStatus().isSucceeded() ? "succeeded" : "failed"); + getLogger() + .info( + "Parallel complete: succeeded={}, failed={}, status={}", + parallelResult.succeeded(), + parallelResult.failed(), + parallelResult.completionStatus().isSucceeded() ? "succeeded" : "failed"); var succeeded = parallelResult.succeeded(); var failed = parallelResult.failed(); - logger.info("Completed: {} succeeded, {} failed", succeeded, failed); + getLogger().info("Completed: {} succeeded, {} failed", succeeded, failed); return new Output(succeeded, failed); } } diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelWithWaitExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelWithWaitExample.java index 06f6510d4..ab755c79c 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelWithWaitExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/parallel/ParallelWithWaitExample.java @@ -2,13 +2,13 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.parallel; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableParallelOperation.parallel; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; import java.time.Duration; import java.util.ArrayList; import java.util.List; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.model.ParallelResult; @@ -38,8 +38,7 @@ public record Output(List deliveries, int success, int faiure) {} @Override public Output handleRequest(Input input) { - var logger = DurableContext.getCurrentContext().getLogger(); - logger.info("Sending notifications to user {}", input.userId()); + getLogger().info("Sending notifications to user {}", input.userId()); var config = ParallelConfig.builder().build(); var futures = new ArrayList>(3); @@ -69,7 +68,7 @@ public Output handleRequest(Input input) { ParallelResult result = parallel.get(); var deliveries = futures.stream().map(DurableFuture::get).toList(); - logger.info("All {} notifications delivered", deliveries.size()); + getLogger().info("All {} notifications delivered", deliveries.size()); // Test replay DurableWaitOperation.wait("wait for finalization", Duration.ofSeconds(5)); return new Output(deliveries, result.succeeded(), result.failed()); diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/ManyAsyncStepsExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/ManyAsyncStepsExample.java index a3274094d..d264349e4 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/ManyAsyncStepsExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/step/ManyAsyncStepsExample.java @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.step; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; import static software.amazon.lambda.durable.operation.DurableStepOperation.stepAsync; @@ -9,7 +10,6 @@ import java.util.ArrayList; import java.util.concurrent.TimeUnit; import software.amazon.lambda.durable.DurableConfig; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.ManyAsyncStepsInput; @@ -34,9 +34,8 @@ public ManyAsyncStepsOutput handleRequest(ManyAsyncStepsInput input) { var startTime = System.nanoTime(); var multiplier = input.multiplier(); var steps = input.steps(); - var logger = DurableContext.getCurrentContext().getLogger(); - logger.info("Starting {} async steps with multiplier {}", steps, multiplier); + getLogger().info("Starting {} async steps with multiplier {}", steps, multiplier); // Create async steps var futures = new ArrayList>(steps); @@ -46,7 +45,7 @@ public ManyAsyncStepsOutput handleRequest(ManyAsyncStepsInput input) { futures.add(future); } - logger.info("All {} async steps created, collecting results", steps); + getLogger().info("All {} async steps created, collecting results", steps); // Collect all results using allOf var results = DurableFuture.allOf(futures); @@ -55,7 +54,7 @@ public ManyAsyncStepsOutput handleRequest(ManyAsyncStepsInput input) { // checkpoint the executionTime so that we can have the same value when replay var executionTimeMs = step("execution-time", Long.class, () -> TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime)); - logger.info("Completed {} steps, total sum: {}, execution time: {}ms", steps, totalSum, executionTimeMs); + getLogger().info("Completed {} steps, total sum: {}, execution time: {}ms", steps, totalSum, executionTimeMs); // Wait 2 seconds to test replay DurableWaitOperation.wait("post-compute-wait", Duration.ofSeconds(2)); diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/vt/ManyAsyncStepsVirtualThreadPoolExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/vt/ManyAsyncStepsVirtualThreadPoolExample.java index a7604f6fd..e9cfbb84e 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/vt/ManyAsyncStepsVirtualThreadPoolExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/vt/ManyAsyncStepsVirtualThreadPoolExample.java @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.vt; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableStepOperation.step; import static software.amazon.lambda.durable.operation.DurableStepOperation.stepAsync; @@ -10,7 +11,6 @@ import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import software.amazon.lambda.durable.DurableConfig; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.ExampleTemplate; @@ -37,9 +37,8 @@ public ManyAsyncStepsOutput handleRequest(ManyAsyncStepsInput input) { var startTime = System.nanoTime(); var multiplier = input.multiplier(); var steps = input.steps(); - var logger = DurableContext.getCurrentContext().getLogger(); - logger.info("Starting {} async steps with multiplier {}", steps, multiplier); + getLogger().info("Starting {} async steps with multiplier {}", steps, multiplier); // Create async steps var futures = new ArrayList>(steps); @@ -49,7 +48,7 @@ public ManyAsyncStepsOutput handleRequest(ManyAsyncStepsInput input) { futures.add(future); } - logger.info("All {} async steps created, collecting results", steps); + getLogger().info("All {} async steps created, collecting results", steps); // Collect all results using allOf var results = DurableFuture.allOf(futures); @@ -58,7 +57,7 @@ public ManyAsyncStepsOutput handleRequest(ManyAsyncStepsInput input) { // checkpoint the executionTime so that we can have the same value when replay var executionTimeMs = step("execution-time", Long.class, () -> TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime)); - logger.info("Completed {} steps, total sum: {}, execution time: {}ms", steps, totalSum, executionTimeMs); + getLogger().info("Completed {} steps, total sum: {}, execution time: {}ms", steps, totalSum, executionTimeMs); // Wait 2 seconds to test replay DurableWaitOperation.wait("post-compute-wait", Duration.ofSeconds(2)); diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAsyncExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAsyncExample.java index db131e0e7..0b6f29ec5 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAsyncExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/operation/wait/WaitAsyncExample.java @@ -2,11 +2,11 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.examples.operation.wait; +import static software.amazon.lambda.durable.logging.DurableLogger.getLogger; import static software.amazon.lambda.durable.operation.DurableStepOperation.stepAsync; import static software.amazon.lambda.durable.operation.DurableWaitOperation.waitAsync; import java.time.Duration; -import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.DurableHandler; import software.amazon.lambda.durable.examples.types.GreetingRequest; @@ -28,8 +28,7 @@ public class WaitAsyncExample extends DurableHandler { @Override public String handleRequest(GreetingRequest input) { - var context = DurableContext.getCurrentContext(); - context.getLogger().info("Starting waitAsync example for {}", input.getName()); + getLogger().info("Starting waitAsync example for {}", input.getName()); // Start a non-blocking wait — returns immediately DurableFuture waitFuture = waitAsync("min-delay", Duration.ofSeconds(5)); @@ -41,7 +40,7 @@ public String handleRequest(GreetingRequest input) { waitFuture.get(); var result = stepFuture.get(); - context.getLogger().info("Both wait and step complete: {}", result); + getLogger().info("Both wait and step complete: {}", result); return result; } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/BaseContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/BaseContextImpl.java index 1753f590e..a8ae159d1 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/BaseContextImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/BaseContextImpl.java @@ -99,12 +99,12 @@ public ExecutionManager getExecutionManager() { /** Returns a durable logger for this context. */ public DurableLogger getLogger() { - return DurableLogger.INSTANCE; + return DurableLogger.getLogger(); } /** Returns a durable logger for this context. */ public DurableLogger getLogger(Logger delegate) { - return new DurableLogger(delegate); + return DurableLogger.getLogger(delegate); } public static void setCurrentContext(BaseContext context) { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/logging/DurableLogger.java b/sdk/src/main/java/software/amazon/lambda/durable/logging/DurableLogger.java index 888c4e248..057ea2892 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/logging/DurableLogger.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/logging/DurableLogger.java @@ -38,6 +38,20 @@ public DurableLogger(Logger delegate) { this.delegate = delegate; } + /** Returns the shared context-aware durable logger. */ + public static DurableLogger getLogger() { + return INSTANCE; + } + + /** + * Returns a context-aware durable logger wrapping the given SLF4J logger. + * + * @param delegate the SLF4J logger to wrap + */ + public static DurableLogger getLogger(Logger delegate) { + return new DurableLogger(delegate); + } + public static SafeCloseable attachContext() { var context = BaseContext.getCurrentContext(); if (context != null) { diff --git a/sdk/src/test/java/software/amazon/lambda/durable/logging/DurableLoggerTest.java b/sdk/src/test/java/software/amazon/lambda/durable/logging/DurableLoggerTest.java index f2ffe85e0..f48fe20aa 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/logging/DurableLoggerTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/logging/DurableLoggerTest.java @@ -48,6 +48,21 @@ void tearDown() throws ReflectiveOperationException { setMdcAdapter(originalMdcAdapter); } + @Test + void getsSharedLogger() { + assertSame(DurableLogger.INSTANCE, DurableLogger.getLogger()); + } + + @Test + void wrapsProvidedLogger() { + var recordingLogger = new RecordingLogger(); + + DurableLogger.getLogger(recordingLogger.delegate()).info("test message"); + + assertEquals(1, recordingLogger.calls().size()); + assertEquals("test message", recordingLogger.calls().get(0).message()); + } + @Test void logsWhenNotReplaying() { var recordingLogger = new RecordingLogger(); From 370eba1b35454cb04051d412067e8b5770e76078 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Mon, 10 Aug 2026 17:52:23 -0700 Subject: [PATCH 37/40] Fix extension replay validation and retry polling --- docs/adr/006-custom-extension-operations.md | 10 ++- docs/advanced/extensions.md | 14 +++- docs/design.md | 7 +- .../lambda/durable/MapIntegrationTest.java | 47 +++++++++++ .../extension/ExtensionContextConfig.java | 15 +++- .../ExtensionContextReplayContext.java | 16 +++- .../operation/DurableMapOperation.java | 84 ++++++++++++++++--- .../primitive/ChildContextPrimitive.java | 42 ++++++---- .../durable/primitive/StepPrimitive.java | 25 ++++-- .../config/ExtensionContextConfigTest.java | 3 + ...DurableMapOperationImplementationTest.java | 1 + .../primitive/ChildContextPrimitiveTest.java | 38 +++++++++ .../StatefulExtensionStepPrimitiveTest.java | 40 ++++++++- 13 files changed, 295 insertions(+), 47 deletions(-) diff --git a/docs/adr/006-custom-extension-operations.md b/docs/adr/006-custom-extension-operations.md index b6aa1827b..3144c01f7 100644 --- a/docs/adr/006-custom-extension-operations.md +++ b/docs/adr/006-custom-extension-operations.md @@ -246,8 +246,9 @@ A stateful extension STEP may return only: The fixed-delay retry checkpoints the supplied state and delay. The normalization-aware retry first serializes and deserializes the state with the configured SerDes, evaluates the delay strategy against that normalized state, then checkpoints the normalized state and selected delay. This keeps first execution and replay behavior consistent for -normalizing serializers. The SDK maps both outcomes onto the fixed STEP lifecycle. Thrown exceptions follow the -normal STEP failure path. Attempt metadata remains available through `StepContext`. +normalizing serializers. The SDK maps both outcomes onto the fixed STEP lifecycle and polls at the computed +retry-ready timestamp. A replayed `PENDING` extension step polls using its checkpointed `nextAttemptTimestamp`. +Thrown exceptions follow the normal STEP failure path. Attempt metadata remains available through `StepContext`. ### Support Context Replay State @@ -274,6 +275,11 @@ behavior: - context failure translation - whether the framework function emits user-function plugin events - whether late child checkpoints are suppressed after parent completion +- whether completed contexts with `replayChildren=false` re-enter the framework function for validation-only replay + +During validation-only replay, `ExtensionContextReplayContext.isValidatingReplay()` is true and `getReplayState()` +contains the checkpointed result. The framework can recreate deterministic child reservations to validate operation +identity without replacing the completed parent checkpoint. Existing configuration classes are not changed. diff --git a/docs/advanced/extensions.md b/docs/advanced/extensions.md index ca7113171..94bda950e 100644 --- a/docs/advanced/extensions.md +++ b/docs/advanced/extensions.md @@ -295,6 +295,9 @@ strategy against that checkpoint-normalized state. Both retry forms checkpoint t remains available through `StepContext.requireCurrentContext()`. Thrown exceptions follow the normal STEP failure path. +After checkpointing a retry, the SDK polls at the computed retry-ready timestamp. When replaying a `PENDING` extension +step, it uses the checkpointed `nextAttemptTimestamp`, preserving the same scheduling behavior across invocations. + `ExtensionStepConfig` owns its exception retry strategy. It returns the fixed-delay retry outcome or a do-not-retry decision, allowing extension libraries to configure exception retries and delivery semantics without depending on the customer-facing config or retry packages: @@ -342,9 +345,14 @@ is exposed as a `null` replay state instead of being deserialized. This also all large-result checkpoints written by earlier SDK versions. `ExtensionContextConfig` directly configures the context serializer and whether the context is virtual. It also -controls framework user-function plugin events and can suppress child checkpoints that finish after the parent. If a -context fails, the SDK first rethrows a deserialized original exception, then calls the configured error handler, and -finally falls back to `ChildContextFailedException`. The handler receives the complete failed context operation +controls framework user-function plugin events, can suppress child checkpoints that finish after the parent, and can +enable validation-only replay with `validateCompletedReplay(true)`. In validation-only replay, the framework callback +is re-entered for a completed context even when `replayChildren` is false. The callback can detect this through +`ExtensionContextReplayContext.isValidatingReplay()`, inspect the checkpointed result through `getReplayState()`, and +recreate deterministic child reservations without emitting a new parent checkpoint. + +If a context fails, the SDK first rethrows a deserialized original exception, then calls the configured error handler, +and finally falls back to `ChildContextFailedException`. The handler receives the complete failed context operation through `failure.operation()`, plus read-only child-operation summaries. ## Explicit child contexts diff --git a/docs/design.md b/docs/design.md index 04980d02e..c8a222cc8 100644 --- a/docs/design.md +++ b/docs/design.md @@ -837,7 +837,9 @@ Key details: deregisters before the step thread starts. - The wrapper deregisters the step thread in `finally`; terminal checkpoint completion re-registers and wakes a context thread waiting on the primitive future. -- For retries, the step sends a RETRY checkpoint and then polls for the READY status before re-executing. If no other threads are active during the retry delay, the execution suspends. +- For retries, the step sends a RETRY checkpoint and polls for READY at the computed retry-ready timestamp before + re-executing. A replayed `PENDING` step uses `stepDetails.nextAttemptTimestamp()` from the checkpoint. If no other + threads are active during the retry delay, the execution suspends. #### WaitPrimitive @@ -888,6 +890,9 @@ Key details: empty payload with `replayChildren=true` and re-executes the child context on replay to reconstruct the result. - Extension contexts use the same empty payload for a `null` replay state. Legacy empty replay payloads are therefore interpreted as `null` instead of being deserialized. +- An extension context can opt completed checkpoints with `replayChildren=false` into validation-only replay. The + framework callback receives the checkpointed result and can recreate deterministic child reservations, while the + completed parent context suppresses new checkpoints. ### In-Process Completion diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/MapIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/MapIntegrationTest.java index c7a5154ba..61f9247c1 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/MapIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/MapIntegrationTest.java @@ -1975,6 +1975,53 @@ void testChangedItemNameFailsCachedReplay() { assertTrue(replay.getError().orElseThrow().errorType().contains("NonDeterministicExecutionException")); } + @Test + void testAddingItemNamerFailsSmallCachedDefaultNamedMapReplay() { + var useItemNamer = new AtomicBoolean(); + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { + var config = MapConfig.builder(); + if (useItemNamer.get()) { + config.itemNamer((item, index) -> "custom-" + item); + } + var result = context.map( + "default-named-map", + List.of("a", "b"), + String.class, + (item, index, ctx) -> item.toUpperCase(), + config.build()); + return String.join(",", result.results()); + }); + + var initial = runner.runUntilComplete("test"); + assertEquals(ExecutionStatus.SUCCEEDED, initial.getStatus()); + assertFalse(Boolean.TRUE.equals( + initial.getOperation("default-named-map").getContextDetails().replayChildren())); + useItemNamer.set(true); + + var replay = runner.run("test"); + + assertEquals(ExecutionStatus.FAILED, replay.getStatus()); + assertTrue(replay.getError().orElseThrow().errorType().contains("NonDeterministicExecutionException")); + } + + @Test + void testChangingItemCountFailsSmallCachedDefaultNamedMapReplay() { + var items = new AtomicReference<>(List.of("a", "b")); + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { + var result = context.map( + "default-named-map", items.get(), String.class, (item, index, ctx) -> item.toUpperCase()); + return String.join(",", result.results()); + }); + + assertEquals(ExecutionStatus.SUCCEEDED, runner.runUntilComplete("test").getStatus()); + items.set(List.of("a", "b", "c")); + + var replay = runner.run("test"); + + assertEquals(ExecutionStatus.FAILED, replay.getStatus()); + assertTrue(replay.getError().orElseThrow().errorType().contains("NonDeterministicExecutionException")); + } + @Test void testEmptyMapDoesNotInvokeItemNamer() { var namerCalls = new AtomicInteger(); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextConfig.java index 86faeb2fa..88a10c897 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextConfig.java @@ -11,6 +11,7 @@ public final class ExtensionContextConfig { private final ExtensionContextErrorHandler errorHandler; private final boolean emitUserFunctionEvents; private final boolean suppressLateChildCheckpoints; + private final boolean validateCompletedReplay; private ExtensionContextConfig(Builder builder) { serDes = builder.serDes; @@ -18,6 +19,7 @@ private ExtensionContextConfig(Builder builder) { errorHandler = builder.errorHandler; emitUserFunctionEvents = builder.emitUserFunctionEvents; suppressLateChildCheckpoints = builder.suppressLateChildCheckpoints; + validateCompletedReplay = builder.validateCompletedReplay; } public SerDes serDes() { @@ -40,13 +42,18 @@ public boolean suppressLateChildCheckpoints() { return suppressLateChildCheckpoints; } + public boolean validateCompletedReplay() { + return validateCompletedReplay; + } + public Builder toBuilder() { return new Builder() .serDes(serDes) .isVirtual(virtual) .errorHandler(errorHandler) .emitUserFunctionEvents(emitUserFunctionEvents) - .suppressLateChildCheckpoints(suppressLateChildCheckpoints); + .suppressLateChildCheckpoints(suppressLateChildCheckpoints) + .validateCompletedReplay(validateCompletedReplay); } public static Builder builder() { @@ -59,6 +66,7 @@ public static final class Builder { private ExtensionContextErrorHandler errorHandler; private boolean emitUserFunctionEvents = true; private boolean suppressLateChildCheckpoints; + private boolean validateCompletedReplay; private Builder() {} @@ -87,6 +95,11 @@ public Builder suppressLateChildCheckpoints(boolean suppressLateChildCheckpoints return this; } + public Builder validateCompletedReplay(boolean validateCompletedReplay) { + this.validateCompletedReplay = validateCompletedReplay; + return this; + } + public ExtensionContextConfig build() { return new ExtensionContextConfig(this); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextReplayContext.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextReplayContext.java index b19c147ae..4428d7e45 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextReplayContext.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionContextReplayContext.java @@ -9,10 +9,12 @@ public final class ExtensionContextReplayContext { private static final ThreadLocal> CURRENT = new ThreadLocal<>(); private final boolean replayingChildren; + private final boolean validatingReplay; private final T replayState; - private ExtensionContextReplayContext(boolean replayingChildren, T replayState) { + private ExtensionContextReplayContext(boolean replayingChildren, boolean validatingReplay, T replayState) { this.replayingChildren = replayingChildren; + this.validatingReplay = validatingReplay; this.replayState = replayState; } @@ -31,6 +33,11 @@ public boolean isReplayingChildren() { return replayingChildren; } + /** Returns whether a completed CONTEXT is re-entering its framework callback only to validate replay. */ + public boolean isValidatingReplay() { + return validatingReplay; + } + /** Returns the checkpointed replay state, or {@code null} on initial execution. */ public T getReplayState() { return replayState; @@ -38,8 +45,13 @@ public T getReplayState() { /** Attaches replay metadata for the duration of an SDK-managed framework callback. */ public static SafeCloseable attach(boolean replayingChildren, T replayState) { + return attach(replayingChildren, false, replayState); + } + + /** Attaches replay and validation metadata for the duration of an SDK-managed framework callback. */ + public static SafeCloseable attach(boolean replayingChildren, boolean validatingReplay, T replayState) { var previous = CURRENT.get(); - CURRENT.set(new ExtensionContextReplayContext<>(replayingChildren, replayState)); + CURRENT.set(new ExtensionContextReplayContext<>(replayingChildren, validatingReplay, replayState)); return () -> { if (previous == null) { CURRENT.remove(); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java index 21d62abf8..aaf2b701c 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java @@ -18,11 +18,14 @@ import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.exception.MapIterationFailedException; +import software.amazon.lambda.durable.exception.NonDeterministicExecutionException; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.extension.ExtensionContext; +import software.amazon.lambda.durable.extension.ExtensionContextConfig; import software.amazon.lambda.durable.extension.ExtensionContextReplayContext; import software.amazon.lambda.durable.extension.ExtensionContextResult; +import software.amazon.lambda.durable.extension.ExtensionOperation; import software.amazon.lambda.durable.model.MapResult; import software.amazon.lambda.durable.model.SafeCloseable; import software.amazon.lambda.durable.serde.SerDes; @@ -100,12 +103,15 @@ public static DurableFuture> mapAsync( var mapConfig = config; var virtualEmptyMap = itemList.isEmpty() && !context.getDurableConfig().shouldCheckpointEmptyMap(); + var parentConfig = parentContextConfig(mapConfig.serDes(), virtualEmptyMap).toBuilder() + .validateCompletedReplay(true) + .build(); return parent.runInChildContextAsync( MAP.getValue(), mapResultType(), () -> executeInChildContext( name, itemList, iterationNames, resultType, function, mapConfig, virtualEmptyMap), - parentContextConfig(mapConfig.serDes(), virtualEmptyMap)); + parentConfig); } private static DurableContext.MapFunction adapt(Function function) { @@ -137,10 +143,14 @@ private static ExtensionContextResult> executeInChildContext } var replay = ExtensionContextReplayContext.>getCurrentContext(); - var replayState = replay.isReplayingChildren() ? replay.getReplayState() : null; - if (replay.isReplayingChildren() && replayState == null) { + var replayingCompletedMap = replay.isReplayingChildren() || replay.isValidatingReplay(); + var replayState = replayingCompletedMap ? replay.getReplayState() : null; + if (replayingCompletedMap && replayState == null) { throw new IllegalStateException("Missing result in completed Map operation"); } + if (replay.isValidatingReplay()) { + return validateCompletedReplay(name, items, iterationNames, resultType, function, config, replayState); + } var coordinator = new OperationConcurrencyCoordinator(config.maxConcurrency(), config.completionConfig()); var registeredItems = @@ -156,6 +166,48 @@ private static ExtensionContextResult> executeInChildContext : ExtensionContextResult.replayChildren(result, strippedResult); } + private static ExtensionContextResult> validateCompletedReplay( + String name, + List items, + List iterationNames, + TypeToken resultType, + DurableContext.MapFunction function, + MapConfig config, + MapResult replayState) { + if (items.size() != replayState.size()) { + throw new NonDeterministicExecutionException(String.format( + "Map item count mismatch for \"%s\". Expected %d, got %d", name, replayState.size(), items.size())); + } + if (config.nestingType() == NestingType.NESTED) { + validateNestedIterations(items, iterationNames, resultType, function, config, replayState); + } + return ExtensionContextResult.completed(replayState); + } + + private static void validateNestedIterations( + List items, + List iterationNames, + TypeToken resultType, + DurableContext.MapFunction function, + MapConfig config, + MapResult replayState) { + var context = ExtensionContext.getCurrentContext(); + var iterationConfig = childContextConfig( + config.serDes(), config.nestingType(), failure -> new MapIterationFailedException(failure.operation())); + for (int index = 0; index < items.size(); index++) { + if (replayState.getItem(index).status() == MapResult.MapResultItem.Status.SKIPPED) { + continue; + } + launchIteration( + context.reserve(iterationNames.get(index)), + items.get(index), + index, + resultType, + function, + iterationConfig); + } + } + private static List> registerItems( OperationConcurrencyCoordinator coordinator, List items, @@ -176,19 +228,29 @@ private static List> registerItem var skipped = replayState != null && replayState.getItem(index).status() == MapResult.MapResultItem.Status.SKIPPED; registeredItems.add(coordinator.register( - () -> reservation.runInChildContextAsync( - MAP_ITERATION.getValue(), - resultType, - () -> ExtensionContextResult.replayChildrenAboveSize( - function.apply(item, itemIndex, DurableContext.requireCurrentContext()), - null, - LARGE_RESULT_THRESHOLD), - iterationConfig), + () -> launchIteration(reservation, item, itemIndex, resultType, function, iterationConfig), skipped)); } return registeredItems; } + private static DurableFuture launchIteration( + ExtensionOperation reservation, + I item, + int index, + TypeToken resultType, + DurableContext.MapFunction function, + ExtensionContextConfig config) { + return reservation.runInChildContextAsync( + MAP_ITERATION.getValue(), + resultType, + () -> ExtensionContextResult.replayChildrenAboveSize( + function.apply(item, index, DurableContext.requireCurrentContext()), + null, + LARGE_RESULT_THRESHOLD), + config); + } + private static List resolveIterationNames(String mapName, List items, MapConfig config) { var namer = config.itemNamer(); var prefix = mapName == null ? "map-iteration-" : mapName + "-iteration-"; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/primitive/ChildContextPrimitive.java b/sdk/src/main/java/software/amazon/lambda/durable/primitive/ChildContextPrimitive.java index 7db0542c8..3f93bd319 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/primitive/ChildContextPrimitive.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/primitive/ChildContextPrimitive.java @@ -52,6 +52,7 @@ public class ChildContextPrimitive extends SerializablePrimitive { private final ExtensionContextFunction extensionFunction; private final ExtensionContextConfig extensionConfig; private final AtomicBoolean replayChildren = new AtomicBoolean(false); + private final AtomicBoolean validatingReplay = new AtomicBoolean(false); private final AtomicReference replayState = new AtomicReference<>(null); private final AtomicReference> cachedOperationResult = new AtomicReference<>(null); @@ -127,19 +128,7 @@ protected void start() { @Override protected void replay(Operation existing) { switch (existing.status()) { - case SUCCEEDED -> { - if (existing.contextDetails() != null - && Boolean.TRUE.equals(existing.contextDetails().replayChildren())) { - replayChildren.set(true); - var result = existing.contextDetails().result(); - if (extensionFunction != null && result != null && !result.isEmpty()) { - replayState.set(deserializeResult(result)); - } - executeChildContext(); - } else { - markAlreadyCompleted(); - } - } + case SUCCEEDED -> replaySucceeded(existing); case FAILED -> markAlreadyCompleted(); case STARTED -> executeChildContext(); default -> @@ -148,6 +137,25 @@ protected void replay(Operation existing) { } } + private void replaySucceeded(Operation existing) { + var details = existing.contextDetails(); + var shouldReplayChildren = details != null && Boolean.TRUE.equals(details.replayChildren()); + var shouldValidateReplay = + extensionFunction != null && extensionConfig.validateCompletedReplay() && !shouldReplayChildren; + if (!shouldReplayChildren && !shouldValidateReplay) { + markAlreadyCompleted(); + return; + } + + replayChildren.set(shouldReplayChildren); + validatingReplay.set(shouldValidateReplay); + var result = details != null ? details.result() : null; + if (extensionFunction != null && result != null && !result.isEmpty()) { + replayState.set(deserializeResult(result)); + } + executeChildContext(); + } + private void executeChildContext() { // The operationId is already globally unique (prefixed by parent context path via // DurableContext.nextOperationId), so we use it directly as the contextId. @@ -191,7 +199,8 @@ private void executeFunction(DurableContextImpl childContext) { return; } - try (var ignoredReplayContext = ExtensionContextReplayContext.attach(replayChildren.get(), replayState.get())) { + try (var ignoredReplayContext = + ExtensionContextReplayContext.attach(replayChildren.get(), validatingReplay.get(), replayState.get())) { var result = extensionConfig.emitUserFunctionEvents() ? runUserFunction(null, extensionFunction::apply) : extensionFunction.apply(); @@ -238,7 +247,10 @@ private String serializeReplayState(T replayState) { } private boolean shouldSkipCheckpoint() { - return replayChildren.get() || isVirtual || parentOperation != null && parentOperation.isOperationCompleted(); + return replayChildren.get() + || validatingReplay.get() + || isVirtual + || parentOperation != null && parentOperation.isOperationCompleted(); } private void cacheSuccessAndComplete(T result) { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/primitive/StepPrimitive.java b/sdk/src/main/java/software/amazon/lambda/durable/primitive/StepPrimitive.java index 14533205d..b61d75af6 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/primitive/StepPrimitive.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/primitive/StepPrimitive.java @@ -3,6 +3,7 @@ package software.amazon.lambda.durable.primitive; import java.time.Duration; +import java.time.Instant; import java.util.concurrent.CompletableFuture; import java.util.function.Function; import software.amazon.awssdk.services.lambda.model.ErrorObject; @@ -63,7 +64,14 @@ protected void start() { protected void replay(Operation existing) { switch (existing.status()) { case SUCCEEDED, FAILED -> markAlreadyCompleted(); - case PENDING -> pollReadyAndResumeExtensionStep(); + case PENDING -> { + var details = existing.stepDetails(); + if (details == null || details.nextAttemptTimestamp() == null) { + throw terminateExecutionWithIllegalDurableOperationException( + "Unexpected PENDING step without nextAttemptTimestamp: " + getOperationId()); + } + pollReadyAndResumeExtensionStep(details.nextAttemptTimestamp()); + } case STARTED -> { if (isAtMostOnce()) { handleExtensionStepFailure( @@ -95,11 +103,11 @@ private T extensionState(Operation existing) { : extensionConfig.initialState(); } - private void pollReadyAndResumeExtensionStep() { - pollForOperationUpdates() + private void pollReadyAndResumeExtensionStep(Instant nextAttemptTimestamp) { + pollForOperationUpdates(nextAttemptTimestamp) .thenCompose(op -> op.status() == OperationStatus.READY ? CompletableFuture.completedFuture(op) - : pollForOperationUpdates()) + : pollForOperationUpdates(nextAttemptTimestamp)) .thenAccept(this::resumeExtensionStep); } @@ -155,14 +163,15 @@ private void handleExtensionStepRetry( update.error(error); } sendOperationUpdate(update); - pollReadyAndExecuteExtensionStep(serializedState.deserialized(), attempt + 1); + pollReadyAndExecuteExtensionStep( + serializedState.deserialized(), attempt + 1, Instant.now().plusSeconds(retryDelaySeconds)); } - private void pollReadyAndExecuteExtensionStep(T state, int attempt) { - pollForOperationUpdates() + private void pollReadyAndExecuteExtensionStep(T state, int attempt, Instant nextAttemptTimestamp) { + pollForOperationUpdates(nextAttemptTimestamp) .thenCompose(op -> op.status() == OperationStatus.READY ? CompletableFuture.completedFuture(op) - : pollForOperationUpdates()) + : pollForOperationUpdates(nextAttemptTimestamp)) .thenRun(() -> executeExtensionStepLogic(state, attempt)); } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionContextConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionContextConfigTest.java index 83cec9aed..ff064b70c 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionContextConfigTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/config/ExtensionContextConfigTest.java @@ -23,6 +23,7 @@ void builderUsesOrdinaryChildContextDefaults() { assertNull(config.errorHandler()); assertTrue(config.emitUserFunctionEvents()); assertFalse(config.suppressLateChildCheckpoints()); + assertFalse(config.validateCompletedReplay()); } @Test @@ -35,6 +36,7 @@ void builderRetainsExtensionPolicies() { .errorHandler(handler) .emitUserFunctionEvents(false) .suppressLateChildCheckpoints(true) + .validateCompletedReplay(true) .build(); assertSame(serDes, config.serDes()); @@ -42,5 +44,6 @@ void builderRetainsExtensionPolicies() { assertEquals(handler, config.errorHandler()); assertFalse(config.emitUserFunctionEvents()); assertTrue(config.suppressLateChildCheckpoints()); + assertTrue(config.validateCompletedReplay()); } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java index 9c299541d..ccf842ceb 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java @@ -72,6 +72,7 @@ void executeBuildsMapAndIterationContextsFromReservations() { assertSame(serDes, parentConfig.getValue().serDes()); assertFalse(parentConfig.getValue().emitUserFunctionEvents()); assertTrue(parentConfig.getValue().suppressLateChildCheckpoints()); + assertTrue(parentConfig.getValue().validateCompletedReplay()); var child = mock(CurrentContext.class); var first = mock(ExtensionOperation.class); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/primitive/ChildContextPrimitiveTest.java b/sdk/src/test/java/software/amazon/lambda/durable/primitive/ChildContextPrimitiveTest.java index dae7f68da..e3e7eb8ed 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/primitive/ChildContextPrimitiveTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/primitive/ChildContextPrimitiveTest.java @@ -185,6 +185,44 @@ void replaySucceededReturnsCachedResult() { assertFalse(functionCalled.get(), "Function should not be called during SUCCEEDED replay"); } + @Test + void extensionCanValidateCompletedReplayWithCachedState() { + when(executionManager.getOperationAndUpdateReplayState("1")) + .thenReturn(Operation.builder() + .id("1") + .name("test-context") + .type(OperationType.CONTEXT) + .subType("AcmeContext") + .status(OperationStatus.SUCCEEDED) + .contextDetails(ContextDetails.builder() + .result("\"cached-value\"") + .build()) + .build()); + var functionCalled = new AtomicBoolean(); + var config = ExtensionContextConfig.builder() + .serDes(SERDES) + .validateCompletedReplay(true) + .build(); + var operation = createExtensionOperation( + "AcmeContext", + () -> { + var replay = ExtensionContextReplayContext.getCurrentContext(); + assertFalse(replay.isReplayingChildren()); + assertTrue(replay.isValidatingReplay()); + assertEquals("cached-value", replay.getReplayState()); + functionCalled.set(true); + return ExtensionContextResult.completed(replay.getReplayState()); + }, + config); + + operation.execute(); + + assertEquals("cached-value", operation.get()); + assertTrue(functionCalled.get()); + verify(executionManager, never()) + .sendOperationUpdate(argThat(update -> update.action() == OperationAction.SUCCEED)); + } + /** Virtual contexts are always executed, even during SUCCEEDED replay. */ @Test void executeVirtualContext() { diff --git a/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java b/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java index 2799c9083..e09961ac2 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/primitive/StatefulExtensionStepPrimitiveTest.java @@ -3,14 +3,19 @@ package software.amazon.lambda.durable.primitive; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.timeout; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import java.time.Duration; +import java.time.Instant; import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; @@ -21,6 +26,7 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.CsvSource; +import org.mockito.ArgumentCaptor; import software.amazon.awssdk.services.lambda.model.ErrorObject; import software.amazon.awssdk.services.lambda.model.Operation; import software.amazon.awssdk.services.lambda.model.OperationAction; @@ -144,14 +150,19 @@ void replayStartedOrReadyResumesWithCheckpointedState(OperationStatus status, in @Test void replayPendingPollsUntilReadyAndResumes() throws Exception { + var nextAttemptTimestamp = Instant.parse("2026-08-10T12:00:00Z"); var pending = operation( OperationStatus.PENDING, - StepDetails.builder().attempt(1).result("5").build()); + StepDetails.builder() + .attempt(1) + .result("5") + .nextAttemptTimestamp(nextAttemptTimestamp) + .build()); var ready = operation( OperationStatus.READY, StepDetails.builder().attempt(1).result("5").build()); when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(pending); - when(executionManager.pollForOperationUpdates(OPERATION_ID)) + when(executionManager.pollForOperationUpdates(OPERATION_ID, nextAttemptTimestamp)) .thenReturn(CompletableFuture.completedFuture(ready)); var called = new CountDownLatch(1); var operation = createOperation(state -> { @@ -162,6 +173,19 @@ void replayPendingPollsUntilReadyAndResumes() throws Exception { operation.execute(); assertTrue(called.await(2, TimeUnit.SECONDS)); + verify(executionManager).pollForOperationUpdates(OPERATION_ID, nextAttemptTimestamp); + } + + @Test + void replayPendingWithoutReadyTimestampFails() { + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)) + .thenReturn(operation( + OperationStatus.PENDING, + StepDetails.builder().attempt(1).result("5").build())); + + var operation = createOperation(ExtensionStepResult::succeed); + + assertThrows(IllegalDurableOperationException.class, operation::execute); } @Test @@ -187,7 +211,8 @@ void replayWithoutCheckpointStateUsesInitialState() throws Exception { @Test void exceptionRetryWithoutStateDoesNotCheckpointPayload() throws Exception { when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(null); - when(executionManager.pollForOperationUpdates(OPERATION_ID)).thenReturn(new CompletableFuture<>()); + when(executionManager.pollForOperationUpdates(eq(OPERATION_ID), any(Instant.class))) + .thenReturn(new CompletableFuture<>()); var retryUpdate = new AtomicReference(); var retrySent = new CountDownLatch(1); when(executionManager.sendOperationUpdate(any())).thenAnswer(invocation -> { @@ -219,7 +244,8 @@ void exceptionRetryWithoutStateDoesNotCheckpointPayload() throws Exception { @Test void retryDelayUsesCheckpointNormalizedState() throws Exception { when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(null); - when(executionManager.pollForOperationUpdates(OPERATION_ID)).thenReturn(new CompletableFuture<>()); + when(executionManager.pollForOperationUpdates(eq(OPERATION_ID), any(Instant.class))) + .thenReturn(new CompletableFuture<>()); var retryUpdate = new AtomicReference(); var retrySent = new CountDownLatch(1); when(executionManager.sendOperationUpdate(any())).thenAnswer(invocation -> { @@ -247,12 +273,18 @@ void retryDelayUsesCheckpointNormalizedState() throws Exception { .build(), durableContext); + var beforeRetry = Instant.now(); operation.execute(); assertTrue(retrySent.await(2, TimeUnit.SECONDS)); + var pollAt = ArgumentCaptor.forClass(Instant.class); + verify(executionManager, timeout(1000)).pollForOperationUpdates(eq(OPERATION_ID), pollAt.capture()); + var afterRetry = Instant.now(); assertEquals("normalized", strategyState.get()); assertEquals("\"raw\"", retryUpdate.get().payload()); assertEquals(7, retryUpdate.get().stepOptions().nextAttemptDelaySeconds()); + assertFalse(pollAt.getValue().isBefore(beforeRetry.plusSeconds(7))); + assertFalse(pollAt.getValue().isAfter(afterRetry.plusSeconds(7))); } @Test From 39b138188a4b12fc167b9d2e84bf62f685a5cf3e Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Tue, 11 Aug 2026 02:54:30 +0000 Subject: [PATCH 38/40] refactor: mark primitive implementations internal --- docs/adr/006-custom-extension-operations.md | 2 +- sdk/pom.xml | 5 +++ .../extension/ExtensionOperationImpl.java | 7 ++++ .../lambda/durable/internal/InternalApi.java | 21 ++++++++++++ .../primitive/SerializablePrimitive.java | 2 +- .../durable/primitive/package-info.java | 13 ++++++++ .../lambda/durable/DurableFutureTest.java | 7 ++-- .../InternalImplementationVisibilityTest.java | 32 +++++++++++++++++++ 8 files changed, 83 insertions(+), 6 deletions(-) create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/internal/InternalApi.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/primitive/package-info.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/InternalImplementationVisibilityTest.java diff --git a/docs/adr/006-custom-extension-operations.md b/docs/adr/006-custom-extension-operations.md index 3144c01f7..00638913f 100644 --- a/docs/adr/006-custom-extension-operations.md +++ b/docs/adr/006-custom-extension-operations.md @@ -82,7 +82,7 @@ DurableContext.step The other primitives follow the same dependency direction through their matching merged operation class. `DurableContextImpl` provides the durable scope and reservation mechanism. Only -`extension.ExtensionOperationImpl` constructs +the internal `extension.ExtensionOperationImpl` constructs the concrete primitive operation engines, so customer APIs and third-party extensions share one backend boundary. Expose each built-in extension family through an independently maintained class: diff --git a/sdk/pom.xml b/sdk/pom.xml index d81047b44..8c2eec82b 100644 --- a/sdk/pom.xml +++ b/sdk/pom.xml @@ -105,6 +105,11 @@ org.apache.maven.plugins maven-javadoc-plugin + + + software.amazon.lambda.durable.internal:software.amazon.lambda.durable.internal.*:software.amazon.lambda.durable.primitive + + attach-javadocs diff --git a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionOperationImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionOperationImpl.java index 546c2d331..788088726 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionOperationImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/extension/ExtensionOperationImpl.java @@ -10,6 +10,7 @@ import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.internal.InternalApi; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.primitive.BasePrimitive; import software.amazon.lambda.durable.primitive.CallbackPrimitive; @@ -19,6 +20,12 @@ import software.amazon.lambda.durable.primitive.WaitPrimitive; import software.amazon.lambda.durable.util.ParameterValidator; +/** + * Internal bridge from extension reservations to checkpoint primitives. + * + * @hidden + */ +@InternalApi public final class ExtensionOperationImpl implements ExtensionOperation { private final DurableContextImpl context; private final String operationId; diff --git a/sdk/src/main/java/software/amazon/lambda/durable/internal/InternalApi.java b/sdk/src/main/java/software/amazon/lambda/durable/internal/InternalApi.java new file mode 100644 index 000000000..8d7a0400e --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/internal/InternalApi.java @@ -0,0 +1,21 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.internal; + +import java.lang.annotation.Documented; +import java.lang.annotation.ElementType; +import java.lang.annotation.Retention; +import java.lang.annotation.RetentionPolicy; +import java.lang.annotation.Target; + +/** + * Marks SDK implementation details that are not part of the supported public API. + * + *

        Types covered by this annotation may change or be removed without compatibility guarantees. + * + * @hidden + */ +@Documented +@Retention(RetentionPolicy.RUNTIME) +@Target({ElementType.PACKAGE, ElementType.TYPE}) +public @interface InternalApi {} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/primitive/SerializablePrimitive.java b/sdk/src/main/java/software/amazon/lambda/durable/primitive/SerializablePrimitive.java index 5a318bbb1..f2baba9bc 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/primitive/SerializablePrimitive.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/primitive/SerializablePrimitive.java @@ -31,7 +31,7 @@ *

      2. Proper thread coordination via future * */ -public abstract class SerializablePrimitive extends BasePrimitive implements DurableFuture { +abstract class SerializablePrimitive extends BasePrimitive implements DurableFuture { private static final Logger logger = LoggerFactory.getLogger(SerializablePrimitive.class); protected record SerializedResult(String serialized, T deserialized) {} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/primitive/package-info.java b/sdk/src/main/java/software/amazon/lambda/durable/primitive/package-info.java new file mode 100644 index 000000000..7cb84d634 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/primitive/package-info.java @@ -0,0 +1,13 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +/** + * Internal checkpoint primitive implementations. + * + *

        These types are SDK implementation details and are not supported for direct use. + * + * @hidden + */ +@InternalApi +package software.amazon.lambda.durable.primitive; + +import software.amazon.lambda.durable.internal.InternalApi; diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java index db57b0471..d93b049da 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableFutureTest.java @@ -11,7 +11,6 @@ import org.junit.jupiter.api.Test; import software.amazon.lambda.durable.context.BaseContextImpl; import software.amazon.lambda.durable.execution.ExecutionManager; -import software.amazon.lambda.durable.primitive.SerializablePrimitive; class DurableFutureTest { @AfterEach @@ -71,7 +70,7 @@ void allOfSingleFutureReturnsSingleResult() { void allOfPropagatesException() { var op1 = mockOperation("first"); @SuppressWarnings("unchecked") - SerializablePrimitive op2 = mock(SerializablePrimitive.class); + DurableFuture op2 = mock(DurableFuture.class); when(op2.get()).thenThrow(new RuntimeException("Step failed")); assertThrows(RuntimeException.class, () -> DurableFuture.allOf(op1, op2)); @@ -106,8 +105,8 @@ void anyOfUsesExecutionManagerWhenCalledFromDurableContext() { } @SuppressWarnings("unchecked") - private SerializablePrimitive mockOperation(T result) { - SerializablePrimitive op = mock(SerializablePrimitive.class); + private DurableFuture mockOperation(T result) { + DurableFuture op = mock(DurableFuture.class); when(op.get()).thenReturn(result); return op; } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/InternalImplementationVisibilityTest.java b/sdk/src/test/java/software/amazon/lambda/durable/InternalImplementationVisibilityTest.java new file mode 100644 index 000000000..866fa6a6e --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/InternalImplementationVisibilityTest.java @@ -0,0 +1,32 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.lang.reflect.Modifier; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.internal.InternalApi; +import software.amazon.lambda.durable.primitive.BasePrimitive; + +class InternalImplementationVisibilityTest { + @Test + void primitivePackageIsMarkedInternal() { + assertTrue(BasePrimitive.class.getPackage().isAnnotationPresent(InternalApi.class)); + } + + @Test + void implementationTypesAreInternalOrPackagePrivate() throws ClassNotFoundException { + assertTrue(Class.forName("software.amazon.lambda.durable.extension.ExtensionOperationImpl") + .isAnnotationPresent(InternalApi.class)); + assertPackagePrivate("software.amazon.lambda.durable.primitive.SerializablePrimitive"); + } + + private void assertPackagePrivate(String className) throws ClassNotFoundException { + var type = Class.forName(className); + assertFalse(Modifier.isPublic(type.getModifiers())); + assertFalse(Modifier.isProtected(type.getModifiers())); + assertFalse(Modifier.isPrivate(type.getModifiers())); + } +} From fecceb3739c8e9344609e01e7f75cb5be29c012a Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Tue, 11 Aug 2026 15:11:31 -0700 Subject: [PATCH 39/40] Fix map replay reservation ordering --- .../operation/DurableMapOperation.java | 9 +-- ...DurableMapOperationImplementationTest.java | 61 +++++++++++++++++++ 2 files changed, 63 insertions(+), 7 deletions(-) diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java index aaf2b701c..a7f8c035a 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/DurableMapOperation.java @@ -195,16 +195,11 @@ private static void validateNestedIterations( var iterationConfig = childContextConfig( config.serDes(), config.nestingType(), failure -> new MapIterationFailedException(failure.operation())); for (int index = 0; index < items.size(); index++) { + var reservation = context.reserve(iterationNames.get(index)); if (replayState.getItem(index).status() == MapResult.MapResultItem.Status.SKIPPED) { continue; } - launchIteration( - context.reserve(iterationNames.get(index)), - items.get(index), - index, - resultType, - function, - iterationConfig); + launchIteration(reservation, items.get(index), index, resultType, function, iterationConfig); } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java index ccf842ceb..7f9a5ebdd 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/DurableMapOperationImplementationTest.java @@ -9,9 +9,12 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.inOrder; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import static software.amazon.lambda.durable.model.ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED; import static software.amazon.lambda.durable.model.OperationSubType.MAP; import static software.amazon.lambda.durable.model.OperationSubType.MAP_ITERATION; @@ -122,6 +125,64 @@ void executeBuildsMapAndIterationContextsFromReservations() { assertSame(failedIteration, failure.getOperation()); } + @Test + void validateCompletedReplayReservesSkippedIterationsBeforeCompletedIterations() { + var context = mock(ExtensionContext.class); + var parent = mock(ExtensionOperation.class); + var config = DurableMapOperation.MapConfig.builder() + .serDes(new JacksonSerDes()) + .build(); + when(context.reserve("map")).thenReturn(parent); + when(parent.runInChildContextAsync( + eq(MAP.getValue()), + any(TypeToken.class), + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) + .thenReturn(mockMapFuture()); + + DurableMapOperation.mapAsync( + context, + "map", + List.of("skipped", "completed"), + TypeToken.get(String.class), + (item, index, child) -> item, + config); + + var function = extensionFunction(); + verify(parent).runInChildContextAsync(eq(MAP.getValue()), any(TypeToken.class), function.capture(), any()); + var child = mock(CurrentContext.class); + var skipped = mock(ExtensionOperation.class); + var completed = mock(ExtensionOperation.class); + when(child.reserve("map-iteration-0")).thenReturn(skipped); + when(child.reserve("map-iteration-1")).thenReturn(completed); + when(completed.runInChildContextAsync( + eq(MAP_ITERATION.getValue()), + eq(TypeToken.get(String.class)), + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class))) + .thenReturn(new CompletedFuture<>("completed")); + var replayState = new MapResult<>( + List.of(MapResult.MapResultItem.skipped(), MapResult.MapResultItem.succeeded("completed")), + MIN_SUCCESSFUL_REACHED); + + MapResult result; + try (var ignoredContext = BaseContextImpl.attachCurrentContext(child); + var ignoredReplay = ExtensionContextReplayContext.attach(false, true, replayState)) { + result = function.getValue().apply().result(); + } + + assertEquals(replayState, result); + var reservations = inOrder(child); + reservations.verify(child).reserve("map-iteration-0"); + reservations.verify(child).reserve("map-iteration-1"); + verify(skipped, never()) + .runInChildContextAsync( + any(String.class), + any(TypeToken.class), + any(ExtensionContextFunction.class), + any(ExtensionContextConfig.class)); + } + @SuppressWarnings({"rawtypes", "unchecked"}) private ArgumentCaptor>> extensionFunction() { return (ArgumentCaptor) ArgumentCaptor.forClass(ExtensionContextFunction.class); From 6d2ff84ffea314da23b4ce06abb7c86e96a787de Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Tue, 11 Aug 2026 20:02:28 -0700 Subject: [PATCH 40/40] Fix awaitFuture suspension race --- .../durable/execution/ExecutionManager.java | 92 ++++++++++++++----- .../execution/ExecutionManagerTest.java | 40 +++++++- 2 files changed, 106 insertions(+), 26 deletions(-) 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 49caa5489..4f45ade7c 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 @@ -71,6 +71,14 @@ public class ExecutionManager implements SafeCloseable { private static final ThreadLocal currentThreadContext = new ThreadLocal<>(); private final CompletableFuture executionExceptionFuture = new CompletableFuture<>(); + enum FutureWaitState { + ACTIVE, + DEREGISTERING, + WAITING, + COMPLETED, + SUSPENDED + } + // ===== Checkpoint Batching ===== private final CheckpointManager checkpointManager; @@ -244,20 +252,11 @@ public T awaitFuture(CompletableFuture future) { CompletableFuture awaitedFuture = future; if (threadContext != null && !future.isDone()) { - var coordinationLock = new Object(); - var deregistered = new boolean[1]; - synchronized (coordinationLock) { - awaitedFuture = future.whenComplete((ignored, throwable) -> { - synchronized (coordinationLock) { - if (deregistered[0] && !isExecutionCompletedExceptionally()) { - registerActiveThread(threadContext.threadId()); - } - } - }); - if (!future.isDone()) { - deregistered[0] = true; - deregisterActiveThread(threadContext.threadId()); - } + var waitState = new AtomicReference<>(FutureWaitState.ACTIVE); + awaitedFuture = future.whenComplete( + (ignored, throwable) -> completeFutureWait(threadContext.threadId(), waitState)); + if (waitState.compareAndSet(FutureWaitState.ACTIVE, FutureWaitState.DEREGISTERING)) { + deregisterActiveThreadForFuture(threadContext.threadId(), waitState); } } @@ -269,6 +268,45 @@ public T awaitFuture(CompletableFuture future) { } } + private void completeFutureWait(String threadId, AtomicReference waitState) { + while (true) { + var current = waitState.get(); + if (current == FutureWaitState.COMPLETED || current == FutureWaitState.SUSPENDED) { + return; + } + if (waitState.compareAndSet(current, FutureWaitState.COMPLETED)) { + if (current == FutureWaitState.WAITING) { + registerActiveThreadIfRunning(threadId); + } + return; + } + } + } + + void deregisterActiveThreadForFuture(String threadId, AtomicReference waitState) { + synchronized (activeThreads) { + removeActiveThread(threadId); + if (!waitState.compareAndSet(FutureWaitState.DEREGISTERING, FutureWaitState.WAITING)) { + if (waitState.get() == FutureWaitState.COMPLETED) { + registerActiveThreadIfRunning(threadId); + } + return; + } + if (activeThreads.isEmpty() + && waitState.compareAndSet(FutureWaitState.WAITING, FutureWaitState.SUSPENDED)) { + suspendForNoActiveThreads(); + } + } + } + + private void registerActiveThreadIfRunning(String threadId) { + synchronized (activeThreads) { + if (!isExecutionCompletedExceptionally()) { + registerActiveThread(threadId); + } + } + } + /** * Registers a thread as active. * @@ -298,21 +336,27 @@ public void deregisterActiveThread(String threadId) { // Add synchronized block to avoid remove then check race condition and make sure that // the suspendExecution is called only once synchronized (activeThreads) { - boolean removed = activeThreads.remove(threadId); - if (removed) { - logger.trace("Deregistered thread '{}' Active threads: {}", threadId, activeThreads.size()); - } else { - logger.warn("Thread '{}' not active, cannot deregister", threadId); - } - + removeActiveThread(threadId); if (activeThreads.isEmpty()) { - logger.info("No active threads remaining - suspending execution"); - preSuspendCheck(); - suspendExecution(); + suspendForNoActiveThreads(); } } } + private void removeActiveThread(String threadId) { + if (activeThreads.remove(threadId)) { + logger.trace("Deregistered thread '{}' Active threads: {}", threadId, activeThreads.size()); + } else { + logger.warn("Thread '{}' not active, cannot deregister", threadId); + } + } + + private void suspendForNoActiveThreads() { + logger.info("No active threads remaining - suspending execution"); + preSuspendCheck(); + suspendExecution(); + } + private void preSuspendCheck() { var hasAnyPendingOperation = operationStorage.values().stream().anyMatch(o -> switch (o.type()) { case STEP -> o.status() == OperationStatus.PENDING; diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java index 72c825e5d..8db1eea59 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java @@ -4,6 +4,7 @@ import static org.junit.jupiter.api.Assertions.*; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.clearInvocations; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.spy; @@ -219,12 +220,16 @@ void isOperationUpdatedSinceLastInvocation_handlesMultipleIds() { void awaitFutureDeregistersAndReregistersCurrentContextThread() { var manager = spy(createManager(List.of(executionOp()))); var deregistered = new CountDownLatch(1); + manager.registerActiveThread("context"); + manager.registerActiveThread("other"); + clearInvocations(manager); doAnswer(invocation -> { + invocation.callRealMethod(); deregistered.countDown(); return null; }) .when(manager) - .deregisterActiveThread("context"); + .deregisterActiveThreadForFuture(any(), any()); var future = new CompletableFuture(); manager.setCurrentThreadContext(new ThreadContext("context", ThreadType.CONTEXT)); CompletableFuture.runAsync(() -> { @@ -238,7 +243,38 @@ void awaitFutureDeregistersAndReregistersCurrentContextThread() { }); assertEquals("done", manager.awaitFuture(future)); - verify(manager).deregisterActiveThread("context"); + verify(manager).deregisterActiveThreadForFuture(any(), any()); verify(manager).registerActiveThread("context"); + assertFalse(manager.isExecutionCompletedExceptionally()); + } + + @Test + void awaitFutureDoesNotSuspendWhenFutureCompletesDuringDeregistration() { + var manager = spy(createManager(List.of(executionOp()))); + var future = new CompletableFuture(); + manager.registerActiveThread("context"); + clearInvocations(manager); + manager.setCurrentThreadContext(new ThreadContext("context", ThreadType.CONTEXT)); + doAnswer(invocation -> { + future.complete("done"); + invocation.callRealMethod(); + return null; + }) + .when(manager) + .deregisterActiveThreadForFuture(any(), any()); + + assertEquals("done", manager.awaitFuture(future)); + verify(manager).registerActiveThread("context"); + assertFalse(manager.isExecutionCompletedExceptionally()); + } + + @Test + void awaitFutureSuspendsWhenFutureRemainsPendingAndNoOtherThreadsAreActive() { + var manager = createManager(List.of(executionOp(), stepOp("step", OperationStatus.PENDING))); + manager.registerActiveThread("context"); + manager.setCurrentThreadContext(new ThreadContext("context", ThreadType.CONTEXT)); + + assertThrows(SuspendExecutionException.class, () -> manager.awaitFuture(new CompletableFuture<>())); + assertTrue(manager.isExecutionCompletedExceptionally()); } }