[EXPERIMENTAL] feat(dag): migrate DAG support to extension SPI by zhongkechen · Pull Request #611 · aws/aws-durable-execution-sdk-java

zhongkechen · 2026-08-10T17:29:27Z

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

Warning

Experimental stacked PR. This work depends on draft PR #607 and is not intended for production use. The DAG API, checkpoint shape, and implementation may change before merge or release.

Issue Link, if available

Stacked on #607.

Description

Migrates the experimental DAG implementation to the extension operation SPI introduced by #607.

Adds static DurableDagOperation.dag(...) and dagAsync(...) entry points following the new Durable*Operation facade style.
Removes DAG-specific methods and internal hooks from DurableContext.
Schedules DAG containers and tasks through ExtensionContext and ExtensionOperation.
Preserves stable name-based DAG task IDs.
Reuses unchanged map, parallel, and wait-for-condition facades through a DAG-internal reserved-context adapter.
Leaves all other Durable*Operation classes and their implementation tests unchanged from the stacked base.
Handles large DAG results with extension child replay state.
Updates examples, conformance handlers, integration tests, and DAG documentation.

Demo/Screenshots

Not applicable. This is an SDK API and execution implementation change.

Checklist

I have filled out every section of the PR template
I have thoroughly tested this change

Testing

Unit Tests

Yes. Added focused coverage for DurableDagOperation and the DAG-internal reserved-context adapter.

Integration Tests

Yes. DAG integration and conformance tests were updated for the static operation facade. The complete Maven reactor passes.

Commands run:

mvn spotless:apply
mvn -pl sdk -Dtest=DurableDagOperationTest,ReservedOperationContextTest test \
  -DargLine=-javaagent:$HOME/.m2/repository/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar
mvn test \
  -DargLine=-javaagent:$HOME/.m2/repository/org/mockito/mockito-core/5.23.0/mockito-core-5.23.0.jar

Cloud example tests remain disabled by default.

Examples

Existing DAG examples were migrated to DurableDagOperation; no additional example was required.

Sign in to view

+        while (!inFlight.isEmpty()) {
+            var it = inFlight.entrySet().iterator();
+            var entry = it.next();
+            it.remove();
+            String name = entry.getKey();
+            try {
+                Object result = entry.getValue().future.get();


Sign in to view

+        return switch (kind) {
+            case PLAIN -> rehydratePlain(raw, taskName, scope);
+            case BATCH -> delegate.deserialize(delegate.serialize(raw), TypeToken.get(MapResult.class));


Sign in to view

+            } catch (UnrecoverableDurableExecutionException e) {
+                throw e;
+            } catch (RuntimeException e) {
+                results.put(name, failed(name, DagTaskError.of(e)));


Sign in to view

+        var cc = unwrap(dcc);
+        int succeeded = countByStatus(results, TaskStatus.SUCCEEDED);
+        int failed = countByStatus(results, TaskStatus.FAILED);
+        if (cc.minSuccessful() != null && succeeded >= cc.minSuccessful()) {
+            return DagCompletionReason.MIN_SUCCESSFUL_REACHED;


Sign in to view

+        try {
+            return delegate.deserialize(delegate.serialize(raw), declared.get());
+        } catch (RuntimeException e) {
+            return raw;


Sign in to view

+    // ── invoke ───────────────────────────────────────────────────────────────
+    <T> TaskHandle<T> invoke(String name, String functionName, Class<T> type, DagPayloadFunction payloadFn);
+
+    <T> TaskHandle<T> invoke(
+            String name, String functionName, Class<T> type, DagPayloadFunction payloadFn, InvokeConfig config);
+
+    // ── callback ─────────────────────────────────────────────────────────────
+    <T> TaskHandle<T> callback(String name, Class<T> type, DagCallbackSubmitter submitter);
+
+    <T> TaskHandle<T> callback(
+            String name, Class<T> type, DagCallbackSubmitter submitter, WaitForCallbackConfig config);
+
+    // ── wait ─────────────────────────────────────────────────────────────────
+    TaskHandle<Void> wait(String name, Duration duration);
+
+    // ── waitForCondition ──────────────────────────────────────────────────────
+    <S> TaskHandle<S> waitForCondition(
+            String name, Class<S> type, DagConditionFunction<S> check, WaitForConditionConfig<S> config);
+
+    // ── runInChildContext ─────────────────────────────────────────────────────
+    <T> TaskHandle<T> runInChildContext(String name, Class<T> type, DagChildFunction<T> fn);
+
+    <T> TaskHandle<T> runInChildContext(String name, TypeToken<T> type, DagChildFunction<T> fn);
+
+    // ── map ──────────────────────────────────────────────────────────────────
+    <I, O> TaskHandle<MapResult<O>> map(String name, Collection<I> items, Class<O> type, MapFunction<I, O> fn);
+
+    <I, O> TaskHandle<MapResult<O>> map(
+            String name, Collection<I> items, Class<O> type, MapFunction<I, O> fn, MapConfig config);
+
+    <I, O> TaskHandle<MapResult<O>> map(
+            String name, Function<Deps, Collection<I>> items, Class<O> type, MapFunction<I, O> fn);
+
+    <I, O> TaskHandle<MapResult<O>> map(
+            String name, Function<Deps, Collection<I>> items, Class<O> type, MapFunction<I, O> fn, MapConfig config);


Sign in to view

+## Completion (threshold only in v1)
+
+`DagConfig.builder().completionConfig(...)` accepts one of six threshold policies:
+`allCompleted`, `allSuccessful`, `firstSuccessful`, `minSuccessful(n)`, `toleratedFailureCount(n)`,
+`toleratedFailurePercentage(p)`. Default (no `completionConfig`) drains the whole reachable graph. `completionReason()`
+reports `ALL_COMPLETED`, `COMPLETED_WITH_FAILURES`, `MIN_SUCCESSFUL_REACHED`, or `FAILURE_TOLERANCE_EXCEEDED`.
+
+> **v2-deferred:** Custom-predicate (result-based) completion is **not** in v1. `DagCompletionConfig` exposes only the
+> threshold factories, and `DagCompletionReason.CUSTOM_COMPLETION_*` are reserved-but-unreachable.
+
+## Results
+
+`DagResult` provides `getResult(TaskHandle<T>) -> Optional<T>` (typed) and `getResult(String) -> Optional<Object>`
+(untyped), `getStatus(...)`, grouped views (`succeeded()`/`failed()`/`skipped()`), counts, `completionReason()`, and
+`throwIfError()` (throws `DagExecutionException` iff `failureCount() > 0`).


…extension-spi

github-actions · 2026-08-20T20:07:46Z

+    private static ExtensionContextResult<DagResult> executeDag(DagContextImpl dag, DagConfig config, SerDes serDes) {
+        var result = DagResultImpl.from(DagExecutor.run(dag.tasks(), ExtensionContext.getCurrentContext(), config));
+        var replayState = serDes instanceof DagResultSerDes dagSerDes ? dagSerDes.replayState(result) : null;
+        return ExtensionContextResult.replayChildrenAboveSize(result, replayState, LARGE_RESULT_THRESHOLD);


Codex AI review

P1: Validate the graph before replaying an offloaded DAG

ReplayChildren re-enters this callback, but the stored replay context is ignored and the current graph is executed. If a completed DAG is redeployed with an added, renamed, or previously skipped task, that task has no checkpoint and its user code can run while replaying an already-successful DAG. Persist a bounded graph signature in replay state and reject mismatches before launching any task; add replay tests for changed and newly runnable tasks.

github-actions · 2026-08-20T20:07:47Z

+    private String selectOffloadPayload(DagResult result) {
+        var payloads = offloadPayloads(result);
+        for (var payload : payloads) {
+            if (payload == null || payload.getBytes(StandardCharsets.UTF_8).length < LARGE_RESULT_THRESHOLD) {
+                return payload;
+            }
+        }
+        return payloads.get(payloads.size() - 1);


Codex AI review

P1: Ensure the final replay payload is actually bounded

When both candidates exceed 256 KiB, this returns the last oversized payload anyway. startedTaskNames is not byte-bounded because explicit maxConcurrency has no upper limit and each name may contain 100 characters, so a wide early-completing DAG can still fail its success checkpoint. Add a compact bounded representation for the started set, or enforce bounds that guarantee the final candidate fits, and test a replay state with thousands of started tasks.

github-actions · 2026-08-20T20:07:49Z

+        while (!inFlight.isEmpty()) {
+            var it = inFlight.entrySet().iterator();
+            var entry = it.next();
+            it.remove();
+            String name = entry.getKey();
+            try {
+                Object result = entry.getValue().future.get();


Codex AI review

P1: Wait for tasks in completion order

The scheduler always blocks on the first-launched future. A long wait registered before a fast task prevents harvesting the fast result, scheduling its dependents, or satisfying firstSuccessful until the earlier task finishes. Wait via completion signals/DurableFuture.anyOf, then harvest all completed entries deterministically; preserve the settled stop set in replay state so early completion remains reproducible.

github-actions · 2026-08-20T20:07:51Z

+            case PLAIN -> rehydratePlain(raw, taskName, scope);
+            case BATCH -> delegate.deserialize(delegate.serialize(raw), TypeToken.get(MapResult.class));


Codex AI review

P1: Preserve map item result types

Deserializing BATCH as raw MapResult.class turns POJO item results into generic maps during the normal serialize-deserialize boundary. A TaskHandle<MapResult<MyPojo>> can therefore throw ClassCastException even on first execution. Record each map task's item TypeToken and rehydrate every successful item with it; cover this with a POJO map-result round trip rather than only strings.

github-actions · 2026-08-20T20:07:52Z

+        TaskExecutor<T> exec = (ctx, operation, deps) -> operation.invokeAsync(
+                INVOKE_SUBTYPE,
+                functionName,
+                payloadFn.apply(deps),


Codex AI review

P1: Convert task-preparation failures into task failures

payloadFn.apply(deps) runs synchronously during scheduler launch, outside the future failure boundary. If it throws, the entire DAG aborts and downstream compensation never runs, despite the documented rule that only a throwing runIf aborts. Execute preparation inside a failure-bearing task operation or record ordinary callback failures as TaskExecution.FAILED while rethrowing unrecoverable execution errors; apply the same treatment to dynamic map inputs and parallel registration.

github-actions · 2026-08-20T20:07:54Z

+                var statuses = depStatuses(task, results);
+                var rule = task.triggerRuleOpt().orElse(defaultRule);
+                if (!TriggerRuleEvaluator.eval(rule, statuses)) {
+                    results.put(name, skipped(name, SkipReason.TRIGGER_RULE));


Codex AI review

P2: Evaluate custom completion after skips

Skips are terminal settlements, but completion is evaluated only after future.get(). An all-skipped graph never invokes its custom predicate, and a final wave of downstream skips is invisible to it. Re-evaluate completion whenever fillReady records a skip, including after the initial fill and before returning with no in-flight tasks.

github-actions · 2026-08-20T20:07:55Z

+        var cc = unwrap(dcc);
+        int succeeded = countByStatus(results, TaskStatus.SUCCEEDED);
+        int failed = countByStatus(results, TaskStatus.FAILED);
+        if (cc.minSuccessful() != null && succeeded >= cc.minSuccessful()) {


Codex AI review

P2: Reject impossible minSuccessful thresholds

When minSuccessful exceeds the registered task count, this condition can never become true and the DAG eventually reports ALL_COMPLETED. The base CompletionConfig contract rejects that configuration. Validate the threshold after registration and before reserving or launching operations, including for empty DAGs.

github-actions · 2026-08-20T20:07:56Z

+        try {
+            return delegate.deserialize(delegate.serialize(raw), declared.get());
+        } catch (RuntimeException e) {
+            return raw;


Codex AI review

P2: Do not silently erase known declared types

For a known task, a deserialization failure currently returns the raw parsed tree, violating TaskHandle<T> and potentially changing the value's runtime type only after replay. Keep generic-tree fallback only for unknown task names and propagate the SerDes failure for declared types.

Suggested change

try {

return delegate.deserialize(delegate.serialize(raw), declared.get());

} catch (RuntimeException e) {

return raw;

return delegate.deserialize(delegate.serialize(raw), declared.get());

github-actions · 2026-08-20T20:07:58Z

+            DagStep3Function<A, B, C, T> fn);
+
+    // ── invoke ───────────────────────────────────────────────────────────────
+    <T> TaskHandle<T> invoke(String name, String functionName, Class<T> type, DagPayloadFunction payloadFn);


Codex AI review

P2: Add TypeToken overloads for generic task results

This Class-only pattern is repeated for invoke, callback, wait-for-condition, and map, so results such as List<MyPojo> cannot be declared without erasing their element type. Add matching TypeToken<T> overloads, as provided by the core operation facades, and carry the full token through registration and aggregate rehydration.

github-actions · 2026-08-20T20:07:59Z

+    public <I, O> TaskHandle<MapResult<O>> map(
+            String name, Collection<I> items, Class<O> type, MapFunction<I, O> fn, MapConfig config) {
+        return map(name, deps -> items, type, fn, config);


Codex AI review

P2: Snapshot static map inputs at registration

This overload captures the caller's mutable collection until the task launches. A predecessor can mutate it on first execution but be skipped on replay, changing map cardinality and causing nondeterminism. Copy the collection immediately; retain the function overload for intentionally dynamic inputs.

Suggested change

public <I, O> TaskHandle<MapResult<O>> map(

String name, Collection<I> items, Class<O> type, MapFunction<I, O> fn, MapConfig config) {

return map(name, deps -> items, type, fn, config);

var itemSnapshot = List.copyOf(items);

return map(name, deps -> itemSnapshot, type, fn, config);

github-actions · 2026-08-20T20:08:01Z

+            if (exec == null) {
+                item = new DagCompletionItemStatus(task.name(), Optional.empty(), Optional.empty(), Optional.empty());


Codex AI review

P2: Report launched tasks as STARTED

Every nonterminal task receives an empty status because this method only sees terminal results, even when the task is present in inFlight. This contradicts DagCompletionItemStatus, which reserves empty status for tasks that have not started, and prevents custom predicates from distinguishing active work. Pass the in-flight names into this snapshot and emit TaskStatus.STARTED for them.

github-actions · 2026-08-20T20:08:03Z

+- `deps.get(handle)` returns the upstream's declared type `T`; `deps.getOptional(handle)` returns `Optional<T>` for
+  non-`ALL_SUCCESS` paths where an upstream may be FAILED/SKIPPED.


Codex AI review

P3: Align the documentation with the published API

Deps has no getOptional method, and get already returns Optional<T>, so this guidance directs customers to nonexistent and incorrectly typed APIs. Update this section and examples accordingly; the completion section should also stop claiming custom completion is unavailable and that throwIfError depends only on failureCount.

github-actions · 2026-08-20T20:08:05Z

Codex AI review

Found 12 correctness and API issues, including replay side effects, completion scheduling regressions, and result type erasure. Static review only; tests were not run per the review constraints.

Reviewed commit b9e9f7cf30480451ab6b4f1fe92bd2c04255b6f4. Workflow run

zhongkechen deployed to ai-pr-review August 10, 2026 17:29 — with GitHub Actions Active

zhongkechen requested a deployment to ai-pr-review-runtime August 10, 2026 17:29 — with GitHub Actions Waiting

zhongkechen had a problem deploying to ai-pr-review-runtime August 10, 2026 17:29 — with GitHub Actions Failure

zhongkechen mentioned this pull request Aug 10, 2026

feat: add custom extension operations API #607

Open

2 tasks

zhongkechen force-pushed the codex/experimental-dag-extension-spi branch 2 times, most recently from ecf0c30 to e25ea9c Compare August 10, 2026 20:05

zhongkechen force-pushed the codex/experimental-dag-extension-spi branch 2 times, most recently from 25cf931 to b08588f Compare August 10, 2026 23:59

zhongkechen and others added 3 commits August 10, 2026 22:19

feat(dag): rebase experimental DAG support onto extension SPI

db49108

refactor(dag): align with durable operation facade style

0a28e84

refactor(dag): publish DAG as separate artifact

818b256

zhongkechen force-pushed the codex/experimental-dag-extension-spi branch from 6a8bbc1 to 818b256 Compare August 11, 2026 05:19

zhongkechen deployed to ai-pr-review August 14, 2026 21:52 — with GitHub Actions Active

zhongkechen had a problem deploying to ai-pr-review-runtime August 14, 2026 21:52 — with GitHub Actions Failure

zhongkechen temporarily deployed to ai-pr-review-runtime August 14, 2026 21:52 — with GitHub Actions Inactive

github-actions Bot reviewed Aug 14, 2026

View reviewed changes

This comment has been minimized.

Sign in to view

Merge codex/extension-operation-refactor into codex/experimental-dag-…

b9e9f7c

…extension-spi

zhongkechen deployed to ai-pr-review August 20, 2026 18:44 — with GitHub Actions Active

zhongkechen had a problem deploying to ai-pr-review-runtime August 20, 2026 19:03 — with GitHub Actions Failure

zhongkechen temporarily deployed to ai-pr-review-runtime August 20, 2026 19:03 — with GitHub Actions Inactive

github-actions Bot reviewed Aug 20, 2026

View reviewed changes

		case PLAIN -> rehydratePlain(raw, taskName, scope);
		case BATCH -> delegate.deserialize(delegate.serialize(raw), TypeToken.get(MapResult.class));

		if (exec == null) {
		item = new DagCompletionItemStatus(task.name(), Optional.empty(), Optional.empty(), Optional.empty());

		- `deps.get(handle)` returns the upstream's declared type `T`; `deps.getOptional(handle)` returns `Optional<T>` for
		non-`ALL_SUCCESS` paths where an upstream may be FAILED/SKIPPED.

Uh oh!

Conversation

zhongkechen commented Aug 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Issue Link, if available

Description

Demo/Screenshots

Checklist

Testing

Unit Tests

Integration Tests

Examples

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment has been minimized.

github-actions Bot Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Aug 20, 2026

Codex AI review

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

zhongkechen commented Aug 10, 2026 •

edited

Loading