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 e26a84003..92e30c66b 100644 --- a/core/src/main/java/com/google/adk/workflow/FunctionNode.java +++ b/core/src/main/java/com/google/adk/workflow/FunctionNode.java @@ -23,10 +23,16 @@ import com.google.adk.agents.Context; import com.google.adk.annotations.Experimental; import com.google.adk.events.Event; +import com.google.common.base.Throwables; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.genai.types.Schema; import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.core.MaybeEmitter; import java.util.Optional; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.FutureTask; import org.jspecify.annotations.Nullable; /** @@ -124,15 +130,42 @@ boolean edgeInputSchemaInferred() { @Override public Flowable runNode(Context context, @Nullable Object nodeInput) { - return Flowable.defer( - () -> { - I input = inputCodec.decode(nodeInput, "input of node '" + name() + "'"); - O result = body.apply(context, input); - encodeDelegatedOutput(context.requireNodeState()); - // After a delegation, NodeRunner drops a returned output. - return Flowable.fromOptional(Optional.ofNullable(output(result))); - }) - .subscribeOn(context.invocationContext().scheduler()); + // Maybe.create rather than defer, so that canceling the node reaches the body while it runs. + return Maybe.create(emitter -> runBody(context, nodeInput, emitter)) + .subscribeOn(context.invocationContext().scheduler()) + .toFlowable(); + } + + /** + * Runs the body on the calling thread and sends its output or failure to {@code emitter}. + * Disposing {@code emitter} interrupts the body while it runs and drops its result. + */ + private void runBody(Context context, @Nullable Object nodeInput, MaybeEmitter emitter) { + Worker worker = + new Worker<>(() -> { + I input = inputCodec.decode(nodeInput, "input of node '" + name() + "'"); + return emitter.isDisposed() ? null : body.apply(context, input); + }); + emitter.setCancellable(worker::cancel); + Object output; + try { + O result = worker.run(); + if (emitter.isDisposed()) { + return; + } + encodeDelegatedOutput(context.requireNodeState()); + // After a delegation, NodeRunner drops a returned output. + output = output(result); + } catch (Exception e) { + // tryOnError, as a canceled run can no longer take a failure. + emitter.tryOnError(e); + return; + } + if (output == null) { + emitter.onComplete(); + } else { + emitter.onSuccess(output); + } } /** @@ -188,6 +221,44 @@ private void encodeDelegatedOutput(NodeExecutionState state) { return outputCodec.encode(result, what); } + /** + * Runs one body on the thread that calls {@link #run}, and {@link #cancel} interrupts that thread + * only while the body runs. When the body returns, the worker clears its own interrupt, so the + * thread's next task does not see it. It keeps an interrupt the caller had set before the body + * ran. A cancel does not wait for a body that ignores the interrupt. + */ + private static final class Worker { + private final FutureTask task; + + Worker(Callable body) { + this.task = new FutureTask<>(body); + } + + @SuppressWarnings("Interruption") // Interrupting the body is how a cancel reaches it. + void cancel() { + task.cancel(/* mayInterruptIfRunning= */ true); + } + + /** Runs the body and returns its result, or {@code null} if it was canceled. */ + @Nullable T run() throws Exception { + boolean callerInterrupted = Thread.currentThread().isInterrupted(); + task.run(); + if (task.isCancelled()) { + if (!callerInterrupted) { + var unused = Thread.interrupted(); + } + return null; + } + try { + return task.get(); // Done, so get does not block. + } catch (ExecutionException e) { + Throwables.throwIfInstanceOf(e.getCause(), Exception.class); + Throwables.throwIfUnchecked(e.getCause()); + throw e; + } + } + } + /** * Fluent builder for {@link FunctionNode}. Any property left unset falls back to the default of * the matching {@link Node} property; name and body are required. 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 9cba186d8..5999d037c 100644 --- a/core/src/main/java/com/google/adk/workflow/NodeBody.java +++ b/core/src/main/java/com/google/adk/workflow/NodeBody.java @@ -24,7 +24,8 @@ * 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. + * 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. */ @Experimental @FunctionalInterface 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 805c00bdb..4597fbbc6 100644 --- a/core/src/test/java/com/google/adk/workflow/FunctionNodeTest.java +++ b/core/src/test/java/com/google/adk/workflow/FunctionNodeTest.java @@ -60,6 +60,7 @@ import com.google.adk.workflow.WorkflowResumeTesting.WorkflowSession; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.util.concurrent.Uninterruptibles; import com.google.genai.types.Content; import com.google.genai.types.FinishReason; import com.google.genai.types.Part; @@ -67,7 +68,10 @@ import com.google.genai.types.Type; import com.google.genai.types.VideoMetadata; import io.reactivex.rxjava3.core.Scheduler; +import io.reactivex.rxjava3.disposables.Disposable; +import io.reactivex.rxjava3.plugins.RxJavaPlugins; import io.reactivex.rxjava3.schedulers.Schedulers; +import io.reactivex.rxjava3.subscribers.TestSubscriber; import java.io.IOException; import java.lang.annotation.Retention; import java.lang.annotation.RetentionPolicy; @@ -84,9 +88,14 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import org.jspecify.annotations.Nullable; @@ -1069,6 +1078,406 @@ public void of_body_runsOnTheRunnersScheduler() { assertThat(result.outputAt("wf@1/where@1")).isEqualTo("app-scheduler"); } + @Test + public void runDisposed_sleepingBody_isInterrupted() throws InterruptedException { + // Arrange: an executor scheduler, which does not interrupt on its own when disposed. + CountDownLatch started = new CountDownLatch(1); + CountDownLatch interrupted = new CountDownLatch(1); + FunctionNode sleeper = + FunctionNodes.of( + "sleeper", + Content.class, + String.class, + (ctx, message) -> { + started.countDown(); + try { + Thread.sleep(60_000); + } catch (InterruptedException e) { + interrupted.countDown(); + } + return "late"; + }); + ExecutorService executor = Executors.newCachedThreadPool(); + try { + Runner runner = runner(chain(sleeper), Schedulers.from(executor)); + Session session = runner.sessionService().createSession(APP, USER).blockingGet(); + Disposable run = + runner.runAsync(USER, session.id(), Content.fromParts(Part.fromText("go"))).subscribe(); + assertThat(started.await(10, SECONDS)).isTrue(); + + run.dispose(); + + assertThat(interrupted.await(10, SECONDS)).isTrue(); + } finally { + executor.shutdownNow(); + } + } + + @Test + public void runDisposed_bodyThatKeepsTheInterrupt_leavesTheThreadUninterruptedForTheNextTask() + throws Exception { + // Arrange: one thread runs every task and keeps an interrupt between them, unlike a pool. + BlockingQueue tasks = new LinkedBlockingQueue<>(); + AtomicBoolean stopped = new AtomicBoolean(); + Thread worker = + new Thread(() -> { + while (!stopped.get()) { + Uninterruptibles.takeUninterruptibly(tasks).run(); + } + }); + worker.setDaemon(true); + worker.start(); + CountDownLatch started = new CountDownLatch(1); + CountDownLatch finished = new CountDownLatch(1); + AtomicBoolean sawInterrupt = new AtomicBoolean(); + FunctionNode spinner = + FunctionNodes.of( + "spinner", + Content.class, + String.class, + (ctx, message) -> { + started.countDown(); + long deadline = System.nanoTime() + SECONDS.toNanos(10); + while (!Thread.currentThread().isInterrupted() && System.nanoTime() < deadline) { + Thread.onSpinWait(); + } + sawInterrupt.set(Thread.currentThread().isInterrupted()); + finished.countDown(); + return "late"; + }); + try { + Runner runner = runner(chain(spinner), Schedulers.from(tasks::add)); + Session session = runner.sessionService().createSession(APP, USER).blockingGet(); + Disposable run = + runner.runAsync(USER, session.id(), Content.fromParts(Part.fromText("go"))).subscribe(); + assertThat(started.await(10, SECONDS)).isTrue(); + + run.dispose(); + assertThat(finished.await(10, SECONDS)).isTrue(); + CompletableFuture nextTaskInterrupted = new CompletableFuture<>(); + tasks.add(() -> nextTaskInterrupted.complete(Thread.currentThread().isInterrupted())); + + assertThat(sawInterrupt.get()).isTrue(); + assertThat(nextTaskInterrupted.get(10, SECONDS)).isFalse(); + } finally { + stopped.set(true); + tasks.add(() -> {}); + } + } + + @Test + public void inlineScheduler_runDisposedFromAnotherThread_interruptsTheBodyAndClearsTheInterrupt() + throws Exception { + // Arrange: a trampoline scheduler runs the body on the caller's thread, which subscribes. + CountDownLatch started = new CountDownLatch(1); + AtomicBoolean sawInterrupt = new AtomicBoolean(); + FunctionNode sleeper = + FunctionNodes.of( + "sleeper", + Content.class, + String.class, + (ctx, message) -> { + started.countDown(); + try { + Thread.sleep(60_000); + } catch (InterruptedException e) { + sawInterrupt.set(true); + Thread.currentThread().interrupt(); + } + return "late"; + }); + Runner runner = runner(chain(sleeper), Schedulers.trampoline()); + Session session = runner.sessionService().createSession(APP, USER).blockingGet(); + TestSubscriber subscriber = new TestSubscriber<>(); + ExecutorService caller = Executors.newSingleThreadExecutor(); + try { + Future callerInterrupted = + caller.submit(() -> { + runner + .runAsync(USER, session.id(), Content.fromParts(Part.fromText("go"))) + .subscribe(subscriber); + return Thread.interrupted(); + }); + assertThat(started.await(10, SECONDS)).isTrue(); + + subscriber.cancel(); + + assertThat(callerInterrupted.get(10, SECONDS)).isFalse(); + assertThat(sawInterrupt.get()).isTrue(); + } finally { + caller.shutdownNow(); + } + } + + @Test + public void inlineScheduler_callerAlreadyInterrupted_keepsTheInterruptAfterTheBodyReturns() + throws Exception { + // Arrange: a trampoline scheduler runs the body on the caller's thread, interrupted beforehand. + AtomicBoolean bodySawInterrupt = new AtomicBoolean(); + FunctionNode quick = + FunctionNodes.of( + "quick", + Content.class, + String.class, + (ctx, message) -> { + bodySawInterrupt.set(Thread.currentThread().isInterrupted()); + return "done"; + }); + Runner runner = runner(chain(quick), Schedulers.trampoline()); + Session session = runner.sessionService().createSession(APP, USER).blockingGet(); + TestSubscriber subscriber = new TestSubscriber<>(); + ExecutorService caller = Executors.newSingleThreadExecutor(); + Future callerInterrupted; + try { + callerInterrupted = + caller.submit(() -> { + Thread.currentThread().interrupt(); + runner + .runAsync(USER, session.id(), Content.fromParts(Part.fromText("go"))) + .subscribe(subscriber); + return Thread.interrupted(); + }); + + assertThat(callerInterrupted.get(10, SECONDS)).isTrue(); + } finally { + caller.shutdownNow(); + } + + subscriber.assertComplete().assertNoErrors(); + assertThat(bodySawInterrupt.get()).isTrue(); + } + + @Test + public void inlineScheduler_callerInterruptedThenRunDisposed_keepsTheCallersInterrupt() + throws Exception { + // Arrange: the caller's interrupt and the cancel's both reach the body on the caller's thread. + CountDownLatch started = new CountDownLatch(1); + CountDownLatch canceled = new CountDownLatch(1); + AtomicBoolean sawCancel = new AtomicBoolean(); + AtomicBoolean cancelInterruptedTheBody = new AtomicBoolean(); + FunctionNode waiter = + FunctionNodes.of( + "waiter", + Content.class, + String.class, + (ctx, message) -> { + boolean callerInterrupt = Thread.interrupted(); + started.countDown(); + sawCancel.set(Uninterruptibles.awaitUninterruptibly(canceled, 10, SECONDS)); + cancelInterruptedTheBody.set(Thread.currentThread().isInterrupted()); + if (callerInterrupt) { + Thread.currentThread().interrupt(); + } + return "late"; + }); + Runner runner = runner(chain(waiter), Schedulers.trampoline()); + Session session = runner.sessionService().createSession(APP, USER).blockingGet(); + TestSubscriber subscriber = new TestSubscriber<>(); + ExecutorService caller = Executors.newSingleThreadExecutor(); + try { + Future callerInterrupted = + caller.submit(() -> { + Thread.currentThread().interrupt(); + runner + .runAsync(USER, session.id(), Content.fromParts(Part.fromText("go"))) + .subscribe(subscriber); + return Thread.interrupted(); + }); + assertThat(started.await(10, SECONDS)).isTrue(); + + subscriber.cancel(); + canceled.countDown(); + + assertThat(callerInterrupted.get(10, SECONDS)).isTrue(); + assertThat(sawCancel.get()).isTrue(); + assertThat(cancelInterruptedTheBody.get()).isTrue(); + } finally { + caller.shutdownNow(); + } + } + + @Test + public void inlineScheduler_bodyInterruptsItself_callerSeesTheInterrupt() throws Exception { + // Arrange: nothing cancels the node, so only the body interrupts the caller's thread. + FunctionNode restorer = + FunctionNodes.of( + "restorer", + Content.class, + String.class, + (ctx, message) -> { + Thread.currentThread().interrupt(); + return "done"; + }); + Runner runner = runner(chain(restorer), Schedulers.trampoline()); + Session session = runner.sessionService().createSession(APP, USER).blockingGet(); + TestSubscriber subscriber = new TestSubscriber<>(); + ExecutorService caller = Executors.newSingleThreadExecutor(); + Future callerInterrupted; + try { + callerInterrupted = + caller.submit(() -> { + runner + .runAsync(USER, session.id(), Content.fromParts(Part.fromText("go"))) + .subscribe(subscriber); + return Thread.interrupted(); + }); + + assertThat(callerInterrupted.get(10, SECONDS)).isTrue(); + } finally { + caller.shutdownNow(); + } + + subscriber.assertComplete().assertNoErrors(); + } + + @Test + public void inlineScheduler_bodyReturns_leavesTheCallerUninterrupted() throws Exception { + // Arrange: RxJava calls the node's cancellable after the body's result is sent. + FunctionNode quick = + FunctionNodes.of("quick", Content.class, String.class, (ctx, message) -> "done"); + Runner runner = runner(chain(quick), Schedulers.trampoline()); + Session session = runner.sessionService().createSession(APP, USER).blockingGet(); + TestSubscriber subscriber = new TestSubscriber<>(); + ExecutorService caller = Executors.newSingleThreadExecutor(); + Future callerInterrupted; + try { + callerInterrupted = + caller.submit(() -> { + runner + .runAsync(USER, session.id(), Content.fromParts(Part.fromText("go"))) + .subscribe(subscriber); + return Thread.interrupted(); + }); + + assertThat(callerInterrupted.get(10, SECONDS)).isFalse(); + } finally { + caller.shutdownNow(); + } + + subscriber.assertComplete().assertNoErrors(); + } + + @Test + public void nodeTimeout_sleepingBody_isInterruptedAndFailsWithNodeTimeoutException() + throws InterruptedException { + CountDownLatch interrupted = new CountDownLatch(1); + FunctionNode sleeper = + FunctionNodes.builder(Content.class, String.class) + .name("sleeper") + .config(NodeConfig.builder().timeout(Duration.ofSeconds(1)).build()) + .body((ctx, message) -> { + try { + Thread.sleep(60_000); + } catch (InterruptedException e) { + interrupted.countDown(); + } + return "late"; + }) + .build(); + ExecutorService executor = Executors.newCachedThreadPool(); + + RunResult result; + try { + result = run(chain(sleeper), Schedulers.from(executor)); + assertThat(interrupted.await(10, SECONDS)).isTrue(); + } finally { + executor.shutdownNow(); + } + + assertThat(result.error()).isInstanceOf(NodeTimeoutException.class); + assertThat(result.errorMessage("NodeTimeoutException")).contains("'sleeper'"); + } + + @Test + public void runDisposed_bodyThatThrowsAfterTheInterrupt_reportsNoUndeliverableError() + throws InterruptedException { + CountDownLatch started = new CountDownLatch(1); + CountDownLatch finished = new CountDownLatch(1); + List undeliverable = new CopyOnWriteArrayList<>(); + var previousErrorHandler = RxJavaPlugins.getErrorHandler(); + FunctionNode sleeper = + FunctionNodes.of( + "sleeper", + Content.class, + String.class, + (ctx, message) -> { + started.countDown(); + try { + Thread.sleep(60_000); + return "late"; + } finally { + finished.countDown(); + } + }); + ExecutorService executor = Executors.newCachedThreadPool(); + RxJavaPlugins.setErrorHandler(undeliverable::add); + try { + Runner runner = runner(chain(sleeper), Schedulers.from(executor)); + Session session = runner.sessionService().createSession(APP, USER).blockingGet(); + Disposable run = + runner.runAsync(USER, session.id(), Content.fromParts(Part.fromText("go"))).subscribe(); + assertThat(started.await(10, SECONDS)).isTrue(); + + run.dispose(); + assertThat(finished.await(10, SECONDS)).isTrue(); + executor.shutdown(); + assertThat(executor.awaitTermination(10, SECONDS)).isTrue(); + + assertThat(undeliverable).isEmpty(); + } finally { + RxJavaPlugins.setErrorHandler(previousErrorHandler); + executor.shutdownNow(); + } + } + + @Test + public void siblingFails_sleepingBody_isInterrupted() throws InterruptedException { + // Arrange: the workflow cancels the nodes still running when one fails. + CountDownLatch started = new CountDownLatch(1); + CountDownLatch interrupted = new CountDownLatch(1); + FunctionNode sleeper = + FunctionNodes.of( + "sleeper", + Content.class, + Void.class, + (ctx, message) -> { + started.countDown(); + try { + Thread.sleep(60_000); + } catch (InterruptedException e) { + interrupted.countDown(); + } + return null; + }); + FunctionNode broken = + FunctionNodes.of( + "broken", + Content.class, + Void.class, + (ctx, message) -> { + if (!started.await(10, SECONDS)) { + throw new IllegalStateException("sleeper never started"); + } + throw new IOException("down"); + }); + Workflow workflow = + Workflow.builder() + .name("wf") + .edges(new Edge(Start.INSTANCE, sleeper), new Edge(Start.INSTANCE, broken)) + .build(); + ExecutorService executor = Executors.newCachedThreadPool(); + + RunResult result; + try { + result = run(workflow, Schedulers.from(executor)); + assertThat(interrupted.await(10, SECONDS)).isTrue(); + } finally { + executor.shutdownNow(); + } + + assertThat(result.errorMessage("IOException")).isEqualTo("down"); + } + @Test public void toBuilder_copy_keepsThePortsAndTheBody() { FunctionNode describe = @@ -3211,11 +3620,7 @@ private static RunResult run(Node root) { /** Runs {@code root} as an app's root node through a {@link Runner} on a "go" message. */ private static RunResult run(Node root, @Nullable Scheduler scheduler) { - Runner runner = - Runner.builder() - .app(App.builder().name(APP).rootNode(root).build()) - .scheduler(scheduler) - .build(); + Runner runner = runner(root, scheduler); Session session = runner.sessionService().createSession(APP, USER).blockingGet(); List events = new ArrayList<>(); AtomicReference error = new AtomicReference<>(); @@ -3226,6 +3631,13 @@ private static RunResult run(Node root, @Nullable Scheduler scheduler) { return new RunResult(ImmutableList.copyOf(events), error.get()); } + private static Runner runner(Node root, @Nullable Scheduler scheduler) { + return Runner.builder() + .app(App.builder().name(APP).rootNode(root).build()) + .scheduler(scheduler) + .build(); + } + /** The events a run streamed and the error it failed with, if any. */ private record RunResult(ImmutableList events, @Nullable Throwable error) {