[python] Prevent Beam state access after runner shutdown

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,

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<BeamFnApi.StateResponse.Builder> handle(BeamFnApi.StateRequest request)
             throws Exception {
-        BeamFnApi.StateKey.TypeCase typeCase = request.getStateKey().getTypeCase();
-        ListState<byte[]> listState;
-        MapState<ByteArrayWrapper, byte[]> mapState;
-
-        switch (typeCase) {
-            case BAG_USER_STATE:
-                listState = keyedStateStore.getListState(request);
-                return CompletableFuture.completedFuture(
-                        bagStateHandler.handle(request, listState));
-            case MULTIMAP_SIDE_INPUT:
-                mapState = keyedStateStore.getMapState(request);
-                return CompletableFuture.completedFuture(mapStateHandler.handle(request, mapState));
-            case MULTIMAP_KEYS_SIDE_INPUT:
-                mapState = operatorStateStore.getMapState(request);
-                return CompletableFuture.completedFuture(mapStateHandler.handle(request, mapState));
-            default:
-                throw new RuntimeException("Unsupported state type: " + typeCase);
+        final Lock readLock = lifecycleLock.readLock();
+        readLock.lock();
+        try {
+            if (closed) {
+                throw new IllegalStateException("Beam state request handler is closed.");
+            }
+
+            BeamFnApi.StateKey.TypeCase typeCase = request.getStateKey().getTypeCase();
+            ListState<byte[]> listState;
+            MapState<ByteArrayWrapper, byte[]> mapState;
+
+            switch (typeCase) {
+                case BAG_USER_STATE:
+                    listState = keyedStateStore.getListState(request);
+                    return CompletableFuture.completedFuture(
+                            bagStateHandler.handle(request, listState));
+                case MULTIMAP_SIDE_INPUT:
+                    mapState = keyedStateStore.getMapState(request);
+                    return CompletableFuture.completedFuture(
+                            mapStateHandler.handle(request, mapState));
+                case MULTIMAP_KEYS_SIDE_INPUT:
+                    mapState = operatorStateStore.getMapState(request);
+                    return CompletableFuture.completedFuture(
+                            mapStateHandler.handle(request, mapState));
+                default:
+                    throw new RuntimeException("Unsupported state type: " + typeCase);
+            }
+        } finally {
+            readLock.unlock();
+        }
+    }
+
+    /**
+     * Stops accepting state requests and waits for requests that are already accessing Flink state
+     * to complete.
+     */
+    @Override
+    public void close() {
+        final Lock writeLock = lifecycleLock.writeLock();
+        writeLock.lock();
+        try {
+            closed = true;
+        } finally {
+            writeLock.unlock();
         }
     }
 

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<byte[]> getListState(BeamFnApi.StateRequest request) {
+                        stateAccessedDuringFactoryClose.set(true);
+                        return null;
+                    }
+
+                    @Override
+                    public MapState<ByteArrayWrapper, byte[]> 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<S> implements BeamStateHandler<S> {
+
+        @Override
+        public BeamFnApi.StateResponse.Builder handle(BeamFnApi.StateRequest request, S state) {
+            return BeamFnApi.StateResponse.newBuilder();
+        }
+
+        @Override
+        public BeamFnApi.StateResponse.Builder handleGet(BeamFnApi.StateRequest request, S state) {
+            return BeamFnApi.StateResponse.newBuilder();
+        }
+
+        @Override
+        public BeamFnApi.StateResponse.Builder handleAppend(
+                BeamFnApi.StateRequest request, S state) {
+            return BeamFnApi.StateResponse.newBuilder();
+        }
+
+        @Override
+        public BeamFnApi.StateResponse.Builder handleClear(
+                BeamFnApi.StateRequest request, S state) {
+            return BeamFnApi.StateResponse.newBuilder();
+        }
+    }
+
+    private static class TestingBeamPythonFunctionRunner extends BeamPythonFunctionRunner {
+
+        private TestingBeamPythonFunctionRunner(
+                ProcessPythonEnvironmentManager environmentManager) {
+            super(
+                    null,
+                    "test-task",
+                    environmentManager,
+                    null,
+                    null,
+                    null,
+                    null,
+                    null,
+                    null,
+                    null,
+                    0.0,
+                    FlinkFnApi.CoderInfoDescriptor.getDefaultInstance(),
+                    FlinkFnApi.CoderInfoDescriptor.getDefaultInstance(),
+                    Collections.emptyMap());
+        }
+
+        @Override
+        protected void buildTransforms(RunnerApi.Components.Builder componentsBuilder) {}
+
+        @Override
+        protected List<TimerReference> getTimers(RunnerApi.Components components) {
+            return Collections.emptyList();
+        }
+
+        @Override
+        protected Optional<RunnerApi.Coder> getOptionalTimerCoderProto() {
+            return Optional.empty();
+        }
+    }
+}