From b701803844deac0fcc6c2247a47de304422abe64 Mon Sep 17 00:00:00 2001 From: Mateusz Krawiec Date: Fri, 13 Nov 2026 12:30:01 -0800 Subject: [PATCH] fix(workflow): wait for a canceled FunctionNode body before a retry PiperOrigin-RevId: 1021240441 --- .../com/google/adk/workflow/FunctionNode.java | 19 +- .../com/google/adk/workflow/NodeBody.java | 5 +- .../adk/workflow/NodeExecutionState.java | 13 + .../com/google/adk/workflow/NodeRunner.java | 11 +- .../google/adk/workflow/FunctionNodeTest.java | 263 ++++++++++++++++++ 5 files changed, 304 insertions(+), 7 deletions(-) diff --git a/core/src/main/java/com/google/adk/workflow/FunctionNode.java b/core/src/main/java/com/google/adk/workflow/FunctionNode.java index 6833fc453..fca0a73d0 100644 --- a/core/src/main/java/com/google/adk/workflow/FunctionNode.java +++ b/core/src/main/java/com/google/adk/workflow/FunctionNode.java @@ -32,6 +32,8 @@ import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Maybe; import io.reactivex.rxjava3.core.MaybeEmitter; +import io.reactivex.rxjava3.schedulers.Schedulers; +import io.reactivex.rxjava3.subjects.CompletableSubject; import java.util.LinkedHashMap; import java.util.Map; import java.util.Optional; @@ -166,7 +168,22 @@ boolean edgeInputSchemaInferred() { @Override public Flowable runNode(Context context, @Nullable Object nodeInput) { // Maybe.create rather than defer, so that canceling the node reaches the body while it runs. - return Maybe.create(emitter -> runBody(context, nodeInput, emitter)) + return Maybe.create( + emitter -> { + // A timeout fails the node only once the body returns, so a retry cannot overlap it. + CompletableSubject returned = CompletableSubject.create(); + NodeExecutionState state = context.requireNodeState(); + // A body run under another node's context must not replace that node's signal. + if (state.node() == this) { + // Not on the body's thread, which may be interrupted, nor the Runner's busy pool. + state.markBodyStarted(returned.observeOn(Schedulers.computation())); + } + try { + runBody(context, nodeInput, emitter); + } finally { + returned.onComplete(); + } + }) .subscribeOn(context.invocationContext().scheduler()) .toFlowable(); } diff --git a/core/src/main/java/com/google/adk/workflow/NodeBody.java b/core/src/main/java/com/google/adk/workflow/NodeBody.java index 5999d037c..90c65c71a 100644 --- a/core/src/main/java/com/google/adk/workflow/NodeBody.java +++ b/core/src/main/java/com/google/adk/workflow/NodeBody.java @@ -24,8 +24,9 @@ * The body of a {@link FunctionNode}, which takes the node run's context and its single input; the * input may be shared with other nodes, so the body must not mutate it. It runs on the {@code * Runner}'s scheduler, so it may block, and a small fixed-size scheduler can deadlock bodies that - * wait for one another. Canceling the node interrupts the body's thread and drops its result; a - * body that ignores the interrupt keeps its thread until it returns, and a retry can overlap it. + * wait for one another. Canceling the node interrupts the body's thread and drops its result; only + * the node's own timeout waits for the body to return, so a body that ignores the interrupt holds + * up the retry and the run. */ @Experimental @FunctionalInterface diff --git a/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java b/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java index 2a8f0bebd..191aca877 100644 --- a/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java +++ b/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java @@ -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.core.Completable; import java.util.Collection; import java.util.Collections; import java.util.LinkedHashMap; @@ -74,6 +75,8 @@ public final class NodeExecutionState { private final AtomicReference outputState = new AtomicReference<>(NO_OUTPUT); // Set once the node drops its own output because a useAsOutput child holds the output. private final AtomicBoolean droppedOwnOutput = new AtomicBoolean(); + private final AtomicReference bodyReturned = + new AtomicReference<>(Completable.complete()); private String eventAuthor; private @Nullable ImmutableList selectedRoutes; private boolean routesEmitted; @@ -298,6 +301,16 @@ void markFinished() { } } + /** Records that a body of this run started; {@code returned} completes once it returns. */ + void markBodyStarted(Completable returned) { + bodyReturned.set(returned); + } + + /** Completes once the body this run started has returned, or at once if none started. */ + Completable bodyReturned() { + return Completable.defer(bodyReturned::get); + } + /** * Returns whether this run's child {@code childKey}, a {@code name@runId} segment, started in * this turn and has not finished. When read before the child runs again, this means a failing diff --git a/core/src/main/java/com/google/adk/workflow/NodeRunner.java b/core/src/main/java/com/google/adk/workflow/NodeRunner.java index 5a16dc449..e36a5baed 100644 --- a/core/src/main/java/com/google/adk/workflow/NodeRunner.java +++ b/core/src/main/java/com/google/adk/workflow/NodeRunner.java @@ -529,16 +529,19 @@ private Completable withTimeout(Completable run) { return run; } Duration limit = timeout.get(); - // Fails on the timer's thread: the timed-out attempt may still hold the scheduler's threads. - Completable timedOut = Completable.error(() -> new NodeTimeoutException(node.name(), limit)); + // The timeout only interrupts a body, so the node fails once the body has returned. + Completable timedOut = + nodeState + .bodyReturned() + .andThen(Completable.error(() -> new NodeTimeoutException(node.name(), limit))); io.opentelemetry.context.Context traceContext = io.opentelemetry.context.Context.current(); - // Off the caller's thread, so a body that blocks its thread cannot hold off the timeout. + // Off the caller's thread, so a body that blocks does not block the caller. return Completable.timer(0, NANOSECONDS, invocationContext.scheduler()) // Unlike subscribeOn, canceling this hop does not interrupt the handling of node events. .andThen(run.compose(Tracing.withContext(traceContext))) // convert saturates, so a timeout too long to count in nanoseconds waits about 292 years. .timeout(NANOSECONDS.convert(limit), NANOSECONDS, timedOut) - // The failure arrives on the timer's thread, so it needs the caller's trace context too. + // The failure arrives on another thread, so it needs the caller's trace context too. .compose(Tracing.withContext(traceContext)); } diff --git a/core/src/test/java/com/google/adk/workflow/FunctionNodeTest.java b/core/src/test/java/com/google/adk/workflow/FunctionNodeTest.java index 71d7ce201..720145bbc 100644 --- a/core/src/test/java/com/google/adk/workflow/FunctionNodeTest.java +++ b/core/src/test/java/com/google/adk/workflow/FunctionNodeTest.java @@ -28,6 +28,7 @@ import static com.google.common.collect.MoreCollectors.onlyElement; import static com.google.common.truth.Truth.assertThat; import static com.google.common.truth.Truth.assertWithMessage; +import static java.util.concurrent.TimeUnit.MILLISECONDS; import static java.util.concurrent.TimeUnit.SECONDS; import static org.junit.Assert.assertThrows; @@ -55,10 +56,12 @@ import com.fasterxml.jackson.databind.ser.std.ToStringSerializer; import com.fasterxml.jackson.databind.util.StdConverter; import com.google.adk.agents.Context; +import com.google.adk.agents.InvocationContext; import com.google.adk.apps.App; import com.google.adk.events.Event; import com.google.adk.events.EventActions; import com.google.adk.events.ToolConfirmation; +import com.google.adk.plugins.BasePlugin; import com.google.adk.runner.Runner; import com.google.adk.sessions.Session; import com.google.adk.testing.TestNode; @@ -76,6 +79,7 @@ import com.google.genai.types.VideoMetadata; import io.reactivex.rxjava3.core.Completable; import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Maybe; import io.reactivex.rxjava3.core.Scheduler; import io.reactivex.rxjava3.disposables.Disposable; import io.reactivex.rxjava3.plugins.RxJavaPlugins; @@ -1375,6 +1379,265 @@ public void nodeTimeout_sleepingBody_isInterruptedAndFailsWithNodeTimeoutExcepti assertThat(result.errorMessage("NodeTimeoutException")).contains("'sleeper'"); } + @Test + public void nodeTimeout_bodyIgnoringTheInterrupt_runFailsOnlyAfterItReturns() throws Exception { + // The body waits for the run to end, at most 2 s, so a run that ends first records it first. + List timeline = Collections.synchronizedList(new ArrayList<>()); + CountDownLatch runEnded = new CountDownLatch(1); + CountDownLatch bodyReturned = new CountDownLatch(1); + AtomicBoolean interrupted = new AtomicBoolean(); + FunctionNode stubborn = + FunctionNodes.builder(Content.class, String.class) + .name("stubborn") + .config(NodeConfig.builder().timeout(Duration.ofMillis(200)).build()) + .body((ctx, message) -> { + var unused = Uninterruptibles.awaitUninterruptibly(runEnded, 2, SECONDS); + interrupted.set(Thread.currentThread().isInterrupted()); + timeline.add("body returned"); + bodyReturned.countDown(); + return "late"; + }) + .build(); + ExecutorService executor = Executors.newCachedThreadPool(); + + RunResult result; + try { + result = run(chain(stubborn), Schedulers.from(executor)); + timeline.add("run ended"); + runEnded.countDown(); + assertThat(bodyReturned.await(10, SECONDS)).isTrue(); + } finally { + executor.shutdownNow(); + } + + assertThat(timeline).containsExactly("body returned", "run ended").inOrder(); + assertThat(interrupted.get()).isTrue(); + assertThat(result.errorMessage("NodeTimeoutException")).contains("'stubborn'"); + } + + @Test + public void nodeTimeoutWithRetry_bodyIgnoringTheInterrupt_retryStartsOnlyAfterItReturns() + throws Exception { + // Attempt 1 waits for attempt 2, at most 2 s, so an overlapping retry records its start first. + AtomicInteger attempts = new AtomicInteger(); + List timeline = Collections.synchronizedList(new ArrayList<>()); + CountDownLatch retryStarted = new CountDownLatch(1); + CountDownLatch firstReturned = new CountDownLatch(1); + AtomicBoolean firstInterrupted = new AtomicBoolean(); + FunctionNode stubborn = + FunctionNodes.builder(Content.class, String.class) + .name("stubborn") + .config( + NodeConfig.builder() + .timeout(Duration.ofMillis(200)) + .retryConfig( + RetryConfig.builder() + .maxAttempts(2) + .initialDelay(Duration.ZERO) + .jitter(0) + .build()) + .build()) + .body((ctx, message) -> { + int attempt = attempts.incrementAndGet(); + timeline.add("start " + attempt); + if (attempt == 1) { + var unused = Uninterruptibles.awaitUninterruptibly(retryStarted, 2, SECONDS); + firstInterrupted.set(Thread.currentThread().isInterrupted()); + timeline.add("end 1"); + firstReturned.countDown(); + return "late"; + } + retryStarted.countDown(); + return "retried"; + }) + .build(); + ExecutorService executor = Executors.newCachedThreadPool(); + + RunResult result; + try { + result = run(chain(stubborn), Schedulers.from(executor)); + assertThat(firstReturned.await(10, SECONDS)).isTrue(); + } finally { + executor.shutdownNow(); + } + + assertThat(timeline).containsExactly("start 1", "end 1", "start 2").inOrder(); + assertThat(firstInterrupted.get()).isTrue(); + result.assertCompleted(); + assertThat(result.outputAt("wf@1/stubborn@1")).isEqualTo("retried"); + } + + @Test + public void nodeTimeoutWithRetry_bodyRunsAnotherNodeUnderItsContext_retryWaitsForTheOuterBody() + throws Exception { + // "inner" runs under "outer"'s context and returns at once; "outer" then outlives its timeout. + AtomicInteger attempts = new AtomicInteger(); + List timeline = Collections.synchronizedList(new ArrayList<>()); + CountDownLatch retryStarted = new CountDownLatch(1); + CountDownLatch firstReturned = new CountDownLatch(1); + FunctionNode inner = + FunctionNodes.of("inner", Object.class, String.class, (ctx, in) -> "inner out"); + FunctionNode outer = + FunctionNodes.builder(Content.class, String.class) + .name("outer") + .config( + NodeConfig.builder() + .timeout(Duration.ofMillis(200)) + .retryConfig( + RetryConfig.builder() + .maxAttempts(2) + .initialDelay(Duration.ZERO) + .jitter(0) + .build()) + .build()) + .body((ctx, message) -> { + int attempt = attempts.incrementAndGet(); + timeline.add("start " + attempt); + if (attempt == 1) { + inner.runNode(ctx, null).blockingSubscribe(); + timeline.add("inner returned"); + var unused = Uninterruptibles.awaitUninterruptibly(retryStarted, 2, SECONDS); + timeline.add("end 1"); + firstReturned.countDown(); + return "late"; + } + retryStarted.countDown(); + return "retried"; + }) + .build(); + ExecutorService executor = Executors.newCachedThreadPool(); + + RunResult result; + try { + result = run(chain(outer), Schedulers.from(executor)); + assertThat(firstReturned.await(10, SECONDS)).isTrue(); + } finally { + executor.shutdownNow(); + } + + assertThat(timeline).containsExactly("start 1", "inner returned", "end 1", "start 2").inOrder(); + result.assertCompleted(); + assertThat(result.outputAt("wf@1/outer@1")).isEqualTo("retried"); + } + + @Test + public void nodeTimeout_oneThreadRunnerWithABusySibling_failsTheRunAndCancelsTheSibling() + throws Exception { + // Once "sleeper" times out, a task it queues and "busy" hold the one thread until released. + // Both nodes have a timeout, so both start through the same hops and "sleeper" starts first. + List timeline = Collections.synchronizedList(new ArrayList<>()); + CountDownLatch release = new CountDownLatch(1); + ExecutorService executor = Executors.newSingleThreadExecutor(); + FunctionNode sleeper = + FunctionNodes.builder(Content.class, String.class) + .name("sleeper") + .config(NodeConfig.builder().timeout(Duration.ofMillis(200)).build()) + .body((ctx, message) -> { + timeline.add("sleeper started"); + try { + Thread.sleep(60_000); + } catch (InterruptedException e) { + executor.execute(() -> Uninterruptibles.awaitUninterruptibly(release)); + } + return "slept"; + }) + .build(); + FunctionNode busy = + FunctionNodes.builder(Content.class, String.class) + .name("busy") + .config(NodeConfig.builder().timeout(Duration.ofHours(1)).build()) + .body((ctx, message) -> { + timeline.add("busy started"); + release.await(); + timeline.add("busy finished"); + return "done"; + }) + .build(); + Workflow workflow = + Workflow.builder() + .name("wf") + .edges(new Edge(Start.INSTANCE, sleeper), new Edge(Start.INSTANCE, busy)) + .build(); + Runner runner = runner(workflow, Schedulers.from(executor)); + Session session = runner.sessionService().createSession(APP, USER).blockingGet(); + TestSubscriber subscriber = new TestSubscriber<>(); + + boolean runEnded; + try { + runner + .runAsync(USER, session.id(), Content.fromParts(Part.fromText("go"))) + .subscribe(subscriber); + runEnded = subscriber.await(5, SECONDS); + } finally { + subscriber.cancel(); + release.countDown(); + executor.shutdownNow(); + } + + assertThat(timeline).isNotEmpty(); + assertThat(timeline.get(0)).isEqualTo("sleeper started"); + assertWithMessage("whether the run ended within 5 s").that(runEnded).isTrue(); + assertThat(timeline).doesNotContain("busy finished"); + subscriber.assertError(NodeTimeoutException.class); + } + + @Test + public void + inlineSchedulerWithInterruptedCaller_nodeTimeout_deliversTheFailureOnAnUninterruptedThread() + throws Exception { + // The body outlives its timeout on the caller's interrupted thread. + AtomicReference errorOnInterruptedThread = new AtomicReference<>(); + FunctionNode stubborn = + FunctionNodes.builder(Content.class, String.class) + .name("stubborn") + .config(NodeConfig.builder().timeout(Duration.ofMillis(200)).build()) + .body((ctx, message) -> { + Uninterruptibles.sleepUninterruptibly(400, MILLISECONDS); + return "late"; + }) + .build(); + BasePlugin observer = + new BasePlugin("observer") { + @Override + public Maybe onEventCallback(InvocationContext context, Event event) { + if (event.errorCode().isPresent()) { + errorOnInterruptedThread.set(Thread.currentThread().isInterrupted()); + } + return Maybe.empty(); + } + }; + Runner runner = + Runner.builder() + .app(App.builder().name(APP).rootNode(chain(stubborn)).plugins(observer).build()) + .scheduler(Schedulers.trampoline()) + .build(); + Session session = runner.sessionService().createSession(APP, USER).blockingGet(); + TestSubscriber subscriber = new TestSubscriber<>(); + ExecutorService caller = Executors.newSingleThreadExecutor(); + + boolean runEnded; + try { + var unused = + caller.submit(() -> { + Thread.currentThread().interrupt(); + runner + .runAsync(USER, session.id(), Content.fromParts(Part.fromText("go"))) + .subscribe(subscriber); + return Thread.interrupted(); + }); + runEnded = subscriber.await(5, SECONDS); + } finally { + subscriber.cancel(); + caller.shutdownNow(); + } + + assertWithMessage("whether the run ended within 5 s").that(runEnded).isTrue(); + subscriber.assertError(NodeTimeoutException.class); + assertWithMessage("whether the error event reached plugins on an interrupted thread") + .that(errorOnInterruptedThread.get()) + .isFalse(); + } + @Test public void runDisposed_bodyThatThrowsAfterTheInterrupt_reportsNoUndeliverableError() throws InterruptedException {