sangkyoonnam commented on code in PR #1348: URL: https://github.com/apache/flink-agents/pull/1348#discussion_r4478157170
########## 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: On recovery, a finished tool's record can be dropped here when the id at this index differs from the first run. I modeled `ToolCallAction.executeParallel` by hand in a local test (not part of this PR). On the first run the subagent submit fails before it reserves its slot, so the tool lands at index 0 as `tool-call:call-2`. On recovery the submit succeeds, index 0 is now `session#call-1`, `saved.matches(id)` is false, and `clearCallResultsFromAndPersist(index)` removes the tool's terminal record. The tool then runs again (`toolCalls` is 2, expected 1). I ran `submitFailingBeforeSlotReservationOnFirstRunKeepsCompletedToolOnRecovery` and `submitPersistenceFailureOnFirstRunKeepsCompletedToolOnRecovery`. Both failed in 3 of 3 runs. The controls `controlNoSubmitFailureKeepsCompletedToolOnRecovery` and `controlRecordedSubmitFailureKeepsCompletedToolOnRecovery` passed in 3 of 3 runs. The test doesn't go through the real `ToolCallAction`, so I haven't shown the first-run state is reachable there. Possible direction (inference): make a failed submit still consume its slot so the first run and recovery agree on indices, or look the record up by id instead of truncating the suffix. ########## 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: A call that completes after the sweep but before this `finish(timeout)` is recorded as a timeout failure. `isDone()` only says whether the outcome is settled, not whether `completion` finished. So `timeout != null` builds `Outcome.failure(timeout)` and, for a started durable call, writes it with `slot.finish(result)`. `cancel()` can't undo it, because `completion.cancel(true)` has no effect on a finished completion. The description says the scheduler "preserves available completions at timeout"; this window is after the sweep. The `finish(timeout)` branch is the same as before `0cecf5f`, and the test still failed on `ee2aeb8`. I injected that window in a local test (not part of this PR): a hook's `preventStart()` waits for the tool to finish, then sleeps 100 ms. `completionBetweenSweepAndTimeoutFinishIsNotRecordedAsTimeout` failed in 3 of 3 runs (`isFailure()` was true, expected false). The control `controlCompletionBeforeDeadlineIsRecordedAsSuccess` passed in 3 of 3 runs. The window is narrow and I injected it by hand. Possible direction (inference): when `timeout != null` but `completion` already finished normally, take the `finish(null)` path, and keep canceled completions as timeouts. -- 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]
