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 7d7519f9c [runtime][java] Keep worker lock ownership when an async
call leaves the interrupt flag set (#1298)
7d7519f9c is described below
commit 7d7519f9cdae66ad7794e9c62b15c2ea393a3a0f
Author: Sangkyoon Nam <[email protected]>
AuthorDate: Wed Oct 21 11:48:15 2026 +0900
[runtime][java] Keep worker lock ownership when an async call leaves the
interrupt flag set (#1298)
Generated-by: Claude Code 2.1.303 (Claude Opus 5.5)
---
.../runtime/async/ContinuationActionExecutor.java | 23 +++++-
.../parallel/ParallelExecutionCoordinator.java | 14 +++-
.../operator/parallel/ParallelExecutionLock.java | 5 ++
.../async/ContinuationActionExecutorTest.java | 71 +++++++++++++++++
.../parallel/ParallelExecutionCoordinatorTest.java | 89 ++++++++++++++++++++--
5 files changed, 192 insertions(+), 10 deletions(-)
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 6132ae039..d61fe390d 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
@@ -101,8 +101,7 @@ public class ContinuationActionExecutor {
try {
return supplier.get();
} finally {
- parallelExecutionLock.acquireByWorker(context.getRecordIndex(),
context.getTaskIndex());
- contextRestorer.restore(context.getKey(), context.getActionTask());
+ reacquireAndRestore(context);
}
}
@@ -132,9 +131,27 @@ public class ContinuationActionExecutor {
try {
return batchExecutor.execute(suppliers, timeout, maxParallelism);
} finally {
+ reacquireAndRestore(context);
+ }
+ }
+
+ /**
+ * Re-acquires worker ownership and restores the runner context. An
interrupt flag left set by
+ * the call (for example by code that re-interrupts after catching {@link
InterruptedException})
+ * must not discard its outcome, so the flag is cleared for the wait and
restored afterward. A
+ * new interrupt during the wait, such as the coordinator's shutdown
interrupts, still aborts
+ * it.
+ */
+ private void reacquireAndRestore(ContinuationContext context) throws
InterruptedException {
+ boolean interrupted = Thread.interrupted();
+ try {
parallelExecutionLock.acquireByWorker(context.getRecordIndex(),
context.getTaskIndex());
- contextRestorer.restore(context.getKey(), context.getActionTask());
+ } finally {
+ if (interrupted) {
+ Thread.currentThread().interrupt();
+ }
}
+ contextRestorer.restore(context.getKey(), context.getActionTask());
}
/**
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
index 993188223..faef8f850 100644
---
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
@@ -191,7 +191,9 @@ public final class ParallelExecutionCoordinator implements
AutoCloseable {
}
// 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.
+ // throwable would only kill this pool thread silently), and the
lock is returned
+ // unless a shutdown interrupt aborted the work's async re-acquire.
+ boolean ownsLock = true;
try {
ParallelExecutionTaskQueue.Node node =
schedule.pollPendingNode();
ParallelExecutionTask work = workFactory.get();
@@ -201,6 +203,12 @@ public final class ParallelExecutionCoordinator implements
AutoCloseable {
// execute() never throws: it captures pull/prepare/run
failures and rethrows at
// commit.
work.execute();
+ if (workerExecutor.isShutdown() &&
!parallelExecutionLock.isHeldByCurrentThread()) {
+ // A shutdown interrupt aborted the work's async
re-acquire, so this worker no
+ // longer owns the lock: abandon the task without
completing or releasing it.
+ ownsLock = false;
+ return;
+ }
schedule.markDone(node);
mailboxDispatcher.execute(() -> completeOnMailbox(node),
"complete action task");
} catch (Throwable t) {
@@ -210,7 +218,9 @@ public final class ParallelExecutionCoordinator implements
AutoCloseable {
},
"throw action task scheduling failure");
} finally {
- parallelExecutionLock.releaseToMain();
+ if (ownsLock) {
+ parallelExecutionLock.releaseToMain();
+ }
}
} finally {
submittedPermits.decrementAndGet();
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
index f5009586c..b82e8d5ee 100644
---
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
@@ -176,6 +176,11 @@ public final class ParallelExecutionLock {
}
}
+ /** Whether the calling thread currently owns the lock. */
+ boolean isHeldByCurrentThread() {
+ return ownerThread == Thread.currentThread();
+ }
+
/** Asserts that the calling thread currently owns the lock. */
public void checkReentrant() {
if (ownerThread != Thread.currentThread()) {
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
index 0872d6ba8..4a196d708 100644
---
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
@@ -25,11 +25,13 @@ import org.junit.jupiter.api.Test;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
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.function.Supplier;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -102,6 +104,75 @@ class ContinuationActionExecutorTest {
}
}
+ @Test
+ void executeAsyncKeepsResultWhenCallLeavesInterruptFlagSet() throws
Throwable {
+ assertThat(
+ runInterruptingCall(
+ () -> {
+ // Code that catches an
InterruptedException and re-interrupts.
+ Thread.currentThread().interrupt();
+ return "result";
+ }))
+ .isEqualTo("result");
+ }
+
+ @Test
+ void executeAsyncKeepsFailureWhenCallLeavesInterruptFlagSet() {
+ IllegalStateException failure = new IllegalStateException("call
failed");
+ assertThatThrownBy(
+ () ->
+ runInterruptingCall(
+ () -> {
+ Thread.currentThread().interrupt();
+ throw failure;
+ }))
+ .isSameAs(failure);
+ }
+
+ /**
+ * Runs {@code call} through {@code executeAsync} on a worker that owns
the lock, and checks
+ * that the worker gets the lock and its context back with the interrupt
flag still set.
+ */
+ private static String runInterruptingCall(Supplier<String> call) throws
Throwable {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ AtomicBoolean restored = new AtomicBoolean();
+ AtomicBoolean interruptedAfterCall = new AtomicBoolean();
+ ContinuationActionExecutor executor =
+ new ContinuationActionExecutor(
+ 1,
+ () -> {},
+ lock,
+ (key, action) -> {
+ lock.checkReentrant();
+ restored.set(true);
+ });
+ ExecutorService workerExecutor = Executors.newSingleThreadExecutor();
+ try {
+ Future<String> workerResult =
+ workerExecutor.submit(
+ () -> {
+ lock.acquireByWorker(0, 0);
+ try {
+ return executor.executeAsync(new
ContinuationContext(), call);
+ } finally {
+
interruptedAfterCall.set(Thread.interrupted());
+ lock.release();
+ }
+ });
+ try {
+ return workerResult.get(5, TimeUnit.SECONDS);
+ } catch (ExecutionException e) {
+ throw e.getCause();
+ } finally {
+ assertThat(restored).isTrue();
+ assertThat(interruptedAfterCall).isTrue();
+ }
+ } finally {
+ workerExecutor.shutdownNow();
+ executor.close();
+ }
+ }
+
@Test
void executeAllAsyncReleasesAndReacquiresLockAndRestoresContext() throws
Exception {
ParallelExecutionLock lock = new ParallelExecutionLock();
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
index 271981df0..76bab05e1 100644
---
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
@@ -34,6 +34,7 @@ import java.util.concurrent.CountDownLatch;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Supplier;
import static org.assertj.core.api.Assertions.assertThat;
@@ -375,6 +376,57 @@ class ParallelExecutionCoordinatorTest {
}
}
+ @Test
+ @Timeout(15)
+ void shutdownDuringAsyncReacquireAbandonsTaskWithoutLockErrors() throws
Exception {
+ ParallelExecutionLock lock = new ParallelExecutionLock();
+ BlockingQueue<ThrowingRunnable<? extends Exception>> mails = new
LinkedBlockingQueue<>();
+ List<String> events = Collections.synchronizedList(new ArrayList<>());
+ AtomicReference<Throwable> uncaught = new AtomicReference<>();
+ CountDownLatch released = new CountDownLatch(1);
+ CountDownLatch mainHoldsLock = new CountDownLatch(1);
+ Supplier<ParallelExecutionTask> works =
+ () ->
+ FakeWork.asyncAbandonedOnInterrupt(
+ "A",
+ events,
+ lock,
+ () -> {
+ Thread.currentThread()
+ .setUncaughtExceptionHandler(
+ (thread, error) ->
+
uncaught.compareAndSet(null, error));
+ released.countDown();
+ await(mainHoldsLock);
+ });
+ ParallelExecutionCoordinator coordinator = newCoordinator(lock, mails,
works);
+ try {
+ lock.acquireByMain();
+ coordinator.addTask(KEY);
+ lock.release();
+ assertThat(released.await(5, TimeUnit.SECONDS)).isTrue();
+
+ // Hold the lock on the main thread so the worker's re-acquire has
to wait, then shut
+ // down, whose interrupt aborts that wait.
+ lock.acquireByMain();
+ try {
+ mainHoldsLock.countDown();
+ waitUntilQueuedWorkers(lock, 1);
+ coordinator.shutdown();
+ coordinator.close();
+ } finally {
+ lock.release();
+ }
+
+
assertThat(events).contains("abandon-A").doesNotContain("commit-A");
+ assertThat(mails).isEmpty();
+ assertThat(uncaught.get()).isNull();
+ } finally {
+ mainHoldsLock.countDown();
+ coordinator.close();
+ }
+ }
+
@Test
@Timeout(15)
void closeWaitsForInFlightWorkerBeforeReturning() throws Exception {
@@ -734,6 +786,17 @@ class ParallelExecutionCoordinatorTest {
}
}
+ private static void waitUntilQueuedWorkers(ParallelExecutionLock lock, int
expected)
+ throws InterruptedException {
+ long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
+ while (lock.getQueuedWorkerCountForTesting() != expected) {
+ assertThat(System.nanoTime())
+ .as("timed out waiting for %d queued workers", 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);
@@ -773,6 +836,7 @@ class ParallelExecutionCoordinatorTest {
@Nullable private final Runnable onExecute;
@Nullable private final ParallelExecutionLock asyncLock;
@Nullable private final RuntimeException commitFailure;
+ private final boolean abandonOnInterrupt;
private volatile boolean done;
private volatile Object assignedKey;
private volatile long assignedRecordIndex = -1L;
@@ -783,16 +847,18 @@ class ParallelExecutionCoordinatorTest {
List<String> events,
@Nullable Runnable onExecute,
@Nullable ParallelExecutionLock asyncLock,
- @Nullable RuntimeException commitFailure) {
+ @Nullable RuntimeException commitFailure,
+ boolean abandonOnInterrupt) {
this.id = id;
this.events = events;
this.onExecute = onExecute;
this.asyncLock = asyncLock;
this.commitFailure = commitFailure;
+ this.abandonOnInterrupt = abandonOnInterrupt;
}
static FakeWork blocking(String id, List<String> events, Runnable
onExecute) {
- return new FakeWork(id, events, onExecute, null, null);
+ return new FakeWork(id, events, onExecute, null, null, false);
}
/**
@@ -802,7 +868,7 @@ class ParallelExecutionCoordinatorTest {
*/
static FakeWork blockingAsync(
String id, List<String> events, ParallelExecutionLock lock,
Runnable onExecute) {
- return new FakeWork(id, events, onExecute, lock, null);
+ return new FakeWork(id, events, onExecute, lock, null, false);
}
static FakeWork blockingAsyncFailingCommit(
@@ -811,12 +877,21 @@ class ParallelExecutionCoordinatorTest {
ParallelExecutionLock lock,
Runnable onExecute,
RuntimeException commitFailure) {
- return new FakeWork(id, events, onExecute, lock, commitFailure);
+ return new FakeWork(id, events, onExecute, lock, commitFailure,
false);
+ }
+
+ /**
+ * Like {@link #blockingAsync}, but an interrupted re-acquire returns
without the lock, as
+ * the operator's work does when it records the failure for commit.
+ */
+ static FakeWork asyncAbandonedOnInterrupt(
+ String id, List<String> events, ParallelExecutionLock lock,
Runnable onExecute) {
+ return new FakeWork(id, events, onExecute, lock, null, true);
}
static FakeWork failingCommit(
String id, List<String> events, RuntimeException
commitFailure) {
- return new FakeWork(id, events, null, null, commitFailure);
+ return new FakeWork(id, events, null, null, commitFailure, false);
}
@Override
@@ -847,6 +922,10 @@ class ParallelExecutionCoordinatorTest {
asyncLock.acquireByWorker(assignedRecordIndex,
assignedTaskIndex);
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
+ if (abandonOnInterrupt) {
+ events.add("abandon-" + id);
+ return;
+ }
throw new AssertionError("Interrupted while
re-acquiring the lock", e);
}
}
