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 8e5e5fdd58682..7791d33744dd0 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 @@ -65,7 +65,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; @@ -164,7 +163,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; @@ -340,6 +339,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 { @@ -734,7 +743,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 3cb259c2eec5f..bfded650aa44c 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 0000000000000..e92457fd482f6 --- /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 0000000000000..2e89868b69d17 --- /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(); + } + } +}