[FLINK-40135][python] Prevent Beam state access after runner shutdown

flink-python/src/main/java/org/apache/flink/python/PythonFunctionRunner.java

@@ -34,6 +34,16 @@ public interface PythonFunctionRunner extends AutoCloseable {
     /** Tear-down the Python function runner. */
     void close() throws Exception;
 
+    /**
+     * Cancels the Python function runner.
+     *
+     * <p>The default implementation keeps compatibility with runners that do not need a separate
+     * cancellation path. Runners with a blocking graceful close should override this method.
+     */
+    default void cancel() throws Exception {
+        close();
+    }
+
     /**
      * Executes the Python function with the input byte array.
      *

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<Throwable> exceptionReference = new AtomicReference<>();
             flushThreadPool.submit(
                     () -> {
                         try {
-                            pythonFunctionRunner.flush();
+                            functionRunner.flush();
                         } catch (Throwable e) {
                             exceptionReference.set(e);
                         } finally {
                             flushThreadFinish.set(true);
                             // interrupt the progress of takeResult to avoid the main thread is
                             // blocked forever.
-                            ((BeamPythonFunctionRunner) pythonFunctionRunner).notifyNoMoreResults();
+                            ((BeamPythonFunctionRunner) functionRunner).notifyNoMoreResults();
                         }
                     });
             Tuple3<String, byte[], Integer> resultTuple;
             while (!flushThreadFinish.get()) {
-                resultTuple = pythonFunctionRunner.takeResult();
+                resultTuple = functionRunner.takeResult();
                 if (resultTuple.f2 != 0) {
                     emitResult(resultTuple);
                     emitResults();