This is an automated email from the ASF dual-hosted git repository.
bowenli86 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push:
new 6de8055d00c [FLINK-40135][python] Prevent Beam state access after
runner shutdown (#28716)
6de8055d00c is described below
commit 6de8055d00c5c7afcc7ec68b8edd5e2c944fd2d2
Author: bowenli86 <[email protected]>
AuthorDate: Fri Oct 23 20:56:34 2026 -0700
[FLINK-40135][python] Prevent Beam state access after runner shutdown
(#28716)
---
.../apache/flink/python/PythonFunctionRunner.java | 10 +
.../AbstractExternalPythonFunctionOperator.java | 34 +-
.../python/beam/BeamPythonFunctionRunner.java | 432 +++++++++-
.../runners/python/beam/PythonSharedResources.java | 29 +-
.../python/beam/state/BeamStateRequestHandler.java | 69 +-
...AbstractExternalPythonFunctionOperatorTest.java | 176 ++++
.../python/beam/BeamPythonFunctionRunnerTest.java | 954 +++++++++++++++++++++
.../beam/state/BeamStateRequestHandlerTest.java | 176 ++++
8 files changed, 1811 insertions(+), 69 deletions(-)
diff --git
a/flink-python/src/main/java/org/apache/flink/python/PythonFunctionRunner.java
b/flink-python/src/main/java/org/apache/flink/python/PythonFunctionRunner.java
index c637c0e659c..708e8c70e8a 100644
---
a/flink-python/src/main/java/org/apache/flink/python/PythonFunctionRunner.java
+++
b/flink-python/src/main/java/org/apache/flink/python/PythonFunctionRunner.java
@@ -34,6 +34,16 @@ public interface PythonFunctionRunner extends AutoCloseable {
/** Tear-down the Python function runner. */
void close() throws Exception;
+ /**
+ * Cancels the Python function runner.
+ *
+ * <p>The default implementation keeps compatibility with runners that do
not need a separate
+ * cancellation path. Runners with a blocking graceful close should
override this method.
+ */
+ default void cancel() throws Exception {
+ close();
+ }
+
/**
* Executes the Python function with the input byte array.
*
diff --git
a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/process/AbstractExternalPythonFunctionOperator.java
b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/process/AbstractExternalPythonFunctionOperator.java
index da50b5335f2..f0367fb6b95 100644
---
a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/process/AbstractExternalPythonFunctionOperator.java
+++
b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/process/AbstractExternalPythonFunctionOperator.java
@@ -27,6 +27,7 @@ import
org.apache.flink.python.env.process.ProcessPythonEnvironmentManager;
import
org.apache.flink.streaming.api.operators.python.AbstractPythonFunctionOperator;
import
org.apache.flink.streaming.api.runners.python.beam.BeamPythonFunctionRunner;
import org.apache.flink.table.functions.python.PythonEnv;
+import org.apache.flink.util.concurrent.ExecutorThreadFactory;
import java.util.HashMap;
import java.util.concurrent.ExecutorService;
@@ -55,19 +56,35 @@ public abstract class
AbstractExternalPythonFunctionOperator<OUT>
super.open();
this.pythonFunctionRunner = createPythonFunctionRunner();
this.pythonFunctionRunner.open(config);
- this.flushThreadPool = Executors.newSingleThreadExecutor();
+ this.flushThreadPool =
+ Executors.newSingleThreadExecutor(
+ new ExecutorThreadFactory("beam-python-bundle-flush"));
}
@Override
public void close() throws Exception {
+ final PythonFunctionRunner functionRunner = pythonFunctionRunner;
+ final ExecutorService currentFlushThreadPool = flushThreadPool;
+ final boolean canceling =
+ getContainingTask().isCanceled() ||
getContainingTask().isFailing();
try {
- if (pythonFunctionRunner != null) {
- pythonFunctionRunner.close();
+ if (canceling && currentFlushThreadPool != null) {
+ currentFlushThreadPool.shutdownNow();
+ }
+
+ if (functionRunner != null) {
+ if (canceling) {
+ functionRunner.cancel();
+ } else {
+ functionRunner.close();
+ }
pythonFunctionRunner = null;
}
- if (flushThreadPool != null) {
- flushThreadPool.shutdown();
+ if (currentFlushThreadPool != null) {
+ if (!canceling) {
+ currentFlushThreadPool.shutdown();
+ }
flushThreadPool = null;
}
} finally {
@@ -78,24 +95,25 @@ public abstract class
AbstractExternalPythonFunctionOperator<OUT>
@Override
protected void invokeFinishBundle() throws Exception {
if (elementCount > 0) {
+ final PythonFunctionRunner functionRunner = pythonFunctionRunner;
AtomicBoolean flushThreadFinish = new AtomicBoolean(false);
AtomicReference<Throwable> exceptionReference = new
AtomicReference<>();
flushThreadPool.submit(
() -> {
try {
- pythonFunctionRunner.flush();
+ functionRunner.flush();
} catch (Throwable e) {
exceptionReference.set(e);
} finally {
flushThreadFinish.set(true);
// interrupt the progress of takeResult to avoid
the main thread is
// blocked forever.
- ((BeamPythonFunctionRunner)
pythonFunctionRunner).notifyNoMoreResults();
+ ((BeamPythonFunctionRunner)
functionRunner).notifyNoMoreResults();
}
});
Tuple3<String, byte[], Integer> resultTuple;
while (!flushThreadFinish.get()) {
- resultTuple = pythonFunctionRunner.takeResult();
+ resultTuple = functionRunner.takeResult();
if (resultTuple.f2 != 0) {
emitResult(resultTuple);
emitResults();
diff --git
a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java
b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java
index 7d89c93cc36..46d0ee1a6f4 100644
---
a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java
+++
b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java
@@ -39,9 +39,11 @@ import org.apache.flink.runtime.state.OperatorStateBackend;
import
org.apache.flink.streaming.api.operators.python.process.timer.TimerRegistration;
import
org.apache.flink.streaming.api.operators.python.process.timer.TimerRegistrationAction;
import
org.apache.flink.streaming.api.runners.python.beam.state.BeamStateRequestHandler;
+import org.apache.flink.util.ExceptionUtils;
import org.apache.flink.util.Preconditions;
import org.apache.flink.util.ShutdownHookUtil;
import org.apache.flink.util.TemporaryClassLoaderContext;
+import org.apache.flink.util.concurrent.ExecutorThreadFactory;
import org.apache.flink.util.function.LongFunctionWithException;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
@@ -55,7 +57,6 @@ import
org.apache.beam.runners.fnexecution.control.RemoteBundle;
import org.apache.beam.runners.fnexecution.control.StageBundleFactory;
import org.apache.beam.runners.fnexecution.control.TimerReceiverFactory;
import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
-import org.apache.beam.runners.fnexecution.state.StateRequestHandler;
import org.apache.beam.sdk.coders.ByteArrayCoder;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.fn.data.FnDataReceiver;
@@ -93,6 +94,11 @@ import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.function.BiConsumer;
import java.util.stream.Collectors;
@@ -165,7 +171,7 @@ public abstract class BeamPythonFunctionRunner implements
PythonFunctionRunner {
private transient StageBundleFactory stageBundleFactory;
/** Handler for state requests. */
- private transient StateRequestHandler stateRequestHandler;
+ private transient BeamStateRequestHandler stateRequestHandler;
/** Handler for bundle progress messages, both during bundle execution and
on its completion. */
private transient BundleProgressHandler progressHandler;
@@ -195,6 +201,32 @@ public abstract class BeamPythonFunctionRunner implements
PythonFunctionRunner {
/** The shared resource among Python operators of the same slot. */
private transient OpaqueMemoryResource<PythonSharedResources>
sharedResources;
+ /** Prevents duplicate teardown from explicit close and the JVM shutdown
hook. */
+ private transient boolean closed;
+
+ /** Coordinates duplicate close calls until teardown has finished. */
+ @Nullable private transient CompletableFuture<Void> closeCompletion;
+
+ /** Whether cancellation should skip waiting for an in-flight bundle
close. */
+ private transient boolean cancelRequested;
+
+ /** Coordinates resource teardown until lease disposal has finished. */
+ @Nullable private transient CompletableFuture<Void>
resourceCloseCompletion;
+
+ /** Coordinates duplicate close and cancel calls while the state handler
drains requests. */
+ @Nullable private transient CompletableFuture<Void>
stateHandlerCloseCompletion;
+
+ /** Owns a detached bundle close after cancellation returns to the task
cleanup path. */
+ @Nullable private transient ExecutorService
cancellationBundleCloseExecutor;
+
+ /**
+ * The bundle close currently owned by a flush thread.
+ *
+ * <p>The operation remains reachable while it is running so a concurrent
graceful close can
+ * wait for the same bundle close without attempting a second one.
+ */
+ @Nullable private transient BundleCloseOperation bundleCloseOperation;
+
private transient Thread shutdownHook;
private transient Environment environment;
@@ -328,45 +360,74 @@ public abstract class BeamPythonFunctionRunner implements
PythonFunctionRunner {
shutdownHook =
ShutdownHookUtil.addShutdownHook(
- this, BeamPythonFunctionRunner.class.getSimpleName(),
LOG);
+ this::cancel,
BeamPythonFunctionRunner.class.getSimpleName(), LOG);
unregisteredTimers = Collections.synchronizedList(new LinkedList<>());
}
@Override
public void close() throws Exception {
+ if (!startClose()) {
+ awaitCloseCompletion();
+ return;
+ }
+
try {
- if (jobBundleFactory != null) {
- jobBundleFactory.close();
+ try {
+ // Normal shutdown finishes this runner's bundle before
tearing down its factory or
+ // shared lease. Cancellation uses the separate non-blocking
path below.
+ flush();
+ } finally {
+ try {
+ closeResources();
+ } finally {
+ try {
+ closeStateRequestHandler();
+ } finally {
+ removeShutdownHook();
+ }
+ }
}
} finally {
- jobBundleFactory = null;
+ completeClose();
}
+ }
+
+ @Override
+ public void cancel() throws Exception {
+ final BundleCloseClaim bundleCloseClaim = requestCancel();
try {
- if (sharedResources != null) {
- sharedResources.close();
- } else {
- // if sharedResources is not null, the close of
environmentManager will be managed
- // in sharedResources,
- // otherwise, we need to close the environmentManager
explicitly
- environmentManager.close();
+ if (bundleCloseClaim != null && bundleCloseClaim.ownsBundleClose) {
+ closeBundleAfterCancellation(bundleCloseClaim.closeOperation);
}
+ // Release only this runner's resource lease. If it is the final
lease, the shared
+ // resource disposer owns factory teardown; otherwise co-located
runners keep using the
+ // shared worker.
+ closeResources();
} finally {
- sharedResources = null;
- }
-
- if (shutdownHook != null) {
- ShutdownHookUtil.removeShutdownHook(
- shutdownHook,
BeamPythonFunctionRunner.class.getSimpleName(), LOG);
- shutdownHook = null;
+ try {
+ // A bundle close already owned by another thread may still
issue callbacks while
+ // it unwinds. Gate those callbacks before state backends are
disposed.
+ closeStateRequestHandler();
+ } finally {
+ try {
+ removeShutdownHook();
+ } finally {
+ completeClose();
+ }
+ }
}
}
@Override
public void process(byte[] data) throws Exception {
- checkInvokeStartBundle();
- mainInputReceiver.accept(WindowedValues.valueInGlobalWindow(data));
+ final FnDataReceiver<WindowedValue<byte[]>> inputReceiver;
+ synchronized (this) {
+ checkInvokeStartBundle();
+ inputReceiver = mainInputReceiver;
+ }
+ inputReceiver.accept(WindowedValues.valueInGlobalWindow(data));
}
@Override
@@ -381,23 +442,34 @@ public abstract class BeamPythonFunctionRunner implements
PythonFunctionRunner {
@Override
public void processTimer(byte[] timerData) throws Exception {
- if (timerInputReceiver == null) {
- checkInvokeStartBundle();
- timerInputReceiver =
- Preconditions.checkNotNull(
-
Iterables.getOnlyElement(remoteBundle.getTimerReceivers().values()),
- "Failed to retrieve main input receiver.");
+ final FnDataReceiver<Timer> inputReceiver;
+ synchronized (this) {
+ if (timerInputReceiver == null) {
+ checkInvokeStartBundle();
+ timerInputReceiver =
+ Preconditions.checkNotNull(
+
Iterables.getOnlyElement(remoteBundle.getTimerReceivers().values()),
+ "Failed to retrieve main input receiver.");
+ }
+ inputReceiver = timerInputReceiver;
}
Timer<byte[]> timerValue = Timer.cleared(timerData, "",
Collections.emptyList());
- timerInputReceiver.accept(timerValue);
+ inputReceiver.accept(timerValue);
}
/** Checks whether to invoke startBundle. */
private void checkInvokeStartBundle() {
+ if (closed) {
+ throw new IllegalStateException("Beam Python function runner is
closed.");
+ }
if (!bundleStarted) {
+ if (bundleCloseOperation != null &&
!bundleCloseOperation.isDone()) {
+ throw new IllegalStateException("The previous Beam bundle is
still closing.");
+ }
startBundle();
bundleStarted = true;
+ bundleCloseOperation = null;
}
}
@@ -442,13 +514,34 @@ public abstract class BeamPythonFunctionRunner implements
PythonFunctionRunner {
@Override
public void flush() throws Exception {
- if (bundleStarted) {
- try {
- finishBundle();
- } finally {
+ final BundleCloseOperation closeOperation;
+ final boolean ownsBundleClose;
+ synchronized (this) {
+ if (cancelRequested) {
+ return;
+ }
+
+ if (bundleStarted) {
+ closeOperation = new BundleCloseOperation(remoteBundle);
+ bundleCloseOperation = closeOperation;
bundleStarted = false;
+ clearBundleReferences();
+ ownsBundleClose = true;
+ } else {
+ closeOperation = bundleCloseOperation;
+ ownsBundleClose = false;
}
}
+
+ if (closeOperation == null) {
+ return;
+ }
+
+ if (ownsBundleClose) {
+ finishBundle(closeOperation);
+ } else {
+ closeOperation.await();
+ }
}
/** Interrupts the progress of takeResult. */
@@ -456,15 +549,274 @@ public abstract class BeamPythonFunctionRunner
implements PythonFunctionRunner {
resultBuffer.add(Tuple2.of(null, new byte[0]));
}
- private void finishBundle() {
+ private void finishBundle(BundleCloseOperation closeOperation) {
+ RuntimeException failure = null;
try {
- remoteBundle.close();
+ closeOperation.remoteBundle.close();
} catch (Throwable t) {
- throw new RuntimeException("Failed to close remote bundle", t);
+ failure = new RuntimeException("Failed to close remote bundle", t);
+ throw failure;
} finally {
- remoteBundle = null;
- mainInputReceiver = null;
- timerInputReceiver = null;
+ closeOperation.complete(failure);
+ if (isCancelRequested()) {
+ try {
+ closeResources();
+ } catch (Exception e) {
+ LOG.warn("Failed to clean up Python resources after
cancellation.", e);
+ }
+ }
+ }
+ }
+
+ private synchronized boolean startClose() {
+ if (closed) {
+ return false;
+ }
+ closed = true;
+ closeCompletion = new CompletableFuture<>();
+ return true;
+ }
+
+ /**
+ * Starts cancellation without waiting for a concurrent bundle close.
+ *
+ * <p>If no flush thread owns the active bundle yet, cancellation detaches
it and becomes
+ * responsible for finishing it asynchronously.
+ */
+ @Nullable
+ private synchronized BundleCloseClaim requestCancel() {
+ cancelRequested = true;
+ closed = true;
+ if (closeCompletion == null) {
+ closeCompletion = new CompletableFuture<>();
+ }
+
+ if (bundleStarted) {
+ final BundleCloseOperation closeOperation = new
BundleCloseOperation(remoteBundle);
+ bundleCloseOperation = closeOperation;
+ bundleStarted = false;
+ clearBundleReferences();
+ return new BundleCloseClaim(closeOperation, true);
+ }
+ if (bundleCloseOperation != null && !bundleCloseOperation.isDone()) {
+ return new BundleCloseClaim(bundleCloseOperation, false);
+ }
+ clearBundleReferences();
+ return null;
+ }
+
+ private synchronized boolean isCancelRequested() {
+ return cancelRequested;
+ }
+
+ private void awaitCloseCompletion() throws Exception {
+ final CompletableFuture<Void> currentCloseCompletion;
+ synchronized (this) {
+ currentCloseCompletion = closeCompletion;
+ }
+ if (currentCloseCompletion != null) {
+ awaitCompletion(currentCloseCompletion);
+ }
+ }
+
+ private void completeClose() {
+ final CompletableFuture<Void> currentCloseCompletion;
+ synchronized (this) {
+ currentCloseCompletion = closeCompletion;
+ }
+ if (currentCloseCompletion != null) {
+ currentCloseCompletion.complete(null);
+ }
+ }
+
+ private void closeResources() throws Exception {
+ final JobBundleFactory ownedJobBundleFactory;
+ final OpaqueMemoryResource<PythonSharedResources>
currentSharedResources;
+ final CompletableFuture<Void> closeCompletion;
+ final boolean ownsClose;
+ synchronized (this) {
+ if (resourceCloseCompletion == null) {
+ closeCompletion = new CompletableFuture<>();
+ resourceCloseCompletion = closeCompletion;
+ ownedJobBundleFactory = jobBundleFactory;
+ currentSharedResources = sharedResources;
+ jobBundleFactory = null;
+ sharedResources = null;
+ ownsClose = true;
+ } else {
+ closeCompletion = resourceCloseCompletion;
+ ownedJobBundleFactory = null;
+ currentSharedResources = null;
+ ownsClose = false;
+ }
+ }
+
+ if (!ownsClose) {
+ // Disposal holds the slot's SharedResources lock. Even an
interrupted cancellation
+ // must wait for that lock to be released before task cleanup can
finish.
+ try {
+ closeCompletion.join();
+ } catch (CompletionException e) {
+ ExceptionUtils.rethrowException(e.getCause());
+ }
+ return;
+ }
+
+ try {
+ if (currentSharedResources != null) {
+ try {
+ currentSharedResources.close();
+ } finally {
+ clearBundleReferences();
+ }
+ } else {
+ try {
+ if (ownedJobBundleFactory != null) {
+ ownedJobBundleFactory.close();
+ }
+ } finally {
+ try {
+ environmentManager.close();
+ } finally {
+ clearBundleReferences();
+ }
+ }
+ }
+ closeCompletion.complete(null);
+ } catch (Throwable t) {
+ closeCompletion.completeExceptionally(t);
+ throw t;
+ }
+ }
+
+ private void closeBundleAfterCancellation(BundleCloseOperation
closeOperation) {
+ final ExecutorService executor;
+ synchronized (this) {
+ if (cancellationBundleCloseExecutor == null) {
+ cancellationBundleCloseExecutor =
+ Executors.newSingleThreadExecutor(
+ new
ExecutorThreadFactory("beam-python-bundle-cancellation"));
+ }
+ executor = cancellationBundleCloseExecutor;
+ }
+ executor.execute(
+ () -> {
+ try {
+ finishBundle(closeOperation);
+ } catch (Throwable t) {
+ LOG.debug("Beam bundle close failed during
cancellation.", t);
+ } finally {
+ executor.shutdown();
+ }
+ });
+ }
+
+ private void closeStateRequestHandler() throws Exception {
+ final BeamStateRequestHandler currentStateRequestHandler;
+ final CompletableFuture<Void> closeCompletion;
+ final boolean ownsClose;
+ synchronized (this) {
+ if (stateHandlerCloseCompletion == null) {
+ closeCompletion = new CompletableFuture<>();
+ stateHandlerCloseCompletion = closeCompletion;
+ currentStateRequestHandler = stateRequestHandler;
+ stateRequestHandler = null;
+ ownsClose = true;
+ } else {
+ closeCompletion = stateHandlerCloseCompletion;
+ currentStateRequestHandler = null;
+ ownsClose = false;
+ }
+ }
+
+ if (ownsClose) {
+ RuntimeException failure = null;
+ try {
+ if (currentStateRequestHandler != null) {
+ currentStateRequestHandler.close();
+ }
+ } catch (RuntimeException e) {
+ failure = e;
+ throw e;
+ } finally {
+ if (failure == null) {
+ closeCompletion.complete(null);
+ } else {
+ closeCompletion.completeExceptionally(failure);
+ }
+ }
+ } else {
+ awaitCompletion(closeCompletion);
+ }
+ }
+
+ private synchronized void clearBundleReferences() {
+ remoteBundle = null;
+ mainInputReceiver = null;
+ timerInputReceiver = null;
+ }
+
+ private void removeShutdownHook() {
+ final Thread currentShutdownHook;
+ synchronized (this) {
+ currentShutdownHook = shutdownHook;
+ shutdownHook = null;
+ }
+ if (currentShutdownHook != null) {
+ ShutdownHookUtil.removeShutdownHook(
+ currentShutdownHook,
BeamPythonFunctionRunner.class.getSimpleName(), LOG);
+ }
+ }
+
+ private static final class BundleCloseOperation {
+
+ private final RemoteBundle remoteBundle;
+ private final CompletableFuture<Void> completion = new
CompletableFuture<>();
+
+ private BundleCloseOperation(RemoteBundle remoteBundle) {
+ this.remoteBundle = remoteBundle;
+ }
+
+ private boolean isDone() {
+ return completion.isDone();
+ }
+
+ private void complete(@Nullable RuntimeException failure) {
+ if (failure == null) {
+ completion.complete(null);
+ } else {
+ completion.completeExceptionally(failure);
+ }
+ }
+
+ private void await() throws Exception {
+ awaitCompletion(completion);
+ }
+ }
+
+ private static final class BundleCloseClaim {
+
+ private final BundleCloseOperation closeOperation;
+ private final boolean ownsBundleClose;
+
+ private BundleCloseClaim(BundleCloseOperation closeOperation, boolean
ownsBundleClose) {
+ this.closeOperation = closeOperation;
+ this.ownsBundleClose = ownsBundleClose;
+ }
+ }
+
+ private static void awaitCompletion(CompletableFuture<Void> completion)
throws Exception {
+ try {
+ completion.get();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw e;
+ } catch (ExecutionException e) {
+ final Throwable cause = e.getCause();
+ if (cause instanceof Exception) {
+ throw (Exception) cause;
+ }
+ throw new RuntimeException(cause);
}
}
@@ -735,7 +1087,7 @@ public abstract class BeamPythonFunctionRunner implements
PythonFunctionRunner {
return new TimerReceiverFactory(stageBundleFactory, timerDataConsumer,
null);
}
- private static StateRequestHandler getStateRequestHandler(
+ private static BeamStateRequestHandler getStateRequestHandler(
@Nullable KeyedStateBackend<?> keyedStateBackend,
@Nullable OperatorStateBackend operatorStateBackend,
@Nullable TypeSerializer<?> keySerializer,
diff --git
a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/PythonSharedResources.java
b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/PythonSharedResources.java
index 293d5a5ad83..9eaf2fbb468 100644
---
a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/PythonSharedResources.java
+++
b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/PythonSharedResources.java
@@ -44,6 +44,8 @@ public final class PythonSharedResources implements
AutoCloseable {
/** Keep track of the PythonEnvironmentManagers of the Python operators in
one slot. */
private final List<PythonEnvironmentManager> environmentManagers;
+ private boolean closed;
+
PythonSharedResources(JobBundleFactory jobBundleFactory, Environment
environment) {
this.jobBundleFactory = jobBundleFactory;
this.environment = environment;
@@ -63,10 +65,31 @@ public final class PythonSharedResources implements
AutoCloseable {
}
@Override
- public void close() throws Exception {
- jobBundleFactory.close();
+ public synchronized void close() throws Exception {
+ if (closed) {
+ return;
+ }
+
+ Exception exception = null;
+ try {
+ jobBundleFactory.close();
+ } catch (Exception e) {
+ exception = e;
+ }
for (PythonEnvironmentManager environmentManager :
environmentManagers) {
- environmentManager.close();
+ try {
+ environmentManager.close();
+ } catch (Exception e) {
+ if (exception == null) {
+ exception = e;
+ } else {
+ exception.addSuppressed(e);
+ }
+ }
+ }
+ closed = true;
+ if (exception != null) {
+ throw exception;
}
}
}
diff --git
a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamStateRequestHandler.java
b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamStateRequestHandler.java
index fb6765dff20..34295b5a519 100644
---
a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamStateRequestHandler.java
+++
b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamStateRequestHandler.java
@@ -37,12 +37,14 @@ import java.util.Collections;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
/**
* The handler for Beam state requests sent from Python side, which does
actual operations on Flink
* state.
*/
-public class BeamStateRequestHandler implements StateRequestHandler {
+public class BeamStateRequestHandler implements StateRequestHandler,
AutoCloseable {
private final BeamStateStore keyedStateStore;
@@ -54,6 +56,10 @@ public class BeamStateRequestHandler implements
StateRequestHandler {
private final BeamFnApi.ProcessBundleRequest.CacheToken cacheToken;
+ private final ReentrantReadWriteLock lifecycleLock = new
ReentrantReadWriteLock(true);
+
+ private boolean closed;
+
public BeamStateRequestHandler(
BeamStateStore keyedStateStore,
BeamStateStore operatorStateStore,
@@ -69,23 +75,50 @@ public class BeamStateRequestHandler implements
StateRequestHandler {
@Override
public CompletionStage<BeamFnApi.StateResponse.Builder>
handle(BeamFnApi.StateRequest request)
throws Exception {
- BeamFnApi.StateKey.TypeCase typeCase =
request.getStateKey().getTypeCase();
- ListState<byte[]> listState;
- MapState<ByteArrayWrapper, byte[]> mapState;
-
- switch (typeCase) {
- case BAG_USER_STATE:
- listState = keyedStateStore.getListState(request);
- return CompletableFuture.completedFuture(
- bagStateHandler.handle(request, listState));
- case MULTIMAP_SIDE_INPUT:
- mapState = keyedStateStore.getMapState(request);
- return
CompletableFuture.completedFuture(mapStateHandler.handle(request, mapState));
- case MULTIMAP_KEYS_SIDE_INPUT:
- mapState = operatorStateStore.getMapState(request);
- return
CompletableFuture.completedFuture(mapStateHandler.handle(request, mapState));
- default:
- throw new RuntimeException("Unsupported state type: " +
typeCase);
+ final Lock readLock = lifecycleLock.readLock();
+ readLock.lock();
+ try {
+ if (closed) {
+ throw new IllegalStateException("Beam state request handler is
closed.");
+ }
+
+ BeamFnApi.StateKey.TypeCase typeCase =
request.getStateKey().getTypeCase();
+ ListState<byte[]> listState;
+ MapState<ByteArrayWrapper, byte[]> mapState;
+
+ switch (typeCase) {
+ case BAG_USER_STATE:
+ listState = keyedStateStore.getListState(request);
+ return CompletableFuture.completedFuture(
+ bagStateHandler.handle(request, listState));
+ case MULTIMAP_SIDE_INPUT:
+ mapState = keyedStateStore.getMapState(request);
+ return CompletableFuture.completedFuture(
+ mapStateHandler.handle(request, mapState));
+ case MULTIMAP_KEYS_SIDE_INPUT:
+ mapState = operatorStateStore.getMapState(request);
+ return CompletableFuture.completedFuture(
+ mapStateHandler.handle(request, mapState));
+ default:
+ throw new RuntimeException("Unsupported state type: " +
typeCase);
+ }
+ } finally {
+ readLock.unlock();
+ }
+ }
+
+ /**
+ * Stops accepting state requests and waits for requests that are already
accessing Flink state
+ * to complete.
+ */
+ @Override
+ public void close() {
+ final Lock writeLock = lifecycleLock.writeLock();
+ writeLock.lock();
+ try {
+ closed = true;
+ } finally {
+ writeLock.unlock();
}
}
diff --git
a/flink-python/src/test/java/org/apache/flink/streaming/api/operators/python/process/AbstractExternalPythonFunctionOperatorTest.java
b/flink-python/src/test/java/org/apache/flink/streaming/api/operators/python/process/AbstractExternalPythonFunctionOperatorTest.java
new file mode 100644
index 00000000000..fa0d307a562
--- /dev/null
+++
b/flink-python/src/test/java/org/apache/flink/streaming/api/operators/python/process/AbstractExternalPythonFunctionOperatorTest.java
@@ -0,0 +1,176 @@
+/*
+ * 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.streaming.api.operators.python.process;
+
+import org.apache.flink.api.java.tuple.Tuple3;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.python.PythonFunctionRunner;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.runtime.tasks.StreamTask;
+import org.apache.flink.table.functions.python.PythonEnv;
+
+import org.junit.jupiter.api.Test;
+
+import java.lang.reflect.Field;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+class AbstractExternalPythonFunctionOperatorTest {
+
+ @Test
+ void testCloseCancelsRunnerAfterStoppingFlushExecutor() throws Exception {
+ final StreamTask<?, ?> containingTask =
org.mockito.Mockito.mock(StreamTask.class);
+ org.mockito.Mockito.when(containingTask.isCanceled()).thenReturn(true);
+ final ExecutorService flushThreadPool =
Executors.newSingleThreadExecutor();
+ final TestingPythonFunctionRunner functionRunner =
+ new TestingPythonFunctionRunner(flushThreadPool);
+ final TestingExternalPythonFunctionOperator operator =
+ createOperator(containingTask, functionRunner,
flushThreadPool);
+
+ try {
+ operator.close();
+
+ assertThat(functionRunner.cancelCalled).isTrue();
+ assertThat(functionRunner.closeCalled).isFalse();
+
assertThat(functionRunner.flushExecutorStoppedBeforeCancel).isTrue();
+ assertThat(flushThreadPool.isShutdown()).isTrue();
+ } finally {
+ flushThreadPool.shutdownNow();
+ }
+ }
+
+ @Test
+ void testCloseClosesRunnerGracefullyForNormalCompletion() throws Exception
{
+ final StreamTask<?, ?> containingTask =
org.mockito.Mockito.mock(StreamTask.class);
+ final ExecutorService flushThreadPool =
Executors.newSingleThreadExecutor();
+ final TestingPythonFunctionRunner functionRunner =
+ new TestingPythonFunctionRunner(flushThreadPool);
+ final TestingExternalPythonFunctionOperator operator =
+ createOperator(containingTask, functionRunner,
flushThreadPool);
+
+ try {
+ operator.close();
+
+ assertThat(functionRunner.closeCalled).isTrue();
+ assertThat(functionRunner.cancelCalled).isFalse();
+ assertThat(flushThreadPool.isShutdown()).isTrue();
+ } finally {
+ flushThreadPool.shutdownNow();
+ }
+ }
+
+ private static TestingExternalPythonFunctionOperator createOperator(
+ StreamTask<?, ?> containingTask,
+ TestingPythonFunctionRunner functionRunner,
+ ExecutorService flushThreadPool)
+ throws ReflectiveOperationException {
+ final TestingExternalPythonFunctionOperator operator =
+ new TestingExternalPythonFunctionOperator(functionRunner);
+ setField(AbstractStreamOperator.class, operator, "container",
containingTask);
+ setField(
+ AbstractExternalPythonFunctionOperator.class,
+ operator,
+ "flushThreadPool",
+ flushThreadPool);
+ return operator;
+ }
+
+ private static void setField(Class<?> owner, Object target, String
fieldName, Object value)
+ throws ReflectiveOperationException {
+ final Field field = owner.getDeclaredField(fieldName);
+ field.setAccessible(true);
+ field.set(target, value);
+ }
+
+ private static class TestingExternalPythonFunctionOperator
+ extends AbstractExternalPythonFunctionOperator<Object> {
+
+ private final PythonFunctionRunner functionRunner;
+
+ private TestingExternalPythonFunctionOperator(PythonFunctionRunner
functionRunner) {
+ super(new Configuration());
+ this.functionRunner = functionRunner;
+ this.pythonFunctionRunner = functionRunner;
+ }
+
+ @Override
+ public PythonEnv getPythonEnv() {
+ return new PythonEnv(PythonEnv.ExecType.PROCESS);
+ }
+
+ @Override
+ public void emitResult(Tuple3<String, byte[], Integer> resultTuple) {}
+
+ @Override
+ public PythonFunctionRunner createPythonFunctionRunner() {
+ return functionRunner;
+ }
+ }
+
+ private static class TestingPythonFunctionRunner implements
PythonFunctionRunner {
+
+ private final ExecutorService flushThreadPool;
+ private final AtomicBoolean closeCalled = new AtomicBoolean();
+ private final AtomicBoolean cancelCalled = new AtomicBoolean();
+ private final AtomicBoolean flushExecutorStoppedBeforeCancel = new
AtomicBoolean();
+
+ private TestingPythonFunctionRunner(ExecutorService flushThreadPool) {
+ this.flushThreadPool = flushThreadPool;
+ }
+
+ @Override
+ public void open(org.apache.flink.configuration.ReadableConfig config)
{}
+
+ @Override
+ public void close() {
+ closeCalled.set(true);
+ }
+
+ @Override
+ public void cancel() {
+ flushExecutorStoppedBeforeCancel.set(flushThreadPool.isShutdown());
+ cancelCalled.set(true);
+ }
+
+ @Override
+ public void process(byte[] data) {}
+
+ @Override
+ public void processTimer(byte[] timerData) {}
+
+ @Override
+ public void drainUnregisteredTimers() {}
+
+ @Override
+ public Tuple3<String, byte[], Integer> pollResult() {
+ return null;
+ }
+
+ @Override
+ public Tuple3<String, byte[], Integer> takeResult() {
+ return null;
+ }
+
+ @Override
+ public void flush() {}
+ }
+}
diff --git
a/flink-python/src/test/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunnerTest.java
b/flink-python/src/test/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunnerTest.java
new file mode 100644
index 00000000000..44979995a82
--- /dev/null
+++
b/flink-python/src/test/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunnerTest.java
@@ -0,0 +1,954 @@
+/*
+ * 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.streaming.api.runners.python.beam;
+
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.MapState;
+import org.apache.flink.fnexecution.v1.FlinkFnApi;
+import org.apache.flink.python.env.PythonDependencyInfo;
+import org.apache.flink.python.env.process.ProcessPythonEnvironmentManager;
+import org.apache.flink.runtime.memory.OpaqueMemoryResource;
+import org.apache.flink.runtime.memory.SharedResources;
+import
org.apache.flink.streaming.api.runners.python.beam.state.BeamStateHandler;
+import
org.apache.flink.streaming.api.runners.python.beam.state.BeamStateRequestHandler;
+import org.apache.flink.streaming.api.runners.python.beam.state.BeamStateStore;
+import org.apache.flink.streaming.api.utils.ByteArrayWrapper;
+
+import org.apache.beam.model.fnexecution.v1.BeamFnApi;
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.runners.fnexecution.control.JobBundleFactory;
+import org.apache.beam.runners.fnexecution.control.RemoteBundle;
+import org.apache.beam.runners.fnexecution.control.StageBundleFactory;
+import org.apache.beam.sdk.fn.data.FnDataReceiver;
+import org.apache.beam.sdk.util.construction.Timer;
+import org.apache.beam.sdk.util.construction.graph.ExecutableStage;
+import org.apache.beam.sdk.util.construction.graph.TimerReference;
+import org.apache.beam.sdk.values.KV;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+
+import java.lang.reflect.Field;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.LockSupport;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+class BeamPythonFunctionRunnerTest {
+
+ @Test
+ void testCloseDrainsStateHandlerAfterStoppingOwnedRequestProduction()
throws Exception {
+ final AtomicBoolean stateAccessedDuringFactoryClose = new
AtomicBoolean();
+ final AtomicBoolean remoteBundleClosed = new AtomicBoolean();
+ final BeamStateRequestHandler stateRequestHandler =
+ createStateRequestHandler(stateAccessedDuringFactoryClose);
+ final JobBundleFactory jobBundleFactory = new
TestingJobBundleFactory(stateRequestHandler);
+ final TestingBeamPythonFunctionRunner runner =
+ new
TestingBeamPythonFunctionRunner(createEnvironmentManager());
+ setField(runner, "jobBundleFactory", jobBundleFactory);
+ setField(runner, "stateRequestHandler", stateRequestHandler);
+ setField(
+ runner,
+ "remoteBundle",
+ new TestingRemoteBundle(stateRequestHandler,
remoteBundleClosed));
+ setField(runner, "bundleStarted", true);
+
+ runner.close();
+
+ assertThat(stateAccessedDuringFactoryClose).isTrue();
+ assertThat(remoteBundleClosed).isTrue();
+ assertStateHandlerClosed(stateRequestHandler);
+ }
+
+ @Test
+ void testCloseDrainsStateHandlerForNonFinalManagedMemoryLease() throws
Exception {
+ final AtomicBoolean stateAccessedDuringBundleClose = new
AtomicBoolean();
+ final AtomicBoolean remoteBundleClosed = new AtomicBoolean();
+ final AtomicBoolean sharedFactoryClosed = new AtomicBoolean();
+ final BeamStateRequestHandler stateRequestHandler =
+ createStateRequestHandler(stateAccessedDuringBundleClose);
+ final TestingBeamPythonFunctionRunner runner =
+ new
TestingBeamPythonFunctionRunner(createEnvironmentManager());
+ final PythonSharedResources pythonSharedResources =
+ new PythonSharedResources(
+ new TrackingJobBundleFactory(sharedFactoryClosed),
+ RunnerApi.Environment.getDefaultInstance());
+ final AtomicInteger sharedResourceLeases = new AtomicInteger(2);
+ final OpaqueMemoryResource<PythonSharedResources> sharedResources =
+ createSharedResourceLease(pythonSharedResources,
sharedResourceLeases);
+ final OpaqueMemoryResource<PythonSharedResources>
remainingSharedResourceLease =
+ createSharedResourceLease(pythonSharedResources,
sharedResourceLeases);
+ setField(runner, "stateRequestHandler", stateRequestHandler);
+ setField(
+ runner,
+ "remoteBundle",
+ new TestingRemoteBundle(stateRequestHandler,
remoteBundleClosed));
+ setField(runner, "bundleStarted", true);
+ setField(runner, "sharedResources", sharedResources);
+
+ runner.close();
+
+ assertThat(stateAccessedDuringBundleClose).isTrue();
+ assertThat(remoteBundleClosed).isTrue();
+ assertThat(sharedFactoryClosed).isFalse();
+ assertStateHandlerClosed(stateRequestHandler);
+
+ remainingSharedResourceLease.close();
+
+ assertThat(sharedFactoryClosed).isTrue();
+ }
+
+ @Test
+ void testCloseDrainsStateHandlerForFinalManagedMemoryLease() throws
Exception {
+ final AtomicBoolean stateAccessedDuringFactoryClose = new
AtomicBoolean();
+ final BeamStateRequestHandler stateRequestHandler =
+ createStateRequestHandler(stateAccessedDuringFactoryClose);
+ final PythonSharedResources pythonSharedResources =
+ new PythonSharedResources(
+ new TestingJobBundleFactory(stateRequestHandler),
+ RunnerApi.Environment.getDefaultInstance());
+ final OpaqueMemoryResource<PythonSharedResources> sharedResources =
+ new OpaqueMemoryResource<>(pythonSharedResources, 1L,
pythonSharedResources::close);
+ final TestingBeamPythonFunctionRunner runner =
+ new
TestingBeamPythonFunctionRunner(createEnvironmentManager());
+ setField(runner, "stateRequestHandler", stateRequestHandler);
+ setField(runner, "sharedResources", sharedResources);
+
+ runner.close();
+
+ assertThat(stateAccessedDuringFactoryClose).isTrue();
+ assertStateHandlerClosed(stateRequestHandler);
+ }
+
+ @Test
+ void testCloseWaitsForConcurrentOwnedBundleFlush() throws Exception {
+ final AtomicBoolean stateAccessedDuringBundleClose = new
AtomicBoolean();
+ final AtomicBoolean ownedFactoryClosed = new AtomicBoolean();
+ final AtomicInteger remoteBundleCloseCalls = new AtomicInteger();
+ final AtomicReference<Thread> closeThread = new AtomicReference<>();
+ final CountDownLatch firstBundleCloseStarted = new CountDownLatch(1);
+ final CountDownLatch releaseFirstBundleClose = new CountDownLatch(1);
+ final CountDownLatch closeStarted = new CountDownLatch(1);
+ final BeamStateRequestHandler stateRequestHandler =
+ createStateRequestHandler(stateAccessedDuringBundleClose);
+ final TestingBeamPythonFunctionRunner runner =
+ new
TestingBeamPythonFunctionRunner(createEnvironmentManager());
+ setField(runner, "jobBundleFactory", new
TrackingJobBundleFactory(ownedFactoryClosed));
+ setField(runner, "stateRequestHandler", stateRequestHandler);
+ setField(
+ runner,
+ "remoteBundle",
+ new BlockingTestingRemoteBundle(
+ stateRequestHandler,
+ remoteBundleCloseCalls,
+ firstBundleCloseStarted,
+ releaseFirstBundleClose));
+ setField(runner, "bundleStarted", true);
+ final ExecutorService executor = Executors.newFixedThreadPool(2);
+
+ try {
+ final Future<?> flushFuture =
+ executor.submit(
+ () -> {
+ runner.flush();
+ return null;
+ });
+ assertThat(firstBundleCloseStarted.await(10,
TimeUnit.SECONDS)).isTrue();
+
+ final Future<?> closeFuture =
+ executor.submit(
+ () -> {
+ closeThread.set(Thread.currentThread());
+ closeStarted.countDown();
+ runner.close();
+ return null;
+ });
+ assertThat(closeStarted.await(10, TimeUnit.SECONDS)).isTrue();
+ assertThreadIsWaitingForBundleClose(
+ closeThread.get(), remoteBundleCloseCalls, closeFuture);
+ assertThat(closeFuture.isDone()).isFalse();
+ assertThat(remoteBundleCloseCalls).hasValue(1);
+ assertThat(ownedFactoryClosed).isFalse();
+
+ releaseFirstBundleClose.countDown();
+ flushFuture.get(10, TimeUnit.SECONDS);
+ closeFuture.get(10, TimeUnit.SECONDS);
+
+ assertThat(remoteBundleCloseCalls).hasValue(1);
+ assertThat(ownedFactoryClosed).isTrue();
+ assertThat(stateAccessedDuringBundleClose).isTrue();
+ assertStateHandlerClosed(stateRequestHandler);
+ } finally {
+ releaseFirstBundleClose.countDown();
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ void testCloseWaitsForConcurrentManagedBundleFlush() throws Exception {
+ final AtomicBoolean stateAccessedDuringBundleClose = new
AtomicBoolean();
+ final AtomicBoolean sharedFactoryClosed = new AtomicBoolean();
+ final AtomicInteger remoteBundleCloseCalls = new AtomicInteger();
+ final AtomicReference<Thread> closeThread = new AtomicReference<>();
+ final CountDownLatch firstBundleCloseStarted = new CountDownLatch(1);
+ final CountDownLatch releaseFirstBundleClose = new CountDownLatch(1);
+ final CountDownLatch closeStarted = new CountDownLatch(1);
+ final BeamStateRequestHandler stateRequestHandler =
+ createStateRequestHandler(stateAccessedDuringBundleClose);
+ final PythonSharedResources pythonSharedResources =
+ new PythonSharedResources(
+ new TrackingJobBundleFactory(sharedFactoryClosed),
+ RunnerApi.Environment.getDefaultInstance());
+ final OpaqueMemoryResource<PythonSharedResources> sharedResources =
+ new OpaqueMemoryResource<>(pythonSharedResources, 1L,
pythonSharedResources::close);
+ final TestingBeamPythonFunctionRunner runner =
+ new
TestingBeamPythonFunctionRunner(createEnvironmentManager());
+ setField(runner, "stateRequestHandler", stateRequestHandler);
+ setField(
+ runner,
+ "remoteBundle",
+ new BlockingTestingRemoteBundle(
+ stateRequestHandler,
+ remoteBundleCloseCalls,
+ firstBundleCloseStarted,
+ releaseFirstBundleClose));
+ setField(runner, "bundleStarted", true);
+ setField(runner, "sharedResources", sharedResources);
+ final ExecutorService executor = Executors.newFixedThreadPool(2);
+
+ try {
+ final Future<?> flushFuture =
+ executor.submit(
+ () -> {
+ runner.flush();
+ return null;
+ });
+ assertThat(firstBundleCloseStarted.await(10,
TimeUnit.SECONDS)).isTrue();
+
+ final Future<?> closeFuture =
+ executor.submit(
+ () -> {
+ closeThread.set(Thread.currentThread());
+ closeStarted.countDown();
+ runner.close();
+ return null;
+ });
+ assertThat(closeStarted.await(10, TimeUnit.SECONDS)).isTrue();
+ assertThreadIsWaitingForBundleClose(
+ closeThread.get(), remoteBundleCloseCalls, closeFuture);
+ assertThat(closeFuture.isDone()).isFalse();
+ assertThat(remoteBundleCloseCalls).hasValue(1);
+ assertThat(sharedFactoryClosed).isFalse();
+
+ releaseFirstBundleClose.countDown();
+ flushFuture.get(10, TimeUnit.SECONDS);
+ closeFuture.get(10, TimeUnit.SECONDS);
+
+ assertThat(remoteBundleCloseCalls).hasValue(1);
+ assertThat(sharedFactoryClosed).isTrue();
+ assertThat(stateAccessedDuringBundleClose).isTrue();
+ assertStateHandlerClosed(stateRequestHandler);
+ } finally {
+ releaseFirstBundleClose.countDown();
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ void testCancelDoesNotWaitForConcurrentManagedBundleFlush() throws
Exception {
+ final AtomicBoolean stateAccessedDuringBundleClose = new
AtomicBoolean();
+ final AtomicBoolean sharedFactoryClosed = new AtomicBoolean();
+ final AtomicInteger remoteBundleCloseCalls = new AtomicInteger();
+ final CountDownLatch firstBundleCloseStarted = new CountDownLatch(1);
+ final CountDownLatch releaseFirstBundleClose = new CountDownLatch(1);
+ final BeamStateRequestHandler stateRequestHandler =
+ createStateRequestHandler(stateAccessedDuringBundleClose);
+ final PythonSharedResources pythonSharedResources =
+ new PythonSharedResources(
+ new TrackingJobBundleFactory(sharedFactoryClosed),
+ RunnerApi.Environment.getDefaultInstance());
+ final OpaqueMemoryResource<PythonSharedResources> sharedResources =
+ new OpaqueMemoryResource<>(pythonSharedResources, 1L,
pythonSharedResources::close);
+ final TestingBeamPythonFunctionRunner runner =
+ new
TestingBeamPythonFunctionRunner(createEnvironmentManager());
+ setField(runner, "stateRequestHandler", stateRequestHandler);
+ setField(
+ runner,
+ "remoteBundle",
+ new BlockingTestingRemoteBundle(
+ stateRequestHandler,
+ remoteBundleCloseCalls,
+ firstBundleCloseStarted,
+ releaseFirstBundleClose));
+ setField(runner, "bundleStarted", true);
+ setField(runner, "sharedResources", sharedResources);
+ final ExecutorService executor = Executors.newFixedThreadPool(2);
+
+ try {
+ final Future<?> flushFuture =
+ executor.submit(
+ () -> {
+ runner.flush();
+ return null;
+ });
+ assertThat(firstBundleCloseStarted.await(10,
TimeUnit.SECONDS)).isTrue();
+
+ final Future<?> cancelFuture =
+ executor.submit(
+ () -> {
+ runner.cancel();
+ return null;
+ });
+ cancelFuture.get(10, TimeUnit.SECONDS);
+
+ assertThat(remoteBundleCloseCalls).hasValue(1);
+ assertThat(sharedFactoryClosed).isTrue();
+ assertThat(stateAccessedDuringBundleClose).isFalse();
+ assertStateHandlerClosed(stateRequestHandler);
+
+ releaseFirstBundleClose.countDown();
+ assertThatThrownBy(() -> flushFuture.get(10, TimeUnit.SECONDS))
+ .hasRootCauseInstanceOf(IllegalStateException.class);
+ assertThat(sharedFactoryClosed).isTrue();
+ assertThat(remoteBundleCloseCalls).hasValue(1);
+ assertThat(stateAccessedDuringBundleClose).isFalse();
+ } finally {
+ releaseFirstBundleClose.countDown();
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ void testCancelRetainsUnclaimedManagedBundleCleanupOwnership() throws
Exception {
+ final AtomicBoolean stateAccessedDuringBundleClose = new
AtomicBoolean();
+ final AtomicBoolean sharedFactoryClosed = new AtomicBoolean();
+ final AtomicInteger remoteBundleCloseCalls = new AtomicInteger();
+ final CountDownLatch bundleCloseStarted = new CountDownLatch(1);
+ final CountDownLatch releaseBundleClose = new CountDownLatch(1);
+ final CountDownLatch bundleCloseFinished = new CountDownLatch(1);
+ final BeamStateRequestHandler stateRequestHandler =
+ createStateRequestHandler(stateAccessedDuringBundleClose);
+ final PythonSharedResources pythonSharedResources =
+ new PythonSharedResources(
+ new TrackingJobBundleFactory(sharedFactoryClosed),
+ RunnerApi.Environment.getDefaultInstance());
+ final AtomicInteger sharedResourceLeases = new AtomicInteger(2);
+ final OpaqueMemoryResource<PythonSharedResources> sharedResources =
+ createSharedResourceLease(pythonSharedResources,
sharedResourceLeases);
+ final OpaqueMemoryResource<PythonSharedResources>
remainingSharedResourceLease =
+ createSharedResourceLease(pythonSharedResources,
sharedResourceLeases);
+ final TestingBeamPythonFunctionRunner runner =
+ new
TestingBeamPythonFunctionRunner(createEnvironmentManager());
+ setField(runner, "stateRequestHandler", stateRequestHandler);
+ setField(
+ runner,
+ "remoteBundle",
+ new BlockingTestingRemoteBundle(
+ stateRequestHandler,
+ remoteBundleCloseCalls,
+ bundleCloseStarted,
+ releaseBundleClose,
+ bundleCloseFinished));
+ setField(runner, "bundleStarted", true);
+ setField(runner, "sharedResources", sharedResources);
+
+ try {
+ runner.cancel();
+ assertThat(bundleCloseStarted.await(10,
TimeUnit.SECONDS)).isTrue();
+ runner.flush();
+
+ assertThat(remoteBundleCloseCalls).hasValue(1);
+ assertThat(sharedFactoryClosed).isFalse();
+ assertStateHandlerClosed(stateRequestHandler);
+
+ releaseBundleClose.countDown();
+ assertThat(bundleCloseFinished.await(10,
TimeUnit.SECONDS)).isTrue();
+ assertRemainingSharedResourceLeases(sharedResourceLeases, 1);
+ assertThat(stateAccessedDuringBundleClose).isFalse();
+ assertThat(sharedFactoryClosed).isFalse();
+
+ remainingSharedResourceLease.close();
+
+ assertThat(sharedFactoryClosed).isTrue();
+ } finally {
+ releaseBundleClose.countDown();
+ remainingSharedResourceLease.close();
+ }
+ }
+
+ @ParameterizedTest
+ @CsvSource({"false, false", "true, false", "false, true", "true, true"})
+ void testCancelWaitsForConcurrentFinalLeaseDisposal(
+ boolean failDisposal, boolean interruptCancel) throws Exception {
+ final AtomicInteger factoryCloseCalls = new AtomicInteger();
+ final AtomicBoolean cancelInterrupted = new AtomicBoolean();
+ final AtomicReference<Thread> cancelThread = new AtomicReference<>();
+ final CountDownLatch resourceCleanupStarted = new CountDownLatch(1);
+ final CountDownLatch releaseResourceCleanup = new CountDownLatch(1);
+ final RuntimeException disposalFailure = new
RuntimeException("Resource disposal failed.");
+ final JobBundleFactory jobBundleFactory =
+ new TrackingJobBundleFactory(new AtomicBoolean()) {
+ @Override
+ public void close() {
+ factoryCloseCalls.incrementAndGet();
+ resourceCleanupStarted.countDown();
+ try {
+ assertThat(releaseResourceCleanup.await(10,
TimeUnit.SECONDS)).isTrue();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException(e);
+ }
+ if (failDisposal) {
+ throw disposalFailure;
+ }
+ }
+ };
+ final PythonSharedResources pythonSharedResources =
+ new PythonSharedResources(
+ jobBundleFactory,
RunnerApi.Environment.getDefaultInstance());
+ final SharedResources registry = new SharedResources();
+ final Object leaseHolder = new Object();
+ final SharedResources.ResourceAndSize<PythonSharedResources> resource =
+ registry.getOrAllocateSharedResource(
+ "python", leaseHolder, ignored ->
pythonSharedResources, 1L);
+ final OpaqueMemoryResource<PythonSharedResources> sharedResources =
+ new OpaqueMemoryResource<>(
+ resource.resourceHandle(),
+ resource.size(),
+ () -> registry.release("python", leaseHolder, ignored
-> {}));
+ final BeamStateRequestHandler stateRequestHandler =
+ createStateRequestHandler(new AtomicBoolean());
+ final TestingBeamPythonFunctionRunner runner =
+ new
TestingBeamPythonFunctionRunner(createEnvironmentManager());
+ final ExecutorService cancellationExecutor =
+ new CleanupAwaitingExecutor(resourceCleanupStarted,
interruptCancel);
+ setField(runner, "stateRequestHandler", stateRequestHandler);
+ setField(
+ runner,
+ "remoteBundle",
+ new TestingRemoteBundle(stateRequestHandler, new
AtomicBoolean()));
+ setField(runner, "bundleStarted", true);
+ setField(runner, "sharedResources", sharedResources);
+ setField(runner, "cancellationBundleCloseExecutor",
cancellationExecutor);
+ final ExecutorService executor = Executors.newSingleThreadExecutor();
+
+ try {
+ final Future<?> cancelFuture =
+ executor.submit(
+ () -> {
+ cancelThread.set(Thread.currentThread());
+ try {
+ runner.cancel();
+ } finally {
+
cancelInterrupted.set(Thread.currentThread().isInterrupted());
+ }
+ return null;
+ });
+ assertThat(resourceCleanupStarted.await(10,
TimeUnit.SECONDS)).isTrue();
+ assertThreadIsWaitingForBundleClose(
+ cancelThread.get(), factoryCloseCalls, cancelFuture);
+ assertThat(cancelFuture.isDone()).isFalse();
+ assertThat(factoryCloseCalls).hasValue(1);
+
+ releaseResourceCleanup.countDown();
+ if (failDisposal) {
+ assertThatThrownBy(() -> cancelFuture.get(10,
TimeUnit.SECONDS))
+ .hasCause(disposalFailure);
+ } else {
+ cancelFuture.get(10, TimeUnit.SECONDS);
+ }
+ assertStateHandlerClosed(stateRequestHandler);
+ assertThat(factoryCloseCalls).hasValue(1);
+ assertThat(cancelInterrupted.get()).isEqualTo(interruptCancel);
+
+ final AtomicBoolean replacementFactoryClosed = new AtomicBoolean();
+ executor.submit(
+ () -> {
+ final Object replacementLeaseHolder = new
Object();
+ registry.getOrAllocateSharedResource(
+ "python",
+ replacementLeaseHolder,
+ ignored ->
+ new PythonSharedResources(
+ new
TrackingJobBundleFactory(
+
replacementFactoryClosed),
+
RunnerApi.Environment.getDefaultInstance()),
+ 1L);
+ registry.release("python",
replacementLeaseHolder, ignored -> {});
+ return null;
+ })
+ .get(10, TimeUnit.SECONDS);
+ assertThat(replacementFactoryClosed).isTrue();
+ assertThat(factoryCloseCalls).hasValue(1);
+ } finally {
+ releaseResourceCleanup.countDown();
+ cancellationExecutor.shutdownNow();
+ executor.shutdownNow();
+ assertThat(cancellationExecutor.awaitTermination(10,
TimeUnit.SECONDS)).isTrue();
+ assertThat(executor.awaitTermination(10,
TimeUnit.SECONDS)).isTrue();
+ }
+ }
+
+ @Test
+ void testConcurrentCloseIsIdempotent() throws Exception {
+ final AtomicInteger stateHandlerCloseCalls = new AtomicInteger();
+ final AtomicReference<Thread> secondCloseThread = new
AtomicReference<>();
+ final CountDownLatch firstStateHandlerCloseStarted = new
CountDownLatch(1);
+ final CountDownLatch releaseFirstStateHandlerClose = new
CountDownLatch(1);
+ final CountDownLatch secondCloseStarted = new CountDownLatch(1);
+ final ProcessPythonEnvironmentManager environmentManager =
createEnvironmentManager();
+ environmentManager.open();
+ final Path environmentDirectory =
Paths.get(environmentManager.getBaseDirectory());
+ final BlockingCloseBeamStateRequestHandler stateRequestHandler =
+ new BlockingCloseBeamStateRequestHandler(
+ stateHandlerCloseCalls,
+ firstStateHandlerCloseStarted,
+ releaseFirstStateHandlerClose);
+ final PythonSharedResources pythonSharedResources =
+ new PythonSharedResources(
+ new TrackingJobBundleFactory(new AtomicBoolean()),
+ RunnerApi.Environment.getDefaultInstance());
+ pythonSharedResources.addPythonEnvironmentManager(environmentManager);
+ final AtomicInteger sharedResourceLeases = new AtomicInteger(2);
+ final OpaqueMemoryResource<PythonSharedResources> sharedResources =
+ createSharedResourceLease(pythonSharedResources,
sharedResourceLeases);
+ final OpaqueMemoryResource<PythonSharedResources>
remainingSharedResourceLease =
+ createSharedResourceLease(pythonSharedResources,
sharedResourceLeases);
+ final TestingBeamPythonFunctionRunner runner =
+ new TestingBeamPythonFunctionRunner(environmentManager);
+ setField(runner, "stateRequestHandler", stateRequestHandler);
+ setField(runner, "sharedResources", sharedResources);
+ final ExecutorService executor = Executors.newFixedThreadPool(2);
+
+ try {
+ final Future<?> firstCloseFuture =
+ executor.submit(
+ () -> {
+ runner.close();
+ return null;
+ });
+ assertThat(firstStateHandlerCloseStarted.await(10,
TimeUnit.SECONDS)).isTrue();
+
+ final Future<?> secondCloseFuture =
+ executor.submit(
+ () -> {
+ secondCloseThread.set(Thread.currentThread());
+ secondCloseStarted.countDown();
+ runner.close();
+ return null;
+ });
+ assertThat(secondCloseStarted.await(10,
TimeUnit.SECONDS)).isTrue();
+ assertThreadIsWaitingForBundleClose(
+ secondCloseThread.get(), stateHandlerCloseCalls,
secondCloseFuture);
+ assertThat(secondCloseFuture.isDone()).isFalse();
+ assertThat(stateHandlerCloseCalls).hasValue(1);
+ assertThat(Files.exists(environmentDirectory)).isTrue();
+
+ releaseFirstStateHandlerClose.countDown();
+ firstCloseFuture.get(10, TimeUnit.SECONDS);
+ secondCloseFuture.get(10, TimeUnit.SECONDS);
+
+ assertThat(stateHandlerCloseCalls).hasValue(1);
+ assertThat(Files.exists(environmentDirectory)).isTrue();
+ } finally {
+ releaseFirstStateHandlerClose.countDown();
+ executor.shutdownNow();
+ remainingSharedResourceLease.close();
+ environmentManager.close();
+ }
+ }
+
+ private static void assertThreadIsWaitingForBundleClose(
+ Thread closeThread, AtomicInteger concurrentCloseCalls, Future<?>
closeFuture) {
+ final long timeoutNanos = TimeUnit.SECONDS.toNanos(10);
+ final long deadlineNanos = System.nanoTime() + timeoutNanos;
+ while (System.nanoTime() < deadlineNanos) {
+ if (closeFuture.isDone()) {
+ assertThat(closeFuture.isDone()).isFalse();
+ }
+ if (concurrentCloseCalls.get() > 1) {
+ assertThat(concurrentCloseCalls).hasValue(1);
+ }
+ if (closeThread.getState() == Thread.State.WAITING) {
+ return;
+ }
+ LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1));
+ }
+ assertThat(closeThread.getState()).isEqualTo(Thread.State.WAITING);
+ }
+
+ private static void assertRemainingSharedResourceLeases(
+ AtomicInteger remainingLeases, int expectedLeases) {
+ final long deadlineNanos = System.nanoTime() +
TimeUnit.SECONDS.toNanos(10);
+ while (System.nanoTime() < deadlineNanos) {
+ if (remainingLeases.get() == expectedLeases) {
+ return;
+ }
+ LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1));
+ }
+ assertThat(remainingLeases).hasValue(expectedLeases);
+ }
+
+ private static OpaqueMemoryResource<PythonSharedResources>
createSharedResourceLease(
+ PythonSharedResources pythonSharedResources, AtomicInteger
remainingLeases) {
+ return new OpaqueMemoryResource<>(
+ pythonSharedResources,
+ 1L,
+ () -> {
+ if (remainingLeases.decrementAndGet() == 0) {
+ pythonSharedResources.close();
+ }
+ });
+ }
+
+ private static BeamStateRequestHandler
createStateRequestHandler(AtomicBoolean stateAccessed) {
+ final BeamStateStore keyedStateStore =
+ new BeamStateStore() {
+ @Override
+ public ListState<byte[]>
getListState(BeamFnApi.StateRequest request) {
+ stateAccessed.set(true);
+ return null;
+ }
+
+ @Override
+ public MapState<ByteArrayWrapper, byte[]> getMapState(
+ BeamFnApi.StateRequest request) {
+ throw new UnsupportedOperationException();
+ }
+ };
+ return new BeamStateRequestHandler(
+ keyedStateStore,
+ BeamStateStore.unsupported(),
+ new NoOpBeamStateHandler<>(),
+ new NoOpBeamStateHandler<>());
+ }
+
+ private static void assertStateHandlerClosed(BeamStateRequestHandler
stateRequestHandler) {
+ assertThatThrownBy(() ->
stateRequestHandler.handle(createBagUserStateRequest()))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessage("Beam state request handler is closed.");
+ }
+
+ private static ProcessPythonEnvironmentManager createEnvironmentManager() {
+ return new ProcessPythonEnvironmentManager(
+ new PythonDependencyInfo(
+ Collections.emptyMap(), null, null,
Collections.emptyMap(), "python"),
+ new String[] {System.getProperty("java.io.tmpdir")},
+ Collections.emptyMap(),
+ new JobID());
+ }
+
+ private static void setField(Object target, String fieldName, Object value)
+ throws ReflectiveOperationException {
+ final Field field =
BeamPythonFunctionRunner.class.getDeclaredField(fieldName);
+ field.setAccessible(true);
+ field.set(target, value);
+ }
+
+ private static BeamFnApi.StateRequest createBagUserStateRequest() {
+ return BeamFnApi.StateRequest.newBuilder()
+ .setStateKey(
+ BeamFnApi.StateKey.newBuilder()
+ .setBagUserState(
+
BeamFnApi.StateKey.BagUserState.getDefaultInstance()))
+ .setGet(BeamFnApi.StateGetRequest.getDefaultInstance())
+ .build();
+ }
+
+ private static class TestingJobBundleFactory implements JobBundleFactory {
+
+ private final BeamStateRequestHandler stateRequestHandler;
+
+ private TestingJobBundleFactory(BeamStateRequestHandler
stateRequestHandler) {
+ this.stateRequestHandler = stateRequestHandler;
+ }
+
+ @Override
+ public StageBundleFactory forStage(ExecutableStage executableStage) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public void close() throws Exception {
+ stateRequestHandler.handle(createBagUserStateRequest());
+ }
+ }
+
+ private static class TestingRemoteBundle implements RemoteBundle {
+
+ private final BeamStateRequestHandler stateRequestHandler;
+ private final AtomicBoolean closed;
+
+ private TestingRemoteBundle(
+ BeamStateRequestHandler stateRequestHandler, AtomicBoolean
closed) {
+ this.stateRequestHandler = stateRequestHandler;
+ this.closed = closed;
+ }
+
+ @Override
+ public String getId() {
+ return "test-bundle";
+ }
+
+ @Override
+ public Map<String, FnDataReceiver> getInputReceivers() {
+ return Collections.emptyMap();
+ }
+
+ @Override
+ public Map<KV<String, String>, FnDataReceiver<Timer>>
getTimerReceivers() {
+ return Collections.emptyMap();
+ }
+
+ @Override
+ public void requestProgress() {}
+
+ @Override
+ public void split(double fractionOfRemainder) {}
+
+ @Override
+ public void close() throws Exception {
+ closed.set(true);
+ stateRequestHandler.handle(createBagUserStateRequest());
+ }
+ }
+
+ private static class BlockingTestingRemoteBundle extends
TestingRemoteBundle {
+
+ private final BeamStateRequestHandler stateRequestHandler;
+ private final AtomicInteger closeCalls;
+ private final CountDownLatch firstCloseStarted;
+ private final CountDownLatch releaseFirstClose;
+ private final CountDownLatch closeFinished;
+
+ private BlockingTestingRemoteBundle(
+ BeamStateRequestHandler stateRequestHandler,
+ AtomicInteger closeCalls,
+ CountDownLatch firstCloseStarted,
+ CountDownLatch releaseFirstClose) {
+ this(
+ stateRequestHandler,
+ closeCalls,
+ firstCloseStarted,
+ releaseFirstClose,
+ new CountDownLatch(0));
+ }
+
+ private BlockingTestingRemoteBundle(
+ BeamStateRequestHandler stateRequestHandler,
+ AtomicInteger closeCalls,
+ CountDownLatch firstCloseStarted,
+ CountDownLatch releaseFirstClose,
+ CountDownLatch closeFinished) {
+ super(stateRequestHandler, new AtomicBoolean());
+ this.stateRequestHandler = stateRequestHandler;
+ this.closeCalls = closeCalls;
+ this.firstCloseStarted = firstCloseStarted;
+ this.releaseFirstClose = releaseFirstClose;
+ this.closeFinished = closeFinished;
+ }
+
+ @Override
+ public void close() throws Exception {
+ try {
+ if (closeCalls.incrementAndGet() == 1) {
+ firstCloseStarted.countDown();
+ awaitUninterruptibly(releaseFirstClose);
+ }
+ stateRequestHandler.handle(createBagUserStateRequest());
+ } finally {
+ closeFinished.countDown();
+ }
+ }
+
+ private static void awaitUninterruptibly(CountDownLatch latch) {
+ boolean interrupted = false;
+ while (true) {
+ try {
+ latch.await();
+ break;
+ } catch (InterruptedException e) {
+ interrupted = true;
+ }
+ }
+ if (interrupted) {
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ private static class BlockingCloseBeamStateRequestHandler extends
BeamStateRequestHandler {
+
+ private final AtomicInteger closeCalls;
+ private final CountDownLatch firstCloseStarted;
+ private final CountDownLatch releaseFirstClose;
+
+ private BlockingCloseBeamStateRequestHandler(
+ AtomicInteger closeCalls,
+ CountDownLatch firstCloseStarted,
+ CountDownLatch releaseFirstClose) {
+ super(
+ BeamStateStore.unsupported(),
+ BeamStateStore.unsupported(),
+ new NoOpBeamStateHandler<>(),
+ new NoOpBeamStateHandler<>());
+ this.closeCalls = closeCalls;
+ this.firstCloseStarted = firstCloseStarted;
+ this.releaseFirstClose = releaseFirstClose;
+ }
+
+ @Override
+ public void close() {
+ if (closeCalls.incrementAndGet() == 1) {
+ firstCloseStarted.countDown();
+ try {
+ releaseFirstClose.await();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while closing
state handler.", e);
+ }
+ }
+ super.close();
+ }
+ }
+
+ private static class TrackingJobBundleFactory implements JobBundleFactory {
+
+ private final AtomicBoolean closed;
+
+ private TrackingJobBundleFactory(AtomicBoolean closed) {
+ this.closed = closed;
+ }
+
+ @Override
+ public StageBundleFactory forStage(ExecutableStage executableStage) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public void close() {
+ closed.set(true);
+ }
+ }
+
+ private static class CleanupAwaitingExecutor extends ThreadPoolExecutor {
+
+ private final CountDownLatch resourceCleanupStarted;
+ private final boolean interruptCancel;
+
+ private CleanupAwaitingExecutor(
+ CountDownLatch resourceCleanupStarted, boolean
interruptCancel) {
+ super(1, 1, 0L, TimeUnit.MILLISECONDS, new
LinkedBlockingQueue<>());
+ this.resourceCleanupStarted = resourceCleanupStarted;
+ this.interruptCancel = interruptCancel;
+ }
+
+ @Override
+ public void execute(Runnable command) {
+ super.execute(command);
+ try {
+ assertThat(resourceCleanupStarted.await(10,
TimeUnit.SECONDS)).isTrue();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException(e);
+ }
+ if (interruptCancel) {
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ private static class NoOpBeamStateHandler<S> implements
BeamStateHandler<S> {
+
+ @Override
+ public BeamFnApi.StateResponse.Builder handle(BeamFnApi.StateRequest
request, S state) {
+ return BeamFnApi.StateResponse.newBuilder();
+ }
+
+ @Override
+ public BeamFnApi.StateResponse.Builder
handleGet(BeamFnApi.StateRequest request, S state) {
+ return BeamFnApi.StateResponse.newBuilder();
+ }
+
+ @Override
+ public BeamFnApi.StateResponse.Builder handleAppend(
+ BeamFnApi.StateRequest request, S state) {
+ return BeamFnApi.StateResponse.newBuilder();
+ }
+
+ @Override
+ public BeamFnApi.StateResponse.Builder handleClear(
+ BeamFnApi.StateRequest request, S state) {
+ return BeamFnApi.StateResponse.newBuilder();
+ }
+ }
+
+ private static class TestingBeamPythonFunctionRunner extends
BeamPythonFunctionRunner {
+
+ private TestingBeamPythonFunctionRunner(
+ ProcessPythonEnvironmentManager environmentManager) {
+ super(
+ null,
+ "test-task",
+ environmentManager,
+ null,
+ null,
+ null,
+ null,
+ null,
+ null,
+ null,
+ 0.0,
+ FlinkFnApi.CoderInfoDescriptor.getDefaultInstance(),
+ FlinkFnApi.CoderInfoDescriptor.getDefaultInstance(),
+ Collections.emptyMap());
+ }
+
+ @Override
+ protected void buildTransforms(RunnerApi.Components.Builder
componentsBuilder) {}
+
+ @Override
+ protected List<TimerReference> getTimers(RunnerApi.Components
components) {
+ return Collections.emptyList();
+ }
+
+ @Override
+ protected Optional<RunnerApi.Coder> getOptionalTimerCoderProto() {
+ return Optional.empty();
+ }
+ }
+}
diff --git
a/flink-python/src/test/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamStateRequestHandlerTest.java
b/flink-python/src/test/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamStateRequestHandlerTest.java
new file mode 100644
index 00000000000..bf99c827dd5
--- /dev/null
+++
b/flink-python/src/test/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamStateRequestHandlerTest.java
@@ -0,0 +1,176 @@
+/*
+ * 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.streaming.api.runners.python.beam.state;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.MapState;
+import org.apache.flink.streaming.api.utils.ByteArrayWrapper;
+
+import org.apache.beam.model.fnexecution.v1.BeamFnApi;
+import org.junit.jupiter.api.Test;
+
+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.concurrent.locks.LockSupport;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+class BeamStateRequestHandlerTest {
+
+ @Test
+ void testRejectsRequestsAfterClose() {
+ final AtomicBoolean stateAccessed = new AtomicBoolean();
+ final BeamStateStore keyedStateStore =
+ new BeamStateStore() {
+ @Override
+ public ListState<byte[]>
getListState(BeamFnApi.StateRequest request) {
+ stateAccessed.set(true);
+ return null;
+ }
+
+ @Override
+ public MapState<ByteArrayWrapper, byte[]> getMapState(
+ BeamFnApi.StateRequest request) {
+ throw new UnsupportedOperationException();
+ }
+ };
+ final BeamStateRequestHandler handler = createHandler(keyedStateStore);
+
+ handler.close();
+
+ assertThatThrownBy(() -> handler.handle(createBagUserStateRequest()))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessage("Beam state request handler is closed.");
+ assertThat(stateAccessed).isFalse();
+ }
+
+ @Test
+ void testCloseWaitsForInFlightRequest() throws Exception {
+ final CountDownLatch stateAccessStarted = new CountDownLatch(1);
+ final CountDownLatch releaseStateAccess = new CountDownLatch(1);
+ final BeamStateStore keyedStateStore =
+ new BeamStateStore() {
+ @Override
+ public ListState<byte[]>
getListState(BeamFnApi.StateRequest request)
+ throws InterruptedException {
+ stateAccessStarted.countDown();
+ releaseStateAccess.await();
+ return null;
+ }
+
+ @Override
+ public MapState<ByteArrayWrapper, byte[]> getMapState(
+ BeamFnApi.StateRequest request) {
+ throw new UnsupportedOperationException();
+ }
+ };
+ final BeamStateRequestHandler handler = createHandler(keyedStateStore);
+ final ExecutorService executor = Executors.newFixedThreadPool(2);
+
+ try {
+ final Future<?> requestFuture =
+ executor.submit(() ->
handler.handle(createBagUserStateRequest()));
+ assertThat(stateAccessStarted.await(10,
TimeUnit.SECONDS)).isTrue();
+
+ final CountDownLatch closeStarted = new CountDownLatch(1);
+ final AtomicReference<Thread> closeThread = new
AtomicReference<>();
+ final Future<?> closeFuture =
+ executor.submit(
+ () -> {
+ closeThread.set(Thread.currentThread());
+ closeStarted.countDown();
+ handler.close();
+ });
+ assertThat(closeStarted.await(10, TimeUnit.SECONDS)).isTrue();
+ assertCloseWaitsForInFlightRequest(closeThread.get(), closeFuture);
+
+ releaseStateAccess.countDown();
+ requestFuture.get(10, TimeUnit.SECONDS);
+ closeFuture.get(10, TimeUnit.SECONDS);
+ } finally {
+ releaseStateAccess.countDown();
+ executor.shutdownNow();
+ }
+ }
+
+ private static void assertCloseWaitsForInFlightRequest(
+ Thread closeThread, Future<?> closeFuture) {
+ final long timeoutNanos = TimeUnit.SECONDS.toNanos(10);
+ final long deadlineNanos = System.nanoTime() + timeoutNanos;
+ while (System.nanoTime() < deadlineNanos) {
+ if (closeFuture.isDone()) {
+ assertThat(closeFuture.isDone()).isFalse();
+ }
+ if (closeThread.getState() == Thread.State.WAITING) {
+ return;
+ }
+ LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1));
+ }
+ assertThat(closeThread.getState()).isEqualTo(Thread.State.WAITING);
+ }
+
+ private static BeamStateRequestHandler createHandler(BeamStateStore
keyedStateStore) {
+ return new BeamStateRequestHandler(
+ keyedStateStore,
+ BeamStateStore.unsupported(),
+ new NoOpBeamStateHandler<>(),
+ new NoOpBeamStateHandler<>());
+ }
+
+ private static BeamFnApi.StateRequest createBagUserStateRequest() {
+ return BeamFnApi.StateRequest.newBuilder()
+ .setStateKey(
+ BeamFnApi.StateKey.newBuilder()
+ .setBagUserState(
+
BeamFnApi.StateKey.BagUserState.getDefaultInstance()))
+ .setGet(BeamFnApi.StateGetRequest.getDefaultInstance())
+ .build();
+ }
+
+ private static class NoOpBeamStateHandler<S> implements
BeamStateHandler<S> {
+
+ @Override
+ public BeamFnApi.StateResponse.Builder handle(BeamFnApi.StateRequest
request, S state) {
+ return BeamFnApi.StateResponse.newBuilder();
+ }
+
+ @Override
+ public BeamFnApi.StateResponse.Builder
handleGet(BeamFnApi.StateRequest request, S state) {
+ return BeamFnApi.StateResponse.newBuilder();
+ }
+
+ @Override
+ public BeamFnApi.StateResponse.Builder handleAppend(
+ BeamFnApi.StateRequest request, S state) {
+ return BeamFnApi.StateResponse.newBuilder();
+ }
+
+ @Override
+ public BeamFnApi.StateResponse.Builder handleClear(
+ BeamFnApi.StateRequest request, S state) {
+ return BeamFnApi.StateResponse.newBuilder();
+ }
+ }
+}
