From dfa7cbc8732b4aab032273663affe43c4254b766 Mon Sep 17 00:00:00 2001 From: bowenli86 Date: Fri, 10 Jul 2026 19:06:55 -0700 Subject: [PATCH 1/8] [FLINK-40135][python] Prevent Beam state access after runner shutdown Generated-by: Codex GPT-5 --- .../python/beam/BeamPythonFunctionRunner.java | 15 +- .../beam/state/BeamStateRequestHandler.java | 69 +++++-- .../beam/BeamPythonFunctionRunnerTest.java | 192 ++++++++++++++++++ .../state/BeamStateRequestHandlerTest.java | 158 ++++++++++++++ 4 files changed, 413 insertions(+), 21 deletions(-) create mode 100644 flink-python/src/test/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunnerTest.java create mode 100644 flink-python/src/test/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamStateRequestHandlerTest.java 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 7d89c93cc36aff..42812a2ef4faa2 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 @@ -55,7 +55,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; @@ -165,7 +164,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; @@ -341,6 +340,16 @@ public void close() throws Exception { } } finally { jobBundleFactory = null; + + // State backends are disposed after the runner is closed. Drain the handler after + // stopping Beam request production so no callback can access disposed state. + try { + if (stateRequestHandler != null) { + stateRequestHandler.close(); + } + } finally { + stateRequestHandler = null; + } } try { @@ -735,7 +744,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/state/BeamStateRequestHandler.java b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamStateRequestHandler.java index fb6765dff202d2..34295b5a519be6 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/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 00000000000000..e92457fd482f68 --- /dev/null +++ b/flink-python/src/test/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunnerTest.java @@ -0,0 +1,192 @@ +/* + * 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.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.core.construction.graph.ExecutableStage; +import org.apache.beam.runners.core.construction.graph.TimerReference; +import org.apache.beam.runners.fnexecution.control.JobBundleFactory; +import org.apache.beam.runners.fnexecution.control.StageBundleFactory; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.Field; +import java.util.Collections; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +class BeamPythonFunctionRunnerTest { + + @Test + void testCloseDrainsStateHandlerAfterStoppingRequestProduction() throws Exception { + final AtomicBoolean stateAccessedDuringFactoryClose = new AtomicBoolean(); + final BeamStateStore keyedStateStore = + new BeamStateStore() { + @Override + public ListState getListState(BeamFnApi.StateRequest request) { + stateAccessedDuringFactoryClose.set(true); + return null; + } + + @Override + public MapState getMapState( + BeamFnApi.StateRequest request) { + throw new UnsupportedOperationException(); + } + }; + final BeamStateRequestHandler stateRequestHandler = + new BeamStateRequestHandler( + keyedStateStore, + BeamStateStore.unsupported(), + new NoOpBeamStateHandler<>(), + new NoOpBeamStateHandler<>()); + final JobBundleFactory jobBundleFactory = new TestingJobBundleFactory(stateRequestHandler); + final TestingBeamPythonFunctionRunner runner = + new TestingBeamPythonFunctionRunner(createEnvironmentManager()); + setField(runner, "jobBundleFactory", jobBundleFactory); + setField(runner, "stateRequestHandler", stateRequestHandler); + + runner.close(); + + assertThat(stateAccessedDuringFactoryClose).isTrue(); + 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 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 00000000000000..2e89868b69d179 --- /dev/null +++ b/flink-python/src/test/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamStateRequestHandlerTest.java @@ -0,0 +1,158 @@ +/* + * 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.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; + +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 Future closeFuture = + executor.submit( + () -> { + closeStarted.countDown(); + handler.close(); + }); + assertThat(closeStarted.await(10, TimeUnit.SECONDS)).isTrue(); + assertThatThrownBy(() -> closeFuture.get(100, TimeUnit.MILLISECONDS)) + .isInstanceOf(TimeoutException.class); + + releaseStateAccess.countDown(); + requestFuture.get(10, TimeUnit.SECONDS); + closeFuture.get(10, TimeUnit.SECONDS); + } finally { + releaseStateAccess.countDown(); + executor.shutdownNow(); + } + } + + 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(); + } + } +} From f68efc1d530036fd7c85617f49e38f2e86a9cd6f Mon Sep 17 00:00:00 2001 From: bowenli86 Date: Fri, 31 Jul 2026 18:10:45 -0700 Subject: [PATCH 2/8] [FLINK-40135][python] Fix managed memory shutdown ordering Generated-by: Codex GPT-5 --- .../python/beam/BeamPythonFunctionRunner.java | 69 ++++--- .../beam/BeamPythonFunctionRunnerTest.java | 184 ++++++++++++++++-- 2 files changed, 209 insertions(+), 44 deletions(-) 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 42812a2ef4faa2..99bf0c30616bbe 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 @@ -335,40 +335,53 @@ public void open(ReadableConfig config) throws Exception { @Override public void close() throws Exception { try { - if (jobBundleFactory != null) { - jobBundleFactory.close(); - } - } finally { - jobBundleFactory = null; - - // State backends are disposed after the runner is closed. Drain the handler after - // stopping Beam request production so no callback can access disposed state. try { - if (stateRequestHandler != null) { - stateRequestHandler.close(); + // A managed-memory JobBundleFactory can be shared by multiple runners. Finish this + // runner's bundle first so it cannot issue state requests after its handler closes. + if (sharedResources != null && bundleStarted) { + flush(); } } finally { - stateRequestHandler = null; - } - } + bundleStarted = false; - 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(); + try { + if (jobBundleFactory != null) { + jobBundleFactory.close(); + } + } finally { + jobBundleFactory = null; + + 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(); + } + } finally { + sharedResources = null; + + // State backends are disposed after the runner is closed. Gate this + // runner's handler only after its bundle and resource teardown can no + // longer use it. + try { + if (stateRequestHandler != null) { + stateRequestHandler.close(); + } + } finally { + stateRequestHandler = null; + } + } + } } } finally { - sharedResources = null; - } - - if (shutdownHook != null) { - ShutdownHookUtil.removeShutdownHook( - shutdownHook, BeamPythonFunctionRunner.class.getSimpleName(), LOG); - shutdownHook = null; + if (shutdownHook != null) { + ShutdownHookUtil.removeShutdownHook( + shutdownHook, BeamPythonFunctionRunner.class.getSimpleName(), LOG); + shutdownHook = null; + } } } 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 index e92457fd482f68..f930599649971a 100644 --- 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 @@ -24,6 +24,7 @@ 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; @@ -31,17 +32,23 @@ import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.core.construction.Timer; import org.apache.beam.runners.core.construction.graph.ExecutableStage; import org.apache.beam.runners.core.construction.graph.TimerReference; 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.values.KV; import org.junit.jupiter.api.Test; import java.lang.reflect.Field; import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -49,13 +56,107 @@ class BeamPythonFunctionRunnerTest { @Test - void testCloseDrainsStateHandlerAfterStoppingRequestProduction() throws Exception { + 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).isFalse(); + 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); + } + + 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) { - stateAccessedDuringFactoryClose.set(true); + stateAccessed.set(true); return null; } @@ -65,21 +166,14 @@ public MapState getMapState( throw new UnsupportedOperationException(); } }; - final BeamStateRequestHandler stateRequestHandler = - new BeamStateRequestHandler( - keyedStateStore, - BeamStateStore.unsupported(), - new NoOpBeamStateHandler<>(), - new NoOpBeamStateHandler<>()); - final JobBundleFactory jobBundleFactory = new TestingJobBundleFactory(stateRequestHandler); - final TestingBeamPythonFunctionRunner runner = - new TestingBeamPythonFunctionRunner(createEnvironmentManager()); - setField(runner, "jobBundleFactory", jobBundleFactory); - setField(runner, "stateRequestHandler", stateRequestHandler); - - runner.close(); + return new BeamStateRequestHandler( + keyedStateStore, + BeamStateStore.unsupported(), + new NoOpBeamStateHandler<>(), + new NoOpBeamStateHandler<>()); + } - assertThat(stateAccessedDuringFactoryClose).isTrue(); + private static void assertStateHandlerClosed(BeamStateRequestHandler stateRequestHandler) { assertThatThrownBy(() -> stateRequestHandler.handle(createBagUserStateRequest())) .isInstanceOf(IllegalStateException.class) .hasMessage("Beam state request handler is closed."); @@ -130,6 +224,64 @@ public void close() throws Exception { } } + 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 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 From a34a8734bbd4101881911004c194b98a4724fb9a Mon Sep 17 00:00:00 2001 From: bowenli86 Date: Fri, 31 Jul 2026 18:12:51 -0700 Subject: [PATCH 3/8] [FLINK-40135][python] Serialize managed bundle shutdown Generated-by: Codex GPT-5 --- .../api/runners/python/beam/BeamPythonFunctionRunner.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 99bf0c30616bbe..10ae01bda6519e 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 @@ -338,7 +338,7 @@ public void close() throws Exception { try { // A managed-memory JobBundleFactory can be shared by multiple runners. Finish this // runner's bundle first so it cannot issue state requests after its handler closes. - if (sharedResources != null && bundleStarted) { + if (sharedResources != null) { flush(); } } finally { @@ -463,7 +463,7 @@ public Tuple3 takeResult() throws Exception { } @Override - public void flush() throws Exception { + public synchronized void flush() throws Exception { if (bundleStarted) { try { finishBundle(); From 4655bca39c2d39124e1fdd070ebf2c40273b96b4 Mon Sep 17 00:00:00 2001 From: bowenli86 Date: Fri, 31 Jul 2026 21:00:35 -0700 Subject: [PATCH 4/8] [FLINK-40135][python] Test concurrent managed bundle shutdown Generated-by: Codex GPT-5 --- .../beam/BeamPythonFunctionRunnerTest.java | 127 ++++++++++++++++++ 1 file changed, 127 insertions(+) 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 index f930599649971a..46af6ef2c9fb17 100644 --- 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 @@ -47,6 +47,12 @@ 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.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; @@ -139,6 +145,73 @@ void testCloseDrainsStateHandlerForFinalManagedMemoryLease() throws Exception { assertStateHandlerClosed(stateRequestHandler); } + @Test + void testCloseWaitsForConcurrentManagedBundleFlush() 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 CountDownLatch closeFlushStarted = 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); + runner.notifyWhenCloseFlushStarts(closeFlushStarted); + 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( + () -> { + runner.close(); + return null; + }); + assertThat(closeFlushStarted.await(10, TimeUnit.SECONDS)).isTrue(); + assertThatThrownBy(() -> closeFuture.get(100, TimeUnit.MILLISECONDS)) + .isInstanceOf(TimeoutException.class); + 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(); + } + } + private static OpaqueMemoryResource createSharedResourceLease( PythonSharedResources pythonSharedResources, AtomicInteger remainingLeases) { return new OpaqueMemoryResource<>( @@ -263,6 +336,35 @@ public void close() throws Exception { } } + private static class BlockingTestingRemoteBundle extends TestingRemoteBundle { + + private final BeamStateRequestHandler stateRequestHandler; + private final AtomicInteger closeCalls; + private final CountDownLatch firstCloseStarted; + private final CountDownLatch releaseFirstClose; + + private BlockingTestingRemoteBundle( + BeamStateRequestHandler stateRequestHandler, + AtomicInteger closeCalls, + CountDownLatch firstCloseStarted, + CountDownLatch releaseFirstClose) { + super(stateRequestHandler, new AtomicBoolean()); + this.stateRequestHandler = stateRequestHandler; + this.closeCalls = closeCalls; + this.firstCloseStarted = firstCloseStarted; + this.releaseFirstClose = releaseFirstClose; + } + + @Override + public void close() throws Exception { + if (closeCalls.incrementAndGet() == 1) { + firstCloseStarted.countDown(); + releaseFirstClose.await(); + } + stateRequestHandler.handle(createBagUserStateRequest()); + } + } + private static class TrackingJobBundleFactory implements JobBundleFactory { private final AtomicBoolean closed; @@ -309,6 +411,9 @@ public BeamFnApi.StateResponse.Builder handleClear( private static class TestingBeamPythonFunctionRunner extends BeamPythonFunctionRunner { + private volatile Thread closingThread; + private volatile CountDownLatch closeFlushStarted; + private TestingBeamPythonFunctionRunner( ProcessPythonEnvironmentManager environmentManager) { super( @@ -328,6 +433,28 @@ private TestingBeamPythonFunctionRunner( Collections.emptyMap()); } + private void notifyWhenCloseFlushStarts(CountDownLatch closeFlushStarted) { + this.closeFlushStarted = closeFlushStarted; + } + + @Override + public void close() throws Exception { + closingThread = Thread.currentThread(); + try { + super.close(); + } finally { + closingThread = null; + } + } + + @Override + public void flush() throws Exception { + if (Thread.currentThread() == closingThread && closeFlushStarted != null) { + closeFlushStarted.countDown(); + } + super.flush(); + } + @Override protected void buildTransforms(RunnerApi.Components.Builder componentsBuilder) {} From be9bdf4f07488fbb65987180e855caf94543ddf4 Mon Sep 17 00:00:00 2001 From: bowenli86 Date: Fri, 31 Jul 2026 21:20:31 -0700 Subject: [PATCH 5/8] [FLINK-40135][python] Make shutdown concurrency test deterministic Generated-by: Codex GPT-5 --- .../beam/BeamPythonFunctionRunnerTest.java | 25 ++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) 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 index 46af6ef2c9fb17..54b0f60a40f336 100644 --- 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 @@ -52,9 +52,10 @@ import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; 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; @@ -150,6 +151,7 @@ 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 closeFlushStarted = new CountDownLatch(1); @@ -189,12 +191,13 @@ void testCloseWaitsForConcurrentManagedBundleFlush() throws Exception { final Future closeFuture = executor.submit( () -> { + closeThread.set(Thread.currentThread()); runner.close(); return null; }); assertThat(closeFlushStarted.await(10, TimeUnit.SECONDS)).isTrue(); - assertThatThrownBy(() -> closeFuture.get(100, TimeUnit.MILLISECONDS)) - .isInstanceOf(TimeoutException.class); + assertCloseWaitsForConcurrentFlush(closeThread.get(), remoteBundleCloseCalls); + assertThat(closeFuture.isDone()).isFalse(); assertThat(remoteBundleCloseCalls).hasValue(1); assertThat(sharedFactoryClosed).isFalse(); @@ -212,6 +215,22 @@ void testCloseWaitsForConcurrentManagedBundleFlush() throws Exception { } } + private static void assertCloseWaitsForConcurrentFlush( + Thread closeThread, AtomicInteger remoteBundleCloseCalls) { + final long timeoutNanos = TimeUnit.SECONDS.toNanos(10); + final long deadlineNanos = System.nanoTime() + timeoutNanos; + while (System.nanoTime() < deadlineNanos) { + if (remoteBundleCloseCalls.get() > 1) { + assertThat(remoteBundleCloseCalls).hasValue(1); + } + if (closeThread.getState() == Thread.State.BLOCKED) { + return; + } + LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1)); + } + assertThat(closeThread.getState()).isEqualTo(Thread.State.BLOCKED); + } + private static OpaqueMemoryResource createSharedResourceLease( PythonSharedResources pythonSharedResources, AtomicInteger remainingLeases) { return new OpaqueMemoryResource<>( From 3105c6e524a1ff1a9edb4349a1d4a6cf9fefea12 Mon Sep 17 00:00:00 2001 From: bowenli86 Date: Tue, 11 Aug 2026 10:09:31 -0700 Subject: [PATCH 6/8] [FLINK-40135][python] Make runner close idempotent Generated-by: Codex GPT-5 --- .../python/beam/BeamPythonFunctionRunner.java | 10 +- .../beam/BeamPythonFunctionRunnerTest.java | 151 ++++++++++++++---- 2 files changed, 127 insertions(+), 34 deletions(-) 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 10ae01bda6519e..f9a72080b75d9a 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 @@ -194,6 +194,9 @@ 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; + private transient Thread shutdownHook; private transient Environment environment; @@ -333,7 +336,12 @@ public void open(ReadableConfig config) throws Exception { } @Override - public void close() throws Exception { + public synchronized void close() throws Exception { + if (closed) { + return; + } + closed = true; + try { try { // A managed-memory JobBundleFactory can be shared by multiple runners. Finish this 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 index 54b0f60a40f336..926d2d6f256abd 100644 --- 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 @@ -43,6 +43,9 @@ 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; @@ -154,7 +157,7 @@ void testCloseWaitsForConcurrentManagedBundleFlush() throws Exception { final AtomicReference closeThread = new AtomicReference<>(); final CountDownLatch firstBundleCloseStarted = new CountDownLatch(1); final CountDownLatch releaseFirstBundleClose = new CountDownLatch(1); - final CountDownLatch closeFlushStarted = new CountDownLatch(1); + final CountDownLatch closeStarted = new CountDownLatch(1); final BeamStateRequestHandler stateRequestHandler = createStateRequestHandler(stateAccessedDuringBundleClose); final PythonSharedResources pythonSharedResources = @@ -176,7 +179,6 @@ void testCloseWaitsForConcurrentManagedBundleFlush() throws Exception { releaseFirstBundleClose)); setField(runner, "bundleStarted", true); setField(runner, "sharedResources", sharedResources); - runner.notifyWhenCloseFlushStarts(closeFlushStarted); final ExecutorService executor = Executors.newFixedThreadPool(2); try { @@ -192,11 +194,12 @@ void testCloseWaitsForConcurrentManagedBundleFlush() throws Exception { executor.submit( () -> { closeThread.set(Thread.currentThread()); + closeStarted.countDown(); runner.close(); return null; }); - assertThat(closeFlushStarted.await(10, TimeUnit.SECONDS)).isTrue(); - assertCloseWaitsForConcurrentFlush(closeThread.get(), remoteBundleCloseCalls); + assertThat(closeStarted.await(10, TimeUnit.SECONDS)).isTrue(); + assertCloseThreadIsBlocked(closeThread.get(), remoteBundleCloseCalls, closeFuture); assertThat(closeFuture.isDone()).isFalse(); assertThat(remoteBundleCloseCalls).hasValue(1); assertThat(sharedFactoryClosed).isFalse(); @@ -215,13 +218,85 @@ void testCloseWaitsForConcurrentManagedBundleFlush() throws Exception { } } - private static void assertCloseWaitsForConcurrentFlush( - Thread closeThread, AtomicInteger remoteBundleCloseCalls) { + @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(); + assertCloseThreadIsBlocked( + 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 assertCloseThreadIsBlocked( + Thread closeThread, AtomicInteger concurrentCloseCalls, Future closeFuture) { final long timeoutNanos = TimeUnit.SECONDS.toNanos(10); final long deadlineNanos = System.nanoTime() + timeoutNanos; while (System.nanoTime() < deadlineNanos) { - if (remoteBundleCloseCalls.get() > 1) { - assertThat(remoteBundleCloseCalls).hasValue(1); + if (closeFuture.isDone()) { + assertThat(closeFuture.isDone()).isFalse(); + } + if (concurrentCloseCalls.get() > 1) { + assertThat(concurrentCloseCalls).hasValue(1); } if (closeThread.getState() == Thread.State.BLOCKED) { return; @@ -384,6 +459,41 @@ public void close() throws Exception { } } + 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; @@ -430,9 +540,6 @@ public BeamFnApi.StateResponse.Builder handleClear( private static class TestingBeamPythonFunctionRunner extends BeamPythonFunctionRunner { - private volatile Thread closingThread; - private volatile CountDownLatch closeFlushStarted; - private TestingBeamPythonFunctionRunner( ProcessPythonEnvironmentManager environmentManager) { super( @@ -452,28 +559,6 @@ private TestingBeamPythonFunctionRunner( Collections.emptyMap()); } - private void notifyWhenCloseFlushStarts(CountDownLatch closeFlushStarted) { - this.closeFlushStarted = closeFlushStarted; - } - - @Override - public void close() throws Exception { - closingThread = Thread.currentThread(); - try { - super.close(); - } finally { - closingThread = null; - } - } - - @Override - public void flush() throws Exception { - if (Thread.currentThread() == closingThread && closeFlushStarted != null) { - closeFlushStarted.countDown(); - } - super.flush(); - } - @Override protected void buildTransforms(RunnerApi.Components.Builder componentsBuilder) {} From a6e99cf7f672086a1e98517e92427607a2c58325 Mon Sep 17 00:00:00 2001 From: bowenli86 Date: Tue, 11 Aug 2026 10:20:44 -0700 Subject: [PATCH 7/8] [FLINK-40135][python] Make state shutdown test deterministic Generated-by: Codex GPT-5 --- .../state/BeamStateRequestHandlerTest.java | 24 ++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) 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 index 2e89868b69d179..bf99c827dd5ac4 100644 --- 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 @@ -30,8 +30,9 @@ import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; 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; @@ -94,15 +95,16 @@ public MapState getMapState( 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(); - assertThatThrownBy(() -> closeFuture.get(100, TimeUnit.MILLISECONDS)) - .isInstanceOf(TimeoutException.class); + assertCloseWaitsForInFlightRequest(closeThread.get(), closeFuture); releaseStateAccess.countDown(); requestFuture.get(10, TimeUnit.SECONDS); @@ -113,6 +115,22 @@ public MapState getMapState( } } + 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, From 35db8a7688127bc09c32daa1b957b3ddffb6f801 Mon Sep 17 00:00:00 2001 From: bowenli86 Date: Tue, 11 Aug 2026 10:59:58 -0700 Subject: [PATCH 8/8] [FLINK-40135][python] Add non-blocking Beam cancellation Generated-by: Codex GPT-5 --- .../flink/python/PythonFunctionRunner.java | 10 + ...bstractExternalPythonFunctionOperator.java | 34 +- .../python/beam/BeamPythonFunctionRunner.java | 415 +++++++++++++++--- .../python/beam/PythonSharedResources.java | 29 +- ...actExternalPythonFunctionOperatorTest.java | 176 ++++++++ .../beam/BeamPythonFunctionRunnerTest.java | 260 ++++++++++- 6 files changed, 840 insertions(+), 84 deletions(-) create mode 100644 flink-python/src/test/java/org/apache/flink/streaming/api/operators/python/process/AbstractExternalPythonFunctionOperatorTest.java 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 c637c0e659c13a..708e8c70e8a22a 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 da50b5335f22c5..f0367fb6b95cfd 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 f9a72080b75d9a..82cf75e60d3088 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; @@ -92,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; @@ -197,6 +202,29 @@ public abstract class BeamPythonFunctionRunner implements PythonFunctionRunner { /** 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; @@ -330,73 +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 synchronized void close() throws Exception { - if (closed) { + public void close() throws Exception { + if (!startClose()) { + awaitCloseCompletion(); return; } - closed = true; try { try { - // A managed-memory JobBundleFactory can be shared by multiple runners. Finish this - // runner's bundle first so it cannot issue state requests after its handler closes. - if (sharedResources != null) { - flush(); - } + // 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 { - bundleStarted = false; - try { - if (jobBundleFactory != null) { - jobBundleFactory.close(); - } + closeResources(); } finally { - jobBundleFactory = null; - 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(); - } + closeStateRequestHandler(); } finally { - sharedResources = null; - - // State backends are disposed after the runner is closed. Gate this - // runner's handler only after its bundle and resource teardown can no - // longer use it. - try { - if (stateRequestHandler != null) { - stateRequestHandler.close(); - } - } finally { - stateRequestHandler = null; - } + removeShutdownHook(); } } } } finally { - if (shutdownHook != null) { - ShutdownHookUtil.removeShutdownHook( - shutdownHook, BeamPythonFunctionRunner.class.getSimpleName(), LOG); - shutdownHook = null; + completeClose(); + } + } + + @Override + public void cancel() throws Exception { + final BundleCloseClaim bundleCloseClaim = requestCancel(); + + try { + 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 { + 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 @@ -411,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; } } @@ -471,14 +511,35 @@ public Tuple3 takeResult() throws Exception { } @Override - public synchronized void flush() throws Exception { - if (bundleStarted) { - try { - finishBundle(); - } finally { + public void flush() throws Exception { + 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. */ @@ -486,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); } } 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 293d5a5ad83d86..9eaf2fbb4688b0 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/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 00000000000000..fa0d307a562ab4 --- /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 index 926d2d6f256abd..b5074fae72d3f0 100644 --- 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 @@ -32,13 +32,13 @@ import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.model.pipeline.v1.RunnerApi; -import org.apache.beam.runners.core.construction.Timer; -import org.apache.beam.runners.core.construction.graph.ExecutableStage; -import org.apache.beam.runners.core.construction.graph.TimerReference; 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; @@ -85,7 +85,7 @@ void testCloseDrainsStateHandlerAfterStoppingOwnedRequestProduction() throws Exc runner.close(); assertThat(stateAccessedDuringFactoryClose).isTrue(); - assertThat(remoteBundleClosed).isFalse(); + assertThat(remoteBundleClosed).isTrue(); assertStateHandlerClosed(stateRequestHandler); } @@ -149,6 +149,70 @@ void testCloseDrainsStateHandlerForFinalManagedMemoryLease() throws Exception { 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(); @@ -199,7 +263,8 @@ void testCloseWaitsForConcurrentManagedBundleFlush() throws Exception { return null; }); assertThat(closeStarted.await(10, TimeUnit.SECONDS)).isTrue(); - assertCloseThreadIsBlocked(closeThread.get(), remoteBundleCloseCalls, closeFuture); + assertThreadIsWaitingForBundleClose( + closeThread.get(), remoteBundleCloseCalls, closeFuture); assertThat(closeFuture.isDone()).isFalse(); assertThat(remoteBundleCloseCalls).hasValue(1); assertThat(sharedFactoryClosed).isFalse(); @@ -218,6 +283,128 @@ void testCloseWaitsForConcurrentManagedBundleFlush() throws Exception { } } + @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(); @@ -267,7 +454,7 @@ void testConcurrentCloseIsIdempotent() throws Exception { return null; }); assertThat(secondCloseStarted.await(10, TimeUnit.SECONDS)).isTrue(); - assertCloseThreadIsBlocked( + assertThreadIsWaitingForBundleClose( secondCloseThread.get(), stateHandlerCloseCalls, secondCloseFuture); assertThat(secondCloseFuture.isDone()).isFalse(); assertThat(stateHandlerCloseCalls).hasValue(1); @@ -287,7 +474,7 @@ void testConcurrentCloseIsIdempotent() throws Exception { } } - private static void assertCloseThreadIsBlocked( + private static void assertThreadIsWaitingForBundleClose( Thread closeThread, AtomicInteger concurrentCloseCalls, Future closeFuture) { final long timeoutNanos = TimeUnit.SECONDS.toNanos(10); final long deadlineNanos = System.nanoTime() + timeoutNanos; @@ -298,12 +485,24 @@ private static void assertCloseThreadIsBlocked( if (concurrentCloseCalls.get() > 1) { assertThat(concurrentCloseCalls).hasValue(1); } - if (closeThread.getState() == Thread.State.BLOCKED) { + if (closeThread.getState() == Thread.State.WAITING) { return; } LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1)); } - assertThat(closeThread.getState()).isEqualTo(Thread.State.BLOCKED); + 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( @@ -436,26 +635,61 @@ private static class BlockingTestingRemoteBundle extends TestingRemoteBundle { 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 { - if (closeCalls.incrementAndGet() == 1) { - firstCloseStarted.countDown(); - releaseFirstClose.await(); + 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(); } - stateRequestHandler.handle(createBagUserStateRequest()); } }