wenjin272 commented on code in PR #1348: URL: https://github.com/apache/flink-agents/pull/1348#discussion_r4479936509
########## runtime/src/main/java/org/apache/flink/agents/runtime/context/DurableCallSlot.java: ########## @@ -0,0 +1,73 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.runtime.context; + +import org.apache.flink.agents.api.context.Outcome; +import org.apache.flink.agents.runtime.actionstate.CallResult; + +/** + * A reserved, indexed call result. Preparing never launches work; completion never moves the + * cursor. + */ +public final class DurableCallSlot<T> { + private final RunnerContextImpl owner; + private final String id; + private final Class<T> type; + private int index; + private boolean recovering; + + public DurableCallSlot(RunnerContextImpl owner, String id, Class<T> type) { + this.owner = owner; + this.id = id; + this.type = type; + } + + public Outcome<T> prepare() throws Exception { + index = owner.getCurrentCallIndex(); + CallResult saved = owner.getCallResultAt(index); + if (saved != null && !saved.matches(id)) { + owner.clearCallResultsFromAndPersist(index); Review Comment: Thanks for catching this. I confirmed the path through the real ToolCallAction and operator recovery: a user-provided deferred `prepare()` or overridden async `submitRequest()` can throw before the durable call is created. Fixed in 56a20ddad using your first direction. The framework now reserves the call's durable slot before invoking preparation, records ordinary preparation failures in that same slot, and skips preparation when replaying a terminal outcome. The same handling covers `awaitResult()` preparation. Framework faults and cancellation still propagate without becoming terminal failure records. `SubagentPreparationRecoveryTest` covers sequential and parallel tool dispatch, with a crash after the following ordinary tool's result is persisted. Recovery replays the preparation failure and reuses the tool result without executing it again. It also covers pending-state persistence failures. The focused suites pass on JDK 17 and 21. ########## runtime/src/main/java/org/apache/flink/agents/runtime/context/AsyncFutureImpl.java: ########## @@ -21,189 +21,406 @@ import org.apache.flink.agents.api.context.DurableCallable; import org.apache.flink.agents.api.context.DurableFuture; import org.apache.flink.agents.api.context.Outcome; +import org.apache.flink.agents.api.context.RunnerContext; import org.apache.flink.agents.plan.utils.CancellationUtils; -import org.apache.flink.util.Preconditions; +import org.apache.flink.agents.runtime.operator.parallel.UncapturableException; +import org.apache.flink.util.function.FunctionWithException; import java.util.ArrayList; import java.util.Collections; import java.util.IdentityHashMap; import java.util.List; import java.util.Set; import java.util.concurrent.Callable; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeoutException; -/** Runtime-owned implementation of a deferred asynchronous call handle. */ -abstract class AsyncFutureImpl<T> implements AsyncFuture<T> { +/** Shared await behavior and local result state for runtime call handles. */ +public abstract class AsyncFutureImpl<T> implements AsyncFuture<T> { private final RunnerContextImpl owner; private boolean resolving; - private boolean done; - private T value; - private Exception error; + private Outcome<T> outcome; - AsyncFutureImpl(RunnerContextImpl owner) { + protected AsyncFutureImpl(RunnerContextImpl owner) { this.owner = owner; } - final RunnerContextImpl getOwner() { + public final RunnerContextImpl getOwner() { return owner; } - final boolean isDone() { - return done; + public final boolean isDone() { + return outcome != null; } - final Outcome<T> getCompletedOutcome() { - if (!done) { + public final Outcome<T> getCompletedOutcome() { + if (outcome == null) { throw new IllegalStateException("Async future has not been resolved"); } - return error == null ? Outcome.success(value) : Outcome.failure(error); + return outcome; } + public RunnerContextImpl getOwnerContext() { + return owner; + } + + protected final void complete(Outcome<T> value) { + outcome = value; + } + + protected boolean cacheResolutionFailures() { + return true; + } + + protected abstract T resolveValue() throws Exception; + @Override public final T await() throws Exception { - if (done) { - return completedValue(); + owner.checkMailboxThread(); + if (isDone()) { + return value(); } if (resolving) { throw new IllegalStateException("An async future cannot await itself recursively"); } - resolving = true; try { - value = resolveValue(); - done = true; - return value; + T result = resolveValue(); + if (!isDone()) { + complete(Outcome.success(result)); + } + return result; } catch (Exception e) { - if (CancellationUtils.isCancellation(e)) { - if (CancellationUtils.isInterruption(e)) { - Thread.currentThread().interrupt(); - } - throw e; + rethrowFatal(e); + if (!isDone() && cacheResolutionFailures()) { + complete(Outcome.failure(e)); } - error = e; - done = true; throw e; } finally { resolving = false; } } - final void complete(Outcome<T> outcome) { - if (done) { - throw new IllegalStateException("Async future has already been resolved"); + protected final T value() throws Exception { + if (outcome.isFailure()) { + throw outcome.getError(); } - value = outcome.getValue(); - error = outcome.getError(); - done = true; + return outcome.getValue(); } - private T completedValue() throws Exception { - if (error != null) { - throw error; + public static void rethrowFatal(Throwable failure) throws Exception { + if (failure instanceof Error) { + throw (Error) failure; + } + if (CancellationUtils.isCancellation(failure) + || UncapturableException.isUncapturable(failure)) { + if (CancellationUtils.isInterruption(failure)) { + Thread.currentThread().interrupt(); + } + throw (Exception) failure; } - return value; } - abstract T resolveValue() throws Exception; + protected static <T> Outcome<T> getOutcome(CompletableFuture<T> future) throws Exception { + // CompletableFuture.get() unwraps an explicitly completed CompletionException. Capture + // the original exception as an Outcome before crossing that transport boundary. + Outcome<T> result; + try { + result = + future.handle( + (value, failure) -> { + if (failure == null) { + return Outcome.success(value); + } + if (failure instanceof Error) { + throw (Error) failure; + } + return Outcome.<T>failure((Exception) failure); + }) + .get(); + } catch (ExecutionException e) { + rethrowFatal(e.getCause()); + throw (Exception) e.getCause(); + } + if (result.isFailure()) { + rethrowFatal(result.getError()); + } + return result; + } + + /** + * Maps a call result without adding a durable boundary, preserving runtime batch scheduling. + */ + public static <S, T> AsyncFuture<T> map( + RunnerContext owner, + AsyncFuture<S> future, + FunctionWithException<S, T, Exception> mapper) + throws Exception { + if (!(owner instanceof RunnerContextImpl)) { + // Other context implementations manage their own waiting and gather behavior. + return () -> mapper.apply(future.await()); + } + RunnerContextImpl runtime = (RunnerContextImpl) owner; + SchedulableFuture<S> source = requireSchedulable(future, runtime); + return new MappedAsyncFuture<>(runtime, source, mapper); + } + + public static <T> SchedulableFuture<T> requireSchedulable( + AsyncFuture<T> future, RunnerContext owner) { + if (!(future instanceof SchedulableFuture)) { + throw new IllegalArgumentException("gather only accepts runtime-managed call futures"); + } + SchedulableFuture<T> participant = (SchedulableFuture<T>) future; + if (participant.getOwnerContext() != null && participant.getOwnerContext() != owner) { + throw new IllegalArgumentException( + "An async future must be gathered by the context that created it"); + } + return participant; + } } -/** Deferred handle for one ordinary callable. */ -class SingleAsyncFuture<T> extends AsyncFutureImpl<T> { +/** Ordinary callable; durable subclasses reuse the same scheduling lifecycle. */ +class SingleAsyncFuture<T> extends AsyncFutureImpl<T> implements SchedulableFuture<T> { private final Callable<T> callable; + private Callable<T> execution; + private DurableCallSlot<T> slot; + private CompletableFuture<T> completion; + private boolean prepared; + private volatile boolean started; + private boolean stopped; SingleAsyncFuture(RunnerContextImpl owner, Callable<T> callable) { super(owner); this.callable = callable; } - Callable<T> getCallable() { - return callable; - } - DurableCallable<T> getDurableCallable() { return null; } @Override - T resolveValue() throws Exception { - return getOwner().resolveAsync(callable); + protected T resolveValue() throws Exception { + return getOwner().resolveAsyncFuture(this); + } + + @Override + public Outcome<T> prepare(RunnerContextImpl context) throws Exception { + if (prepared || isDone()) { + return isDone() ? getCompletedOutcome() : null; + } + execution = callable; + DurableCallable<T> durable = getDurableCallable(); + if (durable != null) { + slot = new DurableCallSlot<>(getOwner(), durable.getId(), durable.getResultClass()); + Outcome<T> recorded = slot.prepare(); + if (recorded != null) { + complete(recorded); + } + if (slot.isRecovering() && durable.reconciler() != null) { + execution = durable.reconciler(); + } + } + prepared = true; + return isDone() ? getCompletedOutcome() : null; + } + + @Override + public CompletableFuture<?> start() { + if (isDone()) { + return CompletableFuture.completedFuture(null); + } + synchronized (this) { + stopped = false; + } + Callable<T> work = + () -> { + synchronized (this) { + if (stopped) { + throw new CancellationException("Call stopped before execution"); + } + started = true; + } + return execution.call(); + }; + completion = getOwner().submitAsync(work); + return completion; + } + + @Override + public Outcome<T> finish(TimeoutException timeout) throws Exception { + if (isDone()) { + return getCompletedOutcome(); + } + if (timeout != null) { Review Comment: Thanks, confirmed and fixed in 8ebe3ad27. After cancellation, the shared scheduler rechecks each completion signal. If it completed and was not cancelled, we use `finish(null)` to preserve the actual result or exception; cancelled or unfinished calls still take the timeout path. Added deterministic race tests in `MixedFutureGatherTest` for successful completion, ordinary failure, fatal Error, and Chat/internal-subagent completion between the timeout sweep and cancellation. They pass on JDK 21, including the latest rerun. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
