Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions core/src/main/java/com/google/adk/workflow/NodeExecutionState.java
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import io.reactivex.rxjava3.subjects.SingleSubject;
import java.util.Collection;
import java.util.Collections;
import java.util.LinkedHashMap;
Expand Down Expand Up @@ -61,6 +62,7 @@ public final class NodeExecutionState {
private final ImmutableList<String> outputFor;
private final Map<String, @Nullable Object> resumeInputs;
private final ConcurrentMap<String, Integer> childRunCounts = new ConcurrentHashMap<>();
private final ConcurrentMap<String, SingleSubject<Context>> childRuns = new ConcurrentHashMap<>();
// The paths of the node runs started but not finished in this turn, shared under one root.
private final Set<String> unfinishedNodePaths;
private final Set<String> interruptIds = new CopyOnWriteArraySet<>();
Expand Down Expand Up @@ -224,6 +226,14 @@ String nextChildRunId(String name) {
return String.valueOf(childRunCounts.merge(name, 1, Integer::sum));
}

/**
* Returns the runs of the children this node runs itself, keyed by {@code name@runId}, which
* calls with the same run ID share.
*/
ConcurrentMap<String, SingleSubject<Context>> childRuns() {
return childRuns;
}

/** Records that this run started; it counts as canceled mid-run until it finishes. */
void markStarted() {
unfinishedNodePaths.add(nodePath);
Expand Down Expand Up @@ -255,6 +265,11 @@ synchronized ImmutableMap<String, ResumeScan.RecoveredNode> recoveredChildren(
return recoveredChildren;
}

/** Returns the parent's state if this run's output is also the parent's, or null. */
@Nullable NodeExecutionState outputParent() {
return outputParent;
}

/** Returns whether an event has carried the node's own or delegated output. */
boolean outputEmitted() {
OutputStatus status = outputState.get().status();
Expand Down
78 changes: 74 additions & 4 deletions core/src/main/java/com/google/adk/workflow/NodeRunner.java
Original file line number Diff line number Diff line change
Expand Up @@ -40,10 +40,12 @@
import io.reactivex.rxjava3.processors.FlowableProcessor;
import io.reactivex.rxjava3.processors.UnicastProcessor;
import io.reactivex.rxjava3.subjects.CompletableSubject;
import io.reactivex.rxjava3.subjects.SingleSubject;
import java.util.Map;
import java.util.Optional;
import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ConcurrentMap;
import java.util.function.Function;
import org.jspecify.annotations.Nullable;

Expand Down Expand Up @@ -218,8 +220,7 @@ private static Maybe<Object> startChild(
parentState.claimOutputDelegation();
}
String runId = config.runId().orElseGet(() -> parentState.nextChildRunId(node.name()));
// Deferred so that a throw from scan or intercept releases the delegation, as a failure does.
return Single.defer(() -> runOrReplayChild(parent, node, nodeInput, runId, config))
return sharedChildRun(parent, node, nodeInput, runId, config)
.doOnError(
error -> {
if (delegatesOutput) {
Expand All @@ -229,6 +230,50 @@ private static Maybe<Object> startChild(
.flatMapMaybe(child -> childResult(parent, child, delegatesOutput, config.raiseOnWait()));
}

/**
* Returns the run of {@code node} under {@code runId} shared by the caller's calls with that run
* ID, starting it if needed. A run that fails or waits for an output it did not produce is not
* shared, so a later call runs the child again. When the call that started a run is canceled, a
* call still waiting on it runs the child itself.
*/
private static Single<Context> sharedChildRun(
Context parent, Node node, @Nullable Object nodeInput, String runId, RunNodeConfig config) {
return Single.defer(
() -> {
ConcurrentMap<String, SingleSubject<Context>> childRuns =
parent.requireNodeState().childRuns();
String childKey = segment(node.name(), runId);
SingleSubject<Context> run = SingleSubject.create();
SingleSubject<Context> shared = childRuns.putIfAbsent(childKey, run);
if (shared != null) {
return shared.onErrorResumeNext(
error ->
error instanceof ChildRunCanceledException
? sharedChildRun(parent, node, nodeInput, runId, config)
: Single.error(error));
}
// Defer so a throw from scan or intercept releases the claimed run instead of leaking it.
return Single.defer(() -> runOrReplayChild(parent, node, nodeInput, runId, config))
.doOnSuccess(
child -> {
if (!keepsItsRun(child)) {
childRuns.remove(childKey, run);
}
run.onSuccess(child);
})
.doOnError(
error -> {
childRuns.remove(childKey, run);
run.onError(error);
})
.doOnDispose(
() -> {
childRuns.remove(childKey, run);
run.onError(new ChildRunCanceledException());
});
});
}

/**
* Runs {@code node} as a child of {@code parent}, or replays the result it recorded in an earlier
* turn under the same run ID.
Expand Down Expand Up @@ -258,6 +303,22 @@ private static Single<Context> runOrReplayChild(
.run(nodeInput);
}

/**
* Returns whether a later call with the same run ID shares this finished run of a child: a failed
* child runs again, as does one that waits for an output and left no outcome.
*/
private static boolean keepsItsRun(Context child) {
NodeExecutionState state = child.requireNodeState();
if (state.failure().isPresent()) {
return false;
}
return !state.interruptIds().isEmpty()
|| state.output().isPresent()
|| state.selectedRoutes().isPresent()
|| child.eventActions().transferToAgent().isPresent()
|| !state.node().waitForOutput();
}

/**
* Returns the child's output, or fails with the child's failure or with an interrupt, which the
* caller takes on as its own.
Expand Down Expand Up @@ -292,9 +353,11 @@ private static Maybe<Object> childResult(
&& (childNode instanceof Workflow || childNode.waitForOutput())) {
return Maybe.error(new NodeInterruptedException());
}
// No event of this run carries a delegate's replayed output, so the caller emits it.
// No event of this run carries a delegate's replayed or shared output, so the caller emits it.
Optional<Object> output = childState.output();
if (delegatesOutput && output.isPresent() && !childState.outputEmitted()) {
boolean emittedForCaller =
childState.outputEmitted() && childState.outputParent() == parentState;
if (delegatesOutput && output.isPresent() && !emittedForCaller) {
parentState.adoptReplayedOutput(output.get());
}
return Maybe.fromOptional(output);
Expand Down Expand Up @@ -532,6 +595,13 @@ private Event.Builder newEvent() {
}
}

/** Ends the waiting calls of a shared child run whose first caller was canceled. */
private static final class ChildRunCanceledException extends RuntimeException {
ChildRunCanceledException() {
super("The call that ran this child was canceled.");
}
}

/**
* An event queued for the runner, paired with the signal that completes once the caller receives
* the event and, for an error event, with the failure it reports.
Expand Down
7 changes: 4 additions & 3 deletions core/src/main/java/com/google/adk/workflow/RunNodeConfig.java
Original file line number Diff line number Diff line change
Expand Up @@ -31,9 +31,10 @@ public abstract class RunNodeConfig {

/**
* Returns the ID of the child's run, which names the run in node paths and branches. A child that
* recorded an output under the same ID in an earlier turn returns it without running again. When
* this ID is empty, the child receives the next number among the caller's runs of nodes with the
* same name, in call order, so concurrent children need explicit IDs.
* already recorded an output under the same ID, in this run of the caller or in an earlier turn,
* returns it without running again; one that finished without an output may run again. When this
* ID is empty, the child receives the next number among the caller's runs of nodes with the same
* name, in call order, so concurrent children need explicit IDs.
*/
public abstract Optional<String> runId();

Expand Down
Loading
Loading