[runtime][java] Avoid async-pool self-deadlock in internal sub-agent await (#1213)

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)

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 <T> BatchExecutionResult<T> 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> T awaitFuture(ContinuationContext context, CompletableFuture<T> future)
+            throws Exception {
+        return future.get();
+    }
+
     /** Stops async work without waiting for worker exit cleanup. */
     public void shutdown() {
         batchExecutor.shutdown();

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.
+ *
+ * <p>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.
+ *
+ * <p>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<T> extends DurableCallable<T> {}

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> T resolveDurableAsync(DurableCallable<T> callable) throws Exception {
+        Callable<T> boundary = executionBoundary(callable);
         if (durableExecutionContext != null) {
             Callable<T> 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> T durableExecuteAsyncWithReconcile(
-            DurableCallable<T> callable, Callable<T> 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 <T> Callable<T> executionBoundary(DurableCallable<T> callable) {
+        if (callable instanceof ContinuationNativeCallable) {
+            return callable::call;
+        }
+        return () -> resolveAsync(callable::call);
     }
 
     @Override
@@ -350,6 +362,17 @@ protected <T> T resolveAsync(Callable<T> 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> T awaitCompletableFuture(CompletableFuture<T> future) throws Exception {
+        if (continuationExecutor != null && continuationContext != null) {
+            return continuationExecutor.awaitFuture(continuationContext, future);
+        }
+        return future.get();
+    }
+
     @Override
     public void close() throws Exception {
         super.close();

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,
      *
      * <p>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.
      *
      * <p>The callable folds its own comprehensible failures into the returned {@link
      * SubagentResult}; an exception escaping {@link DurableCallable#call()} is a system-level