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 c637c0e659c13..708e8c70e8a22 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. + * + *

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 da50b5335f22c..f0367fb6b95cf 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.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 void open() throws Exception { 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 void close() throws Exception { @Override protected void invokeFinishBundle() throws Exception { if (elementCount > 0) { + final PythonFunctionRunner functionRunner = pythonFunctionRunner; AtomicBoolean flushThreadFinish = new AtomicBoolean(false); AtomicReference 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 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 7d89c93cc36af..82cf75e60d308 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 @@ -42,6 +42,7 @@ 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 +56,6 @@ 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 +93,10 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.concurrent.CompletableFuture; +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 +169,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 +199,32 @@ public abstract class BeamPythonFunctionRunner implements PythonFunctionRunner { /** The shared resource among Python operators of the same slot. */ private transient OpaqueMemoryResource 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 closeCompletion; + + /** Whether cancellation should skip waiting for an in-flight bundle close. */ + private transient boolean cancelRequested; + + /** Prevents duplicate resource teardown from close and cancellation. */ + private transient boolean resourcesClosed; + + /** Coordinates duplicate close and cancel calls while the state handler drains requests. */ + @Nullable private transient CompletableFuture 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. + * + *

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 +358,74 @@ public void open(ReadableConfig config) throws Exception { 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> inputReceiver; + synchronized (this) { + checkInvokeStartBundle(); + inputReceiver = mainInputReceiver; + } + inputReceiver.accept(WindowedValues.valueInGlobalWindow(data)); } @Override @@ -381,23 +440,34 @@ public void drainUnregisteredTimers() { @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 inputReceiver; + synchronized (this) { + if (timerInputReceiver == null) { + checkInvokeStartBundle(); + timerInputReceiver = + Preconditions.checkNotNull( + Iterables.getOnlyElement(remoteBundle.getTimerReceivers().values()), + "Failed to retrieve main input receiver."); + } + inputReceiver = timerInputReceiver; } Timer 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 +512,34 @@ public Tuple3 takeResult() throws Exception { @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 +547,249 @@ public void notifyNoMoreResults() { 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. + * + *

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 currentCloseCompletion; + synchronized (this) { + currentCloseCompletion = closeCompletion; + } + if (currentCloseCompletion != null) { + awaitCompletion(currentCloseCompletion); + } + } + + private void completeClose() { + final CompletableFuture currentCloseCompletion; + synchronized (this) { + currentCloseCompletion = closeCompletion; + } + if (currentCloseCompletion != null) { + currentCloseCompletion.complete(null); + } + } + + private void closeResources() throws Exception { + final JobBundleFactory ownedJobBundleFactory; + final OpaqueMemoryResource currentSharedResources; + synchronized (this) { + if (resourcesClosed) { + return; + } + resourcesClosed = true; + ownedJobBundleFactory = jobBundleFactory; + currentSharedResources = sharedResources; + jobBundleFactory = null; + sharedResources = null; + } + + if (currentSharedResources != null) { + try { + currentSharedResources.close(); + } finally { + clearBundleReferences(); + } + } else { + try { + if (ownedJobBundleFactory != null) { + ownedJobBundleFactory.close(); + } + } finally { + try { + environmentManager.close(); + } finally { + clearBundleReferences(); + } + } + } + } + + 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 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 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 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 +1060,7 @@ private TimerReceiverFactory createTimerReceiverFactory() { 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 293d5a5ad83d8..9eaf2fbb4688b 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 environmentManagers; + private boolean closed; + PythonSharedResources(JobBundleFactory jobBundleFactory, Environment environment) { this.jobBundleFactory = jobBundleFactory; this.environment = environment; @@ -63,10 +65,31 @@ synchronized void addPythonEnvironmentManager(PythonEnvironmentManager environme } @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 fb6765dff202d..34295b5a519be 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.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 BeamStateRequestHandler( @Override public CompletionStage handle(BeamFnApi.StateRequest request) throws Exception { - BeamFnApi.StateKey.TypeCase typeCase = request.getStateKey().getTypeCase(); - ListState listState; - MapState 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 listState; + MapState 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 0000000000000..fa0d307a562ab --- /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 { + + 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 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 pollResult() { + return null; + } + + @Override + public Tuple3 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 0000000000000..b5074fae72d3f --- /dev/null +++ b/flink-python/src/test/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunnerTest.java @@ -0,0 +1,809 @@ +/* + * 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.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 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.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 sharedResources = + createSharedResourceLease(pythonSharedResources, sharedResourceLeases); + final OpaqueMemoryResource 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 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 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 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 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 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 sharedResources = + createSharedResourceLease(pythonSharedResources, sharedResourceLeases); + final OpaqueMemoryResource 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(); + } + } + + @Test + void testConcurrentCloseIsIdempotent() throws Exception { + final AtomicInteger stateHandlerCloseCalls = new AtomicInteger(); + final AtomicReference 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 sharedResources = + createSharedResourceLease(pythonSharedResources, sharedResourceLeases); + final OpaqueMemoryResource 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 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 getListState(BeamFnApi.StateRequest request) { + stateAccessed.set(true); + return null; + } + + @Override + public MapState 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 getInputReceivers() { + return Collections.emptyMap(); + } + + @Override + public Map, FnDataReceiver> 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 NoOpBeamStateHandler implements BeamStateHandler { + + @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 getTimers(RunnerApi.Components components) { + return Collections.emptyList(); + } + + @Override + protected Optional 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 0000000000000..bf99c827dd5ac --- /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 getListState(BeamFnApi.StateRequest request) { + stateAccessed.set(true); + return null; + } + + @Override + public MapState 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 getListState(BeamFnApi.StateRequest request) + throws InterruptedException { + stateAccessStarted.countDown(); + releaseStateAccess.await(); + return null; + } + + @Override + public MapState 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 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 implements BeamStateHandler { + + @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(); + } + } +}