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 18556432e [runtime] Bound continuation inputs and prioritize queued
async work by record start order (#1318)
18556432e is described below
commit 18556432e95d0c8b245523f902267aa1cb999d0a
Author: Sangkyoon Nam <[email protected]>
AuthorDate: Fri Oct 23 10:53:08 2026 +0900
[runtime] Bound continuation inputs and prioritize queued async work by
record start order (#1318)
Generated-by: Claude Code 2.1.307 (Claude Opus 5.5)
---
.../agents/api/agents/AgentExecutionOptions.java | 8 +-
.../docs/development/create-agent/actions.md | 3 +
docs/content/docs/operations/configuration.md | 2 +-
.../integration/test/ParallelExecutionE2ETest.java | 10 +-
python/flink_agents/api/core_options.py | 9 +-
.../runtime/operator/ActionExecutionOperator.java | 58 ++++++-
.../runtime/operator/ActionTaskContextManager.java | 2 +-
.../runtime/async/ContinuationActionExecutor.java | 76 ++++++++-
.../agents/runtime/async/ContinuationContext.java | 21 ++-
.../ContinuationActionExecutorPriorityTest.java | 138 +++++++++++++++
.../operator/ActionExecutionOperatorTest.java | 189 +++++++++++++++++++++
11 files changed, 483 insertions(+), 33 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 97a7068d7..a98661d1e 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
@@ -44,10 +44,10 @@ public class AgentExecutionOptions {
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.
+ * Maximum number of input records that may be in flight concurrently on
each operator subtask,
+ * enforced by continuation-based execution and the JDK<21 parallel
execution engine; serial
+ * execution ignores 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);
diff --git a/docs/content/docs/development/create-agent/actions.md
b/docs/content/docs/development/create-agent/actions.md
index 3ee0881ee..d415197a5 100644
--- a/docs/content/docs/development/create-agent/actions.md
+++ b/docs/content/docs/development/create-agent/actions.md
@@ -785,6 +785,9 @@ worker and batch limits. Concurrency between Actions has
additional requirements
- **JDK 21+:** add `--add-exports=java.base/jdk.internal.vm=ALL-UNNAMED` to
the JVM
options, including `env.java.opts.all` on a cluster, to enable
continuation-based execution.
Plans with Python Actions need it too, not only plans whose Actions are all
Java.
+ Queued Java async calls are then prioritized by the order their input
records started
+ processing, then by the order their Actions were first dispatched, so a
record's follow-up
+ calls can run before later records' first calls. Running calls are not
preempted.
- **Java on JDK 11–20:** the experimental parallel execution engine is enabled
by default
for plans whose Actions are all Java. Setting `parallel-execution.enabled`
to `false`,
or including Python Actions, disables it. A Java single-call await then runs
synchronously;
diff --git a/docs/content/docs/operations/configuration.md
b/docs/content/docs/operations/configuration.md
index 8d969236f..a5efff5b9 100644
--- a/docs/content/docs/operations/configuration.md
+++ b/docs/content/docs/operations/configuration.md
@@ -211,7 +211,7 @@ the built-in [parallel tool-call batches]({{< ref
"docs/development/calling-llms
| `async.batch.parallelism` | os cpu count | int | Maximum concurrent calls in
one `ctx.gather(...)` batch, including ordinary, durable, and built-in tool
calls. Values of `1` or less run calls sequentially. Worker-pool capacity can
further limit concurrency. |
| `async.batch.timeout.ms` | -1 (disabled) | long (milliseconds) | Overall
timeout for one `ctx.gather(...)` batch, in milliseconds. Non-positive values
disable it. A timeout does not guarantee that running calls stop and cannot
undo their external effects. See [Tool Use]({{< ref
"docs/development/calling-llms/tool_use#timeouts-and-external-effects" >}}) for
built-in tool-batch behavior. |
| `job-identifier` | none | String | Stable job identifier used to scope Mem0
long-term memory. Keep an explicit value unchanged across restores. If unset,
the initial Flink JobID is stored in operator state and reused when restoring a
checkpoint or savepoint. |
-| `max-in-flight-input-records` | 100 | int | Maximum input records in flight
per Agent operator subtask on the experimental parallel engine. Must be
positive. The limit is ignored by continuation-based execution and plans
containing Python Actions. See [Execution Requirements]({{< ref
"docs/development/create-agent/actions#execution-requirements" >}}). |
+| `max-in-flight-input-records` | 100 | int | Maximum input records in flight
per Agent operator subtask on continuation-based execution and the experimental
parallel engine. Must be positive. When the limit is reached, the operator
admits the next input record only after an in-flight record finishes. Serial
execution ignores the limit. See [Execution Requirements]({{< ref
"docs/development/create-agent/actions#execution-requirements" >}}). |
| `num-async-threads` | os cpu count * 2 | int | Worker threads per async
execution pool on each Agent operator subtask. Must be positive. On the
experimental parallel engine the pool keeps one core worker and grows to the
larger of 10 and twice this value, because a call that waits inline holds its
worker while the work queued behind it draws from the same pool; exhausting
that pool stalls the job until Flink's checkpoint timeout fails it.
`async.batch.parallelism` separately limits con [...]
| `parallel-execution.enabled` | true | boolean | **Experimental.** Enable the
parallel engine for eligible plans with only Java Actions when
continuation-based execution is unavailable. `false` selects serial execution.
Ignored by the continuation engine and plans containing Python Actions. On JDK
11–20, Java `ctx.chat` and internal sub-agent calls need this engine, and an
internal child with a non-Java Action fails the job when the operator opens.
See [Execution Requirements]({{< ref " [...]
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
index c7c23ccd3..5af99608d 100644
---
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
@@ -27,12 +27,10 @@ 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;
@@ -56,8 +54,7 @@ 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.
+ * runs on both execution engines, including the input admission backpressure
scenario.
*
* <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
@@ -1080,11 +1077,6 @@ public class ParallelExecutionE2ETest {
@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(),
diff --git a/python/flink_agents/api/core_options.py
b/python/flink_agents/api/core_options.py
index aa1db965e..653d8b6d4 100644
--- a/python/flink_agents/api/core_options.py
+++ b/python/flink_agents/api/core_options.py
@@ -313,10 +313,11 @@ class AgentExecutionOptions:
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.
+ # Maximum number of input records that may be in flight concurrently on
+ # each operator subtask, enforced by continuation-based execution and the
+ # Java JDK<21 parallel execution engine; serial execution ignores it. At
+ # the cap, admission of further records waits until an in-flight record
+ # retires.
MAX_IN_FLIGHT_INPUT_RECORDS = ConfigOption(
key="max-in-flight-input-records",
config_type=int,
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 5ea003abd..a5bc2e4ef 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
@@ -96,6 +96,7 @@ import javax.annotation.Nullable;
import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.Collections;
+import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
@@ -254,10 +255,9 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
private transient boolean drainAborted;
/**
- * 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.
+ * Number of admitted input records not yet retired, including pending
inputs for busy keys,
+ * compared against {@link #maxInFlightInputRecords} to apply input
backpressure. Maintained on
+ * the mailbox thread and rebuilt from the recovered processing keys and
pending input events.
*/
private transient int inFlightInputRecords;
@@ -266,6 +266,17 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
*/
private int maxInFlightInputRecords;
+ /**
+ * Priority of each active input record by the order it started
processing, keyed by its key.
+ * Continuation execution passes it to the async executor so the work of
records that started
+ * earlier runs first. Set when an input event starts processing and
removed when it retires, so
+ * it also tells whether a key has an active input.
+ */
+ private transient Map<Object, RecordPriority> recordPriorities;
+
+ /** Next record index to assign; see {@link #recordPriorities}. */
+ private transient long nextRecordIndex;
+
// Long-term memory backed by Mem0; non-null only when
LongTermMemoryOptions.Mem0 is configured.
private transient Mem0LongTermMemory ltm;
@@ -386,6 +397,7 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
"%s must be positive, but was %s",
AgentExecutionOptions.MAX_IN_FLIGHT_INPUT_RECORDS.getKey(),
maxInFlightInputRecords);
+ recordPriorities = new HashMap<>();
parallelExecutionWithoutCoroutineEnabled =
shouldEnableParallelExecutionWithoutCoroutine(
@@ -522,6 +534,14 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
processElementInternal(record);
});
} else {
+ if (ContinuationActionExecutor.isContinuationSupported()) {
+ // Mails run while yielding can change the current key, so set
the record's key
+ // after the last yield.
+ while (inFlightInputRecords >= maxInFlightInputRecords) {
+ mailboxExecutor.yield();
+ }
+ super.setKeyContextElement1(record);
+ }
processElementInternal(record);
}
}
@@ -546,10 +566,12 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
try {
eventRouter.getKeySegmentQueue().addKeyToLastSegment(key);
+ // An input that triggers no Actions keeps its key active until
its retirement mail
+ // runs, so check the active inputs rather than the queued Action
tasks.
boolean currentKeyBusy =
parallelExecutionWithoutCoroutineEnabled
?
executionCoordinator.hasOutstanding(getCurrentKey())
- : stateManager.hasMoreActionTasks();
+ : recordPriorities.containsKey(key);
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
@@ -662,6 +684,7 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
freshRecordRound = true;
}
stateManager.addProcessingKey(key);
+ recordPriorities.put(key, new
RecordPriority(nextRecordIndex++));
stateManager.initOrIncSequenceNumber();
if (freshRecordRound) {
// Announced once the record's own bookkeeping is in place
and before anything
@@ -798,6 +821,7 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
notifyRecordFinished(key);
int removedCount = stateManager.removeProcessingKey(key);
inFlightInputRecords--;
+ recordPriorities.remove(key);
checkState(
removedCount == 1,
"Current processing key count for key "
@@ -832,6 +856,7 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
notifyRecordFinished(key);
int removedCount = stateManager.removeProcessingKey(key);
inFlightInputRecords--;
+ recordPriorities.remove(key);
checkState(
removedCount == 1,
"Current processing key count for key "
@@ -863,6 +888,10 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
throws Exception {
// 2. Invoke the action task.
long sequenceNumber = stateManager.getSequenceNumber();
+ // A resumed Action keeps its first-dispatch priority for queued async
work.
+ boolean firstDispatch =
+ !contextManager.hasContexts(actionTask)
+ || !contextManager.hasContinuationContext(actionTask);
contextManager.createAndSetRunnerContext(
actionTask,
contextKey,
@@ -876,6 +905,13 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
pythonBridge.getPythonRunnerContext(),
ltm,
this::createExecutionReportingContext);
+ ContinuationContext continuationContext =
contextManager.getContinuationContext(actionTask);
+ if (firstDispatch && continuationContext != null) {
+ RecordPriority priority = recordPriorities.get(key);
+ if (priority != null) {
+ continuationContext.setPriority(priority.recordIndex,
priority.nextTaskIndex++);
+ }
+ }
notifyActionPrepared(actionTask);
boolean isFinished;
@@ -1049,6 +1085,7 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
// we can proceed to process the next InputEvent.
int removedCount = stateManager.removeProcessingKey(key);
inFlightInputRecords--;
+ recordPriorities.remove(key);
checkState(
removedCount == 1,
"Current processing key count for key "
@@ -1661,6 +1698,7 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
continue;
}
eventRouter.getKeySegmentQueue().addKeyToLastSegment(key);
+ recordPriorities.put(key, new
RecordPriority(nextRecordIndex++));
// 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. Emitted before
// either engine dispatches, so no callback for this record
can precede it.
@@ -2148,4 +2186,14 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
super(message, cause);
}
}
+
+ /** Processing-start order of an active input record and the next index
for its Actions. */
+ private static final class RecordPriority {
+ private final long recordIndex;
+ private long nextTaskIndex;
+
+ private RecordPriority(long recordIndex) {
+ this.recordIndex = recordIndex;
+ }
+ }
}
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 b817bda44..bb29f39da 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
@@ -181,7 +181,7 @@ class ActionTaskContextManager implements AutoCloseable {
requireContexts(actionTask).lifecycleActive = active;
}
- private boolean hasContexts(ActionTask actionTask) {
+ boolean hasContexts(ActionTask actionTask) {
return actionTaskContexts.containsKey(actionTask);
}
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 4c650d853..74305526f 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
@@ -32,11 +32,15 @@ import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
+import java.util.concurrent.Executor;
import java.util.concurrent.Future;
+import java.util.concurrent.FutureTask;
+import java.util.concurrent.PriorityBlockingQueue;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicIntegerArray;
+import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Supplier;
import org.slf4j.Logger;
@@ -46,6 +50,9 @@ import org.slf4j.LoggerFactory;
* Executor for Java actions that supports asynchronous execution using JDK
21+ Continuation API.
*
* <p>This version uses {@code jdk.internal.vm.Continuation} to implement true
async execution.
+ * Queued async work runs in ascending (recordIndex, taskIndex) order of its
{@link
+ * ContinuationContext}, then in submission order, so the follow-up calls of
earlier input records
+ * don't wait behind the first calls of later ones.
*/
public class ContinuationActionExecutor {
@@ -53,8 +60,9 @@ public class ContinuationActionExecutor {
private static final ContinuationScope SCOPE = new
ContinuationScope("FlinkAgentsAction");
- private final ExecutorService asyncExecutor;
+ private final ThreadPoolExecutor asyncExecutor;
private final AsyncExecutorThreadFactory asyncThreadFactory;
+ private final AtomicLong submissionSequence = new AtomicLong();
public ContinuationActionExecutor(int numAsyncThreads) {
this(numAsyncThreads, () -> {}, null, null);
@@ -73,7 +81,13 @@ public class ContinuationActionExecutor {
LOG.info("Initialize fixed thread pool for async task with {}
threads", numAsyncThreads);
this.asyncThreadFactory = new
AsyncExecutorThreadFactory(threadCleanup);
this.asyncExecutor =
- Executors.newFixedThreadPool(numAsyncThreads,
asyncThreadFactory);
+ new ThreadPoolExecutor(
+ numAsyncThreads,
+ numAsyncThreads,
+ 0L,
+ TimeUnit.MILLISECONDS,
+ new PriorityBlockingQueue<>(),
+ asyncThreadFactory);
}
/**
@@ -149,8 +163,8 @@ public class ContinuationActionExecutor {
context.clearAsyncState();
// Submit task to thread pool and store the Future
- Future<?> future =
- asyncExecutor.submit(
+ FutureTask<?> future =
+ new FutureTask<>(
() -> {
try {
T result = supplier.get();
@@ -158,7 +172,9 @@ public class ContinuationActionExecutor {
} catch (Throwable t) {
context.getAsyncExceptionRef().set(t);
}
- });
+ },
+ null);
+ prioritizedExecutor(context).execute(future);
// Store the future reference before yielding (volatile write ensures
visibility)
context.setPendingFuture(future);
@@ -215,6 +231,7 @@ public class ContinuationActionExecutor {
int parallelismLimit = Math.min(Math.max(maxParallelism, 1),
batchSize);
long deadlineNanos = getDeadlineNanos(timeout);
+ Executor executor = prioritizedExecutor(context);
while (completed < batchSize) {
if (System.nanoTime() >= deadlineNanos) {
@@ -240,7 +257,7 @@ public class ContinuationActionExecutor {
} catch (Exception e) {
return Outcome.failure(e);
}
- }, asyncExecutor);
+ }, executor);
}
List<CompletableFuture<?>> inFlight = new
ArrayList<>(parallelismLimit);
@@ -275,6 +292,49 @@ public class ContinuationActionExecutor {
return collectBatchOutcomes(Arrays.asList(slots), started);
}
+ /** Returns an executor that queues work at the priority of {@code
context}. */
+ private Executor prioritizedExecutor(ContinuationContext context) {
+ long recordIndex = context.getRecordIndex();
+ long taskIndex = context.getTaskIndex();
+ return command ->
+ asyncExecutor.execute(
+ new PrioritizedTask(
+ command,
+ recordIndex,
+ taskIndex,
+ submissionSequence.getAndIncrement()));
+ }
+
+ /** Async work ordered by (recordIndex, taskIndex), then by submission. */
+ private static final class PrioritizedTask implements Runnable,
Comparable<PrioritizedTask> {
+ private final Runnable command;
+ private final long recordIndex;
+ private final long taskIndex;
+ private final long sequence;
+
+ private PrioritizedTask(Runnable command, long recordIndex, long
taskIndex, long sequence) {
+ this.command = command;
+ this.recordIndex = recordIndex;
+ this.taskIndex = taskIndex;
+ this.sequence = sequence;
+ }
+
+ @Override
+ public void run() {
+ command.run();
+ }
+
+ @Override
+ public int compareTo(PrioritizedTask other) {
+ int byRecord = Long.compare(recordIndex, other.recordIndex);
+ if (byRecord != 0) {
+ return byRecord;
+ }
+ int byTask = Long.compare(taskIndex, other.taskIndex);
+ return byTask != 0 ? byTask : Long.compare(sequence,
other.sequence);
+ }
+ }
+
private static <T> int countInFlight(
CompletableFuture<Outcome<T>>[] slots, int submittedCount) {
int inFlight = 0;
diff --git
a/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationContext.java
b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationContext.java
index 0523abb5b..78f0bd2d8 100644
---
a/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationContext.java
+++
b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationContext.java
@@ -22,10 +22,16 @@ import jdk.internal.vm.Continuation;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicReference;
-/** Continuation context with JDK 21 continuation state. */
+/**
+ * Continuation context with JDK 21 continuation state and the (recordIndex,
taskIndex) priority
+ * that {@link ContinuationActionExecutor} uses to order the async work
submitted for this context.
+ */
public class ContinuationContext {
private Continuation currentContinuation;
+ // Unprioritized contexts never overtake prioritized work.
+ private long recordIndex = Long.MAX_VALUE;
+ private long taskIndex = Long.MAX_VALUE;
private volatile Future<?> pendingFuture;
private volatile Future<?> pendingBatchFuture;
/** {@link Long#MAX_VALUE} means the pending batch has no deadline. */
@@ -42,6 +48,19 @@ public class ContinuationContext {
this.currentContinuation = currentContinuation;
}
+ public void setPriority(long recordIndex, long taskIndex) {
+ this.recordIndex = recordIndex;
+ this.taskIndex = taskIndex;
+ }
+
+ public long getRecordIndex() {
+ return recordIndex;
+ }
+
+ public long getTaskIndex() {
+ return taskIndex;
+ }
+
public Future<?> getPendingFuture() {
return pendingFuture;
}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutorPriorityTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutorPriorityTest.java
new file mode 100644
index 000000000..c9ed6aaf1
--- /dev/null
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutorPriorityTest.java
@@ -0,0 +1,138 @@
+/*
+ * 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.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.Assumptions.assumeTrue;
+
+/** Tests the order in which the JDK 21 continuation executor runs queued
async work. */
+class ContinuationActionExecutorPriorityTest {
+
+ @BeforeEach
+ void requireContinuations() {
+ assumeTrue(ContinuationActionExecutor.isContinuationSupported());
+ }
+
+ @Test
+ void queuedWorkRunsInRecordThenTaskOrder() throws Exception {
+ ContinuationActionExecutor executor = new
ContinuationActionExecutor(1);
+ try {
+ CountDownLatch blockerStarted = new CountDownLatch(1);
+ CountDownLatch releaseBlocker = new CountDownLatch(1);
+ List<String> runOrder = new CopyOnWriteArrayList<>();
+ List<ContinuationContext> contexts = new ArrayList<>();
+ List<Runnable> actions = new ArrayList<>();
+
+ // Occupies the only async thread so the work submitted below
queues up.
+ ContinuationContext blocker = context(0, 0);
+ contexts.add(blocker);
+ actions.add(
+ () ->
+ executeAsync(
+ executor,
+ blocker,
+ () -> {
+ blockerStarted.countDown();
+ await(releaseBlocker);
+ return "blocker";
+ }));
+ assertThat(executor.executeAction(blocker,
actions.get(0))).isFalse();
+ assertThat(blockerStarted.await(5, TimeUnit.SECONDS)).isTrue();
+
+ // Submitted in the reverse of the expected order. The batch goes
through
+ // executeAllAsync, the rest through executeAsync.
+ ContinuationContext laterRecord = context(2, 0);
+ contexts.add(laterRecord);
+ actions.add(
+ () -> {
+ try {
+ executor.executeAllAsync(
+ laterRecord,
+ List.of(() -> runOrder.add("record 2")),
+ Duration.ZERO,
+ 1);
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ });
+ ContinuationContext laterTask = context(1, 1);
+ contexts.add(laterTask);
+ actions.add(
+ () -> executeAsync(executor, laterTask, () ->
runOrder.add("record 1 task 1")));
+ ContinuationContext earliestTask = context(1, 0);
+ contexts.add(earliestTask);
+ actions.add(
+ () ->
+ executeAsync(
+ executor, earliestTask, () ->
runOrder.add("record 1 task 0")));
+ for (int i = 1; i < actions.size(); i++) {
+ assertThat(executor.executeAction(contexts.get(i),
actions.get(i))).isFalse();
+ }
+
+ releaseBlocker.countDown();
+ for (int i = 0; i < actions.size(); i++) {
+ long deadline = System.nanoTime() +
TimeUnit.SECONDS.toNanos(5);
+ while (!executor.executeAction(contexts.get(i),
actions.get(i))) {
+ assertThat(System.nanoTime()).isLessThan(deadline);
+ Thread.sleep(1);
+ }
+ }
+
+ assertThat(runOrder).containsExactly("record 1 task 0", "record 1
task 1", "record 2");
+ } finally {
+ executor.close();
+ }
+ }
+
+ private static ContinuationContext context(long recordIndex, long
taskIndex) {
+ ContinuationContext context = new ContinuationContext();
+ context.setPriority(recordIndex, taskIndex);
+ return context;
+ }
+
+ private static <T> void executeAsync(
+ ContinuationActionExecutor executor,
+ ContinuationContext context,
+ Supplier<T> supplier) {
+ try {
+ executor.executeAsync(context, supplier);
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ private static void await(CountDownLatch latch) {
+ try {
+ assertThat(latch.await(5, TimeUnit.SECONDS)).isTrue();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new AssertionError(e);
+ }
+ }
+}
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 1a89374f8..98f508061 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
@@ -70,7 +70,9 @@ import
org.apache.flink.agents.runtime.actionstate.KafkaActionStateStore;
import org.apache.flink.agents.runtime.async.BatchExecutionResult;
import org.apache.flink.agents.runtime.async.ContinuationActionExecutor;
import org.apache.flink.agents.runtime.async.ContinuationContext;
+import org.apache.flink.agents.runtime.context.JavaRunnerContextImpl;
import org.apache.flink.agents.runtime.context.RunnerContextImpl;
+import org.apache.flink.agents.runtime.lifecycle.TaskLifecycleListener;
import org.apache.flink.agents.runtime.memory.Mem0LongTermMemory;
import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl;
import
org.apache.flink.agents.runtime.operator.parallel.ParallelExecutionContextLifecycle;
@@ -120,6 +122,7 @@ import java.util.List;
import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Callable;
+import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
@@ -136,6 +139,7 @@ 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.junit.jupiter.api.Assumptions.assumeTrue;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.inOrder;
@@ -706,6 +710,191 @@ public class ActionExecutionOperatorTest {
}
}
+ @Test
+ @Timeout(10)
+ void nonParallelEnginesRetireInputsWithoutTriggeredActions() throws
Exception {
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.PARALLEL_EXECUTION_ENABLED, false);
+ 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();
+
+ // Continuation execution waits at the limit; serial execution
admits both inputs.
+ testHarness.processElement(new StreamRecord<>(1L));
+ testHarness.processElement(new StreamRecord<>(2L));
+ if (ContinuationActionExecutor.isContinuationSupported()) {
+
assertThat(operator.getOperatorStateManager().getProcessingKeys())
+ .containsExactly(2L);
+ } else {
+
assertThat(operator.getOperatorStateManager().getProcessingKeys())
+ .containsExactly(1L, 2L);
+ }
+ operator.waitInFlightEventsFinished();
+
+
assertThat(operator.getOperatorStateManager().getProcessingKeys()).isEmpty();
+ assertThat(testHarness.getRecordOutput()).isEmpty();
+ }
+ }
+
+ @Test
+ @Timeout(10)
+ void
continuationExecutionProcessesRecordAdmittedAtInFlightLimitUnderItsOwnKey()
+ throws Exception {
+ assumeTrue(ContinuationActionExecutor.isContinuationSupported());
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.PARALLEL_EXECUTION_ENABLED, false);
+ config.set(AgentExecutionOptions.MAX_IN_FLIGHT_INPUT_RECORDS, 1);
+
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(
+ TestAgent.getAgentPlanWithConfig(config),
true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>)
testHarness.getOperator();
+ List<Object> recordStartKeys = new CopyOnWriteArrayList<>();
+ operator.addTaskLifecycleListener(
+ new TaskLifecycleListener() {
+ @Override
+ public void onRecordStart(Object key) {
+ recordStartKeys.add(key);
+ }
+ });
+
+ // While the second record waits, the first record's mails set the
current key to 1L.
+ testHarness.processElement(new StreamRecord<>(1L));
+ testHarness.processElement(new StreamRecord<>(2L));
+ operator.waitInFlightEventsFinished();
+
+ assertThat(recordStartKeys).containsExactly(1L, 2L);
+ }
+ }
+
+ @Test
+ @Timeout(10)
+ void nonParallelEnginesKeepKeyBusyUntilNoActionInputRetires() throws
Exception {
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.PARALLEL_EXECUTION_ENABLED, false);
+ config.set(AgentExecutionOptions.MAX_IN_FLIGHT_INPUT_RECORDS, 2);
+ 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 -> 0L,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>)
testHarness.getOperator();
+
+ // The second input queues behind the first, which triggers no
Actions.
+ testHarness.processElement(new StreamRecord<>(1L));
+ testHarness.processElement(new StreamRecord<>(2L));
+
assertThat(operator.getOperatorStateManager().getProcessingKeys()).containsExactly(0L);
+
+ operator.waitInFlightEventsFinished();
+
assertThat(operator.getOperatorStateManager().getProcessingKeys()).isEmpty();
+ }
+ }
+
+ @Test
+ @Timeout(10)
+ void nonParallelEnginesAssignPrioritiesByRecordStartThenFirstDispatch()
throws Exception {
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.PARALLEL_EXECUTION_ENABLED, false);
+
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(
+ TestAgent.getAgentPlanWithConfig(config),
true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>)
testHarness.getOperator();
+ List<String> priorities = new CopyOnWriteArrayList<>();
+ operator.addTaskLifecycleListener(
+ new TaskLifecycleListener() {
+ @Override
+ public void onActionPrepared(ActionTask task) {
+ ContinuationContext context =
+ ((JavaRunnerContextImpl)
task.getRunnerContext())
+ .getContinuationContext();
+ priorities.add(
+ task.action.getName()
+ + " "
+ + context.getRecordIndex()
+ + "/"
+ + context.getTaskIndex());
+ }
+ });
+
+ // action1 handles the input and sends the event that triggers
action2.
+ testHarness.processElement(new StreamRecord<>(1L));
+ testHarness.processElement(new StreamRecord<>(2L));
+ operator.waitInFlightEventsFinished();
+
+ assertThat(priorities)
+ .containsExactlyInAnyOrder(
+ "action1 0/0", "action2 0/1", "action1 1/0",
"action2 1/1");
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(ints = {2, 3})
+ @Timeout(10)
+ void nonParallelEnginesAdvancePastPendingInputWithoutTriggeredActions(int
maxInFlight)
+ throws Exception {
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.PARALLEL_EXECUTION_ENABLED, false);
+ config.set(AgentExecutionOptions.MAX_IN_FLIGHT_INPUT_RECORDS,
maxInFlight);
+ Action action =
+ new Action(
+ "pending-events-single",
+ new JavaFunction(
+ TestAgent.class,
+ "pendingEventsSingleAction",
+ new Class<?>[] {Event.class,
RunnerContext.class}),
+ Collections.singletonList("type ==
EventType.InputEvent && input != 2"));
+ Map<String, Action> actions = new HashMap<>();
+ actions.put(action.getName(), action);
+ AgentPlan agentPlan = new AgentPlan(actions, new HashMap<>(), config);
+
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(agentPlan, true),
+ (KeySelector<Long, Long>) value -> 0L,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>)
testHarness.getOperator();
+
+ // Input 2 triggers no Actions; input 3 must still run after it.
+ testHarness.processElement(new StreamRecord<>(1L));
+ testHarness.processElement(new StreamRecord<>(2L));
+ testHarness.processElement(new StreamRecord<>(3L));
+
assertThat(operator.getOperatorStateManager().getProcessingKeys()).containsExactly(0L);
+
+ operator.endInput();
+ operator.waitInFlightEventsFinished();
+
+ assertThat(testHarness.getRecordOutput())
+ .extracting(StreamRecord::getValue)
+ .containsExactly("single-1", "single-3");
+
assertThat(operator.getOperatorStateManager().getProcessingKeys()).isEmpty();
+ }
+ }
+
@Test
@Timeout(30)
void backpressureBlocksAdmissionAtCapAndReleasesOnRetirement() throws
Exception {
