You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by we...@apache.org on 2026/09/29 02:24:28 UTC
(flink-agents) branch main updated: [api][runtime] Support parallel action execution on JDK < 21 (#1147)
This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new 85e8cd97 [api][runtime] Support parallel action execution on JDK < 21 (#1147)
85e8cd97 is described below
commit 85e8cd972dac35b254d07c0d148ef019a456a708
Author: Yi Zhang <72...@users.noreply.github.com>
AuthorDate: Tue Sep 29 10:24:22 2026 +0800
[api][runtime] Support parallel action execution on JDK < 21 (#1147)
Generated-by: Qoder (Qoder Auto)
Generated-by: Codex (model/version not exposed)
---------
Co-authored-by: Gen Luo <86...@users.noreply.github.com>
Co-authored-by: Haocong Wang <75...@users.noreply.github.com>
---
.../agents/api/agents/AgentExecutionOptions.java | 18 +
docs/content/docs/operations/configuration.md | 2 +
.../integration/test/AsyncExecutionTest.java | 26 +-
.../integration/test/ParallelExecutionE2ETest.java | 1141 ++++++++++++++++++++
.../test/PythonResourceWorkerLifetimeE2ETest.java | 75 ++
.../test/PythonResourceWorkerLifetimeScenario.java | 369 +++++++
.../resources/python/resource_lifetime_fixture.py | 55 +
.../resourceprovider/PythonResourceProvider.java | 29 +-
.../PythonResourceProviderTest.java | 59 +
python/flink_agents/api/core_options.py | 20 +
.../apache/flink/agents/runtime/ResourceCache.java | 12 +-
.../runtime/async/ContinuationActionExecutor.java | 77 +-
.../agents/runtime/async/ContinuationContext.java | 47 +-
.../runtime/context/JavaRunnerContextImpl.java | 5 +-
.../runtime/operator/ActionExecutionOperator.java | 639 ++++++++++-
.../runtime/operator/ActionTaskContextManager.java | 74 +-
.../runtime/operator/DurableExecutionManager.java | 1 +
.../runtime/operator/OperatorStateManager.java | 12 +
.../runtime/operator/PythonBridgeManager.java | 3 +-
.../ParallelExecutionContextRestorer.java} | 17 +-
.../parallel/ParallelExecutionCoordinator.java | 312 ++++++
.../operator/parallel/ParallelExecutionLock.java | 282 +++++
.../operator/parallel/ParallelExecutionTask.java | 58 +
.../parallel/ParallelExecutionTaskQueue.java | 208 ++++
.../python/utils/PythonInterpreterManager.java | 84 +-
.../python/utils/PythonResourceAdapterImpl.java | 5 +-
.../runtime/async/ContinuationActionExecutor.java | 17 +-
.../flink/agents/runtime/ResourceCacheTest.java | 22 +
.../async/ContinuationActionExecutorTest.java | 200 ++++
...vaRunnerContextImplDurableExecuteAsyncTest.java | 2 +-
.../TaskLifecycleListenerNotificationTest.java | 8 +-
...onExecutionOperatorCapacityIntegrationTest.java | 338 ++++++
...cutionOperatorDrainDeadlockIntegrationTest.java | 495 +++++++++
...ActionExecutionOperatorParallelMetricsTest.java | 235 ++++
.../operator/ActionExecutionOperatorTest.java | 1014 ++++++++++++++++-
.../operator/ActionTaskContextManagerTest.java | 168 ++-
.../operator/DurableExecutionManagerTest.java | 23 +
.../parallel/ParallelExecutionCoordinatorTest.java | 881 +++++++++++++++
.../parallel/ParallelExecutionLockTest.java | 441 ++++++++
.../ParallelExecutionPythonResourceTest.java | 297 +++++
.../parallel/ParallelExecutionTaskQueueTest.java | 194 ++++
.../python/utils/PythonInterpreterManagerTest.java | 46 +
42 files changed, 7868 insertions(+), 143 deletions(-)
diff --git a/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java b/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java
index c77d7bda..3d691b50 100644
--- a/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java
+++ b/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java
@@ -34,6 +34,24 @@ public class AgentExecutionOptions {
Integer.class,
Runtime.getRuntime().availableProcessors() * 2);
+ /**
+ * Experimental fallback switch for the JDK<21 parallel execution engine. Only consulted for
+ * pure-Java agents without coroutine support; the JDK 21 coroutine engine and plans containing
+ * Python actions never use the parallel engine regardless of this value (no Python-side
+ * counterpart exists by design). Set to {@code false} to fall back to the serial engine.
+ */
+ public static final ConfigOption<Boolean> PARALLEL_EXECUTION_ENABLED =
+ new ConfigOption<>("parallel-execution.enabled", Boolean.class, true);
+
+ /**
+ * Maximum number of input records that may be in flight concurrently. Only enforced by the
+ * JDK<21 parallel execution engine for pure-Java agents; the JDK 21 coroutine engine and
+ * plans containing Python actions ignore it. Every admitted record consumes one unit of budget;
+ * at the cap, admission of further records blocks until an in-flight record retires.
+ */
+ public static final ConfigOption<Integer> MAX_IN_FLIGHT_INPUT_RECORDS =
+ new ConfigOption<>("max-in-flight-input-records", Integer.class, 100);
+
public static final ConfigOption<Boolean> CHAT_ASYNC =
new ConfigOption<>("chat.async", Boolean.class, true);
diff --git a/docs/content/docs/operations/configuration.md b/docs/content/docs/operations/configuration.md
index f4708670..1a5dad19 100644
--- a/docs/content/docs/operations/configuration.md
+++ b/docs/content/docs/operations/configuration.md
@@ -138,6 +138,8 @@ Here is the list of all built-in core configuration options.
| `tool-call.batch.timeout.ms` | -1 (disabled) | long (milliseconds) | Overall timeout for one parallel tool-call batch. Non-positive disables it. On timeout, completed slots keep their outcome; slots that started but did not finish are recorded as failures; slots that never started executing (for example, queued in a saturated pool) stay pending, so they are re-executed after recovery instead of recording a false failure. Timeout cancellation is best-effort; external side [...]
| `rag.async` | true | boolean | Whether retrieve context asynchronously for built-in context retrieval action. |
| `num-async-threads` | os cpu count * 2 | int | Size of the fixed async executor created once per operator subtask. Chat, RAG, and parallel tool batches on **all keys handled by that subtask** submit work here. Default `tool-call.parallelism` is `os cpu count`, so one full tool batch can occupy up to half of this pool; multiple keys running large batches can saturate it. |
+| `parallel-execution.enabled` | true | boolean | Whether pure-**Java** agents run actions concurrently on the parallel execution engine when the continuation-based engine is unavailable (**JDK < 21**). Same-key input records still commit in input order, and checkpoint (exactly-once) semantics are unchanged. Set to `false` to fall back to the serial engine. Ignored when the JDK 21 coroutine engine is available, and for plans containing **Python** actions, [...]
+| `max-in-flight-input-records` | 100 | int | Maximum number of input records that may be in flight concurrently. Only enforced by the **JDK < 21** parallel execution engine for pure-**Java** agents; the JDK 21 coroutine engine and plans containing **Python** actions ignore it. Every admitted record consumes one unit of budget; at the cap, admission of further records blocks until an in-flight record retires. |
| `job-identifier` | none | String | The unique identifier of job, remaining consistent after restoring from a savepoint. If not set, uses flink job id. |
| `event-log.level` | STANDARD | EventLogLevel | Global default verbosity for the [Event Log]({{< ref "docs/operations/monitoring#event-log" >}}). Valid values: `OFF` (skip event), `STANDARD` (payload may be truncated/summarized to keep logs concise), `VERBOSE` (full payload). Can be overridden per event type — see [Per-event-type log levels]({{< ref "docs/operations/monitoring#per-event-type-log-levels" >}}). |
| `event-log.trace.enabled` | false | boolean | Whether to persist Agent Trace information in the Event Log. When enabled, business Events include trace context and Action/LLM/Parser/Tool lifecycle Events are logged. |
diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionTest.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionTest.java
index a8594d97..561dcd95 100644
--- a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionTest.java
+++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionTest.java
@@ -210,14 +210,15 @@ public class AsyncExecutionTest {
}
/**
- * Tests that async execution on JDK 21+ actually executes tasks in parallel.
+ * Tests that async execution actually executes tasks in parallel.
*
* <p>This test creates multiple tasks that each sleep for a fixed duration. Each async task
* records its start and end timestamps. We verify parallel execution by checking if the
* execution time ranges overlap.
*
- * <p>On JDK 21+: Tasks run in parallel, their execution times overlap On JDK < 21: Tasks run
- * sequentially, no overlap
+ * <p>On JDK 21+: the coroutine engine runs tasks in parallel, their execution times overlap. On
+ * JDK < 21: the blocking-inline parallel execution engine (enabled by default for pure-Java
+ * agents) also runs these multi-key tasks in parallel.
*/
@Test
public void testAsyncExecutionIsActuallyParallel() throws Exception {
@@ -322,15 +323,18 @@ public class AsyncExecutionTest {
expectedOverlaps - 1, overlapCount));
System.out.println("✓ Async execution is PARALLEL (as expected on JDK 21+)");
} else {
- // On JDK < 21, tasks run sequentially - no overlap expected
- Assertions.assertEquals(
- 0,
- overlapCount,
+ // On JDK < 21, the blocking-inline parallel execution engine (enabled by default
+ // for pure-Java agents) also runs these multi-key tasks in parallel.
+ int expectedOverlaps = (numRequests * (numRequests - 1)) / 2;
+ Assertions.assertTrue(
+ overlapCount >= expectedOverlaps - 1, // Allow some tolerance
String.format(
- "On JDK < 21, async tasks should run sequentially (no overlap). "
- + "But found %d overlapping pairs.",
- overlapCount));
- System.out.println("✓ Async execution is SEQUENTIAL (as expected on JDK < 21)");
+ "On JDK < 21, the parallel execution engine should run async tasks "
+ + "in parallel (overlapping). "
+ + "Expected at least %d overlapping pairs, but found %d.",
+ expectedOverlaps - 1, overlapCount));
+ System.out.println(
+ "✓ Async execution is PARALLEL (blocking-inline engine on JDK < 21)");
}
System.out.println("=== Test Passed ===");
diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/ParallelExecutionE2ETest.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/ParallelExecutionE2ETest.java
new file mode 100644
index 00000000..c7c23ccd
--- /dev/null
+++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/ParallelExecutionE2ETest.java
@@ -0,0 +1,1141 @@
+/*
+ * 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.integration.test;
+
+import org.apache.flink.agents.api.AgentsExecutionEnvironment;
+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.annotation.Action;
+import org.apache.flink.agents.api.context.DurableCallable;
+import org.apache.flink.agents.api.context.MemoryObject;
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.runtime.async.ContinuationActionExecutor;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.util.CloseableIterator;
+import org.junit.jupiter.api.Assumptions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+
+import java.io.Serializable;
+import java.lang.reflect.Method;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.catchThrowable;
+
+/**
+ * End-to-end scenarios for parallel action execution, covering fan-out, broadcast, multi-key,
+ * context switch, same-key input order, capacity, failure propagation and backpressure. The suite
+ * runs on both execution engines; only the backpressure scenario is specific to the JDK<21
+ * blocking-inline engine.
+ *
+ * <p>Timing observations are collected in the in-JVM {@link TimingLog} (valid because every
+ * scenario runs on a same-JVM MiniCluster). Its append order is worker completion order, so
+ * order-sensitive assertions always go against the collected job outputs, which carry one {@code
+ * recordId:action} marker per terminal {@link OutputEvent}.
+ */
+public class ParallelExecutionE2ETest {
+
+ private static final int FANOUT_SLEEP_MS = 1500;
+ private static final int FANOUT_HANDLERS = 3;
+
+ private static final int BROADCAST_LONG_SLEEP_MS = 1200;
+ private static final int BROADCAST_SHORT_SLEEP_MS = 100;
+ private static final int BROADCAST_HANDLERS = 3;
+
+ private static final int MULTI_KEY_SLEEP_MS = 1500;
+ private static final int MULTI_KEY_RECORDS = 3;
+ private static final CountDownLatch MULTI_KEY_ALL_STARTED =
+ new CountDownLatch(MULTI_KEY_RECORDS);
+
+ private static final String CTX_OWNER_CONFIG = "context-owner";
+ private static final String CTX_MEMORY_OWNER = "context-owner";
+ private static final int CTX_SLEEP_MS = 1500;
+ private static final int CTX_HANDLERS = 3;
+ /** Independent oracle for the action-specific context configuration. */
+ private static final Map<String, String> CTX_EXPECTED_OWNER_BY_HANDLER =
+ Map.of(
+ "handleAlpha", "ALPHA",
+ "handleBeta", "BETA",
+ "handleGamma", "GAMMA");
+
+ private static final int SAME_KEY = 7;
+ private static final int SAME_KEY_INPUTS = 2;
+ private static final int SAME_KEY_ACTIONS = 2;
+ private static final int SAME_KEY_ASYNC_THREADS = 2;
+ private static final int SAME_KEY_LONG_SLEEP_MS = 1000;
+ private static final int SAME_KEY_SHORT_SLEEP_MS = 100;
+
+ private static final int CAPACITY_ASYNC_THREADS = 2;
+ private static final int CAPACITY_ACTIONS = CAPACITY_ASYNC_THREADS * 2 + 2;
+ private static final int CAPACITY_SLEEP_MS = 350;
+
+ private static final String FAILURE_MARKER = "scenario-async-failure-marker";
+
+ private static final int BACKPRESSURE_ASYNC_THREADS = 2;
+ private static final int BACKPRESSURE_MAX_IN_FLIGHT = 1;
+ private static final int BACKPRESSURE_SLEEP_MS = 700;
+
+ @BeforeEach
+ void clearRegistries() {
+ TimingLog.clear();
+ ContextObservations.clear();
+ }
+
+ /** One timed execution on an async worker. */
+ static final class Span {
+ final String action;
+ final String key;
+ final String recordId;
+ final int slot;
+ final long startMs;
+ final long endMs;
+ final String thread;
+
+ Span(
+ String action,
+ String key,
+ String recordId,
+ int slot,
+ long startMs,
+ long endMs,
+ String thread) {
+ this.action = action;
+ this.key = key;
+ this.recordId = recordId;
+ this.slot = slot;
+ this.startMs = startMs;
+ this.endMs = endMs;
+ this.thread = thread;
+ }
+
+ @Override
+ public String toString() {
+ return String.format(
+ Locale.ROOT,
+ "%s:%s#%d[%d..%d]@%s",
+ recordId,
+ action,
+ slot,
+ startMs,
+ endMs,
+ thread);
+ }
+ }
+
+ /** In-JVM timing registry; append order is worker completion order, not sink output order. */
+ static final class TimingLog {
+ private static final List<Span> SPANS = Collections.synchronizedList(new ArrayList<>());
+
+ static void clear() {
+ SPANS.clear();
+ }
+
+ /** Times the sleep, records the span, and returns the {@code recordId:action} marker. */
+ static String run(String action, Object key, Object recordId, int slot, int sleepMs)
+ throws Exception {
+ long start = System.currentTimeMillis();
+ Thread.sleep(sleepMs);
+ long end = System.currentTimeMillis();
+ SPANS.add(
+ new Span(
+ action,
+ String.valueOf(key),
+ String.valueOf(recordId),
+ slot,
+ start,
+ end,
+ Thread.currentThread().getName()));
+ return recordId + ":" + action;
+ }
+
+ static List<Span> spans() {
+ synchronized (SPANS) {
+ return new ArrayList<>(SPANS);
+ }
+ }
+ }
+
+ private static List<Span> spansOfSlot(int slot) {
+ List<Span> result = new ArrayList<>();
+ for (Span span : TimingLog.spans()) {
+ if (span.slot == slot) {
+ result.add(span);
+ }
+ }
+ return result;
+ }
+
+ private static List<Span> spansOfRecord(String recordId) {
+ List<Span> result = new ArrayList<>();
+ for (Span span : TimingLog.spans()) {
+ if (recordId.equals(span.recordId)) {
+ result.add(span);
+ }
+ }
+ return result;
+ }
+
+ private static Span spanOf(String action) {
+ return TimingLog.spans().stream()
+ .filter(span -> action.equals(span.action))
+ .findFirst()
+ .orElseThrow(() -> new AssertionError("Missing span for action " + action));
+ }
+
+ private static String describe(List<Span> spans) {
+ List<Span> sorted = new ArrayList<>(spans);
+ sorted.sort(Comparator.comparingLong(span -> span.startMs));
+ return sorted.toString();
+ }
+
+ /** Every pair of spans must overlap in time. */
+ private static void assertAllPairsOverlap(String description, List<Span> spans) {
+ assertThat(spans.size())
+ .as("%s needs at least two spans to compare", description)
+ .isGreaterThanOrEqualTo(2);
+ for (int i = 0; i < spans.size(); i++) {
+ for (int j = i + 1; j < spans.size(); j++) {
+ Span a = spans.get(i);
+ Span b = spans.get(j);
+ assertThat(a.startMs < b.endMs && b.startMs < a.endMs)
+ .as(
+ "%s: %s and %s must overlap; spans=%s",
+ description, a, b, describe(spans))
+ .isTrue();
+ }
+ }
+ }
+
+ /** Sum of span durations over the aggregate wall time. */
+ private static void assertSpeedup(String description, List<Span> spans, double minRatio) {
+ long minStart = spans.stream().mapToLong(span -> span.startMs).min().orElseThrow();
+ long maxEnd = spans.stream().mapToLong(span -> span.endMs).max().orElseThrow();
+ long sumDurations = spans.stream().mapToLong(span -> span.endMs - span.startMs).sum();
+ long wall = Math.max(maxEnd - minStart, 1L);
+ assertThat((double) sumDurations / wall)
+ .as("%s: speedup over wall time; spans=%s", description, describe(spans))
+ .isGreaterThanOrEqualTo(minRatio);
+ }
+
+ private static int distinctThreads(List<Span> spans) {
+ return (int) spans.stream().map(span -> span.thread).distinct().count();
+ }
+
+ /** Uniform request POJO for all scenarios; {@code recordId} and {@code key} vary per test. */
+ public static class TestRequest implements Serializable {
+ private static final long serialVersionUID = 1L;
+
+ public final int recordId;
+ public final int key;
+
+ public TestRequest(int recordId, int key) {
+ this.recordId = recordId;
+ this.key = key;
+ }
+ }
+
+ public static class TestRequestKeySelector implements KeySelector<TestRequest, Integer> {
+ @Override
+ public Integer getKey(TestRequest request) {
+ return request.key;
+ }
+ }
+
+ /** Runs a single-parallelism agent job and returns the collected outputs in sink order. */
+ private static List<String> runJob(
+ Agent agent, Consumer<AgentsExecutionEnvironment> configure, TestRequest... inputs)
+ throws Exception {
+ StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
+ // Keep one Flink subtask so observed overlap must come from the JDK<21 action workers.
+ env.setParallelism(1);
+ DataStream<TestRequest> inputStream = env.fromElements(inputs);
+
+ AgentsExecutionEnvironment agentsEnv =
+ AgentsExecutionEnvironment.getExecutionEnvironment(env);
+ configure.accept(agentsEnv);
+
+ DataStream<Object> outputStream =
+ agentsEnv
+ .fromDataStream(inputStream, new TestRequestKeySelector())
+ .apply(agent)
+ .toDataStream();
+
+ CloseableIterator<Object> results = outputStream.collectAsync();
+ agentsEnv.execute();
+
+ List<String> outputs = new ArrayList<>();
+ while (results.hasNext()) {
+ outputs.add(results.next().toString());
+ }
+ results.close();
+ return outputs;
+ }
+
+ /** Serializable body used by {@link #durableTiming} hooks. */
+ @FunctionalInterface
+ public interface ThrowingRunnable extends Serializable {
+ void run() throws Exception;
+ }
+
+ /** One timed {@code durableExecuteAsync} call; returns the {@code recordId:action} marker. */
+ private static String durableTiming(
+ RunnerContext ctx,
+ String id,
+ String action,
+ Object key,
+ Object recordId,
+ int slot,
+ int sleepMs)
+ throws Exception {
+ return durableTiming(ctx, id, action, key, recordId, slot, sleepMs, null);
+ }
+
+ /**
+ * Variant with a {@code beforeTiming} hook running inside the worker callable before the timed
+ * window opens (e.g. a start-alignment barrier), so the hook never inflates the span.
+ */
+ private static String durableTiming(
+ RunnerContext ctx,
+ String id,
+ String action,
+ Object key,
+ Object recordId,
+ int slot,
+ int sleepMs,
+ ThrowingRunnable beforeTiming)
+ throws Exception {
+ return ctx.durableExecuteAsync(
+ new DurableCallable<String>() {
+ @Override
+ public String getId() {
+ return id;
+ }
+
+ @Override
+ public Class<String> getResultClass() {
+ return String.class;
+ }
+
+ @Override
+ public String call() throws Exception {
+ if (beforeTiming != null) {
+ beforeTiming.run();
+ }
+ return TimingLog.run(action, key, recordId, slot, sleepMs);
+ }
+ })
+ .await();
+ }
+
+ /** One {@code durableExecuteAsync} call whose worker body always throws. */
+ private static void durableFailure(RunnerContext ctx, String id, String marker)
+ throws Exception {
+ ctx.durableExecuteAsync(
+ new DurableCallable<String>() {
+ @Override
+ public String getId() {
+ return id;
+ }
+
+ @Override
+ public Class<String> getResultClass() {
+ return String.class;
+ }
+
+ @Override
+ public String call() {
+ throw new IllegalStateException(marker);
+ }
+ })
+ .await();
+ }
+
+ private static Method method(Class<?> agentClass, String name) {
+ try {
+ return agentClass.getMethod(name, Event.class, RunnerContext.class);
+ } catch (NoSuchMethodException e) {
+ throw new IllegalStateException("Missing action method: " + name, e);
+ }
+ }
+
+ @Test
+ @Timeout(30)
+ public void fanoutRunsSiblingActionsInParallel() throws Exception {
+ List<String> outputs =
+ runJob(
+ new FanoutAgent(),
+ agentsEnv ->
+ agentsEnv
+ .getConfig()
+ .set(
+ AgentExecutionOptions.NUM_ASYNC_THREADS,
+ FANOUT_HANDLERS),
+ new TestRequest(1, 1));
+
+ assertThat(outputs).containsExactlyInAnyOrder("1:handleA", "1:handleB", "1:handleC");
+
+ List<Span> spans = TimingLog.spans();
+ assertThat(spans).hasSize(FANOUT_HANDLERS);
+ assertThat(spans)
+ .extracting(span -> span.action)
+ .containsExactlyInAnyOrder("handleA", "handleB", "handleC");
+ assertAllPairsOverlap("fan-out siblings", spans);
+ assertThat(distinctThreads(spans)).isGreaterThanOrEqualTo(2);
+ assertSpeedup("fan-out siblings", spans, 2.0);
+ }
+
+ /** Agent: each handler consumes the input directly and does one durableExecuteAsync. */
+ public static class FanoutAgent extends Agent {
+
+ @Action(InputEvent.EVENT_TYPE)
+ public static void handleA(Event event, RunnerContext ctx) throws Exception {
+ invokeAndEmit(event, ctx, "handleA");
+ }
+
+ @Action(InputEvent.EVENT_TYPE)
+ public static void handleB(Event event, RunnerContext ctx) throws Exception {
+ invokeAndEmit(event, ctx, "handleB");
+ }
+
+ @Action(InputEvent.EVENT_TYPE)
+ public static void handleC(Event event, RunnerContext ctx) throws Exception {
+ invokeAndEmit(event, ctx, "handleC");
+ }
+
+ private static void invokeAndEmit(Event event, RunnerContext ctx, String action)
+ throws Exception {
+ TestRequest request = (TestRequest) InputEvent.fromEvent(event).getInput();
+ String marker =
+ durableTiming(
+ ctx,
+ "fanout-" + action + "-" + request.recordId,
+ action,
+ request.key,
+ request.recordId,
+ 1,
+ FANOUT_SLEEP_MS);
+ ctx.sendEvent(new OutputEvent(marker));
+ }
+ }
+
+ @Test
+ @Timeout(30)
+ public void broadcastRunsAllListenersInParallel() throws Exception {
+ List<String> outputs =
+ runJob(
+ new BroadcastAgent(),
+ agentsEnv ->
+ agentsEnv
+ .getConfig()
+ .set(
+ AgentExecutionOptions.NUM_ASYNC_THREADS,
+ BROADCAST_HANDLERS),
+ new TestRequest(1, 1));
+
+ assertThat(outputs)
+ .containsExactlyInAnyOrder("1:handleAlpha", "1:handleBeta", "1:handleGamma");
+
+ List<Span> spans = TimingLog.spans();
+ assertThat(spans).hasSize(BROADCAST_HANDLERS);
+ assertAllPairsOverlap("broadcast listeners", spans);
+ assertThat(distinctThreads(spans)).isGreaterThanOrEqualTo(2);
+
+ // The short-sleep listeners must finish before the long-sleep one.
+ Span alpha = spanOf("handleAlpha");
+ Span beta = spanOf("handleBeta");
+ Span gamma = spanOf("handleGamma");
+ assertThat(beta.endMs).isLessThan(alpha.endMs);
+ assertThat(gamma.endMs).isLessThan(alpha.endMs);
+ }
+
+ /** Broadcast event consumed by every listener in this scenario. */
+ public static class FanBroadcastEvent extends Event {
+ public static final String EVENT_TYPE = "ParallelExecutionE2ETest.FanBroadcastEvent";
+
+ public final int inputId;
+
+ public FanBroadcastEvent(int inputId) {
+ super(EVENT_TYPE);
+ this.inputId = inputId;
+ }
+ }
+
+ /** Agent: prime sends one FanBroadcastEvent; three listeners share the same EVENT_TYPE. */
+ public static class BroadcastAgent extends Agent {
+
+ public BroadcastAgent() {
+ addAction(new String[] {InputEvent.EVENT_TYPE}, method(BroadcastAgent.class, "prime"));
+ addAction(
+ new String[] {FanBroadcastEvent.EVENT_TYPE},
+ method(BroadcastAgent.class, "handleAlpha"));
+ addAction(
+ new String[] {FanBroadcastEvent.EVENT_TYPE},
+ method(BroadcastAgent.class, "handleBeta"));
+ addAction(
+ new String[] {FanBroadcastEvent.EVENT_TYPE},
+ method(BroadcastAgent.class, "handleGamma"));
+ }
+
+ public static void prime(Event event, RunnerContext ctx) throws Exception {
+ TestRequest request = (TestRequest) InputEvent.fromEvent(event).getInput();
+ ctx.sendEvent(new FanBroadcastEvent(request.recordId));
+ }
+
+ public static void handleAlpha(Event event, RunnerContext ctx) throws Exception {
+ invokeAndEmit(event, ctx, "handleAlpha", BROADCAST_LONG_SLEEP_MS);
+ }
+
+ public static void handleBeta(Event event, RunnerContext ctx) throws Exception {
+ invokeAndEmit(event, ctx, "handleBeta", BROADCAST_SHORT_SLEEP_MS);
+ }
+
+ public static void handleGamma(Event event, RunnerContext ctx) throws Exception {
+ invokeAndEmit(event, ctx, "handleGamma", BROADCAST_SHORT_SLEEP_MS);
+ }
+
+ private static void invokeAndEmit(
+ Event event, RunnerContext ctx, String action, int sleepMs) throws Exception {
+ FanBroadcastEvent in = (FanBroadcastEvent) event;
+ String marker =
+ durableTiming(
+ ctx,
+ "broadcast-" + action + "-" + in.inputId,
+ action,
+ in.inputId,
+ in.inputId,
+ 1,
+ sleepMs);
+ ctx.sendEvent(new OutputEvent(marker));
+ }
+ }
+
+ @Test
+ @Timeout(30)
+ public void multiKeyInputsRunInParallel() throws Exception {
+ List<String> outputs =
+ runJob(
+ new MultiKeyAgent(),
+ agentsEnv ->
+ agentsEnv
+ .getConfig()
+ .set(
+ AgentExecutionOptions.NUM_ASYNC_THREADS,
+ MULTI_KEY_RECORDS),
+ new TestRequest(1, 1),
+ new TestRequest(2, 2),
+ new TestRequest(3, 3));
+
+ assertThat(outputs).containsExactlyInAnyOrder("1:run", "2:run", "3:run");
+
+ List<Span> spans = TimingLog.spans();
+ assertThat(spans).hasSize(MULTI_KEY_RECORDS);
+ assertThat(spans)
+ .extracting(span -> span.recordId)
+ .containsExactlyInAnyOrder("1", "2", "3");
+ assertAllPairsOverlap("multi-key inputs", spans);
+ assertThat(distinctThreads(spans)).isGreaterThanOrEqualTo(2);
+ assertSpeedup("multi-key inputs", spans, 2.0);
+ }
+
+ /** Agent with a single action issuing exactly one durableExecuteAsync call. */
+ public static class MultiKeyAgent extends Agent {
+
+ @Action(InputEvent.EVENT_TYPE)
+ public static void run(Event event, RunnerContext ctx) throws Exception {
+ TestRequest request = (TestRequest) InputEvent.fromEvent(event).getInput();
+ String marker =
+ durableTiming(
+ ctx,
+ "multikey-run-" + request.recordId,
+ "run",
+ request.key,
+ request.recordId,
+ 1,
+ MULTI_KEY_SLEEP_MS,
+ () -> {
+ // Start-alignment barrier; runs before the timed window opens.
+ MULTI_KEY_ALL_STARTED.countDown();
+ if (!MULTI_KEY_ALL_STARTED.await(10, TimeUnit.SECONDS)) {
+ throw new IllegalStateException(
+ "Timed out waiting for all multi-key actions to start");
+ }
+ });
+ ctx.sendEvent(new OutputEvent(marker));
+ }
+ }
+
+ @Test
+ @Timeout(20)
+ public void parallelSiblingsKeepActionScopedContext() throws Exception {
+ List<String> outputs =
+ runJob(
+ new ContextSwitchAgent(),
+ agentsEnv ->
+ agentsEnv
+ .getConfig()
+ .set(AgentExecutionOptions.NUM_ASYNC_THREADS, CTX_HANDLERS),
+ new TestRequest(1, 1));
+
+ assertThat(outputs)
+ .containsExactlyInAnyOrder("1:handleAlpha", "1:handleBeta", "1:handleGamma");
+
+ List<ContextObservation> observations = ContextObservations.snapshot();
+ assertThat(observations).hasSize(CTX_HANDLERS);
+ List<String> handlers = new ArrayList<>();
+ for (ContextObservation observation : observations) {
+ handlers.add(observation.handler);
+ String expectedOwner = CTX_EXPECTED_OWNER_BY_HANDLER.get(observation.handler);
+ assertThat(expectedOwner)
+ .as("missing expected owner for handler %s", observation.handler)
+ .isNotNull();
+ assertThat(observation.beforeConfig).isEqualTo(expectedOwner);
+ assertThat(observation.beforeMemory).isEqualTo(expectedOwner);
+ assertThat(observation.after1Config).isEqualTo(expectedOwner);
+ assertThat(observation.after1Memory).isEqualTo(expectedOwner);
+ assertThat(observation.after2Config).isEqualTo(expectedOwner);
+ assertThat(observation.after2Memory).isEqualTo(expectedOwner);
+ }
+ assertThat(handlers).containsExactlyInAnyOrder("handleAlpha", "handleBeta", "handleGamma");
+
+ List<Span> firstCalls = spansOfSlot(1);
+ List<Span> secondCalls = spansOfSlot(2);
+ assertThat(firstCalls).hasSize(CTX_HANDLERS);
+ assertThat(secondCalls).hasSize(CTX_HANDLERS);
+ assertAllPairsOverlap("context-switch slot 1", firstCalls);
+ assertAllPairsOverlap("context-switch slot 2", secondCalls);
+ assertThat(distinctThreads(firstCalls)).isEqualTo(CTX_HANDLERS);
+ assertThat(distinctThreads(secondCalls)).isEqualTo(CTX_HANDLERS);
+ assertSpeedup("context-switch slot 1", firstCalls, 2.0);
+ assertSpeedup("context-switch slot 2", secondCalls, 2.0);
+ }
+
+ /** Context values observed by one handler around its two yielding calls. */
+ static final class ContextObservation {
+ final String handler;
+ final String beforeConfig;
+ final String beforeMemory;
+ final String after1Config;
+ final String after1Memory;
+ final String after2Config;
+ final String after2Memory;
+
+ ContextObservation(
+ String handler,
+ String beforeConfig,
+ String beforeMemory,
+ String after1Config,
+ String after1Memory,
+ String after2Config,
+ String after2Memory) {
+ this.handler = handler;
+ this.beforeConfig = beforeConfig;
+ this.beforeMemory = beforeMemory;
+ this.after1Config = after1Config;
+ this.after1Memory = after1Memory;
+ this.after2Config = after2Config;
+ this.after2Memory = after2Memory;
+ }
+ }
+
+ /** In-JVM registry of context observations, cleared before each test. */
+ static final class ContextObservations {
+ private static final List<ContextObservation> OBSERVATIONS =
+ Collections.synchronizedList(new ArrayList<>());
+
+ static void clear() {
+ OBSERVATIONS.clear();
+ }
+
+ static void record(ContextObservation observation) {
+ OBSERVATIONS.add(observation);
+ }
+
+ static List<ContextObservation> snapshot() {
+ synchronized (OBSERVATIONS) {
+ return new ArrayList<>(OBSERVATIONS);
+ }
+ }
+ }
+
+ public static class ContextSwitchEvent extends Event {
+ public static final String EVENT_TYPE = "ParallelExecutionE2ETest.ContextSwitchEvent";
+
+ public final int inputId;
+
+ public ContextSwitchEvent(int inputId) {
+ super(EVENT_TYPE);
+ this.inputId = inputId;
+ }
+ }
+
+ /** Agent whose three same-key siblings each own a distinct action-scoped config value. */
+ public static class ContextSwitchAgent extends Agent {
+
+ public ContextSwitchAgent() {
+ addAction(
+ new String[] {InputEvent.EVENT_TYPE},
+ method(ContextSwitchAgent.class, "prime"));
+ addAction(
+ new String[] {ContextSwitchEvent.EVENT_TYPE},
+ method(ContextSwitchAgent.class, "handleAlpha"),
+ ownerConfig("ALPHA"));
+ addAction(
+ new String[] {ContextSwitchEvent.EVENT_TYPE},
+ method(ContextSwitchAgent.class, "handleBeta"),
+ ownerConfig("BETA"));
+ addAction(
+ new String[] {ContextSwitchEvent.EVENT_TYPE},
+ method(ContextSwitchAgent.class, "handleGamma"),
+ ownerConfig("GAMMA"));
+ }
+
+ public static void prime(Event event, RunnerContext ctx) throws Exception {
+ TestRequest request = (TestRequest) InputEvent.fromEvent(event).getInput();
+ ctx.sendEvent(new ContextSwitchEvent(request.recordId));
+ }
+
+ public static void handleAlpha(Event event, RunnerContext ctx) throws Exception {
+ handle(event, ctx, "handleAlpha");
+ }
+
+ public static void handleBeta(Event event, RunnerContext ctx) throws Exception {
+ handle(event, ctx, "handleBeta");
+ }
+
+ public static void handleGamma(Event event, RunnerContext ctx) throws Exception {
+ handle(event, ctx, "handleGamma");
+ }
+
+ private static void handle(Event event, RunnerContext ctx, String action) throws Exception {
+ ContextSwitchEvent input = (ContextSwitchEvent) event;
+ String expectedOwner = CTX_EXPECTED_OWNER_BY_HANDLER.get(action);
+ if (expectedOwner == null) {
+ throw new IllegalStateException("Missing expected owner for handler: " + action);
+ }
+ ctx.getShortTermMemory().set(CTX_MEMORY_OWNER, expectedOwner);
+
+ String[] before = snapshot(ctx);
+ durableTiming(
+ ctx,
+ "context-switch-" + expectedOwner + "-" + input.inputId + "-1",
+ action,
+ input.inputId,
+ input.inputId,
+ 1,
+ CTX_SLEEP_MS);
+ String[] afterFirst = snapshot(ctx);
+ durableTiming(
+ ctx,
+ "context-switch-" + expectedOwner + "-" + input.inputId + "-2",
+ action,
+ input.inputId,
+ input.inputId,
+ 2,
+ CTX_SLEEP_MS);
+ String[] afterSecond = snapshot(ctx);
+
+ assertSnapshot(action, "before", expectedOwner, before);
+ assertSnapshot(action, "after1", expectedOwner, afterFirst);
+ assertSnapshot(action, "after2", expectedOwner, afterSecond);
+
+ ContextObservations.record(
+ new ContextObservation(
+ action,
+ before[0],
+ before[1],
+ afterFirst[0],
+ afterFirst[1],
+ afterSecond[0],
+ afterSecond[1]));
+ // Emitted only after all yielding calls have returned.
+ ctx.sendEvent(new OutputEvent(input.inputId + ":" + action));
+ }
+
+ /** Returns {@code [configOwner, memoryOwner]} as observed through the runner context. */
+ private static String[] snapshot(RunnerContext ctx) throws Exception {
+ Object configOwner = ctx.getActionConfigValue(CTX_OWNER_CONFIG);
+ MemoryObject memoryOwner = ctx.getShortTermMemory().get(CTX_MEMORY_OWNER);
+ Object memoryValue = memoryOwner == null ? null : memoryOwner.getValue();
+ return new String[] {
+ configOwner == null ? null : configOwner.toString(),
+ memoryValue == null ? null : memoryValue.toString()
+ };
+ }
+
+ private static void assertSnapshot(
+ String action, String phase, String expected, String[] snapshot) {
+ if (!expected.equals(snapshot[0]) || !expected.equals(snapshot[1])) {
+ throw new IllegalStateException(
+ "Context mismatch: action="
+ + action
+ + ", phase="
+ + phase
+ + ", expected="
+ + expected
+ + ", configOwner="
+ + snapshot[0]
+ + ", memoryOwner="
+ + snapshot[1]);
+ }
+ }
+
+ private static Map<String, Object> ownerConfig(String owner) {
+ Map<String, Object> config = new HashMap<>();
+ config.put(CTX_OWNER_CONFIG, owner);
+ return config;
+ }
+ }
+
+ @Test
+ @Timeout(25)
+ public void sameKeyInputsCommitInOrder() throws Exception {
+ List<String> outputs =
+ runJob(
+ new SameKeyInputOrderAgent(),
+ agentsEnv ->
+ agentsEnv
+ .getConfig()
+ .set(
+ AgentExecutionOptions.NUM_ASYNC_THREADS,
+ SAME_KEY_ASYNC_THREADS),
+ new TestRequest(1, SAME_KEY),
+ new TestRequest(2, SAME_KEY));
+
+ // Both markers of input 1 must be emitted before both markers of input 2; the order of
+ // siblings within one input is a scheduling detail that varies across Flink versions.
+ assertThat(outputs).hasSize(SAME_KEY_INPUTS * SAME_KEY_ACTIONS);
+ assertThat(outputs.subList(0, SAME_KEY_ACTIONS))
+ .containsExactlyInAnyOrder("1:actionA", "1:actionB");
+ assertThat(outputs.subList(SAME_KEY_ACTIONS, SAME_KEY_INPUTS * SAME_KEY_ACTIONS))
+ .containsExactlyInAnyOrder("2:actionA", "2:actionB");
+
+ List<Span> spans = TimingLog.spans();
+ assertThat(spans).hasSize(SAME_KEY_INPUTS * SAME_KEY_ACTIONS);
+
+ List<Span> firstInput = spansOfRecord("1");
+ List<Span> secondInput = spansOfRecord("2");
+ long firstInputEnd = firstInput.stream().mapToLong(span -> span.endMs).max().orElseThrow();
+ long secondInputStart =
+ secondInput.stream().mapToLong(span -> span.startMs).min().orElseThrow();
+ assertThat(secondInputStart)
+ .as("The second same-key input must wait for the first input to commit")
+ .isGreaterThanOrEqualTo(firstInputEnd);
+ }
+
+ /** Agent with two sibling actions per input, one long and one short. */
+ public static class SameKeyInputOrderAgent extends Agent {
+
+ public SameKeyInputOrderAgent() {
+ addAction(
+ new String[] {InputEvent.EVENT_TYPE},
+ method(SameKeyInputOrderAgent.class, "actionA"));
+ addAction(
+ new String[] {InputEvent.EVENT_TYPE},
+ method(SameKeyInputOrderAgent.class, "actionB"));
+ }
+
+ public static void actionA(Event event, RunnerContext ctx) throws Exception {
+ executeAndEmit(event, ctx, "actionA", SAME_KEY_LONG_SLEEP_MS);
+ }
+
+ public static void actionB(Event event, RunnerContext ctx) throws Exception {
+ executeAndEmit(event, ctx, "actionB", SAME_KEY_SHORT_SLEEP_MS);
+ }
+
+ private static void executeAndEmit(
+ Event event, RunnerContext ctx, String action, int sleepMs) throws Exception {
+ TestRequest request = (TestRequest) InputEvent.fromEvent(event).getInput();
+ String marker =
+ durableTiming(
+ ctx,
+ "same-key-" + request.recordId + "-" + action,
+ action,
+ request.key,
+ request.recordId,
+ 1,
+ sleepMs);
+ ctx.sendEvent(new OutputEvent(marker));
+ }
+ }
+
+ @Test
+ @Timeout(25)
+ public void tasksBeyondWorkerCountCompleteExactlyOnce() throws Exception {
+ List<String> outputs =
+ runJob(
+ new CapacityAgent(),
+ agentsEnv ->
+ agentsEnv
+ .getConfig()
+ .set(
+ AgentExecutionOptions.NUM_ASYNC_THREADS,
+ CAPACITY_ASYNC_THREADS),
+ new TestRequest(1, 1));
+
+ assertThat(outputs)
+ .containsExactlyInAnyOrder(
+ "1:action1",
+ "1:action2",
+ "1:action3",
+ "1:action4",
+ "1:action5",
+ "1:action6");
+
+ List<Span> spans = TimingLog.spans();
+ assertThat(spans).hasSize(CAPACITY_ACTIONS);
+ assertThat(spans)
+ .extracting(span -> span.action)
+ .containsExactlyInAnyOrder(
+ "action1", "action2", "action3", "action4", "action5", "action6");
+ }
+
+ /** Agent registering more sibling actions than async workers. */
+ public static class CapacityAgent extends Agent {
+
+ public CapacityAgent() {
+ for (int i = 1; i <= CAPACITY_ACTIONS; i++) {
+ addAction(
+ new String[] {InputEvent.EVENT_TYPE},
+ method(CapacityAgent.class, "action" + i));
+ }
+ }
+
+ public static void action1(Event event, RunnerContext ctx) throws Exception {
+ executeAndEmit(event, ctx, "action1");
+ }
+
+ public static void action2(Event event, RunnerContext ctx) throws Exception {
+ executeAndEmit(event, ctx, "action2");
+ }
+
+ public static void action3(Event event, RunnerContext ctx) throws Exception {
+ executeAndEmit(event, ctx, "action3");
+ }
+
+ public static void action4(Event event, RunnerContext ctx) throws Exception {
+ executeAndEmit(event, ctx, "action4");
+ }
+
+ public static void action5(Event event, RunnerContext ctx) throws Exception {
+ executeAndEmit(event, ctx, "action5");
+ }
+
+ public static void action6(Event event, RunnerContext ctx) throws Exception {
+ executeAndEmit(event, ctx, "action6");
+ }
+
+ private static void executeAndEmit(Event event, RunnerContext ctx, String action)
+ throws Exception {
+ TestRequest request = (TestRequest) InputEvent.fromEvent(event).getInput();
+ String marker =
+ durableTiming(
+ ctx,
+ "capacity-" + action,
+ action,
+ request.key,
+ request.recordId,
+ 1,
+ CAPACITY_SLEEP_MS);
+ ctx.sendEvent(new OutputEvent(marker));
+ }
+ }
+
+ @Test
+ @Timeout(20)
+ public void asyncFailureFailsTheJob() throws Exception {
+ StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
+ env.setParallelism(1);
+ DataStream<TestRequest> inputStream = env.fromElements(new TestRequest(1, 1));
+
+ AgentsExecutionEnvironment agentsEnv =
+ AgentsExecutionEnvironment.getExecutionEnvironment(env);
+ agentsEnv.getConfig().set(AgentExecutionOptions.NUM_ASYNC_THREADS, 1);
+
+ DataStream<Object> outputStream =
+ agentsEnv
+ .fromDataStream(inputStream, new TestRequestKeySelector())
+ .apply(new FailureAgent())
+ .toDataStream();
+
+ CloseableIterator<Object> results = outputStream.collectAsync();
+ Throwable failure = catchThrowable(agentsEnv::execute);
+ if (failure == null) {
+ results.close();
+ }
+
+ assertThat(failure).isNotNull();
+ assertThat(failure).hasStackTraceContaining(FAILURE_MARKER);
+ }
+
+ @Test
+ @Timeout(20)
+ public void failingSiblingFailsTheJobWithoutHanging() throws Exception {
+ StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
+ env.setParallelism(1);
+ DataStream<TestRequest> inputStream = env.fromElements(new TestRequest(1, 1));
+
+ AgentsExecutionEnvironment agentsEnv =
+ AgentsExecutionEnvironment.getExecutionEnvironment(env);
+ agentsEnv.getConfig().set(AgentExecutionOptions.NUM_ASYNC_THREADS, 2);
+
+ DataStream<Object> outputStream =
+ agentsEnv
+ .fromDataStream(inputStream, new TestRequestKeySelector())
+ .apply(new MultiSiblingFailureAgent())
+ .toDataStream();
+
+ CloseableIterator<Object> results = outputStream.collectAsync();
+ Throwable failure = catchThrowable(agentsEnv::execute);
+ if (failure == null) {
+ results.close();
+ }
+
+ assertThat(failure).isNotNull();
+ assertThat(failure).hasStackTraceContaining(FAILURE_MARKER);
+ }
+
+ /** Agent whose single action always fails on the async worker. */
+ public static class FailureAgent extends Agent {
+
+ public FailureAgent() {
+ addAction(new String[] {InputEvent.EVENT_TYPE}, method(FailureAgent.class, "fail"));
+ }
+
+ public static void fail(Event event, RunnerContext ctx) throws Exception {
+ durableFailure(ctx, "scenario-failure", FAILURE_MARKER);
+ }
+ }
+
+ /** Agent pairing a healthy sibling with a failing one on the same input. */
+ public static class MultiSiblingFailureAgent extends Agent {
+
+ public MultiSiblingFailureAgent() {
+ addAction(
+ new String[] {InputEvent.EVENT_TYPE},
+ method(MultiSiblingFailureAgent.class, "succeed"));
+ addAction(
+ new String[] {InputEvent.EVENT_TYPE},
+ method(MultiSiblingFailureAgent.class, "fail"));
+ }
+
+ public static void succeed(Event event, RunnerContext ctx) throws Exception {
+ String result =
+ ctx.durableExecuteAsync(
+ new DurableCallable<String>() {
+ @Override
+ public String getId() {
+ return "scenario-sibling-success";
+ }
+
+ @Override
+ public Class<String> getResultClass() {
+ return String.class;
+ }
+
+ @Override
+ public String call() {
+ return "sibling-success";
+ }
+ })
+ .await();
+ ctx.sendEvent(new OutputEvent(result));
+ }
+
+ public static void fail(Event event, RunnerContext ctx) throws Exception {
+ durableFailure(ctx, "scenario-sibling-failure", FAILURE_MARKER + "-sibling");
+ }
+ }
+
+ @Test
+ @Timeout(25)
+ public void maxInFlightInputRecordsThrottlesAdmission() throws Exception {
+ // Input admission backpressure is only enforced by the JDK<21 blocking-inline engine.
+ Assumptions.assumeFalse(
+ ContinuationActionExecutor.isContinuationSupported(),
+ "MAX_IN_FLIGHT_INPUT_RECORDS is only enforced by the JDK<21 blocking-inline"
+ + " engine");
+ List<String> outputs =
+ runJob(
+ new BackpressureAgent(),
+ agentsEnv -> {
+ agentsEnv
+ .getConfig()
+ .set(
+ AgentExecutionOptions.NUM_ASYNC_THREADS,
+ BACKPRESSURE_ASYNC_THREADS);
+ agentsEnv
+ .getConfig()
+ .set(
+ AgentExecutionOptions.MAX_IN_FLIGHT_INPUT_RECORDS,
+ BACKPRESSURE_MAX_IN_FLIGHT);
+ },
+ new TestRequest(1, 1),
+ new TestRequest(2, 2));
+
+ assertThat(outputs).containsExactlyInAnyOrder("1:run", "2:run");
+
+ List<Span> spans = TimingLog.spans();
+ spans.sort(Comparator.comparing(span -> span.recordId));
+ assertThat(spans).hasSize(2);
+ assertThat(spans).extracting(span -> span.recordId).containsExactly("1", "2");
+
+ Span first = spans.get(0);
+ Span second = spans.get(1);
+ assertThat(second.startMs)
+ .as("The second different-key input must wait for the first input to retire")
+ .isGreaterThanOrEqualTo(first.endMs);
+ }
+
+ /** Agent with a single action; admission throttling happens before it runs. */
+ public static class BackpressureAgent extends Agent {
+
+ public BackpressureAgent() {
+ addAction(new String[] {InputEvent.EVENT_TYPE}, method(BackpressureAgent.class, "run"));
+ }
+
+ public static void run(Event event, RunnerContext ctx) throws Exception {
+ TestRequest request = (TestRequest) InputEvent.fromEvent(event).getInput();
+ String marker =
+ durableTiming(
+ ctx,
+ "backpressure-" + request.recordId,
+ "run",
+ request.key,
+ request.recordId,
+ 1,
+ BACKPRESSURE_SLEEP_MS);
+ ctx.sendEvent(new OutputEvent(marker));
+ }
+ }
+}
diff --git a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/PythonResourceWorkerLifetimeE2ETest.java b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/PythonResourceWorkerLifetimeE2ETest.java
new file mode 100644
index 00000000..6871e5d5
--- /dev/null
+++ b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/PythonResourceWorkerLifetimeE2ETest.java
@@ -0,0 +1,75 @@
+/*
+ * 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.resource.test;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Native Pemja regression coverage for cached resources outliving managed workers. */
+public class PythonResourceWorkerLifetimeE2ETest {
+
+ @TempDir private Path temporaryDirectory;
+
+ @Test
+ @Timeout(45)
+ public void testCachedPythonResourceOutlivesCreatingWorkerInterpreter() throws Exception {
+ Path outputFile = temporaryDirectory.resolve("scenario-output.log");
+ Path errorFile = temporaryDirectory.resolve("hs_err_pid%p.log");
+
+ List<String> command = new ArrayList<>();
+ command.add(Path.of(System.getProperty("java.home"), "bin", "java").toString());
+ command.add("--add-opens=java.base/java.util=ALL-UNNAMED");
+ command.add("--add-opens=java.base/java.nio=ALL-UNNAMED");
+ command.add("--add-exports=java.base/jdk.internal.vm=ALL-UNNAMED");
+ command.add("-XX:ErrorFile=" + errorFile);
+ command.add("-cp");
+ command.add(
+ System.getProperty(
+ "surefire.test.class.path", System.getProperty("java.class.path")));
+ command.add(PythonResourceWorkerLifetimeScenario.class.getName());
+
+ Process process =
+ new ProcessBuilder(command)
+ .directory(temporaryDirectory.toFile())
+ .redirectErrorStream(true)
+ .redirectOutput(outputFile.toFile())
+ .start();
+
+ boolean finished = process.waitFor(30, TimeUnit.SECONDS);
+ if (!finished) {
+ process.destroyForcibly().waitFor();
+ }
+ String output = Files.exists(outputFile) ? Files.readString(outputFile) : "";
+
+ assertThat(finished).as("native lifetime scenario timed out; output:%n%s", output).isTrue();
+ assertThat(process.exitValue())
+ .as("native lifetime scenario failed; output:%n%s", output)
+ .isZero();
+ assertThat(output).contains(PythonResourceWorkerLifetimeScenario.SUCCESS_MARKER);
+ }
+}
diff --git a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/PythonResourceWorkerLifetimeScenario.java b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/PythonResourceWorkerLifetimeScenario.java
new file mode 100644
index 00000000..91f3bdb9
--- /dev/null
+++ b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/PythonResourceWorkerLifetimeScenario.java
@@ -0,0 +1,369 @@
+/*
+ * 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.resource.test;
+
+import org.apache.flink.agents.api.resource.Resource;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
+import org.apache.flink.agents.api.resource.ResourceType;
+import org.apache.flink.agents.api.resource.python.PythonResourceWrapper;
+import org.apache.flink.agents.plan.resourceprovider.PythonResourceProvider;
+import org.apache.flink.agents.plan.resourceprovider.ResourceProvider;
+import org.apache.flink.agents.runtime.PythonMCPResourceDiscovery;
+import org.apache.flink.agents.runtime.ResourceCache;
+import org.apache.flink.agents.runtime.operator.parallel.ParallelExecutionCoordinator;
+import org.apache.flink.agents.runtime.operator.parallel.ParallelExecutionLock;
+import org.apache.flink.agents.runtime.operator.parallel.ParallelExecutionTask;
+import org.apache.flink.agents.runtime.python.utils.PythonInterpreterManager;
+import org.apache.flink.agents.runtime.python.utils.PythonResourceAdapterImpl;
+import org.apache.flink.util.function.ThrowingRunnable;
+import pemja.core.PythonInterpreter;
+import pemja.core.PythonInterpreterConfig;
+
+import javax.annotation.Nullable;
+
+import java.io.BufferedReader;
+import java.io.InputStreamReader;
+import java.net.URL;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Queue;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
+
+/** Child-process scenario for {@link PythonResourceWorkerLifetimeE2ETest}. */
+public final class PythonResourceWorkerLifetimeScenario {
+
+ static final String SUCCESS_MARKER = "PYTHON_RESOURCE_WORKER_LIFETIME_OK";
+
+ private static final String RESOURCE_NAME = "lifetimeModel";
+ private static final String DEPENDENCY_RESOURCE_NAME = "lifetimeDependency";
+ private static final long WORKER_IDLE_TIMEOUT_MILLIS = 200L;
+
+ private PythonResourceWorkerLifetimeScenario() {}
+
+ public static void main(String[] args) throws Exception {
+ runScenario();
+ System.out.println(SUCCESS_MARKER);
+ }
+
+ private static void runScenario() throws Exception {
+ PythonInterpreterConfig interpreterConfig = createInterpreterConfig();
+ PythonInterpreter ownerInterpreter = new PythonInterpreter(interpreterConfig);
+ ParallelExecutionLock executionLock = new ParallelExecutionLock();
+
+ try (PythonInterpreterManager interpreterManager =
+ new PythonInterpreterManager(
+ ownerInterpreter,
+ () -> new PythonInterpreter(interpreterConfig),
+ 1,
+ executionLock::checkReentrant)) {
+ Map<ResourceType, Map<String, ResourceProvider>> providers = createProviders();
+ ResourceCache resourceCache = new ResourceCache(providers);
+ PythonResourceAdapterImpl adapter =
+ new PythonResourceAdapterImpl(
+ resourceCache.getResourceContext(), interpreterManager, null);
+ adapter.open();
+ PythonMCPResourceDiscovery.discoverPythonMCPResources(
+ providers, adapter, resourceCache);
+
+ try {
+ executeWorkerLifetimeScenario(
+ executionLock, interpreterManager, resourceCache, adapter);
+ } finally {
+ try {
+ resourceCache.close();
+ } finally {
+ adapter.close();
+ }
+ }
+ }
+ }
+
+ private static void executeWorkerLifetimeScenario(
+ ParallelExecutionLock executionLock,
+ PythonInterpreterManager interpreterManager,
+ ResourceCache resourceCache,
+ PythonResourceAdapterImpl adapter)
+ throws Exception {
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails = new LinkedBlockingQueue<>();
+ Queue<ParallelExecutionTask> tasks = new ConcurrentLinkedQueue<>();
+ CountDownLatch residentWorkerStarted = new CountDownLatch(1);
+ CountDownLatch releaseResidentWorker = new CountDownLatch(1);
+ CountDownLatch resourceCreated = new CountDownLatch(1);
+ CountDownLatch resourceReused = new CountDownLatch(1);
+ CountDownLatch surplusWorkerCleanedUp = new CountDownLatch(1);
+
+ ParallelExecutionCoordinator coordinator =
+ new ParallelExecutionCoordinator(
+ executionLock,
+ (mail, description) -> mails.add(mail),
+ tasks::poll,
+ () -> {
+ interpreterManager.releaseCurrentThreadInterpreter();
+ surplusWorkerCleanedUp.countDown();
+ },
+ 2,
+ WORKER_IDLE_TIMEOUT_MILLIS);
+ try {
+ tasks.add(
+ ScenarioTask.blocking(
+ executionLock, residentWorkerStarted, releaseResidentWorker));
+ addTask(executionLock, coordinator, "resident");
+ await(residentWorkerStarted, "resident worker did not start");
+
+ tasks.add(
+ ScenarioTask.running(
+ () -> {
+ resourceCache.getResource(RESOURCE_NAME, ResourceType.CHAT_MODEL);
+ resourceCreated.countDown();
+ }));
+ addTask(executionLock, coordinator, "creator");
+ await(resourceCreated, "surplus worker did not create the Python resource");
+ runOneMail(mails);
+
+ await(
+ surplusWorkerCleanedUp,
+ "surplus worker did not retire and release its interpreter");
+
+ tasks.add(
+ ScenarioTask.running(
+ () -> {
+ Resource cached =
+ resourceCache.getResource(
+ RESOURCE_NAME, ResourceType.CHAT_MODEL);
+ Object result =
+ adapter.callMethod(
+ ((PythonResourceWrapper) cached)
+ .getPythonResource(),
+ "touch",
+ Map.of());
+ if (!"alive".equals(result)) {
+ throw new AssertionError(
+ "Unexpected cached resource result: " + result);
+ }
+ resourceReused.countDown();
+ }));
+ addTask(executionLock, coordinator, "reuser");
+ await(resourceReused, "replacement worker did not reuse the cached resource");
+ runOneMail(mails);
+
+ releaseResidentWorker.countDown();
+ runOneMail(mails);
+ } finally {
+ releaseResidentWorker.countDown();
+ coordinator.close();
+ }
+ }
+
+ private static Map<ResourceType, Map<String, ResourceProvider>> createProviders() {
+ ResourceDescriptor descriptor =
+ new ResourceDescriptor(
+ "resource_lifetime_fixture",
+ "LifetimeChatModelSetup",
+ Map.of("connection", "unused", "model", "unused"));
+ PythonResourceProvider provider =
+ new PythonResourceProvider(RESOURCE_NAME, ResourceType.CHAT_MODEL, descriptor);
+ ResourceDescriptor dependencyDescriptor =
+ new ResourceDescriptor(
+ "resource_lifetime_fixture",
+ "LifetimeDependencySetup",
+ Map.of("connection", "unused", "model", "unused"));
+ PythonResourceProvider dependencyProvider =
+ new PythonResourceProvider(
+ DEPENDENCY_RESOURCE_NAME, ResourceType.CHAT_MODEL, dependencyDescriptor);
+ Map<ResourceType, Map<String, ResourceProvider>> providers = new HashMap<>();
+ providers.put(
+ ResourceType.CHAT_MODEL,
+ Map.of(RESOURCE_NAME, provider, DEPENDENCY_RESOURCE_NAME, dependencyProvider));
+ return providers;
+ }
+
+ private static PythonInterpreterConfig createInterpreterConfig() throws Exception {
+ URL fixtureDirectory =
+ PythonResourceWorkerLifetimeScenario.class.getClassLoader().getResource("python");
+ if (fixtureDirectory == null) {
+ throw new IllegalStateException("Python test fixture directory was not found.");
+ }
+
+ List<String> pythonPaths = new ArrayList<>();
+ pythonPaths.add(Paths.get(fixtureDirectory.toURI()).toString());
+ String inheritedPythonPath = System.getenv("PYTHONPATH");
+ if (inheritedPythonPath != null && !inheritedPythonPath.isBlank()) {
+ pythonPaths.add(inheritedPythonPath);
+ }
+
+ return PythonInterpreterConfig.newBuilder()
+ .setExcType(PythonInterpreterConfig.ExecType.MULTI_THREAD)
+ .setPythonExec(findPythonExecutable())
+ .addPythonPaths(pythonPaths.toArray(new String[0]))
+ .build();
+ }
+
+ private static String findPythonExecutable() throws Exception {
+ String virtualEnvironment = System.getenv("VIRTUAL_ENV");
+ if (virtualEnvironment != null) {
+ Path candidate = Path.of(virtualEnvironment, "bin", "python");
+ if (Files.isExecutable(candidate)) {
+ return candidate.toString();
+ }
+ }
+
+ Process process =
+ new ProcessBuilder("python3", "-c", "import sys; print(sys.executable)")
+ .redirectErrorStream(true)
+ .start();
+ String executable;
+ try (BufferedReader reader =
+ new BufferedReader(new InputStreamReader(process.getInputStream()))) {
+ executable = reader.readLine();
+ }
+ if (!process.waitFor(10, TimeUnit.SECONDS) || process.exitValue() != 0) {
+ throw new IllegalStateException("Unable to resolve the active Python executable.");
+ }
+ if (executable == null || executable.isBlank()) {
+ throw new IllegalStateException("Python executable lookup returned no path.");
+ }
+ return executable.trim();
+ }
+
+ private static void addTask(
+ ParallelExecutionLock executionLock,
+ ParallelExecutionCoordinator coordinator,
+ Object key)
+ throws InterruptedException {
+ executionLock.acquireByMain();
+ try {
+ coordinator.addTask(key);
+ } finally {
+ executionLock.release();
+ }
+ }
+
+ private static void runOneMail(BlockingQueue<ThrowingRunnable<? extends Exception>> mails)
+ throws Exception {
+ ThrowingRunnable<? extends Exception> mail = mails.poll(10, TimeUnit.SECONDS);
+ if (mail == null) {
+ throw new AssertionError("Timed out waiting for mailbox completion.");
+ }
+ mail.run();
+ }
+
+ private static void await(CountDownLatch latch, String message) throws InterruptedException {
+ if (!latch.await(10, TimeUnit.SECONDS)) {
+ throw new AssertionError(message);
+ }
+ }
+
+ private static final class ScenarioTask implements ParallelExecutionTask {
+ private final ThrowingRunnable<? extends Exception> action;
+ private final ParallelExecutionLock executionLock;
+ private final CountDownLatch started;
+ private final CountDownLatch release;
+
+ private long recordIndex;
+ private long taskIndex;
+ private volatile boolean done;
+ private Throwable failure;
+
+ private ScenarioTask(ThrowingRunnable<? extends Exception> action) {
+ this(action, null, null, null);
+ }
+
+ private ScenarioTask(
+ ThrowingRunnable<? extends Exception> action,
+ ParallelExecutionLock executionLock,
+ CountDownLatch started,
+ CountDownLatch release) {
+ this.action = action;
+ this.executionLock = executionLock;
+ this.started = started;
+ this.release = release;
+ }
+
+ private static ScenarioTask running(ThrowingRunnable<? extends Exception> action) {
+ return new ScenarioTask(action);
+ }
+
+ private static ScenarioTask blocking(
+ ParallelExecutionLock executionLock,
+ CountDownLatch started,
+ CountDownLatch release) {
+ return new ScenarioTask(() -> {}, executionLock, started, release);
+ }
+
+ @Override
+ public void setup(Object key, long recordIndex, long taskIndex) {
+ this.recordIndex = recordIndex;
+ this.taskIndex = taskIndex;
+ }
+
+ @Override
+ public void restoreContext() {}
+
+ @Override
+ public void execute() {
+ try {
+ if (executionLock == null) {
+ action.run();
+ } else {
+ started.countDown();
+ executionLock.release();
+ try {
+ await(release, "resident worker was not released");
+ } finally {
+ executionLock.acquireByWorker(recordIndex, taskIndex);
+ }
+ }
+ } catch (Throwable throwable) {
+ failure = throwable;
+ } finally {
+ done = true;
+ }
+ }
+
+ @Override
+ public boolean isDone() {
+ return done;
+ }
+
+ @Override
+ public void commit() throws Exception {
+ if (failure instanceof Exception) {
+ throw (Exception) failure;
+ }
+ if (failure instanceof Error) {
+ throw (Error) failure;
+ }
+ if (failure != null) {
+ throw new RuntimeException(failure);
+ }
+ }
+
+ @Override
+ public void finishGroup(@Nullable ParallelExecutionTask lastCommitted) {}
+ }
+}
diff --git a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/resources/python/resource_lifetime_fixture.py b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/resources/python/resource_lifetime_fixture.py
new file mode 100644
index 00000000..dc67bc68
--- /dev/null
+++ b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/resources/python/resource_lifetime_fixture.py
@@ -0,0 +1,55 @@
+################################################################################
+# 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.
+################################################################################
+
+from typing import Any, Dict
+
+from flink_agents.api.chat_models.chat_model import BaseChatModelSetup
+from flink_agents.api.resource import ResourceType
+
+
+class LifetimeChatModelSetup(BaseChatModelSetup):
+ """Minimal resource used to exercise Pemja handle lifetime."""
+
+ @property
+ def model_kwargs(self) -> Dict[str, Any]:
+ return {}
+
+ def open(self) -> None:
+ self.resource_context.get_resource(
+ "lifetimeDependency", ResourceType.CHAT_MODEL
+ )
+
+ def touch(self) -> str:
+ return "alive"
+
+ def close(self) -> None:
+ pass
+
+
+class LifetimeDependencySetup(BaseChatModelSetup):
+ """Nested Python resource created while the outer resource is opening."""
+
+ @property
+ def model_kwargs(self) -> Dict[str, Any]:
+ return {}
+
+ def open(self) -> None:
+ pass
+
+ def close(self) -> None:
+ pass
diff --git a/plan/src/main/java/org/apache/flink/agents/plan/resourceprovider/PythonResourceProvider.java b/plan/src/main/java/org/apache/flink/agents/plan/resourceprovider/PythonResourceProvider.java
index 54fe1d30..6ba39cc6 100644
--- a/plan/src/main/java/org/apache/flink/agents/plan/resourceprovider/PythonResourceProvider.java
+++ b/plan/src/main/java/org/apache/flink/agents/plan/resourceprovider/PythonResourceProvider.java
@@ -29,6 +29,7 @@ import org.apache.flink.agents.api.resource.ResourceType;
import org.apache.flink.agents.api.resource.python.PythonResourceAdapter;
import org.apache.flink.agents.api.vectorstores.python.PythonCollectionManageableVectorStore;
import org.apache.flink.agents.plan.resource.python.PythonMCPServer;
+import org.apache.flink.util.ExceptionUtils;
import pemja.core.object.PyObject;
import java.lang.reflect.Constructor;
@@ -114,15 +115,25 @@ public class PythonResourceProvider extends ResourceProvider {
}
PyObject pyResource = pythonResourceAdapter.initPythonResource(pyModule, pyClazz, kwargs);
- Constructor<?> constructor =
- clazz.getConstructor(
- PythonResourceAdapter.class,
- PyObject.class,
- ResourceDescriptor.class,
- ResourceContext.class);
- return (Resource)
- constructor.newInstance(
- pythonResourceAdapter, pyResource, descriptor, resourceContext);
+ try {
+ Constructor<?> constructor =
+ clazz.getConstructor(
+ PythonResourceAdapter.class,
+ PyObject.class,
+ ResourceDescriptor.class,
+ ResourceContext.class);
+ return (Resource)
+ constructor.newInstance(
+ pythonResourceAdapter, pyResource, descriptor, resourceContext);
+ } catch (Throwable constructionFailure) {
+ try {
+ pyResource.close();
+ } catch (Throwable closeFailure) {
+ constructionFailure.addSuppressed(closeFailure);
+ }
+ ExceptionUtils.rethrowException(constructionFailure);
+ throw new AssertionError("Unreachable after rethrowing resource construction failure");
+ }
}
@Override
diff --git a/plan/src/test/java/org/apache/flink/agents/plan/resourceprovider/PythonResourceProviderTest.java b/plan/src/test/java/org/apache/flink/agents/plan/resourceprovider/PythonResourceProviderTest.java
index 9dc54824..9199ee7d 100644
--- a/plan/src/test/java/org/apache/flink/agents/plan/resourceprovider/PythonResourceProviderTest.java
+++ b/plan/src/test/java/org/apache/flink/agents/plan/resourceprovider/PythonResourceProviderTest.java
@@ -22,14 +22,19 @@ import org.apache.flink.agents.api.resource.ResourceContext;
import org.apache.flink.agents.api.resource.ResourceDescriptor;
import org.apache.flink.agents.api.resource.ResourceType;
import org.apache.flink.agents.api.resource.python.PythonResourceAdapter;
+import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.EnumSource;
import pemja.core.object.PyObject;
+import java.lang.reflect.InvocationTargetException;
import java.util.Map;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.catchThrowable;
import static org.mockito.ArgumentMatchers.anyMap;
import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
@@ -62,4 +67,58 @@ class PythonResourceProviderTest {
verify(adapter).callMethod(pythonResource, "close", Map.of());
verify(pythonResource).close();
}
+
+ @Test
+ void closesPythonHandleWhenWrapperConstructionFails() throws Exception {
+ PythonResourceAdapter adapter = mock(PythonResourceAdapter.class);
+ PyObject pythonResource = mock(PyObject.class);
+ PythonResourceProvider provider = createProviderWithInvalidWrapperArgument(adapter);
+ when(adapter.initPythonResource(anyString(), anyString(), anyMap()))
+ .thenReturn(pythonResource);
+
+ Throwable failure = catchThrowable(() -> provider.provide(mock(ResourceContext.class)));
+
+ assertThat(failure)
+ .isInstanceOf(InvocationTargetException.class)
+ .hasCauseInstanceOf(IllegalArgumentException.class);
+ assertThat(failure.getCause())
+ .hasMessage(
+ "Unknown structured output strategy 'invalid'. Expected one of: AUTO, NATIVE, PROMPT.");
+ verify(pythonResource).close();
+ }
+
+ @Test
+ void preservesWrapperConstructionFailureWhenPythonHandleCloseFails() throws Exception {
+ PythonResourceAdapter adapter = mock(PythonResourceAdapter.class);
+ PyObject pythonResource = mock(PyObject.class);
+ RuntimeException closeFailure = new RuntimeException("close failed");
+ PythonResourceProvider provider = createProviderWithInvalidWrapperArgument(adapter);
+ when(adapter.initPythonResource(anyString(), anyString(), anyMap()))
+ .thenReturn(pythonResource);
+ doThrow(closeFailure).when(pythonResource).close();
+
+ Throwable failure = catchThrowable(() -> provider.provide(mock(ResourceContext.class)));
+
+ assertThat(failure)
+ .isInstanceOf(InvocationTargetException.class)
+ .hasCauseInstanceOf(IllegalArgumentException.class);
+ assertThat(failure.getCause())
+ .hasMessage(
+ "Unknown structured output strategy 'invalid'. Expected one of: AUTO, NATIVE, PROMPT.");
+ assertThat(failure.getSuppressed()).containsExactly(closeFailure);
+ verify(pythonResource).close();
+ }
+
+ private static PythonResourceProvider createProviderWithInvalidWrapperArgument(
+ PythonResourceAdapter adapter) {
+ ResourceDescriptor descriptor =
+ new ResourceDescriptor(
+ "example.module",
+ "ExampleModel",
+ Map.of("structured_output_strategy", "invalid"));
+ PythonResourceProvider provider =
+ new PythonResourceProvider("model", ResourceType.CHAT_MODEL, descriptor);
+ provider.setPythonResourceAdapter(adapter);
+ return provider;
+ }
}
diff --git a/python/flink_agents/api/core_options.py b/python/flink_agents/api/core_options.py
index 4db30bae..10651b57 100644
--- a/python/flink_agents/api/core_options.py
+++ b/python/flink_agents/api/core_options.py
@@ -292,6 +292,26 @@ class AgentExecutionOptions:
default=os.cpu_count() * 2,
)
+ # Experimental fallback switch for the Java JDK<21 parallel execution engine.
+ # Only consulted for pure-Java agents without coroutine support; plans
+ # containing Python actions never use the parallel engine regardless of
+ # this value.
+ PARALLEL_EXECUTION_ENABLED = ConfigOption(
+ key="parallel-execution.enabled",
+ config_type=bool,
+ default=True,
+ )
+
+ # Maximum number of input records that may be in flight concurrently. Only
+ # enforced by the Java JDK<21 parallel execution engine for pure-Java
+ # agents; the JDK 21 coroutine engine and plans containing Python actions
+ # ignore it.
+ MAX_IN_FLIGHT_INPUT_RECORDS = ConfigOption(
+ key="max-in-flight-input-records",
+ config_type=int,
+ default=100,
+ )
+
CHAT_ASYNC = ConfigOption(
key="chat.async",
config_type=bool,
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java b/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java
index 620c9280..ec6057fe 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java
@@ -163,7 +163,17 @@ public class ResourceCache implements AutoCloseable {
((FunctionTool) resource).setPythonResourceAdapter(pythonResourceAdapter);
}
- resource.open();
+ try {
+ resource.open();
+ } catch (Throwable openFailure) {
+ try {
+ resource.close();
+ } catch (Throwable closeFailure) {
+ openFailure.addSuppressed(closeFailure);
+ }
+ ExceptionUtils.rethrowException(openFailure);
+ throw new AssertionError("Unreachable after rethrowing resource open failure");
+ }
cache.computeIfAbsent(type, k -> new ConcurrentHashMap<>()).put(name, resource);
return resource;
}
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 6a369254..945da425 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
@@ -18,6 +18,10 @@
package org.apache.flink.agents.runtime.async;
import org.apache.flink.agents.api.context.Outcome;
+import org.apache.flink.agents.runtime.operator.parallel.ParallelExecutionContextRestorer;
+import org.apache.flink.agents.runtime.operator.parallel.ParallelExecutionLock;
+
+import javax.annotation.Nullable;
import java.time.Duration;
import java.util.Arrays;
@@ -27,18 +31,39 @@ import java.util.function.Supplier;
import java.util.stream.Collectors;
/**
- * Executor for Java actions that supports asynchronous execution.
+ * Executor for Java actions that supports asynchronous execution (JDK 11 version; JDK 21+ uses a
+ * multi-release variant backed by the Continuation API).
*
- * <p>This is the JDK 11 version that falls back to synchronous execution. On JDK 21+, the
- * Multi-release JAR will use a version that leverages Continuation API for true async execution.
+ * <p>On the parallel path a {@link ParallelExecutionLock} and {@link
+ * ParallelExecutionContextRestorer} are wired in: {@link #executeAsync} releases the lock while the
+ * blocking work runs inline, then re-acquires worker ownership and restores the runner context.
+ * Without a lock it runs synchronously.
*/
public class ContinuationActionExecutor {
+ @Nullable private final ParallelExecutionLock parallelExecutionLock;
+ @Nullable private final ParallelExecutionContextRestorer contextRestorer;
+
/** Creates a new ContinuationActionExecutor. */
- public ContinuationActionExecutor(int numAsyncThreads) {}
+ public ContinuationActionExecutor(int numAsyncThreads) {
+ this(numAsyncThreads, () -> {}, null, null);
+ }
- /** JDK 11 fallback has no worker threads, so the cleanup hook is never needed. */
- public ContinuationActionExecutor(int numAsyncThreads, Runnable threadCleanup) {}
+ /**
+ * Creates a new ContinuationActionExecutor. On the JDK<21 parallel MVP path both {@code
+ * parallelExecutionLock} and {@code contextRestorer} are supplied; otherwise they are {@code
+ * null} and {@link #executeAsync} runs synchronously.
+ *
+ * <p>JDK 11 fallback has no worker threads, so the cleanup hook is never needed.
+ */
+ public ContinuationActionExecutor(
+ int numAsyncThreads,
+ Runnable threadCleanup,
+ @Nullable ParallelExecutionLock parallelExecutionLock,
+ @Nullable ParallelExecutionContextRestorer contextRestorer) {
+ this.parallelExecutionLock = parallelExecutionLock;
+ this.contextRestorer = contextRestorer;
+ }
/**
* Executes the action. In JDK 11, this simply runs the action synchronously.
@@ -61,18 +86,32 @@ public class ContinuationActionExecutor {
* @param <T> the result type
* @return the result of the supplier
*/
- public <T> T executeAsync(ContinuationContext context, Supplier<T> supplier) {
- // JDK 11: Fall back to synchronous execution
- return supplier.get();
+ public <T> T executeAsync(ContinuationContext context, Supplier<T> supplier) throws Exception {
+ if (parallelExecutionLock == null) {
+ // No mailbox lock wired in: fall back to synchronous execution.
+ return supplier.get();
+ }
+
+ parallelExecutionLock.checkReentrant();
+ parallelExecutionLock.release();
+ try {
+ return supplier.get();
+ } finally {
+ parallelExecutionLock.acquireByWorker(context.getRecordIndex(), context.getTaskIndex());
+ contextRestorer.restore(context.getKey(), context.getActionTask());
+ }
}
/**
* Executes all suppliers as one batch. In JDK 11, this falls back to serial execution and
- * captures each supplier's success or failure as an {@link Outcome}.
+ * captures each supplier's success or failure as an {@link Outcome}. When invoked by the
+ * parallel engine, the whole batch runs without the shared operator lock and the task context
+ * is restored after worker ownership is re-acquired.
*
* @param context the continuation context
* @param suppliers the suppliers to execute
* @param timeout ignored in the JDK 11 fallback
+ * @param maxParallelism ignored in the JDK 11 fallback
* @param <T> the result type
* @return outcomes in supplier order
*/
@@ -80,7 +119,23 @@ public class ContinuationActionExecutor {
ContinuationContext context,
List<Callable<T>> suppliers,
Duration timeout,
- int maxParallelism) {
+ int maxParallelism)
+ throws Exception {
+ if (parallelExecutionLock == null || suppliers.isEmpty()) {
+ return executeSuppliers(suppliers);
+ }
+
+ parallelExecutionLock.checkReentrant();
+ parallelExecutionLock.release();
+ try {
+ return executeSuppliers(suppliers);
+ } finally {
+ parallelExecutionLock.acquireByWorker(context.getRecordIndex(), context.getTaskIndex());
+ contextRestorer.restore(context.getKey(), context.getActionTask());
+ }
+ }
+
+ private static <T> BatchExecutionResult<T> executeSuppliers(List<Callable<T>> suppliers) {
List<Outcome<T>> outcomes =
suppliers.stream()
.map(
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationContext.java b/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationContext.java
index 6d0d7e09..f327edfd 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationContext.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationContext.java
@@ -17,10 +17,47 @@
*/
package org.apache.flink.agents.runtime.async;
+import org.apache.flink.agents.runtime.operator.ActionTask;
+
/**
- * Marker context for continuation execution.
- *
- * <p>This class is part of the base (JDK 11) sources. JDK 21 provides a multi-release variant with
- * continuation-specific fields.
+ * Context for continuation execution (base JDK 11 sources). Carries the current key, the {@link
+ * ActionTask}, and its (recordIndex, taskIndex) priority so {@link ContinuationActionExecutor} can
+ * re-acquire the lock at the right priority and restore the runner context on the parallel path.
+ * JDK 21 provides a multi-release variant with continuation-specific state instead.
*/
-public class ContinuationContext {}
+public class ContinuationContext {
+
+ private Object key;
+ private ActionTask actionTask;
+ private long recordIndex;
+ private long taskIndex;
+
+ public void setKey(Object key) {
+ this.key = key;
+ }
+
+ public Object getKey() {
+ return key;
+ }
+
+ public void setActionTask(ActionTask actionTask) {
+ this.actionTask = actionTask;
+ }
+
+ public ActionTask getActionTask() {
+ return actionTask;
+ }
+
+ public void setPriority(long recordIndex, long taskIndex) {
+ this.recordIndex = recordIndex;
+ this.taskIndex = taskIndex;
+ }
+
+ public long getRecordIndex() {
+ return recordIndex;
+ }
+
+ public long getTaskIndex() {
+ return taskIndex;
+ }
+}
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 496fcd7e..5d5dbeca 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
@@ -228,7 +228,7 @@ public class JavaRunnerContextImpl extends RunnerContextImpl {
}
private <T> List<Outcome<T>> executeAllWithoutDurableState(List<DurableCallable<T>> callables)
- throws InterruptedException {
+ throws Exception {
List<Callable<T>> suppliers = new ArrayList<>();
for (DurableCallable<T> callable : callables) {
suppliers.add(callable::call);
@@ -246,7 +246,8 @@ public class JavaRunnerContextImpl extends RunnerContextImpl {
return outcomes;
}
- private <T> BatchExecutionResult<T> executeOutcomeSuppliers(List<Callable<T>> suppliers) {
+ private <T> BatchExecutionResult<T> executeOutcomeSuppliers(List<Callable<T>> suppliers)
+ throws Exception {
if (suppliers.isEmpty()) {
return new BatchExecutionResult<>(List.of(), new boolean[0]);
}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
index ae24256f..9c02520d 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
@@ -19,8 +19,10 @@ package org.apache.flink.agents.runtime.operator;
import org.apache.flink.agents.api.Event;
import org.apache.flink.agents.api.EventContext;
+import org.apache.flink.agents.api.InputEvent;
import org.apache.flink.agents.api.OutputEvent;
import org.apache.flink.agents.api.agents.AgentExecutionOptions;
+import org.apache.flink.agents.api.context.RunnerContext;
import org.apache.flink.agents.api.event.AgentRunBeginEvent;
import org.apache.flink.agents.api.resource.Resource;
import org.apache.flink.agents.api.resource.ResourceType;
@@ -29,9 +31,12 @@ import org.apache.flink.agents.plan.AgentPlan;
import org.apache.flink.agents.plan.JavaFunction;
import org.apache.flink.agents.plan.PythonFunction;
import org.apache.flink.agents.plan.actions.Action;
+import org.apache.flink.agents.plan.resourceprovider.PythonResourceProvider;
import org.apache.flink.agents.runtime.ResourceCache;
import org.apache.flink.agents.runtime.actionstate.ActionState;
import org.apache.flink.agents.runtime.actionstate.ActionStateStore;
+import org.apache.flink.agents.runtime.async.ContinuationActionExecutor;
+import org.apache.flink.agents.runtime.async.ContinuationContext;
import org.apache.flink.agents.runtime.eventlog.EventLogWriter;
import org.apache.flink.agents.runtime.lifecycle.ComponentExecutionListener;
import org.apache.flink.agents.runtime.lifecycle.PythonTaskLifecycleListener;
@@ -42,6 +47,9 @@ import org.apache.flink.agents.runtime.memory.MemoryObjectImpl;
import org.apache.flink.agents.runtime.memory.MemoryUpdateReplayer;
import org.apache.flink.agents.runtime.metrics.BuiltInMetrics;
import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl;
+import org.apache.flink.agents.runtime.operator.parallel.ParallelExecutionCoordinator;
+import org.apache.flink.agents.runtime.operator.parallel.ParallelExecutionLock;
+import org.apache.flink.agents.runtime.operator.parallel.ParallelExecutionTask;
import org.apache.flink.agents.runtime.python.operator.PythonActionTask;
import org.apache.flink.agents.runtime.python.resource.PythonRuntimeResource;
import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor;
@@ -53,6 +61,8 @@ import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.api.common.operators.MailboxExecutor;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.runtime.checkpoint.CheckpointOptions;
+import org.apache.flink.runtime.state.CheckpointStreamFactory;
import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.StateInitializationContext;
import org.apache.flink.runtime.state.StateSnapshotContext;
@@ -61,6 +71,7 @@ import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
import org.apache.flink.streaming.api.operators.BoundedOneInput;
import org.apache.flink.streaming.api.operators.ChainingStrategy;
import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.api.operators.OperatorSnapshotFutures;
import org.apache.flink.streaming.api.operators.Output;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
@@ -69,6 +80,7 @@ import org.apache.flink.streaming.runtime.tasks.StreamTask;
import org.apache.flink.streaming.runtime.tasks.mailbox.MailboxExecutorImpl;
import org.apache.flink.streaming.runtime.tasks.mailbox.MailboxProcessor;
import org.apache.flink.util.ExceptionUtils;
+import org.apache.flink.util.function.ThrowingRunnable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -76,6 +88,7 @@ import javax.annotation.Nullable;
import java.lang.reflect.Field;
import java.util.ArrayList;
+import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
@@ -85,6 +98,9 @@ import java.util.Set;
import java.util.function.IntPredicate;
import static org.apache.flink.agents.api.configuration.AgentConfigOptions.JOB_IDENTIFIER;
+import static org.apache.flink.agents.plan.actions.Utils.supportAsync;
+import static org.apache.flink.util.Preconditions.checkArgument;
+import static org.apache.flink.util.Preconditions.checkNotNull;
import static org.apache.flink.util.Preconditions.checkState;
/**
@@ -105,6 +121,54 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
private static final Logger LOG = LoggerFactory.getLogger(ActionExecutionOperator.class);
+ /** Idle time after which a surplus parallel-execution worker thread retires (min pool is 1). */
+ private static final long WORKER_IDLE_TIMEOUT_MS = 60_000L;
+
+ private static final String INTERNAL_NOOP_INPUT_ACTION_NAME =
+ "__flink_agents_internal_noop_input_action";
+
+ private static final Action INTERNAL_NOOP_INPUT_ACTION = createInternalNoopInputAction();
+
+ private static Action createInternalNoopInputAction() {
+ try {
+ return new Action(
+ INTERNAL_NOOP_INPUT_ACTION_NAME,
+ new JavaFunction(
+ InternalNoopInputAction.class,
+ "run",
+ new Class<?>[] {Event.class, RunnerContext.class}),
+ Collections.singletonList(InputEvent.EVENT_TYPE));
+ } catch (Exception e) {
+ throw new ExceptionInInitializerError(e);
+ }
+ }
+
+ private static boolean isInternalNoopInputAction(Action action) {
+ return INTERNAL_NOOP_INPUT_ACTION_NAME.equals(action.getName());
+ }
+
+ public static final class InternalNoopInputAction {
+ public static void run(Event event, RunnerContext ctx) {}
+
+ private InternalNoopInputAction() {}
+ }
+
+ static boolean shouldEnableParallelExecutionWithoutCoroutine(
+ AgentPlan agentPlan, boolean continuationSupported) {
+ return !continuationSupported
+ && agentPlan.getConfig().get(AgentExecutionOptions.PARALLEL_EXECUTION_ENABLED)
+ && agentPlan.getActions().values().stream()
+ .allMatch(action -> action.getExec() instanceof JavaFunction)
+ && (!usesPythonResource(agentPlan) || supportAsync());
+ }
+
+ /** Whether the plan declares any Python resource. */
+ private static boolean usesPythonResource(AgentPlan agentPlan) {
+ return agentPlan.getResourceProviders().values().stream()
+ .flatMap(providersByName -> providersByName.values().stream())
+ .anyMatch(provider -> provider instanceof PythonResourceProvider);
+ }
+
private final AgentPlan agentPlan;
private transient ResourceCache resourceCache;
@@ -119,6 +183,25 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
private transient ActionTaskContextManager contextManager;
+ private transient boolean parallelExecutionWithoutCoroutineEnabled;
+
+ @Nullable private transient ParallelExecutionLock parallelExecutionLock;
+
+ @Nullable private transient ParallelExecutionCoordinator executionCoordinator;
+
+ /**
+ * Number of input records currently being processed (one per active key), compared against
+ * {@link #maxInFlightInputRecords} to apply input backpressure. Maintained on the mailbox
+ * thread alongside {@code addProcessingKey}/{@code removeProcessingKey} and rebuilt from the
+ * recovered processing keys.
+ */
+ private transient int inFlightInputRecords;
+
+ /**
+ * Configured cap on concurrently processed input records; see {@link #inFlightInputRecords}.
+ */
+ private int maxInFlightInputRecords;
+
// Long-term memory backed by Mem0; non-null only when LongTermMemoryOptions.Mem0 is configured.
private transient Mem0LongTermMemory ltm;
@@ -226,8 +309,36 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
durableExecManager.initRecoveryMarkerState(getOperatorStateBackend());
durableExecManager.initializeKeyedStates(getRuntimeContext());
+ int numAsyncThreads = agentPlan.getConfig().get(AgentExecutionOptions.NUM_ASYNC_THREADS);
+ checkArgument(
+ numAsyncThreads > 0,
+ "%s must be positive, but was %s",
+ AgentExecutionOptions.NUM_ASYNC_THREADS.getKey(),
+ numAsyncThreads);
+
+ maxInFlightInputRecords =
+ agentPlan.getConfig().get(AgentExecutionOptions.MAX_IN_FLIGHT_INPUT_RECORDS);
+ checkArgument(
+ maxInFlightInputRecords > 0,
+ "%s must be positive, but was %s",
+ AgentExecutionOptions.MAX_IN_FLIGHT_INPUT_RECORDS.getKey(),
+ maxInFlightInputRecords);
+
+ parallelExecutionWithoutCoroutineEnabled =
+ shouldEnableParallelExecutionWithoutCoroutine(
+ agentPlan, ContinuationActionExecutor.isContinuationSupported());
+ if (parallelExecutionWithoutCoroutineEnabled) {
+ parallelExecutionLock = new ParallelExecutionLock();
+ } else {
+ // JDK21 cooperative-continuation path: there is no worker pool and no lock. All
+ // engine-shared sections (waitInFlightEventsFinished) must guard on the null lock.
+ parallelExecutionLock = null;
+ }
+
// init PythonActionExecutor and PythonResourceAdapter
pythonBridge = new PythonBridgeManager();
+ // On the parallel engine the execution invariant is "holds the ParallelExecutionLock",
+ // not "is the physical mailbox thread".
pythonBridge.open(
agentPlan,
resourceCache,
@@ -236,10 +347,23 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
getContainingTask().getEnvironment().getTaskManagerInfo().getTmpDirectories(),
getRuntimeContext().getJobInfo().getJobId(),
metricGroup,
- this::checkMailboxThread,
+ parallelExecutionWithoutCoroutineEnabled
+ ? parallelExecutionLock::checkReentrant
+ : this::checkMailboxThread,
jobIdentifier,
getRuntimeContext().getUserCodeClassLoader());
+ if (parallelExecutionWithoutCoroutineEnabled) {
+ executionCoordinator =
+ new ParallelExecutionCoordinator(
+ parallelExecutionLock,
+ mailboxExecutor::execute,
+ Work::new,
+ pythonBridge::releaseCurrentThreadInterpreter,
+ numAsyncThreads,
+ WORKER_IDLE_TIMEOUT_MS);
+ }
+
// Capture the wired Mem0 long-term memory, if any, so it can be plumbed into the Java
// runner context created by ActionTaskContextManager.
ltm = pythonBridge.getLongTermMemory();
@@ -254,11 +378,15 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
registerBuiltInLifecycleListeners();
registerSubagentSetups();
- // init context manager for runner context creation and memory contexts
+ // The continuation executor only borrows the lock on the worker-pool path; on every other
+ // path it is null so executeAsync keeps its synchronous fallback.
contextManager =
new ActionTaskContextManager(
- agentPlan.getConfig().get(AgentExecutionOptions.NUM_ASYNC_THREADS),
- pythonBridge::releaseCurrentThreadInterpreter);
+ this,
+ durableExecManager,
+ numAsyncThreads,
+ pythonBridge::releaseCurrentThreadInterpreter,
+ parallelExecutionLock);
mailboxProcessor = getMailboxProcessor();
@@ -271,17 +399,75 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
// parallelism,
// and {@link tryProcessActionTaskForKey} mails might be lost,
// it is necessary to reprocess all keys to ensure correctness.
- tryResumeProcessActionTasks();
+ if (parallelExecutionWithoutCoroutineEnabled) {
+ runWithParallelExecutionLock(this::tryResumeProcessActionTasks);
+ } else {
+ tryResumeProcessActionTasks();
+ }
}
@Override
public void processWatermark(Watermark mark) throws Exception {
+ if (parallelExecutionWithoutCoroutineEnabled) {
+ runWithParallelExecutionLock(() -> processWatermarkInternal(mark));
+ } else {
+ processWatermarkInternal(mark);
+ }
+ }
+
+ private void processWatermarkInternal(Watermark mark) throws Exception {
eventRouter.getKeySegmentQueue().addWatermark(mark);
eventRouter.processEligibleWatermarks(super::processWatermark);
}
+ /**
+ * On the JDK11 parallel engine, overridden to skip the framework's pre-{@link #processElement}
+ * key switch: the framework performs it on the mailbox thread but <em>outside</em> our mailbox
+ * lock, and a coordinator worker may be executing an action under the lock with a different
+ * current key that an off-lock switch would corrupt. On that engine we re-establish the
+ * record's key under the lock inside {@link #processElement}. On the normal engine there are no
+ * such workers, so we defer to the framework's switch. ({@link
+ * org.apache.flink.streaming.api.operators.Input#setKeyContextElement} default-delegates here,
+ * so this covers the one-input path.)
+ */
+ @Override
+ @SuppressWarnings("rawtypes")
+ public void setKeyContextElement1(StreamRecord record) throws Exception {
+ if (!parallelExecutionWithoutCoroutineEnabled) {
+ super.setKeyContextElement1(record);
+ }
+ }
+
@Override
public void processElement(StreamRecord<IN> record) throws Exception {
+ if (parallelExecutionWithoutCoroutineEnabled) {
+ runWithParallelExecutionLock(
+ () -> {
+ // Soft backpressure BEFORE establishing the key: yields may change the
+ // current key, so the record's key is set after the last yield. Release
+ // around yield (non-reentrant lock); re-acquire in finally for the outer
+ // release.
+ while (inFlightInputRecords >= maxInFlightInputRecords) {
+ parallelExecutionLock.release();
+ try {
+ mailboxExecutor.yield();
+ } finally {
+ parallelExecutionLock.acquireByMain();
+ }
+ }
+ super.setKeyContextElement1(record);
+ processElementInternal(record);
+ });
+ } else {
+ processElementInternal(record);
+ }
+ }
+
+ private void processElementInternal(StreamRecord<IN> record) throws Exception {
+ // Every admitted input record consumes one unit of in-flight budget, regardless of whether
+ // it activates a new key or queues behind an in-flight record for an already-active key.
+ // The matching decrement happens when that record's processing finishes.
+ inFlightInputRecords++;
IN input = record.getValue();
LOG.debug("Receive an element {}", input);
@@ -297,7 +483,11 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
try {
eventRouter.getKeySegmentQueue().addKeyToLastSegment(key);
- if (stateManager.hasMoreActionTasks()) {
+ boolean currentKeyBusy =
+ parallelExecutionWithoutCoroutineEnabled
+ ? executionCoordinator.hasOutstanding(getCurrentKey())
+ : stateManager.hasMoreActionTasks();
+ if (currentKeyBusy) {
// If there are already actions being processed for the current key, the newly
// incoming event should be queued and processed later. Therefore, we add it to
// pendingInputEventsState.
@@ -383,12 +573,26 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
notifyRecordStart(key);
freshRecordRound = false;
}
+ if (parallelExecutionWithoutCoroutineEnabled) {
+ // Add one task node per queued task; a worker pulls+prepares it later.
+ executionCoordinator.addTask(key);
+ }
}
+ } else if (isInputEvent && parallelExecutionWithoutCoroutineEnabled) {
+ stateManager.addActionTask(
+ createActionTask(
+ key,
+ INTERNAL_NOOP_INPUT_ACTION,
+ event,
+ stateManager.getSequenceNumber(),
+ traceContext));
+ executionCoordinator.addTask(key);
}
}
- if (isInputEvent) {
- // If the event is an InputEvent, we submit a new mail to try processing the actions.
+ if (isInputEvent && !parallelExecutionWithoutCoroutineEnabled) {
+ // Kick mail for the normal engine only: it has no workers, so processActionTaskForKey
+ // runs from the mail. The parallel engine's addTask above already dispatched permits.
mailboxExecutor.submit(
() -> tryProcessActionTaskForKey(key, contextKey, traceContext.getInputRunId()),
"process action task");
@@ -441,6 +645,7 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
processEvent(key, contextKey, beginEvent, traceContext);
}
+ /** Kick mail used only on the normal engine to drive queued-task processing for a key. */
private void tryProcessActionTaskForKey(
Object key, String contextKey, @Nullable String inputRunId) {
try {
@@ -457,6 +662,44 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
}
}
+ private void maybeFinishCurrentInput(Object key, @Nullable Work lastCommitted)
+ throws Exception {
+ if (lastCommitted == null) {
+ // This drain committed nothing: the head is still running, or an earlier mail already
+ // drained and retired the input (stale mail). Only the tail-committing mail retires.
+ return;
+ }
+ setCurrentKey(key);
+ if (stateManager.hasMoreActionTasks() || executionCoordinator.hasOutstanding(key)) {
+ return;
+ }
+
+ lastCommitted.actionTask.getRunnerContext().clearSensoryMemory();
+ durableExecManager.updateLastCompletedSequenceNumber(lastCommitted.sequenceNumber);
+ // Mirror the serial path: notify record finished. Noop rounds stay silent.
+ if (!isInternalNoopInputAction(lastCommitted.actionTask.action)) {
+ notifyRecordFinished(key);
+ }
+ int removedCount = stateManager.removeProcessingKey(key);
+ inFlightInputRecords--;
+ checkState(
+ removedCount == 1,
+ "Current processing key count for key "
+ + key
+ + " should be 1, but got "
+ + removedCount);
+ checkState(
+ eventRouter.getKeySegmentQueue().removeKey(key),
+ "Current key" + key + " is missing from the segmentedQueue.");
+ eventRouter.processEligibleWatermarks(super::processWatermark);
+ builtInMetrics.markInputRunCompleted(
+ lastCommitted.actionTask.getTraceContext().getInputRunId());
+ Event pendingInputEvent = pollNextPendingInputEvent();
+ if (pendingInputEvent != null) {
+ processInputEvent(key, pendingInputEvent);
+ }
+ }
+
private void processActionTaskForKey(Object key, String contextKey, @Nullable String inputRunId)
throws Exception {
String currentInputRunId = inputRunId;
@@ -517,6 +760,9 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
// Check if action is already completed
if (actionState != null && actionState.isCompleted()) {
// Action has completed, skip execution and replay memory/events
+ // TODO: unlike the executed path below (and the parallel engine's Work.commit), this
+ // replay path never calls contextManager.removeBundle(actionTask); the bundle lingers
+ // until close. Bounded during recovery, but asymmetric.
LOG.debug(
"Skipping already completed action: {} for key: {}",
actionTask.action.getName(),
@@ -637,6 +883,7 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
// Once all sub-events and actions related to the current InputEvent are completed,
// we can proceed to process the next InputEvent.
int removedCount = stateManager.removeProcessingKey(key);
+ inFlightInputRecords--;
checkState(
removedCount == 1,
"Current processing key count for key "
@@ -670,7 +917,23 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
@VisibleForTesting
public void waitInFlightEventsFinished() throws Exception {
- while (stateManager.hasProcessingKeys()) {
+ checkMailboxThread();
+ while (true) {
+ boolean hasProcessingKeys;
+ if (parallelExecutionLock != null) {
+ // Worker-pool engine: read under the lock so concurrent worker commits settle.
+ parallelExecutionLock.acquireByMain();
+ try {
+ hasProcessingKeys = stateManager.hasProcessingKeys();
+ } finally {
+ parallelExecutionLock.release();
+ }
+ } else {
+ hasProcessingKeys = stateManager.hasProcessingKeys();
+ }
+ if (!hasProcessingKeys) {
+ return;
+ }
mailboxExecutor.yield();
}
}
@@ -688,7 +951,12 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
Throwable firstFailure = null;
for (AutoCloseable closeable :
new AutoCloseable[] {
- resourceCache, contextManager, pythonBridge, eventLogWriter, durableExecManager
+ executionCoordinator,
+ resourceCache,
+ contextManager,
+ pythonBridge,
+ eventLogWriter,
+ durableExecManager
}) {
if (closeable == null) {
continue;
@@ -755,8 +1023,61 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
return getKeyedStateBackend().getKeySerializer();
}
+ @Override
+ public void prepareSnapshotPreBarrier(long checkpointId) throws Exception {
+ if (!parallelExecutionWithoutCoroutineEnabled) {
+ super.prepareSnapshotPreBarrier(checkpointId);
+ return;
+ }
+ // Quiesce pulled-but-uncommitted tasks before the barrier (their results only live in
+ // memory and would be lost on restore): pause dispatch, then yield without the lock
+ // until everything in flight has committed.
+ try {
+ executionCoordinator.startDraining();
+ while (!executionCoordinator.isQuiesced()) {
+ mailboxExecutor.yield();
+ }
+ super.prepareSnapshotPreBarrier(checkpointId);
+ } catch (Exception e) {
+ // A failed pre-barrier never reaches snapshotState's stopDraining; resume dispatch.
+ try {
+ executionCoordinator.stopDraining();
+ } catch (Exception resumeFailure) {
+ e.addSuppressed(resumeFailure);
+ }
+ throw e;
+ }
+ }
+
+ @Override
+ public OperatorSnapshotFutures snapshotState(
+ long checkpointId,
+ long timestamp,
+ CheckpointOptions checkpointOptions,
+ CheckpointStreamFactory storageLocation)
+ throws Exception {
+ if (!parallelExecutionWithoutCoroutineEnabled) {
+ return super.snapshotState(checkpointId, timestamp, checkpointOptions, storageLocation);
+ }
+ try {
+ return super.snapshotState(checkpointId, timestamp, checkpointOptions, storageLocation);
+ } finally {
+ // Resume dispatch only after the synchronous snapshot capture: a worker woken any
+ // earlier could pull a queued task out of keyed state before it is captured.
+ executionCoordinator.stopDraining();
+ }
+ }
+
@Override
public void snapshotState(StateSnapshotContext context) throws Exception {
+ if (parallelExecutionWithoutCoroutineEnabled) {
+ runWithParallelExecutionLock(() -> snapshotStateInternal(context));
+ } else {
+ snapshotStateInternal(context);
+ }
+ }
+
+ private void snapshotStateInternal(StateSnapshotContext context) throws Exception {
durableExecManager.snapshotRecoveryMarker();
durableExecManager.snapshotLastCompletedSequenceNumbers(
getKeyedStateBackend(), context.getCheckpointId());
@@ -766,12 +1087,28 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
@Override
public void notifyCheckpointComplete(long checkpointId) throws Exception {
+ if (parallelExecutionWithoutCoroutineEnabled) {
+ runWithParallelExecutionLock(() -> notifyCheckpointCompleteInternal(checkpointId));
+ } else {
+ notifyCheckpointCompleteInternal(checkpointId);
+ }
+ }
+
+ private void notifyCheckpointCompleteInternal(long checkpointId) throws Exception {
durableExecManager.notifyCheckpointComplete(checkpointId);
super.notifyCheckpointComplete(checkpointId);
}
@Override
public void notifyCheckpointAborted(long checkpointId) throws Exception {
+ if (parallelExecutionWithoutCoroutineEnabled) {
+ runWithParallelExecutionLock(() -> notifyCheckpointAbortedInternal(checkpointId));
+ } else {
+ notifyCheckpointAbortedInternal(checkpointId);
+ }
+ }
+
+ private void notifyCheckpointAbortedInternal(long checkpointId) throws Exception {
durableExecManager.notifyCheckpointAborted(checkpointId);
super.notifyCheckpointAborted(checkpointId);
}
@@ -1000,7 +1337,7 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
@Nullable
private ActionTask pollNextActionTask() throws Exception {
ActionTask actionTask = stateManager.pollNextActionTask();
- if (actionTask != null) {
+ if (actionTask != null && !isInternalNoopInputAction(actionTask.action)) {
builtInMetrics.markActionTaskDequeued(
actionTask.getTraceContext(), actionTask.hasExecutionStartedEventEmitted());
}
@@ -1025,15 +1362,30 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
continue;
}
eventRouter.getKeySegmentQueue().addKeyToLastSegment(key);
- String contextKey = resolveContextKey(key);
- // Align with the task-level replay: re-emit the record start for the resumed
- // round so listeners observe a paired start/finished bracket as well.
- notifyRecordStart(key);
- mailboxExecutor.submit(
- () -> tryProcessActionTaskForKey(key, contextKey, null),
- "process action task");
+ if (parallelExecutionWithoutCoroutineEnabled) {
+ // Restored tasks were never enqueued through processEvent: add one node per
+ // queued task; the worker-completion path finishes the input (no kick mail).
+ setCurrentKey(key);
+ int queued = stateManager.countActionTasks();
+ for (int i = 0; i < queued; i++) {
+ executionCoordinator.addTask(key);
+ }
+ } else {
+ // The normal engine has no workers; the kick mail drives
+ // processActionTaskForKey.
+ String contextKey = resolveContextKey(key);
+ // Align with the task-level replay: re-emit the record start for the resumed
+ // round so listeners observe a paired start/finished bracket as well.
+ notifyRecordStart(key);
+ mailboxExecutor.submit(
+ () -> tryProcessActionTaskForKey(key, contextKey, null),
+ "process action task");
+ }
}
stateManager.replaceProcessingKeys(new ArrayList<>(ownedKeys));
+ // Each recovered key has one active input record occupying in-flight budget; queued
+ // pending records for these keys are added below.
+ inFlightInputRecords = ownedKeys.size();
activeInputRuns = ownedKeys.size();
}
builtInMetrics.restoreActiveInputRuns(activeInputRuns);
@@ -1049,17 +1401,256 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
});
long[] pendingInputEvents = {0L};
+ // Recovered pending input records will each be dequeued, processed, and decremented when
+ // they finish, so they must occupy in-flight budget too — count one unit per pending event.
stateManager.forEachPendingInputEventKey(
getKeyedStateBackend(),
- (key, state) -> {
- for (Event ignored : state.get()) {
- eventRouter.getKeySegmentQueue().addKeyToLastSegment(key);
- pendingInputEvents[0]++;
- }
- });
+ (key, state) ->
+ state.get()
+ .forEach(
+ event -> {
+ inFlightInputRecords++;
+ pendingInputEvents[0]++;
+ eventRouter
+ .getKeySegmentQueue()
+ .addKeyToLastSegment(key);
+ }));
builtInMetrics.restorePendingInputEvents(pendingInputEvents[0]);
}
+ private void runWithParallelExecutionLock(ThrowingRunnable<? extends Exception> action)
+ throws Exception {
+ boolean acquired = false;
+ try {
+ parallelExecutionLock.acquireByMain();
+ acquired = true;
+ action.run();
+ } finally {
+ if (acquired) {
+ parallelExecutionLock.release();
+ }
+ }
+ }
+
+ /**
+ * A single action task's full worker lifecycle on the parallel engine: a pool thread {@code
+ * setup}s and {@code execute}s it under the lock; the mailbox thread later restores and commits
+ * it. All task-specific processing lives here; the operator keeps only mailbox flow.
+ */
+ private final class Work implements ParallelExecutionTask {
+ private Object key;
+ private long recordIndex;
+ private long taskIndex;
+ private String contextKey;
+
+ private ActionTask actionTask;
+ private long sequenceNumber;
+ private boolean replayCompletedAction;
+
+ @Nullable private ActionTask.ActionTaskResult result;
+ @Nullable private Throwable failure;
+
+ @Override
+ public void setup(Object key, long recordIndex, long taskIndex) {
+ this.key = key;
+ this.recordIndex = recordIndex;
+ this.taskIndex = taskIndex;
+ this.contextKey = resolveContextKey(key);
+ }
+
+ @Override
+ public void execute() {
+ try {
+ // Pull + prepare under the lock hold: pairing the FIFO poll with the under-lock
+ // taskIndex keeps the (taskIndex, task) binding correct across concurrent workers.
+ // createAndSetRunnerContext already binds the context on this thread.
+ setCurrentKey(key);
+ ActionTask task = pollNextActionTask();
+ checkState(task != null, "Action task queue was empty.");
+
+ contextManager.createAndSetRunnerContext(
+ task,
+ contextKey,
+ agentPlan,
+ resourceCache,
+ metricGroup,
+ jobIdentifier,
+ parallelExecutionLock::checkReentrant,
+ stateManager.getSensoryMemState(),
+ stateManager.getShortTermMemState(),
+ pythonBridge.getPythonRunnerContext(),
+ ltm,
+ ActionExecutionOperator.this::createComponentListeners);
+
+ // The synthetic noop input action stays lifecycle-silent.
+ boolean lifecycleVisible = !isInternalNoopInputAction(task.action);
+ if (lifecycleVisible) {
+ notifyActionPrepared(task);
+ }
+
+ long seq = stateManager.getSequenceNumber();
+ ActionState state =
+ durableExecManager.maybeGetActionState(key, seq, task.action, task.event);
+ boolean replay = state != null && state.isCompleted();
+ if (!replay) {
+ if (state == null) {
+ durableExecManager.maybeInitActionState(key, seq, task.action, task.event);
+ state =
+ durableExecManager.maybeGetActionState(
+ key, seq, task.action, task.event);
+ }
+ durableExecManager.setupDurableExecutionContext(task, state, seq);
+ }
+
+ // Recording the key and action task on the continuation context uses its
+ // JDK<21-only
+ // internals, so it lives on the operator's JDK<21 parallel MVP path and is never
+ // reached on JDK 21.
+ ContinuationContext context =
+ checkNotNull(
+ contextManager.getContinuationContext(task),
+ "Missing continuation context for action task");
+ context.setKey(key);
+ context.setActionTask(task);
+ context.setPriority(recordIndex, taskIndex);
+
+ this.actionTask = task;
+ this.sequenceNumber = seq;
+ this.replayCompletedAction = replay;
+
+ // Run: replay a durably-completed action from its persisted state, or invoke it.
+ if (replay) {
+ MemoryUpdateReplayer.replay(
+ task.getRunnerContext().getShortTermMemory(),
+ state.getShortTermMemoryUpdates());
+ MemoryUpdateReplayer.replay(
+ task.getRunnerContext().getSensoryMemory(),
+ state.getSensoryMemoryUpdates());
+ if (lifecycleVisible) {
+ notifyActionReused(task);
+ }
+ result = task.new ActionTaskResult(true, state.getOutputEvents(), null);
+ } else {
+ if (lifecycleVisible) {
+ notifyActionStarted(task);
+ }
+ try {
+ if (task instanceof JavaActionTask) {
+ ((JavaActionTask) task).setEventSerializer(eventSerializer);
+ }
+ result =
+ task.invoke(
+ getRuntimeContext().getUserCodeClassLoader(),
+ pythonBridge.getPythonActionExecutor());
+ } catch (Throwable actionFailure) {
+ try {
+ task.getRunnerContext().discardMemoryObservation();
+ } catch (Throwable discardFailure) {
+ if (discardFailure != actionFailure) {
+ actionFailure.addSuppressed(discardFailure);
+ }
+ }
+ if (lifecycleVisible) {
+ try {
+ notifyActionFailed(task, actionFailure);
+ } catch (Throwable notifyFailure) {
+ if (notifyFailure != actionFailure) {
+ actionFailure.addSuppressed(notifyFailure);
+ }
+ }
+ }
+ throw actionFailure;
+ }
+ }
+ } catch (Throwable t) {
+ failure = t;
+ if (t instanceof InterruptedException) {
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ @Override
+ public void restoreContext() {
+ parallelExecutionLock.checkReentrant();
+ if (failure != null) {
+ // Preparation failed before a context existed; commit() will rethrow the failure.
+ return;
+ }
+ contextManager.restore(key, actionTask);
+ }
+
+ @Override
+ public boolean isDone() {
+ return result != null || failure != null;
+ }
+
+ @Override
+ public void commit() throws Exception {
+ if (failure != null) {
+ throw new ActionTaskExecutionException("Failed to execute action task", failure);
+ }
+ checkState(result != null, "Action task completion result must not be null.");
+ // On this engine an action always finishes in one invoke (the base executor runs
+ // durableExecuteAsync inline); an unfinished result would mean the Loom overlay or a
+ // non-Java task leaked in — fail fast rather than run an untested branch.
+ checkState(
+ result.isFinished(),
+ "Parallel-engine action did not finish in one invoke: %s",
+ actionTask.action.getName());
+
+ // The synthetic noop input action stays lifecycle-silent.
+ boolean lifecycleVisible = !isInternalNoopInputAction(actionTask.action);
+ try {
+ contextManager.removeContexts(actionTask);
+ durableExecManager.removeDurableContext(actionTask);
+
+ if (!replayCompletedAction) {
+ if (lifecycleVisible) {
+ notifyActionFinishing(actionTask);
+ }
+ durableExecManager.maybePersistTaskResult(
+ key,
+ sequenceNumber,
+ actionTask.action,
+ actionTask.event,
+ actionTask.getRunnerContext(),
+ result);
+ if (lifecycleVisible) {
+ notifyActionFinished(actionTask);
+ }
+ }
+ } catch (Throwable t) {
+ if (lifecycleVisible) {
+ notifyActionFailed(actionTask, t);
+ }
+ throw new ActionTaskExecutionException("Failed to execute action task", t);
+ }
+
+ for (Event actionOutputEvent : result.getOutputEvents()) {
+ try {
+ processEvent(key, contextKey, actionOutputEvent, actionTask.getTraceContext());
+ } catch (Throwable t) {
+ // Mirror the serial kick-mail wrapper (tryProcessActionTaskForKey): surface
+ // output-event processing failures as ActionTaskExecutionException so callers
+ // observe a uniform cause. The action itself already finished, so it is not
+ // reported failed.
+ throw new ActionTaskExecutionException("Failed to execute action task", t);
+ }
+ }
+
+ if (!isInternalNoopInputAction(actionTask.action)) {
+ builtInMetrics.markActionExecuted(actionTask.action.getName());
+ }
+ actionTask.getRunnerContext().persistMemory();
+ }
+
+ @Override
+ public void finishGroup(@Nullable ParallelExecutionTask lastCommitted) throws Exception {
+ maybeFinishCurrentInput(key, (Work) lastCommitted);
+ }
+ }
+
@VisibleForTesting
DurableExecutionManager getDurableExecutionManager() {
return durableExecManager;
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java
index 30211b11..bafa9314 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java
@@ -32,6 +32,7 @@ import org.apache.flink.agents.runtime.memory.CachedMemoryStore;
import org.apache.flink.agents.runtime.memory.InteranlBaseLongTermMemory;
import org.apache.flink.agents.runtime.memory.MemoryObjectImpl;
import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl;
+import org.apache.flink.agents.runtime.operator.parallel.ParallelExecutionLock;
import org.apache.flink.agents.runtime.python.context.PythonRunnerContextImpl;
import org.apache.flink.api.common.state.MapState;
import org.apache.flink.util.ExceptionUtils;
@@ -63,27 +64,44 @@ import java.util.function.Function;
* <p>The manager is fully constructed in the operator's {@code open()} with the configured
* async-thread count from the agent plan, so it has no separate open step.
*
- * <p>No manager-to-manager references are held here, so cross-cutting data flows in as method
- * parameters. The Python {@link RunnerContextImpl} stays owned by {@link PythonBridgeManager} and
- * the durable-execution context stays on {@link DurableExecutionManager}, and both are passed in
+ * <p>The Python {@link RunnerContextImpl} stays owned by {@link PythonBridgeManager} is passed in
* when a method needs them.
+ *
+ * <p>This manager holds a back-reference to the owning {@link ActionExecutionOperator} (so it can
+ * restore the current Flink key while restoring a runner context) and to the {@link
+ * DurableExecutionManager} (so {@link #restore} can resolve the durable-execution context without
+ * threading it through as a parameter).
*/
class ActionTaskContextManager implements AutoCloseable {
+ private final ActionExecutionOperator<?, ?> operator;
+ private final DurableExecutionManager durableExecManager;
+
private RunnerContextImpl runnerContext;
private final Map<ActionTask, ActionTaskContexts> actionTaskContexts;
private ContinuationActionExecutor continuationActionExecutor;
- ActionTaskContextManager(int numAsyncThreads) {
- this(numAsyncThreads, () -> {});
+ ActionTaskContextManager(
+ ActionExecutionOperator<?, ?> operator,
+ DurableExecutionManager durableExecManager,
+ int numAsyncThreads) {
+ this(operator, durableExecManager, numAsyncThreads, () -> {}, null);
}
- ActionTaskContextManager(int numAsyncThreads, Runnable asyncThreadCleanup) {
+ ActionTaskContextManager(
+ ActionExecutionOperator<?, ?> operator,
+ DurableExecutionManager durableExecManager,
+ int numAsyncThreads,
+ Runnable asyncThreadCleanup,
+ @Nullable ParallelExecutionLock parallelExecutionLock) {
+ this.operator = operator;
+ this.durableExecManager = durableExecManager;
this.actionTaskContexts = new HashMap<>();
this.continuationActionExecutor =
- new ContinuationActionExecutor(numAsyncThreads, asyncThreadCleanup);
+ new ContinuationActionExecutor(
+ numAsyncThreads, asyncThreadCleanup, parallelExecutionLock, this::restore);
}
/**
@@ -96,6 +114,7 @@ class ActionTaskContextManager implements AutoCloseable {
@Nullable private String pythonAwaitableRef;
private List<Event> pendingEvents = new ArrayList<>();
@Nullable private List<ComponentExecutionListener> componentListeners;
+ @Nullable private String contextKey;
}
private boolean hasContexts(ActionTask actionTask) {
@@ -277,6 +296,7 @@ class ActionTaskContextManager implements AutoCloseable {
putMemoryContext(actionTask, memoryContext);
}
+ requireContexts(actionTask).contextKey = contextKey;
context.switchActionContext(
actionTask.action.getName(),
memoryContext,
@@ -307,6 +327,46 @@ class ActionTaskContextManager implements AutoCloseable {
actionTask.setRunnerContext(context);
}
+ /**
+ * Restores the shared runner context to the already-pinned contexts of a task, first
+ * re-establishing its Flink key. Also invoked (as the async executor's restorer) after a worker
+ * re-acquires the lock. Deliberately does not create missing contexts: a missing mapping means
+ * the task was not prepared correctly or its transient execution context was lost.
+ */
+ public void restore(Object key, ActionTask actionTask) {
+ operator.setCurrentKey(key);
+ RunnerContextImpl.MemoryContext memoryContext =
+ Preconditions.checkNotNull(
+ getMemoryContext(actionTask), "Missing memory context for action task");
+ ContinuationContext continuationContext = getContinuationContext(actionTask);
+ RunnerContextImpl context =
+ Preconditions.checkNotNull(
+ actionTask.getRunnerContext(), "Missing runner context for action task");
+
+ context.switchActionContext(
+ actionTask.action.getName(),
+ memoryContext,
+ requireContexts(actionTask).pendingEvents,
+ requireContexts(actionTask).contextKey,
+ actionTask.getObservationId(),
+ MemoryEvent.isMemoryType(actionTask.event.getType()),
+ requireContexts(actionTask).componentListeners);
+ if (context instanceof JavaRunnerContextImpl) {
+ Preconditions.checkNotNull(
+ continuationContext, "Missing continuation context for Java action task");
+ ((JavaRunnerContextImpl) context).setContinuationContext(continuationContext);
+ }
+
+ RunnerContextImpl.DurableExecutionContext durableContext =
+ durableExecManager.getDurableContext(actionTask);
+ if (durableContext == null) {
+ context.clearDurableExecutionContext();
+ } else {
+ context.setDurableExecutionContext(durableContext);
+ }
+ actionTask.setRunnerContext(context);
+ }
+
private void putMemoryContext(
ActionTask actionTask, RunnerContextImpl.MemoryContext memoryContext) {
requireContexts(actionTask).memoryContext = memoryContext;
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/DurableExecutionManager.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/DurableExecutionManager.java
index 49f10995..573fc669 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/DurableExecutionManager.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/DurableExecutionManager.java
@@ -325,6 +325,7 @@ class DurableExecutionManager implements ActionStatePersister, AutoCloseable {
actionTask.event,
actionState,
this);
+ putDurableContext(actionTask, durableContext);
}
actionTask.getRunnerContext().setDurableExecutionContext(durableContext);
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/OperatorStateManager.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/OperatorStateManager.java
index ea0c23f2..e043f3e9 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/OperatorStateManager.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/OperatorStateManager.java
@@ -233,6 +233,18 @@ class OperatorStateManager {
return listStateNotEmpty(actionTasksKState);
}
+ /**
+ * Counts the pending {@link ActionTask}s buffered for the current key without removing any.
+ * Used on recovery to grant one pull credit per already-queued task.
+ */
+ int countActionTasks() throws Exception {
+ int count = 0;
+ for (ActionTask ignored : actionTasksKState.get()) {
+ count++;
+ }
+ return count;
+ }
+
/**
* Removes and returns the next pending {@link ActionTask} for the current key.
*
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
index e8b4d646..f003e12b 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
@@ -186,7 +186,8 @@ class PythonBridgeManager implements AutoCloseable {
new PythonInterpreterManager(
ownerInterpreter,
env::getInterpreter,
- agentPlan.getConfig().get(AgentExecutionOptions.NUM_ASYNC_THREADS));
+ agentPlan.getConfig().get(AgentExecutionOptions.NUM_ASYNC_THREADS),
+ mailboxThreadChecker);
pythonRunnerContext =
new PythonRunnerContextImpl(
metricGroup,
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationContext.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionContextRestorer.java
similarity index 60%
copy from runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationContext.java
copy to runtime/src/main/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionContextRestorer.java
index 6d0d7e09..c84e559c 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationContext.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionContextRestorer.java
@@ -15,12 +15,17 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.flink.agents.runtime.async;
+package org.apache.flink.agents.runtime.operator.parallel;
+
+import org.apache.flink.agents.runtime.operator.ActionTask;
/**
- * Marker context for continuation execution.
- *
- * <p>This class is part of the base (JDK 11) sources. JDK 21 provides a multi-release variant with
- * continuation-specific fields.
+ * Restores the shared runner context to a resuming action task. Called by the async executor after
+ * a worker re-acquires the lock, so the shared context is repointed at the task's key, memory,
+ * continuation, and durable-execution contexts before control returns to user code.
*/
-public class ContinuationContext {}
+@FunctionalInterface
+public interface ParallelExecutionContextRestorer {
+
+ void restore(Object key, ActionTask actionTask);
+}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionCoordinator.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionCoordinator.java
new file mode 100644
index 00000000..69430f14
--- /dev/null
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionCoordinator.java
@@ -0,0 +1,312 @@
+/*
+ * 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.operator.parallel;
+
+import org.apache.flink.agents.runtime.async.AsyncExecutorThreadFactory;
+import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.util.function.ThrowingRunnable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Supplier;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+
+/**
+ * Coordinates worker execution of {@link ParallelExecutionTask}: callers {@link #addTask(Object)
+ * add} one node per queued task, and the coordinator submits one anonymous <em>dispatch permit</em>
+ * per node to an elastic {@link ThreadPoolExecutor}. A pool thread redeems a permit under the lock
+ * (node poll and FIFO task pull share the lock hold, keeping their binding correct — permits never
+ * name a task); commits always run on the mailbox thread in per-key taskIndex order.
+ */
+public final class ParallelExecutionCoordinator implements AutoCloseable {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ParallelExecutionCoordinator.class);
+
+ private static final long WORKER_TERMINATION_TIMEOUT_SECONDS = 180L;
+
+ @FunctionalInterface
+ public interface MailboxDispatcher {
+ void execute(ThrowingRunnable<? extends Exception> command, String description);
+ }
+
+ private final ParallelExecutionTaskQueue schedule;
+ private final ParallelExecutionLock parallelExecutionLock;
+ private final AsyncExecutorThreadFactory workerThreadFactory;
+ private final ThreadPoolExecutor workerExecutor;
+ private final MailboxDispatcher mailboxDispatcher;
+
+ /** Supplies a fresh work; the work pulls and prepares its own task inside {@code execute()}. */
+ private final Supplier<ParallelExecutionTask> workFactory;
+
+ /** Tasks pulled but not yet committed; read by {@link #isQuiesced()} for checkpoint quiesce. */
+ private final AtomicInteger inFlightExecuting = new AtomicInteger();
+
+ /**
+ * When {@code true}, workers finish in-flight tasks but pull no new nodes. Set before a
+ * checkpoint, cleared once the snapshot is taken.
+ */
+ private volatile boolean draining;
+
+ private final Object drainMonitor = new Object();
+
+ /**
+ * Permits submitted and not yet fully redeemed; compared against the pool size by {@link
+ * DispatchQueue#offer} to decide when the pool must grow.
+ */
+ private final AtomicInteger submittedPermits = new AtomicInteger();
+
+ /**
+ * Creates a coordinator with an elastic worker pool: one resident thread, on-demand growth up
+ * to {@code maxWorkers}, and idle retirement after {@code workerIdleTimeoutMillis}. Every
+ * worker participates in the managed async lifecycle and runs {@code threadCleanup} when it
+ * retires or the pool shuts down. Admission is not gated here; the operator bounds outstanding
+ * work via input-record backpressure.
+ */
+ public ParallelExecutionCoordinator(
+ ParallelExecutionLock parallelExecutionLock,
+ MailboxDispatcher mailboxDispatcher,
+ Supplier<ParallelExecutionTask> workFactory,
+ Runnable threadCleanup,
+ int maxWorkers,
+ long workerIdleTimeoutMillis) {
+ this.parallelExecutionLock = checkNotNull(parallelExecutionLock);
+ this.schedule = new ParallelExecutionTaskQueue(parallelExecutionLock);
+ this.mailboxDispatcher = checkNotNull(mailboxDispatcher);
+ this.workFactory = checkNotNull(workFactory);
+ this.workerThreadFactory = new AsyncExecutorThreadFactory(checkNotNull(threadCleanup));
+ DispatchQueue dispatchQueue = new DispatchQueue();
+ this.workerExecutor =
+ new ThreadPoolExecutor(
+ 1,
+ maxWorkers,
+ workerIdleTimeoutMillis,
+ TimeUnit.MILLISECONDS,
+ dispatchQueue,
+ workerThreadFactory,
+ (permit, pool) -> {
+ // Only reachable when offer() returned false to grow the pool but the
+ // addWorker race was lost (already at max), or during shutdown.
+ if (pool.isShutdown()) {
+ submittedPermits.decrementAndGet();
+ } else {
+ dispatchQueue.forceOffer(permit);
+ }
+ });
+ }
+
+ /**
+ * Adds one task node for {@code key} and submits one anonymous dispatch permit (a permit means
+ * "pull one node under the lock", never a specific task, so pool threads may redeem permits in
+ * any order). Must be called on the mailbox thread holding the lock.
+ */
+ public void addTask(Object key) {
+ schedule.addTask(key);
+ submittedPermits.incrementAndGet();
+ workerExecutor.execute(this::runOneTask);
+ }
+
+ /** Whether {@code key} still has outstanding (added but not yet committed) work. */
+ public boolean hasOutstanding(Object key) {
+ return schedule.hasOutstanding(key);
+ }
+
+ /**
+ * Pauses new task dispatch before a checkpoint: workers finish and commit in-flight tasks but
+ * pull no new nodes. Returns only after fencing in-progress pulls: a worker increments {@link
+ * #inFlightExecuting} before it first releases the lock, so after the acquire/release below
+ * {@link #isQuiesced()} can no longer miss a pulled task.
+ */
+ public void startDraining() throws InterruptedException {
+ draining = true;
+ parallelExecutionLock.acquireByMain();
+ parallelExecutionLock.release();
+ }
+
+ /**
+ * Resumes task dispatch after a checkpoint: wakes parked workers, then consumes any dangling
+ * sticky main-grant left by a declining worker (one acquire/release) so woken workers can take
+ * the lock immediately instead of stalling until the next mailbox action.
+ */
+ public void stopDraining() throws InterruptedException {
+ synchronized (drainMonitor) {
+ draining = false;
+ drainMonitor.notifyAll();
+ }
+ parallelExecutionLock.acquireByMain();
+ parallelExecutionLock.release();
+ }
+
+ /** Whether no task is currently pulled-but-uncommitted, i.e. a safe point for a snapshot. */
+ public boolean isQuiesced() {
+ return inFlightExecuting.get() == 0;
+ }
+
+ /**
+ * Redeems one dispatch permit: under the lock, takes the highest-priority pending node, runs a
+ * fresh work, then hands the completion to the mailbox commit drain. During a checkpoint drain
+ * the permit is not forfeited: the thread parks and redeems it afterwards.
+ */
+ private void runOneTask() {
+ try {
+ // None of the interruptible calls below holds the lock when it throws, so one handler
+ // treats every interrupt as the shutdown signal.
+ try {
+ // Fresh pulls acquire at the lowest priority: async resumes always take precedence.
+ parallelExecutionLock.acquireByWorker(Long.MAX_VALUE, Long.MAX_VALUE);
+ while (draining) {
+ // Park without the lock until the drain ends; the node stays pending
+ // meanwhile. No main-thread reservation is left behind, so in-flight tasks
+ // can re-acquire the lock and commit while this worker waits.
+ parallelExecutionLock.release();
+ synchronized (drainMonitor) {
+ while (draining) {
+ drainMonitor.wait();
+ }
+ }
+ parallelExecutionLock.acquireByWorker(Long.MAX_VALUE, Long.MAX_VALUE);
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return;
+ }
+
+ // Holding the lock from here on: every failure is funneled to the mailbox (an uncaught
+ // throwable would only kill this pool thread silently) and the lock is always returned.
+ try {
+ ParallelExecutionTaskQueue.Node node = schedule.pollPendingNode();
+ ParallelExecutionTask work = workFactory.get();
+ work.setup(node.getKey(), node.getRecordIndex(), node.getTaskIndex());
+ node.setWork(work);
+ inFlightExecuting.incrementAndGet();
+ // execute() never throws: it captures pull/prepare/run failures and rethrows at
+ // commit.
+ work.execute();
+ schedule.markDone(node);
+ mailboxDispatcher.execute(() -> completeOnMailbox(node), "complete action task");
+ } catch (Throwable t) {
+ mailboxDispatcher.execute(
+ () -> {
+ throw t;
+ },
+ "throw action task scheduling failure");
+ } finally {
+ parallelExecutionLock.releaseToMain();
+ }
+ } finally {
+ submittedPermits.decrementAndGet();
+ }
+ }
+
+ /**
+ * Mailbox-thread commit drain: commits the completed node's key group in taskIndex order,
+ * stopping at the first not-yet-done head (picked up by a later drain), then runs the group's
+ * follow-up.
+ */
+ private void completeOnMailbox(ParallelExecutionTaskQueue.Node completedNode) throws Exception {
+ parallelExecutionLock.acquireByMain();
+ try {
+ Object key = completedNode.getKey();
+ ParallelExecutionTask lastCommitted = null;
+ while (true) {
+ ParallelExecutionTaskQueue.Node head = schedule.pollCommittableHead(key);
+ if (head == null) {
+ break;
+ }
+ try {
+ head.getWork().restoreContext();
+ head.getWork().commit();
+ lastCommitted = head.getWork();
+ } finally {
+ schedule.removeCommitted(head);
+ inFlightExecuting.decrementAndGet();
+ }
+ }
+ completedNode.getWork().finishGroup(lastCommitted);
+ } finally {
+ parallelExecutionLock.release();
+ }
+ }
+
+ @Override
+ public void close() {
+ // Wait for in-flight workers to terminate before the caller releases resources they may
+ // still be using. Interrupts are re-issued each round: a worker may block on the lock
+ // (e.g. a sticky main-grant nobody consumes while the mailbox waits here) only AFTER the
+ // first interrupt was delivered, so a single shutdownNow is not enough. Waiting for worker
+ // execution is bounded because user code may ignore the interrupt; once the executor has
+ // terminated, wait for every worker's exit cleanup before returning.
+ long deadlineNanos =
+ System.nanoTime() + TimeUnit.SECONDS.toNanos(WORKER_TERMINATION_TIMEOUT_SECONDS);
+ try {
+ do {
+ workerExecutor.shutdownNow();
+ } while (!workerExecutor.awaitTermination(1, TimeUnit.SECONDS)
+ && System.nanoTime() < deadlineNanos);
+ if (workerExecutor.isTerminated()) {
+ // ThreadPoolExecutor may signal termination before the ThreadFactory wrapper has
+ // run its exit cleanup. Join the already-terminating workers before resources used
+ // by that cleanup are closed by the operator.
+ workerThreadFactory.awaitThreadExit();
+ } else {
+ LOG.warn(
+ "Worker pool did not terminate within {}s; proceeding with close.",
+ WORKER_TERMINATION_TIMEOUT_SECONDS);
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ }
+
+ @VisibleForTesting
+ int getWorkerPoolSizeForTesting() {
+ return workerExecutor.getPoolSize();
+ }
+
+ /**
+ * Work queue of anonymous dispatch permits. {@link #offer} reports "full" when every pool
+ * thread is busy and the pool may still grow — the only way to make {@link ThreadPoolExecutor}
+ * add threads beyond its core size; the refused permit is handed to the new thread or re-queued
+ * by the rejection handler. Idle threads simply time out polling this queue and exit.
+ */
+ private final class DispatchQueue extends LinkedBlockingQueue<Runnable> {
+ private static final long serialVersionUID = 1L;
+
+ @Override
+ public boolean offer(Runnable permit) {
+ ThreadPoolExecutor pool = workerExecutor;
+ if (pool != null) {
+ int poolSize = pool.getPoolSize();
+ if (poolSize < pool.getMaximumPoolSize() && submittedPermits.get() > poolSize) {
+ // Every thread is busy and the pool is below its cap: refuse, forcing the
+ // pool to spawn a thread for this permit.
+ return false;
+ }
+ }
+ return super.offer(permit);
+ }
+
+ private boolean forceOffer(Runnable permit) {
+ return super.offer(permit);
+ }
+ }
+}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionLock.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionLock.java
new file mode 100644
index 00000000..f5009586
--- /dev/null
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionLock.java
@@ -0,0 +1,282 @@
+/*
+ * 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.operator.parallel;
+
+import org.apache.flink.annotation.VisibleForTesting;
+
+import java.util.Comparator;
+import java.util.PriorityQueue;
+import java.util.concurrent.locks.Condition;
+import java.util.concurrent.locks.ReentrantLock;
+
+/**
+ * Non-reentrant lock for the JDK<21 parallel execution engine. The main (mailbox) thread has
+ * absolute priority; workers are granted in ascending (recordIndex, taskIndex) order, ties
+ * unordered. {@link #releaseToMain()} reserves the lock for the main thread via a sticky {@code
+ * mailboxGranted} flag even before it starts waiting.
+ */
+public final class ParallelExecutionLock {
+
+ /** Worker grant order: ascending recordIndex, then taskIndex; ties unordered. */
+ private static final Comparator<WorkerNode> WORKER_ORDER =
+ Comparator.comparingLong((WorkerNode node) -> node.recordIndex)
+ .thenComparingLong(node -> node.taskIndex);
+
+ private final ReentrantLock innerLock = new ReentrantLock();
+ private final Condition mailboxCondition = innerLock.newCondition();
+ private final PriorityQueue<WorkerNode> workerQueue = new PriorityQueue<>(WORKER_ORDER);
+
+ /**
+ * Current owner, or {@code null}. Releasers assign ownership directly to the next queued
+ * worker, so the lock is held iff {@code ownerThread != null || mailboxGranted}.
+ */
+ private volatile Thread ownerThread;
+
+ /**
+ * Sticky reservation for the main thread (whose identity may not be known yet); workers block
+ * on it. Read/written under {@link #innerLock} only.
+ */
+ private boolean mailboxGranted;
+
+ /** Whether the main thread is blocked in {@link #acquireByMain()}. Guarded by innerLock. */
+ private boolean mailboxWaiting;
+
+ // -----------------------------------------------------------------------
+ // Lock operations
+ // -----------------------------------------------------------------------
+
+ /** Acquires the lock for the mailbox thread. Mailbox acquisition has the highest priority. */
+ public void acquireByMain() throws InterruptedException {
+ Thread currentThread = Thread.currentThread();
+ innerLock.lockInterruptibly();
+ try {
+ checkNonReentrant();
+ // Fast path 1: lock is free
+ if (ownerThread == null && !mailboxGranted) {
+ ownerThread = currentThread;
+ return;
+ }
+ // Fast path 2: lock is sticky-granted to mailbox
+ if (mailboxGranted) {
+ mailboxGranted = false;
+ ownerThread = currentThread;
+ return;
+ }
+ if (mailboxWaiting) {
+ throw new IllegalStateException(
+ "Only one mailbox thread can wait for the lock at a time.");
+ }
+ // Slow path: wait until granted
+ mailboxWaiting = true;
+ try {
+ while (!mailboxGranted) {
+ mailboxCondition.await();
+ }
+ mailboxGranted = false;
+ ownerThread = currentThread;
+ } catch (InterruptedException e) {
+ // If we were already granted while interrupted, transfer the lock rather than
+ // keeping it.
+ if (mailboxGranted) {
+ mailboxGranted = false;
+ mailboxWaiting = false;
+ transferToNext();
+ }
+ throw e;
+ } finally {
+ mailboxWaiting = false;
+ }
+ } finally {
+ innerLock.unlock();
+ }
+ }
+
+ /**
+ * Acquires the lock for a worker at the given (recordIndex, taskIndex) priority. The main
+ * thread always has absolute priority over workers.
+ */
+ public void acquireByWorker(long recordIndex, long taskIndex) throws InterruptedException {
+ Thread currentThread = Thread.currentThread();
+ innerLock.lockInterruptibly();
+ try {
+ checkNonReentrant();
+ // Fast path: free, no mailbox reservation or waiter, and no queued workers
+ if (ownerThread == null
+ && !mailboxGranted
+ && !mailboxWaiting
+ && workerQueue.isEmpty()) {
+ ownerThread = currentThread;
+ return;
+ }
+ // Slow path: join the priority queue; granted in (recordIndex, taskIndex) order
+ WorkerNode node =
+ new WorkerNode(innerLock.newCondition(), currentThread, recordIndex, taskIndex);
+ workerQueue.add(node);
+ try {
+ while (!node.granted) {
+ node.condition.await();
+ }
+ // The releaser assigned ownership to this thread directly; nothing to claim.
+ } catch (InterruptedException e) {
+ if (node.granted) {
+ // Already made the owner but bailing out: pass the lock on.
+ ownerThread = null;
+ transferToNext();
+ } else {
+ workerQueue.remove(node);
+ }
+ throw e;
+ }
+ } finally {
+ innerLock.unlock();
+ }
+ }
+
+ /**
+ * Releases the lock and transfers ownership mailbox-first, then to the highest-priority worker.
+ */
+ public void release() {
+ innerLock.lock();
+ try {
+ checkReentrant();
+ ownerThread = null;
+ transferToNext();
+ } finally {
+ innerLock.unlock();
+ }
+ }
+
+ /** Releases worker ownership and reserves the lock for the mailbox thread. */
+ public void releaseToMain() {
+ innerLock.lock();
+ try {
+ checkReentrant();
+ ownerThread = null;
+ mailboxGranted = true;
+ if (mailboxWaiting) {
+ mailboxCondition.signal();
+ }
+ } finally {
+ innerLock.unlock();
+ }
+ }
+
+ /** Asserts that the calling thread currently owns the lock. */
+ public void checkReentrant() {
+ if (ownerThread != Thread.currentThread()) {
+ throw new IllegalStateException(
+ "Current thread does not own the lock. owner="
+ + ownerThread
+ + ", current="
+ + Thread.currentThread());
+ }
+ }
+
+ // -----------------------------------------------------------------------
+ // Internal helpers
+ // -----------------------------------------------------------------------
+
+ /**
+ * Transfers ownership main-first, then to the highest-priority queued worker, else frees the
+ * lock. Must be called under {@link #innerLock}.
+ */
+ private void transferToNext() {
+ if (mailboxWaiting) {
+ mailboxGranted = true;
+ mailboxCondition.signal();
+ return;
+ }
+ WorkerNode nextWorker = workerQueue.poll();
+ if (nextWorker != null) {
+ ownerThread = nextWorker.thread;
+ nextWorker.granted = true;
+ nextWorker.condition.signal();
+ }
+ // Otherwise the lock is free: ownerThread is null and no reservation is set.
+ }
+
+ private void checkNonReentrant() {
+ if (ownerThread == Thread.currentThread()) {
+ throw new IllegalStateException(
+ "ParallelExecutionLock is non-reentrant; the owner cannot acquire it again.");
+ }
+ }
+
+ // -----------------------------------------------------------------------
+ // Test-only accessors
+ // -----------------------------------------------------------------------
+
+ @VisibleForTesting
+ boolean isLockedForTesting() {
+ innerLock.lock();
+ try {
+ return ownerThread != null || mailboxGranted;
+ } finally {
+ innerLock.unlock();
+ }
+ }
+
+ @VisibleForTesting
+ boolean isMailboxGrantedForTesting() {
+ innerLock.lock();
+ try {
+ return mailboxGranted;
+ } finally {
+ innerLock.unlock();
+ }
+ }
+
+ @VisibleForTesting
+ boolean isMailboxWaitingForTesting() {
+ innerLock.lock();
+ try {
+ return mailboxWaiting;
+ } finally {
+ innerLock.unlock();
+ }
+ }
+
+ @VisibleForTesting
+ int getQueuedWorkerCountForTesting() {
+ innerLock.lock();
+ try {
+ return workerQueue.size();
+ } finally {
+ innerLock.unlock();
+ }
+ }
+
+ // -----------------------------------------------------------------------
+ // Inner types
+ // -----------------------------------------------------------------------
+
+ private static final class WorkerNode {
+ private final Condition condition;
+ private final Thread thread;
+ private final long recordIndex;
+ private final long taskIndex;
+ private boolean granted;
+
+ private WorkerNode(Condition condition, Thread thread, long recordIndex, long taskIndex) {
+ this.condition = condition;
+ this.thread = thread;
+ this.recordIndex = recordIndex;
+ this.taskIndex = taskIndex;
+ }
+ }
+}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionTask.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionTask.java
new file mode 100644
index 00000000..55913dbe
--- /dev/null
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionTask.java
@@ -0,0 +1,58 @@
+/*
+ * 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.operator.parallel;
+
+import javax.annotation.Nullable;
+
+/**
+ * A self-describing unit of worker execution for the parallel execution engine. A pool thread
+ * {@link #setup} + {@link #execute()}s it (pull, prepare, run; result or failure stored inside);
+ * the mailbox thread later {@link #restoreContext()} + {@link #commit()}s each done key-group head
+ * in order and calls {@link #finishGroup} once per drain.
+ */
+public interface ParallelExecutionTask {
+
+ /**
+ * Binds this work to a key and its coordinator-assigned (recordIndex, taskIndex) priority.
+ * Called once on the worker thread under the lock, before {@link #execute()}.
+ */
+ void setup(Object key, long recordIndex, long taskIndex);
+
+ /** Re-establishes the runner context and Flink key for this work (used by the commit drain). */
+ void restoreContext();
+
+ /**
+ * Runs the action on a worker thread: pulls and prepares the next queued task, invokes the
+ * action, and stores the result or failure inside this work. Never throws; failures are
+ * rethrown by {@link #commit()}.
+ */
+ void execute();
+
+ /** Whether {@link #execute()} has finished (result or failure recorded). */
+ boolean isDone();
+
+ /** Commits this work's result on the mailbox thread; rethrows a recorded execution failure. */
+ void commit() throws Exception;
+
+ /**
+ * Runs the per-key-group follow-up after the commit drain; called once on the completing work.
+ *
+ * @param lastCommitted the last work committed in this drain, or {@code null} if none was.
+ */
+ void finishGroup(@Nullable ParallelExecutionTask lastCommitted) throws Exception;
+}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionTaskQueue.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionTaskQueue.java
new file mode 100644
index 00000000..64ac0e7c
--- /dev/null
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionTaskQueue.java
@@ -0,0 +1,208 @@
+/*
+ * 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.operator.parallel;
+
+import org.apache.flink.annotation.VisibleForTesting;
+
+import javax.annotation.Nullable;
+
+import java.util.ArrayDeque;
+import java.util.Comparator;
+import java.util.Deque;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.PriorityQueue;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * Key-grouped ledger of action task nodes for the parallel execution engine. A node is created on
+ * {@link #addTask(Object)} with its (recordIndex, taskIndex) priority and lives in its key group
+ * until committed; pull takes the highest-priority pending node, commit drains each group's DONE
+ * head in taskIndex order, and an emptied group is dropped ({@code hasOutstanding == false}). Holds
+ * no internal lock: mutators assert the caller owns the {@link ParallelExecutionLock}.
+ */
+final class ParallelExecutionTaskQueue {
+
+ /** Node lifecycle: RUNNING (born, to be pulled) -> DONE (executed) -> COMMITTED (removed). */
+ enum State {
+ RUNNING,
+ DONE,
+ COMMITTED
+ }
+
+ /**
+ * Per-key group: arrival-ordered {@code recordIndex}, next per-key {@code taskIndex}, and its
+ * nodes.
+ */
+ private static final class KeyGroup {
+ private final Object key;
+ private final long recordIndex;
+ private long taskIndexGen;
+ private final Deque<Node> nodes = new ArrayDeque<>();
+
+ private KeyGroup(Object key, long recordIndex) {
+ this.key = key;
+ this.recordIndex = recordIndex;
+ }
+ }
+
+ /** A single action task, ordered within its group by {@code taskIndex}. */
+ static final class Node {
+ private final KeyGroup group;
+ private final long taskIndex;
+ private State state = State.RUNNING;
+ private ParallelExecutionTask work;
+
+ private Node(KeyGroup group, long taskIndex) {
+ this.group = group;
+ this.taskIndex = taskIndex;
+ }
+
+ Object getKey() {
+ return group.key;
+ }
+
+ long getRecordIndex() {
+ return group.recordIndex;
+ }
+
+ long getTaskIndex() {
+ return taskIndex;
+ }
+
+ State getState() {
+ return state;
+ }
+
+ ParallelExecutionTask getWork() {
+ return work;
+ }
+
+ /**
+ * Sets the work; called by the pulling worker under the lock, which also publishes it to
+ * the commit-drain readers.
+ */
+ void setWork(ParallelExecutionTask work) {
+ this.work = work;
+ }
+ }
+
+ /** Worker pull order: ascending group {@code recordIndex}, then per-key {@code taskIndex}. */
+ private static final Comparator<Node> PULL_ORDER =
+ Comparator.comparingLong((Node node) -> node.group.recordIndex)
+ .thenComparingLong(node -> node.taskIndex);
+
+ /** Enforces the "called while holding the mailbox lock" contract on mutating methods. */
+ private final ParallelExecutionLock parallelExecutionLock;
+
+ /**
+ * Active key groups, keyed by the operator key. A group exists iff it has outstanding nodes.
+ */
+ private final Map<Object, KeyGroup> groupsByKey = new HashMap<>();
+
+ /** Not-yet-pulled nodes across all groups, ordered by (recordIndex, taskIndex). */
+ private final PriorityQueue<Node> pendingNodes = new PriorityQueue<>(PULL_ORDER);
+
+ /** Next key arrival sequence to assign. */
+ private long recordIndexGen;
+
+ ParallelExecutionTaskQueue(ParallelExecutionLock parallelExecutionLock) {
+ this.parallelExecutionLock = checkNotNull(parallelExecutionLock);
+ }
+
+ /**
+ * Creates a pending node for {@code key} (lazily creating its arrival-sequenced group) and
+ * registers it in the pull queue; the caller submits the matching dispatch permit.
+ */
+ void addTask(Object key) {
+ parallelExecutionLock.checkReentrant();
+ KeyGroup group = groupsByKey.computeIfAbsent(key, k -> new KeyGroup(k, recordIndexGen++));
+ Node node = new Node(group, group.taskIndexGen++);
+ group.nodes.addLast(node);
+ pendingNodes.add(node);
+ }
+
+ /**
+ * Takes the highest-priority pending node; guaranteed present because the coordinator submits
+ * exactly one dispatch permit per pending node. The node stays in its key group.
+ */
+ Node pollPendingNode() {
+ parallelExecutionLock.checkReentrant();
+ Node node = pendingNodes.poll();
+ checkState(node != null, "No pending node available for a redeemed dispatch permit.");
+ return node;
+ }
+
+ /** Marks a node as executed; called by the owning worker before handing off to the mailbox. */
+ void markDone(Node node) {
+ parallelExecutionLock.checkReentrant();
+ checkState(
+ node.state == State.RUNNING,
+ "Node must be RUNNING to mark DONE, but was: %s",
+ node.state);
+ node.state = State.DONE;
+ }
+
+ /**
+ * Returns {@code key}'s group head iff it is DONE, without removing it; otherwise {@code null}.
+ * The commit drain repeats commit + {@link #removeCommitted} until this returns null.
+ */
+ @Nullable
+ Node pollCommittableHead(Object key) {
+ KeyGroup group = groupsByKey.get(key);
+ if (group == null || group.nodes.isEmpty()) {
+ return null;
+ }
+ Node head = group.nodes.peekFirst();
+ return head.state == State.DONE ? head : null;
+ }
+
+ /**
+ * Removes a committed node (must be its group's head); drops the group once its last node is
+ * removed, so a drained key is no longer outstanding.
+ */
+ void removeCommitted(Node node) {
+ parallelExecutionLock.checkReentrant();
+ checkState(
+ node.state == State.DONE, "Node must be DONE to commit, but was: %s", node.state);
+ node.state = State.COMMITTED;
+ KeyGroup group = node.group;
+ Node head = group.nodes.pollFirst();
+ checkState(head == node, "Committed node was not its key group's head.");
+ if (group.nodes.isEmpty()) {
+ groupsByKey.remove(group.key);
+ }
+ }
+
+ /** Whether {@code key} still has outstanding (added but not yet committed) nodes. */
+ boolean hasOutstanding(Object key) {
+ KeyGroup group = groupsByKey.get(key);
+ return group != null && !group.nodes.isEmpty();
+ }
+
+ @VisibleForTesting
+ int sizeForTesting() {
+ int total = 0;
+ for (KeyGroup group : groupsByKey.values()) {
+ total += group.nodes.size();
+ }
+ return total;
+ }
+}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonInterpreterManager.java b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonInterpreterManager.java
index c007126b..3f042307 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonInterpreterManager.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonInterpreterManager.java
@@ -43,11 +43,13 @@ import java.util.function.Supplier;
* Owns the Pemja interpreters used by one operator subtask.
*
* <p>A {@link PythonInterpreter} is never shared by concurrently executing threads. The owner
- * mailbox thread and Flink Agents' managed Java async workers bind an interpreter to themselves.
- * Calls from every other thread are routed to a bounded set of Java callback workers, each of which
- * owns its interpreter. This distinction is required for threads created by CPython (for example
- * Mem0 or Python's async executor): creating a Pemja interpreter there would attach a second {@code
- * PyThreadState} to the same native thread.
+ * mailbox thread and Flink Agents' managed Java async workers bind an interpreter to themselves for
+ * ordinary calls. A managed worker may also use the owner interpreter to create a handle that must
+ * survive that worker, but only while the operator's execution-access checker guarantees exclusive
+ * access. Calls from every other thread are routed to a bounded set of Java callback workers, each
+ * of which owns its interpreter. This distinction is required for threads created by CPython (for
+ * example Mem0 or Python's async executor): creating a Pemja interpreter there would attach a
+ * second {@code PyThreadState} to the same native thread.
*
* <p>Pemja's {@code MULTI_THREAD} interpreters share the same CPython main interpreter, so Python
* resource objects can still be initialized once and passed as opaque handles to calls made by any
@@ -68,7 +70,9 @@ public final class PythonInterpreterManager implements AutoCloseable {
private final Supplier<PythonInterpreter> interpreterFactory;
private final Consumer<PythonInterpreter> interpreterInitializer;
+ private final PythonInterpreter ownerInterpreter;
private final Thread ownerThread;
+ private final Runnable ownerInterpreterAccessChecker;
private final Map<Thread, PythonInterpreter> interpreters = new ConcurrentHashMap<>();
private final ThreadLocal<PythonInterpreter> threadInterpreter = new ThreadLocal<>();
private final ThreadLocal<Integer> callbackLane;
@@ -85,7 +89,8 @@ public final class PythonInterpreterManager implements AutoCloseable {
ownerInterpreter,
interpreterFactory,
PythonInterpreterManager::initializeInterpreter,
- DEFAULT_CALLBACK_WORKERS);
+ DEFAULT_CALLBACK_WORKERS,
+ PythonInterpreterManager::rejectNonOwnerInterpreterAccess);
}
public PythonInterpreterManager(
@@ -96,7 +101,21 @@ public final class PythonInterpreterManager implements AutoCloseable {
ownerInterpreter,
interpreterFactory,
PythonInterpreterManager::initializeInterpreter,
- callbackWorkerCount);
+ callbackWorkerCount,
+ PythonInterpreterManager::rejectNonOwnerInterpreterAccess);
+ }
+
+ public PythonInterpreterManager(
+ PythonInterpreter ownerInterpreter,
+ Supplier<PythonInterpreter> interpreterFactory,
+ int callbackWorkerCount,
+ Runnable ownerInterpreterAccessChecker) {
+ this(
+ ownerInterpreter,
+ interpreterFactory,
+ PythonInterpreterManager::initializeInterpreter,
+ callbackWorkerCount,
+ ownerInterpreterAccessChecker);
}
PythonInterpreterManager(
@@ -107,7 +126,8 @@ public final class PythonInterpreterManager implements AutoCloseable {
ownerInterpreter,
interpreterFactory,
interpreterInitializer,
- DEFAULT_CALLBACK_WORKERS);
+ DEFAULT_CALLBACK_WORKERS,
+ PythonInterpreterManager::rejectNonOwnerInterpreterAccess);
}
PythonInterpreterManager(
@@ -115,14 +135,30 @@ public final class PythonInterpreterManager implements AutoCloseable {
Supplier<PythonInterpreter> interpreterFactory,
Consumer<PythonInterpreter> interpreterInitializer,
int callbackWorkerCount) {
+ this(
+ ownerInterpreter,
+ interpreterFactory,
+ interpreterInitializer,
+ callbackWorkerCount,
+ PythonInterpreterManager::rejectNonOwnerInterpreterAccess);
+ }
+
+ PythonInterpreterManager(
+ PythonInterpreter ownerInterpreter,
+ Supplier<PythonInterpreter> interpreterFactory,
+ Consumer<PythonInterpreter> interpreterInitializer,
+ int callbackWorkerCount,
+ Runnable ownerInterpreterAccessChecker) {
if (callbackWorkerCount <= 0) {
throw new IllegalArgumentException("Callback worker count must be greater than zero.");
}
this.interpreterFactory = Objects.requireNonNull(interpreterFactory);
this.interpreterInitializer = Objects.requireNonNull(interpreterInitializer);
+ this.ownerInterpreterAccessChecker = Objects.requireNonNull(ownerInterpreterAccessChecker);
this.ownerThread = Thread.currentThread();
PythonInterpreter initializedOwner = initialize(Objects.requireNonNull(ownerInterpreter));
+ this.ownerInterpreter = initializedOwner;
interpreters.put(ownerThread, initializedOwner);
threadInterpreter.set(initializedOwner);
callbackExecutors = createCallbackExecutors(callbackWorkerCount);
@@ -162,6 +198,29 @@ public final class PythonInterpreterManager implements AutoCloseable {
}
}
+ /**
+ * Executes an operation using the owner interpreter without changing the calling Java thread.
+ *
+ * <p>The owner thread may call this method directly. Any other caller must satisfy the supplied
+ * execution-access checker, which requires a managed worker to own the parallel execution lock
+ * in the JDK < 21 execution path. The caller is responsible for that serialization because
+ * Pemja does not synchronize calls to one interpreter.
+ */
+ public <T> T withOwnerInterpreter(Function<PythonInterpreter, T> operation) {
+ lifecycleLock.readLock().lock();
+ try {
+ if (closed) {
+ throw new IllegalStateException("Python interpreter manager is already closed.");
+ }
+ if (Thread.currentThread() != ownerThread) {
+ ownerInterpreterAccessChecker.run();
+ }
+ return Objects.requireNonNull(operation).apply(ownerInterpreter);
+ } finally {
+ lifecycleLock.readLock().unlock();
+ }
+ }
+
public void exec(String code) {
withInterpreter(
interpreter -> {
@@ -174,6 +233,10 @@ public final class PythonInterpreterManager implements AutoCloseable {
return withInterpreter(interpreter -> interpreter.invoke(name, args));
}
+ public Object invokeOnOwner(String name, Object... args) {
+ return withOwnerInterpreter(interpreter -> interpreter.invoke(name, args));
+ }
+
public Object get(String name) {
return withInterpreter(interpreter -> interpreter.get(name));
}
@@ -246,6 +309,11 @@ public final class PythonInterpreterManager implements AutoCloseable {
interpreter.exec(PYTHON_IMPORTS);
}
+ private static void rejectNonOwnerInterpreterAccess() {
+ throw new IllegalStateException(
+ "Owner interpreter access from a non-owner thread requires an execution checker.");
+ }
+
/** Releases the interpreter owned by the current managed Java worker. */
public void releaseCurrentThreadInterpreter() {
if (!AsyncExecutorThreadFactory.isAsyncExecutorThread()) {
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImpl.java b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImpl.java
index ae840b3c..2178ad3d 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImpl.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImpl.java
@@ -141,7 +141,10 @@ public class PythonResourceAdapterImpl implements PythonResourceAdapter, AutoClo
@Override
public PyObject initPythonResource(String module, String clazz, Map<String, Object> kwargs) {
kwargs.put(RESOURCE_CONTEXT_KEY, pythonResourceContext);
- return (PyObject) interpreterManager.invoke(CREATE_RESOURCE, module, clazz, kwargs);
+ // Resource wrappers live in the operator-level cache and may outlive the managed worker
+ // that observes the cache miss. Bind their root handles to the owner interpreter instead of
+ // that worker's disposable interpreter.
+ return (PyObject) interpreterManager.invokeOnOwner(CREATE_RESOURCE, module, clazz, kwargs);
}
@Override
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 e34649ed..f96fd34a 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
@@ -22,6 +22,9 @@ import org.apache.flink.agents.api.context.Outcome;
import jdk.internal.vm.Continuation;
import jdk.internal.vm.ContinuationScope;
+import org.apache.flink.agents.runtime.operator.parallel.ParallelExecutionContextRestorer;
+import org.apache.flink.agents.runtime.operator.parallel.ParallelExecutionLock;
+
import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
@@ -34,6 +37,7 @@ import java.util.concurrent.Future;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicIntegerArray;
import java.util.function.Supplier;
+
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -52,10 +56,19 @@ public class ContinuationActionExecutor {
private final AsyncExecutorThreadFactory asyncThreadFactory;
public ContinuationActionExecutor(int numAsyncThreads) {
- this(numAsyncThreads, () -> {});
+ this(numAsyncThreads, () -> {}, null, null);
}
- public ContinuationActionExecutor(int numAsyncThreads, Runnable threadCleanup) {
+ /**
+ * Constructor kept signature-compatible with the base (JDK<21) executor for multi-release
+ * uniformity. Continuation execution owns its own async thread pool and never uses the mailbox
+ * lock or context restorer, so both are ignored here.
+ */
+ public ContinuationActionExecutor(
+ int numAsyncThreads,
+ Runnable threadCleanup,
+ ParallelExecutionLock parallelExecutionLock,
+ ParallelExecutionContextRestorer contextRestorer) {
LOG.info("Initialize fixed thread pool for async task with {} threads", numAsyncThreads);
this.asyncThreadFactory = new AsyncExecutorThreadFactory(threadCleanup);
this.asyncExecutor =
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/ResourceCacheTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/ResourceCacheTest.java
index d82cc3ec..17ec9d20 100644
--- a/runtime/src/test/java/org/apache/flink/agents/runtime/ResourceCacheTest.java
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/ResourceCacheTest.java
@@ -63,6 +63,7 @@ import java.util.concurrent.atomic.AtomicInteger;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.assertj.core.api.Assertions.catchThrowable;
+import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@@ -289,6 +290,27 @@ public class ResourceCacheTest {
.hasMessageContaining("Resource not found: non-existent");
}
+ @Test
+ public void getResourceClosesUncachedResourceWhenOpenFails() throws Exception {
+ RuntimeException openFailure = new RuntimeException("open failed");
+ RuntimeException closeFailure = new RuntimeException("close failed");
+ Resource resource = mock(Resource.class);
+ ResourceProvider provider = mock(ResourceProvider.class);
+ when(provider.provide(org.mockito.ArgumentMatchers.any())).thenReturn(resource);
+ doThrow(openFailure).when(resource).open();
+ doThrow(closeFailure).when(resource).close();
+ ResourceCache cache =
+ new ResourceCache(Map.of(ResourceType.TOOL, Map.of("tool", provider)));
+
+ assertThatThrownBy(() -> cache.getResource("tool", ResourceType.TOOL))
+ .isSameAs(openFailure)
+ .satisfies(
+ thrown -> assertThat(thrown.getSuppressed()).containsExactly(closeFailure));
+
+ verify(resource).close();
+ assertThat(cachedResources(cache)).isEmpty();
+ }
+
@Test
public void testHasResourceCoversProvidersAndCachedOnlyEntries() throws Exception {
TestAgentWithResources agent = new TestAgentWithResources();
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutorTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutorTest.java
new file mode 100644
index 00000000..fac9c4b1
--- /dev/null
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutorTest.java
@@ -0,0 +1,200 @@
+/*
+ * 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.runtime.operator.parallel.ParallelExecutionContextRestorer;
+import org.apache.flink.agents.runtime.operator.parallel.ParallelExecutionLock;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests the JDK<21 mailbox-lock handoff contract. */
+class ContinuationActionExecutorTest {
+
+ @Test
+ void executeAsyncReleasesAndReacquiresLockAndRestoresContext() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ AtomicBoolean restored = new AtomicBoolean();
+ // The executor drives the restorer once the worker re-acquires the lock after the blocking
+ // callable returns; the context is otherwise opaque to this test.
+ ParallelExecutionContextRestorer restorer =
+ (key, action) -> {
+ lock.checkReentrant();
+ restored.set(true);
+ };
+ ContinuationActionExecutor executor =
+ new ContinuationActionExecutor(1, () -> {}, lock, restorer);
+ ContinuationContext context = new ContinuationContext();
+
+ CountDownLatch supplierStarted = new CountDownLatch(1);
+ CountDownLatch allowSupplierToFinish = new CountDownLatch(1);
+
+ ExecutorService workerExecutor = Executors.newSingleThreadExecutor();
+ try {
+ Future<String> workerResult =
+ workerExecutor.submit(
+ () -> {
+ lock.acquireByWorker(0, 0);
+ try {
+ return executor.executeAsync(
+ context,
+ () -> {
+ supplierStarted.countDown();
+ await(allowSupplierToFinish);
+ return "result";
+ });
+ } finally {
+ lock.release();
+ }
+ });
+
+ assertThat(supplierStarted.await(5, TimeUnit.SECONDS)).isTrue();
+
+ lock.acquireByMain();
+ try {
+ lock.checkReentrant();
+ } finally {
+ lock.release();
+ }
+ assertThat(restored).isFalse();
+
+ allowSupplierToFinish.countDown();
+ assertThat(workerResult.get(5, TimeUnit.SECONDS)).isEqualTo("result");
+ assertThat(restored).isTrue();
+ } finally {
+ workerExecutor.shutdownNow();
+ executor.close();
+ }
+ }
+
+ @Test
+ void executeAllAsyncReleasesAndReacquiresLockAndRestoresContext() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ AtomicBoolean restored = new AtomicBoolean();
+ ParallelExecutionContextRestorer restorer =
+ (key, action) -> {
+ lock.checkReentrant();
+ restored.set(true);
+ };
+ ContinuationActionExecutor executor =
+ new ContinuationActionExecutor(1, () -> {}, lock, restorer);
+ ContinuationContext context = new ContinuationContext();
+ CountDownLatch supplierStarted = new CountDownLatch(1);
+ CountDownLatch allowSupplierToFinish = new CountDownLatch(1);
+ List<Callable<String>> suppliers =
+ List.of(
+ () -> {
+ supplierStarted.countDown();
+ await(allowSupplierToFinish);
+ return "first";
+ },
+ () -> "second");
+
+ ExecutorService workerExecutor = Executors.newSingleThreadExecutor();
+ try {
+ Future<BatchExecutionResult<String>> workerResult =
+ workerExecutor.submit(
+ () -> {
+ lock.acquireByWorker(0, 0);
+ try {
+ return executor.executeAllAsync(
+ context, suppliers, null, suppliers.size());
+ } finally {
+ lock.release();
+ }
+ });
+
+ assertThat(supplierStarted.await(5, TimeUnit.SECONDS)).isTrue();
+
+ lock.acquireByMain();
+ try {
+ lock.checkReentrant();
+ } finally {
+ lock.release();
+ }
+ assertThat(restored).isFalse();
+
+ allowSupplierToFinish.countDown();
+ BatchExecutionResult<String> result = workerResult.get(5, TimeUnit.SECONDS);
+ assertThat(result.getOutcomes().get(0).getValue()).isEqualTo("first");
+ assertThat(result.getOutcomes().get(1).getValue()).isEqualTo("second");
+ assertThat(restored).isTrue();
+ } finally {
+ workerExecutor.shutdownNow();
+ executor.close();
+ }
+ }
+
+ @Test
+ void executeAsyncAbortsBeforeSupplierWhenCallerDoesNotOwnLock() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ AtomicBoolean supplierCalled = new AtomicBoolean();
+ AtomicBoolean restored = new AtomicBoolean();
+ ContinuationActionExecutor executor =
+ new ContinuationActionExecutor(
+ 1, () -> {}, lock, (key, action) -> restored.set(true));
+ try {
+ // executeAsync asserts mailbox-lock ownership before releasing it. Calling without the
+ // lock must fail fast, before the supplier runs and before any restore.
+ assertThatThrownBy(
+ () ->
+ executor.executeAsync(
+ new ContinuationContext(),
+ () -> {
+ supplierCalled.set(true);
+ return "unexpected";
+ }))
+ .isInstanceOf(IllegalStateException.class);
+ assertThat(supplierCalled).isFalse();
+ assertThat(restored).isFalse();
+ } finally {
+ executor.close();
+ }
+ }
+
+ @Test
+ void executeAsyncWithoutMailboxLockKeepsSynchronousFallback() throws Exception {
+ ContinuationActionExecutor executor =
+ new ContinuationActionExecutor(1, () -> {}, null, null);
+ try {
+ assertThat(executor.executeAsync(new ContinuationContext(), () -> "result"))
+ .isEqualTo("result");
+ } finally {
+ executor.close();
+ }
+ }
+
+ private static void await(CountDownLatch latch) {
+ try {
+ assertThat(latch.await(5, TimeUnit.SECONDS)).isTrue();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new AssertionError("Interrupted while waiting for test supplier", e);
+ }
+ }
+}
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 ad94ac7b..401374d0 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
@@ -899,7 +899,7 @@ class JavaRunnerContextImplDurableExecuteAsyncTest {
private ExecutorService batchExecutor = Executors.newFixedThreadPool(4);
private InspectingContinuationActionExecutor() {
- super(1);
+ super(1, () -> {}, null, null);
}
@Override
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/lifecycle/TaskLifecycleListenerNotificationTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/lifecycle/TaskLifecycleListenerNotificationTest.java
index aec715b8..c7711f87 100644
--- a/runtime/src/test/java/org/apache/flink/agents/runtime/lifecycle/TaskLifecycleListenerNotificationTest.java
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/lifecycle/TaskLifecycleListenerNotificationTest.java
@@ -23,8 +23,10 @@ import org.apache.flink.agents.api.EventType;
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.plan.AgentConfiguration;
import org.apache.flink.agents.plan.AgentPlan;
import org.apache.flink.agents.runtime.actionstate.ActionState;
import org.apache.flink.agents.runtime.actionstate.ActionStateStore;
@@ -369,7 +371,11 @@ public class TaskLifecycleListenerNotificationTest {
@Test
void resumedInFlightRecordReEmitsRecordStart() throws Exception {
- AgentPlan plan = new AgentPlan(new SyncAgent());
+ // Force serial: the parallel engine drains in-flight tasks before checkpoint, so no
+ // in-flight record survives the snapshot — this test exercises the serial restore path.
+ AgentConfiguration serialConfig = new AgentConfiguration();
+ serialConfig.set(AgentExecutionOptions.PARALLEL_EXECUTION_ENABLED, false);
+ AgentPlan plan = new AgentPlan(new SyncAgent(), serialConfig);
OperatorSubtaskState snapshot;
// First execution: admit the record but snapshot before its tasks run, so the record
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorCapacityIntegrationTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorCapacityIntegrationTest.java
new file mode 100644
index 00000000..6e796322
--- /dev/null
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorCapacityIntegrationTest.java
@@ -0,0 +1,338 @@
+/*
+ * 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.operator;
+
+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.AgentExecutionOptions;
+import org.apache.flink.agents.api.context.DurableCallable;
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.plan.AgentConfiguration;
+import org.apache.flink.agents.plan.AgentPlan;
+import org.apache.flink.agents.plan.JavaFunction;
+import org.apache.flink.agents.plan.actions.Action;
+import org.apache.flink.agents.runtime.async.ContinuationActionExecutor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assumptions.assumeThat;
+
+/** Component-level integration tests for action backlog beyond the worker count. */
+class ActionExecutionOperatorCapacityIntegrationTest {
+
+ private static final int NUM_ASYNC_THREADS = 2;
+ private static final int WORKER_BACKLOG_ACTIONS = NUM_ASYNC_THREADS * 2;
+ private static final int NUM_ACTIONS = WORKER_BACKLOG_ACTIONS + 2;
+
+ @Test
+ @Timeout(20)
+ void actionsBeyondWorkerCountEventuallyCompleteExactlyOnce() throws Exception {
+ assumeThat(ContinuationActionExecutor.isContinuationSupported()).isFalse();
+ CapacityAgent.reset();
+
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.NUM_ASYNC_THREADS, NUM_ASYNC_THREADS);
+
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory<>(
+ CapacityAgent.getAgentPlan(config), true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+
+ testHarness.processElement(new StreamRecord<>(0L));
+ // Resident workers pull and start the tasks autonomously once the input is admitted;
+ // only NUM_ASYNC_THREADS of them can enter the async stage at once.
+
+ assertThat(CapacityAgent.FIRST_BATCH_STARTED.await(5, TimeUnit.SECONDS)).isTrue();
+ assertThat(CapacityAgent.STARTED_COUNT).hasValue(NUM_ASYNC_THREADS);
+ assertThat(testHarness.getRecordOutput()).isEmpty();
+
+ CapacityAgent.ALLOW_FIRST_BATCH.countDown();
+ operator.waitInFlightEventsFinished();
+
+ List<Object> outputValues = outputValues(testHarness);
+
+ Set<String> expected = expectedActions();
+ assertThat(outputValues)
+ .hasSize(NUM_ACTIONS)
+ .containsExactlyInAnyOrderElementsOf(expected);
+ } finally {
+ CapacityAgent.ALLOW_FIRST_BATCH.countDown();
+ }
+ }
+
+ @Test
+ @Timeout(30)
+ void actionsBeyondWorkerCountRecoverFromCheckpointWithExactlyOnceOutput() throws Exception {
+ assumeThat(ContinuationActionExecutor.isContinuationSupported()).isFalse();
+ CapacityAgent.reset();
+ CapacityAgent.blockActionsAfterFirstBatch();
+
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.NUM_ASYNC_THREADS, NUM_ASYNC_THREADS);
+
+ OperatorSubtaskState snapshot;
+ List<Object> beforeRestore;
+ Thread firstBatchReleaser = null;
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory<>(
+ CapacityAgent.getAgentPlan(config), true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+
+ testHarness.processElement(new StreamRecord<>(0L));
+ assertThat(CapacityAgent.FIRST_BATCH_STARTED.await(5, TimeUnit.SECONDS)).isTrue();
+ assertThat(testHarness.getRecordOutput()).isEmpty();
+
+ firstBatchReleaser =
+ new Thread(
+ () -> {
+ try {
+ Thread.sleep(300L);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ CapacityAgent.ALLOW_FIRST_BATCH.countDown();
+ });
+ firstBatchReleaser.setDaemon(true);
+ firstBatchReleaser.start();
+
+ operator.prepareSnapshotPreBarrier(1L);
+ snapshot = testHarness.snapshot(1L, 1L);
+ assertThat(CapacityAgent.REMAINING_ACTION_STARTED.await(5, TimeUnit.SECONDS))
+ .as(
+ "at least one post-checkpoint callable attempt should start before failover")
+ .isTrue();
+
+ beforeRestore = outputValues(testHarness);
+ assertThat(beforeRestore)
+ .containsExactlyInAnyOrder("capacity-action-1", "capacity-action-2");
+
+ CapacityAgent.ALLOW_REMAINING_ACTIONS.countDown();
+ operator.waitInFlightEventsFinished();
+ } finally {
+ CapacityAgent.ALLOW_FIRST_BATCH.countDown();
+ CapacityAgent.ALLOW_REMAINING_ACTIONS.countDown();
+ if (firstBatchReleaser != null) {
+ firstBatchReleaser.join();
+ }
+ }
+
+ CapacityAgent.allowActionsAfterFirstBatch();
+ CapacityAgent.resetFirstBatchGate();
+ CapacityAgent.ALLOW_FIRST_BATCH.countDown();
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> restored =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory<>(
+ CapacityAgent.getAgentPlan(config), true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ restored.initializeState(snapshot);
+ restored.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) restored.getOperator();
+
+ operator.waitInFlightEventsFinished();
+
+ List<Object> allOutputs = new ArrayList<>(beforeRestore);
+ allOutputs.addAll(outputValues(restored));
+ Set<String> expected = expectedActions();
+ assertThat(allOutputs)
+ .hasSize(NUM_ACTIONS)
+ .containsExactlyInAnyOrderElementsOf(expected);
+ for (String action : expected) {
+ assertThat(allOutputs).filteredOn(action::equals).hasSize(1);
+ }
+
+ assertThat(CapacityAgent.TOTAL_STARTED_COUNT.get()).isGreaterThan(NUM_ACTIONS);
+ assertThat(CapacityAgent.CALL_COUNTS.values())
+ .anySatisfy(counter -> assertThat(counter.get()).isGreaterThan(1));
+ } finally {
+ CapacityAgent.ALLOW_FIRST_BATCH.countDown();
+ CapacityAgent.ALLOW_REMAINING_ACTIONS.countDown();
+ }
+ }
+
+ private static List<Object> outputValues(
+ KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness) {
+ List<Object> outputValues = new ArrayList<>();
+ for (StreamRecord<Object> record : testHarness.getRecordOutput()) {
+ outputValues.add(record.getValue());
+ }
+ return outputValues;
+ }
+
+ private static Set<String> expectedActions() {
+ Set<String> expected = new LinkedHashSet<>();
+ for (int i = 1; i <= NUM_ACTIONS; i++) {
+ expected.add("capacity-action-" + i);
+ }
+ return expected;
+ }
+
+ public static class CapacityAgent {
+
+ private static volatile CountDownLatch FIRST_BATCH_STARTED = new CountDownLatch(0);
+ private static volatile CountDownLatch ALLOW_FIRST_BATCH = new CountDownLatch(0);
+ private static volatile CountDownLatch REMAINING_ACTION_STARTED = new CountDownLatch(0);
+ private static volatile CountDownLatch ALLOW_REMAINING_ACTIONS = new CountDownLatch(0);
+ private static volatile boolean blockActionsAfterFirstBatch;
+ private static final AtomicInteger STARTED_COUNT = new AtomicInteger();
+ private static final AtomicInteger TOTAL_STARTED_COUNT = new AtomicInteger();
+ private static final Map<String, AtomicInteger> CALL_COUNTS = new ConcurrentHashMap<>();
+
+ private static void reset() {
+ resetFirstBatchGate();
+ REMAINING_ACTION_STARTED = new CountDownLatch(0);
+ ALLOW_REMAINING_ACTIONS = new CountDownLatch(0);
+ blockActionsAfterFirstBatch = false;
+ TOTAL_STARTED_COUNT.set(0);
+ CALL_COUNTS.clear();
+ }
+
+ private static void resetFirstBatchGate() {
+ FIRST_BATCH_STARTED = new CountDownLatch(NUM_ASYNC_THREADS);
+ ALLOW_FIRST_BATCH = new CountDownLatch(1);
+ STARTED_COUNT.set(0);
+ }
+
+ private static void blockActionsAfterFirstBatch() {
+ blockActionsAfterFirstBatch = true;
+ REMAINING_ACTION_STARTED = new CountDownLatch(1);
+ ALLOW_REMAINING_ACTIONS = new CountDownLatch(1);
+ }
+
+ private static void allowActionsAfterFirstBatch() {
+ blockActionsAfterFirstBatch = false;
+ ALLOW_REMAINING_ACTIONS.countDown();
+ }
+
+ private static void runAction(Event event, RunnerContext context, int actionNumber)
+ throws Exception {
+ Long input = (Long) InputEvent.fromEvent(event).getInput();
+ String actionId = "capacity-action-" + actionNumber;
+ context.durableExecuteAsync(
+ new DurableCallable<String>() {
+ @Override
+ public String getId() {
+ return actionId + "-" + input;
+ }
+
+ @Override
+ public Class<String> getResultClass() {
+ return String.class;
+ }
+
+ @Override
+ public String call() throws Exception {
+ STARTED_COUNT.incrementAndGet();
+ TOTAL_STARTED_COUNT.incrementAndGet();
+ CALL_COUNTS
+ .computeIfAbsent(
+ actionId, ignored -> new AtomicInteger())
+ .incrementAndGet();
+ if (actionNumber <= NUM_ASYNC_THREADS) {
+ FIRST_BATCH_STARTED.countDown();
+ if (!ALLOW_FIRST_BATCH.await(5, TimeUnit.SECONDS)) {
+ throw new IllegalStateException(
+ "Timed out waiting for first capacity batch");
+ }
+ } else if (blockActionsAfterFirstBatch) {
+ REMAINING_ACTION_STARTED.countDown();
+ if (!ALLOW_REMAINING_ACTIONS.await(30, TimeUnit.SECONDS)) {
+ throw new IllegalStateException(
+ "Timed out waiting for remaining capacity"
+ + " actions");
+ }
+ }
+ return actionId;
+ }
+ })
+ .await();
+ context.sendEvent(new OutputEvent(actionId));
+ }
+
+ public static void capacityAction1(Event event, RunnerContext context) throws Exception {
+ runAction(event, context, 1);
+ }
+
+ public static void capacityAction2(Event event, RunnerContext context) throws Exception {
+ runAction(event, context, 2);
+ }
+
+ public static void capacityAction3(Event event, RunnerContext context) throws Exception {
+ runAction(event, context, 3);
+ }
+
+ public static void capacityAction4(Event event, RunnerContext context) throws Exception {
+ runAction(event, context, 4);
+ }
+
+ public static void capacityAction5(Event event, RunnerContext context) throws Exception {
+ runAction(event, context, 5);
+ }
+
+ public static void capacityAction6(Event event, RunnerContext context) throws Exception {
+ runAction(event, context, 6);
+ }
+
+ private static AgentPlan getAgentPlan(AgentConfiguration config) throws Exception {
+ Map<String, Action> actions = new LinkedHashMap<>();
+ for (int i = 1; i <= NUM_ACTIONS; i++) {
+ String actionName = "capacityAction" + i;
+ Action action =
+ new Action(
+ actionName,
+ new JavaFunction(
+ CapacityAgent.class,
+ actionName,
+ new Class<?>[] {Event.class, RunnerContext.class}),
+ Collections.singletonList(InputEvent.EVENT_TYPE));
+ actions.put(actionName, action);
+ }
+ return new AgentPlan(actions, new LinkedHashMap<>(), config);
+ }
+ }
+}
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorDrainDeadlockIntegrationTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorDrainDeadlockIntegrationTest.java
new file mode 100644
index 00000000..9f70a95a
--- /dev/null
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorDrainDeadlockIntegrationTest.java
@@ -0,0 +1,495 @@
+/*
+ * 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.operator;
+
+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.AgentExecutionOptions;
+import org.apache.flink.agents.api.context.DurableCallable;
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.plan.AgentConfiguration;
+import org.apache.flink.agents.plan.AgentPlan;
+import org.apache.flink.agents.plan.JavaFunction;
+import org.apache.flink.agents.plan.actions.Action;
+import org.apache.flink.agents.runtime.async.ContinuationActionExecutor;
+import org.apache.flink.agents.runtime.operator.parallel.ParallelExecutionCoordinator;
+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.runtime.tasks.mailbox.Mail;
+import org.apache.flink.streaming.runtime.tasks.mailbox.TaskMailbox;
+import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+
+import java.lang.management.ManagementFactory;
+import java.lang.management.ThreadInfo;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.function.BooleanSupplier;
+import java.util.function.Predicate;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.catchThrowable;
+import static org.assertj.core.api.Assumptions.assumeThat;
+
+/**
+ * Component-level regression test for the checkpoint drain quiesce of the JDK<21 parallel
+ * engine: a worker that redeems a queued dispatch permit while the drain is active must park
+ * without holding back the in-flight tasks whose completion the quiesce is waiting for.
+ *
+ * <p>Timeline (three actions, two workers; the harness mailbox thread is the test thread, so the
+ * post-barrier steps are driven from a controller thread):
+ *
+ * <ol>
+ * <li>One input record queues actions A, B and C. The workers pull A and B, which block in their
+ * async sections; C's dispatch permit stays queued behind the saturated pool.
+ * <li>The barrier arrives: the test thread enters {@code prepareSnapshotPreBarrier} and blocks in
+ * the quiesce loop.
+ * <li>The controller releases A; its worker commits A, redeems C's queued permit, sees the drain
+ * and parks.
+ * <li>The controller releases B; B's async resume must re-acquire the lock so the drain can
+ * quiesce. If it cannot, the controller flags the deadlock and interrupts the mailbox thread
+ * so the test reports instead of hanging until the JUnit timeout.
+ * </ol>
+ */
+class ActionExecutionOperatorDrainDeadlockIntegrationTest {
+
+ private static final int NUM_ASYNC_THREADS = 2;
+
+ /** Bound for every phase that must happen for the timeline to be set up. */
+ private static final long SETUP_TIMEOUT_SECONDS = 10;
+
+ /** Time budget for the drain to quiesce after the last in-flight action is released. */
+ private static final long DRAIN_COMPLETION_TIMEOUT_SECONDS = 10;
+
+ @Test
+ @Timeout(120)
+ void drainQuiesceCompletesWhenQueuedPermitIsRedeemedDuringDrain() throws Exception {
+ assumeThat(ContinuationActionExecutor.isContinuationSupported()).isFalse();
+ DrainAgent.reset();
+
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.NUM_ASYNC_THREADS, NUM_ASYNC_THREADS);
+
+ Thread mailboxThread = Thread.currentThread();
+ CountDownLatch drainReturned = new CountDownLatch(1);
+ TimelineObservations observations = new TimelineObservations();
+ Thread controller = null;
+
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory<>(DrainAgent.getAgentPlan(config), true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+
+ // Step 1: workers block inside A's and B's async sections; C's permit stays queued.
+ testHarness.processElement(new StreamRecord<>(0L));
+ assertThat(DrainAgent.A_IN_ASYNC.await(SETUP_TIMEOUT_SECONDS, TimeUnit.SECONDS))
+ .as("action A must reach its async section")
+ .isTrue();
+ assertThat(DrainAgent.B_IN_ASYNC.await(SETUP_TIMEOUT_SECONDS, TimeUnit.SECONDS))
+ .as("action B must reach its async section")
+ .isTrue();
+ assertThat(DrainAgent.C_STARTED.getCount())
+ .as("action C must not have started: its permit is queued behind the workers")
+ .isEqualTo(1);
+
+ controller =
+ new Thread(() -> driveTimeline(observations, drainReturned, mailboxThread));
+ controller.setName("drain-timeline-controller");
+ controller.setDaemon(true);
+ controller.start();
+
+ // Steps 2-4 as seen from the mailbox thread: start the drain and wait for quiesce.
+ Throwable drainFailure = catchThrowable(() -> operator.prepareSnapshotPreBarrier(1L));
+ drainReturned.countDown();
+ controller.join(TimeUnit.SECONDS.toMillis(SETUP_TIMEOUT_SECONDS));
+ // Clear a possible interrupt from the controller before asserting and tearing down.
+ Thread.interrupted();
+
+ assertThat(observations.setupFailure)
+ .as("controller could not establish the timeline")
+ .isNull();
+ assertThat(observations.drainEngaged)
+ .as("the mailbox thread must have entered the drain quiesce loop")
+ .isTrue();
+ assertThat(observations.workerParkedOnQueuedPermit)
+ .as("a worker must have redeemed C's queued permit and parked during the drain")
+ .isTrue();
+ assertThat(observations.deadlockDetected)
+ .as(
+ "the drain must quiesce within "
+ + DRAIN_COMPLETION_TIMEOUT_SECONDS
+ + "s of releasing the last in-flight action")
+ .isFalse();
+ assertThat(drainFailure).as("prepareSnapshotPreBarrier must not fail").isNull();
+
+ // The drain quiesced; the snapshot resumes dispatch, which lets the parked worker
+ // run C.
+ testHarness.snapshot(1L, 1L);
+ operator.waitInFlightEventsFinished();
+
+ List<Object> outputs = new ArrayList<>();
+ for (StreamRecord<Object> record : testHarness.getRecordOutput()) {
+ outputs.add(record.getValue());
+ }
+ assertThat(outputs)
+ .as("every action must complete exactly once across the checkpoint")
+ .containsExactlyInAnyOrder(
+ "drain-action-A", "drain-action-B", "drain-action-C");
+ } finally {
+ DrainAgent.releaseAll();
+ if (controller != null) {
+ controller.join(TimeUnit.SECONDS.toMillis(SETUP_TIMEOUT_SECONDS));
+ }
+ Thread.interrupted();
+ }
+ }
+
+ @Test
+ @Timeout(120)
+ void oldCheckpointAbortDoesNotResumeCurrentCheckpointDrain() throws Exception {
+ assumeThat(ContinuationActionExecutor.isContinuationSupported()).isFalse();
+ DrainAgent.reset();
+
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.NUM_ASYNC_THREADS, NUM_ASYNC_THREADS);
+
+ Thread mailboxThread = Thread.currentThread();
+ CountDownLatch abortExecuted = new CountDownLatch(1);
+ CountDownLatch drainReturned = new CountDownLatch(1);
+ TimelineObservations observations = new TimelineObservations();
+ Thread controller = null;
+
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory<>(DrainAgent.getAgentPlan(config), true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+
+ testHarness.processElement(new StreamRecord<>(0L));
+ assertThat(DrainAgent.A_IN_ASYNC.await(SETUP_TIMEOUT_SECONDS, TimeUnit.SECONDS))
+ .isTrue();
+ assertThat(DrainAgent.B_IN_ASYNC.await(SETUP_TIMEOUT_SECONDS, TimeUnit.SECONDS))
+ .isTrue();
+ assertThat(DrainAgent.C_STARTED.getCount()).isEqualTo(1);
+
+ testHarness
+ .getTaskMailbox()
+ .put(
+ new Mail(
+ () -> {
+ abortExecuted.countDown();
+ operator.notifyCheckpointAborted(1L);
+ },
+ TaskMailbox.MAX_PRIORITY,
+ "abort old checkpoint"));
+ controller =
+ new Thread(
+ () ->
+ driveAbortTimeline(
+ abortExecuted,
+ observations,
+ drainReturned,
+ mailboxThread));
+ controller.setName("abort-during-drain-controller");
+ controller.setDaemon(true);
+ controller.start();
+
+ Throwable drainFailure = catchThrowable(() -> operator.prepareSnapshotPreBarrier(2L));
+ drainReturned.countDown();
+ controller.join(TimeUnit.SECONDS.toMillis(SETUP_TIMEOUT_SECONDS));
+ Thread.interrupted();
+
+ assertThat(observations.setupFailure)
+ .as("controller could not establish the abort-during-drain timeline")
+ .isNull();
+ assertThat(observations.queuedActionAdvancedAfterAbort)
+ .as("the queued action's worker must either park on drain or start")
+ .isTrue();
+ assertThat(observations.queuedActionStartedBeforeSnapshot)
+ .as("an older checkpoint abort must not resume the current checkpoint drain")
+ .isFalse();
+ assertThat(DrainAgent.C_STARTED.getCount())
+ .as("action C must stay queued until the current checkpoint snapshot resumes")
+ .isEqualTo(1);
+ assertThat(observations.deadlockDetected)
+ .as("the current checkpoint drain must still quiesce")
+ .isFalse();
+ assertThat(drainFailure).as("prepareSnapshotPreBarrier must not fail").isNull();
+
+ testHarness.snapshot(2L, 2L);
+ operator.waitInFlightEventsFinished();
+
+ List<Object> outputs = new ArrayList<>();
+ for (StreamRecord<Object> record : testHarness.getRecordOutput()) {
+ outputs.add(record.getValue());
+ }
+ assertThat(outputs)
+ .as("every action must complete exactly once across the checkpoint")
+ .containsExactlyInAnyOrder(
+ "drain-action-A", "drain-action-B", "drain-action-C");
+ } finally {
+ DrainAgent.releaseAll();
+ if (controller != null) {
+ controller.join(TimeUnit.SECONDS.toMillis(SETUP_TIMEOUT_SECONDS));
+ }
+ Thread.interrupted();
+ }
+ }
+
+ private void driveAbortTimeline(
+ CountDownLatch abortExecuted,
+ TimelineObservations observations,
+ CountDownLatch drainReturned,
+ Thread mailboxThread) {
+ try {
+ if (!abortExecuted.await(SETUP_TIMEOUT_SECONDS, TimeUnit.SECONDS)) {
+ observations.setupFailure =
+ new AssertionError("old checkpoint abort mail did not execute");
+ return;
+ }
+
+ DrainAgent.RELEASE_A.countDown();
+ observations.queuedActionAdvancedAfterAbort =
+ awaitCondition(
+ () -> DrainAgent.C_STARTED.getCount() == 0 || workerParkedInDrain());
+ observations.queuedActionStartedBeforeSnapshot = DrainAgent.C_STARTED.getCount() == 0;
+
+ DrainAgent.RELEASE_B.countDown();
+ if (!drainReturned.await(DRAIN_COMPLETION_TIMEOUT_SECONDS, TimeUnit.SECONDS)) {
+ observations.deadlockDetected = true;
+ mailboxThread.interrupt();
+ }
+ } catch (Throwable t) {
+ observations.setupFailure = t;
+ } finally {
+ DrainAgent.releaseAll();
+ }
+ }
+
+ /** Drives the post-barrier timeline off the mailbox thread and records what happened. */
+ private void driveTimeline(
+ TimelineObservations observations, CountDownLatch drainReturned, Thread mailboxThread) {
+ try {
+ // Step 2: wait until the mailbox thread is actually parked in the quiesce loop.
+ observations.drainEngaged =
+ awaitCondition(
+ ActionExecutionOperatorDrainDeadlockIntegrationTest
+ ::mailboxInDrainQuiesce);
+ if (!observations.drainEngaged) {
+ return;
+ }
+
+ // Step 3: free one worker; it must pick up C's queued permit and park on the drain.
+ DrainAgent.RELEASE_A.countDown();
+ observations.workerParkedOnQueuedPermit =
+ awaitCondition(
+ ActionExecutionOperatorDrainDeadlockIntegrationTest
+ ::workerParkedInDrain);
+ if (!observations.workerParkedOnQueuedPermit) {
+ return;
+ }
+
+ // Step 4: release the still-in-flight action; its async resume must get the lock back.
+ DrainAgent.RELEASE_B.countDown();
+ if (!drainReturned.await(DRAIN_COMPLETION_TIMEOUT_SECONDS, TimeUnit.SECONDS)) {
+ observations.deadlockDetected = true;
+ // Break the mailbox out of its blocking yield so the test reports instead of
+ // hanging.
+ mailboxThread.interrupt();
+ }
+ } catch (Throwable t) {
+ observations.setupFailure = t;
+ } finally {
+ DrainAgent.releaseAll();
+ }
+ }
+
+ /** Records what the controller observed; read by the mailbox thread after the drain call. */
+ private static final class TimelineObservations {
+ private volatile boolean drainEngaged;
+ private volatile boolean workerParkedOnQueuedPermit;
+ private volatile boolean queuedActionAdvancedAfterAbort;
+ private volatile boolean queuedActionStartedBeforeSnapshot;
+ private volatile boolean deadlockDetected;
+ private volatile Throwable setupFailure;
+ }
+
+ private static boolean awaitCondition(BooleanSupplier condition) throws InterruptedException {
+ long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(SETUP_TIMEOUT_SECONDS);
+ while (System.nanoTime() < deadline) {
+ if (condition.getAsBoolean()) {
+ return true;
+ }
+ Thread.sleep(20);
+ }
+ return false;
+ }
+
+ /** Whether some thread is blocked in the operator's drain quiesce yield loop. */
+ private static boolean mailboxInDrainQuiesce() {
+ return anyThreadMatches(
+ stack ->
+ hasFrame(
+ stack,
+ ActionExecutionOperator.class.getName(),
+ "prepareSnapshotPreBarrier")
+ && hasFrameMethod(stack, "yield"));
+ }
+
+ /**
+ * Whether some worker thread has redeemed a dispatch permit during the drain and parked on the
+ * drain monitor.
+ */
+ private static boolean workerParkedInDrain() {
+ return anyThreadMatches(
+ stack ->
+ hasFrame(stack, ParallelExecutionCoordinator.class.getName(), "runOneTask")
+ && hasFrame(stack, "java.lang.Object", "wait"));
+ }
+
+ private static boolean anyThreadMatches(Predicate<StackTraceElement[]> test) {
+ for (ThreadInfo info : ManagementFactory.getThreadMXBean().dumpAllThreads(false, false)) {
+ if (info != null && test.test(info.getStackTrace())) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ private static boolean hasFrame(
+ StackTraceElement[] stack, String className, String methodName) {
+ return Arrays.stream(stack)
+ .anyMatch(
+ frame ->
+ frame.getClassName().equals(className)
+ && frame.getMethodName().equals(methodName));
+ }
+
+ private static boolean hasFrameMethod(StackTraceElement[] stack, String methodName) {
+ return Arrays.stream(stack).anyMatch(frame -> frame.getMethodName().equals(methodName));
+ }
+
+ /**
+ * Three actions on one input event: A and B block in their async sections until released; C is
+ * an instant no-op.
+ */
+ public static class DrainAgent {
+
+ private static volatile CountDownLatch A_IN_ASYNC = new CountDownLatch(0);
+ private static volatile CountDownLatch B_IN_ASYNC = new CountDownLatch(0);
+ private static volatile CountDownLatch C_STARTED = new CountDownLatch(0);
+ private static volatile CountDownLatch RELEASE_A = new CountDownLatch(0);
+ private static volatile CountDownLatch RELEASE_B = new CountDownLatch(0);
+
+ private static void reset() {
+ A_IN_ASYNC = new CountDownLatch(1);
+ B_IN_ASYNC = new CountDownLatch(1);
+ C_STARTED = new CountDownLatch(1);
+ RELEASE_A = new CountDownLatch(1);
+ RELEASE_B = new CountDownLatch(1);
+ }
+
+ private static void releaseAll() {
+ RELEASE_A.countDown();
+ RELEASE_B.countDown();
+ }
+
+ public static void drainActionA(Event event, RunnerContext context) throws Exception {
+ runAction(event, context, "A", A_IN_ASYNC, RELEASE_A);
+ }
+
+ public static void drainActionB(Event event, RunnerContext context) throws Exception {
+ runAction(event, context, "B", B_IN_ASYNC, RELEASE_B);
+ }
+
+ public static void drainActionC(Event event, RunnerContext context) throws Exception {
+ runAction(event, context, "C", C_STARTED, null);
+ }
+
+ private static void runAction(
+ Event event,
+ RunnerContext context,
+ String name,
+ CountDownLatch started,
+ CountDownLatch release)
+ throws Exception {
+ Long input = (Long) InputEvent.fromEvent(event).getInput();
+ String actionId = "drain-action-" + name;
+ context.durableExecuteAsync(
+ new DurableCallable<String>() {
+ @Override
+ public String getId() {
+ return actionId + "-" + input;
+ }
+
+ @Override
+ public Class<String> getResultClass() {
+ return String.class;
+ }
+
+ @Override
+ public String call() throws Exception {
+ started.countDown();
+ if (release != null
+ && !release.await(
+ SETUP_TIMEOUT_SECONDS * 6, TimeUnit.SECONDS)) {
+ throw new IllegalStateException(
+ "Timed out waiting for the release of action "
+ + name);
+ }
+ return actionId;
+ }
+ })
+ .await();
+ context.sendEvent(new OutputEvent(actionId));
+ }
+
+ private static AgentPlan getAgentPlan(AgentConfiguration config) throws Exception {
+ Map<String, Action> actions = new LinkedHashMap<>();
+ for (String name : new String[] {"A", "B", "C"}) {
+ String actionName = "drainAction" + name;
+ Action action =
+ new Action(
+ actionName,
+ new JavaFunction(
+ DrainAgent.class,
+ actionName,
+ new Class<?>[] {Event.class, RunnerContext.class}),
+ Collections.singletonList(InputEvent.EVENT_TYPE));
+ actions.put(actionName, action);
+ }
+ return new AgentPlan(actions, new LinkedHashMap<>(), config);
+ }
+ }
+}
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorParallelMetricsTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorParallelMetricsTest.java
new file mode 100644
index 00000000..059af779
--- /dev/null
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorParallelMetricsTest.java
@@ -0,0 +1,235 @@
+/*
+ * 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.operator;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.api.InputEvent;
+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.plan.AgentConfiguration;
+import org.apache.flink.agents.plan.AgentPlan;
+import org.apache.flink.agents.plan.JavaFunction;
+import org.apache.flink.agents.plan.actions.Action;
+import org.apache.flink.agents.runtime.async.ContinuationActionExecutor;
+import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl;
+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.Test;
+import org.junit.jupiter.api.Timeout;
+
+import java.lang.reflect.Field;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assumptions.assumeThat;
+
+/** Component-level tests for metrics bookkeeping on the JDK<21 parallel execution path. */
+class ActionExecutionOperatorParallelMetricsTest {
+
+ private static final int NUM_ASYNC_THREADS = 2;
+ private static final int NUM_ACTIONS = 2;
+
+ @Test
+ @Timeout(20)
+ void completedParallelActionsClearPendingAndSchedulingMetrics() throws Exception {
+ assumeThat(ContinuationActionExecutor.isContinuationSupported()).isFalse();
+ MetricsAgent.reset();
+
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ newHarness(MetricsAgent.getAgentPlan())) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator = operator(testHarness);
+
+ testHarness.processElement(new StreamRecord<>(0L));
+ assertThat(MetricsAgent.ACTIONS_STARTED.await(5, TimeUnit.SECONDS)).isTrue();
+ MetricsAgent.ALLOW_ACTIONS.countDown();
+ operator.waitInFlightEventsFinished();
+
+ FlinkAgentsMetricGroupImpl metricGroup = metricGroup(operator);
+ long pendingTasks = 0L;
+ long schedulingSamples = 0L;
+ for (int i = 1; i <= NUM_ACTIONS; i++) {
+ FlinkAgentsMetricGroupImpl actionMetrics =
+ metricGroup.getSubGroup("action", "metricAction" + i);
+ pendingTasks += gauge(actionMetrics, "numOfPendingActionTasks");
+ schedulingSamples +=
+ actionMetrics.getHistogram("actionSchedulingLatencyMs").getCount();
+ }
+
+ assertThat(List.of(pendingTasks, schedulingSamples))
+ .as("dequeue must clear pending tasks and consume their enqueue timestamps")
+ .containsExactly(0L, (long) NUM_ACTIONS);
+ } finally {
+ MetricsAgent.ALLOW_ACTIONS.countDown();
+ }
+ }
+
+ @Test
+ @Timeout(20)
+ void completedQueuedInputClearsPendingInputMetric() throws Exception {
+ assumeThat(ContinuationActionExecutor.isContinuationSupported()).isFalse();
+ MetricsAgent.reset();
+
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ newHarness(MetricsAgent.getAgentPlan())) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator = operator(testHarness);
+
+ testHarness.processElement(new StreamRecord<>(0L));
+ assertThat(MetricsAgent.ACTIONS_STARTED.await(5, TimeUnit.SECONDS)).isTrue();
+ testHarness.processElement(new StreamRecord<>(0L));
+
+ FlinkAgentsMetricGroupImpl metricGroup = metricGroup(operator);
+ assertThat(gauge(metricGroup, "numOfPendingInputEvents")).isEqualTo(1L);
+
+ MetricsAgent.ALLOW_ACTIONS.countDown();
+ operator.waitInFlightEventsFinished();
+
+ assertThat(gauge(metricGroup, "numOfPendingInputEvents"))
+ .as("consuming a queued input must decrement its pending gauge")
+ .isZero();
+ } finally {
+ MetricsAgent.ALLOW_ACTIONS.countDown();
+ }
+ }
+
+ @Test
+ @Timeout(20)
+ void internalNoopTaskDoesNotParticipateInActionMetrics() throws Exception {
+ assumeThat(ContinuationActionExecutor.isContinuationSupported()).isFalse();
+
+ AgentConfiguration config = parallelConfig();
+ AgentPlan noActionPlan =
+ new AgentPlan(new LinkedHashMap<>(), new LinkedHashMap<>(), config);
+
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ newHarness(noActionPlan)) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator = operator(testHarness);
+
+ testHarness.processElement(new StreamRecord<>(0L));
+ operator.waitInFlightEventsFinished();
+
+ assertThat(testHarness.getRecordOutput()).isEmpty();
+ assertThat(gauge(metricGroup(operator), "numOfPendingInputEvents")).isZero();
+ }
+ }
+
+ private static KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> newHarness(
+ AgentPlan agentPlan) throws Exception {
+ return new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory<>(agentPlan, true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class));
+ }
+
+ @SuppressWarnings("unchecked")
+ private static ActionExecutionOperator<Long, Object> operator(
+ KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness) {
+ return (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+ }
+
+ private static FlinkAgentsMetricGroupImpl metricGroup(
+ ActionExecutionOperator<Long, Object> operator) throws Exception {
+ Field field = ActionExecutionOperator.class.getDeclaredField("metricGroup");
+ field.setAccessible(true);
+ return (FlinkAgentsMetricGroupImpl) field.get(operator);
+ }
+
+ private static long gauge(FlinkAgentsMetricGroupImpl metricGroup, String name) {
+ return ((Number) metricGroup.getGauge(name).getValue()).longValue();
+ }
+
+ private static AgentConfiguration parallelConfig() {
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.NUM_ASYNC_THREADS, NUM_ASYNC_THREADS);
+ return config;
+ }
+
+ public static class MetricsAgent {
+
+ private static volatile CountDownLatch ACTIONS_STARTED = new CountDownLatch(0);
+ private static volatile CountDownLatch ALLOW_ACTIONS = new CountDownLatch(0);
+
+ private static void reset() {
+ ACTIONS_STARTED = new CountDownLatch(NUM_ACTIONS);
+ ALLOW_ACTIONS = new CountDownLatch(1);
+ }
+
+ private static void runAction(Event event, RunnerContext context, int actionNumber)
+ throws Exception {
+ Long input = (Long) InputEvent.fromEvent(event).getInput();
+ context.durableExecuteAsync(
+ new DurableCallable<String>() {
+ @Override
+ public String getId() {
+ return "metric-action-" + actionNumber + "-" + input;
+ }
+
+ @Override
+ public Class<String> getResultClass() {
+ return String.class;
+ }
+
+ @Override
+ public String call() throws Exception {
+ ACTIONS_STARTED.countDown();
+ if (!ALLOW_ACTIONS.await(5, TimeUnit.SECONDS)) {
+ throw new IllegalStateException(
+ "Timed out waiting to release metrics actions");
+ }
+ return "done";
+ }
+ })
+ .await();
+ }
+
+ public static void metricAction1(Event event, RunnerContext context) throws Exception {
+ runAction(event, context, 1);
+ }
+
+ public static void metricAction2(Event event, RunnerContext context) throws Exception {
+ runAction(event, context, 2);
+ }
+
+ private static AgentPlan getAgentPlan() throws Exception {
+ Map<String, Action> actions = new LinkedHashMap<>();
+ for (int i = 1; i <= NUM_ACTIONS; i++) {
+ String actionName = "metricAction" + i;
+ actions.put(
+ actionName,
+ new Action(
+ actionName,
+ new JavaFunction(
+ MetricsAgent.class,
+ actionName,
+ new Class<?>[] {Event.class, RunnerContext.class}),
+ Collections.singletonList(InputEvent.EVENT_TYPE)));
+ }
+ return new AgentPlan(actions, new LinkedHashMap<>(), parallelConfig());
+ }
+ }
+}
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java
index 62167c46..0488b263 100644
--- a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java
@@ -23,6 +23,7 @@ import org.apache.flink.agents.api.Event;
import org.apache.flink.agents.api.EventContext;
import org.apache.flink.agents.api.InputEvent;
import org.apache.flink.agents.api.OutputEvent;
+import org.apache.flink.agents.api.agents.AgentExecutionOptions;
import org.apache.flink.agents.api.configuration.AgentConfigOptions;
import org.apache.flink.agents.api.context.DurableCallable;
import org.apache.flink.agents.api.context.MemoryObject;
@@ -37,6 +38,8 @@ import org.apache.flink.agents.api.logger.EventLoggerFactory;
import org.apache.flink.agents.api.logger.EventLoggerOpenParams;
import org.apache.flink.agents.api.logger.LoggerType;
import org.apache.flink.agents.api.memory.MemorySet;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
+import org.apache.flink.agents.api.resource.ResourceName;
import org.apache.flink.agents.api.resource.ResourceType;
import org.apache.flink.agents.api.trace.ExecutionLifecycleEvents;
import org.apache.flink.agents.api.trace.ExecutionReporter;
@@ -44,9 +47,12 @@ import org.apache.flink.agents.api.trace.ExecutionTraceContext;
import org.apache.flink.agents.plan.AgentConfiguration;
import org.apache.flink.agents.plan.AgentPlan;
import org.apache.flink.agents.plan.JavaFunction;
+import org.apache.flink.agents.plan.PythonFunction;
import org.apache.flink.agents.plan.actions.Action;
import org.apache.flink.agents.plan.actions.ToolCallAction;
+import org.apache.flink.agents.plan.actions.Utils;
import org.apache.flink.agents.plan.resourceprovider.JavaSerializableResourceProvider;
+import org.apache.flink.agents.plan.resourceprovider.PythonResourceProvider;
import org.apache.flink.agents.plan.resourceprovider.ResourceProvider;
import org.apache.flink.agents.plan.tools.FunctionTool;
import org.apache.flink.agents.runtime.ResourceCache;
@@ -57,6 +63,7 @@ import org.apache.flink.agents.runtime.actionstate.ActionStateUtil;
import org.apache.flink.agents.runtime.actionstate.CallResult;
import org.apache.flink.agents.runtime.actionstate.InMemoryActionStateStore;
import org.apache.flink.agents.runtime.actionstate.KafkaActionStateStore;
+import org.apache.flink.agents.runtime.async.ContinuationActionExecutor;
import org.apache.flink.agents.runtime.eventlog.EventLogWriter;
import org.apache.flink.agents.runtime.eventlog.FileEventLogger;
import org.apache.flink.agents.runtime.eventlog.Slf4jEventLogger;
@@ -72,6 +79,7 @@ import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
import org.apache.flink.streaming.api.operators.StreamOperatorStateHandler;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.mailbox.Mail;
import org.apache.flink.streaming.runtime.tasks.mailbox.TaskMailbox;
import org.apache.flink.streaming.util.AbstractStreamOperatorTestHarness;
import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
@@ -79,8 +87,10 @@ import org.apache.flink.util.ExceptionUtils;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
import org.mockito.InOrder;
import org.mockito.MockedConstruction;
+import org.mockito.MockedStatic;
import java.io.IOException;
import java.io.Serializable;
@@ -93,9 +103,12 @@ import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.Callable;
+import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.IntPredicate;
import java.util.stream.Collectors;
@@ -103,10 +116,12 @@ import java.util.stream.Collectors;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.assertj.core.api.Assertions.catchThrowable;
+import static org.junit.jupiter.api.Assumptions.assumeFalse;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.inOrder;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockConstruction;
+import static org.mockito.Mockito.mockStatic;
/** Tests for {@link ActionExecutionOperator}. */
public class ActionExecutionOperatorTest {
@@ -252,6 +267,104 @@ public class ActionExecutionOperatorTest {
}
}
+ @Test
+ void shouldEnableParallelExecutionWithoutCoroutineOnlyForJdk11AllJavaAgents() throws Exception {
+ AgentPlan allJava = TestAgent.getAgentPlan(false);
+ AgentPlan pythonOnly = TestAgent.getPythonOnlyAgentPlan();
+ AgentPlan mixed = TestAgent.getMixedJavaPythonAgentPlan();
+
+ assertThat(
+ ActionExecutionOperator.shouldEnableParallelExecutionWithoutCoroutine(
+ allJava, false))
+ .isTrue();
+ assertThat(
+ ActionExecutionOperator.shouldEnableParallelExecutionWithoutCoroutine(
+ allJava, true))
+ .isFalse();
+ assertThat(
+ ActionExecutionOperator.shouldEnableParallelExecutionWithoutCoroutine(
+ pythonOnly, false))
+ .isFalse();
+ assertThat(
+ ActionExecutionOperator.shouldEnableParallelExecutionWithoutCoroutine(
+ mixed, false))
+ .isFalse();
+
+ // The explicit fallback switch turns the engine off even for eligible plans.
+ AgentConfiguration disabled = new AgentConfiguration();
+ disabled.set(AgentExecutionOptions.PARALLEL_EXECUTION_ENABLED, false);
+ AgentPlan allJavaDisabled = TestAgent.getSingleAsyncAgentPlan(disabled);
+ assertThat(
+ ActionExecutionOperator.shouldEnableParallelExecutionWithoutCoroutine(
+ allJavaDisabled, false))
+ .isFalse();
+ }
+
+ @Test
+ void pythonResourcesRequireAsyncCompatibleFlinkForParallelExecution() throws Exception {
+ AgentPlan allJava = TestAgent.getAgentPlan(false);
+ PythonResourceProvider pythonChatModel =
+ new PythonResourceProvider(
+ "python-chat-model",
+ ResourceType.CHAT_MODEL,
+ new ResourceDescriptor("test.module", "TestChatModel", new HashMap<>()));
+ Map<String, ResourceProvider> chatModels = new HashMap<>();
+ chatModels.put("python-chat-model", pythonChatModel);
+ Map<ResourceType, Map<String, ResourceProvider>> resourceProviders = new HashMap<>();
+ resourceProviders.put(ResourceType.CHAT_MODEL, chatModels);
+ AgentPlan planWithPythonResource =
+ new AgentPlan(allJava.getActions(), resourceProviders, allJava.getConfig());
+
+ try (MockedStatic<Utils> utils = mockStatic(Utils.class)) {
+ utils.when(Utils::supportAsync).thenReturn(false);
+ assertThat(
+ ActionExecutionOperator.shouldEnableParallelExecutionWithoutCoroutine(
+ allJava, false))
+ .isTrue();
+ assertThat(
+ ActionExecutionOperator.shouldEnableParallelExecutionWithoutCoroutine(
+ planWithPythonResource, false))
+ .isFalse();
+
+ utils.when(Utils::supportAsync).thenReturn(true);
+ assertThat(
+ ActionExecutionOperator.shouldEnableParallelExecutionWithoutCoroutine(
+ planWithPythonResource, false))
+ .isTrue();
+ }
+ }
+
+ @Test
+ void pythonChromaResourceCanUseParallelExecutionWithAsyncCompatibleFlink() throws Exception {
+ AgentPlan allJava = TestAgent.getAgentPlan(false);
+ PythonResourceProvider pythonChromaVectorStore =
+ new PythonResourceProvider(
+ "python-chroma-vector-store",
+ ResourceType.VECTOR_STORE,
+ new ResourceDescriptor(
+ "test.module",
+ "TestVectorStore",
+ new HashMap<>(
+ Map.of(
+ "pythonClazz",
+ ResourceName.VectorStore.Python
+ .CHROMA_VECTOR_STORE))));
+ Map<String, ResourceProvider> vectorStores = new HashMap<>();
+ vectorStores.put("python-chroma-vector-store", pythonChromaVectorStore);
+ Map<ResourceType, Map<String, ResourceProvider>> resourceProviders = new HashMap<>();
+ resourceProviders.put(ResourceType.VECTOR_STORE, vectorStores);
+ AgentPlan planWithPythonChroma =
+ new AgentPlan(allJava.getActions(), resourceProviders, allJava.getConfig());
+
+ try (MockedStatic<Utils> utils = mockStatic(Utils.class)) {
+ utils.when(Utils::supportAsync).thenReturn(true);
+ assertThat(
+ ActionExecutionOperator.shouldEnableParallelExecutionWithoutCoroutine(
+ planWithPythonChroma, false))
+ .isTrue();
+ }
+ }
+
@Test
void testSameKeyDataAreProcessedInOrder() throws Exception {
try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
@@ -267,25 +380,15 @@ public class ActionExecutionOperatorTest {
testHarness.processElement(new StreamRecord<>(0L));
// Process input data 2, which has the same key (0)
testHarness.processElement(new StreamRecord<>(0L));
- // Since both pieces of data share the same key, we should consolidate them and process
- // only input data 1.
- // This means we need one mail to execute the action1 action for input data 1.
- assertMailboxSizeAndRun(testHarness.getTaskMailbox(), 1);
- // After executing this mail, we will have another mail to execute the action2 action
- // for input data 1.
- assertMailboxSizeAndRun(testHarness.getTaskMailbox(), 1);
- // Once the above mails are executed, we should get a single output result from input
- // data 1.
+ // Since both pieces of data share the same key, we should consolidate them: input
+ // data 2 is queued behind input data 1. Resident workers run the tasks and the
+ // worker-completion path drains the same-key queue in order; wait for it to finish.
+ operator.waitInFlightEventsFinished();
+ // Once the queued inputs finish in order, we should get two output results.
List<StreamRecord<Object>> recordOutput =
(List<StreamRecord<Object>>) testHarness.getRecordOutput();
- assertThat(recordOutput.size()).isEqualTo(1);
- assertThat(recordOutput.get(0).getValue()).isEqualTo(2L);
-
- // After the processing of input data 1 is finished, we can proceed to process input
- // data 2 and obtain its result.
- operator.waitInFlightEventsFinished();
- recordOutput = (List<StreamRecord<Object>>) testHarness.getRecordOutput();
assertThat(recordOutput.size()).isEqualTo(2);
+ assertThat(recordOutput.get(0).getValue()).isEqualTo(2L);
assertThat(recordOutput.get(1).getValue()).isEqualTo(2L);
}
}
@@ -303,16 +406,14 @@ public class ActionExecutionOperatorTest {
testHarness.processElement(new StreamRecord<>(0L));
// Process input data 2, which has the different key (1)
testHarness.processElement(new StreamRecord<>(1L));
- // Since the two input data items have different keys, they can be processed in
- // parallel.
- // As a result, we should have two separate mails to execute the action1 for each of
- // them.
- assertMailboxSizeAndRun(testHarness.getTaskMailbox(), 2);
- // After these two mails are executed, there should be another two mails — one for each
- // input data item — to execute the corresponding action2.
- assertMailboxSizeAndRun(testHarness.getTaskMailbox(), 2);
- // Once both action2 operations are completed, we should receive two output data items,
- // each corresponding to one of the original inputs.
+ // Since the two input data items have different keys, they can be processed in parallel
+ // by resident workers. Worker completion is asynchronous, so wait for the execution to
+ // drain instead of assuming both completions are queued at the same instant.
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+ operator.waitInFlightEventsFinished();
+ // Once both action chains are completed, we should receive two output data items, each
+ // corresponding to one of the original inputs.
List<StreamRecord<Object>> recordOutput =
(List<StreamRecord<Object>>) testHarness.getRecordOutput();
assertThat(recordOutput.size()).isEqualTo(2);
@@ -321,6 +422,617 @@ public class ActionExecutionOperatorTest {
}
}
+ @Test
+ @Timeout(15)
+ void sameKeySiblingActionsEnterAsyncStageTogether() throws Exception {
+ assumeFalse(ContinuationActionExecutor.isContinuationSupported());
+ TestAgent.resetSiblingParallelFixture();
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.NUM_ASYNC_THREADS, 2);
+
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(
+ TestAgent.getSiblingParallelAgentPlan(config), true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ testHarness.processElement(new StreamRecord<>(0L));
+ // Resident workers pull and run both sibling tasks autonomously once the input is
+ // admitted; wait until both have entered their async stage.
+
+ try {
+ assertThat(TestAgent.BOTH_SIBLING_ASYNC_STARTED.await(5, TimeUnit.SECONDS))
+ .isTrue();
+ } finally {
+ TestAgent.SIBLING_ACTION_A_ALLOW.countDown();
+ TestAgent.SIBLING_ACTION_B_ALLOW.countDown();
+ ((ActionExecutionOperator<Long, Object>) testHarness.getOperator())
+ .waitInFlightEventsFinished();
+ }
+ }
+ }
+
+ @Test
+ @Timeout(15)
+ void sameKeySiblingResultsCommitInSubmissionOrder() throws Exception {
+ assumeFalse(ContinuationActionExecutor.isContinuationSupported());
+ TestAgent.resetSiblingParallelFixture();
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.NUM_ASYNC_THREADS, 2);
+
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(
+ TestAgent.getSiblingParallelAgentPlan(config), true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+ testHarness.processElement(new StreamRecord<>(0L));
+ // Resident workers pull and run both sibling tasks autonomously; wait until both have
+ // entered their async stage.
+ assertThat(TestAgent.BOTH_SIBLING_ASYNC_STARTED.await(5, TimeUnit.SECONDS)).isTrue();
+
+ TestAgent.SIBLING_ACTION_B_ALLOW.countDown();
+ assertThat(TestAgent.SIBLING_B_ACTION_RETURNING.await(5, TimeUnit.SECONDS)).isTrue();
+ testHarness.getTaskMailbox().take(TaskMailbox.MIN_PRIORITY).run();
+ assertThat(testHarness.getRecordOutput()).isEmpty();
+
+ TestAgent.SIBLING_ACTION_A_ALLOW.countDown();
+ operator.waitInFlightEventsFinished();
+ List<StreamRecord<Object>> output =
+ (List<StreamRecord<Object>>) testHarness.getRecordOutput();
+ assertThat(output)
+ .extracting(StreamRecord::getValue)
+ .containsExactly("sibling-A", "sibling-B");
+ } finally {
+ TestAgent.SIBLING_ACTION_A_ALLOW.countDown();
+ TestAgent.SIBLING_ACTION_B_ALLOW.countDown();
+ }
+ }
+
+ @Test
+ @Timeout(15)
+ void staleCompletionMailAfterFullDrainIsNoOp() throws Exception {
+ assumeFalse(ContinuationActionExecutor.isContinuationSupported());
+ TestAgent.resetSiblingParallelFixture();
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.NUM_ASYNC_THREADS, 2);
+
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(
+ TestAgent.getSiblingParallelAgentPlan(config), true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ testHarness.processElement(new StreamRecord<>(0L));
+ assertThat(TestAgent.BOTH_SIBLING_ASYNC_STARTED.await(5, TimeUnit.SECONDS)).isTrue();
+
+ // Release both siblings: each worker marks its node done and dispatches its own
+ // completion mail. A worker's releaseToMain reserves the lock for the mailbox, so
+ // interleave mailbox acquire/release cycles (processWatermark) to let the other worker
+ // re-acquire and finish — without consuming the queued completion mails. Wait until
+ // both mails are queued; at that point both nodes are DONE (markDone precedes the
+ // dispatch).
+ TestAgent.SIBLING_ACTION_A_ALLOW.countDown();
+ TestAgent.SIBLING_ACTION_B_ALLOW.countDown();
+ long watermark = 0;
+ long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
+ while (testHarness.getTaskMailbox().size() < 2) {
+ assertThat(System.nanoTime() < deadline)
+ .as("timed out waiting for both completion mails")
+ .isTrue();
+ testHarness.processWatermark(new Watermark(watermark++));
+ Thread.sleep(5);
+ }
+
+ // The first completion mail finds both nodes DONE, drains and commits both in
+ // taskIndex order, drops the key group, and retires the input.
+ testHarness.getTaskMailbox().take(TaskMailbox.MIN_PRIORITY).run();
+ List<StreamRecord<Object>> output =
+ (List<StreamRecord<Object>>) testHarness.getRecordOutput();
+ assertThat(output)
+ .extracting(StreamRecord::getValue)
+ .containsExactly("sibling-A", "sibling-B");
+
+ // The second (stale) completion mail sees no group and commits nothing; it must be a
+ // no-op rather than attempt to retire the already-retired input.
+ testHarness.getTaskMailbox().take(TaskMailbox.MIN_PRIORITY).run();
+
+ assertThat(
+ ((ActionExecutionOperator<Long, Object>) testHarness.getOperator())
+ .getOperatorStateManager()
+ .getProcessingKeys())
+ .isEmpty();
+ } finally {
+ TestAgent.SIBLING_ACTION_A_ALLOW.countDown();
+ TestAgent.SIBLING_ACTION_B_ALLOW.countDown();
+ }
+ }
+
+ @Test
+ void openRejectsNonPositiveNumAsyncThreads() throws Exception {
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.NUM_ASYNC_THREADS, 0);
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(
+ TestAgent.getAgentPlanWithConfig(config, false), true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ assertThatThrownBy(testHarness::open)
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("num-async-threads");
+ }
+ }
+
+ @Test
+ void openRejectsNonPositiveMaxInFlightInputRecords() throws Exception {
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.MAX_IN_FLIGHT_INPUT_RECORDS, 0);
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(
+ TestAgent.getAgentPlanWithConfig(config, false), true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ assertThatThrownBy(testHarness::open)
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("max-in-flight-input-records");
+ }
+ }
+
+ @Test
+ @Timeout(10)
+ void parallelEngineRetiresInputWithoutTriggeredActions() throws Exception {
+ assumeFalse(ContinuationActionExecutor.isContinuationSupported());
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.NUM_ASYNC_THREADS, 1);
+ config.set(AgentExecutionOptions.MAX_IN_FLIGHT_INPUT_RECORDS, 1);
+ AgentPlan agentPlan = new AgentPlan(new HashMap<>(), new HashMap<>(), config);
+
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(agentPlan, true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+
+ testHarness.processElement(new StreamRecord<>(1L));
+ operator.waitInFlightEventsFinished();
+
+ assertThat(operator.getOperatorStateManager().getProcessingKeys()).isEmpty();
+ assertThat(testHarness.getRecordOutput()).isEmpty();
+
+ testHarness.processElement(new StreamRecord<>(2L));
+ operator.waitInFlightEventsFinished();
+ assertThat(operator.getOperatorStateManager().getProcessingKeys()).isEmpty();
+ assertThat(testHarness.getRecordOutput()).isEmpty();
+ }
+ }
+
+ @Test
+ @Timeout(30)
+ void backpressureBlocksAdmissionAtCapAndReleasesOnRetirement() throws Exception {
+ assumeFalse(ContinuationActionExecutor.isContinuationSupported());
+ TestAgent.resetSingleAsyncFixture();
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.NUM_ASYNC_THREADS, 2);
+ config.set(AgentExecutionOptions.MAX_IN_FLIGHT_INPUT_RECORDS, 2);
+
+ Thread releaser = null;
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(
+ TestAgent.getSingleAsyncAgentPlan(config), true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+
+ // Fill the in-flight cap with two records whose actions park in the async stage.
+ testHarness.processElement(new StreamRecord<>(1L));
+ testHarness.processElement(new StreamRecord<>(2L));
+ assertThat(TestAgent.SINGLE_ASYNC_STARTED.await(5, TimeUnit.SECONDS)).isTrue();
+
+ AtomicBoolean retirementReleased = new AtomicBoolean();
+ releaser =
+ new Thread(
+ () -> {
+ try {
+ Thread.sleep(500);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ retirementReleased.set(true);
+ TestAgent.SINGLE_ASYNC_ALLOW.countDown();
+ });
+ releaser.setDaemon(true);
+ releaser.start();
+
+ // At the cap, admission must block inside processElement until a record retires.
+ testHarness.processElement(new StreamRecord<>(3L));
+ assertThat(retirementReleased.get()).isTrue();
+
+ operator.waitInFlightEventsFinished();
+ List<StreamRecord<Object>> output =
+ (List<StreamRecord<Object>>) testHarness.getRecordOutput();
+ assertThat(output)
+ .extracting(StreamRecord::getValue)
+ .containsExactlyInAnyOrder(2L, 4L, 6L);
+ } finally {
+ TestAgent.SINGLE_ASYNC_ALLOW.countDown();
+ if (releaser != null) {
+ releaser.join();
+ }
+ }
+ }
+
+ @Test
+ @Timeout(30)
+ void restoreRebuildsInFlightBudgetFromProcessingKeysAndPendingEvents() throws Exception {
+ assumeFalse(ContinuationActionExecutor.isContinuationSupported());
+ TestAgent.resetSingleAsyncFixture();
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.NUM_ASYNC_THREADS, 1);
+ config.set(AgentExecutionOptions.MAX_IN_FLIGHT_INPUT_RECORDS, 3);
+
+ // Phase 1: key 1 is pulled and parks in async; key 2 queues one task (single worker is
+ // busy) plus one pending input event (busy key). The drain retires key 1, so the snapshot
+ // captures exactly one processing key and one pending event.
+ OperatorSubtaskState snapshot;
+ Thread releaser = null;
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(
+ TestAgent.getSingleAsyncAgentPlan(config), true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+
+ testHarness.processElement(new StreamRecord<>(1L));
+ assertThat(TestAgent.SINGLE_ASYNC_STARTED.await(5, TimeUnit.SECONDS)).isTrue();
+ testHarness.processElement(new StreamRecord<>(2L));
+ testHarness.processElement(new StreamRecord<>(2L));
+
+ releaser =
+ new Thread(
+ () -> {
+ try {
+ Thread.sleep(300);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ TestAgent.SINGLE_ASYNC_ALLOW.countDown();
+ });
+ releaser.setDaemon(true);
+ releaser.start();
+ operator.prepareSnapshotPreBarrier(1L);
+ assertThat(operator.getOperatorStateManager().getProcessingKeys()).containsExactly(2L);
+ snapshot = testHarness.snapshot(1L, 1L);
+ } finally {
+ TestAgent.SINGLE_ASYNC_ALLOW.countDown();
+ if (releaser != null) {
+ releaser.join();
+ }
+ }
+
+ // Phase 2: restore with cap 2. The rebuilt budget must count the processing key AND the
+ // pending event (2 units): a new record then blocks until the restored work retires. If
+ // either unit were dropped, admission would pass immediately and the assertion fails.
+ TestAgent.resetSingleAsyncFixture();
+ AgentConfiguration restoredConfig = new AgentConfiguration();
+ restoredConfig.set(AgentExecutionOptions.NUM_ASYNC_THREADS, 1);
+ restoredConfig.set(AgentExecutionOptions.MAX_IN_FLIGHT_INPUT_RECORDS, 2);
+ Thread budgetReleaser = null;
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> restored =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(
+ TestAgent.getSingleAsyncAgentPlan(restoredConfig), true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ restored.initializeState(snapshot);
+ restored.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) restored.getOperator();
+ // The restored queued task must be re-dispatched and reach the async stage.
+ assertThat(TestAgent.SINGLE_ASYNC_STARTED.await(5, TimeUnit.SECONDS)).isTrue();
+
+ AtomicBoolean budgetFreed = new AtomicBoolean();
+ budgetReleaser =
+ new Thread(
+ () -> {
+ try {
+ Thread.sleep(500);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ budgetFreed.set(true);
+ TestAgent.SINGLE_ASYNC_ALLOW.countDown();
+ });
+ budgetReleaser.setDaemon(true);
+ budgetReleaser.start();
+
+ restored.processElement(new StreamRecord<>(3L));
+ assertThat(budgetFreed.get()).isTrue();
+
+ operator.waitInFlightEventsFinished();
+ List<StreamRecord<Object>> output =
+ (List<StreamRecord<Object>>) restored.getRecordOutput();
+ assertThat(output)
+ .extracting(StreamRecord::getValue)
+ .containsExactlyInAnyOrder(4L, 4L, 6L);
+ } finally {
+ TestAgent.SINGLE_ASYNC_ALLOW.countDown();
+ if (budgetReleaser != null) {
+ budgetReleaser.join();
+ }
+ }
+ }
+
+ @Test
+ @Timeout(30)
+ void preBarrierFailureResetsDrainingSoDispatchIsNotWedged() throws Exception {
+ assumeFalse(ContinuationActionExecutor.isContinuationSupported());
+ TestAgent.resetSingleAsyncFixture();
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.NUM_ASYNC_THREADS, 1);
+
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(
+ TestAgent.getSingleAsyncAgentPlan(config), true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+
+ // First record is pulled and parks in the async stage; the second queues as a pending
+ // input for the busy key, so a task remains to be dispatched after the drain window.
+ testHarness.processElement(new StreamRecord<>(21L));
+ assertThat(TestAgent.SINGLE_ASYNC_STARTED.await(5, TimeUnit.SECONDS)).isTrue();
+ testHarness.processElement(new StreamRecord<>(21L));
+
+ // Poison the quiesce loop: its first yield runs this mail and prepareSnapshotPreBarrier
+ // fails after startDraining.
+ testHarness
+ .getTaskMailbox()
+ .put(
+ new Mail(
+ () -> {
+ throw new RuntimeException("boom");
+ },
+ 0,
+ "poison mail"));
+ assertThatThrownBy(() -> operator.prepareSnapshotPreBarrier(1L))
+ .hasMessageContaining("boom");
+
+ // Even if the failure is swallowed upstream, dispatch must not stay paused: releasing
+ // the async action must let both inputs finish without any further checkpoint call.
+ TestAgent.SINGLE_ASYNC_ALLOW.countDown();
+ operator.waitInFlightEventsFinished();
+ List<StreamRecord<Object>> output =
+ (List<StreamRecord<Object>>) testHarness.getRecordOutput();
+ assertThat(output).extracting(StreamRecord::getValue).containsExactly(42L, 42L);
+ } finally {
+ TestAgent.SINGLE_ASYNC_ALLOW.countDown();
+ }
+ }
+
+ @Test
+ @Timeout(30)
+ void checkpointDrainsInFlightAsyncActionSoOutputIsNotLostOnRestore() throws Exception {
+ assumeFalse(ContinuationActionExecutor.isContinuationSupported());
+ TestAgent.resetSingleAsyncFixture();
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.NUM_ASYNC_THREADS, 2);
+
+ Thread releaser = null;
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(
+ TestAgent.getSingleAsyncAgentPlan(config), true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+
+ testHarness.processElement(new StreamRecord<>(21L));
+ // A resident worker pulls the task (removing it from the durable queue) and enters the
+ // async stage; its result is not yet committed.
+ assertThat(TestAgent.SINGLE_ASYNC_STARTED.await(5, TimeUnit.SECONDS)).isTrue();
+
+ // Release the async shortly after, so a checkpoint that drains in-flight work can
+ // complete; the un-fixed engine takes the checkpoint immediately, before this release.
+ releaser =
+ new Thread(
+ () -> {
+ try {
+ Thread.sleep(300);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ TestAgent.SINGLE_ASYNC_ALLOW.countDown();
+ });
+ releaser.setDaemon(true);
+ releaser.start();
+
+ // StreamTask calls prepareSnapshotPreBarrier before broadcasting the barrier.
+ operator.prepareSnapshotPreBarrier(1L);
+ testHarness.snapshot(1L, 1L);
+
+ // Invariant: once the checkpoint is taken, the in-flight action must already have
+ // committed and emitted its output. Otherwise its output — whose input was consumed
+ // before the barrier — is permanently lost on restore. Fails on the un-fixed engine.
+ List<StreamRecord<Object>> output =
+ (List<StreamRecord<Object>>) testHarness.getRecordOutput();
+ assertThat(output).extracting(StreamRecord::getValue).containsExactly(42L);
+ } finally {
+ if (TestAgent.SINGLE_ASYNC_ALLOW != null) {
+ TestAgent.SINGLE_ASYNC_ALLOW.countDown();
+ }
+ if (releaser != null) {
+ releaser.join();
+ }
+ }
+ }
+
+ @Test
+ @Timeout(30)
+ void checkpointQuiesceSnapshotAndRestoreDoesNotReplayTransientSchedule() throws Exception {
+ assumeFalse(ContinuationActionExecutor.isContinuationSupported());
+ TestAgent.resetSingleAsyncFixture();
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.NUM_ASYNC_THREADS, 2);
+
+ OperatorSubtaskState snapshot;
+ Thread releaser = null;
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(
+ TestAgent.getSingleAsyncAgentPlan(config), true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+
+ testHarness.processElement(new StreamRecord<>(31L));
+ assertThat(TestAgent.SINGLE_ASYNC_STARTED.await(5, TimeUnit.SECONDS)).isTrue();
+
+ releaser =
+ new Thread(
+ () -> {
+ try {
+ Thread.sleep(300);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ TestAgent.SINGLE_ASYNC_ALLOW.countDown();
+ });
+ releaser.setDaemon(true);
+ releaser.start();
+
+ operator.prepareSnapshotPreBarrier(1L);
+ snapshot = testHarness.snapshot(1L, 1L);
+
+ List<StreamRecord<Object>> output =
+ (List<StreamRecord<Object>>) testHarness.getRecordOutput();
+ assertThat(output).extracting(StreamRecord::getValue).containsExactly(62L);
+ assertThat(operator.getOperatorStateManager().getProcessingKeys()).isEmpty();
+ } finally {
+ TestAgent.SINGLE_ASYNC_ALLOW.countDown();
+ if (releaser != null) {
+ releaser.join();
+ }
+ }
+
+ TestAgent.resetSingleAsyncFixture();
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> restored =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(
+ TestAgent.getSingleAsyncAgentPlan(config), true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ restored.initializeState(snapshot);
+ restored.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) restored.getOperator();
+
+ operator.waitInFlightEventsFinished();
+ assertThat(restored.getRecordOutput()).isEmpty();
+
+ TestAgent.SINGLE_ASYNC_ALLOW.countDown();
+ restored.processElement(new StreamRecord<>(32L));
+ operator.waitInFlightEventsFinished();
+
+ List<StreamRecord<Object>> output =
+ (List<StreamRecord<Object>>) restored.getRecordOutput();
+ assertThat(output).extracting(StreamRecord::getValue).containsExactly(64L);
+ } finally {
+ TestAgent.SINGLE_ASYNC_ALLOW.countDown();
+ }
+ }
+
+ @Test
+ @Timeout(15)
+ void defaultSingleActionPendingEventsDrainOncePerInput() throws Exception {
+ assumeFalse(ContinuationActionExecutor.isContinuationSupported());
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.NUM_ASYNC_THREADS, 2);
+
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(
+ TestAgent.getPendingEventsSingleActionPlan(config), true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+
+ testHarness.processElement(new StreamRecord<>(10L));
+ testHarness.processElement(new StreamRecord<>(10L));
+ operator.waitInFlightEventsFinished();
+
+ List<StreamRecord<Object>> output =
+ (List<StreamRecord<Object>>) testHarness.getRecordOutput();
+ assertThat(output)
+ .extracting(StreamRecord::getValue)
+ .containsExactly("single-10", "single-10");
+ }
+ }
+
+ @Test
+ @Timeout(15)
+ void sameKeySiblingPendingEventsStayBufferedUntilCommitDrainReachesTheirTask()
+ throws Exception {
+ assumeFalse(ContinuationActionExecutor.isContinuationSupported());
+ TestAgent.resetSiblingParallelFixture();
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.NUM_ASYNC_THREADS, 2);
+
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(
+ TestAgent.getSiblingParallelAgentPlan(config), true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+
+ testHarness.processElement(new StreamRecord<>(0L));
+ assertThat(TestAgent.BOTH_SIBLING_ASYNC_STARTED.await(5, TimeUnit.SECONDS)).isTrue();
+
+ TestAgent.SIBLING_ACTION_B_ALLOW.countDown();
+ assertThat(TestAgent.SIBLING_B_ACTION_RETURNING.await(5, TimeUnit.SECONDS)).isTrue();
+ testHarness.getTaskMailbox().take(TaskMailbox.MIN_PRIORITY).run();
+ assertThat(testHarness.getRecordOutput()).isEmpty();
+
+ TestAgent.SIBLING_ACTION_A_ALLOW.countDown();
+ operator.waitInFlightEventsFinished();
+ List<StreamRecord<Object>> output =
+ (List<StreamRecord<Object>>) testHarness.getRecordOutput();
+ assertThat(output)
+ .extracting(StreamRecord::getValue)
+ .containsExactly("sibling-A", "sibling-B");
+ } finally {
+ TestAgent.SIBLING_ACTION_A_ALLOW.countDown();
+ TestAgent.SIBLING_ACTION_B_ALLOW.countDown();
+ }
+ }
+
@Test
void testRestoreOnlyResumesKeysOwnedByCurrentSubtask() throws Exception {
final int maxParallelism = 4;
@@ -339,7 +1051,12 @@ public class ActionExecutionOperatorTest {
0)) {
testHarness.open();
testHarness.processElement(new StreamRecord<>(key));
- assertThat(testHarness.getTaskMailbox().size()).isEqualTo(1);
+ // The input is now in flight (a processing key), so it is captured by the snapshot.
+ assertThat(
+ ((ActionExecutionOperator<Long, Object>) testHarness.getOperator())
+ .getOperatorStateManager()
+ .getProcessingKeys())
+ .containsExactly(key);
snapshot = testHarness.snapshot(1L, 1L);
}
@@ -381,8 +1098,6 @@ public class ActionExecutionOperatorTest {
ownerHarness.open();
nonOwnerHarness.open();
- assertThat(ownerHarness.getTaskMailbox().size()).isEqualTo(1);
- assertThat(nonOwnerHarness.getTaskMailbox().size()).isZero();
assertThat(
((ActionExecutionOperator<Long, Object>) ownerHarness.getOperator())
.getOperatorStateManager()
@@ -416,7 +1131,12 @@ public class ActionExecutionOperatorTest {
restoredOwnerHarness.initializeState(secondRestoreOwnerState);
restoredOwnerHarness.open();
- assertThat(restoredOwnerHarness.getTaskMailbox().size()).isEqualTo(1);
+ assertThat(
+ ((ActionExecutionOperator<Long, Object>)
+ restoredOwnerHarness.getOperator())
+ .getOperatorStateManager()
+ .getProcessingKeys())
+ .containsExactly(key);
}
}
}
@@ -433,9 +1153,15 @@ public class ActionExecutionOperatorTest {
new Class<?>[] {Event.class, RunnerContext.class}),
Collections.singletonList(InputEvent.EVENT_TYPE),
Collections.singletonMap("mode", "strict"));
+ // Force serial: the parallel engine drains all in-flight tasks before checkpoint, so no
+ // pending task survives the snapshot — this test exercises the serial restore contract.
+ AgentConfiguration serialConfig = new AgentConfiguration();
+ serialConfig.set(AgentExecutionOptions.PARALLEL_EXECUTION_ENABLED, false);
AgentPlan agentPlan =
new AgentPlan(
- Collections.singletonMap(configuredAction.getName(), configuredAction));
+ Collections.singletonMap(configuredAction.getName(), configuredAction),
+ new HashMap<>(),
+ serialConfig);
OperatorSubtaskState snapshot;
try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
@@ -445,7 +1171,12 @@ public class ActionExecutionOperatorTest {
TypeInformation.of(Long.class))) {
testHarness.open();
testHarness.processElement(new StreamRecord<>(key));
- assertThat(testHarness.getTaskMailbox().size()).isEqualTo(1);
+ // The input is now in flight (a processing key), so it is captured by the snapshot.
+ assertThat(
+ ((ActionExecutionOperator<Long, Object>) testHarness.getOperator())
+ .getOperatorStateManager()
+ .getProcessingKeys())
+ .containsExactly(key);
snapshot = testHarness.snapshot(1L, 1L);
}
@@ -470,7 +1201,14 @@ public class ActionExecutionOperatorTest {
@Test
void testRestoredActionUsesSameTextualContextKeyForLtmWriteAndCleanup() throws Exception {
+ // Force serial: the parallel engine drains in-flight tasks before checkpoint, so the
+ // failing action would execute and fail during drain rather than surviving to restore.
AgentPlan agentPlan = TestAgent.getFailedActionAfterLtmAgentPlan();
+ AgentConfiguration serialConfig = new AgentConfiguration();
+ serialConfig.set(AgentExecutionOptions.PARALLEL_EXECUTION_ENABLED, false);
+ agentPlan =
+ new AgentPlan(
+ agentPlan.getActions(), agentPlan.getResourceProviders(), serialConfig);
OperatorSubtaskState snapshot;
long key = 1L << 32;
String expectedContextKey = "4294967296";
@@ -482,7 +1220,12 @@ public class ActionExecutionOperatorTest {
TypeInformation.of(Long.class))) {
testHarness.open();
testHarness.processElement(new StreamRecord<>(key));
- assertThat(testHarness.getTaskMailbox().size()).isEqualTo(1);
+ // The input is now in flight (a processing key), so it is captured by the snapshot.
+ assertThat(
+ ((ActionExecutionOperator<Long, Object>) testHarness.getOperator())
+ .getOperatorStateManager()
+ .getProcessingKeys())
+ .containsExactly(key);
snapshot = testHarness.snapshot(1L, 1L);
}
@@ -523,7 +1266,7 @@ public class ActionExecutionOperatorTest {
assertThatThrownBy(() -> operator.waitInFlightEventsFinished())
.hasCauseInstanceOf(ActionExecutionOperator.ActionTaskExecutionException.class)
.rootCause()
- .hasMessageContaining("Expected to be running on the task mailbox thread");
+ .hasMessageContaining("Current thread does not own the lock");
}
}
@@ -673,10 +1416,13 @@ public class ActionExecutionOperatorTest {
@Test
void testFailedActionAfterLtmDiscardsCurrentKeyBeforeRethrowing() throws Exception {
RecordingMem0LongTermMemory ltm = new RecordingMem0LongTermMemory();
+ AgentPlan agentPlan = TestAgent.getFailedActionAfterLtmAgentPlan();
+ // Verify the serial drain path: the failing action discards its observation exactly once,
+ // and the following same-key action never executes.
+ agentPlan.getConfig().set(AgentExecutionOptions.PARALLEL_EXECUTION_ENABLED, false);
try (KeyedOneInputStreamOperatorTestHarness<String, Long, Object> testHarness =
new KeyedOneInputStreamOperatorTestHarness<>(
- new ActionExecutionOperatorFactory(
- TestAgent.getFailedActionAfterLtmAgentPlan(), true),
+ new ActionExecutionOperatorFactory(agentPlan, true),
(KeySelector<Long, String>) String::valueOf,
TypeInformation.of(String.class))) {
testHarness.open();
@@ -708,10 +1454,13 @@ public class ActionExecutionOperatorTest {
RuntimeException discardFailure = new RuntimeException("discard failed");
ltm.failDrainWith(discardFailure);
+ AgentPlan agentPlan = TestAgent.getFailedActionAfterLtmAgentPlan();
+ // Verify on the serial path that a discard failure is suppressed rather than replacing
+ // the original action failure.
+ agentPlan.getConfig().set(AgentExecutionOptions.PARALLEL_EXECUTION_ENABLED, false);
try (KeyedOneInputStreamOperatorTestHarness<String, Long, Object> testHarness =
new KeyedOneInputStreamOperatorTestHarness<>(
- new ActionExecutionOperatorFactory(
- TestAgent.getFailedActionAfterLtmAgentPlan(), true),
+ new ActionExecutionOperatorFactory(agentPlan, true),
(KeySelector<Long, String>) String::valueOf,
TypeInformation.of(String.class))) {
testHarness.open();
@@ -1523,7 +2272,8 @@ public class ActionExecutionOperatorTest {
actionStateStore.clearPruneCalls();
testHarness.processElement(new StreamRecord<>(5L));
- assertThat(testHarness.getTaskMailbox().size()).isEqualTo(1);
+ // The second input is now in flight (a processing key) but not yet completed.
+ assertThat(operator.getOperatorStateManager().getProcessingKeys()).containsExactly(5L);
testHarness.snapshot(1L, 1L);
testHarness.notifyOfCompletedCheckpoint(1L);
@@ -3081,6 +3831,14 @@ public class ActionExecutionOperatorTest {
NESTED_MEMORY_ACTION_CALL_COUNTER =
new java.util.concurrent.atomic.AtomicInteger(0);
+ public static volatile CountDownLatch BOTH_SIBLING_ASYNC_STARTED;
+ public static volatile CountDownLatch SIBLING_ACTION_A_ALLOW;
+ public static volatile CountDownLatch SIBLING_ACTION_B_ALLOW;
+ public static volatile CountDownLatch SIBLING_B_ACTION_RETURNING;
+
+ public static volatile CountDownLatch SINGLE_ASYNC_STARTED;
+ public static volatile CountDownLatch SINGLE_ASYNC_ALLOW;
+
public static class MiddleEvent extends Event {
public static final String EVENT_TYPE = "MiddleEvent";
@@ -3147,8 +3905,8 @@ public class ActionExecutionOperatorTest {
public static void action3(MiddleEvent event, RunnerContext context) {
// To test disallows memory access from non-mailbox threads.
+ ExecutorService executor = Executors.newSingleThreadExecutor();
try {
- ExecutorService executor = Executors.newSingleThreadExecutor();
Future<Long> future =
executor.submit(
() -> (Long) context.getShortTermMemory().get("tmp").getValue());
@@ -3156,6 +3914,8 @@ public class ActionExecutionOperatorTest {
context.sendEvent(new OutputEvent(tmp * 2));
} catch (Exception e) {
ExceptionUtils.rethrow(e);
+ } finally {
+ executor.shutdownNow();
}
}
@@ -3196,6 +3956,95 @@ public class ActionExecutionOperatorTest {
};
}
+ public static void siblingActionA(Event event, RunnerContext context) {
+ try {
+ runSiblingAction("sibling-A", event, context, SIBLING_ACTION_A_ALLOW, false);
+ } catch (Exception e) {
+ ExceptionUtils.rethrow(e);
+ }
+ }
+
+ public static void siblingActionB(Event event, RunnerContext context) {
+ try {
+ runSiblingAction("sibling-B", event, context, SIBLING_ACTION_B_ALLOW, true);
+ } catch (Exception e) {
+ ExceptionUtils.rethrow(e);
+ }
+ }
+
+ private static void runSiblingAction(
+ String actionId,
+ Event event,
+ RunnerContext context,
+ CountDownLatch allowLatch,
+ boolean recordReturn)
+ throws Exception {
+ Long inputData = (Long) InputEvent.fromEvent(event).getInput();
+ Long result =
+ context.durableExecuteAsync(
+ durableCallable(
+ actionId,
+ Long.class,
+ () -> {
+ BOTH_SIBLING_ASYNC_STARTED.countDown();
+ awaitLatch(allowLatch);
+ return inputData;
+ }))
+ .await();
+ context.sendEvent(new OutputEvent(actionId));
+ if (recordReturn) {
+ SIBLING_B_ACTION_RETURNING.countDown();
+ }
+ }
+
+ private static void awaitLatch(CountDownLatch latch) throws Exception {
+ if (!latch.await(5, TimeUnit.SECONDS)) {
+ throw new IllegalStateException("Timed out waiting for sibling action release");
+ }
+ }
+
+ public static void resetSiblingParallelFixture() {
+ BOTH_SIBLING_ASYNC_STARTED = new CountDownLatch(2);
+ SIBLING_ACTION_A_ALLOW = new CountDownLatch(1);
+ SIBLING_ACTION_B_ALLOW = new CountDownLatch(1);
+ SIBLING_B_ACTION_RETURNING = new CountDownLatch(1);
+ }
+
+ public static void resetSingleAsyncFixture() {
+ SINGLE_ASYNC_STARTED = new CountDownLatch(1);
+ SINGLE_ASYNC_ALLOW = new CountDownLatch(1);
+ }
+
+ /**
+ * Single latch-controlled async action: signals when it enters the async stage, then blocks
+ * until released, so a test can hold the task in-flight (pulled from the durable queue,
+ * result not yet committed) across a checkpoint.
+ */
+ public static void singleAsyncAction(Event event, RunnerContext context) {
+ Long inputData = (Long) InputEvent.fromEvent(event).getInput();
+ try {
+ Long result =
+ context.durableExecuteAsync(
+ durableCallable(
+ "single-async",
+ Long.class,
+ () -> {
+ SINGLE_ASYNC_STARTED.countDown();
+ awaitLatch(SINGLE_ASYNC_ALLOW);
+ return inputData * 2;
+ }))
+ .await();
+ context.sendEvent(new OutputEvent(result));
+ } catch (Exception e) {
+ ExceptionUtils.rethrow(e);
+ }
+ }
+
+ public static void pendingEventsSingleAction(Event event, RunnerContext context) {
+ Long inputData = (Long) InputEvent.fromEvent(event).getInput();
+ context.sendEvent(new OutputEvent("single-" + inputData));
+ }
+
private static <T> DurableCallable<T> reconcilableDurableCallable(
String id,
Class<T> resultClass,
@@ -3499,6 +4348,91 @@ public class ActionExecutionOperatorTest {
return null;
}
+ public static AgentPlan getSiblingParallelAgentPlan(AgentConfiguration config) {
+ try {
+ Action siblingA =
+ new Action(
+ "sibling-A",
+ new JavaFunction(
+ TestAgent.class,
+ "siblingActionA",
+ new Class<?>[] {Event.class, RunnerContext.class}),
+ Collections.singletonList(InputEvent.EVENT_TYPE));
+ Action siblingB =
+ new Action(
+ "sibling-B",
+ new JavaFunction(
+ TestAgent.class,
+ "siblingActionB",
+ new Class<?>[] {Event.class, RunnerContext.class}),
+ Collections.singletonList(InputEvent.EVENT_TYPE));
+ Map<String, Action> actions = new HashMap<>();
+ actions.put(siblingA.getName(), siblingA);
+ actions.put(siblingB.getName(), siblingB);
+ return new AgentPlan(actions, new HashMap<>(), config);
+ } catch (Exception e) {
+ ExceptionUtils.rethrow(e);
+ }
+ return null;
+ }
+
+ public static AgentPlan getSingleAsyncAgentPlan(AgentConfiguration config) {
+ try {
+ Action action =
+ new Action(
+ "single-async",
+ new JavaFunction(
+ TestAgent.class,
+ "singleAsyncAction",
+ new Class<?>[] {Event.class, RunnerContext.class}),
+ Collections.singletonList(InputEvent.EVENT_TYPE));
+ Map<String, Action> actions = new HashMap<>();
+ actions.put(action.getName(), action);
+ return new AgentPlan(actions, new HashMap<>(), config);
+ } catch (Exception e) {
+ ExceptionUtils.rethrow(e);
+ }
+ return null;
+ }
+
+ public static AgentPlan getPendingEventsSingleActionPlan(AgentConfiguration config) {
+ try {
+ Action action =
+ new Action(
+ "pending-events-single",
+ new JavaFunction(
+ TestAgent.class,
+ "pendingEventsSingleAction",
+ new Class<?>[] {Event.class, RunnerContext.class}),
+ Collections.singletonList(InputEvent.EVENT_TYPE));
+ Map<String, Action> actions = new HashMap<>();
+ actions.put(action.getName(), action);
+ return new AgentPlan(actions, new HashMap<>(), config);
+ } catch (Exception e) {
+ ExceptionUtils.rethrow(e);
+ }
+ return null;
+ }
+
+ public static AgentPlan getPythonOnlyAgentPlan() throws Exception {
+ Action pythonAction =
+ new Action(
+ "pythonAction",
+ new PythonFunction("module", "pythonAction"),
+ Collections.singletonList(InputEvent.EVENT_TYPE));
+ Map<String, Action> actions = new HashMap<>();
+ actions.put(pythonAction.getName(), pythonAction);
+ return new AgentPlan(actions, new HashMap<>());
+ }
+
+ public static AgentPlan getMixedJavaPythonAgentPlan() throws Exception {
+ AgentPlan javaPlan = getAgentPlan(false);
+ AgentPlan pythonPlan = getPythonOnlyAgentPlan();
+ Map<String, Action> actions = new HashMap<>(javaPlan.getActions());
+ actions.putAll(pythonPlan.getActions());
+ return new AgentPlan(actions, new HashMap<>());
+ }
+
public static AgentPlan getAgentPlanWithConfig(AgentConfiguration config) {
return getAgentPlanWithConfig(config, false);
}
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java
index d95ab33f..214a2be2 100644
--- a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java
@@ -20,6 +20,7 @@ package org.apache.flink.agents.runtime.operator;
import org.apache.flink.agents.api.Event;
import org.apache.flink.agents.api.EventContext;
import org.apache.flink.agents.api.InputEvent;
+import org.apache.flink.agents.api.OutputEvent;
import org.apache.flink.agents.api.trace.ExecutionLifecycleEvents;
import org.apache.flink.agents.api.trace.ExecutionReporter;
import org.apache.flink.agents.plan.AgentPlan;
@@ -70,7 +71,7 @@ class ActionTaskContextManagerTest {
*/
@Test
void closeClosesContinuationExecutorWhenRunnerContextFails() throws Exception {
- ActionTaskContextManager mgr = new ActionTaskContextManager(1);
+ ActionTaskContextManager mgr = newManager();
RunnerContextImpl failingContext = mock(RunnerContextImpl.class);
ContinuationActionExecutor continuationExecutor = mock(ContinuationActionExecutor.class);
doThrow(new IllegalStateException("runner context close failed"))
@@ -92,7 +93,7 @@ class ActionTaskContextManagerTest {
/** The first failure is rethrown and the later one is attached as suppressed, never dropped. */
@Test
void closeReportsFirstFailureWithLaterOneSuppressed() throws Exception {
- ActionTaskContextManager mgr = new ActionTaskContextManager(1);
+ ActionTaskContextManager mgr = newManager();
RunnerContextImpl failingContext = mock(RunnerContextImpl.class);
ContinuationActionExecutor failingExecutor = mock(ContinuationActionExecutor.class);
doThrow(new IllegalStateException("runner context close failed"))
@@ -124,7 +125,7 @@ class ActionTaskContextManagerTest {
@Test
void perTaskContextsAreIsolatedAcrossPutGetRemove() throws Exception {
- try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) {
+ try (ActionTaskContextManager mgr = newManager()) {
Action action = TestActions.noopAction();
ActionTask t1 = new JavaActionTask("k", new InputEvent(1L), action, 1L);
ActionTask t2 = new JavaActionTask("k", new InputEvent(2L), action, 1L);
@@ -156,7 +157,7 @@ class ActionTaskContextManagerTest {
@Test
void createOrGetRunnerContextThrowsWhenPythonContextRequestedButNull() throws Exception {
- try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) {
+ try (ActionTaskContextManager mgr = newManager()) {
assertThatThrownBy(
() ->
mgr.createOrGetRunnerContext(
@@ -175,7 +176,7 @@ class ActionTaskContextManagerTest {
@Test
void createAndSetRunnerContextBuildsFreshMemoryContextOnFirstCall() throws Exception {
- try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) {
+ try (ActionTaskContextManager mgr = newManager()) {
ActionTask t =
new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction(), 1L);
invokeCreateAndSetRunnerContext(mgr, t);
@@ -186,9 +187,79 @@ class ActionTaskContextManagerTest {
}
}
+ @Test
+ void createAndSetRunnerContextPinsFreshContextsAndRestoresTaskSpecificContexts()
+ throws Exception {
+ try (ActionTaskContextManager mgr = newManager()) {
+ Action action = TestActions.noopAction();
+ ActionTask taskA = new JavaActionTask("a", new InputEvent(1L), action, 1L);
+ ActionTask taskB = new JavaActionTask("b", new InputEvent(2L), action, 2L);
+
+ invokeCreateAndSetRunnerContext(mgr, taskA);
+ RunnerContextImpl.MemoryContext memoryA = taskA.getRunnerContext().getMemoryContext();
+ ContinuationContext continuationA =
+ ((JavaRunnerContextImpl) taskA.getRunnerContext()).getContinuationContext();
+
+ invokeCreateAndSetRunnerContext(mgr, taskB);
+ RunnerContextImpl.MemoryContext memoryB = taskB.getRunnerContext().getMemoryContext();
+ ContinuationContext continuationB =
+ ((JavaRunnerContextImpl) taskB.getRunnerContext()).getContinuationContext();
+
+ mgr.restore("a", taskA);
+ assertThat(taskA.getRunnerContext().getMemoryContext()).isSameAs(memoryA);
+ assertThat(((JavaRunnerContextImpl) taskA.getRunnerContext()).getContinuationContext())
+ .isSameAs(continuationA);
+
+ mgr.restore("b", taskB);
+ assertThat(taskB.getRunnerContext().getMemoryContext()).isSameAs(memoryB);
+ assertThat(((JavaRunnerContextImpl) taskB.getRunnerContext()).getContinuationContext())
+ .isSameAs(continuationB);
+ }
+ }
+
+ @Test
+ void restoreRestoresOrClearsDurableContext() throws Exception {
+ DurableExecutionManager durableManager =
+ new DurableExecutionManager(new InMemoryActionStateStore(false));
+ try (ActionTaskContextManager mgr = newManager(durableManager)) {
+ Action action = TestActions.noopAction();
+ InputEvent event = new InputEvent(1L);
+ ActionTask taskA = new JavaActionTask("a", event, action, 1L);
+ ActionTask taskB = new JavaActionTask("b", new InputEvent(2L), action, 2L);
+
+ invokeCreateAndSetRunnerContext(mgr, taskA);
+ ActionState actionState = new ActionState(event);
+ durableManager.setupDurableExecutionContext(taskA, actionState, 0L);
+ RunnerContextImpl.DurableExecutionContext durableContext =
+ taskA.getRunnerContext().getDurableExecutionContext();
+
+ invokeCreateAndSetRunnerContext(mgr, taskB);
+ mgr.restore("b", taskB);
+ assertThat(taskB.getRunnerContext().getDurableExecutionContext()).isNull();
+
+ mgr.restore("a", taskA);
+ assertThat(taskA.getRunnerContext().getDurableExecutionContext())
+ .isSameAs(durableContext);
+ }
+ }
+
+ @Test
+ void restoreReestablishesRunnerContextFromKeyAndActionTask() throws Exception {
+ try (ActionTaskContextManager mgr = newManager()) {
+ ActionTask task =
+ new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction(), 1L);
+ invokeCreateAndSetRunnerContext(mgr, task);
+
+ // The continuation executor extracts key and action task from its continuation context
+ // and hands them to the restorer directly; the manager never reads the context itself.
+ mgr.restore("k", task);
+ assertThat(task.getRunnerContext().getMemoryContext()).isNotNull();
+ }
+ }
+
@Test
void createAndSetRunnerContextReusesExistingMemoryContext() throws Exception {
- try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) {
+ try (ActionTaskContextManager mgr = newManager()) {
Action action = TestActions.noopAction();
ActionTask from = new JavaActionTask("k", new InputEvent(1L), action, 1L);
ActionTask to = new JavaActionTask("k", new InputEvent(2L), action, 1L);
@@ -217,7 +288,7 @@ class ActionTaskContextManagerTest {
@Test
void sameKeyTasksSwitchLtmWithDistinctObservationIds() throws Exception {
- try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) {
+ try (ActionTaskContextManager mgr = newManager()) {
Action action = TestActions.noopAction();
ActionTask suspended = new JavaActionTask("k", new InputEvent(1L), action, 1L);
ActionTask sibling = new JavaActionTask("k", new InputEvent(1L), action, 1L);
@@ -234,7 +305,7 @@ class ActionTaskContextManagerTest {
@Test
void transferContextsCopiesMemoryAndContinuationToNewTask() throws Exception {
- try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) {
+ try (ActionTaskContextManager mgr = newManager()) {
Action action = TestActions.noopAction();
ActionTask from = new JavaActionTask("k", new InputEvent(1L), action, 1L);
ActionTask to = new JavaActionTask("k", new InputEvent(2L), action, 1L);
@@ -275,7 +346,7 @@ class ActionTaskContextManagerTest {
@Test
void componentListenersFollowActionExecutionAcrossContinuationTasks() throws Exception {
- try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) {
+ try (ActionTaskContextManager mgr = newManager()) {
Action action = TestActions.noopAction();
ActionTask from = new JavaActionTask("k", new InputEvent(1L), action, 1L);
ActionTask to =
@@ -309,7 +380,7 @@ class ActionTaskContextManagerTest {
@Test
void removingContextsDropsComponentListeners() throws Exception {
- try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) {
+ try (ActionTaskContextManager mgr = newManager()) {
ActionTask task =
new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction(), 1L);
List<RecordingComponentListener> created = new ArrayList<>();
@@ -340,7 +411,7 @@ class ActionTaskContextManagerTest {
@Test
void activeExecutionReportsDoNotEnterActionTaskState() throws Exception {
- try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) {
+ try (ActionTaskContextManager mgr = newManager()) {
ActionTask task =
new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction(), 1L);
invokeCreateAndSetRunnerContext(
@@ -366,9 +437,68 @@ class ActionTaskContextManagerTest {
}
}
+ @Test
+ void transferContextsSharesPendingEventsBufferWithGeneratedTask() throws Exception {
+ try (ActionTaskContextManager mgr = newManager()) {
+ Action action = TestActions.noopAction();
+ ActionTask from = new JavaActionTask("k", new InputEvent(1L), action, 1L);
+ ActionTask to = new JavaActionTask("k", new InputEvent(1L), action, 1L);
+
+ invokeCreateAndSetRunnerContext(mgr, from);
+ OutputEvent bufferedBeforeYield = new OutputEvent("before-yield");
+ from.getRunnerContext().sendEvent(bufferedBeforeYield);
+ List<Event> liveBuffer = from.getRunnerContext().getPendingEvents();
+
+ mgr.transferContexts(from, to, new DurableExecutionManager(null));
+ invokeCreateAndSetRunnerContext(mgr, to);
+
+ assertThat(to.getRunnerContext().getPendingEvents()).isSameAs(liveBuffer);
+ assertThat(to.getRunnerContext().drainEvents(null))
+ .containsExactly(bufferedBeforeYield);
+ assertThat(liveBuffer).isEmpty();
+ }
+ }
+
+ @Test
+ void transferContextsRecreatesIdentityContinuationContextsAfterCleanup() throws Exception {
+ DurableExecutionManager durableManager =
+ new DurableExecutionManager(new InMemoryActionStateStore(false));
+ try (ActionTaskContextManager mgr = newManager(durableManager)) {
+ Action action = TestActions.noopAction();
+ InputEvent event = new InputEvent(1L);
+ ActionTask task = new JavaActionTask("k", event, action, 1L);
+
+ invokeCreateAndSetRunnerContext(mgr, task);
+ RunnerContextImpl.MemoryContext memoryContext =
+ task.getRunnerContext().getMemoryContext();
+ ContinuationContext continuationContext =
+ ((JavaRunnerContextImpl) task.getRunnerContext()).getContinuationContext();
+ List<Event> pendingEvents = task.getRunnerContext().getPendingEvents();
+
+ durableManager.setupDurableExecutionContext(task, new ActionState(event), 0L);
+ RunnerContextImpl.DurableExecutionContext durableContext =
+ task.getRunnerContext().getDurableExecutionContext();
+
+ // Python generator continuations return `this` as the generated task, so from==to.
+ // Verify the remove-then-transfer-to-self cycle preserves all context instances.
+ mgr.removeContexts(task);
+ durableManager.removeDurableContext(task);
+ mgr.transferContexts(task, task, durableManager);
+ invokeCreateAndSetRunnerContext(mgr, task);
+
+ assertThat(task.getRunnerContext().getMemoryContext()).isSameAs(memoryContext);
+ assertThat(((JavaRunnerContextImpl) task.getRunnerContext()).getContinuationContext())
+ .isSameAs(continuationContext);
+ assertThat(task.getRunnerContext().getPendingEvents()).isSameAs(pendingEvents);
+ assertThat(durableManager.getDurableContext(task)).isSameAs(durableContext);
+ assertThat(task.getRunnerContext().getDurableExecutionContext())
+ .isSameAs(durableContext);
+ }
+ }
+
@Test
void transferContextsRoutesDurableContextThroughManager() throws Exception {
- try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) {
+ try (ActionTaskContextManager mgr = newManager()) {
Action action = TestActions.noopAction();
InputEvent event = new InputEvent(1L);
ActionTask from = new JavaActionTask("k", event, action, 1L);
@@ -406,7 +536,7 @@ class ActionTaskContextManagerTest {
@Test
void closeIsIdempotent() throws Exception {
// Not using try-with-resources here because we want to call close() explicitly twice.
- ActionTaskContextManager mgr = new ActionTaskContextManager(1);
+ ActionTaskContextManager mgr = newManager();
ActionTask t = new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction(), 1L);
invokeCreateAndSetRunnerContext(mgr, t);
@@ -491,4 +621,16 @@ class ActionTaskContextManagerTest {
private static AgentPlan newEmptyAgentPlan() {
return new AgentPlan(new HashMap<>(), new HashMap<>());
}
+
+ private static ActionTaskContextManager newManager() {
+ return newManager(new DurableExecutionManager(null));
+ }
+
+ private static ActionTaskContextManager newManager(DurableExecutionManager durableExecManager) {
+ // The operator is only touched to re-establish the current key during a restore; a mock is
+ // sufficient for these contract tests. The continuation executor runs synchronously here
+ // (null mailbox lock), so the parallel-MVP hand-off is exercised elsewhere.
+ return new ActionTaskContextManager(
+ mock(ActionExecutionOperator.class), durableExecManager, 1);
+ }
}
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/DurableExecutionManagerTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/DurableExecutionManagerTest.java
index ef58ef20..6c5b6e51 100644
--- a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/DurableExecutionManagerTest.java
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/DurableExecutionManagerTest.java
@@ -21,6 +21,7 @@ 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.plan.AgentConfiguration;
+import org.apache.flink.agents.plan.AgentPlan;
import org.apache.flink.agents.plan.actions.Action;
import org.apache.flink.agents.runtime.actionstate.ActionState;
import org.apache.flink.agents.runtime.actionstate.ActionStateKeyEncoder;
@@ -109,6 +110,28 @@ class DurableExecutionManagerTest {
dem.close();
}
+ @Test
+ void setupDurableExecutionContextPinsFreshContext() {
+ InMemoryActionStateStore store = new InMemoryActionStateStore(false);
+ DurableExecutionManager manager = new DurableExecutionManager(store);
+ Action action = TestActions.noopAction();
+ Event event = new InputEvent(1L);
+ ActionTask task = new JavaActionTask("key", event, action, 1L);
+ RunnerContextImpl runnerContext =
+ new RunnerContextImpl(
+ null,
+ () -> {},
+ new AgentPlan(new HashMap<>(), new HashMap<>()),
+ null,
+ "job");
+ task.setRunnerContext(runnerContext);
+
+ manager.setupDurableExecutionContext(task, new ActionState(event), 0L);
+
+ assertThat(manager.getDurableContext(task))
+ .isSameAs(task.getRunnerContext().getDurableExecutionContext());
+ }
+
@Test
@SuppressWarnings("unchecked")
void noStoreModeSnapshotAndNotifyKeepCheckpointMapEmpty() throws Exception {
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionCoordinatorTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionCoordinatorTest.java
new file mode 100644
index 00000000..271981df
--- /dev/null
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionCoordinatorTest.java
@@ -0,0 +1,881 @@
+/*
+ * 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.operator.parallel;
+
+import org.apache.flink.agents.runtime.async.AsyncExecutorThreadFactory;
+import org.apache.flink.util.function.ThrowingRunnable;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+
+import javax.annotation.Nullable;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Queue;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Supplier;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/**
+ * Tests the coordinator's {@link ParallelExecutionTask} lifecycle: node-driven pull, eager per-key
+ * taskIndex assignment at add time, worker execution, the mailbox-thread in-order commit drain, and
+ * its follow-up. Works are supplied by a FIFO factory standing in for the operator's lazy {@code
+ * Work}.
+ */
+class ParallelExecutionCoordinatorTest {
+
+ private static final Object KEY = "k0";
+
+ @Test
+ void addTaskRequiresMailboxLock() {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails = new LinkedBlockingQueue<>();
+ try (ParallelExecutionCoordinator coordinator = newCoordinator(lock, mails, () -> null)) {
+ assertThatThrownBy(() -> coordinator.addTask(KEY))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("does not own");
+ }
+ }
+
+ @Test
+ @Timeout(15)
+ void workerRunsSetsUpThenExecutesCommitsAndFinishes() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails = new LinkedBlockingQueue<>();
+ List<String> events = Collections.synchronizedList(new ArrayList<>());
+ Queue<ParallelExecutionTask> works = new ConcurrentLinkedQueue<>();
+ try (ParallelExecutionCoordinator coordinator = newCoordinator(lock, mails, works::poll)) {
+ lock.acquireByMain();
+ // The worker owns the mailbox lock while it runs the action.
+ works.add(FakeWork.blocking("A", events, lock::checkReentrant));
+ coordinator.addTask(KEY);
+ lock.release();
+
+ runOneMail(mails);
+
+ // The work is set up + executed on the worker, then restored + committed and its group
+ // finished on the mailbox thread; finishGroup receives the just-committed work as last.
+ assertThat(events).containsExactly("execute-A", "restore-A", "commit-A", "finish-A->A");
+ assertThat(coordinator.hasOutstanding(KEY)).isFalse();
+ }
+ }
+
+ @Test
+ @Timeout(15)
+ void workerRunsAsManagedAsyncExecutorThread() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails = new LinkedBlockingQueue<>();
+ List<String> events = Collections.synchronizedList(new ArrayList<>());
+ Queue<ParallelExecutionTask> works = new ConcurrentLinkedQueue<>();
+ AtomicBoolean managedWorker = new AtomicBoolean();
+
+ try (ParallelExecutionCoordinator coordinator = newCoordinator(lock, mails, works::poll)) {
+ works.add(
+ FakeWork.blocking(
+ "A",
+ events,
+ () ->
+ managedWorker.set(
+ AsyncExecutorThreadFactory.isAsyncExecutorThread())));
+ lock.acquireByMain();
+ coordinator.addTask(KEY);
+ lock.release();
+
+ runOneMail(mails);
+
+ assertThat(managedWorker)
+ .as("parallel action workers must participate in managed-worker lifecycle")
+ .isTrue();
+ }
+ }
+
+ @Test
+ @Timeout(15)
+ void outOfOrderCompletionCommitsInKeyOrder() throws Exception {
+ // Both same-key works run on separate workers; each releases the lock inside execute()
+ // (mimicking the durable-async segment) so they can block concurrently under the real lock.
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails = new LinkedBlockingQueue<>();
+ List<String> events = Collections.synchronizedList(new ArrayList<>());
+ Queue<ParallelExecutionTask> works = new ConcurrentLinkedQueue<>();
+ CountDownLatch aStarted = new CountDownLatch(1);
+ CountDownLatch bStarted = new CountDownLatch(1);
+ CountDownLatch allowA = new CountDownLatch(1);
+ CountDownLatch allowB = new CountDownLatch(1);
+ try (ParallelExecutionCoordinator coordinator =
+ newCoordinator(2, lock, mails, works::poll)) {
+ lock.acquireByMain();
+ // Add A, wait until it is running (taskIndex 0), then add B (taskIndex 1): the
+ // sequential
+ // add fixes the pull order deterministically.
+ works.add(
+ FakeWork.blockingAsync(
+ "A",
+ events,
+ lock,
+ () -> {
+ aStarted.countDown();
+ await(allowA);
+ }));
+ coordinator.addTask(KEY);
+ lock.release();
+ assertThat(aStarted.await(5, TimeUnit.SECONDS)).isTrue();
+
+ works.add(
+ FakeWork.blockingAsync(
+ "B",
+ events,
+ lock,
+ () -> {
+ bStarted.countDown();
+ await(allowB);
+ }));
+ lock.acquireByMain();
+ coordinator.addTask(KEY);
+ lock.release();
+ assertThat(bStarted.await(5, TimeUnit.SECONDS)).isTrue();
+
+ // B finishes first, but its drain must stop at the still-running head A.
+ allowB.countDown();
+ runOneMail(mails);
+ assertThat(committed(events)).isEmpty();
+ assertThat(coordinator.hasOutstanding(KEY)).isTrue();
+
+ // Once A finishes, the contiguous done prefix commits A then B in key-taskIndex order.
+ allowA.countDown();
+ runOneMail(mails);
+ assertThat(committed(events)).containsExactly("A", "B");
+ assertThat(coordinator.hasOutstanding(KEY)).isFalse();
+ } finally {
+ allowA.countDown();
+ allowB.countDown();
+ }
+ }
+
+ @Test
+ @Timeout(15)
+ void commitFailureEscapesMailbox() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails = new LinkedBlockingQueue<>();
+ List<String> events = Collections.synchronizedList(new ArrayList<>());
+ Queue<ParallelExecutionTask> works = new ConcurrentLinkedQueue<>();
+ RuntimeException commitFailure = new RuntimeException("commit failure");
+ try (ParallelExecutionCoordinator coordinator = newCoordinator(lock, mails, works::poll)) {
+ lock.acquireByMain();
+ works.add(FakeWork.failingCommit("A", events, commitFailure));
+ coordinator.addTask(KEY);
+ lock.release();
+
+ assertThatThrownBy(() -> runOneMail(mails)).isSameAs(commitFailure);
+ }
+ }
+
+ @Test
+ @Timeout(15)
+ void commitDrainFailureAmongSiblingsDoesNotLeakLockOrWedgeGroup() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails = new LinkedBlockingQueue<>();
+ List<String> events = Collections.synchronizedList(new ArrayList<>());
+ Queue<ParallelExecutionTask> works = new ConcurrentLinkedQueue<>();
+ CountDownLatch aStarted = new CountDownLatch(1);
+ CountDownLatch bStarted = new CountDownLatch(1);
+ CountDownLatch allowA = new CountDownLatch(1);
+ CountDownLatch allowB = new CountDownLatch(1);
+ RuntimeException commitFailure = new RuntimeException("commit failure");
+
+ try (ParallelExecutionCoordinator coordinator =
+ newCoordinator(2, lock, mails, works::poll)) {
+ lock.acquireByMain();
+ works.add(
+ FakeWork.blockingAsync(
+ "A",
+ events,
+ lock,
+ () -> {
+ aStarted.countDown();
+ await(allowA);
+ }));
+ coordinator.addTask(KEY);
+ lock.release();
+ assertThat(aStarted.await(5, TimeUnit.SECONDS)).isTrue();
+
+ works.add(
+ FakeWork.blockingAsyncFailingCommit(
+ "B",
+ events,
+ lock,
+ () -> {
+ bStarted.countDown();
+ await(allowB);
+ },
+ commitFailure));
+ lock.acquireByMain();
+ coordinator.addTask(KEY);
+ lock.release();
+ assertThat(bStarted.await(5, TimeUnit.SECONDS)).isTrue();
+
+ allowB.countDown();
+ runOneMail(mails);
+ assertThat(committed(events)).isEmpty();
+ assertThat(coordinator.hasOutstanding(KEY)).isTrue();
+
+ allowA.countDown();
+ assertThatThrownBy(() -> runOneMail(mails)).isSameAs(commitFailure);
+
+ lock.acquireByMain();
+ lock.release();
+ assertThat(coordinator.hasOutstanding(KEY)).isFalse();
+ assertThat(committed(events)).containsExactly("A");
+ } finally {
+ allowA.countDown();
+ allowB.countDown();
+ }
+ }
+
+ @Test
+ @Timeout(15)
+ void hasOutstandingReflectsOutstandingWork() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails = new LinkedBlockingQueue<>();
+ List<String> events = Collections.synchronizedList(new ArrayList<>());
+ Queue<ParallelExecutionTask> works = new ConcurrentLinkedQueue<>();
+ CountDownLatch started = new CountDownLatch(1);
+ CountDownLatch allow = new CountDownLatch(1);
+ try (ParallelExecutionCoordinator coordinator = newCoordinator(lock, mails, works::poll)) {
+ // An inactive key has no outstanding work.
+ assertThat(coordinator.hasOutstanding(KEY)).isFalse();
+
+ lock.acquireByMain();
+ works.add(
+ FakeWork.blockingAsync(
+ "A",
+ events,
+ lock,
+ () -> {
+ started.countDown();
+ await(allow);
+ }));
+ coordinator.addTask(KEY);
+ lock.release();
+
+ // Outstanding from the moment a task is added until the commit drain removes it.
+ assertThat(coordinator.hasOutstanding(KEY)).isTrue();
+ assertThat(started.await(5, TimeUnit.SECONDS)).isTrue();
+ assertThat(coordinator.hasOutstanding(KEY)).isTrue();
+
+ allow.countDown();
+ runOneMail(mails);
+ assertThat(coordinator.hasOutstanding(KEY)).isFalse();
+ } finally {
+ allow.countDown();
+ }
+ }
+
+ @Test
+ @Timeout(15)
+ void isQuiescedTracksInFlightWork() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails = new LinkedBlockingQueue<>();
+ List<String> events = Collections.synchronizedList(new ArrayList<>());
+ Queue<ParallelExecutionTask> works = new ConcurrentLinkedQueue<>();
+ CountDownLatch started = new CountDownLatch(1);
+ CountDownLatch allow = new CountDownLatch(1);
+ try (ParallelExecutionCoordinator coordinator = newCoordinator(lock, mails, works::poll)) {
+ // No in-flight work: quiesced.
+ assertThat(coordinator.isQuiesced()).isTrue();
+
+ lock.acquireByMain();
+ works.add(
+ FakeWork.blockingAsync(
+ "A",
+ events,
+ lock,
+ () -> {
+ started.countDown();
+ await(allow);
+ }));
+ coordinator.addTask(KEY);
+ lock.release();
+
+ assertThat(started.await(5, TimeUnit.SECONDS)).isTrue();
+ // Pulled and executing but not yet committed: not quiesced.
+ assertThat(coordinator.isQuiesced()).isFalse();
+
+ allow.countDown();
+ runOneMail(mails); // commit drain removes the node
+ assertThat(coordinator.isQuiesced()).isTrue();
+ } finally {
+ allow.countDown();
+ }
+ }
+
+ @Test
+ @Timeout(15)
+ void drainingPausesNewDispatchUntilStopped() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails = new LinkedBlockingQueue<>();
+ List<String> events = Collections.synchronizedList(new ArrayList<>());
+ Queue<ParallelExecutionTask> works = new ConcurrentLinkedQueue<>();
+ CountDownLatch started = new CountDownLatch(1);
+ CountDownLatch allow = new CountDownLatch(1);
+ try (ParallelExecutionCoordinator coordinator = newCoordinator(lock, mails, works::poll)) {
+ coordinator.startDraining();
+
+ works.add(
+ FakeWork.blockingAsync(
+ "A",
+ events,
+ lock,
+ () -> {
+ started.countDown();
+ await(allow);
+ }));
+ lock.acquireByMain();
+ coordinator.addTask(KEY);
+ lock.release();
+
+ // While draining, the worker declines the pulled credit; the task is not started, but
+ // its node stays pending (outstanding), so its task stays in the durable queue.
+ assertThat(started.await(1, TimeUnit.SECONDS)).isFalse();
+ assertThat(events).isEmpty();
+ assertThat(coordinator.hasOutstanding(KEY)).isTrue();
+
+ // Resuming dispatch alone must revive the worker — no other mailbox lock cycle runs.
+ coordinator.stopDraining();
+ assertThat(started.await(5, TimeUnit.SECONDS)).isTrue();
+ assertThat(coordinator.isQuiesced()).isFalse();
+
+ allow.countDown();
+ runOneMail(mails);
+ } finally {
+ allow.countDown();
+ }
+ }
+
+ @Test
+ @Timeout(15)
+ void closeWaitsForInFlightWorkerBeforeReturning() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails = new LinkedBlockingQueue<>();
+ List<String> events = Collections.synchronizedList(new ArrayList<>());
+ CountDownLatch workerEntered = new CountDownLatch(1);
+ CountDownLatch allowFinish = new CountDownLatch(1);
+ // Emulates user code that does not react to the shutdown interrupt.
+ Supplier<ParallelExecutionTask> works =
+ () ->
+ FakeWork.blocking(
+ "A",
+ events,
+ () -> {
+ workerEntered.countDown();
+ boolean released = false;
+ while (!released) {
+ try {
+ released = allowFinish.await(5, TimeUnit.SECONDS);
+ } catch (InterruptedException ignored) {
+ // keep waiting
+ }
+ }
+ });
+ ParallelExecutionCoordinator coordinator = newCoordinator(lock, mails, works);
+ try {
+ lock.acquireByMain();
+ coordinator.addTask(KEY);
+ lock.release();
+ assertThat(workerEntered.await(5, TimeUnit.SECONDS)).isTrue();
+
+ Thread releaser =
+ new Thread(
+ () -> {
+ try {
+ Thread.sleep(400L);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ allowFinish.countDown();
+ });
+ releaser.setDaemon(true);
+ releaser.start();
+
+ // close() must block until the in-flight worker has terminated, so callers can
+ // safely release resources the worker may still be using.
+ coordinator.close();
+ assertThat(coordinator.getWorkerPoolSizeForTesting()).isZero();
+ releaser.join();
+ } finally {
+ allowFinish.countDown();
+ coordinator.close();
+ }
+ }
+
+ @Test
+ @Timeout(15)
+ void startDrainingFencesInProgressPullSoQuiesceCannotMissIt() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails = new LinkedBlockingQueue<>();
+ List<String> events = Collections.synchronizedList(new ArrayList<>());
+ CountDownLatch factoryEntered = new CountDownLatch(1);
+ CountDownLatch allowFactory = new CountDownLatch(1);
+ CountDownLatch drainerDone = new CountDownLatch(1);
+ AtomicBoolean quiescedAfterFence = new AtomicBoolean();
+ // The factory runs inside the worker's pull section: after the draining check, before
+ // inFlightExecuting is incremented — exactly the TOCTOU window.
+ Supplier<ParallelExecutionTask> works =
+ () -> {
+ factoryEntered.countDown();
+ await(allowFactory);
+ return FakeWork.blocking("A", events, () -> {});
+ };
+ try (ParallelExecutionCoordinator coordinator = newCoordinator(lock, mails, works)) {
+ lock.acquireByMain();
+ coordinator.addTask(KEY);
+ lock.release();
+ assertThat(factoryEntered.await(5, TimeUnit.SECONDS)).isTrue();
+
+ // Start draining while the worker sits in the window; the fence must block until the
+ // worker's lock hold (which includes the increment) is over.
+ Thread drainer =
+ new Thread(
+ () -> {
+ try {
+ coordinator.startDraining();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ quiescedAfterFence.set(coordinator.isQuiesced());
+ drainerDone.countDown();
+ });
+ drainer.start();
+ // Give the drainer time to reach the fence (pre-fix it would finish immediately).
+ Thread.sleep(300L);
+ allowFactory.countDown();
+ assertThat(drainerDone.await(5, TimeUnit.SECONDS)).isTrue();
+
+ // The pulled-but-uncommitted task must be visible to the quiesce check.
+ assertThat(quiescedAfterFence.get()).isFalse();
+
+ coordinator.stopDraining();
+ runOneMail(mails);
+ assertThat(coordinator.isQuiesced()).isTrue();
+ } finally {
+ allowFactory.countDown();
+ }
+ }
+
+ @Test
+ @Timeout(15)
+ void workerFailureSurfacesOnMailboxAndKeepsLockUsable() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails = new LinkedBlockingQueue<>();
+ // A non-RuntimeException failure in the pull/handoff section: without the mailbox funnel
+ // it would silently kill the pool thread and leak the lock.
+ AssertionError boom = new AssertionError("worker boom");
+ try (ParallelExecutionCoordinator coordinator =
+ newCoordinator(
+ lock,
+ mails,
+ () -> {
+ throw boom;
+ })) {
+ lock.acquireByMain();
+ coordinator.addTask(KEY);
+ lock.release();
+
+ // The failure arrives as a throwing mail and fails the mailbox (i.e. the job).
+ assertThatThrownBy(() -> runOneMail(mails)).isSameAs(boom);
+
+ // The lock was handed back on the failure path: the mailbox thread can still use it.
+ lock.acquireByMain();
+ lock.release();
+ }
+ }
+
+ @Test
+ @Timeout(15)
+ void growsWorkerPoolWithBacklogUpToMax() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails = new LinkedBlockingQueue<>();
+ List<String> events = Collections.synchronizedList(new ArrayList<>());
+ Queue<ParallelExecutionTask> works = new ConcurrentLinkedQueue<>();
+ CountDownLatch aStarted = new CountDownLatch(1);
+ CountDownLatch bStarted = new CountDownLatch(1);
+ CountDownLatch cStarted = new CountDownLatch(1);
+ CountDownLatch allowAll = new CountDownLatch(1);
+ try (ParallelExecutionCoordinator coordinator =
+ newCoordinator(3, lock, mails, works::poll)) {
+ // The pool is lazy: no thread exists before the first task.
+ assertThat(coordinator.getWorkerPoolSizeForTesting()).isZero();
+
+ // Every live thread is busy when the next task arrives, so each add grows the pool.
+ works.add(
+ FakeWork.blockingAsync(
+ "A",
+ events,
+ lock,
+ () -> {
+ aStarted.countDown();
+ await(allowAll);
+ }));
+ lock.acquireByMain();
+ coordinator.addTask(KEY);
+ lock.release();
+ assertThat(aStarted.await(5, TimeUnit.SECONDS)).isTrue();
+ assertThat(coordinator.getWorkerPoolSizeForTesting()).isEqualTo(1);
+
+ works.add(
+ FakeWork.blockingAsync(
+ "B",
+ events,
+ lock,
+ () -> {
+ bStarted.countDown();
+ await(allowAll);
+ }));
+ lock.acquireByMain();
+ coordinator.addTask(KEY);
+ lock.release();
+ assertThat(bStarted.await(5, TimeUnit.SECONDS)).isTrue();
+ assertThat(coordinator.getWorkerPoolSizeForTesting()).isEqualTo(2);
+
+ works.add(
+ FakeWork.blockingAsync(
+ "C",
+ events,
+ lock,
+ () -> {
+ cStarted.countDown();
+ await(allowAll);
+ }));
+ lock.acquireByMain();
+ coordinator.addTask(KEY);
+ lock.release();
+ assertThat(cStarted.await(5, TimeUnit.SECONDS)).isTrue();
+ assertThat(coordinator.getWorkerPoolSizeForTesting()).isEqualTo(3);
+
+ allowAll.countDown();
+ runOneMail(mails);
+ runOneMail(mails);
+ runOneMail(mails);
+ assertThat(committed(events)).containsExactly("A", "B", "C");
+ } finally {
+ allowAll.countDown();
+ }
+ }
+
+ @Test
+ @Timeout(15)
+ void shrinksIdlePoolToOneCoreThreadAndGrowsAgainOnDemand() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails = new LinkedBlockingQueue<>();
+ List<String> events = Collections.synchronizedList(new ArrayList<>());
+ Queue<ParallelExecutionTask> works = new ConcurrentLinkedQueue<>();
+ CountDownLatch aStarted = new CountDownLatch(1);
+ CountDownLatch bStarted = new CountDownLatch(1);
+ CountDownLatch allowFirst = new CountDownLatch(1);
+ CountDownLatch cStarted = new CountDownLatch(1);
+ CountDownLatch dStarted = new CountDownLatch(1);
+ CountDownLatch allowSecond = new CountDownLatch(1);
+ CountDownLatch idleWorkerCleanedUp = new CountDownLatch(1);
+ try (ParallelExecutionCoordinator coordinator =
+ newCoordinator(2, 150L, lock, mails, works::poll, idleWorkerCleanedUp::countDown)) {
+ // Grow to the cap of 2.
+ works.add(
+ FakeWork.blockingAsync(
+ "A",
+ events,
+ lock,
+ () -> {
+ aStarted.countDown();
+ await(allowFirst);
+ }));
+ lock.acquireByMain();
+ coordinator.addTask(KEY);
+ lock.release();
+ assertThat(aStarted.await(5, TimeUnit.SECONDS)).isTrue();
+ works.add(
+ FakeWork.blockingAsync(
+ "B",
+ events,
+ lock,
+ () -> {
+ bStarted.countDown();
+ await(allowFirst);
+ }));
+ lock.acquireByMain();
+ coordinator.addTask(KEY);
+ lock.release();
+ assertThat(bStarted.await(5, TimeUnit.SECONDS)).isTrue();
+ assertThat(coordinator.getWorkerPoolSizeForTesting()).isEqualTo(2);
+
+ allowFirst.countDown();
+ runOneMail(mails);
+ runOneMail(mails);
+
+ // With nothing to do, the surplus thread times out and exits; the core thread does
+ // not, even across several keep-alive periods.
+ waitUntilPoolSize(coordinator, 1);
+ assertThat(idleWorkerCleanedUp.await(5, TimeUnit.SECONDS))
+ .as("idle worker must run its exit cleanup before retiring")
+ .isTrue();
+ Thread.sleep(500L);
+ assertThat(coordinator.getWorkerPoolSizeForTesting()).isEqualTo(1);
+
+ // New demand grows the pool again.
+ works.add(
+ FakeWork.blockingAsync(
+ "C",
+ events,
+ lock,
+ () -> {
+ cStarted.countDown();
+ await(allowSecond);
+ }));
+ lock.acquireByMain();
+ coordinator.addTask(KEY);
+ lock.release();
+ assertThat(cStarted.await(5, TimeUnit.SECONDS)).isTrue();
+ works.add(
+ FakeWork.blockingAsync(
+ "D",
+ events,
+ lock,
+ () -> {
+ dStarted.countDown();
+ await(allowSecond);
+ }));
+ lock.acquireByMain();
+ coordinator.addTask(KEY);
+ lock.release();
+ assertThat(dStarted.await(5, TimeUnit.SECONDS)).isTrue();
+ assertThat(coordinator.getWorkerPoolSizeForTesting()).isEqualTo(2);
+
+ allowSecond.countDown();
+ runOneMail(mails);
+ runOneMail(mails);
+ } finally {
+ allowFirst.countDown();
+ allowSecond.countDown();
+ }
+ }
+
+ private static ParallelExecutionCoordinator newCoordinator(
+ ParallelExecutionLock lock,
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails,
+ Supplier<ParallelExecutionTask> workFactory) {
+ return newCoordinator(1, lock, mails, workFactory);
+ }
+
+ private static ParallelExecutionCoordinator newCoordinator(
+ int maxWorkers,
+ ParallelExecutionLock lock,
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails,
+ Supplier<ParallelExecutionTask> workFactory) {
+ // A practically-infinite keep-alive keeps the pool size deterministic for these tests.
+ return newCoordinator(maxWorkers, 3_600_000L, lock, mails, workFactory);
+ }
+
+ private static ParallelExecutionCoordinator newCoordinator(
+ int maxWorkers,
+ long workerIdleTimeoutMillis,
+ ParallelExecutionLock lock,
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails,
+ Supplier<ParallelExecutionTask> workFactory) {
+ return newCoordinator(
+ maxWorkers, workerIdleTimeoutMillis, lock, mails, workFactory, () -> {});
+ }
+
+ private static ParallelExecutionCoordinator newCoordinator(
+ int maxWorkers,
+ long workerIdleTimeoutMillis,
+ ParallelExecutionLock lock,
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails,
+ Supplier<ParallelExecutionTask> workFactory,
+ Runnable threadCleanup) {
+ return new ParallelExecutionCoordinator(
+ lock,
+ (mail, description) -> mails.add(mail),
+ workFactory,
+ threadCleanup,
+ maxWorkers,
+ workerIdleTimeoutMillis);
+ }
+
+ private static void waitUntilPoolSize(ParallelExecutionCoordinator coordinator, int expected)
+ throws InterruptedException {
+ long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
+ while (coordinator.getWorkerPoolSizeForTesting() != expected) {
+ assertThat(System.nanoTime())
+ .as("timed out waiting for pool size %d", expected)
+ .isLessThan(deadline);
+ Thread.sleep(10L);
+ }
+ }
+
+ private static void runOneMail(BlockingQueue<ThrowingRunnable<? extends Exception>> mails)
+ throws Exception {
+ ThrowingRunnable<? extends Exception> mail = mails.poll(5, TimeUnit.SECONDS);
+ assertThat(mail).isNotNull();
+ mail.run();
+ }
+
+ /** Extracts, in order, the ids of works whose {@link ParallelExecutionTask#commit()} ran. */
+ private static List<String> committed(List<String> events) {
+ List<String> out = new ArrayList<>();
+ for (String event : events) {
+ if (event.startsWith("commit-")) {
+ out.add(event.substring("commit-".length()));
+ }
+ }
+ return out;
+ }
+
+ private static void await(CountDownLatch latch) {
+ try {
+ assertThat(latch.await(5, TimeUnit.SECONDS)).isTrue();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new AssertionError("Interrupted while waiting for test latch", e);
+ }
+ }
+
+ /**
+ * A controllable {@link ParallelExecutionTask}: records each lifecycle step into a shared log,
+ * can block or run arbitrary code inside {@link #execute()} — optionally releasing the lock
+ * while blocked, mimicking the durable-async segment — and can fail its {@link #commit()}. The
+ * coordinator binds its scheduling priority via {@link #setup(Object, long, long)}.
+ */
+ private static final class FakeWork implements ParallelExecutionTask {
+ private final String id;
+ private final List<String> events;
+ @Nullable private final Runnable onExecute;
+ @Nullable private final ParallelExecutionLock asyncLock;
+ @Nullable private final RuntimeException commitFailure;
+ private volatile boolean done;
+ private volatile Object assignedKey;
+ private volatile long assignedRecordIndex = -1L;
+ private volatile long assignedTaskIndex = -1L;
+
+ private FakeWork(
+ String id,
+ List<String> events,
+ @Nullable Runnable onExecute,
+ @Nullable ParallelExecutionLock asyncLock,
+ @Nullable RuntimeException commitFailure) {
+ this.id = id;
+ this.events = events;
+ this.onExecute = onExecute;
+ this.asyncLock = asyncLock;
+ this.commitFailure = commitFailure;
+ }
+
+ static FakeWork blocking(String id, List<String> events, Runnable onExecute) {
+ return new FakeWork(id, events, onExecute, null, null);
+ }
+
+ /**
+ * Like {@link #blocking}, but releases {@code lock} around {@code onExecute} and
+ * re-acquires it at the work's assigned priority — the worker-side protocol of a
+ * durable-async segment.
+ */
+ static FakeWork blockingAsync(
+ String id, List<String> events, ParallelExecutionLock lock, Runnable onExecute) {
+ return new FakeWork(id, events, onExecute, lock, null);
+ }
+
+ static FakeWork blockingAsyncFailingCommit(
+ String id,
+ List<String> events,
+ ParallelExecutionLock lock,
+ Runnable onExecute,
+ RuntimeException commitFailure) {
+ return new FakeWork(id, events, onExecute, lock, commitFailure);
+ }
+
+ static FakeWork failingCommit(
+ String id, List<String> events, RuntimeException commitFailure) {
+ return new FakeWork(id, events, null, null, commitFailure);
+ }
+
+ @Override
+ public void setup(Object key, long recordIndex, long taskIndex) {
+ this.assignedKey = key;
+ this.assignedRecordIndex = recordIndex;
+ this.assignedTaskIndex = taskIndex;
+ }
+
+ @Override
+ public void restoreContext() {
+ events.add("restore-" + id);
+ }
+
+ @Override
+ public void execute() {
+ events.add("execute-" + id);
+ if (asyncLock != null) {
+ // Durable-async protocol: give up the lock while blocked, then resume at the
+ // work's assigned (recordIndex, taskIndex) priority.
+ asyncLock.release();
+ try {
+ if (onExecute != null) {
+ onExecute.run();
+ }
+ } finally {
+ try {
+ asyncLock.acquireByWorker(assignedRecordIndex, assignedTaskIndex);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new AssertionError("Interrupted while re-acquiring the lock", e);
+ }
+ }
+ } else if (onExecute != null) {
+ onExecute.run();
+ }
+ done = true;
+ }
+
+ @Override
+ public boolean isDone() {
+ return done;
+ }
+
+ @Override
+ public void commit() {
+ if (commitFailure != null) {
+ throw commitFailure;
+ }
+ events.add("commit-" + id);
+ }
+
+ @Override
+ public void finishGroup(@Nullable ParallelExecutionTask lastCommitted) {
+ events.add(
+ "finish-"
+ + id
+ + "->"
+ + (lastCommitted == null ? "none" : ((FakeWork) lastCommitted).id));
+ }
+ }
+}
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionLockTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionLockTest.java
new file mode 100644
index 00000000..6f93c1ca
--- /dev/null
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionLockTest.java
@@ -0,0 +1,441 @@
+/*
+ * 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.operator.parallel;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.function.BooleanSupplier;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.junit.jupiter.api.Assertions.fail;
+
+/** Tests for {@link ParallelExecutionLock}. */
+class ParallelExecutionLockTest {
+
+ // -----------------------------------------------------------------------
+ // Basic acquire / release / check
+ // -----------------------------------------------------------------------
+
+ @Test
+ void acquireReleaseAndCheckUseCurrentOwnerOnly() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+
+ lock.acquireByMain();
+ lock.checkReentrant();
+ lock.release();
+
+ assertThatThrownBy(lock::checkReentrant)
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("does not own");
+ }
+
+ // -----------------------------------------------------------------------
+ // Non-reentrant guard
+ // -----------------------------------------------------------------------
+
+ @Test
+ void acquireIsNonReentrant() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+
+ lock.acquireByMain();
+ try {
+ assertThatThrownBy(lock::acquireByMain)
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("non-reentrant");
+ assertThatThrownBy(() -> lock.acquireByWorker(0, 0))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("non-reentrant");
+ } finally {
+ lock.release();
+ }
+ }
+
+ // -----------------------------------------------------------------------
+ // Non-owner release guard
+ // -----------------------------------------------------------------------
+
+ @Test
+ void nonOwnerCannotRelease() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ lock.acquireByMain();
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ try {
+ Future<?> result =
+ executor.submit(
+ () ->
+ assertThatThrownBy(lock::release)
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("does not own"));
+ result.get(5, TimeUnit.SECONDS);
+ } finally {
+ executor.shutdownNow();
+ lock.release();
+ }
+ }
+
+ @Test
+ void nonOwnerCannotReleaseToMailbox() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ lock.acquireByMain();
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ try {
+ Future<?> result =
+ executor.submit(
+ () ->
+ assertThatThrownBy(lock::releaseToMain)
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("does not own"));
+ result.get(5, TimeUnit.SECONDS);
+ } finally {
+ executor.shutdownNow();
+ lock.release();
+ }
+ }
+
+ // -----------------------------------------------------------------------
+ // releaseToMain sticky grant
+ // -----------------------------------------------------------------------
+
+ @Test
+ void releaseToMainKeepsStickyGrantUntilMainAcquires() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ try {
+ Future<?> worker =
+ executor.submit(
+ () -> {
+ lock.acquireByWorker(0, 0);
+ lock.releaseToMain();
+ return null;
+ });
+
+ worker.get(5, TimeUnit.SECONDS);
+ // Lock is still held (sticky grant pending)
+ assertThat(lock.isMailboxGrantedForTesting()).isTrue();
+ assertThat(lock.isLockedForTesting()).isTrue();
+
+ // Mailbox acquires via fast path — flag must be cleared
+ lock.acquireByMain();
+ try {
+ lock.checkReentrant();
+ assertThat(lock.isMailboxGrantedForTesting()).isFalse();
+ } finally {
+ lock.release();
+ }
+ } finally {
+ executor.shutdownNow();
+ }
+ }
+
+ // -----------------------------------------------------------------------
+ // Worker FIFO order
+ // -----------------------------------------------------------------------
+
+ @Test
+ void equalPriorityWorkersAllAcquireInArbitraryOrder() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ ExecutorService executor = Executors.newFixedThreadPool(3);
+ List<Integer> acquisitionOrder = Collections.synchronizedList(new ArrayList<>());
+
+ lock.acquireByMain();
+ try {
+ Future<?> worker1 = submitWorker(lock, executor, acquisitionOrder, 1);
+ waitUntil(() -> lock.getQueuedWorkerCountForTesting() == 1);
+
+ Future<?> worker2 = submitWorker(lock, executor, acquisitionOrder, 2);
+ waitUntil(() -> lock.getQueuedWorkerCountForTesting() == 2);
+
+ Future<?> worker3 = submitWorker(lock, executor, acquisitionOrder, 3);
+ waitUntil(() -> lock.getQueuedWorkerCountForTesting() == 3);
+
+ lock.release();
+
+ worker1.get(5, TimeUnit.SECONDS);
+ worker2.get(5, TimeUnit.SECONDS);
+ worker3.get(5, TimeUnit.SECONDS);
+ // Equal-priority workers have no defined grant order; all must acquire exactly once.
+ assertThat(acquisitionOrder).containsExactlyInAnyOrder(1, 2, 3);
+ } finally {
+ executor.shutdownNow();
+ }
+ }
+
+ // -----------------------------------------------------------------------
+ // Mailbox priority over queued workers
+ // -----------------------------------------------------------------------
+
+ @Test
+ void mailboxAcquireHasPriorityOverQueuedWorkers() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ ExecutorService executor = Executors.newFixedThreadPool(2);
+ List<String> acquisitionOrder = Collections.synchronizedList(new ArrayList<>());
+
+ lock.acquireByMain();
+ try {
+ Future<?> worker =
+ executor.submit(
+ () -> {
+ lock.acquireByWorker(0, 0);
+ try {
+ acquisitionOrder.add("worker");
+ } finally {
+ lock.release();
+ }
+ return null;
+ });
+ waitUntil(() -> lock.getQueuedWorkerCountForTesting() == 1);
+
+ Future<?> mailbox =
+ executor.submit(
+ () -> {
+ lock.acquireByMain();
+ try {
+ acquisitionOrder.add("mailbox");
+ } finally {
+ lock.release();
+ }
+ return null;
+ });
+ waitUntil(lock::isMailboxWaitingForTesting);
+
+ lock.release();
+
+ mailbox.get(5, TimeUnit.SECONDS);
+ worker.get(5, TimeUnit.SECONDS);
+ assertThat(acquisitionOrder).containsExactly("mailbox", "worker");
+ } finally {
+ executor.shutdownNow();
+ }
+ }
+
+ // -----------------------------------------------------------------------
+ // Worker interrupt cleanup
+ // -----------------------------------------------------------------------
+
+ @Test
+ void interruptedWorkerIsRemovedFromQueue() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+
+ lock.acquireByMain();
+ try {
+ Future<?> worker =
+ executor.submit(
+ () -> {
+ lock.acquireByWorker(0, 0);
+ lock.release();
+ return null;
+ });
+ waitUntil(() -> lock.getQueuedWorkerCountForTesting() == 1);
+
+ worker.cancel(true);
+ waitUntil(() -> lock.getQueuedWorkerCountForTesting() == 0);
+ assertThat(worker.isCancelled()).isTrue();
+ } finally {
+ lock.release();
+ executor.shutdownNow();
+ }
+ }
+
+ // -----------------------------------------------------------------------
+ // Worker acquireByWorker fast path (no contention)
+ // -----------------------------------------------------------------------
+
+ @Test
+ void workerAcquiresImmediatelyWhenLockFree() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ try {
+ Future<?> worker =
+ executor.submit(
+ () -> {
+ lock.acquireByWorker(0, 0);
+ lock.release();
+ return null;
+ });
+ worker.get(5, TimeUnit.SECONDS);
+ assertThat(lock.isLockedForTesting()).isFalse();
+ } finally {
+ executor.shutdownNow();
+ }
+ }
+
+ // -----------------------------------------------------------------------
+ // Worker priority order
+ // -----------------------------------------------------------------------
+
+ @Test
+ void workersAcquireInPriorityOrder() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ ExecutorService executor = Executors.newFixedThreadPool(3);
+ List<Integer> acquisitionOrder = Collections.synchronizedList(new ArrayList<>());
+
+ lock.acquireByMain();
+ try {
+ Future<?> worker1 = submitWorker(lock, executor, acquisitionOrder, 1, 5L);
+ waitUntil(() -> lock.getQueuedWorkerCountForTesting() == 1);
+
+ Future<?> worker2 = submitWorker(lock, executor, acquisitionOrder, 2, 1L);
+ waitUntil(() -> lock.getQueuedWorkerCountForTesting() == 2);
+
+ Future<?> worker3 = submitWorker(lock, executor, acquisitionOrder, 3, 3L);
+ waitUntil(() -> lock.getQueuedWorkerCountForTesting() == 3);
+
+ lock.release();
+
+ worker1.get(5, TimeUnit.SECONDS);
+ worker2.get(5, TimeUnit.SECONDS);
+ worker3.get(5, TimeUnit.SECONDS);
+ // Smaller priority value is granted first despite later arrival.
+ assertThat(acquisitionOrder).containsExactly(2, 3, 1);
+ } finally {
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ void samePriorityWorkersAllAcquireInArbitraryOrder() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ ExecutorService executor = Executors.newFixedThreadPool(3);
+ List<Integer> acquisitionOrder = Collections.synchronizedList(new ArrayList<>());
+
+ lock.acquireByMain();
+ try {
+ Future<?> worker1 = submitWorker(lock, executor, acquisitionOrder, 1, 7L);
+ waitUntil(() -> lock.getQueuedWorkerCountForTesting() == 1);
+
+ Future<?> worker2 = submitWorker(lock, executor, acquisitionOrder, 2, 7L);
+ waitUntil(() -> lock.getQueuedWorkerCountForTesting() == 2);
+
+ Future<?> worker3 = submitWorker(lock, executor, acquisitionOrder, 3, 7L);
+ waitUntil(() -> lock.getQueuedWorkerCountForTesting() == 3);
+
+ lock.release();
+
+ worker1.get(5, TimeUnit.SECONDS);
+ worker2.get(5, TimeUnit.SECONDS);
+ worker3.get(5, TimeUnit.SECONDS);
+ // Equal-priority workers have no defined grant order; all must acquire exactly once.
+ assertThat(acquisitionOrder).containsExactlyInAnyOrder(1, 2, 3);
+ } finally {
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ void mailboxAcquireHasPriorityOverHighPriorityWorker() throws Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ ExecutorService executor = Executors.newFixedThreadPool(2);
+ List<String> acquisitionOrder = Collections.synchronizedList(new ArrayList<>());
+
+ lock.acquireByMain();
+ try {
+ Future<?> worker =
+ executor.submit(
+ () -> {
+ // Even the smallest priority value must not overtake the mailbox.
+ lock.acquireByWorker(Long.MIN_VALUE, Long.MIN_VALUE);
+ try {
+ acquisitionOrder.add("worker");
+ } finally {
+ lock.release();
+ }
+ return null;
+ });
+ waitUntil(() -> lock.getQueuedWorkerCountForTesting() == 1);
+
+ Future<?> mailbox =
+ executor.submit(
+ () -> {
+ lock.acquireByMain();
+ try {
+ acquisitionOrder.add("mailbox");
+ } finally {
+ lock.release();
+ }
+ return null;
+ });
+ waitUntil(lock::isMailboxWaitingForTesting);
+
+ lock.release();
+
+ mailbox.get(5, TimeUnit.SECONDS);
+ worker.get(5, TimeUnit.SECONDS);
+ assertThat(acquisitionOrder).containsExactly("mailbox", "worker");
+ } finally {
+ executor.shutdownNow();
+ }
+ }
+
+ // -----------------------------------------------------------------------
+ // Helpers
+ // -----------------------------------------------------------------------
+
+ private static Future<?> submitWorker(
+ ParallelExecutionLock lock,
+ ExecutorService executor,
+ List<Integer> acquisitionOrder,
+ int workerId) {
+ return executor.submit(
+ () -> {
+ lock.acquireByWorker(0, 0);
+ try {
+ acquisitionOrder.add(workerId);
+ } finally {
+ lock.release();
+ }
+ return null;
+ });
+ }
+
+ private static Future<?> submitWorker(
+ ParallelExecutionLock lock,
+ ExecutorService executor,
+ List<Integer> acquisitionOrder,
+ int workerId,
+ long priority) {
+ return executor.submit(
+ () -> {
+ lock.acquireByWorker(priority, 0L);
+ try {
+ acquisitionOrder.add(workerId);
+ } finally {
+ lock.release();
+ }
+ return null;
+ });
+ }
+
+ private static void waitUntil(BooleanSupplier condition) throws Exception {
+ long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
+ while (!condition.getAsBoolean()) {
+ if (System.nanoTime() > deadline) {
+ fail("Condition was not met before timeout.");
+ }
+ Thread.sleep(10L);
+ }
+ }
+}
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionPythonResourceTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionPythonResourceTest.java
new file mode 100644
index 00000000..c43da10d
--- /dev/null
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionPythonResourceTest.java
@@ -0,0 +1,297 @@
+/*
+ * 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.operator.parallel;
+
+import org.apache.flink.agents.api.resource.Resource;
+import org.apache.flink.agents.api.resource.ResourceContext;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
+import org.apache.flink.agents.api.resource.ResourceType;
+import org.apache.flink.agents.api.vectorstores.python.PythonVectorStore;
+import org.apache.flink.agents.plan.resourceprovider.PythonResourceProvider;
+import org.apache.flink.agents.plan.resourceprovider.ResourceProvider;
+import org.apache.flink.agents.runtime.PythonMCPResourceDiscovery;
+import org.apache.flink.agents.runtime.ResourceCache;
+import org.apache.flink.agents.runtime.python.utils.PythonInterpreterManager;
+import org.apache.flink.agents.runtime.python.utils.PythonResourceAdapterImpl;
+import org.apache.flink.util.function.ThrowingRunnable;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+import pemja.core.PythonInterpreter;
+import pemja.core.object.PyObject;
+
+import javax.annotation.Nullable;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Queue;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/** Regression tests for Python resource use on parallel action workers. */
+class ParallelExecutionPythonResourceTest {
+
+ private static final Object KEY = "key";
+ private static final String CREATE_RESOURCE = "python_java_utils.create_resource";
+ private static final String FROM_JAVA_RESOURCE = "python_java_utils.from_java_resource";
+
+ @Test
+ @Timeout(15)
+ void pythonObjectsReturnedToJavaAreConsumedAndClosedOnTheirOwningWorker() throws Exception {
+ PythonInterpreter ownerInterpreter = mock(PythonInterpreter.class);
+ PythonInterpreter actionInterpreter = mock(PythonInterpreter.class);
+ PyObject vectorStoreHandle = mock(PyObject.class);
+ PyObject pythonDocument = mock(PyObject.class);
+ AtomicReference<Thread> resultThread = new AtomicReference<>();
+ AtomicReference<Thread> accessThread = new AtomicReference<>();
+ AtomicReference<Thread> closeThread = new AtomicReference<>();
+ AtomicReference<Thread> interpreterCloseThread = new AtomicReference<>();
+
+ doAnswer(
+ invocation -> {
+ resultThread.set(Thread.currentThread());
+ return List.of(pythonDocument);
+ })
+ .when(actionInterpreter)
+ .invoke("python_java_utils.call_method", vectorStoreHandle, "get", Map.of());
+ when(pythonDocument.getAttr("content"))
+ .thenAnswer(
+ invocation -> {
+ accessThread.compareAndSet(null, Thread.currentThread());
+ return "content";
+ });
+ when(pythonDocument.getAttr("metadata", Map.class))
+ .thenAnswer(
+ invocation -> {
+ accessThread.compareAndSet(null, Thread.currentThread());
+ return Map.of("source", "test");
+ });
+ when(pythonDocument.getAttr("id"))
+ .thenAnswer(
+ invocation -> {
+ accessThread.compareAndSet(null, Thread.currentThread());
+ return "doc-1";
+ });
+ doAnswer(
+ invocation -> {
+ closeThread.set(Thread.currentThread());
+ return null;
+ })
+ .when(pythonDocument)
+ .close();
+ doAnswer(
+ invocation -> {
+ interpreterCloseThread.set(Thread.currentThread());
+ return null;
+ })
+ .when(actionInterpreter)
+ .close();
+
+ try (PythonInterpreterManager interpreterManager =
+ new PythonInterpreterManager(ownerInterpreter, () -> actionInterpreter, 1)) {
+ PythonResourceAdapterImpl adapter =
+ new PythonResourceAdapterImpl(
+ mock(ResourceContext.class), interpreterManager, null);
+ PythonVectorStore vectorStore =
+ new PythonVectorStore(
+ adapter,
+ vectorStoreHandle,
+ new ResourceDescriptor("test.module", "TestVectorStore", Map.of()),
+ mock(ResourceContext.class));
+
+ runOnParallelWorker(
+ () -> vectorStore.get(null, null, null, null, Map.of()),
+ interpreterManager::releaseCurrentThreadInterpreter);
+
+ assertThat(resultThread.get()).isNotNull();
+ assertThat(accessThread.get())
+ .as("PyObject attributes must be read on the interpreter-owning worker")
+ .isSameAs(resultThread.get());
+ assertThat(closeThread.get())
+ .as("PyObject references must be closed on the interpreter-owning worker")
+ .isSameAs(resultThread.get());
+ assertThat(interpreterCloseThread.get())
+ .as("worker interpreter must be closed before the manager")
+ .isSameAs(resultThread.get());
+ }
+ }
+
+ @Test
+ @Timeout(15)
+ void pythonResourceOpenCanReenterResourceCacheOnParallelWorker() throws Exception {
+ PythonInterpreter ownerInterpreter = mock(PythonInterpreter.class);
+ PythonInterpreter actionInterpreter = mock(PythonInterpreter.class);
+ PyObject modelHandle = mock(PyObject.class);
+ PythonResourceProvider modelProvider =
+ new PythonResourceProvider(
+ "model",
+ ResourceType.CHAT_MODEL,
+ new ResourceDescriptor("test.module", "TestChatModel", Map.of()));
+ Resource connection =
+ new Resource() {
+ @Override
+ public ResourceType getResourceType() {
+ return ResourceType.CHAT_MODEL_CONNECTION;
+ }
+ };
+ ResourceProvider connectionProvider =
+ new ResourceProvider("connection", ResourceType.CHAT_MODEL_CONNECTION) {
+ @Override
+ public Resource provide(ResourceContext resourceContext) {
+ return connection;
+ }
+ };
+ Map<ResourceType, Map<String, ResourceProvider>> providers = new HashMap<>();
+ providers.put(ResourceType.CHAT_MODEL, Map.of("model", modelProvider));
+ providers.put(ResourceType.CHAT_MODEL_CONNECTION, Map.of("connection", connectionProvider));
+ ResourceCache resourceCache = new ResourceCache(providers);
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+
+ try (PythonInterpreterManager interpreterManager =
+ new PythonInterpreterManager(
+ ownerInterpreter, () -> actionInterpreter, 1, lock::checkReentrant)) {
+ PythonResourceAdapterImpl adapter =
+ new PythonResourceAdapterImpl(
+ resourceCache.getResourceContext(), interpreterManager, null);
+ PythonMCPResourceDiscovery.discoverPythonMCPResources(
+ providers, adapter, resourceCache);
+
+ Map<String, Object> createArguments = new HashMap<>();
+ createArguments.put("resource_context", null);
+ when(ownerInterpreter.invoke(
+ CREATE_RESOURCE, "test.module", "TestChatModel", createArguments))
+ .thenReturn(modelHandle);
+ Map<String, Object> conversionArguments = new HashMap<>();
+ conversionArguments.put("j_resource", connection);
+ conversionArguments.put("j_resource_adapter", null);
+ conversionArguments.put("resource_context", null);
+ when(actionInterpreter.invoke(
+ FROM_JAVA_RESOURCE,
+ ResourceType.CHAT_MODEL_CONNECTION.getValue(),
+ conversionArguments))
+ .thenReturn(new Object());
+ doAnswer(
+ invocation -> {
+ assertThat(Thread.holdsLock(resourceCache))
+ .as(
+ "Python resource open must run on the worker that"
+ + " owns the ResourceCache monitor")
+ .isTrue();
+ assertThat(
+ adapter.getResource(
+ "connection",
+ ResourceType.CHAT_MODEL_CONNECTION
+ .getValue()))
+ .isNotNull();
+ return null;
+ })
+ .when(actionInterpreter)
+ .invoke("python_java_utils.call_method", modelHandle, "open", Map.of());
+
+ runOnParallelWorker(
+ lock,
+ () -> resourceCache.getResource("model", ResourceType.CHAT_MODEL),
+ interpreterManager::releaseCurrentThreadInterpreter);
+
+ verify(ownerInterpreter)
+ .invoke(CREATE_RESOURCE, "test.module", "TestChatModel", createArguments);
+ verify(actionInterpreter, never())
+ .invoke(CREATE_RESOURCE, "test.module", "TestChatModel", createArguments);
+ resourceCache.close();
+ verify(modelHandle).close();
+ }
+ }
+
+ private static void runOnParallelWorker(
+ ThrowingRunnable<? extends Exception> action, Runnable threadCleanup) throws Exception {
+ runOnParallelWorker(new ParallelExecutionLock(), action, threadCleanup);
+ }
+
+ private static void runOnParallelWorker(
+ ParallelExecutionLock lock,
+ ThrowingRunnable<? extends Exception> action,
+ Runnable threadCleanup)
+ throws Exception {
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails = new LinkedBlockingQueue<>();
+ Queue<ParallelExecutionTask> works = new ConcurrentLinkedQueue<>();
+ works.add(new TestTask(action));
+
+ try (ParallelExecutionCoordinator coordinator =
+ new ParallelExecutionCoordinator(
+ lock,
+ (mail, description) -> mails.add(mail),
+ works::poll,
+ threadCleanup,
+ 1,
+ 3_600_000L)) {
+ lock.acquireByMain();
+ coordinator.addTask(KEY);
+ lock.release();
+
+ ThrowingRunnable<? extends Exception> mail = mails.poll(5, TimeUnit.SECONDS);
+ assertThat(mail).isNotNull();
+ mail.run();
+ }
+ }
+
+ private static final class TestTask implements ParallelExecutionTask {
+ private final ThrowingRunnable<? extends Exception> action;
+ private volatile boolean done;
+
+ private TestTask(ThrowingRunnable<? extends Exception> action) {
+ this.action = action;
+ }
+
+ @Override
+ public void setup(Object key, long recordIndex, long taskIndex) {}
+
+ @Override
+ public void restoreContext() {}
+
+ @Override
+ public void execute() {
+ try {
+ action.run();
+ done = true;
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ @Override
+ public boolean isDone() {
+ return done;
+ }
+
+ @Override
+ public void commit() {}
+
+ @Override
+ public void finishGroup(@Nullable ParallelExecutionTask lastCommitted) {}
+ }
+}
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionTaskQueueTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionTaskQueueTest.java
new file mode 100644
index 00000000..1a23681f
--- /dev/null
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionTaskQueueTest.java
@@ -0,0 +1,194 @@
+/*
+ * 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.operator.parallel;
+
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link ParallelExecutionTaskQueue}. */
+class ParallelExecutionTaskQueueTest {
+
+ /**
+ * A queue whose lock is held by the test thread for the whole test, so every lock-asserted call
+ * can be driven single-threaded.
+ */
+ private static ParallelExecutionTaskQueue newSchedule() {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ try {
+ lock.acquireByMain();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new AssertionError("Interrupted while acquiring the test lock", e);
+ }
+ return new ParallelExecutionTaskQueue(lock);
+ }
+
+ @Test
+ void addTaskCreatesGroupAndAssignsRecordIndexAndPerKeyTaskIndex() {
+ ParallelExecutionTaskQueue schedule = newSchedule();
+ schedule.addTask("a");
+ schedule.addTask("a");
+
+ ParallelExecutionTaskQueue.Node first = schedule.pollPendingNode();
+ assertThat(first.getKey()).isEqualTo("a");
+ assertThat(first.getRecordIndex()).isEqualTo(0L);
+ assertThat(first.getTaskIndex()).isEqualTo(0L);
+
+ ParallelExecutionTaskQueue.Node second = schedule.pollPendingNode();
+ assertThat(second.getRecordIndex()).isEqualTo(0L);
+ assertThat(second.getTaskIndex()).isEqualTo(1L);
+ }
+
+ @Test
+ void pollPendingServesLowestRecordIndexThenTaskIndex() {
+ ParallelExecutionTaskQueue schedule = newSchedule();
+ // Key group recordIndex is assigned on the first addTask: "a" first (recordIndex 0), then
+ // "b" (recordIndex 1). The interleaved third add for "a" still sorts by (recordIndex,
+ // taskIndex), not add order.
+ schedule.addTask("a");
+ schedule.addTask("b");
+ schedule.addTask("a");
+
+ ParallelExecutionTaskQueue.Node n0 = schedule.pollPendingNode();
+ assertThat(n0.getKey()).isEqualTo("a");
+ assertThat(n0.getRecordIndex()).isEqualTo(0L);
+ assertThat(n0.getTaskIndex()).isEqualTo(0L);
+ ParallelExecutionTaskQueue.Node n1 = schedule.pollPendingNode();
+ assertThat(n1.getKey()).isEqualTo("a");
+ assertThat(n1.getTaskIndex()).isEqualTo(1L);
+ ParallelExecutionTaskQueue.Node n2 = schedule.pollPendingNode();
+ assertThat(n2.getKey()).isEqualTo("b");
+ assertThat(n2.getRecordIndex()).isEqualTo(1L);
+ assertThat(n2.getTaskIndex()).isEqualTo(0L);
+ }
+
+ @Test
+ void hasOutstandingStaysTrueFromAddTaskThroughCommit() {
+ ParallelExecutionTaskQueue schedule = newSchedule();
+ assertThat(schedule.hasOutstanding("a")).isFalse();
+
+ schedule.addTask("a");
+ assertThat(schedule.hasOutstanding("a")).isTrue();
+
+ ParallelExecutionTaskQueue.Node node = schedule.pollPendingNode();
+ // Node pulled off the pull queue but still in its group: must remain outstanding.
+ assertThat(schedule.hasOutstanding("a")).isTrue();
+
+ schedule.markDone(node);
+ assertThat(schedule.hasOutstanding("a")).isTrue();
+
+ // Committing the last node drops the group, so the key is no longer outstanding.
+ schedule.removeCommitted(node);
+ assertThat(schedule.hasOutstanding("a")).isFalse();
+ }
+
+ @Test
+ void hasOutstandingCountsEachNode() {
+ ParallelExecutionTaskQueue schedule = newSchedule();
+ schedule.addTask("a");
+ schedule.addTask("a");
+ ParallelExecutionTaskQueue.Node a0 = schedule.pollPendingNode();
+ ParallelExecutionTaskQueue.Node a1 = schedule.pollPendingNode();
+ schedule.markDone(a0);
+ schedule.markDone(a1);
+
+ schedule.removeCommitted(a0);
+ assertThat(schedule.hasOutstanding("a")).isTrue();
+ schedule.removeCommitted(a1);
+ assertThat(schedule.hasOutstanding("a")).isFalse();
+ }
+
+ @Test
+ void groupIsDroppedWhenLastNodeCommitted() {
+ ParallelExecutionTaskQueue schedule = newSchedule();
+ schedule.addTask("a");
+ ParallelExecutionTaskQueue.Node node = schedule.pollPendingNode();
+ schedule.markDone(node);
+ schedule.removeCommitted(node);
+
+ assertThat(schedule.hasOutstanding("a")).isFalse();
+ assertThat(schedule.sizeForTesting()).isZero();
+
+ // Re-adding the same key mints a fresh group with the next arrival recordIndex.
+ schedule.addTask("a");
+ assertThat(schedule.pollPendingNode().getRecordIndex()).isEqualTo(1L);
+ }
+
+ @Test
+ void nodesAreBornRunning() {
+ ParallelExecutionTaskQueue schedule = newSchedule();
+ schedule.addTask("a");
+ assertThat(schedule.pollPendingNode().getState())
+ .isEqualTo(ParallelExecutionTaskQueue.State.RUNNING);
+ }
+
+ @Test
+ void markDoneRequiresRunningNode() {
+ ParallelExecutionTaskQueue schedule = newSchedule();
+ schedule.addTask("a");
+ ParallelExecutionTaskQueue.Node node = schedule.pollPendingNode();
+ schedule.markDone(node);
+ // Already DONE: a second markDone is rejected.
+ assertThatThrownBy(() -> schedule.markDone(node))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("RUNNING");
+ }
+
+ @Test
+ void removeCommittedRequiresDoneNodeAndRemovesIt() {
+ ParallelExecutionTaskQueue schedule = newSchedule();
+ schedule.addTask("a");
+ schedule.addTask("a");
+ ParallelExecutionTaskQueue.Node a0 = schedule.pollPendingNode();
+
+ assertThatThrownBy(() -> schedule.removeCommitted(a0))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("DONE");
+
+ schedule.markDone(a0);
+ schedule.removeCommitted(a0);
+ assertThat(a0.getState()).isEqualTo(ParallelExecutionTaskQueue.State.COMMITTED);
+ assertThat(schedule.sizeForTesting()).isEqualTo(1);
+ }
+
+ @Test
+ void pollCommittableHeadStopsAtHeadNotYetDone() {
+ ParallelExecutionTaskQueue schedule = newSchedule();
+ schedule.addTask("a");
+ schedule.addTask("a");
+ ParallelExecutionTaskQueue.Node a0 = schedule.pollPendingNode();
+ ParallelExecutionTaskQueue.Node a1 = schedule.pollPendingNode();
+
+ // Out-of-order completion: a1 is DONE, but the drain must stop at the not-yet-done head a0.
+ schedule.markDone(a1);
+ assertThat(schedule.pollCommittableHead("a")).isNull();
+
+ // Once a0 is DONE, the mailbox drain commits the contiguous done prefix a0, a1 in order.
+ schedule.markDone(a0);
+ ParallelExecutionTaskQueue.Node head = schedule.pollCommittableHead("a");
+ assertThat(head).isSameAs(a0);
+ schedule.removeCommitted(head);
+ head = schedule.pollCommittableHead("a");
+ assertThat(head).isSameAs(a1);
+ schedule.removeCommitted(head);
+ assertThat(schedule.pollCommittableHead("a")).isNull();
+ assertThat(schedule.sizeForTesting()).isZero();
+ }
+}
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonInterpreterManagerTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonInterpreterManagerTest.java
index 5df8d9f2..9df1ec6b 100644
--- a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonInterpreterManagerTest.java
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonInterpreterManagerTest.java
@@ -38,6 +38,7 @@ import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
/** Defect-oriented concurrency tests for {@link PythonInterpreterManager}. */
class PythonInterpreterManagerTest {
@@ -72,6 +73,51 @@ class PythonInterpreterManagerTest {
}
}
+ @Test
+ void authorizedNonOwnerThreadCanInvokeOwnerInterpreter() throws Exception {
+ PythonInterpreter owner = mock(PythonInterpreter.class);
+ AtomicInteger accessChecks = new AtomicInteger();
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ when(owner.invoke("owner-call")).thenReturn("result");
+
+ try (PythonInterpreterManager manager =
+ new PythonInterpreterManager(
+ owner,
+ () -> {
+ throw new AssertionError(
+ "owner invocation created a worker interpreter");
+ },
+ 1,
+ accessChecks::incrementAndGet)) {
+ assertThat(executor.submit(() -> manager.invokeOnOwner("owner-call")).get())
+ .isEqualTo("result");
+
+ assertThat(accessChecks).hasValue(1);
+ verify(owner).invoke("owner-call");
+ } finally {
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ void nonOwnerThreadCannotInvokeOwnerInterpreterWithoutAccess() throws Exception {
+ PythonInterpreter owner = mock(PythonInterpreter.class);
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+
+ try (PythonInterpreterManager manager =
+ new PythonInterpreterManager(owner, () -> mock(PythonInterpreter.class))) {
+ Future<?> invocation = executor.submit(() -> manager.invokeOnOwner("owner-call"));
+
+ assertThatThrownBy(() -> invocation.get(5, TimeUnit.SECONDS))
+ .hasCauseInstanceOf(IllegalStateException.class)
+ .hasRootCauseMessage(
+ "Owner interpreter access from a non-owner thread requires an execution checker.");
+ verify(owner, never()).invoke("owner-call");
+ } finally {
+ executor.shutdownNow();
+ }
+ }
+
@Test
void boundsInterpretersCreatedForTransientUnmanagedThreads() throws Exception {
PythonInterpreter owner = mock(PythonInterpreter.class);