diff --git a/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py b/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py index 92c0e154a..f0456d870 100644 --- a/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py +++ b/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py @@ -246,17 +246,47 @@ def invoke() -> tuple: assert store.call_results[0].exception_payload is None -def _run_async(result: Any) -> object: - iterator = result.__await__() - value = None +def _drain(iterator: Any) -> object: while True: try: - iterator.send(value) - value = None + iterator.send(None) except StopIteration as e: # noqa: PERF203 return e.value +def _run_async(result: Any) -> object: + return _drain(result.__await__()) + + +_WORKER_START_TIMEOUT_SECONDS = 30.0 + + +def _await_workers(started: list[threading.Event]) -> None: + deadline = time.monotonic() + _WORKER_START_TIMEOUT_SECONDS + for index, event in enumerate(started): + remaining = max(deadline - time.monotonic(), 0.0) + assert event.wait(remaining), f"async worker {index} never started" + + +def _run_async_after_workers_started( + result: Any, started: list[threading.Event] +) -> object: + """Run a batch, pausing after submission until its workers have begun. + + A batch deadline is wall-clock from the first submission, so a test asserting + that a slot was *started* when the timeout hits cannot leave thread startup + racing that deadline: under CI load the worker has not begun yet and the slot + is persisted as PENDING instead of FAILED. One generator step submits the + batch; waiting there until every worker reports it began makes "started but + unfinished" hold by construction, and the deadline has certainly passed by the + time driving resumes. + """ + iterator = result.__await__() + iterator.send(None) + _await_workers(started) + return _drain(iterator) + + def _preload_pending( j_runner_context: _FakeJavaRunnerContext, func: Callable[..., Any], @@ -1052,27 +1082,32 @@ def fail_call() -> str: def test_flink_runner_context_gather_timeout_keeps_completed_results() -> None: j_runner_context = _FakeJavaRunnerContext() - config = AgentConfiguration({"async.batch.timeout.ms": 100}) + config = AgentConfiguration( + {"async.batch.timeout.ms": 100, "async.batch.parallelism": 2} + ) ctx = _create_runner_context(j_runner_context, config=config) # The slow worker blocks on an event instead of racing a sleep against the - # deadline. Once its worker thread starts, _mark_started_on_run has already - # flipped the started flag, so the timeout must record it as FAILED (started - # but unfinished) rather than PENDING. The event guarantees it never - # finishes before the deadline, and the finally block releases it so the - # pool can shut down. + # deadline, and announces itself so the driver can wait for it. Once its worker + # thread starts, _mark_started_on_run has already flipped the started flag, so + # the timeout must record it as FAILED (started but unfinished) rather than + # PENDING. The event guarantees it never finishes before the deadline, and the + # finally block releases it so the pool can shut down. release = threading.Event() + slow_started = threading.Event() def slow_call() -> str: + slow_started.set() release.wait(5) return "slow" try: - outcomes = _run_async( + outcomes = _run_async_after_workers_started( ctx.gather( ctx.durable_execute_async(lambda: "fast"), ctx.durable_execute_async(slow_call), - ) + ), + [slow_started], ) finally: release.set() @@ -1096,22 +1131,26 @@ def test_flink_runner_context_gather_timeout_leaves_unsubmitted_slots_pending() # The two blocking calls saturate the parallelism budget and hold their # worker threads (started flag flipped) until the deadline, so they fail; the # remaining two suppliers are never submitted and stay pending. Blocking on an - # event instead of sleeping makes "started" deterministic rather than racing - # thread startup against a short wall-clock timeout. + # event instead of sleeping, and waiting for both workers to announce + # themselves, makes "started" deterministic rather than racing thread startup + # against a short wall-clock timeout. release = threading.Event() + started_events = {"one": threading.Event(), "two": threading.Event()} def blocking_call(value: str) -> str: + started_events[value].set() release.wait(5) return value try: - outcomes = _run_async( + outcomes = _run_async_after_workers_started( ctx.gather( ctx.durable_execute_async(blocking_call, "one"), ctx.durable_execute_async(blocking_call, "two"), ctx.durable_execute_async(_call_value, "three"), ctx.durable_execute_async(_call_value, "four"), - ) + ), + [started_events["one"], started_events["two"]], ) finally: release.set() @@ -1141,12 +1180,15 @@ def test_flink_runner_context_gather_timeout_leaves_queued_slots_pending() -> No # Blocking on an event (rather than sleeping) keeps the two running workers # deterministically "started but unfinished" until the deadline, while the two - # queued suppliers never begin. + # queued suppliers never begin. Waiting for both workers to announce themselves + # removes the race between thread startup and the wall-clock deadline. release = threading.Event() + started_events = {"one": threading.Event(), "two": threading.Event()} executions: list[str] = [] def blocking_call(value: str) -> str: executions.append(value) + started_events[value].set() release.wait(5) return value @@ -1155,13 +1197,14 @@ def queued_call(value: str) -> str: return value try: - outcomes = _run_async( + outcomes = _run_async_after_workers_started( ctx.gather( ctx.durable_execute_async(blocking_call, "one"), ctx.durable_execute_async(blocking_call, "two"), ctx.durable_execute_async(queued_call, "three"), ctx.durable_execute_async(queued_call, "four"), - ) + ), + [started_events["one"], started_events["two"]], ) finally: release.set() @@ -1532,15 +1575,19 @@ def durable() -> str: try: # The cached durable child takes no worker; the ordinary child occupies the - # only worker, leaving the new durable child queued until the timeout. + # only worker, leaving the new durable child queued until the timeout. The + # driver waits for that worker to announce itself, so "the only worker is + # occupied when the deadline trips" does not race thread startup. assert ctx.durable_execute(lambda: "cached", durable_id="cached") == "cached" j_ctx.current_call_index = 0 - outcomes = _run_async(ctx.gather( - ctx.durable_execute_async(lambda: "cached", durable_id="cached"), - ctx.execute_async(ordinary), - ctx.durable_execute_async(durable, durable_id="pending"), - )) - assert started.is_set() + outcomes = _run_async_after_workers_started( + ctx.gather( + ctx.durable_execute_async(lambda: "cached", durable_id="cached"), + ctx.execute_async(ordinary), + ctx.durable_execute_async(durable, durable_id="pending"), + ), + [started], + ) assert outcomes[0].value == "cached" assert isinstance(outcomes[1].error, TimeoutError) assert isinstance(outcomes[2].error, TimeoutError) diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java b/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java index 37769a072..6132ae039 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java @@ -25,6 +25,7 @@ import java.time.Duration; import java.util.List; import java.util.concurrent.Callable; +import java.util.concurrent.CompletableFuture; import java.util.function.Supplier; /** @@ -136,6 +137,16 @@ public BatchExecutionResult executeAllAsync( } } + /** + * Awaits a {@link CompletableFuture} by blocking the calling thread. On JDK 21+ the + * multi-release variant yields the continuation instead, avoiding async-pool exhaustion (issue + * #1213). + */ + public T awaitFuture(ContinuationContext context, CompletableFuture future) + throws Exception { + return future.get(); + } + /** Stops async work without waiting for worker exit cleanup. */ public void shutdown() { batchExecutor.shutdown(); diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationNativeCallable.java b/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationNativeCallable.java new file mode 100644 index 000000000..d9bd89e36 --- /dev/null +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationNativeCallable.java @@ -0,0 +1,34 @@ +/* + * 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.async; + +import org.apache.flink.agents.api.context.DurableCallable; + +/** + * Marks a {@link DurableCallable} that suspends by yielding the caller's continuation, and must + * therefore run inline on the mailbox thread instead of being dispatched to an async-pool worker. + * + *

A waiting caller that occupies a pool worker holds a resource its own callee may need to make + * progress, which self-deadlocks once concurrent or nested waits exhaust the pool (issue #1213). A + * callable marked here awaits by yielding, so the worker pool stays free for the callee. + * + *

Only meaningful on JDK 21+, where the continuation the callable yields is the caller's; the + * callers that produce such callables gate on {@code requireMailboxSuspension} beforehand. + */ +public interface ContinuationNativeCallable extends DurableCallable {} diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java b/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java index 27b402969..d4f541637 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java @@ -27,6 +27,7 @@ import org.apache.flink.agents.runtime.async.BatchExecutionResult; import org.apache.flink.agents.runtime.async.ContinuationActionExecutor; import org.apache.flink.agents.runtime.async.ContinuationContext; +import org.apache.flink.agents.runtime.async.ContinuationNativeCallable; import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl; import java.time.Duration; @@ -35,6 +36,7 @@ import java.util.Collections; import java.util.List; import java.util.concurrent.Callable; +import java.util.concurrent.CompletableFuture; import java.util.function.Supplier; import java.util.stream.Collectors; @@ -71,19 +73,29 @@ public ContinuationContext getContinuationContext() { @Override protected T resolveDurableAsync(DurableCallable callable) throws Exception { + Callable boundary = executionBoundary(callable); if (durableExecutionContext != null) { Callable reconcileCallable = callable.reconciler(); if (reconcileCallable != null) { - return durableExecuteAsyncWithReconcile(callable, reconcileCallable); + return durableExecuteWithReconcile(callable, reconcileCallable, boundary); } } - return durableExecuteCompletionOnly(callable, () -> resolveAsync(callable::call)); + return durableExecuteCompletionOnly(callable, boundary); } - private T durableExecuteAsyncWithReconcile( - DurableCallable callable, Callable reconcileCallable) throws Exception { - return durableExecuteWithReconcile( - callable, reconcileCallable, () -> resolveAsync(callable::call)); + /** + * Chooses where one durable async call executes, on both the completion-only and the reconcile + * path. A {@link ContinuationNativeCallable} suspends by yielding the caller's continuation, so + * it runs inline on the mailbox thread; handing it to an async-pool worker would make the + * waiting caller hold a worker its own callee needs to progress (issue #1213), and the yield + * itself would fail off the mailbox thread. Any other callable is offloaded to the pool, which + * lets the continuation yield while it runs. + */ + private Callable executionBoundary(DurableCallable callable) { + if (callable instanceof ContinuationNativeCallable) { + return callable::call; + } + return () -> resolveAsync(callable::call); } @Override @@ -350,6 +362,17 @@ protected T resolveAsync(Callable callable) throws Exception { } } + /** + * Awaits a {@link CompletableFuture} by yielding the continuation, without occupying an + * async-pool worker. Falls back to blocking {@code get()} when continuations are unavailable. + */ + public T awaitCompletableFuture(CompletableFuture future) throws Exception { + if (continuationExecutor != null && continuationContext != null) { + return continuationExecutor.awaitFuture(continuationContext, future); + } + return future.get(); + } + @Override public void close() throws Exception { super.close(); diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseDeferredSubagentSetup.java b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseDeferredSubagentSetup.java index 56b7cedf8..7d840cd94 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseDeferredSubagentSetup.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseDeferredSubagentSetup.java @@ -58,8 +58,10 @@ public SubagentFuture submit(RunnerContext ctx, Object prompt, String sessionId, * *

Called exactly once per invocation, when the deferred handle is first resolved, on the * mailbox thread. Implementations may therefore perform the mailbox-confined part of issuing - * the request here; the returned callable's {@link DurableCallable#call()} carries only the - * part that runs off the mailbox thread. + * the request here; the returned callable's {@link DurableCallable#call()} then carries the + * wait, normally on an async-pool worker off the mailbox thread. A callable marked {@link + * org.apache.flink.agents.runtime.async.ContinuationNativeCallable} instead runs inline on the + * mailbox thread and releases it by suspending the caller's continuation. * *

The callable folds its own comprehensible failures into the returned {@link * SubagentResult}; an exception escaping {@link DurableCallable#call()} is a system-level diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/InternalSubagentSetup.java b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/InternalSubagentSetup.java index 9a660fe90..774205210 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/InternalSubagentSetup.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/InternalSubagentSetup.java @@ -33,6 +33,7 @@ import org.apache.flink.agents.plan.actions.Action; import org.apache.flink.agents.runtime.ResourceCache; import org.apache.flink.agents.runtime.async.ContinuationActionExecutor; +import org.apache.flink.agents.runtime.async.ContinuationNativeCallable; import org.apache.flink.agents.runtime.condition.ActionMatcher; import org.apache.flink.agents.runtime.context.JavaRunnerContextImpl; import org.apache.flink.agents.runtime.context.RunnerContextImpl; @@ -54,8 +55,10 @@ *

Execution mode is the deferred one, inherited from {@link BaseDeferredSubagentSetup}: {@link * #submit} returns a deferred handle whose request is prepared on first resolve. {@link #prepare} * then runs the mailbox-confined bootstrap — registering the call status and sending one {@code - * InternalSubagentCallEvent} — and returns the durable callable that waits off the mailbox thread - * for the child to quiesce, so the operator can dispatch the child agent's actions in between. + * InternalSubagentCallEvent} — and returns the durable callable that waits for the child to + * quiesce. Unlike the other deferred setups, that callable runs inline on the mailbox thread and + * releases it by yielding the caller's continuation, so the operator can dispatch the child agent's + * actions in between while no async-pool worker is held. * *

Orchestration state lives here, not in the operator: the per-call quiesce statuses, the * per-scope child resource caches, and the per-key session index used to clean up when a record @@ -158,8 +161,15 @@ public List matchActions(Event event) { } /** - * The mailbox-confined bootstrap plus the off-mailbox wait. Sending the bootstrap event must - * happen before the wait releases the mailbox, so it runs here, on the resolving thread. + * The mailbox-confined bootstrap plus the wait. Sending the bootstrap event must happen before + * the wait releases the mailbox, so it runs here, on the resolving thread. + * + *

The returned callable is {@link ContinuationNativeCallable}, so it runs inline inside the + * caller's continuation rather than on an async-pool worker, and its wait suspends by yielding + * that continuation. A worker-dispatched wait would hold a pool worker for the whole child call + * while the child's own durable calls queue behind it for a worker from the same finite pool — + * a self-deadlock at {@code num-async-threads=1} and under enough nested or concurrent waits at + * any size (issue #1213). */ @Override protected DurableCallable prepare( @@ -171,7 +181,7 @@ protected DurableCallable prepare( throw new IllegalStateException( "Failed to bootstrap internal sub-agent call for scope " + scope, e); } - return new DurableCallable() { + return new ContinuationNativeCallable() { @Override public String getId() { return sessionId + "#" + callId; @@ -184,11 +194,11 @@ public Class getResultClass() { @Override public SubagentResult call() { - // Runs off the mailbox thread, so the operator can dispatch the child's - // actions; failures converge into a failed SubagentResult like the other - // deferred setups. + // Runs inline on the mailbox thread; the wait yields the continuation, so the + // operator can dispatch the child's actions without a pool worker being held. + // Failures converge into a failed SubagentResult like the other deferred setups. try { - return SubagentResult.ok(awaitSubagentCall(sessionId, callId)); + return SubagentResult.ok(awaitSubagentCallYielding(ctx, sessionId, callId)); } catch (Exception e) { return SubagentResult.error(e); } @@ -251,6 +261,29 @@ public void bootstrap(RunnerContext ctx, String sessionId, String callId, Object * and returns its accumulated output. Must be invoked off the mailbox thread. */ public List awaitSubagentCall(String sessionId, String callId) throws Exception { + try { + return requireCallStatus(sessionId, callId).getResponseFuture().get(); + } catch (java.util.concurrent.ExecutionException e) { + throw unwrapChildFailure(e); + } + } + + /** + * Awaits the identified call by yielding the caller's continuation, so the wait holds no + * async-pool worker and the child can draw from the pool to make progress (issue #1213). + * + *

The wait runs inline on the mailbox thread, so the blocking {@link #awaitSubagentCall} is + * not an option here — it would pin the mailbox and stall the very child it waits for. Only a + * Java action reaches {@link #prepare}; Python bootstraps through {@code + * RunnerContextImpl.bootstrapSubagentCallForScope} and blocks on its own async worker instead. + */ + private List awaitSubagentCallYielding( + RunnerContext ctx, String sessionId, String callId) throws Exception { + InternalSubagentCallStatus cs = requireCallStatus(sessionId, callId); + return ((JavaRunnerContextImpl) ctx).awaitCompletableFuture(cs.getResponseFuture()); + } + + private InternalSubagentCallStatus requireCallStatus(String sessionId, String callId) { InternalSubagentCallStatus cs = getCallStatus(sessionId, callId); if (cs == null) { throw new IllegalStateException( @@ -259,17 +292,16 @@ public List awaitSubagentCall(String sessionId, String callId) throws Ex + ", callId=" + callId); } - try { - return cs.getResponseFuture().get(); - } catch (java.util.concurrent.ExecutionException e) { - // Surface the child's failure directly, so the caller's SubagentResult carries - // the failing action's exception rather than the future's plumbing. - Throwable cause = e.getCause(); - if (cause instanceof Exception) { - throw (Exception) cause; - } - throw e; - } + return cs; + } + + /** + * Surfaces the child's failure directly, so the caller's {@link SubagentResult} carries the + * failing action's exception rather than the future's plumbing. + */ + private static Exception unwrapChildFailure(java.util.concurrent.ExecutionException e) { + Throwable cause = e.getCause(); + return cause instanceof Exception ? (Exception) cause : e; } /** The quiesce status of the identified call, or {@code null} when this setup owns none. */ diff --git a/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java index eb0b6003c..ef7e4f662 100644 --- a/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java +++ b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java @@ -370,6 +370,44 @@ private long getDeadlineNanos(Duration timeout) { : System.nanoTime() + timeout.toNanos(); } + /** + * Yields the current continuation until {@code future} completes, without occupying an + * async-pool worker. The caller must be running inside a continuation on the mailbox thread; + * this method yields that continuation and lets {@link #executeAction} resume it once the + * future is done. + * + *

This is the non-blocking alternative to submitting a blocking {@code future.get()} to the + * pool via {@link #executeAsync}. It prevents the deadlock where a parent awaiting a child + * holds the only pool worker the child needs to make progress (issue #1213). + */ + public T awaitFuture(ContinuationContext context, CompletableFuture future) + throws Exception { + context.clearAsyncState(); + // Bridge to Future so executeAction's hasPendingAsync gate drives the resumption. + context.setPendingFuture(future.thenApply(v -> null)); + // Re-checked before every yield, and resumption is polled by the operator's mail loop + // rather than signalled by a completion callback, so a future already done here — or one + // completing between the two lines above — simply never suspends. + while (!future.isDone()) { + Continuation.yield(SCOPE); + } + return joinCompletableFuture(future); + } + + private static T joinCompletableFuture(CompletableFuture future) throws Exception { + try { + return future.join(); + } catch (CompletionException e) { + Throwable cause = e.getCause(); + if (cause instanceof Exception) { + throw (Exception) cause; + } else if (cause instanceof Error) { + throw (Error) cause; + } + throw e; + } + } + /** Stops async work without waiting for worker exit cleanup. */ public void shutdown() { asyncExecutor.shutdownNow(); diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java index ed996a297..67ce7c7ec 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java @@ -125,6 +125,29 @@ void testDurableExecuteAsyncReconcilableSuccessCall() throws Exception { assertTrue(persisted.isSuccess()); } + @Test + void testDurableExecuteAsyncContinuationNativeReconcilableRunsInline() throws Exception { + InspectingContinuationActionExecutor executor = new InspectingContinuationActionExecutor(); + JavaRunnerContextImpl context = createContext(new ActionState(null), executor); + TestContinuationNativeReconcilableCallable callable = + new TestContinuationNativeReconcilableCallable<>( + "recon-inline", + String.class, + () -> "ok", + () -> fail("reconcile should not be called on initial async execution")); + + String result = context.durableExecuteAsync(callable).await(); + + assertEquals("ok", result); + assertEquals(1, callable.getCallCount()); + assertEquals(0, callable.getReconcileCount()); + assertEquals(0, executor.getExecuteAsyncCallCount()); + assertEquals(2, persistCallCount.get()); + CallResult persisted = + context.getDurableExecutionContext().getActionState().getCallResults().get(0); + assertTrue(persisted.isSuccess()); + } + @Test void testDurableExecuteAsyncReconcilableReplaySuccess() throws Exception { InspectingContinuationActionExecutor executor = new InspectingContinuationActionExecutor(); diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/context/TestDurableCallables.java b/runtime/src/test/java/org/apache/flink/agents/runtime/context/TestDurableCallables.java index 46ad2cedb..c1311cfd5 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/context/TestDurableCallables.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/context/TestDurableCallables.java @@ -19,6 +19,7 @@ package org.apache.flink.agents.runtime.context; import org.apache.flink.agents.api.context.DurableCallable; +import org.apache.flink.agents.runtime.async.ContinuationNativeCallable; import java.util.concurrent.Callable; @@ -108,3 +109,20 @@ int getReconcileCount() { return reconcileCount; } } + +/** + * Reconciliable fixture that is also marked {@link ContinuationNativeCallable}, so the runner + * context must run its {@link #call()} inline on the caller's thread rather than dispatch it to an + * async-pool worker. + */ +class TestContinuationNativeReconcilableCallable extends TestReconcilableCallable + implements ContinuationNativeCallable { + + TestContinuationNativeReconcilableCallable( + String id, + Class resultClass, + Callable callSupplier, + Callable reconcileSupplier) { + super(id, resultClass, callSupplier, reconcileSupplier); + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/InternalSubagentDeadlockTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/InternalSubagentDeadlockTest.java new file mode 100644 index 000000000..f2b0acd4f --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/InternalSubagentDeadlockTest.java @@ -0,0 +1,233 @@ +/* + * 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.subagent; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.InputEvent; +import org.apache.flink.agents.api.OutputEvent; +import org.apache.flink.agents.api.agents.Agent; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; +import org.apache.flink.agents.api.context.DurableCallable; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.api.subagent.SubagentFutures; +import org.apache.flink.agents.api.subagent.SubagentResult; +import org.apache.flink.agents.api.subagent.SubagentSetup; +import org.apache.flink.agents.plan.AgentConfiguration; +import org.apache.flink.agents.plan.AgentPlan; +import org.apache.flink.agents.runtime.operator.ActionExecutionOperator; +import org.apache.flink.agents.runtime.operator.ActionExecutionOperatorFactory; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.api.condition.EnabledForJreRange; +import org.junit.jupiter.api.condition.JRE; + +import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Reproduces the deadlock described in GitHub issue #1213: with {@code num-async-threads=1}, a + * parent awaiting an internal sub-agent holds the single async-pool worker via a blocking {@code + * responseFuture.get()}, while the child's own durable async call needs a worker from the same pool + * to make progress. The result is a self-deadlock — neither side can advance. + * + *

The fix must let the parent's await yield its worker so the child can proceed, regardless of + * pool size. This test times out under the broken behavior and completes promptly once fixed. + */ +@EnabledForJreRange(min = JRE.JAVA_21) +public class InternalSubagentDeadlockTest { + + private static final AtomicBoolean CHILD_DURABLE_EXECUTED = new AtomicBoolean(false); + private static final AtomicInteger DURABLE_EXECUTIONS = new AtomicInteger(); + + @BeforeEach + void resetCounters() { + CHILD_DURABLE_EXECUTED.set(false); + DURABLE_EXECUTIONS.set(0); + } + + /** + * Child agent whose action performs a durable async call before emitting output. The call + * itself is trivial (returns immediately), but it still requires an async-pool worker to + * execute. Under the deadlock, the parent holds the only worker and this call never starts. + */ + public static class DurableAsyncChildAgent extends Agent { + public DurableAsyncChildAgent() throws Exception { + addAction( + new String[] {InputEvent.EVENT_TYPE}, + DurableAsyncChildAgent.class.getMethod( + "handle", Event.class, RunnerContext.class)); + } + + @SuppressWarnings("unused") + public static void handle(Event event, RunnerContext ctx) throws Exception { + String result = ctx.durableExecuteAsync(durableCall()).await(); + ctx.sendEvent(new OutputEvent(result)); + } + } + + /** + * Child that itself awaits a grandchild performing a durable call, so two waits are nested + * while the grandchild's call still needs a worker from the same single-thread pool. + */ + public static class NestingDurableChildAgent extends Agent { + public NestingDurableChildAgent() throws Exception { + addResource("grandchild", ResourceType.AGENT, new DurableAsyncChildAgent()); + addAction( + new String[] {InputEvent.EVENT_TYPE}, + NestingDurableChildAgent.class.getMethod( + "relay", Event.class, RunnerContext.class)); + } + + @SuppressWarnings("unused") + public static void relay(Event event, RunnerContext ctx) throws Exception { + SubagentSetup grandchild = + (SubagentSetup) ctx.getResource("grandchild", ResourceType.AGENT); + SubagentResult result = + grandchild.submit(ctx, InputEvent.fromEvent(event).getInput()).await(); + ctx.sendEvent(new OutputEvent("relayed:" + firstOutput(result))); + } + } + + private static DurableCallable durableCall() { + return new DurableCallable() { + @Override + public String getId() { + return "child-durable"; + } + + @Override + public Class getResultClass() { + return String.class; + } + + @Override + public String call() { + CHILD_DURABLE_EXECUTED.set(true); + DURABLE_EXECUTIONS.incrementAndGet(); + return "durable-result"; + } + }; + } + + private static String firstOutput(SubagentResult result) { + return String.valueOf(((List) result.getResult()).get(0)); + } + + @SuppressWarnings("unused") + public static void callChild(Event event, RunnerContext ctx) throws Exception { + SubagentSetup child = (SubagentSetup) ctx.getResource("child", ResourceType.AGENT); + SubagentResult result = child.submit(ctx, InputEvent.fromEvent(event).getInput()).await(); + ctx.sendEvent(new OutputEvent("parent:" + firstOutput(result))); + } + + @SuppressWarnings("unused") + public static void callChildTwiceCombined(Event event, RunnerContext ctx) throws Exception { + SubagentSetup child = (SubagentSetup) ctx.getResource("child", ResourceType.AGENT); + SubagentFutures combined = child.submit(ctx, "a").combine(child.submit(ctx, "b")); + List results = combined.awaitAll(); + ctx.sendEvent( + new OutputEvent( + "parent:" + + firstOutput(results.get(0)) + + "," + + firstOutput(results.get(1)))); + } + + /** + * With {@code num-async-threads=1}, the parent's await must not hold the only worker. If it + * does, the child's durable call can never execute and this test times out. + */ + @Test + @Timeout(30) + void singleAsyncThreadDoesNotDeadlockWhenChildUsesDurableCall() throws Exception { + List output = run(plan("callChild", new DurableAsyncChildAgent())); + + assertThat(CHILD_DURABLE_EXECUTED).isTrue(); + assertThat(output).containsExactly("parent:durable-result"); + } + + /** + * Nesting multiplies the waits in flight: a child awaiting its own grandchild must not consume + * a worker either, or one nested call deadlocks the single-thread pool. + */ + @Test + @Timeout(30) + void singleAsyncThreadDoesNotDeadlockOnNestedCalls() throws Exception { + List output = run(plan("callChild", new NestingDurableChildAgent())); + + assertThat(CHILD_DURABLE_EXECUTED).isTrue(); + assertThat(output).containsExactly("parent:relayed:durable-result"); + } + + /** + * Concurrent calls combined into one wait: {@code awaitAll} prepares both calls (dispatching + * both children) before awaiting either, so both waits are in flight at once and neither may + * hold a worker. + */ + @Test + @Timeout(30) + void singleAsyncThreadDoesNotDeadlockOnCombinedCalls() throws Exception { + List output = run(plan("callChildTwiceCombined", new DurableAsyncChildAgent())); + + assertThat(DURABLE_EXECUTIONS.get()).isEqualTo(2); + assertThat(output).containsExactly("parent:durable-result,durable-result"); + } + + // --- helpers --- + + /** Builds a plan whose single caller action is {@code callerAction}, with one async worker. */ + private static AgentPlan plan(String callerAction, Agent child) throws Exception { + AgentConfiguration config = new AgentConfiguration(); + config.set(AgentExecutionOptions.NUM_ASYNC_THREADS, 1); + + Agent parent = new Agent(); + parent.addResource("child", ResourceType.AGENT, child); + parent.addAction( + new String[] {InputEvent.EVENT_TYPE}, + InternalSubagentDeadlockTest.class.getMethod( + callerAction, Event.class, RunnerContext.class)); + return new AgentPlan(parent, config); + } + + @SuppressWarnings("unchecked") + private static List run(AgentPlan plan) throws Exception { + try (KeyedOneInputStreamOperatorTestHarness harness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>(plan, true), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + harness.open(); + harness.processElement(new StreamRecord<>(1L)); + ((ActionExecutionOperator) harness.getOperator()) + .waitInFlightEventsFinished(); + return ((List>) harness.getRecordOutput()) + .stream().map(StreamRecord::getValue).collect(Collectors.toList()); + } + } +}